class_docstr: Test the `pixel_control_loss` op.
docstr: Defines example data and expected result for the op.
docstr_tok: ["Defines","example","data","and","expected","result","for","the","op","."]
code: def setUp(self):
      """Defines example data and expected result for the op."""
      super().setUp()
  
      # Observation shape is (2,2,3) (i.e., height 2, width 2, and 3 channels).
      # We will use no cropping, and a cell size of 1. We have num_actions = 3,
      # meaning our Q values should be (2,2,3). We will set the Q value equal to
      # the observation.
      self.seq_length = 3
      self.discount = 0.9
      self.cell_size = 1
  
      # Observations.
      obs1 = np.array([[[1, 2, 3], [3, 4, 5]], [[5, 6, 7], [7, 8, 9]]])
      obs2 = np.array([[[7, 8, 9], [1, 2, 3]], [[3, 4, 5], [5, 6, 7]]])
      obs3 = np.array([[[5, 6, 7], [7, 8, 9]], [[1, 2, 3], [3, 4, 5]]])
      obs4 = np.array([[[3, 4, 5], [5, 6, 7]], [[7, 8, 9], [1, 2, 3]]])
  
      # Actions.
      action1 = 0
      action2 = 1
      action3 = 2
  
      # Compute loss for constant discount.
      qa_tm1 = obs3[:, :, action3]
      reward3 = np.mean(np.abs(obs4 - obs3), axis=2)
      qmax_t = np.amax(obs4, axis=2)
      target = reward3 + self.discount * qmax_t
      error3 = target - qa_tm1
  
      qa_tm1 = obs2[:, :, action2]
      reward2 = np.mean(np.abs(obs3 - obs2), axis=2)
      target = reward2 + self.discount * target
      error2 = target - qa_tm1
  
      qa_tm1 = obs1[:, :, action1]
      reward1 = np.mean(np.abs(obs2 - obs1), axis=2)
      target = reward1 + self.discount * target
      error1 = target - qa_tm1
  
      # Compute loss for episode termination with discount 0.
      qa_tm1 = obs1[:, :, action1]
      reward1 = np.mean(np.abs(obs2 - obs1), axis=2)
      target = reward1 + 0.0 * target
      error1_term = target - qa_tm1
  
      self.error = (
          np.sum(
              np.square(error1) + np.square(error2) + np.square(error3)
          )
          * 0.5
      )
      self.error_term = (
          np.sum(
              np.square(error1_term)
              + np.square(error2)
              + np.square(error3)
          )
          * 0.5
      )
  
      self.observations = np.stack(
          [obs1, obs2, obs3, obs4], axis=0
      ).astype(np.float32)
      self.action_values = self.observations
      self.actions = np.array([action1, action2, action3])
  
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 losses
label_desc: Use image node_id 1 for calling the PixelControlLossTest obj's underlying member method code with example usage: obj.setUp() without return types