fix naming of joints and links to also be prepended by scene-{i}

mani_skill2/agents/utils.py

@@ -7,6 +7,8 @@
 import sapien.physx as physx
 from gymnasium import spaces
 
+from mani_skill2.utils.structs.articulation import Articulation
+
 
 def get_joint_indices(
     articulation: physx.PhysxArticulation, joint_names: Sequence[str]
@@ -31,9 +33,7 @@ def get_joints(articulation: physx.PhysxArticulation, joint_names: Sequence[str]
     return [joints[idx] for idx in joint_indices]
 
 
-def get_active_joints(
-    articulation: physx.PhysxArticulation, joint_names: Sequence[str]
-):
+def get_active_joints(articulation: Articulation, joint_names: Sequence[str]):
     joints = articulation.get_active_joints()
     joint_indices = get_active_joint_indices(articulation, joint_names)
     return [joints[idx] for idx in joint_indices]

mani_skill2/envs/tasks/open_cabinet_drawer.py

@@ -63,50 +63,52 @@ def _load_actors(self):
         self.ground = build_tesselated_square_floor(self._scene)
 
         cabinets = []
+        self.cabinet_heights = []
         for i, model_id in enumerate(model_ids):
             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.cabinet_heights.append(
+                metadata.bbox.bounds[1, 2] - metadata.bbox.bounds[0, 2]
+            )
             # self.cabinet = cabinet
             # self.cabinet_metadata = metadata
             cabinets.append(cabinet)
         self.cabinet = Articulation.merge_articulations(cabinets, name="cabinet")
         self.cabinet_metadata = metadata
 
     def _initialize_actors(self):
-        height = (
-            self.cabinet_metadata.bbox.bounds[0, 2]
-            - self.cabinet_metadata.bbox.bounds[1, 2]
-        )
-        self.cabinet.set_pose(Pose.create_from_pq(p=[0, 0, -height / 2]))
-        qlimits = self.cabinet.get_qlimits()  # [N, self.cabinet.max_dof, 2]
-        qpos = qlimits[:, :, 0]
-        self.cabinet.set_qpos(
-            qpos
-        )  # close all the cabinets. We know beforehand that lower qlimit means "closed" for these assets.
-
-        # initialize robot
-        if self.robot_uid == "panda":
-            self.agent.robot.set_qpos(self.agent.robot.qpos * 0)
-            self.agent.robot.set_pose(Pose.create_from_pq(p=[-1, 0, 0]))
-        elif self.robot_uid == "mobile_panda_single_arm":
-            center = np.array([0, 0.8])
-            dist = self._episode_rng.uniform(1.6, 1.8)
-            theta = self._episode_rng.uniform(0.9 * np.pi, 1.1 * np.pi)
-            direction = np.array([np.cos(theta), np.sin(theta)])
-            xy = center + direction * dist
-
-            # Base orientation
-            noise_ori = self._episode_rng.uniform(-0.05 * np.pi, 0.05 * np.pi)
-            ori = (theta - np.pi) + noise_ori
-
-            h = 1e-4
-            arm_qpos = np.array([0, 0, 0, -1.5, 0, 3, 0.78, 0.02, 0.02])
-
-            qpos = np.hstack([xy, ori, h, arm_qpos])
-            self.agent.reset(qpos)
+        with torch.device(self.device):
+            xyz = torch.zeros((self.num_envs, 3))
+            xyz[:, 2] = torch.tensor(self.cabinet_heights) / 2
+            self.cabinet.set_pose(Pose.create_from_pq(p=xyz))
+            qlimits = self.cabinet.get_qlimits()  # [N, self.cabinet.max_dof, 2]
+            qpos = qlimits[:, :, 0]
+            self.cabinet.set_qpos(
+                qpos
+            )  # close all the cabinets. We know beforehand that lower qlimit means "closed" for these assets.
+            # initialize robot
+            if self.robot_uid == "panda":
+                self.agent.robot.set_qpos(self.agent.robot.qpos * 0)
+                self.agent.robot.set_pose(Pose.create_from_pq(p=[-1, 0, 0]))
+            elif self.robot_uid == "mobile_panda_single_arm":
+                center = np.array([0, 0.8])
+                dist = self._episode_rng.uniform(1.6, 1.8)
+                theta = self._episode_rng.uniform(0.9 * np.pi, 1.1 * np.pi)
+                direction = np.array([np.cos(theta), np.sin(theta)])
+                xy = center + direction * dist
+
+                # Base orientation
+                noise_ori = self._episode_rng.uniform(-0.05 * np.pi, 0.05 * np.pi)
+                ori = (theta - np.pi) + noise_ori
+
+                h = 1e-4
+                arm_qpos = np.array([0, 0, 0, -1.5, 0, 3, 0.78, 0.02, 0.02])
+
+                qpos = np.hstack([xy, ori, h, arm_qpos])
+                self.agent.reset(qpos)
 
     def _get_obs_extra(self):
         return OrderedDict()

mani_skill2/utils/building/articulation_builder.py

@@ -39,6 +39,52 @@ def set_scene_mask(
         """
         self.scene_mask = scene_mask
 
+    def build_entities(self, fix_root_link=None, name_prefix=""):
+        entities = []
+        links = []
+        for b in self.link_builders:
+            b._check()
+            b.physx_body_type = "link"
+
+            entity = sapien.Entity()
+
+            link_component = b.build_physx_component(
+                links[b.parent.index] if b.parent else None
+            )
+
+            entity.add_component(link_component)
+            if b.visual_records:
+                entity.add_component(b.build_render_component())
+            entity.name = b.name
+
+            link_component.name = f"{name_prefix}{b.name}"
+            link_component.joint.name = f"{name_prefix}{b.joint_record.name}"
+            link_component.joint.type = b.joint_record.joint_type
+            link_component.joint.pose_in_child = b.joint_record.pose_in_child
+            link_component.joint.pose_in_parent = b.joint_record.pose_in_parent
+
+            if link_component.joint.type in [
+                "revolute",
+                "prismatic",
+                "revolute_unwrapped",
+            ]:
+                link_component.joint.limit = np.array(b.joint_record.limits).flatten()
+                link_component.joint.set_drive_property(0, b.joint_record.damping)
+
+            if link_component.joint.type == "continuous":
+                link_component.joint.limit = [-np.inf, np.inf]
+                link_component.joint.set_drive_property(0, b.joint_record.damping)
+
+            links.append(link_component)
+            entities.append(entity)
+
+        if fix_root_link is not None:
+            entities[0].components[0].joint.type = (
+                "fixed" if fix_root_link else "undefined"
+            )
+        entities[0].pose = self.initial_pose
+        return entities
+
     def build(self, name=None, fix_root_link=None):
         assert self.scene is not None
         if name is not None:
@@ -61,14 +107,16 @@ def build(self, name=None, fix_root_link=None):
         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
-            links: List[sapien.Entity] = self.build_entities()
+            links: List[sapien.Entity] = self.build_entities(
+                name_prefix=f"scene-{scene_idx}_"
+            )
             if fix_root_link is not None:
                 links[0].components[0].joint.type = (
                     "fixed" if fix_root_link else "undefined"
                 )
             links[0].pose = self.initial_pose
             for l in links:
-                l.name = f"scene-{scene_idx}_{l.name}"
+                # l.name = f"scene-{scene_idx}_{l.name}"
                 scene.add_entity(l)
             articulation: physx.PhysxArticulation = l.components[0].articulation
             articulation.name = f"scene-{scene_idx}_{self.name}"

mani_skill2/utils/structs/articulation.py

@@ -165,11 +165,11 @@ def set_state(self, state: Array):
     def max_dof(self) -> int:
         return max([obj.dof for obj in self._objs])
 
-    def bbox(self):
-        import ipdb
-
-        ipdb.set_trace()
-        self._objs[0]
+    # def bbox(self):
+    # you should avoid calling this too often
+    # the bounding box can change easily if the qpos changes.
+    # in theory this can be implemented by precomputing the meshes for each link, then transforming the meshes by the link poses
+    # however there will have to be some for loops if this articulation is managing different articulations with different links
 
     # -------------------------------------------------------------------------- #
     # Functions from physx.PhysxArticulation
@@ -323,12 +323,14 @@ def qf(self):
             return torch.from_numpy(self._objs[0].qf[None, :])
 
     @qf.setter
-    def qf(self, arg1):
+    def qf(self, arg1: torch.Tensor):
         if physx.is_gpu_enabled():
             arg1 = to_tensor(arg1)
             self.px.cuda_articulation_qf[self._data_index, : self.max_dof] = arg1
         else:
-            self._objs[0].qf = arg1
+            if len(arg1.shape) == 2:
+                arg1 = arg1[0]
+            self._objs[0].qf = to_numpy(arg1)
 
     @cached_property
     def qlimits(self):
@@ -354,12 +356,14 @@ def qpos(self):
             return torch.from_numpy(self._objs[0].qpos[None, :])
 
     @qpos.setter
-    def qpos(self, arg1):
+    def qpos(self, arg1: torch.Tensor):
         if physx.is_gpu_enabled():
             arg1 = to_tensor(arg1)
             self.px.cuda_articulation_qpos[self._data_index, : self.max_dof] = arg1
         else:
-            self._objs[0].qpos = arg1
+            if len(arg1.shape) == 2:
+                arg1 = arg1[0]
+            self._objs[0].qpos = to_numpy(arg1)
 
     @property
     def qvel(self):
@@ -369,12 +373,14 @@ def qvel(self):
             return torch.from_numpy(self._objs[0].qvel[None, :])
 
     @qvel.setter
-    def qvel(self, arg1):
+    def qvel(self, arg1: torch.Tensor):
         if physx.is_gpu_enabled():
             arg1 = to_tensor(arg1)
             self.px.cuda_articulation_qvel[self._data_index, : self.max_dof] = arg1
         else:
-            self._objs[0].qvel = arg1
+            if len(arg1.shape) == 2:
+                arg1 = arg1[0]
+            self._objs[0].qvel = to_numpy(arg1)
 
     @property
     def root_angular_velocity(self) -> torch.Tensor:

mani_skill2/utils/structs/joint.py

@@ -45,13 +45,17 @@ def create(
         articulation: Articulation,
     ):
         px: physx.PhysxSystem = articulation.px
-        shared_name = physx_joints[0].name
+        shared_name = "_".join(physx_joints[0].name.split("_")[1:])
         child_link = None
         parent_link = None
         if physx_joints[0].child_link is not None:
-            child_link = articulation.link_map[physx_joints[0].child_link.name]
+            child_link = articulation.link_map[
+                "_".join(physx_joints[0].child_link.name.split("_")[1:])
+            ]
         if physx_joints[0].parent_link is not None:
-            parent_link = articulation.link_map[physx_joints[0].parent_link.name]
+            parent_link = articulation.link_map[
+                "_".join(physx_joints[0].parent_link.name.split("_")[1:])
+            ]
         return cls(
             articulation=articulation,
             _objs=physx_joints,

mani_skill2/utils/structs/link.py

@@ -33,7 +33,7 @@ def create(
         physx_links: List[physx.PhysxArticulationLinkComponent],
         articulation: Articulation,
     ):
-        shared_name = physx_links[0].name
+        shared_name = "_".join(physx_links[0].name.split("_")[1:])
         return cls(
             articulation=articulation,
             _objs=physx_links,

manualtest/visual_all_envs_cpu.py

@@ -8,7 +8,7 @@
 
 if __name__ == "__main__":
     # , "StackCube-v1", "PickCube-v1", "PushCube-v1", "PickSingleYCB-v1"
-    num_envs = 4
+    num_envs = 8
     sapien.physx.set_gpu_memory_config(
         found_lost_pairs_capacity=2**26,
         max_rigid_patch_count=2**19,