work

mani_skill2/utils/building/articulation_builder.py

@@ -11,7 +11,7 @@
 )
 from sapien.wrapper.articulation_builder import LinkBuilder
 
-from mani_skill2.utils.sapien_utils import to_tensor
+from mani_skill2.utils import sapien_utils
 from mani_skill2.utils.structs.articulation import Articulation
 
 if TYPE_CHECKING:
@@ -111,10 +111,12 @@ def build(self, name=None, fix_root_link=None):
                 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 = to_tensor(self.scene_mask)
+            self.scene_mask = sapien_utils.to_tensor(self.scene_mask)
         else:
             # if scene mask is none, set it here
-            self.scene_mask = to_tensor(torch.ones((self.scene.num_envs), dtype=bool))
+            self.scene_mask = sapien_utils.to_tensor(
+                torch.ones((self.scene.num_envs), dtype=bool)
+            )
 
         articulations = []
 

mani_skill2/utils/building/articulations.py

@@ -14,12 +14,12 @@
 
 from mani_skill2 import ASSET_DIR, PACKAGE_ASSET_DIR
 from mani_skill2.envs.scene import ManiSkillScene
+from mani_skill2.utils import sapien_utils
 from mani_skill2.utils.geometry.trimesh_utils import (
     get_articulation_meshes,
     merge_meshes,
 )
 from mani_skill2.utils.io_utils import load_json
-from mani_skill2.utils.sapien_utils import apply_urdf_config, parse_urdf_config
 
 
 @dataclass
@@ -105,8 +105,8 @@ def build_preprocessed_partnet_mobility_articulation(
     # loader.multiple_collisions_decomposition="coacd"
     loader.disable_self_collisions = True
     urdf_path = MODEL_DBS["PartnetMobility"]["model_urdf_paths"][model_id]
-    urdf_config = parse_urdf_config(urdf_config or {}, scene)
-    apply_urdf_config(loader, urdf_config)
+    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)
     metadata = ArticulationMetadata(
         joints=dict(), links=dict(), movable_links=[], bbox=None, scale=loader.scale

mani_skill2/utils/common.py

@@ -7,7 +7,7 @@
 import torch
 from gymnasium import spaces
 
-from mani_skill2.utils.sapien_utils import to_tensor
+from mani_skill2.utils import sapien_utils
 from mani_skill2.utils.structs.types import Array, Device
 
 from .logging_utils import logger
@@ -209,28 +209,28 @@ def flatten_state_dict(
             if state.size == 0:
                 state = None
             if use_torch:
-                state = to_tensor(state)
+                state = sapien_utils.to_tensor(state)
         elif isinstance(value, (tuple, list)):
             state = None if len(value) == 0 else value
             if use_torch:
-                state = to_tensor(state)
+                state = sapien_utils.to_tensor(state)
         elif isinstance(value, (bool, np.bool_, int, np.int32, np.int64)):
             # x = np.array(1) > 0 is np.bool_ instead of ndarray
             state = int(value)
             if use_torch:
-                state = to_tensor(state)
+                state = sapien_utils.to_tensor(state)
         elif isinstance(value, (float, np.float32, np.float64)):
             state = np.float32(value)
             if use_torch:
-                state = to_tensor(state)
+                state = sapien_utils.to_tensor(state)
         elif isinstance(value, np.ndarray):
             if value.ndim > 2:
                 raise AssertionError(
                     "The dimension of {} should not be more than 2.".format(key)
                 )
             state = value if value.size > 0 else None
             if use_torch:
-                state = to_tensor(state)
+                state = sapien_utils.to_tensor(state)
 
         else:
             is_torch_tensor = False

mani_skill2/utils/structs/actor.py

@@ -10,7 +10,7 @@
 import sapien.render
 import torch
 
-from mani_skill2.utils.sapien_utils import to_numpy, to_tensor
+from mani_skill2.utils import sapien_utils
 from mani_skill2.utils.structs.base import BaseStruct, PhysxRigidDynamicComponentStruct
 from mani_skill2.utils.structs.pose import Pose, to_sapien_pose, vectorize_pose
 from mani_skill2.utils.structs.types import Array
@@ -130,12 +130,12 @@ def get_state(self):
 
     def set_state(self, state: Array):
         if physx.is_gpu_enabled():
-            state = to_tensor(state)
+            state = sapien_utils.to_tensor(state)
             self.set_pose(Pose.create(state[:, :7]))
             self.set_linear_velocity(state[:, 7:10])
             self.set_angular_velocity(state[:, 10:13])
         else:
-            state = to_numpy(state[0])
+            state = sapien_utils.to_numpy(state[0])
             self.set_pose(sapien.Pose(state[0:3], state[3:7]))
             if self.px_body_type == "dynamic":
                 self.set_linear_velocity(state[7:10])

mani_skill2/utils/structs/articulation.py

@@ -11,13 +11,8 @@
 import torch
 import trimesh
 
+from mani_skill2.utils import sapien_utils
 from mani_skill2.utils.geometry.trimesh_utils import get_component_meshes, merge_meshes
-from mani_skill2.utils.sapien_utils import (
-    compute_total_impulse,
-    get_articulation_contacts,
-    to_numpy,
-    to_tensor,
-)
 from mani_skill2.utils.structs.base import BaseStruct
 from mani_skill2.utils.structs.joint import Joint
 from mani_skill2.utils.structs.link import Link
@@ -201,15 +196,15 @@ def get_state(self):
 
     def set_state(self, state: Array):
         if physx.is_gpu_enabled():
-            state = to_tensor(state)
+            state = sapien_utils.to_tensor(state)
             self.set_root_pose(Pose.create(state[:, :7]))
             self.set_root_linear_velocity(state[:, 7:10])
             self.set_root_angular_velocity(state[:, 10:13])
             # TODO (stao): Handle get/set state for envs with different DOFs. Perhaps need to let user set a padding ahead of time to ensure state is the same?
             self.set_qpos(state[:, 13 : 13 + self.max_dof])
             self.set_qvel(state[:, 13 + self.max_dof :])
         else:
-            state = to_numpy(state[0])
+            state = sapien_utils.to_numpy(state[0])
             self.set_root_pose(sapien.Pose(state[0:3], state[3:7]))
             self.set_root_linear_velocity(state[7:10])
             self.set_root_angular_velocity(state[10:13])
@@ -283,13 +278,16 @@ def get_net_contact_forces(self, link_names: Union[List[str], Tuple[str]]):
             )
         else:
 
-            body_contacts = get_articulation_contacts(
+            body_contacts = sapien_utils.get_articulation_contacts(
                 self.px.get_contacts(),
                 self._objs[0],
                 included_links=[self.link_map[k]._objs[0] for k in link_names],
             )
             net_force = (
-                to_tensor(compute_total_impulse(body_contacts)) / self._scene.timestep
+                sapien_utils.to_tensor(
+                    sapien_utils.compute_total_impulse(body_contacts)
+                )
+                / self._scene.timestep
             )
             return net_force[None, :]
 
@@ -451,12 +449,12 @@ def qf(self):
     @qf.setter
     def qf(self, arg1: torch.Tensor):
         if physx.is_gpu_enabled():
-            arg1 = to_tensor(arg1)
+            arg1 = sapien_utils.to_tensor(arg1)
             self.px.cuda_articulation_qf.torch()[
                 self._data_index[self._scene._reset_mask], : self.max_dof
             ] = arg1
         else:
-            arg1 = to_numpy(arg1)
+            arg1 = sapien_utils.to_numpy(arg1)
             if len(arg1.shape) == 2:
                 arg1 = arg1[0]
             self._objs[0].qf = arg1
@@ -489,12 +487,12 @@ def qpos(self):
     @qpos.setter
     def qpos(self, arg1: torch.Tensor):
         if physx.is_gpu_enabled():
-            arg1 = to_tensor(arg1)
+            arg1 = sapien_utils.to_tensor(arg1)
             self.px.cuda_articulation_qpos.torch()[
                 self._data_index[self._scene._reset_mask], : self.max_dof
             ] = arg1
         else:
-            arg1 = to_numpy(arg1)
+            arg1 = sapien_utils.to_numpy(arg1)
             if len(arg1.shape) == 2:
                 arg1 = arg1[0]
             self._objs[0].qpos = arg1
@@ -511,12 +509,12 @@ def qvel(self):
     @qvel.setter
     def qvel(self, arg1: torch.Tensor):
         if physx.is_gpu_enabled():
-            arg1 = to_tensor(arg1)
+            arg1 = sapien_utils.to_tensor(arg1)
             self.px.cuda_articulation_qvel.torch()[
                 self._data_index[self._scene._reset_mask], : self.max_dof
             ] = arg1
         else:
-            arg1 = to_numpy(arg1)
+            arg1 = sapien_utils.to_numpy(arg1)
             if len(arg1.shape) == 2:
                 arg1 = arg1[0]
             self._objs[0].qvel = arg1
@@ -528,12 +526,12 @@ def root_angular_velocity(self) -> torch.Tensor:
     @root_angular_velocity.setter
     def root_angular_velocity(self, arg1: Array) -> None:
         if physx.is_gpu_enabled():
-            arg1 = to_tensor(arg1)
+            arg1 = sapien_utils.to_tensor(arg1)
             self.px.cuda_rigid_body_data.torch()[
                 self.root._body_data_index[self._scene._reset_mask], 10:13
             ] = arg1
         else:
-            arg1 = to_numpy(arg1)
+            arg1 = sapien_utils.to_numpy(arg1)
             if len(arg1.shape) == 2:
                 arg1 = arg1[0]
             self._objs[0].set_root_angular_velocity(arg1)
@@ -545,12 +543,12 @@ def root_linear_velocity(self) -> torch.Tensor:
     @root_linear_velocity.setter
     def root_linear_velocity(self, arg1: Array) -> None:
         if physx.is_gpu_enabled():
-            arg1 = to_tensor(arg1)
+            arg1 = sapien_utils.to_tensor(arg1)
             self.px.cuda_rigid_body_data.torch()[
                 self.root._body_data_index[self._scene._reset_mask], 7:10
             ] = arg1
         else:
-            arg1 = to_numpy(arg1)
+            arg1 = sapien_utils.to_numpy(arg1)
             if len(arg1.shape) == 2:
                 arg1 = arg1[0]
             self._objs[0].set_root_linear_velocity(arg1)
@@ -590,7 +588,7 @@ def set_joint_drive_targets(
         TODO (stao): can we use joint indices for the CPU sim as well? Some global joint indices?
         """
         if physx.is_gpu_enabled():
-            targets = to_tensor(targets)
+            targets = sapien_utils.to_tensor(targets)
             if joint_indices not in self._cached_joint_target_indices:
                 self._cached_joint_target_indices[joint_indices] = torch.meshgrid(
                     self._data_index, joint_indices, indexing="ij"
@@ -613,7 +611,7 @@ def set_joint_drive_velocity_targets(
         TODO (stao): can we use joint indices for the CPU sim as well? Some global joint indices?
         """
         if physx.is_gpu_enabled():
-            targets = to_tensor(targets)
+            targets = sapien_utils.to_tensor(targets)
             if joint_indices not in self._cached_joint_target_indices:
                 self._cached_joint_target_indices[joint_indices] = torch.meshgrid(
                     self._data_index, joint_indices, indexing="ij"

mani_skill2/utils/structs/base.py

@@ -7,12 +7,7 @@
 import sapien.physx as physx
 import torch
 
-from mani_skill2.utils.sapien_utils import (
-    compute_total_impulse,
-    get_actor_contacts,
-    to_numpy,
-    to_tensor,
-)
+from mani_skill2.utils import sapien_utils
 from mani_skill2.utils.structs.decorators import before_gpu_init
 from mani_skill2.utils.structs.types import Array
 
@@ -94,11 +89,14 @@ def get_net_contact_forces(self):
                 / self._scene.timestep
             )
         else:
-            body_contacts = get_actor_contacts(
+            body_contacts = sapien_utils.get_actor_contacts(
                 self.px.get_contacts(), self._bodies[0].entity
             )
             net_force = (
-                to_tensor(compute_total_impulse(body_contacts)) / self._scene.timestep
+                sapien_utils.to_tensor(
+                    sapien_utils.compute_total_impulse(body_contacts)
+                )
+                / self._scene.timestep
             )
             return net_force[None, :]
 
@@ -299,12 +297,12 @@ def angular_velocity(self) -> torch.Tensor:
     @angular_velocity.setter
     def angular_velocity(self, arg1: Array):
         if physx.is_gpu_enabled():
-            arg1 = to_tensor(arg1)
+            arg1 = sapien_utils.to_tensor(arg1)
             self._body_data[
                 self._body_data_index[self._scene._reset_mask], 10:13
             ] = arg1
         else:
-            arg1 = to_numpy(arg1)
+            arg1 = sapien_utils.to_numpy(arg1)
             if len(arg1.shape) == 2:
                 arg1 = arg1[0]
             self._bodies[0].angular_velocity = arg1
@@ -368,10 +366,10 @@ def linear_velocity(self) -> torch.Tensor:
     @linear_velocity.setter
     def linear_velocity(self, arg1: Array):
         if physx.is_gpu_enabled():
-            arg1 = to_tensor(arg1)
+            arg1 = sapien_utils.to_tensor(arg1)
             self._body_data[self._body_data_index[self._scene._reset_mask], 7:10] = arg1
         else:
-            arg1 = to_numpy(arg1)
+            arg1 = sapien_utils.to_numpy(arg1)
             if len(arg1.shape) == 2:
                 arg1 = arg1[0]
             self._bodies[0].linear_velocity = arg1

mani_skill2/utils/structs/joint.py

@@ -8,7 +8,7 @@
 import sapien.physx as physx
 import torch
 
-from mani_skill2.utils.sapien_utils import to_tensor
+from mani_skill2.utils import sapien_utils
 from mani_skill2.utils.structs.base import BaseStruct
 from mani_skill2.utils.structs.decorators import before_gpu_init
 from mani_skill2.utils.structs.link import Link
@@ -194,7 +194,7 @@ def drive_target(self) -> torch.Tensor:
     @drive_target.setter
     def drive_target(self, arg1: Array) -> None:
         if physx.is_gpu_enabled():
-            arg1 = to_tensor(arg1)
+            arg1 = sapien_utils.to_tensor(arg1)
             raise NotImplementedError(
                 "Setting drive targets of individual joints is not implemented yet."
             )
@@ -217,7 +217,7 @@ def drive_velocity_target(self) -> torch.Tensor:
     @drive_velocity_target.setter
     def drive_velocity_target(self, arg1: Array) -> None:
         if physx.is_gpu_enabled():
-            arg1 = to_tensor(arg1)
+            arg1 = sapien_utils.to_tensor(arg1)
             raise NotImplementedError(
                 "Cannot set drive velocity targets at the moment in GPU simulation"
             )