code: def compute_gae_for_episode(
      episode: SingleAgentEpisode,
      config: AlgorithmConfig,
      module: RLModule,
  ):
      """Adds GAE to a trajectory."""
      # TODO (simon): All of this can be batched over multiple episodes.
      # This should increase performance.
      # TODO (sven): Shall do postprocessing in the training_step or
      # in the env_runner? Here we could batch over episodes as we have
      # them now in the training_step.
      episode = compute_bootstrap_value(episode, module)
  
      vf_preds = episode.get_extra_model_outputs(SampleBatch.VF_PREDS)
      rewards = episode.get_rewards()
  
      # TODO (simon): In case of recurrent models sequeeze out time dimension.
  
      episode = compute_advantages(
          episode,
          last_r=episode.extra_model_outputs[
              SampleBatch.VALUES_BOOTSTRAPPED
          ][-1],
          gamma=config["gamma"],
          lambda_=config["lambda"],
          use_gae=config["use_gae"],
          use_critic=config.get("use_critic", True),
          vf_preds=vf_preds,
          rewards=rewards,
      )
  
      # TODO (simon): Add dimension in case of recurrent model.
      return episode
  
file_path: ray/python/ray/rllib/evaluation/postprocessing_v2.py
imports: from typing import List
  import numpy
  import tree
  from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
  from ray.rllib.core.models.base import STATE_IN
  from ray.rllib.core.rl_module.rl_module import RLModule
  from ray.rllib.env.single_agent_episode import SingleAgentEpisode
  from ray.rllib.evaluation.postprocessing import discount_cumsum
  from ray.rllib.policy.sample_batch import concat_samples, SampleBatch
  from ray.rllib.utils.annotations import DeveloperAPI
  from ray.rllib.utils.framework import try_import_tf
  from ray.rllib.utils.nested_dict import NestedDict
  from ray.rllib.utils.numpy import convert_to_numpy
  from ray.rllib.utils.torch_utils import convert_to_torch_tensor
  from ray.rllib.utils.typing import TensorType
label_desc: Use image node_id 2 for calling a global function with example usage: compute_gae_for_episode(episode, config, module) and returns: episode