From 5214cae36ad62c71dcceb76a6586ed9df1a5458d Mon Sep 17 00:00:00 2001 From: Peter Du Date: Mon, 19 Oct 2026 17:51:55 +0000 Subject: [PATCH 1/2] Add SoftMimicGen deformable trajectory generation --- README.md | 4 +- autodata_core/__init__.py | 2 + autodata_core/algorithms.py | 139 ++++++++- autodata_core/data_generator.py | 40 ++- autodata_core/datagen_info.py | 6 + autodata_core/deformable_transforms.py | 151 ++++++++++ autodata_core/pool.py | 25 +- autodata_core/selection_strategy.py | 65 ++++- autodata_interfaces/datastream/datastream.py | 21 ++ autodata_interfaces/env/__init__.py | 2 + .../env/isaaclab_env_interface.py | 80 +++++- autodata_interfaces/env/reset_request.py | 23 ++ autodata_interfaces/env/scene_state.py | 36 +++ .../tasks/generation_policy_spec.py | 4 + autodata_interfaces/tasks/subtask_spec.py | 18 ++ autodata_tests/core/test_algorithms.py | 18 +- autodata_tests/core/test_datagen_info.py | 11 + .../core/test_deformable_transforms.py | 58 ++++ autodata_tests/core/test_pool.py | 72 +++++ .../core/test_selection_strategy.py | 66 +++++ .../interfaces/datastream/test_datastream.py | 23 ++ .../interfaces/env/test_env_loop.py | 161 +++++++++++ autodata_tests/interfaces/mocks.py | 4 + .../tasks/test_generation_policy_spec.py | 4 +- .../interfaces/tasks/test_subtask_spec.py | 11 + .../utils/test_thin_plate_spline.py | 86 ++++++ autodata_utils/thin_plate_spline.py | 263 ++++++++++++++++++ scripts/annotate_demos.py | 1 + scripts/generate_dataset.py | 5 +- 29 files changed, 1372 insertions(+), 27 deletions(-) create mode 100644 autodata_core/deformable_transforms.py create mode 100644 autodata_interfaces/env/reset_request.py create mode 100644 autodata_interfaces/env/scene_state.py create mode 100644 autodata_tests/core/test_deformable_transforms.py create mode 100644 autodata_tests/interfaces/env/test_env_loop.py create mode 100644 autodata_tests/utils/test_thin_plate_spline.py create mode 100644 autodata_utils/thin_plate_spline.py diff --git a/README.md b/README.md index 7bde5cd..cd483a4 100644 --- a/README.md +++ b/README.md @@ -198,8 +198,8 @@ If you use AutoData in your research, please cite: ``` Depending on the generation algorithm used, please also cite the original -[MimicGen](https://arxiv.org/abs/2310.17596), [DexMimicGen](https://arxiv.org/abs/2410.24185), or -[SkillMimicGen](https://arxiv.org/abs/2410.18907) work. Isaac Lab users should also cite the +[MimicGen](https://arxiv.org/abs/2310.17596), [DexMimicGen](https://arxiv.org/abs/2410.24185), +[SkillMimicGen](https://arxiv.org/abs/2410.18907), or [SoftMimicGen](https://arxiv.org/abs/2603.25725) work. Isaac Lab users should also cite the [Isaac Lab paper](https://arxiv.org/abs/2511.04831). ## Acknowledgements diff --git a/autodata_core/__init__.py b/autodata_core/__init__.py index edf52bf..4e3dc92 100644 --- a/autodata_core/__init__.py +++ b/autodata_core/__init__.py @@ -16,6 +16,7 @@ GenerationAlgorithm, MimicGen, SkillGen, + SoftMimicGen, get_algorithm, iter_algorithms, ) @@ -33,6 +34,7 @@ "GenerationResult", "MimicGen", "MultiWaypoint", + "SoftMimicGen", "SkillGen", "Waypoint", "WaypointSequence", diff --git a/autodata_core/algorithms.py b/autodata_core/algorithms.py index 65d6daf..9814430 100644 --- a/autodata_core/algorithms.py +++ b/autodata_core/algorithms.py @@ -7,14 +7,15 @@ behavior diverges between Mimic, DexMimicGen, SkillGen, etc. Subclasses self-register at import time via :class:`_AlgorithmMeta`. -The single behavioral hook is :meth:`GenerationAlgorithm.plan_subtask_trajectory`, called by -``DataGenerator`` every time an EEF needs a new executable trajectory. The default implementation -generates a subtask trajectory and merges an interpolation segment (the Mimic / DexMimicGen path). -SkillGen overrides this to insert a motion-planned transit ahead of the merge. +The main behavioral hook is :meth:`GenerationAlgorithm.plan_subtask_trajectory`, called by +``DataGenerator`` every time an EEF needs a new executable trajectory. Object-state adaptation is +routed through :meth:`GenerationAlgorithm.transform_source_eef_poses`. Rigid algorithms use the +default object-pose transform while SoftMimicGen overrides it with nodal TPS registration. """ from __future__ import annotations +import torch from collections.abc import Iterator from typing import TYPE_CHECKING, Any @@ -22,6 +23,7 @@ from autodata_core.data_generator import DataGenerator, _EEFGenerationState from autodata_core.waypoint import Waypoint from autodata_interfaces.datastream.datastream import Datastream + from autodata_interfaces.tasks.subtask_spec import Subtask REGISTERED_ALGORITHMS: dict[str, type[GenerationAlgorithm]] = {} @@ -71,6 +73,45 @@ def validate_setup(self, datastream: Datastream) -> None: or ``datastream.get_env()`` (escape hatch) to enforce algorithm-specific invariants. """ + def is_deformable_subtask(self, subtask: Subtask) -> bool: + """Return whether ``subtask`` uses deformable nodal state instead of a rigid pose.""" + + return False + + def transform_source_eef_poses( + self, + *, + data_generator: DataGenerator, + eef_name: str, + subtask_ind: int, + subtask_object_name: str | None, + subtask_object_pose: torch.Tensor | None, + src_subtask_object_pose: torch.Tensor | None, + subtask_object_nodal_positions: torch.Tensor | None, + src_subtask_object_nodal_positions: torch.Tensor | None, + src_eef_poses: torch.Tensor, + use_delta_transform: torch.Tensor | None, + coord_transform_scheme: Any, + runtime_subtask_constraints_dict: dict, + ) -> torch.Tensor: + """Adapt source EEF poses to the current subtask object state. + + The default implementation performs the existing rigid-object transform. Algorithms using + another object representation override this method. + """ + + return data_generator._apply_subtask_transform( + eef_name=eef_name, + subtask_ind=subtask_ind, + subtask_object_name=subtask_object_name, + subtask_object_pose=subtask_object_pose, + src_subtask_object_pose=src_subtask_object_pose, + src_eef_poses=src_eef_poses, + use_delta_transform=use_delta_transform, + coord_transform_scheme=coord_transform_scheme, + runtime_subtask_constraints_dict=runtime_subtask_constraints_dict, + ) + def plan_subtask_trajectory( self, *, @@ -165,6 +206,96 @@ class DexMimicGen(GenerationAlgorithm): supports_coordination = True +class SoftMimicGen(GenerationAlgorithm): + """MimicGen for deformable reference objects. + + Deformable subtasks select source segments using nodal state and warp their EEF trajectories + with a thin-plate-spline transform. Rigid subtasks retain the normal MimicGen transform. + """ + + name = "softmimicgen" + expected_eef_count = (1, 2) + requires_motion_planner = False + uses_subtask_start_signals = False + supports_coordination = False + + def is_deformable_subtask(self, subtask: Subtask) -> bool: + return bool(getattr(subtask.algo_params, "object_soft", False)) + + def validate_setup(self, datastream: Datastream) -> None: + live_nodal_positions = datastream.get_object_nodal_positions() + for eef_name in datastream.get_eef_names(): + for subtask_index, subtask in enumerate(datastream.get_subtasks(eef_name)): + if not self.is_deformable_subtask(subtask): + continue + assert ( + subtask.object_ref + ), f"SoftMimicGen deformable subtask {eef_name}[{subtask_index}] requires object_ref" + assert subtask.object_ref in live_nodal_positions, ( + f"Deformable object {subtask.object_ref!r} is not present in the live scene; " + f"available: {sorted(live_nodal_positions)}" + ) + live_node_count = live_nodal_positions[subtask.object_ref].shape[-2] + for demo_index, datagen_info in enumerate(datastream.source_pool.datagen_infos): + assert ( + datagen_info.object_nodal_positions is not None + ), f"Source demo {demo_index} lacks object_nodal_position annotations" + assert ( + subtask.object_ref in datagen_info.object_nodal_positions + ), f"Source demo {demo_index} lacks nodal positions for {subtask.object_ref!r}" + source_node_count = datagen_info.object_nodal_positions[subtask.object_ref].shape[-2] + assert source_node_count == live_node_count, ( + f"Source demo {demo_index} has {source_node_count} nodes for " + f"{subtask.object_ref!r}, live object has {live_node_count}" + ) + + def transform_source_eef_poses( + self, + *, + data_generator: DataGenerator, + eef_name: str, + subtask_ind: int, + subtask_object_name: str | None, + subtask_object_pose: torch.Tensor | None, + src_subtask_object_pose: torch.Tensor | None, + subtask_object_nodal_positions: torch.Tensor | None, + src_subtask_object_nodal_positions: torch.Tensor | None, + src_eef_poses: torch.Tensor, + use_delta_transform: torch.Tensor | None, + coord_transform_scheme: Any, + runtime_subtask_constraints_dict: dict, + ) -> torch.Tensor: + subtask = data_generator.datastream.get_subtask(eef_name, subtask_ind) + if not self.is_deformable_subtask(subtask): + return super().transform_source_eef_poses( + data_generator=data_generator, + eef_name=eef_name, + subtask_ind=subtask_ind, + subtask_object_name=subtask_object_name, + subtask_object_pose=subtask_object_pose, + src_subtask_object_pose=src_subtask_object_pose, + subtask_object_nodal_positions=subtask_object_nodal_positions, + src_subtask_object_nodal_positions=src_subtask_object_nodal_positions, + src_eef_poses=src_eef_poses, + use_delta_transform=use_delta_transform, + coord_transform_scheme=coord_transform_scheme, + runtime_subtask_constraints_dict=runtime_subtask_constraints_dict, + ) + + assert subtask_object_nodal_positions is not None, "current deformable nodal state is missing" + assert src_subtask_object_nodal_positions is not None, "source deformable nodal state is missing" + from autodata_core.deformable_transforms import transform_source_data_segment_using_nodal_registration + + return transform_source_data_segment_using_nodal_registration( + src_eef_poses=src_eef_poses, + src_obj_nodal_pos=src_subtask_object_nodal_positions, + tgt_obj_nodal_pos=subtask_object_nodal_positions, + use_rotation_transform=subtask.algo_params.use_rotation_transform, + bend_coef=subtask.algo_params.bend_coef, + rot_coef=subtask.algo_params.rot_coef, + ) + + class SkillGen(GenerationAlgorithm): """Single-arm SkillGen — motion-planned transit followed by skill replay per subtask. diff --git a/autodata_core/data_generator.py b/autodata_core/data_generator.py index 59481cb..a46b4ce 100644 --- a/autodata_core/data_generator.py +++ b/autodata_core/data_generator.py @@ -28,6 +28,7 @@ transform_source_data_segment_using_object_pose, ) from autodata_core.waypoint import MultiWaypoint, Waypoint, WaypointSequence, WaypointTrajectory +from autodata_interfaces.env.reset_request import EnvResetRequest from autodata_interfaces.tasks.subtask_constraint_spec import SubTaskConstraintCoordinationScheme, SubTaskConstraintType if TYPE_CHECKING: @@ -209,8 +210,10 @@ def select_source_demo( eef_name: str, eef_pose: torch.Tensor, object_pose: torch.Tensor | None, + object_nodal_positions: torch.Tensor | None, src_demo_current_subtask_boundaries: np.ndarray, subtask_object_name: str | None, + subtask_object_soft: bool, selection_strategy_name: str, selection_strategy_kwargs: dict | None = None, ) -> int: @@ -228,7 +231,12 @@ def select_source_demo( eef_pose=src_ep.eef_pose[eef_name][start_ind:end_ind], object_poses=( {subtask_object_name: src_ep.object_poses[subtask_object_name][start_ind:end_ind]} - if subtask_object_name is not None + if subtask_object_name is not None and not subtask_object_soft + else None + ), + object_nodal_positions=( + {subtask_object_name: src_ep.object_nodal_positions[subtask_object_name][start_ind:end_ind]} + if subtask_object_name is not None and subtask_object_soft else None ), subtask_term_signals=None, @@ -243,6 +251,7 @@ def select_source_demo( eef_pose=eef_pose, object_pose=object_pose, src_subtask_datagen_infos=src_subtask_datagen_infos, + object_nodal_positions=object_nodal_positions, **kwargs, ) @@ -265,9 +274,15 @@ def generate_eef_subtask_trajectory( # Subtask.object_ref is empty string when no object is involved; normalize to None so the # rest of the pipeline can keep using the upstream `is not None` convention. subtask_object_name = subtasks[subtask_ind].object_ref or None + subtask_object_soft = self.algorithm.is_deformable_subtask(subtasks[subtask_ind]) subtask_object_pose = ( self.datastream.get_object_poses(env_ids=[env_id])[subtask_object_name][0] - if subtask_object_name is not None + if subtask_object_name is not None and not subtask_object_soft + else None + ) + subtask_object_nodal_positions = ( + self.datastream.get_object_nodal_positions(env_ids=[env_id])[subtask_object_name][0] + if subtask_object_name is not None and subtask_object_soft else None ) @@ -291,8 +306,10 @@ def generate_eef_subtask_trajectory( eef_name=eef_name, eef_pose=self.datastream.get_robot_eef_pose(env_ids=[env_id], eef_name=eef_name)[0], object_pose=subtask_object_pose, + object_nodal_positions=subtask_object_nodal_positions, src_demo_current_subtask_boundaries=all_randomized_subtask_boundaries[eef_name][:, subtask_ind], subtask_object_name=subtask_object_name, + subtask_object_soft=subtask_object_soft, selection_strategy_name=subtasks[subtask_ind].selection_strategy, selection_strategy_kwargs=subtasks[subtask_ind].selection_strategy_kwargs, ) @@ -325,7 +342,14 @@ def generate_eef_subtask_trajectory( if channel_name == eef_name or channel_name not in eef_names } src_subtask_object_pose = ( - src_ep.object_poses[subtask_object_name][selected_boundary[0]] if subtask_object_name is not None else None + src_ep.object_poses[subtask_object_name][selected_boundary[0]] + if subtask_object_name is not None and not subtask_object_soft + else None + ) + src_subtask_object_nodal_positions = ( + src_ep.object_nodal_positions[subtask_object_name][selected_boundary[0]] + if subtask_object_name is not None and subtask_object_soft + else None ) if is_first_subtask or policy.transform_first_robot_pose: @@ -343,12 +367,15 @@ def generate_eef_subtask_trajectory( for channel_name, channel_tensor in src_subtask_passthrough_actions.items() } - transformed_eef_poses = self._apply_subtask_transform( + transformed_eef_poses = self.algorithm.transform_source_eef_poses( + data_generator=self, eef_name=eef_name, subtask_ind=subtask_ind, subtask_object_name=subtask_object_name, subtask_object_pose=subtask_object_pose, src_subtask_object_pose=src_subtask_object_pose, + subtask_object_nodal_positions=subtask_object_nodal_positions, + src_subtask_object_nodal_positions=src_subtask_object_nodal_positions, src_eef_poses=src_eef_poses, use_delta_transform=use_delta_transform, coord_transform_scheme=coord_transform_scheme, @@ -662,8 +689,9 @@ async def _reset_environment_for_generation( # Recorder + reset queue stay on env. The initial scene state # snapshot is read through the Datastream interface. self.datastream.get_env().recorder_manager.reset(env_ids=env_id_tensor) - await env_reset_queue.put(env_id) - await env_reset_queue.join() + completion = asyncio.get_running_loop().create_future() + await env_reset_queue.put(EnvResetRequest(env_id=env_id, completion=completion)) + await completion return env_id_tensor, self.datastream.get_scene_state(is_relative=True) def _build_runtime_subtask_constraints(self) -> dict: diff --git a/autodata_core/datagen_info.py b/autodata_core/datagen_info.py index 4027841..959b86d 100644 --- a/autodata_core/datagen_info.py +++ b/autodata_core/datagen_info.py @@ -22,6 +22,8 @@ class DatagenInfo: Attributes: eef_pose: ``{eef_name: [T, 4, 4]}`` recorded EEF poses [m, rad]. object_poses: ``{object_name: [T, 4, 4]}`` recorded object poses [m, rad]. + object_nodal_positions: ``{object_name: [T, N, 3]}`` recorded deformable-object + nodal positions [m]. subtask_term_signals: ``{subtask_name: [T]}`` binary completion flag per step. subtask_start_signals: ``{subtask_name: [T]}`` binary start flag per step; required by SkillGen. target_eef_pose: ``{eef_name: [T, 4, 4]}`` controller target poses [m, rad]. @@ -34,6 +36,7 @@ def __init__( self, eef_pose: dict[str, torch.Tensor] | None = None, object_poses: dict[str, torch.Tensor] | None = None, + object_nodal_positions: dict[str, torch.Tensor] | None = None, subtask_term_signals: dict[str, Any] | None = None, subtask_start_signals: dict[str, Any] | None = None, target_eef_pose: dict[str, torch.Tensor] | None = None, @@ -41,6 +44,7 @@ def __init__( ) -> None: self.eef_pose = eef_pose self.object_poses = dict(object_poses) if object_poses is not None else None + self.object_nodal_positions = dict(object_nodal_positions) if object_nodal_positions is not None else None self.subtask_term_signals = dict(subtask_term_signals) if subtask_term_signals is not None else None self.subtask_start_signals = dict(subtask_start_signals) if subtask_start_signals is not None else None self.target_eef_pose = target_eef_pose @@ -53,6 +57,8 @@ def to_dict(self) -> dict[str, Any]: out["eef_pose"] = self.eef_pose if self.object_poses is not None: out["object_poses"] = deepcopy(self.object_poses) + if self.object_nodal_positions is not None: + out["object_nodal_positions"] = deepcopy(self.object_nodal_positions) if self.subtask_start_signals is not None: out["subtask_start_signals"] = deepcopy(self.subtask_start_signals) if self.subtask_term_signals is not None: diff --git a/autodata_core/deformable_transforms.py b/autodata_core/deformable_transforms.py new file mode 100644 index 0000000..f760ae6 --- /dev/null +++ b/autodata_core/deformable_transforms.py @@ -0,0 +1,151 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Non-rigid transforms used by SoftMimicGen.""" + +from __future__ import annotations + +import numpy as np +import torch + +import isaaclab.utils.math as PoseUtils + +from autodata_utils import thin_plate_spline as tps + + +def transform_source_data_segment_using_nodal_registration( + src_eef_poses: torch.Tensor, + src_obj_nodal_pos: torch.Tensor, + tgt_obj_nodal_pos: torch.Tensor, + *, + use_rotation_transform: bool = True, + bend_coef: float = 0.1, + rot_coef: float = 1e-3, +) -> torch.Tensor: + """Warp an EEF pose sequence using TPS registration between deformable-object nodes. + + Args: + src_eef_poses: Source EEF poses shaped ``(T, 4, 4)``. + src_obj_nodal_pos: Source object node positions shaped ``(N, 3)`` [m]. + tgt_obj_nodal_pos: Target object node positions shaped ``(N, 3)`` [m]. + use_rotation_transform: Transform EEF rotations using the local TPS Jacobian. + bend_coef: TPS bending regularization coefficient. + rot_coef: TPS affine-rotation regularization coefficient. + + Returns: + Warped EEF poses shaped ``(T, 4, 4)``. + """ + + assert src_eef_poses.ndim == 3 and src_eef_poses.shape[-2:] == ( + 4, + 4, + ), f"src_eef_poses must have shape (T, 4, 4), got {tuple(src_eef_poses.shape)}" + assert ( + src_obj_nodal_pos.ndim == 2 and src_obj_nodal_pos.shape[1] == 3 + ), f"source nodal positions must have shape (N, 3), got {tuple(src_obj_nodal_pos.shape)}" + assert ( + tgt_obj_nodal_pos.ndim == 2 and tgt_obj_nodal_pos.shape[1] == 3 + ), f"target nodal positions must have shape (N, 3), got {tuple(tgt_obj_nodal_pos.shape)}" + assert src_obj_nodal_pos.shape[0] == tgt_obj_nodal_pos.shape[0], ( + "corresponding-node TPS requires source and target objects to have the same node count: " + f"{src_obj_nodal_pos.shape[0]} != {tgt_obj_nodal_pos.shape[0]}" + ) + + src_nodal_np = src_obj_nodal_pos.detach().cpu().numpy() + tgt_nodal_np = tgt_obj_nodal_pos.detach().cpu().numpy() + linear, translation, weights = tps.fit_reduced( + source_points=src_nodal_np, + target_points=tgt_nodal_np, + bend_coefficient=bend_coef, + rotation_coefficient=rot_coef, + ) + + src_eef_pos, src_eef_rot = PoseUtils.unmake_pose(src_eef_poses) + src_eef_pos_np = src_eef_pos.detach().cpu().numpy() + transformed_pos_np = tps.evaluate( + query_points=src_eef_pos_np, + linear=linear, + translation=translation, + weights=weights, + source_points=src_nodal_np, + ) + + if use_rotation_transform: + jacobians = tps.gradient( + query_points=src_eef_pos_np, + linear=linear, + translation=translation, + weights=weights, + source_points=src_nodal_np, + ) + src_eef_rot_np = src_eef_rot.detach().cpu().numpy() + transformed_rot_np = np.empty_like(src_eef_rot_np) + for pose_index, jacobian in enumerate(jacobians): + transformed_rotation = jacobian @ src_eef_rot_np[pose_index] + u_mat, _, vt_mat = np.linalg.svd(transformed_rotation) + transformed_rotation = u_mat @ vt_mat + if np.linalg.det(transformed_rotation) < 0: + vt_mat[-1, :] *= -1 + transformed_rotation = u_mat @ vt_mat + transformed_rot_np[pose_index] = transformed_rotation + else: + transformed_rot_np = src_eef_rot.detach().cpu().numpy() + + transformed_pos = torch.as_tensor( + transformed_pos_np, + device=src_eef_poses.device, + dtype=src_eef_poses.dtype, + ) + transformed_rot = torch.as_tensor( + transformed_rot_np, + device=src_eef_poses.device, + dtype=src_eef_poses.dtype, + ) + assert torch.isfinite(transformed_pos).all(), "TPS produced non-finite EEF positions" + assert torch.isfinite(transformed_rot).all(), "TPS produced non-finite EEF rotations" + return PoseUtils.make_pose(pos=transformed_pos, rot=transformed_rot) + + +def nodal_registration_cost( + src_obj_nodal_pos: torch.Tensor, + tgt_obj_nodal_pos: torch.Tensor, + *, + bend_coef: float = 0.1, + rot_reg: float = 1e-3, +) -> float: + """Return the TPS fitting cost between two corresponding deformable node sets. + + Args: + src_obj_nodal_pos: Source object node positions shaped ``(N, 3)`` [m]. + tgt_obj_nodal_pos: Target object node positions shaped ``(N, 3)`` [m]. + bend_coef: TPS bending regularization coefficient. + rot_reg: TPS affine-rotation regularization coefficient. + + Returns: + Scalar TPS residual-plus-bending cost. + """ + + assert src_obj_nodal_pos.shape == tgt_obj_nodal_pos.shape, ( + "registration-cost TPS requires matching nodal shapes: " + f"{tuple(src_obj_nodal_pos.shape)} != {tuple(tgt_obj_nodal_pos.shape)}" + ) + assert ( + src_obj_nodal_pos.ndim == 2 and src_obj_nodal_pos.shape[1] == 3 + ), f"nodal positions must have shape (N, 3), got {tuple(src_obj_nodal_pos.shape)}" + + src_nodal_np = src_obj_nodal_pos.detach().cpu().numpy() + tgt_nodal_np = tgt_obj_nodal_pos.detach().cpu().numpy() + linear, translation, weights = tps.fit( + source_points=src_nodal_np, + target_points=tgt_nodal_np, + bend_coefficient=bend_coef, + rotation_coefficient=rot_reg, + ) + return tps.cost( + linear=linear, + translation=translation, + weights=weights, + source_points=src_nodal_np, + target_points=tgt_nodal_np, + bend_coefficient=bend_coef, + ) diff --git a/autodata_core/pool.py b/autodata_core/pool.py index 102b8a2..19ed6fd 100644 --- a/autodata_core/pool.py +++ b/autodata_core/pool.py @@ -124,8 +124,30 @@ def _add_episode(self, episode: EpisodeData) -> None: raise ValueError("Episode lacks 'datagen_info' obs annotations") eef_pose = ep_grp["obs"]["datagen_info"]["eef_pose"] - object_poses_dict = ep_grp["obs"]["datagen_info"]["object_pose"] + object_poses_dict = ep_grp["obs"]["datagen_info"].get("object_pose") + object_nodal_positions_dict = ep_grp["obs"]["datagen_info"].get("object_nodal_position") target_eef_pose = ep_grp["obs"]["datagen_info"]["target_eef_pose"] + + # Legacy single-arm datasets use names such as "robot0". With exactly one + # recorded and one configured EEF, the correspondence is unambiguous. + # Normalize both pose mappings in memory without modifying the source episode. + expected_eefs = set(self.task_descriptor.get_eef_names()) + recorded_eefs = set(eef_pose) + assert recorded_eefs == set(target_eef_pose), ( + f"Source eef_pose names {sorted(recorded_eefs)} do not match " + f"target_eef_pose names {sorted(target_eef_pose)}." + ) + if recorded_eefs != expected_eefs: + assert len(recorded_eefs) == len(expected_eefs) == 1, ( + f"Source EEF names {sorted(recorded_eefs)} do not match configured names " + f"{sorted(expected_eefs)}. Automatic name matching requires exactly one " + "source and one configured EEF; multi-arm datasets require matching names." + ) + source_eef = next(iter(recorded_eefs)) + configured_eef = next(iter(expected_eefs)) + eef_pose = {configured_eef: eef_pose[source_eef]} + target_eef_pose = {configured_eef: target_eef_pose[source_eef]} + subtask_term_signals_dict = ep_grp["obs"]["datagen_info"]["subtask_term_signals"] subtask_start_signals_dict = ep_grp["obs"]["datagen_info"].get("subtask_start_signals") @@ -134,6 +156,7 @@ def _add_episode(self, episode: EpisodeData) -> None: ep_datagen_info = DatagenInfo( eef_pose=eef_pose, object_poses=object_poses_dict, + object_nodal_positions=object_nodal_positions_dict, subtask_start_signals=subtask_start_signals_dict, subtask_term_signals=subtask_term_signals_dict, target_eef_pose=target_eef_pose, diff --git a/autodata_core/selection_strategy.py b/autodata_core/selection_strategy.py index df1442f..c8aefe1 100644 --- a/autodata_core/selection_strategy.py +++ b/autodata_core/selection_strategy.py @@ -10,6 +10,7 @@ from __future__ import annotations import abc +import numpy as np import torch from typing import Any @@ -51,6 +52,7 @@ def select_source_demo( eef_pose: torch.Tensor, object_pose: torch.Tensor | None, src_subtask_datagen_infos: list, + object_nodal_positions: torch.Tensor | None = None, ) -> int: """Return the index of the source demo whose subtask segment best fits the current scene. @@ -58,6 +60,7 @@ def select_source_demo( eef_pose: Current 4x4 EEF pose [m, rad]. object_pose: Current 4x4 pose of the subtask's reference object [m, rad], or None. src_subtask_datagen_infos: Per-source-demo :class:`DatagenInfo` slices covering this subtask. + object_nodal_positions: Current deformable-object nodes shaped ``(N, 3)`` [m], or None. """ raise NotImplementedError @@ -67,7 +70,13 @@ class RandomStrategy(SelectionStrategy): NAME = "random" - def select_source_demo(self, eef_pose, object_pose, src_subtask_datagen_infos) -> int: + def select_source_demo( + self, + eef_pose, + object_pose, + src_subtask_datagen_infos, + object_nodal_positions=None, + ) -> int: n_src_demo = len(src_subtask_datagen_infos) return torch.randint(0, n_src_demo, (1,)).item() @@ -85,6 +94,7 @@ def select_source_demo( pos_weight: float = 1.0, rot_weight: float = 1.0, nn_k: int = 3, + object_nodal_positions: torch.Tensor | None = None, ) -> int: src_object_poses = [] for di in src_subtask_datagen_infos: @@ -126,6 +136,7 @@ def select_source_demo( pos_weight: float = 1.0, rot_weight: float = 1.0, nn_k: int = 3, + object_nodal_positions: torch.Tensor | None = None, ) -> int: src_eef_poses = [] src_object_poses = [] @@ -165,3 +176,55 @@ def select_source_demo( nn_k = min(nn_k, len(dists)) rand_k = torch.randint(0, nn_k, (1,)).item() return torch.argsort(dists)[:nn_k][rand_k] + + +class RegistrationCostStrategy(SelectionStrategy): + """Pick a source segment using deformable-object TPS registration cost.""" + + NAME = "registration_cost" + + def select_source_demo( + self, + eef_pose: torch.Tensor, + object_pose: torch.Tensor | None, + src_subtask_datagen_infos: list, + object_nodal_positions: torch.Tensor | None = None, + bend_coef: float = 0.1, + rot_reg: float = 1e-3, + nn_k: int = 3, + ) -> int: + """Choose uniformly among the ``nn_k`` source shapes with lowest TPS cost.""" + + from autodata_core.deformable_transforms import nodal_registration_cost + + assert object_nodal_positions is not None, "registration_cost requires current object nodal positions" + assert nn_k >= 1, f"nn_k must be at least 1, got {nn_k}" + + valid_candidates: list[tuple[int, float]] = [] + for demo_index, datagen_info in enumerate(src_subtask_datagen_infos): + assert ( + datagen_info.object_nodal_positions is not None + ), "registration_cost requires source object nodal positions" + source_objects = list(datagen_info.object_nodal_positions.values()) + assert ( + len(source_objects) == 1 + ), f"registration_cost expects exactly one source object, got {len(source_objects)}" + source_nodal_positions = source_objects[0][0] + try: + cost = nodal_registration_cost( + source_nodal_positions, + object_nodal_positions, + bend_coef=bend_coef, + rot_reg=rot_reg, + ) + except (AssertionError, np.linalg.LinAlgError, ValueError): + continue + if np.isfinite(cost): + valid_candidates.append((demo_index, cost)) + + assert valid_candidates, "registration_cost could not compute a finite TPS cost for any source demo" + costs_tensor = torch.tensor([cost for _, cost in valid_candidates], dtype=torch.float32) + nn_k = min(nn_k, len(valid_candidates)) + rand_k = torch.randint(0, nn_k, (1,)).item() + selected_candidate = torch.argsort(costs_tensor)[:nn_k][rand_k].item() + return valid_candidates[selected_candidate][0] diff --git a/autodata_interfaces/datastream/datastream.py b/autodata_interfaces/datastream/datastream.py index 9733fa7..12923d0 100644 --- a/autodata_interfaces/datastream/datastream.py +++ b/autodata_interfaces/datastream/datastream.py @@ -271,6 +271,27 @@ def get_object_poses(self, env_ids: Sequence[int] | None = None) -> dict[str, to object_pose_matrix[obj_name] = reference_inverse @ object_pose_matrix[obj_name] return object_pose_matrix + def get_object_nodal_positions(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor]: + """Get all deformable-object nodal positions in the env-relative frame. + + Args: + env_ids: Environments to read; ``None`` reads all. + + Returns: + Mapping from deformable-object name to a tensor shaped + ``(len(env_ids), num_nodes, 3)`` containing positions [m]. + """ + + index: slice | Sequence[int] = slice(None) if env_ids is None else env_ids + scene = self.env.scene + env_origins = scene.env_origins[index] + object_nodal_positions: dict[str, torch.Tensor] = {} + for obj_name, obj in scene.deformable_objects.items(): + nodal_pos_w = obj.data.nodal_pos_w + nodal_pos_w = nodal_pos_w.torch if hasattr(nodal_pos_w, "torch") else as_torch(nodal_pos_w) + object_nodal_positions[obj_name] = nodal_pos_w[index] - env_origins.unsqueeze(1) + return object_nodal_positions + def get_subtask_term_signals( self, env_ids: Sequence[int] | None = None, obs_group: str = "subtask_terms" ) -> dict[str, torch.Tensor]: diff --git a/autodata_interfaces/env/__init__.py b/autodata_interfaces/env/__init__.py index a673b88..392fc38 100644 --- a/autodata_interfaces/env/__init__.py +++ b/autodata_interfaces/env/__init__.py @@ -11,9 +11,11 @@ setup_env_config, setup_output_paths, ) +from autodata_interfaces.env.reset_request import EnvResetRequest __all__ = [ "EnvironmentProfile", + "EnvResetRequest", "apply_env_profile", "get_env_name_from_dataset", "setup_output_paths", diff --git a/autodata_interfaces/env/isaaclab_env_interface.py b/autodata_interfaces/env/isaaclab_env_interface.py index 3a8f916..435cc8d 100644 --- a/autodata_interfaces/env/isaaclab_env_interface.py +++ b/autodata_interfaces/env/isaaclab_env_interface.py @@ -34,6 +34,8 @@ from autodata_interfaces.env.env_profile import EnvironmentProfile, convert_event_params from autodata_interfaces.env.pink_ik_action import PinkInverseKinematicsActionSharedUrdf +from autodata_interfaces.env.reset_request import EnvResetRequest +from autodata_interfaces.env.scene_state import get_scene_state from autodata_interfaces.tasks.generation_policy_spec import GenerationPolicy @@ -264,7 +266,7 @@ def env_loop( Args: env: The environment to run the main step loop on. - env_reset_queue: The asyncio queue carrying per-env reset requests. + env_reset_queue: The asyncio queue carrying :class:`EnvResetRequest` instances. env_action_queue: The asyncio queue carrying ``(env_id, action)`` pairs to execute. asyncio_event_loop: The main asyncio event loop. generation_policy_params: The task descriptor's generation config (source of truth). @@ -277,13 +279,17 @@ def env_loop( """ num_trials = generation_policy_params.num_trials guarantee_success = generation_policy_params.guarantee_success + reset_settling_steps = generation_policy_params.reset_settling_steps + assert reset_settling_steps >= 0, "reset_settling_steps must be non-negative" env_id_tensor = torch.tensor([0], dtype=torch.int64, device=env.device) + settling_requests: dict[int, EnvResetRequest] = {} + settling_steps_remaining: dict[int, int] = {} prev_num_attempts = 0 # simulate environment -- run everything in inference mode with contextlib.suppress(KeyboardInterrupt) and torch.inference_mode(): while True: # check if any environment needs to be reset while waiting for actions - while env_action_queue.qsize() != env.num_envs: + while env_action_queue.qsize() + len(settling_requests) != env.num_envs: asyncio_event_loop.run_until_complete(asyncio.sleep(0)) if data_gen_tasks is not None and data_gen_tasks.done(): exc = data_gen_tasks.exception() @@ -291,25 +297,83 @@ def env_loop( raise exc return False while not env_reset_queue.empty(): - env_id_tensor[0] = env_reset_queue.get_nowait() + request = env_reset_queue.get_nowait() + assert isinstance( + request, EnvResetRequest + ), f"env_reset_queue entries must be EnvResetRequest instances, got {type(request).__name__}" + assert ( + 0 <= request.env_id < env.num_envs + ), f"Reset environment index {request.env_id} is outside [0, {env.num_envs})." + assert ( + request.env_id not in settling_requests + ), f"Environment {request.env_id} already has a reset settling request in progress." + env_id_tensor[0] = request.env_id env.reset(env_ids=env_id_tensor) - env_reset_queue.task_done() + + if reset_settling_steps > 0: + settling_requests[request.env_id] = request + settling_steps_remaining[request.env_id] = reset_settling_steps + else: + request.completion.set_result(None) + env_reset_queue.task_done() + + expected_action_count = env.num_envs - len(settling_requests) + assert env_action_queue.qsize() <= expected_action_count, ( + f"Received {env_action_queue.qsize()} queued actions for only {expected_action_count} " + "non-settling environments." + ) actions = torch.zeros(env.action_space.shape) - # batch-fetch all per-env actions in one gather instead of sequential blocking calls - get_tasks = [env_action_queue.get() for _ in range(env.num_envs)] - results = asyncio_event_loop.run_until_complete(asyncio.gather(*get_tasks)) + # Settling environments use the zero action. Batch-fetch one real action for every + # other environment so their generation trajectories continue without interruption. + action_count = env.num_envs - len(settling_requests) + get_tasks = [env_action_queue.get() for _ in range(action_count)] + results = asyncio_event_loop.run_until_complete(asyncio.gather(*get_tasks)) if get_tasks else [] + action_env_ids: set[int] = set() for env_id, action in results: + assert ( + env_id not in settling_requests + ), f"Environment {env_id} supplied an action while its reset was still settling." + assert env_id not in action_env_ids, f"Received multiple actions for environment {env_id}." + action_env_ids.add(env_id) actions[env_id] = action + expected_action_env_ids = set(range(env.num_envs)) - settling_requests.keys() + assert ( + action_env_ids == expected_action_env_ids + ), f"Expected actions for environments {sorted(expected_action_env_ids)}, got {sorted(action_env_ids)}." # perform action on environment env.step(actions) # mark done so the data generators can continue with the step results - for _ in range(env.num_envs): + for _ in range(action_count): env_action_queue.task_done() + settled_env_ids: list[int] = [] + for env_id in settling_steps_remaining: + settling_steps_remaining[env_id] -= 1 + if settling_steps_remaining[env_id] == 0: + settled_env_ids.append(env_id) + + if settled_env_ids: + settled_env_ids_tensor = torch.tensor(settled_env_ids, dtype=torch.int64, device=env.device) + # Discard temporary settle frames, then capture the settled state as the generated + # episode's initial state before releasing the corresponding generators. + env.recorder_manager.reset(env_ids=settled_env_ids_tensor) + # Settling is not a reset: Arena correctly ignores post-reset callbacks here. + # Snapshot the actual settled state through the public recorder API instead. + env.recorder_manager.add_to_episodes( + "initial_state", + get_scene_state(env.scene, is_relative=True, env_ids=settled_env_ids_tensor), + env_ids=settled_env_ids_tensor, + ) + for env_id in settled_env_ids: + request = settling_requests.pop(env_id) + settling_steps_remaining.pop(env_id) + request.completion.set_result(None) + env_reset_queue.task_done() + if prev_num_attempts != stats["num_attempts"]: prev_num_attempts = stats["num_attempts"] num_success = stats["num_success"] diff --git a/autodata_interfaces/env/reset_request.py b/autodata_interfaces/env/reset_request.py new file mode 100644 index 0000000..f6a7e51 --- /dev/null +++ b/autodata_interfaces/env/reset_request.py @@ -0,0 +1,23 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Reset coordination between asynchronous generators and the synchronous environment loop.""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass + + +@dataclass(frozen=True) +class EnvResetRequest: + """Request an environment reset and wait for data-generation preparation to finish. + + Args: + env_id: Environment index to reset. + completion: Future resolved by the environment loop after reset settling and recorder + initialization are complete. + """ + + env_id: int + completion: asyncio.Future[None] diff --git a/autodata_interfaces/env/scene_state.py b/autodata_interfaces/env/scene_state.py new file mode 100644 index 0000000..29fe0fa --- /dev/null +++ b/autodata_interfaces/env/scene_state.py @@ -0,0 +1,36 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Scene-state reads with optional environment selection.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + + +def get_scene_state(scene: Any, is_relative: bool = True, env_ids: Sequence[int] | None = None) -> dict: + """Return scene state in world or per-environment coordinates. + + Isaac Lab owns the coordinate conversion for every supported asset type. + Selection is applied afterwards to retain the requested environment order. + + Args: + scene: Isaac Lab interactive scene. + is_relative: Whether positions are relative to each environment origin. + env_ids: Environment rows to include, in the requested order; None selects all. + + Returns: + Nested scene-state mapping containing cloned tensors. + """ + + state = scene.get_state(is_relative=is_relative) + if env_ids is None: + return state + + def select_envs(value): + if isinstance(value, dict): + return {key: select_envs(item) for key, item in value.items()} + return value[env_ids] + + return select_envs(state) diff --git a/autodata_interfaces/tasks/generation_policy_spec.py b/autodata_interfaces/tasks/generation_policy_spec.py index 0883ed3..70f4853 100644 --- a/autodata_interfaces/tasks/generation_policy_spec.py +++ b/autodata_interfaces/tasks/generation_policy_spec.py @@ -23,6 +23,9 @@ class GenerationPolicy: use_skillgen: Whether SkillGen is used to generate motion trajectories. use_navigation_controller: Whether a navigation controller generates loco-manipulation trajectories. + reset_settling_steps: Number of ordinary batched environment steps to run after reset + before generation begins. During these steps, the resetting environment receives zero + actions while other environments continue executing their generated actions. max_navigation_hold_steps: Maximum consecutive simulation steps spent waiting for a navigation goal at the end of its recorded segment. Reaching the limit fails the attempt. @@ -49,6 +52,7 @@ class GenerationPolicy: task_name: str | None = None use_skillgen: bool = False use_navigation_controller: bool = False + reset_settling_steps: int = 0 max_navigation_hold_steps: int = 1000 # --- segment stitching --- diff --git a/autodata_interfaces/tasks/subtask_spec.py b/autodata_interfaces/tasks/subtask_spec.py index 3a14e2f..160f4a9 100644 --- a/autodata_interfaces/tasks/subtask_spec.py +++ b/autodata_interfaces/tasks/subtask_spec.py @@ -44,6 +44,23 @@ class SkillGenSubtaskAlgoParams(SubtaskAlgoParams): subtask_start_offset_range: tuple[int, int] = (0, 0) +@dataclass +class SoftMimicGenSubtaskAlgoParams(SubtaskAlgoParams): + """SoftMimicGen-specific subtask parameters. + + Args: + object_soft: Whether the reference object is deformable. + use_rotation_transform: Whether to transform EEF rotations with the local TPS Jacobian. + bend_coef: TPS bending regularization coefficient. + rot_coef: TPS affine-rotation regularization coefficient. + """ + + object_soft: bool = False + use_rotation_transform: bool = True + bend_coef: float = 0.1 + rot_coef: float = 1e-3 + + @dataclass(kw_only=True) class Subtask: """Configuration object used to specify subtasks used in data generation. @@ -94,6 +111,7 @@ class Subtask: "mimicgen": MimicGenSubtaskAlgoParams, "dexmimicgen": DexMimicGenSubtaskAlgoParams, "skillgen": SkillGenSubtaskAlgoParams, + "softmimicgen": SoftMimicGenSubtaskAlgoParams, } """Maps the ``algo:`` discriminator in a YAML task config to the corresponding :class:`SubtaskAlgoParams` subclass. diff --git a/autodata_tests/core/test_algorithms.py b/autodata_tests/core/test_algorithms.py index b3b0867..fb74702 100644 --- a/autodata_tests/core/test_algorithms.py +++ b/autodata_tests/core/test_algorithms.py @@ -11,6 +11,7 @@ GenerationAlgorithm, MimicGen, SkillGen, + SoftMimicGen, get_algorithm, iter_algorithms, ) @@ -20,6 +21,7 @@ def test_registry_contents(): assert REGISTERED_ALGORITHMS == { "mimicgen": MimicGen, "dexmimicgen": DexMimicGen, + "softmimicgen": SoftMimicGen, "skillgen": SkillGen, } @@ -31,10 +33,13 @@ def test_base_class_not_registered(): def test_iter_algorithms(): - assert set(iter_algorithms()) == {MimicGen, DexMimicGen, SkillGen} + assert set(iter_algorithms()) == {MimicGen, DexMimicGen, SoftMimicGen, SkillGen} -@pytest.mark.parametrize("name, cls", [("mimicgen", MimicGen), ("dexmimicgen", DexMimicGen)]) +@pytest.mark.parametrize( + "name, cls", + [("mimicgen", MimicGen), ("dexmimicgen", DexMimicGen), ("softmimicgen", SoftMimicGen)], +) def test_get_algorithm_no_kwargs(name, cls): assert isinstance(get_algorithm(name), cls) @@ -61,6 +66,15 @@ def test_dexmimicgen_attributes(): assert algo.requires_motion_planner is False +def test_softmimicgen_attributes(): + algo = SoftMimicGen() + assert algo.name == "softmimicgen" + assert algo.expected_eef_count == (1, 2) + assert algo.requires_motion_planner is False + assert algo.uses_subtask_start_signals is False + assert algo.supports_coordination is False + + def test_skillgen_attributes(): algo = SkillGen(motion_planners={0: object()}) assert algo.name == "skillgen" diff --git a/autodata_tests/core/test_datagen_info.py b/autodata_tests/core/test_datagen_info.py index e487eb3..306c9f2 100644 --- a/autodata_tests/core/test_datagen_info.py +++ b/autodata_tests/core/test_datagen_info.py @@ -12,6 +12,7 @@ def test_defaults_all_none(): di = DatagenInfo() assert di.eef_pose is None assert di.object_poses is None + assert di.object_nodal_positions is None assert di.subtask_term_signals is None assert di.subtask_start_signals is None assert di.target_eef_pose is None @@ -33,6 +34,7 @@ def test_to_dict_key_names(): di = DatagenInfo( eef_pose={}, object_poses={}, + object_nodal_positions={}, subtask_term_signals={}, subtask_start_signals={}, target_eef_pose={}, @@ -41,6 +43,7 @@ def test_to_dict_key_names(): assert set(di.to_dict()) == { "eef_pose", "object_poses", + "object_nodal_positions", "subtask_term_signals", "subtask_start_signals", "target_eef_pose", @@ -72,3 +75,11 @@ def test_to_dict_deepcopies_object_poses(): out = di.to_dict() out["object_poses"]["cube"] += 1.0 # mutate the returned copy assert torch.count_nonzero(di.object_poses["cube"]) == 0 # original untouched + + +def test_to_dict_deepcopies_object_nodal_positions(): + nodes = {"rope": torch.zeros(2, 8, 3)} + di = DatagenInfo(object_nodal_positions=nodes) + out = di.to_dict() + out["object_nodal_positions"]["rope"] += 1.0 + assert torch.count_nonzero(di.object_nodal_positions["rope"]) == 0 diff --git a/autodata_tests/core/test_deformable_transforms.py b/autodata_tests/core/test_deformable_transforms.py new file mode 100644 index 0000000..c860495 --- /dev/null +++ b/autodata_tests/core/test_deformable_transforms.py @@ -0,0 +1,58 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for SoftMimicGen deformable transforms.""" + +import torch + +import pytest + +from autodata_core.deformable_transforms import ( + nodal_registration_cost, + transform_source_data_segment_using_nodal_registration, +) + +_NODES = torch.tensor( + [ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + [1.0, 1.0, 0.0], + [1.0, 0.0, 1.0], + ], + dtype=torch.float64, +) + + +def _poses() -> torch.Tensor: + poses = torch.eye(4, dtype=torch.float64).repeat(2, 1, 1) + poses[0, :3, 3] = torch.tensor([0.25, 0.25, 0.25], dtype=torch.float64) + poses[1, :3, 3] = torch.tensor([0.75, 0.25, 0.25], dtype=torch.float64) + return poses + + +def test_nodal_tps_translation_transforms_positions_and_preserves_rotations(): + translation = torch.tensor([0.2, -0.1, 0.3], dtype=torch.float64) + source_poses = _poses() + transformed = transform_source_data_segment_using_nodal_registration( + source_poses, + _NODES, + _NODES + translation, + ) + assert torch.allclose(transformed[:, :3, 3], source_poses[:, :3, 3] + translation, atol=1e-6) + assert torch.allclose(transformed[:, :3, :3], source_poses[:, :3, :3], atol=1e-6) + assert torch.allclose(torch.linalg.det(transformed[:, :3, :3]), torch.ones(2, dtype=torch.float64)) + + +def test_nodal_registration_cost_is_near_zero_for_identical_nodes(): + assert nodal_registration_cost(_NODES, _NODES) == pytest.approx(0.0, abs=1e-8) + + +def test_nodal_tps_rejects_mismatched_node_counts_before_fitting(): + with pytest.raises(AssertionError, match="same node count"): + transform_source_data_segment_using_nodal_registration( + _poses(), + _NODES, + _NODES[:-1], + ) diff --git a/autodata_tests/core/test_pool.py b/autodata_tests/core/test_pool.py index 6ca0edd..70f5490 100644 --- a/autodata_tests/core/test_pool.py +++ b/autodata_tests/core/test_pool.py @@ -86,6 +86,27 @@ def _skillgen_task() -> TaskDescriptor: }) +def _softmimicgen_task() -> TaskDescriptor: + return TaskDescriptor.from_dict({ + "name": "rope", + "algo": "softmimicgen", + "subtasks": { + "franka": [ + { + "object_ref": "rope", + "subtask_term_signal": "grasp", + "algo_params": {"object_soft": True}, + }, + { + "object_ref": "rope", + "subtask_term_signal": "", + "algo_params": {"object_soft": True}, + }, + ] + }, + }) + + def _step_signal(length: int, edge: int) -> torch.Tensor: """A per-step step-function: False before ``edge``, True from ``edge`` on.""" sig = torch.ones(length, dtype=torch.bool) @@ -157,6 +178,18 @@ def _skillgen_episode( ) +def _softmimicgen_episode(actions_len: int = 10, term_edge: int = 4) -> types.SimpleNamespace: + return _episode( + actions_len=actions_len, + datagen_info={ + "eef_pose": {"franka": torch.zeros(actions_len, 4, 4)}, + "object_nodal_position": {"rope": torch.zeros(actions_len, 12, 3)}, + "target_eef_pose": {"franka": torch.zeros(actions_len, 4, 4)}, + "subtask_term_signals": {"grasp": _step_signal(actions_len, term_edge)}, + }, + ) + + # --------------------------------------------------------------------------------------------------- # __init__: subtask signal-name / offset-range tables built from the task descriptor # --------------------------------------------------------------------------------------------------- @@ -192,6 +225,16 @@ def test_add_episode_populates_datagen_info_and_passthrough(): assert di.passthrough_action["franka"].shape == (8, 1) +def test_add_episode_accepts_deformable_state_without_rigid_object_pose(): + pool = _pool(_softmimicgen_task()) + pool._add_episode(_softmimicgen_episode(actions_len=8, term_edge=3)) + di = pool.datagen_infos[0] + assert di.object_poses is None + assert set(di.object_nodal_positions) == {"rope"} + assert di.object_nodal_positions["rope"].shape == (8, 12, 3) + assert pool.subtask_boundaries["franka"] == [[(0, 4), (4, 8)]] + + def test_add_episode_missing_datagen_info_error(): pool = _pool(_mimicgen_task()) bad = types.SimpleNamespace(data={"actions": torch.zeros(5, 7), "obs": {}}) @@ -266,3 +309,32 @@ def test_add_episode_skillgen_rejects_start_offset_that_empties_subtask(): pool = _pool(task, uses_start_signals=True) with pytest.raises(AssertionError, match="empty in worst case"): pool._add_episode(_skillgen_episode(actions_len=10, start_edges=(1, 6), term_edges=(4, 8))) + + +def test_add_episode_maps_legacy_single_arm_names_without_mutating_source(): + pool = _pool(_softmimicgen_task()) + episode = _softmimicgen_episode() + info = episode.data["obs"]["datagen_info"] + for key in ("eef_pose", "target_eef_pose"): + info[key]["robot0"] = info[key].pop("franka") + + pool._add_episode(episode) + recorded = pool.datagen_infos[0] + for key in ("eef_pose", "target_eef_pose"): + assert set(info[key]) == {"robot0"} + assert set(getattr(recorded, key)) == {"franka"} + torch.testing.assert_close(getattr(recorded, key)["franka"], info[key]["robot0"]) + + +def test_add_episode_rejects_inconsistent_pose_names(): + episode = _softmimicgen_episode() + info = episode.data["obs"]["datagen_info"] + info["eef_pose"]["robot0"] = info["eef_pose"].pop("franka") + with pytest.raises(AssertionError, match="do not match target_eef_pose"): + _pool(_softmimicgen_task())._add_episode(episode) + + +def test_add_episode_rejects_ambiguous_arm_name_mapping(): + episode = _dexmimicgen_episode() + with pytest.raises(AssertionError, match="multi-arm datasets require matching names"): + _pool(_softmimicgen_task())._add_episode(episode) diff --git a/autodata_tests/core/test_selection_strategy.py b/autodata_tests/core/test_selection_strategy.py index 3cfbaaf..d102269 100644 --- a/autodata_tests/core/test_selection_strategy.py +++ b/autodata_tests/core/test_selection_strategy.py @@ -13,6 +13,7 @@ NearestNeighborObjectStrategy, NearestNeighborRobotDistanceStrategy, RandomStrategy, + RegistrationCostStrategy, make_selection_strategy, ) @@ -34,6 +35,7 @@ def test_registry_contents(): "random", "nearest_neighbor_object", "nearest_neighbor_robot_distance", + "registration_cost", } @@ -109,3 +111,67 @@ def test_nearest_neighbor_robot_distance_transforms_source_eef_into_current_obje ) assert int(index) == 0 + + +def test_registration_cost_picks_lowest_cost(monkeypatch): + import autodata_core.deformable_transforms as deformable_transforms + + monkeypatch.setattr( + deformable_transforms, + "nodal_registration_cost", + lambda source, target, **kwargs: float(torch.linalg.vector_norm(source - target)), + ) + current = torch.zeros(6, 3) + infos = [DatagenInfo(object_nodal_positions={"rope": torch.full((1, 6, 3), value)}) for value in (2.0, 0.0, 1.0)] + index = RegistrationCostStrategy().select_source_demo( + None, + None, + infos, + object_nodal_positions=current, + nn_k=1, + ) + assert int(index) == 1 + + +def test_registration_cost_skips_failed_and_nonfinite_candidates(monkeypatch): + import autodata_core.deformable_transforms as deformable_transforms + + def registration_cost(source, target, **kwargs): + del target, kwargs + source_value = source[0, 0].item() + if source_value == 0.0: + raise ValueError("invalid TPS input") + if source_value == 1.0: + return float("nan") + return 1.0 + + monkeypatch.setattr(deformable_transforms, "nodal_registration_cost", registration_cost) + current = torch.zeros(6, 3) + infos = [DatagenInfo(object_nodal_positions={"rope": torch.full((1, 6, 3), value)}) for value in (0.0, 1.0, 2.0)] + + index = RegistrationCostStrategy().select_source_demo( + None, + None, + infos, + object_nodal_positions=current, + nn_k=3, + ) + + assert index == 2 + + +def test_registration_cost_rejects_all_invalid_candidates(monkeypatch): + import autodata_core.deformable_transforms as deformable_transforms + + monkeypatch.setattr(deformable_transforms, "nodal_registration_cost", lambda source, target, **kwargs: float("inf")) + current = torch.zeros(6, 3) + infos = [DatagenInfo(object_nodal_positions={"rope": torch.zeros(1, 6, 3)})] + + with pytest.raises(AssertionError, match="could not compute a finite TPS cost"): + RegistrationCostStrategy().select_source_demo( + None, + None, + infos, + object_nodal_positions=current, + nn_k=1, + ) diff --git a/autodata_tests/interfaces/datastream/test_datastream.py b/autodata_tests/interfaces/datastream/test_datastream.py index 56486fa..36f36b7 100644 --- a/autodata_tests/interfaces/datastream/test_datastream.py +++ b/autodata_tests/interfaces/datastream/test_datastream.py @@ -57,9 +57,15 @@ def _env(env_origins: torch.Tensor | None = None) -> MockEnv: root_quat_w=torch.tensor(_IDENTITY_QUAT_XYZW), ) ) + rope = MockAsset( + MockArticulationData( + nodal_pos_w=torch.tensor([[[1.0, 1.0, 0.05], [1.2, 1.0, 0.05]]]), + ) + ) scene = MockScene( assets={"robot": robot}, rigid_objects={"cube_1": cube}, + deformable_objects={"rope": rope}, env_origins=torch.tensor([[0.0, 0.0, 0.0]]) if env_origins is None else env_origins, state=_SCENE_STATE_SENTINEL, ) @@ -232,6 +238,23 @@ def test_get_object_poses_subtracts_env_origin(): assert torch.allclose(pose[0, :3, 3], torch.tensor([0.0, 0.0, 0.05])) +def test_get_object_nodal_positions_is_env_relative(): + task = TaskDescriptor.from_dict(_task_dict()) + embodiment_adapter = _embodiment_adapter() + env = _env(env_origins=torch.tensor([[1.0, 1.0, 0.0]])) + pool = DataGenInfoPool( + task_descriptor=task, + embodiment_adapter=embodiment_adapter, + device=env.device, + uses_start_signals=False, + ) + datastream = Datastream(env=env, task_descriptor=task, embodiment_adapter=embodiment_adapter, source_pool=pool) + nodes = datastream.get_object_nodal_positions()["rope"] + assert nodes.shape == (1, 2, 3) + assert torch.allclose(nodes[0, 0], torch.tensor([0.0, 0.0, 0.05])) + assert torch.allclose(nodes[0, 1], torch.tensor([0.2, 0.0, 0.05])) + + def test_articulated_object_poses_use_controller_frame(): """A moving controller frame must transform articulated roots just like rigid objects.""" datastream, _, adapter, env = _datastream() diff --git a/autodata_tests/interfaces/env/test_env_loop.py b/autodata_tests/interfaces/env/test_env_loop.py new file mode 100644 index 0000000..6e497f4 --- /dev/null +++ b/autodata_tests/interfaces/env/test_env_loop.py @@ -0,0 +1,161 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for reset settling in the synchronous data-generation environment loop.""" + +from __future__ import annotations + +import asyncio +import contextlib +import torch +from types import SimpleNamespace + +import pytest + +from autodata_interfaces.env.isaaclab_env_interface import env_loop +from autodata_interfaces.env.reset_request import EnvResetRequest +from autodata_interfaces.tasks.generation_policy_spec import GenerationPolicy + + +class _MockRecorderManager: + def __init__(self) -> None: + self.reset_env_ids: list[tuple[int, ...]] = [] + self.post_reset_env_ids: list[tuple[int, ...]] = [] + self.initial_states: list[tuple[tuple[int, ...], dict]] = [] + + def reset(self, env_ids: torch.Tensor) -> None: + self.reset_env_ids.append(tuple(env_ids.tolist())) + + def record_post_reset(self, env_ids: torch.Tensor) -> None: + # Arena ignores this callback outside an actual reset. + self.post_reset_env_ids.append(tuple(env_ids.tolist())) + + def add_to_episodes(self, key: str, value: dict, env_ids: torch.Tensor) -> None: + assert key == "initial_state" + self.initial_states.append((tuple(env_ids.tolist()), value)) + + +class _MockSimulation: + def is_stopped(self) -> bool: + return False + + +class _MockScene: + def __init__(self, num_envs: int) -> None: + self.env_origins = torch.arange(num_envs * 3, dtype=torch.float32).reshape(num_envs, 3) + self.nodes = torch.zeros((num_envs, 5, 3)) + + def get_state(self, is_relative: bool) -> dict: + nodes = self.nodes.clone() + if not is_relative: + nodes += self.env_origins[:, None] + return {"deformable_object": {"object": {"nodal_position": nodes}}} + + +class _MockEnv: + def __init__(self, num_envs: int) -> None: + self.num_envs = num_envs + self.device = "cpu" + self.action_space = SimpleNamespace(shape=(num_envs, 1)) + self.recorder_manager = _MockRecorderManager() + self.sim = _MockSimulation() + self.scene = _MockScene(num_envs) + self.reset_env_ids: list[tuple[int, ...]] = [] + self.actions: list[torch.Tensor] = [] + + def reset(self, env_ids: torch.Tensor) -> None: + self.reset_env_ids.append(tuple(env_ids.tolist())) + + def step(self, actions: torch.Tensor) -> None: + self.actions.append(actions.clone()) + self.scene.nodes[:, :, 2] += 0.01 + + +def _close_event_loop(loop: asyncio.AbstractEventLoop, tasks: asyncio.Future) -> None: + if not tasks.done(): + tasks.cancel() + with contextlib.suppress(asyncio.CancelledError): + loop.run_until_complete(tasks) + loop.close() + asyncio.set_event_loop(None) + + +def test_reset_without_settling_completes_without_environment_steps(): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + env = _MockEnv(num_envs=1) + reset_queue: asyncio.Queue = asyncio.Queue() + action_queue: asyncio.Queue = asyncio.Queue() + completion = loop.create_future() + reset_queue.put_nowait(EnvResetRequest(env_id=0, completion=completion)) + + async def wait_for_reset() -> None: + await completion + + tasks = asyncio.gather(wait_for_reset()) + try: + completed = env_loop( + env=env, + env_reset_queue=reset_queue, + env_action_queue=action_queue, + asyncio_event_loop=loop, + generation_policy_params=GenerationPolicy(reset_settling_steps=0), + stats={"num_success": 0, "num_failures": 0, "num_attempts": 0}, + data_gen_tasks=tasks, + ) + assert completed is False + assert env.reset_env_ids == [(0,)] + assert env.actions == [] + assert env.recorder_manager.reset_env_ids == [] + assert env.recorder_manager.post_reset_env_ids == [] + loop.run_until_complete(asyncio.wait_for(reset_queue.join(), timeout=0.1)) + finally: + _close_event_loop(loop, tasks) + + +@pytest.mark.parametrize("settling_env_id", [0, 1]) +def test_settling_env_uses_zero_actions_while_other_env_continues(settling_env_id): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + env = _MockEnv(num_envs=2) + running_env_id = 1 - settling_env_id + reset_queue: asyncio.Queue = asyncio.Queue() + action_queue: asyncio.Queue = asyncio.Queue() + completion = loop.create_future() + reset_queue.put_nowait(EnvResetRequest(env_id=settling_env_id, completion=completion)) + + async def provide_running_env_actions() -> None: + for value in (3.0, 4.0): + await action_queue.put((running_env_id, torch.tensor([value]))) + await action_queue.join() + + async def wait_for_reset() -> None: + await completion + + tasks = asyncio.gather(provide_running_env_actions(), wait_for_reset()) + try: + completed = env_loop( + env=env, + env_reset_queue=reset_queue, + env_action_queue=action_queue, + asyncio_event_loop=loop, + generation_policy_params=GenerationPolicy(reset_settling_steps=2), + stats={"num_success": 0, "num_failures": 0, "num_attempts": 0}, + data_gen_tasks=tasks, + ) + assert completed is False + assert env.reset_env_ids == [(settling_env_id,)] + assert len(env.actions) == 2 + assert [action[settling_env_id, 0].item() for action in env.actions] == [0.0, 0.0] + assert [action[running_env_id, 0].item() for action in env.actions] == [3.0, 4.0] + assert env.recorder_manager.reset_env_ids == [(settling_env_id,)] + assert env.recorder_manager.post_reset_env_ids == [] + assert len(env.recorder_manager.initial_states) == 1 + ids, state = env.recorder_manager.initial_states[0] + assert ids == (settling_env_id,) + nodes = state["deformable_object"]["object"]["nodal_position"] + torch.testing.assert_close(nodes, env.scene.nodes[[settling_env_id]]) + torch.testing.assert_close(nodes[:, :, 2], torch.full((1, 5), 0.02)) + loop.run_until_complete(asyncio.wait_for(reset_queue.join(), timeout=0.1)) + finally: + _close_event_loop(loop, tasks) diff --git a/autodata_tests/interfaces/mocks.py b/autodata_tests/interfaces/mocks.py index 83359ae..0443977 100644 --- a/autodata_tests/interfaces/mocks.py +++ b/autodata_tests/interfaces/mocks.py @@ -18,11 +18,13 @@ def __init__( joint_names: list[str] | None = None, root_pos_w: torch.Tensor | None = None, root_quat_w: torch.Tensor | None = None, + nodal_pos_w: torch.Tensor | None = None, ) -> None: self.joint_pos = joint_pos self.joint_names = joint_names self.root_pos_w = root_pos_w self.root_quat_w = root_quat_w + self.nodal_pos_w = nodal_pos_w class MockAsset: @@ -39,11 +41,13 @@ def __init__( self, assets: dict[str, MockAsset] | None = None, rigid_objects: dict[str, MockAsset] | None = None, + deformable_objects: dict[str, MockAsset] | None = None, env_origins: torch.Tensor | None = None, state: Any = None, ) -> None: self._assets = assets or {} self.rigid_objects = rigid_objects or {} + self.deformable_objects = deformable_objects or {} self.env_origins = env_origins self._state = state diff --git a/autodata_tests/interfaces/tasks/test_generation_policy_spec.py b/autodata_tests/interfaces/tasks/test_generation_policy_spec.py index 8c78c7e..c4a430f 100644 --- a/autodata_tests/interfaces/tasks/test_generation_policy_spec.py +++ b/autodata_tests/interfaces/tasks/test_generation_policy_spec.py @@ -20,6 +20,7 @@ def test_generation_policy_defaults(): assert p.task_name is None assert p.use_skillgen is False assert p.use_navigation_controller is False + assert p.reset_settling_steps == 0 assert p.select_src_per_subtask is False assert p.select_src_per_arm is False assert p.transform_first_robot_pose is False @@ -27,11 +28,12 @@ def test_generation_policy_defaults(): def test_generation_policy_from_kwargs(): - p = GenerationPolicy(name="run", seed=42, num_trials=100, use_skillgen=True) + p = GenerationPolicy(name="run", seed=42, num_trials=100, use_skillgen=True, reset_settling_steps=10) assert p.name == "run" assert p.seed == 42 assert p.num_trials == 100 assert p.use_skillgen is True + assert p.reset_settling_steps == 10 # Unspecified fields keep their defaults. assert p.guarantee_success is True assert p.interpolate_from_last_target_pose is True diff --git a/autodata_tests/interfaces/tasks/test_subtask_spec.py b/autodata_tests/interfaces/tasks/test_subtask_spec.py index 54ff041..ff98aaf 100644 --- a/autodata_tests/interfaces/tasks/test_subtask_spec.py +++ b/autodata_tests/interfaces/tasks/test_subtask_spec.py @@ -10,6 +10,7 @@ DexMimicGenSubtaskAlgoParams, MimicGenSubtaskAlgoParams, SkillGenSubtaskAlgoParams, + SoftMimicGenSubtaskAlgoParams, Subtask, SubtaskAlgoParams, ) @@ -81,6 +82,7 @@ def test_algo_params_registry_contents(): "mimicgen": MimicGenSubtaskAlgoParams, "dexmimicgen": DexMimicGenSubtaskAlgoParams, "skillgen": SkillGenSubtaskAlgoParams, + "softmimicgen": SoftMimicGenSubtaskAlgoParams, } for cls in ALGO_PARAMS_REGISTRY.values(): assert issubclass(cls, SubtaskAlgoParams) @@ -91,6 +93,15 @@ def test_skillgen_algo_params_default_and_override(): assert SkillGenSubtaskAlgoParams(subtask_start_offset_range=(1, 3)).subtask_start_offset_range == (1, 3) +def test_softmimicgen_algo_params_defaults_and_override(): + defaults = SoftMimicGenSubtaskAlgoParams() + assert defaults.object_soft is False + assert defaults.use_rotation_transform is True + assert defaults.bend_coef == 0.1 + assert defaults.rot_coef == 1e-3 + assert SoftMimicGenSubtaskAlgoParams(object_soft=True).object_soft is True + + @pytest.mark.parametrize("cls", [MimicGenSubtaskAlgoParams, DexMimicGenSubtaskAlgoParams]) def test_plain_algo_params_have_no_extra_fields(cls): # MimicGen / DexMimicGen carry no params beyond the shared ones on Subtask. diff --git a/autodata_tests/utils/test_thin_plate_spline.py b/autodata_tests/utils/test_thin_plate_spline.py new file mode 100644 index 0000000..7683ffc --- /dev/null +++ b/autodata_tests/utils/test_thin_plate_spline.py @@ -0,0 +1,86 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Regression tests for AutoData's three-dimensional TPS implementation.""" + +import numpy as np + +from autodata_utils import thin_plate_spline as tps + +_SOURCE_POINTS = np.array([ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + [1.0, 1.0, 0.0], + [1.0, 0.0, 1.0], +]) +_TARGET_POINTS = _SOURCE_POINTS.copy() +_TARGET_POINTS[0] += [0.10, -0.05, 0.02] +_TARGET_POINTS[1] += [0.20, 0.10, -0.03] +_TARGET_POINTS[4] += [-0.10, 0.15, 0.05] +_QUERY_POINTS = np.array([ + [0.25, 0.25, 0.25], + [0.75, 0.25, 0.25], +]) + + +def test_reduced_fit_matches_previous_tps_reference_outputs(): + linear, translation, weights = tps.fit_reduced( + source_points=_SOURCE_POINTS, + target_points=_TARGET_POINTS, + bend_coefficient=0.1, + rotation_coefficient=1e-3, + ) + + transformed = tps.evaluate( + query_points=_QUERY_POINTS, + linear=linear, + translation=translation, + weights=weights, + source_points=_SOURCE_POINTS, + ) + expected_transformed = np.array([ + [0.303313771813, 0.251697797185, 0.258325640304], + [0.821598859240, 0.310808443454, 0.249183096591], + ]) + np.testing.assert_allclose(transformed, expected_transformed, rtol=1e-10, atol=1e-10) + + jacobians = tps.gradient( + query_points=_QUERY_POINTS, + linear=linear, + translation=translation, + weights=weights, + source_points=_SOURCE_POINTS, + ) + expected_jacobians = np.array([ + [ + [1.032732261752, -0.163893736190, -0.132353662016], + [0.116302039264, 1.049445488374, 0.002135377114], + [-0.016366130876, 0.012051710831, 0.996281673744], + ], + [ + [1.032732261752, -0.235557062625, -0.167147111810], + [0.116302039264, 1.050479586533, -0.052135339689], + [-0.016366130876, 0.047883374048, 1.013678398641], + ], + ]) + np.testing.assert_allclose(jacobians, expected_jacobians, rtol=1e-10, atol=1e-10) + + +def test_regularized_fit_cost_matches_previous_tps_reference_output(): + linear, translation, weights = tps.fit( + source_points=_SOURCE_POINTS, + target_points=_TARGET_POINTS, + bend_coefficient=0.1, + rotation_coefficient=1e-3, + ) + registration_cost = tps.cost( + linear=linear, + translation=translation, + weights=weights, + source_points=_SOURCE_POINTS, + target_points=_TARGET_POINTS, + bend_coefficient=0.1, + ) + np.testing.assert_allclose(registration_cost, 0.0073160059705457805, rtol=1e-10, atol=1e-10) diff --git a/autodata_utils/thin_plate_spline.py b/autodata_utils/thin_plate_spline.py new file mode 100644 index 0000000..01235a6 --- /dev/null +++ b/autodata_utils/thin_plate_spline.py @@ -0,0 +1,263 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Three-dimensional thin-plate-spline fitting and evaluation. + +This utility module contains the small numerical surface SoftMimicGen needs for deformable-object +registration. It intentionally avoids the unrelated robotics stack that accompanied the former +Rapprentice dependency. +""" + +from __future__ import annotations + +import numpy as np + +_POINT_DIMENSION = 3 + + +def _points(points: np.ndarray, name: str) -> np.ndarray: + """Return finite three-dimensional points as double-precision values.""" + + points = np.asarray(points, dtype=np.float64) + assert ( + points.ndim == 2 and points.shape[1] == _POINT_DIMENSION + ), f"{name} must have shape (N, {_POINT_DIMENSION}), got {points.shape}" + assert np.isfinite(points).all(), f"{name} contains non-finite values" + return points + + +def _corresponding_points(source_points: np.ndarray, target_points: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Validate and return corresponding source and target points.""" + + source_points = _points(source_points, "source_points") + target_points = _points(target_points, "target_points") + assert ( + source_points.shape == target_points.shape + ), f"source_points and target_points must have matching shapes, got {source_points.shape} and {target_points.shape}" + assert ( + source_points.shape[0] >= _POINT_DIMENSION + 1 + ), f"TPS requires at least {_POINT_DIMENSION + 1} corresponding points, got {source_points.shape[0]}" + return source_points, target_points + + +def _kernel_matrix(left_points: np.ndarray, right_points: np.ndarray) -> np.ndarray: + """Return the three-dimensional TPS kernel ``K(r) = -r``.""" + + differences = left_points[:, None, :] - right_points[None, :, :] + return -np.linalg.norm(differences, axis=-1) + + +def _transform_parameters( + linear: np.ndarray, + translation: np.ndarray, + weights: np.ndarray, + source_point_count: int, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Validate and return TPS transform parameters.""" + + linear = np.asarray(linear, dtype=np.float64) + translation = np.asarray(translation, dtype=np.float64) + weights = np.asarray(weights, dtype=np.float64) + assert linear.shape == ( + _POINT_DIMENSION, + _POINT_DIMENSION, + ), f"linear must have shape ({_POINT_DIMENSION}, {_POINT_DIMENSION}), got {linear.shape}" + assert translation.shape == ( + _POINT_DIMENSION, + ), f"translation must have shape ({_POINT_DIMENSION},), got {translation.shape}" + assert weights.shape == ( + source_point_count, + _POINT_DIMENSION, + ), f"weights must have shape ({source_point_count}, {_POINT_DIMENSION}), got {weights.shape}" + return linear, translation, weights + + +def fit( + source_points: np.ndarray, + target_points: np.ndarray, + bend_coefficient: float, + rotation_coefficient: float, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Fit a regularized TPS transform between corresponding points. + + Args: + source_points: Source control points shaped ``(N, 3)``. + target_points: Target control points shaped ``(N, 3)``. + bend_coefficient: TPS bending regularization coefficient. + rotation_coefficient: Affine-rotation regularization coefficient. + + Returns: + A ``(linear, translation, weights)`` tuple describing the fitted transform. + """ + + source_points, target_points = _corresponding_points(source_points, target_points) + point_count, dimension = source_points.shape + kernel = _kernel_matrix(source_points, source_points) + rotation_ratio = bend_coefficient / rotation_coefficient if rotation_coefficient > 0 else 0.0 + + system = np.zeros((point_count + dimension + 1, point_count + dimension + 1)) + system[:point_count, :point_count] = kernel + diagonal = np.arange(point_count) + system[diagonal, diagonal] += bend_coefficient + system[:point_count, point_count : point_count + dimension] = source_points + system[:point_count, point_count + dimension] = 1.0 + system[point_count : point_count + dimension, :point_count] = source_points.T + system[point_count + dimension, :point_count] = 1.0 + system[ + point_count : point_count + dimension, + point_count : point_count + dimension, + ] = rotation_ratio * np.eye(dimension) + + right_hand_side = np.empty((point_count + dimension + 1, dimension)) + right_hand_side[:point_count] = target_points + right_hand_side[point_count : point_count + dimension] = rotation_ratio * np.eye(dimension) + right_hand_side[point_count + dimension] = 0.0 + + solution = np.linalg.solve(system, right_hand_side) + weights = solution[:point_count] + linear = solution[point_count : point_count + dimension] + translation = solution[point_count + dimension] + return linear, translation, weights + + +def fit_reduced( + source_points: np.ndarray, + target_points: np.ndarray, + bend_coefficient: float, + rotation_coefficient: float, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Fit a TPS transform in the null space of its affine constraints. + + Args: + source_points: Source control points shaped ``(N, 3)``. + target_points: Target control points shaped ``(N, 3)``. + bend_coefficient: TPS bending regularization coefficient. + rotation_coefficient: Affine-rotation regularization coefficient. + + Returns: + A ``(linear, translation, weights)`` tuple describing the fitted transform. + """ + + source_points, target_points = _corresponding_points(source_points, target_points) + point_count, dimension = source_points.shape + affine_column_count = dimension + 1 + affine_basis = np.column_stack((source_points, np.ones(point_count))) + left_singular_vectors, _, _ = np.linalg.svd(affine_basis, full_matrices=True) + null_space = left_singular_vectors[:, affine_column_count:] + + kernel = _kernel_matrix(source_points, source_points) + design = np.column_stack((source_points, np.ones(point_count), kernel @ null_space)) + system = design.T @ design + system[affine_column_count:, affine_column_count:] += bend_coefficient * null_space.T @ kernel @ null_space + right_hand_side = design.T @ target_points + system[:dimension, :dimension] += rotation_coefficient * np.eye(dimension) + right_hand_side[:dimension, :dimension] += rotation_coefficient * np.eye(dimension) + + solution = np.linalg.solve(system, right_hand_side) + linear = solution[:dimension] + translation = solution[dimension] + weights = null_space @ solution[affine_column_count:] + return linear, translation, weights + + +def evaluate( + query_points: np.ndarray, + linear: np.ndarray, + translation: np.ndarray, + weights: np.ndarray, + source_points: np.ndarray, +) -> np.ndarray: + """Evaluate a fitted TPS transform at query points. + + Args: + query_points: Points to transform, shaped ``(M, 3)``. + linear: Affine linear component, shaped ``(3, 3)``. + translation: Affine translation component, shaped ``(3,)``. + weights: Non-rigid kernel weights, shaped ``(N, 3)``. + source_points: Source control points, shaped ``(N, 3)``. + + Returns: + Transformed points shaped ``(M, 3)``. + """ + + query_points = _points(query_points, "query_points") + source_points = _points(source_points, "source_points") + linear, translation, weights = _transform_parameters( + linear, + translation, + weights, + source_points.shape[0], + ) + kernel = _kernel_matrix(query_points, source_points) + return kernel @ weights + query_points @ linear + translation[None, :] + + +def gradient( + query_points: np.ndarray, + linear: np.ndarray, + translation: np.ndarray, + weights: np.ndarray, + source_points: np.ndarray, +) -> np.ndarray: + """Evaluate the local Jacobian of a fitted TPS transform. + + Args: + query_points: Points at which to evaluate the Jacobian, shaped ``(M, 3)``. + linear: Affine linear component, shaped ``(3, 3)``. + translation: Affine translation component, shaped ``(3,)``. + weights: Non-rigid kernel weights, shaped ``(N, 3)``. + source_points: Source control points, shaped ``(N, 3)``. + + Returns: + Transform Jacobians shaped ``(M, 3, 3)``. + """ + + query_points = _points(query_points, "query_points") + source_points = _points(source_points, "source_points") + linear, _, weights = _transform_parameters(linear, translation, weights, source_points.shape[0]) + differences = query_points[:, None, :] - source_points[None, :, :] + distances = np.linalg.norm(differences, axis=-1) + directions = np.divide( + differences, + distances[:, :, None], + out=np.zeros_like(differences), + where=distances[:, :, None] != 0.0, + ) + kernel_gradient = np.einsum("mna,ng->mga", directions, weights) + return linear.T[None, :, :] - kernel_gradient + + +def cost( + linear: np.ndarray, + translation: np.ndarray, + weights: np.ndarray, + source_points: np.ndarray, + target_points: np.ndarray, + bend_coefficient: float, +) -> float: + """Return the residual-plus-bending cost of a fitted TPS transform. + + Args: + linear: Affine linear component, shaped ``(3, 3)``. + translation: Affine translation component, shaped ``(3,)``. + weights: Non-rigid kernel weights, shaped ``(N, 3)``. + source_points: Source control points, shaped ``(N, 3)``. + target_points: Target control points, shaped ``(N, 3)``. + bend_coefficient: TPS bending regularization coefficient. + + Returns: + Scalar residual-plus-bending cost. + """ + + source_points, target_points = _corresponding_points(source_points, target_points) + linear, translation, weights = _transform_parameters( + linear, + translation, + weights, + source_points.shape[0], + ) + kernel = _kernel_matrix(source_points, source_points) + predicted_points = kernel @ weights + source_points @ linear + translation[None, :] + residual_cost = np.square(predicted_points - target_points).sum() + bending_cost = bend_coefficient * np.sum(weights * (kernel @ weights)) + return float(residual_cost + bending_cost) diff --git a/scripts/annotate_demos.py b/scripts/annotate_demos.py index b0d6b2b..ea522d6 100644 --- a/scripts/annotate_demos.py +++ b/scripts/annotate_demos.py @@ -159,6 +159,7 @@ def record_pre_step(self): assert _datastream is not None, "Datastream must be initialized before recording." datagen_info = { "object_pose": _datastream.get_object_poses(), + "object_nodal_position": _datastream.get_object_nodal_positions(), "eef_pose": _datastream.embodiment_adapter.get_eef_poses(env_ids=None), "target_eef_pose": _datastream.action_to_target_eef_pose(self._env.action_manager.action), } diff --git a/scripts/generate_dataset.py b/scripts/generate_dataset.py index d6f5167..2cbbbd2 100755 --- a/scripts/generate_dataset.py +++ b/scripts/generate_dataset.py @@ -8,7 +8,7 @@ python scripts/generate_dataset.py \\ --env_name \\ - [--alg {mimicgen|dexmimicgen|skillgen}] \\ + [--alg {mimicgen|dexmimicgen|skillgen|softmimicgen}] \\ --task_descriptor \\ --embodiment \\ --env_profile \\ @@ -25,6 +25,7 @@ * ``skillgen`` — single-arm SkillGen. SkillGen depends on a motion-planner interface; until the planner code is ported into this repo, the CLI satisfies that interface with the upstream Arena ``CuroboPlanner``. +* ``softmimicgen`` — one- or two-arm MimicGen with deformable-object nodal registration. The CLI composes a :class:`Datastream` from the task descriptor YAML, the embodiment YAML, the live env, and the HDF5 source dataset, then hands it to :class:`DataGenerator`. @@ -44,7 +45,7 @@ # Hardcoded to keep argparse importable without pulling in the heavy core package. # Add new algorithms here when registering them in autodata_core.algorithms. -_ALG_CHOICES = ["mimicgen", "dexmimicgen", "skillgen"] +_ALG_CHOICES = ["mimicgen", "dexmimicgen", "skillgen", "softmimicgen"] parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) parser.add_argument( From ba2658e88deb29d0fdf3d7976905bfd20a41784d Mon Sep 17 00:00:00 2001 From: Peter Du Date: Wed, 21 Oct 2026 10:47:10 -0700 Subject: [PATCH 2/2] address comments --- autodata_core/algorithms.py | 14 ++++++- autodata_core/selection_strategy.py | 11 ++++- .../env/isaaclab_env_interface.py | 16 +++++++- .../tasks/generation_policy_spec.py | 4 +- autodata_tests/core/test_algorithms.py | 40 +++++++++++++++++++ .../interfaces/env/test_env_loop.py | 32 ++++++++++++++- 6 files changed, 109 insertions(+), 8 deletions(-) diff --git a/autodata_core/algorithms.py b/autodata_core/algorithms.py index 9814430..9e09267 100644 --- a/autodata_core/algorithms.py +++ b/autodata_core/algorithms.py @@ -19,6 +19,8 @@ from collections.abc import Iterator from typing import TYPE_CHECKING, Any +from autodata_core.selection_strategy import make_selection_strategy + if TYPE_CHECKING: from autodata_core.data_generator import DataGenerator, _EEFGenerationState from autodata_core.waypoint import Waypoint @@ -226,7 +228,17 @@ def validate_setup(self, datastream: Datastream) -> None: live_nodal_positions = datastream.get_object_nodal_positions() for eef_name in datastream.get_eef_names(): for subtask_index, subtask in enumerate(datastream.get_subtasks(eef_name)): - if not self.is_deformable_subtask(subtask): + is_deformable = self.is_deformable_subtask(subtask) + strategy = make_selection_strategy(subtask.selection_strategy) + supports_object = ( + strategy.supports_deformable_objects if is_deformable else strategy.supports_rigid_objects + ) + object_type = "deformable" if is_deformable else "rigid" + assert supports_object, ( + f"Selection strategy {subtask.selection_strategy!r} does not support {object_type} reference" + f" objects for SoftMimicGen subtask {eef_name}[{subtask_index}]." + ) + if not is_deformable: continue assert ( subtask.object_ref diff --git a/autodata_core/selection_strategy.py b/autodata_core/selection_strategy.py index c8aefe1..9368c16 100644 --- a/autodata_core/selection_strategy.py +++ b/autodata_core/selection_strategy.py @@ -42,9 +42,15 @@ def __new__(mcs, name: str, bases: tuple, class_dict: dict) -> type: class SelectionStrategy(metaclass=SelectionStrategyMeta): - """ABC for source-demo selection strategies.""" + """ABC for source-demo selection strategies. + + Strategies declare which reference-object representations they can select from. + Strategies that do not inspect the object state can support both representations. + """ NAME: str + supports_rigid_objects: bool = True + supports_deformable_objects: bool = False @abc.abstractmethod def select_source_demo( @@ -69,6 +75,7 @@ class RandomStrategy(SelectionStrategy): """Uniform random pick over source demos.""" NAME = "random" + supports_deformable_objects = True def select_source_demo( self, @@ -182,6 +189,8 @@ class RegistrationCostStrategy(SelectionStrategy): """Pick a source segment using deformable-object TPS registration cost.""" NAME = "registration_cost" + supports_rigid_objects = False + supports_deformable_objects = True def select_source_demo( self, diff --git a/autodata_interfaces/env/isaaclab_env_interface.py b/autodata_interfaces/env/isaaclab_env_interface.py index 435cc8d..4f3ae8e 100644 --- a/autodata_interfaces/env/isaaclab_env_interface.py +++ b/autodata_interfaces/env/isaaclab_env_interface.py @@ -264,6 +264,9 @@ def env_loop( config) rather than from ``env.cfg.datagen_config``: ``num_trials`` interpreted as successes when ``guarantee_success`` is set, otherwise as total attempts. + Reset settling uses ``env.cfg.idle_action`` when configured, or zeros otherwise. A configured + idle action must be a finite vector matching one environment's action dimensions. + Args: env: The environment to run the main step loop on. env_reset_queue: The asyncio queue carrying :class:`EnvResetRequest` instances. @@ -281,6 +284,15 @@ def env_loop( guarantee_success = generation_policy_params.guarantee_success reset_settling_steps = generation_policy_params.reset_settling_steps assert reset_settling_steps >= 0, "reset_settling_steps must be non-negative" + idle_actions = torch.zeros(env.action_space.shape, dtype=torch.float32, device=env.device) + configured_idle_action = getattr(env.cfg, "idle_action", None) + if reset_settling_steps > 0 and configured_idle_action is not None: + idle_action = torch.as_tensor(configured_idle_action, dtype=idle_actions.dtype, device=env.device) + assert ( + idle_action.shape == env.action_space.shape[1:] + ), f"env.cfg.idle_action must have shape {env.action_space.shape[1:]}, got {tuple(idle_action.shape)}." + assert torch.isfinite(idle_action).all(), "env.cfg.idle_action must contain only finite values." + idle_actions[:] = idle_action env_id_tensor = torch.tensor([0], dtype=torch.int64, device=env.device) settling_requests: dict[int, EnvResetRequest] = {} settling_steps_remaining: dict[int, int] = {} @@ -323,9 +335,9 @@ def env_loop( "non-settling environments." ) - actions = torch.zeros(env.action_space.shape) + actions = idle_actions.clone() - # Settling environments use the zero action. Batch-fetch one real action for every + # Settling environments use the idle action. Batch-fetch one real action for every # other environment so their generation trajectories continue without interruption. action_count = env.num_envs - len(settling_requests) get_tasks = [env_action_queue.get() for _ in range(action_count)] diff --git a/autodata_interfaces/tasks/generation_policy_spec.py b/autodata_interfaces/tasks/generation_policy_spec.py index 70f4853..7efa53b 100644 --- a/autodata_interfaces/tasks/generation_policy_spec.py +++ b/autodata_interfaces/tasks/generation_policy_spec.py @@ -24,8 +24,8 @@ class GenerationPolicy: use_navigation_controller: Whether a navigation controller generates loco-manipulation trajectories. reset_settling_steps: Number of ordinary batched environment steps to run after reset - before generation begins. During these steps, the resetting environment receives zero - actions while other environments continue executing their generated actions. + before generation begins. During these steps, the resetting environment receives + ``env.cfg.idle_action`` (zeros when absent) while other environments execute their generated actions. max_navigation_hold_steps: Maximum consecutive simulation steps spent waiting for a navigation goal at the end of its recorded segment. Reaching the limit fails the attempt. diff --git a/autodata_tests/core/test_algorithms.py b/autodata_tests/core/test_algorithms.py index fb74702..c08b5d4 100644 --- a/autodata_tests/core/test_algorithms.py +++ b/autodata_tests/core/test_algorithms.py @@ -3,6 +3,9 @@ """Unit tests for :mod:`autodata_core.algorithms`.""" +import torch +from types import SimpleNamespace + import pytest from autodata_core.algorithms import ( @@ -15,6 +18,8 @@ get_algorithm, iter_algorithms, ) +from autodata_core.datagen_info import DatagenInfo +from autodata_interfaces.tasks.subtask_spec import SoftMimicGenSubtaskAlgoParams, Subtask def test_registry_contents(): @@ -75,6 +80,41 @@ def test_softmimicgen_attributes(): assert algo.supports_coordination is False +@pytest.mark.parametrize( + "object_soft, strategy, compatible", + [ + (False, "random", True), + (False, "nearest_neighbor_object", True), + (False, "nearest_neighbor_robot_distance", True), + (False, "registration_cost", False), + (True, "random", True), + (True, "nearest_neighbor_object", False), + (True, "nearest_neighbor_robot_distance", False), + (True, "registration_cost", True), + ], +) +def test_softmimicgen_validates_selection_strategy_at_setup(object_soft, strategy, compatible): + subtask = Subtask( + object_ref="object", + selection_strategy=strategy, + algo_params=SoftMimicGenSubtaskAlgoParams(object_soft=object_soft), + ) + nodes = torch.zeros(1, 6, 3) + datastream = SimpleNamespace( + get_eef_names=lambda: ["franka"], + get_subtasks=lambda _eef: [subtask], + get_object_nodal_positions=lambda: {"object": nodes}, + source_pool=SimpleNamespace(datagen_infos=[DatagenInfo(object_nodal_positions={"object": nodes})]), + ) + algorithm = SoftMimicGen() + if compatible: + algorithm.validate_setup(datastream) + else: + object_type = "deformable" if object_soft else "rigid" + with pytest.raises(AssertionError, match=f"{strategy!r} does not support {object_type} reference objects"): + algorithm.validate_setup(datastream) + + def test_skillgen_attributes(): algo = SkillGen(motion_planners={0: object()}) assert algo.name == "skillgen" diff --git a/autodata_tests/interfaces/env/test_env_loop.py b/autodata_tests/interfaces/env/test_env_loop.py index 6e497f4..4a9c0a8 100644 --- a/autodata_tests/interfaces/env/test_env_loop.py +++ b/autodata_tests/interfaces/env/test_env_loop.py @@ -56,6 +56,7 @@ class _MockEnv: def __init__(self, num_envs: int) -> None: self.num_envs = num_envs self.device = "cpu" + self.cfg = SimpleNamespace() self.action_space = SimpleNamespace(shape=(num_envs, 1)) self.recorder_manager = _MockRecorderManager() self.sim = _MockSimulation() @@ -114,10 +115,13 @@ async def wait_for_reset() -> None: @pytest.mark.parametrize("settling_env_id", [0, 1]) -def test_settling_env_uses_zero_actions_while_other_env_continues(settling_env_id): +@pytest.mark.parametrize("idle_action", [None, [2.0]], ids=["default_zero", "configured"]) +def test_settling_env_uses_idle_action_while_other_env_continues(settling_env_id, idle_action): loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) env = _MockEnv(num_envs=2) + if idle_action is not None: + env.cfg.idle_action = idle_action running_env_id = 1 - settling_env_id reset_queue: asyncio.Queue = asyncio.Queue() action_queue: asyncio.Queue = asyncio.Queue() @@ -146,7 +150,8 @@ async def wait_for_reset() -> None: assert completed is False assert env.reset_env_ids == [(settling_env_id,)] assert len(env.actions) == 2 - assert [action[settling_env_id, 0].item() for action in env.actions] == [0.0, 0.0] + expected_idle = 0.0 if idle_action is None else idle_action[0] + assert [action[settling_env_id, 0].item() for action in env.actions] == [expected_idle, expected_idle] assert [action[running_env_id, 0].item() for action in env.actions] == [3.0, 4.0] assert env.recorder_manager.reset_env_ids == [(settling_env_id,)] assert env.recorder_manager.post_reset_env_ids == [] @@ -159,3 +164,26 @@ async def wait_for_reset() -> None: loop.run_until_complete(asyncio.wait_for(reset_queue.join(), timeout=0.1)) finally: _close_event_loop(loop, tasks) + + +@pytest.mark.parametrize("idle_action", [[1.0, 2.0], [float("nan")]], ids=["wrong_shape", "nonfinite"]) +def test_settling_rejects_invalid_idle_action_before_reset(idle_action): + loop = asyncio.new_event_loop() + env = _MockEnv(num_envs=1) + env.cfg.idle_action = idle_action + reset_queue: asyncio.Queue = asyncio.Queue() + reset_queue.put_nowait(EnvResetRequest(env_id=0, completion=loop.create_future())) + try: + with pytest.raises(AssertionError, match="env.cfg.idle_action"): + env_loop( + env=env, + env_reset_queue=reset_queue, + env_action_queue=asyncio.Queue(), + asyncio_event_loop=loop, + generation_policy_params=GenerationPolicy(reset_settling_steps=2), + stats={"num_success": 0, "num_failures": 0, "num_attempts": 0}, + ) + assert env.reset_env_ids == [] + assert env.actions == [] + finally: + loop.close()