docstr: Tests that truncated GAE is consistent with scalar or array lambda_.
docstr_tok: ["Tests","that","truncated","GAE","is","consistent","with","scalar","or","array","lambda_","."]
code: def test_array_lambda(self):
      """Tests that truncated GAE is consistent with scalar or array lambda_."""
      scalar_lambda_fn = self.variant(
          jax.vmap(
              multistep.truncated_generalized_advantage_estimation,
              in_axes=(0, 0, None, 0),
              out_axes=0,
          )
      )
      array_lambda_fn = self.variant(
          jax.vmap(multistep.truncated_generalized_advantage_estimation)
      )
      scalar_lambda_result = scalar_lambda_fn(
          self.r_t, self.discount_t, 0.9, self.v_t
      )
      array_lambda_result = array_lambda_fn(
          self.r_t, self.discount_t, self.array_lambda, self.v_t
      )
      np.testing.assert_allclose(
          scalar_lambda_result, array_lambda_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 3 for calling the TruncatedGeneralizedAdvantageEstimationTest obj's underlying member method code with example usage: obj.test_array_lambda() without return types