# Buggy code
def postprocess_trajectory(
    policy: Policy,
    sample_batch: SampleBatch,
    other_agent_batches: Optional[Dict[AgentID, SampleBatch]] = None,
    episode: Optional[MultiAgentEpisode] = None,
) -> SampleBatch:
    """Postprocesses a trajectory and returns the processed trajectory.

    The trajectory contains only data from one episode and from one agent.
    - If  `config.batch_mode=truncate_episodes` (default), sample_batch may
    contain a truncated (at-the-end) episode, in case the
    `config.rollout_fragment_length` was reached by the sampler.
    - If `config.batch_mode=complete_episodes`, sample_batch will contain
    exactly one episode (no matter how long).
    New columns can be added to sample_batch and existing ones may be altered.

    Args:
        policy (Policy): The Policy used to generate the trajectory
            (`sample_batch`)
        sample_batch (SampleBatch): The SampleBatch to postprocess.
        other_agent_batches (Optional[Dict[PolicyID, SampleBatch]]): Optional
            dict of AgentIDs mapping to other agents' trajectory data (from the
            same episode). NOTE: The other agents use the same policy.
        episode (Optional[MultiAgentEpisode]): Optional multi-agent episode
            object in which the agents operated.

    Returns:
        SampleBatch: The postprocessed, modified SampleBatch (or a new one).
    """
    if not policy.config["vtrace"]:
        sample_batch = postprocess_ppo_gae(
            policy, sample_batch, other_agent_batches, episode
        )

    # TODO: (sven) remove this del once we have trajectory view API fully in
    #  place.
    del sample_batch.data["new_obs"]  # not used, so save some bandwidth

    return sample_batch


# Traceback (tensorflow.python.framework.errors_impl.InvalidArgumentError)
2020-06-21 13:09:32,571	ERROR trial_runner.py:524 -- Trial PPO_TestingGym_f28cf_00000: Error processing event.
Traceback (most recent call last):
File "/home/user/anaconda3/envs/RLlibTesting/lib/python3.8/site-packages/ray/tune/trial_runner.py", line 472, in _process_trial
result = self.trial_executor.fetch_result(trial)
File "/home/user/anaconda3/envs/RLlibTesting/lib/python3.8/site-packages/ray/tune/ray_trial_executor.py", line 430, in fetch_result
result = ray.get(trial_future[0], DEFAULT_GET_TIMEOUT)
File "/home/user/anaconda3/envs/RLlibTesting/lib/python3.8/site-packages/ray/worker.py", line 1478, in get
raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(InvalidArgumentError): ray::PPO.train() (pid=20426, ip=192.168.2.105)
File "/home/user/anaconda3/envs/RLlibTesting/lib/python3.8/site-packages/tensorflow/python/client/session.py", line 1349, in _run_fn
return self._call_tf_sessionrun(options, feed_dict, fetch_list,
File "/home/user/anaconda3/envs/RLlibTesting/lib/python3.8/site-packages/tensorflow/python/client/session.py", line 1441, in _call_tf_sessionrun
return tf_session.TF_SessionRun_wrapper(self._session, options, feed_dict,
tensorflow.python.framework.errors_impl.InvalidArgumentError: Incompatible shapes: [120,64] vs. [6,64]
[[{{node default_policy_1/tower_1/model_1/lstm/while/lstm_cell/add_6}}]]

During handling of the above exception, another exception occurred:

ray::PPO.train() (pid=20426, ip=192.168.2.105)
File "python/ray/_raylet.pyx", line 443, in ray._raylet.execute_task
File "python/ray/_raylet.pyx", line 446, in ray._raylet.execute_task
File "python/ray/_raylet.pyx", line 447, in ray._raylet.execute_task
File "python/ray/_raylet.pyx", line 401, in ray._raylet.execute_task.function_executor
File "/home/user/anaconda3/envs/RLlibTesting/lib/python3.8/site-packages/ray/rllib/agents/trainer.py", line 520, in train
raise e
File "/home/user/anaconda3/envs/RLlibTesting/lib/python3.8/site-packages/ray/rllib/agents/trainer.py", line 506, in train
result = Trainable.train(self)
File "/home/user/anaconda3/envs/RLlibTesting/lib/python3.8/site-packages/ray/tune/trainable.py", line 317, in train
result = self._train()
File "/home/user/anaconda3/envs/RLlibTesting/lib/python3.8/site-packages/ray/rllib/agents/trainer_template.py", line 137, in _train
return self._train_exec_impl()
File "/home/user/anaconda3/envs/RLlibTesting/lib/python3.8/site-packages/ray/rllib/agents/trainer_template.py", line 175, in _train_exec_impl
res = next(self.train_exec_impl)
File "/home/user/anaconda3/envs/RLlibTesting/lib/python3.8/site-packages/ray/util/iter.py", line 731, in __next__
return next(self.built_iterator)
File "/home/user/anaconda3/envs/RLlibTesting/lib/python3.8/site-packages/ray/util/iter.py", line 744, in apply_foreach
for item in it:
File "/home/user/anaconda3/envs/RLlibTesting/lib/python3.8/site-packages/ray/util/iter.py", line 744, in apply_foreach
for item in it:
File "/home/user/anaconda3/envs/RLlibTesting/lib/python3.8/site-packages/ray/util/iter.py", line 814, in apply_filter
for item in it:
File "/home/user/anaconda3/envs/RLlibTesting/lib/python3.8/site-packages/ray/util/iter.py", line 814, in apply_filter
for item in it:
File "/home/user/anaconda3/envs/RLlibTesting/lib/python3.8/site-packages/ray/util/iter.py", line 744, in apply_foreach
for item in it:
File "/home/user/anaconda3/envs/RLlibTesting/lib/python3.8/site-packages/ray/util/iter.py", line 744, in apply_foreach
for item in it:
File "/home/user/anaconda3/envs/RLlibTesting/lib/python3.8/site-packages/ray/util/iter.py", line 752, in apply_foreach
result = fn(item)
File "/home/user/anaconda3/envs/RLlibTesting/lib/python3.8/site-packages/ray/rllib/execution/train_ops.py", line 204, in __call__
batch_fetches = optimizer.optimize(
File "/home/user/anaconda3/envs/RLlibTesting/lib/python3.8/site-packages/ray/rllib/execution/multi_gpu_impl.py", line 257, in optimize
return sess.run(fetches, feed_dict=feed_dict)
File "/home/user/anaconda3/envs/RLlibTesting/lib/python3.8/site-packages/tensorflow/python/client/session.py", line 957, in run
result = self._run(None, fetches, feed_dict, options_ptr,
File "/home/user/anaconda3/envs/RLlibTesting/lib/python3.8/site-packages/tensorflow/python/client/session.py", line 1180, in _run
results = self._do_run(handle, final_targets, final_fetches,
File "/home/user/anaconda3/envs/RLlibTesting/lib/python3.8/site-packages/tensorflow/python/client/session.py", line 1358, in _do_run
return self._do_call(_run_fn, feeds, fetches, targets, options,
File "/home/user/anaconda3/envs/RLlibTesting/lib/python3.8/site-packages/tensorflow/python/client/session.py", line 1384, in _do_call
raise type(e)(node_def, op, message)
tensorflow.python.framework.errors_impl.InvalidArgumentError: Incompatible shapes: [120,64] vs. [6,64]
[[node default_policy_1/tower_1/model_1/lstm/while/lstm_cell/add_6 (defined at mnt/c/Users/user/Desktop/RLlib_Issue/rnn_model.py:81) ]]

# Fixed code
def postprocess_trajectory(
    policy: Policy,
    sample_batch: SampleBatch,
    other_agent_batches: Optional[Dict[AgentID, SampleBatch]] = None,
    episode: Optional[MultiAgentEpisode] = None,
) -> SampleBatch:
    """Postprocesses a trajectory and returns the processed trajectory.

    The trajectory contains only data from one episode and from one agent.
    - If  `config.batch_mode=truncate_episodes` (default), sample_batch may
    contain a truncated (at-the-end) episode, in case the
    `config.rollout_fragment_length` was reached by the sampler.
    - If `config.batch_mode=complete_episodes`, sample_batch will contain
    exactly one episode (no matter how long).
    New columns can be added to sample_batch and existing ones may be altered.

    Args:
        policy (Policy): The Policy used to generate the trajectory
            (`sample_batch`)
        sample_batch (SampleBatch): The SampleBatch to postprocess.
        other_agent_batches (Optional[Dict[PolicyID, SampleBatch]]): Optional
            dict of AgentIDs mapping to other agents' trajectory data (from the
            same episode). NOTE: The other agents use the same policy.
        episode (Optional[MultiAgentEpisode]): Optional multi-agent episode
            object in which the agents operated.

    Returns:
        SampleBatch: The postprocessed, modified SampleBatch (or a new one).
    """
    if not policy.config["vtrace"]:
        sample_batch = compute_gae_for_sample_batch(
            policy, sample_batch, other_agent_batches, episode
        )

    # TODO: (sven) remove this del once we have trajectory view API fully in
    #  place.
    del sample_batch.data["new_obs"]  # not used, so save some bandwidth

    return sample_batch