docstr: Tests that the rewards are valid.
docstr_tok: ["Tests","that","the","rewards","are","valid","."]
code: def test_rewards_validity(self, batch_size):
      """Tests that the rewards are valid."""
      env = wheel_py_environment.WheelPyEnvironment(
          delta=0.5,
          mu_base=[1.2, 1.0, 1.0, 1.0, 1.0],
          std_base=0.01 * np.ones(5),
          mu_high=50.0,
          std_high=0.01,
          batch_size=batch_size,
      )
      time_step = env.reset()
      time_step = env.step(np.arange(batch_size))
      self.assertEqual(time_step.reward.shape, (batch_size,))
  
file_path: agents/tf_agents/bandits/environments/wheel_py_environment_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import numpy
  import tensorflow
  from tf_agents.bandits.environments import wheel_py_environment
label_desc: Use image node_id 6 for calling the WheelBanditPyEnvironmentTest obj's underlying member method code with example usage: obj.test_rewards_validity(batch_size) without return types