def testRewardCheckerSizeMismatch(self):
    obs_spec = array_spec.BoundedArraySpec((2, 3), np.int32, (- 10), 10)
    env = random_py_environment.RandomPyEnvironment(obs_spec, reward_fn=(lambda *_: 1.0), batch_size=5)
    env.reset()
    env._done = False
    with self.assertRaises(ValueError):
        env.step([0])
def testRewardCheckerSizeMismatch(self):
    obs_spec = array_spec.BoundedArraySpec((2, 3), np.int32, (- 10), 10)
    env = random_py_environment.RandomPyEnvironment(obs_spec, reward_fn=(lambda *_: 1.0), batch_size=5)
    env.reset()
    env._done = False
    with self.assertRaises(ValueError):
        env.step([0])