docstr: Tests for a single element.
docstr_tok: ["Tests","for","a","single","element","."]
code: def test_pixel_control_rewards(self):
      """Tests for a single element."""
      pixel_control_rewards = self.variant(
          functools.partial(
              general_value_functions.pixel_control_rewards,
              cell_size=self.cell_size,
          )
      )
      # Test pseudo rewards.
      for i in range(self.batch_size):
          rs = pixel_control_rewards(self.obs[:, i])
          np.testing.assert_allclose(self.expected[:, i], rs, rtol=1e-5)
  
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 2 for calling the PixelControlTest obj's underlying member method code with example usage: obj.test_pixel_control_rewards() without return types