From ade373c251d79b16eb7ec5119e2a4152656033ce Mon Sep 17 00:00:00 2001 From: Feishi Wang Date: Tue, 22 Apr 2025 01:30:07 -0700 Subject: [PATCH] infra: add contact force sensor cfg; isaaclab: support it --- metasim/cfg/scenario.py | 3 +- metasim/cfg/sensors/__init__.py | 2 + metasim/cfg/sensors/base_sensor.py | 13 ++++++ metasim/cfg/sensors/contact.py | 25 ++++++++++ metasim/sim/base.py | 1 + metasim/sim/isaaclab/env_overwriter.py | 6 ++- metasim/sim/isaaclab/isaaclab.py | 17 ++++++- metasim/sim/isaaclab/isaaclab_helper.py | 62 +++++++++++++++++++++++++ metasim/utils/state.py | 13 ++++++ 9 files changed, 138 insertions(+), 4 deletions(-) create mode 100644 metasim/cfg/sensors/base_sensor.py create mode 100644 metasim/cfg/sensors/contact.py diff --git a/metasim/cfg/scenario.py b/metasim/cfg/scenario.py index a156f2053..2f109c79b 100644 --- a/metasim/cfg/scenario.py +++ b/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() diff --git a/metasim/cfg/sensors/__init__.py b/metasim/cfg/sensors/__init__.py index 7e091a946..b2d23ef9d 100644 --- a/metasim/cfg/sensors/__init__.py +++ b/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 diff --git a/metasim/cfg/sensors/base_sensor.py b/metasim/cfg/sensors/base_sensor.py new file mode 100644 index 000000000..e376038da --- /dev/null +++ b/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""" diff --git a/metasim/cfg/sensors/contact.py b/metasim/cfg/sensors/contact.py new file mode 100644 index 000000000..c3783f167 --- /dev/null +++ b/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. + """ diff --git a/metasim/sim/base.py b/metasim/sim/base.py index 285990f52..fd6d2a27c 100644 --- a/metasim/sim/base.py +++ b/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()} diff --git a/metasim/sim/isaaclab/env_overwriter.py b/metasim/sim/isaaclab/env_overwriter.py index 2cac35c30..a68567d14 100644 --- a/metasim/sim/isaaclab/env_overwriter.py +++ b/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 diff --git a/metasim/sim/isaaclab/isaaclab.py b/metasim/sim/isaaclab/isaaclab.py index 942063cc3..39093198d 100644 --- a/metasim/sim/isaaclab/isaaclab.py +++ b/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: diff --git a/metasim/sim/isaaclab/isaaclab_helper.py b/metasim/sim/isaaclab/isaaclab_helper.py index 12b55d5ac..3a0971f96 100644 --- a/metasim/sim/isaaclab/isaaclab_helper.py +++ b/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]: diff --git a/metasim/utils/state.py b/metasim/utils/state.py index fe0297bab..940239265 100644 --- a/metasim/utils/state.py +++ b/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]: