docstr: Tests for a full batch.
docstr_tok: ["Tests","for","a","full","batch","."]
code: def test_clipped_surrogate_pg_loss_batch(self):
      """Tests for a full batch."""
      get_ratios = jax.vmap(
          distributions.categorical_importance_sampling_ratios
      )
      prob_ratios = get_ratios(
          self.logits, self.old_logits, self.actions
      )
      batched_fn_variant = self.variant(
          jax.vmap(
              functools.partial(
                  policy_gradients.clipped_surrogate_pg_loss,
                  epsilon=self.epsilon,
              )
          )
      )
      actual = batched_fn_variant(prob_ratios, self.advantages)
      np.testing.assert_allclose(actual, self.expected, atol=1e-4)
  
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 distributions
  from rlax._src import policy_gradients
label_desc: Use image node_id 2 for calling the ClippedSurrogatePGLossTest obj's underlying member method code with example usage: obj.test_clipped_surrogate_pg_loss_batch() without return types