Source code for embodichain.lab.sim.sensors.contact_sensor
# ----------------------------------------------------------------------------
# Copyright (c) 2021-2026 DexForce Technology Co., Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ----------------------------------------------------------------------------
from __future__ import annotations
import uuid
from typing import TYPE_CHECKING, Sequence
import dexsim
import numpy as np
import torch
import warp as wp
from tensordict import TensorDict
from embodichain.lab.sim.sensors import BaseSensor, SensorCfg
from embodichain.lab.sim.sensors._warp.contact import scatter_contact_rows
from embodichain.lab.sim.sensors.contact_history import ContactHistory
from embodichain.utils import configclass, logger
from embodichain.utils.device_utils import current_warp_stream
if TYPE_CHECKING:
from dexsim.scene import (
ContactActorInfo,
ContactBuffer,
ContactQuery,
ContactQueryCapabilities,
)
from embodichain.lab.sim.sim_manager import SimulationManager
__all__ = [
"ArticulationContactFilterCfg",
"ContactSensor",
"ContactSensorCfg",
]
[docs]
@configclass
class ContactSensorCfg(SensorCfg):
"""Configuration class for contact sensors.
This class defines the configuration for contact sensors that detect
collisions between rigid bodies and articulation links.
"""
rigid_uid_list: list[str] = []
"""rigid body contact filter configs"""
articulation_cfg_list: list[ArticulationContactFilterCfg] = []
"""articulation link contact filter configs"""
filter_need_both_actor: bool = True
"""Whether to filter contact only when both actors are in the filter list."""
max_contacts_per_env: int = 64
"""Maximum number of contacts per environment the sensor can handle."""
sensor_type: str = "ContactSensor"
[docs]
@configclass
class ArticulationContactFilterCfg:
"""Configuration for filtering contacts from an articulation's links.
This class defines which articulation and which links to monitor
for contact events.
"""
articulation_uid: str = ""
"""Articulation unique identifier."""
link_name_list: list[str] = []
"""link names in the articulation whose contacts need to be filtered."""
[docs]
@classmethod
def from_dict(
cls, init_dict: dict[str, str | list[str]]
) -> "ArticulationContactFilterCfg":
"""Initialize the configuration from a dictionary.
Args:
init_dict: Dictionary containing configuration parameters.
Returns:
ArticulationContactFilterCfg: The initialized configuration.
"""
cfg = cls()
for key, value in init_dict.items():
if hasattr(cfg, key):
setattr(cfg, key, value)
else:
logger.log_warning(f"Key '{key}' not found in {cls.__name__}.")
return cfg
[docs]
class ContactSensor(BaseSensor):
"""Sensor to get contacts from rigid bodies and articulation links.
The sensor preserves every backend-emitted contact row that passes the
backend's positive-impulse filter. Geometry-only Newton solvers retain all
candidate contact rows.
"""
SUPPORTED_DATA_TYPES = [
"position",
"normal",
"friction",
"impulse",
"distance",
"user_ids",
"is_valid",
]
[docs]
def __init__(
self,
config: ContactSensorCfg,
device: torch.device = torch.device("cpu"),
*,
owner: "SimulationManager | None" = None,
) -> None:
if owner is None:
from embodichain.lab.sim.sim_manager import SimulationManager
owner = SimulationManager.get_instance()
self._sim = owner
self.item_user_ids: torch.Tensor | None = None
"""Backend-neutral actor IDs selected by the contact query."""
self.item_env_ids: torch.Tensor | None = None
"""Environment IDs of the selected contact actors."""
self.item_user_env_ids_map: torch.Tensor | None = None
"""Compatibility map from contact actor ID to environment ID."""
self._visualizer: dexsim.models.PointCloud | None = None
"""contact point visualizer. Default to None"""
self.device = device
self.cfg = config
self._query: ContactQuery | None = None
self._metadata_actor_table: tuple[ContactActorInfo, ...] | None = None
self._metadata_selected_ids: tuple[int, ...] | None = None
self._query_dropped_count = torch.zeros(1, dtype=torch.int32, device=device)
self._scatter_dropped_count = torch.zeros(
owner.num_envs, dtype=torch.int32, device=device
)
self._sample_scatter_dropped_count = torch.zeros_like(
self._scatter_dropped_count
)
self._histories: dict[str, ContactHistory] = {}
self._sample_graph = None
self._sample_graph_key = None
self._sample_stream = None
self._sample_torch_stream = None
self._num_contacts_per_env: torch.Tensor | None = None
"""Number of contacts per environment."""
super().__init__(config, device, num_instances=owner.num_envs)
@property
def max_total_contacts(self) -> int:
"""Get the maximum total number of contacts across all environments.
Returns:
int: Maximum total number of contacts.
"""
return self.cfg.max_contacts_per_env * self.num_instances
@property
def total_current_contacts(self) -> int:
"""Get the current total number of contacts across all environments.
Note:
This method returns the total number of contacts detected in the most recent update.
Returns:
int: Total number of contacts.
"""
assert self._num_contacts_per_env is not None
return int(self._num_contacts_per_env.sum().item())
@property
def contact_capabilities(self) -> "ContactQueryCapabilities":
"""Capabilities reported by the active DexSim contact binding."""
assert self._query is not None
return self._query.capabilities
def _build_sensor_from_config(
self,
config: ContactSensorCfg,
device: torch.device,
) -> None:
result = self._sim.spawn_result
if result is None:
raise RuntimeError("ContactSensor requires SimulationManager.prepare().")
targets: list[object] = []
for uid in config.rigid_uid_list:
try:
handles = self._sim._spawn_scene.handles(uid)
except KeyError as exc:
raise KeyError(f"Contact rigid-body UID not found: {uid!r}.") from exc
if not handles:
raise RuntimeError(f"Contact rigid body {uid!r} is not materialized.")
targets.extend(handles)
for articulation_cfg in config.articulation_cfg_list:
uid = articulation_cfg.articulation_uid
try:
handles = self._sim._spawn_scene.handles(uid)
except KeyError as exc:
raise KeyError(f"Contact articulation UID not found: {uid!r}.") from exc
if not handles:
raise RuntimeError(f"Contact articulation {uid!r} is not materialized.")
if not articulation_cfg.link_name_list:
targets.extend(handles)
continue
for handle in handles:
available = set(handle.get_link_names())
missing = set(articulation_cfg.link_name_list) - available
if missing:
raise ValueError(
f"Contact articulation {uid!r} has no links: {sorted(missing)}."
)
targets.extend(
(handle, link_name) for link_name in articulation_cfg.link_name_list
)
if not targets:
raise ValueError(
"ContactSensor requires at least one rigid or link target."
)
self._query = result.create_contact_query(
targets,
match="all" if config.filter_need_both_actor else "any",
capacity=self.max_total_contacts,
capacity_per_env=config.max_contacts_per_env,
device=device,
frame="arena",
)
self._sync_filter_actor_metadata()
num_envs = self.num_instances
self._num_contacts_per_env = torch.zeros(
num_envs, dtype=torch.int32, device=device
)
self._data_buffer = TensorDict(
{
"position": torch.zeros(
(num_envs, config.max_contacts_per_env, 3), device=device
),
"normal": torch.zeros(
(num_envs, config.max_contacts_per_env, 3), device=device
),
"friction": torch.zeros(
(num_envs, config.max_contacts_per_env, 3), device=device
),
"impulse": torch.zeros(
(num_envs, config.max_contacts_per_env), device=device
),
"distance": torch.zeros(
(num_envs, config.max_contacts_per_env), device=device
),
"user_ids": torch.zeros(
(num_envs, config.max_contacts_per_env, 2),
dtype=torch.int32,
device=device,
),
"is_valid": torch.zeros(
(num_envs, config.max_contacts_per_env),
dtype=torch.bool,
device=device,
),
},
batch_size=[num_envs, config.max_contacts_per_env],
device=device,
)
def _sync_filter_actor_metadata(self) -> None:
assert self._query is not None
# ContactQuery replaces these immutable tables when topology refreshes.
# Avoid visiting every actor and copying metadata on every physics step.
actors = self._query.actors
actor_ids = self._query.selected_actor_ids
if (
actors is self._metadata_actor_table
and actor_ids is self._metadata_selected_ids
):
return
self.item_user_ids = torch.as_tensor(
actor_ids, dtype=torch.int32, device=self.device
)
self.item_env_ids = torch.as_tensor(
[self._query.actor_info(actor_id).env_id for actor_id in actor_ids],
dtype=torch.int32,
device=self.device,
)
self.item_user_env_ids_map = torch.full(
(max(actor_ids, default=-1) + 1,),
-1,
dtype=torch.int32,
device=self.device,
)
if actor_ids:
self.item_user_env_ids_map[self.item_user_ids.to(torch.long)] = (
self.item_env_ids
)
self._metadata_actor_table = actors
self._metadata_selected_ids = actor_ids
[docs]
def update(self, **kwargs) -> None:
"""Update the sensor state based on the current simulation state.
This method is called periodically to ensure the sensor data is up-to-date.
Args:
**kwargs: Additional keyword arguments for sensor update.
"""
buffer = self._fetch_contacts()
if not self._histories:
self._query_dropped_count.copy_(buffer.dropped_count_device)
self._scatter_dropped_count.copy_(self._sample_scatter_dropped_count)
def _fetch_contacts(self) -> ContactBuffer:
buffer = self._read_contacts()
self._scatter_contacts(buffer)
return buffer
def _read_contacts(self) -> ContactBuffer:
assert self._query is not None
if self.device.type == "cuda":
buffer = self._query.fetch_async()
else:
buffer = self._query.fetch()
buffer.device_counts[0] = buffer.count
buffer.device_counts[1] = buffer.dropped_count
self._sync_filter_actor_metadata()
return buffer
def _scatter_contacts(self, buffer: ContactBuffer) -> None:
assert self._num_contacts_per_env is not None
self._num_contacts_per_env.zero_()
self._data_buffer["is_valid"].zero_()
self._sample_scatter_dropped_count.zero_()
wp.launch(
kernel=scatter_contact_rows,
dim=buffer.capacity,
inputs=[
wp.from_torch(buffer.data),
wp.from_torch(buffer.actor_ids),
wp.from_torch(buffer.env_ids),
wp.from_torch(self._num_contacts_per_env),
self.cfg.max_contacts_per_env,
*[
wp.from_torch(self._data_buffer[name])
for name in (
"position",
"normal",
"friction",
"impulse",
"distance",
"user_ids",
"is_valid",
)
],
wp.from_torch(buffer.count_device),
wp.from_torch(self._sample_scatter_dropped_count),
],
device=str(self.device),
stream=current_warp_stream(self.device),
)
@property
def dropped_contacts(self) -> int:
"""Rows lost across the current control interval, or latest standalone update.
This diagnostic reads device counters only when requested. Query overflow
and per-environment scatter overflow are counted separately and summed.
"""
return int(
(self._query_dropped_count.sum() + self._scatter_dropped_count.sum()).item()
)
[docs]
def get_arena_pose(self, to_matrix: bool = False) -> torch.Tensor | None:
"""Not used.
Args:
to_matrix: If True, return the pose as a 4x4 transformation matrix.
Returns:
A tensor representing the pose of the sensor in the arena frame.
"""
logger.log_error("`get_arena_pose` for contact sensor is not implemented yet.")
return None
[docs]
def get_local_pose(self, to_matrix: bool = False) -> torch.Tensor | None:
"""Get the local pose of the camera.
Args:
to_matrix (bool): If True, return the pose as a 4x4 matrix. If False, return as a quaternion.
Returns:
torch.Tensor: The local pose of the camera.
"""
logger.log_error("`get_local_pose` for contact sensor is not implemented yet.")
[docs]
def set_local_pose(
self, pose: torch.Tensor, env_ids: Sequence[int] | None = None
) -> None:
"""Set the local pose of the camera.
Note: The pose should be in the OpenGL coordinate system, which means the Y is up and Z is forward.
Args:
pose (torch.Tensor): The local pose to set, should be a 4x4 transformation matrix.
env_ids (Sequence[int] | None): The environment IDs to set the pose for. If None, set for all environments.
"""
logger.log_error("`set_local_pose` for contact sensor is not implemented yet.")
[docs]
def get_data(self) -> TensorDict:
"""Retrieve data from the sensor.
Returns:
Batched contact data. ``position`` is in the Arena frame;
``normal`` points from actor 0 toward actor 1; ``friction`` and
``impulse`` are impulses; ``distance`` is signed separation;
``user_ids`` contains backend-neutral contact actor IDs; and
``is_valid`` marks rows populated by the latest update. Values in
invalid rows are unspecified and may come from an earlier update.
"""
return self._data_buffer
[docs]
def get_actor_info(self, actor_id: int) -> "ContactActorInfo":
"""Resolve an ID from the ``user_ids`` field to its Spawn identity."""
assert self._query is not None
return self._query.actor_info(actor_id)
[docs]
def get_actor_ids(
self, uid: str, link_names: Sequence[str] | None = None
) -> torch.Tensor:
"""Resolve a scene entity's ordered contact IDs for every environment.
Args:
uid: Robot, articulation or rigid-object UID in the Spawn scene.
link_names: Ordered articulation links; None selects the rigid actor.
Returns:
IDs shaped (environments, selected bodies). Global actors are broadcast.
"""
assert self._query is not None
handles = self._sim._spawn_scene.handles(uid)
paths = {handle.path for handle in handles}
names = (None,) if link_names is None else tuple(link_names)
actors = {
(actor.env_id, actor.link_name): actor.actor_id
for actor in self._query.actors
if actor.path in paths
}
rows = []
for env_id in range(self.num_instances):
row = []
for name in names:
key = (env_id, name) if (env_id, name) in actors else (-1, name)
if key not in actors:
raise ValueError(
f"No contact actor for {uid}/{name} in row {env_id}."
)
row.append(actors[key])
rows.append(row)
return torch.tensor(rows, device=self.device, dtype=torch.int32).reshape(
self.num_instances, len(names)
)
[docs]
def create_history(
self,
name: str,
actor_ids: torch.Tensor,
*,
counterpart_ids: torch.Tensor | None = None,
include_unknown_counterpart: bool = False,
force_threshold: float = 0.0,
) -> ContactHistory:
"""Register a sensor-owned reduction with explicit counterpart filtering.
Args:
name: Unique history name within the sensor.
actor_ids: Body IDs to monitor, in per-environment order.
counterpart_ids: Allowed counterparts; None accepts all contacts.
include_unknown_counterpart: Also accept unidentified query counterparts.
force_threshold: Minimum contact force in newtons for timing events.
Returns:
History that automatically enables substep sampling and selected-row reset.
"""
if name in self._histories:
raise ValueError(f"Contact history {name!r} already exists.")
history = ContactHistory(
actor_ids,
counterpart_ids=counterpart_ids,
include_unknown_counterpart=include_unknown_counterpart,
force_threshold=force_threshold,
)
self._histories[name] = history
return history
@property
def requires_substep_update(self) -> bool:
"""Whether this sensor participates in the physics-substep lifecycle."""
return bool(self._histories)
[docs]
def begin_control_step(self) -> None:
"""Start a new interval without discarding persistent contact timing."""
self._query_dropped_count.zero_()
self._scatter_dropped_count.zero_()
for history in self._histories.values():
history.begin_control_step()
[docs]
def update_physics_step(self, dt: float) -> None:
"""Sample contacts and advance registered histories after physics.
Args:
dt: Elapsed physics time in seconds.
"""
if dt <= 0:
raise ValueError("Contact sampling dt must be positive.")
buffer = self._read_contacts()
if self.device.type != "cuda" or torch.cuda.is_current_stream_capturing():
self._sample_contacts(buffer, dt)
return
key = (
id(buffer),
dt,
self.cfg.max_contacts_per_env,
tuple(
(
id(h),
h.force_threshold,
h.include_unknown_counterpart,
h.counterpart_ids is None,
)
for h in self._histories.values()
),
)
if self._sample_stream is None:
self._sample_stream = wp.get_stream(str(self.device))
self._sample_torch_stream = wp.stream_to_torch(self._sample_stream)
caller = torch.cuda.current_stream(self.device)
same_stream = caller.cuda_stream == self._sample_torch_stream.cuda_stream
if not same_stream:
self._sample_torch_stream.wait_stream(caller)
with (
wp.ScopedStream(self._sample_stream),
torch.cuda.stream(self._sample_torch_stream),
):
if key != self._sample_graph_key:
self._sample_graph = None
self._sample_graph_key = key
# This real sample also compiles/warms every kernel. Do not
# execute history updates twice while preparing capture.
self._sample_contacts(buffer, dt)
else:
if self._sample_graph is None:
with wp.ScopedCapture(stream=self._sample_stream) as capture:
self._sample_contacts(buffer, dt)
self._sample_graph = capture.graph
wp.capture_launch(self._sample_graph, stream=self._sample_stream)
if not same_stream:
caller.wait_stream(self._sample_torch_stream)
def _sample_contacts(self, buffer: ContactBuffer, dt: float) -> None:
self._scatter_contacts(buffer)
self._query_dropped_count.add_(buffer.dropped_count_device)
self._scatter_dropped_count.add_(self._sample_scatter_dropped_count)
for history in self._histories.values():
history._update_from_query(buffer, dt)
[docs]
def reset(self, env_ids: Sequence[int] | None = None) -> None:
"""Clear contact samples and histories only for the selected rows.
Args:
env_ids: Rows to reset. None selects every row.
"""
ids = (
None
if env_ids is None
else torch.as_tensor(env_ids, dtype=torch.long).to(
self.device, non_blocking=True
)
)
for value in (
self._data_buffer["is_valid"],
self._num_contacts_per_env,
self._scatter_dropped_count,
self._sample_scatter_dropped_count,
):
if ids is None:
value.zero_()
else:
value.index_fill_(0, ids, 0)
if env_ids is None:
self._query_dropped_count.zero_()
for history in self._histories.values():
history.reset(ids)
[docs]
def filter_by_user_ids(
self, item_user_ids: torch.Tensor, env_ids: Sequence[int] | None = None
) -> TensorDict:
"""Filter contact report by backend-neutral contact actor IDs.
Args:
item_user_ids: Actor IDs from this sensor's ``user_ids`` field.
env_ids: Environment IDs to filter. If None, filter all environments.
Returns:
data: A TensorDict containing only the filtered contacts for the specified environments.
"""
env_ids_tensor = (
torch.arange(self.num_instances, device=self.device)
if env_ids is None
else torch.as_tensor(list(env_ids), dtype=torch.long, device=self.device)
)
item_user_ids = item_user_ids.to(device=self.device, dtype=torch.int32)
# Flatten data across all specified environments
env_data = {
"position": self._data_buffer["position"][env_ids_tensor].flatten(0, 1),
"normal": self._data_buffer["normal"][env_ids_tensor].flatten(0, 1),
"friction": self._data_buffer["friction"][env_ids_tensor].flatten(0, 1),
"impulse": self._data_buffer["impulse"][env_ids_tensor].flatten(0, 1),
"distance": self._data_buffer["distance"][env_ids_tensor].flatten(0, 1),
"user_ids": self._data_buffer["user_ids"][env_ids_tensor].flatten(0, 1),
"is_valid": self._data_buffer["is_valid"][env_ids_tensor].flatten(0, 1),
}
# Create valid mask (only slots up to _num_contacts_per_env are valid)
num_envs_to_filter = len(env_ids_tensor)
valid_mask = (
torch.arange(self.cfg.max_contacts_per_env, device=self.device).expand(
num_envs_to_filter, -1
)
< self._num_contacts_per_env[env_ids_tensor][:, None]
)
valid_mask = valid_mask.flatten()
# Create user ID filter mask
user_ids_flat = env_data["user_ids"]
filter0_mask = torch.isin(user_ids_flat[:, 0], item_user_ids)
filter1_mask = torch.isin(user_ids_flat[:, 1], item_user_ids)
if self.cfg.filter_need_both_actor:
filter_mask = torch.logical_and(filter0_mask, filter1_mask)
else:
filter_mask = torch.logical_or(filter0_mask, filter1_mask)
# Combine valid and user ID filters
combined_mask = torch.logical_and(valid_mask, filter_mask)
if not bool(combined_mask.any()):
# Return empty TensorDict if no matches
return TensorDict(
{
"position": torch.empty((0, 3), device=self.device),
"normal": torch.empty((0, 3), device=self.device),
"friction": torch.empty((0, 3), device=self.device),
"impulse": torch.empty((0,), device=self.device),
"distance": torch.empty((0,), device=self.device),
"user_ids": torch.empty(
(0, 2), dtype=torch.int32, device=self.device
),
"is_valid": torch.empty((0,), dtype=torch.bool, device=self.device),
},
batch_size=[0],
device=self.device,
)
# Extract filtered data using the combined mask
filtered_data = {key: value[combined_mask] for key, value in env_data.items()}
return TensorDict(
filtered_data,
batch_size=[filtered_data["position"].shape[0]],
device=self.device,
)
def set_contact_point_visibility(
self,
visible: bool = True,
rgba: Sequence[float] | None = None,
point_size: float = 3.0,
env_ids: Sequence[int] | None = None,
) -> None:
if env_ids is None:
env_ids = range(self.num_instances)
if visible:
# Convert env_ids to tensor if needed
env_ids_tensor = (
torch.tensor(env_ids, device=self.device)
if not isinstance(env_ids, torch.Tensor)
else env_ids
)
# Get number of contacts for each environment
num_contacts = self._num_contacts_per_env[env_ids_tensor]
# Create mask for valid contacts across all environments
# Shape: [num_envs, max_contacts_per_env]
contact_mask = torch.arange(
self.cfg.max_contacts_per_env, device=self.device
).unsqueeze(0) < num_contacts.unsqueeze(1)
if not contact_mask.any():
# No contacts to visualize
if isinstance(self._visualizer, dexsim.models.PointCloud):
self._visualizer.clear()
return
# Extract contact positions for all specified environments
# Shape: [num_envs, max_contacts_per_env, 3]
contact_position_arena = self._data_buffer["position"][env_ids_tensor]
# Get arena offsets and broadcast to match positions shape
# Shape: [num_envs, 1, 3] -> [num_envs, max_contacts_per_env, 3]
contact_offsets = self._sim.arena_offsets[env_ids_tensor].unsqueeze(1)
# Convert to world coordinates and apply mask in one go
contact_position_world = (contact_position_arena + contact_offsets)[
contact_mask
]
if self._visualizer is None:
# create new visualizer
temp_str = uuid.uuid4().hex
self._visualizer = self._sim.get_env().create_point_cloud(name=temp_str)
else:
# update existing visualizer points
self._visualizer.clear()
rgba = rgba if rgba is not None else (0.8, 0.2, 0.2, 1.0)
if len(rgba) != 4:
logger.log_error(
f"Invalid rgba {rgba}, should be a sequence of 4 floats."
)
rgba = np.array(
[
rgba[0],
rgba[1],
rgba[2],
rgba[3],
]
)
self._visualizer.add_points(
points=contact_position_world.to("cpu").numpy(), color=rgba
)
self._visualizer.set_point_size(point_size)
else:
if isinstance(self._visualizer, dexsim.models.PointCloud):
self._visualizer.clear()