code: def general_off_policy_returns_from_action_values(
      q_t: Array,
      a_t: Array,
      r_t: Array,
      discount_t: Array,
      c_t: Array,
      pi_t: Array,
      stop_target_gradients: bool = False,
  ) -> Array:
      """Calculates targets for various off-policy correction algorithms.
  
      Given a window of experience of length `K`, generated by a behaviour policy μ,
      for each time-step `t` we can estimate the return `G_t` from that step
      onwards, under some target policy π, using the rewards in the trajectory, the
      actions selected by μ and the action-values under π, according to equation:
  
        Gₜ = rₜ₊₁ + γₜ₊₁ * (E[q(aₜ₊₁)] - cₜ * q(aₜ₊₁) + cₜ * Gₜ₊₁),
  
      where, depending on the choice of `c_t`, the algorithm implements:
  
        Importance Sampling             c_t = π(x_t, a_t) / μ(x_t, a_t),
        Harutyunyan's et al. Q(lambda)  c_t = λ,
        Precup's et al. Tree-Backup     c_t = π(x_t, a_t),
        Munos' et al. Retrace           c_t = λ min(1, π(x_t, a_t) / μ(x_t, a_t)).
  
      See "Safe and Efficient Off-Policy Reinforcement Learning" by Munos et al.
      (https://arxiv.org/abs/1606.02647).
  
      Args:
        q_t: Q-values at times [1, ..., K - 1].
        a_t: action index at times [1, ..., K - 1].
        r_t: reward at times [1, ..., K - 1].
        discount_t: discount at times [1, ..., K - 1].
        c_t: importance weights at times [1, ..., K - 1].
        pi_t: target policy probs at times [1, ..., K - 1].
        stop_target_gradients: bool indicating whether or not to apply stop gradient
          to targets.
  
      Returns:
        Off-policy estimates of the generalized returns from states visited at times
        [0, ..., K - 1].
      """
      chex.assert_rank(
          [q_t, a_t, r_t, discount_t, c_t, pi_t], [2, 1, 1, 1, 1, 2]
      )
      chex.assert_type(
          [q_t, a_t, r_t, discount_t, c_t, pi_t],
          [float, int, float, float, float, float],
      )
      chex.assert_equal_shape(
          [q_t[..., 0], a_t, r_t, discount_t, c_t, pi_t[..., 0]]
      )
  
      # Get the expected values and the values of actually selected actions.
      exp_q_t = (pi_t * q_t).sum(axis=-1)
      # The generalized returns are independent of Q-values and cs at the final
      # state.
      q_a_t = base.batched_index(q_t, a_t)[:-1]
      c_t = c_t[:-1]
  
      return general_off_policy_returns_from_q_and_v(
          q_a_t, exp_q_t, r_t, discount_t, c_t, 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 6 for calling a global function with example usage: general_off_policy_returns_from_action_values(q_t, a_t, r_t, discount_t, c_t, pi_t, stop_target_gradients) and returns: general_off_policy_returns_from_q_and_v