infra: add contact force sensor cfg; isaaclab: support it

metasim/cfg/scenario.py

@@ -19,7 +19,7 @@
 from .render import RenderCfg
 from .robots.base_robot_cfg import BaseRobotCfg
 from .scenes.base_scene_cfg import SceneCfg
-from .sensors import BaseCameraCfg, PinholeCameraCfg
+from .sensors import BaseCameraCfg, BaseSensorCfg, PinholeCameraCfg
 from .tasks.base_task_cfg import BaseTaskCfg
 
 
@@ -49,6 +49,7 @@ class ScenarioCfg:
     lights: list[BaseLightCfg] = [DistantLightCfg()]
     objects: list[BaseObjCfg] = []
     cameras: list[BaseCameraCfg] = [PinholeCameraCfg()]
+    sensors: list[BaseSensorCfg] = []
     checker: BaseChecker = EmptyChecker()
     render: RenderCfg = RenderCfg()
     random: RandomizationCfg = RandomizationCfg()

metasim/cfg/sensors/__init__.py

@@ -2,4 +2,6 @@
 
 """Sub-module containing the camera configuration."""
 
+from .base_sensor import BaseSensorCfg
 from .cameras import BaseCameraCfg, PinholeCameraCfg
+from .contact import ContactForceSensorCfg

metasim/cfg/sensors/base_sensor.py

@@ -0,0 +1,13 @@
+"""Sub-module containing the base sensor configuration."""
+
+from dataclasses import MISSING
+
+from metasim.utils.configclass import configclass
+
+
+@configclass
+class BaseSensorCfg:
+    """Base sensor configuration."""
+
+    name: str = MISSING
+    """Sensor name"""

metasim/cfg/sensors/contact.py

@@ -0,0 +1,25 @@
+"""Sub-module containing the contact force sensor configuration."""
+
+from __future__ import annotations
+
+from dataclasses import MISSING
+
+from metasim.cfg.sensors import BaseSensorCfg
+from metasim.utils.configclass import configclass
+
+
+@configclass
+class ContactForceSensorCfg(BaseSensorCfg):
+    """Contact force sensor cfg."""
+
+    base_link: str | tuple[str, str] = MISSING
+    """Body link to feel the contact force.
+        If a ``str``, the sensor will be attached to the root link of the object specified by the name.
+        If a ``tuple[str, str]``, the sensor will be attached to the body link specified by the second str of the object specified by the first str.
+    """
+    source_link: str | tuple[str, str] | None = None
+    """Body link to feel the contact force from.
+        If ``None``, the sensor will feel the contact force from all the source links.
+        If a ``str``, the sensor will only feel the contact force from the root link of the object specified by the name.
+        If a ``tuple[str, str]``, the sensor will only feel the contact force from the body link specified by the second str of the object specified by the first str.
+    """

metasim/sim/base.py

@@ -27,6 +27,7 @@ def __init__(self, scenario: ScenarioCfg):
         self.task = scenario.task
         self.robot = scenario.robot
         self.cameras = scenario.cameras
+        self.sensors = scenario.sensors
         self.objects = scenario.objects
         self.checker = scenario.checker
         self.object_dict = {obj.name: obj for obj in self.objects + [self.robot] + self.checker.get_debug_viewers()}

metasim/sim/isaaclab/env_overwriter.py

@@ -5,7 +5,7 @@
 from metasim.cfg.scenario import ScenarioCfg
 from metasim.utils.camera_util import get_cam_params
 
-from .isaaclab_helper import add_lights, add_objects, add_robot, get_pose
+from .isaaclab_helper import add_lights, add_objects, add_robot, add_sensors, get_pose
 
 try:
     from .empty_env import EmptyEnv
@@ -34,6 +34,7 @@ def __init__(self, scenario: ScenarioCfg):
         self.task = scenario.task
         self.robot = scenario.robot
         self.cameras = scenario.cameras
+        self.sensors = scenario.sensors
         self.objects = scenario.objects
         self.scene = scenario.scene
         self.checker = scenario.checker
@@ -351,6 +352,9 @@ def _setup_scene(self, env: "EmptyEnv") -> None:
                 )
             )
 
+        ## Add sensors
+        add_sensors(env, self.sensors)
+
     def _pre_physics_step(self, env: "EmptyEnv", actions: torch.Tensor) -> None:
         ## TODO: Clip action or not?
         env.actions = actions

metasim/sim/isaaclab/isaaclab.py

@@ -9,9 +9,10 @@
 
 from metasim.cfg.objects import ArticulationObjCfg, BaseObjCfg, PrimitiveFrameCfg, RigidObjCfg
 from metasim.cfg.scenario import ScenarioCfg
+from metasim.cfg.sensors import ContactForceSensorCfg
 from metasim.sim import BaseSimHandler, EnvWrapper, IdentityEnvWrapper
 from metasim.types import Action, EnvState, Extra, Obs, Reward, Success, TimeOut
-from metasim.utils.state import CameraState, ObjectState, RobotState, TensorState
+from metasim.utils.state import CameraState, ContactForceState, ObjectState, RobotState, TensorState
 
 from .env_overwriter import IsaaclabEnvOverwriter
 from .isaaclab_helper import get_pose
@@ -357,7 +358,19 @@ def get_states(self, env_ids: list[int] | None = None) -> TensorState:
             depth_data = camera_inst.data.output.get("depth", None)
             camera_states[camera.name] = CameraState(rgb=rgb_data, depth=depth_data)
 
-        return TensorState(objects=object_states, robots=robot_states, cameras=camera_states)
+        sensor_states = {}
+        for sensor in self.sensors:
+            if isinstance(sensor, ContactForceSensorCfg):
+                sensor_inst = self.env.scene.sensors[sensor.name]
+                if sensor.source_link is None:
+                    force = sensor_inst.data.net_forces_w.squeeze(1)
+                else:
+                    force = sensor_inst.data.force_matrix_w.squeeze((1, 2))
+                sensor_states[sensor.name] = ContactForceState(force=force)
+            else:
+                raise ValueError(f"Unknown sensor type: {type(sensor)}")
+
+        return TensorState(objects=object_states, robots=robot_states, cameras=camera_states, sensors=sensor_states)
 
     def get_pos(self, obj_name: str, env_ids: list[int] | None = None) -> torch.FloatTensor:
         if env_ids is None:

metasim/sim/isaaclab/isaaclab_helper.py

@@ -14,6 +14,7 @@
     RigidObjCfg,
 )
 from metasim.cfg.robots import BaseRobotCfg
+from metasim.cfg.sensors import BaseSensorCfg, ContactForceSensorCfg
 
 try:
     from .empty_env import EmptyEnv
@@ -150,6 +151,7 @@ def add_robot(env: "EmptyEnv", robot: BaseRobotCfg) -> None:
     cfg = ArticulationCfg(
         spawn=sim_utils.UsdFileCfg(
             usd_path=robot.usd_path,
+            activate_contact_sensors=True,  # TODO: only activate when contact sensor is added
             rigid_props=sim_utils.RigidBodyPropertiesCfg(),
             articulation_props=sim_utils.ArticulationRootPropertiesCfg(),
         ),
@@ -208,6 +210,66 @@ def add_lights(env: "EmptyEnv", lights: list[BaseLightCfg]) -> None:
             _add_light(env, light, f"/World/envs/env_0/lights/light_{i}")
 
 
+def _add_contact_force_sensor(env: "EmptyEnv", sensor: ContactForceSensorCfg) -> None:
+    try:
+        import omni.isaac.core.utils.prims as prim_utils
+        from omni.isaac.lab.sensors import ContactSensor, ContactSensorCfg
+    except ModuleNotFoundError:
+        import isaacsim.core.utils.prims as prim_utils
+        from isaaclab.sensors import ContactSensor, ContactSensorCfg
+
+    if isinstance(sensor, ContactForceSensorCfg):
+        _base_prim_regex_path = (
+            f"/World/envs/env_0/{sensor.base_link}"
+            if isinstance(sensor.base_link, str)
+            else f"/World/envs/env_0/{sensor.base_link[0]}/.*{sensor.base_link[1]}"  # TODO: improve the regex
+        )
+        _base_prim_paths = prim_utils.find_matching_prim_paths(_base_prim_regex_path)
+        if len(_base_prim_paths) == 0:
+            log.error(f"Base link {sensor.base_link} of cotact force sensor not found")
+            return
+        if len(_base_prim_paths) > 1:
+            log.warning(
+                f"Multiple base links found for contact force sensor {sensor.name}, using the first one: {_base_prim_paths[0]}"
+            )
+        base_prim_path = _base_prim_paths[0]
+        log.info(f"Base prim path: {base_prim_path}")
+        if sensor.source_link is not None:
+            _source_prim_regex_path = (
+                f"/World/envs/env_0/{sensor.source_link}"
+                if isinstance(sensor.source_link, str)
+                else f"/World/envs/env_0/{sensor.source_link[0]}/.*{sensor.source_link[1]}"  # TODO: improve the regex
+            )
+            _source_prim_paths = prim_utils.find_matching_prim_paths(_source_prim_regex_path)
+            if len(_source_prim_paths) == 0:
+                log.error(f"Source link {sensor.source_link} of cotact force sensor not found")
+                return
+            if len(_source_prim_paths) > 1:
+                log.warning(
+                    f"Multiple source links found for contact force sensor {sensor.name}, using the first one: {_source_prim_paths[0]}"
+                )
+            source_prim_path = _source_prim_paths[0]
+        else:
+            source_prim_path = None
+
+        env.scene.sensors[sensor.name] = ContactSensor(
+            ContactSensorCfg(
+                prim_path=base_prim_path.replace("env_0", "env_.*"),  # HACK: this is so hacky
+                filter_prim_paths_expr=[source_prim_path.replace("env_0", "env_.*")]  # HACK: this is so hacky
+                if source_prim_path is not None
+                else [],
+                history_length=6,  # XXX: hard-coded
+                update_period=0.0,  # XXX: hard-coded
+            )
+        )
+
+
+def add_sensors(env: "EmptyEnv", sensors: list[BaseSensorCfg]) -> None:
+    for sensor in sensors:
+        if isinstance(sensor, ContactForceSensorCfg):
+            _add_contact_force_sensor(env, sensor)
+
+
 def get_pose(
     env: "EmptyEnv", obj_name: str, obj_subpath: str | None = None, env_ids: list[int] | None = None
 ) -> tuple[torch.FloatTensor, torch.FloatTensor]:

metasim/utils/state.py

@@ -16,6 +16,17 @@
     pass
 
 
+@dataclass
+class ContactForceState:
+    """State of a single contact force sensor."""
+
+    force: torch.Tensor
+    """Contact force. Shape is (num_envs, 3)."""
+
+
+SensorState = ContactForceState
+
+
 @dataclass
 class ObjectState:
     """State of a single object."""
@@ -74,6 +85,8 @@ class TensorState:
     """States of all robots."""
     cameras: dict[str, CameraState]
     """States of all cameras."""
+    sensors: dict[str, SensorState]
+    """States of all sensors."""
 
 
 def _dof_tensor_to_dict(dof_tensor: torch.Tensor, joint_names: list[str]) -> dict[str, float]: