docstr: Tests for a full batch.
docstr_tok: ["Tests","for","a","full","batch","."]
code: def test_lambda_q_estimate(self):
      """Tests for a full batch."""
      lambda_ = 0.8
      vtrace_td_error_and_advantage = self.variant(
          jax.vmap(
              functools.partial(
                  vtrace.vtrace_td_error_and_advantage,
                  clip_rho_threshold=self._clip_rho_threshold,
                  lambda_=lambda_,
              )
          )
      )
      # Get function arguments.
      r_t, discount_t, rho_tm1, v_tm1, bootstrap_value = self._inputs
      v_t = np.concatenate(
          [v_tm1[:, 1:], bootstrap_value[:, None]], axis=1
      )
      # Compute vtrace output.
      vtrace_output = vtrace_td_error_and_advantage(
          v_tm1, v_t, r_t, discount_t, rho_tm1
      )
      expected_vs = vtrace_output.errors + v_tm1
      clipped_rho_tm1 = np.minimum(self._clip_rho_threshold, rho_tm1)
      vs_from_q = v_tm1 + clipped_rho_tm1 * (
          vtrace_output.q_estimate - v_tm1
      )
      # Test output.
      np.testing.assert_allclose(expected_vs, vs_from_q, rtol=1e-3)
  
file_path: rlax/rlax/_src/vtrace_test.py
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  import chex
  import jax
  import numpy
  from rlax._src import distributions
  from rlax._src import vtrace
label_desc: Use image node_id 4 for calling the VTraceTest obj's underlying member method code with example usage: obj.test_lambda_q_estimate() without return types