# Buggy code
def _gather_rollout(self, rollout_number, last_observation):
    assert 0 < self._path_lengths[rollout_number] <= self._max_episode_length
    traj = TrajectoryBatch(
        env_spec=self._envs[rollout_number].spec,
        observations=np.asarray(self._observations[rollout_number]),
        last_observations=np.asarray([last_observation]),
        actions=np.asarray(self._actions[rollout_number]),
        rewards=np.asarray(self._rewards[rollout_number]),
        step_types=np.asarray(self._step_types[rollout_number], dtype=StepType),
        env_infos=self._env_infos[rollout_number],
        agent_infos=self._agent_infos[rollout_number],
        lengths=np.asarray([self._path_lengths[rollout_number]], dtype="l"),
    )
    self._completed_rollouts.append(traj)
    self._observations[rollout_number] = []
    self._actions[rollout_number] = []
    self._rewards[rollout_number] = []
    self._step_types[rollout_number] = []
    self._path_lengths[rollout_number] = 0
    self._prev_obs[rollout_number] = self._envs[rollout_number].reset()


# Traceback (ValueError)
---------------------------------------  --------------
Sampling  [####################################]  100%
2020-07-24 12:28:46 | [trpo_gym_tf_cartpole] epoch #3 | Optimizing policy...
2020-07-24 12:28:46 | [trpo_gym_tf_cartpole] epoch #3 | Computing loss before
2020-07-24 12:28:46 | [trpo_gym_tf_cartpole] epoch #3 | Computing KL before
2020-07-24 12:28:46 | [trpo_gym_tf_cartpole] epoch #3 | Optimizing
2020-07-24 12:28:46 | [trpo_gym_tf_cartpole] epoch #3 | Start CG optimization: #parameters: 1282, #inputs: 47, #subsample_inputs: 47
2020-07-24 12:28:46 | [trpo_gym_tf_cartpole] epoch #3 | computing loss before
2020-07-24 12:28:46 | [trpo_gym_tf_cartpole] epoch #3 | computing gradient
2020-07-24 12:28:46 | [trpo_gym_tf_cartpole] epoch #3 | gradient computed
2020-07-24 12:28:46 | [trpo_gym_tf_cartpole] epoch #3 | computing descent direction
2020-07-24 12:28:46 | [trpo_gym_tf_cartpole] epoch #3 | descent direction computed
2020-07-24 12:28:46 | [trpo_gym_tf_cartpole] epoch #3 | backtrack iters: 2
2020-07-24 12:28:46 | [trpo_gym_tf_cartpole] epoch #3 | optimization finished
2020-07-24 12:28:46 | [trpo_gym_tf_cartpole] epoch #3 | Computing KL after
2020-07-24 12:28:46 | [trpo_gym_tf_cartpole] epoch #3 | Computing loss after
2020-07-24 12:28:46 | [trpo_gym_tf_cartpole] epoch #3 | Fitting baseline...
2020-07-24 12:28:46 | [trpo_gym_tf_cartpole] epoch #3 | Saving snapshot...
2020-07-24 12:28:46 | [trpo_gym_tf_cartpole] epoch #3 | Saved
2020-07-24 12:28:46 | [trpo_gym_tf_cartpole] epoch #3 | Time 7.64 s
2020-07-24 12:28:46 | [trpo_gym_tf_cartpole] epoch #3 | EpochTime 0.93 s
---------------------------------------  --------------
Evaluation/AverageDiscountedReturn          54.6411
Evaluation/AverageReturn                    85.1489
Evaluation/Iteration                         3
Evaluation/MaxReturn                       188
Evaluation/MinReturn                        21
Evaluation/NumTrajs                         47
Evaluation/StdReturn                        37.8738
Evaluation/TerminationRate                   1
Extras/EpisodeRewardMean                    70.39
LinearFeatureBaseline/ExplainedVariance      0.402734
TotalEnvSteps                            16156
policy/Entropy                               1.41126
policy/KL                                    0.00784782
policy/KLBefore                              0
policy/LossAfter                            -0.101578
policy/LossBefore                           -0.0811697
policy/Perplexity                            4.10113
policy/dLoss                                 0.0204084
---------------------------------------  --------------
Traceback (most recent call last):
File "garage/examples/tf/trpo_gym_tf_cartpole.py", line 48, in <module>
trpo_gym_tf_cartpole()
File "/home/csidrane/Documents/NASA/garage/src/garage/experiment/experiment.py", line 362, in __call__
result = self.function(ctxt, **kwargs)
File "garage/examples/tf/trpo_gym_tf_cartpole.py", line 45, in trpo_gym_tf_cartpole
runner.train(n_epochs=120, batch_size=4000)
File "/home/csidrane/Documents/NASA/garage/src/garage/experiment/local_runner.py", line 500, in train
average_return = self._algo.train(self)
File "/home/csidrane/Documents/NASA/garage/src/garage/tf/algos/npo.py", line 185, in train
runner.step_path = runner.obtain_samples(runner.step_itr)
File "/home/csidrane/Documents/NASA/garage/src/garage/experiment/local_runner.py", line 358, in obtain_samples
env_update)
File "/home/csidrane/Documents/NASA/garage/src/garage/experiment/local_runner.py", line 325, in obtain_trajectories
env_update=env_update)
File "/home/csidrane/Documents/NASA/garage/src/garage/sampler/ray_sampler.py", line 164, in obtain_samples
ready_worker_id, trajectory_batch = ray.get(result)
File "/home/csidrane/anaconda3/envs/garage/lib/python3.6/site-packages/ray/worker.py", line 1474, in get
raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(ValueError): ray::SamplerWorker.rollout() (pid=17540, ip=171.64.160.86)
File "python/ray/_raylet.pyx", line 446, in ray._raylet.execute_task
File "python/ray/_raylet.pyx", line 400, in ray._raylet.execute_task.function_executor
File "/home/csidrane/Documents/NASA/garage/src/garage/sampler/ray_sampler.py", line 299, in rollout
return (self.worker_id, self.inner_worker.rollout())
File "/home/csidrane/Documents/NASA/garage/src/garage/tf/samplers/worker.py", line 137, in rollout
return self._inner_worker.rollout()
File "/home/csidrane/Documents/NASA/garage/src/garage/sampler/default_worker.py", line 181, in rollout
return self.collect_rollout()
File "/home/csidrane/Documents/NASA/garage/src/garage/sampler/default_worker.py", line 169, in collect_rollout
dtype='i'))
File "/home/csidrane/Documents/NASA/garage/src/garage/_dtypes.py", line 213, in __new__
format(inferred_batch_size, key, val.shape[0]))
ValueError: Each entry in env_infos must have a batch dimension of length 200, but got key TimeLimit.truncated with batch size 1 instead.

# Fixed code
def _gather_rollout(self, rollout_number, last_observation):
    assert 0 < self._path_lengths[rollout_number] <= self._max_episode_length
    env_infos = self._env_infos[rollout_number]
    agent_infos = self._agent_infos[rollout_number]
    for k, v in env_infos.items():
        env_infos[k] = np.asarray(v)
    for k, v in agent_infos.items():
        agent_infos[k] = np.asarray(v)
    traj = TrajectoryBatch(
        env_spec=self._envs[rollout_number].spec,
        observations=np.asarray(self._observations[rollout_number]),
        last_observations=np.asarray([last_observation]),
        actions=np.asarray(self._actions[rollout_number]),
        rewards=np.asarray(self._rewards[rollout_number]),
        step_types=np.asarray(self._step_types[rollout_number], dtype=StepType),
        env_infos=dict(env_infos),
        agent_infos=dict(agent_infos),
        lengths=np.asarray([self._path_lengths[rollout_number]], dtype="l"),
    )
    self._completed_rollouts.append(traj)
    self._observations[rollout_number] = []
    self._actions[rollout_number] = []
    self._rewards[rollout_number] = []
    self._step_types[rollout_number] = []
    self._path_lengths[rollout_number] = 0
    self._prev_obs[rollout_number] = self._envs[rollout_number].reset()
    self._env_infos[rollout_number] = collections.defaultdict(list)
    self._agent_infos[rollout_number] = collections.defaultdict(list)