code: def setUp(self):
      super().setUp()
  
      self.q_values = jnp.array([[0.0, -1.0, 1.0], [1.0, -1.0, 0]])
      self.policy_logits = jnp.array([[1.0, 1.0, 1.0], [1.0, 1.0, 4.0]])
  
      # baseline = \sum_a pi_a * Q_a = 0.
      # -\sum_a pi_a * relu(Q_a - baseline)
      # negative sign as it's a loss term and loss needs to be minimized.
      self.expected_policy_loss = (0.0 + 0.0) / 2
  
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 1 for calling the QPGLossTest obj's underlying member method code with example usage: obj.setUp() without return types