def testRewardFnCalled(self):

    def reward_fn(unused_step_type, action, unused_observation):
        return action
    action_spec = array_spec.BoundedArraySpec((1,), np.int32, (- 10), 10)
    observation_spec = array_spec.BoundedArraySpec((1,), np.int32, (- 10), 10)
    env = random_py_environment.RandomPyEnvironment(observation_spec, action_spec, reward_fn=reward_fn)
    time_step = env.step(1)
    self.assertEqual(0.0, time_step.reward)
    time_step = env.step(1)
    self.assertEqual(1, time_step.reward)
def testRewardFnCalled(self):

    def reward_fn(unused_step_type, action, unused_observation):
        return action
    action_spec = array_spec.BoundedArraySpec((1,), np.int32, (- 10), 10)
    observation_spec = array_spec.BoundedArraySpec((1,), np.int32, (- 10), 10)
    env = random_py_environment.RandomPyEnvironment(observation_spec, action_spec, reward_fn=reward_fn)
    time_step = env.step(1)
    self.assertEqual(0.0, time_step.reward)
    time_step = env.step(1)
    self.assertEqual(1, time_step.reward)