code: def n_step_bootstrapped_returns(
      r_t: Array,
      discount_t: Array,
      v_t: Array,
      n: int,
      lambda_t: Numeric = 1.0,
      stop_target_gradients: bool = False,
  ) -> Array:
      """Computes strided n-step bootstrapped return targets over a sequence.
  
      The returns are computed according to the below equation iterated `n` times:
  
         Gₜ = rₜ₊₁ + γₜ₊₁ [(1 - λₜ₊₁) vₜ₊₁ + λₜ₊₁ Gₜ₊₁].
  
      When lambda_t == 1. (default), this reduces to
  
         Gₜ = rₜ₊₁ + γₜ₊₁ * (rₜ₊₂ + γₜ₊₂ * (... * (rₜ₊ₙ + γₜ₊ₙ * vₜ₊ₙ ))).
  
      Args:
        r_t: rewards at times [1, ..., T].
        discount_t: discounts at times [1, ..., T].
        v_t: state or state-action values to bootstrap from at time [1, ...., T].
        n: number of steps over which to accumulate reward before bootstrapping.
        lambda_t: lambdas at times [1, ..., T]. Shape is [], or [T-1].
        stop_target_gradients: bool indicating whether or not to apply stop gradient
          to targets.
  
      Returns:
        estimated bootstrapped returns at times [0, ...., T-1]
      """
      chex.assert_rank(
          [r_t, discount_t, v_t, lambda_t], [1, 1, 1, {0, 1}]
      )
      chex.assert_type([r_t, discount_t, v_t, lambda_t], float)
      chex.assert_equal_shape([r_t, discount_t, v_t])
      seq_len = r_t.shape[0]
  
      # Maybe change scalar lambda to an array.
      lambda_t = jnp.ones_like(discount_t) * lambda_t
  
      # Shift bootstrap values by n and pad end of sequence with last value v_t[-1].
      pad_size = min(n - 1, seq_len)
      targets = jnp.concatenate(
          [v_t[n - 1 :], jnp.array([v_t[-1]] * pad_size)]
      )
  
      # Pad sequences. Shape is now (T + n - 1,).
      r_t = jnp.concatenate([r_t, jnp.zeros(n - 1)])
      discount_t = jnp.concatenate([discount_t, jnp.ones(n - 1)])
      lambda_t = jnp.concatenate([lambda_t, jnp.ones(n - 1)])
      v_t = jnp.concatenate([v_t, jnp.array([v_t[-1]] * (n - 1))])
  
      # Work backwards to compute n-step returns.
      for i in reversed(range(n)):
          r_ = r_t[i : i + seq_len]
          discount_ = discount_t[i : i + seq_len]
          lambda_ = lambda_t[i : i + seq_len]
          v_ = v_t[i : i + seq_len]
          targets = r_ + discount_ * (
              (1.0 - lambda_) * v_ + lambda_ * targets
          )
  
      return jax.lax.select(
          stop_target_gradients, jax.lax.stop_gradient(targets), targets
      )
  
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 2 for calling a global function with example usage: n_step_bootstrapped_returns(r_t, discount_t, v_t, n, lambda_t, stop_target_gradients) and returns: jax