Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
77 changes: 56 additions & 21 deletions genesis/engine/sensors/contact_force.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,13 @@
from genesis.options.sensors import Contact as ContactSensorOptions
from genesis.options.sensors import ContactForce as ContactForceSensorOptions
from genesis.utils.geom import inv_transform_by_quat, qd_inv_transform_by_quat, transform_by_quat
from genesis.utils.misc import concat_with_tensor, make_tensor_field, qd_to_torch, tensor_to_array
from genesis.utils.misc import (
append_filter_links_idx,
concat_with_tensor,
make_tensor_field,
qd_to_torch,
tensor_to_array,
)
from genesis.utils.ring_buffer import TensorRingBuffer

from .base_sensor import RigidSensorMetadataMixin, RigidSensorMixin, Sensor, SimpleSensor, SimpleSensorMetadata
Expand All @@ -23,13 +29,35 @@
from .sensor_manager import SensorManager


@qd.func
def _func_link_is_filtered(filter_links_idx: qd.types.ndarray(), i_s: int, link: int):
"""Whether ``link`` is in sensor ``i_s``'s filter list. ``-1`` padding entries never match a real link index."""
is_filtered = False
for i_f in range(filter_links_idx.shape[-1]):
if filter_links_idx[i_s, i_f] == link:
is_filtered = True
return is_filtered


def _drop_filtered_counterpart_contacts(
is_a: torch.Tensor, is_b: torch.Tensor, link_a, link_b, shared_metadata
) -> None:
filt = shared_metadata.filtered_sensor_idx
if filt.numel() == 0:
return
sub_filter = shared_metadata.filter_links_idx[filt][None, :, None, :]
is_a[:, filt, :] &= ~(link_b[:, None, :, None] == sub_filter).any(dim=-1)
is_b[:, filt, :] &= ~(link_a[:, None, :, None] == sub_filter).any(dim=-1)


@qd.kernel
def _kernel_get_contacts_forces(
contact_forces: qd.types.ndarray(),
link_a: qd.types.ndarray(),
link_b: qd.types.ndarray(),
links_quat: qd.types.ndarray(),
sensors_link_idx: qd.types.ndarray(),
filter_links_idx: qd.types.ndarray(), # (n_sensors, max_filter_links); unused slots are -1
output: qd.types.ndarray(),
):
for i_c, i_s, i_b in qd.ndrange(link_a.shape[-1], sensors_link_idx.shape[-1], output.shape[-1]):
Expand All @@ -51,10 +79,15 @@ def _kernel_get_contacts_forces(
force_a = qd_inv_transform_by_quat(-force_vec, quat_a)
force_b = qd_inv_transform_by_quat(force_vec, quat_b)

if contact_data_link_a == sensors_link_idx[i_s]:
# Accumulate the force on whichever side is the sensor link, unless the counterpart link is filtered.
if contact_data_link_a == sensors_link_idx[i_s] and not _func_link_is_filtered(
filter_links_idx, i_s, contact_data_link_b
):
for j in qd.static(range(3)):
output[j_s + j, i_b] += force_a[j]
if contact_data_link_b == sensors_link_idx[i_s]:
if contact_data_link_b == sensors_link_idx[i_s] and not _func_link_is_filtered(
filter_links_idx, i_s, contact_data_link_a
):
for j in qd.static(range(3)):
output[j_s + j, i_b] += force_b[j]

Expand Down Expand Up @@ -108,15 +141,10 @@ def build(self):
self._shared_metadata.expanded_links_idx, link_idx, expand=(1,), dim=0
)

num_sensors, cur_num_filter_links = self._shared_metadata.filter_links_idx.shape
max_num_filter_links = max(cur_num_filter_links, len(self._options.filter_link_idx))
filter_links_idx = torch.full((num_sensors + 1, max_num_filter_links), -1, dtype=gs.tc_int, device=gs.device)
filter_links_idx[:num_sensors, :cur_num_filter_links] = self._shared_metadata.filter_links_idx
filter_links_idx[num_sensors, : len(self._options.filter_link_idx)] = torch.tensor(
self._options.filter_link_idx, dtype=gs.tc_int, device=gs.device
num_sensors = self._shared_metadata.filter_links_idx.shape[0]
self._shared_metadata.filter_links_idx = append_filter_links_idx(
self._shared_metadata.filter_links_idx, self._options.filter_link_idx
)
self._shared_metadata.filter_links_idx = filter_links_idx

if len(self._options.filter_link_idx) > 0:
self._shared_metadata.filtered_sensor_idx = concat_with_tensor(
self._shared_metadata.filtered_sensor_idx, num_sensors, expand=(1,), dim=0
Expand Down Expand Up @@ -151,18 +179,9 @@ def _update_raw_data(cls, shared_context: None, shared_metadata: ContactSensorMe

is_contact_a = link_a[..., None, :] == shared_metadata.expanded_links_idx[..., None]
is_contact_b = link_b[..., None, :] == shared_metadata.expanded_links_idx[..., None]
_drop_filtered_counterpart_contacts(is_contact_a, is_contact_b, link_a, link_b, shared_metadata)
# Float-valued contact count per sensor (intermediate cache is float; bool projection in `_post_process`).
result = (is_contact_a | is_contact_b).sum(dim=-1).to(dtype=gs.tc_float)
# Apply the (more expensive) filter-aware update only on sensors that declared a filter; other sensors keep the
# cheap aggregate result above.
if shared_metadata.filtered_sensor_idx.numel() > 0:
filt = shared_metadata.filtered_sensor_idx
sub_filter = shared_metadata.filter_links_idx[filt][None, :, None, :]
filtered_a = (link_b[:, None, :, None] == sub_filter).any(dim=-1)
filtered_b = (link_a[:, None, :, None] == sub_filter).any(dim=-1)
sub_is_a = is_contact_a[:, filt, :]
sub_is_b = is_contact_b[:, filt, :]
result[:, filt] = ((sub_is_a & ~filtered_a) | (sub_is_b & ~filtered_b)).sum(dim=-1).to(dtype=gs.tc_float)
raw_data_T[:] = result.T

@classmethod
Expand Down Expand Up @@ -203,6 +222,11 @@ class ContactForceSensorMetadata(RigidSensorMetadataMixin, SimpleSensorMetadata)

min_force: torch.Tensor = make_tensor_field((0, 3))
max_force: torch.Tensor = make_tensor_field((0, 3))
# (num_force_sensors, max_num_filter_links); unused slots are -1.
filter_links_idx: torch.Tensor = make_tensor_field((0, 0), dtype_factory=lambda: gs.tc_int)
# Indices into links_idx of sensors that have at least one filter link. Lets the GT update skip the
# 4D contact-vs-filter comparison for the (typically larger) subset of sensors with no filter.
filtered_sensor_idx: torch.Tensor = make_tensor_field((0,), dtype_factory=lambda: gs.tc_int)


class ContactForceSensor(
Expand Down Expand Up @@ -238,6 +262,15 @@ def build(self):
self._shared_metadata.max_force, self._options.max_force, expand=(1, 3)
)

num_sensors = self._shared_metadata.filter_links_idx.shape[0]
self._shared_metadata.filter_links_idx = append_filter_links_idx(
self._shared_metadata.filter_links_idx, self._options.filter_link_idx
)
if len(self._options.filter_link_idx) > 0:
self._shared_metadata.filtered_sensor_idx = concat_with_tensor(
self._shared_metadata.filtered_sensor_idx, num_sensors, expand=(1,), dim=0
)

def _get_return_format(self) -> tuple[int, ...]:
return (3,)

Expand Down Expand Up @@ -276,6 +309,7 @@ def _update_raw_data(
# Forces are aggregated BEFORE moving them in local frame for efficiency.
force_mask_a = link_a[:, None] == shared_metadata.links_idx[None, :, None]
force_mask_b = link_b[:, None] == shared_metadata.links_idx[None, :, None]
_drop_filtered_counterpart_contacts(force_mask_a, force_mask_b, link_a, link_b, shared_metadata)
force_mask = force_mask_b.to(dtype=gs.tc_float) - force_mask_a.to(dtype=gs.tc_float)
sensors_force = (force_mask[..., None] * force[:, None]).sum(dim=2)
sensors_quat = links_quat[:, shared_metadata.links_idx]
Expand All @@ -290,6 +324,7 @@ def _update_raw_data(
link_b.contiguous(),
links_quat.contiguous(),
shared_metadata.links_idx,
shared_metadata.filter_links_idx,
raw_data_T,
)

Expand Down
67 changes: 50 additions & 17 deletions genesis/engine/sensors/kinematic_tactile.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
from genesis.options.sensors import ContactDepthProbe as ContactDepthProbeOptions
from genesis.options.sensors import ContactProbe as ContactProbeOptions
from genesis.options.sensors import KinematicTaxel as KinematicTaxelOptions
from genesis.utils.misc import concat_with_tensor, make_tensor_field, tensor_to_array
from genesis.utils.misc import append_filter_links_idx, concat_with_tensor, make_tensor_field, tensor_to_array
from genesis.utils.raycast_qd import (
closest_point_on_triangle,
get_triangle_vertices,
Expand All @@ -23,6 +23,7 @@
from .raycaster import RaycastContext

from .base_sensor import RigidSensorMetadataMixin, RigidSensorMixin, SimpleSensor, SimpleSensorMetadata
from .contact_force import _func_link_is_filtered
from .probe import (
ProbeSensorMetadataMixin,
ProbeSensorMixin,
Expand Down Expand Up @@ -52,6 +53,7 @@ def _func_query_contact_depth_penetration(
probe_pos: qd.types.vector(3),
probe_radius_gt: float,
probe_radius_m: float,
filter_links_idx: qd.types.ndarray(),
geoms_info: array_class.GeomsInfo,
geoms_state: array_class.GeomsState,
sensor_geoms_idx: qd.types.ndarray(),
Expand All @@ -60,22 +62,28 @@ def _func_query_contact_depth_penetration(
):
"""
Max probe penetration from SDF over the sensor link's unique opposing geoms, dual-radius.

Returns ``(max_pen_gt, max_pen_m)`` from a single SDF pass: both penetrations come from the same ``sd`` per geom
via ``pen = radius - sd``. Callers that do not need the noised-radius branch pass ``probe_radius_m ==
probe_radius_gt`` and ignore the second return. Geoms whose owning link is in sensor ``i_s``'s
``filter_links_idx`` list are skipped.
"""
max_pen_gt = gs.qd_float(0.0)
max_pen_m = gs.qd_float(0.0)

n_g = sensor_n_geoms[i_b, i_s]
for i_g_ in range(n_g):
i_g = sensor_geoms_idx[i_b, i_s, i_g_]
g_pos = geoms_state.pos[i_g, i_b]
g_quat = geoms_state.quat[i_g, i_b]
sd = sdf.sdf_func_world_local(geoms_info, sdf_info, probe_pos, i_g, g_pos, g_quat)
pen_gt = probe_radius_gt - sd
if pen_gt > max_pen_gt:
max_pen_gt = pen_gt
pen_m = probe_radius_m - sd
if pen_m > max_pen_m:
max_pen_m = pen_m
if not _func_link_is_filtered(filter_links_idx, i_s, geoms_info.link_idx[i_g]):
g_pos = geoms_state.pos[i_g, i_b]
g_quat = geoms_state.quat[i_g, i_b]
sd = sdf.sdf_func_world_local(geoms_info, sdf_info, probe_pos, i_g, g_pos, g_quat)
pen_gt = probe_radius_gt - sd
if pen_gt > max_pen_gt:
max_pen_gt = pen_gt
pen_m = probe_radius_m - sd
if pen_m > max_pen_m:
max_pen_m = pen_m

return max_pen_gt, max_pen_m

Expand Down Expand Up @@ -177,6 +185,7 @@ def _func_query_contact_depth(
probe_pos: qd.types.vector(3),
probe_radius_gt: float,
probe_radius_m: float,
filter_links_idx: qd.types.ndarray(),
geoms_info: array_class.GeomsInfo,
geoms_state: array_class.GeomsState,
rigid_global_info: array_class.RigidGlobalInfo,
Expand All @@ -192,8 +201,9 @@ def _func_query_contact_depth(
Iterates the per-(env, sensor) deduplicated opposing-geom list built by ``_kernel_build_sensor_geom_idx``;
every geom in that list contacts the sensor's tracked link, so the reported contact link is recovered as
``geoms_info.link_idx[i_g]`` (the link owning the opposing geom). AABB pre-filter expands by
``max(probe_radius_gt, probe_radius_m)`` so neither branch is
silently skipped. Callers without a noised radius pass ``probe_radius_m == probe_radius_gt``.
``max(probe_radius_gt, probe_radius_m)`` so neither branch is silently skipped. Callers without a noised radius
pass ``probe_radius_m == probe_radius_gt``. Geoms whose owning link is in sensor ``i_s``'s ``filter_links_idx``
list are skipped.
"""
max_pen_gt = gs.qd_float(0.0)
contact_link_gt = gs.qd_int(-1)
Expand All @@ -206,7 +216,9 @@ def _func_query_contact_depth(
n_g = sensor_n_geoms[i_b, i_s]
for i_g_ in range(n_g):
i_g = sensor_geoms_idx[i_b, i_s, i_g_]
if func_point_in_geom_aabb(geoms_state, i_g, i_b, probe_pos, aabb_expansion):
if not _func_link_is_filtered(filter_links_idx, i_s, geoms_info.link_idx[i_g]) and func_point_in_geom_aabb(
geoms_state, i_g, i_b, probe_pos, aabb_expansion
):
g_pos = geoms_state.pos[i_g, i_b]
g_quat = geoms_state.quat[i_g, i_b]
sd = sdf.sdf_func_world_local(geoms_info, sdf_info, probe_pos, i_g, g_pos, g_quat)
Expand Down Expand Up @@ -297,6 +309,7 @@ def _kernel_kinematic_taxel(
shear_scalar: qd.types.ndarray(),
twist_scalar: qd.types.ndarray(),
links_idx: qd.types.ndarray(),
filter_links_idx: qd.types.ndarray(),
sensor_cache_start: qd.types.ndarray(),
sensor_probe_start: qd.types.ndarray(),
n_probes_per_sensor: qd.types.ndarray(),
Expand Down Expand Up @@ -363,6 +376,7 @@ def _kernel_kinematic_taxel(
probe_pos,
probe_radius,
probe_radius_m,
filter_links_idx,
geoms_info,
geoms_state,
rigid_global_info,
Expand Down Expand Up @@ -429,6 +443,7 @@ def _kernel_contact_depth_probe(
probe_radii_noise: qd.types.ndarray(),
probe_gains: qd.types.ndarray(),
links_idx: qd.types.ndarray(),
filter_links_idx: qd.types.ndarray(),
sensor_cache_start: qd.types.ndarray(),
sensor_probe_start: qd.types.ndarray(),
sensor_geoms_idx: qd.types.ndarray(),
Expand Down Expand Up @@ -475,6 +490,7 @@ def _kernel_contact_depth_probe(
probe_pos,
probe_radius,
probe_radius_m,
filter_links_idx,
geoms_info,
geoms_state,
sensor_geoms_idx,
Expand All @@ -493,6 +509,7 @@ def _kernel_contact_depth_probe(
@qd.kernel
def _kernel_build_sensor_candidate_geom_mask(
sensor_link_idx: qd.types.ndarray(),
filter_links_idx: qd.types.ndarray(),
sensor_contacts_idx: qd.types.ndarray(),
sensor_n_contacts: qd.types.ndarray(),
collider_state: array_class.ColliderState,
Expand All @@ -505,7 +522,8 @@ def _kernel_build_sensor_candidate_geom_mask(
to skip triangles whose owning geom isn't in the sensor's current contact list. Only the geom on the side opposite
the sensor link is marked (mirroring the SDF path's ``i_g = <other geom>`` selection); marking the sensor's own
geom would let the BVH closest-point test latch onto the sensor's own surface, pinning the reported depth to
``probe_radius`` regardless of the pressing object.
``probe_radius`` regardless of the pressing object. A contact whose counterpart link is in sensor ``i_s``'s
``filter_links_idx`` list is skipped, so its opposing geom never becomes a BVH candidate.
"""
n_batches = sensor_n_contacts.shape[0]
n_sensors = sensor_n_contacts.shape[1]
Expand All @@ -517,9 +535,11 @@ def _kernel_build_sensor_candidate_geom_mask(
n_c = sensor_n_contacts[i_b, i_s]
for i_c_ in range(n_c):
i_c = sensor_contacts_idx[i_b, i_s, i_c_]
if collider_state.contact_data.link_a[i_c, i_b] == link:
c_link_a = collider_state.contact_data.link_a[i_c, i_b]
c_link_b = collider_state.contact_data.link_b[i_c, i_b]
if c_link_a == link and not _func_link_is_filtered(filter_links_idx, i_s, c_link_b):
sensor_candidate_geom_mask[i_b, i_s, collider_state.contact_data.geom_b[i_c, i_b]] = True
if collider_state.contact_data.link_b[i_c, i_b] == link:
if c_link_b == link and not _func_link_is_filtered(filter_links_idx, i_s, c_link_a):
sensor_candidate_geom_mask[i_b, i_s, collider_state.contact_data.geom_a[i_c, i_b]] = True


Expand Down Expand Up @@ -889,6 +909,12 @@ class KinematicTactileSensorMixin(ContactDepthQuerySensorMixin, ProbeSensorMixin
subclasses add their own metadata.
"""

def build(self):
super().build()
self._shared_metadata.filter_links_idx = append_filter_links_idx(
self._shared_metadata.filter_links_idx, self._options.filter_link_idx
)


@dataclass
class ContactDepthProbeMetadata(
Expand All @@ -898,7 +924,8 @@ class ContactDepthProbeMetadata(
RigidSensorMetadataMixin,
SimpleSensorMetadata,
):
pass
# (n_sensors, max_num_filter_links); unused slots and empty filters are -1. See `append_filter_links_idx`.
filter_links_idx: torch.Tensor = make_tensor_field((0, 0), dtype_factory=lambda: gs.tc_int)


class ContactDepthProbeSensor(
Expand Down Expand Up @@ -958,6 +985,7 @@ def _update_current_timestep_data(
shared_metadata.probe_radii_noise,
shared_metadata.probe_gains,
shared_metadata.links_idx,
shared_metadata.filter_links_idx,
shared_metadata.sensor_cache_start,
shared_metadata.sensor_probe_start,
shared_metadata.sensor_geoms_idx,
Expand All @@ -982,6 +1010,7 @@ def _update_current_timestep_data(
shared_metadata.sensor_candidate_geom_mask = torch.zeros(mask_shape, dtype=gs.tc_bool, device=gs.device)
_kernel_build_sensor_candidate_geom_mask(
shared_metadata.links_idx,
shared_metadata.filter_links_idx,
shared_metadata.sensor_contacts_idx,
shared_metadata.sensor_n_contacts,
solver.collider._collider_state,
Expand Down Expand Up @@ -1129,6 +1158,8 @@ class KinematicTaxelMetadata(
normal_exponent: torch.Tensor = make_tensor_field((0,))
shear_scalar: torch.Tensor = make_tensor_field((0,))
twist_scalar: torch.Tensor = make_tensor_field((0,))
# (n_sensors, max_num_filter_links); unused slots and empty filters are -1. See `append_filter_links_idx`.
filter_links_idx: torch.Tensor = make_tensor_field((0, 0), dtype_factory=lambda: gs.tc_int)


class KinematicTaxelSensor(
Expand Down Expand Up @@ -1215,6 +1246,7 @@ def _update_current_timestep_data(
shared_metadata.shear_scalar,
shared_metadata.twist_scalar,
shared_metadata.links_idx,
shared_metadata.filter_links_idx,
shared_metadata.sensor_cache_start,
shared_metadata.sensor_probe_start,
shared_metadata.n_probes_per_sensor,
Expand Down Expand Up @@ -1244,6 +1276,7 @@ def _update_current_timestep_data(
shared_metadata.sensor_candidate_geom_mask = torch.zeros(mask_shape, dtype=gs.tc_bool, device=gs.device)
_kernel_build_sensor_candidate_geom_mask(
shared_metadata.links_idx,
shared_metadata.filter_links_idx,
shared_metadata.sensor_contacts_idx,
shared_metadata.sensor_n_contacts,
solver.collider._collider_state,
Expand Down
Loading
Loading