diff --git a/docs/source/overview/sim/atomic_actions/builtin_actions.md b/docs/source/overview/sim/atomic_actions/builtin_actions.md index fd46a99bc..b351e405a 100644 --- a/docs/source/overview/sim/atomic_actions/builtin_actions.md +++ b/docs/source/overview/sim/atomic_actions/builtin_actions.md @@ -677,10 +677,15 @@ 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`, 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; each output sample is grounded with IK instead of being interpolated only in @@ -698,7 +703,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 +716,32 @@ 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 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 @@ -806,8 +830,24 @@ 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` 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. + +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 eec111ded..e0888bec0 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 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)`. | @@ -193,6 +194,52 @@ 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"] +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. + ### 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..0c3e52938 100644 --- a/embodichain/lab/sim/atomic_actions/affordance.py +++ b/embodichain/lab/sim/atomic_actions/affordance.py @@ -16,11 +16,18 @@ 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" +_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" + @dataclass class Affordance: @@ -49,6 +56,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 +261,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 parent-joint rotation 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 revolute-joint origin metadata.""" 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.""" + """Parent revolute-joint axis, signed toward articulation geometry.""" joint_name: str | None = None """Optional stable articulation-joint name associated with the axis.""" @@ -275,11 +296,47 @@ 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: + """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] + + 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. + + 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 geometry['target_link_revolute_axis_origin']; " + "the target link's revolute joint origin is missing." + ) + 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 +377,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 parent-joint translation geometry. The positive translation-axis direction denotes approaching and pushing the articulated part closed. Pulling moves in the opposite direction. @@ -337,7 +394,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.""" + """Parent prismatic-joint axis, signed toward articulation geometry.""" joint_name: str | None = None """Optional stable prismatic-joint name associated with the link.""" @@ -360,6 +417,16 @@ 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: + """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] + @dataclass class OpenDoorAffordance(AntipodalAffordance): @@ -558,15 +625,18 @@ def from_articulation( @dataclass class PressAffordance(Affordance): - """Explicit target-local contact point and 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 expressed in the target object's local frame.""" + """Parent prismatic-joint axis, signed toward articulation 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 +648,26 @@ 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: + """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 + 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 +703,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 joint geometry." + ) local_press_position = torch.tensor( configured_position, dtype=torch.float32, @@ -653,6 +745,183 @@ def _validate_press_position( return tuple(float(component) for component in position) +def _infer_articulation_neighborhood_axis( + geometry: Mapping[str, Any], + *, + field_name: str, + axis_key: str, +) -> tuple[torch.Tensor, torch.Tensor] | None: + """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 neighborhood-center offset disambiguates the + sign of the normalized joint axis supplied in target-link coordinates. + + Args: + 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 joint-axis metadata is present, otherwise + ``None`` so legacy and non-articulation affordances retain their axis. + + Raises: + 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 + 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 or not has_axis: + raise ValueError( + 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( + 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." + ) + 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( + 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 + direction_score = torch.dot(center_offset, joint_axis) + offset_tolerance = max(1.0e-8, float(target_radius.item()) * 1.0e-6) + if abs(float(direction_score.item())) <= offset_tolerance: + raise ValueError( + f"{field_name} direction is ambiguous because the articulation " + "neighborhood-center offset is orthogonal to the joint axis." + ) + + 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: + """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 _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], + *, + 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, +) -> 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..713578203 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,364 @@ 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: 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/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/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/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/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..0e3efd0c9 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,36 @@ _EXPECTED_GRASP_SAMPLES = 1_000 +def _drawer_handle_point_cloud_geometry() -> dict[str, torch.Tensor]: + """Return target-local clouds and a prismatic axis oriented toward +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), + "target_link_prismatic_joint_axis": torch.tensor( + [0.0, 1.0, 0.0], + dtype=torch.float32, + ), + } + + class _NeverObserveProvider: """Reject dynamic observations during configuration decoding/compilation.""" @@ -533,6 +563,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 +611,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 +658,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 +704,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..d378d5dc4 100644 --- a/tests/sim/atomic_actions/test_actions.py +++ b/tests/sim/atomic_actions/test_actions.py @@ -124,12 +124,45 @@ DUAL_ARM_DOF = 2 * ARM_DOF DUAL_ROBOT_DOF = DUAL_ARM_DOF + 2 * HAND_DOF 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], + [-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 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 + 1.5 * AUTOMATIC_TWIST_JOINT_AXIS_UNIT + ) + return { + "target_link_point_cloud": target_points, + "articulation_point_cloud": torch.cat( + (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(), + } + + class _StubGraspPoseGenerator(ParallelJawGraspPoseGenerator): """Deterministic planning-service double used by atomic-action tests.""" @@ -2388,13 +2421,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 +2461,14 @@ 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.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) + ) 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 +2505,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_joint_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,15 +2521,15 @@ 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, ) assert torch.allclose( twisted[:, -1, :3, 3], - torch.tensor([1.0, 1.0, 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 1fdd5fc87..280247605 100644 --- a/tests/sim/atomic_actions/test_affordance.py +++ b/tests/sim/atomic_actions/test_affordance.py @@ -36,6 +36,105 @@ SlideAffordance, TwistAffordance, ) +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( + [ + [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" +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" + + +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, + ), + 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, + *, + 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 +254,414 @@ 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("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, + 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]), + ) + + ObjectSemantics( + affordance=affordance, + 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.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_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( + ("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, + ) + + ObjectSemantics( + affordance=affordance, + 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( + "present_fields", + ( + ("target",), + ("articulation",), + ("axis",), + ("target", "articulation"), + ("target", "axis"), + ("articulation", "axis"), + ), + ) + @pytest.mark.parametrize("kind", ("slide", "press", "twist")) + def test_incomplete_joint_axis_inference_metadata_is_rejected( + self, + 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( + kind, + fallback_axis=torch.tensor([1.0, 0.0, 0.0]), + ) + + with pytest.raises(ValueError, match=axis_key): + ObjectSemantics( + affordance=affordance, + geometry=geometry, + 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 = _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]), + ) + + with pytest.raises(ValueError, match="degenerate target-link"): + ObjectSemantics( + affordance=affordance, + geometry=geometry, + entity_id="degenerate-geometry-target", + ) + + @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( + kind, + fallback_axis=torch.tensor([1.0, 0.0, 0.0]), + ) + + with pytest.raises(ValueError, match="direction is ambiguous"): + ObjectSemantics( + affordance=affordance, + geometry=_axis_geometry(perpendicular_offset), + entity_id=f"ambiguous-{kind}-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( + tuple(float(value) for value in -1.5 * PRISMATIC_JOINT_AXIS_UNIT) + ), + entity_id="automatic-press-target", + ) + + assert torch.allclose( + affordance.press_axis, + -PRISMATIC_JOINT_AXIS_UNIT, + atol=1.0e-6, + ) + 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) + ) + + 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), + twist_axis=torch.tensor([1.0, 0.0, 0.0]), + ) + + ObjectSemantics( + affordance=affordance, + geometry=_axis_geometry( + tuple(float(value) for value in -1.5 * REVOLUTE_JOINT_AXIS_UNIT) + ), + entity_id="automatic-twist-target", + ) + + assert torch.allclose( + affordance.twist_axis, + -REVOLUTE_JOINT_AXIS_UNIT, + atol=1.0e-6, + ) + assert affordance.axis_origin == pytest.approx( + 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( + 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), + 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.allclose( + affordance.twist_axis, + -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( + 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), + 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.allclose( + affordance.twist_axis, + -REVOLUTE_JOINT_AXIS_UNIT, + atol=1.0e-6, + ) + 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, + ) -> 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 +904,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..ffe2c8a68 100644 --- a/tests/sim/atomic_actions/test_tutorial_utils.py +++ b/tests/sim/atomic_actions/test_tutorial_utils.py @@ -130,6 +130,77 @@ "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]) + + +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), + ) + ), + "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(), + } + + +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]: @@ -155,6 +226,73 @@ 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", + (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", + (-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.5, -1.0, -1.0), + (1.0 / 3.0, -2.0 / 3.0, -2.0 / 3.0), + ), + ), +) +def test_articulation_tutorial_semantics_resolve_signed_parent_joint_axis( + 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.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.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) + ) + + def test_should_wait_for_tutorial_input_is_disabled_for_headless_modes() -> None: assert ( should_wait_for_tutorial_input( @@ -474,17 +612,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") @@ -704,6 +831,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: diff --git a/tests/sim/objects/test_articulation.py b/tests/sim/objects/test_articulation.py index 03f07983c..bc018bc6b 100644 --- a/tests/sim/objects/test_articulation.py +++ b/tests/sim/objects/test_articulation.py @@ -41,6 +41,76 @@ 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(): @@ -96,6 +166,667 @@ 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