docstr: Tests for a full batch.
docstr_tok: ["Tests","for","a","full","batch","."]
code: def test_dpg_loss_batch(self):
      """Tests for a full batch."""
      dpg = self.variant(
          jax.vmap(
              functools.partial(
                  policy_gradients.dpg_loss, dqda_clipping=1.0
              )
          )
      )
      # Actor and critic function approximators.
      actor = lambda s_t: jnp.matmul(s_t, self.w_s) + self.b_s
      critic = lambda a_t: jnp.squeeze(jnp.matmul(a_t, self.w) + self.b)
      # Compute loss.
      a_t = actor(self.s_t)
      dqda = jax.vmap(jax.grad(critic))(a_t)
      # Test outputs.
      actual = np.sum(dpg(a_t, dqda), axis=1)
      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 DpgLossTest obj's underlying member method code with example usage: obj.test_dpg_loss_batch() without return types