Refactor scene mask to scene idxs (#242)

* move out scene masks in favor of scene idxs

* Update link.py

mani_skill/envs/tasks/assembling_kits.py

@@ -92,7 +92,7 @@ def _load_scene(self):
             self.goal_rot = np.zeros((self.num_envs,))
 
             for i, eps_idx in enumerate(eps_idxs):
-                scene_mask = [i]
+                scene_idxs = [i]
                 episode = self._episodes[eps_idx]
 
                 # get the kit builder and the goal positions/rotations of all other objects
@@ -102,15 +102,15 @@ def _load_scene(self):
                     object_goal_rot,
                 ) = self._get_kit_builder_and_goals(episode["kit"])
                 kit = (
-                    kit_builder.set_scene_idxs(scene_mask)
+                    kit_builder.set_scene_idxs(scene_idxs)
                     .set_initial_pose(sapien.Pose([0, 0, 0.01]))
                     .build_static(f"kit_{i}")
                 )
                 kits.append(kit)
                 # create the object to place and make it dynamic
                 obj_to_place = (
                     self._get_object_builder(episode["obj_to_place"])
-                    .set_scene_idxs(scene_mask)
+                    .set_scene_idxs(scene_idxs)
                     .build(f"obj_{i}")
                 )
                 self.object_ids.append(episode["obj_to_place"])
@@ -119,7 +119,7 @@ def _load_scene(self):
                 # create all other objects and leave them as static as they do not need to be manipulated
                 other_objs = [
                     self._get_object_builder(obj_id, static=True)
-                    .set_scene_idxs(scene_mask)
+                    .set_scene_idxs(scene_idxs)
                     .set_initial_pose(
                         sapien.Pose(
                             object_goal_pos[obj_id],

mani_skill/envs/tasks/dexterity/rotate_single_object_in_hand.py

@@ -110,9 +110,8 @@ def _load_scene(self):
                         base_color=np.array([255, 255, 255, 255]) / 255,
                     ),
                 )
-                scene_mask = np.zeros(self.num_envs, dtype=bool)
-                scene_mask[i] = True
-                builder.set_scene_mask(scene_mask)
+                scene_idxs = [i]
+                builder.set_scene_idxs(scene_idxs)
                 actors.append(builder.build(name=f"cube-{i}"))
                 obj_heights.append(half_size)
             self.obj = Actor.merge(actors, name="cube")
@@ -128,9 +127,8 @@ def _load_scene(self):
                 builder, obj_height = build_actor_ycb(
                     model_id, self._scene, name=model_id, return_builder=True
                 )
-                scene_mask = np.zeros(self.num_envs, dtype=bool)
-                scene_mask[i] = True
-                builder.set_scene_mask(scene_mask)
+                scene_idxs = [i]
+                builder.set_scene_idxs(scene_idxs)
                 actors.append(builder.build(name=f"{model_id}-{i}"))
                 obj_heights.append(obj_height)
             self.obj = Actor.merge(actors, name="ycb_object")

mani_skill/envs/tasks/dexterity/rotate_valve.py

@@ -110,21 +110,20 @@ def _load_articulations(self):
         valves: List[Articulation] = []
         capsule_lens = []
         for i, valve_angles in enumerate(valve_angles_list):
-            scene_mask = np.zeros(self.num_envs, dtype=bool)
-            scene_mask[i] = True
+            scene_idxs = [i]
             if self.difficulty_level < 3:
                 valve, capsule_len = build_robel_valve(
                     self._scene,
                     valve_angles=valve_angles,
-                    scene_mask=scene_mask,
+                    scene_idxs=scene_idxs,
                     name=f"valve_station_{i}",
                 )
             else:
                 scales = self._main_rng.randn(2) * 0.1 + 1
                 valve, capsule_len = build_robel_valve(
                     self._scene,
                     valve_angles=valve_angles,
-                    scene_mask=scene_mask,
+                    scene_idxs=scene_idxs,
                     name=f"valve_station_{i}",
                     radius_scale=scales[0],
                     capsule_radius_scale=scales[1],

mani_skill/envs/tasks/open_cabinet_drawer.py

@@ -83,7 +83,7 @@ def _load_cabinets(self, joint_types: List[str]):
             scene_mask = np.zeros(self.num_envs, dtype=bool)
             scene_mask[i] = True
             cabinet, metadata = build_preprocessed_partnet_mobility_articulation(
-                self._scene, model_id, name=f"{model_id}-{i}", scene_mask=scene_mask
+                self._scene, model_id, name=f"{model_id}-{i}", scene_idxs=scene_mask
             )
             # TODO (stao): since we processed the assets we know that the bounds[0, 1] is the actual height to set the object at
             # but in future we will store a visual origin offset so we can place them by using the actual bbox height / 2

mani_skill/envs/tasks/peg_insertion_side.py

@@ -112,8 +112,7 @@ def _load_scene(self):
             boxes = []
 
             for i in range(self.num_envs):
-                scene_mask = np.zeros((self.num_envs), dtype=bool)
-                scene_mask[i] = True
+                scene_idxs = [i]
                 length = lengths[i]
                 radius = radii[i]
                 builder = self._scene.create_actor_builder()
@@ -140,7 +139,7 @@ def _load_scene(self):
                     half_size=[length / 2, radius, radius],
                     material=mat,
                 )
-                builder.set_scene_mask(scene_mask)
+                builder.set_scene_idxs(scene_idxs)
                 peg = builder.build(f"peg_{i}")
 
                 # box with hole
@@ -153,7 +152,7 @@ def _load_scene(self):
                 builder = _build_box_with_hole(
                     self._scene, inner_radius, outer_radius, depth, center=centers[i]
                 )
-                builder.set_scene_mask(scene_mask)
+                builder.set_scene_idxs(scene_idxs)
                 box = builder.build_kinematic(f"box_with_hole_{i}")
 
                 pegs.append(peg)

mani_skill/utils/building/actor_builder.py

@@ -28,23 +28,10 @@ class ActorBuilder(SAPIENActorBuilder):
     def __init__(self):
         super().__init__()
         self.initial_pose = Pose.create(self.initial_pose)
-        self.scene_mask = None
         self.scene_idxs = None
         self._allow_overlapping_plane_collisions = False
         self._plane_collision_poses = set()
 
-    def set_scene_mask(
-        self,
-        scene_mask: Optional[
-            Union[List[bool], Sequence[bool], torch.Tensor, np.ndarray]
-        ] = None,
-    ):
-        """
-        Set a scene mask so that the actor builder builds the actor only in a subset of the environments
-        """
-        self.scene_mask = scene_mask
-        return self
-
     def set_scene_idxs(
         self,
         scene_idxs: Optional[
@@ -187,16 +174,11 @@ def build(self, name):
         self.set_name(name)
 
         num_actors = self.scene.num_envs
-        if self.scene_mask is not None:
-            num_actors = np.sum(num_actors)
-            self.scene_mask = sapien_utils.to_tensor(self.scene_mask).to(bool)
-            self.scene_idxs = self.scene_mask.argwhere()
-        elif self.scene_idxs is not None:
-            self.scene_mask = torch.zeros((self.scene.num_envs), dtype=bool)
-            self.scene_mask[self.scene_idxs] = True
+        if self.scene_idxs is not None:
+            pass
         else:
-            self.scene_mask = torch.ones((self.scene.num_envs), dtype=bool)
             self.scene_idxs = torch.arange((self.scene.num_envs), dtype=int)
+        num_actors = len(self.scene_idxs)
         initial_pose = Pose.create(self.initial_pose)
         initial_pose_b = initial_pose.raw_pose.shape[0]
         assert initial_pose_b == 1 or initial_pose_b == num_actors
@@ -217,7 +199,7 @@ def build(self, name):
             sub_scene.add_entity(entity)
             entities.append(entity)
             i += 1
-        actor = Actor._create_from_entities(entities, self.scene, self.scene_mask)
+        actor = Actor.create_from_entities(entities, self.scene, self.scene_idxs)
 
         # if it is a static body type and this is a GPU sim but we are given a single initial pose, we repeat it for the purposes of observations
         if (

mani_skill/utils/building/articulation_builder.py

@@ -25,23 +25,24 @@ class ArticulationBuilder(SapienArticulationBuilder):
 
     def __init__(self):
         super().__init__()
-        self.scene_mask = None
         self.name = None
+        self.scene_idxs = None
 
     def set_name(self, name: str):
         self.name = name
         return self
 
-    def set_scene_mask(
+    def set_scene_idxs(
         self,
-        scene_mask: Optional[
-            Union[List[bool], Sequence[bool], torch.Tensor, np.ndarray]
+        scene_idxs: Optional[
+            Union[List[int], Sequence[int], torch.Tensor, np.ndarray]
         ] = None,
     ):
         """
-        Set a scene mask so that the articulation builder builds the articulation only in a subset of the environments
+        Set a list of scene indices to build this object in. Cannot be used in conjunction with scene mask
         """
-        self.scene_mask = scene_mask
+        self.scene_idxs = scene_idxs
+        return self
 
     def create_link_builder(self, parent: LinkBuilder = None):
         if self.link_builders:
@@ -105,25 +106,15 @@ def build(self, name=None, fix_root_link=None, build_mimic_joints=True):
         if name is not None:
             self.set_name(name)
         assert self.name is not None
-
-        num_arts = self.scene.num_envs
-        if self.scene_mask is not None:
-            assert (
-                len(self.scene_mask) == self.scene.num_envs
-            ), "Scene mask size is not correct. Must be the same as the number of sub scenes"
-            num_arts = np.sum(num_arts)
-            self.scene_mask = sapien_utils.to_tensor(self.scene_mask)
+        if self.scene_idxs is not None:
+            pass
         else:
-            # if scene mask is none, set it here
-            self.scene_mask = sapien_utils.to_tensor(
-                torch.ones((self.scene.num_envs), dtype=bool)
-            )
+            self.scene_idxs = torch.arange((self.scene.num_envs), dtype=int)
 
         articulations = []
 
-        for scene_idx, scene in enumerate(self.scene.sub_scenes):
-            if self.scene_mask is not None and self.scene_mask[scene_idx] == False:
-                continue
+        for scene_idx in self.scene_idxs:
+            sub_scene = self.scene.sub_scenes[scene_idx]
             links: List[sapien.Entity] = self.build_entities(
                 name_prefix=f"scene-{scene_idx}-{self.name}_"
             )
@@ -176,13 +167,13 @@ def build(self, name=None, fix_root_link=None, build_mimic_joints=True):
                         )
 
             for l in links:
-                scene.add_entity(l)
+                sub_scene.add_entity(l)
             articulation: physx.PhysxArticulation = l.components[0].articulation
             articulation.name = f"scene-{scene_idx}_{self.name}"
             articulations.append(articulation)
 
-        articulation = Articulation._create_from_physx_articulations(
-            articulations, self.scene, self.scene_mask
+        articulation = Articulation.create_from_physx_articulations(
+            articulations, self.scene, self.scene_idxs
         )
         self.scene.articulations[self.name] = articulation
         return articulation

mani_skill/utils/building/articulations.py

@@ -79,7 +79,7 @@ def build_preprocessed_partnet_mobility_articulation(
     name: str,
     fix_root_link=True,
     urdf_config: dict = None,
-    scene_mask=None,
+    scene_idxs=None,
 ):
     """
     Builds a physx.PhysxArticulation object into the scene and returns metadata containing annotations of the object's links and joints.
@@ -107,7 +107,7 @@ def build_preprocessed_partnet_mobility_articulation(
     urdf_path = MODEL_DBS["PartnetMobility"]["model_urdf_paths"][model_id]
     urdf_config = sapien_utils.parse_urdf_config(urdf_config or {}, scene)
     sapien_utils.apply_urdf_config(loader, urdf_config)
-    articulation = loader.load(str(urdf_path), name=name, scene_mask=scene_mask)
+    articulation = loader.load(str(urdf_path), name=name, scene_idxs=scene_idxs)
     metadata = ArticulationMetadata(
         joints=dict(), links=dict(), movable_links=[], bbox=None, scale=loader.scale
     )
@@ -169,7 +169,7 @@ def build_robel_valve(
     name: str,
     radius_scale: float = 1.0,
     capsule_radius_scale: float = 1.0,
-    scene_mask=None,
+    scene_idxs=None,
 ):
     # Size and geometry of valve are based on the original setting of Robel benchmark, unit: m
     # Ref: https://github.com/google-research/robel
@@ -182,7 +182,7 @@ def build_robel_valve(
     bearing_height = 0.032
 
     builder = scene.create_articulation_builder()
-    builder.set_scene_mask(scene_mask)
+    builder.set_scene_idxs(scene_idxs)
 
     # Mount link
     mount_builder = builder.create_link_builder(parent=None)

mani_skill/utils/building/urdf_loader.py

@@ -51,14 +51,15 @@ def load(
         srdf_file=None,
         package_dir=None,
         name=None,
-        scene_mask=None,
+        scene_idxs=None,
     ) -> Articulation:
         """
         Args:
             urdf_file: filename for URDL file
             srdf_file: SRDF for urdf_file. If srdf_file is None, it defaults to the ".srdf" file with the same as the urdf file
             package_dir: base directory used to resolve asset files in the URDF file. If an asset path starts with "package://", "package://" is simply removed from the file name
             name (str): name of the created articulation
+            scene_idxs (list[int]): the ids of the scenes to build the objects in
         Returns:
             returns a single Articulation loaded from the URDF file. It throws an error if multiple objects exists
         """
@@ -75,7 +76,7 @@ def load(
 
         articulations: List[Articulation] = []
         for b in articulation_builders:
-            b.set_scene_mask(scene_mask)
+            b.set_scene_idxs(scene_idxs)
             b.disable_self_collisions = self.disable_self_collisions
             articulations.append(b.build())
 

mani_skill/utils/structs/actor.py

@@ -34,18 +34,18 @@ class Actor(PhysxRigidDynamicComponentStruct, BaseStruct[sapien.Entity]):
 
     # track the initial pose of the actor builder for this actor. Necessary to ensure the actor is reset correctly once
     # gpu system is initialized
-    _builder_initial_pose: sapien.Pose = None
+    _builder_initial_pose: Pose = None
     name: str = None
 
     def __hash__(self):
         return self._objs[0].__hash__()
 
     @classmethod
-    def _create_from_entities(
+    def create_from_entities(
         cls,
         entities: List[sapien.Entity],
         scene: ManiSkillScene,
-        scene_mask: torch.Tensor,
+        scene_idxs: torch.Tensor,
     ):
 
         shared_name = "_".join(entities[0].name.split("_")[1:])
@@ -69,7 +69,7 @@ def _create_from_entities(
         return cls(
             _objs=entities,
             _scene=scene,
-            _scene_mask=scene_mask,
+            _scene_idxs=scene_idxs,
             px_body_type=px_body_type,
             _bodies=bodies,
             _body_data_name="cuda_rigid_body_data"
@@ -94,19 +94,20 @@ def merge(cls, actors: List["Actor"], name: str = None):
         objs = []
         scene = actors[0]._scene
         _builder_initial_poses = []
-        merged_scene_mask = actors[0]._scene_mask.clone()
+        merged_scene_idxs = []
         num_objs_per_actor = actors[0]._num_objs
         for actor in actors:
             objs += actor._objs
-            merged_scene_mask[actor._scene_mask] = True
+            merged_scene_idxs.append(actor._scene_idxs)
             _builder_initial_poses.append(actor._builder_initial_pose.raw_pose)
             del scene.actors[actor.name]
             assert (
                 actor._num_objs == num_objs_per_actor
             ), "Each given actor must have the same number of managed objects"
         # TODO (stao): Can we support e.g. each Actor having len(actor._objs) > 1? It would mean fetching pose data or any kind of data is highly uintuitive
         # we definitely cannot permit some actors to have more objs than others, otherwise the data is ragged.
-        merged_actor = Actor._create_from_entities(objs, scene, merged_scene_mask)
+        merged_scene_idxs = torch.concat(merged_scene_idxs)
+        merged_actor = Actor.create_from_entities(objs, scene, merged_scene_idxs)
         merged_actor.name = name
         merged_actor._builder_initial_pose = Pose.create(
             torch.vstack(_builder_initial_poses)