Merge pull request #638 from ARISE-Initiative/seg-bodies-in-arena

Enable instance segmentation to include bodies in arena

Author: kevin-thankyou-lin

robosuite/demos/demo_segmentation.py

@@ -57,18 +57,17 @@ def segmentation_to_rgb(seg_im, random_colors=False):
     parser.add_argument("--video-path", type=str, default="/tmp/video.mp4", help="Path to video file")
     parser.add_argument("--random-colors", action="store_true", help="Radnomize segmentation colors")
     parser.add_argument("--segmentation-level", type=str, default="element", help="instance, class, or element")
+    parser.add_argument("--env-name", type=str, default="TwoArmHandover", help="Environment name")
+    parser.add_argument("--camera", type=str, default="frontview", help="Camera name")
     args = parser.parse_args()
 
     # Create dict to hold options that will be passed to env creation call
     options = {}
 
     # Choose environment and add it to options
-    options["env_name"] = "TwoArmHandover"
+    options["env_name"] = args.env_name
     options["robots"] = ["Panda", "Panda"]
 
-    # Choose camera
-    camera = "frontview"
-
     # Choose segmentation type
     segmentation_level = args.segmentation_level  # Options are {instance, class, element}
 
@@ -80,7 +79,7 @@ def segmentation_to_rgb(seg_im, random_colors=False):
         ignore_done=True,
         use_camera_obs=True,
         control_freq=20,
-        camera_names=camera,
+        camera_names=args.camera,
         camera_segmentations=segmentation_level,
         camera_heights=512,
         camera_widths=512,
@@ -98,7 +97,7 @@ def segmentation_to_rgb(seg_im, random_colors=False):
         action = 0.5 * np.random.uniform(low, high)
         obs, reward, done, _ = env.step(action)
 
-        video_img = obs[f"{camera}_segmentation_{segmentation_level}"].squeeze(-1)[::-1]
+        video_img = obs[f"{args.camera}_segmentation_{segmentation_level}"].squeeze(-1)[::-1]
         np.savetxt("/tmp/seg_{}.txt".format(i), video_img, fmt="%.2f")
         video_img = segmentation_to_rgb(video_img, args.random_colors)
         video_writer.append_data(video_img)

robosuite/models/tasks/task.py

@@ -1,11 +1,41 @@
+import xml.etree.ElementTree as ET
 from copy import deepcopy
 
+import mujoco
+
 from robosuite.models.objects import MujocoObject
 from robosuite.models.robots import RobotModel
 from robosuite.models.world import MujocoWorldBase
 from robosuite.utils.mjcf_utils import get_ids
 
 
+def get_subtree_geom_ids_by_group(model: mujoco.MjModel, body_id: int, group: int) -> list[int]:
+    """Get all geoms belonging to a subtree starting at a given body, filtered by group.
+
+    Args:
+        model: MuJoCo model.
+        body_id: ID of body where subtree starts.
+        group: Group ID to filter geoms.
+
+    Returns:
+        A list containing all subtree geom ids in the specified group.
+
+    Adapted from https://github.com/kevinzakka/mink/blob/main/mink/utils.py
+    """
+
+    def gather_geoms(body_id: int) -> list[int]:
+        geoms: list[int] = []
+        geom_start = model.body_geomadr[body_id]
+        geom_end = geom_start + model.body_geomnum[body_id]
+        geoms.extend(geom_id for geom_id in range(geom_start, geom_end) if model.geom_group[geom_id] == group)
+        children = [i for i in range(model.nbody) if model.body_parentid[i] == body_id]
+        for child_id in children:
+            geoms.extend(gather_geoms(child_id))
+        return geoms
+
+    return gather_geoms(body_id)
+
+
 class Task(MujocoWorldBase):
     """
     Creates MJCF model for a task performed.
@@ -106,15 +136,32 @@ def generate_id_mappings(self, sim):
         for robot in self.mujoco_robots:
             models += [robot] + robot.models
 
+        worldbody = self.mujoco_arena.root.find("worldbody")
+        exclude_bodies = ["table"]
+        top_level_bodies = [
+            body.attrib.get("name")
+            for body in worldbody.findall("body")
+            if body.attrib.get("name") not in exclude_bodies
+        ]
+        models.extend(top_level_bodies)
+
         # Parse all mujoco models from robots and objects
         for model in models:
-            # Grab model class name and visual IDs
-            cls = str(type(model)).split("'")[1].split(".")[-1]
-            inst = model.name
-            id_groups = [
-                get_ids(sim=sim, elements=model.visual_geoms + model.contact_geoms, element_type="geom"),
-                get_ids(sim=sim, elements=model.sites, element_type="site"),
-            ]
+            if isinstance(model, str):
+                body_name = model
+                visual_group_number = 1
+                body_id = sim.model.body_name2id(body_name)
+                inst, cls = body_name, body_name
+                geom_ids = get_subtree_geom_ids_by_group(sim.model, body_id, visual_group_number)
+                id_groups = [geom_ids, []]
+            else:
+                # Grab model class name and visual IDs
+                cls = str(type(model)).split("'")[1].split(".")[-1]
+                inst = model.name
+                id_groups = [
+                    get_ids(sim=sim, elements=model.visual_geoms + model.contact_geoms, element_type="geom"),
+                    get_ids(sim=sim, elements=model.sites, element_type="site"),
+                ]
             group_types = ("geom", "site")
             ids_to_instances = (self._geom_ids_to_instances, self._site_ids_to_instances)
             ids_to_classes = (self._geom_ids_to_classes, self._site_ids_to_classes)