# Buggy code
def postprocess_ppo_gae(
    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).
    """

    # Trajectory is actually complete -> last r=0.0.
    if sample_batch[SampleBatch.DONES][-1]:
        last_r = 0.0
    # Trajectory has been truncated -> last r=VF estimate of last obs.
    else:
        # Input dict is provided to us automatically via the Model's
        # requirements. It's a single-timestep (last one in trajectory)
        # input_dict.
        if policy.config["_use_trajectory_view_api"]:
            # Create an input dict according to the Model's requirements.
            input_dict = policy.model.get_input_dict(sample_batch, index="last")
            last_r = policy._value(**input_dict)
        # TODO: (sven) Remove once trajectory view API is all-algo default.
        else:
            next_state = []
            for i in range(policy.num_state_tensors()):
                next_state.append(sample_batch["state_out_{}".format(i)][-1])
            last_r = policy._value(
                sample_batch[SampleBatch.NEXT_OBS][-1],
                sample_batch[SampleBatch.ACTIONS][-1],
                sample_batch[SampleBatch.REWARDS][-1],
                *next_state,
            )

    # Adds the policy logits, VF preds, and advantages to the batch,
    # using GAE ("generalized advantage estimation") or not.
    batch = compute_advantages(
        sample_batch,
        last_r,
        policy.config["gamma"],
        policy.config["lambda"],
        use_gae=policy.config["use_gae"],
        use_critic=policy.config.get("use_critic", True),
    )

    return 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_ppo_gae(
    policy: Policy,
    sample_batch: SampleBatch,
    other_agent_batches: Optional[Dict[AgentID, SampleBatch]] = None,
    episode: Optional[MultiAgentEpisode] = None,
) -> SampleBatch:
    # Stub serving backward compatibility.
    deprecation_warning(
        old="rllib.agents.ppo.ppo_tf_policy.postprocess_ppo_gae",
        new="rllib.evaluation.postprocessing.compute_gae_for_sample_batch",
        error=False,
    )

    return compute_gae_for_sample_batch(
        policy, sample_batch, other_agent_batches, episode
    )