From afc7e38f990f33bef5a1e982fa5306aeeac5b6ff Mon Sep 17 00:00:00 2001 From: matafela Date: Tue, 1 Sep 2026 15:41:25 +0800 Subject: [PATCH 1/8] update --- .../sim/atomic_actions/builtin_actions.md | 42 ++- docs/source/overview/sim/sim_articulation.md | 29 ++ .../lab/sim/atomic_actions/affordance.py | 238 +++++++++++- embodichain/lab/sim/atomic_actions/core.py | 3 +- .../sim/atomic_actions/primitives/twist.py | 2 +- embodichain/lab/sim/objects/articulation.py | 271 +++++++++++++ .../integrations/_configured_services.py | 38 +- .../task_program/integrations/configured.py | 15 +- .../open_drawer/task_program/integration.yaml | 1 - scripts/tutorials/atomic_action/press.py | 13 +- scripts/tutorials/atomic_action/slide.py | 9 +- scripts/tutorials/atomic_action/twist.py | 13 +- .../task_program/test_task_vertical_slices.py | 68 +++- tests/sim/atomic_actions/test_actions.py | 49 ++- tests/sim/atomic_actions/test_affordance.py | 303 ++++++++++++++- .../sim/atomic_actions/test_tutorial_utils.py | 129 +++++++ tests/sim/objects/test_articulation.py | 356 ++++++++++++++++++ 17 files changed, 1495 insertions(+), 84 deletions(-) diff --git a/docs/source/overview/sim/atomic_actions/builtin_actions.md b/docs/source/overview/sim/atomic_actions/builtin_actions.md index fd46a99bc..c67f6867b 100644 --- a/docs/source/overview/sim/atomic_actions/builtin_actions.md +++ b/docs/source/overview/sim/atomic_actions/builtin_actions.md @@ -677,10 +677,13 @@ migration. ## `Press` Plans **close hand -> approach target -> contact -> press along axis -> return -to the approach pose**. `PressAffordance` is entity-free and stores an explicit -target-local surface `press_position` and `press_axis`. `PressGoal.target_pose` -is either a pose snapshot or `SceneEntityPose`, which resolves through the -current `PlanningContext.scene` and participates in dynamic-goal recovery. +to the approach pose**. For an articulation link, +`Articulation.sample_initial_point_clouds()` stores target-link-local target and +whole-articulation clouds in `ObjectSemantics.geometry`. `PressAffordance` +derives `press_axis` from that geometry and can derive the outer-surface +`press_position` when it is omitted. `PressGoal.target_pose` is either a pose +snapshot or `SceneEntityPose`, which resolves through the current +`PlanningContext.scene` and participates in dynamic-goal recovery. The contact, press, and retract segments use axis-aligned Cartesian keyframes; each output sample is grounded with IK instead of being interpolated only in @@ -698,7 +701,7 @@ right-handed orthonormal rotation even for vertical or oblique press axes. `PressOptions` controls hand-close interpolation, approach distance, press distance, and an optional target-local `press_position`. An options-level -position overrides the affordance's explicit surface point. The bound +position overrides the affordance's resolved surface point. The bound `primary.grasp` endpoint must provide `grasp`; both endpoints come from the generic `ActionBinding`, and the action keeps the gripper closed for all arm motion segments. Applications that require force/contact confirmation must @@ -711,13 +714,25 @@ verify it externally. ## `Slide` Plans a grasped linear interaction for one articulation link. The entity-free -`SlideAffordance` stores the link-local grasp mesh, `translation_axis`, and +`SlideAffordance` stores the link-local grasp mesh and resolves its +`translation_axis` from initial articulation point-cloud geometry, plus optional joint name/limits. `SlideGoal.target_pose` supplies the link pose as a snapshot or `SceneEntityPose`. The positive axis direction means approach and push/close; pull/open uses its negative direction. The affordance inherits -`AntipodalAffordance` and selects a grasp with `get_best_grasp_poses()`. The grasp -approach direction is the link-frame translation axis transformed by the -current link rotation. +`AntipodalAffordance` and selects a grasp with `get_best_grasp_poses()`. The +grasp approach direction is the resolved link-frame translation axis +transformed by the current link rotation. + +Axis inference samples the target link and the merged articulation surface at +`ArticulationCfg.init_qpos`, expressed in the target link's initial local frame. +Both clouds use Open3D uniform surface sampling. Sampling the whole articulation +as one merged mesh preserves triangle-area weighting instead of giving every +link an equal point budget, which would over-represent tiny decorative links. +The target cloud center and twice its distribution radius define a spherical +neighborhood in the full cloud. The largest signed component of the +neighborhood-center offset selects exactly one of `+X`, `-X`, `+Y`, `-Y`, `+Z`, +or `-Z`. A complete point-cloud geometry snapshot overrides any legacy explicit +axis fallback. With `direction="pull"`, the sequence is **approach -> reach -> close -> pull -> open**. With `direction="push"`, it is **approach -> reach -> close -> push -> open @@ -806,8 +821,13 @@ normalized against the resolved hinge limits and passed as the goal's Plans **approach -> reach -> close -> twist -> open -> retract** for an articulation link or a rigid object. The entity-free `TwistAffordance` stores an -explicit local `grasp_position`, `twist_axis`, and `axis_origin`, plus optional -joint name/limits. `TwistGoal.target_pose` supplies the grounded target pose. +explicit local `grasp_position`, plus optional joint name/limits. For an +articulation link, the same initial target-neighborhood geometry used by +`Slide` and `Press` resolves `twist_axis` and sets `axis_origin` to the sampled +target-link cloud centroid. Complete point-cloud geometry overrides both legacy +fallback values, including when the centroid is not the link-frame origin. A +rigid object without articulation context may still provide explicit +compatibility values. `TwistGoal.target_pose` supplies the grounded target pose. The grasp frame's z-axis follows the world-transformed twist axis; an adaptive reference completes a right-handed orthonormal frame. Twist keyframes rotate diff --git a/docs/source/overview/sim/sim_articulation.md b/docs/source/overview/sim/sim_articulation.md index eec111ded..d5012546c 100644 --- a/docs/source/overview/sim/sim_articulation.md +++ b/docs/source/overview/sim/sim_articulation.md @@ -181,6 +181,7 @@ State data is accessed via getter methods that return batched tensors (`N` envir | :--- | :--- | :--- | | `get_local_pose(to_matrix=False)` | `(N, 7)` or `(N, 4, 4)` | Root link pose `[x, y, z, qw, qx, qy, qz]` or a 4x4 matrix. | | `get_link_pose(link_name, to_matrix=False)` | `(N, 7)` or `(N, 4, 4)` | Specific link pose `[x, y, z, qw, qx, qy, qz]` or a 4x4 matrix. | +| `sample_initial_point_clouds(target_link_name)` | `Dict[str, Tensor]` | Uniformly sample the target link and merged articulation surface at `init_qpos`, both in the target link's initial local frame. | | `get_qpos(target=False)` | `(N, dof)` | Current joint positions (or joint targets if `target=True`). | | `get_qvel(target=False)` | `(N, dof)` | Current joint velocities (or velocity targets if `target=True`). | | `get_joint_drive()` | `Tuple[Tensor, ...]` | Returns `(stiffness, damping, max_effort, max_velocity, friction, armature)`, each shaped `(N, dof)`. | @@ -193,6 +194,34 @@ print(f"Current Joint Positions: {articulation.get_qpos()}") print(f"End Effector Pose: {articulation.get_link_pose('ee_link')}") ``` +### Initial Link-Local Point Clouds + +`sample_initial_point_clouds()` uses the configured initial joint positions and +forward kinematics, rather than the articulation's mutable runtime state. It +transforms every link mesh into the requested target link's initial frame, +merges the meshes, and uses Open3D uniform surface sampling on the combined +triangle mesh. Sampling the merged mesh makes each surface's representation +proportional to triangle area instead of assigning the same point count to +every link. The returned float32 tensors are moved back to the articulation +device and are ready to store in an atomic action's `ObjectSemantics.geometry`: + +```python +geometry = articulation.sample_initial_point_clouds( + "button_cap", + articulation_point_count=100_000, + target_point_count=5_000, +) +target_points = geometry["target_link_point_cloud"] +articulation_points = geometry["articulation_point_cloud"] +``` + +The shown point counts are the defaults. Open3D draws random uniform samples +from the target mesh and merged articulation mesh independently, so repeated +calls are not expected to be bitwise identical and consumers should rely on +the spatial distribution rather than point-for-point correspondence. The +method requires a built kinematic chain and currently supports unit +`body_scale`. + ### Visual Appearance Asset materials are wrapped automatically during articulation construction. Materials are organized by environment and link: diff --git a/embodichain/lab/sim/atomic_actions/affordance.py b/embodichain/lab/sim/atomic_actions/affordance.py index 90282d300..4caa3e1da 100644 --- a/embodichain/lab/sim/atomic_actions/affordance.py +++ b/embodichain/lab/sim/atomic_actions/affordance.py @@ -16,11 +16,15 @@ from __future__ import annotations +from collections.abc import Mapping from dataclasses import dataclass, field from typing import Any, ClassVar import torch +_TARGET_LINK_POINT_CLOUD_KEY = "target_link_point_cloud" +_ARTICULATION_POINT_CLOUD_KEY = "articulation_point_cloud" + @dataclass class Affordance: @@ -49,6 +53,17 @@ def get_batch_size(self) -> int: """Return the batch size of this affordance data.""" return 1 + def resolve_from_object_geometry(self, geometry: Mapping[str, Any]) -> None: + """Resolve geometry-derived fields after object semantics are assembled. + + Subclasses may override this hook when their derived semantic values + require metadata owned by ``ObjectSemantics.geometry``. + + Args: + geometry: Non-affordance object geometry metadata. + """ + del geometry + @dataclass class AntipodalAffordance(Affordance): @@ -243,18 +258,21 @@ def __post_init__(self) -> None: @dataclass class TwistAffordance(Affordance): - """Target-local grasp point and rotation-axis geometry for twisting.""" + """Target-local grasp point and inferred rotation-axis geometry.""" grasp_position: tuple[float, float, float] = field(kw_only=True) """Explicit target-local center of the gripper contact region.""" - axis_origin: tuple[float, float, float] = field(kw_only=True) - """Explicit point on the rotation axis in the target-local frame.""" + axis_origin: tuple[float, float, float] | None = field( + default=None, + kw_only=True, + ) + """Fallback axis point, overridden by the target point-cloud centroid.""" twist_axis: torch.Tensor = field( default_factory=lambda: torch.tensor([0.0, 1.0, 0.0]) ) - """Twist axis expressed in the target object's local frame.""" + """Twist axis in the target frame, overridden by articulation geometry.""" joint_name: str | None = None """Optional stable articulation-joint name associated with the axis.""" @@ -275,11 +293,41 @@ def __post_init__(self) -> None: self.grasp_position = _validate_local_point( self.grasp_position, "TwistAffordance.grasp_position" ) - self.axis_origin = _validate_local_point( - self.axis_origin, "TwistAffordance.axis_origin" - ) + if self.axis_origin is not None: + self.axis_origin = _validate_local_point( + self.axis_origin, "TwistAffordance.axis_origin" + ) _validate_joint_metadata(self.joint_name, self.joint_limits) + def resolve_from_object_geometry(self, geometry: Mapping[str, Any]) -> None: + """Infer the twist axis and its target-local origin from point clouds.""" + resolved = _infer_articulation_neighborhood_axis( + geometry, + field_name="TwistAffordance.twist_axis", + ) + if resolved is not None: + self.twist_axis, target_points = resolved + self.axis_origin = tuple( + float(component) for component in target_points.mean(dim=0) + ) + + def require_axis_origin(self) -> tuple[float, float, float]: + """Return the explicit or geometry-derived local rotation-axis point. + + Returns: + Resolved target-local rotation-axis origin. + + Raises: + ValueError: If neither a fallback nor articulation geometry supplied + the rotation-axis origin. + """ + if self.axis_origin is None: + raise ValueError( + "TwistAffordance.axis_origin must be provided explicitly or " + "resolved from articulation point-cloud geometry." + ) + return self.axis_origin + def get_grasp_pose(self, target_pose: torch.Tensor) -> torch.Tensor: """Construct a deterministic world grasp pose from local geometry. @@ -320,7 +368,7 @@ def get_grasp_pose(self, target_pose: torch.Tensor) -> torch.Tensor: @dataclass class SlideAffordance(AntipodalAffordance): - """Target-local antipodal grasp and translation-axis geometry. + """Target-local antipodal grasp and inferred translation-axis geometry. The positive translation-axis direction denotes approaching and pushing the articulated part closed. Pulling moves in the opposite direction. @@ -337,7 +385,7 @@ class SlideAffordance(AntipodalAffordance): translation_axis: torch.Tensor = field( default_factory=lambda: torch.tensor([0.0, 1.0, 0.0]) ) - """Approach and push/close direction in the articulation-link frame.""" + """Approach/push axis, overridden by articulation point-cloud geometry.""" joint_name: str | None = None """Optional stable prismatic-joint name associated with the link.""" @@ -360,6 +408,15 @@ def __post_init__(self) -> None: self.translation_axis = self.translation_axis.clone() _validate_joint_metadata(self.joint_name, self.joint_limits) + def resolve_from_object_geometry(self, geometry: Mapping[str, Any]) -> None: + """Infer the signed Cartesian translation axis from local point clouds.""" + resolved = _infer_articulation_neighborhood_axis( + geometry, + field_name="SlideAffordance.translation_axis", + ) + if resolved is not None: + self.translation_axis = resolved[0] + @dataclass class OpenDoorAffordance(AntipodalAffordance): @@ -558,15 +615,18 @@ def from_articulation( @dataclass class PressAffordance(Affordance): - """Explicit target-local contact point and pressing direction.""" + """Target-local contact point and inferred pressing direction.""" press_axis: torch.Tensor = field( default_factory=lambda: torch.tensor([0.0, 0.0, 1.0]) ) - """Press direction expressed in the target object's local frame.""" + """Press direction in the target frame, overridden by point-cloud geometry.""" - press_position: tuple[float, float, float] = field(kw_only=True) - """Explicit local-frame point on the pressable contact surface.""" + press_position: tuple[float, float, float] | None = field( + default=None, + kw_only=True, + ) + """Local contact point; inferred from articulation geometry when omitted.""" def __post_init__(self) -> None: if ( @@ -578,9 +638,25 @@ def __post_init__(self) -> None: if torch.linalg.vector_norm(self.press_axis) <= 1.0e-6: raise ValueError("PressAffordance.press_axis must be non-zero.") self.press_axis = self.press_axis.clone() - self.press_position = _validate_local_point( - self.press_position, "PressAffordance.press_position" + if self.press_position is not None: + self.press_position = _validate_local_point( + self.press_position, "PressAffordance.press_position" + ) + + def resolve_from_object_geometry(self, geometry: Mapping[str, Any]) -> None: + """Infer the press axis and optional outer-surface contact point.""" + resolved = _infer_articulation_neighborhood_axis( + geometry, + field_name="PressAffordance.press_axis", ) + if resolved is None: + return + self.press_axis, target_points = resolved + if self.press_position is None: + self.press_position = _outer_surface_center( + target_points, + self.press_axis, + ) def get_press_pose( self, @@ -616,6 +692,11 @@ def get_press_pose( configured_position = ( self.press_position if configured_position is None else configured_position ) + if configured_position is None: + raise ValueError( + "PressAffordance.press_position must be provided explicitly or " + "resolved from articulation point-cloud geometry." + ) local_press_position = torch.tensor( configured_position, dtype=torch.float32, @@ -653,6 +734,133 @@ def _validate_press_position( return tuple(float(component) for component in position) +def _infer_articulation_neighborhood_axis( + geometry: Mapping[str, Any], + *, + field_name: str, +) -> tuple[torch.Tensor, torch.Tensor] | None: + """Infer one signed Cartesian axis from target-centered local geometry. + + The target-link point-cloud center defines a spherical neighborhood in the + complete articulation cloud. Its radius is twice the target cloud's maximum + distance from that center. The largest signed component of the neighborhood + center offset selects exactly one of ``+/-X``, ``+/-Y``, or ``+/-Z``. + + Args: + geometry: Object geometry containing both point-cloud entries. + field_name: Affordance field name used in validation errors. + + Returns: + ``(axis, target_points)`` when point-cloud geometry is present, otherwise + ``None`` for legacy or non-articulation affordances. + + Raises: + TypeError: If geometry or a point cloud has an invalid type. + ValueError: If point-cloud metadata is incomplete, malformed, or + geometrically degenerate. + """ + if not isinstance(geometry, Mapping): + raise TypeError("geometry must be a mapping.") + has_target = _TARGET_LINK_POINT_CLOUD_KEY in geometry + has_articulation = _ARTICULATION_POINT_CLOUD_KEY in geometry + if not has_target and not has_articulation: + return None + if not has_target or not has_articulation: + raise ValueError( + f"{field_name} inference requires both " + f"{_TARGET_LINK_POINT_CLOUD_KEY!r} and " + f"{_ARTICULATION_POINT_CLOUD_KEY!r}." + ) + + target_points = _validate_local_point_cloud( + geometry[_TARGET_LINK_POINT_CLOUD_KEY], + field_name=f"geometry[{_TARGET_LINK_POINT_CLOUD_KEY!r}]", + ) + articulation_points = _validate_local_point_cloud( + geometry[_ARTICULATION_POINT_CLOUD_KEY], + field_name=f"geometry[{_ARTICULATION_POINT_CLOUD_KEY!r}]", + ) + if target_points.device != articulation_points.device: + raise ValueError( + "Articulation and target-link point clouds must share a device." + ) + target_points = target_points.to(dtype=torch.float32) + articulation_points = articulation_points.to(dtype=torch.float32) + + target_center = target_points.mean(dim=0) + target_distances = torch.linalg.vector_norm( + target_points - target_center, + dim=1, + ) + target_radius = target_distances.max() + if float(target_radius.item()) <= 1.0e-8: + raise ValueError( + f"{field_name} cannot be inferred from a degenerate target-link " + "point cloud." + ) + + neighborhood_radius = target_radius * 2.0 + neighborhood_mask = ( + torch.linalg.vector_norm( + articulation_points - target_center, + dim=1, + ) + <= neighborhood_radius + ) + if not bool(neighborhood_mask.any().item()): + raise ValueError(f"{field_name} point-cloud neighborhood is empty.") + neighborhood_center = articulation_points[neighborhood_mask].mean(dim=0) + center_offset = neighborhood_center - target_center + axis_index = int(torch.argmax(torch.abs(center_offset)).item()) + dominant_offset = center_offset[axis_index] + offset_tolerance = max(1.0e-8, float(target_radius.item()) * 1.0e-6) + if abs(float(dominant_offset.item())) <= offset_tolerance: + raise ValueError( + f"{field_name} cannot be inferred because the articulation " + "neighborhood is centered on the target link." + ) + + axis = torch.zeros(3, dtype=torch.float32, device=target_points.device) + axis[axis_index] = 1.0 if float(dominant_offset.item()) > 0.0 else -1.0 + return axis, target_points + + +def _validate_local_point_cloud(value: Any, *, field_name: str) -> torch.Tensor: + """Validate one non-empty finite floating point cloud.""" + if not isinstance(value, torch.Tensor): + raise TypeError(f"{field_name} must be a torch.Tensor.") + if ( + not value.is_floating_point() + or value.dim() != 2 + or value.shape[1:] != (3,) + or value.shape[0] == 0 + or not bool(torch.isfinite(value).all().item()) + ): + raise ValueError( + f"{field_name} must be a non-empty finite floating tensor with " + "shape (N, 3)." + ) + return value + + +def _outer_surface_center( + target_points: torch.Tensor, + inward_axis: torch.Tensor, +) -> tuple[float, float, float]: + """Return the sampled outer-surface center opposite an inward axis.""" + axis = inward_axis.to(device=target_points.device, dtype=torch.float32) + projections = torch.matmul(target_points, axis) + minimum = projections.min() + radius = torch.linalg.vector_norm( + target_points - target_points.mean(dim=0), + dim=1, + ).max() + tolerance = max(1.0e-5, float(radius.item()) * 1.0e-4) + surface_points = target_points[projections <= minimum + tolerance] + center = surface_points.mean(dim=0) + return tuple(float(component) for component in center) + + def _orthogonal_xy_from_z(z_axis: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Complete normalized z axes into right-handed orthonormal frames.""" basis = torch.eye(3, dtype=z_axis.dtype, device=z_axis.device) diff --git a/embodichain/lab/sim/atomic_actions/core.py b/embodichain/lab/sim/atomic_actions/core.py index 404781920..69baaa9c2 100644 --- a/embodichain/lab/sim/atomic_actions/core.py +++ b/embodichain/lab/sim/atomic_actions/core.py @@ -99,7 +99,7 @@ class ObjectSemantics: """Affordance data describing supported interactions.""" geometry: dict[str, Any] - """Non-affordance geometric metadata.""" + """Non-affordance metadata used to resolve geometry-derived affordance data.""" entity_id: str """Stable scene identifier used by snapshot grounding and object identity.""" @@ -121,6 +121,7 @@ def __post_init__(self) -> None: raise ValueError("label must be a non-empty string.") if not isinstance(self.entity_id, str) or not self.entity_id.strip(): raise ValueError("entity_id must be a non-empty string.") + self.affordance.resolve_from_object_geometry(self.geometry) self.affordance.object_label = self.label diff --git a/embodichain/lab/sim/atomic_actions/primitives/twist.py b/embodichain/lab/sim/atomic_actions/primitives/twist.py index 4a0a10b24..73a0d5532 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/twist.py +++ b/embodichain/lab/sim/atomic_actions/primitives/twist.py @@ -211,7 +211,7 @@ def _plan( link_pose, grasp_xpos, affordance.twist_axis, - affordance.axis_origin, + affordance.require_axis_origin(), options.twist_angle, options.twist_waypoint_count, ) diff --git a/embodichain/lab/sim/objects/articulation.py b/embodichain/lab/sim/objects/articulation.py index 6499dfa15..ac63d7463 100644 --- a/embodichain/lab/sim/objects/articulation.py +++ b/embodichain/lab/sim/objects/articulation.py @@ -21,6 +21,7 @@ import torch import dexsim import numpy as np +import open3d as o3d from dataclasses import dataclass from functools import cached_property @@ -1197,6 +1198,276 @@ def get_link_vert_face(self, link_name: str) -> Tuple[torch.Tensor, torch.Tensor verts, faces = self.body_data.link_vert_face[link_name] return verts, faces + def sample_initial_point_clouds( + self, + target_link_name: str, + *, + articulation_point_count: int = 100_000, + target_point_count: int = 5_000, + ) -> dict[str, torch.Tensor]: + """Sample initial articulation geometry in a target-link frame. + + All link meshes are transformed with forward kinematics evaluated at + :attr:`ArticulationCfg.init_qpos`, merged into one triangle mesh, and + sampled proportionally to triangle area. The target link is sampled + separately so consumers can derive a target-centered neighborhood. + Both returned clouds use the target link's initial local frame. + + Args: + target_link_name: Link whose initial local frame defines the output. + articulation_point_count: Number of points sampled from the merged + articulation surface. + target_point_count: Number of points sampled from the target-link + surface. + + Returns: + Geometry metadata containing ``articulation_point_cloud`` with + shape ``(articulation_point_count, 3)`` and + ``target_link_point_cloud`` with shape + ``(target_point_count, 3)``. + + Raises: + TypeError: If a name or point count has the wrong type. + ValueError: If the target link, point counts, initial joint state, + FK output, or mesh geometry is invalid. + RuntimeError: If the articulation has no kinematic chain. + """ + if type(target_link_name) is not str: + raise TypeError("target_link_name must be a string.") + if not target_link_name or target_link_name != target_link_name.strip(): + raise ValueError("target_link_name must be non-empty.") + if target_link_name not in self.link_names: + raise ValueError( + f"Unknown articulation link {target_link_name!r}. Available " + f"links: {list(self.link_names)}." + ) + for value, field_name in ( + (articulation_point_count, "articulation_point_count"), + (target_point_count, "target_point_count"), + ): + if type(value) is not int: + raise TypeError(f"{field_name} must be an integer.") + if value <= 0: + raise ValueError(f"{field_name} must be positive.") + if self.pk_chain is None: + raise RuntimeError( + "Initial point-cloud sampling requires cfg.build_pk_chain=True." + ) + body_scale = torch.as_tensor(self.cfg.body_scale, dtype=torch.float32) + if body_scale.shape != (3,) or not torch.allclose( + body_scale, + torch.ones(3, dtype=torch.float32), + ): + raise ValueError( + "Initial point-cloud sampling currently requires unit body_scale." + ) + + initial_qpos = torch.as_tensor( + self.cfg.init_qpos, + dtype=torch.float32, + device=self.device, + ) + if initial_qpos.shape != (self.dof,) or not torch.isfinite(initial_qpos).all(): + raise ValueError( + "ArticulationCfg.init_qpos must be a finite vector matching " + f"the articulation DoF ({self.dof})." + ) + backend_joint_names = list(self.joint_names) + kinematic_joint_names = list(self.pk_chain.get_joint_parameter_names()) + if backend_joint_names != kinematic_joint_names: + if len(backend_joint_names) != len(kinematic_joint_names) or set( + backend_joint_names + ) != set(kinematic_joint_names): + raise ValueError( + "Initial point-cloud sampling requires matching simulator " + "and kinematic-chain joint names." + ) + qpos_by_name = dict(zip(backend_joint_names, initial_qpos)) + initial_qpos = torch.stack( + [qpos_by_name[name] for name in kinematic_joint_names] + ) + link_names = list(self.link_names) + initial_link_poses = self.compute_fk( + initial_qpos.unsqueeze(0), + link_names=link_names, + ) + if ( + initial_link_poses.shape != (1, len(link_names), 4, 4) + or not torch.isfinite(initial_link_poses).all() + ): + raise ValueError( + "Initial FK must return finite poses with shape " + f"(1, {len(link_names)}, 4, 4)." + ) + initial_link_poses = initial_link_poses[0].to( + device=self.device, + dtype=torch.float32, + ) + target_index = link_names.index(target_link_name) + target_from_root = torch.linalg.inv(initial_link_poses[target_index]) + + merged_vertices: list[torch.Tensor] = [] + merged_triangles: list[torch.Tensor] = [] + vertex_offset = 0 + target_vertices: torch.Tensor | None = None + target_triangles: torch.Tensor | None = None + for link_index, link_name in enumerate(link_names): + vertices, triangles = self.get_link_vert_face(link_name) + vertices = torch.as_tensor( + vertices, + dtype=torch.float32, + device=self.device, + ) + triangles = torch.as_tensor( + triangles, + dtype=torch.long, + device=self.device, + ) + self._validate_point_cloud_mesh(vertices, triangles, link_name=link_name) + if vertices.shape[0] == 0: + continue + + target_from_link = torch.matmul( + target_from_root, + initial_link_poses[link_index], + ) + transformed_vertices = ( + torch.matmul( + vertices, + target_from_link[:3, :3].transpose(0, 1), + ) + + target_from_link[:3, 3] + ) + merged_vertices.append(transformed_vertices) + if triangles.shape[0] > 0: + merged_triangles.append(triangles + vertex_offset) + if link_name == target_link_name: + target_vertices = transformed_vertices + target_triangles = triangles + vertex_offset += vertices.shape[0] + + if target_vertices is None or target_vertices.shape[0] == 0: + raise ValueError( + f"Target link {target_link_name!r} has no point-cloud geometry." + ) + if not merged_vertices: + raise ValueError("Articulation has no point-cloud geometry.") + articulation_vertices = torch.cat(merged_vertices, dim=0) + articulation_triangles = ( + torch.cat(merged_triangles, dim=0) + if merged_triangles + else torch.empty((0, 3), dtype=torch.long, device=self.device) + ) + assert target_triangles is not None + return { + "target_link_point_cloud": self._sample_mesh_surface_points( + target_vertices, + target_triangles, + target_point_count, + ), + "articulation_point_cloud": self._sample_mesh_surface_points( + articulation_vertices, + articulation_triangles, + articulation_point_count, + ), + } + + @staticmethod + def _validate_point_cloud_mesh( + vertices: torch.Tensor, + triangles: torch.Tensor, + *, + link_name: str, + ) -> None: + """Validate one link-local mesh used for surface sampling.""" + if vertices.dim() != 2 or vertices.shape[1:] != (3,): + raise ValueError(f"Link {link_name!r} vertices must have shape (N, 3).") + if not torch.isfinite(vertices).all(): + raise ValueError(f"Link {link_name!r} vertices must be finite.") + if triangles.dim() != 2 or triangles.shape[1:] != (3,): + raise ValueError(f"Link {link_name!r} triangles must have shape (M, 3).") + if triangles.shape[0] == 0: + return + if vertices.shape[0] == 0: + raise ValueError( + f"Link {link_name!r} triangles cannot reference an empty mesh." + ) + if bool((triangles < 0).any().item()) or int(triangles.max().item()) >= len( + vertices + ): + raise ValueError( + f"Link {link_name!r} triangles reference invalid vertices." + ) + + @staticmethod + def _sample_mesh_surface_points( + vertices: torch.Tensor, + triangles: torch.Tensor, + point_count: int, + ) -> torch.Tensor: + """Uniformly sample a triangle mesh with Open3D's CPU sampler.""" + if vertices.shape[0] == 0: + raise ValueError("Cannot sample an empty mesh.") + if triangles.shape[0] == 0: + indices = ( + torch.linspace( + 0, + vertices.shape[0] - 1, + point_count, + device=vertices.device, + ) + .round() + .to(torch.long) + ) + return vertices[indices] + + face_vertices = vertices[triangles] + face_areas = 0.5 * torch.linalg.vector_norm( + torch.cross( + face_vertices[:, 1] - face_vertices[:, 0], + face_vertices[:, 2] - face_vertices[:, 0], + dim=1, + ), + dim=1, + ) + valid_faces = face_areas > torch.finfo(vertices.dtype).eps + if not bool(valid_faces.any().item()): + indices = ( + torch.linspace( + 0, + vertices.shape[0] - 1, + point_count, + device=vertices.device, + ) + .round() + .to(torch.long) + ) + return vertices[indices] + + mesh = o3d.geometry.TriangleMesh( + vertices=o3d.utility.Vector3dVector( + vertices.detach().to(device="cpu", dtype=torch.float64).numpy() + ), + triangles=o3d.utility.Vector3iVector( + triangles[valid_faces] + .detach() + .to(device="cpu", dtype=torch.int32) + .numpy() + ), + ) + point_cloud = mesh.sample_points_uniformly(number_of_points=point_count) + sampled_points = np.asarray(point_cloud.points).copy() + if sampled_points.shape != (point_count, 3): + raise RuntimeError( + "Open3D surface sampling returned an unexpected point-cloud shape " + f"{sampled_points.shape}; expected ({point_count}, 3)." + ) + return torch.tensor( + sampled_points, + device=vertices.device, + dtype=vertices.dtype, + ) + def get_link_pose( self, link_name: str, env_ids: Sequence[int] | None = None, to_matrix=False ) -> torch.Tensor: diff --git a/embodichain/lab/task_program/integrations/_configured_services.py b/embodichain/lab/task_program/integrations/_configured_services.py index c210780fe..2dd73e7dc 100644 --- a/embodichain/lab/task_program/integrations/_configured_services.py +++ b/embodichain/lab/task_program/integrations/_configured_services.py @@ -102,7 +102,7 @@ def _identifier(value: object, *, field_name: str) -> str: def _axis(value: tuple[float, float, float]) -> tuple[float, float, float]: - """Validate one finite non-zero three-dimensional axis.""" + """Validate one legacy finite non-zero three-dimensional axis fallback.""" if type(value) is not tuple or len(value) != 3: raise TypeError("translation_axis must be an exact three-value tuple.") normalized = tuple(float(item) for item in value) @@ -289,16 +289,18 @@ def lower( @dataclass(frozen=True, slots=True) class _ArticulationLinkSlideLowererFactory(RegisteredSemanticLowererFactory): - """Create Slide semantics from one configured articulation-link mesh.""" + """Create Slide semantics from one configured articulation-link geometry.""" call_id: ClassVar[str] = _ARTICULATION_LINK_SLIDE_CALL_ID - revision: ClassVar[str] = "2" + revision: ClassVar[str] = "3" target_descriptor: ClassVar[SkillDescriptor] = Slide.descriptor() articulation_id: str articulation_simulation_uid: str link_entity_id: str - translation_axis: tuple[float, float, float] + translation_axis: tuple[float, float, float] | None = None + """Optional compatibility fallback superseded by complete point-cloud geometry.""" + target_pose_mode: str = "live" def __post_init__(self) -> None: @@ -308,7 +310,8 @@ def __post_init__(self) -> None: "link_entity_id", ): _identifier(getattr(self, field_name), field_name=field_name) - object.__setattr__(self, "translation_axis", _axis(self.translation_axis)) + if self.translation_axis is not None: + object.__setattr__(self, "translation_axis", _axis(self.translation_axis)) _slide_target_pose_mode(self.target_pose_mode) def create( @@ -356,10 +359,27 @@ def create( if not callable(get_link_vert_face): raise TypeError("Articulation must provide get_link_vert_face().") vertices, triangles = get_link_vert_face(native_link_name) + sample_initial_point_clouds = getattr( + articulation, + "sample_initial_point_clouds", + None, + ) + if not callable(sample_initial_point_clouds): + raise TypeError("Articulation must provide sample_initial_point_clouds().") + geometry = sample_initial_point_clouds(native_link_name) + if not isinstance(geometry, dict): + raise TypeError("Articulation point-cloud geometry must be a dict.") + affordance_kwargs: dict[str, object] = {} + if self.translation_axis is not None: + affordance_kwargs["translation_axis"] = torch.tensor( + self.translation_axis, + dtype=torch.float32, + device=engine.device, + ) semantics = ObjectSemantics( label="articulation_link", entity_id=self.link_entity_id, - geometry={}, + geometry=geometry, affordance=SlideAffordance( mesh_vertices=torch.as_tensor( vertices, @@ -371,11 +391,7 @@ def create( dtype=torch.long, device=engine.device, ), - translation_axis=torch.tensor( - self.translation_axis, - dtype=torch.float32, - device=engine.device, - ), + **affordance_kwargs, ), ) return _ArticulationLinkSlideLowerer( diff --git a/embodichain/lab/task_program/integrations/configured.py b/embodichain/lab/task_program/integrations/configured.py index 06753bbeb..2062c2a35 100644 --- a/embodichain/lab/task_program/integrations/configured.py +++ b/embodichain/lab/task_program/integrations/configured.py @@ -1750,10 +1750,9 @@ def _decode_registered_lowerer( "articulation_id", "articulation_simulation_uid", "link_entity_id", - "translation_axis", } ), - optional=frozenset({"target_pose_mode"}), + optional=frozenset({"target_pose_mode", "translation_axis"}), ) target_pose_mode = _identifier( config.get("target_pose_mode", "live"), @@ -1776,10 +1775,14 @@ def _decode_registered_lowerer( config["link_entity_id"], path=f"{path}.link_entity_id", ), - translation_axis=_finite_tuple( - config["translation_axis"], - path=f"{path}.translation_axis", - expected_length=3, + translation_axis=( + _finite_tuple( + config["translation_axis"], + path=f"{path}.translation_axis", + expected_length=3, + ) + if "translation_axis" in config + else None ), target_pose_mode=target_pose_mode, ) diff --git a/embodichain_tasks/configs/tasks/manipulation/open_drawer/task_program/integration.yaml b/embodichain_tasks/configs/tasks/manipulation/open_drawer/task_program/integration.yaml index 2f6ff9932..f5f4495e2 100644 --- a/embodichain_tasks/configs/tasks/manipulation/open_drawer/task_program/integration.yaml +++ b/embodichain_tasks/configs/tasks/manipulation/open_drawer/task_program/integration.yaml @@ -42,5 +42,4 @@ runtime_services: articulation_id: drawer articulation_simulation_uid: drawer link_entity_id: drawer_handle - translation_axis: [0.0, 1.0, 0.0] target_pose_mode: snapshot diff --git a/scripts/tutorials/atomic_action/press.py b/scripts/tutorials/atomic_action/press.py index 8e0400776..e7e7c5afd 100644 --- a/scripts/tutorials/atomic_action/press.py +++ b/scripts/tutorials/atomic_action/press.py @@ -130,19 +130,14 @@ def create_button_semantics( ) -> tuple[ObjectSemantics, torch.Tensor]: """Create press semantics for an articulation-link or rigid button.""" if isinstance(target, Articulation): - vertices, _ = target.get_link_vert_face(BUTTON_LINK_NAME) target_pose = target.get_link_pose(BUTTON_LINK_NAME, to_matrix=True) - press_axis = torch.tensor([0.0, 0.0, -1.0], device=target.device) - affordance = PressAffordance( - # button_cap's local -z direction matches the prismatic joint's - # inward press direction in this asset. - press_axis=press_axis, - press_position=_surface_center(vertices, press_axis), - ) + geometry = target.sample_initial_point_clouds(BUTTON_LINK_NAME) + affordance = PressAffordance() label = "microwave_start_button" else: vertices = target.get_vertices(env_ids=[0], scale=True)[0] target_pose = target.get_local_pose(to_matrix=True) + geometry = {} press_axis = torch.tensor([-1.0, 0.0, 0.0], device=target.device) affordance = PressAffordance( press_axis=press_axis, @@ -152,7 +147,7 @@ def create_button_semantics( return ( ObjectSemantics( label=label, - geometry={}, + geometry=geometry, entity_id=BUTTON_SCENE_ENTITY_ID, affordance=affordance, ), diff --git a/scripts/tutorials/atomic_action/slide.py b/scripts/tutorials/atomic_action/slide.py index 9780d77a9..ca7b5053b 100644 --- a/scripts/tutorials/atomic_action/slide.py +++ b/scripts/tutorials/atomic_action/slide.py @@ -65,7 +65,6 @@ HANDLE_LINK_NAME = "large_handle_bar" DRAWER_POSITION = (-1.1, 0.0, 0.0) DRAWER_ORIENTATION = (0.0, 0.0, 90.0) # degrees -TRANSLATION_AXIS = (0.0, 1.0, 0.0) # handle-link frame, approach/push direction TRAJECTORY_SAMPLE_COUNT = 140 HAND_INTERP_STEPS = 12 POST_TRAJECTORY_STEPS = 240 @@ -115,18 +114,14 @@ def create_drawer_semantics(drawer: Articulation) -> ObjectSemantics: Pure target-local semantics for the handle's pull/push affordance. """ vertices, triangles = drawer.get_link_vert_face(HANDLE_LINK_NAME) + geometry = drawer.sample_initial_point_clouds(HANDLE_LINK_NAME) return ObjectSemantics( label="drawer_large_handle", - geometry={}, + geometry=geometry, entity_id=HANDLE_SCENE_ENTITY_ID, affordance=SlideAffordance( mesh_vertices=torch.as_tensor(vertices), mesh_triangles=torch.as_tensor(triangles), - translation_axis=torch.tensor( - TRANSLATION_AXIS, - dtype=torch.float32, - device=drawer.device, - ), ), ) diff --git a/scripts/tutorials/atomic_action/twist.py b/scripts/tutorials/atomic_action/twist.py index 88cccab65..22ec41289 100644 --- a/scripts/tutorials/atomic_action/twist.py +++ b/scripts/tutorials/atomic_action/twist.py @@ -67,7 +67,6 @@ RIGID_KNOB_POSITION = (-0.7, -0.00, 0.70) RIGID_KNOB_SIZE = (0.05, 0.05, 0.05) KNOB_SCENE_ENTITY_ID = "twist-target" -KNOB_AXIS_ORIGIN = (0.0, 0.0, 0.0) def parse_arguments() -> argparse.Namespace: @@ -124,26 +123,26 @@ def create_knob_semantics( if isinstance(target, Articulation): vertices, _ = target.get_link_vert_face(KNOB_LINK_NAME) target_pose = target.get_link_pose(KNOB_LINK_NAME, to_matrix=True) + geometry = target.sample_initial_point_clouds(KNOB_LINK_NAME) affordance = TwistAffordance( grasp_position=_mesh_center(vertices), - # The cap_1 revolute axis passes through its link-frame origin. - axis_origin=KNOB_AXIS_ORIGIN, - twist_axis=torch.tensor([0.0, 0.0, -1.0], device=target.device), ) label = "microwave_power_knob" else: vertices = target.get_vertices(env_ids=[0], scale=True)[0] target_pose = target.get_local_pose(to_matrix=True) + geometry = {} + mesh_center = _mesh_center(vertices) affordance = TwistAffordance( - grasp_position=_mesh_center(vertices), - axis_origin=KNOB_AXIS_ORIGIN, + grasp_position=mesh_center, + axis_origin=mesh_center, twist_axis=torch.tensor([-1.0, 0.0, 0.0], device=target.device), ) label = "rigid_knob" return ( ObjectSemantics( label=label, - geometry={}, + geometry=geometry, entity_id=KNOB_SCENE_ENTITY_ID, affordance=affordance, ), diff --git a/tests/gym/envs/task_program/test_task_vertical_slices.py b/tests/gym/envs/task_program/test_task_vertical_slices.py index 3702748c1..cde3947f5 100644 --- a/tests/gym/envs/task_program/test_task_vertical_slices.py +++ b/tests/gym/envs/task_program/test_task_vertical_slices.py @@ -101,6 +101,32 @@ _EXPECTED_GRASP_SAMPLES = 1_000 +def _drawer_handle_point_cloud_geometry() -> dict[str, torch.Tensor]: + """Return target-local clouds whose articulation neighborhood lies along +Y.""" + target_points = torch.tensor( + [ + [-0.10, 0.0, 0.0], + [0.10, 0.0, 0.0], + [0.0, 0.0, -0.05], + [0.0, 0.0, 0.05], + ], + dtype=torch.float32, + ) + body_points = torch.tensor( + [ + [-0.05, 0.15, -0.05], + [0.05, 0.15, -0.05], + [-0.05, 0.15, 0.05], + [0.05, 0.15, 0.05], + ], + dtype=torch.float32, + ) + return { + "target_link_point_cloud": target_points, + "articulation_point_cloud": torch.cat((target_points, body_points), dim=0), + } + + class _NeverObserveProvider: """Reject dynamic observations during configuration decoding/compilation.""" @@ -533,6 +559,11 @@ def get_link_vert_face(name: str) -> tuple[torch.Tensor, torch.Tensor]: torch.tensor([[0, 1, 2]]), ) + @staticmethod + def sample_initial_point_clouds(name: str) -> dict[str, torch.Tensor]: + assert name == _OPEN_DRAWER_HANDLE_LINK_NAME + return _drawer_handle_point_cloud_geometry() + drawer_ref = SceneArticulationRef(_OPEN_DRAWER_ENTITY_ID) registry = SceneRegistry( ( @@ -576,6 +607,35 @@ def get_link_vert_face(name: str) -> tuple[torch.Tensor, torch.Tensor]: assert type(first[0]) is type(second[0]) assert type(first[0]).call_id == _OPEN_DRAWER_CALL_ID assert first[0] is not second[0] + first_affordance = first[0]._semantics.affordance + assert isinstance(first_affordance, SlideAffordance) + assert torch.equal( + first_affordance.translation_axis, + torch.tensor([0.0, 1.0, 0.0]), + ) + + +def test_open_drawer_lowerer_keeps_optional_legacy_axis_compatibility() -> None: + """Configured Slide accepts an old axis while new deployments omit it.""" + from embodichain.lab.task_program.integrations.configured import ( + _decode_registered_lowerer, + ) + + config = { + "kind": "articulation_link_slide", + "articulation_id": _OPEN_DRAWER_ENTITY_ID, + "articulation_simulation_uid": _OPEN_DRAWER_ENTITY_ID, + "link_entity_id": _OPEN_DRAWER_HANDLE_ID, + } + automatic = _decode_registered_lowerer(config, path="runtime_services.lowerer") + legacy = _decode_registered_lowerer( + {**config, "translation_axis": [0.0, 1.0, 0.0]}, + path="runtime_services.lowerer", + ) + + assert automatic.revision == "3" + assert automatic.translation_axis is None + assert legacy.translation_axis == (0.0, 1.0, 0.0) def test_open_drawer_lowerer_accepts_only_canonical_payload() -> None: @@ -594,13 +654,13 @@ def test_open_drawer_lowerer_accepts_only_canonical_payload() -> None: ObjectSemantics( label="drawer_handle", entity_id=_OPEN_DRAWER_HANDLE_ID, - geometry={}, + geometry=_drawer_handle_point_cloud_geometry(), affordance=SlideAffordance( mesh_vertices=torch.tensor( [[-0.1, 0.0, 0.0], [0.1, 0.0, 0.0], [0.0, 0.0, 0.0]] ), mesh_triangles=torch.tensor([[0, 1, 2]]), - translation_axis=torch.tensor([0.0, 1.0, 0.0]), + translation_axis=torch.tensor([-1.0, 0.0, 0.0]), ), ), _OPEN_DRAWER_HANDLE_ID, @@ -640,13 +700,13 @@ def test_open_drawer_lowerer_owns_a_snapshot_of_the_current_target_pose() -> Non ObjectSemantics( label="drawer_handle", entity_id=_OPEN_DRAWER_HANDLE_ID, - geometry={}, + geometry=_drawer_handle_point_cloud_geometry(), affordance=SlideAffordance( mesh_vertices=torch.tensor( [[-0.1, 0.0, 0.0], [0.1, 0.0, 0.0], [0.0, 0.0, 0.0]] ), mesh_triangles=torch.tensor([[0, 1, 2]]), - translation_axis=torch.tensor([0.0, 1.0, 0.0]), + translation_axis=torch.tensor([-1.0, 0.0, 0.0]), ), ), _OPEN_DRAWER_HANDLE_ID, diff --git a/tests/sim/atomic_actions/test_actions.py b/tests/sim/atomic_actions/test_actions.py index ad0619b02..344b903c0 100644 --- a/tests/sim/atomic_actions/test_actions.py +++ b/tests/sim/atomic_actions/test_actions.py @@ -124,12 +124,38 @@ DUAL_ARM_DOF = 2 * ARM_DOF DUAL_ROBOT_DOF = DUAL_ARM_DOF + 2 * HAND_DOF DOOR_ENTITY_ID = "door" +AUTOMATIC_TWIST_AXIS_ORIGIN = torch.tensor([1.0, 0.0, 0.0]) +AUTOMATIC_TWIST_TARGET_OFFSETS = torch.tensor( + [ + [1.0, 0.0, 0.0], + [-1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, -1.0, 0.0], + [0.0, 0.0, 1.0], + [0.0, 0.0, -1.0], + ] +) ActionT = TypeVar("ActionT", bound=AtomicAction) _ACTION_ENGINES: dict[int, AtomicActionEngine] = {} _GRASP_GENERATORS: dict[int, _StubGraspPoseGenerator] = {} +def _automatic_twist_geometry() -> dict[str, torch.Tensor]: + """Build target-local geometry with a non-origin center and +Z neighbor.""" + target_points = AUTOMATIC_TWIST_AXIS_ORIGIN + AUTOMATIC_TWIST_TARGET_OFFSETS + articulation_neighbor = AUTOMATIC_TWIST_AXIS_ORIGIN + torch.tensor( + [0.25, -0.125, 1.5] + ) + return { + "target_link_point_cloud": target_points, + "articulation_point_cloud": torch.cat( + (target_points, articulation_neighbor.unsqueeze(0)), + dim=0, + ), + } + + class _StubGraspPoseGenerator(ParallelJawGraspPoseGenerator): """Deterministic planning-service double used by atomic-action tests.""" @@ -2388,13 +2414,11 @@ def test_press_closes_hand_without_changing_projected_attachment() -> None: def test_twist_plans_six_segments_from_articulation_link() -> None: affordance = TwistAffordance( - grasp_position=(0.0, 0.0, 0.0), - axis_origin=(0.0, 0.0, 0.0), - twist_axis=torch.tensor([0.0, 1.0, 0.0]), + grasp_position=(2.0, 0.0, 0.0), ) semantics = ObjectSemantics( affordance=affordance, - geometry={}, + geometry=_automatic_twist_geometry(), label="knob", entity_id="knob", ) @@ -2430,6 +2454,10 @@ def test_twist_plans_six_segments_from_articulation_link() -> None: assert torch.all( trajectory.positions[:, plan.segment("open").stop - 1, ARM_DOF:] == 0.0 ) + assert torch.equal(affordance.twist_axis, torch.tensor([0.0, 0.0, 1.0])) + assert affordance.require_axis_origin() == pytest.approx( + tuple(float(value) for value in AUTOMATIC_TWIST_AXIS_ORIGIN) + ) first_target = generator.robot.compute_ik.call_args_list[0].kwargs["pose"] grasp_pose = affordance.get_grasp_pose(torch.eye(4).repeat(NUM_ENVS, 1, 1)) expected_pre_grasp_position = ( @@ -2466,8 +2494,15 @@ def test_twist_plans_from_explicit_rigid_object_pose_snapshot() -> None: assert plan.plan_success.tolist() == [True, True] -def test_twist_rotates_grasp_about_explicit_axis_origin() -> None: +def test_twist_rotates_grasp_about_geometry_derived_axis_origin() -> None: action = _bind_action(_motion_generator(), Twist()) + affordance = TwistAffordance(grasp_position=(2.0, 0.0, 0.0)) + ObjectSemantics( + affordance=affordance, + geometry=_automatic_twist_geometry(), + label="knob", + entity_id="knob", + ) target_pose = torch.eye(4).repeat(NUM_ENVS, 1, 1) grasp_pose = target_pose.clone() grasp_pose[:, 0, 3] = 2.0 @@ -2475,8 +2510,8 @@ def test_twist_rotates_grasp_about_explicit_axis_origin() -> None: twisted = action._twisted_grasp_poses( target_pose, grasp_pose, - torch.tensor([0.0, 0.0, 1.0]), - (1.0, 0.0, 0.0), + affordance.twist_axis, + affordance.require_axis_origin(), math.pi / 2, 4, ) diff --git a/tests/sim/atomic_actions/test_affordance.py b/tests/sim/atomic_actions/test_affordance.py index 1fdd5fc87..4f51eed98 100644 --- a/tests/sim/atomic_actions/test_affordance.py +++ b/tests/sim/atomic_actions/test_affordance.py @@ -36,6 +36,81 @@ SlideAffordance, TwistAffordance, ) +from embodichain.lab.sim.atomic_actions.core import ObjectSemantics + +POINT_CLOUD_CENTER = torch.tensor([2.0, -3.0, 4.0]) +TARGET_POINT_OFFSETS = torch.tensor( + [ + [1.0, 0.0, 0.0], + [-1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, -1.0, 0.0], + [0.0, 0.0, 1.0], + [0.0, 0.0, -1.0], + ] +) +TARGET_LINK_POINT_CLOUD_KEY = "target_link_point_cloud" +ARTICULATION_POINT_CLOUD_KEY = "articulation_point_cloud" + + +def _axis_geometry(neighbor_offset: tuple[float, float, float]) -> dict[str, object]: + """Build target-local clouds with one neighbor and one distant outlier.""" + target_points = POINT_CLOUD_CENTER + TARGET_POINT_OFFSETS + neighbor = POINT_CLOUD_CENTER + torch.tensor(neighbor_offset) + distant_outlier = POINT_CLOUD_CENTER + torch.tensor([8.0, 8.0, 8.0]) + return { + TARGET_LINK_POINT_CLOUD_KEY: target_points.clone(), + ARTICULATION_POINT_CLOUD_KEY: torch.cat( + ( + target_points, + neighbor.unsqueeze(0), + distant_outlier.unsqueeze(0), + ), + dim=0, + ), + } + + +def _axis_affordance( + kind: str, + *, + fallback_axis: torch.Tensor, + press_position: tuple[float, float, float] | None = (0.0, 0.0, 0.0), +) -> tuple[Affordance, str]: + """Construct one axis-bearing affordance with a legacy fallback axis.""" + if kind == "slide": + return ( + SlideAffordance( + mesh_vertices=torch.tensor( + [ + [-0.1, -0.1, 0.0], + [0.1, -0.1, 0.0], + [0.0, 0.1, 0.0], + ] + ), + mesh_triangles=torch.tensor([[0, 1, 2]]), + translation_axis=fallback_axis, + ), + "translation_axis", + ) + if kind == "press": + return ( + PressAffordance( + press_axis=fallback_axis, + press_position=press_position, + ), + "press_axis", + ) + if kind == "twist": + return ( + TwistAffordance( + grasp_position=(0.0, 0.0, 0.0), + axis_origin=(0.0, 0.0, 0.0), + twist_axis=fallback_axis, + ), + "twist_axis", + ) + raise ValueError(f"Unsupported affordance kind: {kind!r}.") class TestAffordance: @@ -155,11 +230,229 @@ def test_rejects_invalid_internal_axis(self, internal_axis): AxisAlignAffordance(internal_axis=internal_axis) +class TestArticulationGeometryAxisInference: + @pytest.mark.parametrize( + ("kind", "axis_field"), + ( + ("slide", "translation_axis"), + ("press", "press_axis"), + ("twist", "twist_axis"), + ), + ) + @pytest.mark.parametrize( + ("neighbor_offset", "expected_axis"), + ( + ((1.5, 0.25, -0.125), (1.0, 0.0, 0.0)), + ((-1.5, 0.25, -0.125), (-1.0, 0.0, 0.0)), + ((0.25, 1.5, -0.125), (0.0, 1.0, 0.0)), + ((0.25, -1.5, -0.125), (0.0, -1.0, 0.0)), + ((0.25, -0.125, 1.5), (0.0, 0.0, 1.0)), + ((0.25, -0.125, -1.5), (0.0, 0.0, -1.0)), + ), + ids=( + "positive_x", + "negative_x", + "positive_y", + "negative_y", + "positive_z", + "negative_z", + ), + ) + def test_object_geometry_ignores_outlier_and_quantizes_offset_to_signed_axis( + self, + kind: str, + axis_field: str, + neighbor_offset: tuple[float, float, float], + expected_axis: tuple[float, float, float], + ) -> None: + affordance, actual_axis_field = _axis_affordance( + kind, + fallback_axis=torch.tensor([1.0, 1.0, 1.0]), + ) + + ObjectSemantics( + affordance=affordance, + geometry=_axis_geometry(neighbor_offset), + entity_id=f"{kind}-target", + ) + + assert actual_axis_field == axis_field + actual_axis = getattr(affordance, axis_field) + assert torch.equal(actual_axis, torch.tensor(expected_axis)) + assert torch.count_nonzero(actual_axis).item() == 1 + assert actual_axis.abs().sum().item() == pytest.approx(1.0) + + @pytest.mark.parametrize("kind", ("slide", "press", "twist")) + def test_complete_geometry_overrides_explicit_fallback_axis( + self, kind: str + ) -> None: + affordance, axis_field = _axis_affordance( + kind, + fallback_axis=torch.tensor([1.0, 0.0, 0.0]), + ) + + ObjectSemantics( + affordance=affordance, + geometry=_axis_geometry((0.25, -0.125, -1.5)), + entity_id=f"{kind}-target", + ) + + assert torch.equal( + getattr(affordance, axis_field), + torch.tensor([0.0, 0.0, -1.0]), + ) + + @pytest.mark.parametrize("kind", ("slide", "press", "twist")) + def test_empty_geometry_preserves_explicit_fallback_axis(self, kind: str) -> None: + fallback_axis = torch.tensor([-1.0, 0.0, 0.0]) + affordance, axis_field = _axis_affordance( + kind, + fallback_axis=fallback_axis, + ) + + ObjectSemantics( + affordance=affordance, + geometry={}, + entity_id=f"rigid-{kind}-target", + ) + + assert torch.equal(getattr(affordance, axis_field), fallback_axis) + + @pytest.mark.parametrize( + "geometry", + ( + {TARGET_LINK_POINT_CLOUD_KEY: POINT_CLOUD_CENTER.unsqueeze(0)}, + {ARTICULATION_POINT_CLOUD_KEY: POINT_CLOUD_CENTER.unsqueeze(0)}, + ), + ) + def test_partial_point_cloud_geometry_is_rejected( + self, + geometry: dict[str, object], + ) -> None: + affordance, _ = _axis_affordance( + "slide", + fallback_axis=torch.tensor([1.0, 0.0, 0.0]), + ) + + with pytest.raises(ValueError, match="requires both"): + ObjectSemantics( + affordance=affordance, + geometry=geometry, + entity_id="partial-geometry-target", + ) + + def test_degenerate_target_point_cloud_is_rejected(self) -> None: + geometry = { + TARGET_LINK_POINT_CLOUD_KEY: POINT_CLOUD_CENTER.repeat(3, 1), + ARTICULATION_POINT_CLOUD_KEY: POINT_CLOUD_CENTER.unsqueeze(0), + } + affordance, _ = _axis_affordance( + "twist", + fallback_axis=torch.tensor([1.0, 0.0, 0.0]), + ) + + with pytest.raises(ValueError, match="degenerate target-link"): + ObjectSemantics( + affordance=affordance, + geometry=geometry, + entity_id="degenerate-geometry-target", + ) + + def test_centered_articulation_neighborhood_is_rejected(self) -> None: + target_points = POINT_CLOUD_CENTER + TARGET_POINT_OFFSETS + geometry = { + TARGET_LINK_POINT_CLOUD_KEY: target_points, + ARTICULATION_POINT_CLOUD_KEY: target_points.clone(), + } + affordance, _ = _axis_affordance( + "press", + fallback_axis=torch.tensor([1.0, 0.0, 0.0]), + ) + + with pytest.raises(ValueError, match="centered on the target link"): + ObjectSemantics( + affordance=affordance, + geometry=geometry, + entity_id="centered-geometry-target", + ) + + def test_press_without_position_uses_outer_surface_opposite_inferred_axis( + self, + ) -> None: + affordance = PressAffordance( + press_axis=torch.tensor([1.0, 0.0, 0.0]), + press_position=None, + ) + + ObjectSemantics( + affordance=affordance, + geometry=_axis_geometry((0.25, -0.125, -1.5)), + entity_id="automatic-press-target", + ) + + assert torch.equal( + affordance.press_axis, + torch.tensor([0.0, 0.0, -1.0]), + ) + expected_surface_center = POINT_CLOUD_CENTER + torch.tensor([0.0, 0.0, 1.0]) + assert affordance.press_position == pytest.approx( + tuple(float(value) for value in expected_surface_center) + ) + + def test_twist_geometry_uses_non_origin_target_center_as_axis_origin(self) -> None: + affordance = TwistAffordance( + grasp_position=(0.0, 0.0, 0.0), + axis_origin=(9.0, 8.0, 7.0), + twist_axis=torch.tensor([1.0, 0.0, 0.0]), + ) + + ObjectSemantics( + affordance=affordance, + geometry=_axis_geometry((0.25, -0.125, -1.5)), + entity_id="automatic-twist-target", + ) + + assert torch.equal( + affordance.twist_axis, + torch.tensor([0.0, 0.0, -1.0]), + ) + assert affordance.axis_origin == pytest.approx( + tuple(float(value) for value in POINT_CLOUD_CENTER) + ) + + def test_twist_without_point_cloud_geometry_preserves_axis_origin_fallback( + self, + ) -> None: + affordance = TwistAffordance( + grasp_position=(0.0, 0.0, 0.0), + axis_origin=(9.0, 8.0, 7.0), + twist_axis=torch.tensor([1.0, 0.0, 0.0]), + ) + + ObjectSemantics( + affordance=affordance, + geometry={}, + entity_id="rigid-twist-target", + ) + + assert torch.equal( + affordance.twist_axis, + torch.tensor([1.0, 0.0, 0.0]), + ) + assert affordance.axis_origin == pytest.approx((9.0, 8.0, 7.0)) + + class TestTwistAffordance: - def test_requires_explicit_grasp_position_and_axis_origin(self): + def test_requires_explicit_grasp_position(self): with pytest.raises(TypeError, match="grasp_position"): TwistAffordance() # type: ignore[call-arg] + def test_requires_axis_origin_before_use_without_point_cloud_geometry(self): + affordance = TwistAffordance(grasp_position=(0.0, 0.0, 0.0)) + + with pytest.raises(ValueError, match="provided explicitly or resolved"): + affordance.require_axis_origin() + @pytest.mark.parametrize( "twist_axis", ( @@ -402,9 +695,11 @@ def test_rejects_explicit_non_revolute_hinge(self): class TestPressAffordance: - def test_requires_explicit_surface_press_position(self): - with pytest.raises(TypeError, match="press_position"): - PressAffordance() # type: ignore[call-arg] + def test_requires_surface_position_before_pose_without_point_cloud_geometry(self): + affordance = PressAffordance() + + with pytest.raises(ValueError, match="provided explicitly or resolved"): + affordance.get_press_pose(torch.eye(4).unsqueeze(0)) @pytest.mark.parametrize( "press_axis", diff --git a/tests/sim/atomic_actions/test_tutorial_utils.py b/tests/sim/atomic_actions/test_tutorial_utils.py index 73ec4d748..0b8c8701a 100644 --- a/tests/sim/atomic_actions/test_tutorial_utils.py +++ b/tests/sim/atomic_actions/test_tutorial_utils.py @@ -132,6 +132,71 @@ ) +def _tutorial_axis_geometry( + neighbor_offset: tuple[float, float, float], +) -> dict[str, torch.Tensor]: + """Build non-origin target-local geometry for tutorial semantics tests.""" + center = torch.tensor([2.0, -3.0, 4.0]) + target_points = center + torch.tensor( + [ + [1.0, 0.0, 0.0], + [-1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, -1.0, 0.0], + [0.0, 0.0, 1.0], + [0.0, 0.0, -1.0], + ] + ) + return { + "target_link_point_cloud": target_points, + "articulation_point_cloud": torch.cat( + ( + target_points, + (center + torch.tensor(neighbor_offset)).unsqueeze(0), + ) + ), + } + + +class _TutorialArticulation: + """Minimal articulation surface used by automatic-axis tutorial tests.""" + + def __init__(self, geometry: dict[str, torch.Tensor]) -> None: + self.device = torch.device("cpu") + self.geometry = geometry + self.sampled_link_name: str | None = None + + def get_link_vert_face(self, link_name: str) -> tuple[torch.Tensor, torch.Tensor]: + del link_name + return ( + self.geometry["target_link_point_cloud"], + torch.tensor( + [ + [0, 2, 4], + [2, 1, 4], + [1, 3, 4], + [3, 0, 4], + [2, 0, 5], + [1, 2, 5], + [3, 1, 5], + [0, 3, 5], + ] + ), + ) + + def get_link_pose(self, link_name: str, *, to_matrix: bool) -> torch.Tensor: + del link_name + assert to_matrix is True + return torch.eye(4).unsqueeze(0) + + def sample_initial_point_clouds( + self, + target_link_name: str, + ) -> dict[str, torch.Tensor]: + self.sampled_link_name = target_link_name + return self.geometry + + def _run_obstacle_animation(*, pace_wall_time: bool) -> tuple[MagicMock, MagicMock]: obstacle = MagicMock() adapter = MagicMock() @@ -155,6 +220,70 @@ def _run_obstacle_animation(*, pace_wall_time: bool) -> tuple[MagicMock, MagicMo return obstacle, adapter +@pytest.mark.parametrize( + ( + "module_name", + "factory_name", + "link_name", + "axis_field", + "neighbor_offset", + "expected_axis", + ), + ( + ( + "slide", + "create_drawer_semantics", + "large_handle_bar", + "translation_axis", + (0.25, 1.5, -0.125), + (0.0, 1.0, 0.0), + ), + ( + "press", + "create_button_semantics", + "button_cap", + "press_axis", + (0.25, -0.125, -1.5), + (0.0, 0.0, -1.0), + ), + ( + "twist", + "create_knob_semantics", + "cap_1", + "twist_axis", + (0.25, -0.125, -1.5), + (0.0, 0.0, -1.0), + ), + ), +) +def test_articulation_tutorial_semantics_resolve_axis_from_initial_point_clouds( + module_name: str, + factory_name: str, + link_name: str, + axis_field: str, + neighbor_offset: tuple[float, float, float], + expected_axis: tuple[float, float, float], +) -> None: + module = importlib.import_module(f"scripts.tutorials.atomic_action.{module_name}") + geometry = _tutorial_axis_geometry(neighbor_offset) + articulation = _TutorialArticulation(geometry) + + with patch.object(module, "Articulation", _TutorialArticulation): + result = getattr(module, factory_name)(articulation) + + semantics = result[0] if isinstance(result, tuple) else result + assert articulation.sampled_link_name == link_name + assert semantics.geometry is geometry + assert torch.equal( + getattr(semantics.affordance, axis_field), + torch.tensor(expected_axis), + ) + if module_name == "press": + assert semantics.affordance.press_position == pytest.approx((2.0, -3.0, 5.0)) + elif module_name == "twist": + assert semantics.affordance.axis_origin == pytest.approx((2.0, -3.0, 4.0)) + + def test_should_wait_for_tutorial_input_is_disabled_for_headless_modes() -> None: assert ( should_wait_for_tutorial_input( diff --git a/tests/sim/objects/test_articulation.py b/tests/sim/objects/test_articulation.py index 03f07983c..41213eee9 100644 --- a/tests/sim/objects/test_articulation.py +++ b/tests/sim/objects/test_articulation.py @@ -41,6 +41,66 @@ ART_PATH = "SlidingBoxDrawer/SlidingBoxDrawer.urdf" NUM_ARENAS = 10 +POINT_CLOUD_TOLERANCE = 1.0e-6 + + +def _make_point_cloud_articulation( + *, + initial_qpos: tuple[float, ...] = (0.25,), + backend_joint_names: tuple[str, ...] = ("joint",), + kinematic_joint_names: tuple[str, ...] | None = None, + link_meshes: dict[str, tuple[torch.Tensor, torch.Tensor]] | None = None, + link_poses: torch.Tensor | None = None, + body_scale: tuple[float, float, float] = (1.0, 1.0, 1.0), +) -> tuple[Articulation, list[tuple[torch.Tensor, list[str]]]]: + """Build a pure-Python articulation double for point-cloud sampling.""" + if link_meshes is None: + link_meshes = { + "target": ( + torch.tensor( + ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)), + dtype=torch.float32, + ), + torch.tensor(((0, 1, 2),), dtype=torch.long), + ) + } + link_names = list(link_meshes) + if link_poses is None: + link_poses = torch.eye(4, dtype=torch.float32).repeat(1, len(link_names), 1, 1) + if kinematic_joint_names is None: + kinematic_joint_names = backend_joint_names + + articulation = object.__new__(Articulation) + articulation.device = torch.device("cpu") + articulation.cfg = SimpleNamespace( + init_qpos=initial_qpos, + body_scale=body_scale, + ) + articulation._data = SimpleNamespace( + dof=len(backend_joint_names), + link_names=link_names, + link_vert_face=link_meshes, + ) + articulation._entities = [ + SimpleNamespace( + get_actived_joint_names=lambda: list(backend_joint_names), + ) + ] + articulation.pk_chain = SimpleNamespace( + get_joint_parameter_names=lambda: list(kinematic_joint_names), + ) + fk_calls: list[tuple[torch.Tensor, list[str]]] = [] + + def compute_fk( + qpos: torch.Tensor, + *, + link_names: list[str], + ) -> torch.Tensor: + fk_calls.append((qpos.clone(), list(link_names))) + return link_poses.clone() + + articulation.compute_fk = compute_fk # type: ignore[method-assign] + return articulation, fk_calls def test_get_qf_returns_all_articulation_joint_efforts(): @@ -96,6 +156,302 @@ def test_get_parent_joint_chain_returns_backend_neutral_child_to_root_values(): assert chain[0].origin_pose[0, 3].item() == 0.0 +@pytest.mark.no_sim +class TestInitialPointCloudSampling: + """Pure CPU coverage for initial articulation surface sampling.""" + + def test_reorders_initial_qpos_into_kinematic_joint_order(self): + articulation, fk_calls = _make_point_cloud_articulation( + initial_qpos=(2.0, 1.0), + backend_joint_names=("joint_b", "joint_a"), + kinematic_joint_names=("joint_a", "joint_b"), + ) + + articulation.sample_initial_point_clouds( + "target", + articulation_point_count=5, + target_point_count=3, + ) + + assert len(fk_calls) == 1 + fk_qpos, fk_link_names = fk_calls[0] + assert torch.equal(fk_qpos, torch.tensor(((1.0, 2.0),))) + assert fk_link_names == ["target"] + + def test_returns_both_clouds_in_target_link_initial_frame(self): + triangle = torch.tensor(((0, 1, 2),), dtype=torch.long) + link_meshes = { + "body": ( + torch.tensor( + ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)), + dtype=torch.float32, + ), + triangle, + ), + "target": ( + torch.tensor( + ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)), + dtype=torch.float32, + ), + triangle, + ), + } + initial_link_poses = torch.eye(4, dtype=torch.float32).repeat(1, 2, 1, 1) + initial_link_poses[0, 0, :3, 3] = torch.tensor((4.0, 0.0, 0.0)) + initial_link_poses[0, 1, :3, :3] = torch.tensor( + ((0.0, -1.0, 0.0), (1.0, 0.0, 0.0), (0.0, 0.0, 1.0)) + ) + initial_link_poses[0, 1, :3, 3] = torch.tensor((1.0, 2.0, 0.0)) + articulation, _ = _make_point_cloud_articulation( + link_meshes=link_meshes, + link_poses=initial_link_poses, + ) + + geometry = articulation.sample_initial_point_clouds( + "target", + articulation_point_count=256, + target_point_count=32, + ) + + assert set(geometry) == { + "target_link_point_cloud", + "articulation_point_cloud", + } + target_points = geometry["target_link_point_cloud"] + assert target_points.shape == (32, 3) + assert torch.allclose( + target_points[:, 2], + torch.zeros(32), + atol=POINT_CLOUD_TOLERANCE, + ) + assert bool((target_points[:, :2] >= -POINT_CLOUD_TOLERANCE).all()) + assert bool( + ( + target_points[:, 0] + target_points[:, 1] <= 1.0 + POINT_CLOUD_TOLERANCE + ).all() + ) + + articulation_points = geometry["articulation_point_cloud"] + assert articulation_points.shape == (256, 3) + body_mask = articulation_points[:, 0] < -0.5 + assert bool(body_mask.any()) + assert bool((~body_mask).any()) + body_points = articulation_points[body_mask] + assert bool( + ( + (body_points[:, 0] >= -2.0 - POINT_CLOUD_TOLERANCE) + & (body_points[:, 0] <= -1.0 + POINT_CLOUD_TOLERANCE) + & (body_points[:, 1] >= -4.0 - POINT_CLOUD_TOLERANCE) + & (body_points[:, 1] <= -3.0 + POINT_CLOUD_TOLERANCE) + ).all() + ) + sampled_target_points = articulation_points[~body_mask] + assert bool((sampled_target_points[:, :2] >= -POINT_CLOUD_TOLERANCE).all()) + assert bool( + ( + sampled_target_points[:, 0] + sampled_target_points[:, 1] + <= 1.0 + POINT_CLOUD_TOLERANCE + ).all() + ) + + def test_sampling_consumes_open3d_rng_without_resetting_it(self): + import open3d as o3d + + articulation, _ = _make_point_cloud_articulation() + o3d.utility.random.seed(7) + + first = articulation.sample_initial_point_clouds( + "target", + articulation_point_count=23, + target_point_count=17, + ) + second = articulation.sample_initial_point_clouds( + "target", + articulation_point_count=23, + target_point_count=17, + ) + + assert not torch.equal( + first["target_link_point_cloud"], + second["target_link_point_cloud"], + ) + assert not torch.equal( + first["articulation_point_cloud"], + second["articulation_point_cloud"], + ) + + def test_surface_sampling_preserves_torch_dtype_and_owns_its_data(self): + vertices = torch.tensor( + ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)), + dtype=torch.float64, + ) + triangles = torch.tensor(((0, 1, 2),), dtype=torch.long) + original_vertices = vertices.clone() + + points = Articulation._sample_mesh_surface_points(vertices, triangles, 11) + + assert points.device == vertices.device + assert points.dtype == vertices.dtype + assert points.shape == (11, 3) + points.zero_() + assert torch.equal(vertices, original_vertices) + + @pytest.mark.parametrize( + "triangles", + ( + torch.empty((0, 3), dtype=torch.long), + torch.tensor(((0, 1, 2),), dtype=torch.long), + ), + ) + def test_surface_sampling_falls_back_to_vertices_without_valid_faces( + self, + triangles: torch.Tensor, + ): + vertices = torch.tensor( + ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (2.0, 0.0, 0.0)), + dtype=torch.float32, + ) + + points = Articulation._sample_mesh_surface_points(vertices, triangles, 7) + + assert points.shape == (7, 3) + assert set(points[:, 0].tolist()) <= {0.0, 1.0, 2.0} + + def test_merged_surface_sampling_is_weighted_by_face_area(self): + triangle = torch.tensor(((0, 1, 2),), dtype=torch.long) + articulation, _ = _make_point_cloud_articulation( + link_meshes={ + "small": ( + torch.tensor( + ((-10.0, 0.0, 0.0), (-9.0, 0.0, 0.0), (-10.0, 1.0, 0.0)), + dtype=torch.float32, + ), + triangle, + ), + "target": ( + torch.tensor( + ((0.0, 0.0, 0.0), (2.0, 0.0, 0.0), (0.0, 2.0, 0.0)), + dtype=torch.float32, + ), + triangle, + ), + } + ) + + geometry = articulation.sample_initial_point_clouds( + "target", + articulation_point_count=5_000, + target_point_count=3, + ) + + articulation_points = geometry["articulation_point_cloud"] + small_face_fraction = float((articulation_points[:, 0] < -9.0).float().mean()) + assert small_face_fraction == pytest.approx(0.2, abs=0.03) + + @pytest.mark.parametrize( + ("target_link_name", "error_type", "message"), + ( + (None, TypeError, "target_link_name must be a string"), + (" ", ValueError, "target_link_name must be non-empty"), + ("missing", ValueError, "Unknown articulation link"), + ), + ) + def test_rejects_invalid_target_link_names( + self, + target_link_name: object, + error_type: type[Exception], + message: str, + ): + articulation, _ = _make_point_cloud_articulation() + + with pytest.raises(error_type, match=message): + articulation.sample_initial_point_clouds(target_link_name) # type: ignore[arg-type] + + @pytest.mark.parametrize( + ("field_name", "value", "error_type", "message"), + ( + ( + "articulation_point_count", + True, + TypeError, + "articulation_point_count must be an integer", + ), + ( + "target_point_count", + 0, + ValueError, + "target_point_count must be positive", + ), + ), + ) + def test_rejects_invalid_point_counts( + self, + field_name: str, + value: object, + error_type: type[Exception], + message: str, + ): + articulation, _ = _make_point_cloud_articulation() + + with pytest.raises(error_type, match=message): + articulation.sample_initial_point_clouds( + "target", + **{field_name: value}, # type: ignore[arg-type] + ) + + def test_requires_a_kinematic_chain(self): + articulation, _ = _make_point_cloud_articulation() + articulation.pk_chain = None + + with pytest.raises(RuntimeError, match="cfg.build_pk_chain=True"): + articulation.sample_initial_point_clouds("target") + + def test_rejects_non_unit_body_scale(self): + articulation, _ = _make_point_cloud_articulation( + body_scale=(1.0, 2.0, 1.0), + ) + + with pytest.raises(ValueError, match="requires unit body_scale"): + articulation.sample_initial_point_clouds("target") + + def test_rejects_mismatched_joint_name_sets(self): + articulation, _ = _make_point_cloud_articulation( + backend_joint_names=("backend_joint",), + kinematic_joint_names=("kinematic_joint",), + ) + + with pytest.raises(ValueError, match="matching simulator and kinematic-chain"): + articulation.sample_initial_point_clouds("target") + + @pytest.mark.parametrize( + "initial_qpos", + ( + (), + (float("nan"),), + ), + ) + def test_rejects_invalid_initial_qpos(self, initial_qpos: tuple[float, ...]): + articulation, _ = _make_point_cloud_articulation( + initial_qpos=initial_qpos, + ) + + with pytest.raises(ValueError, match="finite vector matching"): + articulation.sample_initial_point_clouds("target") + + def test_rejects_invalid_mesh_indices(self): + articulation, _ = _make_point_cloud_articulation( + link_meshes={ + "target": ( + torch.zeros((3, 3), dtype=torch.float32), + torch.tensor(((0, 1, 3),), dtype=torch.long), + ) + } + ) + + with pytest.raises(ValueError, match="triangles reference invalid vertices"): + articulation.sample_initial_point_clouds("target") + + def _link_static_friction(art: Articulation, link_name: str, env_idx: int = 0) -> float: return art._entities[env_idx].get_physical_attr(link_name).static_friction From f963be6e5d79a82d30330a24106bb800ffec48a9 Mon Sep 17 00:00:00 2001 From: matafela Date: Tue, 1 Sep 2026 16:19:38 +0800 Subject: [PATCH 2/8] fix pour tutorial --- scripts/tutorials/atomic_action/pour.py | 19 +++++++++++++++---- .../sim/atomic_actions/test_tutorial_utils.py | 4 ++++ 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/scripts/tutorials/atomic_action/pour.py b/scripts/tutorials/atomic_action/pour.py index eec54b058..fced481d5 100644 --- a/scripts/tutorials/atomic_action/pour.py +++ b/scripts/tutorials/atomic_action/pour.py @@ -38,6 +38,7 @@ PourGoal, PourOptions, ) +from embodichain.lab.sim.planners import ToppraPlanOptions, TrajectorySampleMethod from embodichain.utils import logger from scripts.tutorials.atomic_action.axis_align import ( create_align_object, @@ -64,6 +65,7 @@ POUR_SAMPLE_INTERVAL = 80 HAND_INTERP_STEPS = 12 POST_TRAJECTORY_STEPS = 240 +PICK_MOTION_SAMPLE_COUNT = PICK_SAMPLE_INTERVAL - HAND_INTERP_STEPS OBJ_POSITION = (-0.5, 0.0, 0.0) RECORD_LOOK_AT = ( (-1.5, 0.2, 1.2), @@ -87,6 +89,18 @@ def parse_arguments() -> argparse.Namespace: return parser.parse_args() +def _create_pick_motion_policy() -> MotionPolicy: + """Create the PickUp policy with an explicit valid TOPPRA sample count.""" + return MotionPolicy( + strategy="motion_gen", + sample_count=PICK_SAMPLE_INTERVAL, + plan_opts=ToppraPlanOptions( + sample_method=TrajectorySampleMethod.QUANTITY, + sample_interval=PICK_MOTION_SAMPLE_COUNT, + ), + ) + + def main() -> None: """Plan and replay PickUp followed by Pour.""" args = parse_arguments() @@ -132,10 +146,7 @@ def main() -> None: "pick_up", GraspGoal(semantics), control_parts=control_parts, - motion_policy=MotionPolicy( - strategy="motion_gen", - sample_count=PICK_SAMPLE_INTERVAL, - ), + motion_policy=_create_pick_motion_policy(), skill_options=PickUpOptions( approach_direction=torch.tensor( APPROACH_DIRECTION, diff --git a/tests/sim/atomic_actions/test_tutorial_utils.py b/tests/sim/atomic_actions/test_tutorial_utils.py index 0b8c8701a..7fa8922fe 100644 --- a/tests/sim/atomic_actions/test_tutorial_utils.py +++ b/tests/sim/atomic_actions/test_tutorial_utils.py @@ -833,6 +833,10 @@ def test_pour_tutorial_uses_configured_pickup_and_local_rotation_axis() -> None: assert configured_args.rotate_angle == pytest.approx(-1.25) assert module.APPROACH_DIRECTION == pytest.approx((-0.707, 0.0, -0.707)) assert module.POUR_INTERNAL_AXIS == (1.0, 0.0, 0.0) + pick_policy = module._create_pick_motion_policy() + assert pick_policy.sample_count == module.PICK_SAMPLE_INTERVAL + assert pick_policy.plan_opts.sample_method is module.TrajectorySampleMethod.QUANTITY + assert pick_policy.plan_opts.sample_interval == module.PICK_MOTION_SAMPLE_COUNT def test_replay_timed_trajectory_uses_arrival_intervals() -> None: From 1df12a94d87af50ff4ad30f68f52e98df9eeba5a Mon Sep 17 00:00:00 2001 From: matafela Date: Wed, 2 Sep 2026 10:31:01 +0800 Subject: [PATCH 3/8] update --- scripts/tutorials/atomic_action/tutorial_utils.py | 5 ++++- tests/sim/atomic_actions/test_tutorial_utils.py | 11 ----------- 2 files changed, 4 insertions(+), 12 deletions(-) diff --git a/scripts/tutorials/atomic_action/tutorial_utils.py b/scripts/tutorials/atomic_action/tutorial_utils.py index eac646336..af8e288ce 100644 --- a/scripts/tutorials/atomic_action/tutorial_utils.py +++ b/scripts/tutorials/atomic_action/tutorial_utils.py @@ -93,7 +93,7 @@ finger_thickness=0.01, palm_depth=0.096, ) -DEFAULT_GRIPPER_CLOSE_QPOS = 0.024 +DEFAULT_GRIPPER_CLOSE_QPOS = 0.036 DEFAULT_TUTORIAL_LIGHT_POS = (1.0, 0.0, 3.0) _FRANKA_TUTORIAL_BASE_ROTATION = (0.0, 0.0, 180.0) _DEFAULT_GRIPPER_TCP_Z = 0.17 @@ -924,12 +924,15 @@ def create_ur5_gripper_robot_cfg( }, "drive_pros": { "stiffness": { + "arm": 5e4, GRIPPER_HAND_JOINT_PATTERN: 1e3, }, "damping": { + "arm": 5e3, GRIPPER_HAND_JOINT_PATTERN: 1e2, }, "max_effort": { + "arm": 1e6, GRIPPER_HAND_JOINT_PATTERN: 1e4, }, }, diff --git a/tests/sim/atomic_actions/test_tutorial_utils.py b/tests/sim/atomic_actions/test_tutorial_utils.py index 7fa8922fe..f65eff02d 100644 --- a/tests/sim/atomic_actions/test_tutorial_utils.py +++ b/tests/sim/atomic_actions/test_tutorial_utils.py @@ -603,17 +603,6 @@ def test_dual_franka_mount_preserves_single_arm_facing_direction() -> None: ) -def test_hand_commands_use_pgi_open_limit() -> None: - robot = MagicMock() - robot.device = torch.device("cpu") - robot.get_qpos_limits.return_value = torch.tensor([[[0.0, 0.04]]]) - - hand_open, hand_close = get_hand_open_close_qpos(robot) - - assert torch.allclose(hand_open, torch.tensor([0.0])) - assert torch.allclose(hand_close, torch.tensor([0.024])) - - def test_hand_commands_cover_all_six_robotiq_joints_with_mimic_directions() -> None: robot = MagicMock() robot.device = torch.device("cpu") From 5dc39b3a3cbb883bf0a426052eb3a7e58a41479c Mon Sep 17 00:00:00 2001 From: matafela Date: Wed, 2 Sep 2026 11:41:00 +0800 Subject: [PATCH 4/8] update twist axis_origin --- .../sim/atomic_actions/builtin_actions.md | 18 +- docs/source/overview/sim/sim_articulation.md | 13 +- .../lab/sim/atomic_actions/affordance.py | 42 +++- embodichain/lab/sim/objects/articulation.py | 45 ++++- tests/sim/atomic_actions/test_actions.py | 14 +- tests/sim/atomic_actions/test_affordance.py | 92 ++++++++- .../sim/atomic_actions/test_tutorial_utils.py | 6 +- tests/sim/objects/test_articulation.py | 182 ++++++++++++++++++ 8 files changed, 387 insertions(+), 25 deletions(-) diff --git a/docs/source/overview/sim/atomic_actions/builtin_actions.md b/docs/source/overview/sim/atomic_actions/builtin_actions.md index c67f6867b..461475701 100644 --- a/docs/source/overview/sim/atomic_actions/builtin_actions.md +++ b/docs/source/overview/sim/atomic_actions/builtin_actions.md @@ -823,11 +823,19 @@ Plans **approach -> reach -> close -> twist -> open -> retract** for an articulation link or a rigid object. The entity-free `TwistAffordance` stores an explicit local `grasp_position`, plus optional joint name/limits. For an articulation link, the same initial target-neighborhood geometry used by -`Slide` and `Press` resolves `twist_axis` and sets `axis_origin` to the sampled -target-link cloud centroid. Complete point-cloud geometry overrides both legacy -fallback values, including when the centroid is not the link-frame origin. A -rigid object without articulation context may still provide explicit -compatibility values. `TwistGoal.target_pose` supplies the grounded target pose. +`Slide` and `Press` resolves `twist_axis`. The independent +`target_link_revolute_axis_origin` geometry entry sets `axis_origin` to the +nearest parent revolute joint's initial origin, expressed in the target link's +initial local frame. Resolution walks through fixed parent joints. The sampled +target-link centroid is only a neighborhood/contact reference; it is never used +as the rotation origin. + +Complete point-cloud geometry overrides the legacy `twist_axis` fallback but +does not overwrite an explicit `axis_origin` when revolute-joint metadata is +absent. If neither source supplies an origin, planning reports the missing +rotation-axis point. A rigid object without articulation context may still +provide explicit compatibility values. `TwistGoal.target_pose` supplies the +grounded target pose. The grasp frame's z-axis follows the world-transformed twist axis; an adaptive reference completes a right-handed orthonormal frame. Twist keyframes rotate diff --git a/docs/source/overview/sim/sim_articulation.md b/docs/source/overview/sim/sim_articulation.md index d5012546c..45619c379 100644 --- a/docs/source/overview/sim/sim_articulation.md +++ b/docs/source/overview/sim/sim_articulation.md @@ -181,7 +181,7 @@ State data is accessed via getter methods that return batched tensors (`N` envir | :--- | :--- | :--- | | `get_local_pose(to_matrix=False)` | `(N, 7)` or `(N, 4, 4)` | Root link pose `[x, y, z, qw, qx, qy, qz]` or a 4x4 matrix. | | `get_link_pose(link_name, to_matrix=False)` | `(N, 7)` or `(N, 4, 4)` | Specific link pose `[x, y, z, qw, qx, qy, qz]` or a 4x4 matrix. | -| `sample_initial_point_clouds(target_link_name)` | `Dict[str, Tensor]` | Uniformly sample the target link and merged articulation surface at `init_qpos`, both in the target link's initial local frame. | +| `sample_initial_point_clouds(target_link_name)` | `Dict[str, Tensor]` | Uniformly sample the target link and merged articulation surface at `init_qpos` in the target link's initial local frame, and include the nearest parent revolute-joint origin when available. | | `get_qpos(target=False)` | `(N, dof)` | Current joint positions (or joint targets if `target=True`). | | `get_qvel(target=False)` | `(N, dof)` | Current joint velocities (or velocity targets if `target=True`). | | `get_joint_drive()` | `Tuple[Tensor, ...]` | Returns `(stiffness, damping, max_effort, max_velocity, friction, armature)`, each shaped `(N, dof)`. | @@ -213,6 +213,7 @@ geometry = articulation.sample_initial_point_clouds( ) target_points = geometry["target_link_point_cloud"] articulation_points = geometry["articulation_point_cloud"] +revolute_origin = geometry.get("target_link_revolute_axis_origin") ``` The shown point counts are the defaults. Open3D draws random uniform samples @@ -222,6 +223,16 @@ the spatial distribution rather than point-for-point correspondence. The method requires a built kinematic chain and currently supports unit `body_scale`. +When the target link has a revolute ancestor, the returned +`target_link_revolute_axis_origin` tensor has shape `(3,)`. It is the first +revolute joint found while walking upward from the target link (through any +fixed parent joints), evaluated at the initial joint configuration and +expressed in the target link's initial local frame. If no revolute ancestor +exists, the key is omitted without preventing point-cloud sampling. Atomic +`Twist` uses this metadata as its rotation origin. The target-cloud centroid is +used only to center neighborhood/contact calculations and is not a substitute +for a physical joint origin. + ### Visual Appearance Asset materials are wrapped automatically during articulation construction. Materials are organized by environment and link: diff --git a/embodichain/lab/sim/atomic_actions/affordance.py b/embodichain/lab/sim/atomic_actions/affordance.py index 4caa3e1da..a84d1911b 100644 --- a/embodichain/lab/sim/atomic_actions/affordance.py +++ b/embodichain/lab/sim/atomic_actions/affordance.py @@ -24,6 +24,7 @@ _TARGET_LINK_POINT_CLOUD_KEY = "target_link_point_cloud" _ARTICULATION_POINT_CLOUD_KEY = "articulation_point_cloud" +_TARGET_LINK_REVOLUTE_AXIS_ORIGIN_KEY = "target_link_revolute_axis_origin" @dataclass @@ -267,7 +268,7 @@ class TwistAffordance(Affordance): default=None, kw_only=True, ) - """Fallback axis point, overridden by the target point-cloud centroid.""" + """Fallback axis point, overridden by revolute-joint origin metadata.""" twist_axis: torch.Tensor = field( default_factory=lambda: torch.tensor([0.0, 1.0, 0.0]) @@ -300,16 +301,20 @@ def __post_init__(self) -> None: _validate_joint_metadata(self.joint_name, self.joint_limits) def resolve_from_object_geometry(self, geometry: Mapping[str, Any]) -> None: - """Infer the twist axis and its target-local origin from point clouds.""" + """Resolve the twist axis and target-local revolute-joint origin.""" resolved = _infer_articulation_neighborhood_axis( geometry, field_name="TwistAffordance.twist_axis", ) if resolved is not None: - self.twist_axis, target_points = resolved - self.axis_origin = tuple( - float(component) for component in target_points.mean(dim=0) - ) + self.twist_axis = resolved[0] + + resolved_origin = _resolve_geometry_local_point( + geometry, + key=_TARGET_LINK_REVOLUTE_AXIS_ORIGIN_KEY, + ) + if resolved_origin is not None: + self.axis_origin = resolved_origin def require_axis_origin(self) -> tuple[float, float, float]: """Return the explicit or geometry-derived local rotation-axis point. @@ -324,7 +329,8 @@ def require_axis_origin(self) -> tuple[float, float, float]: if self.axis_origin is None: raise ValueError( "TwistAffordance.axis_origin must be provided explicitly or " - "resolved from articulation point-cloud geometry." + "resolved from geometry['target_link_revolute_axis_origin']; " + "the target link's revolute joint origin is missing." ) return self.axis_origin @@ -843,6 +849,28 @@ def _validate_local_point_cloud(value: Any, *, field_name: str) -> torch.Tensor: return value +def _resolve_geometry_local_point( + geometry: Mapping[str, Any], + *, + key: str, +) -> tuple[float, float, float] | None: + """Resolve an optional finite target-local point tensor from geometry.""" + if key not in geometry: + return None + value = geometry[key] + if not isinstance(value, torch.Tensor): + raise TypeError(f"geometry[{key!r}] must be a torch.Tensor.") + if ( + not value.is_floating_point() + or value.shape != (3,) + or not bool(torch.isfinite(value).all().item()) + ): + raise ValueError( + f"geometry[{key!r}] must be a finite floating tensor with shape (3,)." + ) + return tuple(float(component) for component in value) + + def _outer_surface_center( target_points: torch.Tensor, inward_axis: torch.Tensor, diff --git a/embodichain/lab/sim/objects/articulation.py b/embodichain/lab/sim/objects/articulation.py index ac63d7463..f09ba0709 100644 --- a/embodichain/lab/sim/objects/articulation.py +++ b/embodichain/lab/sim/objects/articulation.py @@ -1224,12 +1224,14 @@ def sample_initial_point_clouds( Geometry metadata containing ``articulation_point_cloud`` with shape ``(articulation_point_count, 3)`` and ``target_link_point_cloud`` with shape - ``(target_point_count, 3)``. + ``(target_point_count, 3)``. If the target link has a revolute + ancestor, ``target_link_revolute_axis_origin`` contains that + nearest joint's origin in the target link's initial local frame. Raises: TypeError: If a name or point count has the wrong type. ValueError: If the target link, point counts, initial joint state, - FK output, or mesh geometry is invalid. + FK output, parent-joint geometry, or mesh geometry is invalid. RuntimeError: If the articulation has no kinematic chain. """ if type(target_link_name) is not str: @@ -1305,6 +1307,38 @@ def sample_initial_point_clouds( ) target_index = link_names.index(target_link_name) target_from_root = torch.linalg.inv(initial_link_poses[target_index]) + target_link_revolute_axis_origin: torch.Tensor | None = None + for joint in self.get_parent_joint_chain(target_link_name): + if joint.joint_type != "revolute": + continue + if joint.parent_link_name not in link_names: + raise ValueError( + f"Revolute joint {joint.name!r} parent link " + f"{joint.parent_link_name!r} is not an articulation link." + ) + joint_origin_pose = torch.as_tensor( + joint.origin_pose, + dtype=torch.float32, + device=self.device, + ) + if ( + joint_origin_pose.shape != (4, 4) + or not torch.isfinite(joint_origin_pose).all() + ): + raise ValueError( + f"Revolute joint {joint.name!r} origin pose must be finite " + "with shape (4, 4)." + ) + parent_index = link_names.index(joint.parent_link_name) + target_from_joint = torch.matmul( + torch.matmul( + target_from_root, + initial_link_poses[parent_index], + ), + joint_origin_pose, + ) + target_link_revolute_axis_origin = target_from_joint[:3, 3].clone() + break merged_vertices: list[torch.Tensor] = [] merged_triangles: list[torch.Tensor] = [] @@ -1359,7 +1393,7 @@ def sample_initial_point_clouds( else torch.empty((0, 3), dtype=torch.long, device=self.device) ) assert target_triangles is not None - return { + geometry = { "target_link_point_cloud": self._sample_mesh_surface_points( target_vertices, target_triangles, @@ -1371,6 +1405,11 @@ def sample_initial_point_clouds( articulation_point_count, ), } + if target_link_revolute_axis_origin is not None: + geometry["target_link_revolute_axis_origin"] = ( + target_link_revolute_axis_origin + ) + return geometry @staticmethod def _validate_point_cloud_mesh( diff --git a/tests/sim/atomic_actions/test_actions.py b/tests/sim/atomic_actions/test_actions.py index 344b903c0..bfa05fbf2 100644 --- a/tests/sim/atomic_actions/test_actions.py +++ b/tests/sim/atomic_actions/test_actions.py @@ -124,7 +124,8 @@ DUAL_ARM_DOF = 2 * ARM_DOF DUAL_ROBOT_DOF = DUAL_ARM_DOF + 2 * HAND_DOF DOOR_ENTITY_ID = "door" -AUTOMATIC_TWIST_AXIS_ORIGIN = torch.tensor([1.0, 0.0, 0.0]) +AUTOMATIC_TWIST_TARGET_CENTER = torch.tensor([1.0, 0.0, 0.0]) +AUTOMATIC_TWIST_AXIS_ORIGIN = torch.tensor([0.25, -0.5, 0.75]) AUTOMATIC_TWIST_TARGET_OFFSETS = torch.tensor( [ [1.0, 0.0, 0.0], @@ -142,9 +143,9 @@ def _automatic_twist_geometry() -> dict[str, torch.Tensor]: - """Build target-local geometry with a non-origin center and +Z neighbor.""" - target_points = AUTOMATIC_TWIST_AXIS_ORIGIN + AUTOMATIC_TWIST_TARGET_OFFSETS - articulation_neighbor = AUTOMATIC_TWIST_AXIS_ORIGIN + torch.tensor( + """Build geometry whose target centroid and revolute origin are distinct.""" + target_points = AUTOMATIC_TWIST_TARGET_CENTER + AUTOMATIC_TWIST_TARGET_OFFSETS + articulation_neighbor = AUTOMATIC_TWIST_TARGET_CENTER + torch.tensor( [0.25, -0.125, 1.5] ) return { @@ -153,6 +154,7 @@ def _automatic_twist_geometry() -> dict[str, torch.Tensor]: (target_points, articulation_neighbor.unsqueeze(0)), dim=0, ), + "target_link_revolute_axis_origin": AUTOMATIC_TWIST_AXIS_ORIGIN.clone(), } @@ -2494,7 +2496,7 @@ def test_twist_plans_from_explicit_rigid_object_pose_snapshot() -> None: assert plan.plan_success.tolist() == [True, True] -def test_twist_rotates_grasp_about_geometry_derived_axis_origin() -> None: +def test_twist_rotates_grasp_about_geometry_derived_joint_axis_origin() -> None: action = _bind_action(_motion_generator(), Twist()) affordance = TwistAffordance(grasp_position=(2.0, 0.0, 0.0)) ObjectSemantics( @@ -2518,7 +2520,7 @@ def test_twist_rotates_grasp_about_geometry_derived_axis_origin() -> None: assert torch.allclose( twisted[:, -1, :3, 3], - torch.tensor([1.0, 1.0, 0.0]).expand(NUM_ENVS, -1), + torch.tensor([-0.25, 1.25, 0.0]).expand(NUM_ENVS, -1), atol=1.0e-6, ) diff --git a/tests/sim/atomic_actions/test_affordance.py b/tests/sim/atomic_actions/test_affordance.py index 4f51eed98..de2bd1f39 100644 --- a/tests/sim/atomic_actions/test_affordance.py +++ b/tests/sim/atomic_actions/test_affordance.py @@ -39,6 +39,7 @@ from embodichain.lab.sim.atomic_actions.core import ObjectSemantics POINT_CLOUD_CENTER = torch.tensor([2.0, -3.0, 4.0]) +REVOLUTE_AXIS_ORIGIN = torch.tensor([0.75, -0.5, 0.25]) TARGET_POINT_OFFSETS = torch.tensor( [ [1.0, 0.0, 0.0], @@ -51,6 +52,7 @@ ) TARGET_LINK_POINT_CLOUD_KEY = "target_link_point_cloud" ARTICULATION_POINT_CLOUD_KEY = "articulation_point_cloud" +TARGET_LINK_REVOLUTE_AXIS_ORIGIN_KEY = "target_link_revolute_axis_origin" def _axis_geometry(neighbor_offset: tuple[float, float, float]) -> dict[str, object]: @@ -68,6 +70,7 @@ def _axis_geometry(neighbor_offset: tuple[float, float, float]) -> dict[str, obj ), dim=0, ), + TARGET_LINK_REVOLUTE_AXIS_ORIGIN_KEY: REVOLUTE_AXIS_ORIGIN.clone(), } @@ -399,7 +402,7 @@ def test_press_without_position_uses_outer_surface_opposite_inferred_axis( tuple(float(value) for value in expected_surface_center) ) - def test_twist_geometry_uses_non_origin_target_center_as_axis_origin(self) -> None: + def test_twist_geometry_uses_revolute_joint_axis_origin(self) -> None: affordance = TwistAffordance( grasp_position=(0.0, 0.0, 0.0), axis_origin=(9.0, 8.0, 7.0), @@ -417,8 +420,93 @@ def test_twist_geometry_uses_non_origin_target_center_as_axis_origin(self) -> No torch.tensor([0.0, 0.0, -1.0]), ) assert affordance.axis_origin == pytest.approx( - tuple(float(value) for value in POINT_CLOUD_CENTER) + tuple(float(value) for value in REVOLUTE_AXIS_ORIGIN) ) + assert not torch.allclose( + torch.tensor(affordance.require_axis_origin()), + POINT_CLOUD_CENTER, + ) + + def test_twist_complete_clouds_without_joint_origin_preserve_fallback( + self, + ) -> None: + geometry = _axis_geometry((0.25, -0.125, -1.5)) + geometry.pop(TARGET_LINK_REVOLUTE_AXIS_ORIGIN_KEY) + affordance = TwistAffordance( + grasp_position=(0.0, 0.0, 0.0), + axis_origin=(9.0, 8.0, 7.0), + twist_axis=torch.tensor([1.0, 0.0, 0.0]), + ) + + ObjectSemantics( + affordance=affordance, + geometry=geometry, + entity_id="automatic-twist-target-with-origin-fallback", + ) + + assert torch.equal( + affordance.twist_axis, + torch.tensor([0.0, 0.0, -1.0]), + ) + assert affordance.axis_origin == pytest.approx((9.0, 8.0, 7.0)) + + def test_twist_complete_clouds_without_joint_origin_do_not_use_centroid( + self, + ) -> None: + geometry = _axis_geometry((0.25, -0.125, -1.5)) + geometry.pop(TARGET_LINK_REVOLUTE_AXIS_ORIGIN_KEY) + affordance = TwistAffordance( + grasp_position=(0.0, 0.0, 0.0), + twist_axis=torch.tensor([1.0, 0.0, 0.0]), + ) + + ObjectSemantics( + affordance=affordance, + geometry=geometry, + entity_id="automatic-twist-target-without-joint-origin", + ) + + assert torch.equal( + affordance.twist_axis, + torch.tensor([0.0, 0.0, -1.0]), + ) + assert affordance.axis_origin is None + with pytest.raises( + ValueError, + match="target_link_revolute_axis_origin", + ): + affordance.require_axis_origin() + + @pytest.mark.parametrize( + ("axis_origin", "exception_type"), + ( + ((0.0, 0.0, 0.0), TypeError), + (torch.tensor([0, 0, 0]), ValueError), + (torch.tensor([0.0, 0.0]), ValueError), + (torch.tensor([0.0, float("nan"), 0.0]), ValueError), + ), + ) + def test_twist_rejects_invalid_revolute_joint_axis_origin_metadata( + self, + axis_origin: object, + exception_type: type[Exception], + ) -> None: + geometry = _axis_geometry((0.25, -0.125, -1.5)) + geometry[TARGET_LINK_REVOLUTE_AXIS_ORIGIN_KEY] = axis_origin + affordance = TwistAffordance( + grasp_position=(0.0, 0.0, 0.0), + axis_origin=(9.0, 8.0, 7.0), + ) + + with pytest.raises( + exception_type, + match="target_link_revolute_axis_origin", + ): + ObjectSemantics( + affordance=affordance, + geometry=geometry, + entity_id="twist-target-with-invalid-joint-origin", + ) def test_twist_without_point_cloud_geometry_preserves_axis_origin_fallback( self, diff --git a/tests/sim/atomic_actions/test_tutorial_utils.py b/tests/sim/atomic_actions/test_tutorial_utils.py index f65eff02d..2cc806306 100644 --- a/tests/sim/atomic_actions/test_tutorial_utils.py +++ b/tests/sim/atomic_actions/test_tutorial_utils.py @@ -130,6 +130,7 @@ "move_end_effector", "move_joints", ) +TUTORIAL_REVOLUTE_AXIS_ORIGIN = torch.tensor([0.75, -0.5, 0.25]) def _tutorial_axis_geometry( @@ -155,6 +156,7 @@ def _tutorial_axis_geometry( (center + torch.tensor(neighbor_offset)).unsqueeze(0), ) ), + "target_link_revolute_axis_origin": TUTORIAL_REVOLUTE_AXIS_ORIGIN.clone(), } @@ -281,7 +283,9 @@ def test_articulation_tutorial_semantics_resolve_axis_from_initial_point_clouds( if module_name == "press": assert semantics.affordance.press_position == pytest.approx((2.0, -3.0, 5.0)) elif module_name == "twist": - assert semantics.affordance.axis_origin == pytest.approx((2.0, -3.0, 4.0)) + assert semantics.affordance.axis_origin == pytest.approx( + tuple(float(value) for value in TUTORIAL_REVOLUTE_AXIS_ORIGIN) + ) def test_should_wait_for_tutorial_input_is_disabled_for_headless_modes() -> None: diff --git a/tests/sim/objects/test_articulation.py b/tests/sim/objects/test_articulation.py index 41213eee9..a6dad2958 100644 --- a/tests/sim/objects/test_articulation.py +++ b/tests/sim/objects/test_articulation.py @@ -52,6 +52,7 @@ def _make_point_cloud_articulation( link_meshes: dict[str, tuple[torch.Tensor, torch.Tensor]] | None = None, link_poses: torch.Tensor | None = None, body_scale: tuple[float, float, float] = (1.0, 1.0, 1.0), + parent_joint_chain: tuple[object, ...] = (), ) -> tuple[Articulation, list[tuple[torch.Tensor, list[str]]]]: """Build a pure-Python articulation double for point-cloud sampling.""" if link_meshes is None: @@ -99,7 +100,16 @@ def compute_fk( fk_calls.append((qpos.clone(), list(link_names))) return link_poses.clone() + def get_parent_joint_chain( + link_name: str, + ) -> tuple[object, ...]: + assert link_name in link_names + return parent_joint_chain + articulation.compute_fk = compute_fk # type: ignore[method-assign] + articulation.get_parent_joint_chain = ( # type: ignore[method-assign] + get_parent_joint_chain + ) return articulation, fk_calls @@ -178,6 +188,178 @@ def test_reorders_initial_qpos_into_kinematic_joint_order(self): assert torch.equal(fk_qpos, torch.tensor(((1.0, 2.0),))) assert fk_link_names == ["target"] + def test_uses_nearest_revolute_ancestor_after_fixed_descendants(self): + triangle = torch.tensor(((0, 1, 2),), dtype=torch.long) + link_meshes = { + link_name: ( + torch.tensor( + ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)), + dtype=torch.float32, + ), + triangle, + ) + for link_name in ("base", "near_parent", "moving", "target") + } + fixed_joint = ArticulationJointKinematics( + name="target_fixed", + joint_type="fixed", + parent_link_name="moving", + child_link_name="target", + origin_pose=torch.eye(4), + axis=torch.zeros(3), + ) + near_origin_pose = torch.eye(4) + near_origin_pose[:3, 3] = torch.tensor((1.0, 2.0, 3.0)) + near_joint = ArticulationJointKinematics( + name="near_hinge", + joint_type="revolute", + parent_link_name="near_parent", + child_link_name="moving", + origin_pose=near_origin_pose, + axis=torch.tensor((0.0, 0.0, 1.0)), + ) + far_origin_pose = torch.eye(4) + far_origin_pose[:3, 3] = torch.tensor((9.0, 9.0, 9.0)) + far_joint = ArticulationJointKinematics( + name="far_hinge", + joint_type="revolute", + parent_link_name="base", + child_link_name="near_parent", + origin_pose=far_origin_pose, + axis=torch.tensor((1.0, 0.0, 0.0)), + ) + articulation, _ = _make_point_cloud_articulation( + link_meshes=link_meshes, + parent_joint_chain=(fixed_joint, near_joint, far_joint), + ) + + geometry = articulation.sample_initial_point_clouds( + "target", + articulation_point_count=16, + target_point_count=8, + ) + + assert torch.equal( + geometry["target_link_revolute_axis_origin"], + near_origin_pose[:3, 3], + ) + + def test_transforms_revolute_origin_from_rotated_parent_to_target_frame(self): + triangle = torch.tensor(((0, 1, 2),), dtype=torch.long) + link_meshes = { + link_name: ( + torch.tensor( + ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)), + dtype=torch.float32, + ), + triangle, + ) + for link_name in ("parent", "target") + } + root_from_parent = torch.eye(4) + root_from_parent[:3, :3] = torch.tensor( + ((0.0, -1.0, 0.0), (1.0, 0.0, 0.0), (0.0, 0.0, 1.0)) + ) + root_from_parent[:3, 3] = torch.tensor((3.0, 4.0, 5.0)) + root_from_target = torch.eye(4) + root_from_target[:3, :3] = torch.tensor( + ((1.0, 0.0, 0.0), (0.0, 0.0, -1.0), (0.0, 1.0, 0.0)) + ) + root_from_target[:3, 3] = torch.tensor((0.5, -1.0, 2.0)) + joint_origin_pose = torch.eye(4) + joint_origin_pose[:3, :3] = torch.tensor( + ((0.0, 0.0, 1.0), (0.0, 1.0, 0.0), (-1.0, 0.0, 0.0)) + ) + joint_origin_pose[:3, 3] = torch.tensor((1.0, 2.0, 3.0)) + link_poses = torch.stack((root_from_parent, root_from_target)).unsqueeze(0) + joint = ArticulationJointKinematics( + name="target_hinge", + joint_type="revolute", + parent_link_name="parent", + child_link_name="target", + origin_pose=joint_origin_pose, + axis=torch.tensor((0.0, 1.0, 0.0)), + ) + articulation, _ = _make_point_cloud_articulation( + link_meshes=link_meshes, + link_poses=link_poses, + parent_joint_chain=(joint,), + ) + + geometry = articulation.sample_initial_point_clouds( + "target", + articulation_point_count=16, + target_point_count=8, + ) + + target_from_joint = ( + torch.linalg.inv(root_from_target) @ root_from_parent @ joint_origin_pose + ) + assert torch.allclose( + geometry["target_link_revolute_axis_origin"], + target_from_joint[:3, 3], + atol=POINT_CLOUD_TOLERANCE, + ) + + def test_omits_revolute_origin_when_parent_chain_has_only_fixed_joints(self): + fixed_joint = ArticulationJointKinematics( + name="target_fixed", + joint_type="fixed", + parent_link_name="parent", + child_link_name="target", + origin_pose=torch.eye(4), + axis=torch.zeros(3), + ) + articulation, _ = _make_point_cloud_articulation( + parent_joint_chain=(fixed_joint,), + ) + + geometry = articulation.sample_initial_point_clouds( + "target", + articulation_point_count=5, + target_point_count=3, + ) + + assert "target_link_revolute_axis_origin" not in geometry + + def test_rejects_revolute_ancestor_with_unknown_parent_link(self): + joint = ArticulationJointKinematics( + name="target_hinge", + joint_type="revolute", + parent_link_name="missing_parent", + child_link_name="target", + origin_pose=torch.eye(4), + axis=torch.tensor((0.0, 0.0, 1.0)), + ) + articulation, _ = _make_point_cloud_articulation( + parent_joint_chain=(joint,), + ) + + with pytest.raises(ValueError, match="parent link.*not an articulation link"): + articulation.sample_initial_point_clouds( + "target", + articulation_point_count=5, + target_point_count=3, + ) + + def test_rejects_revolute_ancestor_with_invalid_origin_pose(self): + joint = SimpleNamespace( + name="target_hinge", + joint_type="revolute", + parent_link_name="target", + origin_pose=torch.full((4, 4), float("nan")), + ) + articulation, _ = _make_point_cloud_articulation( + parent_joint_chain=(joint,), + ) + + with pytest.raises(ValueError, match="origin pose must be finite"): + articulation.sample_initial_point_clouds( + "target", + articulation_point_count=5, + target_point_count=3, + ) + def test_returns_both_clouds_in_target_link_initial_frame(self): triangle = torch.tensor(((0, 1, 2),), dtype=torch.long) link_meshes = { From 2c1f53a0136a9844af83c446b9ad2df018738960 Mon Sep 17 00:00:00 2001 From: matafela Date: Wed, 2 Sep 2026 12:13:30 +0800 Subject: [PATCH 5/8] update --- .../sim/atomic_actions/builtin_actions.md | 44 ++- docs/source/overview/sim/sim_articulation.md | 27 +- .../lab/sim/atomic_actions/affordance.py | 97 +++++-- embodichain/lab/sim/objects/articulation.py | 59 +++- .../task_program/test_task_vertical_slices.py | 6 +- tests/sim/atomic_actions/test_actions.py | 17 +- tests/sim/atomic_actions/test_affordance.py | 273 +++++++++++++----- .../sim/atomic_actions/test_tutorial_utils.py | 23 +- tests/sim/objects/test_articulation.py | 209 +++++++++++++- 9 files changed, 594 insertions(+), 161 deletions(-) diff --git a/docs/source/overview/sim/atomic_actions/builtin_actions.md b/docs/source/overview/sim/atomic_actions/builtin_actions.md index 461475701..b351e405a 100644 --- a/docs/source/overview/sim/atomic_actions/builtin_actions.md +++ b/docs/source/overview/sim/atomic_actions/builtin_actions.md @@ -679,10 +679,12 @@ migration. Plans **close hand -> approach target -> contact -> press along axis -> return to the approach pose**. For an articulation link, `Articulation.sample_initial_point_clouds()` stores target-link-local target and -whole-articulation clouds in `ObjectSemantics.geometry`. `PressAffordance` -derives `press_axis` from that geometry and can derive the outer-surface -`press_position` when it is omitted. `PressGoal.target_pose` is either a pose -snapshot or `SceneEntityPose`, which resolves through the current +whole-articulation clouds in `ObjectSemantics.geometry`, together with the +nearest parent prismatic joint axis transformed into the same frame. +`PressAffordance` uses `target_link_prismatic_joint_axis` for its axis direction; +the point-cloud neighborhood selects only the sign. It can also derive the +outer-surface `press_position` when it is omitted. `PressGoal.target_pose` is +either a pose snapshot or `SceneEntityPose`, which resolves through the current `PlanningContext.scene` and participates in dynamic-goal recovery. The contact, press, and retract segments use axis-aligned Cartesian keyframes; @@ -728,11 +730,18 @@ Axis inference samples the target link and the merged articulation surface at Both clouds use Open3D uniform surface sampling. Sampling the whole articulation as one merged mesh preserves triangle-area weighting instead of giving every link an equal point budget, which would over-represent tiny decorative links. -The target cloud center and twice its distribution radius define a spherical -neighborhood in the full cloud. The largest signed component of the -neighborhood-center offset selects exactly one of `+X`, `-X`, `+Y`, `-Y`, `+Z`, -or `-Z`. A complete point-cloud geometry snapshot overrides any legacy explicit -axis fallback. +The sampler also walks the target's parent chain and transforms the nearest +prismatic joint axis into the target link's initial local frame. The target +cloud center and twice its distribution radius define a spherical neighborhood +in the full cloud. The dot product between the neighborhood-center offset and +the normalized joint axis selects its sign. It does not quantize an oblique +joint axis to a Cartesian basis direction; an offset perpendicular to the joint +axis is rejected as directionally ambiguous. + +Automatic resolution requires all three entries: +`target_link_point_cloud`, `articulation_point_cloud`, and +`target_link_prismatic_joint_axis`. With none of those inference entries, an +explicit compatibility axis is preserved; a partial set is rejected. With `direction="pull"`, the sequence is **approach -> reach -> close -> pull -> open**. With `direction="push"`, it is **approach -> reach -> close -> push -> open @@ -823,19 +832,22 @@ Plans **approach -> reach -> close -> twist -> open -> retract** for an articulation link or a rigid object. The entity-free `TwistAffordance` stores an explicit local `grasp_position`, plus optional joint name/limits. For an articulation link, the same initial target-neighborhood geometry used by -`Slide` and `Press` resolves `twist_axis`. The independent +`Slide` and `Press` signs the normalized `target_link_revolute_joint_axis` to +resolve `twist_axis`. The independent `target_link_revolute_axis_origin` geometry entry sets `axis_origin` to the nearest parent revolute joint's initial origin, expressed in the target link's initial local frame. Resolution walks through fixed parent joints. The sampled target-link centroid is only a neighborhood/contact reference; it is never used as the rotation origin. -Complete point-cloud geometry overrides the legacy `twist_axis` fallback but -does not overwrite an explicit `axis_origin` when revolute-joint metadata is -absent. If neither source supplies an origin, planning reports the missing -rotation-axis point. A rigid object without articulation context may still -provide explicit compatibility values. `TwistGoal.target_pose` supplies the -grounded target pose. +Automatic axis resolution requires the revolute joint-axis entry and both point +clouds. When all three axis-inference entries are absent, the legacy +`twist_axis` fallback is preserved; a partial set is rejected. Axis-origin +resolution is independent, so missing origin metadata does not overwrite an +explicit `axis_origin`. If neither source supplies an origin, planning reports +the missing rotation-axis point. A rigid object without articulation context +may still provide explicit compatibility values. `TwistGoal.target_pose` +supplies the grounded target pose. The grasp frame's z-axis follows the world-transformed twist axis; an adaptive reference completes a right-handed orthonormal frame. Twist keyframes rotate diff --git a/docs/source/overview/sim/sim_articulation.md b/docs/source/overview/sim/sim_articulation.md index 45619c379..e0888bec0 100644 --- a/docs/source/overview/sim/sim_articulation.md +++ b/docs/source/overview/sim/sim_articulation.md @@ -181,7 +181,7 @@ State data is accessed via getter methods that return batched tensors (`N` envir | :--- | :--- | :--- | | `get_local_pose(to_matrix=False)` | `(N, 7)` or `(N, 4, 4)` | Root link pose `[x, y, z, qw, qx, qy, qz]` or a 4x4 matrix. | | `get_link_pose(link_name, to_matrix=False)` | `(N, 7)` or `(N, 4, 4)` | Specific link pose `[x, y, z, qw, qx, qy, qz]` or a 4x4 matrix. | -| `sample_initial_point_clouds(target_link_name)` | `Dict[str, Tensor]` | Uniformly sample the target link and merged articulation surface at `init_qpos` in the target link's initial local frame, and include the nearest parent revolute-joint origin when available. | +| `sample_initial_point_clouds(target_link_name)` | `Dict[str, Tensor]` | Uniformly sample target and merged articulation surfaces at `init_qpos`, and include nearest parent prismatic/revolute axes and the revolute origin in the target link's initial local frame when available. | | `get_qpos(target=False)` | `(N, dof)` | Current joint positions (or joint targets if `target=True`). | | `get_qvel(target=False)` | `(N, dof)` | Current joint velocities (or velocity targets if `target=True`). | | `get_joint_drive()` | `Tuple[Tensor, ...]` | Returns `(stiffness, damping, max_effort, max_velocity, friction, armature)`, each shaped `(N, dof)`. | @@ -213,6 +213,8 @@ geometry = articulation.sample_initial_point_clouds( ) target_points = geometry["target_link_point_cloud"] articulation_points = geometry["articulation_point_cloud"] +prismatic_axis = geometry.get("target_link_prismatic_joint_axis") +revolute_axis = geometry.get("target_link_revolute_joint_axis") revolute_origin = geometry.get("target_link_revolute_axis_origin") ``` @@ -223,15 +225,20 @@ the spatial distribution rather than point-for-point correspondence. The method requires a built kinematic chain and currently supports unit `body_scale`. -When the target link has a revolute ancestor, the returned -`target_link_revolute_axis_origin` tensor has shape `(3,)`. It is the first -revolute joint found while walking upward from the target link (through any -fixed parent joints), evaluated at the initial joint configuration and -expressed in the target link's initial local frame. If no revolute ancestor -exists, the key is omitted without preventing point-cloud sampling. Atomic -`Twist` uses this metadata as its rotation origin. The target-cloud centroid is -used only to center neighborhood/contact calculations and is not a substitute -for a physical joint origin. +The sampler walks the immediate-parent-first joint chain and records the nearest +ancestor of each supported type. `target_link_prismatic_joint_axis` and +`target_link_revolute_joint_axis` are normalized `(3,)` axes transformed by the +initial parent-link and joint-origin rotations into the target link's initial +local frame. `target_link_revolute_axis_origin` is the matching nearest +revolute-joint origin in that frame. A key is omitted when the corresponding +joint type has no ancestor; this does not prevent point-cloud sampling. + +`Slide` and `Press` consume the prismatic axis, while `Twist` consumes the +revolute axis and origin. The point-cloud neighborhood only chooses between the +stored axis and its negation by projecting the neighborhood-center offset onto +the axis. It never replaces an oblique physical axis with a Cartesian basis +direction. The target-cloud centroid centers neighborhood/contact calculations +and is not a substitute for the revolute joint origin. ### Visual Appearance diff --git a/embodichain/lab/sim/atomic_actions/affordance.py b/embodichain/lab/sim/atomic_actions/affordance.py index a84d1911b..0c3e52938 100644 --- a/embodichain/lab/sim/atomic_actions/affordance.py +++ b/embodichain/lab/sim/atomic_actions/affordance.py @@ -24,6 +24,8 @@ _TARGET_LINK_POINT_CLOUD_KEY = "target_link_point_cloud" _ARTICULATION_POINT_CLOUD_KEY = "articulation_point_cloud" +_TARGET_LINK_PRISMATIC_JOINT_AXIS_KEY = "target_link_prismatic_joint_axis" +_TARGET_LINK_REVOLUTE_JOINT_AXIS_KEY = "target_link_revolute_joint_axis" _TARGET_LINK_REVOLUTE_AXIS_ORIGIN_KEY = "target_link_revolute_axis_origin" @@ -259,7 +261,7 @@ def __post_init__(self) -> None: @dataclass class TwistAffordance(Affordance): - """Target-local grasp point and inferred rotation-axis geometry.""" + """Target-local grasp point and parent-joint rotation geometry.""" grasp_position: tuple[float, float, float] = field(kw_only=True) """Explicit target-local center of the gripper contact region.""" @@ -273,7 +275,7 @@ class TwistAffordance(Affordance): twist_axis: torch.Tensor = field( default_factory=lambda: torch.tensor([0.0, 1.0, 0.0]) ) - """Twist axis in the target frame, overridden by articulation geometry.""" + """Parent revolute-joint axis, signed toward articulation geometry.""" joint_name: str | None = None """Optional stable articulation-joint name associated with the axis.""" @@ -301,10 +303,11 @@ def __post_init__(self) -> None: _validate_joint_metadata(self.joint_name, self.joint_limits) def resolve_from_object_geometry(self, geometry: Mapping[str, Any]) -> None: - """Resolve the twist axis and target-local revolute-joint origin.""" + """Resolve the target-local revolute axis, sign, and joint origin.""" resolved = _infer_articulation_neighborhood_axis( geometry, field_name="TwistAffordance.twist_axis", + axis_key=_TARGET_LINK_REVOLUTE_JOINT_AXIS_KEY, ) if resolved is not None: self.twist_axis = resolved[0] @@ -374,7 +377,7 @@ def get_grasp_pose(self, target_pose: torch.Tensor) -> torch.Tensor: @dataclass class SlideAffordance(AntipodalAffordance): - """Target-local antipodal grasp and inferred translation-axis geometry. + """Target-local antipodal grasp and parent-joint translation geometry. The positive translation-axis direction denotes approaching and pushing the articulated part closed. Pulling moves in the opposite direction. @@ -391,7 +394,7 @@ class SlideAffordance(AntipodalAffordance): translation_axis: torch.Tensor = field( default_factory=lambda: torch.tensor([0.0, 1.0, 0.0]) ) - """Approach/push axis, overridden by articulation point-cloud geometry.""" + """Parent prismatic-joint axis, signed toward articulation geometry.""" joint_name: str | None = None """Optional stable prismatic-joint name associated with the link.""" @@ -415,10 +418,11 @@ def __post_init__(self) -> None: _validate_joint_metadata(self.joint_name, self.joint_limits) def resolve_from_object_geometry(self, geometry: Mapping[str, Any]) -> None: - """Infer the signed Cartesian translation axis from local point clouds.""" + """Resolve the target-local prismatic axis and neighborhood sign.""" resolved = _infer_articulation_neighborhood_axis( geometry, field_name="SlideAffordance.translation_axis", + axis_key=_TARGET_LINK_PRISMATIC_JOINT_AXIS_KEY, ) if resolved is not None: self.translation_axis = resolved[0] @@ -621,12 +625,12 @@ def from_articulation( @dataclass class PressAffordance(Affordance): - """Target-local contact point and inferred pressing direction.""" + """Target-local contact point and parent-joint pressing geometry.""" press_axis: torch.Tensor = field( default_factory=lambda: torch.tensor([0.0, 0.0, 1.0]) ) - """Press direction in the target frame, overridden by point-cloud geometry.""" + """Parent prismatic-joint axis, signed toward articulation geometry.""" press_position: tuple[float, float, float] | None = field( default=None, @@ -650,10 +654,11 @@ def __post_init__(self) -> None: ) def resolve_from_object_geometry(self, geometry: Mapping[str, Any]) -> None: - """Infer the press axis and optional outer-surface contact point.""" + """Resolve the target-local prismatic axis, sign, and contact point.""" resolved = _infer_articulation_neighborhood_axis( geometry, field_name="PressAffordance.press_axis", + axis_key=_TARGET_LINK_PRISMATIC_JOINT_AXIS_KEY, ) if resolved is None: return @@ -701,7 +706,7 @@ def get_press_pose( if configured_position is None: raise ValueError( "PressAffordance.press_position must be provided explicitly or " - "resolved from articulation point-cloud geometry." + "resolved from articulation joint geometry." ) local_press_position = torch.tensor( configured_position, @@ -744,38 +749,42 @@ def _infer_articulation_neighborhood_axis( geometry: Mapping[str, Any], *, field_name: str, + axis_key: str, ) -> tuple[torch.Tensor, torch.Tensor] | None: - """Infer one signed Cartesian axis from target-centered local geometry. + """Resolve one signed joint axis from target-centered local geometry. The target-link point-cloud center defines a spherical neighborhood in the complete articulation cloud. Its radius is twice the target cloud's maximum - distance from that center. The largest signed component of the neighborhood - center offset selects exactly one of ``+/-X``, ``+/-Y``, or ``+/-Z``. + distance from that center. The neighborhood-center offset disambiguates the + sign of the normalized joint axis supplied in target-link coordinates. Args: - geometry: Object geometry containing both point-cloud entries. + geometry: Object geometry containing joint-axis and point-cloud entries. field_name: Affordance field name used in validation errors. + axis_key: Geometry key containing the target-local joint axis. Returns: - ``(axis, target_points)`` when point-cloud geometry is present, otherwise - ``None`` for legacy or non-articulation affordances. + ``(axis, target_points)`` when joint-axis metadata is present, otherwise + ``None`` so legacy and non-articulation affordances retain their axis. Raises: - TypeError: If geometry or a point cloud has an invalid type. - ValueError: If point-cloud metadata is incomplete, malformed, or - geometrically degenerate. + TypeError: If geometry, a point cloud, or the joint axis has an invalid + type. + ValueError: If geometry metadata is incomplete, malformed, + geometrically degenerate, or directionally ambiguous. """ if not isinstance(geometry, Mapping): raise TypeError("geometry must be a mapping.") has_target = _TARGET_LINK_POINT_CLOUD_KEY in geometry has_articulation = _ARTICULATION_POINT_CLOUD_KEY in geometry - if not has_target and not has_articulation: + has_axis = axis_key in geometry + if not has_target and not has_articulation and not has_axis: return None - if not has_target or not has_articulation: + if not has_target or not has_articulation or not has_axis: raise ValueError( - f"{field_name} inference requires both " - f"{_TARGET_LINK_POINT_CLOUD_KEY!r} and " - f"{_ARTICULATION_POINT_CLOUD_KEY!r}." + f"{field_name} inference requires " + f"{_TARGET_LINK_POINT_CLOUD_KEY!r}, " + f"{_ARTICULATION_POINT_CLOUD_KEY!r}, and {axis_key!r}." ) target_points = _validate_local_point_cloud( @@ -790,8 +799,16 @@ def _infer_articulation_neighborhood_axis( raise ValueError( "Articulation and target-link point clouds must share a device." ) + joint_axis = _validate_geometry_axis( + geometry[axis_key], + field_name=f"geometry[{axis_key!r}]", + ) + if joint_axis.device != target_points.device: + raise ValueError("Joint axis and point clouds must share a device.") target_points = target_points.to(dtype=torch.float32) articulation_points = articulation_points.to(dtype=torch.float32) + joint_axis = joint_axis.to(dtype=torch.float32) + joint_axis = joint_axis / torch.linalg.vector_norm(joint_axis) target_center = target_points.mean(dim=0) target_distances = torch.linalg.vector_norm( @@ -817,18 +834,17 @@ def _infer_articulation_neighborhood_axis( raise ValueError(f"{field_name} point-cloud neighborhood is empty.") neighborhood_center = articulation_points[neighborhood_mask].mean(dim=0) center_offset = neighborhood_center - target_center - axis_index = int(torch.argmax(torch.abs(center_offset)).item()) - dominant_offset = center_offset[axis_index] + direction_score = torch.dot(center_offset, joint_axis) offset_tolerance = max(1.0e-8, float(target_radius.item()) * 1.0e-6) - if abs(float(dominant_offset.item())) <= offset_tolerance: + if abs(float(direction_score.item())) <= offset_tolerance: raise ValueError( - f"{field_name} cannot be inferred because the articulation " - "neighborhood is centered on the target link." + f"{field_name} direction is ambiguous because the articulation " + "neighborhood-center offset is orthogonal to the joint axis." ) - axis = torch.zeros(3, dtype=torch.float32, device=target_points.device) - axis[axis_index] = 1.0 if float(dominant_offset.item()) > 0.0 else -1.0 - return axis, target_points + if float(direction_score.item()) < 0.0: + joint_axis = -joint_axis + return joint_axis, target_points def _validate_local_point_cloud(value: Any, *, field_name: str) -> torch.Tensor: @@ -849,6 +865,23 @@ def _validate_local_point_cloud(value: Any, *, field_name: str) -> torch.Tensor: return value +def _validate_geometry_axis(value: Any, *, field_name: str) -> torch.Tensor: + """Validate one finite, non-zero floating geometry axis.""" + if not isinstance(value, torch.Tensor): + raise TypeError(f"{field_name} must be a torch.Tensor.") + if ( + not value.is_floating_point() + or value.shape != (3,) + or not bool(torch.isfinite(value).all().item()) + ): + raise ValueError( + f"{field_name} must be a finite floating tensor with shape (3,)." + ) + if float(torch.linalg.vector_norm(value).item()) <= 1.0e-6: + raise ValueError(f"{field_name} must be non-zero.") + return value + + def _resolve_geometry_local_point( geometry: Mapping[str, Any], *, diff --git a/embodichain/lab/sim/objects/articulation.py b/embodichain/lab/sim/objects/articulation.py index f09ba0709..713578203 100644 --- a/embodichain/lab/sim/objects/articulation.py +++ b/embodichain/lab/sim/objects/articulation.py @@ -1227,6 +1227,9 @@ def sample_initial_point_clouds( ``(target_point_count, 3)``. If the target link has a revolute ancestor, ``target_link_revolute_axis_origin`` contains that nearest joint's origin in the target link's initial local frame. + ``target_link_prismatic_joint_axis`` and + ``target_link_revolute_joint_axis`` contain the normalized axes of + the nearest ancestors of each type in the same frame when present. Raises: TypeError: If a name or point count has the wrong type. @@ -1308,12 +1311,19 @@ def sample_initial_point_clouds( target_index = link_names.index(target_link_name) target_from_root = torch.linalg.inv(initial_link_poses[target_index]) target_link_revolute_axis_origin: torch.Tensor | None = None + target_link_joint_axes: dict[str, torch.Tensor] = {} + joint_axis_geometry_keys = { + "prismatic": "target_link_prismatic_joint_axis", + "revolute": "target_link_revolute_joint_axis", + } for joint in self.get_parent_joint_chain(target_link_name): - if joint.joint_type != "revolute": + joint_type = joint.joint_type + geometry_key = joint_axis_geometry_keys.get(joint_type) + if geometry_key is None or geometry_key in target_link_joint_axes: continue if joint.parent_link_name not in link_names: raise ValueError( - f"Revolute joint {joint.name!r} parent link " + f"{joint_type.capitalize()} joint {joint.name!r} parent link " f"{joint.parent_link_name!r} is not an articulation link." ) joint_origin_pose = torch.as_tensor( @@ -1326,9 +1336,25 @@ def sample_initial_point_clouds( or not torch.isfinite(joint_origin_pose).all() ): raise ValueError( - f"Revolute joint {joint.name!r} origin pose must be finite " + f"{joint_type.capitalize()} joint {joint.name!r} origin pose " + "must be finite " "with shape (4, 4)." ) + joint_axis = torch.as_tensor( + joint.axis, + dtype=torch.float32, + device=self.device, + ) + if ( + joint_axis.shape != (3,) + or not torch.isfinite(joint_axis).all() + or torch.linalg.vector_norm(joint_axis) + <= torch.finfo(joint_axis.dtype).eps + ): + raise ValueError( + f"{joint_type.capitalize()} joint {joint.name!r} axis must " + "be finite and nonzero with shape (3,)." + ) parent_index = link_names.index(joint.parent_link_name) target_from_joint = torch.matmul( torch.matmul( @@ -1337,8 +1363,30 @@ def sample_initial_point_clouds( ), joint_origin_pose, ) - target_link_revolute_axis_origin = target_from_joint[:3, 3].clone() - break + target_link_joint_axis = torch.matmul( + target_from_joint[:3, :3], + joint_axis, + ) + target_link_joint_axis_norm = torch.linalg.vector_norm( + target_link_joint_axis + ) + if ( + not torch.isfinite(target_link_joint_axis).all() + or not torch.isfinite(target_link_joint_axis_norm) + or target_link_joint_axis_norm + <= torch.finfo(target_link_joint_axis.dtype).eps + ): + raise ValueError( + f"{joint_type.capitalize()} joint {joint.name!r} axis must " + "transform to a finite, nonzero target-link vector." + ) + target_link_joint_axes[geometry_key] = ( + target_link_joint_axis / target_link_joint_axis_norm + ) + if joint_type == "revolute": + target_link_revolute_axis_origin = target_from_joint[:3, 3].clone() + if len(target_link_joint_axes) == len(joint_axis_geometry_keys): + break merged_vertices: list[torch.Tensor] = [] merged_triangles: list[torch.Tensor] = [] @@ -1409,6 +1457,7 @@ def sample_initial_point_clouds( geometry["target_link_revolute_axis_origin"] = ( target_link_revolute_axis_origin ) + geometry.update(target_link_joint_axes) return geometry @staticmethod diff --git a/tests/gym/envs/task_program/test_task_vertical_slices.py b/tests/gym/envs/task_program/test_task_vertical_slices.py index cde3947f5..0e3efd0c9 100644 --- a/tests/gym/envs/task_program/test_task_vertical_slices.py +++ b/tests/gym/envs/task_program/test_task_vertical_slices.py @@ -102,7 +102,7 @@ def _drawer_handle_point_cloud_geometry() -> dict[str, torch.Tensor]: - """Return target-local clouds whose articulation neighborhood lies along +Y.""" + """Return target-local clouds and a prismatic axis oriented toward +Y.""" target_points = torch.tensor( [ [-0.10, 0.0, 0.0], @@ -124,6 +124,10 @@ def _drawer_handle_point_cloud_geometry() -> dict[str, torch.Tensor]: return { "target_link_point_cloud": target_points, "articulation_point_cloud": torch.cat((target_points, body_points), dim=0), + "target_link_prismatic_joint_axis": torch.tensor( + [0.0, 1.0, 0.0], + dtype=torch.float32, + ), } diff --git a/tests/sim/atomic_actions/test_actions.py b/tests/sim/atomic_actions/test_actions.py index bfa05fbf2..d378d5dc4 100644 --- a/tests/sim/atomic_actions/test_actions.py +++ b/tests/sim/atomic_actions/test_actions.py @@ -126,6 +126,10 @@ DOOR_ENTITY_ID = "door" AUTOMATIC_TWIST_TARGET_CENTER = torch.tensor([1.0, 0.0, 0.0]) AUTOMATIC_TWIST_AXIS_ORIGIN = torch.tensor([0.25, -0.5, 0.75]) +AUTOMATIC_TWIST_JOINT_AXIS = torch.tensor([2.0, 2.0, 1.0]) +AUTOMATIC_TWIST_JOINT_AXIS_UNIT = AUTOMATIC_TWIST_JOINT_AXIS / torch.linalg.vector_norm( + AUTOMATIC_TWIST_JOINT_AXIS +) AUTOMATIC_TWIST_TARGET_OFFSETS = torch.tensor( [ [1.0, 0.0, 0.0], @@ -145,8 +149,8 @@ def _automatic_twist_geometry() -> dict[str, torch.Tensor]: """Build geometry whose target centroid and revolute origin are distinct.""" target_points = AUTOMATIC_TWIST_TARGET_CENTER + AUTOMATIC_TWIST_TARGET_OFFSETS - articulation_neighbor = AUTOMATIC_TWIST_TARGET_CENTER + torch.tensor( - [0.25, -0.125, 1.5] + articulation_neighbor = ( + AUTOMATIC_TWIST_TARGET_CENTER + 1.5 * AUTOMATIC_TWIST_JOINT_AXIS_UNIT ) return { "target_link_point_cloud": target_points, @@ -154,6 +158,7 @@ def _automatic_twist_geometry() -> dict[str, torch.Tensor]: (target_points, articulation_neighbor.unsqueeze(0)), dim=0, ), + "target_link_revolute_joint_axis": AUTOMATIC_TWIST_JOINT_AXIS.clone(), "target_link_revolute_axis_origin": AUTOMATIC_TWIST_AXIS_ORIGIN.clone(), } @@ -2456,7 +2461,11 @@ def test_twist_plans_six_segments_from_articulation_link() -> None: assert torch.all( trajectory.positions[:, plan.segment("open").stop - 1, ARM_DOF:] == 0.0 ) - assert torch.equal(affordance.twist_axis, torch.tensor([0.0, 0.0, 1.0])) + assert torch.allclose( + affordance.twist_axis, + AUTOMATIC_TWIST_JOINT_AXIS_UNIT, + atol=1.0e-6, + ) assert affordance.require_axis_origin() == pytest.approx( tuple(float(value) for value in AUTOMATIC_TWIST_AXIS_ORIGIN) ) @@ -2520,7 +2529,7 @@ def test_twist_rotates_grasp_about_geometry_derived_joint_axis_origin() -> None: assert torch.allclose( twisted[:, -1, :3, 3], - torch.tensor([-0.25, 1.25, 0.0]).expand(NUM_ENVS, -1), + torch.tensor([5.0 / 12.0, 17.0 / 12.0, 1.0 / 3.0]).expand(NUM_ENVS, -1), atol=1.0e-6, ) diff --git a/tests/sim/atomic_actions/test_affordance.py b/tests/sim/atomic_actions/test_affordance.py index de2bd1f39..280247605 100644 --- a/tests/sim/atomic_actions/test_affordance.py +++ b/tests/sim/atomic_actions/test_affordance.py @@ -39,6 +39,14 @@ from embodichain.lab.sim.atomic_actions.core import ObjectSemantics POINT_CLOUD_CENTER = torch.tensor([2.0, -3.0, 4.0]) +PRISMATIC_JOINT_AXIS = torch.tensor([2.0, 2.0, 1.0]) +REVOLUTE_JOINT_AXIS = torch.tensor([-1.0, 2.0, 2.0]) +PRISMATIC_JOINT_AXIS_UNIT = PRISMATIC_JOINT_AXIS / torch.linalg.vector_norm( + PRISMATIC_JOINT_AXIS +) +REVOLUTE_JOINT_AXIS_UNIT = REVOLUTE_JOINT_AXIS / torch.linalg.vector_norm( + REVOLUTE_JOINT_AXIS +) REVOLUTE_AXIS_ORIGIN = torch.tensor([0.75, -0.5, 0.25]) TARGET_POINT_OFFSETS = torch.tensor( [ @@ -52,6 +60,8 @@ ) TARGET_LINK_POINT_CLOUD_KEY = "target_link_point_cloud" ARTICULATION_POINT_CLOUD_KEY = "articulation_point_cloud" +TARGET_LINK_PRISMATIC_JOINT_AXIS_KEY = "target_link_prismatic_joint_axis" +TARGET_LINK_REVOLUTE_JOINT_AXIS_KEY = "target_link_revolute_joint_axis" TARGET_LINK_REVOLUTE_AXIS_ORIGIN_KEY = "target_link_revolute_axis_origin" @@ -70,10 +80,21 @@ def _axis_geometry(neighbor_offset: tuple[float, float, float]) -> dict[str, obj ), dim=0, ), + TARGET_LINK_PRISMATIC_JOINT_AXIS_KEY: PRISMATIC_JOINT_AXIS.clone(), + TARGET_LINK_REVOLUTE_JOINT_AXIS_KEY: REVOLUTE_JOINT_AXIS.clone(), TARGET_LINK_REVOLUTE_AXIS_ORIGIN_KEY: REVOLUTE_AXIS_ORIGIN.clone(), } +def _joint_axis_metadata(kind: str) -> tuple[str, torch.Tensor]: + """Return the geometry key and raw parent-joint axis for an affordance.""" + if kind in ("slide", "press"): + return TARGET_LINK_PRISMATIC_JOINT_AXIS_KEY, PRISMATIC_JOINT_AXIS + if kind == "twist": + return TARGET_LINK_REVOLUTE_JOINT_AXIS_KEY, REVOLUTE_JOINT_AXIS + raise ValueError(f"Unsupported affordance kind: {kind!r}.") + + def _axis_affordance( kind: str, *, @@ -242,32 +263,18 @@ class TestArticulationGeometryAxisInference: ("twist", "twist_axis"), ), ) - @pytest.mark.parametrize( - ("neighbor_offset", "expected_axis"), - ( - ((1.5, 0.25, -0.125), (1.0, 0.0, 0.0)), - ((-1.5, 0.25, -0.125), (-1.0, 0.0, 0.0)), - ((0.25, 1.5, -0.125), (0.0, 1.0, 0.0)), - ((0.25, -1.5, -0.125), (0.0, -1.0, 0.0)), - ((0.25, -0.125, 1.5), (0.0, 0.0, 1.0)), - ((0.25, -0.125, -1.5), (0.0, 0.0, -1.0)), - ), - ids=( - "positive_x", - "negative_x", - "positive_y", - "negative_y", - "positive_z", - "negative_z", - ), - ) - def test_object_geometry_ignores_outlier_and_quantizes_offset_to_signed_axis( + @pytest.mark.parametrize("direction", (1.0, -1.0), ids=("positive", "negative")) + def test_uses_parent_joint_axis_and_neighborhood_only_selects_sign( self, kind: str, axis_field: str, - neighbor_offset: tuple[float, float, float], - expected_axis: tuple[float, float, float], + direction: float, ) -> None: + _, raw_joint_axis = _joint_axis_metadata(kind) + normalized_joint_axis = raw_joint_axis / torch.linalg.vector_norm( + raw_joint_axis + ) + neighbor_offset = normalized_joint_axis * (1.5 * direction) affordance, actual_axis_field = _axis_affordance( kind, fallback_axis=torch.tensor([1.0, 1.0, 1.0]), @@ -275,39 +282,63 @@ def test_object_geometry_ignores_outlier_and_quantizes_offset_to_signed_axis( ObjectSemantics( affordance=affordance, - geometry=_axis_geometry(neighbor_offset), + geometry=_axis_geometry(tuple(float(value) for value in neighbor_offset)), entity_id=f"{kind}-target", ) assert actual_axis_field == axis_field actual_axis = getattr(affordance, axis_field) - assert torch.equal(actual_axis, torch.tensor(expected_axis)) - assert torch.count_nonzero(actual_axis).item() == 1 - assert actual_axis.abs().sum().item() == pytest.approx(1.0) + assert torch.allclose( + actual_axis, + normalized_joint_axis * direction, + atol=1.0e-6, + ) + assert torch.linalg.vector_norm(actual_axis).item() == pytest.approx(1.0) + assert torch.count_nonzero(actual_axis).item() == 3 @pytest.mark.parametrize("kind", ("slide", "press", "twist")) - def test_complete_geometry_overrides_explicit_fallback_axis( - self, kind: str - ) -> None: + def test_empty_geometry_preserves_explicit_fallback_axis(self, kind: str) -> None: + fallback_axis = torch.tensor([-1.0, 0.0, 0.0]) affordance, axis_field = _axis_affordance( kind, - fallback_axis=torch.tensor([1.0, 0.0, 0.0]), + fallback_axis=fallback_axis, ) ObjectSemantics( affordance=affordance, - geometry=_axis_geometry((0.25, -0.125, -1.5)), - entity_id=f"{kind}-target", + geometry={}, + entity_id=f"rigid-{kind}-target", ) - assert torch.equal( - getattr(affordance, axis_field), - torch.tensor([0.0, 0.0, -1.0]), - ) + assert torch.equal(getattr(affordance, axis_field), fallback_axis) - @pytest.mark.parametrize("kind", ("slide", "press", "twist")) - def test_empty_geometry_preserves_explicit_fallback_axis(self, kind: str) -> None: - fallback_axis = torch.tensor([-1.0, 0.0, 0.0]) + @pytest.mark.parametrize( + ("kind", "unrelated_axis_key", "unrelated_axis"), + ( + ( + "slide", + TARGET_LINK_REVOLUTE_JOINT_AXIS_KEY, + REVOLUTE_JOINT_AXIS, + ), + ( + "press", + TARGET_LINK_REVOLUTE_JOINT_AXIS_KEY, + REVOLUTE_JOINT_AXIS, + ), + ( + "twist", + TARGET_LINK_PRISMATIC_JOINT_AXIS_KEY, + PRISMATIC_JOINT_AXIS, + ), + ), + ) + def test_unrelated_joint_axis_metadata_preserves_fallback_axis( + self, + kind: str, + unrelated_axis_key: str, + unrelated_axis: torch.Tensor, + ) -> None: + fallback_axis = torch.tensor([-1.0, 0.5, 0.25]) affordance, axis_field = _axis_affordance( kind, fallback_axis=fallback_axis, @@ -315,40 +346,111 @@ def test_empty_geometry_preserves_explicit_fallback_axis(self, kind: str) -> Non ObjectSemantics( affordance=affordance, - geometry={}, - entity_id=f"rigid-{kind}-target", + geometry={unrelated_axis_key: unrelated_axis.clone()}, + entity_id=f"unrelated-{kind}-joint-axis-target", ) assert torch.equal(getattr(affordance, axis_field), fallback_axis) + def test_twist_origin_metadata_alone_preserves_axis_fallback(self) -> None: + fallback_axis = torch.tensor([-1.0, 0.5, 0.25]) + affordance = TwistAffordance( + grasp_position=(0.0, 0.0, 0.0), + axis_origin=(9.0, 8.0, 7.0), + twist_axis=fallback_axis, + ) + + ObjectSemantics( + affordance=affordance, + geometry={ + TARGET_LINK_REVOLUTE_AXIS_ORIGIN_KEY: REVOLUTE_AXIS_ORIGIN.clone() + }, + entity_id="twist-origin-only-target", + ) + + assert torch.equal(affordance.twist_axis, fallback_axis) + assert affordance.axis_origin == pytest.approx( + tuple(float(value) for value in REVOLUTE_AXIS_ORIGIN) + ) + @pytest.mark.parametrize( - "geometry", + "present_fields", ( - {TARGET_LINK_POINT_CLOUD_KEY: POINT_CLOUD_CENTER.unsqueeze(0)}, - {ARTICULATION_POINT_CLOUD_KEY: POINT_CLOUD_CENTER.unsqueeze(0)}, + ("target",), + ("articulation",), + ("axis",), + ("target", "articulation"), + ("target", "axis"), + ("articulation", "axis"), ), ) - def test_partial_point_cloud_geometry_is_rejected( + @pytest.mark.parametrize("kind", ("slide", "press", "twist")) + def test_incomplete_joint_axis_inference_metadata_is_rejected( self, - geometry: dict[str, object], + kind: str, + present_fields: tuple[str, ...], ) -> None: + axis_key, raw_joint_axis = _joint_axis_metadata(kind) + complete_geometry = _axis_geometry((1.0, 1.0, 0.5)) + geometry: dict[str, object] = {} + if "target" in present_fields: + geometry[TARGET_LINK_POINT_CLOUD_KEY] = complete_geometry[ + TARGET_LINK_POINT_CLOUD_KEY + ] + if "articulation" in present_fields: + geometry[ARTICULATION_POINT_CLOUD_KEY] = complete_geometry[ + ARTICULATION_POINT_CLOUD_KEY + ] + if "axis" in present_fields: + geometry[axis_key] = raw_joint_axis.clone() affordance, _ = _axis_affordance( - "slide", + kind, fallback_axis=torch.tensor([1.0, 0.0, 0.0]), ) - with pytest.raises(ValueError, match="requires both"): + with pytest.raises(ValueError, match=axis_key): ObjectSemantics( affordance=affordance, geometry=geometry, - entity_id="partial-geometry-target", + entity_id=f"incomplete-{kind}-geometry-target", + ) + + @pytest.mark.parametrize("kind", ("slide", "press", "twist")) + @pytest.mark.parametrize( + ("joint_axis", "exception_type"), + ( + ((1.0, 0.0, 0.0), TypeError), + (torch.tensor([1, 0, 0]), ValueError), + (torch.tensor([1.0, 0.0]), ValueError), + (torch.tensor([float("nan"), 0.0, 0.0]), ValueError), + (torch.zeros(3), ValueError), + ), + ) + def test_invalid_parent_joint_axis_metadata_is_rejected( + self, + kind: str, + joint_axis: object, + exception_type: type[Exception], + ) -> None: + axis_key, _ = _joint_axis_metadata(kind) + geometry = _axis_geometry((1.0, 1.0, 0.5)) + geometry[axis_key] = joint_axis + affordance, _ = _axis_affordance( + kind, + fallback_axis=torch.tensor([1.0, 0.0, 0.0]), + ) + + with pytest.raises(exception_type, match=axis_key): + ObjectSemantics( + affordance=affordance, + geometry=geometry, + entity_id=f"invalid-{kind}-joint-axis-target", ) def test_degenerate_target_point_cloud_is_rejected(self) -> None: - geometry = { - TARGET_LINK_POINT_CLOUD_KEY: POINT_CLOUD_CENTER.repeat(3, 1), - ARTICULATION_POINT_CLOUD_KEY: POINT_CLOUD_CENTER.unsqueeze(0), - } + geometry = _axis_geometry((0.5, -1.0, -1.0)) + geometry[TARGET_LINK_POINT_CLOUD_KEY] = POINT_CLOUD_CENTER.repeat(3, 1) + geometry[ARTICULATION_POINT_CLOUD_KEY] = POINT_CLOUD_CENTER.unsqueeze(0) affordance, _ = _axis_affordance( "twist", fallback_axis=torch.tensor([1.0, 0.0, 0.0]), @@ -361,22 +463,29 @@ def test_degenerate_target_point_cloud_is_rejected(self) -> None: entity_id="degenerate-geometry-target", ) - def test_centered_articulation_neighborhood_is_rejected(self) -> None: - target_points = POINT_CLOUD_CENTER + TARGET_POINT_OFFSETS - geometry = { - TARGET_LINK_POINT_CLOUD_KEY: target_points, - ARTICULATION_POINT_CLOUD_KEY: target_points.clone(), - } + @pytest.mark.parametrize( + ("kind", "perpendicular_offset"), + ( + ("slide", (1.0, -1.0, 0.0)), + ("press", (1.0, -1.0, 0.0)), + ("twist", (2.0, 1.0, 0.0)), + ), + ) + def test_neighborhood_offset_perpendicular_to_joint_axis_is_ambiguous( + self, + kind: str, + perpendicular_offset: tuple[float, float, float], + ) -> None: affordance, _ = _axis_affordance( - "press", + kind, fallback_axis=torch.tensor([1.0, 0.0, 0.0]), ) - with pytest.raises(ValueError, match="centered on the target link"): + with pytest.raises(ValueError, match="direction is ambiguous"): ObjectSemantics( affordance=affordance, - geometry=geometry, - entity_id="centered-geometry-target", + geometry=_axis_geometry(perpendicular_offset), + entity_id=f"ambiguous-{kind}-geometry-target", ) def test_press_without_position_uses_outer_surface_opposite_inferred_axis( @@ -389,15 +498,18 @@ def test_press_without_position_uses_outer_surface_opposite_inferred_axis( ObjectSemantics( affordance=affordance, - geometry=_axis_geometry((0.25, -0.125, -1.5)), + geometry=_axis_geometry( + tuple(float(value) for value in -1.5 * PRISMATIC_JOINT_AXIS_UNIT) + ), entity_id="automatic-press-target", ) - assert torch.equal( + assert torch.allclose( affordance.press_axis, - torch.tensor([0.0, 0.0, -1.0]), + -PRISMATIC_JOINT_AXIS_UNIT, + atol=1.0e-6, ) - expected_surface_center = POINT_CLOUD_CENTER + torch.tensor([0.0, 0.0, 1.0]) + expected_surface_center = POINT_CLOUD_CENTER + torch.tensor([0.5, 0.5, 0.0]) assert affordance.press_position == pytest.approx( tuple(float(value) for value in expected_surface_center) ) @@ -411,13 +523,16 @@ def test_twist_geometry_uses_revolute_joint_axis_origin(self) -> None: ObjectSemantics( affordance=affordance, - geometry=_axis_geometry((0.25, -0.125, -1.5)), + geometry=_axis_geometry( + tuple(float(value) for value in -1.5 * REVOLUTE_JOINT_AXIS_UNIT) + ), entity_id="automatic-twist-target", ) - assert torch.equal( + assert torch.allclose( affordance.twist_axis, - torch.tensor([0.0, 0.0, -1.0]), + -REVOLUTE_JOINT_AXIS_UNIT, + atol=1.0e-6, ) assert affordance.axis_origin == pytest.approx( tuple(float(value) for value in REVOLUTE_AXIS_ORIGIN) @@ -430,7 +545,9 @@ def test_twist_geometry_uses_revolute_joint_axis_origin(self) -> None: def test_twist_complete_clouds_without_joint_origin_preserve_fallback( self, ) -> None: - geometry = _axis_geometry((0.25, -0.125, -1.5)) + geometry = _axis_geometry( + tuple(float(value) for value in -1.5 * REVOLUTE_JOINT_AXIS_UNIT) + ) geometry.pop(TARGET_LINK_REVOLUTE_AXIS_ORIGIN_KEY) affordance = TwistAffordance( grasp_position=(0.0, 0.0, 0.0), @@ -444,16 +561,19 @@ def test_twist_complete_clouds_without_joint_origin_preserve_fallback( entity_id="automatic-twist-target-with-origin-fallback", ) - assert torch.equal( + assert torch.allclose( affordance.twist_axis, - torch.tensor([0.0, 0.0, -1.0]), + -REVOLUTE_JOINT_AXIS_UNIT, + atol=1.0e-6, ) assert affordance.axis_origin == pytest.approx((9.0, 8.0, 7.0)) def test_twist_complete_clouds_without_joint_origin_do_not_use_centroid( self, ) -> None: - geometry = _axis_geometry((0.25, -0.125, -1.5)) + geometry = _axis_geometry( + tuple(float(value) for value in -1.5 * REVOLUTE_JOINT_AXIS_UNIT) + ) geometry.pop(TARGET_LINK_REVOLUTE_AXIS_ORIGIN_KEY) affordance = TwistAffordance( grasp_position=(0.0, 0.0, 0.0), @@ -466,9 +586,10 @@ def test_twist_complete_clouds_without_joint_origin_do_not_use_centroid( entity_id="automatic-twist-target-without-joint-origin", ) - assert torch.equal( + assert torch.allclose( affordance.twist_axis, - torch.tensor([0.0, 0.0, -1.0]), + -REVOLUTE_JOINT_AXIS_UNIT, + atol=1.0e-6, ) assert affordance.axis_origin is None with pytest.raises( diff --git a/tests/sim/atomic_actions/test_tutorial_utils.py b/tests/sim/atomic_actions/test_tutorial_utils.py index 2cc806306..ffe2c8a68 100644 --- a/tests/sim/atomic_actions/test_tutorial_utils.py +++ b/tests/sim/atomic_actions/test_tutorial_utils.py @@ -130,6 +130,8 @@ "move_end_effector", "move_joints", ) +TUTORIAL_PRISMATIC_JOINT_AXIS = torch.tensor([2.0, 2.0, 1.0]) +TUTORIAL_REVOLUTE_JOINT_AXIS = torch.tensor([-1.0, 2.0, 2.0]) TUTORIAL_REVOLUTE_AXIS_ORIGIN = torch.tensor([0.75, -0.5, 0.25]) @@ -156,6 +158,8 @@ def _tutorial_axis_geometry( (center + torch.tensor(neighbor_offset)).unsqueeze(0), ) ), + "target_link_prismatic_joint_axis": TUTORIAL_PRISMATIC_JOINT_AXIS.clone(), + "target_link_revolute_joint_axis": TUTORIAL_REVOLUTE_JOINT_AXIS.clone(), "target_link_revolute_axis_origin": TUTORIAL_REVOLUTE_AXIS_ORIGIN.clone(), } @@ -237,28 +241,28 @@ def _run_obstacle_animation(*, pace_wall_time: bool) -> tuple[MagicMock, MagicMo "create_drawer_semantics", "large_handle_bar", "translation_axis", - (0.25, 1.5, -0.125), - (0.0, 1.0, 0.0), + (1.0, 1.0, 0.5), + (2.0 / 3.0, 2.0 / 3.0, 1.0 / 3.0), ), ( "press", "create_button_semantics", "button_cap", "press_axis", - (0.25, -0.125, -1.5), - (0.0, 0.0, -1.0), + (-1.0, -1.0, -0.5), + (-2.0 / 3.0, -2.0 / 3.0, -1.0 / 3.0), ), ( "twist", "create_knob_semantics", "cap_1", "twist_axis", - (0.25, -0.125, -1.5), - (0.0, 0.0, -1.0), + (0.5, -1.0, -1.0), + (1.0 / 3.0, -2.0 / 3.0, -2.0 / 3.0), ), ), ) -def test_articulation_tutorial_semantics_resolve_axis_from_initial_point_clouds( +def test_articulation_tutorial_semantics_resolve_signed_parent_joint_axis( module_name: str, factory_name: str, link_name: str, @@ -276,12 +280,13 @@ def test_articulation_tutorial_semantics_resolve_axis_from_initial_point_clouds( semantics = result[0] if isinstance(result, tuple) else result assert articulation.sampled_link_name == link_name assert semantics.geometry is geometry - assert torch.equal( + assert torch.allclose( getattr(semantics.affordance, axis_field), torch.tensor(expected_axis), + atol=1.0e-6, ) if module_name == "press": - assert semantics.affordance.press_position == pytest.approx((2.0, -3.0, 5.0)) + assert semantics.affordance.press_position == pytest.approx((2.5, -2.5, 4.0)) elif module_name == "twist": assert semantics.affordance.axis_origin == pytest.approx( tuple(float(value) for value in TUTORIAL_REVOLUTE_AXIS_ORIGIN) diff --git a/tests/sim/objects/test_articulation.py b/tests/sim/objects/test_articulation.py index a6dad2958..bc018bc6b 100644 --- a/tests/sim/objects/test_articulation.py +++ b/tests/sim/objects/test_articulation.py @@ -243,6 +243,146 @@ def test_uses_nearest_revolute_ancestor_after_fixed_descendants(self): geometry["target_link_revolute_axis_origin"], near_origin_pose[:3, 3], ) + assert torch.equal( + geometry["target_link_revolute_joint_axis"], + near_joint.axis, + ) + + def test_uses_nearest_prismatic_axis_after_fixed_descendant(self): + triangle = torch.tensor(((0, 1, 2),), dtype=torch.long) + link_meshes = { + link_name: ( + torch.tensor( + ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)), + dtype=torch.float32, + ), + triangle, + ) + for link_name in ("base", "near_parent", "moving", "target") + } + fixed_joint = ArticulationJointKinematics( + name="target_fixed", + joint_type="fixed", + parent_link_name="moving", + child_link_name="target", + origin_pose=torch.eye(4), + axis=torch.zeros(3), + ) + near_origin_pose = torch.eye(4) + near_origin_pose[:3, :3] = torch.tensor( + ((0.0, -1.0, 0.0), (1.0, 0.0, 0.0), (0.0, 0.0, 1.0)) + ) + near_joint = ArticulationJointKinematics( + name="near_slider", + joint_type="prismatic", + parent_link_name="near_parent", + child_link_name="moving", + origin_pose=near_origin_pose, + axis=torch.tensor((4.0, 0.0, 0.0)), + ) + far_joint = ArticulationJointKinematics( + name="far_slider", + joint_type="prismatic", + parent_link_name="base", + child_link_name="near_parent", + origin_pose=torch.eye(4), + axis=torch.tensor((0.0, 0.0, 2.0)), + ) + articulation, _ = _make_point_cloud_articulation( + link_meshes=link_meshes, + parent_joint_chain=(fixed_joint, near_joint, far_joint), + ) + + geometry = articulation.sample_initial_point_clouds( + "target", + articulation_point_count=16, + target_point_count=8, + ) + + assert torch.allclose( + geometry["target_link_prismatic_joint_axis"], + torch.tensor((0.0, 1.0, 0.0)), + atol=POINT_CLOUD_TOLERANCE, + ) + assert "target_link_revolute_joint_axis" not in geometry + assert "target_link_revolute_axis_origin" not in geometry + + def test_collects_both_nearest_joint_types_from_one_parent_chain(self): + triangle = torch.tensor(((0, 1, 2),), dtype=torch.long) + link_meshes = { + link_name: ( + torch.tensor( + ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)), + dtype=torch.float32, + ), + triangle, + ) + for link_name in ("root", "base", "near_parent", "hinge_child", "target") + } + near_prismatic = ArticulationJointKinematics( + name="near_slider", + joint_type="prismatic", + parent_link_name="hinge_child", + child_link_name="target", + origin_pose=torch.eye(4), + axis=torch.tensor((4.0, 0.0, 0.0)), + ) + near_revolute_origin = torch.eye(4) + near_revolute_origin[:3, 3] = torch.tensor((1.0, 2.0, 3.0)) + near_revolute = ArticulationJointKinematics( + name="near_hinge", + joint_type="revolute", + parent_link_name="near_parent", + child_link_name="hinge_child", + origin_pose=near_revolute_origin, + axis=torch.tensor((0.0, 5.0, 0.0)), + ) + far_prismatic = ArticulationJointKinematics( + name="far_slider", + joint_type="prismatic", + parent_link_name="base", + child_link_name="near_parent", + origin_pose=torch.eye(4), + axis=torch.tensor((0.0, 0.0, 2.0)), + ) + far_revolute_origin = torch.eye(4) + far_revolute_origin[:3, 3] = torch.tensor((9.0, 9.0, 9.0)) + far_revolute = ArticulationJointKinematics( + name="far_hinge", + joint_type="revolute", + parent_link_name="root", + child_link_name="base", + origin_pose=far_revolute_origin, + axis=torch.tensor((0.0, 0.0, 7.0)), + ) + articulation, _ = _make_point_cloud_articulation( + link_meshes=link_meshes, + parent_joint_chain=( + near_prismatic, + near_revolute, + far_prismatic, + far_revolute, + ), + ) + + geometry = articulation.sample_initial_point_clouds( + "target", + articulation_point_count=16, + target_point_count=8, + ) + + assert torch.equal( + geometry["target_link_prismatic_joint_axis"], + torch.tensor((1.0, 0.0, 0.0)), + ) + assert torch.equal( + geometry["target_link_revolute_joint_axis"], + torch.tensor((0.0, 1.0, 0.0)), + ) + assert torch.equal( + geometry["target_link_revolute_axis_origin"], + near_revolute_origin[:3, 3], + ) def test_transforms_revolute_origin_from_rotated_parent_to_target_frame(self): triangle = torch.tensor(((0, 1, 2),), dtype=torch.long) @@ -278,7 +418,7 @@ def test_transforms_revolute_origin_from_rotated_parent_to_target_frame(self): parent_link_name="parent", child_link_name="target", origin_pose=joint_origin_pose, - axis=torch.tensor((0.0, 1.0, 0.0)), + axis=torch.tensor((0.0, 4.0, 0.0)), ) articulation, _ = _make_point_cloud_articulation( link_meshes=link_meshes, @@ -300,6 +440,16 @@ def test_transforms_revolute_origin_from_rotated_parent_to_target_frame(self): target_from_joint[:3, 3], atol=POINT_CLOUD_TOLERANCE, ) + expected_axis = target_from_joint[:3, :3] @ joint.axis + expected_axis = expected_axis / torch.linalg.vector_norm(expected_axis) + assert torch.allclose( + geometry["target_link_revolute_joint_axis"], + expected_axis, + atol=POINT_CLOUD_TOLERANCE, + ) + assert torch.linalg.vector_norm( + geometry["target_link_revolute_joint_axis"] + ).item() == pytest.approx(1.0) def test_omits_revolute_origin_when_parent_chain_has_only_fixed_joints(self): fixed_joint = ArticulationJointKinematics( @@ -320,12 +470,19 @@ def test_omits_revolute_origin_when_parent_chain_has_only_fixed_joints(self): target_point_count=3, ) - assert "target_link_revolute_axis_origin" not in geometry + assert set(geometry) == { + "target_link_point_cloud", + "articulation_point_cloud", + } - def test_rejects_revolute_ancestor_with_unknown_parent_link(self): + @pytest.mark.parametrize("joint_type", ("prismatic", "revolute")) + def test_rejects_joint_axis_metadata_with_unknown_parent_link( + self, + joint_type: str, + ): joint = ArticulationJointKinematics( - name="target_hinge", - joint_type="revolute", + name="target_joint", + joint_type=joint_type, parent_link_name="missing_parent", child_link_name="target", origin_pose=torch.eye(4), @@ -342,12 +499,17 @@ def test_rejects_revolute_ancestor_with_unknown_parent_link(self): target_point_count=3, ) - def test_rejects_revolute_ancestor_with_invalid_origin_pose(self): + @pytest.mark.parametrize("joint_type", ("prismatic", "revolute")) + def test_rejects_joint_axis_metadata_with_invalid_origin_pose( + self, + joint_type: str, + ): joint = SimpleNamespace( - name="target_hinge", - joint_type="revolute", + name="target_joint", + joint_type=joint_type, parent_link_name="target", origin_pose=torch.full((4, 4), float("nan")), + axis=torch.tensor((0.0, 0.0, 1.0)), ) articulation, _ = _make_point_cloud_articulation( parent_joint_chain=(joint,), @@ -360,6 +522,37 @@ def test_rejects_revolute_ancestor_with_invalid_origin_pose(self): target_point_count=3, ) + @pytest.mark.parametrize( + ("joint_type", "axis"), + ( + ("prismatic", torch.zeros(3)), + ("revolute", torch.tensor((float("nan"), 0.0, 1.0))), + ("prismatic", torch.zeros(2)), + ), + ) + def test_rejects_invalid_joint_axes( + self, + joint_type: str, + axis: torch.Tensor, + ): + joint = SimpleNamespace( + name="target_joint", + joint_type=joint_type, + parent_link_name="target", + origin_pose=torch.eye(4), + axis=axis, + ) + articulation, _ = _make_point_cloud_articulation( + parent_joint_chain=(joint,), + ) + + with pytest.raises(ValueError, match="axis must be finite and nonzero"): + articulation.sample_initial_point_clouds( + "target", + articulation_point_count=5, + target_point_count=3, + ) + def test_returns_both_clouds_in_target_link_initial_frame(self): triangle = torch.tensor(((0, 1, 2),), dtype=torch.long) link_meshes = { From 02eb80ec8c3da50c8ae40bd8f1cf2e8978ab2b13 Mon Sep 17 00:00:00 2001 From: matafela Date: Wed, 2 Sep 2026 14:35:14 +0800 Subject: [PATCH 6/8] update --- agent_context/MAP.yaml | 7 + .../topics/atomic-actions/atomic-actions.md | 15 + .../simulation-system/simulation-system.md | 8 + .../embodichain.lab.sim.atomic_actions.rst | 32 + .../sim/atomic_actions/builtin_actions.md | 12 +- docs/source/overview/sim/sim_articulation.md | 54 +- .../lab/sim/atomic_actions/__init__.py | 10 + .../_articulation_geometry_keys.py | 27 + .../lab/sim/atomic_actions/affordance.py | 12 +- .../atomic_actions/articulation_geometry.py | 713 ++++++++++++++ embodichain/lab/sim/objects/articulation.py | 412 +------- embodichain/lab/sim/objects/robot.py | 10 + .../integrations/_configured_services.py | 16 +- scripts/tutorials/atomic_action/press.py | 9 +- scripts/tutorials/atomic_action/slide.py | 9 +- scripts/tutorials/atomic_action/twist.py | 9 +- .../task_program/test_task_vertical_slices.py | 74 +- .../test_articulation_geometry.py | 899 ++++++++++++++++++ .../sim/atomic_actions/test_tutorial_utils.py | 54 +- tests/sim/objects/test_articulation.py | 777 +-------------- tests/sim/objects/test_robot.py | 27 +- 21 files changed, 1996 insertions(+), 1190 deletions(-) create mode 100644 embodichain/lab/sim/atomic_actions/_articulation_geometry_keys.py create mode 100644 embodichain/lab/sim/atomic_actions/articulation_geometry.py create mode 100644 tests/sim/atomic_actions/test_articulation_geometry.py diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index 8fae060e0..3e9e8e98f 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -626,6 +626,11 @@ topics: - RigidObjectSceneProvider - create_simulation_atomic_action_engine - ObjectSemantics + - ArticulationAffordanceGeometry + - ArticulationGeometryProvider + - ArticulationJointGeometry + - sample_initial_articulation_geometry + - articulation geometry adapter - entity_id - frozen ObjectSemantics - canonical entity identity @@ -736,6 +741,8 @@ topics: - embodichain/lab/sim/atomic_actions/goals.py - embodichain/lab/sim/atomic_actions/effects.py - embodichain/lab/sim/atomic_actions/affordance.py + - embodichain/lab/sim/atomic_actions/_articulation_geometry_keys.py + - embodichain/lab/sim/atomic_actions/articulation_geometry.py - embodichain/lab/sim/atomic_actions/bindings.py - embodichain/lab/sim/atomic_actions/control.py - embodichain/lab/sim/atomic_actions/invocation.py diff --git a/agent_context/topics/atomic-actions/atomic-actions.md b/agent_context/topics/atomic-actions/atomic-actions.md index df027c9cd..f7458c813 100644 --- a/agent_context/topics/atomic-actions/atomic-actions.md +++ b/agent_context/topics/atomic-actions/atomic-actions.md @@ -53,6 +53,7 @@ a physical effect. | Concern | Source of truth | |---|---| | Core action and goal contracts | `atomic_actions/core.py`, `goals.py`, `affordance.py` | +| Articulation affordance geometry adapter | `atomic_actions/articulation_geometry.py` | | Invocation, binding, and policies | `invocation.py`, `bindings.py`, `policies.py`, `control.py` | | Robot/task/scene state | `state.py`, `scene.py` | | Plans and runtime commands | `plans.py`, `runtime_commands.py` | @@ -65,6 +66,19 @@ a physical effect. | Task Program lowering/execution | `embodichain/lab/task_program/compiler/`, `runtime/` | | Gym lifecycle bridge | `embodichain/lab/gym/envs/task_program/bridge.py` | +## Articulation geometry boundary + +`sample_initial_articulation_geometry()` adapts domain-neutral articulation +facts into the geometry consumed by link affordances. Callers pass an +`ArticulationGeometryProvider`, the explicitly named initial joint state, and +body scale. The adapter owns Open3D surface sampling, target-link-frame mesh +transforms, nearest prismatic/revolute ancestor geometry, and the private +`ObjectSemantics.geometry` keys. + +The adapter returns `ArticulationAffordanceGeometry`; convert it with +`to_object_geometry()` only when constructing `ObjectSemantics`. Keep random +sampling and semantic geometry keys out of `objects/articulation.py`. + ## Direct resolution path ### 1. Construct the engine @@ -221,6 +235,7 @@ an action implementation. | Change | Owning location | |---|---| | Goal/options or action behavior | matching file under `atomic_actions/primitives/` plus `goals.py` when shared | +| Articulation-link sampling or affordance geometry metadata | `atomic_actions/articulation_geometry.py`; keep `objects/articulation.py` limited to meshes, FK, and topology | | Engine catalog or lifecycle | `engine.py` | | Session transitions or recovery | `execution.py` | | Verification value contracts | `verification.py` | diff --git a/agent_context/topics/simulation-system/simulation-system.md b/agent_context/topics/simulation-system/simulation-system.md index 606eb1e9c..24a403db3 100644 --- a/agent_context/topics/simulation-system/simulation-system.md +++ b/agent_context/topics/simulation-system/simulation-system.md @@ -96,6 +96,12 @@ for integrations that need link ancestry. It returns immediate-parent-first origin, axis, and optional limits. Consumers must not reach into `BatchEntity._entities` or retain backend-native joint-info objects. +`Articulation` also exposes deterministic link meshes through +`get_link_vert_face()` and named-state FK through `compute_fk()` with +`qpos_joint_names`. Stochastic surface sampling and Atomic Action geometry keys +do not belong to the simulation object; use +`atomic_actions.sample_initial_articulation_geometry()` for that adaptation. + ## Configuration Flow `SimulationManagerCfg` owns window size, headless mode, rendering, GPU/CPU @@ -141,6 +147,8 @@ corresponding robot/sensor module. Scene composition belongs in `SimulationManager.flush_cleanup_queue()`. - Resolve articulation ancestry through `get_parent_joint_chain()`; keep DexSim topology access encapsulated by `Articulation`. +- Keep articulation mesh access, FK, and topology domain-neutral. Perform + affordance sampling and semantic-key conversion in the Atomic Action adapter. ## Common Failure Modes diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.atomic_actions.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.atomic_actions.rst index 1bc85a991..147ebae2d 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.atomic_actions.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.atomic_actions.rst @@ -41,6 +41,15 @@ embodichain.lab.sim.atomic_actions ActionPlan CompiledTrajectory + .. rubric:: Articulation geometry adaptation + + .. autosummary:: + + ArticulationAffordanceGeometry + ArticulationGeometryProvider + ArticulationJointGeometry + sample_initial_articulation_geometry + .. rubric:: Semantic resource contracts .. autosummary:: @@ -145,6 +154,29 @@ embodichain.lab.sim.atomic_actions .. currentmodule:: embodichain.lab.sim.atomic_actions +Articulation geometry adaptation +-------------------------------- + +The adapter converts deterministic articulation meshes, FK, and parent-joint +topology into sampled Atomic Action affordance geometry. Initial state and scale +are explicit inputs; the simulation object does not own sampling or semantic +geometry keys. + +.. currentmodule:: embodichain.lab.sim.atomic_actions.articulation_geometry + +.. autoclass:: ArticulationGeometryProvider + :members: + +.. autoclass:: ArticulationJointGeometry + :members: + +.. autoclass:: ArticulationAffordanceGeometry + :members: + +.. autofunction:: sample_initial_articulation_geometry + +.. currentmodule:: embodichain.lab.sim.atomic_actions + Semantic resource contracts --------------------------- diff --git a/docs/source/overview/sim/atomic_actions/builtin_actions.md b/docs/source/overview/sim/atomic_actions/builtin_actions.md index b351e405a..fff4ce566 100644 --- a/docs/source/overview/sim/atomic_actions/builtin_actions.md +++ b/docs/source/overview/sim/atomic_actions/builtin_actions.md @@ -678,9 +678,10 @@ migration. Plans **close hand -> approach target -> contact -> press along axis -> return to the approach pose**. For an articulation link, -`Articulation.sample_initial_point_clouds()` stores target-link-local target and -whole-articulation clouds in `ObjectSemantics.geometry`, together with the -nearest parent prismatic joint axis transformed into the same frame. +`sample_initial_articulation_geometry()` builds typed target-link-local target +and whole-articulation geometry. Its `to_object_geometry()` conversion stores +the two clouds and the nearest parent prismatic joint axis in +`ObjectSemantics.geometry`. `PressAffordance` uses `target_link_prismatic_joint_axis` for its axis direction; the point-cloud neighborhood selects only the sign. It can also derive the outer-surface `press_position` when it is omitted. `PressGoal.target_pose` is @@ -725,8 +726,9 @@ push/close; pull/open uses its negative direction. The affordance inherits grasp approach direction is the resolved link-frame translation axis transformed by the current link rotation. -Axis inference samples the target link and the merged articulation surface at -`ArticulationCfg.init_qpos`, expressed in the target link's initial local frame. +The Atomic Action articulation-geometry adapter samples the target link and the +merged articulation surface at `ArticulationCfg.init_qpos`, expressed in the +target link's initial local frame. Both clouds use Open3D uniform surface sampling. Sampling the whole articulation as one merged mesh preserves triangle-area weighting instead of giving every link an equal point budget, which would over-represent tiny decorative links. diff --git a/docs/source/overview/sim/sim_articulation.md b/docs/source/overview/sim/sim_articulation.md index e0888bec0..455946172 100644 --- a/docs/source/overview/sim/sim_articulation.md +++ b/docs/source/overview/sim/sim_articulation.md @@ -181,7 +181,6 @@ State data is accessed via getter methods that return batched tensors (`N` envir | :--- | :--- | :--- | | `get_local_pose(to_matrix=False)` | `(N, 7)` or `(N, 4, 4)` | Root link pose `[x, y, z, qw, qx, qy, qz]` or a 4x4 matrix. | | `get_link_pose(link_name, to_matrix=False)` | `(N, 7)` or `(N, 4, 4)` | Specific link pose `[x, y, z, qw, qx, qy, qz]` or a 4x4 matrix. | -| `sample_initial_point_clouds(target_link_name)` | `Dict[str, Tensor]` | Uniformly sample target and merged articulation surfaces at `init_qpos`, and include nearest parent prismatic/revolute axes and the revolute origin in the target link's initial local frame when available. | | `get_qpos(target=False)` | `(N, dof)` | Current joint positions (or joint targets if `target=True`). | | `get_qvel(target=False)` | `(N, dof)` | Current joint velocities (or velocity targets if `target=True`). | | `get_joint_drive()` | `Tuple[Tensor, ...]` | Returns `(stiffness, damping, max_effort, max_velocity, friction, armature)`, each shaped `(N, dof)`. | @@ -194,51 +193,18 @@ print(f"Current Joint Positions: {articulation.get_qpos()}") print(f"End Effector Pose: {articulation.get_link_pose('ee_link')}") ``` -### Initial Link-Local Point Clouds +### Atomic Action Geometry Adapter -`sample_initial_point_clouds()` uses the configured initial joint positions and -forward kinematics, rather than the articulation's mutable runtime state. It -transforms every link mesh into the requested target link's initial frame, -merges the meshes, and uses Open3D uniform surface sampling on the combined -triangle mesh. Sampling the merged mesh makes each surface's representation -proportional to triangle area instead of assigning the same point count to -every link. The returned float32 tensors are moved back to the articulation -device and are ready to store in an atomic action's `ObjectSemantics.geometry`: +`Articulation` exposes deterministic, domain-neutral facts through APIs such as +`get_link_vert_face()`, `compute_fk()`, and `get_parent_joint_chain()`. Atomic +Action integrations compose those facts with initial-state mesh transforms, +Open3D surface sampling, and affordance-specific geometry metadata through +`sample_initial_articulation_geometry()` in +`embodichain.lab.sim.atomic_actions`. The adapter returns a typed value; call +`to_object_geometry()` only at the `ObjectSemantics.geometry` boundary. -```python -geometry = articulation.sample_initial_point_clouds( - "button_cap", - articulation_point_count=100_000, - target_point_count=5_000, -) -target_points = geometry["target_link_point_cloud"] -articulation_points = geometry["articulation_point_cloud"] -prismatic_axis = geometry.get("target_link_prismatic_joint_axis") -revolute_axis = geometry.get("target_link_revolute_joint_axis") -revolute_origin = geometry.get("target_link_revolute_axis_origin") -``` - -The shown point counts are the defaults. Open3D draws random uniform samples -from the target mesh and merged articulation mesh independently, so repeated -calls are not expected to be bitwise identical and consumers should rely on -the spatial distribution rather than point-for-point correspondence. The -method requires a built kinematic chain and currently supports unit -`body_scale`. - -The sampler walks the immediate-parent-first joint chain and records the nearest -ancestor of each supported type. `target_link_prismatic_joint_axis` and -`target_link_revolute_joint_axis` are normalized `(3,)` axes transformed by the -initial parent-link and joint-origin rotations into the target link's initial -local frame. `target_link_revolute_axis_origin` is the matching nearest -revolute-joint origin in that frame. A key is omitted when the corresponding -joint type has no ancestor; this does not prevent point-cloud sampling. - -`Slide` and `Press` consume the prismatic axis, while `Twist` consumes the -revolute axis and origin. The point-cloud neighborhood only chooses between the -stored axis and its negation by projecting the neighborhood-center offset onto -the axis. It never replaces an oblique physical axis with a Cartesian basis -direction. The target-cloud centroid centers neighborhood/contact calculations -and is not a substitute for the revolute joint origin. +Keeping this stochastic conversion in the Atomic Action layer means simulation +objects do not own private point-cloud keys or affordance interpretation. ### Visual Appearance diff --git a/embodichain/lab/sim/atomic_actions/__init__.py b/embodichain/lab/sim/atomic_actions/__init__.py index 5ce887e8a..a00cbbe41 100644 --- a/embodichain/lab/sim/atomic_actions/__init__.py +++ b/embodichain/lab/sim/atomic_actions/__init__.py @@ -28,6 +28,12 @@ from __future__ import annotations +from .articulation_geometry import ( + ArticulationAffordanceGeometry, + ArticulationGeometryProvider, + ArticulationJointGeometry, + sample_initial_articulation_geometry, +) from .affordance import ( Affordance, AntipodalAffordance, @@ -253,6 +259,9 @@ "ActionPlanningServices", "Affordance", "AntipodalAffordance", + "ArticulationAffordanceGeometry", + "ArticulationGeometryProvider", + "ArticulationJointGeometry", "ArticulationJointState", "AssembleAffordance", "AssembleGoal", @@ -401,6 +410,7 @@ "SkillResourceSlot", "StateDelta", "SimulationExecutionAdapter", + "sample_initial_articulation_geometry", "TaskState", "TimedCommandSequence", "TimedTerminalAcceptance", diff --git a/embodichain/lab/sim/atomic_actions/_articulation_geometry_keys.py b/embodichain/lab/sim/atomic_actions/_articulation_geometry_keys.py new file mode 100644 index 000000000..a7992db11 --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/_articulation_geometry_keys.py @@ -0,0 +1,27 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Internal keys for articulation affordance geometry.""" + +from __future__ import annotations + +_TARGET_LINK_POINT_CLOUD_KEY = "target_link_point_cloud" +_ARTICULATION_POINT_CLOUD_KEY = "articulation_point_cloud" +_TARGET_LINK_PRISMATIC_JOINT_AXIS_KEY = "target_link_prismatic_joint_axis" +_TARGET_LINK_REVOLUTE_JOINT_AXIS_KEY = "target_link_revolute_joint_axis" +_TARGET_LINK_REVOLUTE_AXIS_ORIGIN_KEY = "target_link_revolute_axis_origin" + +__all__: list[str] = [] diff --git a/embodichain/lab/sim/atomic_actions/affordance.py b/embodichain/lab/sim/atomic_actions/affordance.py index 0c3e52938..739a09279 100644 --- a/embodichain/lab/sim/atomic_actions/affordance.py +++ b/embodichain/lab/sim/atomic_actions/affordance.py @@ -22,11 +22,13 @@ import torch -_TARGET_LINK_POINT_CLOUD_KEY = "target_link_point_cloud" -_ARTICULATION_POINT_CLOUD_KEY = "articulation_point_cloud" -_TARGET_LINK_PRISMATIC_JOINT_AXIS_KEY = "target_link_prismatic_joint_axis" -_TARGET_LINK_REVOLUTE_JOINT_AXIS_KEY = "target_link_revolute_joint_axis" -_TARGET_LINK_REVOLUTE_AXIS_ORIGIN_KEY = "target_link_revolute_axis_origin" +from ._articulation_geometry_keys import ( + _ARTICULATION_POINT_CLOUD_KEY, + _TARGET_LINK_POINT_CLOUD_KEY, + _TARGET_LINK_PRISMATIC_JOINT_AXIS_KEY, + _TARGET_LINK_REVOLUTE_AXIS_ORIGIN_KEY, + _TARGET_LINK_REVOLUTE_JOINT_AXIS_KEY, +) @dataclass diff --git a/embodichain/lab/sim/atomic_actions/articulation_geometry.py b/embodichain/lab/sim/atomic_actions/articulation_geometry.py new file mode 100644 index 000000000..7023674aa --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/articulation_geometry.py @@ -0,0 +1,713 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Adapt deterministic articulation facts into sampled affordance geometry.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Protocol + +import torch + +from ._articulation_geometry_keys import ( + _ARTICULATION_POINT_CLOUD_KEY, + _TARGET_LINK_POINT_CLOUD_KEY, + _TARGET_LINK_PRISMATIC_JOINT_AXIS_KEY, + _TARGET_LINK_REVOLUTE_AXIS_ORIGIN_KEY, + _TARGET_LINK_REVOLUTE_JOINT_AXIS_KEY, +) + + +class ArticulationJointGeometry(Protocol): + """Structural joint geometry consumed by the articulation adapter. + + Attributes: + name: Stable joint name. + joint_type: Normalized joint type such as ``fixed``, ``prismatic``, or + ``revolute``. + parent_link_name: Stable parent-link name. + child_link_name: Stable child-link name. + origin_pose: Joint-frame pose in the parent-link frame, shape ``(4, 4)``. + axis: Joint axis in the joint frame, shape ``(3,)``. + """ + + name: str + joint_type: str + parent_link_name: str + child_link_name: str + origin_pose: torch.Tensor + axis: torch.Tensor + + +class ArticulationGeometryProvider(Protocol): + """Deterministic articulation facts required for geometry adaptation. + + Implementations provide raw link meshes, FK, and immediate-parent-first + joint topology. Initial configuration and scale are explicit adapter inputs + so this protocol does not depend on an ``ArticulationCfg`` or PK-chain API. + + Attributes: + device: Device on which geometry tensors are assembled. + link_names: Stable articulation link names. + """ + + device: torch.device + link_names: Sequence[str] + + def get_link_vert_face( + self, + link_name: str, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return one link-local triangle mesh. + + Args: + link_name: Stable link name. + + Returns: + Link-local vertices and triangle indices. + """ + ... + + def compute_fk( + self, + qpos: torch.Tensor, + *, + link_names: Sequence[str], + qpos_joint_names: Sequence[str], + ) -> torch.Tensor: + """Return root-frame link poses for the supplied named joint state. + + Args: + qpos: Joint positions with shape ``(B, J)``. + link_names: Links whose poses should be returned. + qpos_joint_names: Names corresponding to the last ``qpos`` axis. + + Returns: + Root-frame poses with shape ``(B, L, 4, 4)``. + """ + ... + + def get_parent_joint_chain( + self, + link_name: str, + ) -> tuple[ArticulationJointGeometry, ...]: + """Return parent joints ordered from the link toward the root. + + Args: + link_name: Link at which to begin traversal. + + Returns: + Immediate-parent-first structural joint geometry. + """ + ... + + +@dataclass(frozen=True, slots=True, eq=False) +class ArticulationAffordanceGeometry: + """Owned sampled geometry for articulation-link affordances. + + The point clouds and optional joint data are expressed in the target link's + initial local frame. Joint axes are normalized. Use + :meth:`to_object_geometry` at the ``ObjectSemantics`` boundary so the + Atomic Action-specific string-key protocol remains in this module. + + Args: + target_link_point_cloud: Sampled target-link surface, shape ``(N, 3)``. + articulation_point_cloud: Sampled merged-articulation surface, shape + ``(M, 3)``. + prismatic_joint_axis: Optional nearest prismatic ancestor axis. + revolute_joint_axis: Optional nearest revolute ancestor axis. + revolute_axis_origin: Optional matching revolute-joint origin. + """ + + target_link_point_cloud: torch.Tensor + articulation_point_cloud: torch.Tensor + prismatic_joint_axis: torch.Tensor | None = None + revolute_joint_axis: torch.Tensor | None = None + revolute_axis_origin: torch.Tensor | None = None + + def __post_init__(self) -> None: + """Validate and own all tensor fields.""" + target_points = _owned_point_cloud( + self.target_link_point_cloud, + field_name="target_link_point_cloud", + ) + articulation_points = _owned_point_cloud( + self.articulation_point_cloud, + field_name="articulation_point_cloud", + ) + if target_points.device != articulation_points.device: + raise ValueError( + "Articulation affordance point clouds must share a device." + ) + + prismatic_axis = _owned_axis( + self.prismatic_joint_axis, + field_name="prismatic_joint_axis", + ) + revolute_axis = _owned_axis( + self.revolute_joint_axis, + field_name="revolute_joint_axis", + ) + revolute_origin = _owned_point( + self.revolute_axis_origin, + field_name="revolute_axis_origin", + ) + if (revolute_axis is None) != (revolute_origin is None): + raise ValueError( + "revolute_joint_axis and revolute_axis_origin must be provided " + "together." + ) + for value, field_name in ( + (prismatic_axis, "prismatic_joint_axis"), + (revolute_axis, "revolute_joint_axis"), + (revolute_origin, "revolute_axis_origin"), + ): + if value is not None and value.device != target_points.device: + raise ValueError( + f"{field_name} and articulation point clouds must share a device." + ) + + object.__setattr__(self, "target_link_point_cloud", target_points) + object.__setattr__(self, "articulation_point_cloud", articulation_points) + object.__setattr__(self, "prismatic_joint_axis", prismatic_axis) + object.__setattr__(self, "revolute_joint_axis", revolute_axis) + object.__setattr__(self, "revolute_axis_origin", revolute_origin) + + def to_object_geometry(self) -> dict[str, torch.Tensor]: + """Return an owned ``ObjectSemantics.geometry`` dictionary. + + Returns: + A new real dictionary containing cloned point clouds and whichever + optional joint entries are available. + """ + geometry = { + _TARGET_LINK_POINT_CLOUD_KEY: self.target_link_point_cloud.clone(), + _ARTICULATION_POINT_CLOUD_KEY: self.articulation_point_cloud.clone(), + } + if self.prismatic_joint_axis is not None: + geometry[_TARGET_LINK_PRISMATIC_JOINT_AXIS_KEY] = ( + self.prismatic_joint_axis.clone() + ) + if self.revolute_joint_axis is not None: + geometry[_TARGET_LINK_REVOLUTE_JOINT_AXIS_KEY] = ( + self.revolute_joint_axis.clone() + ) + if self.revolute_axis_origin is not None: + geometry[_TARGET_LINK_REVOLUTE_AXIS_ORIGIN_KEY] = ( + self.revolute_axis_origin.clone() + ) + return geometry + + +def sample_initial_articulation_geometry( + provider: ArticulationGeometryProvider, + target_link_name: str, + *, + initial_qpos: torch.Tensor | Sequence[float], + initial_qpos_joint_names: Sequence[str], + body_scale: torch.Tensor | Sequence[float], + articulation_point_count: int = 100_000, + target_point_count: int = 5_000, +) -> ArticulationAffordanceGeometry: + """Sample initial articulation geometry for Atomic Action affordances. + + The adapter evaluates FK at an explicitly named initial joint state, + transforms every raw link mesh into the target link's initial local frame, + merges those meshes, and samples the target and merged surfaces with + Open3D. It also transforms the nearest prismatic and revolute ancestor + joint geometry into that frame. + + Args: + provider: Structural source of deterministic link meshes, FK, and + parent-joint topology. + target_link_name: Link whose initial local frame defines the result. + initial_qpos: Initial joint positions with shape ``(J,)``. + initial_qpos_joint_names: Names corresponding to ``initial_qpos``. + body_scale: Configured articulation scale. Only unit scale is currently + supported because raw meshes and FK must share one metric frame. + articulation_point_count: Merged-articulation surface sample count. + target_point_count: Target-link surface sample count. + + Returns: + Owned typed affordance geometry in the target link's initial frame. + + Raises: + TypeError: If a name, joint-name sequence, or point count has the wrong + type. + ValueError: If the target, initial state, scale, FK output, topology, or + mesh geometry is invalid. + """ + link_names = _validate_adapter_inputs( + provider, + target_link_name=target_link_name, + articulation_point_count=articulation_point_count, + target_point_count=target_point_count, + ) + device = torch.device(provider.device) + initial_qpos, qpos_joint_names = _validate_initial_state( + initial_qpos, + initial_qpos_joint_names, + device=device, + ) + _validate_body_scale(body_scale) + + initial_link_poses = provider.compute_fk( + initial_qpos.unsqueeze(0), + link_names=link_names, + qpos_joint_names=qpos_joint_names, + ) + if ( + not isinstance(initial_link_poses, torch.Tensor) + or initial_link_poses.shape != (1, len(link_names), 4, 4) + or not bool(torch.isfinite(initial_link_poses).all().item()) + ): + raise ValueError( + "Initial FK must return finite poses with shape " + f"(1, {len(link_names)}, 4, 4)." + ) + initial_link_poses = initial_link_poses[0].to( + device=device, + dtype=torch.float32, + ) + target_index = link_names.index(target_link_name) + target_from_root = torch.linalg.inv(initial_link_poses[target_index]) + + prismatic_axis, revolute_axis, revolute_origin = _resolve_joint_geometry( + provider, + target_link_name=target_link_name, + link_names=link_names, + initial_link_poses=initial_link_poses, + target_from_root=target_from_root, + device=device, + ) + ( + target_vertices, + target_triangles, + articulation_vertices, + articulation_triangles, + ) = _merge_initial_link_meshes( + provider, + target_link_name=target_link_name, + link_names=link_names, + initial_link_poses=initial_link_poses, + target_from_root=target_from_root, + device=device, + ) + + return ArticulationAffordanceGeometry( + target_link_point_cloud=_sample_mesh_surface_points( + target_vertices, + target_triangles, + target_point_count, + ), + articulation_point_cloud=_sample_mesh_surface_points( + articulation_vertices, + articulation_triangles, + articulation_point_count, + ), + prismatic_joint_axis=prismatic_axis, + revolute_joint_axis=revolute_axis, + revolute_axis_origin=revolute_origin, + ) + + +def _validate_adapter_inputs( + provider: ArticulationGeometryProvider, + *, + target_link_name: str, + articulation_point_count: int, + target_point_count: int, +) -> list[str]: + """Validate names and point counts shared by adapter stages.""" + if type(target_link_name) is not str: + raise TypeError("target_link_name must be a string.") + if not target_link_name or target_link_name != target_link_name.strip(): + raise ValueError("target_link_name must be non-empty.") + link_names = list(provider.link_names) + if target_link_name not in link_names: + raise ValueError( + f"Unknown articulation link {target_link_name!r}. Available links: " + f"{link_names}." + ) + for value, field_name in ( + (articulation_point_count, "articulation_point_count"), + (target_point_count, "target_point_count"), + ): + if type(value) is not int: + raise TypeError(f"{field_name} must be an integer.") + if value <= 0: + raise ValueError(f"{field_name} must be positive.") + return link_names + + +def _validate_initial_state( + initial_qpos: torch.Tensor | Sequence[float], + initial_qpos_joint_names: Sequence[str], + *, + device: torch.device, +) -> tuple[torch.Tensor, tuple[str, ...]]: + """Validate one explicitly named initial joint state.""" + if isinstance(initial_qpos_joint_names, (str, bytes)) or not isinstance( + initial_qpos_joint_names, + Sequence, + ): + raise TypeError("initial_qpos_joint_names must be a sequence of names.") + joint_names = tuple(initial_qpos_joint_names) + if any( + type(name) is not str or not name or name != name.strip() + for name in joint_names + ): + raise ValueError( + "initial_qpos_joint_names must contain non-empty string names." + ) + if len(set(joint_names)) != len(joint_names): + raise ValueError("initial_qpos_joint_names must not contain duplicates.") + qpos = torch.as_tensor( + initial_qpos, + dtype=torch.float32, + device=device, + ) + if qpos.shape != (len(joint_names),) or not bool(torch.isfinite(qpos).all().item()): + raise ValueError( + "initial_qpos must be a finite vector matching its joint-name sequence." + ) + return qpos, joint_names + + +def _validate_body_scale(body_scale: torch.Tensor | Sequence[float]) -> None: + """Require the raw mesh and FK frames to use unit body scale.""" + scale = torch.as_tensor( + body_scale, + dtype=torch.float32, + device="cpu", + ) + if scale.shape != (3,) or not torch.allclose( + scale, + torch.ones(3, dtype=torch.float32), + ): + raise ValueError( + "Initial articulation geometry sampling currently requires unit " + "body_scale." + ) + + +def _resolve_joint_geometry( + provider: ArticulationGeometryProvider, + *, + target_link_name: str, + link_names: list[str], + initial_link_poses: torch.Tensor, + target_from_root: torch.Tensor, + device: torch.device, +) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + """Resolve nearest prismatic/revolute axes and the revolute origin.""" + resolved_axes: dict[str, torch.Tensor] = {} + revolute_origin: torch.Tensor | None = None + for joint in provider.get_parent_joint_chain(target_link_name): + joint_type = joint.joint_type + if joint_type not in ("prismatic", "revolute") or joint_type in resolved_axes: + continue + if joint.parent_link_name not in link_names: + raise ValueError( + f"{joint_type.capitalize()} joint {joint.name!r} parent link " + f"{joint.parent_link_name!r} is not an articulation link." + ) + joint_origin_pose = torch.as_tensor( + joint.origin_pose, + dtype=torch.float32, + device=device, + ) + if joint_origin_pose.shape != (4, 4) or not bool( + torch.isfinite(joint_origin_pose).all().item() + ): + raise ValueError( + f"{joint_type.capitalize()} joint {joint.name!r} origin pose " + "must be finite with shape (4, 4)." + ) + joint_axis = torch.as_tensor( + joint.axis, + dtype=torch.float32, + device=device, + ) + if ( + joint_axis.shape != (3,) + or not bool(torch.isfinite(joint_axis).all().item()) + or float(torch.linalg.vector_norm(joint_axis).item()) + <= torch.finfo(joint_axis.dtype).eps + ): + raise ValueError( + f"{joint_type.capitalize()} joint {joint.name!r} axis must be " + "finite and nonzero with shape (3,)." + ) + parent_index = link_names.index(joint.parent_link_name) + target_from_joint = target_from_root @ initial_link_poses[parent_index] + target_from_joint = target_from_joint @ joint_origin_pose + target_axis = target_from_joint[:3, :3] @ joint_axis + axis_norm = torch.linalg.vector_norm(target_axis) + if ( + not bool(torch.isfinite(target_axis).all().item()) + or not bool(torch.isfinite(axis_norm).item()) + or float(axis_norm.item()) <= torch.finfo(target_axis.dtype).eps + ): + raise ValueError( + f"{joint_type.capitalize()} joint {joint.name!r} axis must " + "transform to a finite, nonzero target-link vector." + ) + resolved_axes[joint_type] = target_axis / axis_norm + if joint_type == "revolute": + revolute_origin = target_from_joint[:3, 3].clone() + if len(resolved_axes) == 2: + break + return ( + resolved_axes.get("prismatic"), + resolved_axes.get("revolute"), + revolute_origin, + ) + + +def _merge_initial_link_meshes( + provider: ArticulationGeometryProvider, + *, + target_link_name: str, + link_names: list[str], + initial_link_poses: torch.Tensor, + target_from_root: torch.Tensor, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Transform and merge raw link meshes in the target-link frame.""" + merged_vertices: list[torch.Tensor] = [] + merged_triangles: list[torch.Tensor] = [] + vertex_offset = 0 + target_vertices: torch.Tensor | None = None + target_triangles: torch.Tensor | None = None + for link_index, link_name in enumerate(link_names): + vertices, triangles = provider.get_link_vert_face(link_name) + vertices = torch.as_tensor( + vertices, + dtype=torch.float32, + device=device, + ) + triangles = torch.as_tensor( + triangles, + device=device, + ) + if ( + triangles.dtype == torch.bool + or triangles.is_floating_point() + or triangles.is_complex() + ): + raise ValueError(f"Link {link_name!r} triangles must use integer indices.") + triangles = triangles.to(dtype=torch.long) + _validate_point_cloud_mesh(vertices, triangles, link_name=link_name) + if vertices.shape[0] == 0: + continue + + target_from_link = target_from_root @ initial_link_poses[link_index] + transformed_vertices = ( + vertices @ target_from_link[:3, :3].transpose(0, 1) + + target_from_link[:3, 3] + ) + merged_vertices.append(transformed_vertices) + if triangles.shape[0] > 0: + merged_triangles.append(triangles + vertex_offset) + if link_name == target_link_name: + target_vertices = transformed_vertices + target_triangles = triangles + vertex_offset += vertices.shape[0] + + if target_vertices is None or target_vertices.shape[0] == 0: + raise ValueError( + f"Target link {target_link_name!r} has no point-cloud geometry." + ) + if not merged_vertices: + raise ValueError("Articulation has no point-cloud geometry.") + articulation_vertices = torch.cat(merged_vertices, dim=0) + articulation_triangles = ( + torch.cat(merged_triangles, dim=0) + if merged_triangles + else torch.empty((0, 3), dtype=torch.long, device=device) + ) + assert target_triangles is not None + return ( + target_vertices, + target_triangles, + articulation_vertices, + articulation_triangles, + ) + + +def _validate_point_cloud_mesh( + vertices: torch.Tensor, + triangles: torch.Tensor, + *, + link_name: str, +) -> None: + """Validate one link-local mesh used for surface sampling.""" + if vertices.dim() != 2 or vertices.shape[1:] != (3,): + raise ValueError(f"Link {link_name!r} vertices must have shape (N, 3).") + if not bool(torch.isfinite(vertices).all().item()): + raise ValueError(f"Link {link_name!r} vertices must be finite.") + if triangles.dim() != 2 or triangles.shape[1:] != (3,): + raise ValueError(f"Link {link_name!r} triangles must have shape (M, 3).") + if triangles.shape[0] == 0: + return + if vertices.shape[0] == 0: + raise ValueError( + f"Link {link_name!r} triangles cannot reference an empty mesh." + ) + if bool((triangles < 0).any().item()) or int(triangles.max().item()) >= len( + vertices + ): + raise ValueError(f"Link {link_name!r} triangles reference invalid vertices.") + + +def _sample_mesh_surface_points( + vertices: torch.Tensor, + triangles: torch.Tensor, + point_count: int, +) -> torch.Tensor: + """Uniformly sample a triangle mesh with Open3D's CPU sampler.""" + if type(point_count) is not int: + raise TypeError("point_count must be an integer.") + if point_count <= 0: + raise ValueError("point_count must be positive.") + if vertices.shape[0] == 0: + raise ValueError("Cannot sample an empty mesh.") + if triangles.shape[0] == 0: + indices = ( + torch.linspace( + 0, + vertices.shape[0] - 1, + point_count, + device=vertices.device, + ) + .round() + .to(torch.long) + ) + return vertices[indices].clone() + + face_vertices = vertices[triangles] + face_areas = 0.5 * torch.linalg.vector_norm( + torch.cross( + face_vertices[:, 1] - face_vertices[:, 0], + face_vertices[:, 2] - face_vertices[:, 0], + dim=1, + ), + dim=1, + ) + valid_faces = face_areas > torch.finfo(vertices.dtype).eps + if not bool(valid_faces.any().item()): + indices = ( + torch.linspace( + 0, + vertices.shape[0] - 1, + point_count, + device=vertices.device, + ) + .round() + .to(torch.long) + ) + return vertices[indices].clone() + + import numpy as np + import open3d as o3d + + mesh = o3d.geometry.TriangleMesh( + vertices=o3d.utility.Vector3dVector( + vertices.detach().to(device="cpu", dtype=torch.float64).numpy() + ), + triangles=o3d.utility.Vector3iVector( + triangles[valid_faces].detach().to(device="cpu", dtype=torch.int32).numpy() + ), + ) + point_cloud = mesh.sample_points_uniformly(number_of_points=point_count) + sampled_points = np.asarray(point_cloud.points).copy() + if sampled_points.shape != (point_count, 3): + raise RuntimeError( + "Open3D surface sampling returned an unexpected point-cloud shape " + f"{sampled_points.shape}; expected ({point_count}, 3)." + ) + return torch.tensor( + sampled_points, + device=vertices.device, + dtype=vertices.dtype, + ) + + +def _owned_point_cloud(value: object, *, field_name: str) -> torch.Tensor: + """Validate and clone one floating point cloud.""" + if not isinstance(value, torch.Tensor): + raise TypeError(f"{field_name} must be a torch.Tensor.") + if ( + not value.is_floating_point() + or value.dim() != 2 + or value.shape[1:] != (3,) + or value.shape[0] == 0 + or not bool(torch.isfinite(value).all().item()) + ): + raise ValueError( + f"{field_name} must be a non-empty finite floating tensor with " + "shape (N, 3)." + ) + return value.to(dtype=torch.float32).clone() + + +def _owned_axis( + value: object | None, + *, + field_name: str, +) -> torch.Tensor | None: + """Validate, normalize, and clone one optional axis.""" + if value is None: + return None + axis = _owned_point(value, field_name=field_name) + assert axis is not None + axis_norm = torch.linalg.vector_norm(axis) + if float(axis_norm.item()) <= 1.0e-6: + raise ValueError(f"{field_name} must be non-zero.") + return axis / axis_norm + + +def _owned_point( + value: object | None, + *, + field_name: str, +) -> torch.Tensor | None: + """Validate and clone one optional local point.""" + if value is None: + return None + if not isinstance(value, torch.Tensor): + raise TypeError(f"{field_name} must be a torch.Tensor.") + if ( + not value.is_floating_point() + or value.shape != (3,) + or not bool(torch.isfinite(value).all().item()) + ): + raise ValueError( + f"{field_name} must be a finite floating tensor with shape (3,)." + ) + return value.to(dtype=torch.float32).clone() + + +__all__ = [ + "ArticulationAffordanceGeometry", + "ArticulationGeometryProvider", + "ArticulationJointGeometry", + "sample_initial_articulation_geometry", +] diff --git a/embodichain/lab/sim/objects/articulation.py b/embodichain/lab/sim/objects/articulation.py index 713578203..9940f6108 100644 --- a/embodichain/lab/sim/objects/articulation.py +++ b/embodichain/lab/sim/objects/articulation.py @@ -21,7 +21,6 @@ import torch import dexsim import numpy as np -import open3d as o3d from dataclasses import dataclass from functools import cached_property @@ -1198,364 +1197,6 @@ def get_link_vert_face(self, link_name: str) -> Tuple[torch.Tensor, torch.Tensor verts, faces = self.body_data.link_vert_face[link_name] return verts, faces - def sample_initial_point_clouds( - self, - target_link_name: str, - *, - articulation_point_count: int = 100_000, - target_point_count: int = 5_000, - ) -> dict[str, torch.Tensor]: - """Sample initial articulation geometry in a target-link frame. - - All link meshes are transformed with forward kinematics evaluated at - :attr:`ArticulationCfg.init_qpos`, merged into one triangle mesh, and - sampled proportionally to triangle area. The target link is sampled - separately so consumers can derive a target-centered neighborhood. - Both returned clouds use the target link's initial local frame. - - Args: - target_link_name: Link whose initial local frame defines the output. - articulation_point_count: Number of points sampled from the merged - articulation surface. - target_point_count: Number of points sampled from the target-link - surface. - - Returns: - Geometry metadata containing ``articulation_point_cloud`` with - shape ``(articulation_point_count, 3)`` and - ``target_link_point_cloud`` with shape - ``(target_point_count, 3)``. If the target link has a revolute - ancestor, ``target_link_revolute_axis_origin`` contains that - nearest joint's origin in the target link's initial local frame. - ``target_link_prismatic_joint_axis`` and - ``target_link_revolute_joint_axis`` contain the normalized axes of - the nearest ancestors of each type in the same frame when present. - - Raises: - TypeError: If a name or point count has the wrong type. - ValueError: If the target link, point counts, initial joint state, - FK output, parent-joint geometry, or mesh geometry is invalid. - RuntimeError: If the articulation has no kinematic chain. - """ - if type(target_link_name) is not str: - raise TypeError("target_link_name must be a string.") - if not target_link_name or target_link_name != target_link_name.strip(): - raise ValueError("target_link_name must be non-empty.") - if target_link_name not in self.link_names: - raise ValueError( - f"Unknown articulation link {target_link_name!r}. Available " - f"links: {list(self.link_names)}." - ) - for value, field_name in ( - (articulation_point_count, "articulation_point_count"), - (target_point_count, "target_point_count"), - ): - if type(value) is not int: - raise TypeError(f"{field_name} must be an integer.") - if value <= 0: - raise ValueError(f"{field_name} must be positive.") - if self.pk_chain is None: - raise RuntimeError( - "Initial point-cloud sampling requires cfg.build_pk_chain=True." - ) - body_scale = torch.as_tensor(self.cfg.body_scale, dtype=torch.float32) - if body_scale.shape != (3,) or not torch.allclose( - body_scale, - torch.ones(3, dtype=torch.float32), - ): - raise ValueError( - "Initial point-cloud sampling currently requires unit body_scale." - ) - - initial_qpos = torch.as_tensor( - self.cfg.init_qpos, - dtype=torch.float32, - device=self.device, - ) - if initial_qpos.shape != (self.dof,) or not torch.isfinite(initial_qpos).all(): - raise ValueError( - "ArticulationCfg.init_qpos must be a finite vector matching " - f"the articulation DoF ({self.dof})." - ) - backend_joint_names = list(self.joint_names) - kinematic_joint_names = list(self.pk_chain.get_joint_parameter_names()) - if backend_joint_names != kinematic_joint_names: - if len(backend_joint_names) != len(kinematic_joint_names) or set( - backend_joint_names - ) != set(kinematic_joint_names): - raise ValueError( - "Initial point-cloud sampling requires matching simulator " - "and kinematic-chain joint names." - ) - qpos_by_name = dict(zip(backend_joint_names, initial_qpos)) - initial_qpos = torch.stack( - [qpos_by_name[name] for name in kinematic_joint_names] - ) - link_names = list(self.link_names) - initial_link_poses = self.compute_fk( - initial_qpos.unsqueeze(0), - link_names=link_names, - ) - if ( - initial_link_poses.shape != (1, len(link_names), 4, 4) - or not torch.isfinite(initial_link_poses).all() - ): - raise ValueError( - "Initial FK must return finite poses with shape " - f"(1, {len(link_names)}, 4, 4)." - ) - initial_link_poses = initial_link_poses[0].to( - device=self.device, - dtype=torch.float32, - ) - target_index = link_names.index(target_link_name) - target_from_root = torch.linalg.inv(initial_link_poses[target_index]) - target_link_revolute_axis_origin: torch.Tensor | None = None - target_link_joint_axes: dict[str, torch.Tensor] = {} - joint_axis_geometry_keys = { - "prismatic": "target_link_prismatic_joint_axis", - "revolute": "target_link_revolute_joint_axis", - } - for joint in self.get_parent_joint_chain(target_link_name): - joint_type = joint.joint_type - geometry_key = joint_axis_geometry_keys.get(joint_type) - if geometry_key is None or geometry_key in target_link_joint_axes: - continue - if joint.parent_link_name not in link_names: - raise ValueError( - f"{joint_type.capitalize()} joint {joint.name!r} parent link " - f"{joint.parent_link_name!r} is not an articulation link." - ) - joint_origin_pose = torch.as_tensor( - joint.origin_pose, - dtype=torch.float32, - device=self.device, - ) - if ( - joint_origin_pose.shape != (4, 4) - or not torch.isfinite(joint_origin_pose).all() - ): - raise ValueError( - f"{joint_type.capitalize()} joint {joint.name!r} origin pose " - "must be finite " - "with shape (4, 4)." - ) - joint_axis = torch.as_tensor( - joint.axis, - dtype=torch.float32, - device=self.device, - ) - if ( - joint_axis.shape != (3,) - or not torch.isfinite(joint_axis).all() - or torch.linalg.vector_norm(joint_axis) - <= torch.finfo(joint_axis.dtype).eps - ): - raise ValueError( - f"{joint_type.capitalize()} joint {joint.name!r} axis must " - "be finite and nonzero with shape (3,)." - ) - parent_index = link_names.index(joint.parent_link_name) - target_from_joint = torch.matmul( - torch.matmul( - target_from_root, - initial_link_poses[parent_index], - ), - joint_origin_pose, - ) - target_link_joint_axis = torch.matmul( - target_from_joint[:3, :3], - joint_axis, - ) - target_link_joint_axis_norm = torch.linalg.vector_norm( - target_link_joint_axis - ) - if ( - not torch.isfinite(target_link_joint_axis).all() - or not torch.isfinite(target_link_joint_axis_norm) - or target_link_joint_axis_norm - <= torch.finfo(target_link_joint_axis.dtype).eps - ): - raise ValueError( - f"{joint_type.capitalize()} joint {joint.name!r} axis must " - "transform to a finite, nonzero target-link vector." - ) - target_link_joint_axes[geometry_key] = ( - target_link_joint_axis / target_link_joint_axis_norm - ) - if joint_type == "revolute": - target_link_revolute_axis_origin = target_from_joint[:3, 3].clone() - if len(target_link_joint_axes) == len(joint_axis_geometry_keys): - break - - merged_vertices: list[torch.Tensor] = [] - merged_triangles: list[torch.Tensor] = [] - vertex_offset = 0 - target_vertices: torch.Tensor | None = None - target_triangles: torch.Tensor | None = None - for link_index, link_name in enumerate(link_names): - vertices, triangles = self.get_link_vert_face(link_name) - vertices = torch.as_tensor( - vertices, - dtype=torch.float32, - device=self.device, - ) - triangles = torch.as_tensor( - triangles, - dtype=torch.long, - device=self.device, - ) - self._validate_point_cloud_mesh(vertices, triangles, link_name=link_name) - if vertices.shape[0] == 0: - continue - - target_from_link = torch.matmul( - target_from_root, - initial_link_poses[link_index], - ) - transformed_vertices = ( - torch.matmul( - vertices, - target_from_link[:3, :3].transpose(0, 1), - ) - + target_from_link[:3, 3] - ) - merged_vertices.append(transformed_vertices) - if triangles.shape[0] > 0: - merged_triangles.append(triangles + vertex_offset) - if link_name == target_link_name: - target_vertices = transformed_vertices - target_triangles = triangles - vertex_offset += vertices.shape[0] - - if target_vertices is None or target_vertices.shape[0] == 0: - raise ValueError( - f"Target link {target_link_name!r} has no point-cloud geometry." - ) - if not merged_vertices: - raise ValueError("Articulation has no point-cloud geometry.") - articulation_vertices = torch.cat(merged_vertices, dim=0) - articulation_triangles = ( - torch.cat(merged_triangles, dim=0) - if merged_triangles - else torch.empty((0, 3), dtype=torch.long, device=self.device) - ) - assert target_triangles is not None - geometry = { - "target_link_point_cloud": self._sample_mesh_surface_points( - target_vertices, - target_triangles, - target_point_count, - ), - "articulation_point_cloud": self._sample_mesh_surface_points( - articulation_vertices, - articulation_triangles, - articulation_point_count, - ), - } - if target_link_revolute_axis_origin is not None: - geometry["target_link_revolute_axis_origin"] = ( - target_link_revolute_axis_origin - ) - geometry.update(target_link_joint_axes) - return geometry - - @staticmethod - def _validate_point_cloud_mesh( - vertices: torch.Tensor, - triangles: torch.Tensor, - *, - link_name: str, - ) -> None: - """Validate one link-local mesh used for surface sampling.""" - if vertices.dim() != 2 or vertices.shape[1:] != (3,): - raise ValueError(f"Link {link_name!r} vertices must have shape (N, 3).") - if not torch.isfinite(vertices).all(): - raise ValueError(f"Link {link_name!r} vertices must be finite.") - if triangles.dim() != 2 or triangles.shape[1:] != (3,): - raise ValueError(f"Link {link_name!r} triangles must have shape (M, 3).") - if triangles.shape[0] == 0: - return - if vertices.shape[0] == 0: - raise ValueError( - f"Link {link_name!r} triangles cannot reference an empty mesh." - ) - if bool((triangles < 0).any().item()) or int(triangles.max().item()) >= len( - vertices - ): - raise ValueError( - f"Link {link_name!r} triangles reference invalid vertices." - ) - - @staticmethod - def _sample_mesh_surface_points( - vertices: torch.Tensor, - triangles: torch.Tensor, - point_count: int, - ) -> torch.Tensor: - """Uniformly sample a triangle mesh with Open3D's CPU sampler.""" - if vertices.shape[0] == 0: - raise ValueError("Cannot sample an empty mesh.") - if triangles.shape[0] == 0: - indices = ( - torch.linspace( - 0, - vertices.shape[0] - 1, - point_count, - device=vertices.device, - ) - .round() - .to(torch.long) - ) - return vertices[indices] - - face_vertices = vertices[triangles] - face_areas = 0.5 * torch.linalg.vector_norm( - torch.cross( - face_vertices[:, 1] - face_vertices[:, 0], - face_vertices[:, 2] - face_vertices[:, 0], - dim=1, - ), - dim=1, - ) - valid_faces = face_areas > torch.finfo(vertices.dtype).eps - if not bool(valid_faces.any().item()): - indices = ( - torch.linspace( - 0, - vertices.shape[0] - 1, - point_count, - device=vertices.device, - ) - .round() - .to(torch.long) - ) - return vertices[indices] - - mesh = o3d.geometry.TriangleMesh( - vertices=o3d.utility.Vector3dVector( - vertices.detach().to(device="cpu", dtype=torch.float64).numpy() - ), - triangles=o3d.utility.Vector3iVector( - triangles[valid_faces] - .detach() - .to(device="cpu", dtype=torch.int32) - .numpy() - ), - ) - point_cloud = mesh.sample_points_uniformly(number_of_points=point_count) - sampled_points = np.asarray(point_cloud.points).copy() - if sampled_points.shape != (point_count, 3): - raise RuntimeError( - "Open3D surface sampling returned an unexpected point-cloud shape " - f"{sampled_points.shape}; expected ({point_count}, 3)." - ) - return torch.tensor( - sampled_points, - device=vertices.device, - dtype=vertices.dtype, - ) - def get_link_pose( self, link_name: str, env_ids: Sequence[int] | None = None, to_matrix=False ) -> torch.Tensor: @@ -2539,10 +2180,12 @@ def _set_default_joint_drive(self) -> None: def compute_fk( self, qpos: torch.Tensor | np.ndarray | None, - link_names: str | list[str] | tuple[str] | None = None, + link_names: str | list[str] | tuple[str, ...] | None = None, end_link_name: str | None = None, root_link_name: str | None = None, to_dict: bool = False, + *, + qpos_joint_names: Sequence[str] | None = None, **kwargs, ) -> Union[torch.Tensor, dict[str, "pk.Transform3d"]]: """Compute the forward kinematics (FK) for the given joint positions. @@ -2555,6 +2198,9 @@ def compute_fk( end_link_name (str, optional): Name of the end link for which FK is computed. If None, all links are considered. root_link_name (str, optional): Name of the root link for which FK is computed. Defaults to None. to_dict (bool, optional): If True, returns the FK result as a dictionary of Transform3d objects. Defaults to False. + qpos_joint_names: Optional names corresponding to the last dimension + of ``qpos``. When supplied, the values are reordered into the + kinematic chain's parameter order before FK. **kwargs: Additional keyword arguments for customization. Raises: @@ -2571,6 +2217,52 @@ def compute_fk( if self.pk_chain is None: logger.log_error("pk_chain is not initialized for this articulation.") + if qpos_joint_names is not None: + if end_link_name is not None or root_link_name is not None: + raise ValueError( + "qpos_joint_names cannot be combined with serial-chain FK." + ) + if isinstance(qpos_joint_names, (str, bytes)) or not isinstance( + qpos_joint_names, + Sequence, + ): + raise TypeError("qpos_joint_names must be a sequence of names.") + supplied_joint_names = tuple(qpos_joint_names) + if any( + type(name) is not str or not name or name != name.strip() + for name in supplied_joint_names + ): + raise ValueError( + "qpos_joint_names must contain non-empty string names." + ) + if len(set(supplied_joint_names)) != len(supplied_joint_names): + raise ValueError("qpos_joint_names must not contain duplicates.") + if qpos is None: + raise ValueError("qpos is required when qpos_joint_names is supplied.") + qpos_tensor = torch.as_tensor(qpos, device=self.device) + if qpos_tensor.dim() not in (1, 2) or qpos_tensor.shape[-1] != len( + supplied_joint_names + ): + raise ValueError("qpos last dimension must match qpos_joint_names.") + kinematic_joint_names = tuple(self.pk_chain.get_joint_parameter_names()) + if len(supplied_joint_names) != len(kinematic_joint_names) or set( + supplied_joint_names + ) != set(kinematic_joint_names): + raise ValueError( + "qpos_joint_names must match the kinematic-chain joint names." + ) + if kinematic_joint_names: + qpos_by_name = { + name: qpos_tensor[..., index] + for index, name in enumerate(supplied_joint_names) + } + qpos = torch.stack( + [qpos_by_name[name] for name in kinematic_joint_names], + dim=-1, + ) + else: + qpos = qpos_tensor[..., :0].clone() + # Adapt link_names to work with get_frame_indices if link_names is not None: if isinstance(link_names, str): diff --git a/embodichain/lab/sim/objects/robot.py b/embodichain/lab/sim/objects/robot.py index 7b8a1340e..9de9599e4 100644 --- a/embodichain/lab/sim/objects/robot.py +++ b/embodichain/lab/sim/objects/robot.py @@ -770,6 +770,8 @@ def compute_fk( root_link_name: str | None = None, env_ids: Sequence[int] | None = None, to_matrix: bool = False, + *, + qpos_joint_names: Sequence[str] | None = None, ) -> torch.Tensor: """Compute the forward kinematics of the robot given joint positions and optionally a specific part name. The output pose will be in the local arena frame. @@ -782,6 +784,8 @@ def compute_fk( root_link_name (str | None): The name of the root link to compute the FK for. If None, the default root link is used. env_ids (Sequence[int] | None): The environment ids to compute the FK for. If None, all environments are used. to_matrix (bool): If True, returns the transformation in the form of a 4x4 matrix. + qpos_joint_names: Optional names corresponding to the last dimension + of ``qpos`` for full-articulation FK. Returns: torch.Tensor: The forward kinematics result with shape (num_envs, 7) or (num_envs, 4, 4) if `to_matrix` is True. @@ -794,6 +798,12 @@ def compute_fk( link_names=link_names, end_link_name=end_link_name, root_link_name=root_link_name, + qpos_joint_names=qpos_joint_names, + ) + + if qpos_joint_names is not None: + raise ValueError( + "qpos_joint_names cannot be combined with named robot control-part FK." ) if not self._solvers: diff --git a/embodichain/lab/task_program/integrations/_configured_services.py b/embodichain/lab/task_program/integrations/_configured_services.py index 2dd73e7dc..0ea67d8d7 100644 --- a/embodichain/lab/task_program/integrations/_configured_services.py +++ b/embodichain/lab/task_program/integrations/_configured_services.py @@ -53,6 +53,7 @@ SlideAffordance, SlideGoal, SlideOptions, + sample_initial_articulation_geometry, ) from embodichain.lab.task_program.semantics import ( BinaryEffectEvidenceQuery, @@ -359,16 +360,13 @@ def create( if not callable(get_link_vert_face): raise TypeError("Articulation must provide get_link_vert_face().") vertices, triangles = get_link_vert_face(native_link_name) - sample_initial_point_clouds = getattr( + geometry = sample_initial_articulation_geometry( articulation, - "sample_initial_point_clouds", - None, - ) - if not callable(sample_initial_point_clouds): - raise TypeError("Articulation must provide sample_initial_point_clouds().") - geometry = sample_initial_point_clouds(native_link_name) - if not isinstance(geometry, dict): - raise TypeError("Articulation point-cloud geometry must be a dict.") + native_link_name, + initial_qpos=articulation.cfg.init_qpos, + initial_qpos_joint_names=articulation.joint_names, + body_scale=articulation.cfg.body_scale, + ).to_object_geometry() affordance_kwargs: dict[str, object] = {} if self.translation_axis is not None: affordance_kwargs["translation_axis"] = torch.tensor( diff --git a/scripts/tutorials/atomic_action/press.py b/scripts/tutorials/atomic_action/press.py index e7e7c5afd..e9bbcbf53 100644 --- a/scripts/tutorials/atomic_action/press.py +++ b/scripts/tutorials/atomic_action/press.py @@ -37,6 +37,7 @@ PressAffordance, PressGoal, PressOptions, + sample_initial_articulation_geometry, ) from embodichain.lab.sim.cfg import ( ArticulationCfg, @@ -131,7 +132,13 @@ def create_button_semantics( """Create press semantics for an articulation-link or rigid button.""" if isinstance(target, Articulation): target_pose = target.get_link_pose(BUTTON_LINK_NAME, to_matrix=True) - geometry = target.sample_initial_point_clouds(BUTTON_LINK_NAME) + geometry = sample_initial_articulation_geometry( + target, + BUTTON_LINK_NAME, + initial_qpos=target.cfg.init_qpos, + initial_qpos_joint_names=target.joint_names, + body_scale=target.cfg.body_scale, + ).to_object_geometry() affordance = PressAffordance() label = "microwave_start_button" else: diff --git a/scripts/tutorials/atomic_action/slide.py b/scripts/tutorials/atomic_action/slide.py index ca7b5053b..3a70cd517 100644 --- a/scripts/tutorials/atomic_action/slide.py +++ b/scripts/tutorials/atomic_action/slide.py @@ -40,6 +40,7 @@ SlideAffordance, SlideGoal, SlideOptions, + sample_initial_articulation_geometry, ) from embodichain.lab.sim.cfg import ( ArticulationCfg, @@ -114,7 +115,13 @@ def create_drawer_semantics(drawer: Articulation) -> ObjectSemantics: Pure target-local semantics for the handle's pull/push affordance. """ vertices, triangles = drawer.get_link_vert_face(HANDLE_LINK_NAME) - geometry = drawer.sample_initial_point_clouds(HANDLE_LINK_NAME) + geometry = sample_initial_articulation_geometry( + drawer, + HANDLE_LINK_NAME, + initial_qpos=drawer.cfg.init_qpos, + initial_qpos_joint_names=drawer.joint_names, + body_scale=drawer.cfg.body_scale, + ).to_object_geometry() return ObjectSemantics( label="drawer_large_handle", geometry=geometry, diff --git a/scripts/tutorials/atomic_action/twist.py b/scripts/tutorials/atomic_action/twist.py index 22ec41289..bc366bdb5 100644 --- a/scripts/tutorials/atomic_action/twist.py +++ b/scripts/tutorials/atomic_action/twist.py @@ -37,6 +37,7 @@ TwistAffordance, TwistGoal, TwistOptions, + sample_initial_articulation_geometry, ) from embodichain.lab.sim.cfg import ( ArticulationCfg, @@ -123,7 +124,13 @@ def create_knob_semantics( if isinstance(target, Articulation): vertices, _ = target.get_link_vert_face(KNOB_LINK_NAME) target_pose = target.get_link_pose(KNOB_LINK_NAME, to_matrix=True) - geometry = target.sample_initial_point_clouds(KNOB_LINK_NAME) + geometry = sample_initial_articulation_geometry( + target, + KNOB_LINK_NAME, + initial_qpos=target.cfg.init_qpos, + initial_qpos_joint_names=target.joint_names, + body_scale=target.cfg.body_scale, + ).to_object_geometry() affordance = TwistAffordance( grasp_position=_mesh_center(vertices), ) diff --git a/tests/gym/envs/task_program/test_task_vertical_slices.py b/tests/gym/envs/task_program/test_task_vertical_slices.py index 0e3efd0c9..839d9c248 100644 --- a/tests/gym/envs/task_program/test_task_vertical_slices.py +++ b/tests/gym/envs/task_program/test_task_vertical_slices.py @@ -24,6 +24,7 @@ import json from pathlib import Path from types import SimpleNamespace +from unittest.mock import call, patch import pytest import torch @@ -48,6 +49,7 @@ ) from embodichain.lab.sim.atomic_actions import ( Affordance, + ArticulationAffordanceGeometry, AtomicActionEngine, DynamicCollisionMode, EntityState, @@ -131,6 +133,16 @@ def _drawer_handle_point_cloud_geometry() -> dict[str, torch.Tensor]: } +def _drawer_handle_articulation_geometry() -> ArticulationAffordanceGeometry: + """Return the typed adapter result used by the configured lowerer.""" + geometry = _drawer_handle_point_cloud_geometry() + return ArticulationAffordanceGeometry( + target_link_point_cloud=geometry["target_link_point_cloud"], + articulation_point_cloud=geometry["articulation_point_cloud"], + prismatic_joint_axis=geometry["target_link_prismatic_joint_axis"], + ) + + class _NeverObserveProvider: """Reject dynamic observations during configuration decoding/compilation.""" @@ -554,6 +566,11 @@ def test_open_drawer_config_owns_registered_lowerer_factory() -> None: class Drawer: link_names = (_OPEN_DRAWER_HANDLE_LINK_NAME,) + joint_names = ("drawer_joint",) + cfg = SimpleNamespace( + init_qpos=(0.0,), + body_scale=(1.0, 1.0, 1.0), + ) @staticmethod def get_link_vert_face(name: str) -> tuple[torch.Tensor, torch.Tensor]: @@ -563,11 +580,6 @@ def get_link_vert_face(name: str) -> tuple[torch.Tensor, torch.Tensor]: torch.tensor([[0, 1, 2]]), ) - @staticmethod - def sample_initial_point_clouds(name: str) -> dict[str, torch.Tensor]: - assert name == _OPEN_DRAWER_HANDLE_LINK_NAME - return _drawer_handle_point_cloud_geometry() - drawer_ref = SceneArticulationRef(_OPEN_DRAWER_ENTITY_ID) registry = SceneRegistry( ( @@ -589,24 +601,46 @@ def sample_initial_point_clouds(name: str) -> dict[str, torch.Tensor]: robot=robot, device=torch.device("cpu"), ) + drawer = Drawer() simulation = SimpleNamespace( - get_articulation=lambda identifier: ( - Drawer() if identifier == "drawer" else None + get_articulation=lambda identifier: drawer if identifier == "drawer" else None + ) + sampled_geometry = _drawer_handle_articulation_geometry() + + with patch( + "embodichain.lab.task_program.integrations._configured_services." + "sample_initial_articulation_geometry", + return_value=sampled_geometry, + ) as sampler: + first = registration.create_registered_semantic_lowerers( + simulation=simulation, + robot=robot, + scene_registry=registry, + engine=engine, + ) + second = registration.create_registered_semantic_lowerers( + simulation=simulation, + robot=robot, + scene_registry=registry, + engine=engine, ) - ) - first = registration.create_registered_semantic_lowerers( - simulation=simulation, - robot=robot, - scene_registry=registry, - engine=engine, - ) - second = registration.create_registered_semantic_lowerers( - simulation=simulation, - robot=robot, - scene_registry=registry, - engine=engine, - ) + assert sampler.call_args_list == [ + call( + drawer, + _OPEN_DRAWER_HANDLE_LINK_NAME, + initial_qpos=drawer.cfg.init_qpos, + initial_qpos_joint_names=drawer.joint_names, + body_scale=drawer.cfg.body_scale, + ), + call( + drawer, + _OPEN_DRAWER_HANDLE_LINK_NAME, + initial_qpos=drawer.cfg.init_qpos, + initial_qpos_joint_names=drawer.joint_names, + body_scale=drawer.cfg.body_scale, + ), + ] assert type(first[0]) is type(second[0]) assert type(first[0]).call_id == _OPEN_DRAWER_CALL_ID diff --git a/tests/sim/atomic_actions/test_articulation_geometry.py b/tests/sim/atomic_actions/test_articulation_geometry.py new file mode 100644 index 000000000..450e398f9 --- /dev/null +++ b/tests/sim/atomic_actions/test_articulation_geometry.py @@ -0,0 +1,899 @@ +# ---------------------------------------------------------------------------- +# 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 + +from types import SimpleNamespace + +import pytest +import torch + +import embodichain.lab.sim.atomic_actions.articulation_geometry as articulation_geometry_module +from embodichain.lab.sim.atomic_actions.articulation_geometry import ( + ArticulationAffordanceGeometry, + sample_initial_articulation_geometry, +) +from embodichain.lab.sim.objects import ArticulationJointKinematics + +POINT_CLOUD_TOLERANCE = 1.0e-6 + + +class _ArticulationGeometryProvider: + """Minimal structural provider for initial articulation geometry sampling.""" + + def __init__( + self, + *, + initial_qpos: tuple[float, ...], + backend_joint_names: tuple[str, ...], + link_meshes: dict[str, tuple[torch.Tensor, torch.Tensor]], + link_poses: torch.Tensor, + body_scale: tuple[float, float, float], + parent_joint_chain: tuple[object, ...], + ) -> None: + self.device = torch.device("cpu") + self.initial_qpos = initial_qpos + self.initial_qpos_joint_names = backend_joint_names + self.body_scale = body_scale + self.link_names = list(link_meshes) + self._link_meshes = link_meshes + self._link_poses = link_poses + self._parent_joint_chain = parent_joint_chain + self.fk_calls: list[tuple[torch.Tensor, list[str], tuple[str, ...]]] = [] + + def compute_fk( + self, + qpos: torch.Tensor, + *, + link_names: list[str], + qpos_joint_names: tuple[str, ...], + ) -> torch.Tensor: + """Return configured initial link poses and record the FK request.""" + self.fk_calls.append((qpos.clone(), list(link_names), tuple(qpos_joint_names))) + return self._link_poses.clone() + + def get_link_vert_face( + self, + link_name: str, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return one configured link-local triangle mesh.""" + return self._link_meshes[link_name] + + def get_parent_joint_chain( + self, + link_name: str, + ) -> tuple[object, ...]: + """Return the configured target-to-root parent-joint chain.""" + assert link_name in self.link_names + return self._parent_joint_chain + + +def _make_point_cloud_articulation( + *, + initial_qpos: tuple[float, ...] = (0.25,), + backend_joint_names: tuple[str, ...] = ("joint",), + link_meshes: dict[str, tuple[torch.Tensor, torch.Tensor]] | None = None, + link_poses: torch.Tensor | None = None, + body_scale: tuple[float, float, float] = (1.0, 1.0, 1.0), + parent_joint_chain: tuple[object, ...] = (), +) -> tuple[ + _ArticulationGeometryProvider, + list[tuple[torch.Tensor, list[str], tuple[str, ...]]], +]: + """Build a pure-Python structural provider for the geometry adapter.""" + if link_meshes is None: + link_meshes = { + "target": ( + torch.tensor( + ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)), + dtype=torch.float32, + ), + torch.tensor(((0, 1, 2),), dtype=torch.long), + ) + } + if link_poses is None: + link_poses = torch.eye(4, dtype=torch.float32).repeat( + 1, + len(link_meshes), + 1, + 1, + ) + articulation = _ArticulationGeometryProvider( + initial_qpos=initial_qpos, + backend_joint_names=backend_joint_names, + link_meshes=link_meshes, + link_poses=link_poses, + body_scale=body_scale, + parent_joint_chain=parent_joint_chain, + ) + return articulation, articulation.fk_calls + + +@pytest.mark.no_sim +class TestInitialArticulationGeometrySampling: + """Pure CPU coverage for the Atomic Action articulation geometry adapter.""" + + def test_passes_initial_joint_state_and_names_to_provider(self): + articulation, fk_calls = _make_point_cloud_articulation( + initial_qpos=(2.0, 1.0), + backend_joint_names=("joint_b", "joint_a"), + ) + + sample_initial_articulation_geometry( + articulation, + "target", + initial_qpos=articulation.initial_qpos, + initial_qpos_joint_names=articulation.initial_qpos_joint_names, + body_scale=articulation.body_scale, + articulation_point_count=5, + target_point_count=3, + ) + + assert len(fk_calls) == 1 + fk_qpos, fk_link_names, fk_qpos_joint_names = fk_calls[0] + assert torch.equal(fk_qpos, torch.tensor(((2.0, 1.0),))) + assert fk_link_names == ["target"] + assert fk_qpos_joint_names == ("joint_b", "joint_a") + + def test_uses_nearest_revolute_ancestor_after_fixed_descendants(self): + triangle = torch.tensor(((0, 1, 2),), dtype=torch.long) + link_meshes = { + link_name: ( + torch.tensor( + ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)), + dtype=torch.float32, + ), + triangle, + ) + for link_name in ("base", "near_parent", "moving", "target") + } + fixed_joint = ArticulationJointKinematics( + name="target_fixed", + joint_type="fixed", + parent_link_name="moving", + child_link_name="target", + origin_pose=torch.eye(4), + axis=torch.zeros(3), + ) + near_origin_pose = torch.eye(4) + near_origin_pose[:3, 3] = torch.tensor((1.0, 2.0, 3.0)) + near_joint = ArticulationJointKinematics( + name="near_hinge", + joint_type="revolute", + parent_link_name="near_parent", + child_link_name="moving", + origin_pose=near_origin_pose, + axis=torch.tensor((0.0, 0.0, 1.0)), + ) + far_origin_pose = torch.eye(4) + far_origin_pose[:3, 3] = torch.tensor((9.0, 9.0, 9.0)) + far_joint = ArticulationJointKinematics( + name="far_hinge", + joint_type="revolute", + parent_link_name="base", + child_link_name="near_parent", + origin_pose=far_origin_pose, + axis=torch.tensor((1.0, 0.0, 0.0)), + ) + articulation, _ = _make_point_cloud_articulation( + link_meshes=link_meshes, + parent_joint_chain=(fixed_joint, near_joint, far_joint), + ) + + geometry = sample_initial_articulation_geometry( + articulation, + "target", + initial_qpos=articulation.initial_qpos, + initial_qpos_joint_names=articulation.initial_qpos_joint_names, + body_scale=articulation.body_scale, + articulation_point_count=16, + target_point_count=8, + ) + + assert torch.equal( + geometry.revolute_axis_origin, + near_origin_pose[:3, 3], + ) + assert torch.equal( + geometry.revolute_joint_axis, + near_joint.axis, + ) + + def test_uses_nearest_prismatic_axis_after_fixed_descendant(self): + triangle = torch.tensor(((0, 1, 2),), dtype=torch.long) + link_meshes = { + link_name: ( + torch.tensor( + ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)), + dtype=torch.float32, + ), + triangle, + ) + for link_name in ("base", "near_parent", "moving", "target") + } + fixed_joint = ArticulationJointKinematics( + name="target_fixed", + joint_type="fixed", + parent_link_name="moving", + child_link_name="target", + origin_pose=torch.eye(4), + axis=torch.zeros(3), + ) + near_origin_pose = torch.eye(4) + near_origin_pose[:3, :3] = torch.tensor( + ((0.0, -1.0, 0.0), (1.0, 0.0, 0.0), (0.0, 0.0, 1.0)) + ) + near_joint = ArticulationJointKinematics( + name="near_slider", + joint_type="prismatic", + parent_link_name="near_parent", + child_link_name="moving", + origin_pose=near_origin_pose, + axis=torch.tensor((4.0, 0.0, 0.0)), + ) + far_joint = ArticulationJointKinematics( + name="far_slider", + joint_type="prismatic", + parent_link_name="base", + child_link_name="near_parent", + origin_pose=torch.eye(4), + axis=torch.tensor((0.0, 0.0, 2.0)), + ) + articulation, _ = _make_point_cloud_articulation( + link_meshes=link_meshes, + parent_joint_chain=(fixed_joint, near_joint, far_joint), + ) + + geometry = sample_initial_articulation_geometry( + articulation, + "target", + initial_qpos=articulation.initial_qpos, + initial_qpos_joint_names=articulation.initial_qpos_joint_names, + body_scale=articulation.body_scale, + articulation_point_count=16, + target_point_count=8, + ) + + assert torch.allclose( + geometry.prismatic_joint_axis, + torch.tensor((0.0, 1.0, 0.0)), + atol=POINT_CLOUD_TOLERANCE, + ) + assert geometry.revolute_joint_axis is None + assert geometry.revolute_axis_origin is None + + def test_collects_both_nearest_joint_types_from_one_parent_chain(self): + triangle = torch.tensor(((0, 1, 2),), dtype=torch.long) + link_meshes = { + link_name: ( + torch.tensor( + ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)), + dtype=torch.float32, + ), + triangle, + ) + for link_name in ("root", "base", "near_parent", "hinge_child", "target") + } + near_prismatic = ArticulationJointKinematics( + name="near_slider", + joint_type="prismatic", + parent_link_name="hinge_child", + child_link_name="target", + origin_pose=torch.eye(4), + axis=torch.tensor((4.0, 0.0, 0.0)), + ) + near_revolute_origin = torch.eye(4) + near_revolute_origin[:3, 3] = torch.tensor((1.0, 2.0, 3.0)) + near_revolute = ArticulationJointKinematics( + name="near_hinge", + joint_type="revolute", + parent_link_name="near_parent", + child_link_name="hinge_child", + origin_pose=near_revolute_origin, + axis=torch.tensor((0.0, 5.0, 0.0)), + ) + far_prismatic = ArticulationJointKinematics( + name="far_slider", + joint_type="prismatic", + parent_link_name="base", + child_link_name="near_parent", + origin_pose=torch.eye(4), + axis=torch.tensor((0.0, 0.0, 2.0)), + ) + far_revolute_origin = torch.eye(4) + far_revolute_origin[:3, 3] = torch.tensor((9.0, 9.0, 9.0)) + far_revolute = ArticulationJointKinematics( + name="far_hinge", + joint_type="revolute", + parent_link_name="root", + child_link_name="base", + origin_pose=far_revolute_origin, + axis=torch.tensor((0.0, 0.0, 7.0)), + ) + articulation, _ = _make_point_cloud_articulation( + link_meshes=link_meshes, + parent_joint_chain=( + near_prismatic, + near_revolute, + far_prismatic, + far_revolute, + ), + ) + + geometry = sample_initial_articulation_geometry( + articulation, + "target", + initial_qpos=articulation.initial_qpos, + initial_qpos_joint_names=articulation.initial_qpos_joint_names, + body_scale=articulation.body_scale, + articulation_point_count=16, + target_point_count=8, + ) + + assert torch.equal( + geometry.prismatic_joint_axis, + torch.tensor((1.0, 0.0, 0.0)), + ) + assert torch.equal( + geometry.revolute_joint_axis, + torch.tensor((0.0, 1.0, 0.0)), + ) + assert torch.equal( + geometry.revolute_axis_origin, + near_revolute_origin[:3, 3], + ) + + def test_transforms_revolute_origin_from_rotated_parent_to_target_frame(self): + triangle = torch.tensor(((0, 1, 2),), dtype=torch.long) + link_meshes = { + link_name: ( + torch.tensor( + ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)), + dtype=torch.float32, + ), + triangle, + ) + for link_name in ("parent", "target") + } + root_from_parent = torch.eye(4) + root_from_parent[:3, :3] = torch.tensor( + ((0.0, -1.0, 0.0), (1.0, 0.0, 0.0), (0.0, 0.0, 1.0)) + ) + root_from_parent[:3, 3] = torch.tensor((3.0, 4.0, 5.0)) + root_from_target = torch.eye(4) + root_from_target[:3, :3] = torch.tensor( + ((1.0, 0.0, 0.0), (0.0, 0.0, -1.0), (0.0, 1.0, 0.0)) + ) + root_from_target[:3, 3] = torch.tensor((0.5, -1.0, 2.0)) + joint_origin_pose = torch.eye(4) + joint_origin_pose[:3, :3] = torch.tensor( + ((0.0, 0.0, 1.0), (0.0, 1.0, 0.0), (-1.0, 0.0, 0.0)) + ) + joint_origin_pose[:3, 3] = torch.tensor((1.0, 2.0, 3.0)) + link_poses = torch.stack((root_from_parent, root_from_target)).unsqueeze(0) + joint = ArticulationJointKinematics( + name="target_hinge", + joint_type="revolute", + parent_link_name="parent", + child_link_name="target", + origin_pose=joint_origin_pose, + axis=torch.tensor((0.0, 4.0, 0.0)), + ) + articulation, _ = _make_point_cloud_articulation( + link_meshes=link_meshes, + link_poses=link_poses, + parent_joint_chain=(joint,), + ) + + geometry = sample_initial_articulation_geometry( + articulation, + "target", + initial_qpos=articulation.initial_qpos, + initial_qpos_joint_names=articulation.initial_qpos_joint_names, + body_scale=articulation.body_scale, + articulation_point_count=16, + target_point_count=8, + ) + + target_from_joint = ( + torch.linalg.inv(root_from_target) @ root_from_parent @ joint_origin_pose + ) + assert torch.allclose( + geometry.revolute_axis_origin, + target_from_joint[:3, 3], + atol=POINT_CLOUD_TOLERANCE, + ) + expected_axis = target_from_joint[:3, :3] @ joint.axis + expected_axis = expected_axis / torch.linalg.vector_norm(expected_axis) + assert torch.allclose( + geometry.revolute_joint_axis, + expected_axis, + atol=POINT_CLOUD_TOLERANCE, + ) + assert torch.linalg.vector_norm( + geometry.revolute_joint_axis + ).item() == pytest.approx(1.0) + + def test_omits_revolute_origin_when_parent_chain_has_only_fixed_joints(self): + fixed_joint = ArticulationJointKinematics( + name="target_fixed", + joint_type="fixed", + parent_link_name="parent", + child_link_name="target", + origin_pose=torch.eye(4), + axis=torch.zeros(3), + ) + articulation, _ = _make_point_cloud_articulation( + parent_joint_chain=(fixed_joint,), + ) + + geometry = sample_initial_articulation_geometry( + articulation, + "target", + initial_qpos=articulation.initial_qpos, + initial_qpos_joint_names=articulation.initial_qpos_joint_names, + body_scale=articulation.body_scale, + articulation_point_count=5, + target_point_count=3, + ) + + assert isinstance(geometry, ArticulationAffordanceGeometry) + first_object_geometry = geometry.to_object_geometry() + second_object_geometry = geometry.to_object_geometry() + assert first_object_geometry is not second_object_geometry + assert set(first_object_geometry) == { + "target_link_point_cloud", + "articulation_point_cloud", + } + assert geometry.prismatic_joint_axis is None + assert geometry.revolute_joint_axis is None + assert geometry.revolute_axis_origin is None + + @pytest.mark.parametrize("joint_type", ("prismatic", "revolute")) + def test_rejects_joint_axis_metadata_with_unknown_parent_link( + self, + joint_type: str, + ): + joint = ArticulationJointKinematics( + name="target_joint", + joint_type=joint_type, + parent_link_name="missing_parent", + child_link_name="target", + origin_pose=torch.eye(4), + axis=torch.tensor((0.0, 0.0, 1.0)), + ) + articulation, _ = _make_point_cloud_articulation( + parent_joint_chain=(joint,), + ) + + with pytest.raises(ValueError, match="parent link.*not an articulation link"): + sample_initial_articulation_geometry( + articulation, + "target", + initial_qpos=articulation.initial_qpos, + initial_qpos_joint_names=articulation.initial_qpos_joint_names, + body_scale=articulation.body_scale, + articulation_point_count=5, + target_point_count=3, + ) + + @pytest.mark.parametrize("joint_type", ("prismatic", "revolute")) + def test_rejects_joint_axis_metadata_with_invalid_origin_pose( + self, + joint_type: str, + ): + joint = SimpleNamespace( + name="target_joint", + joint_type=joint_type, + parent_link_name="target", + origin_pose=torch.full((4, 4), float("nan")), + axis=torch.tensor((0.0, 0.0, 1.0)), + ) + articulation, _ = _make_point_cloud_articulation( + parent_joint_chain=(joint,), + ) + + with pytest.raises(ValueError, match="origin pose must be finite"): + sample_initial_articulation_geometry( + articulation, + "target", + initial_qpos=articulation.initial_qpos, + initial_qpos_joint_names=articulation.initial_qpos_joint_names, + body_scale=articulation.body_scale, + articulation_point_count=5, + target_point_count=3, + ) + + @pytest.mark.parametrize( + ("joint_type", "axis"), + ( + ("prismatic", torch.zeros(3)), + ("revolute", torch.tensor((float("nan"), 0.0, 1.0))), + ("prismatic", torch.zeros(2)), + ), + ) + def test_rejects_invalid_joint_axes( + self, + joint_type: str, + axis: torch.Tensor, + ): + joint = SimpleNamespace( + name="target_joint", + joint_type=joint_type, + parent_link_name="target", + origin_pose=torch.eye(4), + axis=axis, + ) + articulation, _ = _make_point_cloud_articulation( + parent_joint_chain=(joint,), + ) + + with pytest.raises(ValueError, match="axis must be finite and nonzero"): + sample_initial_articulation_geometry( + articulation, + "target", + initial_qpos=articulation.initial_qpos, + initial_qpos_joint_names=articulation.initial_qpos_joint_names, + body_scale=articulation.body_scale, + articulation_point_count=5, + target_point_count=3, + ) + + def test_returns_both_clouds_in_target_link_initial_frame(self): + triangle = torch.tensor(((0, 1, 2),), dtype=torch.long) + link_meshes = { + "body": ( + torch.tensor( + ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)), + dtype=torch.float32, + ), + triangle, + ), + "target": ( + torch.tensor( + ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)), + dtype=torch.float32, + ), + triangle, + ), + } + initial_link_poses = torch.eye(4, dtype=torch.float32).repeat(1, 2, 1, 1) + initial_link_poses[0, 0, :3, 3] = torch.tensor((4.0, 0.0, 0.0)) + initial_link_poses[0, 1, :3, :3] = torch.tensor( + ((0.0, -1.0, 0.0), (1.0, 0.0, 0.0), (0.0, 0.0, 1.0)) + ) + initial_link_poses[0, 1, :3, 3] = torch.tensor((1.0, 2.0, 0.0)) + articulation, _ = _make_point_cloud_articulation( + link_meshes=link_meshes, + link_poses=initial_link_poses, + ) + + geometry = sample_initial_articulation_geometry( + articulation, + "target", + initial_qpos=articulation.initial_qpos, + initial_qpos_joint_names=articulation.initial_qpos_joint_names, + body_scale=articulation.body_scale, + articulation_point_count=256, + target_point_count=32, + ) + + assert set(geometry.to_object_geometry()) == { + "target_link_point_cloud", + "articulation_point_cloud", + } + target_points = geometry.target_link_point_cloud + assert target_points.shape == (32, 3) + assert torch.allclose( + target_points[:, 2], + torch.zeros(32), + atol=POINT_CLOUD_TOLERANCE, + ) + assert bool((target_points[:, :2] >= -POINT_CLOUD_TOLERANCE).all()) + assert bool( + ( + target_points[:, 0] + target_points[:, 1] <= 1.0 + POINT_CLOUD_TOLERANCE + ).all() + ) + + articulation_points = geometry.articulation_point_cloud + assert articulation_points.shape == (256, 3) + body_mask = articulation_points[:, 0] < -0.5 + assert bool(body_mask.any()) + assert bool((~body_mask).any()) + body_points = articulation_points[body_mask] + assert bool( + ( + (body_points[:, 0] >= -2.0 - POINT_CLOUD_TOLERANCE) + & (body_points[:, 0] <= -1.0 + POINT_CLOUD_TOLERANCE) + & (body_points[:, 1] >= -4.0 - POINT_CLOUD_TOLERANCE) + & (body_points[:, 1] <= -3.0 + POINT_CLOUD_TOLERANCE) + ).all() + ) + sampled_target_points = articulation_points[~body_mask] + assert bool((sampled_target_points[:, :2] >= -POINT_CLOUD_TOLERANCE).all()) + assert bool( + ( + sampled_target_points[:, 0] + sampled_target_points[:, 1] + <= 1.0 + POINT_CLOUD_TOLERANCE + ).all() + ) + + def test_sampling_consumes_open3d_rng_without_resetting_it(self): + import open3d as o3d + + articulation, _ = _make_point_cloud_articulation() + o3d.utility.random.seed(7) + + first = sample_initial_articulation_geometry( + articulation, + "target", + initial_qpos=articulation.initial_qpos, + initial_qpos_joint_names=articulation.initial_qpos_joint_names, + body_scale=articulation.body_scale, + articulation_point_count=23, + target_point_count=17, + ) + second = sample_initial_articulation_geometry( + articulation, + "target", + initial_qpos=articulation.initial_qpos, + initial_qpos_joint_names=articulation.initial_qpos_joint_names, + body_scale=articulation.body_scale, + articulation_point_count=23, + target_point_count=17, + ) + + assert not torch.equal( + first.target_link_point_cloud, + second.target_link_point_cloud, + ) + assert not torch.equal( + first.articulation_point_cloud, + second.articulation_point_cloud, + ) + + def test_surface_sampling_preserves_torch_dtype_and_owns_its_data(self): + vertices = torch.tensor( + ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)), + dtype=torch.float64, + ) + triangles = torch.tensor(((0, 1, 2),), dtype=torch.long) + original_vertices = vertices.clone() + + points = articulation_geometry_module._sample_mesh_surface_points( + vertices, triangles, 11 + ) + + assert points.device == vertices.device + assert points.dtype == vertices.dtype + assert points.shape == (11, 3) + points.zero_() + assert torch.equal(vertices, original_vertices) + + @pytest.mark.parametrize( + "triangles", + ( + torch.empty((0, 3), dtype=torch.long), + torch.tensor(((0, 1, 2),), dtype=torch.long), + ), + ) + def test_surface_sampling_falls_back_to_vertices_without_valid_faces( + self, + triangles: torch.Tensor, + ): + vertices = torch.tensor( + ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (2.0, 0.0, 0.0)), + dtype=torch.float32, + ) + + points = articulation_geometry_module._sample_mesh_surface_points( + vertices, triangles, 7 + ) + + assert points.shape == (7, 3) + assert set(points[:, 0].tolist()) <= {0.0, 1.0, 2.0} + + def test_merged_surface_sampling_is_weighted_by_face_area(self): + triangle = torch.tensor(((0, 1, 2),), dtype=torch.long) + articulation, _ = _make_point_cloud_articulation( + link_meshes={ + "small": ( + torch.tensor( + ((-10.0, 0.0, 0.0), (-9.0, 0.0, 0.0), (-10.0, 1.0, 0.0)), + dtype=torch.float32, + ), + triangle, + ), + "target": ( + torch.tensor( + ((0.0, 0.0, 0.0), (2.0, 0.0, 0.0), (0.0, 2.0, 0.0)), + dtype=torch.float32, + ), + triangle, + ), + } + ) + + geometry = sample_initial_articulation_geometry( + articulation, + "target", + initial_qpos=articulation.initial_qpos, + initial_qpos_joint_names=articulation.initial_qpos_joint_names, + body_scale=articulation.body_scale, + articulation_point_count=5_000, + target_point_count=3, + ) + + articulation_points = geometry.articulation_point_cloud + small_face_fraction = float((articulation_points[:, 0] < -9.0).float().mean()) + assert small_face_fraction == pytest.approx(0.2, abs=0.03) + + @pytest.mark.parametrize( + ("target_link_name", "error_type", "message"), + ( + (None, TypeError, "target_link_name must be a string"), + (" ", ValueError, "target_link_name must be non-empty"), + ("missing", ValueError, "Unknown articulation link"), + ), + ) + def test_rejects_invalid_target_link_names( + self, + target_link_name: object, + error_type: type[Exception], + message: str, + ): + articulation, _ = _make_point_cloud_articulation() + + with pytest.raises(error_type, match=message): + sample_initial_articulation_geometry( + articulation, + target_link_name, # type: ignore[arg-type] + initial_qpos=articulation.initial_qpos, + initial_qpos_joint_names=articulation.initial_qpos_joint_names, + body_scale=articulation.body_scale, + ) + + @pytest.mark.parametrize( + ("field_name", "value", "error_type", "message"), + ( + ( + "articulation_point_count", + True, + TypeError, + "articulation_point_count must be an integer", + ), + ( + "target_point_count", + 0, + ValueError, + "target_point_count must be positive", + ), + ), + ) + def test_rejects_invalid_point_counts( + self, + field_name: str, + value: object, + error_type: type[Exception], + message: str, + ): + articulation, _ = _make_point_cloud_articulation() + + with pytest.raises(error_type, match=message): + sample_initial_articulation_geometry( + articulation, + "target", + initial_qpos=articulation.initial_qpos, + initial_qpos_joint_names=articulation.initial_qpos_joint_names, + body_scale=articulation.body_scale, + **{field_name: value}, # type: ignore[arg-type] + ) + + def test_rejects_non_unit_body_scale(self): + articulation, _ = _make_point_cloud_articulation( + body_scale=(1.0, 2.0, 1.0), + ) + + with pytest.raises(ValueError, match="requires unit body_scale"): + sample_initial_articulation_geometry( + articulation, + "target", + initial_qpos=articulation.initial_qpos, + initial_qpos_joint_names=articulation.initial_qpos_joint_names, + body_scale=articulation.body_scale, + ) + + def test_rejects_initial_joint_name_count_mismatch(self): + articulation, _ = _make_point_cloud_articulation( + initial_qpos=(0.1, 0.2), + backend_joint_names=("joint",), + ) + + with pytest.raises(ValueError, match="matching its joint-name sequence"): + sample_initial_articulation_geometry( + articulation, + "target", + initial_qpos=articulation.initial_qpos, + initial_qpos_joint_names=articulation.initial_qpos_joint_names, + body_scale=articulation.body_scale, + ) + + @pytest.mark.parametrize( + "initial_qpos", + ( + (), + (float("nan"),), + ), + ) + def test_rejects_invalid_initial_qpos(self, initial_qpos: tuple[float, ...]): + articulation, _ = _make_point_cloud_articulation( + initial_qpos=initial_qpos, + ) + + with pytest.raises(ValueError, match="finite vector matching"): + sample_initial_articulation_geometry( + articulation, + "target", + initial_qpos=articulation.initial_qpos, + initial_qpos_joint_names=articulation.initial_qpos_joint_names, + body_scale=articulation.body_scale, + ) + + def test_rejects_invalid_mesh_indices(self): + articulation, _ = _make_point_cloud_articulation( + link_meshes={ + "target": ( + torch.zeros((3, 3), dtype=torch.float32), + torch.tensor(((0, 1, 3),), dtype=torch.long), + ) + } + ) + + with pytest.raises(ValueError, match="triangles reference invalid vertices"): + sample_initial_articulation_geometry( + articulation, + "target", + initial_qpos=articulation.initial_qpos, + initial_qpos_joint_names=articulation.initial_qpos_joint_names, + body_scale=articulation.body_scale, + ) + + @pytest.mark.parametrize( + "triangles", + ( + torch.tensor(((0.0, 1.0, 2.0),), dtype=torch.float32), + torch.tensor(((False, True, True),), dtype=torch.bool), + ), + ) + def test_rejects_non_integer_mesh_indices(self, triangles: torch.Tensor): + articulation, _ = _make_point_cloud_articulation( + link_meshes={ + "target": ( + torch.zeros((3, 3), dtype=torch.float32), + triangles, + ) + } + ) + + with pytest.raises(ValueError, match="triangles must use integer indices"): + sample_initial_articulation_geometry( + articulation, + "target", + initial_qpos=articulation.initial_qpos, + initial_qpos_joint_names=articulation.initial_qpos_joint_names, + body_scale=articulation.body_scale, + ) diff --git a/tests/sim/atomic_actions/test_tutorial_utils.py b/tests/sim/atomic_actions/test_tutorial_utils.py index ffe2c8a68..9a97ae1a4 100644 --- a/tests/sim/atomic_actions/test_tutorial_utils.py +++ b/tests/sim/atomic_actions/test_tutorial_utils.py @@ -30,7 +30,10 @@ import pytest import torch -from embodichain.lab.sim.atomic_actions import TimedTrajectory +from embodichain.lab.sim.atomic_actions import ( + ArticulationAffordanceGeometry, + TimedTrajectory, +) from embodichain.lab.sim.solvers import BaseSolver from scripts.tutorials.atomic_action.dynamic_obstacle_recovery import ( _animate_obstacle_to_pose, @@ -164,13 +167,30 @@ def _tutorial_axis_geometry( } +def _tutorial_articulation_geometry( + geometry: dict[str, torch.Tensor], +) -> ArticulationAffordanceGeometry: + """Build the typed adapter result consumed by tutorial semantics.""" + return ArticulationAffordanceGeometry( + target_link_point_cloud=geometry["target_link_point_cloud"], + articulation_point_cloud=geometry["articulation_point_cloud"], + prismatic_joint_axis=geometry["target_link_prismatic_joint_axis"], + revolute_joint_axis=geometry["target_link_revolute_joint_axis"], + revolute_axis_origin=geometry["target_link_revolute_axis_origin"], + ) + + class _TutorialArticulation: """Minimal articulation surface used by automatic-axis tutorial tests.""" def __init__(self, geometry: dict[str, torch.Tensor]) -> None: self.device = torch.device("cpu") self.geometry = geometry - self.sampled_link_name: str | None = None + self.cfg = SimpleNamespace( + init_qpos=(0.0,), + body_scale=(1.0, 1.0, 1.0), + ) + self.joint_names = ("target_joint",) def get_link_vert_face(self, link_name: str) -> tuple[torch.Tensor, torch.Tensor]: del link_name @@ -195,13 +215,6 @@ def get_link_pose(self, link_name: str, *, to_matrix: bool) -> torch.Tensor: assert to_matrix is True return torch.eye(4).unsqueeze(0) - def sample_initial_point_clouds( - self, - target_link_name: str, - ) -> dict[str, torch.Tensor]: - self.sampled_link_name = target_link_name - return self.geometry - def _run_obstacle_animation(*, pace_wall_time: bool) -> tuple[MagicMock, MagicMock]: obstacle = MagicMock() @@ -272,14 +285,31 @@ def test_articulation_tutorial_semantics_resolve_signed_parent_joint_axis( ) -> None: module = importlib.import_module(f"scripts.tutorials.atomic_action.{module_name}") geometry = _tutorial_axis_geometry(neighbor_offset) + sampled_geometry = _tutorial_articulation_geometry(geometry) articulation = _TutorialArticulation(geometry) - with patch.object(module, "Articulation", _TutorialArticulation): + with ( + patch.object(module, "Articulation", _TutorialArticulation), + patch.object( + module, + "sample_initial_articulation_geometry", + return_value=sampled_geometry, + ) as sampler, + ): result = getattr(module, factory_name)(articulation) semantics = result[0] if isinstance(result, tuple) else result - assert articulation.sampled_link_name == link_name - assert semantics.geometry is geometry + sampler.assert_called_once_with( + articulation, + link_name, + initial_qpos=articulation.cfg.init_qpos, + initial_qpos_joint_names=articulation.joint_names, + body_scale=articulation.cfg.body_scale, + ) + expected_geometry = sampled_geometry.to_object_geometry() + assert set(semantics.geometry) == set(expected_geometry) + for key, expected_value in expected_geometry.items(): + assert torch.equal(semantics.geometry[key], expected_value) assert torch.allclose( getattr(semantics.affordance, axis_field), torch.tensor(expected_axis), diff --git a/tests/sim/objects/test_articulation.py b/tests/sim/objects/test_articulation.py index bc018bc6b..f4e51a6ae 100644 --- a/tests/sim/objects/test_articulation.py +++ b/tests/sim/objects/test_articulation.py @@ -41,76 +41,6 @@ ART_PATH = "SlidingBoxDrawer/SlidingBoxDrawer.urdf" NUM_ARENAS = 10 -POINT_CLOUD_TOLERANCE = 1.0e-6 - - -def _make_point_cloud_articulation( - *, - initial_qpos: tuple[float, ...] = (0.25,), - backend_joint_names: tuple[str, ...] = ("joint",), - kinematic_joint_names: tuple[str, ...] | None = None, - link_meshes: dict[str, tuple[torch.Tensor, torch.Tensor]] | None = None, - link_poses: torch.Tensor | None = None, - body_scale: tuple[float, float, float] = (1.0, 1.0, 1.0), - parent_joint_chain: tuple[object, ...] = (), -) -> tuple[Articulation, list[tuple[torch.Tensor, list[str]]]]: - """Build a pure-Python articulation double for point-cloud sampling.""" - if link_meshes is None: - link_meshes = { - "target": ( - torch.tensor( - ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)), - dtype=torch.float32, - ), - torch.tensor(((0, 1, 2),), dtype=torch.long), - ) - } - link_names = list(link_meshes) - if link_poses is None: - link_poses = torch.eye(4, dtype=torch.float32).repeat(1, len(link_names), 1, 1) - if kinematic_joint_names is None: - kinematic_joint_names = backend_joint_names - - articulation = object.__new__(Articulation) - articulation.device = torch.device("cpu") - articulation.cfg = SimpleNamespace( - init_qpos=initial_qpos, - body_scale=body_scale, - ) - articulation._data = SimpleNamespace( - dof=len(backend_joint_names), - link_names=link_names, - link_vert_face=link_meshes, - ) - articulation._entities = [ - SimpleNamespace( - get_actived_joint_names=lambda: list(backend_joint_names), - ) - ] - articulation.pk_chain = SimpleNamespace( - get_joint_parameter_names=lambda: list(kinematic_joint_names), - ) - fk_calls: list[tuple[torch.Tensor, list[str]]] = [] - - def compute_fk( - qpos: torch.Tensor, - *, - link_names: list[str], - ) -> torch.Tensor: - fk_calls.append((qpos.clone(), list(link_names))) - return link_poses.clone() - - def get_parent_joint_chain( - link_name: str, - ) -> tuple[object, ...]: - assert link_name in link_names - return parent_joint_chain - - articulation.compute_fk = compute_fk # type: ignore[method-assign] - articulation.get_parent_joint_chain = ( # type: ignore[method-assign] - get_parent_joint_chain - ) - return articulation, fk_calls def test_get_qf_returns_all_articulation_joint_efforts(): @@ -123,6 +53,52 @@ def test_get_qf_returns_all_articulation_joint_efforts(): assert torch.equal(actual_qf, expected_qf) +@pytest.mark.no_sim +def test_compute_fk_reorders_named_qpos_into_kinematic_joint_order(): + articulation = object.__new__(Articulation) + captured: dict[str, object] = {} + expected_result = {"target": object()} + + def forward_kinematics( + *, + th: torch.Tensor, + frame_indices: object, + ) -> dict[str, object]: + captured["qpos"] = th.clone() + captured["frame_indices"] = frame_indices + return expected_result + + articulation.pk_chain = SimpleNamespace( + get_joint_parameter_names=lambda: ["joint_a", "joint_b"], + forward_kinematics=forward_kinematics, + ) + + result = articulation.compute_fk( + torch.tensor(((2.0, 1.0),)), + qpos_joint_names=("joint_b", "joint_a"), + to_dict=True, + ) + + assert result is expected_result + assert torch.equal(captured["qpos"], torch.tensor(((1.0, 2.0),))) + assert captured["frame_indices"] is None + + +@pytest.mark.no_sim +def test_compute_fk_rejects_named_qpos_for_serial_chain(): + articulation = object.__new__(Articulation) + articulation.pk_chain = SimpleNamespace( + get_joint_parameter_names=lambda: ["joint"], + ) + + with pytest.raises(ValueError, match="cannot be combined with serial-chain FK"): + articulation.compute_fk( + torch.tensor((0.0,)), + end_link_name="target", + qpos_joint_names=("joint",), + ) + + @pytest.mark.no_sim def test_get_parent_joint_chain_returns_backend_neutral_child_to_root_values(): articulation = object.__new__(Articulation) @@ -166,667 +142,6 @@ def test_get_parent_joint_chain_returns_backend_neutral_child_to_root_values(): assert chain[0].origin_pose[0, 3].item() == 0.0 -@pytest.mark.no_sim -class TestInitialPointCloudSampling: - """Pure CPU coverage for initial articulation surface sampling.""" - - def test_reorders_initial_qpos_into_kinematic_joint_order(self): - articulation, fk_calls = _make_point_cloud_articulation( - initial_qpos=(2.0, 1.0), - backend_joint_names=("joint_b", "joint_a"), - kinematic_joint_names=("joint_a", "joint_b"), - ) - - articulation.sample_initial_point_clouds( - "target", - articulation_point_count=5, - target_point_count=3, - ) - - assert len(fk_calls) == 1 - fk_qpos, fk_link_names = fk_calls[0] - assert torch.equal(fk_qpos, torch.tensor(((1.0, 2.0),))) - assert fk_link_names == ["target"] - - def test_uses_nearest_revolute_ancestor_after_fixed_descendants(self): - triangle = torch.tensor(((0, 1, 2),), dtype=torch.long) - link_meshes = { - link_name: ( - torch.tensor( - ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)), - dtype=torch.float32, - ), - triangle, - ) - for link_name in ("base", "near_parent", "moving", "target") - } - fixed_joint = ArticulationJointKinematics( - name="target_fixed", - joint_type="fixed", - parent_link_name="moving", - child_link_name="target", - origin_pose=torch.eye(4), - axis=torch.zeros(3), - ) - near_origin_pose = torch.eye(4) - near_origin_pose[:3, 3] = torch.tensor((1.0, 2.0, 3.0)) - near_joint = ArticulationJointKinematics( - name="near_hinge", - joint_type="revolute", - parent_link_name="near_parent", - child_link_name="moving", - origin_pose=near_origin_pose, - axis=torch.tensor((0.0, 0.0, 1.0)), - ) - far_origin_pose = torch.eye(4) - far_origin_pose[:3, 3] = torch.tensor((9.0, 9.0, 9.0)) - far_joint = ArticulationJointKinematics( - name="far_hinge", - joint_type="revolute", - parent_link_name="base", - child_link_name="near_parent", - origin_pose=far_origin_pose, - axis=torch.tensor((1.0, 0.0, 0.0)), - ) - articulation, _ = _make_point_cloud_articulation( - link_meshes=link_meshes, - parent_joint_chain=(fixed_joint, near_joint, far_joint), - ) - - geometry = articulation.sample_initial_point_clouds( - "target", - articulation_point_count=16, - target_point_count=8, - ) - - assert torch.equal( - geometry["target_link_revolute_axis_origin"], - near_origin_pose[:3, 3], - ) - assert torch.equal( - geometry["target_link_revolute_joint_axis"], - near_joint.axis, - ) - - def test_uses_nearest_prismatic_axis_after_fixed_descendant(self): - triangle = torch.tensor(((0, 1, 2),), dtype=torch.long) - link_meshes = { - link_name: ( - torch.tensor( - ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)), - dtype=torch.float32, - ), - triangle, - ) - for link_name in ("base", "near_parent", "moving", "target") - } - fixed_joint = ArticulationJointKinematics( - name="target_fixed", - joint_type="fixed", - parent_link_name="moving", - child_link_name="target", - origin_pose=torch.eye(4), - axis=torch.zeros(3), - ) - near_origin_pose = torch.eye(4) - near_origin_pose[:3, :3] = torch.tensor( - ((0.0, -1.0, 0.0), (1.0, 0.0, 0.0), (0.0, 0.0, 1.0)) - ) - near_joint = ArticulationJointKinematics( - name="near_slider", - joint_type="prismatic", - parent_link_name="near_parent", - child_link_name="moving", - origin_pose=near_origin_pose, - axis=torch.tensor((4.0, 0.0, 0.0)), - ) - far_joint = ArticulationJointKinematics( - name="far_slider", - joint_type="prismatic", - parent_link_name="base", - child_link_name="near_parent", - origin_pose=torch.eye(4), - axis=torch.tensor((0.0, 0.0, 2.0)), - ) - articulation, _ = _make_point_cloud_articulation( - link_meshes=link_meshes, - parent_joint_chain=(fixed_joint, near_joint, far_joint), - ) - - geometry = articulation.sample_initial_point_clouds( - "target", - articulation_point_count=16, - target_point_count=8, - ) - - assert torch.allclose( - geometry["target_link_prismatic_joint_axis"], - torch.tensor((0.0, 1.0, 0.0)), - atol=POINT_CLOUD_TOLERANCE, - ) - assert "target_link_revolute_joint_axis" not in geometry - assert "target_link_revolute_axis_origin" not in geometry - - def test_collects_both_nearest_joint_types_from_one_parent_chain(self): - triangle = torch.tensor(((0, 1, 2),), dtype=torch.long) - link_meshes = { - link_name: ( - torch.tensor( - ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)), - dtype=torch.float32, - ), - triangle, - ) - for link_name in ("root", "base", "near_parent", "hinge_child", "target") - } - near_prismatic = ArticulationJointKinematics( - name="near_slider", - joint_type="prismatic", - parent_link_name="hinge_child", - child_link_name="target", - origin_pose=torch.eye(4), - axis=torch.tensor((4.0, 0.0, 0.0)), - ) - near_revolute_origin = torch.eye(4) - near_revolute_origin[:3, 3] = torch.tensor((1.0, 2.0, 3.0)) - near_revolute = ArticulationJointKinematics( - name="near_hinge", - joint_type="revolute", - parent_link_name="near_parent", - child_link_name="hinge_child", - origin_pose=near_revolute_origin, - axis=torch.tensor((0.0, 5.0, 0.0)), - ) - far_prismatic = ArticulationJointKinematics( - name="far_slider", - joint_type="prismatic", - parent_link_name="base", - child_link_name="near_parent", - origin_pose=torch.eye(4), - axis=torch.tensor((0.0, 0.0, 2.0)), - ) - far_revolute_origin = torch.eye(4) - far_revolute_origin[:3, 3] = torch.tensor((9.0, 9.0, 9.0)) - far_revolute = ArticulationJointKinematics( - name="far_hinge", - joint_type="revolute", - parent_link_name="root", - child_link_name="base", - origin_pose=far_revolute_origin, - axis=torch.tensor((0.0, 0.0, 7.0)), - ) - articulation, _ = _make_point_cloud_articulation( - link_meshes=link_meshes, - parent_joint_chain=( - near_prismatic, - near_revolute, - far_prismatic, - far_revolute, - ), - ) - - geometry = articulation.sample_initial_point_clouds( - "target", - articulation_point_count=16, - target_point_count=8, - ) - - assert torch.equal( - geometry["target_link_prismatic_joint_axis"], - torch.tensor((1.0, 0.0, 0.0)), - ) - assert torch.equal( - geometry["target_link_revolute_joint_axis"], - torch.tensor((0.0, 1.0, 0.0)), - ) - assert torch.equal( - geometry["target_link_revolute_axis_origin"], - near_revolute_origin[:3, 3], - ) - - def test_transforms_revolute_origin_from_rotated_parent_to_target_frame(self): - triangle = torch.tensor(((0, 1, 2),), dtype=torch.long) - link_meshes = { - link_name: ( - torch.tensor( - ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)), - dtype=torch.float32, - ), - triangle, - ) - for link_name in ("parent", "target") - } - root_from_parent = torch.eye(4) - root_from_parent[:3, :3] = torch.tensor( - ((0.0, -1.0, 0.0), (1.0, 0.0, 0.0), (0.0, 0.0, 1.0)) - ) - root_from_parent[:3, 3] = torch.tensor((3.0, 4.0, 5.0)) - root_from_target = torch.eye(4) - root_from_target[:3, :3] = torch.tensor( - ((1.0, 0.0, 0.0), (0.0, 0.0, -1.0), (0.0, 1.0, 0.0)) - ) - root_from_target[:3, 3] = torch.tensor((0.5, -1.0, 2.0)) - joint_origin_pose = torch.eye(4) - joint_origin_pose[:3, :3] = torch.tensor( - ((0.0, 0.0, 1.0), (0.0, 1.0, 0.0), (-1.0, 0.0, 0.0)) - ) - joint_origin_pose[:3, 3] = torch.tensor((1.0, 2.0, 3.0)) - link_poses = torch.stack((root_from_parent, root_from_target)).unsqueeze(0) - joint = ArticulationJointKinematics( - name="target_hinge", - joint_type="revolute", - parent_link_name="parent", - child_link_name="target", - origin_pose=joint_origin_pose, - axis=torch.tensor((0.0, 4.0, 0.0)), - ) - articulation, _ = _make_point_cloud_articulation( - link_meshes=link_meshes, - link_poses=link_poses, - parent_joint_chain=(joint,), - ) - - geometry = articulation.sample_initial_point_clouds( - "target", - articulation_point_count=16, - target_point_count=8, - ) - - target_from_joint = ( - torch.linalg.inv(root_from_target) @ root_from_parent @ joint_origin_pose - ) - assert torch.allclose( - geometry["target_link_revolute_axis_origin"], - target_from_joint[:3, 3], - atol=POINT_CLOUD_TOLERANCE, - ) - expected_axis = target_from_joint[:3, :3] @ joint.axis - expected_axis = expected_axis / torch.linalg.vector_norm(expected_axis) - assert torch.allclose( - geometry["target_link_revolute_joint_axis"], - expected_axis, - atol=POINT_CLOUD_TOLERANCE, - ) - assert torch.linalg.vector_norm( - geometry["target_link_revolute_joint_axis"] - ).item() == pytest.approx(1.0) - - def test_omits_revolute_origin_when_parent_chain_has_only_fixed_joints(self): - fixed_joint = ArticulationJointKinematics( - name="target_fixed", - joint_type="fixed", - parent_link_name="parent", - child_link_name="target", - origin_pose=torch.eye(4), - axis=torch.zeros(3), - ) - articulation, _ = _make_point_cloud_articulation( - parent_joint_chain=(fixed_joint,), - ) - - geometry = articulation.sample_initial_point_clouds( - "target", - articulation_point_count=5, - target_point_count=3, - ) - - assert set(geometry) == { - "target_link_point_cloud", - "articulation_point_cloud", - } - - @pytest.mark.parametrize("joint_type", ("prismatic", "revolute")) - def test_rejects_joint_axis_metadata_with_unknown_parent_link( - self, - joint_type: str, - ): - joint = ArticulationJointKinematics( - name="target_joint", - joint_type=joint_type, - parent_link_name="missing_parent", - child_link_name="target", - origin_pose=torch.eye(4), - axis=torch.tensor((0.0, 0.0, 1.0)), - ) - articulation, _ = _make_point_cloud_articulation( - parent_joint_chain=(joint,), - ) - - with pytest.raises(ValueError, match="parent link.*not an articulation link"): - articulation.sample_initial_point_clouds( - "target", - articulation_point_count=5, - target_point_count=3, - ) - - @pytest.mark.parametrize("joint_type", ("prismatic", "revolute")) - def test_rejects_joint_axis_metadata_with_invalid_origin_pose( - self, - joint_type: str, - ): - joint = SimpleNamespace( - name="target_joint", - joint_type=joint_type, - parent_link_name="target", - origin_pose=torch.full((4, 4), float("nan")), - axis=torch.tensor((0.0, 0.0, 1.0)), - ) - articulation, _ = _make_point_cloud_articulation( - parent_joint_chain=(joint,), - ) - - with pytest.raises(ValueError, match="origin pose must be finite"): - articulation.sample_initial_point_clouds( - "target", - articulation_point_count=5, - target_point_count=3, - ) - - @pytest.mark.parametrize( - ("joint_type", "axis"), - ( - ("prismatic", torch.zeros(3)), - ("revolute", torch.tensor((float("nan"), 0.0, 1.0))), - ("prismatic", torch.zeros(2)), - ), - ) - def test_rejects_invalid_joint_axes( - self, - joint_type: str, - axis: torch.Tensor, - ): - joint = SimpleNamespace( - name="target_joint", - joint_type=joint_type, - parent_link_name="target", - origin_pose=torch.eye(4), - axis=axis, - ) - articulation, _ = _make_point_cloud_articulation( - parent_joint_chain=(joint,), - ) - - with pytest.raises(ValueError, match="axis must be finite and nonzero"): - articulation.sample_initial_point_clouds( - "target", - articulation_point_count=5, - target_point_count=3, - ) - - def test_returns_both_clouds_in_target_link_initial_frame(self): - triangle = torch.tensor(((0, 1, 2),), dtype=torch.long) - link_meshes = { - "body": ( - torch.tensor( - ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)), - dtype=torch.float32, - ), - triangle, - ), - "target": ( - torch.tensor( - ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)), - dtype=torch.float32, - ), - triangle, - ), - } - initial_link_poses = torch.eye(4, dtype=torch.float32).repeat(1, 2, 1, 1) - initial_link_poses[0, 0, :3, 3] = torch.tensor((4.0, 0.0, 0.0)) - initial_link_poses[0, 1, :3, :3] = torch.tensor( - ((0.0, -1.0, 0.0), (1.0, 0.0, 0.0), (0.0, 0.0, 1.0)) - ) - initial_link_poses[0, 1, :3, 3] = torch.tensor((1.0, 2.0, 0.0)) - articulation, _ = _make_point_cloud_articulation( - link_meshes=link_meshes, - link_poses=initial_link_poses, - ) - - geometry = articulation.sample_initial_point_clouds( - "target", - articulation_point_count=256, - target_point_count=32, - ) - - assert set(geometry) == { - "target_link_point_cloud", - "articulation_point_cloud", - } - target_points = geometry["target_link_point_cloud"] - assert target_points.shape == (32, 3) - assert torch.allclose( - target_points[:, 2], - torch.zeros(32), - atol=POINT_CLOUD_TOLERANCE, - ) - assert bool((target_points[:, :2] >= -POINT_CLOUD_TOLERANCE).all()) - assert bool( - ( - target_points[:, 0] + target_points[:, 1] <= 1.0 + POINT_CLOUD_TOLERANCE - ).all() - ) - - articulation_points = geometry["articulation_point_cloud"] - assert articulation_points.shape == (256, 3) - body_mask = articulation_points[:, 0] < -0.5 - assert bool(body_mask.any()) - assert bool((~body_mask).any()) - body_points = articulation_points[body_mask] - assert bool( - ( - (body_points[:, 0] >= -2.0 - POINT_CLOUD_TOLERANCE) - & (body_points[:, 0] <= -1.0 + POINT_CLOUD_TOLERANCE) - & (body_points[:, 1] >= -4.0 - POINT_CLOUD_TOLERANCE) - & (body_points[:, 1] <= -3.0 + POINT_CLOUD_TOLERANCE) - ).all() - ) - sampled_target_points = articulation_points[~body_mask] - assert bool((sampled_target_points[:, :2] >= -POINT_CLOUD_TOLERANCE).all()) - assert bool( - ( - sampled_target_points[:, 0] + sampled_target_points[:, 1] - <= 1.0 + POINT_CLOUD_TOLERANCE - ).all() - ) - - def test_sampling_consumes_open3d_rng_without_resetting_it(self): - import open3d as o3d - - articulation, _ = _make_point_cloud_articulation() - o3d.utility.random.seed(7) - - first = articulation.sample_initial_point_clouds( - "target", - articulation_point_count=23, - target_point_count=17, - ) - second = articulation.sample_initial_point_clouds( - "target", - articulation_point_count=23, - target_point_count=17, - ) - - assert not torch.equal( - first["target_link_point_cloud"], - second["target_link_point_cloud"], - ) - assert not torch.equal( - first["articulation_point_cloud"], - second["articulation_point_cloud"], - ) - - def test_surface_sampling_preserves_torch_dtype_and_owns_its_data(self): - vertices = torch.tensor( - ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)), - dtype=torch.float64, - ) - triangles = torch.tensor(((0, 1, 2),), dtype=torch.long) - original_vertices = vertices.clone() - - points = Articulation._sample_mesh_surface_points(vertices, triangles, 11) - - assert points.device == vertices.device - assert points.dtype == vertices.dtype - assert points.shape == (11, 3) - points.zero_() - assert torch.equal(vertices, original_vertices) - - @pytest.mark.parametrize( - "triangles", - ( - torch.empty((0, 3), dtype=torch.long), - torch.tensor(((0, 1, 2),), dtype=torch.long), - ), - ) - def test_surface_sampling_falls_back_to_vertices_without_valid_faces( - self, - triangles: torch.Tensor, - ): - vertices = torch.tensor( - ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (2.0, 0.0, 0.0)), - dtype=torch.float32, - ) - - points = Articulation._sample_mesh_surface_points(vertices, triangles, 7) - - assert points.shape == (7, 3) - assert set(points[:, 0].tolist()) <= {0.0, 1.0, 2.0} - - def test_merged_surface_sampling_is_weighted_by_face_area(self): - triangle = torch.tensor(((0, 1, 2),), dtype=torch.long) - articulation, _ = _make_point_cloud_articulation( - link_meshes={ - "small": ( - torch.tensor( - ((-10.0, 0.0, 0.0), (-9.0, 0.0, 0.0), (-10.0, 1.0, 0.0)), - dtype=torch.float32, - ), - triangle, - ), - "target": ( - torch.tensor( - ((0.0, 0.0, 0.0), (2.0, 0.0, 0.0), (0.0, 2.0, 0.0)), - dtype=torch.float32, - ), - triangle, - ), - } - ) - - geometry = articulation.sample_initial_point_clouds( - "target", - articulation_point_count=5_000, - target_point_count=3, - ) - - articulation_points = geometry["articulation_point_cloud"] - small_face_fraction = float((articulation_points[:, 0] < -9.0).float().mean()) - assert small_face_fraction == pytest.approx(0.2, abs=0.03) - - @pytest.mark.parametrize( - ("target_link_name", "error_type", "message"), - ( - (None, TypeError, "target_link_name must be a string"), - (" ", ValueError, "target_link_name must be non-empty"), - ("missing", ValueError, "Unknown articulation link"), - ), - ) - def test_rejects_invalid_target_link_names( - self, - target_link_name: object, - error_type: type[Exception], - message: str, - ): - articulation, _ = _make_point_cloud_articulation() - - with pytest.raises(error_type, match=message): - articulation.sample_initial_point_clouds(target_link_name) # type: ignore[arg-type] - - @pytest.mark.parametrize( - ("field_name", "value", "error_type", "message"), - ( - ( - "articulation_point_count", - True, - TypeError, - "articulation_point_count must be an integer", - ), - ( - "target_point_count", - 0, - ValueError, - "target_point_count must be positive", - ), - ), - ) - def test_rejects_invalid_point_counts( - self, - field_name: str, - value: object, - error_type: type[Exception], - message: str, - ): - articulation, _ = _make_point_cloud_articulation() - - with pytest.raises(error_type, match=message): - articulation.sample_initial_point_clouds( - "target", - **{field_name: value}, # type: ignore[arg-type] - ) - - def test_requires_a_kinematic_chain(self): - articulation, _ = _make_point_cloud_articulation() - articulation.pk_chain = None - - with pytest.raises(RuntimeError, match="cfg.build_pk_chain=True"): - articulation.sample_initial_point_clouds("target") - - def test_rejects_non_unit_body_scale(self): - articulation, _ = _make_point_cloud_articulation( - body_scale=(1.0, 2.0, 1.0), - ) - - with pytest.raises(ValueError, match="requires unit body_scale"): - articulation.sample_initial_point_clouds("target") - - def test_rejects_mismatched_joint_name_sets(self): - articulation, _ = _make_point_cloud_articulation( - backend_joint_names=("backend_joint",), - kinematic_joint_names=("kinematic_joint",), - ) - - with pytest.raises(ValueError, match="matching simulator and kinematic-chain"): - articulation.sample_initial_point_clouds("target") - - @pytest.mark.parametrize( - "initial_qpos", - ( - (), - (float("nan"),), - ), - ) - def test_rejects_invalid_initial_qpos(self, initial_qpos: tuple[float, ...]): - articulation, _ = _make_point_cloud_articulation( - initial_qpos=initial_qpos, - ) - - with pytest.raises(ValueError, match="finite vector matching"): - articulation.sample_initial_point_clouds("target") - - def test_rejects_invalid_mesh_indices(self): - articulation, _ = _make_point_cloud_articulation( - link_meshes={ - "target": ( - torch.zeros((3, 3), dtype=torch.float32), - torch.tensor(((0, 1, 3),), dtype=torch.long), - ) - } - ) - - with pytest.raises(ValueError, match="triangles reference invalid vertices"): - articulation.sample_initial_point_clouds("target") - - def _link_static_friction(art: Articulation, link_name: str, env_idx: int = 0) -> float: return art._entities[env_idx].get_physical_attr(link_name).static_friction diff --git a/tests/sim/objects/test_robot.py b/tests/sim/objects/test_robot.py index e6e050f91..f423f33bd 100644 --- a/tests/sim/objects/test_robot.py +++ b/tests/sim/objects/test_robot.py @@ -18,13 +18,14 @@ import os from types import SimpleNamespace +from unittest.mock import patch import numpy as np import pytest import torch from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.objects import Robot +from embodichain.lab.sim.objects import Articulation, Robot from embodichain.lab.sim.robots.dexforce_w1 import DexforceW1Cfg from embodichain.data import get_data_path @@ -65,6 +66,30 @@ def test_get_qf_selects_control_part_joint_efforts(): assert torch.equal(actual_qf, full_qf[:, [3, 1]]) +@pytest.mark.no_sim +def test_compute_fk_forwards_named_joint_state_to_articulation(): + robot = object.__new__(Robot) + qpos = torch.tensor(((2.0, 1.0),)) + expected = torch.eye(4).reshape(1, 1, 4, 4) + + with patch.object(Articulation, "compute_fk", return_value=expected) as compute_fk: + result = robot.compute_fk( + qpos, + link_names=["target"], + env_ids=(), + qpos_joint_names=("joint_b", "joint_a"), + ) + + assert result is expected + compute_fk.assert_called_once_with( + qpos=qpos, + link_names=["target"], + end_link_name=None, + root_link_name=None, + qpos_joint_names=("joint_b", "joint_a"), + ) + + # Base test class for CPU and CUDA class BaseRobotTest: @classmethod From e5271f7272c20e6aeb8ae2bc6e0f58162eb4ce3f Mon Sep 17 00:00:00 2001 From: matafela Date: Wed, 2 Sep 2026 15:54:49 +0800 Subject: [PATCH 7/8] add enable gravity --- agent_context/MAP.yaml | 3 ++ .../topics/robot-system/robot-system.md | 3 +- .../simulation-system/simulation-system.md | 6 +++ docs/source/overview/sim/sim_articulation.md | 5 ++- embodichain/lab/sim/cfg.py | 6 +++ embodichain/lab/sim/objects/articulation.py | 19 +++++++++ embodichain/lab/sim/robots/dual_arm.py | 1 + scripts/tutorials/atomic_action/hand_over.py | 2 +- tests/sim/objects/test_articulation.py | 41 +++++++++++++++++++ tests/sim/objects/test_dual_arm.py | 9 ++++ tests/sim/test_cfg.py | 12 ++++++ 11 files changed, 104 insertions(+), 3 deletions(-) diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index 3e9e8e98f..278bbb500 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -33,6 +33,8 @@ topics: - simulation lifecycle - ArticulationJointKinematics - get_parent_joint_chain + - ArticulationCfg + - enable_gravity paths: - topics/simulation-system/simulation-system.md source_of_truth: @@ -253,6 +255,7 @@ topics: - dexforce_w1 - gripper - arm + - enable_gravity paths: - topics/robot-system/robot-system.md source_of_truth: diff --git a/agent_context/topics/robot-system/robot-system.md b/agent_context/topics/robot-system/robot-system.md index bf7e4ce9a..3fc36e140 100644 --- a/agent_context/topics/robot-system/robot-system.md +++ b/agent_context/topics/robot-system/robot-system.md @@ -30,7 +30,8 @@ Inheritance chain: ``` ObjectBaseCfg uid, init_pos, init_rot, init_local_pose └─ ArticulationCfg fpath, drive_pros, attrs, link_attrs, fix_base, - │ disable_self_collision, init_qpos, body_scale, + │ disable_self_collision, enable_gravity, init_qpos, + │ body_scale, │ build_pk_chain, use_usd_properties └─ RobotCfg control_parts, urdf_cfg, solver_cfg, drive_pros (override default to "force") ├─ DexforceW1Cfg version, hand_versions, with_default_eef diff --git a/agent_context/topics/simulation-system/simulation-system.md b/agent_context/topics/simulation-system/simulation-system.md index 24a403db3..7ecd402b9 100644 --- a/agent_context/topics/simulation-system/simulation-system.md +++ b/agent_context/topics/simulation-system/simulation-system.md @@ -70,6 +70,12 @@ then advances the world for the requested number of physics steps. Each environment control step normally calls it with `sim_steps_per_control`. +`ArticulationCfg.enable_gravity` defaults to `True`. During articulation +construction, `Articulation` applies this explicit runtime flag to every native +entity before the first physics update, including when +`use_usd_properties=True`. Use `Articulation.set_gravity(...)` to change the +flag later for all or selected environment indices. + ## Module Boundaries | Area | Owner | Routed topic | diff --git a/docs/source/overview/sim/sim_articulation.md b/docs/source/overview/sim/sim_articulation.md index 455946172..6e5118654 100644 --- a/docs/source/overview/sim/sim_articulation.md +++ b/docs/source/overview/sim/sim_articulation.md @@ -18,11 +18,14 @@ Articulations are configured using the {class}`~cfg.ArticulationCfg` dataclass. | `init_qpos` | `List[float]` | `None` | Initial joint positions. | | `qpos_limits` | `Tensor` / `Dict[str, List[float]]` | `None` | Override joint position limits. Replaces asset limits and may either tighten or expand the range. | | `body_scale` | `List[float]` | `[1.0, 1.0, 1.0]` | Scaling factors for the articulation links. | -| `disable_self_collisions` | `bool` | `True` | Whether to disable self-collisions. | +| `disable_self_collision` | `bool` | `True` | Whether to disable self-collisions. | +| `enable_gravity` | `bool` | `True` | Whether gravity affects the articulation. This runtime flag also applies when `use_usd_properties=True`. | | `drive_pros` | `JointDrivePropertiesCfg` | `drive_type="none"` | Default drive properties. | | `attrs` | `RigidBodyAttributesCfg` | `...` | Default rigid body attributes applied to all links. | | `link_attrs` | `dict[str, LinkPhysicsOverrideCfg]` | `None` | Optional per-link overrides keyed by group name; each group matches link names via regex. | +At runtime, call `articulation.set_gravity(...)` to change gravity for every +environment or for a selected set of environment indices. ### Per-link physics (`link_attrs`) diff --git a/embodichain/lab/sim/cfg.py b/embodichain/lab/sim/cfg.py index c36fabbcd..3335af461 100644 --- a/embodichain/lab/sim/cfg.py +++ b/embodichain/lab/sim/cfg.py @@ -1634,6 +1634,12 @@ class ArticulationCfg(ObjectBaseCfg): disable_self_collision: bool = True """Whether to enable or disable self-collisions.""" + enable_gravity: bool = True + """Whether gravity is enabled for the articulation. + + This runtime flag is applied regardless of :attr:`use_usd_properties`. + """ + init_qpos: torch.Tensor | np.ndarray | Sequence[float] = None """Initial joint positions of the articulation. diff --git a/embodichain/lab/sim/objects/articulation.py b/embodichain/lab/sim/objects/articulation.py index 9940f6108..d92825163 100644 --- a/embodichain/lab/sim/objects/articulation.py +++ b/embodichain/lab/sim/objects/articulation.py @@ -654,6 +654,10 @@ def __init__( # Store all indices for batch operations self._all_indices = torch.arange(len(entities), dtype=torch.int32) + # Apply gravity before the first physics update. Unlike asset-backed + # physical properties, this config field is an explicit runtime flag. + self.set_gravity(self.cfg.enable_gravity) + if device.type == "cuda": self._world.update(0.001) @@ -2727,6 +2731,21 @@ def set_self_collision( ArticulationFlag.DISABLE_SELF_COLLISION, not enable ) + def set_gravity( + self, + enable: bool = True, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set whether gravity is enabled for the articulation. + + Args: + enable: Whether to enable gravity. Defaults to True. + env_ids: Environment indices. If None, all environments are used. + """ + local_env_ids = self._all_indices if env_ids is None else env_ids + for env_idx in local_env_ids: + self._entities[env_idx].enable_gravity(bool(enable)) + def destroy(self) -> None: env = self._world.get_env() arenas = env.get_all_arenas() diff --git a/embodichain/lab/sim/robots/dual_arm.py b/embodichain/lab/sim/robots/dual_arm.py index c42c85396..6ddd63b28 100644 --- a/embodichain/lab/sim/robots/dual_arm.py +++ b/embodichain/lab/sim/robots/dual_arm.py @@ -410,6 +410,7 @@ def _populate_dual_cfg( cfg.min_velocity_iters = base_cfg.min_velocity_iters cfg.fix_base = base_cfg.fix_base cfg.disable_self_collision = base_cfg.disable_self_collision + cfg.enable_gravity = base_cfg.enable_gravity cfg.sleep_threshold = base_cfg.sleep_threshold diff --git a/scripts/tutorials/atomic_action/hand_over.py b/scripts/tutorials/atomic_action/hand_over.py index fc110e525..f34d7116b 100644 --- a/scripts/tutorials/atomic_action/hand_over.py +++ b/scripts/tutorials/atomic_action/hand_over.py @@ -68,7 +68,7 @@ ) VERTICAL_OBJECT_SCALE = (0.56, 0.56, 0.56) HORIZONTAL_OBJECT_SCALE = (2.0, 2.0, 2.0) -GRIPPER_TCP_Z = 0.17 +GRIPPER_TCP_Z = 0.16 SUPPORT_SURFACE_Z = 0.50 SUPPORT_SURFACE_SIZE = (0.8, 1.2, 0.02) SUPPORT_SURFACE_CENTER = ( diff --git a/tests/sim/objects/test_articulation.py b/tests/sim/objects/test_articulation.py index f4e51a6ae..6b3288191 100644 --- a/tests/sim/objects/test_articulation.py +++ b/tests/sim/objects/test_articulation.py @@ -43,6 +43,47 @@ NUM_ARENAS = 10 +class _GravityEntity: + """Record native gravity calls for an articulation test double.""" + + def __init__(self) -> None: + self.calls: list[bool] = [] + + def enable_gravity(self, flag: bool) -> None: + self.calls.append(flag) + + +@pytest.mark.no_sim +@pytest.mark.parametrize("enable", [True, False]) +def test_set_gravity_updates_only_selected_environments(enable: bool) -> None: + """Runtime gravity updates are dispatched to the requested native entities.""" + articulation = object.__new__(Articulation) + articulation._entities = [_GravityEntity() for _ in range(3)] + articulation._all_indices = torch.arange(3, dtype=torch.int32) + + articulation.set_gravity(enable, env_ids=(2, 0)) + + assert articulation._entities[0].calls == [enable] + assert articulation._entities[1].calls == [] + assert articulation._entities[2].calls == [enable] + + +@pytest.mark.no_sim +def test_set_gravity_updates_all_environments_by_default() -> None: + """Omitting environment indices applies gravity to every native entity.""" + articulation = object.__new__(Articulation) + articulation._entities = [_GravityEntity() for _ in range(3)] + articulation._all_indices = torch.arange(3, dtype=torch.int32) + + articulation.set_gravity(False) + + assert [entity.calls for entity in articulation._entities] == [ + [False], + [False], + [False], + ] + + def test_get_qf_returns_all_articulation_joint_efforts(): expected_qf = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=torch.float32) articulation = object.__new__(Articulation) diff --git a/tests/sim/objects/test_dual_arm.py b/tests/sim/objects/test_dual_arm.py index d4d9febf5..ae32f02c9 100644 --- a/tests/sim/objects/test_dual_arm.py +++ b/tests/sim/objects/test_dual_arm.py @@ -162,6 +162,15 @@ def test_build_dual_arm_dual_part_toggle(): assert "dual_arm" not in cfg.control_parts +def test_build_dual_arm_preserves_gravity_setting() -> None: + base = URRobotCfg.from_dict({"robot_type": "ur5", "enable_gravity": False}) + mounts = resolve_mounts({"preset": "side_by_side", "separation": 0.6}) + + cfg = build_dual_arm_cfg(base, mounts) + + assert cfg.enable_gravity is False + + # --------------------------------------------------------------------------- # # DualArmRobotCfg from_dict + round-trip # --------------------------------------------------------------------------- # diff --git a/tests/sim/test_cfg.py b/tests/sim/test_cfg.py index c9cfc28fe..3bc44ec53 100644 --- a/tests/sim/test_cfg.py +++ b/tests/sim/test_cfg.py @@ -40,6 +40,18 @@ def test_articulation_cfg_partial_drive_properties_preserve_no_drive() -> None: assert articulation_cfg.drive_pros.drive_type == "none" +def test_articulation_cfg_enables_gravity_by_default() -> None: + """Articulations opt into gravity unless explicitly configured otherwise.""" + assert ArticulationCfg().enable_gravity is True + + +def test_articulation_cfg_parses_disabled_gravity() -> None: + """Dictionary configuration can disable articulation gravity.""" + articulation_cfg = ArticulationCfg.from_dict({"enable_gravity": False}) + + assert articulation_cfg.enable_gravity is False + + def test_robot_cfg_defaults_to_force_joint_drive() -> None: """Robots retain force-based joint drives by default.""" robot_cfg = RobotCfg() From df3ae5986e530def0f356bf9ed0e6b572a9bbb80 Mon Sep 17 00:00:00 2001 From: matafela Date: Wed, 2 Sep 2026 16:13:58 +0800 Subject: [PATCH 8/8] update --- .../topics/atomic-actions/atomic-actions.md | 5 ++ .../sim/atomic_actions/builtin_actions.md | 3 ++ .../atomic_actions/articulation_geometry.py | 42 +++++++++++---- .../test_articulation_geometry.py | 51 +++++++++++++++++++ 4 files changed, 91 insertions(+), 10 deletions(-) diff --git a/agent_context/topics/atomic-actions/atomic-actions.md b/agent_context/topics/atomic-actions/atomic-actions.md index f7458c813..084d5f4de 100644 --- a/agent_context/topics/atomic-actions/atomic-actions.md +++ b/agent_context/topics/atomic-actions/atomic-actions.md @@ -75,6 +75,11 @@ body scale. The adapter owns Open3D surface sampling, target-link-frame mesh transforms, nearest prismatic/revolute ancestor geometry, and the private `ObjectSemantics.geometry` keys. +Every non-empty provider link mesh must expose at least one non-degenerate +triangle surface. The adapter rejects vertices-only or fully degenerate links +before merging so the target-link and whole-articulation clouds cannot disagree +about whether a link contributes geometry. + The adapter returns `ArticulationAffordanceGeometry`; convert it with `to_object_geometry()` only when constructing `ObjectSemantics`. Keep random sampling and semantic geometry keys out of `objects/articulation.py`. diff --git a/docs/source/overview/sim/atomic_actions/builtin_actions.md b/docs/source/overview/sim/atomic_actions/builtin_actions.md index fff4ce566..635e35ba6 100644 --- a/docs/source/overview/sim/atomic_actions/builtin_actions.md +++ b/docs/source/overview/sim/atomic_actions/builtin_actions.md @@ -732,6 +732,9 @@ target link's initial local frame. Both clouds use Open3D uniform surface sampling. Sampling the whole articulation as one merged mesh preserves triangle-area weighting instead of giving every link an equal point budget, which would over-represent tiny decorative links. +To keep the target and merged clouds geometrically consistent, every non-empty +link mesh must contain at least one non-degenerate triangle surface; the adapter +rejects vertices-only and fully degenerate link meshes before merging. The sampler also walks the target's parent chain and transforms the nearest prismatic joint axis into the target link's initial local frame. The target cloud center and twice its distribution radius define a spherical neighborhood diff --git a/embodichain/lab/sim/atomic_actions/articulation_geometry.py b/embodichain/lab/sim/atomic_actions/articulation_geometry.py index 7023674aa..ae61ab790 100644 --- a/embodichain/lab/sim/atomic_actions/articulation_geometry.py +++ b/embodichain/lab/sim/atomic_actions/articulation_geometry.py @@ -80,6 +80,10 @@ def get_link_vert_face( Returns: Link-local vertices and triangle indices. + + .. attention:: + A non-empty link mesh must contain at least one non-degenerate + triangle surface. """ ... @@ -566,6 +570,11 @@ def _validate_point_cloud_mesh( if triangles.dim() != 2 or triangles.shape[1:] != (3,): raise ValueError(f"Link {link_name!r} triangles must have shape (M, 3).") if triangles.shape[0] == 0: + if vertices.shape[0] > 0: + raise ValueError( + f"Link {link_name!r} must contain at least one non-degenerate " + "triangle surface." + ) return if vertices.shape[0] == 0: raise ValueError( @@ -575,6 +584,28 @@ def _validate_point_cloud_mesh( vertices ): raise ValueError(f"Link {link_name!r} triangles reference invalid vertices.") + if not bool(_valid_triangle_mask(vertices, triangles).any().item()): + raise ValueError( + f"Link {link_name!r} must contain at least one non-degenerate " + "triangle surface." + ) + + +def _valid_triangle_mask( + vertices: torch.Tensor, + triangles: torch.Tensor, +) -> torch.Tensor: + """Return the non-degenerate triangle mask used by validation and sampling.""" + face_vertices = vertices[triangles] + face_areas = 0.5 * torch.linalg.vector_norm( + torch.cross( + face_vertices[:, 1] - face_vertices[:, 0], + face_vertices[:, 2] - face_vertices[:, 0], + dim=1, + ), + dim=1, + ) + return face_areas > torch.finfo(vertices.dtype).eps def _sample_mesh_surface_points( @@ -602,16 +633,7 @@ def _sample_mesh_surface_points( ) return vertices[indices].clone() - face_vertices = vertices[triangles] - face_areas = 0.5 * torch.linalg.vector_norm( - torch.cross( - face_vertices[:, 1] - face_vertices[:, 0], - face_vertices[:, 2] - face_vertices[:, 0], - dim=1, - ), - dim=1, - ) - valid_faces = face_areas > torch.finfo(vertices.dtype).eps + valid_faces = _valid_triangle_mask(vertices, triangles) if not bool(valid_faces.any().item()): indices = ( torch.linspace( diff --git a/tests/sim/atomic_actions/test_articulation_geometry.py b/tests/sim/atomic_actions/test_articulation_geometry.py index 450e398f9..0809c796d 100644 --- a/tests/sim/atomic_actions/test_articulation_geometry.py +++ b/tests/sim/atomic_actions/test_articulation_geometry.py @@ -707,6 +707,57 @@ def test_surface_sampling_falls_back_to_vertices_without_valid_faces( assert points.shape == (7, 3) assert set(points[:, 0].tolist()) <= {0.0, 1.0, 2.0} + @pytest.mark.parametrize( + "target_triangles", + ( + pytest.param( + torch.empty((0, 3), dtype=torch.long), + id="no-faces", + ), + pytest.param( + torch.tensor(((0, 1, 2),), dtype=torch.long), + id="degenerate-faces", + ), + ), + ) + def test_rejects_mixed_mesh_when_target_has_no_valid_triangle_surface( + self, + target_triangles: torch.Tensor, + ) -> None: + valid_triangle = torch.tensor(((0, 1, 2),), dtype=torch.long) + articulation, _ = _make_point_cloud_articulation( + link_meshes={ + "body": ( + torch.tensor( + ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)), + dtype=torch.float32, + ), + valid_triangle, + ), + "target": ( + torch.tensor( + ((10.0, 0.0, 0.0), (11.0, 0.0, 0.0), (12.0, 0.0, 0.0)), + dtype=torch.float32, + ), + target_triangles, + ), + } + ) + + with pytest.raises( + ValueError, + match="Link 'target' must contain at least one non-degenerate triangle", + ): + sample_initial_articulation_geometry( + articulation, + "target", + initial_qpos=articulation.initial_qpos, + initial_qpos_joint_names=articulation.initial_qpos_joint_names, + body_scale=articulation.body_scale, + articulation_point_count=32, + target_point_count=8, + ) + def test_merged_surface_sampling_is_weighted_by_face_area(self): triangle = torch.tensor(((0, 1, 2),), dtype=torch.long) articulation, _ = _make_point_cloud_articulation(