code: def setUp(self):
      super().setUp()
      self.cell_size = 2
      time, self.batch_size = 3, 2
      height, width, channels = 4, 4, 3
  
      shape = (time, self.batch_size, height, width, channels)
      hw = np.matmul(
          np.arange(0, 1, 0.25)[:, None], np.arange(0, 1, 0.25)[None, :]
      )
      hwc = np.stack([hw, hw + 0.1, hw + 0.2], axis=-1)
      bhwc = np.stack([hwc, hwc + 0.1], axis=0)
      tbhwc = np.stack([bhwc, bhwc + 0.05, bhwc + 0.1], axis=0)
      assert tbhwc.shape == shape
      self.obs = tbhwc
  
      self.expected = 0.05 * np.ones((2, 2, 2, 2), dtype=np.float32)
  
file_path: rlax/rlax/_src/general_value_functions_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 general_value_functions
label_desc: Use image node_id 1 for calling the PixelControlTest obj's underlying member method code with example usage: obj.setUp() without return types