From 02cbc0057154b1ba9ebc56240ff86e06a09ab4d8 Mon Sep 17 00:00:00 2001 From: josiah_wong <84cremebrule@gmail.com> Date: Sat, 7 Nov 2020 10:59:38 -0800 Subject: [PATCH 01/46] add new papers (#118) --- docs/references.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/docs/references.md b/docs/references.md index e5dd9fe06a..9d2e60c453 100644 --- a/docs/references.md +++ b/docs/references.md @@ -13,6 +13,7 @@ A list of references of projects and papers that use **robosuite**. If you would - [Long-Horizon Visual Planning with Goal-Conditioned Hierarchical Predictors](https://arxiv.org/abs/2006.13205). Karl Pertsch, Oleh Rybkin, Frederik Ebert, Chelsea Finn, Dinesh Jayaraman, Sergey Levine - [Balance Between Efficient and Effective Learning: Dense2Sparse Reward Shaping for Robot Manipulation with Environment Uncertainty](https://arxiv.org/abs/2003.02740). Yongle Luo, Kun Dong, Lili Zhao, Zhiyong Sun, Chao Zhou, Bo Song - [Hierarchical 6-DoF Grasping with Approaching Direction Selection](http://rllab.snu.ac.kr/publications/papers/2020_icra_gads.pdf). Yunho Choi, Hogun Kee, Kyungjae Lee, JaeGoo Choy, Junhong Min, Sohee Lee, and Songhwai Oh +- [Conservative Safety Critics for Exploration](https://arxiv.org/abs/2010.14497). Homanga Bharadhwaj, Aviral Kumar, Nicholas Rhinehart, Sergey Levine, Florian Shkurti, Animesh Garg ## Imitation Learning and Batch (Offline) Reinforcement Learning @@ -21,6 +22,8 @@ A list of references of projects and papers that use **robosuite**. If you would - [To Follow or not to Follow: Selective Imitation Learning from Observations](https://arxiv.org/abs/1912.07670). Youngwoon Lee, Edward S. Hu, Zhengyu Yang, Joseph J. Lim - [Learning Robot Skills with Temporal Variational Inference](https://arxiv.org/abs/2006.16232). Tanmay Shankar, Abhinav Gupta - [Residual Learning from Demonstration](https://arxiv.org/abs/2008.07682). Todor Davchev, Kevin Sebastian Luck, Michael Burke, Franziska Meier, Stefan Schaal, Subramanian Ramamoorthy +- [Variational Imitation Learning with Diverse-quality Demonstrations](https://proceedings.icml.cc/static/paper_files/icml/2020/577-Paper.pdf). Voot Tangkaratt, Bo Han, Mohammad Emtiyaz Khan, Masashi Sugiyama + ## Benchmarks From 81011ac3f93c21c4771c2462a5c67639b503fcd2 Mon Sep 17 00:00:00 2001 From: josiah_wong <84cremebrule@gmail.com> Date: Sat, 7 Nov 2020 11:13:03 -0800 Subject: [PATCH 02/46] fix env.model.get_xml() --> env.sim.model.get_xml() so saved xml is up-to-date --- robosuite/scripts/collect_human_demonstrations.py | 2 +- robosuite/scripts/tune_camera.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/robosuite/scripts/collect_human_demonstrations.py b/robosuite/scripts/collect_human_demonstrations.py index 8e44917735..b954f4b7cf 100644 --- a/robosuite/scripts/collect_human_demonstrations.py +++ b/robosuite/scripts/collect_human_demonstrations.py @@ -75,7 +75,7 @@ def collect_human_trajectory(env, device, arm, env_configuration): # applied because the data collector wrapper only starts recording # after the first action has been played. initial_mjstate = env.sim.get_state().flatten() - xml_str = env.model.get_xml() + xml_str = env.sim.model.get_xml() env.reset_from_xml_string(xml_str) env.sim.reset() env.sim.set_state_from_flattened(initial_mjstate) diff --git a/robosuite/scripts/tune_camera.py b/robosuite/scripts/tune_camera.py index 41e6e556ce..4acd1c6caa 100644 --- a/robosuite/scripts/tune_camera.py +++ b/robosuite/scripts/tune_camera.py @@ -249,7 +249,7 @@ def print_command(char, info): ) env.reset() initial_mjstate = env.sim.get_state().flatten() - xml = env.model.get_xml() + xml = env.sim.model.get_xml() # add mocap body to camera to be able to move it around xml = modify_xml_for_camera_movement(xml, camera_name=CAMERA_NAME) From 1045ee4d134b8fbdab72a90a4f1902fc252e895f Mon Sep 17 00:00:00 2001 From: josiah_wong <84cremebrule@gmail.com> Date: Sat, 7 Nov 2020 11:22:07 -0800 Subject: [PATCH 03/46] remove redundant proprioceptive obs from object obs in handover and two arm lift tasks --- robosuite/environments/two_arm_handover.py | 4 ---- robosuite/environments/two_arm_lift.py | 4 ---- 2 files changed, 8 deletions(-) diff --git a/robosuite/environments/two_arm_handover.py b/robosuite/environments/two_arm_handover.py index 1bae2d8d78..c776aeb298 100644 --- a/robosuite/environments/two_arm_handover.py +++ b/robosuite/environments/two_arm_handover.py @@ -526,8 +526,6 @@ def _get_observation(self): di["hammer_quat"] = np.array(self._hammer_quat) di["handle_xpos"] = np.array(self._handle_xpos) - di[pr0 + "eef_xpos"] = np.array(self._eef0_xpos) - di[pr1 + "eef_xpos"] = np.array(self._eef1_xpos) di[pr0 + "gripper_to_handle"] = np.array(self._gripper_0_to_handle) di[pr1 + "gripper_to_handle"] = np.array(self._gripper_1_to_handle) @@ -536,8 +534,6 @@ def _get_observation(self): di["hammer_pos"], di["hammer_quat"], di["handle_xpos"], - di[pr0 + "eef_xpos"], - di[pr1 + "eef_xpos"], di[pr0 + "gripper_to_handle"], di[pr1 + "gripper_to_handle"], ] diff --git a/robosuite/environments/two_arm_lift.py b/robosuite/environments/two_arm_lift.py index 8e73e79d89..1ae9002cde 100644 --- a/robosuite/environments/two_arm_lift.py +++ b/robosuite/environments/two_arm_lift.py @@ -449,8 +449,6 @@ def _get_observation(self): di["cube_pos"] = cube_pos di["cube_quat"] = cube_quat - di[pr0 + "eef_xpos"] = self._eef0_xpos - di[pr1 + "eef_xpos"] = self._eef1_xpos di["handle_0_xpos"] = np.array(self._handle_0_xpos) di["handle_1_xpos"] = np.array(self._handle_1_xpos) di[pr0 + "gripper_to_handle"] = np.array(self._gripper_0_to_handle) @@ -460,8 +458,6 @@ def _get_observation(self): [ di["cube_pos"], di["cube_quat"], - di[pr0 + "eef_xpos"], - di[pr1 + "eef_xpos"], di["handle_0_xpos"], di["handle_1_xpos"], di[pr0 + "gripper_to_handle"], From 4a7e250a15c3728e8b2a406772b437ce281ba4d6 Mon Sep 17 00:00:00 2001 From: josiah_wong <84cremebrule@gmail.com> Date: Sun, 8 Nov 2020 11:52:22 -0800 Subject: [PATCH 04/46] update find_contacts to take in str or list of str --- robosuite/environments/base.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/robosuite/environments/base.py b/robosuite/environments/base.py index 71628e1ba9..1378c2f0d4 100644 --- a/robosuite/environments/base.py +++ b/robosuite/environments/base.py @@ -398,12 +398,17 @@ def find_contacts(self, geoms_1, geoms_2): Finds contact between two geom groups. Args: - geoms_1 (list of str): a list of geom names - geoms_2 (list of str): another list of geom names + geoms_1 (str or list of str): an individual geom name or list of geom names + geoms_2 (str or list of str): another individual geom name or list of geom names Returns: generator: iterator of all contacts between @geoms_1 and @geoms_2 """ + # Check if either geoms_1 or geoms_2 is a string, convert to list if so + if type(geoms_1) is str: + geoms_1 = [geoms_1] + if type(geoms_2) is str: + geoms_2 = [geoms_2] for contact in self.sim.data.contact[0 : self.sim.data.ncon]: # check contact geom in geoms c1_in_g1 = self.sim.model.geom_id2name(contact.geom1) in geoms_1 From 40cc21500dc01c34a72ec206484c9542a9140632 Mon Sep 17 00:00:00 2001 From: josiah_wong <84cremebrule@gmail.com> Date: Sun, 8 Nov 2020 12:05:36 -0800 Subject: [PATCH 05/46] fix _check_gripper_contact (#119), add more accurate checks for bimanual case --- robosuite/environments/robot_env.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/robosuite/environments/robot_env.py b/robosuite/environments/robot_env.py index 6e79fa4630..b5f7274ec5 100644 --- a/robosuite/environments/robot_env.py +++ b/robosuite/environments/robot_env.py @@ -418,16 +418,18 @@ def _get_observation(self): def _check_gripper_contact(self): """ - Checks whether each gripper is in contact with an object. + Checks whether each gripper is in contact with an object. Note that each entry in the return list corresponds + to each robot in the env.robots list. This means that a bimanual robot will have its own nested list within the + return list at the corresponding entry. Returns: - list of bool: True if the specific gripper is in contact with an object + list of bool: True if the specific gripper is in contact with an object. """ collisions = [False] * self.num_robots for idx, robot in enumerate(self.robots): for contact in self.sim.data.contact[: self.sim.data.ncon]: # Single arm case - if robot.arm_type == "single": + if robot.robot_model.arm_type == "single": if ( self.sim.model.geom_id2name(contact.geom1) in robot.gripper.contact_geoms @@ -438,15 +440,17 @@ def _check_gripper_contact(self): break # Bimanual case else: - for arm in robot.arms: + # Set this list entry to be its own list of bools + collisions[idx] = [False, False] + for i, arm in enumerate(robot.arms): if ( self.sim.model.geom_id2name(contact.geom1) in robot.gripper[arm].contact_geoms or self.sim.model.geom_id2name(contact.geom2) in robot.gripper[arm].contact_geoms ): - collisions[idx] = True - break + collisions[idx][i] = True + continue return collisions def _check_arm_contact(self): From 9c5efeb66d0b874c207b55db79b6b823d9c4d25a Mon Sep 17 00:00:00 2001 From: josiah_wong <84cremebrule@gmail.com> Date: Wed, 18 Nov 2020 19:58:59 -0800 Subject: [PATCH 06/46] fix controller limits bug --- robosuite/controllers/joint_pos.py | 2 +- robosuite/controllers/joint_tor.py | 2 +- robosuite/controllers/joint_vel.py | 2 +- robosuite/controllers/osc.py | 4 ++-- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/robosuite/controllers/joint_pos.py b/robosuite/controllers/joint_pos.py index 51d59e59b5..e0a2a52c75 100644 --- a/robosuite/controllers/joint_pos.py +++ b/robosuite/controllers/joint_pos.py @@ -120,7 +120,7 @@ def __init__(self, self.output_min = self.nums2array(output_min, self.control_dim) # limits - self.position_limits = qpos_limits + self.position_limits = np.array(qpos_limits) if qpos_limits is not None else qpos_limits # kp kd self.kp = self.nums2array(kp, self.control_dim) diff --git a/robosuite/controllers/joint_tor.py b/robosuite/controllers/joint_tor.py index bf2c90046d..cf1e9fc998 100644 --- a/robosuite/controllers/joint_tor.py +++ b/robosuite/controllers/joint_tor.py @@ -87,7 +87,7 @@ def __init__(self, self.output_min = self.nums2array(output_min, self.control_dim) # limits (if not specified, set them to actuator limits by default) - self.torque_limits = torque_limits if torque_limits is not None else self.actuator_limits + self.torque_limits = np.array(torque_limits) if torque_limits is not None else self.actuator_limits # control frequency self.control_freq = policy_freq diff --git a/robosuite/controllers/joint_vel.py b/robosuite/controllers/joint_vel.py index b8f440ec39..9598bb329d 100644 --- a/robosuite/controllers/joint_vel.py +++ b/robosuite/controllers/joint_vel.py @@ -104,7 +104,7 @@ def __init__(self, self.last_joint_vel = np.zeros(self.joint_dim) # limits - self.velocity_limits = velocity_limits + self.velocity_limits = np.array(velocity_limits) if velocity_limits is not None else None # control frequency self.control_freq = policy_freq diff --git a/robosuite/controllers/osc.py b/robosuite/controllers/osc.py index 7b5c5c69fd..b7831db4f1 100644 --- a/robosuite/controllers/osc.py +++ b/robosuite/controllers/osc.py @@ -173,8 +173,8 @@ def __init__(self, self.control_dim += 6 # limits - self.position_limits = position_limits - self.orientation_limits = orientation_limits + self.position_limits = np.array(position_limits) if position_limits is not None else position_limits + self.orientation_limits = np.array(orientation_limits) if orientation_limits is not None else orientation_limits # control frequency self.control_freq = policy_freq From 5b63732bd82130528b669baa43c3fc25e91ba048 Mon Sep 17 00:00:00 2001 From: josiah_wong <84cremebrule@gmail.com> Date: Wed, 18 Nov 2020 22:18:19 -0800 Subject: [PATCH 07/46] standardize groups for collision / visual geoms, standardize group colors for collision geoms, refactor generated / xml object parsing, fix naming conventions in two_arm_lift --- robosuite/environments/door.py | 3 +- robosuite/environments/lift.py | 3 +- robosuite/environments/nut_assembly.py | 5 +- robosuite/environments/pick_place.py | 19 +- robosuite/environments/stack.py | 6 +- robosuite/environments/two_arm_handover.py | 3 +- robosuite/environments/two_arm_lift.py | 72 +- robosuite/environments/two_arm_peg_in_hole.py | 4 +- robosuite/environments/wipe.py | 12 +- robosuite/models/arenas/arena.py | 8 +- robosuite/models/arenas/bins_arena.py | 2 - robosuite/models/arenas/wipe_arena.py | 5 +- robosuite/models/assets/arenas/bins_arena.xml | 24 +- robosuite/models/assets/arenas/pegs_arena.xml | 6 +- .../models/assets/arenas/table_arena.xml | 2 +- .../grippers/jaco_three_finger_gripper.xml | 20 +- .../models/assets/grippers/panda_gripper.xml | 8 +- .../assets/grippers/rethink_gripper.xml | 16 +- .../assets/grippers/robotiq_gripper_140.xml | 22 +- .../assets/grippers/robotiq_gripper_85.xml | 22 +- .../assets/grippers/robotiq_gripper_s.xml | 8 +- .../models/assets/grippers/wiping_gripper.xml | 70 +- robosuite/models/assets/objects/bottle.xml | 4 +- .../models/assets/objects/bread-visual.xml | 11 +- robosuite/models/assets/objects/bread.xml | 8 +- .../models/assets/objects/can-visual.xml | 9 +- robosuite/models/assets/objects/can.xml | 8 +- .../models/assets/objects/cereal-visual.xml | 15 +- robosuite/models/assets/objects/cereal.xml | 16 +- robosuite/models/assets/objects/door.xml | 2 +- robosuite/models/assets/objects/door_lock.xml | 16 +- robosuite/models/assets/objects/lemon.xml | 8 +- .../models/assets/objects/milk-visual.xml | 11 +- robosuite/models/assets/objects/milk.xml | 8 +- .../models/assets/objects/plate-with-hole.xml | 16 +- robosuite/models/assets/objects/round-nut.xml | 20 +- .../models/assets/objects/square-nut.xml | 12 +- robosuite/models/base.py | 25 + robosuite/models/grippers/gripper_model.py | 4 + robosuite/models/objects/generated_objects.py | 945 ++++++++++-------- robosuite/models/objects/objects.py | 297 ++++-- robosuite/models/objects/xml_objects.py | 50 +- robosuite/models/robots/robot_model.py | 5 +- robosuite/models/tasks/manipulation_task.py | 38 +- robosuite/utils/mjcf_utils.py | 35 +- 45 files changed, 1077 insertions(+), 826 deletions(-) diff --git a/robosuite/environments/door.py b/robosuite/environments/door.py index 44588be062..585570f57f 100644 --- a/robosuite/environments/door.py +++ b/robosuite/environments/door.py @@ -283,8 +283,7 @@ def _load_model(self): self.model = ManipulationTask( mujoco_arena=self.mujoco_arena, mujoco_robots=[robot.robot_model for robot in self.robots], - mujoco_objects=self.mujoco_objects, - visual_objects=None, + mujoco_objects=self.mujoco_objects, initializer=self.placement_initializer, ) self.model.place_objects() diff --git a/robosuite/environments/lift.py b/robosuite/environments/lift.py index bb5d1b52e0..62acec0516 100644 --- a/robosuite/environments/lift.py +++ b/robosuite/environments/lift.py @@ -318,8 +318,7 @@ def _load_model(self): self.model = ManipulationTask( mujoco_arena=self.mujoco_arena, mujoco_robots=[robot.robot_model for robot in self.robots], - mujoco_objects=self.mujoco_objects, - visual_objects=None, + mujoco_objects=self.mujoco_objects, initializer=self.placement_initializer, ) self.model.place_objects() diff --git a/robosuite/environments/nut_assembly.py b/robosuite/environments/nut_assembly.py index 5628f6c260..0c1c4aa3de 100644 --- a/robosuite/environments/nut_assembly.py +++ b/robosuite/environments/nut_assembly.py @@ -466,8 +466,7 @@ def _load_model(self): self.model = ManipulationTask( mujoco_arena=self.mujoco_arena, mujoco_robots=[robot.robot_model for robot in self.robots], - mujoco_objects=self.mujoco_objects, - visual_objects=None, + mujoco_objects=self.mujoco_objects, initializer=self.placement_initializer, ) @@ -495,7 +494,7 @@ def _get_reference(self): self.obj_body_id[obj_str] = self.sim.model.body_name2id(obj_str) geom_ids = [] for j in range(self.ngeoms[i]): - geom_ids.append(self.sim.model.geom_name2id(obj_str + "-{}".format(j))) + geom_ids.append(self.sim.model.geom_name2id(obj_str + "_{}".format(j))) self.obj_geom_id[obj_str] = geom_ids # information of objects diff --git a/robosuite/environments/pick_place.py b/robosuite/environments/pick_place.py index 35d960c672..163b497cf8 100644 --- a/robosuite/environments/pick_place.py +++ b/robosuite/environments/pick_place.py @@ -430,7 +430,7 @@ def _get_placement_initializer(self): bin_y_half = self.mujoco_arena.table_full_size[1] / 2 - 0.05 # each object should just be sampled in the bounds of the bin (with some tolerance) - for obj_name in self.mujoco_objects: + for obj_name in self.collision_objects: self.placement_initializer.sample_on_top( obj_name, @@ -523,7 +523,7 @@ def _load_model(self): visual_ob_name = ("Visual" + self.item_names[j] + "0") visual_ob = self.vis_inits[j]( name=visual_ob_name, - joints=[], # no free joint for visual objects + joints=[], # no free joint for visual objects ) lst.append((visual_ob_name, visual_ob)) self.visual_objects = OrderedDict(lst) @@ -537,7 +537,8 @@ def _load_model(self): ) lst.append((ob_name, ob)) - self.mujoco_objects = OrderedDict(lst) + self.collision_objects = OrderedDict(lst) + self.mujoco_objects = OrderedDict(**self.visual_objects, **self.collision_objects) self.n_objects = len(self.mujoco_objects) # task includes arena, robot, and objects of interest @@ -545,8 +546,7 @@ def _load_model(self): self.model = ManipulationTask( mujoco_arena=self.mujoco_arena, mujoco_robots=[robot.robot_model for robot in self.robots], - mujoco_objects=self.mujoco_objects, - visual_objects=self.visual_objects, + mujoco_objects=self.mujoco_objects, initializer=self.placement_initializer, ) @@ -579,7 +579,7 @@ def _get_reference(self): for i in range(len(self.ob_inits)): obj_str = str(self.item_names[i]) + "0" self.obj_body_id[obj_str] = self.sim.model.body_name2id(obj_str) - self.obj_geom_id[obj_str] = self.sim.model.geom_name2id(obj_str) + self.obj_geom_id[obj_str] = self.sim.model.geom_name2id(obj_str + "_0") # for checking distance to / contact with objects we want to pick up self.target_object_body_ids = list(map(int, self.obj_body_id.values())) @@ -616,10 +616,11 @@ def _reset_internal(self): # Loop through all objects and reset their positions for i, (obj_name, _) in enumerate(self.mujoco_objects.items()): - self.sim.data.set_joint_qpos(obj_name + "_jnt0", np.concatenate([np.array(obj_pos[i]), np.array(obj_quat[i])])) + if "visual" not in obj_name.lower(): + self.sim.data.set_joint_qpos(obj_name + "_jnt0", np.concatenate([np.array(obj_pos[i]), np.array(obj_quat[i])])) - # information of objects - self.object_names = list(self.mujoco_objects.keys()) + # information of (collision) objects + self.object_names = list(self.collision_objects.keys()) self.object_site_ids = [ self.sim.model.site_name2id(ob_name) for ob_name in self.object_names ] diff --git a/robosuite/environments/stack.py b/robosuite/environments/stack.py index 8950bf925c..f40e2f84ee 100644 --- a/robosuite/environments/stack.py +++ b/robosuite/environments/stack.py @@ -378,9 +378,9 @@ def _load_model(self): # task includes arena, robot, and objects of interest self.model = ManipulationTask( - self.mujoco_arena, - [robot.robot_model for robot in self.robots], - self.mujoco_objects, + mujoco_arena=self.mujoco_arena, + mujoco_robots=[robot.robot_model for robot in self.robots], + mujoco_objects=self.mujoco_objects, initializer=self.placement_initializer, ) self.model.place_objects() diff --git a/robosuite/environments/two_arm_handover.py b/robosuite/environments/two_arm_handover.py index c776aeb298..6d57f6a10c 100644 --- a/robosuite/environments/two_arm_handover.py +++ b/robosuite/environments/two_arm_handover.py @@ -349,8 +349,7 @@ def _load_model(self): self.model = ManipulationTask( mujoco_arena=self.mujoco_arena, mujoco_robots=[robot.robot_model for robot in self.robots], - mujoco_objects=self.mujoco_objects, - visual_objects=None, + mujoco_objects=self.mujoco_objects, initializer=self.placement_initializer, ) self.model.place_objects() diff --git a/robosuite/environments/two_arm_lift.py b/robosuite/environments/two_arm_lift.py index 1ae9002cde..a152c7fdc9 100644 --- a/robosuite/environments/two_arm_lift.py +++ b/robosuite/environments/two_arm_lift.py @@ -258,8 +258,8 @@ def reward(self, action=None): r_lift = min(max(elevation - 0.05, 0), 0.15) reward += 10. * direction_coef * r_lift - _gripper_0_to_handle = self._gripper_0_to_handle - _gripper_1_to_handle = self._gripper_1_to_handle + _gripper0_to_handle0 = self._gripper0_to_handle0 + _gripper1_to_handle1 = self._gripper1_to_handle1 # gh stands for gripper-handle # When grippers are far away, tell them to be closer @@ -268,48 +268,48 @@ def reward(self, action=None): if self.env_configuration == "bimanual": _contacts_0_lf = len(list( self.find_contacts( - self.robots[0].gripper["left"].important_geoms["left_finger"], self.pot.handle_2_geoms() + self.robots[0].gripper["left"].important_geoms["left_finger"], self.pot.handle1_geoms ) )) > 0 _contacts_0_rf = len(list( self.find_contacts( - self.robots[0].gripper["left"].important_geoms["right_finger"], self.pot.handle_2_geoms() + self.robots[0].gripper["left"].important_geoms["right_finger"], self.pot.handle1_geoms ) )) > 0 _contacts_1_lf = len(list( self.find_contacts( - self.robots[0].gripper["right"].important_geoms["left_finger"], self.pot.handle_1_geoms() + self.robots[0].gripper["right"].important_geoms["left_finger"], self.pot.handle0_geoms ) )) > 0 _contacts_1_rf = len(list( self.find_contacts( - self.robots[0].gripper["right"].important_geoms["right_finger"], self.pot.handle_1_geoms() + self.robots[0].gripper["right"].important_geoms["right_finger"], self.pot.handle0_geoms ) )) > 0 # Multi single arm setting else: _contacts_0_lf = len(list( self.find_contacts( - self.robots[0].gripper.important_geoms["left_finger"], self.pot.handle_2_geoms() + self.robots[0].gripper.important_geoms["left_finger"], self.pot.handle1_geoms ) )) > 0 _contacts_0_rf = len(list( self.find_contacts( - self.robots[0].gripper.important_geoms["right_finger"], self.pot.handle_2_geoms() + self.robots[0].gripper.important_geoms["right_finger"], self.pot.handle1_geoms ) )) > 0 _contacts_1_lf = len(list( self.find_contacts( - self.robots[1].gripper.important_geoms["left_finger"], self.pot.handle_1_geoms() + self.robots[1].gripper.important_geoms["left_finger"], self.pot.handle0_geoms ) )) > 0 _contacts_1_rf = len(list( self.find_contacts( - self.robots[1].gripper.important_geoms["right_finger"], self.pot.handle_1_geoms() + self.robots[1].gripper.important_geoms["right_finger"], self.pot.handle0_geoms ) )) > 0 - _g0h_dist = np.linalg.norm(_gripper_0_to_handle) - _g1h_dist = np.linalg.norm(_gripper_1_to_handle) + _g0h_dist = np.linalg.norm(_gripper0_to_handle0) + _g1h_dist = np.linalg.norm(_gripper1_to_handle1) # Grasping reward if _contacts_0_lf and _contacts_0_rf: @@ -367,15 +367,15 @@ def _load_model(self): self.mujoco_arena.set_origin([0, 0, 0]) # initialize objects of interest - self.pot = PotWithHandlesObject(name="pot") - self.mujoco_objects = OrderedDict([("pot", self.pot)]) + self.pot_name = "pot" + self.pot = PotWithHandlesObject(name=self.pot_name) + self.mujoco_objects = OrderedDict([(self.pot_name, self.pot)]) # task includes arena, robot, and objects of interest self.model = ManipulationTask( mujoco_arena=self.mujoco_arena, mujoco_robots=[robot.robot_model for robot in self.robots], - mujoco_objects=self.mujoco_objects, - visual_objects=None, + mujoco_objects=self.mujoco_objects, initializer=self.placement_initializer, ) self.model.place_objects() @@ -389,11 +389,11 @@ def _get_reference(self): super()._get_reference() # Additional object references from this env - self.pot_body_id = self.sim.model.body_name2id("pot") - self.handle_1_site_id = self.sim.model.site_name2id("pot_handle_1") - self.handle_0_site_id = self.sim.model.site_name2id("pot_handle_2") + self.pot_body_id = self.sim.model.body_name2id(self.pot_name) + self.handle1_site_id = self.sim.model.site_name2id(f"{self.pot_name}_handle0") + self.handle0_site_id = self.sim.model.site_name2id(f"{self.pot_name}_handle1") self.table_top_id = self.sim.model.site_name2id("table_top") - self.pot_center_id = self.sim.model.site_name2id("pot_center") + self.pot_center_id = self.sim.model.site_name2id(f"{self.pot_name}_center") def _reset_internal(self): """ @@ -449,19 +449,19 @@ def _get_observation(self): di["cube_pos"] = cube_pos di["cube_quat"] = cube_quat - di["handle_0_xpos"] = np.array(self._handle_0_xpos) - di["handle_1_xpos"] = np.array(self._handle_1_xpos) - di[pr0 + "gripper_to_handle"] = np.array(self._gripper_0_to_handle) - di[pr1 + "gripper_to_handle"] = np.array(self._gripper_1_to_handle) + di["handle0_xpos"] = np.array(self._handle0_xpos) + di["handle1_xpos"] = np.array(self._handle1_xpos) + di[pr0 + "gripper_to_handle0"] = np.array(self._gripper0_to_handle0) + di[pr1 + "gripper_to_handle1"] = np.array(self._gripper1_to_handle1) di["object-state"] = np.concatenate( [ di["cube_pos"], di["cube_quat"], - di["handle_0_xpos"], - di["handle_1_xpos"], - di[pr0 + "gripper_to_handle"], - di[pr1 + "gripper_to_handle"], + di["handle0_xpos"], + di["handle1_xpos"], + di[pr0 + "gripper_to_handle0"], + di[pr1 + "gripper_to_handle1"], ] ) @@ -513,24 +513,24 @@ def _check_robot_configuration(self, robots): "instead, got {}.".format(self.env_configuration, is_bimanual, check_bimanual(robot))) @property - def _handle_0_xpos(self): + def _handle0_xpos(self): """ Grab the position of the left (blue) hammer handle. Returns: np.array: (x,y,z) position of handle """ - return self.sim.data.site_xpos[self.handle_0_site_id] + return self.sim.data.site_xpos[self.handle0_site_id] @property - def _handle_1_xpos(self): + def _handle1_xpos(self): """ Grab the position of the right (green) hammer handle. Returns: np.array: (x,y,z) position of handle """ - return self.sim.data.site_xpos[self.handle_1_site_id] + return self.sim.data.site_xpos[self.handle1_site_id] @property def _pot_quat(self): @@ -579,21 +579,21 @@ def _eef1_xpos(self): return np.array(self.sim.data.site_xpos[self.robots[1].eef_site_id]) @property - def _gripper_0_to_handle(self): + def _gripper0_to_handle0(self): """ Calculate vector from the left gripper to the left pot handle. Returns: np.array: (dx,dy,dz) distance vector between handle and EEF0 """ - return self._handle_0_xpos - self._eef0_xpos + return self._handle0_xpos - self._eef0_xpos @property - def _gripper_1_to_handle(self): + def _gripper1_to_handle1(self): """ Calculate vector from the right gripper to the right pot handle. Returns: np.array: (dx,dy,dz) distance vector between handle and EEF0 """ - return self._handle_1_xpos - self._eef1_xpos + return self._handle1_xpos - self._eef1_xpos diff --git a/robosuite/environments/two_arm_peg_in_hole.py b/robosuite/environments/two_arm_peg_in_hole.py index 9cf25454f2..ebd344452c 100644 --- a/robosuite/environments/two_arm_peg_in_hole.py +++ b/robosuite/environments/two_arm_peg_in_hole.py @@ -314,13 +314,13 @@ def _load_model(self): ) # Load hole object - self.hole_obj = self.hole.get_collision(site=True) + self.hole_obj = self.hole.get_object_subtree(site=True) self.hole_obj.set("quat", "0 0 0.707 0.707") self.hole_obj.set("pos", "0.11 0 0.17") self.model.merge_asset(self.hole) # Load peg object - self.peg_obj = self.peg.get_collision(site=True) + self.peg_obj = self.peg.get_object_subtree(site=True) self.peg_obj.set("pos", array_to_string((0, 0, self.peg_length))) self.model.merge_asset(self.peg) diff --git a/robosuite/environments/wipe.py b/robosuite/environments/wipe.py index 8e4e0c7715..1b8e631b1e 100644 --- a/robosuite/environments/wipe.py +++ b/robosuite/environments/wipe.py @@ -422,7 +422,7 @@ def PointInRectangle(X, Y, Z, W, P): for new_sensor_active_id in new_sensors_active_ids: # Grab relevant sensor id info sensor_name = self.model.arena.sensor_site_names['contact_' + str(new_sensor_active_id) + '_sensor'] - new_sensor_active_geom_id = self.sim.model.geom_name2id(sensor_name) + new_sensor_active_geom_id = self.sim.model.geom_name2id(sensor_name + "_vis") # Make this sensor transparent since we wiped it (alpha = 0) self.sim.model.geom_rgba[new_sensor_active_geom_id] = [0, 0, 0, 0] # Add this sensor the wiped list @@ -526,10 +526,12 @@ def _load_model(self): self.mujoco_objects = OrderedDict() # task includes arena, robot, and objects of interest - self.model = ManipulationTask(self.mujoco_arena, - [robot.robot_model for robot in self.robots], - self.mujoco_objects, - initializer=self.placement_initializer) + self.model = ManipulationTask( + mujoco_arena=self.mujoco_arena, + mujoco_robots=[robot.robot_model for robot in self.robots], + mujoco_objects=self.mujoco_objects, + initializer=self.placement_initializer + ) self.model.place_objects() def _reset_internal(self): diff --git a/robosuite/models/arenas/arena.py b/robosuite/models/arenas/arena.py index 8bb94d367a..5a68643ddb 100644 --- a/robosuite/models/arenas/arena.py +++ b/robosuite/models/arenas/arena.py @@ -1,13 +1,17 @@ import numpy as np from robosuite.models.base import MujocoXML -from robosuite.utils.mjcf_utils import array_to_string, string_to_array -from robosuite.utils.mjcf_utils import new_geom, new_body, new_joint +from robosuite.utils.mjcf_utils import array_to_string, string_to_array, \ + new_geom, new_body, new_joint, ENVIRONMENT_COLLISION_COLOR class Arena(MujocoXML): """Base arena class.""" + def __init__(self, fname): + super().__init__(fname) + self.recolor_collision_geoms(ENVIRONMENT_COLLISION_COLOR) + def set_origin(self, offset): """ Applies a constant offset to all objects. diff --git a/robosuite/models/arenas/bins_arena.py b/robosuite/models/arenas/bins_arena.py index 3b3f1f5ad1..cfa57af188 100644 --- a/robosuite/models/arenas/bins_arena.py +++ b/robosuite/models/arenas/bins_arena.py @@ -17,8 +17,6 @@ class BinsArena(Arena): def __init__( self, bin1_pos=(0.1, -0.5, 0.8), table_full_size=(0.39, 0.49, 0.82), table_friction=(1, 0.005, 0.0001) ): - """ - """ super().__init__(xml_path_completion("arenas/bins_arena.xml")) self.table_full_size = np.array(table_full_size) diff --git a/robosuite/models/arenas/wipe_arena.py b/robosuite/models/arenas/wipe_arena.py index 7e03862b55..92b65f227b 100644 --- a/robosuite/models/arenas/wipe_arena.py +++ b/robosuite/models/arenas/wipe_arena.py @@ -94,9 +94,10 @@ def configure_location(self): density=1, material=dirt, friction=friction, + obj_type="visual", ) self.merge_asset(square2) - visual_c = square2.get_visual(site=True) + visual_c = square2.get_object_subtree(site=True) visual_c.set("pos", array_to_string([pos[0], pos[1], self.table_half_size[2]])) visual_c.find("site").set("pos", [0, 0, 0.005]) visual_c.find("site").set("rgba", array_to_string([0, 0, 0, 0])) @@ -128,7 +129,7 @@ def reset_arena(self, sim): pos = self.sample_start_pos() # Get IDs to the body, geom, and site of each sensor body_id = sim.model.body_name2id(sensor_name) - geom_id = sim.model.geom_name2id(sensor_name) + geom_id = sim.model.geom_name2id(sensor_name + "_vis") # Determine new position for this sensor position = np.array([pos[0], pos[1], self.table_half_size[2]]) # Set the current sensor (body) to this new position diff --git a/robosuite/models/assets/arenas/bins_arena.xml b/robosuite/models/assets/arenas/bins_arena.xml index d26ed7f1f7..31d0d7c8b0 100644 --- a/robosuite/models/assets/arenas/bins_arena.xml +++ b/robosuite/models/assets/arenas/bins_arena.xml @@ -26,15 +26,15 @@ - + - + - + - + - + @@ -43,19 +43,19 @@ - + - + - + - + - + - + - + diff --git a/robosuite/models/assets/arenas/pegs_arena.xml b/robosuite/models/assets/arenas/pegs_arena.xml index 5abcce09f0..dcde695c6c 100644 --- a/robosuite/models/assets/arenas/pegs_arena.xml +++ b/robosuite/models/assets/arenas/pegs_arena.xml @@ -33,7 +33,7 @@ - + @@ -43,11 +43,11 @@ - + - + diff --git a/robosuite/models/assets/arenas/table_arena.xml b/robosuite/models/assets/arenas/table_arena.xml index 1433d1f2d5..1810ded344 100644 --- a/robosuite/models/assets/arenas/table_arena.xml +++ b/robosuite/models/assets/arenas/table_arena.xml @@ -29,7 +29,7 @@ - + diff --git a/robosuite/models/assets/grippers/jaco_three_finger_gripper.xml b/robosuite/models/assets/grippers/jaco_three_finger_gripper.xml index 88c0a62c07..ad4aa3dc12 100644 --- a/robosuite/models/assets/grippers/jaco_three_finger_gripper.xml +++ b/robosuite/models/assets/grippers/jaco_three_finger_gripper.xml @@ -53,7 +53,7 @@ - + @@ -67,14 +67,14 @@ - + - - + + @@ -82,14 +82,14 @@ - + - - + + @@ -97,14 +97,14 @@ - + - - + + diff --git a/robosuite/models/assets/grippers/panda_gripper.xml b/robosuite/models/assets/grippers/panda_gripper.xml index c9cf8becc5..5a1e508f4c 100644 --- a/robosuite/models/assets/grippers/panda_gripper.xml +++ b/robosuite/models/assets/grippers/panda_gripper.xml @@ -32,20 +32,20 @@ - + - + - + - + diff --git a/robosuite/models/assets/grippers/rethink_gripper.xml b/robosuite/models/assets/grippers/rethink_gripper.xml index b6835dde59..3d1c6cd082 100644 --- a/robosuite/models/assets/grippers/rethink_gripper.xml +++ b/robosuite/models/assets/grippers/rethink_gripper.xml @@ -28,7 +28,7 @@ - + @@ -37,24 +37,26 @@ - - + + - + + - - + + - + + diff --git a/robosuite/models/assets/grippers/robotiq_gripper_140.xml b/robosuite/models/assets/grippers/robotiq_gripper_140.xml index 379b597ac6..40520351a1 100644 --- a/robosuite/models/assets/grippers/robotiq_gripper_140.xml +++ b/robosuite/models/assets/grippers/robotiq_gripper_140.xml @@ -50,7 +50,7 @@ - + @@ -63,45 +63,45 @@ - + - + - + - + - + - + - + - + - + - + diff --git a/robosuite/models/assets/grippers/robotiq_gripper_85.xml b/robosuite/models/assets/grippers/robotiq_gripper_85.xml index 665b98d293..1e1bab659a 100644 --- a/robosuite/models/assets/grippers/robotiq_gripper_85.xml +++ b/robosuite/models/assets/grippers/robotiq_gripper_85.xml @@ -45,50 +45,50 @@ - + - + - + - + - + - + - + - + - + - + - + diff --git a/robosuite/models/assets/grippers/robotiq_gripper_s.xml b/robosuite/models/assets/grippers/robotiq_gripper_s.xml index 9bca33fe95..ae124bf920 100644 --- a/robosuite/models/assets/grippers/robotiq_gripper_s.xml +++ b/robosuite/models/assets/grippers/robotiq_gripper_s.xml @@ -85,7 +85,7 @@ - + @@ -109,7 +109,7 @@ - + @@ -135,7 +135,7 @@ - + @@ -160,7 +160,7 @@ - + diff --git a/robosuite/models/assets/grippers/wiping_gripper.xml b/robosuite/models/assets/grippers/wiping_gripper.xml index 51a51c7ef3..65a098082d 100644 --- a/robosuite/models/assets/grippers/wiping_gripper.xml +++ b/robosuite/models/assets/grippers/wiping_gripper.xml @@ -6,50 +6,50 @@ - - - - - - - + + + + + + + - + - - - - - - + + + + + + - - - - - - - + + + + + + + - + - - - - - - + + + + + + - - - - - - - + + + + + + + diff --git a/robosuite/models/assets/objects/bottle.xml b/robosuite/models/assets/objects/bottle.xml index 7882543c4f..5b83fb2f0d 100644 --- a/robosuite/models/assets/objects/bottle.xml +++ b/robosuite/models/assets/objects/bottle.xml @@ -6,8 +6,8 @@ - - + + diff --git a/robosuite/models/assets/objects/bread-visual.xml b/robosuite/models/assets/objects/bread-visual.xml index b33041f12a..cfc2a21ad9 100644 --- a/robosuite/models/assets/objects/bread-visual.xml +++ b/robosuite/models/assets/objects/bread-visual.xml @@ -1,13 +1,8 @@ - - - - - - + + - - + diff --git a/robosuite/models/assets/objects/bread.xml b/robosuite/models/assets/objects/bread.xml index 46c5e71438..a5796c20d2 100644 --- a/robosuite/models/assets/objects/bread.xml +++ b/robosuite/models/assets/objects/bread.xml @@ -6,12 +6,8 @@ - - - - - - + + diff --git a/robosuite/models/assets/objects/can-visual.xml b/robosuite/models/assets/objects/can-visual.xml index 4acfc27711..49665439d2 100644 --- a/robosuite/models/assets/objects/can-visual.xml +++ b/robosuite/models/assets/objects/can-visual.xml @@ -1,11 +1,8 @@ - - - - + + - - + diff --git a/robosuite/models/assets/objects/can.xml b/robosuite/models/assets/objects/can.xml index aa192841ae..4628509951 100644 --- a/robosuite/models/assets/objects/can.xml +++ b/robosuite/models/assets/objects/can.xml @@ -6,12 +6,8 @@ - - - - - - + + diff --git a/robosuite/models/assets/objects/cereal-visual.xml b/robosuite/models/assets/objects/cereal-visual.xml index 97275ea7e6..db8e7d9977 100644 --- a/robosuite/models/assets/objects/cereal-visual.xml +++ b/robosuite/models/assets/objects/cereal-visual.xml @@ -1,18 +1,13 @@ - - - - - - + + - - + - + - + diff --git a/robosuite/models/assets/objects/cereal.xml b/robosuite/models/assets/objects/cereal.xml index 03db569923..4a2107800d 100644 --- a/robosuite/models/assets/objects/cereal.xml +++ b/robosuite/models/assets/objects/cereal.xml @@ -1,21 +1,17 @@ - + - + - - + + - - - - - + - + diff --git a/robosuite/models/assets/objects/door.xml b/robosuite/models/assets/objects/door.xml index 886b43610a..2041c63b49 100644 --- a/robosuite/models/assets/objects/door.xml +++ b/robosuite/models/assets/objects/door.xml @@ -13,7 +13,7 @@ - + diff --git a/robosuite/models/assets/objects/door_lock.xml b/robosuite/models/assets/objects/door_lock.xml index 5096f78c33..7f9a2ade7d 100644 --- a/robosuite/models/assets/objects/door_lock.xml +++ b/robosuite/models/assets/objects/door_lock.xml @@ -13,22 +13,22 @@ - + - - + + - + - - - - + + + + diff --git a/robosuite/models/assets/objects/lemon.xml b/robosuite/models/assets/objects/lemon.xml index a045d5caf9..9874d8cc6c 100644 --- a/robosuite/models/assets/objects/lemon.xml +++ b/robosuite/models/assets/objects/lemon.xml @@ -6,12 +6,8 @@ - - - - - - + + diff --git a/robosuite/models/assets/objects/milk-visual.xml b/robosuite/models/assets/objects/milk-visual.xml index 66eb125853..1670572737 100644 --- a/robosuite/models/assets/objects/milk-visual.xml +++ b/robosuite/models/assets/objects/milk-visual.xml @@ -1,13 +1,8 @@ - - - - - - + + - - + diff --git a/robosuite/models/assets/objects/milk.xml b/robosuite/models/assets/objects/milk.xml index 56373fee20..15d66d93b1 100644 --- a/robosuite/models/assets/objects/milk.xml +++ b/robosuite/models/assets/objects/milk.xml @@ -6,12 +6,8 @@ - - - - - - + + diff --git a/robosuite/models/assets/objects/plate-with-hole.xml b/robosuite/models/assets/objects/plate-with-hole.xml index 2f02e82664..29ff23acd5 100644 --- a/robosuite/models/assets/objects/plate-with-hole.xml +++ b/robosuite/models/assets/objects/plate-with-hole.xml @@ -5,17 +5,11 @@ - - - - - - - - + + + + + diff --git a/robosuite/models/assets/objects/round-nut.xml b/robosuite/models/assets/objects/round-nut.xml index dc92805347..f7984049e9 100644 --- a/robosuite/models/assets/objects/round-nut.xml +++ b/robosuite/models/assets/objects/round-nut.xml @@ -2,16 +2,16 @@ - - - - - - - - - - + + + + + + + + + + diff --git a/robosuite/models/assets/objects/square-nut.xml b/robosuite/models/assets/objects/square-nut.xml index 036c7fc5e3..721c551cf5 100644 --- a/robosuite/models/assets/objects/square-nut.xml +++ b/robosuite/models/assets/objects/square-nut.xml @@ -2,12 +2,12 @@ - - - - - - + + + + + + diff --git a/robosuite/models/base.py b/robosuite/models/base.py index ab951086b2..4adff3528e 100644 --- a/robosuite/models/base.py +++ b/robosuite/models/base.py @@ -5,6 +5,7 @@ import numpy as np from robosuite.utils import XMLError +from robosuite.utils.mjcf_utils import array_to_string class MujocoXML(object): @@ -279,3 +280,27 @@ def _add_prefix_recursively(self, root, tags, prefix): for child in root: if child.tag in tags: self._add_prefix_recursively(child, tags, prefix) + + def recolor_collision_geoms(self, rgba): + """ + Utility method to recolor all collision geoms (where collision geoms are defined as being part of group 0). + + Args: + rgba (4-array): (R, G, B, A) values to assign to all geoms with this group. + """ + for body in self.worldbody: + self._recolor_collision_geoms_recursively(body, rgba) + + def _recolor_collision_geoms_recursively(self, root, rgba): + """ + Iteratively searches through all children nodes in "root" element to find all geoms belonging to group 0 and set + the corresponding rgba value to the specified @rgba argument. + + Args: + root (ET.Element): Root of the xml element tree to start recursively searching through + rgba (4-array): (R, G, B, A) values to assign to all geoms with this group. + """ + for child in root: + if child.tag == "geom" and child.get("group") in {None, "0"}: + child.set("rgba", array_to_string(rgba)) + self._recolor_collision_geoms_recursively(child, rgba) \ No newline at end of file diff --git a/robosuite/models/grippers/gripper_model.py b/robosuite/models/grippers/gripper_model.py index d6455fb87a..10139f711e 100644 --- a/robosuite/models/grippers/gripper_model.py +++ b/robosuite/models/grippers/gripper_model.py @@ -2,6 +2,7 @@ Defines the base class of all grippers """ from robosuite.models.base import MujocoXML +from robosuite.utils.mjcf_utils import GRIPPER_COLLISION_COLOR import numpy as np @@ -26,6 +27,9 @@ def __init__(self, fname, idn=0): # Update all xml element prefixes self.add_prefix(self.naming_prefix) + # Update collision geom colors + self.recolor_collision_geoms(GRIPPER_COLLISION_COLOR) + # Set public attributes with prefixes appended to values self.joints = [self.naming_prefix + joint for joint in self._joints] self.actuators = [self.naming_prefix + actuator for actuator in self._actuators] diff --git a/robosuite/models/objects/generated_objects.py b/robosuite/models/objects/generated_objects.py index 735ca4e90b..3ee8546a81 100644 --- a/robosuite/models/objects/generated_objects.py +++ b/robosuite/models/objects/generated_objects.py @@ -1,16 +1,263 @@ import numpy as np from robosuite.models.objects import MujocoGeneratedObject -from robosuite.utils.mjcf_utils import new_body, new_geom, new_site, array_to_string -from robosuite.utils.mjcf_utils import RED, GREEN, BLUE, CustomMaterial +from robosuite.utils.mjcf_utils import new_body, new_geom, new_site, array_to_string, add_to_dict +from robosuite.utils.mjcf_utils import RED, GREEN, BLUE, CYAN, OBJECT_COLLISION_COLOR, CustomMaterial +import robosuite.utils.transform_utils as T from collections.abc import Iterable -# Define custom colors -CYAN = [0, 1, 1, 1] +from copy import deepcopy +import xml.etree.ElementTree as ET -class HammerObject(MujocoGeneratedObject): +class CompositeObject(MujocoGeneratedObject): + """ + An object constructed out of basic geoms to make more intricate shapes. + + Note that by default, specifying None for a specific geom element will usually set a value to the mujoco defaults. + + Args: + total_size (list): (x, y, z) half-size in each dimension for the bounding box for + this Composite object + + object_name (str): Name of overall object + + geom_types (list): list of geom types in the composite. Must correspond + to MuJoCo geom primitives, such as "box" or "capsule". + + geom_locations (list): list of geom locations in the composite. Each + location should be a list or tuple of 3 elements and all + locations are relative to the lower left corner of the total box + (e.g. (0, 0, 0) corresponds to this corner). + + geom_sizes (list): list of geom sizes ordered the same as @geom_locations + + geom_names (list): list of geom names ordered the same as @geom_locations. The + names will get appended with an underscore to the passed name in @get_collision + and @get_visual + + geom_rgbas (list): list of geom colors ordered the same as @geom_locations. If + passed as an argument, @rgba is ignored. + + geom_materials (list of CustomTexture): list of custom textures to use for this object material + + geom_frictions (list): list of geom frictions to use for each geom. + + geom_quats (list): list of (w, x, y, z) quaternions for each geom. + + joints (list): list of joints to use for each geom. Note that each entry should be its own list of dictionaries, + where each dictionary specifies the specific joint attributes necessary. + See http://www.mujoco.org/book/XMLreference.html#joint for reference. + + rgba (list): (r, g, b, a) default values to use if geom-specific @geom_rgbas isn't specified for a given element + + density (float or list of float): either single value to use for all geom densities or geom-specific values + + solref (list or list of list): parameters used for the mujoco contact solver. Can be single set of values or + element-specific values. See http://www.mujoco.org/book/modeling.html#CSolver for details. + + solimp (list or list of list): parameters used for the mujoco contact solver. Can be single set of values or + element-specific values. See http://www.mujoco.org/book/modeling.html#CSolver for details. + + locations_relative_to_center (bool): If true, @geom_locations will be considered relative to the center of the + overall object bounding box defined by @total_size. Else, the corner of this bounding box is considered the + origin. + + sites (list of dict): list of sites to add to this object. Each dict should specify the appropriate attributes + for the given site. See http://www.mujoco.org/book/XMLreference.html#site for reference. + + obj_types (str or list of str): either single obj_type for all geoms or geom-specific type. Choices are + {"collision", "visual", "all"} + """ + + def __init__( + self, + total_size, + object_name, + geom_types, + geom_locations, + geom_sizes, + geom_names=None, + geom_rgbas=None, + geom_materials=None, + geom_frictions=None, + geom_quats=None, + joints="default", + rgba=None, + density=100., + solref=(0.02, 1.), + solimp=(0.9, 0.95, 0.001), + locations_relative_to_center=False, + sites=None, + obj_types="all", + duplicate_collision_geoms=True, + ): + super().__init__( + name=object_name, + joints=joints, + rgba=rgba, + duplicate_collision_geoms=duplicate_collision_geoms, + ) + + n_geoms = len(geom_types) + self.total_size = np.array(total_size) + self.geom_types = np.array(geom_types) + self.geom_locations = np.array(geom_locations) + self.geom_sizes = deepcopy(geom_sizes) + self.geom_names = list(geom_names) if geom_names is not None else None + self.geom_rgbas = list(geom_rgbas) if geom_rgbas is not None else None + self.geom_materials = list(geom_materials) if geom_materials is not None else None + self.geom_frictions = list(geom_frictions) if geom_frictions is not None else None + self.density = [density] * n_geoms if density is None or type(density) in {float, int} else list(density) + self.solref = [solref] * n_geoms if solref is None or type(solref[0]) in {float, int} else list(solref) + self.solimp = [solimp] * n_geoms if obj_types is None or type(solimp[0]) in {float, int} else list(solimp) + self.rgba = rgba # override superclass setting of this variable + self.locations_relative_to_center = locations_relative_to_center + self.geom_quats = deepcopy(geom_quats) if geom_quats is not None else None + self.sites = sites + self.obj_types = [obj_types] * n_geoms if obj_types is None or type(obj_types) is str else list(obj_types) + + def get_bottom_offset(self): + return np.array([0., 0., -self.total_size[2]]) + + def get_top_offset(self): + return np.array([0., 0., self.total_size[2]]) + + def get_horizontal_radius(self): + return np.linalg.norm(self.total_size[:2], 2) + + def _size_to_cartesian_half_lengths(self, geom_type, geom_size): + """ + converts from geom size specification to x, y, and z half-length bounding box + """ + if geom_type in ['box', 'ellipsoid']: + return geom_size + if geom_type == 'sphere': + # size is radius + return [geom_size[0], geom_size[0], geom_size[0]] + if geom_type == 'capsule': + # size is radius, half-length of cylinder part + return [geom_size[0], geom_size[0], geom_size[0] + geom_size[1]] + if geom_type == 'cylinder': + # size is radius, half-length + return [geom_size[0], geom_size[0], geom_size[1]] + raise Exception("unsupported geom type!") + + def get_object_subtree(self, site=None): + obj = new_body() + obj.set("name", self.name) + + for i in range(self.geom_locations.shape[0]): + # geom type + geom_type = self.geom_types[i] + # get cartesian size from size spec + size = self.geom_sizes[i] + cartesian_size = self._size_to_cartesian_half_lengths(geom_type, size) + if self.locations_relative_to_center: + # no need to convert + pos = self.geom_locations[i] + else: + # use geom location to convert to position coordinate (the origin is the + # center of the composite object) + loc = self.geom_locations[i] + pos = [ + (-self.total_size[0] + cartesian_size[0]) + loc[0], + (-self.total_size[1] + cartesian_size[1]) + loc[1], + (-self.total_size[2] + cartesian_size[2]) + loc[2], + ] + + # geom name + if self.geom_names is not None: + geom_name = "{}_{}".format(self.name, self.geom_names[i]) + else: + geom_name = "{}_{}".format(self.name, i) + + # geom rgba + if self.geom_rgbas is not None and self.geom_rgbas[i] is not None: + geom_rgba = self.geom_rgbas[i] + else: + geom_rgba = self.rgba + + # geom friction + if self.geom_frictions is not None and self.geom_frictions[i] is not None: + geom_friction = array_to_string(self.geom_frictions[i]) + else: + geom_friction = array_to_string(np.array([1., 0.005, 0.0001])) # mujoco default + + # Define base geom attributes + geom_attr = { + "size": size, + "pos": pos, + "name": geom_name, + "geom_type": geom_type, + } + + # Optionally define quat if specified + if self.geom_quats is not None: + geom_attr['quat'] = array_to_string(self.geom_quats[i]) + + # Add collision geom if necessary + if self.obj_types[i] in {"collision", "all"}: + col_geom_attr = deepcopy(geom_attr) + col_geom_attr.update(self.get_collision_attrib_template()) + if self.density[i] is not None: + col_geom_attr['density'] = str(self.density[i]) + col_geom_attr['friction'] = geom_friction + col_geom_attr['solref'] = array_to_string(self.solref[i]) + col_geom_attr['solimp'] = array_to_string(self.solimp[i]) + col_geom_attr['rgba'] = OBJECT_COLLISION_COLOR + obj.append(new_geom(**col_geom_attr)) + + # Add visual geom if necessary + if self.obj_types[i] in {"visual", "all"}: + vis_geom_attr = deepcopy(geom_attr) + vis_geom_attr.update(self.get_visual_attrib_template()) + vis_geom_attr["name"] += "_vis" + if self.geom_materials is not None: + vis_geom_attr['material'] = self.geom_materials[i] + vis_geom_attr["rgba"] = geom_rgba + obj.append(new_geom(**vis_geom_attr)) + + # add top level site if requested + if site: + # add a site as well + site_element_attr = self.get_site_attrib_template() + site_element_attr["rgba"] = "1 0 0 0" + site_element_attr["name"] = self.name + obj.append(ET.Element("site", attrib=site_element_attr)) + + # Also create specific sites requested + if self.sites is not None: + # add relevant sites + for s in self.sites: + obj.append(ET.Element("site", attrib=s)) + + return obj + + def get_bounding_box_size(self): + return np.array(self.total_size) + + def in_box(self, position, object_position): + """ + Checks whether the object is contained within this CompositeObject. + Useful for when the CompositeObject has holes and the object should + be within one of the holes. Makes an approximation by treating the + object as a point, and the CompositeBoxObject as an axis-aligned grid. + Args: + position: 3D body position of CompositeObject + object_position: 3D position of object to test for insertion + """ + ub = position + self.total_size + lb = position - self.total_size + + # fudge factor for the z-check, since after insertion the object falls to table + lb[2] -= 0.01 + + return np.all(object_position > lb) and np.all(object_position < ub) + + +class HammerObject(CompositeObject): """ Generates a Hammer object with a cylindrical or box-shaped handle, cubic head, cylindrical face and triangular claw (used in Handover task) @@ -44,10 +291,6 @@ class HammerObject(MujocoGeneratedObject): rgba_claw (4-array or None): If specified, sets handle rgba values - joints (list of dict): array of dictionaries, where each dictionary corresponds to a joint that will be created - for this object. The dictionary should specify the joint attributes (type, pos, etc.) according to - the MuJoCo xml specification. - Raises: ValueError: [Invalid handle shape] """ @@ -66,10 +309,9 @@ def __init__( rgba_head=None, rgba_face=None, rgba_claw=None, - joints=None, ): - # Run super() init - super().__init__(name=name, joints=joints) + # Set name + self.name = name # Set handle type and density ratio self.handle_shape = handle_shape @@ -95,6 +337,9 @@ def __init__( self.rgba_face = rgba_face if rgba_face is not None else BLUE self.rgba_claw = rgba_claw if rgba_claw is not None else GREEN + # Create dictionary of values to create geoms for composite object and run super init + super().__init__(**self._get_geom_attrs()) + # Define materials we want to use for this object tex_attrib = { "type": "cube", @@ -132,117 +377,101 @@ def get_top_offset(self): def get_horizontal_radius(self): return self.head_halfsize + 0.5 * self.handle_length - @property - def handle_distance(self): + def _get_geom_attrs(self): """ - Calculates how wide the handle is + Creates geom elements that will be passed to superclass CompositeObject constructor Returns: - float: handle diameter - """ - return 2.0 * self.handle_radius - - def get_collision(self, site=None): - # Create new body - main_body = new_body() - main_body.set("name", self.name) - - # Define handle and append to the main body - if self.handle_shape == "cylinder": - main_body.append( - new_geom( - geom_type="cylinder", - name="hammer_handle", - size=[self.handle_radius, self.handle_length / 2.0], - pos=(0, 0, 0), - rgba=None if self.use_texture else self.rgba_handle, - group=1, - density=str(self.handle_density), - friction=array_to_string((self.handle_friction, 0.005, 0.0001)), - material="wood_mat" if self.use_texture else None, - ) - ) - elif self.handle_shape == "box": - main_body.append( - new_geom( - geom_type="box", - name="hammer_handle", - size=[self.handle_radius, self.handle_radius, self.handle_length / 2.0], - pos=(0, 0, 0), - rgba=None if self.use_texture else self.rgba_handle, - group=1, - density=str(self.handle_density), - friction=array_to_string((self.handle_friction, 0.005, 0.0001)), - material="wood_mat" if self.use_texture else None, - ) - ) - else: - # Raise error - raise ValueError("Error loading hammer: Handle type must either be 'box' or 'cylinder', got {}.".format( - self.handle_shape - )) - - # Define head and append to the main body - main_body.append( - new_geom( - geom_type="box", - name="hammer_head", - size=[self.head_halfsize * 2, self.head_halfsize, self.head_halfsize], - pos=(0, 0, self.handle_length / 2.0 + self.head_halfsize), - rgba=None if self.use_texture else self.rgba_head, - group=1, - density=str(self.handle_density * self.head_density_ratio), - material="metal_mat" if self.use_texture else None, - ) + dict: args to be used by CompositeObject to generate geoms + """ + full_size = np.array(( + 3.2 * self.head_halfsize, + self.head_halfsize, + self.handle_length + 2 * self.head_halfsize + )) + # Initialize dict of obj args that we'll pass to the CompositeObject constructor + obj_args = { + "total_size": full_size / 2.0, + "object_name": self.name, + "locations_relative_to_center": True, + "obj_types": "all", + } + + # Add handle component + assert self.handle_shape in {"cylinder", "box"},\ + "Error loading hammer: Handle type must either be 'box' or 'cylinder', got {}.".format(self.handle_shape) + add_to_dict( + dic=obj_args, + geom_types="cylinder" if self.handle_shape == "cylinder" else "box", + geom_locations=(0, 0, 0), + geom_quats=(1, 0, 0, 0), + geom_sizes=np.array([self.handle_radius, self.handle_length / 2.0]) if self.handle_shape == "cylinder" else\ + np.array([self.handle_radius, self.handle_radius, self.handle_length / 2.0]), + geom_names="handle", + geom_rgbas=None if self.use_texture else self.rgba_handle, + geom_materials="wood_mat" if self.use_texture else None, + geom_frictions=(self.handle_friction, 0.005, 0.0001), + density=self.handle_density, ) - # Define face (and neck) and append to the main body - main_body.append( - new_geom( - geom_type="cylinder", - name="hammer_neck", - size=[self.head_halfsize * 0.8, self.head_halfsize * 0.2], - pos=(self.head_halfsize * 2.2, 0, self.handle_length / 2.0 + self.head_halfsize), - quat=array_to_string([0.707106, 0, 0.707106, 0]), - rgba=None if self.use_texture else self.rgba_face, - group=1, - density=str(self.handle_density * self.head_density_ratio), - material="metal_mat" if self.use_texture else None, - ) + # Add head component + add_to_dict( + dic=obj_args, + geom_types="box", + geom_locations=(0, 0, self.handle_length / 2.0 + self.head_halfsize), + geom_quats=(1, 0, 0, 0), + geom_sizes=np.array([self.head_halfsize * 2, self.head_halfsize, self.head_halfsize]), + geom_names="head", + geom_rgbas=None if self.use_texture else self.rgba_head, + geom_materials="metal_mat" if self.use_texture else None, + geom_frictions=None, + density=self.handle_density * self.head_density_ratio, ) - main_body.append( - new_geom( - geom_type="cylinder", - name="hammer_face", - size=[self.head_halfsize, self.head_halfsize * 0.4], - pos=(self.head_halfsize * 2.8, 0, self.handle_length / 2.0 + self.head_halfsize), - quat=array_to_string([0.707106, 0, 0.707106, 0]), - rgba=None if self.use_texture else self.rgba_face, - group=1, - density=str(self.handle_density * self.head_density_ratio), - material="metal_mat" if self.use_texture else None, - ) + + # Add neck component + add_to_dict( + dic=obj_args, + geom_types="cylinder", + geom_locations=(self.head_halfsize * 2.2, 0, self.handle_length / 2.0 + self.head_halfsize), + geom_quats=(0.707106, 0, 0.707106, 0), + geom_sizes=np.array([self.head_halfsize * 0.8, self.head_halfsize * 0.2]), + geom_names="neck", + geom_rgbas=None if self.use_texture else self.rgba_face, + geom_materials="metal_mat" if self.use_texture else None, + geom_frictions=None, + density=self.handle_density * self.head_density_ratio, ) - # Define claw and append to the main body - main_body.append( - new_geom( - geom_type="box", - name="hammer_claw", - size=[self.head_halfsize * 0.7072, self.head_halfsize * 0.95, self.head_halfsize * 0.7072], - pos=(-self.head_halfsize * 2, 0, self.handle_length / 2.0 + self.head_halfsize), - quat=array_to_string([0.9238795, 0, 0.3826834, 0]), - rgba=None if self.use_texture else self.rgba_claw, - group=1, - density=str(self.handle_density * self.head_density_ratio), - material="metal_mat" if self.use_texture else None, - ) + # Add face component + add_to_dict( + dic=obj_args, + geom_types="cylinder", + geom_locations=(self.head_halfsize * 2.8, 0, self.handle_length / 2.0 + self.head_halfsize), + geom_quats=(0.707106, 0, 0.707106, 0), + geom_sizes=np.array([self.head_halfsize, self.head_halfsize * 0.4]), + geom_names="face", + geom_rgbas=None if self.use_texture else self.rgba_face, + geom_materials="metal_mat" if self.use_texture else None, + geom_frictions=None, + density=self.handle_density * self.head_density_ratio, ) - return main_body + # Add claw component + add_to_dict( + dic=obj_args, + geom_types="box", + geom_locations=(-self.head_halfsize * 2, 0, self.handle_length / 2.0 + self.head_halfsize), + geom_quats=(0.9238795, 0, 0.3826834, 0), + geom_sizes=np.array([self.head_halfsize * 0.7072, self.head_halfsize * 0.95, self.head_halfsize * 0.7072]), + geom_names="claw", + geom_rgbas=None if self.use_texture else self.rgba_claw, + geom_materials="metal_mat" if self.use_texture else None, + geom_frictions=None, + density=self.handle_density * self.head_density_ratio, + ) - def get_visual(self, site=None): - return self.get_collision(site) + # Return this dict + return obj_args @property def init_quat(self): @@ -296,14 +525,14 @@ def all_geoms(self): return self.handle_geoms + self.head_geoms + self.face_geoms + self.claw_geoms -class PotWithHandlesObject(MujocoGeneratedObject): +class PotWithHandlesObject(CompositeObject): """ Generates the Pot object with side handles (used in TwoArmLift) Args: name (str): Name of this Pot object - body_half_size (None or 3-array of float): If specified, defines the (x,y,z) half-dimensions of the main pot + body_half_size (3-array of float): If specified, defines the (x,y,z) half-dimensions of the main pot body. Otherwise, defaults to [0.07, 0.07, 0.07] handle_radius (float): Determines the pot handle radius @@ -312,53 +541,64 @@ class PotWithHandlesObject(MujocoGeneratedObject): handle_width (float): Determines the pot handle width + handle_friction (float): Friction value to use for pot handles. Defauls to 1.0 + + density (float): Density value to use for all geoms. Defaults to 1000 + use_texture (bool): If true, geoms will be defined by realistic textures and rgba values will be ignored rgba_body (4-array or None): If specified, sets pot body rgba values - rgba_handle_1 (4-array or None): If specified, sets handle 1 rgba values + rgba_handle_0 (4-array or None): If specified, sets handle 0 rgba values - rgba_handle_2 (4-array or None): If specified, sets handle 2 rgba values + rgba_handle_1 (4-array or None): If specified, sets handle 1 rgba values solid_handle (bool): If true, uses a single geom to represent the handle thickness (float): How thick to make the pot body walls - - joints (list of dict): array of dictionaries, where each dictionary corresponds to a joint that will be created - for this object. The dictionary should specify the joint attributes (type, pos, etc.) according to - the MuJoCo xml specification. """ def __init__( self, name, - body_half_size=None, + body_half_size=(0.07, 0.07, 0.07), handle_radius=0.01, handle_length=0.09, handle_width=0.09, + handle_friction=1.0, + density=1000, use_texture=True, rgba_body=None, + rgba_handle_0=None, rgba_handle_1=None, - rgba_handle_2=None, solid_handle=False, thickness=0.01, # For body - joints=None, ): - super().__init__(name=name, joints=joints) - if body_half_size: - self.body_half_size = body_half_size - else: - self.body_half_size = np.array([0.07, 0.07, 0.07]) + # Set name + self.name = name + + # Set object attributes + self.body_half_size = np.array(body_half_size) self.thickness = thickness self.handle_radius = handle_radius self.handle_length = handle_length self.handle_width = handle_width + self.handle_friction = handle_friction + self.density = density self.use_texture = use_texture self.rgba_body = np.array(rgba_body) if rgba_body else RED - self.rgba_handle_1 = np.array(rgba_handle_1) if rgba_handle_1 else GREEN - self.rgba_handle_2 = np.array(rgba_handle_2) if rgba_handle_2 else BLUE + self.rgba_handle_0 = np.array(rgba_handle_0) if rgba_handle_0 else GREEN + self.rgba_handle_1 = np.array(rgba_handle_1) if rgba_handle_1 else BLUE self.solid_handle = solid_handle + # Geoms to be filled when generated + self.handle0_geoms = None + self.handle1_geoms = None + self.pot_base = None + + # Create dictionary of values to create geoms for composite object and run super init + super().__init__(**self._get_geom_attrs()) + # Define materials we want to use for this object tex_attrib = { "type": "cube", @@ -378,14 +618,14 @@ def __init__( greenwood = CustomMaterial( texture="WoodGreen", tex_name="greenwood", - mat_name="handle1_mat", + mat_name="handle0_mat", tex_attrib=tex_attrib, mat_attrib=mat_attrib, ) bluewood = CustomMaterial( texture="WoodBlue", tex_name="bluewood", - mat_name="handle2_mat", + mat_name="handle1_mat", tex_attrib=tex_attrib, mat_attrib=mat_attrib, ) @@ -413,252 +653,163 @@ def handle_distance(self): """ return self.body_half_size[1] * 2 + self.handle_length * 2 - def get_collision(self, site=None): - main_body = new_body() - main_body.set("name", self.name) - - for geom in five_sided_box( - size=self.body_half_size, - rgba=None if self.use_texture else self.rgba_body, - group=1, - thickness=self.thickness, - material="pot_mat" if self.use_texture else None, - ): - main_body.append(geom) - handle_z = self.body_half_size[2] - self.handle_radius - handle_1_center = [0, self.body_half_size[1] + self.handle_length, handle_z] - handle_2_center = [ - 0, - -1 * (self.body_half_size[1] + self.handle_length), - handle_z, - ] - # the bar on handle horizontal to body - main_bar_size = [ - self.handle_width / 2 + self.handle_radius, - self.handle_radius, - self.handle_radius, - ] - side_bar_size = [self.handle_radius, self.handle_length / 2, self.handle_radius] - handle_1 = new_body(name="handle_1") - if self.solid_handle: - handle_1.append( - new_geom( - geom_type="box", - name="handle_1", - pos=[0, self.body_half_size[1] + self.handle_length / 2, handle_z], - size=[ - self.handle_width / 2, - self.handle_length / 2, - self.handle_radius, - ], - rgba=None if self.use_texture else self.rgba_handle_1, - group=1, - material="handle1_mat" if self.use_texture else None, - ) - ) - else: - handle_1.append( - new_geom( - geom_type="box", - name="handle_1_c", - pos=handle_1_center, - size=main_bar_size, - rgba=None if self.use_texture else self.rgba_handle_1, - group=1, - material="handle1_mat" if self.use_texture else None, - ) - ) - handle_1.append( - new_geom( - geom_type="box", - name="handle_1_+", # + for positive x - pos=[ - self.handle_width / 2, - self.body_half_size[1] + self.handle_length / 2, - handle_z, - ], - size=side_bar_size, - rgba=None if self.use_texture else self.rgba_handle_1, - group=1, - material="handle1_mat" if self.use_texture else None, - ) - ) - handle_1.append( - new_geom( - geom_type="box", - name="handle_1_-", - pos=[ - -self.handle_width / 2, - self.body_half_size[1] + self.handle_length / 2, - handle_z, - ], - size=side_bar_size, - rgba=None if self.use_texture else self.rgba_handle_1, - group=1, - material="handle1_mat" if self.use_texture else None, - ) - ) + def _get_geom_attrs(self): + """ + Creates geom elements that will be passed to superclass CompositeObject constructor - handle_2 = new_body(name="handle_2") - if self.solid_handle: - handle_2.append( - new_geom( - geom_type="box", - name="handle_2", - pos=[0, -self.body_half_size[1] - self.handle_length / 2, handle_z], - size=[ - self.handle_width / 2, - self.handle_length / 2, - self.handle_radius, - ], - rgba=None if self.use_texture else self.rgba_handle_2, - group=1, - material="handle2_mat" if self.use_texture else None, - ) - ) - else: - handle_2.append( - new_geom( - geom_type="box", - name="handle_2_c", - pos=handle_2_center, - size=main_bar_size, - rgba=None if self.use_texture else self.rgba_handle_2, - group=1, - material="handle2_mat" if self.use_texture else None, - ) - ) - handle_2.append( - new_geom( - geom_type="box", - name="handle_2_+", # + for positive x - pos=[ - self.handle_width / 2, - -self.body_half_size[1] - self.handle_length / 2, - handle_z, - ], - size=side_bar_size, - rgba=None if self.use_texture else self.rgba_handle_2, - group=1, - material="handle2_mat" if self.use_texture else None, - ) - ) - handle_2.append( - new_geom( - geom_type="box", - name="handle_2_-", - pos=[ - -self.handle_width / 2, - -self.body_half_size[1] - self.handle_length / 2, - handle_z, - ], - size=side_bar_size, - rgba=None if self.use_texture else self.rgba_handle_2, - group=1, - material="handle2_mat" if self.use_texture else None, - ) - ) + Returns: + dict: args to be used by CompositeObject to generate geoms + """ + full_size = np.array(( + self.body_half_size, + self.body_half_size + self.handle_length * 2, + self.body_half_size, + )) + # Initialize dict of obj args that we'll pass to the CompositeObject constructor + obj_args = { + "total_size": full_size / 2.0, + "object_name": self.name, + "locations_relative_to_center": True, + "obj_types": "all", + "sites": [], + } - main_body.append(handle_1) - main_body.append(handle_2) - main_body.append( - new_site( - name="pot_handle_1", - rgba=self.rgba_handle_1, - pos=handle_1_center - np.array([0, 0.005, 0]), - size=[0.005], - ) - ) - main_body.append( - new_site( - name="pot_handle_2", - rgba=self.rgba_handle_2, - pos=handle_2_center + np.array([0, 0.005, 0]), - size=[0.005], - ) + # Initialize geom lists + self.handle0_geoms = [] + self.handle1_geoms = [] + + # Add main pot body + # Base geom + name = f"base" + self.pot_base = [f"{self.name}_{name}"] + add_to_dict( + dic=obj_args, + geom_types="box", + geom_locations=(0, 0, -self.body_half_size[2] + self.thickness / 2), + geom_quats=(1, 0, 0, 0), + geom_sizes=np.array([self.body_half_size[0], self.body_half_size[1], self.thickness / 2]), + geom_names=name, + geom_rgbas=None if self.use_texture else self.rgba_body, + geom_materials="pot_mat" if self.use_texture else None, + geom_frictions=None, + density=self.density, ) - main_body.append( - new_site( - name="pot_center", - pos=[0, 0, 0], - rgba=[1, 0, 0, 0], + + # Walls + x_off = np.array([0, -(self.body_half_size[0] - self.thickness / 2), + 0, self.body_half_size[0] - self.thickness / 2]) + y_off = np.array([-(self.body_half_size[1] - self.thickness / 2), + 0, self.body_half_size[1] - self.thickness / 2, 0]) + w_vals = np.array([self.body_half_size[1], self.body_half_size[0], + self.body_half_size[1], self.body_half_size[0]]) + r_vals = np.array([np.pi / 2, 0, -np.pi / 2, np.pi]) + for i, (x, y, w, r) in enumerate(zip(x_off, y_off, w_vals, r_vals)): + add_to_dict( + dic=obj_args, + geom_types="box", + geom_locations=(x, y, 0), + geom_quats=T.convert_quat(T.axisangle2quat(np.array([0, 0, r])), to="wxyz"), + geom_sizes=np.array([self.thickness / 2, w, self.body_half_size[2]]), + geom_names=f"body{i}", + geom_rgbas=None if self.use_texture else self.rgba_body, + geom_materials="pot_mat" if self.use_texture else None, + geom_frictions=None, + density=self.density, ) - ) - return main_body + # Add handles + main_bar_size = np.array([ + self.handle_width / 2 + self.handle_radius, + self.handle_radius, + self.handle_radius, + ]) + side_bar_size = np.array([self.handle_radius, self.handle_length / 2, self.handle_radius]) + handle_z = self.body_half_size[2] - self.handle_radius + for i, (handle_side, rgba) in enumerate(zip([1.0, -1.0], [self.rgba_handle_0, self.rgba_handle_1])): + handle_center = np.array((0, handle_side * (self.body_half_size[1] + self.handle_length), handle_z)) + # Get reference to relevant geom list + g_list = getattr(self, f"handle{i}_geoms") + # Solid handle case + name = f"handle{i}" + g_list.append(f"{self.name}_{name}") + if self.solid_handle: + add_to_dict( + dic=obj_args, + geom_types="box", + geom_locations=handle_center, + geom_quats=(1, 0, 0, 0), + geom_sizes=np.array([self.handle_width / 2, self.handle_length / 2, self.handle_radius]), + geom_names=name, + geom_rgbas=None if self.use_texture else rgba, + geom_materials=f"handle{i}_mat" if self.use_texture else None, + geom_frictions=(self.handle_friction, 0.005, 0.0001), + density=self.density, + ) + # Hollow handle case + else: + # Center bar + name = f"handle{i}_c" + g_list.append(f"{self.name}_{name}") + add_to_dict( + dic=obj_args, + geom_types="box", + geom_locations=handle_center, + geom_quats=(1, 0, 0, 0), + geom_sizes=main_bar_size, + geom_names=name, + geom_rgbas=None if self.use_texture else rgba, + geom_materials=f"handle{i}_mat" if self.use_texture else None, + geom_frictions=(self.handle_friction, 0.005, 0.0001), + density=self.density, + ) + # Side bars + for bar_side, suffix in zip([-1., 1.], ["-", "+"]): + name = f"handle{i}_{suffix}" + g_list.append(f"{self.name}_{name}") + add_to_dict( + dic=obj_args, + geom_types="box", + geom_locations=( + bar_side * self.handle_width / 2, + handle_side * (self.body_half_size[1] + self.handle_length / 2), + handle_z + ), + geom_quats=(1, 0, 0, 0), + geom_sizes=side_bar_size, + geom_names=name, + geom_rgbas=None if self.use_texture else rgba, + geom_materials=f"handle{i}_mat" if self.use_texture else None, + geom_frictions=(self.handle_friction, 0.005, 0.0001), + density=self.density, + ) + # Add relevant site + handle_site = self.get_site_attrib_template() + handle_site.update({ + "name": f"{self.name}_handle{i}", + "pos": array_to_string(handle_center - handle_side * np.array([0, 0.005, 0])), + "size": "0.005", + "rgba": rgba, + }) + obj_args["sites"].append(handle_site) + + # Add pot body site + pot_site = self.get_site_attrib_template() + pot_site.update({ + "name": f"{self.name}_center", + "size": "0.005", + }) + obj_args["sites"].append(pot_site) + + # Return this dict + return obj_args + @property def handle_geoms(self): """ Returns: list of str: geom names corresponding to both handles """ - return self.handle_1_geoms() + self.handle_2_geoms() - - def handle_1_geoms(self): - """ - Returns: - list of str: geom names corresponding to handle 1 - """ - if self.solid_handle: - return ["handle_1"] - return ["handle_1_c", "handle_1_+", "handle_1_-"] - - def handle_2_geoms(self): - """ - Returns: - list of str: geom names corresponding to handle 2 - """ - if self.solid_handle: - return ["handle_2"] - return ["handle_2_c", "handle_2_+", "handle_2_-"] - - def get_visual(self, site=None): - return self.get_collision(site) - - -def five_sided_box(size, rgba, group, thickness, material): - """ - Procedurally generates a five-sided (open) box - - Args: - size (3-array of float): the (x,y,z) half-size desired for the box - rgba (4-array of float): rgba color for this box - group (int): Mujoco group to assign these geoms - thickness (float): wall thickness - material (str): material for this box - - Returns: - list: array of geoms corresponding to the 5 sides of the generated box - """ - geoms = [] - x, y, z = size - r = thickness / 2 - geoms.append( - new_geom( - geom_type="box", size=[x, y, r], pos=[0, 0, -z + r], rgba=rgba, group=group, material=material - ) - ) - geoms.append( - new_geom( - geom_type="box", size=[x, r, z], pos=[0, -y + r, 0], rgba=rgba, group=group, material=material - ) - ) - geoms.append( - new_geom( - geom_type="box", size=[x, r, z], pos=[0, y - r, 0], rgba=rgba, group=group, material=material - ) - ) - geoms.append( - new_geom( - geom_type="box", size=[r, y, z], pos=[x - r, 0, 0], rgba=rgba, group=group, material=material - ) - ) - geoms.append( - new_geom( - geom_type="box", size=[r, y, z], pos=[-x + r, 0, 0], rgba=rgba, group=group, material=material - ) - ) - return geoms + return self.handle0_geoms() + self.handle1_geoms() def _get_size(size, @@ -720,7 +871,9 @@ def __init__( solref=None, solimp=None, material=None, - joints=None, + joints="default", + obj_type="all", + duplicate_collision_geoms=True, ): size = _get_size(size, size_max, @@ -737,6 +890,8 @@ def __init__( solimp=solimp, material=material, joints=joints, + obj_type=obj_type, + duplicate_collision_geoms=duplicate_collision_geoms, ) def sanity_check(self): @@ -757,13 +912,8 @@ def get_top_offset(self): def get_horizontal_radius(self): return np.linalg.norm(self.size[0:2], 2) - # returns a copy, Returns xml body node - def get_collision(self, site=False): - return self._get_collision(site=site, ob_type="box") - - # returns a copy, Returns xml body node - def get_visual(self, site=False): - return self._get_visual(site=site, ob_type="box") + def get_object_subtree(self, site=False): + return self._get_object_subtree(site=site, ob_type="box") class CylinderObject(MujocoGeneratedObject): @@ -786,7 +936,9 @@ def __init__( solref=None, solimp=None, material=None, - joints=None, + joints="default", + obj_type="all", + duplicate_collision_geoms=True, ): size = _get_size(size, size_max, @@ -803,6 +955,8 @@ def __init__( solimp=solimp, material=material, joints=joints, + obj_type=obj_type, + duplicate_collision_geoms=duplicate_collision_geoms, ) def sanity_check(self): @@ -823,13 +977,8 @@ def get_top_offset(self): def get_horizontal_radius(self): return self.size[0] - # returns a copy, Returns xml body node - def get_collision(self, site=False): - return self._get_collision(site=site, ob_type="cylinder") - - # returns a copy, Returns xml body node - def get_visual(self, site=False): - return self._get_visual(site=site, ob_type="cylinder") + def get_object_subtree(self, site=False): + return self._get_object_subtree(site=site, ob_type="cylinder") class BallObject(MujocoGeneratedObject): @@ -852,7 +1001,9 @@ def __init__( solref=None, solimp=None, material=None, - joints=None, + joints="default", + obj_type="all", + duplicate_collision_geoms=True, ): size = _get_size(size, size_max, @@ -869,6 +1020,8 @@ def __init__( solimp=solimp, material=material, joints=joints, + obj_type=obj_type, + duplicate_collision_geoms=duplicate_collision_geoms, ) def sanity_check(self): @@ -889,13 +1042,8 @@ def get_top_offset(self): def get_horizontal_radius(self): return self.size[0] - # returns a copy, Returns xml body node - def get_collision(self, site=False): - return self._get_collision(site=site, ob_type="sphere") - - # returns a copy, Returns xml body node - def get_visual(self, site=False): - return self._get_visual(site=site, ob_type="sphere") + def get_object_subtree(self, site=False): + return self._get_object_subtree(site=site, ob_type="sphere") class CapsuleObject(MujocoGeneratedObject): @@ -918,7 +1066,9 @@ def __init__( solref=None, solimp=None, material=None, - joints=None, + joints="default", + obj_type="all", + duplicate_collision_geoms=True, ): size = _get_size(size, size_max, @@ -935,6 +1085,8 @@ def __init__( solimp=solimp, material=material, joints=joints, + obj_type=obj_type, + duplicate_collision_geoms=duplicate_collision_geoms, ) def sanity_check(self): @@ -955,10 +1107,5 @@ def get_top_offset(self): def get_horizontal_radius(self): return self.size[0] - # returns a copy, Returns xml body node - def get_collision(self, site=False): - return self._get_collision(site=site, ob_type="capsule") - - # returns a copy, Returns xml body node - def get_visual(self, site=False): - return self._get_visual(site=site, ob_type="capsule") + def get_object_subtree(self, site=False): + return self._get_object_subtree(site=site, ob_type="capsule") diff --git a/robosuite/models/objects/objects.py b/robosuite/models/objects/objects.py index 28ac086688..38136172a7 100644 --- a/robosuite/models/objects/objects.py +++ b/robosuite/models/objects/objects.py @@ -1,8 +1,19 @@ import copy +from copy import deepcopy import xml.etree.ElementTree as ET from robosuite.models.base import MujocoXML -from robosuite.utils.mjcf_utils import string_to_array, array_to_string, CustomMaterial +from robosuite.utils.mjcf_utils import string_to_array, array_to_string, CustomMaterial, OBJECT_COLLISION_COLOR + + +# Dict mapping geom type string keywords to group number +GEOMTYPE2GROUP = { + "collision": {0}, # If we want to use a geom for physics, but NOT visualize + "visual": {1}, # If we want to use a geom for visualization, but NOT physics + "all": {0, 1}, # If we want to use a geom for BOTH physics + visualization +} + +GEOM_GROUPS = GEOMTYPE2GROUP.keys() class MujocoObject: @@ -15,10 +26,24 @@ class MujocoObject: 2) can be swapped between different tasks Typical methods return copy so the caller can all joints/attributes as wanted + + Args: + obj_type (str): Geom elements to generate / extract for this object. Must be one of: + + :`'collision'`: Only collision geoms are returned (this corresponds to group 0 geoms) + :`'visual'`: Only visual geoms are returned (this corresponds to group 1 geoms) + :`'all'`: All geoms are returned + + duplicate_collision_geoms (bool): If set, will guarantee that each collision geom has a + visual geom copy + """ - def __init__(self): + def __init__(self, obj_type="all", duplicate_collision_geoms=True): self.asset = ET.Element("asset") + assert obj_type in GEOM_GROUPS, "object type must be one in {}, got: {} instead.".format(GEOM_GROUPS, obj_type) + self.obj_type = obj_type + self.duplicate_collision_geoms = duplicate_collision_geoms def get_bottom_offset(self): """ @@ -56,34 +81,18 @@ def get_horizontal_radius(self): float: radius """ raise NotImplementedError - # return 2 - - def get_collision(self, site=False): - """ - Returns a ET.Element - It is a subtree that defines all collision related fields - of this object. - Return should be a copy. - Must be defined by subclass. - - Args: - site (bool): Add a site (with name @name when applicable) to the returned body - Returns: - ET.Element: body - """ - raise NotImplementedError + def get_object_subtree(self, site=False): - def get_visual(self, site=False): """ Returns a ET.Element - It is a subtree that defines all visualization related fields + It is a subtree that defines all collision and / or visualization related fields of this object. Return should be a copy. Must be defined by subclass. Args: - site (bool): Add a site (with name @name when applicable) to the returned body + site (bool): If set, add a site (with name @name when applicable) to the returned body Returns: ET.Element: body @@ -103,6 +112,7 @@ def get_site_attrib_template(): "size": "0.002 0.002 0.002", "rgba": "1 0 0 1", "type": "sphere", + "group": "0", } @@ -112,14 +122,30 @@ class MujocoXMLObject(MujocoXML, MujocoObject): Args: fname (str): XML File path - name (None or str): Name of this MujocoXMLObject + + name (str): Name of this MujocoXMLObject + joints (list of dict): each dictionary corresponds to a joint that will be created for this object. The dictionary should specify the joint attributes (type, pos, etc.) according to the MuJoCo xml specification. + + obj_type (str): Geom elements to generate / extract for this object. Must be one of: + + :`'collision'`: Only collision geoms are returned (this corresponds to group 0 geoms) + :`'visual'`: Only visual geoms are returned (this corresponds to group 1 geoms) + :`'all'`: All geoms are returned + + duplicate_collision_geoms (bool): If set, will guarantee that each collision geom has a + visual geom copy """ - def __init__(self, fname, name=None, joints=None): + def __init__(self, fname, name, joints=None, obj_type="all", duplicate_collision_geoms=True): MujocoXML.__init__(self, fname) + # Set obj type and duplicate args + assert obj_type in GEOM_GROUPS, "object type must be one in {}, got: {} instead.".format(GEOM_GROUPS, obj_type) + self.obj_type = obj_type + self.duplicate_collision_geoms = duplicate_collision_geoms + # Set name self.name = name # joints for this object @@ -142,41 +168,92 @@ def get_horizontal_radius(self): ) return string_to_array(horizontal_radius_site.get("pos"))[0] - def get_collision(self, site=False): - - collision = copy.deepcopy(self.worldbody.find("./body/body[@name='collision']")) - collision.attrib.pop("name") - if self.name is not None: - collision.attrib["name"] = self.name - geoms = collision.findall("geom") - if len(geoms) == 1: - geoms[0].set("name", self.name) + def get_object_subtree(self, site=False): + # Parse object + obj = copy.deepcopy(self.worldbody.find("./body/body[@name='object']")) + # Rename this top level object body + obj.attrib.pop("name") + obj.attrib["name"] = self.name + # Get all geom_pairs in this tree + geom_pairs = self._get_geoms(obj) + + # Define a temp function so we don't duplicate so much code + obj_type = self.obj_type + + def _should_keep(el): + return int(el.get("group")) in GEOMTYPE2GROUP[obj_type] + + # Loop through each of these pairs and modify them according to @elements arg + for i, (parent, element) in enumerate(geom_pairs): + # Delete non-relevant geoms and rename remaining ones + if not _should_keep(element): + parent.remove(element) else: - for i in range(len(geoms)): - geoms[i].set("name", "{}-{}".format(self.name, i)) + g_name = element.get("name") + modded_name = f"{self.name}_{g_name}" if g_name is not None else f"{self.name}_{i}" + element.set("name", modded_name) + # Also optionally duplicate collision geoms if requested (and this is a collision geom) + if self.duplicate_collision_geoms and int(element.get("group")) in {None, 0}: + parent.append(self._duplicate_visual_from_collision(element)) + # Also manually set the visual appearances to the original collision model + element.set("rgba", array_to_string(OBJECT_COLLISION_COLOR)) + if element.get("material") is not None: + del element.attrib["material"] + # Lastly, add the site if requested if site: # add a site as well template = self.get_site_attrib_template() template["rgba"] = "1 0 0 0" - if self.name is not None: - template["name"] = self.name - collision.append(ET.Element("site", attrib=template)) - return collision + template["name"] = self.name + obj.append(ET.Element("site", attrib=template)) - def get_visual(self, site=False): + return obj - visual = copy.deepcopy(self.worldbody.find("./body/body[@name='visual']")) - visual.attrib.pop("name") - if self.name is not None: - visual.attrib["name"] = self.name - if site: - # add a site as well - template = self.get_site_attrib_template() - template["rgba"] = "1 0 0 0" - if self.name is not None: - template["name"] = self.name - visual.append(ET.Element("site", attrib=template)) - return visual + @staticmethod + def _duplicate_visual_from_collision(element): + """ + Helper function to duplicate a geom element to be a visual element. Namely, this corresponds to the + following attribute requirements: group=1, conaffinity/contype=0, no mass, name appended with "_visual" + + Args: + element (ET.Element): element to duplicate as a visual geom + + Returns: + element (ET.Element): duplicated element + """ + # Copy element + vis_element = deepcopy(element) + # Modify for visual-specific attributes (group=1, conaffinity/contype=0, no mass, update name) + vis_element.set("group", "1") + vis_element.set("conaffinity", "0") + vis_element.set("contype", "0") + vis_element.set("mass", "1e-8") + vis_element.set("name", vis_element.get("name") + "_visual") + return vis_element + + def _get_geoms(self, root, parent=None): + """ + Helper function to recursively search through element tree starting at @root and returns + a list of (parent, child) tuples where the child is a geom element + + Args: + root (ET.Element): Root of xml element tree to start recursively searching through + parent (ET.Element): Parent of the root element tree. Should not be used externally; only set + during the recursive call + + Returns: + list: array of (parent, child) tuples where the child element is a geom type + """ + # Initialize return array + geom_pairs = [] + # If the parent exists and this is a geom element, we add this current (parent, element) combo to the output + if parent is not None and root.tag == "geom": + geom_pairs.append((parent, root)) + # Loop through all children elements recursively and add to pairs + for child in root: + geom_pairs += self._get_geoms(child, parent=root) + # Return all found pairs + return geom_pairs class MujocoGeneratedObject(MujocoObject): @@ -212,8 +289,19 @@ class MujocoGeneratedObject(MujocoObject): Note that specifying a custom texture in this way automatically overrides any rgba values set - joints (list of dict): each dictionary corresponds to a joint that will be created for this object. The - dictionary should specify the joint attributes (type, pos, etc.) according to the MuJoCo xml specification. + joints (None or str or list of dict): Joints for this object. If None, no joint will be created. If "default", + a single joint will be crated. Else, should be a list of dict, where each dictionary corresponds to a joint + that will be created for this object. The dictionary should specify the joint attributes (type, pos, etc.) + according to the MuJoCo xml specification. + + obj_type (str): Geom elements to generate / extract for this object. Must be one of: + + :`'collision'`: Only collision geoms are returned (this corresponds to group 0 geoms) + :`'visual'`: Only visual geoms are returned (this corresponds to group 1 geoms) + :`'all'`: All geoms are returned + + duplicate_collision_geoms (bool): If set, will guarantee that each collision geom has a + visual geom copy """ def __init__( @@ -226,9 +314,11 @@ def __init__( solref=None, solimp=None, material=None, - joints=None, + joints="default", + obj_type="all", + duplicate_collision_geoms=True, ): - super().__init__() + super().__init__(obj_type=obj_type, duplicate_collision_geoms=duplicate_collision_geoms) self.name = name @@ -276,8 +366,10 @@ def __init__( self.append_material(material) # joints for this object - if joints is None: + if joints == "default": self.joints = [{'type': 'free'}] # default free joint + elif joints is None: + self.joints = [] else: self.joints = joints @@ -299,8 +391,7 @@ def get_collision_attrib_template(): Returns: dict: Initial template with `'pos'` and `'group'` already specified """ - # TODO: collision group should be 0, but this removes the generated obj like the cube when we only render visual meshes - return {"pos": "0 0 0", "group": "1"} + return {"group": "0", "rgba": array_to_string(OBJECT_COLLISION_COLOR)} @staticmethod def get_visual_attrib_template(): @@ -310,7 +401,7 @@ def get_visual_attrib_template(): Returns: dict: Initial template with `'conaffinity'`, `'contype'`, and `'group'` already specified """ - return {"conaffinity": "0", "contype": "0", "group": "1"} + return {"conaffinity": "0", "contype": "0", "mass": "1e-8", "group": "1"} def append_material(self, material): """ @@ -333,50 +424,48 @@ def append_material(self, material): self.asset.append(ET.Element("texture", attrib=material.tex_attrib)) self.asset.append(ET.Element("material", attrib=material.mat_attrib)) - def _get_collision(self, site=False, ob_type="box"): - main_body = ET.Element("body") - main_body.set("name", self.name) - template = self.get_collision_attrib_template() - template["name"] = self.name - template["type"] = ob_type - if self.material == "default": - template["rgba"] = "0.5 0.5 0.5 1" # mujoco default - template["material"] = "{}_mat".format(self.name) - elif self.material is not None: - template["material"] = self.material.mat_attrib["name"] - else: - template["rgba"] = array_to_string(self.rgba) - template["size"] = array_to_string(self.size) - template["density"] = str(self.density) - template["friction"] = array_to_string(self.friction) - template["solref"] = array_to_string(self.solref) - template["solimp"] = array_to_string(self.solimp) - main_body.append(ET.Element("geom", attrib=template)) - if site: - # add a site as well - template = self.get_site_attrib_template() - template["name"] = self.name - main_body.append(ET.Element("site", attrib=template)) - return main_body - - def _get_visual(self, site=False, ob_type="box"): - main_body = ET.Element("body") - main_body.set("name", self.name) - template = self.get_visual_attrib_template() - template["name"] = self.name - template["type"] = ob_type - if self.material == "default": - template["rgba"] = "0.5 0.5 0.5 1" # mujoco default - template["material"] = "{}_mat".format(self.name) - elif self.material is not None: - template["material"] = self.material.mat_attrib["name"] - else: - template["rgba"] = array_to_string(self.rgba) - template["size"] = array_to_string(self.size) - main_body.append(ET.Element("geom", attrib=template)) + def get_object_subtree(self, site=False): + raise NotImplementedError + + def _get_object_subtree(self, site=False, ob_type="box"): + # Create element tree + obj = ET.Element("body") + obj.set("name", self.name) + + # Get base element attributes + element_attr = { + "name": self.name, + "type": ob_type, + "size": array_to_string(self.size) + } + + # Add collision geom if necessary + if self.obj_type in {"collision", "all"}: + col_element_attr = deepcopy(element_attr) + col_element_attr.update(self.get_collision_attrib_template()) + # col_element_attr["name"] += "_col" + col_element_attr["density"] = str(self.density) + col_element_attr["friction"] = array_to_string(self.friction) + col_element_attr["solref"] = array_to_string(self.solref) + col_element_attr["solimp"] = array_to_string(self.solimp) + obj.append(ET.Element("geom", attrib=col_element_attr)) + # Add visual geom if necessary + if self.obj_type in {"visual", "all"}: + vis_element_attr = deepcopy(element_attr) + vis_element_attr.update(self.get_visual_attrib_template()) + vis_element_attr["name"] += "_vis" + if self.material == "default": + vis_element_attr["rgba"] = "0.5 0.5 0.5 1" # mujoco default + vis_element_attr["material"] = "{}_mat".format(self.name) + elif self.material is not None: + vis_element_attr["material"] = self.material.mat_attrib["name"] + else: + vis_element_attr["rgba"] = array_to_string(self.rgba) + obj.append(ET.Element("geom", attrib=vis_element_attr)) + # Lastly, add site as well if requested if site: # add a site as well - template = self.get_site_attrib_template() - template["name"] = self.name - main_body.append(ET.Element("site", attrib=template)) - return main_body + site_element_attr = self.get_site_attrib_template() + site_element_attr["name"] = self.name + obj.append(ET.Element("site", attrib=site_element_attr)) + return obj diff --git a/robosuite/models/objects/xml_objects.py b/robosuite/models/objects/xml_objects.py index f424f899d6..c3d470adde 100644 --- a/robosuite/models/objects/xml_objects.py +++ b/robosuite/models/objects/xml_objects.py @@ -9,7 +9,8 @@ class BottleObject(MujocoXMLObject): """ def __init__(self, name=None, joints=None): - super().__init__(xml_path_completion("objects/bottle.xml"), name=name, joints=joints) + super().__init__(xml_path_completion("objects/bottle.xml"), + name=name, joints=joints, obj_type="all", duplicate_collision_geoms=True) class CanObject(MujocoXMLObject): @@ -18,7 +19,8 @@ class CanObject(MujocoXMLObject): """ def __init__(self, name=None, joints=None): - super().__init__(xml_path_completion("objects/can.xml"), name=name, joints=joints) + super().__init__(xml_path_completion("objects/can.xml"), + name=name, joints=joints, obj_type="all", duplicate_collision_geoms=True) class LemonObject(MujocoXMLObject): @@ -27,7 +29,8 @@ class LemonObject(MujocoXMLObject): """ def __init__(self, name=None, joints=None): - super().__init__(xml_path_completion("objects/lemon.xml"), name=name, joints=joints) + super().__init__(xml_path_completion("objects/lemon.xml"), + name=name, joints=joints, obj_type="all", duplicate_collision_geoms=True) class MilkObject(MujocoXMLObject): @@ -36,7 +39,8 @@ class MilkObject(MujocoXMLObject): """ def __init__(self, name=None, joints=None): - super().__init__(xml_path_completion("objects/milk.xml"), name=name, joints=joints) + super().__init__(xml_path_completion("objects/milk.xml"), + name=name, joints=joints, obj_type="all", duplicate_collision_geoms=True) class BreadObject(MujocoXMLObject): @@ -45,7 +49,8 @@ class BreadObject(MujocoXMLObject): """ def __init__(self, name=None, joints=None): - super().__init__(xml_path_completion("objects/bread.xml"), name=name, joints=joints) + super().__init__(xml_path_completion("objects/bread.xml"), + name=name, joints=joints, obj_type="all", duplicate_collision_geoms=True) class CerealObject(MujocoXMLObject): @@ -54,7 +59,8 @@ class CerealObject(MujocoXMLObject): """ def __init__(self, name=None, joints=None): - super().__init__(xml_path_completion("objects/cereal.xml"), name=name, joints=joints) + super().__init__(xml_path_completion("objects/cereal.xml"), + name=name, joints=joints, obj_type="all", duplicate_collision_geoms=True) class SquareNutObject(MujocoXMLObject): @@ -63,7 +69,8 @@ class SquareNutObject(MujocoXMLObject): """ def __init__(self, name=None, joints=None): - super().__init__(xml_path_completion("objects/square-nut.xml"), name=name, joints=joints) + super().__init__(xml_path_completion("objects/square-nut.xml"), + name=name, joints=joints, obj_type="all", duplicate_collision_geoms=True) class RoundNutObject(MujocoXMLObject): @@ -72,7 +79,8 @@ class RoundNutObject(MujocoXMLObject): """ def __init__(self, name=None, joints=None): - super().__init__(xml_path_completion("objects/round-nut.xml"), name=name, joints=joints) + super().__init__(xml_path_completion("objects/round-nut.xml"), + name=name, joints=joints, obj_type="all", duplicate_collision_geoms=True) class MilkVisualObject(MujocoXMLObject): @@ -84,7 +92,8 @@ class MilkVisualObject(MujocoXMLObject): """ def __init__(self, name=None, joints=None): - super().__init__(xml_path_completion("objects/milk-visual.xml"), name=name, joints=joints) + super().__init__(xml_path_completion("objects/milk-visual.xml"), + name=name, joints=joints, obj_type="visual", duplicate_collision_geoms=True) class BreadVisualObject(MujocoXMLObject): @@ -96,7 +105,8 @@ class BreadVisualObject(MujocoXMLObject): """ def __init__(self, name=None, joints=None): - super().__init__(xml_path_completion("objects/bread-visual.xml"), name=name, joints=joints) + super().__init__(xml_path_completion("objects/bread-visual.xml"), + name=name, joints=joints, obj_type="visual", duplicate_collision_geoms=True) class CerealVisualObject(MujocoXMLObject): @@ -108,7 +118,8 @@ class CerealVisualObject(MujocoXMLObject): """ def __init__(self, name=None, joints=None): - super().__init__(xml_path_completion("objects/cereal-visual.xml"), name=name, joints=joints) + super().__init__(xml_path_completion("objects/cereal-visual.xml"), + name=name, joints=joints, obj_type="visual", duplicate_collision_geoms=True) class CanVisualObject(MujocoXMLObject): @@ -120,7 +131,8 @@ class CanVisualObject(MujocoXMLObject): """ def __init__(self, name=None, joints=None): - super().__init__(xml_path_completion("objects/can-visual.xml"), name=name, joints=joints) + super().__init__(xml_path_completion("objects/can-visual.xml"), + name=name, joints=joints, obj_type="visual", duplicate_collision_geoms=True) class PlateWithHoleObject(MujocoXMLObject): @@ -129,7 +141,8 @@ class PlateWithHoleObject(MujocoXMLObject): """ def __init__(self, name=None, joints=None): - super().__init__(xml_path_completion("objects/plate-with-hole.xml"), name=name, joints=joints) + super().__init__(xml_path_completion("objects/plate-with-hole.xml"), + name=name, joints=joints, obj_type="all", duplicate_collision_geoms=True) class DoorObject(MujocoXMLObject): @@ -145,7 +158,8 @@ def __init__(self, name=None, joints=None, friction=None, damping=None, lock=Fal xml_path = "objects/door.xml" if lock: xml_path = "objects/door_lock.xml" - super().__init__(xml_path_completion(xml_path), name=name, joints=joints) + super().__init__(xml_path_completion(xml_path), + name=name, joints=joints, obj_type="all", duplicate_collision_geoms=True) self.lock = lock self.friction = friction self.damping = damping @@ -161,8 +175,8 @@ def _set_door_friction(self, friction): Args: friction (3-tuple of float): friction parameters to override the ones specified in the XML """ - collision = self.worldbody.find("./body/body[@name='collision']") - node = collision.find("./body[@name='frame']") + obj = self.worldbody.find("./body/body[@name='object']") + node = obj.find("./body[@name='frame']") node = node.find("./body[@name='door']") hinge = node.find("./joint[@name='door_hinge']") hinge.set("frictionloss", array_to_string(np.array([friction]))) @@ -174,8 +188,8 @@ def _set_door_damping(self, damping): Args: damping (float): damping parameter to override the ones specified in the XML """ - collision = self.worldbody.find("./body/body[@name='collision']") - node = collision.find("./body[@name='frame']") + obj = self.worldbody.find("./body/body[@name='object']") + node = obj.find("./body[@name='frame']") node = node.find("./body[@name='door']") hinge = node.find("./joint[@name='door_hinge']") hinge.set("damping", array_to_string(np.array([damping]))) diff --git a/robosuite/models/robots/robot_model.py b/robosuite/models/robots/robot_model.py index 9524f969ef..202ad69b5c 100644 --- a/robosuite/models/robots/robot_model.py +++ b/robosuite/models/robots/robot_model.py @@ -2,7 +2,7 @@ from robosuite.models.base import MujocoXML from robosuite.utils import XMLError -from robosuite.utils.mjcf_utils import array_to_string +from robosuite.utils.mjcf_utils import array_to_string, ROBOT_COLLISION_COLOR from robosuite.utils.transform_utils import euler2mat, mat2quat import numpy as np @@ -90,6 +90,9 @@ def __init__(self, fname, idn=0, bottom_offset=(0, 0, 0)): # Update all xml element prefixes self.add_prefix(self.naming_prefix) + # Recolor all collision geoms appropriately + self.recolor_collision_geoms(ROBOT_COLLISION_COLOR) + # key: gripper name and value: gripper model self.grippers = OrderedDict() diff --git a/robosuite/models/tasks/manipulation_task.py b/robosuite/models/tasks/manipulation_task.py index 1a7f0ef946..fb318f5662 100644 --- a/robosuite/models/tasks/manipulation_task.py +++ b/robosuite/models/tasks/manipulation_task.py @@ -3,7 +3,7 @@ from robosuite.models.world import MujocoWorldBase from robosuite.models.tasks import UniformRandomSampler -from robosuite.models.objects import MujocoGeneratedObject, MujocoXMLObject +from robosuite.models.objects import MujocoObject from robosuite.utils.mjcf_utils import new_joint, array_to_string @@ -22,9 +22,6 @@ class ManipulationTask(MujocoWorldBase): mujoco_objects (OrderedDict of MujocoObject): a list of MJCF models of physical objects - visual_objects (OrderedDict of MujocoObject): a list of MJCF models of visual-only objects that do not - participate in collisions - initializer (ObjectPositionSampler): placement sampler to initialize object positions. Raises: @@ -35,8 +32,7 @@ def __init__( self, mujoco_arena, mujoco_robots, - mujoco_objects, - visual_objects=None, + mujoco_objects, initializer=None, ): super().__init__() @@ -48,23 +44,16 @@ def __init__( if initializer is None: initializer = UniformRandomSampler() - if visual_objects is None: - visual_objects = collections.OrderedDict() - assert isinstance(mujoco_objects, collections.OrderedDict) - assert isinstance(visual_objects, collections.OrderedDict) mujoco_objects = deepcopy(mujoco_objects) - visual_objects = deepcopy(visual_objects) # xml manifestations of all objects self.objects = [] self.merge_objects(mujoco_objects) - self.merge_objects(visual_objects, is_visual=True) - merged_objects = collections.OrderedDict(**mujoco_objects, **visual_objects) + merged_objects = collections.OrderedDict(**mujoco_objects) self.mujoco_objects = mujoco_objects - self.visual_objects = visual_objects self.initializer = initializer self.initializer.setup(merged_objects, self.table_top_offset, self.table_size) @@ -90,32 +79,29 @@ def merge_arena(self, mujoco_arena): self.table_size = mujoco_arena.table_full_size self.merge(mujoco_arena) - def merge_objects(self, mujoco_objects, is_visual=False): + def merge_objects(self, mujoco_objects): """ Adds object models to the MJCF model. Args: - mujoco_objects (OrderedDict or MujocoObject): objects to merge into this MJCF model - is_visual (bool): Whether the object is a visual object or not + mujoco_objects (OrderedDict of MujocoObject): objects to merge into this MJCF model """ - if not is_visual: - self.max_horizontal_radius = 0 + self.max_horizontal_radius = 0 for obj_name, obj_mjcf in mujoco_objects.items(): - assert(isinstance(obj_mjcf, MujocoGeneratedObject) or isinstance(obj_mjcf, MujocoXMLObject)) + assert(isinstance(obj_mjcf, MujocoObject)) self.merge_asset(obj_mjcf) # Load object - if is_visual: - obj = obj_mjcf.get_visual(site=False) - else: - obj = obj_mjcf.get_collision(site=True) + obj = obj_mjcf.get_object_subtree(site=(obj_mjcf.obj_type != "visual")) for i, joint in enumerate(obj_mjcf.joints): - obj.append(new_joint(name="{}_jnt{}".format(obj_name, i), **joint)) + if "name" not in joint: + joint["name"] = "{}_jnt{}".format(obj_name, i) + obj.append(new_joint(**joint)) self.objects.append(obj) self.worldbody.append(obj) - if not is_visual: + if obj_mjcf.obj_type != "visual": self.max_horizontal_radius = max( self.max_horizontal_radius, obj_mjcf.get_horizontal_radius() ) diff --git a/robosuite/utils/mjcf_utils.py b/robosuite/utils/mjcf_utils.py index 369fa166fe..5f2fbf6407 100644 --- a/robosuite/utils/mjcf_utils.py +++ b/robosuite/utils/mjcf_utils.py @@ -12,6 +12,11 @@ RED = [1, 0, 0, 1] GREEN = [0, 1, 0, 1] BLUE = [0, 0, 1, 1] +CYAN = [0, 1, 1, 1] +ROBOT_COLLISION_COLOR = [0, 0.5, 0, 1] +GRIPPER_COLLISION_COLOR = [0, 0, 0.5, 1] +OBJECT_COLLISION_COLOR = [0.5, 0, 0, 1] +ENVIRONMENT_COLLISION_COLOR = [0.5, 0.5, 0, 1] TEXTURES = { "WoodRed": "red-wood.png", @@ -111,7 +116,7 @@ def new_joint(**kwargs): Creates a joint tag with attributes specified by @**kwargs. Args: - **kwargs (dict): Specified attributes for the new joint + **kwargs: Specified attributes for the new joint Returns: ET.Element: new joint xml element @@ -128,7 +133,7 @@ def new_actuator(joint, act_type="actuator", **kwargs): joint (str): type of actuator transmission. see all types here: http://mujoco.org/book/modeling.html#actuator act_type (str): actuator type. Defaults to "actuator" - **kwargs (dict): Any additional specified attributes for the new joint + **kwargs: Any additional specified attributes for the new joint Returns: ET.Element: new actuator xml element @@ -151,7 +156,7 @@ def new_site(name, rgba=RED, pos=(0, 0, 0), size=(0.005,), **kwargs): rgba (4-array): (r,g,b,a) color and transparency. Defaults to solid red. pos (3-array): (x,y,z) 3d position of the site. size (array of float): site size (sites are spherical by default). - **kwargs (dict): Any additional specified attributes for the new site + **kwargs: Any additional specified attributes for the new site Returns: ET.Element: new site xml element @@ -184,7 +189,7 @@ def new_geom(geom_type, size, pos=(0, 0, 0), rgba=RED, group=0, **kwargs): rgba (4-array): (r,g,b,a) color and transparency. Defaults to solid red. group (int): the integrer group that the geom belongs to. useful for separating visual and physical elements. - **kwargs (dict): Any additional specified attributes for the new geom + **kwargs: Any additional specified attributes for the new geom Returns: ET.Element: new geom xml element @@ -209,7 +214,7 @@ def new_body(name=None, pos=None, **kwargs): Args: name (str): body name. pos (3-array): (x,y,z) 3d position of the body frame. - **kwargs (dict): Any additional specified attributes for the new body + **kwargs: Any additional specified attributes for the new body Returns: ET.Element: new body xml element @@ -230,7 +235,7 @@ def new_inertial(name=None, pos=(0, 0, 0), mass=None, **kwargs): name (str): [NOT USED] pos (3-array): (x,y,z) 3d position of the inertial frame. mass (float): The mass of inertial - **kwargs (dict): Any additional specified attributes for the new inertial element + **kwargs: Any additional specified attributes for the new inertial element Returns: ET.Element: new inertial xml element @@ -280,6 +285,24 @@ def postprocess_model_xml(xml_str): return ET.tostring(root, encoding="utf8").decode("utf8") +def add_to_dict(dic, **kwargs): + """ + Helper function to add key-values to dictionary @dic where each entry is its own array (list). + Args: + dic (dict): Dictionary to which new key / value pairs will be added. If the key already exists, + will append the value to that key entry + + Returns: + dict: Modified dictionary + """ + for k, v in kwargs.items(): + if k in dic: + dic[k].append(v) + else: + dic[k] = [v] + return dic + + class CustomMaterial(object): """ Simple class to instantiate the necessary parameters to define an appropriate texture / material combo From 6b3798e80b7ae031e91e2875e0b878225a15739f Mon Sep 17 00:00:00 2001 From: josiah_wong <84cremebrule@gmail.com> Date: Thu, 3 Dec 2020 11:00:43 -0800 Subject: [PATCH 08/46] fix ringbuffer so it is accurate even during warmstart --- robosuite/utils/control_utils.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/robosuite/utils/control_utils.py b/robosuite/utils/control_utils.py index 70378d278e..234dfa5cbb 100644 --- a/robosuite/utils/control_utils.py +++ b/robosuite/utils/control_utils.py @@ -290,6 +290,9 @@ def __init__(self, dim, length): self.dim = dim self.length = length + # Variable so that initial average values are accurate + self._size = 0 + # Save pointer to current place in the buffer self.ptr = 0 @@ -303,9 +306,11 @@ def push(self, value): Args: value (int or float or array): Value(s) to push into the array (taken as a single new element) """ - # Add value, then increment pointer + # Add value, then increment pointer (and size if necessary) self.buf[self.ptr] = np.array(value) self.ptr = (self.ptr + 1) % self.length + if self._size < self.length: + self._size += 1 def clear(self): """ @@ -313,6 +318,7 @@ def clear(self): """ self.buf = np.zeros((self.length, self.dim)) self.ptr = 0 + self._size = 0 @property def average(self): @@ -322,7 +328,7 @@ def average(self): Returns: float or np.array: Averaged value of all elements in buffer """ - return np.mean(self.buf, axis=0) + return np.mean(self.buf[:self._size], axis=0) class DeltaBuffer(Buffer): From de99886c31076d9cf6b1caf12d9bbbf2b152464d Mon Sep 17 00:00:00 2001 From: josiah_wong <84cremebrule@gmail.com> Date: Thu, 3 Dec 2020 12:30:20 -0800 Subject: [PATCH 09/46] overhaul all (robot, gripper, object) model classes / hierarchy, remove redundancies, fix misc. bugs, standardize all envs and class hierarchy --- robosuite/__init__.py | 19 +- robosuite/environments/base.py | 119 ++- .../environments/{ => manipulation}/door.py | 130 ++-- .../environments/{ => manipulation}/lift.py | 134 ++-- .../manipulation/manipulation_env.py | 308 ++++++++ .../{ => manipulation}/nut_assembly.py | 365 ++++----- .../{ => manipulation}/pick_place.py | 375 ++++------ .../manipulation/single_arm_env.py | 69 ++ .../environments/{ => manipulation}/stack.py | 164 ++-- .../environments/manipulation/two_arm_env.py | 130 ++++ .../{ => manipulation}/two_arm_handover.py | 261 ++----- .../{ => manipulation}/two_arm_lift.py | 207 ++---- .../{ => manipulation}/two_arm_peg_in_hole.py | 143 ++-- .../environments/{ => manipulation}/wipe.py | 242 +++--- robosuite/environments/robot_env.py | 211 ++---- robosuite/models/arenas/wipe_arena.py | 81 +- robosuite/models/assets/base.xml | 2 +- .../grippers/jaco_three_finger_gripper.xml | 11 +- .../models/assets/grippers/null_gripper.xml | 2 +- .../models/assets/grippers/panda_gripper.xml | 2 +- .../assets/grippers/rethink_gripper.xml | 4 +- .../assets/grippers/robotiq_gripper_140.xml | 4 +- .../assets/grippers/robotiq_gripper_85.xml | 4 +- .../assets/grippers/robotiq_gripper_s.xml | 28 +- .../models/assets/grippers/wiping_gripper.xml | 2 +- .../models/assets/objects/bread-visual.xml | 4 +- .../models/assets/objects/can-visual.xml | 4 +- .../models/assets/objects/cereal-visual.xml | 4 +- robosuite/models/assets/objects/door.xml | 8 +- robosuite/models/assets/objects/door_lock.xml | 8 +- .../models/assets/objects/milk-visual.xml | 4 +- robosuite/models/assets/objects/round-nut.xml | 7 +- .../models/assets/objects/square-nut.xml | 7 +- .../models/assets/robots/baxter/robot.xml | 3 +- .../models/assets/robots/panda/robot.xml | 147 ++-- robosuite/models/base.py | 238 +++++- robosuite/models/grippers/gripper_model.py | 162 ++-- .../grippers/jaco_three_finger_gripper.py | 39 +- robosuite/models/grippers/null_gripper.py | 16 - robosuite/models/grippers/panda_gripper.py | 28 +- robosuite/models/grippers/rethink_gripper.py | 29 +- .../models/grippers/robotiq_140_gripper.py | 42 +- .../models/grippers/robotiq_85_gripper.py | 46 +- .../grippers/robotiq_three_finger_gripper.py | 56 +- robosuite/models/grippers/wiping_gripper.py | 18 +- robosuite/models/objects/generated_objects.py | 702 ++++++++++++------ robosuite/models/objects/objects.py | 375 +++++----- robosuite/models/objects/xml_objects.py | 127 +++- robosuite/models/robots/__init__.py | 10 +- robosuite/models/robots/baxter_robot.py | 147 ---- robosuite/models/robots/iiwa_robot.py | 80 -- robosuite/models/robots/jaco_robot.py | 80 -- robosuite/models/robots/kinova3_robot.py | 80 -- .../models/robots/manipulators/__init__.py | 8 + .../robots/manipulators/baxter_robot.py | 79 ++ .../models/robots/manipulators/iiwa_robot.py | 40 + .../models/robots/manipulators/jaco_robot.py | 40 + .../robots/manipulators/kinova3_robot.py | 40 + .../robots/manipulators/manipulator_model.py | 159 ++++ .../robots/{ => manipulators}/panda_robot.py | 43 +- .../robots/manipulators/sawyer_robot.py | 40 + .../models/robots/manipulators/ur5e_robot.py | 40 + robosuite/models/robots/robot_model.py | 325 ++------ robosuite/models/robots/sawyer_robot.py | 80 -- robosuite/models/robots/ur5e_robot.py | 80 -- robosuite/models/tasks/__init__.py | 1 + robosuite/models/tasks/manipulation_task.py | 117 +-- robosuite/models/tasks/placement_sampler.py | 463 +++++------- robosuite/models/tasks/task.py | 78 ++ robosuite/robots/__init__.py | 18 +- robosuite/robots/bimanual.py | 147 +--- robosuite/robots/manipulator.py | 143 ++++ robosuite/robots/robot.py | 135 +++- robosuite/robots/single_arm.py | 151 +--- robosuite/utils/mjcf_utils.py | 351 +++++++-- robosuite/utils/robot_utils.py | 16 + 76 files changed, 4097 insertions(+), 3985 deletions(-) rename robosuite/environments/{ => manipulation}/door.py (83%) rename robosuite/environments/{ => manipulation}/lift.py (79%) create mode 100644 robosuite/environments/manipulation/manipulation_env.py rename robosuite/environments/{ => manipulation}/nut_assembly.py (67%) rename robosuite/environments/{ => manipulation}/pick_place.py (69%) create mode 100644 robosuite/environments/manipulation/single_arm_env.py rename robosuite/environments/{ => manipulation}/stack.py (77%) create mode 100644 robosuite/environments/manipulation/two_arm_env.py rename robosuite/environments/{ => manipulation}/two_arm_handover.py (72%) rename robosuite/environments/{ => manipulation}/two_arm_lift.py (74%) rename robosuite/environments/{ => manipulation}/two_arm_peg_in_hole.py (80%) rename robosuite/environments/{ => manipulation}/wipe.py (78%) delete mode 100644 robosuite/models/robots/baxter_robot.py delete mode 100644 robosuite/models/robots/iiwa_robot.py delete mode 100644 robosuite/models/robots/jaco_robot.py delete mode 100644 robosuite/models/robots/kinova3_robot.py create mode 100644 robosuite/models/robots/manipulators/__init__.py create mode 100644 robosuite/models/robots/manipulators/baxter_robot.py create mode 100644 robosuite/models/robots/manipulators/iiwa_robot.py create mode 100644 robosuite/models/robots/manipulators/jaco_robot.py create mode 100644 robosuite/models/robots/manipulators/kinova3_robot.py create mode 100644 robosuite/models/robots/manipulators/manipulator_model.py rename robosuite/models/robots/{ => manipulators}/panda_robot.py (53%) create mode 100644 robosuite/models/robots/manipulators/sawyer_robot.py create mode 100644 robosuite/models/robots/manipulators/ur5e_robot.py delete mode 100644 robosuite/models/robots/sawyer_robot.py delete mode 100644 robosuite/models/robots/ur5e_robot.py create mode 100644 robosuite/models/tasks/task.py create mode 100644 robosuite/robots/manipulator.py create mode 100644 robosuite/utils/robot_utils.py diff --git a/robosuite/__init__.py b/robosuite/__init__.py index cec6df6a44..e68e8f3556 100644 --- a/robosuite/__init__.py +++ b/robosuite/__init__.py @@ -1,14 +1,15 @@ from robosuite.environments.base import make -from robosuite.environments.lift import Lift -from robosuite.environments.stack import Stack -from robosuite.environments.nut_assembly import NutAssembly -from robosuite.environments.pick_place import PickPlace -from robosuite.environments.door import Door -from robosuite.environments.wipe import Wipe -from robosuite.environments.two_arm_lift import TwoArmLift -from robosuite.environments.two_arm_peg_in_hole import TwoArmPegInHole -from robosuite.environments.two_arm_handover import TwoArmHandover +# Manipulation environments +from robosuite.environments.manipulation.lift import Lift +from robosuite.environments.manipulation.stack import Stack +from robosuite.environments.manipulation.nut_assembly import NutAssembly +from robosuite.environments.manipulation.pick_place import PickPlace +from robosuite.environments.manipulation.door import Door +from robosuite.environments.manipulation.wipe import Wipe +from robosuite.environments.manipulation.two_arm_lift import TwoArmLift +from robosuite.environments.manipulation.two_arm_peg_in_hole import TwoArmPegInHole +from robosuite.environments.manipulation.two_arm_handover import TwoArmHandover from robosuite.environments import ALL_ENVIRONMENTS from robosuite.controllers import ALL_CONTROLLERS, load_controller_config diff --git a/robosuite/environments/base.py b/robosuite/environments/base.py index 1378c2f0d4..d55c76684b 100644 --- a/robosuite/environments/base.py +++ b/robosuite/environments/base.py @@ -3,6 +3,9 @@ from mujoco_py import load_model_from_xml from robosuite.utils import SimulationError, XMLError, MujocoPyRenderer +from robosuite.models.base import MujocoModel + +import numpy as np REGISTERED_ENVS = {} @@ -44,7 +47,7 @@ def __new__(meta, name, bases, class_dict): cls = super().__new__(meta, name, bases, class_dict) # List all environments that should not be registered here. - _unregistered_envs = ["MujocoEnv", "RobotEnv"] + _unregistered_envs = ["MujocoEnv", "RobotEnv", "ManipulationEnv"] if cls.__name__ not in _unregistered_envs: register_env(cls) @@ -56,6 +59,12 @@ class MujocoEnv(metaclass=EnvMeta): Initializes a Mujoco Environment. Args: + use_indicator_object (bool): if True, sets up an indicator object that + is useful for debugging. + + env_visualization (bool): True if visualizing sites for the arena / objects in this environment. Useful for + teleoperation. + has_renderer (bool): If true, render the simulation state in a viewer instead of headless mode. @@ -88,6 +97,8 @@ class MujocoEnv(metaclass=EnvMeta): def __init__( self, + use_indicator_object=False, + env_visualization=False, has_renderer=False, has_offscreen_renderer=True, render_camera="frontview", @@ -110,6 +121,12 @@ def __init__( self.render_visual_mesh = render_visual_mesh self.viewer = None + # whether to use indicator object or not + self.use_indicator_object = use_indicator_object + + # Whether to render visual sites in this env + self.env_visualization = env_visualization + # Simulation-specific attributes self.control_freq = control_freq self.horizon = horizon @@ -158,7 +175,15 @@ def _get_reference(self): index or a list of indices that point to the corresponding elements in a flatten array, which is how MuJoCo stores physical simulation data. """ - pass + # Indicator object references + if self.use_indicator_object: + ind_qpos = self.sim.model.get_joint_qpos_addr("pos_indicator") + self._ref_indicator_pos_low, self._ref_indicator_pos_high = ind_qpos + + ind_qvel = self.sim.model.get_joint_qvel_addr("pos_indicator") + self._ref_indicator_vel_low, self._ref_indicator_vel_high = ind_qvel + + self.indicator_id = self.sim.model.body_name2id("pos_indicator") def _initialize_sim(self, xml_string=None): """ @@ -226,6 +251,10 @@ def _reset_internal(self): self.sim._render_context_offscreen.vopt.geomgroup[0] = (1 if self.render_collision_mesh else 0) self.sim._render_context_offscreen.vopt.geomgroup[1] = (1 if self.render_visual_mesh else 0) + # Set visuals for objects + for obj in self.model.mujoco_objects: + obj.set_sites_visibility(sim=self.sim, visible=self.env_visualization) + # additional housekeeping self.sim_state_initial = self.sim.get_state() self._get_reference() @@ -316,6 +345,10 @@ def _post_action(self, action): # done if number of elapsed timesteps is greater than horizon self.done = (self.timestep >= self.horizon) and not self.ignore_done + + # Run any necessary visualization + self._visualization() + return reward, self.done, {} def reward(self, action): @@ -351,6 +384,37 @@ def observation_spec(self): observation = self._get_observation() return observation + def move_indicator(self, pos): + """ + Sets 3d position of indicator object to @pos. + + Args: + pos (3-tuple): (x,y,z) values to place the indicator within the env + """ + if self.use_indicator_object: + index = self._ref_indicator_pos_low + self.sim.data.qpos[index:index + 3] = pos + + def clear_objects(self, object_names): + """ + Clears objects with the name @object_names out of the task space. This is useful + for supporting task modes with single types of objects, as in + @self.single_object_mode without changing the model definition. + + Args: + object_names (str or list of str): Name of object(s) to remove from the task workspace + """ + object_names = {object_names} if type(object_names) is str else set(object_names) + for obj in self.model.mujoco_objects: + if obj.name in object_names: + self.sim.data.set_joint_qpos(obj.joints[0], np.array((10, 10, 10, 1, 0, 0, 0))) + + def _visualization(self): + """ + Do any needed visualization here + """ + pass + @property def action_spec(self): """ @@ -393,31 +457,66 @@ def reset_from_xml_string(self, xml_string): # Turn off deterministic reset self.deterministic_reset = False - def find_contacts(self, geoms_1, geoms_2): + def check_contact(self, geoms_1, geoms_2=None): """ Finds contact between two geom groups. Args: - geoms_1 (str or list of str): an individual geom name or list of geom names - geoms_2 (str or list of str): another individual geom name or list of geom names + geoms_1 (str or list of str or MujocoModel): an individual geom name or list of geom names or a model. If + a MujocoModel is specified, the geoms checked will be its contact_geoms + geoms_2 (str or list of str or MujocoModel or None): another individual geom name or list of geom names. + If a MujocoModel is specified, the geoms checked will be its contact_geoms. If None, will check + any collision with @geoms_1 to any other geom in the environment Returns: - generator: iterator of all contacts between @geoms_1 and @geoms_2 + bool: True if any geom in @geoms_1 is in contact with any geom in @geoms_2. """ # Check if either geoms_1 or geoms_2 is a string, convert to list if so if type(geoms_1) is str: geoms_1 = [geoms_1] + elif isinstance(geoms_1, MujocoModel): + geoms_1 = geoms_1.contact_geoms if type(geoms_2) is str: geoms_2 = [geoms_2] - for contact in self.sim.data.contact[0 : self.sim.data.ncon]: + elif isinstance(geoms_2, MujocoModel): + geoms_2 = geoms_2.contact_geoms + for contact in self.sim.data.contact[: self.sim.data.ncon]: # check contact geom in geoms c1_in_g1 = self.sim.model.geom_id2name(contact.geom1) in geoms_1 - c2_in_g2 = self.sim.model.geom_id2name(contact.geom2) in geoms_2 + c2_in_g2 = self.sim.model.geom_id2name(contact.geom2) in geoms_2 if geoms_2 is not None else True # check contact geom in geoms (flipped) c2_in_g1 = self.sim.model.geom_id2name(contact.geom2) in geoms_1 - c1_in_g2 = self.sim.model.geom_id2name(contact.geom1) in geoms_2 + c1_in_g2 = self.sim.model.geom_id2name(contact.geom1) in geoms_2 if geoms_2 is not None else True if (c1_in_g1 and c2_in_g2) or (c1_in_g2 and c2_in_g1): - yield contact + return True + return False + + def get_contacts(self, model): + """ + Checks for any contacts with @model (as defined by @model's contact_geoms) and returns the set of + geom names currently in contact with that model (excluding the geoms that are part of the model itself). + + Args: + model (MujocoModel): Model to check contacts for. + + Returns: + set: Unique geoms that are actively in contact with this model. + + Raises: + AssertionError: [Invalid input type] + """ + # Make sure model is MujocoModel type + assert isinstance(model, MujocoModel), \ + "Inputted model must be of type MujocoModel; got type {} instead!".format(type(model)) + contact_set = set() + for contact in self.sim.data.contact[: self.sim.data.ncon]: + # check contact geom in geoms; add to contact set if match is found + g1, g2 = self.sim.model.geom_id2name(contact.geom1), self.sim.model.geom_id2name(contact.geom2) + if g1 in model.contact_geoms and g2 not in model.contact_geoms: + contact_set.add(g2) + elif g2 in model.contact_geoms and g1 not in model.contact_geoms: + contact_set.add(g1) + return contact_set def _check_success(self): """ diff --git a/robosuite/environments/door.py b/robosuite/environments/manipulation/door.py similarity index 83% rename from robosuite/environments/door.py rename to robosuite/environments/manipulation/door.py index 585570f57f..13a158cdae 100644 --- a/robosuite/environments/door.py +++ b/robosuite/environments/manipulation/door.py @@ -1,15 +1,14 @@ from collections import OrderedDict import numpy as np -from robosuite.environments.robot_env import RobotEnv -from robosuite.robots import SingleArm +from robosuite.environments.manipulation.single_arm_env import SingleArmEnv from robosuite.models.arenas import TableArena from robosuite.models.objects import DoorObject from robosuite.models.tasks import ManipulationTask, UniformRandomSampler -class Door(RobotEnv): +class Door(SingleArmEnv): """ This class corresponds to the door opening task for a single robot arm. @@ -18,6 +17,9 @@ class Door(RobotEnv): (e.g: "Sawyer" would generate one arm; ["Panda", "Panda", "Sawyer"] would generate three robot arms) Note: Must be a single single-arm robot! + env_configuration (str): Specifies how to position the robots within the environment (default is "default"). + For most single arm environments, this argument has no impact on the robot setup. + controller_configs (str or list of dict): If set, contains relevant controller parameters for creating a custom controller. Else, uses the default controller for this specific task. Should either be single dict if same controller is to be used for all robots or else it should be a list of the same length as @@ -61,13 +63,20 @@ class Door(RobotEnv): reward_shaping (bool): if True, use dense rewards. - placement_initializer (ObjectPositionSampler instance): if provided, will + placement_initializer (ObjectPositionSampler): if provided, will be used to place objects on every reset, else a UniformRandomSampler is used by default. use_indicator_object (bool): if True, sets up an indicator object that is useful for debugging. + robot_visualizations (bool or list of bool): True if using robot visualization. + Useful for teleoperation. Should either be single bool if robot visualization is to be used for all + robots or else it should be a list of the same length as "robots" param + + env_visualization (bool): True if visualizing sites for the arena / objects in this environment. Useful for + teleoperation. + has_renderer (bool): If true, render the simulation state in a viewer instead of headless mode. @@ -119,6 +128,7 @@ class Door(RobotEnv): def __init__( self, robots, + env_configuration="default", controller_configs=None, gripper_types="default", gripper_visualizations=False, @@ -130,6 +140,8 @@ def __init__( reward_shaping=False, placement_initializer=None, use_indicator_object=False, + robot_visualizations=False, + env_visualization=False, has_renderer=False, has_offscreen_renderer=True, render_camera="frontview", @@ -144,9 +156,6 @@ def __init__( camera_widths=256, camera_depths=False, ): - # First, verify that only one robot is being inputted - self._check_robot_configuration(robots) - # settings for table top (hardcoded since it's not an essential part of the environment) self.table_full_size = (0.8, 0.3, 0.05) self.table_offset = (-0.2, -0.35, 0.8) @@ -166,19 +175,24 @@ def __init__( self.placement_initializer = UniformRandomSampler( x_range=[0.07, 0.09], y_range=[-0.01, 0.01], - ensure_object_boundary_in_range=False, rotation=(-np.pi / 2. - 0.25, -np.pi / 2.), rotation_axis='z', + ensure_object_boundary_in_range=False, + ensure_valid_placement=True, + reference_pos=self.table_offset, ) super().__init__( robots=robots, + env_configuration=env_configuration, controller_configs=controller_configs, gripper_types=gripper_types, gripper_visualizations=gripper_visualizations, initialization_noise=initialization_noise, use_camera_obs=use_camera_obs, use_indicator_object=use_indicator_object, + robot_visualizations=robot_visualizations, + env_visualization=env_visualization, has_renderer=has_renderer, has_offscreen_renderer=has_offscreen_renderer, render_camera=render_camera, @@ -249,44 +263,38 @@ def _load_model(self): """ super()._load_model() - # Verify the correct robot has been loaded - assert isinstance(self.robots[0], SingleArm), \ - "Error: Expected one single-armed robot! Got {} type instead.".format(type(self.robots[0])) - # Adjust base pose accordingly xpos = self.robots[0].robot_model.base_xpos_offset["table"](self.table_full_size[0]) self.robots[0].robot_model.set_base_xpos(xpos) # load model for table top workspace - self.mujoco_arena = TableArena( + mujoco_arena = TableArena( table_full_size=self.table_full_size, table_offset=self.table_offset, ) if self.use_indicator_object: - self.mujoco_arena.add_pos_indicator() + mujoco_arena.add_pos_indicator() # Arena always gets set to zero origin - self.mujoco_arena.set_origin([0, 0, 0]) + mujoco_arena.set_origin([0, 0, 0]) # initialize objects of interest - door = DoorObject( + self.door = DoorObject( name="Door", friction=0.0, damping=0.1, lock=self.use_latch, - joints=[], # ensures that door object does not have a free joint ) - self.mujoco_objects = OrderedDict([("Door", door)]) - self.n_objects = len(self.mujoco_objects) + + # Add door to initializer + self.placement_initializer.add_objects(self.door) # task includes arena, robot, and objects of interest self.model = ManipulationTask( - mujoco_arena=self.mujoco_arena, + mujoco_arena=mujoco_arena, mujoco_robots=[robot.robot_model for robot in self.robots], - mujoco_objects=self.mujoco_objects, - initializer=self.placement_initializer, + mujoco_objects=self.door, ) - self.model.place_objects() def _get_reference(self): """ @@ -298,20 +306,13 @@ def _get_reference(self): # Additional object references from this env self.object_body_ids = dict() - self.object_body_ids["door"] = self.sim.model.body_name2id("door") - self.object_body_ids["frame"] = self.sim.model.body_name2id("frame") - self.object_body_ids["latch"] = self.sim.model.body_name2id("latch") - self.door_handle_site_id = self.sim.model.site_name2id("door_handle") - self.hinge_qpos_addr = self.sim.model.get_joint_qpos_addr("door_hinge") + self.object_body_ids["door"] = self.sim.model.body_name2id(self.door.door_body) + self.object_body_ids["frame"] = self.sim.model.body_name2id(self.door.frame_body) + self.object_body_ids["latch"] = self.sim.model.body_name2id(self.door.latch_body) + self.door_handle_site_id = self.sim.model.site_name2id(self.door.important_sites["handle"]) + self.hinge_qpos_addr = self.sim.model.get_joint_qpos_addr(self.door.joints[0]) if self.use_latch: - self.handle_qpos_addr = self.sim.model.get_joint_qpos_addr("latch_joint") - - self.l_finger_geom_ids = [ - self.sim.model.geom_name2id(x) for x in self.robots[0].gripper.important_geoms["left_finger"] - ] - self.r_finger_geom_ids = [ - self.sim.model.geom_name2id(x) for x in self.robots[0].gripper.important_geoms["right_finger"] - ] + self.handle_qpos_addr = self.sim.model.get_joint_qpos_addr(self.door.joints[1]) def _reset_internal(self): """ @@ -322,11 +323,14 @@ def _reset_internal(self): # Reset all object positions using initializer sampler if we're not directly loading from an xml if not self.deterministic_reset: - # Sample from the placement initializer for the Door object - door_pos, door_quat = self.model.place_objects() - door_body_id = self.sim.model.body_name2id("Door") - self.sim.model.body_pos[door_body_id] = door_pos[0] - self.sim.model.body_quat[door_body_id] = door_quat[0] + # Sample from the placement initializer for all objects + object_placements = self.placement_initializer.sample() + + # We know we're only setting a single object (the door), so specifically set its pose + door_pos, door_quat, _ = object_placements[self.door.name] + door_body_id = self.sim.model.body_name2id(self.door.root_body) + self.sim.model.body_pos[door_body_id] = door_pos + self.sim.model.body_quat[door_body_id] = door_quat def _get_observation(self): """ @@ -394,47 +398,17 @@ def _visualization(self): """ Do any needed visualization here. Overrides superclass implementations. """ + # Run super call first + super()._visualization() - # color the gripper site appropriately based on distance to door handle + # Color the gripper visualization site according to its distance to the cube if self.robots[0].gripper_visualization: - # get distance to door handle - dist = np.sum( - np.square( - self._handle_xpos - - self.sim.data.get_site_xpos(self.robots[0].gripper.visualization_sites["grip_site"]) - ) + self._visualize_gripper_to_target( + gripper=self.robots[0].gripper, + target=self.door.important_sites["handle"], + target_type="site" ) - # set RGBA for the EEF site here - max_dist = 0.1 - scaled = (1.0 - min(dist / max_dist, 1.)) ** 15 - rgba = np.zeros(4) - rgba[0] = 1 - scaled - rgba[1] = scaled - rgba[3] = 0.5 - - self.sim.model.site_rgba[self.robots[0].eef_site_id] = rgba - - def _check_robot_configuration(self, robots): - """ - Sanity check to make sure the inputted robots and configuration is acceptable - - Args: - robots (str or list of str): Robots to instantiate within this env - """ - if type(robots) is list: - assert len(robots) == 1, "Error: Only one robot should be inputted for this task!" - - @property - def _eef_xpos(self): - """ - Grabs End Effector position - - Returns: - np.array: End effector(x,y,z) - """ - return np.array(self.sim.data.site_xpos[self.robots[0].eef_site_id]) - @property def _handle_xpos(self): """ diff --git a/robosuite/environments/lift.py b/robosuite/environments/manipulation/lift.py similarity index 79% rename from robosuite/environments/lift.py rename to robosuite/environments/manipulation/lift.py index 62acec0516..1db37591b7 100644 --- a/robosuite/environments/lift.py +++ b/robosuite/environments/manipulation/lift.py @@ -4,15 +4,14 @@ from robosuite.utils.transform_utils import convert_quat from robosuite.utils.mjcf_utils import CustomMaterial -from robosuite.environments.robot_env import RobotEnv -from robosuite.robots import SingleArm +from robosuite.environments.manipulation.single_arm_env import SingleArmEnv from robosuite.models.arenas import TableArena from robosuite.models.objects import BoxObject from robosuite.models.tasks import ManipulationTask, UniformRandomSampler -class Lift(RobotEnv): +class Lift(SingleArmEnv): """ This class corresponds to the lifting task for a single robot arm. @@ -21,6 +20,9 @@ class Lift(RobotEnv): (e.g: "Sawyer" would generate one arm; ["Panda", "Panda", "Sawyer"] would generate three robot arms) Note: Must be a single single-arm robot! + env_configuration (str): Specifies how to position the robots within the environment (default is "default"). + For most single arm environments, this argument has no impact on the robot setup. + controller_configs (str or list of dict): If set, contains relevant controller parameters for creating a custom controller. Else, uses the default controller for this specific task. Should either be single dict if same controller is to be used for all robots or else it should be a list of the same length as @@ -66,13 +68,20 @@ class Lift(RobotEnv): reward_shaping (bool): if True, use dense rewards. - placement_initializer (ObjectPositionSampler instance): if provided, will + placement_initializer (ObjectPositionSampler): if provided, will be used to place objects on every reset, else a UniformRandomSampler is used by default. use_indicator_object (bool): if True, sets up an indicator object that is useful for debugging. + robot_visualizations (bool or list of bool): True if using robot visualization. + Useful for teleoperation. Should either be single bool if robot visualization is to be used for all + robots or else it should be a list of the same length as "robots" param + + env_visualization (bool): True if visualizing sites for the arena / objects in this environment. Useful for + teleoperation. + has_renderer (bool): If true, render the simulation state in a viewer instead of headless mode. @@ -124,6 +133,7 @@ class Lift(RobotEnv): def __init__( self, robots, + env_configuration="default", controller_configs=None, gripper_types="default", gripper_visualizations=False, @@ -136,6 +146,8 @@ def __init__( reward_shaping=False, placement_initializer=None, use_indicator_object=False, + robot_visualizations=False, + env_visualization=False, has_renderer=False, has_offscreen_renderer=True, render_camera="frontview", @@ -150,12 +162,10 @@ def __init__( camera_widths=256, camera_depths=False, ): - # First, verify that only one robot is being inputted - self._check_robot_configuration(robots) - # settings for table top self.table_full_size = table_full_size self.table_friction = table_friction + self.table_offset = np.array((0, 0, 0.8)) # reward configuration self.reward_scale = reward_scale @@ -171,19 +181,24 @@ def __init__( self.placement_initializer = UniformRandomSampler( x_range=[-0.03, 0.03], y_range=[-0.03, 0.03], - ensure_object_boundary_in_range=False, rotation=None, + ensure_object_boundary_in_range=False, + ensure_valid_placement=True, + reference_pos=self.table_offset, z_offset=0.01, ) super().__init__( robots=robots, + env_configuration=env_configuration, controller_configs=controller_configs, gripper_types=gripper_types, gripper_visualizations=gripper_visualizations, initialization_noise=initialization_noise, use_camera_obs=use_camera_obs, use_indicator_object=use_indicator_object, + robot_visualizations=robot_visualizations, + env_visualization=env_visualization, has_renderer=has_renderer, has_offscreen_renderer=has_offscreen_renderer, render_camera=render_camera, @@ -241,19 +256,7 @@ def reward(self, action=None): reward += reaching_reward # grasping reward - touch_left_finger = False - touch_right_finger = False - for i in range(self.sim.data.ncon): - c = self.sim.data.contact[i] - if c.geom1 in self.l_finger_geom_ids and c.geom2 == self.cube_geom_id: - touch_left_finger = True - if c.geom1 == self.cube_geom_id and c.geom2 in self.l_finger_geom_ids: - touch_left_finger = True - if c.geom1 in self.r_finger_geom_ids and c.geom2 == self.cube_geom_id: - touch_right_finger = True - if c.geom1 == self.cube_geom_id and c.geom2 in self.r_finger_geom_ids: - touch_right_finger = True - if touch_left_finger and touch_right_finger: + if self._check_grasp(gripper=self.robots[0].gripper, object_geoms=self.cube): reward += 0.25 # Scale reward if requested @@ -268,25 +271,21 @@ def _load_model(self): """ super()._load_model() - # Verify the correct robot has been loaded - assert isinstance(self.robots[0], SingleArm), \ - "Error: Expected one single-armed robot! Got {} type instead.".format(type(self.robots[0])) - # Adjust base pose accordingly xpos = self.robots[0].robot_model.base_xpos_offset["table"](self.table_full_size[0]) self.robots[0].robot_model.set_base_xpos(xpos) # load model for table top workspace - self.mujoco_arena = TableArena( + mujoco_arena = TableArena( table_full_size=self.table_full_size, table_friction=self.table_friction, - table_offset=(0, 0, 0.8), + table_offset=self.table_offset, ) if self.use_indicator_object: - self.mujoco_arena.add_pos_indicator() + mujoco_arena.add_pos_indicator() # Arena always gets set to zero origin - self.mujoco_arena.set_origin([0, 0, 0]) + mujoco_arena.set_origin([0, 0, 0]) # initialize objects of interest tex_attrib = { @@ -304,24 +303,23 @@ def _load_model(self): tex_attrib=tex_attrib, mat_attrib=mat_attrib, ) - cube = BoxObject( + self.cube = BoxObject( name="cube", size_min=[0.020, 0.020, 0.020], # [0.015, 0.015, 0.015], size_max=[0.022, 0.022, 0.022], # [0.018, 0.018, 0.018]) rgba=[1, 0, 0, 1], material=redwood, ) - self.mujoco_objects = OrderedDict([("cube", cube)]) - self.n_objects = len(self.mujoco_objects) + + # Add cube to initializer + self.placement_initializer.add_objects(self.cube) # task includes arena, robot, and objects of interest self.model = ManipulationTask( - mujoco_arena=self.mujoco_arena, + mujoco_arena=mujoco_arena, mujoco_robots=[robot.robot_model for robot in self.robots], - mujoco_objects=self.mujoco_objects, - initializer=self.placement_initializer, + mujoco_objects=self.cube, ) - self.model.place_objects() def _get_reference(self): """ @@ -332,14 +330,7 @@ def _get_reference(self): super()._get_reference() # Additional object references from this env - self.cube_body_id = self.sim.model.body_name2id("cube") - self.l_finger_geom_ids = [ - self.sim.model.geom_name2id(x) for x in self.robots[0].gripper.important_geoms["left_finger"] - ] - self.r_finger_geom_ids = [ - self.sim.model.geom_name2id(x) for x in self.robots[0].gripper.important_geoms["right_finger"] - ] - self.cube_geom_id = self.sim.model.geom_name2id("cube") + self.cube_body_id = self.sim.model.body_name2id(self.cube.root_body) def _reset_internal(self): """ @@ -351,11 +342,11 @@ def _reset_internal(self): if not self.deterministic_reset: # Sample from the placement initializer for all objects - obj_pos, obj_quat = self.model.place_objects() + object_placements = self.placement_initializer.sample() # Loop through all objects and reset their positions - for i, (obj_name, _) in enumerate(self.mujoco_objects.items()): - self.sim.data.set_joint_qpos(obj_name + "_jnt0", np.concatenate([np.array(obj_pos[i]), np.array(obj_quat[i])])) + for obj_pos, obj_quat, obj in object_placements.values(): + self.sim.data.set_joint_qpos(obj.joints[0], np.concatenate([np.array(obj_pos), np.array(obj_quat)])) def _get_observation(self): """ @@ -399,6 +390,17 @@ def _get_observation(self): return di + def _visualization(self): + """ + Visualize the gripper site proportional to the distance to the cube + """ + # Run super call first + super()._visualization() + + # Color the gripper visualization site according to its distance to the cube + if self.robots[0].gripper_visualization: + self._visualize_gripper_to_target(gripper=self.robots[0].gripper, target=self.cube) + def _check_success(self): """ Check if cube has been lifted. @@ -407,43 +409,7 @@ def _check_success(self): bool: True if cube has been lifted """ cube_height = self.sim.data.body_xpos[self.cube_body_id][2] - table_height = self.mujoco_arena.table_offset[2] + table_height = self.model.mujoco_arena.table_offset[2] # cube is higher than the table top above a margin return cube_height > table_height + 0.04 - - def _visualization(self): - """ - Do any needed visualization here. Overrides superclass implementations. - """ - - # color the gripper site appropriately based on distance to cube - if self.robots[0].gripper_visualization: - # get distance to cube - cube_site_id = self.sim.model.site_name2id("cube") - dist = np.sum( - np.square( - self.sim.data.site_xpos[cube_site_id] - - self.sim.data.get_site_xpos(self.robots[0].gripper.visualization_sites["grip_site"]) - ) - ) - - # set RGBA for the EEF site here - max_dist = 0.1 - scaled = (1.0 - min(dist / max_dist, 1.)) ** 15 - rgba = np.zeros(4) - rgba[0] = 1 - scaled - rgba[1] = scaled - rgba[3] = 0.5 - - self.sim.model.site_rgba[self.robots[0].eef_site_id] = rgba - - def _check_robot_configuration(self, robots): - """ - Sanity check to make sure the inputted robots and configuration is acceptable - - Args: - robots (str or list of str): Robots to instantiate within this env - """ - if type(robots) is list: - assert len(robots) == 1, "Error: Only one robot should be inputted for this task!" diff --git a/robosuite/environments/manipulation/manipulation_env.py b/robosuite/environments/manipulation/manipulation_env.py new file mode 100644 index 0000000000..6df3fbfe00 --- /dev/null +++ b/robosuite/environments/manipulation/manipulation_env.py @@ -0,0 +1,308 @@ +import numpy as np + +from robosuite.environments.robot_env import RobotEnv +from robosuite.models.grippers import GripperModel +from robosuite.models.base import MujocoModel +from robosuite.robots import Manipulator, ROBOT_CLASS_MAPPING + + +class ManipulationEnv(RobotEnv): + """ + Initializes a manipulation-specific robot environment in Mujoco. + + Args: + robots: Specification for specific robot arm(s) to be instantiated within this env + (e.g: "Sawyer" would generate one arm; ["Panda", "Panda", "Sawyer"] would generate three robot arms) + + env_configuration (str): Specifies how to position the robot(s) within the environment. Default is "default", + which should be interpreted accordingly by any subclasses. + + controller_configs (str or list of dict): If set, contains relevant controller parameters for creating a + custom controller. Else, uses the default controller for this specific task. Should either be single + dict if same controller is to be used for all robots or else it should be a list of the same length as + "robots" param + + gripper_types (None or str or list of str): type of gripper, used to instantiate + gripper models from gripper factory. Default is "default", which is the default grippers(s) associated + with the robot(s) the 'robots' specification. None removes the gripper, and any other (valid) model + overrides the default gripper. Should either be single str if same gripper type is to be used for all + robots or else it should be a list of the same length as "robots" param + + gripper_visualizations (bool or list of bool): True if using gripper visualization. + Useful for teleoperation. Should either be single bool if gripper visualization is to be used for all + robots or else it should be a list of the same length as "robots" param + + initialization_noise (dict or list of dict): Dict containing the initialization noise parameters. + The expected keys and corresponding value types are specified below: + + :`'magnitude'`: The scale factor of uni-variate random noise applied to each of a robot's given initial + joint positions. Setting this value to `None` or 0.0 results in no noise being applied. + If "gaussian" type of noise is applied then this magnitude scales the standard deviation applied, + If "uniform" type of noise is applied then this magnitude sets the bounds of the sampling range + :`'type'`: Type of noise to apply. Can either specify "gaussian" or "uniform" + + Should either be single dict if same noise value is to be used for all robots or else it should be a + list of the same length as "robots" param + + :Note: Specifying "default" will automatically use the default noise settings. + Specifying None will automatically create the required dict with "magnitude" set to 0.0. + + use_camera_obs (bool): if True, every observation includes rendered image(s) + + use_indicator_object (bool): if True, sets up an indicator object that + is useful for debugging. + + robot_visualizations (bool or list of bool): True if using robot visualization. + Useful for teleoperation. Should either be single bool if robot visualization is to be used for all + robots or else it should be a list of the same length as "robots" param + + env_visualization (bool): True if visualizing sites for the arena / objects in this environment. Useful for + teleoperation. + + has_renderer (bool): If true, render the simulation state in + a viewer instead of headless mode. + + has_offscreen_renderer (bool): True if using off-screen rendering + + render_camera (str): Name of camera to render if `has_renderer` is True. Setting this value to 'None' + will result in the default angle being applied, which is useful as it can be dragged / panned by + the user using the mouse + + render_collision_mesh (bool): True if rendering collision meshes in camera. False otherwise. + + render_visual_mesh (bool): True if rendering visual meshes in camera. False otherwise. + + control_freq (float): how many control signals to receive in every second. This sets the amount of + simulation time that passes between every action input. + + horizon (int): Every episode lasts for exactly @horizon timesteps. + + ignore_done (bool): True if never terminating the environment (ignore @horizon). + + hard_reset (bool): If True, re-loads model, sim, and render object upon a reset call, else, + only calls sim.reset and resets all robosuite-internal variables + + camera_names (str or list of str): name of camera to be rendered. Should either be single str if + same name is to be used for all cameras' rendering or else it should be a list of cameras to render. + + :Note: At least one camera must be specified if @use_camera_obs is True. + + :Note: To render all robots' cameras of a certain type (e.g.: "robotview" or "eye_in_hand"), use the + convention "all-{name}" (e.g.: "all-robotview") to automatically render all camera images from each + robot's camera list). + + camera_heights (int or list of int): height of camera frame. Should either be single int if + same height is to be used for all cameras' frames or else it should be a list of the same length as + "camera names" param. + + camera_widths (int or list of int): width of camera frame. Should either be single int if + same width is to be used for all cameras' frames or else it should be a list of the same length as + "camera names" param. + + camera_depths (bool or list of bool): True if rendering RGB-D, and RGB otherwise. Should either be single + bool if same depth setting is to be used for all cameras or else it should be a list of the same length as + "camera names" param. + + Raises: + ValueError: [Camera obs require offscreen renderer] + ValueError: [Camera name must be specified to use camera obs] + """ + + def __init__( + self, + robots, + env_configuration="default", + controller_configs=None, + gripper_types="default", + gripper_visualizations=False, + initialization_noise=None, + use_camera_obs=True, + use_indicator_object=False, + robot_visualizations=False, + env_visualization=False, + has_renderer=False, + has_offscreen_renderer=True, + render_camera="frontview", + render_collision_mesh=False, + render_visual_mesh=True, + control_freq=10, + horizon=1000, + ignore_done=False, + hard_reset=True, + camera_names="agentview", + camera_heights=256, + camera_widths=256, + camera_depths=False, + ): + # Robot info + robots = list(robots) if type(robots) is list or type(robots) is tuple else [robots] + num_robots = len(robots) + + # Gripper + gripper_types = self._input2list(gripper_types, num_robots) + gripper_visualizations = self._input2list(gripper_visualizations, num_robots) + + # Robot configurations to pass to super call + robot_configs = [ + { + "gripper_type": gripper_types[idx], + "gripper_visualization": gripper_visualizations[idx], + } + for idx in range(num_robots) + ] + + # Run superclass init + super().__init__( + robots=robots, + env_configuration=env_configuration, + controller_configs=controller_configs, + initialization_noise=initialization_noise, + use_camera_obs=use_camera_obs, + use_indicator_object=use_indicator_object, + robot_visualizations=robot_visualizations, + env_visualization=env_visualization, + has_renderer=has_renderer, + has_offscreen_renderer=has_offscreen_renderer, + render_camera=render_camera, + render_collision_mesh=render_collision_mesh, + render_visual_mesh=render_visual_mesh, + control_freq=control_freq, + horizon=horizon, + ignore_done=ignore_done, + hard_reset=hard_reset, + camera_names=camera_names, + camera_heights=camera_heights, + camera_widths=camera_widths, + camera_depths=camera_depths, + robot_configs=robot_configs, + ) + + def _visualization(self): + """ + Do any needed visualization here + """ + # Run superclass method first + super()._visualization() + # Loop over robot grippers to visualize them independently + for robot in self.robots: + robot.visualize_gripper() + + def _check_grasp(self, gripper, object_geoms): + """ + Checks whether the specified gripper as defined by @gripper is grasping the specified object in the environment. + + By default, this will return True if at least one geom in both the "left_fingerpad" and "right_fingerpad" geom + groups are in contact with any geom specified by @object_geoms. Custom gripper geom groups can be + specified with @gripper as well. + + Args: + gripper (GripperModel or str or list of str or list of list of str): If a MujocoModel, this is specific + gripper to check for grasping (as defined by "left_fingerpad" and "right_fingerpad" geom groups). Otherwise, + this sets custom gripper geom groups which together define a grasp. This can be a string + (one group of single gripper geom), a list of string (multiple groups of single gripper geoms) or a + list of list of string (multiple groups of multiple gripper geoms). At least one geom from each group + must be in contact with any geom in @object_geoms for this method to return True. + object_geoms (str or list of str or MujocoModel): If a MujocoModel is inputted, will check for any + collisions with the model's contact_geoms. Otherwise, this should be specific geom name(s) composing + the object to check for contact. + + Returns: + bool: True if the gripper is grasping the given object + """ + # Convert object, gripper geoms into standardized form + if isinstance(object_geoms, MujocoModel): + o_geoms = object_geoms.contact_geoms + else: + o_geoms = [object_geoms] if type(object_geoms) is str else object_geoms + if isinstance(gripper, GripperModel): + g_geoms = [gripper.important_geoms["left_fingerpad"], gripper.important_geoms["right_fingerpad"]] + elif type(gripper) is str: + g_geoms = [[gripper]] + else: + # Parse each element in the gripper_geoms list accordingly + g_geoms = [[g_group] if type(g_group) is str else g_group for g_group in gripper] + + # Search for collisions between each gripper geom group and the object geoms group + for g_group in g_geoms: + if not self.check_contact(g_group, o_geoms): + return False + return True + + def _gripper_to_target(self, gripper, target, target_type="body", return_distance=False): + """ + Calculates the (x,y,z) Cartesian distance (target_pos - gripper_pos) from the specified @gripper to the + specified @target. If @return_distance is set, will return the Euclidean (scalar) distance instead. + + Args: + gripper (MujocoModel): Gripper model to update grip site rgb + target (MujocoModel or str): Either a site / geom / body name, or a model that serves as the target. + If a model is given, then the root body will be used as the target. + target_type (str): One of {"body", "geom", or "site"}, corresponding to the type of element @target + refers to. + return_distance (bool): If set, will return Euclidean distance instead of Cartesian distance + + Returns: + np.array or float: (Cartesian or Euclidean) distance from gripper to target + """ + # Get gripper and target positions + gripper_pos = self.sim.data.get_site_xpos(gripper.important_sites["grip_site"]) + # If target is MujocoModel, grab the correct body as the target and find the target position + if isinstance(target, MujocoModel): + target_pos = self.sim.data.get_body_xpos(target.root_body) + elif target_type == "body": + target_pos = self.sim.data.get_body_xpos(target) + elif target_type == "site": + target_pos = self.sim.data.get_site_xpos(target) + else: + target_pos = self.sim.data.get_geom_xpos(target) + # Calculate distance + diff = target_pos - gripper_pos + # Return appropriate value + return np.linalg.norm(diff) if return_distance else diff + + def _visualize_gripper_to_target(self, gripper, target, target_type="body"): + """ + Colors the grip visualization site proportional to the Euclidean distance to the specified @target. + Colors go from red --> green as the gripper gets closer. + + Args: + gripper (MujocoModel): Gripper model to update grip site rgb + target (MujocoModel or str): Either a site / geom / body name, or a model that serves as the target. + If a model is given, then the root body will be used as the target. + target_type (str): One of {"body", "geom", or "site"}, corresponding to the type of element @target + refers to. + """ + # Get gripper and target positions + gripper_pos = self.sim.data.get_site_xpos(gripper.important_sites["grip_site"]) + # If target is MujocoModel, grab the correct body as the target and find the target position + if isinstance(target, MujocoModel): + target_pos = self.sim.data.get_body_xpos(target.root_body) + elif target_type == "body": + target_pos = self.sim.data.get_body_xpos(target) + elif target_type == "site": + target_pos = self.sim.data.get_site_xpos(target) + else: + target_pos = self.sim.data.get_geom_xpos(target) + # color the gripper site appropriately based on (squared) distance to target + dist = np.sum(np.square((target_pos - gripper_pos))) + max_dist = 0.1 + scaled = (1.0 - min(dist / max_dist, 1.)) ** 15 + rgba = np.zeros(3) + rgba[0] = 1 - scaled + rgba[1] = scaled + self.sim.model.site_rgba[self.sim.model.site_name2id(gripper.important_sites["grip_site"])][:3] = rgba + + def _check_robot_configuration(self, robots): + """ + Sanity check to make sure inputted robots and the corresponding requested task/configuration combo is legal. + Should be implemented in every specific task module + + Args: + robots (str or list of str): Inputted requested robots at the task-level environment + """ + # Make sure all inputted robots are a manipulation robot + if type(robots) is str: + robots = [robots] + for robot in robots: + assert issubclass(ROBOT_CLASS_MAPPING[robot], Manipulator),\ + "Only manipulator robots supported for manipulation environment!" diff --git a/robosuite/environments/nut_assembly.py b/robosuite/environments/manipulation/nut_assembly.py similarity index 67% rename from robosuite/environments/nut_assembly.py rename to robosuite/environments/manipulation/nut_assembly.py index 0c1c4aa3de..7cf620a606 100644 --- a/robosuite/environments/nut_assembly.py +++ b/robosuite/environments/manipulation/nut_assembly.py @@ -3,15 +3,14 @@ import numpy as np import robosuite.utils.transform_utils as T -from robosuite.environments.robot_env import RobotEnv -from robosuite.robots import SingleArm +from robosuite.environments.manipulation.single_arm_env import SingleArmEnv from robosuite.models.arenas import PegsArena from robosuite.models.objects import SquareNutObject, RoundNutObject -from robosuite.models.tasks import ManipulationTask, SequentialCompositeSampler +from robosuite.models.tasks import ManipulationTask, SequentialCompositeSampler, UniformRandomSampler -class NutAssembly(RobotEnv): +class NutAssembly(SingleArmEnv): """ This class corresponds to the nut assembly task for a single robot arm. @@ -20,6 +19,9 @@ class NutAssembly(RobotEnv): (e.g: "Sawyer" would generate one arm; ["Panda", "Panda", "Sawyer"] would generate three robot arms) Note: Must be a single single-arm robot! + env_configuration (str): Specifies how to position the robots within the environment (default is "default"). + For most single arm environments, this argument has no impact on the robot setup. + controller_configs (str or list of dict): If set, contains relevant controller parameters for creating a custom controller. Else, uses the default controller for this specific task. Should either be single dict if same controller is to be used for all robots or else it should be a list of the same length as @@ -65,7 +67,7 @@ class NutAssembly(RobotEnv): reward_shaping (bool): if True, use dense rewards. - placement_initializer (ObjectPositionSampler instance): if provided, will + placement_initializer (ObjectPositionSampler): if provided, will be used to place objects on every reset, else a UniformRandomSampler is used by default. @@ -88,6 +90,13 @@ class NutAssembly(RobotEnv): use_indicator_object (bool): if True, sets up an indicator object that is useful for debugging. + robot_visualizations (bool or list of bool): True if using robot visualization. + Useful for teleoperation. Should either be single bool if robot visualization is to be used for all + robots or else it should be a list of the same length as "robots" param + + env_visualization (bool): True if visualizing sites for the arena / objects in this environment. Useful for + teleoperation. + has_renderer (bool): If true, render the simulation state in a viewer instead of headless mode. @@ -140,6 +149,7 @@ class NutAssembly(RobotEnv): def __init__( self, robots, + env_configuration="default", controller_configs=None, gripper_types="default", gripper_visualizations=False, @@ -154,6 +164,8 @@ def __init__( single_object_mode=0, nut_type=None, use_indicator_object=False, + robot_visualizations=False, + env_visualization=False, has_renderer=False, has_offscreen_renderer=True, render_camera="frontview", @@ -168,9 +180,6 @@ def __init__( camera_widths=256, camera_depths=False, ): - # First, verify that only one robot is being inputted - self._check_robot_configuration(robots) - # task settings self.single_object_mode = single_object_mode self.nut_to_id = {"square": 0, "round": 1} @@ -186,6 +195,7 @@ def __init__( # settings for table top self.table_full_size = table_full_size self.table_friction = table_friction + self.table_offset = np.array((0, 0, 0.82)) # reward configuration self.reward_scale = reward_scale @@ -198,38 +208,20 @@ def __init__( if placement_initializer: self.placement_initializer = placement_initializer else: - # treat sampling of each type of nut differently since we require different - # sampling ranges for each + # treat sampling of each type of nut differently since we require different sampling ranges for each self.placement_initializer = SequentialCompositeSampler() - self.placement_initializer.sample_on_top( - "SquareNut0", - surface_name="table", - x_range=[-0.115, -0.11], - y_range=[0.11, 0.225], - rotation=None, - rotation_axis='z', - z_offset=0.02, - ensure_object_boundary_in_range=False, - ) - self.placement_initializer.sample_on_top( - "RoundNut0", - surface_name="table", - x_range=[-0.115, -0.11], - y_range=[-0.225, -0.11], - rotation=None, - rotation_axis='z', - z_offset=0.02, - ensure_object_boundary_in_range=False, - ) super().__init__( robots=robots, + env_configuration=env_configuration, controller_configs=controller_configs, gripper_types=gripper_types, gripper_visualizations=gripper_visualizations, initialization_noise=initialization_noise, use_camera_obs=use_camera_obs, use_indicator_object=use_indicator_object, + robot_visualizations=robot_visualizations, + env_visualization=env_visualization, has_renderer=has_renderer, has_offscreen_renderer=has_offscreen_renderer, render_camera=render_camera, @@ -306,75 +298,60 @@ def staged_rewards(self): hover_mult = 0.7 # filter out objects that are already on the correct pegs - names_to_reach = [] - objs_to_reach = [] - geoms_to_grasp = [] - geoms_by_array = [] - - for i in range(len(self.ob_inits)): + active_nuts = [] + for i, nut in enumerate(self.nuts): if self.objects_on_pegs[i]: continue - obj_str = str(self.item_names[i]) + "0" - names_to_reach.append(obj_str) - objs_to_reach.append(self.obj_body_id[obj_str]) - geoms_to_grasp.extend(self.obj_geom_id[obj_str]) - geoms_by_array.append(self.obj_geom_id[obj_str]) + active_nuts.append(nut) - ### reaching reward governed by distance to closest object ### + # reaching reward governed by distance to closest object r_reach = 0. - if len(objs_to_reach): - # reaching reward via minimum distance to the handles of the objects (the last geom of each nut) - geom_ids = [elem[-1] for elem in geoms_by_array] - target_geom_pos = self.sim.data.geom_xpos[geom_ids] - gripper_site_pos = self.sim.data.site_xpos[self.robots[0].eef_site_id] - dists = np.linalg.norm( - target_geom_pos - gripper_site_pos.reshape(1, -1), axis=1 - ) + if active_nuts: + # reaching reward via minimum distance to the handles of the objects + dists = [ + self._gripper_to_target( + gripper=self.robots[0].gripper, + target=active_nut.important_sites["handle"], + target_type="site", + return_distance=True, + ) for active_nut in active_nuts + ] r_reach = (1 - np.tanh(10.0 * min(dists))) * reach_mult - ### grasping reward for touching any objects of interest ### - touch_left_finger = False - touch_right_finger = False - for i in range(self.sim.data.ncon): - c = self.sim.data.contact[i] - if c.geom1 in geoms_to_grasp: - if c.geom2 in self.l_finger_geom_ids: - touch_left_finger = True - if c.geom2 in self.r_finger_geom_ids: - touch_right_finger = True - elif c.geom2 in geoms_to_grasp: - if c.geom1 in self.l_finger_geom_ids: - touch_left_finger = True - if c.geom1 in self.r_finger_geom_ids: - touch_right_finger = True - has_grasp = touch_left_finger and touch_right_finger - r_grasp = int(has_grasp) * grasp_mult - - ### lifting reward for picking up an object ### + # grasping reward for touching any objects of interest + r_grasp = int(self._check_grasp( + gripper=self.robots[0].gripper, + object_geoms=[g for active_nut in active_nuts for g in active_nut.contact_geoms]) + ) * grasp_mult + + # lifting reward for picking up an object r_lift = 0. table_pos = np.array(self.sim.data.body_xpos[self.table_body_id]) - if len(objs_to_reach) and r_grasp > 0.: + if active_nuts and r_grasp > 0.: z_target = table_pos[2] + 0.2 - object_z_locs = self.sim.data.body_xpos[objs_to_reach][:, 2] + object_z_locs = self.sim.data.body_xpos[[self.obj_body_id[active_nut.name] + for active_nut in active_nuts]][:, 2] z_dists = np.maximum(z_target - object_z_locs, 0.) r_lift = grasp_mult + (1 - np.tanh(15.0 * min(z_dists))) * ( lift_mult - grasp_mult ) - ### hover reward for getting object above peg ### + # hover reward for getting object above peg r_hover = 0. - if len(objs_to_reach): - r_hovers = np.zeros(len(objs_to_reach)) - for i in range(len(objs_to_reach)): - if names_to_reach[i].startswith(self.item_names[0]): - peg_pos = np.array(self.sim.data.body_xpos[self.peg1_body_id])[:2] - elif names_to_reach[i].startswith(self.item_names[1]): - peg_pos = np.array(self.sim.data.body_xpos[self.peg2_body_id])[:2] - else: - raise Exception( - "Got invalid object to reach: {}".format(names_to_reach[i]) - ) - ob_xy = self.sim.data.body_xpos[objs_to_reach[i]][:2] + if active_nuts: + r_hovers = np.zeros(len(active_nuts)) + peg_body_ids = [self.peg1_body_id, self.peg2_body_id] + for i, nut in enumerate(active_nuts): + valid_obj = False + peg_pos = None + for nut_name, idn in self.nut_to_id.items(): + if nut_name in nut.name.lower(): + peg_pos = np.array(self.sim.data.body_xpos[peg_body_ids[idn]])[:2] + valid_obj = True + break + if not valid_obj: + raise Exception("Got invalid object to reach: {}".format(nut.name)) + ob_xy = self.sim.data.body_xpos[self.obj_body_id[nut.name]][:2] dist = np.linalg.norm(peg_pos - ob_xy) r_hovers[i] = r_lift + (1 - np.tanh(10.0 * dist)) * ( hover_mult - lift_mult @@ -393,86 +370,79 @@ def on_peg(self, obj_pos, peg_id): if ( abs(obj_pos[0] - peg_pos[0]) < 0.03 and abs(obj_pos[1] - peg_pos[1]) < 0.03 - and obj_pos[2] < self.mujoco_arena.table_offset[2] + 0.05 + and obj_pos[2] < self.table_offset[2] + 0.05 ): res = True return res - def clear_objects(self, obj): - """ - Clears objects without the name @obj out of the task space. This is useful - for supporting task modes with single types of objects, as in - @self.single_object_mode without changing the model definition. - - Args: - obj (str): Name of object to keep in the task space - """ - for obj_name, obj_mjcf in self.mujoco_objects.items(): - if obj_name == obj: - continue - else: - sim_state = self.sim.get_state() - # print(self.sim.model.get_joint_qpos_addr(obj_name)) - sim_state.qpos[self.sim.model.get_joint_qpos_addr(obj_name + "_jnt0")[0]] = 10 - self.sim.set_state(sim_state) - self.sim.forward() - def _load_model(self): """ Loads an xml model, puts it in self.model """ super()._load_model() - # Verify the correct robot has been loaded - assert isinstance(self.robots[0], SingleArm), \ - "Error: Expected one single-armed robot! Got {} type instead.".format(type(self.robots[0])) - # Adjust base pose accordingly xpos = self.robots[0].robot_model.base_xpos_offset["table"](self.table_full_size[0]) self.robots[0].robot_model.set_base_xpos(xpos) # load model for table top workspace - self.mujoco_arena = PegsArena( + mujoco_arena = PegsArena( table_full_size=self.table_full_size, table_friction=self.table_friction, - table_offset=(0, 0, 0.82) + table_offset=self.table_offset, ) if self.use_indicator_object: - self.mujoco_arena.add_pos_indicator() + mujoco_arena.add_pos_indicator() # Arena always gets set to zero origin - self.mujoco_arena.set_origin([0, 0, 0]) - - # define mujoco objects - self.ob_inits = [SquareNutObject, RoundNutObject] - self.item_names = ["SquareNut", "RoundNut"] - self.item_names_org = list(self.item_names) - self.obj_to_use = (self.item_names[1] + "{}").format(0) - self.ngeoms = [5, 9] - - lst = [] - for i in range(len(self.ob_inits)): - ob_name = (self.item_names[i] + "0") - ob = self.ob_inits[i]( - name=ob_name, - joints=[dict(type="free", damping="0.0005")], # damp the free joint for each object - ) - lst.append((ob_name, ob)) - - self.mujoco_objects = OrderedDict(lst) - self.n_objects = len(self.mujoco_objects) + mujoco_arena.set_origin([0, 0, 0]) + + # define nuts + self.nuts = [] + + # We need to check if custom placement initializer has been specified. This is true if the sampler is either + # NOT a SequentialCompsiteSampler OR if the sampler already has sub-samplers in it + is_custom_sampler = not isinstance(self.placement_initializer, SequentialCompositeSampler) or \ + len(self.placement_initializer.samplers) > 0 + + for i, (nut_cls, nut_name, default_y_range) in enumerate(zip( + (SquareNutObject, RoundNutObject), + ("SquareNut", "RoundNut"), + ([0.11, 0.225], [-0.225, -0.11]), + )): + nut = nut_cls(name=nut_name) + self.nuts.append(nut) + # Add this nut to the placement initializer + if is_custom_sampler: + if isinstance(self.placement_initializer, SequentialCompositeSampler): + # assumes user has already specified two samplers so we add nuts to them + self.placement_initializer.samplers[i].add_objects(nut) + else: + # This is assumed to be a flat sampler, so we just add all nuts to this sampler + self.placement_initializer.add_objects(nut) + else: + # Default sampler; must define sub-samplers + self.placement_initializer.append_sampler( + sampler=UniformRandomSampler( + mujoco_objects=nut, + x_range=[-0.115, -0.11], + y_range=default_y_range, + rotation=None, + rotation_axis='z', + ensure_object_boundary_in_range=False, + ensure_valid_placement=True, + reference_pos=self.table_offset, + z_offset=0.02, + ) + ) # task includes arena, robot, and objects of interest self.model = ManipulationTask( - mujoco_arena=self.mujoco_arena, + mujoco_arena=mujoco_arena, mujoco_robots=[robot.robot_model for robot in self.robots], - mujoco_objects=self.mujoco_objects, - initializer=self.placement_initializer, + mujoco_objects=self.nuts, ) - # set positions of objects - self.model.place_objects() - def _get_reference(self): """ Sets up references to important components. A reference is typically an @@ -489,35 +459,15 @@ def _get_reference(self): self.peg1_body_id = self.sim.model.body_name2id("peg1") self.peg2_body_id = self.sim.model.body_name2id("peg2") - for i in range(len(self.ob_inits)): - obj_str = str(self.item_names[i]) + "0" - self.obj_body_id[obj_str] = self.sim.model.body_name2id(obj_str) - geom_ids = [] - for j in range(self.ngeoms[i]): - geom_ids.append(self.sim.model.geom_name2id(obj_str + "_{}".format(j))) - self.obj_geom_id[obj_str] = geom_ids + for nut in self.nuts: + self.obj_body_id[nut.name] = self.sim.model.body_name2id(nut.root_body) + self.obj_geom_id[nut.name] = [self.sim.model.geom_name2id(g) for g in nut.contact_geoms] # information of objects - self.object_names = list(self.mujoco_objects.keys()) - self.object_site_ids = [ - self.sim.model.site_name2id(ob_name) for ob_name in self.object_names - ] - - # id of grippers for contact checking - self.l_finger_geom_ids = [ - self.sim.model.geom_name2id(x) for x in self.robots[0].gripper.important_geoms["left_finger"] - ] - self.r_finger_geom_ids = [ - self.sim.model.geom_name2id(x) for x in self.robots[0].gripper.important_geoms["right_finger"] - ] - # self.sim.data.contact # list, geom1, geom2 - self.collision_check_geom_names = self.sim.model._geom_name2id.keys() - self.collision_check_geom_ids = [ - self.sim.model._geom_name2id[k] for k in self.collision_check_geom_names - ] + self.object_site_ids = [self.sim.model.site_name2id(nut.important_sites["handle"]) for nut in self.nuts] # keep track of which objects are on their corresponding pegs - self.objects_on_pegs = np.zeros(len(self.ob_inits)) + self.objects_on_pegs = np.zeros(len(self.nuts)) def _reset_internal(self): """ @@ -529,25 +479,21 @@ def _reset_internal(self): if not self.deterministic_reset: # Sample from the placement initializer for all objects - obj_pos, obj_quat = self.model.place_objects() + object_placements = self.placement_initializer.sample() # Loop through all objects and reset their positions - for i, (obj_name, _) in enumerate(self.mujoco_objects.items()): - self.sim.data.set_joint_qpos(obj_name + "_jnt0", np.concatenate([np.array(obj_pos[i]), np.array(obj_quat[i])])) - - # information of objects - self.object_names = list(self.mujoco_objects.keys()) - self.object_site_ids = [ - self.sim.model.site_name2id(ob_name) for ob_name in self.object_names - ] + for obj_pos, obj_quat, obj in object_placements.values(): + self.sim.data.set_joint_qpos(obj.joints[0], np.concatenate([np.array(obj_pos), np.array(obj_quat)])) # Move objects out of the scene depending on the mode + nut_names = {nut.name for nut in self.nuts} if self.single_object_mode == 1: - self.obj_to_use = (random.choice(self.item_names) + "{}").format(0) - self.clear_objects(self.obj_to_use) + self.obj_to_use = random.choice(list(nut_names)) elif self.single_object_mode == 2: - self.obj_to_use = (self.item_names[self.nut_id] + "{}").format(0) - self.clear_objects(self.obj_to_use) + self.obj_to_use = self.nuts[self.nut_id].name + if self.single_object_mode in {1, 2}: + nut_names.remove(self.obj_to_use) + self.clear_objects(list(nut_names)) def _get_observation(self): """ @@ -581,13 +527,13 @@ def _get_observation(self): gripper_pose = T.pose2mat((di[pr + "eef_pos"], di[pr + "eef_quat"])) world_pose_in_gripper = T.pose_inv(gripper_pose) - for i in range(len(self.item_names_org)): + for i, nut in enumerate(self.nuts): if self.single_object_mode == 2 and self.nut_id != i: # skip observations continue - obj_str = str(self.item_names_org[i]) + "0" + obj_str = nut.name obj_pos = np.array(self.sim.data.body_xpos[self.obj_body_id[obj_str]]) obj_quat = T.convert_quat( self.sim.data.body_xquat[self.obj_body_id[obj_str]], to="xyzw" @@ -608,14 +554,14 @@ def _get_observation(self): if self.single_object_mode == 1: # zero out other objs - for obj_str, obj_mjcf in self.mujoco_objects.items(): - if obj_str == self.obj_to_use: + for obj in self.model.mujoco_objects: + if obj.name == self.obj_to_use: continue else: - di["{}_pos".format(obj_str)] *= 0.0 - di["{}_quat".format(obj_str)] *= 0.0 - di["{}_to_{}eef_pos".format(obj_str, pr)] *= 0.0 - di["{}_to_{}eef_quat".format(obj_str, pr)] *= 0.0 + di["{}_pos".format(obj.name)] *= 0.0 + di["{}_quat".format(obj.name)] *= 0.0 + di["{}_to_{}eef_pos".format(obj.name, pr)] *= 0.0 + di["{}_to_{}eef_quat".format(obj.name, pr)] *= 0.0 di["object-state"] = np.concatenate([di[k] for k in object_state_keys]) @@ -630,8 +576,8 @@ def _check_success(self): """ # remember objects that are on the correct pegs gripper_site_pos = self.sim.data.site_xpos[self.robots[0].eef_site_id] - for i in range(len(self.ob_inits)): - obj_str = str(self.item_names[i]) + "0" + for i, nut in enumerate(self.nuts): + obj_str = nut.name obj_pos = self.sim.data.body_xpos[self.obj_body_id[obj_str]] dist = np.linalg.norm(gripper_site_pos - obj_pos) r_reach = 1 - np.tanh(10.0 * dist) @@ -641,47 +587,30 @@ def _check_success(self): return np.sum(self.objects_on_pegs) > 0 # need one object on peg # returns True if all objects are on correct pegs - return np.sum(self.objects_on_pegs) == len(self.ob_inits) + return np.sum(self.objects_on_pegs) == len(self.nuts) def _visualization(self): """ Do any needed visualization here. Overrides superclass implementations. """ - - # color the gripper site appropriately based on distance to cube + # color the gripper site appropriately based on distance to closest nut if self.robots[0].gripper_visualization: # find closest object - square_dist = lambda x: np.sum( - np.square(x - self.sim.data.get_site_xpos(self.robots[0].gripper.visualization_sites["grip_site"])) + dists = [ + self._gripper_to_target( + gripper=self.robots[0].gripper, + target=nut.important_sites["handle"], + target_type="site", + return_distance=True, + ) for nut in self.nuts + ] + closest_nut_id = np.argmin(dists) + # Visualize the distance to this target + self._visualize_gripper_to_target( + gripper=self.robots[0].gripper, + target=self.nuts[closest_nut_id].important_sites["handle"], + target_type="site", ) - dists = np.array(list(map(square_dist, self.sim.data.site_xpos))) - dists[self.robots[0].eef_site_id] = np.inf # make sure we don't pick the same site - dists[self.robots[0].eef_cylinder_id] = np.inf - ob_dists = dists[ - self.object_site_ids - ] # filter out object sites we care about - min_dist = np.min(ob_dists) - ob_id = np.argmin(ob_dists) - - # set RGBA for the EEF site here - max_dist = 0.1 - scaled = (1.0 - min(min_dist / max_dist, 1.)) ** 15 - rgba = np.zeros(4) - rgba[0] = 1 - scaled - rgba[1] = scaled - rgba[3] = 0.5 - - self.sim.model.site_rgba[self.robots[0].eef_site_id] = rgba - - def _check_robot_configuration(self, robots): - """ - Sanity check to make sure the inputted robots and configuration is acceptable - - Args: - robots (str or list of str): Robots to instantiate within this env - """ - if type(robots) is list: - assert len(robots) == 1, "Error: Only one robot should be inputted for this task!" class NutAssemblySingle(NutAssembly): diff --git a/robosuite/environments/pick_place.py b/robosuite/environments/manipulation/pick_place.py similarity index 69% rename from robosuite/environments/pick_place.py rename to robosuite/environments/manipulation/pick_place.py index 163b497cf8..519b52607b 100644 --- a/robosuite/environments/pick_place.py +++ b/robosuite/environments/manipulation/pick_place.py @@ -3,8 +3,7 @@ import numpy as np import robosuite.utils.transform_utils as T -from robosuite.environments.robot_env import RobotEnv -from robosuite.robots import SingleArm +from robosuite.environments.manipulation.single_arm_env import SingleArmEnv from robosuite.models.arenas import BinsArena from robosuite.models.objects import ( @@ -19,10 +18,10 @@ CerealVisualObject, CanVisualObject, ) -from robosuite.models.tasks import ManipulationTask, SequentialCompositeSampler +from robosuite.models.tasks import ManipulationTask, SequentialCompositeSampler, UniformRandomSampler -class PickPlace(RobotEnv): +class PickPlace(SingleArmEnv): """ This class corresponds to the pick place task for a single robot arm. @@ -31,6 +30,9 @@ class PickPlace(RobotEnv): (e.g: "Sawyer" would generate one arm; ["Panda", "Panda", "Sawyer"] would generate three robot arms) Note: Must be a single single-arm robot! + env_configuration (str): Specifies how to position the robots within the environment (default is "default"). + For most single arm environments, this argument has no impact on the robot setup. + controller_configs (str or list of dict): If set, contains relevant controller parameters for creating a custom controller. Else, uses the default controller for this specific task. Should either be single dict if same controller is to be used for all robots or else it should be a list of the same length as @@ -99,6 +101,13 @@ class PickPlace(RobotEnv): use_indicator_object (bool): if True, sets up an indicator object that is useful for debugging. + robot_visualizations (bool or list of bool): True if using robot visualization. + Useful for teleoperation. Should either be single bool if robot visualization is to be used for all + robots or else it should be a list of the same length as "robots" param + + env_visualization (bool): True if visualizing sites for the arena / objects in this environment. Useful for + teleoperation. + has_renderer (bool): If true, render the simulation state in a viewer instead of headless mode. @@ -151,6 +160,7 @@ class PickPlace(RobotEnv): def __init__( self, robots, + env_configuration="default", controller_configs=None, gripper_types="default", gripper_visualizations=False, @@ -166,6 +176,8 @@ def __init__( single_object_mode=0, object_type=None, use_indicator_object=False, + robot_visualizations=False, + env_visualization=False, has_renderer=False, has_offscreen_renderer=True, render_camera="frontview", @@ -180,12 +192,10 @@ def __init__( camera_widths=256, camera_depths=False, ): - # First, verify that only one robot is being inputted - self._check_robot_configuration(robots) - # task settings self.single_object_mode = single_object_mode self.object_to_id = {"milk": 0, "bread": 1, "cereal": 2, "can": 3} + self.obj_names = ["Milk", "Bread", "Cereal", "Can"] if object_type is not None: assert ( object_type in self.object_to_id.keys() @@ -214,12 +224,15 @@ def __init__( super().__init__( robots=robots, + env_configuration=env_configuration, controller_configs=controller_configs, gripper_types=gripper_types, gripper_visualizations=gripper_visualizations, initialization_noise=initialization_noise, use_camera_obs=use_camera_obs, use_indicator_object=use_indicator_object, + robot_visualizations=robot_visualizations, + env_visualization=env_visualization, has_renderer=has_renderer, has_offscreen_renderer=has_offscreen_renderer, render_camera=render_camera, @@ -296,79 +309,66 @@ def staged_rewards(self): hover_mult = 0.7 # filter out objects that are already in the correct bins - objs_to_reach = [] - geoms_to_grasp = [] - target_bin_placements = [] - for i in range(len(self.ob_inits)): + active_objs = [] + for i, obj in enumerate(self.objects): if self.objects_in_bins[i]: continue - obj_str = str(self.item_names[i]) + "0" - objs_to_reach.append(self.obj_body_id[obj_str]) - geoms_to_grasp.append(self.obj_geom_id[obj_str]) - target_bin_placements.append(self.target_bin_placements[i]) - target_bin_placements = np.array(target_bin_placements) + active_objs.append(obj) - ### reaching reward governed by distance to closest object ### + # reaching reward governed by distance to closest object r_reach = 0. - if len(objs_to_reach): + if active_objs: # get reaching reward via minimum distance to a target object - target_object_pos = self.sim.data.body_xpos[objs_to_reach] - gripper_site_pos = self.sim.data.site_xpos[self.robots[0].eef_site_id] - dists = np.linalg.norm( - target_object_pos - gripper_site_pos.reshape(1, -1), axis=1 - ) + dists = [ + self._gripper_to_target( + gripper=self.robots[0].gripper, + target=active_obj.root_body, + target_type="body", + return_distance=True, + ) for active_obj in active_objs + ] r_reach = (1 - np.tanh(10.0 * min(dists))) * reach_mult - ### grasping reward for touching any objects of interest ### - touch_left_finger = False - touch_right_finger = False - for i in range(self.sim.data.ncon): - c = self.sim.data.contact[i] - if c.geom1 in geoms_to_grasp: - bin_id = geoms_to_grasp.index(c.geom1) - if c.geom2 in self.l_finger_geom_ids: - touch_left_finger = True - if c.geom2 in self.r_finger_geom_ids: - touch_right_finger = True - elif c.geom2 in geoms_to_grasp: - bin_id = geoms_to_grasp.index(c.geom2) - if c.geom1 in self.l_finger_geom_ids: - touch_left_finger = True - if c.geom1 in self.r_finger_geom_ids: - touch_right_finger = True - has_grasp = touch_left_finger and touch_right_finger - r_grasp = int(has_grasp) * grasp_mult - - ### lifting reward for picking up an object ### + # grasping reward for touching any objects of interest + r_grasp = int(self._check_grasp( + gripper=self.robots[0].gripper, + object_geoms=[g for active_obj in active_objs for g in active_obj.contact_geoms]) + ) * grasp_mult + + # lifting reward for picking up an object r_lift = 0. - if len(objs_to_reach) and r_grasp > 0.: + if active_objs and r_grasp > 0.: z_target = self.bin2_pos[2] + 0.25 - object_z_locs = self.sim.data.body_xpos[objs_to_reach][:, 2] + object_z_locs = self.sim.data.body_xpos[[self.obj_body_id[active_obj.name] + for active_obj in active_objs]][:, 2] z_dists = np.maximum(z_target - object_z_locs, 0.) r_lift = grasp_mult + (1 - np.tanh(15.0 * min(z_dists))) * ( lift_mult - grasp_mult ) - ### hover reward for getting object above bin ### + # hover reward for getting object above bin r_hover = 0. - if len(objs_to_reach): + if active_objs: + target_bin_ids = [self.object_to_id[active_obj.name.lower()] for active_obj in active_objs] # segment objects into left of the bins and above the bins - object_xy_locs = self.sim.data.body_xpos[objs_to_reach][:, :2] + object_xy_locs = self.sim.data.body_xpos[[self.obj_body_id[active_obj.name] + for active_obj in active_objs]][:, :2] y_check = ( - np.abs(object_xy_locs[:, 1] - target_bin_placements[:, 1]) + np.abs(object_xy_locs[:, 1] - self.target_bin_placements[target_bin_ids, 1]) < self.bin_size[1] / 4. ) x_check = ( - np.abs(object_xy_locs[:, 0] - target_bin_placements[:, 0]) + np.abs(object_xy_locs[:, 0] - self.target_bin_placements[target_bin_ids, 0]) < self.bin_size[0] / 4. ) objects_above_bins = np.logical_and(x_check, y_check) objects_not_above_bins = np.logical_not(objects_above_bins) dists = np.linalg.norm( - target_bin_placements[:, :2] - object_xy_locs, axis=1 + self.target_bin_placements[target_bin_ids, :2] - object_xy_locs, axis=1 ) - # objects to the left get r_lift added to hover reward, those on the right get max(r_lift) added (to encourage dropping) - r_hover_all = np.zeros(len(objs_to_reach)) + # objects to the left get r_lift added to hover reward, + # those on the right get max(r_lift) added (to encourage dropping) + r_hover_all = np.zeros(len(active_objs)) r_hover_all[objects_above_bins] = lift_mult + ( 1 - np.tanh(10.0 * dists[objects_above_bins]) ) * (hover_mult - lift_mult) @@ -400,25 +400,6 @@ def not_in_bin(self, obj_pos, bin_id): res = False return res - def clear_objects(self, obj): - """ - Clears objects without the name @obj out of the task space. This is useful - for supporting task modes with single types of objects, as in - @self.single_object_mode without changing the model definition. - - Args: - obj (str): Name of object to keep in the task space - """ - for obj_name, obj_mjcf in self.mujoco_objects.items(): - if obj_name == obj: - continue - else: - sim_state = self.sim.get_state() - # print(self.sim.model.get_joint_qpos_addr(obj_name)) - sim_state.qpos[self.sim.model.get_joint_qpos_addr(obj_name + "_jnt0")[0]] = 10 - self.sim.set_state(sim_state) - self.sim.forward() - def _get_placement_initializer(self): """ Helper function for defining placement initializer and object sampling bounds. @@ -426,26 +407,27 @@ def _get_placement_initializer(self): self.placement_initializer = SequentialCompositeSampler() # can sample anywhere in bin - bin_x_half = self.mujoco_arena.table_full_size[0] / 2 - 0.05 - bin_y_half = self.mujoco_arena.table_full_size[1] / 2 - 0.05 + bin_x_half = self.model.mujoco_arena.table_full_size[0] / 2 - 0.05 + bin_y_half = self.model.mujoco_arena.table_full_size[1] / 2 - 0.05 # each object should just be sampled in the bounds of the bin (with some tolerance) - for obj_name in self.collision_objects: - - self.placement_initializer.sample_on_top( - obj_name, - surface_name="table", + self.placement_initializer.append_sampler( + sampler=UniformRandomSampler( + mujoco_objects=self.objects, x_range=[-bin_x_half, bin_x_half], y_range=[-bin_y_half, bin_y_half], rotation=None, rotation_axis='z', - z_offset=0., ensure_object_boundary_in_range=True, + ensure_valid_placement=True, + reference_pos=self.bin1_pos, + z_offset=0., ) + ) # each visual object should just be at the center of each target bin index = 0 - for obj_name in self.visual_objects: + for vis_obj in self.visual_objects: # get center of target bin bin_x_low = self.bin2_pos[0] @@ -464,17 +446,19 @@ def _get_placement_initializer(self): # placement is relative to object bin, so compute difference and send to placement initializer rel_center = bin_center - self.bin1_pos[:2] - self.placement_initializer.sample_on_top( - obj_name, - surface_name="table", - x_range=[rel_center[0], rel_center[0]], - y_range=[rel_center[1], rel_center[1]], - rotation=0., - rotation_axis='z', - z_offset=self.bin2_pos[2] - self.bin1_pos[2], - ensure_object_boundary_in_range=False, + self.placement_initializer.append_sampler( + sampler=UniformRandomSampler( + mujoco_objects=vis_obj, + x_range=[rel_center[0], rel_center[0]], + y_range=[rel_center[1], rel_center[1]], + rotation=0., + rotation_axis='z', + ensure_object_boundary_in_range=False, + ensure_valid_placement=False, + reference_pos=self.bin1_pos, + z_offset=self.bin2_pos[2] - self.bin1_pos[2], + ) ) - index += 1 def _load_model(self): @@ -483,77 +467,51 @@ def _load_model(self): """ super()._load_model() - # Verify the correct robot has been loaded - assert isinstance(self.robots[0], SingleArm), \ - "Error: Expected one single-armed robot! Got {} type instead.".format(type(self.robots[0])) - # Adjust base pose accordingly xpos = self.robots[0].robot_model.base_xpos_offset["bins"] self.robots[0].robot_model.set_base_xpos(xpos) # load model for table top workspace - self.mujoco_arena = BinsArena( + mujoco_arena = BinsArena( bin1_pos=self.bin1_pos, table_full_size=self.table_full_size, table_friction=self.table_friction ) if self.use_indicator_object: - self.mujoco_arena.add_pos_indicator() + mujoco_arena.add_pos_indicator() # Arena always gets set to zero origin - self.mujoco_arena.set_origin([0, 0, 0]) + mujoco_arena.set_origin([0, 0, 0]) # store some arena attributes - self.bin_size = self.mujoco_arena.table_full_size - - # define mujoco objects - self.ob_inits = [MilkObject, BreadObject, CerealObject, CanObject] - self.vis_inits = [ - MilkVisualObject, - BreadVisualObject, - CerealVisualObject, - CanVisualObject, - ] - self.item_names = ["Milk", "Bread", "Cereal", "Can"] - self.item_names_org = list(self.item_names) - self.obj_to_use = (self.item_names[0] + "{}").format(0) - - lst = [] - for j in range(len(self.vis_inits)): - visual_ob_name = ("Visual" + self.item_names[j] + "0") - visual_ob = self.vis_inits[j]( - name=visual_ob_name, - joints=[], # no free joint for visual objects - ) - lst.append((visual_ob_name, visual_ob)) - self.visual_objects = OrderedDict(lst) - - lst = [] - for i in range(len(self.ob_inits)): - ob_name = (self.item_names[i] + "0") - ob = self.ob_inits[i]( - name=ob_name, - joints=[dict(type="free", damping="0.0005")], # damp the free joint for each object - ) - lst.append((ob_name, ob)) + self.bin_size = mujoco_arena.table_full_size + + self.objects = [] + self.visual_objects = [] + for vis_obj_cls, obj_name in zip( + (MilkVisualObject, BreadVisualObject, CerealVisualObject, CanVisualObject), + self.obj_names, + ): + vis_name = "Visual" + obj_name + vis_obj = vis_obj_cls(name=vis_name) + self.visual_objects.append(vis_obj) - self.collision_objects = OrderedDict(lst) - self.mujoco_objects = OrderedDict(**self.visual_objects, **self.collision_objects) - self.n_objects = len(self.mujoco_objects) + for obj_cls, obj_name in zip( + (MilkObject, BreadObject, CerealObject, CanObject), + self.obj_names, + ): + obj = obj_cls(name=obj_name) + self.objects.append(obj) # task includes arena, robot, and objects of interest - self._get_placement_initializer() self.model = ManipulationTask( - mujoco_arena=self.mujoco_arena, + mujoco_arena=mujoco_arena, mujoco_robots=[robot.robot_model for robot in self.robots], - mujoco_objects=self.mujoco_objects, - initializer=self.placement_initializer, + mujoco_objects=self.visual_objects + self.objects, ) - # set positions of objects - self.model.place_objects() - - # self.model.place_visual() + # Generate placement initializer + self._get_placement_initializer() def _get_reference(self): """ @@ -567,31 +525,18 @@ def _get_reference(self): self.obj_body_id = {} self.obj_geom_id = {} - # id of grippers for contact checking - self.l_finger_geom_ids = [ - self.sim.model.geom_name2id(x) for x in self.robots[0].gripper.important_geoms["left_finger"] - ] - self.r_finger_geom_ids = [ - self.sim.model.geom_name2id(x) for x in self.robots[0].gripper.important_geoms["right_finger"] - ] - # object-specific ids - for i in range(len(self.ob_inits)): - obj_str = str(self.item_names[i]) + "0" - self.obj_body_id[obj_str] = self.sim.model.body_name2id(obj_str) - self.obj_geom_id[obj_str] = self.sim.model.geom_name2id(obj_str + "_0") - - # for checking distance to / contact with objects we want to pick up - self.target_object_body_ids = list(map(int, self.obj_body_id.values())) - self.contact_with_object_geom_ids = list(map(int, self.obj_geom_id.values())) + for obj in (self.visual_objects + self.objects): + self.obj_body_id[obj.name] = self.sim.model.body_name2id(obj.root_body) + self.obj_geom_id[obj.name] = [self.sim.model.geom_name2id(g) for g in obj.contact_geoms] # keep track of which objects are in their corresponding bins - self.objects_in_bins = np.zeros(len(self.ob_inits)) + self.objects_in_bins = np.zeros(len(self.objects)) # target locations in bin for each object type - self.target_bin_placements = np.zeros((len(self.ob_inits), 3)) - for j in range(len(self.ob_inits)): - bin_id = j + self.target_bin_placements = np.zeros((len(self.objects), 3)) + for i, obj in enumerate(self.objects): + bin_id = i bin_x_low = self.bin2_pos[0] bin_y_low = self.bin2_pos[1] if bin_id == 0 or bin_id == 2: @@ -600,7 +545,7 @@ def _get_reference(self): bin_y_low -= self.bin_size[1] / 2. bin_x_low += self.bin_size[0] / 4. bin_y_low += self.bin_size[1] / 4. - self.target_bin_placements[j, :] = [bin_x_low, bin_y_low, self.bin2_pos[2]] + self.target_bin_placements[i, :] = [bin_x_low, bin_y_low, self.bin2_pos[2]] def _reset_internal(self): """ @@ -612,30 +557,31 @@ def _reset_internal(self): if not self.deterministic_reset: # Sample from the placement initializer for all objects - obj_pos, obj_quat = self.model.place_objects() + object_placements = self.placement_initializer.sample() # Loop through all objects and reset their positions - for i, (obj_name, _) in enumerate(self.mujoco_objects.items()): - if "visual" not in obj_name.lower(): - self.sim.data.set_joint_qpos(obj_name + "_jnt0", np.concatenate([np.array(obj_pos[i]), np.array(obj_quat[i])])) - - # information of (collision) objects - self.object_names = list(self.collision_objects.keys()) - self.object_site_ids = [ - self.sim.model.site_name2id(ob_name) for ob_name in self.object_names - ] + for obj_pos, obj_quat, obj in object_placements.values(): + # Set the visual object body locations + if "visual" in obj.name.lower(): + self.sim.model.body_pos[self.obj_body_id[obj.name]] = obj_pos + self.sim.model.body_quat[self.obj_body_id[obj.name]] = obj_quat + else: + # Set the collision object joints + self.sim.data.set_joint_qpos(obj.joints[0], np.concatenate([np.array(obj_pos), np.array(obj_quat)])) # Set the bins to the desired position self.sim.model.body_pos[self.sim.model.body_name2id("bin1")] = self.bin1_pos self.sim.model.body_pos[self.sim.model.body_name2id("bin2")] = self.bin2_pos # Move objects out of the scene depending on the mode + obj_names = {obj.name for obj in self.objects} if self.single_object_mode == 1: - self.obj_to_use = (random.choice(self.item_names) + "{}").format(0) - self.clear_objects(self.obj_to_use) + self.obj_to_use = random.choice(list(obj_names)) elif self.single_object_mode == 2: - self.obj_to_use = (self.item_names[self.object_id] + "{}").format(0) - self.clear_objects(self.obj_to_use) + self.obj_to_use = self.objects[self.object_id].name + if self.single_object_mode in {1, 2}: + obj_names.remove(self.obj_to_use) + self.clear_objects(list(obj_names)) def _get_observation(self): """ @@ -669,13 +615,13 @@ def _get_observation(self): gripper_pose = T.pose2mat((di[pr + "eef_pos"], di[pr + "eef_quat"])) world_pose_in_gripper = T.pose_inv(gripper_pose) - for i in range(len(self.item_names_org)): + for i, obj in enumerate(self.objects): if self.single_object_mode == 2 and self.object_id != i: # Skip adding to observations continue - obj_str = str(self.item_names_org[i]) + "0" + obj_str = obj.name obj_pos = np.array(self.sim.data.body_xpos[self.obj_body_id[obj_str]]) obj_quat = T.convert_quat( self.sim.data.body_xquat[self.obj_body_id[obj_str]], to="xyzw" @@ -697,14 +643,14 @@ def _get_observation(self): if self.single_object_mode == 1: # Zero out other objects observations - for obj_str, obj_mjcf in self.mujoco_objects.items(): - if obj_str == self.obj_to_use: + for obj in self.objects: + if obj.name == self.obj_to_use: continue else: - di["{}_pos".format(obj_str)] *= 0.0 - di["{}_quat".format(obj_str)] *= 0.0 - di["{}_to_{}eef_pos".format(obj_str, pr)] *= 0.0 - di["{}_to_{}eef_quat".format(obj_str, pr)] *= 0.0 + di["{}_pos".format(obj.name)] *= 0.0 + di["{}_quat".format(obj.name)] *= 0.0 + di["{}_to_{}eef_pos".format(obj.name, pr)] *= 0.0 + di["{}_to_{}eef_quat".format(obj.name, pr)] *= 0.0 di["object-state"] = np.concatenate([di[k] for k in object_state_keys]) @@ -719,61 +665,42 @@ def _check_success(self): """ # remember objects that are in the correct bins gripper_site_pos = self.sim.data.site_xpos[self.robots[0].eef_site_id] - for i in range(len(self.ob_inits)): - obj_str = str(self.item_names[i]) + "0" + for i, obj in enumerate(self.objects): + obj_str = obj.name obj_pos = self.sim.data.body_xpos[self.obj_body_id[obj_str]] dist = np.linalg.norm(gripper_site_pos - obj_pos) r_reach = 1 - np.tanh(10.0 * dist) - self.objects_in_bins[i] = int( - (not self.not_in_bin(obj_pos, i)) and r_reach < 0.6 - ) + self.objects_in_bins[i] = int((not self.not_in_bin(obj_pos, i)) and r_reach < 0.6) # returns True if a single object is in the correct bin - if self.single_object_mode == 1 or self.single_object_mode == 2: + if self.single_object_mode in {1, 2}: return np.sum(self.objects_in_bins) > 0 # returns True if all objects are in correct bins - return np.sum(self.objects_in_bins) == len(self.ob_inits) + return np.sum(self.objects_in_bins) == len(self.objects) def _visualization(self): """ Do any needed visualization here. Overrides superclass implementations. """ - - # color the gripper site appropriately based on distance to cube + # color the gripper site appropriately based on distance to closest nut if self.robots[0].gripper_visualization: # find closest object - square_dist = lambda x: np.sum( - np.square(x - self.sim.data.get_site_xpos(self.robots[0].gripper.visualization_sites["grip_site"])) + dists = [ + self._gripper_to_target( + gripper=self.robots[0].gripper, + target=obj.root_body, + target_type="body", + return_distance=True, + ) for obj in self.objects + ] + closest_obj_id = np.argmin(dists) + # Visualize the distance to this target + self._visualize_gripper_to_target( + gripper=self.robots[0].gripper, + target=self.objects[closest_obj_id].root_body, + target_type="body", ) - dists = np.array(list(map(square_dist, self.sim.data.site_xpos))) - dists[self.robots[0].eef_site_id] = np.inf # make sure we don't pick the same site - dists[self.robots[0].eef_cylinder_id] = np.inf - ob_dists = dists[ - self.object_site_ids - ] # filter out object sites we care about - min_dist = np.min(ob_dists) - ob_id = np.argmin(ob_dists) - - # set RGBA for the EEF site here - max_dist = 0.1 - scaled = (1.0 - min(min_dist / max_dist, 1.)) ** 15 - rgba = np.zeros(4) - rgba[0] = 1 - scaled - rgba[1] = scaled - rgba[3] = 0.5 - - self.sim.model.site_rgba[self.robots[0].eef_site_id] = rgba - - def _check_robot_configuration(self, robots): - """ - Sanity check to make sure the inputted robots and configuration is acceptable - - Args: - robots (str or list of str): Robots to instantiate within this env - """ - if type(robots) is list: - assert len(robots) == 1, "Error: Only one robot should be inputted for this task!" class PickPlaceSingle(PickPlace): diff --git a/robosuite/environments/manipulation/single_arm_env.py b/robosuite/environments/manipulation/single_arm_env.py new file mode 100644 index 0000000000..5d11f8703f --- /dev/null +++ b/robosuite/environments/manipulation/single_arm_env.py @@ -0,0 +1,69 @@ +import numpy as np + +from robosuite.environments.manipulation.manipulation_env import ManipulationEnv +from robosuite.robots import SingleArm +from robosuite.utils.transform_utils import mat2quat + + +class SingleArmEnv(ManipulationEnv): + """ + A manipulation environment intended for a single robot arm. + """ + def _load_model(self): + """ + Verifies correct robot model is loaded + """ + super()._load_model() + + # Verify the correct robot has been loaded + assert isinstance(self.robots[0], SingleArm), \ + "Error: Expected one single-armed robot! Got {} type instead.".format(type(self.robots[0])) + + def _check_robot_configuration(self, robots): + """ + Sanity check to make sure the inputted robots and configuration is acceptable + + Args: + robots (str or list of str): Robots to instantiate within this env + """ + super()._check_robot_configuration(robots) + if type(robots) is list: + assert len(robots) == 1, "Error: Only one robot should be inputted for this task!" + + @property + def _eef_xpos(self): + """ + Grabs End Effector position + + Returns: + np.array: End effector(x,y,z) + """ + return np.array(self.sim.data.site_xpos[self.robots[0].eef_site_id]) + + @property + def _eef_xmat(self): + """ + End Effector orientation as a rotation matrix + Note that this draws the orientation from the "ee" site, NOT the gripper site, since the gripper + orientations are inconsistent! + + Returns: + np.array: (3,3) End Effector orientation matrix + """ + pf = self.robots[0].robot_model.naming_prefix + if self.env_configuration == "bimanual": + return np.array(self.sim.data.site_xmat[self.sim.model.site_name2id(pf + "right_ee")]).reshape(3, 3) + else: + return np.array(self.sim.data.site_xmat[self.sim.model.site_name2id(pf + "ee")]).reshape(3, 3) + + @property + def _eef_xquat(self): + """ + End Effector orientation as a (x,y,z,w) quaternion + Note that this draws the orientation from the "ee" site, NOT the gripper site, since the gripper + orientations are inconsistent! + + Returns: + np.array: (x,y,z,w) End Effector quaternion + """ + return mat2quat(self._eef_xmat) diff --git a/robosuite/environments/stack.py b/robosuite/environments/manipulation/stack.py similarity index 77% rename from robosuite/environments/stack.py rename to robosuite/environments/manipulation/stack.py index f40e2f84ee..f338d9b78d 100644 --- a/robosuite/environments/stack.py +++ b/robosuite/environments/manipulation/stack.py @@ -4,15 +4,14 @@ from robosuite.utils.transform_utils import convert_quat from robosuite.utils.mjcf_utils import CustomMaterial -from robosuite.environments.robot_env import RobotEnv -from robosuite.robots import SingleArm +from robosuite.environments.manipulation.single_arm_env import SingleArmEnv from robosuite.models.arenas import TableArena from robosuite.models.objects import BoxObject from robosuite.models.tasks import ManipulationTask, UniformRandomSampler -class Stack(RobotEnv): +class Stack(SingleArmEnv): """ This class corresponds to the stacking task for a single robot arm. @@ -21,6 +20,9 @@ class Stack(RobotEnv): (e.g: "Sawyer" would generate one arm; ["Panda", "Panda", "Sawyer"] would generate three robot arms) Note: Must be a single single-arm robot! + env_configuration (str): Specifies how to position the robots within the environment (default is "default"). + For most single arm environments, this argument has no impact on the robot setup. + controller_configs (str or list of dict): If set, contains relevant controller parameters for creating a custom controller. Else, uses the default controller for this specific task. Should either be single dict if same controller is to be used for all robots or else it should be a list of the same length as @@ -66,13 +68,20 @@ class Stack(RobotEnv): reward_shaping (bool): if True, use dense rewards. - placement_initializer (ObjectPositionSampler instance): if provided, will + placement_initializer (ObjectPositionSampler): if provided, will be used to place objects on every reset, else a UniformRandomSampler is used by default. use_indicator_object (bool): if True, sets up an indicator object that is useful for debugging. + robot_visualizations (bool or list of bool): True if using robot visualization. + Useful for teleoperation. Should either be single bool if robot visualization is to be used for all + robots or else it should be a list of the same length as "robots" param + + env_visualization (bool): True if visualizing sites for the arena / objects in this environment. Useful for + teleoperation. + has_renderer (bool): If true, render the simulation state in a viewer instead of headless mode. @@ -124,6 +133,7 @@ class Stack(RobotEnv): def __init__( self, robots, + env_configuration="default", controller_configs=None, gripper_types="default", gripper_visualizations=False, @@ -136,6 +146,8 @@ def __init__( reward_shaping=False, placement_initializer=None, use_indicator_object=False, + robot_visualizations=False, + env_visualization=False, has_renderer=False, has_offscreen_renderer=True, render_camera="frontview", @@ -150,15 +162,10 @@ def __init__( camera_widths=256, camera_depths=False, ): - """ - - """ - # First, verify that only one robot is being inputted - self._check_robot_configuration(robots) - # settings for table top self.table_full_size = table_full_size self.table_friction = table_friction + self.table_offset = np.array((0, 0, 0.8)) # reward configuration self.reward_scale = reward_scale @@ -174,18 +181,24 @@ def __init__( self.placement_initializer = UniformRandomSampler( x_range=[-0.08, 0.08], y_range=[-0.08, 0.08], - ensure_object_boundary_in_range=False, rotation=None, + ensure_object_boundary_in_range=False, + ensure_valid_placement=True, + reference_pos=self.table_offset, + z_offset=0.01, ) super().__init__( robots=robots, + env_configuration=env_configuration, controller_configs=controller_configs, gripper_types=gripper_types, gripper_visualizations=gripper_visualizations, initialization_noise=initialization_noise, use_camera_obs=use_camera_obs, use_indicator_object=use_indicator_object, + robot_visualizations=robot_visualizations, + env_visualization=env_visualization, has_renderer=has_renderer, has_offscreen_renderer=has_offscreen_renderer, render_camera=render_camera, @@ -256,42 +269,21 @@ def staged_rewards(self): - (float): reward for lifting and aligning - (float): reward for stacking """ - # reaching is successful when the gripper site is close to - # the center of the cube + # reaching is successful when the gripper site is close to the center of the cube cubeA_pos = self.sim.data.body_xpos[self.cubeA_body_id] cubeB_pos = self.sim.data.body_xpos[self.cubeB_body_id] gripper_site_pos = self.sim.data.site_xpos[self.robots[0].eef_site_id] dist = np.linalg.norm(gripper_site_pos - cubeA_pos) r_reach = (1 - np.tanh(10.0 * dist)) * 0.25 - # collision checking - touch_left_finger = False - touch_right_finger = False - touch_cubeA_cubeB = False - - for i in range(self.sim.data.ncon): - c = self.sim.data.contact[i] - if c.geom1 in self.l_finger_geom_ids and c.geom2 == self.cubeA_geom_id: - touch_left_finger = True - if c.geom1 == self.cubeA_geom_id and c.geom2 in self.l_finger_geom_ids: - touch_left_finger = True - if c.geom1 in self.r_finger_geom_ids and c.geom2 == self.cubeA_geom_id: - touch_right_finger = True - if c.geom1 == self.cubeA_geom_id and c.geom2 in self.r_finger_geom_ids: - touch_right_finger = True - if c.geom1 == self.cubeA_geom_id and c.geom2 == self.cubeB_geom_id: - touch_cubeA_cubeB = True - if c.geom1 == self.cubeB_geom_id and c.geom2 == self.cubeA_geom_id: - touch_cubeA_cubeB = True - - # additional grasping reward - if touch_left_finger and touch_right_finger: + # grasping reward + grasping_cubeA = self._check_grasp(gripper=self.robots[0].gripper, object_geoms=self.cubeA) + if grasping_cubeA: r_reach += 0.25 - # lifting is successful when the cube is above the table top - # by a margin + # lifting is successful when the cube is above the table top by a margin cubeA_height = cubeA_pos[2] - table_height = self.mujoco_arena.table_offset[2] + table_height = self.table_offset[2] cubeA_lifted = cubeA_height > table_height + 0.04 r_lift = 1.0 if cubeA_lifted else 0.0 @@ -302,11 +294,10 @@ def staged_rewards(self): ) r_lift += 0.5 * (1 - np.tanh(horiz_dist)) - # stacking is successful when the block is lifted and - # the gripper is not holding the object + # stacking is successful when the block is lifted and the gripper is not holding the object r_stack = 0 - not_touching = not touch_left_finger and not touch_right_finger - if not_touching and r_lift > 0 and touch_cubeA_cubeB: + cubeA_touching_cubeB = self.check_contact(self.cubeA, self.cubeB) + if not grasping_cubeA and r_lift > 0 and cubeA_touching_cubeB: r_stack = 2.0 return r_reach, r_lift, r_stack @@ -317,25 +308,21 @@ def _load_model(self): """ super()._load_model() - # Verify the correct robot has been loaded - assert isinstance(self.robots[0], SingleArm), \ - "Error: Expected one single-armed robot! Got {} type instead.".format(type(self.robots[0])) - # Adjust base pose accordingly xpos = self.robots[0].robot_model.base_xpos_offset["table"](self.table_full_size[0]) self.robots[0].robot_model.set_base_xpos(xpos) # load model for table top workspace - self.mujoco_arena = TableArena( + mujoco_arena = TableArena( table_full_size=self.table_full_size, table_friction=self.table_friction, - table_offset=(0, 0, 0.8), + table_offset=self.table_offset, ) if self.use_indicator_object: - self.mujoco_arena.add_pos_indicator() + mujoco_arena.add_pos_indicator() # Arena always gets set to zero origin - self.mujoco_arena.set_origin([0, 0, 0]) + mujoco_arena.set_origin([0, 0, 0]) # initialize objects of interest tex_attrib = { @@ -360,30 +347,30 @@ def _load_model(self): tex_attrib=tex_attrib, mat_attrib=mat_attrib, ) - cubeA = BoxObject( + self.cubeA = BoxObject( name="cubeA", size_min=[0.02, 0.02, 0.02], size_max=[0.02, 0.02, 0.02], rgba=[1, 0, 0, 1], material=redwood, ) - cubeB = BoxObject( + self.cubeB = BoxObject( name="cubeB", size_min=[0.025, 0.025, 0.025], size_max=[0.025, 0.025, 0.025], rgba=[0, 1, 0, 1], material=greenwood, ) - self.mujoco_objects = OrderedDict([("cubeA", cubeA), ("cubeB", cubeB)]) + + # Add cubes to initializer + self.placement_initializer.add_objects([self.cubeA, self.cubeB]) # task includes arena, robot, and objects of interest self.model = ManipulationTask( - mujoco_arena=self.mujoco_arena, + mujoco_arena=mujoco_arena, mujoco_robots=[robot.robot_model for robot in self.robots], - mujoco_objects=self.mujoco_objects, - initializer=self.placement_initializer, + mujoco_objects=[self.cubeA, self.cubeB], ) - self.model.place_objects() def _get_reference(self): """ @@ -394,24 +381,8 @@ def _get_reference(self): super()._get_reference() # Additional object references from this env - self.cubeA_body_id = self.sim.model.body_name2id("cubeA") - self.cubeB_body_id = self.sim.model.body_name2id("cubeB") - - # information of objects - self.object_names = list(self.mujoco_objects.keys()) - self.object_site_ids = [ - self.sim.model.site_name2id(ob_name) for ob_name in self.object_names - ] - - # id of grippers for contact checking - self.l_finger_geom_ids = [ - self.sim.model.geom_name2id(x) for x in self.robots[0].gripper.important_geoms["left_finger"] - ] - self.r_finger_geom_ids = [ - self.sim.model.geom_name2id(x) for x in self.robots[0].gripper.important_geoms["right_finger"] - ] - self.cubeA_geom_id = self.sim.model.geom_name2id("cubeA") - self.cubeB_geom_id = self.sim.model.geom_name2id("cubeB") + self.cubeA_body_id = self.sim.model.body_name2id(self.cubeA.root_body) + self.cubeB_body_id = self.sim.model.body_name2id(self.cubeB.root_body) def _reset_internal(self): """ @@ -423,11 +394,11 @@ def _reset_internal(self): if not self.deterministic_reset: # Sample from the placement initializer for all objects - obj_pos, obj_quat = self.model.place_objects() + object_placements = self.placement_initializer.sample() # Loop through all objects and reset their positions - for i, (obj_name, _) in enumerate(self.mujoco_objects.items()): - self.sim.data.set_joint_qpos(obj_name + "_jnt0", np.concatenate([np.array(obj_pos[i]), np.array(obj_quat[i])])) + for obj_pos, obj_quat, obj in object_placements.values(): + self.sim.data.set_joint_qpos(obj.joints[0], np.concatenate([np.array(obj_pos), np.array(obj_quat)])) def _get_observation(self): """ @@ -504,38 +475,9 @@ def _visualization(self): """ Do any needed visualization here. Overrides superclass implementations. """ + # Run super call first + super()._visualization() - # color the gripper site appropriately based on distance to cube + # Color the gripper visualization site according to its distance to cubeA if self.robots[0].gripper_visualization: - # find closest object - square_dist = lambda x: np.sum( - np.square(x - self.sim.data.get_site_xpos(self.robots[0].gripper.visualization_sites["grip_site"])) - ) - dists = np.array(list(map(square_dist, self.sim.data.site_xpos))) - dists[self.robots[0].eef_site_id] = np.inf # make sure we don't pick the same site - dists[self.robots[0].eef_cylinder_id] = np.inf - ob_dists = dists[ - self.object_site_ids - ] # filter out object sites we care about - min_dist = np.min(ob_dists) - - # set RGBA for the EEF site here - max_dist = 0.1 - scaled = (1.0 - min(min_dist / max_dist, 1.)) ** 15 - rgba = np.zeros(4) - rgba[0] = 1 - scaled - rgba[1] = scaled - rgba[3] = 0.5 - - self.sim.model.site_rgba[self.robots[0].eef_site_id] = rgba - - def _check_robot_configuration(self, robots): - """ - Sanity check to make sure the inputted robots and configuration is acceptable - - Args: - robots (str or list of str): Robots to instantiate within this env - """ - if type(robots) is list: - assert len(robots) == 1, "Error: Only one robot should be inputted for this task!" - + self._visualize_gripper_to_target(gripper=self.robots[0].gripper, target=self.cubeA) diff --git a/robosuite/environments/manipulation/two_arm_env.py b/robosuite/environments/manipulation/two_arm_env.py new file mode 100644 index 0000000000..1b2cb79a3d --- /dev/null +++ b/robosuite/environments/manipulation/two_arm_env.py @@ -0,0 +1,130 @@ +import numpy as np + +from robosuite.environments.manipulation.manipulation_env import ManipulationEnv +from robosuite.utils.robot_utils import check_bimanual +from robosuite.utils.transform_utils import mat2quat + + +class TwoArmEnv(ManipulationEnv): + """ + A manipulation environment intended for two robot arms. + """ + def _check_robot_configuration(self, robots): + """ + Sanity check to make sure the inputted robots and configuration is acceptable + + Args: + robots (str or list of str): Robots to instantiate within this env + """ + super()._check_robot_configuration(robots) + robots = robots if type(robots) == list or type(robots) == tuple else [robots] + # If default config is used, set env_configuration accordingly + if self.env_configuration == "default": + self.env_configuration = "bimanual" if check_bimanual(robots[0]) else "single-arm-opposed" + + if self.env_configuration == "single-arm-opposed" or self.env_configuration == "single-arm-parallel": + # Specifically two robots should be inputted! + is_bimanual = False + if type(robots) is not list or len(robots) != 2: + raise ValueError("Error: Exactly two single-armed robots should be inputted " + "for this task configuration!") + elif self.env_configuration == "bimanual": + is_bimanual = True + # Specifically one robot should be inputted! + if type(robots) is list and len(robots) != 1: + raise ValueError("Error: Exactly one bimanual robot should be inputted " + "for this task configuration!") + else: + # This is an unknown env configuration, print error + raise ValueError("Error: Unknown environment configuration received. Only 'bimanual'," + "'single-arm-parallel', and 'single-arm-opposed' are supported. Got: {}" + .format(self.env_configuration)) + + # Lastly, check to make sure all inputted robot names are of their correct type (bimanual / not bimanual) + for robot in robots: + if check_bimanual(robot) != is_bimanual: + raise ValueError("Error: For {} configuration, expected bimanual check to return {}; " + "instead, got {}.".format(self.env_configuration, is_bimanual, check_bimanual(robot))) + + @property + def _eef0_xpos(self): + """ + Grab the position of Robot 0's end effector. + + Returns: + np.array: (x,y,z) position of EEF0 + """ + if self.env_configuration == "bimanual": + return np.array(self.sim.data.site_xpos[self.robots[0].eef_site_id["right"]]) + else: + return np.array(self.sim.data.site_xpos[self.robots[0].eef_site_id]) + + @property + def _eef1_xpos(self): + """ + Grab the position of Robot 1's end effector. + + Returns: + np.array: (x,y,z) position of EEF1 + """ + if self.env_configuration == "bimanual": + return np.array(self.sim.data.site_xpos[self.robots[0].eef_site_id["left"]]) + else: + return np.array(self.sim.data.site_xpos[self.robots[1].eef_site_id]) + + @property + def _eef0_xmat(self): + """ + End Effector 0 orientation as a rotation matrix + Note that this draws the orientation from the "ee" site, NOT the gripper site, since the gripper + orientations are inconsistent! + + Returns: + np.array: (3,3) orientation matrix for EEF0 + """ + pf = self.robots[0].robot_model.naming_prefix + if self.env_configuration == "bimanual": + return np.array(self.sim.data.site_xmat[self.sim.model.site_name2id(pf + "right_ee")]).reshape(3, 3) + else: + return np.array(self.sim.data.site_xmat[self.sim.model.site_name2id(pf + "ee")]).reshape(3, 3) + + @property + def _eef1_xmat(self): + """ + End Effector 1 orientation as a rotation matrix + Note that this draws the orientation from the "ee" site, NOT the gripper site, since the gripper + orientations are inconsistent! + + Returns: + np.array: (3,3) orientation matrix for EEF1 + """ + if self.env_configuration == "bimanual": + pf = self.robots[0].robot_model.naming_prefix + return np.array(self.sim.data.site_xmat[self.sim.model.site_name2id(pf + "left_ee")]).reshape(3, 3) + else: + pf = self.robots[1].robot_model.naming_prefix + return np.array(self.sim.data.site_xmat[self.sim.model.site_name2id(pf + "ee")]).reshape(3, 3) + + @property + def _eef0_xquat(self): + """ + End Effector 0 orientation as a (x,y,z,w) quaternion + Note that this draws the orientation from the "ee" site, NOT the gripper site, since the gripper + orientations are inconsistent! + + Returns: + np.array: (x,y,z,w) quaternion for EEF0 + """ + return mat2quat(self._eef0_xmat) + + @property + def _eef1_xquat(self): + """ + End Effector 1 orientation as a (x,y,z,w) quaternion + Note that this draws the orientation from the "ee" site, NOT the gripper site, since the gripper + orientations are inconsistent! + + Returns: + np.array: (x,y,z,w) quaternion for EEF1 + """ + return mat2quat(self._eef1_xmat) diff --git a/robosuite/environments/two_arm_handover.py b/robosuite/environments/manipulation/two_arm_handover.py similarity index 72% rename from robosuite/environments/two_arm_handover.py rename to robosuite/environments/manipulation/two_arm_handover.py index 6d57f6a10c..a401aa93c7 100644 --- a/robosuite/environments/two_arm_handover.py +++ b/robosuite/environments/manipulation/two_arm_handover.py @@ -1,17 +1,16 @@ from collections import OrderedDict import numpy as np -from robosuite.environments.robot_env import RobotEnv +from robosuite.environments.manipulation.two_arm_env import TwoArmEnv from robosuite.models.arenas import TableArena from robosuite.models.objects import HammerObject from robosuite.models.tasks import ManipulationTask, UniformRandomSampler -from robosuite.models.robots import check_bimanual import robosuite.utils.transform_utils as T -class TwoArmHandover(RobotEnv): +class TwoArmHandover(TwoArmEnv): """ This class corresponds to the handover task for two robot arms. @@ -27,7 +26,10 @@ class TwoArmHandover(RobotEnv): :`'single-arm-parallel'`: Only applicable for multi single arm setups. Sets up the (two) single armed robots next to each other on the -x side of the table :`'single-arm-opposed'`: Only applicable for multi single arm setups. Sets up the (two) single armed - robots opposed from each others on the opposite +/-y sides of the table (Default option) + robots opposed from each others on the opposite +/-y sides of the table. + + Note that "default" corresponds to either "bimanual" if a bimanual robot is used or "single-arm-opposed" if two + single-arm robots are used. controller_configs (str or list of dict): If set, contains relevant controller parameters for creating a custom controller. Else, uses the default controller for this specific task. Should either be single @@ -76,13 +78,20 @@ class TwoArmHandover(RobotEnv): reward_shaping (bool): if True, use dense rewards. - placement_initializer (ObjectPositionSampler instance): if provided, will + placement_initializer (ObjectPositionSampler): if provided, will be used to place objects on every reset, else a UniformRandomSampler is used by default. use_indicator_object (bool): if True, sets up an indicator object that is useful for debugging. + robot_visualizations (bool or list of bool): True if using robot visualization. + Useful for teleoperation. Should either be single bool if robot visualization is to be used for all + robots or else it should be a list of the same length as "robots" param + + env_visualization (bool): True if visualizing sites for the arena / objects in this environment. Useful for + teleoperation. + has_renderer (bool): If true, render the simulation state in a viewer instead of headless mode. @@ -150,6 +159,8 @@ def __init__( reward_shaping=False, placement_initializer=None, use_indicator_object=False, + robot_visualizations=False, + env_visualization=False, has_renderer=False, has_offscreen_renderer=True, render_camera="frontview", @@ -164,10 +175,6 @@ def __init__( camera_widths=256, camera_depths=False, ): - # First, verify that correct number of robots are being inputted - self.env_configuration = env_configuration - self._check_robot_configuration(robots) - # Task settings self.prehensile = prehensile @@ -195,19 +202,24 @@ def __init__( self.placement_initializer = UniformRandomSampler( x_range=[-0.1, 0.1], y_range=[-0.05, 0.05], - ensure_object_boundary_in_range=False, rotation=None, rotation_axis=rotation_axis, + ensure_object_boundary_in_range=False, + ensure_valid_placement=True, + reference_pos=self.table_offset, ) super().__init__( robots=robots, + env_configuration=env_configuration, controller_configs=controller_configs, gripper_types=gripper_types, gripper_visualizations=gripper_visualizations, initialization_noise=initialization_noise, use_camera_obs=use_camera_obs, use_indicator_object=use_indicator_object, + robot_visualizations=robot_visualizations, + env_visualization=env_visualization, has_renderer=has_renderer, has_offscreen_renderer=has_offscreen_renderer, render_camera=render_camera, @@ -330,29 +342,29 @@ def _load_model(self): robot.robot_model.set_base_xpos(xpos) # load model for table top workspace - self.mujoco_arena = TableArena( + mujoco_arena = TableArena( table_full_size=self.table_true_size, table_friction=self.table_friction, table_offset=self.table_offset ) if self.use_indicator_object: - self.mujoco_arena.add_pos_indicator() + mujoco_arena.add_pos_indicator() # Arena always gets set to zero origin - self.mujoco_arena.set_origin([0, 0, 0]) + mujoco_arena.set_origin([0, 0, 0]) # initialize objects of interest self.hammer = HammerObject(name="hammer") - self.mujoco_objects = OrderedDict([("hammer", self.hammer)]) + + # Add hammer to initializer + self.placement_initializer.add_objects(self.hammer) # task includes arena, robot, and objects of interest self.model = ManipulationTask( - mujoco_arena=self.mujoco_arena, + mujoco_arena=mujoco_arena, mujoco_robots=[robot.robot_model for robot in self.robots], - mujoco_objects=self.mujoco_objects, - initializer=self.placement_initializer, + mujoco_objects=self.hammer, ) - self.model.place_objects() def _get_reference(self): """ @@ -363,11 +375,8 @@ def _get_reference(self): super()._get_reference() # Hammer object references from this env - self.hammer_body_id = self.sim.model.body_name2id("hammer") - self.hammer_handle_geom_id = self.sim.model.geom_name2id("hammer_handle") - self.hammer_head_geom_id = self.sim.model.geom_name2id("hammer_head") - self.hammer_face_geom_id = self.sim.model.geom_name2id("hammer_face") - self.hammer_claw_geom_id = self.sim.model.geom_name2id("hammer_claw") + self.hammer_body_id = self.sim.model.body_name2id(self.hammer.root_body) + self.hammer_handle_geom_id = self.sim.model.geom_name2id(self.hammer.handle_geoms[0]) # General env references self.table_top_id = self.sim.model.site_name2id("table_top") @@ -382,39 +391,38 @@ def _reset_internal(self): if not self.deterministic_reset: # Sample from the placement initializer for all objects - obj_pos, obj_quat = self.model.place_objects() + object_placements = self.placement_initializer.sample() # Loop through all objects and reset their positions - for i, (obj_name, _) in enumerate(self.mujoco_objects.items()): + for obj_pos, obj_quat, obj in object_placements.values(): # If prehensile, set the object normally if self.prehensile: - self.sim.data.set_joint_qpos(obj_name + "_jnt0", - np.concatenate([np.array(obj_pos[i]), np.array(obj_quat[i])])) + self.sim.data.set_joint_qpos(obj.joints[0], np.concatenate([np.array(obj_pos), np.array(obj_quat)])) # Else, set the object in the hand of the robot and loop a few steps to guarantee the robot is grasping # the object initially else: eef_rot_quat = T.mat2quat(T.euler2mat([np.pi - T.mat2euler(self._eef0_xmat)[2], 0, 0])) - obj_quat[i] = T.quat_multiply(obj_quat[i], eef_rot_quat) + obj_quat = T.quat_multiply(obj_quat, eef_rot_quat) for j in range(100): # Set object in hand - self.sim.data.set_joint_qpos(obj_name + "_jnt0", - np.concatenate([self._eef0_xpos, np.array(obj_quat[i])])) + self.sim.data.set_joint_qpos(obj.joints[0], + np.concatenate([self._eef0_xpos, np.array(obj_quat)])) # Close gripper (action = 1) and prevent arm from moving if self.env_configuration == 'bimanual': # Execute no-op action with gravity compensation torques = np.concatenate([self.robots[0].controller["right"].torque_compensation, self.robots[0].controller["left"].torque_compensation]) - self.sim.data.ctrl[self.robots[0]._ref_joint_torq_actuator_indexes] = torques + self.sim.data.ctrl[self.robots[0]._ref_joint_actuator_indexes] = torques # Execute gripper action - self.robots[0].grip_action([1], "right") + self.robots[0].grip_action(gripper=self.robots[0].gripper["right"], gripper_action=[1]) else: # Execute no-op action with gravity compensation - self.sim.data.ctrl[self.robots[0]._ref_joint_torq_actuator_indexes] =\ + self.sim.data.ctrl[self.robots[0]._ref_joint_actuator_indexes] =\ self.robots[0].controller.torque_compensation - self.sim.data.ctrl[self.robots[1]._ref_joint_torq_actuator_indexes] = \ + self.sim.data.ctrl[self.robots[1]._ref_joint_actuator_indexes] = \ self.robots[1].controller.torque_compensation # Execute gripper action - self.robots[0].grip_action([1]) + self.robots[0].grip_action(gripper=self.robots[0].gripper, gripper_action=[1]) # Take forward step self.sim.step() @@ -438,54 +446,10 @@ def _get_task_info(self): table_height = self.sim.data.site_xpos[self.table_top_id][2] # Check if any Arm's gripper is grasping the hammer handle - - # Single bimanual robot setting - if self.env_configuration == "bimanual": - _contacts_0_lf = len(list( - self.find_contacts( - self.robots[0].gripper["left"].important_geoms["left_finger"], self.hammer.all_geoms - ) - )) > 0 - _contacts_0_rf = len(list( - self.find_contacts( - self.robots[0].gripper["left"].important_geoms["right_finger"], self.hammer.all_geoms - ) - )) > 0 - _contacts_1_lf = len(list( - self.find_contacts( - self.robots[0].gripper["right"].important_geoms["left_finger"], self.hammer.handle_geoms - ) - )) > 0 - _contacts_1_rf = len(list( - self.find_contacts( - self.robots[0].gripper["right"].important_geoms["right_finger"], self.hammer.handle_geoms - ) - )) > 0 - # Multi single arm setting - else: - _contacts_0_lf = len(list( - self.find_contacts( - self.robots[0].gripper.important_geoms["left_finger"], self.hammer.all_geoms - ) - )) > 0 - _contacts_0_rf = len(list( - self.find_contacts( - self.robots[0].gripper.important_geoms["right_finger"], self.hammer.all_geoms - ) - )) > 0 - _contacts_1_lf = len(list( - self.find_contacts( - self.robots[1].gripper.important_geoms["left_finger"], self.hammer.handle_geoms - ) - )) > 0 - _contacts_1_rf = len(list( - self.find_contacts( - self.robots[1].gripper.important_geoms["right_finger"], self.hammer.handle_geoms - ) - )) > 0 - - arm0_grasp_any = True if _contacts_0_lf and _contacts_0_rf else False - arm1_grasp_handle = True if _contacts_1_lf and _contacts_1_rf else False + (g0, g1) = (self.robots[0].gripper["right"], self.robots[0].gripper["left"]) if \ + self.env_configuration == "bimanual" else (self.robots[0].gripper, self.robots[1].gripper) + arm0_grasp_any = self._check_grasp(gripper=g0, object_geoms=self.hammer) + arm1_grasp_handle = self._check_grasp(gripper=g1, object_geoms=self.hammer.handle_geoms) # Return all relevant values return arm0_grasp_any, arm1_grasp_handle, hammer_height, table_height @@ -556,38 +520,6 @@ def _check_success(self): hammer_height - table_height > self.height_threshold \ else False - def _check_robot_configuration(self, robots): - """ - Sanity check to make sure the inputted robots and configuration is acceptable - - Args: - robots (str or list of str): Robots to instantiate within this env - """ - robots = robots if type(robots) == list or type(robots) == tuple else [robots] - if self.env_configuration == "single-arm-opposed" or self.env_configuration == "single-arm-parallel": - # Specifically two robots should be inputted! - is_bimanual = False - if type(robots) is not list or len(robots) != 2: - raise ValueError("Error: Exactly two single-armed robots should be inputted " - "for this task configuration!") - elif self.env_configuration == "bimanual": - is_bimanual = True - # Specifically one robot should be inputted! - if type(robots) is list and len(robots) != 1: - raise ValueError("Error: Exactly one bimanual robot should be inputted " - "for this task configuration!") - else: - # This is an unknown env configuration, print error - raise ValueError("Error: Unknown environment configuration received. Only 'bimanual'," - "'single-arm-parallel', and 'single-arm-opposed' are supported. Got: {}" - .format(self.env_configuration)) - - # Lastly, check to make sure all inputted robot names are of their correct type (bimanual / not bimanual) - for robot in robots: - if check_bimanual(robot) != is_bimanual: - raise ValueError("Error: For {} configuration, expected bimanual check to return {}; " - "instead, got {}.".format(self.env_configuration, is_bimanual, check_bimanual(robot))) - @property def _handle_xpos(self): """ @@ -598,36 +530,6 @@ def _handle_xpos(self): """ return self.sim.data.geom_xpos[self.hammer_handle_geom_id] - @property - def _head_xpos(self): - """ - Grab the position of the hammer head. - - Returns: - np.array: (x,y,z) position of head - """ - return self.sim.data.geom_xpos[self.hammer_head_geom_id] - - @property - def _face_xpos(self): - """ - Grab the position of the hammer face. - - Returns: - np.array: (x,y,z) position of face - """ - return self.sim.data.geom_xpos[self.hammer_face_geom_id] - - @property - def _claw_xpos(self): - """ - Grab the position of the hammer claw. - - Returns: - np.array: (x,y,z) position of claw - """ - return self.sim.data.geom_xpos[self.hammer_claw_geom_id] - @property def _hammer_pos(self): """ @@ -661,75 +563,6 @@ def _hammer_angle(self): z_rotated = np.matmul(mat, z_unit) return np.pi/2 - np.arccos(np.dot(z_unit, z_rotated)) - @property - def _world_quat(self): - """ - Grab the world orientation - - Returns: - np.array: (x,y,z,w) world quaternion - """ - return T.convert_quat(np.array([1, 0, 0, 0]), to="xyzw") - - @property - def _eef0_xpos(self): - """ - Grab the position of Robot 0's end effector. - - Returns: - np.array: (x,y,z) position of EEF0 - """ - if self.env_configuration == "bimanual": - return np.array(self.sim.data.site_xpos[self.robots[0].eef_site_id["right"]]) - else: - return np.array(self.sim.data.site_xpos[self.robots[0].eef_site_id]) - - @property - def _eef1_xpos(self): - """ - Grab the position of Robot 1's end effector. - - Returns: - np.array: (x,y,z) position of EEF1 - """ - if self.env_configuration == "bimanual": - return np.array(self.sim.data.site_xpos[self.robots[0].eef_site_id["left"]]) - else: - return np.array(self.sim.data.site_xpos[self.robots[1].eef_site_id]) - - @property - def _eef0_xmat(self): - """ - End Effector 0 orientation as a rotation matrix - Note that this draws the orientation from the "ee" site, NOT the gripper site, since the gripper - orientations are inconsistent! - - Returns: - np.array: (3,3) orientation matrix for EEF0 - """ - pf = self.robots[0].robot_model.naming_prefix - if self.env_configuration == "bimanual": - return np.array(self.sim.data.site_xmat[self.sim.model.site_name2id(pf + "right_ee")]).reshape(3, 3) - else: - return np.array(self.sim.data.site_xmat[self.sim.model.site_name2id(pf + "ee")]).reshape(3, 3) - - @property - def _eef1_xmat(self): - """ - End Effector 1 orientation as a rotation matrix - Note that this draws the orientation from the "right_/left_hand" body, NOT the gripper site, since the gripper - orientations are inconsistent! - - Returns: - np.array: (3,3) orientation matrix for EEF1 - """ - if self.env_configuration == "bimanual": - pf = self.robots[0].robot_model.naming_prefix - return np.array(self.sim.data.site_xmat[self.sim.model.site_name2id(pf + "left_ee")]).reshape(3, 3) - else: - pf = self.robots[1].robot_model.naming_prefix - return np.array(self.sim.data.site_xmat[self.sim.model.site_name2id(pf + "ee")]).reshape(3, 3) - @property def _gripper_0_to_handle(self): """ diff --git a/robosuite/environments/two_arm_lift.py b/robosuite/environments/manipulation/two_arm_lift.py similarity index 74% rename from robosuite/environments/two_arm_lift.py rename to robosuite/environments/manipulation/two_arm_lift.py index a152c7fdc9..5aed546de5 100644 --- a/robosuite/environments/two_arm_lift.py +++ b/robosuite/environments/manipulation/two_arm_lift.py @@ -1,17 +1,16 @@ from collections import OrderedDict import numpy as np -from robosuite.environments.robot_env import RobotEnv +from robosuite.environments.manipulation.two_arm_env import TwoArmEnv from robosuite.models.arenas import TableArena from robosuite.models.objects import PotWithHandlesObject from robosuite.models.tasks import ManipulationTask, UniformRandomSampler -from robosuite.models.robots import check_bimanual import robosuite.utils.transform_utils as T -class TwoArmLift(RobotEnv): +class TwoArmLift(TwoArmEnv): """ This class corresponds to the lifting task for two robot arms. @@ -27,7 +26,10 @@ class TwoArmLift(RobotEnv): :`'single-arm-parallel'`: Only applicable for multi single arm setups. Sets up the (two) single armed robots next to each other on the -x side of the table :`'single-arm-opposed'`: Only applicable for multi single arm setups. Sets up the (two) single armed - robots opposed from each others on the opposite +/-y sides of the table (Default option) + robots opposed from each others on the opposite +/-y sides of the table. + + Note that "default" corresponds to either "bimanual" if a bimanual robot is used or "single-arm-opposed" if two + single-arm robots are used. controller_configs (str or list of dict): If set, contains relevant controller parameters for creating a custom controller. Else, uses the default controller for this specific task. Should either be single @@ -74,13 +76,20 @@ class TwoArmLift(RobotEnv): reward_shaping (bool): if True, use dense rewards. - placement_initializer (ObjectPositionSampler instance): if provided, will + placement_initializer (ObjectPositionSampler): if provided, will be used to place objects on every reset, else a UniformRandomSampler is used by default. use_indicator_object (bool): if True, sets up an indicator object that is useful for debugging. + robot_visualizations (bool or list of bool): True if using robot visualization. + Useful for teleoperation. Should either be single bool if robot visualization is to be used for all + robots or else it should be a list of the same length as "robots" param + + env_visualization (bool): True if visualizing sites for the arena / objects in this environment. Useful for + teleoperation. + has_renderer (bool): If true, render the simulation state in a viewer instead of headless mode. @@ -147,6 +156,8 @@ def __init__( reward_shaping=False, placement_initializer=None, use_indicator_object=False, + robot_visualizations=False, + env_visualization=False, has_renderer=False, has_offscreen_renderer=True, render_camera="frontview", @@ -161,13 +172,10 @@ def __init__( camera_widths=256, camera_depths=False, ): - # First, verify that correct number of robots are being inputted - self.env_configuration = env_configuration - self._check_robot_configuration(robots) - # settings for table top self.table_full_size = table_full_size self.table_friction = table_friction + self.table_offset = np.array((0, 0, 0.8)) # reward configuration self.reward_scale = reward_scale @@ -184,17 +192,22 @@ def __init__( x_range=[-0.03, 0.03], y_range=[-0.03, 0.03], ensure_object_boundary_in_range=False, - rotation=(-np.pi / 3, np.pi / 3), + ensure_valid_placement=True, + reference_pos=self.table_offset, + rotation=(np.pi + -np.pi / 3, np.pi + np.pi / 3), ) super().__init__( robots=robots, + env_configuration=env_configuration, controller_configs=controller_configs, gripper_types=gripper_types, gripper_visualizations=gripper_visualizations, initialization_noise=initialization_noise, use_camera_obs=use_camera_obs, use_indicator_object=use_indicator_object, + robot_visualizations=robot_visualizations, + env_visualization=env_visualization, has_renderer=has_renderer, has_offscreen_renderer=has_offscreen_renderer, render_camera=render_camera, @@ -264,61 +277,21 @@ def reward(self, action=None): # gh stands for gripper-handle # When grippers are far away, tell them to be closer - # Single bimanual robot setting - if self.env_configuration == "bimanual": - _contacts_0_lf = len(list( - self.find_contacts( - self.robots[0].gripper["left"].important_geoms["left_finger"], self.pot.handle1_geoms - ) - )) > 0 - _contacts_0_rf = len(list( - self.find_contacts( - self.robots[0].gripper["left"].important_geoms["right_finger"], self.pot.handle1_geoms - ) - )) > 0 - _contacts_1_lf = len(list( - self.find_contacts( - self.robots[0].gripper["right"].important_geoms["left_finger"], self.pot.handle0_geoms - ) - )) > 0 - _contacts_1_rf = len(list( - self.find_contacts( - self.robots[0].gripper["right"].important_geoms["right_finger"], self.pot.handle0_geoms - ) - )) > 0 - # Multi single arm setting - else: - _contacts_0_lf = len(list( - self.find_contacts( - self.robots[0].gripper.important_geoms["left_finger"], self.pot.handle1_geoms - ) - )) > 0 - _contacts_0_rf = len(list( - self.find_contacts( - self.robots[0].gripper.important_geoms["right_finger"], self.pot.handle1_geoms - ) - )) > 0 - _contacts_1_lf = len(list( - self.find_contacts( - self.robots[1].gripper.important_geoms["left_finger"], self.pot.handle0_geoms - ) - )) > 0 - _contacts_1_rf = len(list( - self.find_contacts( - self.robots[1].gripper.important_geoms["right_finger"], self.pot.handle0_geoms - ) - )) > 0 + # Get contacts + (g0, g1) = (self.robots[0].gripper["right"], self.robots[0].gripper["left"]) if \ + self.env_configuration == "bimanual" else (self.robots[0].gripper, self.robots[1].gripper) + _g0h_dist = np.linalg.norm(_gripper0_to_handle0) _g1h_dist = np.linalg.norm(_gripper1_to_handle1) # Grasping reward - if _contacts_0_lf and _contacts_0_rf: + if self._check_grasp(gripper=g0, object_geoms=self.pot.handle0_geoms): reward += 0.25 # Reaching reward reward += 0.5 * (1 - np.tanh(10.0 * _g0h_dist)) # Grasping reward - if _contacts_1_lf and _contacts_1_rf: + if self._check_grasp(gripper=g1, object_geoms=self.pot.handle1_geoms): reward += 0.25 # Reaching reward reward += 0.5 * (1 - np.tanh(10.0 * _g1h_dist)) @@ -355,30 +328,29 @@ def _load_model(self): robot.robot_model.set_base_xpos(xpos) # load model for table top workspace - self.mujoco_arena = TableArena( + mujoco_arena = TableArena( table_full_size=self.table_full_size, table_friction=self.table_friction, - table_offset=(0, 0, 0.8), + table_offset=self.table_offset, ) if self.use_indicator_object: - self.mujoco_arena.add_pos_indicator() + mujoco_arena.add_pos_indicator() # Arena always gets set to zero origin - self.mujoco_arena.set_origin([0, 0, 0]) + mujoco_arena.set_origin([0, 0, 0]) # initialize objects of interest - self.pot_name = "pot" - self.pot = PotWithHandlesObject(name=self.pot_name) - self.mujoco_objects = OrderedDict([(self.pot_name, self.pot)]) + self.pot = PotWithHandlesObject(name="pot") + + # Add pot to initializer + self.placement_initializer.add_objects(self.pot) # task includes arena, robot, and objects of interest self.model = ManipulationTask( - mujoco_arena=self.mujoco_arena, + mujoco_arena=mujoco_arena, mujoco_robots=[robot.robot_model for robot in self.robots], - mujoco_objects=self.mujoco_objects, - initializer=self.placement_initializer, + mujoco_objects=self.pot, ) - self.model.place_objects() def _get_reference(self): """ @@ -389,11 +361,11 @@ def _get_reference(self): super()._get_reference() # Additional object references from this env - self.pot_body_id = self.sim.model.body_name2id(self.pot_name) - self.handle1_site_id = self.sim.model.site_name2id(f"{self.pot_name}_handle0") - self.handle0_site_id = self.sim.model.site_name2id(f"{self.pot_name}_handle1") + self.pot_body_id = self.sim.model.body_name2id(self.pot.root_body) + self.handle0_site_id = self.sim.model.site_name2id(self.pot.important_sites["handle0"]) + self.handle1_site_id = self.sim.model.site_name2id(self.pot.important_sites["handle1"]) self.table_top_id = self.sim.model.site_name2id("table_top") - self.pot_center_id = self.sim.model.site_name2id(f"{self.pot_name}_center") + self.pot_center_id = self.sim.model.site_name2id(self.pot.important_sites["center"]) def _reset_internal(self): """ @@ -405,11 +377,11 @@ def _reset_internal(self): if not self.deterministic_reset: # Sample from the placement initializer for all objects - obj_pos, obj_quat = self.model.place_objects() + object_placements = self.placement_initializer.sample() # Loop through all objects and reset their positions - for i, (obj_name, _) in enumerate(self.mujoco_objects.items()): - self.sim.data.set_joint_qpos(obj_name + "_jnt0", np.concatenate([np.array(obj_pos[i]), np.array(obj_quat[i])])) + for obj_pos, obj_quat, obj in object_placements.values(): + self.sim.data.set_joint_qpos(obj.joints[0], np.concatenate([np.array(obj_pos), np.array(obj_quat)])) def _get_observation(self): """ @@ -467,6 +439,25 @@ def _get_observation(self): return di + def _visualization(self): + """ + Visualize each gripper site proportional to the distance to their respective handles + """ + # Run super call first + super()._visualization() + + # Color the gripper visualization site according to its distance to the cube + handles = [self.pot.important_sites[f"handle{i}"] for i in range(2)] + if self.env_configuration == "bimanual": + visualizations = [self.robots[0].gripper_visualization[arm] for arm in self.robots[0].arms] + grippers = [self.robots[0].gripper[arm] for arm in self.robots[0].arms] + else: + visualizations = [robot.gripper_visualization for robot in self.robots] + grippers = [robot.gripper for robot in self.robots] + for vis, gripper, handle in zip(visualizations, grippers, handles): + if vis: + self._visualize_gripper_to_target(gripper=gripper, target=handle, target_type="site") + def _check_success(self): """ Check if pot is successfully lifted @@ -480,38 +471,6 @@ def _check_success(self): # cube is higher than the table top above a margin return pot_bottom_height > table_height + 0.10 - def _check_robot_configuration(self, robots): - """ - Sanity check to make sure the inputted robots and configuration is acceptable - - Args: - robots (str or list of str): Robots to instantiate within this env - """ - robots = robots if type(robots) == list or type(robots) == tuple else [robots] - if self.env_configuration == "single-arm-opposed" or self.env_configuration == "single-arm-parallel": - # Specifically two robots should be inputted! - is_bimanual = False - if type(robots) is not list or len(robots) != 2: - raise ValueError("Error: Exactly two single-armed robots should be inputted " - "for this task configuration!") - elif self.env_configuration == "bimanual": - is_bimanual = True - # Specifically one robot should be inputted! - if type(robots) is list and len(robots) != 1: - raise ValueError("Error: Exactly one bimanual robot should be inputted " - "for this task configuration!") - else: - # This is an unknown env configuration, print error - raise ValueError("Error: Unknown environment configuration received. Only 'bimanual'," - "'single-arm-parallel', and 'single-arm-opposed' are supported. Got: {}" - .format(self.env_configuration)) - - # Lastly, check to make sure all inputted robot names are of their correct type (bimanual / not bimanual) - for robot in robots: - if check_bimanual(robot) != is_bimanual: - raise ValueError("Error: For {} configuration, expected bimanual check to return {}; " - "instead, got {}.".format(self.env_configuration, is_bimanual, check_bimanual(robot))) - @property def _handle0_xpos(self): """ @@ -542,42 +501,6 @@ def _pot_quat(self): """ return T.convert_quat(self.sim.data.body_xquat[self.pot_body_id], to="xyzw") - @property - def _world_quat(self): - """ - Grab the world orientation - - Returns: - np.array: (x,y,z,w) world quaternion - """ - return T.convert_quat(np.array([1, 0, 0, 0]), to="xyzw") - - @property - def _eef0_xpos(self): - """ - Grab the position of Robot 0's end effector. - - Returns: - np.array: (x,y,z) position of EEF0 - """ - if self.env_configuration == "bimanual": - return np.array(self.sim.data.site_xpos[self.robots[0].eef_site_id["left"]]) - else: - return np.array(self.sim.data.site_xpos[self.robots[0].eef_site_id]) - - @property - def _eef1_xpos(self): - """ - Grab the position of Robot 1's end effector. - - Returns: - np.array: (x,y,z) position of EEF1 - """ - if self.env_configuration == "bimanual": - return np.array(self.sim.data.site_xpos[self.robots[0].eef_site_id["right"]]) - else: - return np.array(self.sim.data.site_xpos[self.robots[1].eef_site_id]) - @property def _gripper0_to_handle0(self): """ diff --git a/robosuite/environments/two_arm_peg_in_hole.py b/robosuite/environments/manipulation/two_arm_peg_in_hole.py similarity index 80% rename from robosuite/environments/two_arm_peg_in_hole.py rename to robosuite/environments/manipulation/two_arm_peg_in_hole.py index ebd344452c..e1ce1a9f00 100644 --- a/robosuite/environments/two_arm_peg_in_hole.py +++ b/robosuite/environments/manipulation/two_arm_peg_in_hole.py @@ -1,16 +1,15 @@ import numpy as np import robosuite.utils.transform_utils as T -from robosuite.environments.robot_env import RobotEnv -from robosuite.utils.mjcf_utils import CustomMaterial, array_to_string +from robosuite.environments.manipulation.two_arm_env import TwoArmEnv +from robosuite.utils.mjcf_utils import CustomMaterial, array_to_string, find_elements from robosuite.models.objects import CylinderObject, PlateWithHoleObject from robosuite.models.arenas import EmptyArena -from robosuite.models import MujocoWorldBase -from robosuite.models.robots import check_bimanual +from robosuite.models.tasks import ManipulationTask -class TwoArmPegInHole(RobotEnv): +class TwoArmPegInHole(TwoArmEnv): """ This class corresponds to the peg-in-hole task for two robot arms. @@ -26,15 +25,17 @@ class TwoArmPegInHole(RobotEnv): :`'single-arm-parallel'`: Only applicable for multi single arm setups. Sets up the (two) single armed robots next to each other on the -x side of the table :`'single-arm-opposed'`: Only applicable for multi single arm setups. Sets up the (two) single armed - robots opposed from each others on the opposite +/-y sides of the table (Default option) + robots opposed from each others on the opposite +/-y sides of the table. + + Note that "default" corresponds to either "bimanual" if a bimanual robot is used or "single-arm-opposed" if two + single-arm robots are used. controller_configs (str or list of dict): If set, contains relevant controller parameters for creating a custom controller. Else, uses the default controller for this specific task. Should either be single dict if same controller is to be used for all robots or else it should be a list of the same length as "robots" param - gripper_types (str or list of str): type of gripper, used to instantiate - gripper models from gripper factory. + gripper_types (str or list of str): type of gripper, used to instantiate gripper models from gripper factory. For this environment, setting a value other than the default (None) will raise an AssertionError, as this environment is not meant to be used with any gripper at all. @@ -77,6 +78,13 @@ class TwoArmPegInHole(RobotEnv): use_indicator_object (bool): if True, sets up an indicator object that is useful for debugging. + robot_visualizations (bool or list of bool): True if using robot visualization. + Useful for teleoperation. Should either be single bool if robot visualization is to be used for all + robots or else it should be a list of the same length as "robots" param + + env_visualization (bool): True if visualizing sites for the arena / objects in this environment. Useful for + teleoperation. + has_renderer (bool): If true, render the simulation state in a viewer instead of headless mode. @@ -143,6 +151,8 @@ def __init__( peg_radius=(0.015, 0.03), peg_length=0.13, use_indicator_object=False, + robot_visualizations=False, + env_visualization=False, has_renderer=False, has_offscreen_renderer=True, render_camera="frontview", @@ -157,10 +167,6 @@ def __init__( camera_widths=256, camera_depths=False, ): - # First, verify that correct number of robots are being inputted - self.env_configuration = env_configuration - self._check_robot_configuration(robots) - # Assert that the gripper type is None assert gripper_types is None, "Tried to specify gripper other than None in TwoArmPegInHole environment!" @@ -177,12 +183,15 @@ def __init__( super().__init__( robots=robots, + env_configuration=env_configuration, controller_configs=controller_configs, gripper_types=gripper_types, gripper_visualizations=gripper_visualizations, initialization_noise=initialization_noise, use_camera_obs=use_camera_obs, use_indicator_object=use_indicator_object, + robot_visualizations=robot_visualizations, + env_visualization=env_visualization, has_renderer=has_renderer, has_offscreen_renderer=has_offscreen_renderer, render_camera=render_camera, @@ -198,7 +207,7 @@ def __init__( camera_depths=camera_depths, ) - def reward(self, action): + def reward(self, action=None): """ Reward function for the task. @@ -241,8 +250,7 @@ def reward(self, action): reward += 1 - np.tanh(np.abs(t)) reward += cos - # if we're not reward shaping, we need to scale our sparse reward so that the max reward is identical - # to its dense version + # if we're not reward shaping, scale sparse reward so that the max reward is identical to its dense version else: reward *= 5.0 @@ -278,18 +286,15 @@ def _load_model(self): robot.robot_model.set_base_xpos(xpos) # Add arena and robot - self.model = MujocoWorldBase() - self.mujoco_arena = EmptyArena() + mujoco_arena = EmptyArena() if self.use_indicator_object: - self.mujoco_arena.add_pos_indicator() - self.model.merge(self.mujoco_arena) - for robot in self.robots: - self.model.merge(robot.robot_model) + mujoco_arena.add_pos_indicator() + + # Arena always gets set to zero origin + mujoco_arena.set_origin([0, 0, 0]) # initialize objects of interest - self.hole = PlateWithHoleObject( - name="hole", - ) + self.hole = PlateWithHoleObject(name="hole") tex_attrib = { "type": "cube", } @@ -311,30 +316,40 @@ def _load_model(self): size_max=(self.peg_radius[1], self.peg_length), material=greenwood, rgba=[0, 1, 0, 1], + joints=None, ) # Load hole object - self.hole_obj = self.hole.get_object_subtree(site=True) - self.hole_obj.set("quat", "0 0 0.707 0.707") - self.hole_obj.set("pos", "0.11 0 0.17") - self.model.merge_asset(self.hole) + hole_obj = self.hole.get_obj() + hole_obj.set("quat", "0 0 0.707 0.707") + hole_obj.set("pos", "0.11 0 0.17") # Load peg object - self.peg_obj = self.peg.get_object_subtree(site=True) - self.peg_obj.set("pos", array_to_string((0, 0, self.peg_length))) - self.model.merge_asset(self.peg) + peg_obj = self.peg.get_obj() + peg_obj.set("pos", array_to_string((0, 0, self.peg_length))) - # Depending on env configuration, append appropriate objects to arms + # Append appropriate objects to arms if self.env_configuration == "bimanual": - self.model.worldbody.find(".//body[@name='{}']" - .format(self.robots[0].robot_model.eef_name["left"])).append(self.hole_obj) - self.model.worldbody.find(".//body[@name='{}']" - .format(self.robots[0].robot_model.eef_name["right"])).append(self.peg_obj) + r_eef, l_eef = [self.robots[0].robot_model.eef_name[arm] for arm in self.robots[0].arms] + r_model, l_model = [self.robots[0].robot_model, self.robots[0].robot_model] else: - self.model.worldbody.find(".//body[@name='{}']" - .format(self.robots[1].robot_model.eef_name)).append(self.hole_obj) - self.model.worldbody.find(".//body[@name='{}']" - .format(self.robots[0].robot_model.eef_name)).append(self.peg_obj) + r_eef, l_eef = [robot.robot_model.eef_name for robot in self.robots] + r_model, l_model = [self.robots[0].robot_model, self.robots[1].robot_model] + r_body = find_elements(root=r_model.worldbody, tags="body", attribs={"name": r_eef}, return_first=True) + l_body = find_elements(root=l_model.worldbody, tags="body", attribs={"name": l_eef}, return_first=True) + r_body.append(peg_obj) + l_body.append(hole_obj) + + # task includes arena, robot, and objects of interest + # We don't add peg and hole directly since they were already appended to the robots + self.model = ManipulationTask( + mujoco_arena=mujoco_arena, + mujoco_robots=[robot.robot_model for robot in self.robots], + ) + + # Make sure to add relevant assets from peg and hole objects + self.model.merge_asset(self.hole) + self.model.merge_asset(self.peg) def _get_reference(self): """ @@ -345,8 +360,8 @@ def _get_reference(self): super()._get_reference() # Additional object references from this env - self.hole_body_id = self.sim.model.body_name2id("hole") - self.peg_body_id = self.sim.model.body_name2id("peg") + self.hole_body_id = self.sim.model.body_name2id(self.hole.root_body) + self.peg_body_id = self.sim.model.body_name2id(self.peg.root_body) def _reset_internal(self): """ @@ -463,9 +478,7 @@ def _compute_orientation(self): return ( t, d, - abs( - np.dot(hole_normal, v) / np.linalg.norm(hole_normal) / np.linalg.norm(v) - ), + abs(np.dot(hole_normal, v) / np.linalg.norm(hole_normal) / np.linalg.norm(v)), ) def _peg_pose_in_hole_frame(self): @@ -477,13 +490,13 @@ def _peg_pose_in_hole_frame(self): np.array: (4,4) matrix corresponding to the pose of the peg in the hole frame """ # World frame - peg_pos_in_world = self.sim.data.get_body_xpos("peg") - peg_rot_in_world = self.sim.data.get_body_xmat("peg").reshape((3, 3)) + peg_pos_in_world = self.sim.data.get_body_xpos(self.peg.root_body) + peg_rot_in_world = self.sim.data.get_body_xmat(self.peg.root_body).reshape((3, 3)) peg_pose_in_world = T.make_pose(peg_pos_in_world, peg_rot_in_world) # World frame - hole_pos_in_world = self.sim.data.get_body_xpos("hole") - hole_rot_in_world = self.sim.data.get_body_xmat("hole").reshape((3, 3)) + hole_pos_in_world = self.sim.data.get_body_xpos(self.hole.root_body) + hole_rot_in_world = self.sim.data.get_body_xmat(self.hole.root_body).reshape((3, 3)) hole_pose_in_world = T.make_pose(hole_pos_in_world, hole_rot_in_world) world_pose_in_hole = T.pose_inv(hole_pose_in_world) @@ -492,35 +505,3 @@ def _peg_pose_in_hole_frame(self): peg_pose_in_world, world_pose_in_hole ) return peg_pose_in_hole - - def _check_robot_configuration(self, robots): - """ - Sanity check to make sure the inputted robots and configuration is acceptable - - Args: - robots (str or list of str): Robots to instantiate within this env - """ - robots = robots if type(robots) == list or type(robots) == tuple else [robots] - if self.env_configuration == "single-arm-opposed" or self.env_configuration == "single-arm-parallel": - # Specifically two robots should be inputted! - is_bimanual = False - if type(robots) is not list or len(robots) != 2: - raise ValueError("Error: Exactly two single-armed robots should be inputted " - "for this task configuration!") - elif self.env_configuration == "bimanual": - is_bimanual = True - # Specifically one robot should be inputted! - if type(robots) is list and len(robots) != 1: - raise ValueError("Error: Exactly one bimanual robot should be inputted " - "for this task configuration!") - else: - # This is an unknown env configuration, print error - raise ValueError("Error: Unknown environment configuration received. Only 'bimanual'," - "'single-arm-parallel', and 'single-arm-opposed' are supported. Got: {}" - .format(self.env_configuration)) - - # Lastly, check to make sure all inputted robot names are of their correct type (bimanual / not bimanual) - for robot in robots: - if check_bimanual(robot) != is_bimanual: - raise ValueError("Error: For {} configuration, expected bimanual check to return {}; " - "instead, got {}.".format(self.env_configuration, is_bimanual, check_bimanual(robot))) diff --git a/robosuite/environments/wipe.py b/robosuite/environments/manipulation/wipe.py similarity index 78% rename from robosuite/environments/wipe.py rename to robosuite/environments/manipulation/wipe.py index 1b8e631b1e..5a39fb2d8f 100644 --- a/robosuite/environments/wipe.py +++ b/robosuite/environments/manipulation/wipe.py @@ -1,8 +1,7 @@ from collections import OrderedDict import numpy as np -from robosuite.environments.robot_env import RobotEnv -from robosuite.robots import SingleArm +from robosuite.environments.manipulation.single_arm_env import SingleArmEnv from robosuite.models.arenas import WipeArena from robosuite.models.tasks import ManipulationTask, UniformRandomSampler @@ -30,11 +29,11 @@ "line_width": 0.04, # Width of the line to wipe (diameter of the pegs) "two_clusters": False, # if the dirt to wipe is one continuous line or two "coverage_factor": 0.6, # how much of the table surface we cover - "num_sensors": 50, # How many particles of dirt to generate in the environment + "num_markers": 50, # How many particles of dirt to generate in the environment # settings for thresholds "contact_threshold": 3, # Minimum eef force to qualify as contact [N] - "touch_threshold": 5, # force threshold (N) to overcome to change the color of the sensor (wipe the peg) + "touch_threshold": 5, # force threshold (N) to overcome to change the color of the marker (wipe the peg) "pressure_threshold_max": 70, # maximum force allowed (N) # misc settings @@ -45,7 +44,7 @@ } -class Wipe(RobotEnv): +class Wipe(SingleArmEnv): """ This class corresponds to the Wiping task for a single robot arm @@ -54,6 +53,9 @@ class Wipe(RobotEnv): (e.g: "Sawyer" would generate one arm; ["Panda", "Panda", "Sawyer"] would generate three robot arms) Note: Must be a single single-arm robot! + env_configuration (str): Specifies how to position the robots within the environment (default is "default"). + For most single arm environments, this argument has no impact on the robot setup. + controller_configs (str or list of dict): If set, contains relevant controller parameters for creating a custom controller. Else, uses the default controller for this specific task. Should either be single dict if same controller is to be used for all robots or else it should be a list of the same length as @@ -93,13 +95,16 @@ class Wipe(RobotEnv): reward_shaping (bool): if True, use dense rewards. - placement_initializer (ObjectPositionSampler instance): if provided, will - be used to place objects on every reset, else a UniformRandomSampler - is used by default. - use_indicator_object (bool): if True, sets up an indicator object that is useful for debugging. + robot_visualizations (bool or list of bool): True if using robot visualization. + Useful for teleoperation. Should either be single bool if robot visualization is to be used for all + robots or else it should be a list of the same length as "robots" param + + env_visualization (bool): True if visualizing sites for the arena / objects in this environment. Useful for + teleoperation. + has_renderer (bool): If true, render the simulation state in a viewer instead of headless mode. @@ -157,6 +162,7 @@ class Wipe(RobotEnv): def __init__( self, robots, + env_configuration="default", controller_configs=None, gripper_types="WipingGripper", gripper_visualizations=False, @@ -165,8 +171,9 @@ def __init__( use_object_obs=True, reward_scale=1.0, reward_shaping=True, - placement_initializer=None, use_indicator_object=False, + robot_visualizations=False, + env_visualization=False, has_renderer=False, has_offscreen_renderer=True, render_camera="frontview", @@ -182,9 +189,6 @@ def __init__( camera_depths=False, task_config=None, ): - # First, verify that only one robot is being inputted - self._check_robot_configuration(robots) - # Assert that the gripper type is None assert gripper_types == "WipingGripper",\ "Tried to specify gripper other than WipingGripper in Wipe environment!" @@ -214,15 +218,16 @@ def __init__( # settings for table top self.table_full_size = self.task_config['table_full_size'] - self.table_offset = self.task_config['table_offset'] - self.table_friction = self.task_config['table_friction'] - self.table_friction_std = self.task_config['table_friction_std'] self.table_height = self.task_config['table_height'] self.table_height_std = self.task_config['table_height_std'] + delta_height = min(0, np.random.normal(self.table_height, self.table_height_std)) # sample variation in height + self.table_offset = np.array(self.task_config['table_offset']) + np.array((0, 0, delta_height)) + self.table_friction = self.task_config['table_friction'] + self.table_friction_std = self.task_config['table_friction_std'] self.line_width = self.task_config['line_width'] self.two_clusters = self.task_config['two_clusters'] self.coverage_factor = self.task_config['coverage_factor'] - self.num_sensors = self.task_config['num_sensors'] + self.num_markers = self.task_config['num_markers'] # settings for thresholds self.contact_threshold = self.task_config['contact_threshold'] @@ -238,7 +243,7 @@ def __init__( # Scale reward if desired (see reward method for details) self.reward_normalization_factor = horizon / \ - (self.num_sensors * self.unit_wiped_reward + + (self.num_markers * self.unit_wiped_reward + horizon * (self.wipe_contact_reward + self.task_complete_reward)) # ee resets @@ -246,7 +251,7 @@ def __init__( self.ee_torque_bias = np.zeros(3) # set other wipe-specific attributes - self.wiped_sensors = [] + self.wiped_markers = [] self.collisions = 0 self.f_excess = 0 self.metadata = [] @@ -255,24 +260,17 @@ def __init__( # whether to include and use ground-truth object states self.use_object_obs = use_object_obs - # object placement initializer - if placement_initializer: - self.placement_initializer = placement_initializer - else: - self.placement_initializer = UniformRandomSampler( - x_range=[0, 0.2], - y_range=[0, 0.2], - ensure_object_boundary_in_range=False, - rotation=None) - super().__init__( robots=robots, + env_configuration=env_configuration, controller_configs=controller_configs, gripper_types=gripper_types, gripper_visualizations=gripper_visualizations, initialization_noise=initialization_noise, use_camera_obs=use_camera_obs, use_indicator_object=use_indicator_object, + robot_visualizations=robot_visualizations, + env_visualization=env_visualization, has_renderer=has_renderer, has_offscreen_renderer=has_offscreen_renderer, render_camera=render_camera, @@ -288,7 +286,7 @@ def __init__( camera_depths=camera_depths, ) - def reward(self, action): + def reward(self, action=None): """ Reward function for the task. @@ -316,7 +314,7 @@ def reward(self, action): Note that the final per-step reward is normalized given the theoretical best episode return and then scaled: reward_scale * (horizon / - (num_sensors * unit_wiped_reward + horizon * (wipe_contact_reward + task_complete_reward))) + (num_markers * unit_wiped_reward + horizon * (wipe_contact_reward + task_complete_reward))) Args: action (np array): [NOT USED] @@ -329,18 +327,18 @@ def reward(self, action): total_force_ee = np.linalg.norm(np.array(self.robots[0].recent_ee_forcetorques.current[:3])) # Neg Reward from collisions of the arm with the table - if self._check_arm_contact()[0]: + if self.check_contact(self.robots[0].robot_model): if self.reward_shaping: reward = self.arm_limit_collision_penalty self.collisions += 1 - elif self._check_q_limits()[0]: + elif self.robots[0].check_q_limits(): if self.reward_shaping: reward = self.arm_limit_collision_penalty self.collisions += 1 else: # If the arm is not colliding or in joint limits, we check if we are wiping # (we don't want to reward wiping if there are unsafe situations) - sensors_active_ids = [] + active_markers = [] # Current 3D location of the corners of the wiping tool in world frame c_geoms = self.robots[0].gripper.important_geoms["corners"] @@ -380,75 +378,65 @@ def PointInRectangle(X, Y, Z, W, P): # Only go into this computation if there are contact points if self.sim.data.ncon != 0: - # Check each sensor that is still active - for sensor_name in self.model.arena.sensor_names: + # Check each marker that is still active + for marker in self.model.mujoco_arena.markers: - # Current sensor 3D location in world frame - # sensor_pos = np.array( - # self.sim.data.body_xpos[self.sim.model.site_bodyid[self.sim.model.site_name2id(self.model.arena.sensor_site_names[sensor_name])]]) - sensor_pos = np.array( - self.sim.data.site_xpos[ - self.sim.model.site_name2id(self.model.arena.sensor_site_names[sensor_name])]) + # Current marker 3D location in world frame + marker_pos = np.array(self.sim.data.body_xpos[self.sim.model.body_name2id(marker.root_body)]) # We use the second tool corner as point on the plane and define the vector connecting - # the sensor position to that point - v = sensor_pos - corner2_pos + # the marker position to that point + v = marker_pos - corner2_pos - # Shortest distance between the center of the sensor and the plane + # Shortest distance between the center of the marker and the plane dist = np.dot(v, n) - # Projection of the center of the sensor onto the plane - projected_point = np.array(sensor_pos) - dist * n + # Projection of the center of the marker onto the plane + projected_point = np.array(marker_pos) - dist * n - # Positive distances means the center of the sensor is over the plane - # The plane is aligned with the bottom of the wiper and pointing up, so the sensor would be over it + # Positive distances means the center of the marker is over the plane + # The plane is aligned with the bottom of the wiper and pointing up, so the marker would be over it if dist > 0.0: # Distance smaller than this threshold means we are close to the plane on the upper part if dist < 0.02: # Write touching points and projected point in coordinates of the plane pp_2 = np.array( [np.dot(projected_point - corner2_pos, v1), np.dot(projected_point - corner2_pos, v2)]) - # Check if sensor is within the tool center: + # Check if marker is within the tool center: if PointInRectangle(pp[0], pp[1], pp[2], pp[3], pp_2): - parts = sensor_name.split('_') - sensors_active_ids += [int(parts[1])] - - # Obtain the list of currently active (wiped) sensors that where not wiped before - # These are the sensors we are wiping at this step - lall = np.where(np.isin(sensors_active_ids, self.wiped_sensors, invert=True)) - new_sensors_active_ids = np.array(sensors_active_ids)[lall] - - # Loop through all new sensors we are wiping at this step - for new_sensor_active_id in new_sensors_active_ids: - # Grab relevant sensor id info - sensor_name = self.model.arena.sensor_site_names['contact_' + str(new_sensor_active_id) + '_sensor'] - new_sensor_active_geom_id = self.sim.model.geom_name2id(sensor_name + "_vis") - # Make this sensor transparent since we wiped it (alpha = 0) - self.sim.model.geom_rgba[new_sensor_active_geom_id] = [0, 0, 0, 0] - # Add this sensor the wiped list - self.wiped_sensors += [new_sensor_active_id] + active_markers.append(marker) + + # Obtain the list of currently active (wiped) markers that where not wiped before + # These are the markers we are wiping at this step + lall = np.where(np.isin(active_markers, self.wiped_markers, invert=True)) + new_active_markers = np.array(active_markers)[lall] + + # Loop through all new markers we are wiping at this step + for new_active_marker in new_active_markers: + # Grab relevant marker id info + new_active_marker_geom_id = self.sim.model.geom_name2id(new_active_marker.visual_geoms[0]) + # Make this marker transparent since we wiped it (alpha = 0) + self.sim.model.geom_rgba[new_active_marker_geom_id][3] = 0 + # Add this marker the wiped list + self.wiped_markers.append(new_active_marker) # Add reward if we're using the dense reward if self.reward_shaping: reward += self.unit_wiped_reward # Additional reward components if using dense rewards if self.reward_shaping: - # If we haven't wiped all the sensors yet, add a smooth reward for getting closer + # If we haven't wiped all the markers yet, add a smooth reward for getting closer # to the centroid of the dirt to wipe - if len(self.wiped_sensors) < len(self.model.arena.sensor_names): + if len(self.wiped_markers) < self.num_markers: mean_distance_to_things_to_wipe = 0 - num_non_wiped_sensors = 0 - for sensor_name in self.model.arena.sensor_names: - parts = sensor_name.split('_') - sensor_id = int(parts[1]) - if sensor_id not in self.wiped_sensors: - sensor_pos = np.array( - self.sim.data.site_xpos[ - self.sim.model.site_name2id(self.model.arena.sensor_site_names[sensor_name])]) - gripper_position = np.array(self.sim.data.site_xpos[self.robots[0].eef_site_id]) - mean_distance_to_things_to_wipe += np.linalg.norm(gripper_position - sensor_pos) - num_non_wiped_sensors += 1 - mean_distance_to_things_to_wipe /= max(1, num_non_wiped_sensors) + num_non_wiped_markers = 0 + for marker in self.model.mujoco_arena.markers: + if marker not in self.wiped_markers: + marker_pos = np.array( + self.sim.data.body_xpos[self.sim.model.body_name2id(marker.root_body)]) + mean_distance_to_things_to_wipe += np.linalg.norm(marker_pos - self._eef_xpos) + num_non_wiped_markers += 1 + mean_distance_to_things_to_wipe /= max(1, num_non_wiped_markers) reward += self.distance_multiplier * ( 1 - np.tanh(self.distance_th_multiplier * mean_distance_to_things_to_wipe)) @@ -465,24 +453,25 @@ def PointInRectangle(X, Y, Z, W, P): reward -= self.ee_accel_penalty * np.mean(abs(self.robots[0].recent_ee_acc.current)) # Final reward if all wiped - if len(self.wiped_sensors) == len(self.model.arena.sensor_names): + if len(self.wiped_markers) == self.num_markers: reward += self.task_complete_reward # Printing results if self.print_results: - string_to_print = 'Process {pid}, timestep {ts:>4}: reward: {rw:8.4f} wiped sensors: {ws:>3} collisions: {sc:>3} f-excess: {fe:>3}'.format( - pid=id(multiprocessing.current_process()), - ts=self.timestep, - rw=reward, - ws=len(self.wiped_sensors), - sc=self.collisions, - fe=self.f_excess) + string_to_print = 'Process {pid}, timestep {ts:>4}: reward: {rw:8.4f}' \ + 'wiped markers: {ws:>3} collisions: {sc:>3} f-excess: {fe:>3}'.format( + pid=id(multiprocessing.current_process()), + ts=self.timestep, + rw=reward, + ws=len(self.wiped_markers), + sc=self.collisions, + fe=self.f_excess) print(string_to_print) # If we're scaling our reward, we normalize the per-step rewards given the theoretical best episode return # This is equivalent to scaling the reward by: # reward_scale * (horizon / - # (num_sensors * unit_wiped_reward + horizon * (wipe_contact_reward + task_complete_reward))) + # (num_markers * unit_wiped_reward + horizon * (wipe_contact_reward + task_complete_reward))) if self.reward_scale: reward *= self.reward_scale * self.reward_normalization_factor return reward @@ -493,10 +482,6 @@ def _load_model(self): """ super()._load_model() - # Verify the correct robot has been loaded - assert isinstance(self.robots[0], SingleArm), \ - "Error: Expected one single-armed robot! Got {} type instead.".format(type(self.robots[0])) - # Adjust base pose accordingly xpos = self.robots[0].robot_model.base_xpos_offset["table"](self.table_full_size[0]) self.robots[0].robot_model.set_base_xpos(xpos) @@ -504,47 +489,38 @@ def _load_model(self): # Get robot's contact geoms self.robot_contact_geoms = self.robots[0].robot_model.contact_geoms - # Delta goes down - delta_height = min(0, np.random.normal(self.table_height, self.table_height_std)) - - self.mujoco_arena = WipeArena( + mujoco_arena = WipeArena( table_full_size=self.table_full_size, table_friction=self.table_friction, - table_offset=np.array(self.table_offset) + np.array((0, 0, delta_height)), + table_offset=self.table_offset, table_friction_std=self.table_friction_std, coverage_factor=self.coverage_factor, - num_sensors=self.num_sensors, + num_markers=self.num_markers, line_width=self.line_width, two_clusters=self.two_clusters ) if self.use_indicator_object: - self.mujoco_arena.add_pos_indicator() + mujoco_arena.add_pos_indicator() # Arena always gets set to zero origin - self.mujoco_arena.set_origin([0, 0, 0]) - - self.mujoco_objects = OrderedDict() + mujoco_arena.set_origin([0, 0, 0]) # task includes arena, robot, and objects of interest self.model = ManipulationTask( - mujoco_arena=self.mujoco_arena, + mujoco_arena=mujoco_arena, mujoco_robots=[robot.robot_model for robot in self.robots], - mujoco_objects=self.mujoco_objects, - initializer=self.placement_initializer ) - self.model.place_objects() def _reset_internal(self): super()._reset_internal() # inherited class should reset positions of objects (only if we're not using a deterministic reset) if not self.deterministic_reset: - self.model.place_objects() - self.mujoco_arena.reset_arena(self.sim) + self.model.mujoco_arena.reset_arena(self.sim) # Reset all internal vars for this wipe task self.timestep = 0 - self.wiped_sensors = [] + self.wiped_markers = [] self.collisions = 0 self.f_excess = 0 @@ -581,23 +557,19 @@ def _get_observation(self): # object information in the observation if self.use_object_obs: - gripper_site_pos = np.array(self.sim.data.site_xpos[self.robots[0].eef_site_id]) + gripper_site_pos = self._eef_xpos # position of objects to wipe acc = np.array([]) - for sensor_name in self.model.arena.sensor_names: - parts = sensor_name.split('_') - sensor_id = int(parts[1]) - sensor_pos = np.array( - self.sim.data.site_xpos[ - self.sim.model.site_name2id(self.model.arena.sensor_site_names[sensor_name])]) - di['sensor' + str(sensor_id) + '_pos'] = sensor_pos - acc = np.concatenate([acc, di['sensor' + str(sensor_id) + '_pos']]) - di['sensor' + str(sensor_id) + '_wiped'] = [0, 1][sensor_id in self.wiped_sensors] - acc = np.concatenate([acc, [di['sensor' + str(sensor_id) + '_wiped']]]) + for i, marker in enumerate(self.model.mujoco_arena.markers): + marker_pos = np.array(self.sim.data.body_xpos[self.sim.model.body_name2id(marker.root_body)]) + di[f'marker{i}_pos'] = marker_pos + acc = np.concatenate([acc, di[f'marker{i}_pos']]) + di[f'marker{i}_wiped'] = [0, 1][marker in self.wiped_markers] + acc = np.concatenate([acc, [di[f'marker{i}_wiped']]]) # proprioception if self.use_robot_obs: - di['gripper_to_sensor' + str(sensor_id)] = gripper_site_pos - sensor_pos - acc = np.concatenate([acc, di['gripper_to_sensor' + str(sensor_id)]]) + di[f'gripper_to_marker{i}'] = gripper_site_pos - marker_pos + acc = np.concatenate([acc, di[f'gripper_to_marker{i}']]) di['object-state'] = acc return di @@ -609,7 +581,7 @@ def _check_success(self): Returns: bool: True if completed task """ - return True if len(self.wiped_sensors) == len(self.model.arena.sensor_names) else False + return True if len(self.wiped_markers) == self.num_markers else False def _check_terminated(self): """ @@ -626,7 +598,7 @@ def _check_terminated(self): terminated = False # Prematurely terminate if contacting the table with the arm - if self._check_arm_contact()[0]: + if self.check_contact(self.robots[0].robot_model): if self.print_results: print(40 * '-' + " COLLIDED " + 40 * '-') terminated = True @@ -638,7 +610,7 @@ def _check_terminated(self): terminated = True # Prematurely terminate if contacting the table with the arm - if self._check_q_limits()[0]: + if self.robots[0].check_q_limits(): if self.print_results: print(40 * '-' + " JOINT LIMIT " + 40 * '-') terminated = True @@ -662,10 +634,10 @@ def _post_action(self, action): reward, done, info = super()._post_action(action) if self.get_info: - info['add_vals'] = ['nwipedsensors', 'colls', 'percent_viapoints_', 'f_excess'] - info['nwipedsensors'] = len(self.wiped_sensors) + info['add_vals'] = ['nwipedmarkers', 'colls', 'percent_viapoints_', 'f_excess'] + info['nwipedmarkers'] = len(self.wiped_markers) info['colls'] = self.collisions - info['percent_viapoints_'] = len(self.wiped_sensors) / len(self.model.arena.sensor_names) + info['percent_viapoints_'] = len(self.wiped_markers) / self.num_markers info['f_excess'] = self.f_excess # allow episode to finish early if allowed @@ -674,16 +646,6 @@ def _post_action(self, action): return reward, done, info - def _check_robot_configuration(self, robots): - """ - Sanity check to make sure the inputted robots and configuration is acceptable - - Args: - robots (str or list of str): Robots to instantiate within this env - """ - if type(robots) is list: - assert len(robots) == 1, "Error: Only one robot should be inputted for this task!" - @property def _has_gripper_contact(self): """ diff --git a/robosuite/environments/robot_env.py b/robosuite/environments/robot_env.py index b5f7274ec5..7b2d07a42f 100644 --- a/robosuite/environments/robot_env.py +++ b/robosuite/environments/robot_env.py @@ -2,10 +2,7 @@ from robosuite.environments.base import MujocoEnv -from robosuite.robots.single_arm import SingleArm -from robosuite.robots.bimanual import Bimanual -from robosuite.models.robots import check_bimanual - +from robosuite.robots import ROBOT_CLASS_MAPPING from robosuite.controllers import reset_controllers @@ -14,24 +11,16 @@ class RobotEnv(MujocoEnv): Initializes a robot environment in Mujoco. Args: - robots: Specification for specific robot arm(s) to be instantiated within this env - (e.g: "Sawyer" would generate one arm; ["Panda", "Panda", "Sawyer"] would generate three robot arms) + robots: Specification for specific robot(s) to be instantiated within this env + + env_configuration (str): Specifies how to position the robot(s) within the environment. Default is "default", + which should be interpreted accordingly by any subclasses. controller_configs (str or list of dict): If set, contains relevant controller parameters for creating a custom controller. Else, uses the default controller for this specific task. Should either be single dict if same controller is to be used for all robots or else it should be a list of the same length as "robots" param - gripper_types (None or str or list of str): type of gripper, used to instantiate - gripper models from gripper factory. Default is "default", which is the default grippers(s) associated - with the robot(s) the 'robots' specification. None removes the gripper, and any other (valid) model - overrides the default gripper. Should either be single str if same gripper type is to be used for all - robots or else it should be a list of the same length as "robots" param - - gripper_visualizations (bool or list of bool): True if using gripper visualization. - Useful for teleoperation. Should either be single bool if gripper visualization is to be used for all - robots or else it should be a list of the same length as "robots" param - initialization_noise (dict or list of dict): Dict containing the initialization noise parameters. The expected keys and corresponding value types are specified below: @@ -52,6 +41,13 @@ class RobotEnv(MujocoEnv): use_indicator_object (bool): if True, sets up an indicator object that is useful for debugging. + robot_visualizations (bool or list of bool): True if using robot visualization. + Useful for teleoperation. Should either be single bool if robot visualization is to be used for all + robots or else it should be a list of the same length as "robots" param + + env_visualization (bool): True if visualizing sites for the arena / objects in this environment. Useful for + teleoperation. + has_renderer (bool): If true, render the simulation state in a viewer instead of headless mode. @@ -96,6 +92,8 @@ class RobotEnv(MujocoEnv): bool if same depth setting is to be used for all cameras or else it should be a list of the same length as "camera names" param. + robot_configs (list of dict): Per-robot configurations set from any subclass initializers. + Raises: ValueError: [Camera obs require offscreen renderer] ValueError: [Camera name must be specified to use camera obs] @@ -104,12 +102,13 @@ class RobotEnv(MujocoEnv): def __init__( self, robots, + env_configuration="default", controller_configs=None, - gripper_types="default", - gripper_visualizations=False, initialization_noise=None, use_camera_obs=True, use_indicator_object=False, + robot_visualizations=False, + env_visualization=False, has_renderer=False, has_offscreen_renderer=True, render_camera="frontview", @@ -123,7 +122,12 @@ def __init__( camera_heights=256, camera_widths=256, camera_depths=False, + robot_configs=None, ): + # First, verify that correct number of robots are being inputted + self.env_configuration = env_configuration + self._check_robot_configuration(robots) + # Robot robots = list(robots) if type(robots) is list or type(robots) is tuple else [robots] self.num_robots = len(robots) @@ -137,9 +141,8 @@ def __init__( # Initialization Noise initialization_noise = self._input2list(initialization_noise, self.num_robots) - # Gripper - gripper_types = self._input2list(gripper_types, self.num_robots) - gripper_visualizations = self._input2list(gripper_visualizations, self.num_robots) + # Visualization + robot_visualizations = self._input2list(robot_visualizations, self.num_robots) # Observations -- Ground truth = object_obs, Image data = camera_obs self.use_camera_obs = use_camera_obs @@ -160,23 +163,26 @@ def __init__( if self.use_camera_obs and self.camera_names is None: raise ValueError("Must specify at least one camera name when using camera obs") - # Robot configurations + # Robot configurations -- update from subclass configs + if robot_configs is None: + robot_configs = [{} for _ in range(self.num_robots)] self.robot_configs = [ - { - "controller_config": controller_configs[idx], - "initialization_noise": initialization_noise[idx], - "gripper_type": gripper_types[idx], - "gripper_visualization": gripper_visualizations[idx], - "control_freq": control_freq - } - for idx in range(self.num_robots) + dict( + **{ + "controller_config": controller_configs[idx], + "initialization_noise": initialization_noise[idx], + "robot_visualization": robot_visualizations[idx], + "control_freq": control_freq + }, + **robot_config, + ) + for idx, robot_config in enumerate(robot_configs) ] - # whether to use indicator object or not - self.use_indicator_object = use_indicator_object - # Run superclass init super().__init__( + use_indicator_object=use_indicator_object, + env_visualization=env_visualization, has_renderer=has_renderer, has_offscreen_renderer=self.has_offscreen_renderer, render_camera=render_camera, @@ -215,17 +221,6 @@ def action_dim(self): """ return self._action_dim - def move_indicator(self, pos): - """ - Sets 3d position of indicator object to @pos. - - Args: - pos (3-tuple): (x,y,z) values to place the indicator within the env - """ - if self.use_indicator_object: - index = self._ref_indicator_pos_low - self.sim.data.qpos[index : index + 3] = pos - @staticmethod def _input2list(inp, length): """ @@ -263,16 +258,6 @@ def _get_reference(self): robot.reset_sim(self.sim) robot.setup_references() - # Indicator object references - if self.use_indicator_object: - ind_qpos = self.sim.model.get_joint_qpos_addr("pos_indicator") - self._ref_indicator_pos_low, self._ref_indicator_pos_high = ind_qpos - - ind_qvel = self.sim.model.get_joint_qvel_addr("pos_indicator") - self._ref_indicator_vel_low, self._ref_indicator_vel_high = ind_qvel - - self.indicator_id = self.sim.model.body_name2id("pos_indicator") - def _reset_internal(self): """ Resets simulation internal configurations. @@ -356,25 +341,6 @@ def _pre_action(self, action, policy_step=False): ] = self.sim.data.qfrc_bias[ self._ref_indicator_vel_low: self._ref_indicator_vel_high] - def _post_action(self, action): - """ - Run any necessary visualization after running the action - - Args: - action (np.array): Action being passed during this timestep - - Returns: - 3-tuple: - - - (float) reward from the environment - - (bool) whether the current episode is completed or not - - (dict) empty dict to be filled with information by subclassed method - - """ - ret = super()._post_action(action) - self._visualization() - return ret - def _get_observation(self): """ Returns an OrderedDict containing observations [(name_string, np.array), ...]. @@ -416,92 +382,15 @@ def _get_observation(self): return di - def _check_gripper_contact(self): - """ - Checks whether each gripper is in contact with an object. Note that each entry in the return list corresponds - to each robot in the env.robots list. This means that a bimanual robot will have its own nested list within the - return list at the corresponding entry. - - Returns: - list of bool: True if the specific gripper is in contact with an object. - """ - collisions = [False] * self.num_robots - for idx, robot in enumerate(self.robots): - for contact in self.sim.data.contact[: self.sim.data.ncon]: - # Single arm case - if robot.robot_model.arm_type == "single": - if ( - self.sim.model.geom_id2name(contact.geom1) - in robot.gripper.contact_geoms - or self.sim.model.geom_id2name(contact.geom2) - in robot.gripper.contact_geoms - ): - collisions[idx] = True - break - # Bimanual case - else: - # Set this list entry to be its own list of bools - collisions[idx] = [False, False] - for i, arm in enumerate(robot.arms): - if ( - self.sim.model.geom_id2name(contact.geom1) - in robot.gripper[arm].contact_geoms - or self.sim.model.geom_id2name(contact.geom2) - in robot.gripper[arm].contact_geoms - ): - collisions[idx][i] = True - continue - return collisions - - def _check_arm_contact(self): - """ - Checks whether each robot arm is in contact with an object. - - Returns: - list of bool: True if the specific gripper is in contact with an object - """ - collisions = [False] * self.num_robots - for idx, robot in enumerate(self.robots): - for contact in self.sim.data.contact[: self.sim.data.ncon]: - # Single arm case and Bimanual case are the same - if ( - self.sim.model.geom_id2name(contact.geom1) - in robot.robot_model.contact_geoms - or self.sim.model.geom_id2name(contact.geom2) - in robot.robot_model.contact_geoms - ): - collisions[idx] = True - break - return collisions - - def _check_q_limits(self): - """ - Check if each robot arm is either very close or at the joint limits - - Returns: - list of bool: True if the specific arm is near its joint limits - """ - joint_limits = [False] * self.num_robots - tolerance = 0.1 - for idx, robot in enumerate(self.robots): - for (qidx, (q, q_limits)) in enumerate( - zip( - self.sim.data.qpos[robot._ref_joint_pos_indexes], - self.sim.model.jnt_range[robot._ref_joint_indexes] - ) - ): - if q_limits[0] != q_limits[1] and not (q_limits[0] + tolerance < q < q_limits[1] - tolerance): - print("Joint limit reached in joint " + str(qidx)) - joint_limits[idx] = True - return joint_limits - def _visualization(self): """ Do any needed visualization here """ - # Loop over robot grippers to visualize them independently + # Run superclass method first + super()._visualization() + # Loop over robots to visualize them independently for robot in self.robots: - robot.visualize_gripper() + robot.visualize() def _load_robots(self): """ @@ -510,19 +399,7 @@ def _load_robots(self): # Loop through robots and instantiate Robot object for each for idx, (name, config) in enumerate(zip(self.robot_names, self.robot_configs)): # Create the robot instance - if not check_bimanual(name): - self.robots[idx] = SingleArm( - robot_type=name, - idn=idx, - **config - ) - else: - self.robots[idx] = Bimanual( - robot_type=name, - idn=idx, - **config - ) - + self.robots[idx] = ROBOT_CLASS_MAPPING[name](robot_type=name, idn=idx, **config) # Now, load the robot models self.robots[idx].load_model() diff --git a/robosuite/models/arenas/wipe_arena.py b/robosuite/models/arenas/wipe_arena.py index 92b65f227b..61f63271e2 100644 --- a/robosuite/models/arenas/wipe_arena.py +++ b/robosuite/models/arenas/wipe_arena.py @@ -1,12 +1,12 @@ import numpy as np from robosuite.models.arenas import TableArena -from robosuite.utils.mjcf_utils import array_to_string, CustomMaterial +from robosuite.utils.mjcf_utils import CustomMaterial, find_elements from robosuite.models.objects import CylinderObject class WipeArena(TableArena): """ - Workspace that contains an empty table with tactile sensors on its surface. + Workspace that contains an empty table with visual markers on its surface. Args: table_full_size (3-tuple): (L,W,H) full dimensions of the table @@ -14,7 +14,7 @@ class WipeArena(TableArena): table_offset (3-tuple): (x,y,z) offset from center of arena when placing table. Note that the z value sets the upper limit of the table coverage_factor (float): Fraction of table that will be sampled for dirt placement - num_sensors (int): Number of dirt (peg) particles to generate in a path on the table + num_markers (int): Number of dirt (peg) particles to generate in a path on the table table_friction_std (float): Standard deviation to sample for the peg friction line_width (float): Diameter of dirt path trace two_clusters (bool): If set, will generate two separate dirt paths with half the number of sensors in each @@ -26,7 +26,7 @@ def __init__( table_friction=(0.01, 0.005, 0.0001), table_offset=(0, 0, 0.8), coverage_factor=0.9, - num_sensors=10, + num_markers=10, table_friction_std=0, line_width=0.02, two_clusters=False @@ -34,10 +34,9 @@ def __init__( # Tactile table-specific features self.table_friction_std = table_friction_std self.line_width = line_width - self.sensor_names = [] - self.sensor_site_names = {} + self.markers = [] self.coverage_factor = coverage_factor - self.num_sensors = num_sensors + self.num_markers = num_markers self.two_clusters = two_clusters # Attribute to hold current direction of sampled dirt path @@ -55,12 +54,6 @@ def configure_location(self): # Run superclass first super().configure_location() - # Determine peg friction - friction = max(0.001, np.random.normal(self.table_friction[0], self.table_friction_std)) - - # Grab reference to the table body in the xml - table_subtree = self.worldbody.find(".//body[@name='{}']".format("table")) - # Define start position for drawing the line pos = self.sample_start_pos() @@ -82,61 +75,59 @@ def configure_location(self): ) # Define line(s) drawn on table - for i in range(self.num_sensors): + for i in range(self.num_markers): # If we're using two clusters, we resample the starting position and direction at the halfway point - if self.two_clusters and i == int(np.floor(self.num_sensors / 2)): + if self.two_clusters and i == int(np.floor(self.num_markers / 2)): pos = self.sample_start_pos() - square_name2 = 'contact_'+str(i) - square2 = CylinderObject( - name=square_name2, + marker_name = f'contact{i}' + marker = CylinderObject( + name=marker_name, size=[self.line_width / 2, 0.001], rgba=[1, 1, 1, 1], - density=1, material=dirt, - friction=friction, obj_type="visual", + joints=None, ) - self.merge_asset(square2) - visual_c = square2.get_object_subtree(site=True) - visual_c.set("pos", array_to_string([pos[0], pos[1], self.table_half_size[2]])) - visual_c.find("site").set("pos", [0, 0, 0.005]) - visual_c.find("site").set("rgba", array_to_string([0, 0, 0, 0])) - table_subtree.append(visual_c) - - sensor_name = square_name2 + "_sensor" - sensor_site_name = square_name2 - self.sensor_names += [sensor_name] - self.sensor_site_names[sensor_name] = sensor_site_name + # Manually add this object to the arena xml + self.merge_asset(marker) + table = find_elements(root=self.worldbody, tags="body", attribs={"name": "table"}, return_first=True) + table.append(marker.get_obj()) + + # Add this marker to our saved list of all markers + self.markers.append(marker) # Add to the current dirt path pos = self.sample_path_pos(pos) def reset_arena(self, sim): """ - Reset the tactile sensor locations in the environment. Requires @sim (MjSim) reference to be passed in so that + Reset the visual marker locations in the environment. Requires @sim (MjSim) reference to be passed in so that the Mujoco sim can be directly modified Args: - sim (MjSim): Simulation instance containing this arena and tactile sensors + sim (MjSim): Simulation instance containing this arena and visual markers """ - # Sample new initial position and direction for generated sensor paths + # Sample new initial position and direction for generated marker paths pos = self.sample_start_pos() - # Loop through all sensor collision body / site pairs - for i, (_, sensor_name) in enumerate(self.sensor_site_names.items()): + # Loop through all visual markers + for i, marker in enumerate(self.markers): # If we're using two clusters, we resample the starting position and direction at the halfway point - if self.two_clusters and i == int(np.floor(self.num_sensors / 2)): + if self.two_clusters and i == int(np.floor(self.num_markers / 2)): pos = self.sample_start_pos() - # Get IDs to the body, geom, and site of each sensor - body_id = sim.model.body_name2id(sensor_name) - geom_id = sim.model.geom_name2id(sensor_name + "_vis") - # Determine new position for this sensor + # Get IDs to the body, geom, and site of each marker + body_id = sim.model.body_name2id(marker.root_body) + geom_id = sim.model.geom_name2id(marker.visual_geoms[0]) + site_id = sim.model.site_name2id(marker.sites[0]) + # Determine new position for this marker position = np.array([pos[0], pos[1], self.table_half_size[2]]) - # Set the current sensor (body) to this new position + # Set the current marker (body) to this new position sim.model.body_pos[body_id] = position - # Reset the sensor visualization -- setting geom rgba to all 1's - sim.model.geom_rgba[geom_id] = [1, 1, 1, 1] - # Sample next values in local sensor trajectory + # Reset the marker visualization -- setting geom rgba alpha value to 1 + sim.model.geom_rgba[geom_id][3] = 1 + # Hide the default visualization site + sim.model.site_rgba[site_id][3] = 0 + # Sample next values in local marker trajectory pos = self.sample_path_pos(pos) def sample_start_pos(self): diff --git a/robosuite/models/assets/base.xml b/robosuite/models/assets/base.xml index 7f863e6062..3de7de2565 100644 --- a/robosuite/models/assets/base.xml +++ b/robosuite/models/assets/base.xml @@ -1,6 +1,6 @@ - +