Packaging and import fixing

mani_skill2/agents/base_agent.py

@@ -11,13 +11,8 @@
 
 from mani_skill2 import format_path
 from mani_skill2.sensors.base_sensor import BaseSensor, BaseSensorConfig
-from mani_skill2.utils.sapien_utils import (
-    apply_urdf_config,
-    check_urdf_config,
-    parse_urdf_config,
-)
-from mani_skill2.utils.structs.actor import Actor
-from mani_skill2.utils.structs.articulation import Articulation
+from mani_skill2.utils import sapien_utils
+from mani_skill2.utils.structs import Actor, Articulation
 
 from .controllers.base_controller import (
     BaseController,
@@ -98,11 +93,11 @@ def _load_articulation(self):
 
         urdf_path = format_path(str(self.urdf_path))
 
-        urdf_config = parse_urdf_config(self.urdf_config, self.scene)
-        check_urdf_config(urdf_config)
+        urdf_config = sapien_utils.parse_urdf_config(self.urdf_config, self.scene)
+        sapien_utils.check_urdf_config(urdf_config)
 
         # TODO(jigu): support loading multiple convex collision shapes
-        apply_urdf_config(loader, urdf_config)
+        sapien_utils.apply_urdf_config(loader, urdf_config)
         self.robot: Articulation = loader.load(urdf_path)
         assert self.robot is not None, f"Fail to load URDF from {urdf_path}"
 

mani_skill2/agents/controllers/base_controller.py

@@ -16,8 +16,8 @@
     get_active_joint_indices,
     get_joints_by_names,
 )
+from mani_skill2.utils import sapien_utils
 from mani_skill2.utils.common import clip_and_scale_action, normalize_action_space
-from mani_skill2.utils.sapien_utils import to_tensor
 from mani_skill2.utils.structs.articulation import Articulation
 from mani_skill2.utils.structs.joint import Joint
 from mani_skill2.utils.structs.types import Array
@@ -37,8 +37,10 @@ class BaseController:
     """the action space. If the number of parallel environments is > 1, this action space is also batched"""
     single_action_space: spaces.Space
     """The unbatched version of the action space which is also typically already normalized by this class"""
+    """The batched version of the action space which is also typically already normalized by this class"""
     _original_single_action_space: spaces.Space
     """The unbatched, original action space without any additional processing like normalization"""
+    """The batched, original action space without any additional processing like normalization"""
 
     def __init__(
         self,
@@ -162,8 +164,8 @@ def _clip_and_scale_action_space(self):
             self._original_single_action_space.low,
             self._original_single_action_space.high,
         )
-        self.action_space_low = to_tensor(low)
-        self.action_space_high = to_tensor(high)
+        self.action_space_low = sapien_utils.to_tensor(low)
+        self.action_space_high = sapien_utils.to_tensor(high)
 
     def _clip_and_scale_action(self, action):
         return clip_and_scale_action(
@@ -251,7 +253,7 @@ def reset(self):
 
     def set_action(self, action: Dict[str, np.ndarray]):
         for uid, controller in self.controllers.items():
-            controller.set_action(to_tensor(action[uid]))
+            controller.set_action(sapien_utils.to_tensor(action[uid]))
 
     def before_simulation_step(self):
         if physx.is_gpu_enabled():

mani_skill2/agents/controllers/pd_ee_pose.py

@@ -12,12 +12,12 @@
 import torch
 from gymnasium import spaces
 
+from mani_skill2.utils import sapien_utils
 from mani_skill2.utils.common import clip_and_scale_action
 from mani_skill2.utils.geometry.rotation_conversions import (
     euler_angles_to_matrix,
     matrix_to_quaternion,
 )
-from mani_skill2.utils.sapien_utils import get_obj_by_name, to_numpy, to_tensor
 from mani_skill2.utils.structs.pose import Pose, vectorize_pose
 from mani_skill2.utils.structs.types import Array
 
@@ -60,7 +60,7 @@ def suppress_stdout_stderr():
         self.qmask[self.joint_indices] = 1
 
         if self.config.ee_link:
-            self.ee_link = get_obj_by_name(
+            self.ee_link = sapien_utils.get_obj_by_name(
                 self.articulation.get_links(), self.config.ee_link
             )
         else:
@@ -112,12 +112,14 @@ def compute_ik(self, target_pose: Pose, action: Array, max_iterations=100):
             result, success, error = self.pmodel.compute_inverse_kinematics(
                 self.ee_link_idx,
                 target_pose.sp,
-                initial_qpos=to_numpy(self.articulation.get_qpos()).squeeze(0),
+                initial_qpos=sapien_utils.to_numpy(
+                    self.articulation.get_qpos()
+                ).squeeze(0),
                 active_qmask=self.qmask,
                 max_iterations=max_iterations,
             )
         if success:
-            return to_tensor([result[self.joint_indices]])
+            return sapien_utils.to_tensor([result[self.joint_indices]])
         else:
             return None
 

mani_skill2/agents/controllers/pd_joint_pos.py

@@ -5,7 +5,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
 
 from .base_controller import BaseController, ControllerConfig
@@ -67,7 +67,7 @@ def set_drive_targets(self, targets):
 
     def set_action(self, action: Array):
         action = self._preprocess_action(action)
-        action = to_tensor(action)
+        action = sapien_utils.to_tensor(action)
         self._step = 0
         self._start_qpos = self.qpos
         if self.config.use_delta:

mani_skill2/agents/robots/_template/template_robot.py

@@ -2,12 +2,6 @@
 from mani_skill2.agents.base_agent import BaseAgent
 from mani_skill2.agents.controllers import *
 from mani_skill2.sensors.camera import CameraConfig
-from mani_skill2.utils.sapien_utils import (
-    compute_total_impulse,
-    get_actor_contacts,
-    get_obj_by_name,
-    get_pairwise_contact_impulse,
-)
 
 
 class TemplateRobot(BaseAgent):

mani_skill2/agents/robots/allegro_hand/allegro.py

@@ -8,7 +8,7 @@
 from mani_skill2.agents.base_agent import BaseAgent
 from mani_skill2.agents.controllers import *
 from mani_skill2.agents.registration import register_agent
-from mani_skill2.utils.sapien_utils import get_obj_by_name, get_objs_by_names
+from mani_skill2.utils import sapien_utils
 from mani_skill2.utils.structs.pose import vectorize_pose
 
 
@@ -73,10 +73,10 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
     def _after_init(self):
-        self.tip_links: List[sapien.Entity] = get_objs_by_names(
+        self.tip_links: List[sapien.Entity] = sapien_utils.get_objs_by_names(
             self.robot.get_links(), self.tip_link_names
         )
-        self.palm_link: sapien.Entity = get_obj_by_name(
+        self.palm_link: sapien.Entity = sapien_utils.get_obj_by_name(
             self.robot.get_links(), self.palm_link_name
         )
 

mani_skill2/agents/robots/allegro_hand/allegro_touch.py

@@ -9,12 +9,7 @@
 from mani_skill2 import PACKAGE_ASSET_DIR
 from mani_skill2.agents.registration import register_agent
 from mani_skill2.agents.robots.allegro_hand.allegro import AllegroHandRight
-from mani_skill2.utils.sapien_utils import (
-    compute_total_impulse,
-    get_actors_contacts,
-    get_multiple_pairwise_contacts,
-    get_objs_by_names,
-)
+from mani_skill2.utils import sapien_utils
 from mani_skill2.utils.structs.actor import Actor
 
 
@@ -59,7 +54,7 @@ def __init__(self, *args, **kwargs):
 
     def _after_init(self):
         super()._after_init()
-        self.fsr_links: List[Actor] = get_objs_by_names(
+        self.fsr_links: List[Actor] = sapien_utils.get_objs_by_names(
             self.robot.get_links(),
             self.palm_fsr_link_names + self.finger_fsr_link_names,
         )
@@ -100,7 +95,8 @@ def get_fsr_obj_impulse(self, obj: Actor = None):
             )
             sorted_contacts = [obj_contacts[link] for link in internal_fsr_links]
             contact_forces = [
-                compute_total_impulse(contact) for contact in sorted_contacts
+                sapien_utils.compute_total_impulse(contact)
+                for contact in sorted_contacts
             ]
 
             return np.stack(contact_forces)
@@ -133,10 +129,11 @@ def get_fsr_impulse(self):
         else:
             internal_fsr_links = [link._bodies[0].entity for link in self.fsr_links]
             contacts = self.scene.get_contacts()
-            contact_map = get_actors_contacts(contacts, internal_fsr_links)
+            contact_map = sapien_utils.get_actors_contacts(contacts, internal_fsr_links)
             sorted_contacts = [contact_map[link] for link in internal_fsr_links]
             contact_forces = [
-                compute_total_impulse(contact) for contact in sorted_contacts
+                sapien_utils.compute_total_impulse(contact)
+                for contact in sorted_contacts
             ]
 
             contact_impulse = torch.from_numpy(

mani_skill2/agents/robots/anymal/__init__.py

@@ -0,0 +1 @@
+from .anymal_c import ANYmalC

mani_skill2/agents/robots/anymal/anymal_c.py

@@ -1,14 +1,11 @@
+# TODO (stao): Anymal may not be modelled correctly or efficiently at the moment
 import torch
 
 from mani_skill2 import PACKAGE_ASSET_DIR, format_path
 from mani_skill2.agents.base_agent import BaseAgent
 from mani_skill2.agents.controllers import *
 from mani_skill2.agents.registration import register_agent
-from mani_skill2.utils.sapien_utils import (
-    apply_urdf_config,
-    check_urdf_config,
-    parse_urdf_config,
-)
+from mani_skill2.utils import sapien_utils
 from mani_skill2.utils.structs.articulation import Articulation
 
 
@@ -82,11 +79,11 @@ def _load_articulation(self):
 
         urdf_path = format_path(str(self.urdf_path))
 
-        urdf_config = parse_urdf_config(self.urdf_config, self.scene)
-        check_urdf_config(urdf_config)
+        urdf_config = sapien_utils.parse_urdf_config(self.urdf_config, self.scene)
+        sapien_utils.check_urdf_config(urdf_config)
 
         # TODO(jigu): support loading multiple convex collision shapes
-        apply_urdf_config(loader, urdf_config)
+        sapien_utils.apply_urdf_config(loader, urdf_config)
         loader.disable_self_collisions = True
         self.robot: Articulation = loader.load(urdf_path)
         assert self.robot is not None, f"Fail to load URDF from {urdf_path}"

mani_skill2/agents/robots/dclaw/dclaw.py

@@ -8,7 +8,7 @@
 from mani_skill2.agents.controllers import *
 from mani_skill2.agents.registration import register_agent
 from mani_skill2.agents.utils import get_active_joint_indices
-from mani_skill2.utils.sapien_utils import get_objs_by_names
+from mani_skill2.utils import sapien_utils
 from mani_skill2.utils.structs.joint import Joint
 from mani_skill2.utils.structs.link import Link
 from mani_skill2.utils.structs.pose import vectorize_pose
@@ -52,7 +52,7 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
     def _after_init(self):
-        self.tip_links: List[Link] = get_objs_by_names(
+        self.tip_links: List[Link] = sapien_utils.get_objs_by_names(
             self.robot.get_links(), self.tip_link_names
         )
         self.root_joints: List[Joint] = [

mani_skill2/agents/robots/fetch/fetch.py

@@ -11,14 +11,8 @@
 from mani_skill2.agents.controllers import *
 from mani_skill2.agents.registration import register_agent
 from mani_skill2.sensors.camera import CameraConfig
+from mani_skill2.utils import sapien_utils
 from mani_skill2.utils.common import compute_angle_between, np_compute_angle_between
-from mani_skill2.utils.sapien_utils import (
-    compute_total_impulse,
-    get_actor_contacts,
-    get_obj_by_name,
-    get_pairwise_contact_impulse,
-    to_tensor,
-)
 from mani_skill2.utils.structs.actor import Actor
 from mani_skill2.utils.structs.base import BaseStruct
 from mani_skill2.utils.structs.joint import Joint
@@ -332,19 +326,23 @@ def controller_configs(self):
         return deepcopy_dict(controller_configs)
 
     def _after_init(self):
-        self.finger1_link: Link = get_obj_by_name(
+        self.finger1_link: Link = sapien_utils.get_obj_by_name(
             self.robot.get_links(), "l_gripper_finger_link"
         )
-        self.finger2_link: Link = get_obj_by_name(
+        self.finger2_link: Link = sapien_utils.get_obj_by_name(
             self.robot.get_links(), "r_gripper_finger_link"
         )
-        self.tcp: Link = get_obj_by_name(self.robot.get_links(), self.ee_link_name)
+        self.tcp: Link = sapien_utils.get_obj_by_name(
+            self.robot.get_links(), self.ee_link_name
+        )
 
-        self.base_link: Link = get_obj_by_name(self.robot.get_links(), "base_link")
-        self.l_wheel_link: Link = get_obj_by_name(
+        self.base_link: Link = sapien_utils.get_obj_by_name(
+            self.robot.get_links(), "base_link"
+        )
+        self.l_wheel_link: Link = sapien_utils.get_obj_by_name(
             self.robot.get_links(), "l_wheel_link"
         )
-        self.r_wheel_link: Link = get_obj_by_name(
+        self.r_wheel_link: Link = sapien_utils.get_obj_by_name(
             self.robot.get_links(), "r_wheel_link"
         )
         for link in [self.l_wheel_link, self.r_wheel_link]:
@@ -354,11 +352,11 @@ def _after_init(self):
                 cg[2] |= FETCH_UNIQUE_COLLISION_BIT
                 cs.set_collision_groups(cg)
 
-        self.torso_lift_link: Link = get_obj_by_name(
+        self.torso_lift_link: Link = sapien_utils.get_obj_by_name(
             self.robot.get_links(), "torso_lift_link"
         )
 
-        self.head_camera_link: Link = get_obj_by_name(
+        self.head_camera_link: Link = sapien_utils.get_obj_by_name(
             self.robot.get_links(), "head_camera_link"
         )
 
@@ -405,25 +403,27 @@ def is_grasping(self, object: Actor = None, min_impulse=1e-6, max_angle=85):
             contacts = self.scene.get_contacts()
 
             if object is None:
-                finger1_contacts = get_actor_contacts(
+                finger1_contacts = sapien_utils.get_actor_contacts(
                     contacts, self.finger1_link._bodies[0].entity
                 )
-                finger2_contacts = get_actor_contacts(
+                finger2_contacts = sapien_utils.get_actor_contacts(
                     contacts, self.finger2_link._bodies[0].entity
                 )
                 return (
-                    np.linalg.norm(compute_total_impulse(finger1_contacts))
+                    np.linalg.norm(sapien_utils.compute_total_impulse(finger1_contacts))
                     >= min_impulse
-                    and np.linalg.norm(compute_total_impulse(finger2_contacts))
+                    and np.linalg.norm(
+                        sapien_utils.compute_total_impulse(finger2_contacts)
+                    )
                     >= min_impulse
                 )
             else:
-                limpulse = get_pairwise_contact_impulse(
+                limpulse = sapien_utils.get_pairwise_contact_impulse(
                     contacts,
                     self.finger1_link._bodies[0].entity,
                     object._bodies[0].entity,
                 )
-                rimpulse = get_pairwise_contact_impulse(
+                rimpulse = sapien_utils.get_pairwise_contact_impulse(
                     contacts,
                     self.finger2_link._bodies[0].entity,
                     object._bodies[0].entity,