# Buggy code
def step(self, action):
    """Call step on wrapped env.

    This method is necessary to suppress a deprecated warning
    thrown by gym.Wrapper.

    Args:
        action (np.ndarray): An action provided by the agent.

    Returns:
        np.ndarray: Agent's observation of the current environment
        float: Amount of reward returned after previous action
        bool: Whether the episode has ended, in which case further step()
            calls will return undefined results
        dict: Contains auxiliary diagnostic information (helpful for
            debugging, and sometimes learning)

    """
    observation, reward, done, info = self.env.step(action)
    # gym envs that are wrapped in TimeLimit wrapper modify
    # the done/termination signal to be true whenever a time
    # limit expiration occurs. The following statement sets
    # the done signal to be True only if caused by an
    # environment termination, and not a time limit
    # termination. The time limit termination signal
    # will be saved inside env_infos as
    # 'BulletEnv.TimeLimitTerminated'
    if "TimeLimit.truncated" in info:
        info["BulletEnv.TimeLimitTerminated"] = done  # done = True always
        done = not info["TimeLimit.truncated"]
    return observation, reward, done, info


# 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 step(self, action):
    """Call step on wrapped env.

    This method is necessary to suppress a deprecated warning
    thrown by gym.Wrapper.

    Args:
        action (np.ndarray): An action provided by the agent.

    Returns:
        np.ndarray: Agent's observation of the current environment
        float: Amount of reward returned after previous action
        bool: Whether the episode has ended, in which case further step()
            calls will return undefined results
        dict: Contains auxiliary diagnostic information (helpful for
            debugging, and sometimes learning)

    """
    observation, reward, done, info = self.env.step(action)
    # gym envs that are wrapped in TimeLimit wrapper modify
    # the done/termination signal to be true whenever a time
    # limit expiration occurs. The following statement sets
    # the done signal to be True only if caused by an
    # environment termination, and not a time limit
    # termination. The time limit termination signal
    # will be saved inside env_infos as
    # 'BulletEnv.TimeLimitTerminated'
    if "TimeLimit.truncated" in info:
        info["BulletEnv.TimeLimitTerminated"] = done  # done = True always
        done = not info["TimeLimit.truncated"]
    else:
        info["TimeLimit.truncated"] = False
        info["BulletEnv.TimeLimitTerminated"] = False
    return observation, reward, done, info