# Buggy code
def postprocess_advantages(
    policy, sample_batch, other_agent_batches=None, episode=None
):
    completed = sample_batch[SampleBatch.DONES][-1]
    if completed:
        last_r = 0.0
    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,
        )
    return compute_advantages(
        sample_batch,
        last_r,
        policy.config["gamma"],
        policy.config["lambda"],
        policy.config["use_gae"],
        policy.config["use_critic"],
    )


# 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_advantages(
    policy, sample_batch, other_agent_batches=None, episode=None
):
    # Stub serving backward compatibility.
    deprecation_warning(
        old="rllib.agents.a3c.a3c_tf_policy.postprocess_advantages",
        new="rllib.evaluation.postprocessing.compute_gae_for_sample_batch",
        error=False,
    )

    return compute_gae_for_sample_batch(
        policy, sample_batch, other_agent_batches, episode
    )