code: def truncated_generalized_advantage_estimation(
      r_t: Array,
      discount_t: Array,
      lambda_: Union[Array, Scalar],
      values: Array,
      stop_target_gradients: bool = False,
  ) -> Array:
      """Computes truncated generalized advantage estimates for a sequence length k.
  
      The advantages are computed in a backwards fashion according to the equation:
      Âₜ = δₜ + (γλ) * δₜ₊₁ + ... + ... + (γλ)ᵏ⁻ᵗ⁺¹ * δₖ₋₁
      where δₜ = rₜ₊₁ + γₜ₊₁ * v(sₜ₊₁) - v(sₜ).
  
      See Proximal Policy Optimization Algorithms, Schulman et al.:
      https://arxiv.org/abs/1707.06347
  
      Note: This paper uses a different notation than the RLax standard
      convention that follows Sutton & Barto. We use rₜ₊₁ to denote the reward
      received after acting in state sₜ, while the PPO paper uses rₜ.
  
      Args:
        r_t: Sequence of rewards at times [1, k]
        discount_t: Sequence of discounts at times [1, k]
        lambda_: Mixing parameter; a scalar or sequence of lambda_t at times [1, k]
        values: Sequence of values under π at times [0, k]
        stop_target_gradients: bool indicating whether or not to apply stop gradient
          to targets.
  
      Returns:
        Multistep truncated generalized advantage estimation at times [0, k-1].
      """
      chex.assert_rank([r_t, values, discount_t], 1)
      chex.assert_type([r_t, values, discount_t], float)
      lambda_ = (
          jnp.ones_like(discount_t) * lambda_
      )  # If scalar, make into vector.
  
      delta_t = r_t + discount_t * values[1:] - values[:-1]
  
      # Iterate backwards to calculate advantages.
      def _body(acc, xs):
          deltas, discounts, lambda_ = xs
          acc = deltas + discounts * lambda_ * acc
          return acc, acc
  
      _, advantage_t = jax.lax.scan(
          _body, 0.0, (delta_t, discount_t, lambda_), reverse=True
      )
  
      return jax.lax.select(
          stop_target_gradients,
          jax.lax.stop_gradient(advantage_t),
          advantage_t,
      )
  
file_path: rlax/rlax/_src/multistep.py
imports: from typing import Union
  import chex
  import jax
  import jax.numpy
  from rlax._src import base
label_desc: Use image node_id 5 for calling a global function with example usage: truncated_generalized_advantage_estimation(r_t, discount_t, lambda_, values, stop_target_gradients) and returns: jax, acc,  acc