| import unittest |
| import numpy as np |
| import gym |
| import rlbench.gym |
|
|
|
|
| class TestGym(unittest.TestCase): |
|
|
| def test_state_env(self): |
| env = gym.make('reach_target-state-v0') |
| env.reset() |
| obs, _, _, _ = env.step(env.action_space.sample()) |
| self.assertEqual(env.observation_space.shape, |
| obs.shape) |
| env.close() |
|
|
| def test_vision_env(self): |
| env = gym.make('reach_target-vision-v0') |
| env.reset() |
| obs, _, _, _ = env.step(env.action_space.sample()) |
| self.assertEqual(env.observation_space['state'].shape, |
| obs['state'].shape) |
| self.assertEqual(env.observation_space['left_shoulder_rgb'].shape, |
| obs['left_shoulder_rgb'].shape) |
| self.assertEqual(env.observation_space['right_shoulder_rgb'].shape, |
| obs['right_shoulder_rgb'].shape) |
| self.assertEqual(env.observation_space['wrist_rgb'].shape, |
| obs['wrist_rgb'].shape) |
| self.assertEqual(env.observation_space['front_rgb'].shape, |
| obs['front_rgb'].shape) |
| env.close() |
|
|
| def test_env_render(self): |
| env = gym.make('reach_target-vision-v0', render_mode='rgb_array') |
| env.reset() |
| obs, _, _, _ = env.step(env.action_space.sample()) |
| img = env.render('rgb_array') |
| self.assertGreater(np.mean(img), 0) |
| env.close() |
|
|
|
|
