def _unflatten(flat_data):
  """Converts a flat dict of numpy arrays to the batch tuple structure."""
  o_tm1 = {
      'penalties': flat_data['penalties_tm1'],
      'position': flat_data['position_tm1'],
      'velocity': flat_data['velocity_tm1'],
  }
  a_tm1 = flat_data['action_tm1']
  r_t = flat_data['reward_t']
  d_t = flat_data['discount_t']
  o_t = {
      'penalties': flat_data['penalties_t'],
      'position': flat_data['position_t'],
      'velocity': flat_data['velocity_t'],
  }
  return (o_tm1, a_tm1, r_t, d_t, o_t)