def test_swap_arm(self):
# Checks if the environment can be setup with each arm
action_mode = ActionMode(ArmActionMode.DELTA_JOINT_VELOCITY)
for robot_config in environment.SUPPORTED_ROBOTS.keys():
with self.subTest(robot_config=robot_config):
self.env = environment.Environment(
action_mode, headless=True,
robot_configuration=robot_config)
self.env.launch()
self.env.shutdown()
def get_task(self, task_class, arm_action_mode, obs_config=None):
if obs_config is None:
obs_config = ObservationConfig()
obs_config.set_all(False)
obs_config.set_all_low_dim(True)
obs_config.right_shoulder_camera.rgb = True
action_mode = ActionMode(arm_action_mode)
self.env = environment.Environment(
action_mode, ASSET_DIR, obs_config, headless=True)
self.env.launch()
return self.env.get_task(task_class)
def test_get_stored_demos_images_without_init_sim(self):
obs_config = ObservationConfig()
obs_config.set_all(False)
obs_config.set_all_low_dim(True)
obs_config.right_shoulder_camera.rgb = True
action_mode = ActionMode()
self.env = environment.Environment(
action_mode, ASSET_DIR, obs_config, headless=True)
demos = self.env.get_demos('reach_target', 2)
self.assertEqual(len(demos), 2)
self.assertGreater(len(demos[0]), 0)
self.assertIsInstance(demos[0][0].right_shoulder_rgb, np.ndarray)
self.assertIsNone(demos[0][0].left_shoulder_rgb)