def testStepSavesCurrentTimeStep(self):
obs_spec = BoundedTensorSpec((1, ), torch.int32)
action_spec = BoundedTensorSpec((1, ), torch.int64)
random_env = RandomAlfEnvironment(observation_spec=obs_spec,
action_spec=action_spec)
random_env.reset()
time_step = random_env.step(action=torch.ones((1, )))
current_time_step = random_env.current_time_step()
nest.map_structure(self.assertEqual, time_step, current_time_step)
def testRewardCheckerBatchSizeOne(self):
# Ensure batch size 1 with scalar reward works
obs_spec = BoundedTensorSpec((2, 3), torch.int32, -10, 10)
action_spec = BoundedTensorSpec((1, ), torch.int64)
env = RandomAlfEnvironment(obs_spec,
action_spec,
reward_fn=lambda *_: np.array([1.0]),
batch_size=1)
env._done = False
env.reset()
action = torch.tensor([0], dtype=torch.int64)
time_step = env.step(action)
self.assertEqual(time_step.reward, 1.0)
def testCustomRewardFn(self):
obs_spec = BoundedTensorSpec((2, 3), torch.int32, -10, 10)
action_spec = BoundedTensorSpec((1, ), torch.int64)
batch_size = 3
env = RandomAlfEnvironment(obs_spec,
action_spec,
reward_fn=lambda *_: np.ones(batch_size),
batch_size=batch_size)
env._done = False
env.reset()
action = torch.ones(batch_size)
time_step = env.step(action)
self.assertSequenceAlmostEqual([1.0] * 3, time_step.reward)
def testRewardCheckerSizeMismatch(self):
# Ensure custom scalar reward with batch_size greater than 1 raises
# ValueError
obs_spec = BoundedTensorSpec((2, 3), torch.int32, -10, 10)
action_spec = BoundedTensorSpec((1, ), torch.int64)
env = RandomAlfEnvironment(obs_spec,
action_spec,
reward_fn=lambda *_: np.array([1.0]),
batch_size=5)
env.reset()
env._done = False
action = torch.tensor(0, dtype=torch.int64)
with self.assertRaises(ValueError):
env.step(action)
def testRendersImage(self):
action_spec = BoundedTensorSpec((1, ), torch.int64, -10, 10)
observation_spec = BoundedTensorSpec((1, ), torch.int32, -10, 10)
env = RandomAlfEnvironment(observation_spec,
action_spec,
render_size=(4, 4, 3))
env.reset()
img = env.render()
self.assertTrue(np.all(img < 256))
self.assertTrue(np.all(img >= 0))
self.assertEqual((4, 4, 3), img.shape)
self.assertEqual(np.uint8, img.dtype)
