code: def convert_episode_to_batch_major(episode):
      
      episode_batch = {}
      for key in episode.keys():
          val = np.array(episode[key]).copy()
          # make inputs batch-major instead of time-major
          episode_batch[key] = val.swapaxes(0, 1)
  
      return episode_batch
docs: Converts an episode to have the batch dimension in the major (first)
      dimension.