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 3f6b8b6..9d0efd6 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 dc91e69..197cb2f 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, @@ -623,8 +650,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..c501ba7 100644 --- a/autodata_core/pool.py +++ b/autodata_core/pool.py @@ -124,7 +124,8 @@ 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"] 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 +135,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_examples/envs/__init__.py b/autodata_examples/envs/__init__.py new file mode 100644 index 0000000..67bc286 --- /dev/null +++ b/autodata_examples/envs/__init__.py @@ -0,0 +1,62 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Example simulation environments used by AutoData workflows.""" + +import argparse +from typing import Any + + +def register_environment_for_run( + *, + env_name: str, + enable_cameras: bool, + num_envs: int, + device: str, + seed: int, +) -> dict[str, dict[str, Any]]: + """Register the AutoData-owned environment selected for the current run. + + Args: + env_name: Gym ID of the environment selected for this run. + enable_cameras: Whether Arena environments should include their camera sensors. + num_envs: Number of parallel environments in this run. + device: Simulation device for this run. + seed: Environment seed for this run. + + Returns: + Gym constructor kwargs keyed by registered environment ID. + """ + + from .isaac_lab_arena import build_and_register_arena_environment, is_arena_environment + + if not is_arena_environment(env_name): + return {} + return build_and_register_arena_environment( + enable_cameras=enable_cameras, + num_envs=num_envs, + device=device, + seed=seed, + ) + + +def register_environments() -> list[str]: + """External callback used by unmodified Isaac Lab scripts. + + Isaac Lab invokes this function without arguments after starting Isaac Sim. The + selected ``--task`` determines whether the callback needs to compile an AutoData + Arena environment or only import AutoData's regular Isaac Lab registrations. + + Returns: + Command-line arguments not consumed by Arena environment registration. + """ + + from .isaac_lab_arena import is_arena_environment, register_environment_from_cli + + parser = argparse.ArgumentParser(add_help=False) + parser.add_argument("--task", type=str) + args, _ = parser.parse_known_args() + env_name = args.task.split(":")[-1] if args.task is not None else None + if env_name is not None and not is_arena_environment(env_name): + return [] + return register_environment_from_cli() diff --git a/autodata_examples/envs/isaac_lab_arena/__init__.py b/autodata_examples/envs/isaac_lab_arena/__init__.py new file mode 100644 index 0000000..d33c1b6 --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from .registration import build_and_register_arena_environment, is_arena_environment, register_environment_from_cli + +__all__ = ["build_and_register_arena_environment", "is_arena_environment", "register_environment_from_cli"] diff --git a/autodata_examples/envs/isaac_lab_arena/franka_rope/__init__.py b/autodata_examples/envs/isaac_lab_arena/franka_rope/__init__.py new file mode 100644 index 0000000..2574837 --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/__init__.py @@ -0,0 +1,14 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +from . import embodiment as _embodiment # noqa: F401 +from .environment import FRANKA_ROPE_ARENA_ENV_ID, FrankaRopeArenaEnvironment, FrankaRopeArenaEnvironmentCfg + +__all__ = [ + "FRANKA_ROPE_ARENA_ENV_ID", + "FrankaRopeArenaEnvironment", + "FrankaRopeArenaEnvironmentCfg", +] diff --git a/autodata_examples/envs/isaac_lab_arena/franka_rope/assets/Rope.usd b/autodata_examples/envs/isaac_lab_arena/franka_rope/assets/Rope.usd new file mode 100644 index 0000000..fac1654 Binary files /dev/null and b/autodata_examples/envs/isaac_lab_arena/franka_rope/assets/Rope.usd differ diff --git a/autodata_examples/envs/isaac_lab_arena/franka_rope/embodiment.py b/autodata_examples/envs/isaac_lab_arena/franka_rope/embodiment.py new file mode 100644 index 0000000..70b1364 --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/embodiment.py @@ -0,0 +1,183 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +from collections.abc import Callable + +import isaaclab.envs.mdp as mdp_isaac_lab +import isaaclab.sim as sim_utils +from isaaclab.managers import EventTermCfg as EventTerm +from isaaclab.managers import ObservationGroupCfg as ObsGroup +from isaaclab.managers import ObservationTermCfg as ObsTerm +from isaaclab.managers import SceneEntityCfg +from isaaclab.sensors import CameraCfg +from isaaclab.utils.configclass import configclass +from isaaclab_arena.assets.register import register_asset, register_retargeter +from isaaclab_arena.assets.retargeter_library import RetargetterBase +from isaaclab_arena.embodiments.common.arm_mode import ArmMode +from isaaclab_arena.embodiments.franka.franka import FrankaIKEmbodiment +from isaaclab_arena.utils.cameras import ArenaCameraCfg +from isaaclab_arena.utils.pose import Pose + +from . import mdp + +_FRANKA_ROPE_INIT_JOINT_POS: dict[str, float] = { + "panda_joint1": 0.0444, + "panda_joint2": -0.1894, + "panda_joint3": -0.1107, + "panda_joint4": -2.5148, + "panda_joint5": 0.0044, + "panda_joint6": 2.3775, + "panda_joint7": 2.43, + "panda_finger_joint.*": 0.04, +} + +_FRANKA_ROPE_INIT_JOINT_POSE = [ + 0.0444, + -0.1894, + -0.1107, + -2.5148, + 0.0044, + 2.3775, + 2.43, + 0.04, + 0.04, +] + + +@configclass +class FrankaRopeCameraCfg(ArenaCameraCfg): + """Wrist and fixed-view cameras matching the Isaac Lab rope environment.""" + + robot0_eye_in_hand_image = CameraCfg( + prim_path="{ENV_REGEX_NS}/Robot/panda_hand/robot0_eye_in_hand_image", + update_period=0.0, + height=128, + width=128, + data_types=["rgb"], + spawn=sim_utils.PinholeCameraCfg( + focal_length=24.0, + focus_distance=400.0, + horizontal_aperture=20.955, + clipping_range=(0.1, 2.0), + ), + offset=CameraCfg.OffsetCfg( + pos=(0.13, 0.0, -0.15), + rot=(0.03701, 0.03701, -0.70614, -0.70614), + convention="ros", + ), + ) + + agentview_image = CameraCfg( + prim_path="{ENV_REGEX_NS}/agentview_image", + update_period=0.0, + height=128, + width=128, + data_types=["rgb"], + spawn=sim_utils.PinholeCameraCfg( + focal_length=14.0, + focus_distance=400.0, + horizontal_aperture=20.955, + clipping_range=(0.1, 2.0), + ), + offset=CameraCfg.OffsetCfg( + pos=(0.5, 0.5, 0.4), + rot=(0.0, 0.38268, 0.92388, 0.0), + convention="opengl", + ), + ) + + +@configclass +class FrankaRopeObservationsCfg: + """Arena-standard Franka observations plus deformable rope nodal state.""" + + @configclass + class PolicyCfg(ObsGroup): + """Policy observations recorded into demonstration datasets.""" + + actions = ObsTerm(func=mdp_isaac_lab.last_action) + joint_pos = ObsTerm(func=mdp_isaac_lab.joint_pos_rel) + joint_vel = ObsTerm(func=mdp_isaac_lab.joint_vel_rel) + eef_pos = ObsTerm(func=mdp.ee_frame_pos) + eef_quat = ObsTerm(func=mdp.ee_frame_quat) + gripper_pos = ObsTerm(func=mdp.gripper_pos) + object_nodal_pos = ObsTerm(func=mdp.object_nodal_pos) + + def __post_init__(self) -> None: + self.enable_corruption = False + self.concatenate_terms = False + + policy: PolicyCfg = PolicyCfg() + + +@configclass +class FrankaRopeEventsCfg: + """Reset the Franka exactly as the equivalent Isaac Lab rope environment.""" + + reset_all = EventTerm(func=mdp_isaac_lab.reset_scene_to_default, mode="reset") + randomize_franka_joint_state = EventTerm( + func=mdp_isaac_lab.reset_joints_by_offset, + mode="reset", + params={ + "position_range": (-0.02, 0.02), + "velocity_range": (0.0, 0.0), + "asset_cfg": SceneEntityCfg("robot", joint_names=["panda_joint.*"]), + }, + ) + + +@register_asset +class FrankaRopeIKEmbodiment(FrankaIKEmbodiment): + """Arena Franka IK embodiment with rope-specific pose, observations, and cameras.""" + + name = "franka_rope_ik" + + def __init__( + self, + enable_cameras: bool = False, + initial_pose: Pose | None = None, + concatenate_observation_terms: bool = False, + arm_mode: ArmMode | None = None, + ) -> None: + super().__init__( + enable_cameras=enable_cameras, + initial_pose=initial_pose, + initial_joint_pose=list(_FRANKA_ROPE_INIT_JOINT_POSE), + concatenate_observation_terms=concatenate_observation_terms, + arm_mode=arm_mode, + ) + self.scene_config.robot.init_state.joint_pos = dict(_FRANKA_ROPE_INIT_JOINT_POS) + self.observation_config = FrankaRopeObservationsCfg() + self.observation_config.policy.concatenate_terms = concatenate_observation_terms + self.camera_config = FrankaRopeCameraCfg() + self.event_config = FrankaRopeEventsCfg() + self.reward_config = None + + +@register_retargeter +class FrankaRopeKeyboardRetargeter(RetargetterBase): + """Use the standard relative-pose keyboard pipeline for the rope embodiment.""" + + device = "keyboard" + embodiment = FrankaRopeIKEmbodiment.name + + def get_pipeline_builder(self, embodiment: object) -> Callable | None: + """Return no custom pipeline; Isaac Lab handles relative-pose keyboard input.""" + + del embodiment + + +@register_retargeter +class FrankaRopeSpaceMouseRetargeter(RetargetterBase): + """Use the standard relative-pose SpaceMouse pipeline for the rope embodiment.""" + + device = "spacemouse" + embodiment = FrankaRopeIKEmbodiment.name + + def get_pipeline_builder(self, embodiment: object) -> Callable | None: + """Return no custom pipeline; Isaac Lab handles relative-pose SpaceMouse input.""" + + del embodiment diff --git a/autodata_examples/envs/isaac_lab_arena/franka_rope/environment.py b/autodata_examples/envs/isaac_lab_arena/franka_rope/environment.py new file mode 100644 index 0000000..55b5358 --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/environment.py @@ -0,0 +1,108 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from isaaclab.devices import DevicesCfg, Se3KeyboardCfg, Se3SpaceMouseCfg +from isaaclab_arena.assets.register import register_environment +from isaaclab_arena.environments.arena_environment_factory import ArenaEnvironmentCfg, ArenaEnvironmentFactory +from isaaclab_arena.utils.pose import Pose + +from .rope_asset import FrankaRopeAsset +from .task import FrankaRopeTask + +if TYPE_CHECKING: + from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment + +FRANKA_ROPE_ARENA_ENV_ID = "Isaac-Rope-Franka-Arena-IK-Rel-v0" + + +@dataclass +class FrankaRopeArenaEnvironmentCfg(ArenaEnvironmentCfg): + """Configure the AutoData-owned Arena Franka rope environment.""" + + enable_cameras: bool = False + embodiment: str = "franka_rope_ik" + teleop_device: str | None = None + + +def _configure_franka_rope_env(env_cfg: Any) -> Any: + """Apply rope-specific simulation timing and legacy teleop devices.""" + + env_cfg.decimation = 5 + env_cfg.sim.dt = 0.01 + env_cfg.sim.render_interval = 2 + env_cfg.sim.render.antialiasing_mode = "DLSS" + env_cfg.num_rerenders_on_reset = 1 + env_cfg.scene.replicate_physics = False + env_cfg.teleop_devices = DevicesCfg( + devices={ + "keyboard": Se3KeyboardCfg( + pos_sensitivity=0.05, + rot_sensitivity=0.2, + ), + "spacemouse": Se3SpaceMouseCfg( + pos_sensitivity=0.2, + rot_sensitivity=0.5, + ), + }, + ) + return env_cfg + + +@register_environment +class FrankaRopeArenaEnvironment(ArenaEnvironmentFactory[FrankaRopeArenaEnvironmentCfg]): + """Compose the rope scene, Franka embodiment, and rope task through Arena.""" + + name = FRANKA_ROPE_ARENA_ENV_ID + _legacy_argparse_cfg_type = FrankaRopeArenaEnvironmentCfg + + def build(self, cfg: FrankaRopeArenaEnvironmentCfg) -> IsaacLabArenaEnvironment: + """Build an Arena environment description from ``cfg``.""" + + from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment + from isaaclab_arena.scene.scene import Scene + + table = self.asset_registry.get_asset_by_name("table")() + table.set_initial_pose( + Pose( + position_xyz=(0.5, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.707, 0.707), + ), + ) + rope = FrankaRopeAsset() + ground_plane = self.asset_registry.get_asset_by_name("ground_plane")() + ground_plane.set_initial_pose(Pose(position_xyz=(0.0, 0.0, -1.05))) + light = self.asset_registry.get_asset_by_name("light")() + light.set_intensity(3000.0) + light.set_color((0.75, 0.75, 0.75)) + + embodiment = self.asset_registry.get_asset_by_name(cfg.embodiment)( + enable_cameras=cfg.enable_cameras, + ) + teleop_device = ( + self.device_registry.get_device_by_name(cfg.teleop_device)() if cfg.teleop_device is not None else None + ) + + scene = Scene( + assets=[ + table, + rope, + ground_plane, + light, + ], + ) + task = FrankaRopeTask() + + return IsaacLabArenaEnvironment( + name=self.name, + embodiment=embodiment, + scene=scene, + task=task, + teleop_device=teleop_device, + env_cfg_callback=_configure_franka_rope_env, + ) diff --git a/autodata_examples/envs/isaac_lab_arena/franka_rope/mdp/__init__.py b/autodata_examples/envs/isaac_lab_arena/franka_rope/mdp/__init__.py new file mode 100644 index 0000000..5185858 --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/mdp/__init__.py @@ -0,0 +1,19 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from .events import reset_rope_end_tracking, reset_rope_nodal_state +from .observations import ee_frame_pos, ee_frame_quat, gripper_pos, object_grasped, object_nodal_pos +from .terminations import rope_below_minimum, rope_ends_close_tracked + +__all__ = [ + "ee_frame_pos", + "ee_frame_quat", + "gripper_pos", + "object_grasped", + "object_nodal_pos", + "reset_rope_end_tracking", + "reset_rope_nodal_state", + "rope_below_minimum", + "rope_ends_close_tracked", +] diff --git a/autodata_examples/envs/isaac_lab_arena/franka_rope/mdp/events.py b/autodata_examples/envs/isaac_lab_arena/franka_rope/mdp/events.py new file mode 100644 index 0000000..2297c31 --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/mdp/events.py @@ -0,0 +1,106 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +import isaaclab.utils.math as math_utils +from isaaclab.assets import DeformableObject +from isaaclab.managers import SceneEntityCfg + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv, ManagerBasedRLEnv + + +# Mesh-resolution-specific nodes covering one half of the 549-node rope simulation mesh. +# fmt: off +_ROPE_PARTIAL_NODE_IDS = ( + 222, 223, 224, 225, 226, 227, 228, 229, 230, 231, 232, 233, + 270, 271, 272, 273, 274, 275, 276, 277, 278, 279, 280, 281, + 282, 283, 284, 285, 286, 287, 288, 289, 290, 291, 296, 297, + 298, 299, 300, 301, 302, 303, 304, 305, 306, 307, 308, 309, + 310, 311, 312, 313, 314, 315, 316, 317, 318, 319, 320, 321, + 322, 323, 324, 325, 326, 327, 328, 329, 331, 332, 333, 334, + 335, 336, 338, 339, 340, 341, 342, 343, 344, 345, 346, 347, + 348, 349, 350, 351, 352, 353, 354, 355, 356, 357, 358, 359, + 360, 361, 362, 363, 364, 365, 366, 367, 368, 369, 370, 371, + 372, 373, 374, 375, 376, 377, 378, 379, 380, 381, 382, 383, + 384, 385, 386, 387, 388, 389, 390, 391, 392, 393, 394, 395, + 396, 397, 398, 399, 400, 401, 402, 403, 404, 405, 406, 407, + 408, 409, 410, 411, 412, 413, 414, 415, 416, 417, 418, 419, + 420, 421, 422, 423, 424, 425, 426, 427, 428, 429, 430, 431, + 432, 433, 434, 435, 436, 437, 438, 439, 440, 441, 442, 443, + 444, 445, 446, 447, 448, 449, 450, 451, 452, 453, 454, 455, + 456, 457, 458, 459, 460, 461, 462, 463, 464, 465, 466, 467, + 468, 469, 470, 471, 472, 473, 474, 475, 476, 477, 478, 479, + 480, 481, 482, 483, 484, 485, 486, 487, 488, 489, 490, 491, + 492, 493, 494, 495, 496, 497, 498, 499, 500, 501, 502, 503, + 504, 505, 506, 507, 508, 509, 510, 511, 512, 513, 514, 515, + 516, 517, 518, 519, 520, 528, 529, 530, 531, 532, 533, 534, + 535, 536, 537, 538, 539, 540, 541, 542, 543, 544, 545, 546, + 547, 548, +) +# fmt: on + + +def _sample_pose_delta( + pose_range: dict[str, tuple[float, float]], + count: int, + device: str, +) -> tuple[torch.Tensor, torch.Tensor]: + """Sample translation and rotation deltas from named pose ranges.""" + + range_list = [pose_range.get(key, (0.0, 0.0)) for key in ("x", "y", "z", "roll", "pitch", "yaw")] + ranges = torch.tensor(range_list, device=device) + samples = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], (count, 6), device=device) + quaternion = math_utils.quat_from_euler_xyz(samples[:, 3], samples[:, 4], samples[:, 5]) + return samples[:, :3], quaternion + + +def reset_rope_end_tracking(env: ManagerBasedRLEnv, env_ids: torch.Tensor | None = None) -> None: + """Discard cached rope endpoint indices for reset environments.""" + + if not hasattr(env, "_rope_end_indices"): + return + if env_ids is None: + env._rope_end_indices.clear() + return + for env_id in env_ids.tolist(): + env._rope_end_indices.pop(int(env_id), None) + + +def reset_rope_nodal_state( + env: ManagerBasedEnv, + env_ids: torch.Tensor, + pose_range: dict[str, tuple[float, float]], + partial_pose_range: dict[str, tuple[float, float]], + asset_cfg: SceneEntityCfg = SceneEntityCfg("object"), +) -> None: + """Reset and randomize the rope, including a transform of one rope half. + + Args: + env: Environment containing the deformable rope. + env_ids: Environments to reset. + pose_range: Translation [m] and rotation [rad] ranges applied to all nodes. + partial_pose_range: Additional translation [m] and rotation [rad] ranges applied to one rope half. + asset_cfg: Deformable rope entity. + """ + + rope: DeformableObject = env.scene[asset_cfg.name] + nodal_state = rope.data.default_nodal_state_w.torch[env_ids].clone() + + position, quaternion = _sample_pose_delta(pose_range, len(env_ids), rope.device) + nodal_state[..., :3] = rope.transform_nodal_pos(nodal_state[..., :3], position, quaternion) + + assert max(_ROPE_PARTIAL_NODE_IDS) < nodal_state.shape[1], ( + f"Rope mesh has {nodal_state.shape[1]} simulation nodes, but the reset mapping expects at least " + f"{max(_ROPE_PARTIAL_NODE_IDS) + 1}." + ) + node_ids = torch.tensor(_ROPE_PARTIAL_NODE_IDS, device=rope.device) + position, quaternion = _sample_pose_delta(partial_pose_range, len(env_ids), rope.device) + nodal_state[..., node_ids, :3] = rope.transform_nodal_pos(nodal_state[..., node_ids, :3], position, quaternion) + + rope.write_nodal_state_to_sim_index(nodal_state, env_ids=env_ids) diff --git a/autodata_examples/envs/isaac_lab_arena/franka_rope/mdp/observations.py b/autodata_examples/envs/isaac_lab_arena/franka_rope/mdp/observations.py new file mode 100644 index 0000000..98900ca --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/mdp/observations.py @@ -0,0 +1,93 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +from isaaclab.assets import Articulation, DeformableObject +from isaaclab.managers import SceneEntityCfg +from isaaclab.sensors import FrameTransformer + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + + +def ee_frame_pos( + env: ManagerBasedRLEnv, + ee_frame_cfg: SceneEntityCfg = SceneEntityCfg("ee_frame"), +) -> torch.Tensor: + """Return the end-effector position in the environment frame [m].""" + + ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] + return ee_frame.data.target_pos_w.torch[:, 0, :] - env.scene.env_origins + + +def ee_frame_quat( + env: ManagerBasedRLEnv, + ee_frame_cfg: SceneEntityCfg = SceneEntityCfg("ee_frame"), +) -> torch.Tensor: + """Return the end-effector orientation quaternion in the world frame.""" + + ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] + return ee_frame.data.target_quat_w.torch[:, 0, :] + + +def gripper_pos( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntityCfg = SceneEntityCfg("robot", joint_names=["panda_finger_joint.*"]), +) -> torch.Tensor: + """Return the two Franka finger positions using the SoftMimicGen sign convention [m].""" + + robot: Articulation = env.scene[robot_cfg.name] + joint_pos = robot.data.joint_pos.torch[:, robot_cfg.joint_ids] + return torch.cat((joint_pos[:, 0:1], -joint_pos[:, 1:2]), dim=1) + + +def object_grasped( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntityCfg, + ee_frame_cfg: SceneEntityCfg, + object_cfg: SceneEntityCfg, + distance_threshold: float = 0.015, + gripper_open_value: float = 0.04, + gripper_threshold: float = 0.005, +) -> torch.Tensor: + """Return whether the gripper is closed around a nearby rope node. + + Args: + env: Environment containing the robot and rope. + robot_cfg: Franka articulation and finger-joint selection. + ee_frame_cfg: End-effector frame sensor. + object_cfg: Deformable rope entity. + distance_threshold: Maximum end-effector-to-node distance [m]. + gripper_open_value: Finger position when fully open [m]. + gripper_threshold: Required deviation from the open position [m]. + + Returns: + Boolean tensor with one value per environment. + """ + + robot: Articulation = env.scene[robot_cfg.name] + ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] + rope: DeformableObject = env.scene[object_cfg.name] + + nodal_pos_w = rope.data.nodal_pos_w.torch + end_effector_pos_w = ee_frame.data.target_pos_w.torch[:, 0, :] + minimum_distance = torch.linalg.vector_norm(nodal_pos_w - end_effector_pos_w.unsqueeze(1), dim=-1).amin(dim=1) + + finger_pos = robot.data.joint_pos.torch[:, robot_cfg.joint_ids] + gripper_closed = (torch.abs(finger_pos - gripper_open_value) > gripper_threshold).all(dim=1) + return (minimum_distance < distance_threshold) & gripper_closed + + +def object_nodal_pos( + env: ManagerBasedRLEnv, + object_cfg: SceneEntityCfg = SceneEntityCfg("object"), +) -> torch.Tensor: + """Return rope simulation-node positions in the environment frame [m].""" + + rope: DeformableObject = env.scene[object_cfg.name] + return rope.data.nodal_pos_w.torch - env.scene.env_origins.unsqueeze(1) diff --git a/autodata_examples/envs/isaac_lab_arena/franka_rope/mdp/terminations.py b/autodata_examples/envs/isaac_lab_arena/franka_rope/mdp/terminations.py new file mode 100644 index 0000000..ac42349 --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/mdp/terminations.py @@ -0,0 +1,78 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +from isaaclab.assets import DeformableObject +from isaaclab.managers import SceneEntityCfg + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + + +def rope_below_minimum( + env: ManagerBasedRLEnv, + minimum_height: float, + object_cfg: SceneEntityCfg = SceneEntityCfg("object"), +) -> torch.Tensor: + """Return whether the rope center is below a minimum world height [m].""" + + rope: DeformableObject = env.scene[object_cfg.name] + return rope.data.root_pos_w.torch[:, 2] < minimum_height + + +def rope_ends_close_tracked( + env: ManagerBasedRLEnv, + object_cfg: SceneEntityCfg = SceneEntityCfg("object"), + distance_threshold: float = 0.08, + distance_threshold_min: float = 0.02, + num_end_nodes: int = 3, +) -> torch.Tensor: + """Return whether the two tracked rope ends are close enough for success. + + Endpoint nodes are discovered from the farthest node pair at the start of each + episode and then tracked by simulation-node index. + + Args: + env: Environment containing the rope. + object_cfg: Deformable rope entity. + distance_threshold: Maximum distance between endpoint neighborhoods [m]. + distance_threshold_min: Minimum separation used to reject collapsed false positives [m]. + num_end_nodes: Number of nearby nodes used for each endpoint neighborhood. + + Returns: + Boolean tensor with one value per environment. + """ + + rope: DeformableObject = env.scene[object_cfg.name] + nodal_positions = rope.data.nodal_pos_w.torch + num_envs, num_nodes, _ = nodal_positions.shape + assert num_end_nodes <= num_nodes, f"Requested {num_end_nodes} endpoint nodes from a {num_nodes}-node rope." + + if not hasattr(env, "_rope_end_indices"): + env._rope_end_indices = {} + + result = torch.zeros(num_envs, dtype=torch.bool, device=env.device) + for env_id in range(num_envs): + env_nodal_pos = nodal_positions[env_id] + if env_id not in env._rope_end_indices: + farthest_flat_index = torch.cdist(env_nodal_pos, env_nodal_pos).argmax() + env._rope_end_indices[env_id] = ( + int((farthest_flat_index // num_nodes).item()), + int((farthest_flat_index % num_nodes).item()), + ) + + end_1_id, end_2_id = env._rope_end_indices[env_id] + end_1_distances = torch.linalg.vector_norm(env_nodal_pos - env_nodal_pos[end_1_id], dim=1) + end_2_distances = torch.linalg.vector_norm(env_nodal_pos - env_nodal_pos[end_2_id], dim=1) + end_1_node_ids = torch.topk(end_1_distances, num_end_nodes, largest=False).indices + end_2_node_ids = torch.topk(end_2_distances, num_end_nodes, largest=False).indices + + minimum_end_distance = torch.cdist(env_nodal_pos[end_1_node_ids], env_nodal_pos[end_2_node_ids]).amin() + result[env_id] = distance_threshold_min < minimum_end_distance < distance_threshold + + return result diff --git a/autodata_examples/envs/isaac_lab_arena/franka_rope/rope_asset.py b/autodata_examples/envs/isaac_lab_arena/franka_rope/rope_asset.py new file mode 100644 index 0000000..1eb21cd --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/rope_asset.py @@ -0,0 +1,39 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +from pathlib import Path + +from isaaclab.assets import DeformableObjectCfg +from isaaclab.sim.spawners.from_files.from_files_cfg import UsdFileCfg +from isaaclab_arena.assets.asset import Asset + +_ROPE_USD_PATH = Path(__file__).resolve().parent / "assets" / "Rope.usd" + + +class FrankaRopeAsset(Asset): + """Deformable rope asset that can participate in an Arena :class:`Scene`.""" + + def __init__(self, name: str = "object") -> None: + super().__init__(name=name, tags=["deformable"]) + self.object_cfg = DeformableObjectCfg( + prim_path="{ENV_REGEX_NS}/Object", + init_state=DeformableObjectCfg.InitialStateCfg( + pos=(0.5, 0.0, 0.02), + rot=(0.0, 0.0, 0.707, 0.707), + ), + spawn=UsdFileCfg(usd_path=str(_ROPE_USD_PATH)), + debug_vis=False, + ) + + def get_object_cfg(self) -> tuple[str, DeformableObjectCfg]: + """Return the scene key and Isaac Lab deformable-object configuration.""" + + return self.name, self.object_cfg + + def get_event_cfg(self) -> tuple[str, None]: + """Return no asset-level event; the Arena task owns nodal reset semantics.""" + + return self.name, None diff --git a/autodata_examples/envs/isaac_lab_arena/franka_rope/task.py b/autodata_examples/envs/isaac_lab_arena/franka_rope/task.py new file mode 100644 index 0000000..634f71c --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/task.py @@ -0,0 +1,149 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +import math +from typing import Never + +import isaaclab.envs.mdp as mdp_isaac_lab +from isaaclab.envs.common import ViewerCfg +from isaaclab.managers import EventTermCfg as EventTerm +from isaaclab.managers import ObservationGroupCfg as ObsGroup +from isaaclab.managers import ObservationTermCfg as ObsTerm +from isaaclab.managers import SceneEntityCfg +from isaaclab.managers import TerminationTermCfg as DoneTerm +from isaaclab.utils.configclass import configclass +from isaaclab_arena.assets.register import register_task +from isaaclab_arena.embodiments.common.arm_mode import ArmMode +from isaaclab_arena.metrics.metric_base import MetricBase +from isaaclab_arena.metrics.success_rate import SuccessRateMetric +from isaaclab_arena.tasks.task_base import TaskBase + +from . import mdp + + +@configclass +class FrankaRopeTaskObservationsCfg: + """Automatic SoftMimicGen subtask signals.""" + + @configclass + class SubtaskCfg(ObsGroup): + """Signals consumed by AutoData annotation.""" + + grasp = ObsTerm( + func=mdp.object_grasped, + params={ + "robot_cfg": SceneEntityCfg("robot", joint_names=["panda_finger_joint.*"]), + "ee_frame_cfg": SceneEntityCfg("ee_frame"), + "object_cfg": SceneEntityCfg("object"), + }, + ) + + def __post_init__(self) -> None: + self.enable_corruption = False + self.concatenate_terms = False + + subtask_terms: SubtaskCfg = SubtaskCfg() + + +@configclass +class FrankaRopeTaskEventsCfg: + """Rope randomization and endpoint tracking events.""" + + reset_rope_end_tracking = EventTerm( + func=mdp.reset_rope_end_tracking, + mode="reset", + ) + + reset_object_position = EventTerm( + func=mdp.reset_rope_nodal_state, + mode="reset", + params={ + "pose_range": { + "x": (0.0, 0.1), + "y": (0.0, 0.0), + "z": (0.0, 0.0), + "roll": (0.0, 0.0), + "pitch": (0.0, 0.0), + "yaw": (0.0, 0.0), + }, + "partial_pose_range": { + "x": (0.0, 0.1), + "y": (0.0, 0.1), + "z": (0.0, 0.0), + "roll": (0.0, 0.0), + "pitch": (0.0, 0.0), + "yaw": (-math.pi / 6.0, math.pi / 6.0), + }, + "asset_cfg": SceneEntityCfg("object"), + }, + ) + + +@configclass +class FrankaRopeTaskTerminationsCfg: + """Rope success, failure, and time-limit terms.""" + + time_out = DoneTerm(func=mdp_isaac_lab.time_out, time_out=True) + object_dropping = DoneTerm( + func=mdp.rope_below_minimum, + params={ + "minimum_height": -0.05, + "object_cfg": SceneEntityCfg("object"), + }, + ) + success = DoneTerm(func=mdp.rope_ends_close_tracked) + + +@register_task +class FrankaRopeTask(TaskBase): + """Bring the two ends of a deformable rope together with a Franka arm.""" + + def __init__(self, episode_length_s: float = 10.0) -> None: + super().__init__( + episode_length_s=episode_length_s, + task_description="Grasp and manipulate the deformable rope until its ends meet.", + ) + self.observation_cfg = FrankaRopeTaskObservationsCfg() + self.events_cfg = FrankaRopeTaskEventsCfg() + self.termination_cfg = FrankaRopeTaskTerminationsCfg() + + def get_scene_cfg(self) -> None: + """Return no task scene additions; the environment owns all assets.""" + + def get_observation_cfg(self) -> FrankaRopeTaskObservationsCfg: + """Return automatic subtask-signal observations.""" + + return self.observation_cfg + + def get_termination_cfg(self) -> FrankaRopeTaskTerminationsCfg: + """Return task termination terms.""" + + return self.termination_cfg + + def get_events_cfg(self) -> FrankaRopeTaskEventsCfg: + """Return deformable reset events.""" + + return self.events_cfg + + def get_mimic_env_cfg(self, arm_mode: ArmMode) -> Never: + """Reject Arena Mimic mode; SoftMimicGen is provided by AutoData.""" + + del arm_mode + raise NotImplementedError("FrankaRopeTask uses AutoData SoftMimicGen, not Arena Mimic mode.") + + def get_metrics(self) -> list[MetricBase]: + """Return success-rate evaluation.""" + + return [SuccessRateMetric()] + + def get_viewer_cfg(self) -> ViewerCfg: + """Return the rope workspace camera framing.""" + + return ViewerCfg( + eye=(1.0, 1.0, 0.5), + lookat=(0.0, 0.0, 0.0), + origin_type="env", + ) diff --git a/autodata_examples/envs/isaac_lab_arena/registration.py b/autodata_examples/envs/isaac_lab_arena/registration.py new file mode 100644 index 0000000..e66aeb7 --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/registration.py @@ -0,0 +1,128 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +import argparse +import gymnasium as gym +from functools import partial +from typing import Any + +from isaaclab_arena.assets.registries import EnvironmentRegistry + +from .franka_rope import FRANKA_ROPE_ARENA_ENV_ID + + +def is_arena_environment(env_name: str) -> bool: + """Return whether ``env_name`` is owned by AutoData's Arena package.""" + + return env_name == FRANKA_ROPE_ARENA_ENV_ID + + +def _bind_gym_make_kwargs(env_name: str, env_kwargs: dict[str, Any]) -> None: + """Bind Arena-only constructor arguments for callers that use plain ``gym.make``. + + Isaac Lab's external-callback API can register an environment but cannot return + keyword arguments to the later ``gym.make`` call. Bind those arguments into the + registered entry point instead. + """ + + env_spec = gym.spec(env_name) + entry_point = env_spec.entry_point + assert entry_point is not None, f"Environment {env_name!r} has no Gym entry point." + if isinstance(entry_point, str): + from gymnasium.envs.registration import load_env_creator + + entry_point = load_env_creator(entry_point) + env_spec.entry_point = partial(entry_point, **env_kwargs) + + +def build_and_register_arena_environment( + *, + enable_cameras: bool, + num_envs: int, + device: str, + seed: int, +) -> dict[str, dict[str, Any]]: + """Build and register AutoData's typed Arena environment for the current run. + + Args: + enable_cameras: Whether to include environment camera sensors. + num_envs: Number of parallel environments in this run. + device: Simulation device for this run. + seed: Environment seed for this run. + + Returns: + Gym constructor kwargs keyed by registered environment ID. + """ + + assert ( + FRANKA_ROPE_ARENA_ENV_ID not in gym.registry + ), f"Environment {FRANKA_ROPE_ARENA_ENV_ID!r} is already registered in this process." + + # ArenaEnvBuilder imports simulator modules and must be loaded only after AppLauncher starts Isaac Sim. + from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder + from isaaclab_arena.environments.arena_env_builder_cfg import ArenaEnvBuilderCfg + + environment_registry = EnvironmentRegistry() + environment_factory_type = environment_registry.get_component_by_name(FRANKA_ROPE_ARENA_ENV_ID) + environment_cfg_type = environment_registry.get_environment_cfg_type(environment_factory_type) + arena_environment = environment_factory_type().build(environment_cfg_type(enable_cameras=enable_cameras)) + + builder = ArenaEnvBuilder( + arena_environment, + ArenaEnvBuilderCfg( + num_envs=num_envs, + env_spacing=2.5, + seed=seed, + solve_relations=False, + device=device, + ), + ) + env_name, _, env_kwargs = builder.build_registered() + return {env_name: env_kwargs} + + +def register_environment_from_cli() -> list[str]: + """Register the Arena rope environment through Isaac Lab's callback API. + + Isaac Lab invokes external registration callbacks without arguments. This adapter + reads the standard Lab and Arena options already present in ``sys.argv``, builds the + typed AutoData environment, and binds Arena's additional constructor arguments for + the script's later plain ``gym.make`` call. + + Returns: + Command-line arguments not consumed by the environment registration parser. + """ + + from isaaclab.app import AppLauncher + from isaaclab_arena.cli.isaaclab_arena_cli import ( + add_isaac_lab_cli_args, + add_isaaclab_arena_cli_args, + arena_env_builder_cfg_from_argparse, + ) + from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder + from isaaclab_arena_environments.cli import add_environment_cli_args, build_environment_from_cli + + parser = argparse.ArgumentParser(add_help=False) + parser.add_argument("--task", type=str, default=FRANKA_ROPE_ARENA_ENV_ID) + AppLauncher.add_app_launcher_args(parser) + add_isaac_lab_cli_args(parser) + add_isaaclab_arena_cli_args(parser) + + environment_factory_type = EnvironmentRegistry().get_component_by_name(FRANKA_ROPE_ARENA_ENV_ID) + add_environment_cli_args(parser, environment_factory_type) + parser.set_defaults(env_spacing=2.5, solve_relations=False) + args, remaining_args = parser.parse_known_args() + + requested_env_name = args.task.split(":")[-1] + assert ( + requested_env_name == FRANKA_ROPE_ARENA_ENV_ID + ), f"Arena callback expected {FRANKA_ROPE_ARENA_ENV_ID!r}, got {requested_env_name!r}." + + arena_environment = build_environment_from_cli(environment_factory_type, args) + builder = ArenaEnvBuilder(arena_environment, arena_env_builder_cfg_from_argparse(args)) + env_name, _, env_kwargs = builder.build_registered() + _bind_gym_make_kwargs(env_name, env_kwargs) + return remaining_args diff --git a/autodata_examples/tasks/franka_rope_softmimicgen.yaml b/autodata_examples/tasks/franka_rope_softmimicgen.yaml new file mode 100644 index 0000000..682618b --- /dev/null +++ b/autodata_examples/tasks/franka_rope_softmimicgen.yaml @@ -0,0 +1,59 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +name: franka_rope +description: Grasp and manipulate a deformable rope until its ends meet. +algo: softmimicgen + +generation_policy: + name: franka_rope + seed: 1 + num_trials: 10 + guarantee_success: true + keep_failed: true + source_dataset_path: null + generation_path: null + task_name: null + use_skillgen: false + use_navigation_controller: false + reset_settling_steps: 10 + select_src_per_subtask: false + select_src_per_arm: false + transform_first_robot_pose: false + interpolate_from_last_target_pose: true + +subtasks: + franka: + - object_ref: object + description: Approach and grasp the deformable rope. + subtask_term_signal: grasp + selection_strategy: registration_cost + selection_strategy_kwargs: + nn_k: 3 + subtask_term_offset_range: [10, 15] + action_noise: 0.0 + num_interpolation_steps: 10 + num_fixed_steps: 0 + apply_noise_during_interpolation: false + algo_params: + object_soft: true + use_rotation_transform: true + bend_coef: 0.1 + rot_coef: 0.001 + + # Final subtask: no term signal (end-of-trajectory) and no term offset. + - object_ref: object + description: Manipulate the rope until its ends meet. + selection_strategy: registration_cost + selection_strategy_kwargs: + nn_k: 3 + subtask_term_offset_range: [0, 0] + action_noise: 0.0 + num_interpolation_steps: 5 + num_fixed_steps: 0 + apply_noise_during_interpolation: false + algo_params: + object_soft: true + use_rotation_transform: true + bend_coef: 0.1 + rot_coef: 0.001 diff --git a/autodata_interfaces/datastream/datastream.py b/autodata_interfaces/datastream/datastream.py index 38b7109..b2d9cba 100644 --- a/autodata_interfaces/datastream/datastream.py +++ b/autodata_interfaces/datastream/datastream.py @@ -12,6 +12,7 @@ from autodata_core.pool import DataGenInfoPool from autodata_interfaces.embodiments.embodiment_adapter import EmbodimentAdapter +from autodata_interfaces.env.scene_state import get_scene_state from autodata_interfaces.tasks.generation_policy_spec import GenerationPolicy from autodata_interfaces.tasks.subtask_constraint_spec import SubtaskConstraint from autodata_interfaces.tasks.subtask_spec import Subtask, SubtaskAlgoParams @@ -255,6 +256,27 @@ def get_object_poses(self, env_ids: Sequence[int] | None = None) -> dict[str, to object_pose_matrix[obj_name] = pose_math.make_pose(pos_rel, pose_math.matrix_from_quat(quat)) 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]: @@ -322,7 +344,8 @@ def get_scene_state(self, is_relative: bool = True) -> dict: hatch for callers that need the full dict (e.g. recorder ``initial_state``). """ - return self.env.scene.get_state(is_relative=is_relative) + # TODO: Remove this custom method once upstream Lab fixes the nodal position relative scene state bug. + return get_scene_state(self.env.scene, is_relative=is_relative) # ------------------------------------------------------------------ # Collision-world source 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 1011b0f..64ab5e9 100644 --- a/autodata_interfaces/env/isaaclab_env_interface.py +++ b/autodata_interfaces/env/isaaclab_env_interface.py @@ -22,7 +22,6 @@ import isaaclab.sim as sim_utils from isaaclab.assets import RigidObjectCfg -from isaaclab.envs.mdp.recorders.recorders_cfg import ActionStateRecorderManagerCfg from isaaclab.managers import DatasetExportMode, EventTermCfg, SceneEntityCfg from isaaclab.managers.recorder_manager import RecorderManagerBaseCfg from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, ISAACLAB_NUCLEUS_DIR @@ -31,6 +30,8 @@ from isaaclab_tasks.utils.parse_cfg import parse_env_cfg from autodata_interfaces.env.env_profile import EnvironmentProfile, convert_event_params +from autodata_interfaces.env.recorders import make_action_state_recorder_manager_cfg +from autodata_interfaces.env.reset_request import EnvResetRequest from autodata_interfaces.tasks.generation_policy_spec import GenerationPolicy @@ -192,7 +193,7 @@ def setup_env_config( # Setup recorders if recorder_cfg is None: - env_cfg.recorders = ActionStateRecorderManagerCfg() + env_cfg.recorders = make_action_state_recorder_manager_cfg() else: env_cfg.recorders = recorder_cfg env_cfg.recorders.dataset_export_dir_path = output_dir @@ -228,7 +229,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). @@ -241,13 +242,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() @@ -255,25 +260,77 @@ 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) + env.recorder_manager.record_post_reset(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/recorders.py b/autodata_interfaces/env/recorders.py new file mode 100644 index 0000000..3fec801 --- /dev/null +++ b/autodata_interfaces/env/recorders.py @@ -0,0 +1,46 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +from collections.abc import Sequence + +from isaaclab.envs.mdp.recorders.recorders_cfg import ActionStateRecorderManagerCfg +from isaaclab.managers.recorder_manager import RecorderTerm + +from autodata_interfaces.env.scene_state import get_scene_state + +# TODO: Remove this custom recorder once upstream Lab fixes the nodal position relative scene state bug. + + +class InitialStateRecorder(RecorderTerm): + """Record the correctly env-relative initial scene state after reset.""" + + def record_post_reset(self, env_ids: Sequence[int] | None): + """Return initial state for the reset environments.""" + + def select_envs(value): + if isinstance(value, dict): + return {key: select_envs(item) for key, item in value.items()} + return value if env_ids is None else value[env_ids] + + return "initial_state", select_envs(get_scene_state(self._env.scene, is_relative=True)) + + +class PostStepStatesRecorder(RecorderTerm): + """Record correctly env-relative scene state after each environment step.""" + + def record_post_step(self): + """Return state for every environment.""" + + return "states", get_scene_state(self._env.scene, is_relative=True) + + +def make_action_state_recorder_manager_cfg() -> ActionStateRecorderManagerCfg: + """Return Isaac Lab's action-state recorder with corrected state terms.""" + + cfg = ActionStateRecorderManagerCfg() + cfg.record_initial_state.class_type = InitialStateRecorder + cfg.record_post_step_states.class_type = PostStepStatesRecorder + return cfg 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..e012c1d --- /dev/null +++ b/autodata_interfaces/env/scene_state.py @@ -0,0 +1,41 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Scene-state reads with correct per-environment coordinate conversion.""" + +from __future__ import annotations + +from typing import Any + + +# TODO: Remove this custom method once upstream Lab fixes the nodal position relative scene state bug. +def get_scene_state(scene: Any, is_relative: bool = True) -> dict: + """Return scene state in world or per-environment coordinates. + + Isaac Lab's relative-state conversion indexes deformable nodal positions as if their second + dimension were XYZ. Read world state first and apply the environment origin along the actual + ``(environment, node, xyz)`` layout here. + + Args: + scene: Isaac Lab interactive scene. + is_relative: Whether positions are relative to each environment origin. + + Returns: + Nested scene-state mapping containing cloned tensors. + """ + + state = scene.get_state(is_relative=False) + if not is_relative: + return state + + env_origins = scene.env_origins + for asset_state in state.get("articulation", {}).values(): + if "root_pose" in asset_state: + asset_state["root_pose"][:, :3] -= env_origins + for asset_state in state.get("deformable_object", {}).values(): + if "nodal_position" in asset_state: + asset_state["nodal_position"] -= env_origins.unsqueeze(1) + for asset_state in state.get("rigid_object", {}).values(): + if "root_pose" in asset_state: + asset_state["root_pose"][:, :3] -= env_origins + return state diff --git a/autodata_interfaces/tasks/generation_policy_spec.py b/autodata_interfaces/tasks/generation_policy_spec.py index 67cf8c2..0b23c87 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. Segment stitching parameters: select_src_per_subtask: If True, re-select a source demo for every subtask. If False, the @@ -47,6 +50,7 @@ class GenerationPolicy: task_name: str | None = None use_skillgen: bool = False use_navigation_controller: bool = False + reset_settling_steps: int = 0 # --- segment stitching --- select_src_per_subtask: bool = False 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 d37e319..988a451 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..c0a1f8a 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": {}}) 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/e2e/test_softmimicgen_data_generation.py b/autodata_tests/e2e/test_softmimicgen_data_generation.py new file mode 100644 index 0000000..a3d3103 --- /dev/null +++ b/autodata_tests/e2e/test_softmimicgen_data_generation.py @@ -0,0 +1,66 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os +import tempfile + +import pytest + +from autodata_tests.utils.constants import TestPaths +from autodata_tests.utils.subprocess import run_subprocess +from autodata_tests.utils.utils import assert_valid_dataset + +HEADLESS = True +GENERATION_NUM_TRIALS = 1 + + +def _run_franka_rope_softmimicgen(num_envs: int, device: str) -> None: + """Run SoftMimicGen data generation for the Franka rope task.""" + + with tempfile.TemporaryDirectory() as temp_dir: + output_file = os.path.join(temp_dir, "generated.hdf5") + + args = [ + TestPaths.python_path, + TestPaths.generate_dataset_script, + "--env_name", + "Isaac-Rope-Franka-Arena-IK-Rel-v0", + "--alg", + "softmimicgen", + "--task_descriptor", + os.path.join(TestPaths.tasks_dir, "franka_rope_softmimicgen.yaml"), + "--embodiment", + os.path.join(TestPaths.embodiments_dir, "franka_ik_rel.yaml"), + "--input_file", + os.path.join(TestPaths.test_data_dir, "annotated_dataset_franka_rope_softmimicgen.hdf5"), + "--output_file", + output_file, + "--generation_num_trials", + str(GENERATION_NUM_TRIALS), + "--num_envs", + str(num_envs), + "--device", + device, + "--viz", + "none" if HEADLESS else "kit", + ] + run_subprocess(args) + + assert_valid_dataset(output_file, min_num_demos=GENERATION_NUM_TRIALS) + + +@pytest.mark.with_subprocess +def test_franka_rope_softmimicgen_data_generation_single_env_cuda(): + """SoftMimicGen generation for the Franka rope task on a single env on GPU.""" + _run_franka_rope_softmimicgen(num_envs=1, device="cuda") + + +@pytest.mark.with_subprocess +def test_franka_rope_softmimicgen_data_generation_multi_env_cuda(): + """SoftMimicGen generation for the Franka rope task on multiple parallel envs on GPU.""" + _run_franka_rope_softmimicgen(num_envs=10, device="cuda") + + +if __name__ == "__main__": + test_franka_rope_softmimicgen_data_generation_single_env_cuda() + test_franka_rope_softmimicgen_data_generation_multi_env_cuda() diff --git a/autodata_tests/interfaces/datastream/test_datastream.py b/autodata_tests/interfaces/datastream/test_datastream.py index 98ab94a..995672e 100644 --- a/autodata_tests/interfaces/datastream/test_datastream.py +++ b/autodata_tests/interfaces/datastream/test_datastream.py @@ -56,9 +56,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, ) @@ -225,6 +231,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_get_robot_root_pose(): datastream, _, _, _ = _datastream() pose = datastream.get_robot_root_pose(env_ids=[0]) 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..39bf021 --- /dev/null +++ b/autodata_tests/interfaces/env/test_env_loop.py @@ -0,0 +1,131 @@ +# 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 + +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, ...]] = [] + + 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: + self.post_reset_env_ids.append(tuple(env_ids.tolist())) + + +class _MockSimulation: + def is_stopped(self) -> bool: + return False + + +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.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()) + + +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) + + +def test_settling_env_uses_zero_actions_while_other_env_continues(): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + env = _MockEnv(num_envs=2) + 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 provide_running_env_actions() -> None: + for value in (3.0, 4.0): + await action_queue.put((1, 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 == [(0,)] + assert len(env.actions) == 2 + assert env.actions[0][:, 0].tolist() == [0.0, 3.0] + assert env.actions[1][:, 0].tolist() == [0.0, 4.0] + assert env.recorder_manager.reset_env_ids == [(0,)] + assert env.recorder_manager.post_reset_env_ids == [(0,)] + 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/env/test_franka_rope_arena.py b/autodata_tests/interfaces/env/test_franka_rope_arena.py new file mode 100644 index 0000000..3f2f770 --- /dev/null +++ b/autodata_tests/interfaces/env/test_franka_rope_arena.py @@ -0,0 +1,182 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Configuration tests for the AutoData-owned Arena Franka rope environment.""" + +import gymnasium as gym +import sys +from functools import partial +from types import SimpleNamespace + +from isaaclab_arena.assets.registries import EnvironmentRegistry + +import autodata_examples.envs as env_registration +import autodata_examples.envs.isaac_lab_arena as arena_registration +from autodata_examples.envs.isaac_lab_arena import registration +from autodata_examples.envs.isaac_lab_arena.franka_rope import mdp +from autodata_examples.envs.isaac_lab_arena.franka_rope.embodiment import FrankaRopeIKEmbodiment +from autodata_examples.envs.isaac_lab_arena.franka_rope.environment import ( + FRANKA_ROPE_ARENA_ENV_ID, + FrankaRopeArenaEnvironment, + FrankaRopeArenaEnvironmentCfg, +) +from autodata_examples.envs.isaac_lab_arena.franka_rope.rope_asset import FrankaRopeAsset +from autodata_examples.envs.isaac_lab_arena.franka_rope.task import FrankaRopeTask + + +def test_arena_environment_factory_has_typed_config(): + assert FrankaRopeArenaEnvironment.name == FRANKA_ROPE_ARENA_ENV_ID + assert FrankaRopeArenaEnvironment._legacy_argparse_cfg_type is FrankaRopeArenaEnvironmentCfg + + +def test_arena_environment_factory_is_registered(): + registry = EnvironmentRegistry() + assert registry.get_component_by_name(FRANKA_ROPE_ARENA_ENV_ID) is FrankaRopeArenaEnvironment + assert registry.get_environment_cfg_type(FrankaRopeArenaEnvironment) is FrankaRopeArenaEnvironmentCfg + + +def test_arena_registration_uses_run_configuration_and_preserves_make_kwargs(monkeypatch): + from isaaclab_arena.environments import arena_env_builder + + captured = {} + variation_recorder = object() + + class FakeArenaEnvBuilder: + def __init__(self, arena_environment, cfg): + captured["arena_environment"] = arena_environment + captured["cfg"] = cfg + + def build_registered(self): + return FRANKA_ROPE_ARENA_ENV_ID, object(), {"variation_recorder": variation_recorder} + + monkeypatch.delitem(gym.registry, FRANKA_ROPE_ARENA_ENV_ID, raising=False) + monkeypatch.setattr(arena_env_builder, "ArenaEnvBuilder", FakeArenaEnvBuilder) + monkeypatch.setattr( + FrankaRopeArenaEnvironment, + "build", + lambda _self, cfg: SimpleNamespace( + name=FRANKA_ROPE_ARENA_ENV_ID, + embodiment=SimpleNamespace(enable_cameras=cfg.enable_cameras), + ), + ) + + make_kwargs = registration.build_and_register_arena_environment( + enable_cameras=True, + num_envs=8, + device="cuda:1", + seed=23, + ) + + builder_cfg = captured["cfg"] + assert captured["arena_environment"].name == FRANKA_ROPE_ARENA_ENV_ID + assert captured["arena_environment"].embodiment.enable_cameras is True + assert builder_cfg.num_envs == 8 + assert builder_cfg.env_spacing == 2.5 + assert builder_cfg.seed == 23 + assert builder_cfg.solve_relations is False + assert builder_cfg.device == "cuda:1" + assert make_kwargs == {FRANKA_ROPE_ARENA_ENV_ID: {"variation_recorder": variation_recorder}} + + +def test_regular_isaac_lab_run_does_not_build_arena(monkeypatch): + arena_registration_called = False + + def register_arena_environment(**_kwargs): + nonlocal arena_registration_called + arena_registration_called = True + return {} + + monkeypatch.setattr(arena_registration, "build_and_register_arena_environment", register_arena_environment) + + make_kwargs = env_registration.register_environment_for_run( + env_name="Isaac-Regular-Lab-Env-v0", + enable_cameras=False, + num_envs=4, + device="cuda:0", + seed=17, + ) + + assert make_kwargs == {} + assert arena_registration_called is False + + +def test_external_callback_routes_arena_task(monkeypatch): + callback_called = False + + def register_arena_environment_from_cli(): + nonlocal callback_called + callback_called = True + return ["hydra.option=value"] + + monkeypatch.setattr(arena_registration, "register_environment_from_cli", register_arena_environment_from_cli) + monkeypatch.setattr(sys, "argv", ["replay_demos.py", "--task", FRANKA_ROPE_ARENA_ENV_ID]) + + remaining_args = env_registration.register_environments() + + assert callback_called is True + assert remaining_args == ["hydra.option=value"] + + +def test_external_callback_skips_arena_for_regular_lab_task(monkeypatch): + callback_called = False + + def register_arena_environment_from_cli(): + nonlocal callback_called + callback_called = True + return [] + + monkeypatch.setattr(arena_registration, "register_environment_from_cli", register_arena_environment_from_cli) + monkeypatch.setattr(sys, "argv", ["replay_demos.py", "--task", "Isaac-Regular-Lab-Env-v0"]) + + assert env_registration.register_environments() == [] + assert callback_called is False + + +def test_external_callback_binds_make_kwargs_without_gym_deepcopy(): + test_env_name = "AutoData-Test-Arena-Callback-v0" + variation_recorder = object() + gym.register(id=test_env_name, entry_point=lambda **_kwargs: None) + try: + registration._bind_gym_make_kwargs( # noqa: SLF001 + test_env_name, + {"variation_recorder": variation_recorder}, + ) + + env_spec = gym.spec(test_env_name) + assert isinstance(env_spec.entry_point, partial) + assert env_spec.entry_point.keywords["variation_recorder"] is variation_recorder + assert "variation_recorder" not in env_spec.kwargs + finally: + del gym.registry[test_env_name] + + +def test_rope_asset_uses_deformable_scene_key(): + rope = FrankaRopeAsset() + name, object_cfg = rope.get_object_cfg() + assert name == "object" + assert object_cfg.prim_path == "{ENV_REGEX_NS}/Object" + assert object_cfg.spawn.usd_path.endswith("/isaac_lab_arena/franka_rope/assets/Rope.usd") + + +def test_rope_mdp_terms_are_owned_by_arena_package(): + arena_module_prefix = "autodata_examples.envs.isaac_lab_arena.franka_rope.mdp" + assert mdp.object_nodal_pos.__module__.startswith(arena_module_prefix) + assert mdp.reset_rope_nodal_state.__module__.startswith(arena_module_prefix) + assert mdp.rope_ends_close_tracked.__module__.startswith(arena_module_prefix) + + +def test_rope_task_exposes_annotation_and_success_terms(): + task = FrankaRopeTask() + assert task.get_observation_cfg().subtask_terms.grasp is not None + assert task.get_termination_cfg().success is not None + assert task.get_events_cfg().reset_object_position is not None + + +def test_rope_embodiment_preserves_arena_observation_names(): + embodiment = FrankaRopeIKEmbodiment(enable_cameras=False) + assert embodiment.name == "franka_rope_ik" + assert embodiment.observation_config.policy.eef_pos is not None + assert embodiment.observation_config.policy.eef_quat is not None + assert embodiment.observation_config.policy.object_nodal_pos is not None + assert embodiment.event_config.randomize_franka_joint_state.params["position_range"] == (-0.02, 0.02) + assert embodiment.reward_config is None 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 2cfcb4a..7c258a2 100644 --- a/autodata_tests/interfaces/tasks/test_generation_policy_spec.py +++ b/autodata_tests/interfaces/tasks/test_generation_policy_spec.py @@ -18,6 +18,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 @@ -25,11 +26,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/test_data/annotated_dataset_franka_rope_softmimicgen.hdf5 b/autodata_tests/test_data/annotated_dataset_franka_rope_softmimicgen.hdf5 new file mode 100644 index 0000000..2c25c2e --- /dev/null +++ b/autodata_tests/test_data/annotated_dataset_franka_rope_softmimicgen.hdf5 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6fa2b5db4752242eb91cec522053ad1a040741764dcb71ebab6f3abe80e4992c +size 30514205 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 8b64959..cbe0929 100644 --- a/scripts/annotate_demos.py +++ b/scripts/annotate_demos.py @@ -91,7 +91,9 @@ import contextlib # noqa: E402 import gymnasium as gym # noqa: E402 import math # noqa: E402 +import numpy as np # noqa: E402 import os # noqa: E402 +import random # noqa: E402 import torch # noqa: E402 from collections.abc import Callable # noqa: E402 @@ -103,6 +105,7 @@ from isaaclab.utils.datasets import EpisodeData, HDF5DatasetFileHandler # noqa: E402 from autodata_core.pool import DataGenInfoPool # noqa: E402 +from autodata_examples.envs import register_environment_for_run # noqa: E402 from autodata_interfaces.datastream import Datastream # noqa: E402 from autodata_interfaces.embodiments import embodiment_adapter_from_yaml # noqa: E402 from autodata_interfaces.env import get_env_name_from_dataset, setup_env_config, setup_output_paths # noqa: E402 @@ -145,6 +148,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), } @@ -220,6 +224,18 @@ def main() -> int: task_descriptor = TaskDescriptor.from_yaml(args_cli.task_descriptor) generation_policy = task_descriptor.get_generation_policy() + random.seed(generation_policy.seed) + np.random.seed(generation_policy.seed) + torch.manual_seed(generation_policy.seed) + + env_make_kwargs = register_environment_for_run( + env_name=env_name, + enable_cameras=args_cli.enable_cameras, + num_envs=1, + device=args_cli.device, + seed=generation_policy.seed, + ) + # Start signals are required only by SkillGen. annotate_start_signals = generation_policy.use_skillgen @@ -254,7 +270,7 @@ def main() -> int: # Only export episodes we explicitly mark successful (i.e. fully annotated). env_cfg.recorders.dataset_export_mode = DatasetExportMode.EXPORT_SUCCEEDED_ONLY - env = gym.make(env_name, cfg=env_cfg).unwrapped + env = gym.make(env_name, cfg=env_cfg, **env_make_kwargs.get(env_name, {})).unwrapped try: # Create the Datastream embodiment_adapter = embodiment_adapter_from_yaml(args_cli.embodiment) diff --git a/scripts/generate_dataset.py b/scripts/generate_dataset.py index 7f3f311..2b7ce38 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 \\ @@ -24,6 +24,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`. @@ -37,7 +38,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( @@ -119,6 +120,7 @@ from autodata_core import DataGenerator, get_algorithm # noqa: E402 from autodata_core.algorithms import REGISTERED_ALGORITHMS # noqa: E402 +from autodata_examples.envs import register_environment_for_run # noqa: E402 from autodata_interfaces.datastream import Datastream # noqa: E402 from autodata_interfaces.embodiments import embodiment_adapter_from_yaml # noqa: E402 from autodata_interfaces.env import ( # noqa: E402 @@ -286,6 +288,18 @@ def main() -> None: if args_cli.generation_num_trials is not None: generation_policy_params.num_trials = args_cli.generation_num_trials + random.seed(generation_policy_params.seed) + np.random.seed(generation_policy_params.seed) + torch.manual_seed(generation_policy_params.seed) + + env_make_kwargs = register_environment_for_run( + env_name=env_name, + enable_cameras=args_cli.enable_cameras, + num_envs=args_cli.num_envs, + device=args_cli.device, + seed=generation_policy_params.seed, + ) + # The algorithm's start-signal expectation is a class attribute, so we resolve it before # instantiating (the Datastream needs it, and SkillGen can only be instantiated once the # planners exist, which in turn need the Datastream). @@ -306,11 +320,7 @@ def main() -> None: env_profile=env_profile, ) - env = gym.make(env_name, cfg=env_cfg).unwrapped - - random.seed(generation_policy_params.seed) - np.random.seed(generation_policy_params.seed) - torch.manual_seed(generation_policy_params.seed) + env = gym.make(env_name, cfg=env_cfg, **env_make_kwargs.get(env_name, {})).unwrapped env.reset()