code: def discounted_returns(
      r_t: Array,
      discount_t: Array,
      v_t: Array,
      stop_target_gradients: bool = False,
  ) -> Array:
      """Calculates a discounted return from a trajectory.
  
      The returns are computed recursively, from `G_{T-1}` to `G_0`, according to:
  
        Gₜ = rₜ₊₁ + γₜ₊₁ Gₜ₊₁.
  
      See "Reinforcement Learning: An Introduction" by Sutton and Barto.
      (http://incompleteideas.net/sutton/book/ebook/node61.html).
  
      Args:
        r_t: reward sequence at time t.
        discount_t: discount sequence at time t.
        v_t: value sequence or scalar at time t.
        stop_target_gradients: bool indicating whether or not to apply stop gradient
          to targets.
  
      Returns:
        Discounted returns.
      """
      chex.assert_rank([r_t, discount_t, v_t], [1, 1, {0, 1}])
      chex.assert_type([r_t, discount_t, v_t], float)
  
      # If scalar make into vector.
      bootstrapped_v = jnp.ones_like(discount_t) * v_t
      return lambda_returns(
          r_t,
          discount_t,
          bootstrapped_v,
          lambda_=1.0,
          stop_target_gradients=stop_target_gradients,
      )
  
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 3 for calling a global function with example usage: discounted_returns(r_t, discount_t, v_t, stop_target_gradients) and returns: lambda_returns