| import unittest |
| from os import path |
|
|
| import numpy as np |
| from pyrep.const import RenderMode |
|
|
| from rlbench import Environment |
| from rlbench.action_modes.action_mode import MoveArmThenGripper |
| from rlbench.action_modes.arm_action_modes import JointVelocity |
| from rlbench.action_modes.gripper_action_modes import Discrete |
| from rlbench.observation_config import ObservationConfig, CameraConfig |
| from rlbench.sim2real.domain_randomization import RandomizeEvery, \ |
| VisualRandomizationConfig |
| from rlbench.tasks.open_microwave import OpenMicrowave |
|
|
| ASSET_DIR = path.join(path.dirname(path.abspath(__file__)), 'assets', 'textures') |
|
|
|
|
| class TestDomainRandomizaionEnvironment(unittest.TestCase): |
|
|
| def tearDown(self): |
| self.env.shutdown() |
|
|
| def get_task(self, randomize_every, frequency, visual_config=None, |
| dynamics_config=None): |
| cam_config = CameraConfig(render_mode=RenderMode.OPENGL) |
| obs_config = ObservationConfig(right_shoulder_camera=cam_config) |
| obs_config.set_all(False) |
| obs_config.set_all_low_dim(True) |
| obs_config.right_shoulder_camera.rgb = True |
| mode = MoveArmThenGripper(JointVelocity(), Discrete()) |
| self.env = Environment( |
| mode, ASSET_DIR, obs_config, True, True, 'panda', |
| randomize_every, frequency, visual_config, dynamics_config) |
| self.env.launch() |
| return self.env.get_task(OpenMicrowave) |
|
|
| def test_visual_randomize_every_2_episodes(self): |
| visual_config = VisualRandomizationConfig(ASSET_DIR) |
| task = self.get_task(RandomizeEvery.EPISODE, 2, visual_config) |
| _, obs1 = task.reset() |
| _, obs2 = task.reset() |
| _, obs3 = task.reset() |
| |
| self.assertLess( |
| np.mean(np.abs(obs1.right_shoulder_rgb - obs2.right_shoulder_rgb)), |
| 5) |
| |
| self.assertGreater( |
| np.mean(np.abs(obs2.right_shoulder_rgb - obs3.right_shoulder_rgb)), |
| 5) |
|
|
| def test_visual_randomize_every_2_transitions(self): |
| visual_config = VisualRandomizationConfig(ASSET_DIR) |
| task = self.get_task(RandomizeEvery.TRANSITION, 2, visual_config) |
| _, _ = task.reset() |
| a = ([0.0] * 7) + [1.0] |
| obs1, _, _ = task.step(a) |
| obs2, _, _ = task.step(a) |
| obs3, _, _ = task.step(a) |
| |
| self.assertLess( |
| np.mean(np.abs(obs1.right_shoulder_rgb - obs2.right_shoulder_rgb)), |
| 5) |
| |
| self.assertGreater( |
| np.mean(np.abs(obs2.right_shoulder_rgb - obs3.right_shoulder_rgb)), |
| 5) |
|
|
