docstr: Tests truncated GAE.
  
  Tests that truncated GAE yields same output as importance corrected
  td errors with dummy ratios.
  
  Args:
    lambda_: a lambda to use in GAE.
docstr_tok: ["Tests","truncated","GAE",".","Tests","that","truncated","GAE","yields","same","output","as","importance","corrected","td","errors","with","dummy","ratios",".","Args",":","lambda_",":","a","lambda","to","use","in","GAE","."]
code: def test_gae_as_special_case_of_importance_corrected_td_errors(
      self, lambda_
  ):
      """Tests truncated GAE.
  
      Tests that truncated GAE yields same output as importance corrected
      td errors with dummy ratios.
  
      Args:
        lambda_: a lambda to use in GAE.
      """
      batched_gae_fn_variant = self.variant(
          jax.vmap(
              multistep.truncated_generalized_advantage_estimation,
              in_axes=(0, 0, None, 0),
              out_axes=0,
          )
      )
      gae_result = batched_gae_fn_variant(
          self.r_t, self.discount_t, lambda_, self.v_t
      )
  
      batched_ictd_errors_fn_variant = self.variant(
          jax.vmap(multistep.importance_corrected_td_errors)
      )
      ictd_errors_result = batched_ictd_errors_fn_variant(
          self.r_t,
          self.discount_t,
          self.dummy_rho_tm1,
          jnp.ones_like(self.discount_t) * lambda_,
          self.v_t,
      )
      np.testing.assert_allclose(
          gae_result, ictd_errors_result, atol=1e-3
      )
  
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 4 for calling the TruncatedGeneralizedAdvantageEstimationTest obj's underlying member method code with example usage: obj.test_gae_as_special_case_of_importance_corrected_td_errors(lambda_) without return types