docstr: Tests equivalence to computing the error from a the lambda-return.
docstr_tok: ["Tests","equivalence","to","computing","the","error","from","a","the","lambda-return","."]
code: def test_importance_corrected_td_errors_batch(self):
      """Tests equivalence to computing the error from a the lambda-return."""
      # Vmap and optionally compile.
      lambda_returns = self.variant(jax.vmap(multistep.lambda_returns))
      td_errors = self.variant(
          jax.vmap(multistep.importance_corrected_td_errors)
      )
      # Compute multistep td-error with recursion on deltas.
      td_direct = td_errors(
          self.r_t,
          self.discount_t,
          self.rho_tm1,
          np.ones_like(self.discount_t),
          self.values,
      )
      # Compute off-policy corrected return, and derive td-error from it.
      ls_ = np.concatenate(
          (self.rho_tm1[:, 1:], [[1.0], [1.0]]), axis=1
      )
      td_from_returns = self.rho_tm1 * (
          lambda_returns(
              self.r_t, self.discount_t, self.values[:, 1:], ls_
          )
          - self.values[:, :-1]
      )
      # Check equivalence.
      np.testing.assert_allclose(td_direct, td_from_returns, rtol=1e-5)
  
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  import chex
  import jax
  import jax.numpy
  import numpy
  from rlax._src import multistep
label_desc: Use image node_id 2 for calling the TDErrorTest obj's underlying member method code with example usage: obj.test_importance_corrected_td_errors_batch() without return types