def test_box(self):
env = Gym(env_name="CartPole-v0")
obs_preprocessor = env.get_obs_preprocessor()
obs = env.reset()
state = obs_preprocessor(obs)
self.assertTrue(state.has_float_features_only)
self.assertEqual(state.float_features.shape, (1, obs.shape[0]))
self.assertEqual(state.float_features.dtype, torch.float32)
self.assertEqual(state.float_features.device, torch.device("cpu"))
npt.assert_array_almost_equal(obs, state.float_features.squeeze(0))
def test_box_cuda(self):
env = Gym(env_name="CartPole-v0")
device = torch.device("cuda")
obs_preprocessor = env.get_obs_preprocessor(device=device)
obs = env.reset()
state = obs_preprocessor(obs)
self.assertTrue(state.has_float_features_only)
self.assertEqual(state.float_features.shape, (1, obs.shape[0]))
self.assertEqual(state.float_features.dtype, torch.float32)
# `device` doesn't have index. So we need this.
x = torch.zeros(1, device=device)
self.assertEqual(state.float_features.device, x.device)
npt.assert_array_almost_equal(obs, state.float_features.cpu().squeeze(0))