docstr: Tests truncated GAE for a full batch.
docstr_tok: ["Tests","truncated","GAE","for","a","full","batch","."]
code: def test_truncated_gae(self, lambda_):
      """Tests truncated GAE for a full batch."""
      batched_advantage_fn_variant = self.variant(
          jax.vmap(
              multistep.truncated_generalized_advantage_estimation,
              in_axes=(0, 0, None, 0),
              out_axes=0,
          )
      )
      actual = batched_advantage_fn_variant(
          self.r_t, self.discount_t, lambda_, self.v_t
      )
      np.testing.assert_allclose(
          self.expected[lambda_], actual, 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 2 for calling the TruncatedGeneralizedAdvantageEstimationTest obj's underlying member method code with example usage: obj.test_truncated_gae(lambda_) without return types