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..e797c8d 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..19ed6fd 100644 --- a/autodata_core/pool.py +++ b/autodata_core/pool.py @@ -124,8 +124,30 @@ def _add_episode(self, episode: EpisodeData) -> None: raise ValueError("Episode lacks 'datagen_info' obs annotations") eef_pose = ep_grp["obs"]["datagen_info"]["eef_pose"] - object_poses_dict = ep_grp["obs"]["datagen_info"]["object_pose"] + object_poses_dict = ep_grp["obs"]["datagen_info"].get("object_pose") + object_nodal_positions_dict = ep_grp["obs"]["datagen_info"].get("object_nodal_position") target_eef_pose = ep_grp["obs"]["datagen_info"]["target_eef_pose"] + + # Legacy single-arm datasets use names such as "robot0". With exactly one + # recorded and one configured EEF, the correspondence is unambiguous. + # Normalize both pose mappings in memory without modifying the source episode. + expected_eefs = set(self.task_descriptor.get_eef_names()) + recorded_eefs = set(eef_pose) + assert recorded_eefs == set(target_eef_pose), ( + f"Source eef_pose names {sorted(recorded_eefs)} do not match " + f"target_eef_pose names {sorted(target_eef_pose)}." + ) + if recorded_eefs != expected_eefs: + assert len(recorded_eefs) == len(expected_eefs) == 1, ( + f"Source EEF names {sorted(recorded_eefs)} do not match configured names " + f"{sorted(expected_eefs)}. Automatic name matching requires exactly one " + "source and one configured EEF; multi-arm datasets require matching names." + ) + source_eef = next(iter(recorded_eefs)) + configured_eef = next(iter(expected_eefs)) + eef_pose = {configured_eef: eef_pose[source_eef]} + target_eef_pose = {configured_eef: target_eef_pose[source_eef]} + subtask_term_signals_dict = ep_grp["obs"]["datagen_info"]["subtask_term_signals"] subtask_start_signals_dict = ep_grp["obs"]["datagen_info"].get("subtask_start_signals") @@ -134,6 +156,7 @@ def _add_episode(self, episode: EpisodeData) -> None: ep_datagen_info = DatagenInfo( eef_pose=eef_pose, object_poses=object_poses_dict, + object_nodal_positions=object_nodal_positions_dict, subtask_start_signals=subtask_start_signals_dict, subtask_term_signals=subtask_term_signals_dict, target_eef_pose=target_eef_pose, diff --git a/autodata_core/selection_strategy.py b/autodata_core/selection_strategy.py index df1442f..c8aefe1 100644 --- a/autodata_core/selection_strategy.py +++ b/autodata_core/selection_strategy.py @@ -10,6 +10,7 @@ from __future__ import annotations import abc +import numpy as np import torch from typing import Any @@ -51,6 +52,7 @@ def select_source_demo( eef_pose: torch.Tensor, object_pose: torch.Tensor | None, src_subtask_datagen_infos: list, + object_nodal_positions: torch.Tensor | None = None, ) -> int: """Return the index of the source demo whose subtask segment best fits the current scene. @@ -58,6 +60,7 @@ def select_source_demo( eef_pose: Current 4x4 EEF pose [m, rad]. object_pose: Current 4x4 pose of the subtask's reference object [m, rad], or None. src_subtask_datagen_infos: Per-source-demo :class:`DatagenInfo` slices covering this subtask. + object_nodal_positions: Current deformable-object nodes shaped ``(N, 3)`` [m], or None. """ raise NotImplementedError @@ -67,7 +70,13 @@ class RandomStrategy(SelectionStrategy): NAME = "random" - def select_source_demo(self, eef_pose, object_pose, src_subtask_datagen_infos) -> int: + def select_source_demo( + self, + eef_pose, + object_pose, + src_subtask_datagen_infos, + object_nodal_positions=None, + ) -> int: n_src_demo = len(src_subtask_datagen_infos) return torch.randint(0, n_src_demo, (1,)).item() @@ -85,6 +94,7 @@ def select_source_demo( pos_weight: float = 1.0, rot_weight: float = 1.0, nn_k: int = 3, + object_nodal_positions: torch.Tensor | None = None, ) -> int: src_object_poses = [] for di in src_subtask_datagen_infos: @@ -126,6 +136,7 @@ def select_source_demo( pos_weight: float = 1.0, rot_weight: float = 1.0, nn_k: int = 3, + object_nodal_positions: torch.Tensor | None = None, ) -> int: src_eef_poses = [] src_object_poses = [] @@ -165,3 +176,55 @@ def select_source_demo( nn_k = min(nn_k, len(dists)) rand_k = torch.randint(0, nn_k, (1,)).item() return torch.argsort(dists)[:nn_k][rand_k] + + +class RegistrationCostStrategy(SelectionStrategy): + """Pick a source segment using deformable-object TPS registration cost.""" + + NAME = "registration_cost" + + def select_source_demo( + self, + eef_pose: torch.Tensor, + object_pose: torch.Tensor | None, + src_subtask_datagen_infos: list, + object_nodal_positions: torch.Tensor | None = None, + bend_coef: float = 0.1, + rot_reg: float = 1e-3, + nn_k: int = 3, + ) -> int: + """Choose uniformly among the ``nn_k`` source shapes with lowest TPS cost.""" + + from autodata_core.deformable_transforms import nodal_registration_cost + + assert object_nodal_positions is not None, "registration_cost requires current object nodal positions" + assert nn_k >= 1, f"nn_k must be at least 1, got {nn_k}" + + valid_candidates: list[tuple[int, float]] = [] + for demo_index, datagen_info in enumerate(src_subtask_datagen_infos): + assert ( + datagen_info.object_nodal_positions is not None + ), "registration_cost requires source object nodal positions" + source_objects = list(datagen_info.object_nodal_positions.values()) + assert ( + len(source_objects) == 1 + ), f"registration_cost expects exactly one source object, got {len(source_objects)}" + source_nodal_positions = source_objects[0][0] + try: + cost = nodal_registration_cost( + source_nodal_positions, + object_nodal_positions, + bend_coef=bend_coef, + rot_reg=rot_reg, + ) + except (AssertionError, np.linalg.LinAlgError, ValueError): + continue + if np.isfinite(cost): + valid_candidates.append((demo_index, cost)) + + assert valid_candidates, "registration_cost could not compute a finite TPS cost for any source demo" + costs_tensor = torch.tensor([cost for _, cost in valid_candidates], dtype=torch.float32) + nn_k = min(nn_k, len(valid_candidates)) + rand_k = torch.randint(0, nn_k, (1,)).item() + selected_candidate = torch.argsort(costs_tensor)[:nn_k][rand_k].item() + return valid_candidates[selected_candidate][0] diff --git a/autodata_examples/envs/__init__.py b/autodata_examples/envs/__init__.py new file mode 100644 index 0000000..d829881 --- /dev/null +++ b/autodata_examples/envs/__init__.py @@ -0,0 +1,65 @@ +# 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, + presets: str | None = None, +) -> 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. + presets: Optional Arena physics backend preset; None preserves the environment default. + + 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, + presets=presets, + ) + + +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/assets/RopeNewton.usda b/autodata_examples/envs/isaac_lab_arena/franka_rope/assets/RopeNewton.usda new file mode 100644 index 0000000..87de490 --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/assets/RopeNewton.usda @@ -0,0 +1,19 @@ +#usda 1.0 +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +( + defaultPrim = "World" +) + +# Keep the shared simulation topology, round visual mesh, and material. +# Newton uses the simulation tetrahedra directly; the PhysX collision mesh is redundant. +def Xform "World" ( + prepend references = @Rope.usd@ +) +{ + over "collision_mesh" ( + active = false + ) + { + } +} 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..b831c80 --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/embodiment.py @@ -0,0 +1,216 @@ +# 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.physics_backend import PhysicsBackend +from isaaclab_arena.utils.pose import Pose +from isaaclab_physx.renderers import IsaacRtxRendererCfg, IsaacRtxRendererGlobalSettingsCfg + +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"], + renderer_cfg=IsaacRtxRendererCfg(global_settings=IsaacRtxRendererGlobalSettingsCfg(antialiasing_mode="DLSS")), + 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"], + renderer_cfg=IsaacRtxRendererCfg(global_settings=IsaacRtxRendererGlobalSettingsCfg(antialiasing_mode="DLSS")), + 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 + + def _configure_physics_backend(self, backend: PhysicsBackend) -> None: + """Apply Newton gravity compensation and mesh contact settings to the Franka.""" + if backend is PhysicsBackend.NEWTON: + from isaaclab_newton.renderers import NewtonWarpRendererCfg + from isaaclab_newton.sim.schemas import ( + MujocoRigidBodyCfg, + NewtonMaterialPropertiesCfg, + NewtonSDFCollisionPropertiesCfg, + ) + + from .newton_cfg import spawn_rope_franka_newton + + spawn = self.scene_config.robot.spawn + # Match the shared high-PD Franka's PhysX-only disable_gravity setting. + spawn.rigid_props = MujocoRigidBodyCfg(gravcomp=1.0) + spawn.make_uninstanceable = True + spawn.func = spawn_rope_franka_newton + # Full-surface soft contacts require SDFs for the mesh colliders. + spawn.collision_props = NewtonSDFCollisionPropertiesCfg(sdf_max_resolution=64) + spawn.physics_material = NewtonMaterialPropertiesCfg( + static_friction=3.0, dynamic_friction=3.0, restitution=0.0 + ) + # Maintain pose tracking under the deformable solver's contact feedback. + self.scene_config.robot.actuators["panda_shoulder"].stiffness = 800.0 + self.scene_config.robot.actuators["panda_forearm"].stiffness = 800.0 + # Read rigid poses and deformable triangles directly from Newton state. + for camera_name in self.camera_config.camera_names(): + getattr(self.camera_config, camera_name).renderer_cfg = NewtonWarpRendererCfg() + + +@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..6f6f633 --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/environment.py @@ -0,0 +1,129 @@ +# 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, rope_newton_spawn_cfg +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 timing and asset/solver settings for Arena's selected backend.""" + + from isaaclab_newton.physics import NewtonCfg + + env_cfg.decimation = 5 + env_cfg.sim.dt = 0.01 + env_cfg.sim.render_interval = 2 + env_cfg.num_rerenders_on_reset = 1 + if isinstance(env_cfg.sim.physics, NewtonCfg): + from isaaclab_newton.sim.schemas import MujocoCollisionCfg + + from .newton_cfg import rope_newton_physics_cfg, spawn_rope_table_newton + + env_cfg.sim.physics = rope_newton_physics_cfg() + # Refresh Kit once per control step; the solver still advances at 1 ms. + env_cfg.sim.render_interval = env_cfg.decimation + env_cfg.scene.object.spawn = rope_newton_spawn_cfg() + # A discontinuous half-mesh reset can invert the connecting tetrahedra. + # Blend the same sampled poses through the middle half of the rope. + env_cfg.events.reset_object_position.params["transition_length"] = 0.4 + env_cfg.scene.replicate_physics = True + env_cfg.scene.table.spawn.make_uninstanceable = True + env_cfg.scene.table.spawn.func = spawn_rope_table_newton + # A firm, critically damped table contact in MJWarp. Priority keeps the + # table response from being softened by the gripper's default material. + # These MuJoCo parameters do not change VBD's rope contact properties. + env_cfg.scene.table.spawn.collision_props = MujocoCollisionCfg(priority=1, solref=(0.004, 1.0)) + else: + 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..b2cf0f1 --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/mdp/events.py @@ -0,0 +1,127 @@ +# 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"), + transition_length: float | None = None, +) -> 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. + transition_length: Length [m] over which to blend the half-rope transform. + None preserves the original piecewise reset used by PhysX. + """ + + 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) + if transition_length is None: + nodal_state[..., node_ids, :3] = rope.transform_nodal_pos(nodal_state[..., node_ids, :3], position, quaternion) + else: + assert transition_length > 0.0 + nodes = nodal_state[..., :3] + center = nodes.mean(dim=1, keepdim=True) + half_center = nodes[:, node_ids].mean(dim=1, keepdim=True) + axis = torch.nn.functional.normalize(half_center - center, dim=-1) + distance = ((nodes - center) * axis).sum(dim=-1, keepdim=True) + fraction = (distance / transition_length + 0.5).clamp(0.0, 1.0) + # A C2-continuous transition avoids tearing/inverting the tetrahedra at + # the boundary of the two halves. End sections retain their rigid poses. + weight = fraction**3 * (10.0 - 15.0 * fraction + 6.0 * fraction**2) + transformed = ( + math_utils.quat_apply(quaternion[:, None, :].expand(-1, nodes.shape[1], -1), nodes - half_center) + + half_center + + position[:, None, :] + ) + nodal_state[..., :3] = nodes + weight * (transformed - nodes) + + 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/newton_cfg.py b/autodata_examples/envs/isaac_lab_arena/franka_rope/newton_cfg.py new file mode 100644 index 0000000..6722bec --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/newton_cfg.py @@ -0,0 +1,141 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Newton rigid--deformable physics for the shared Franka rope task.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from isaaclab.sim.spawners.from_files import UsdFileCfg, spawn_from_usd +from isaaclab.sim.utils import clone +from isaaclab_contrib.coupling import CouplerEntryCfg, CouplerProxyCfg, CouplerProxyMappingCfg +from isaaclab_newton.physics import ( + MJWarpSolverCfg, + NewtonCfg, + NewtonCollisionPipelineCfg, + NewtonSoftContactCfg, + VBDSolverCfg, +) + +if TYPE_CHECKING: + from pxr import Usd + + +@clone +def spawn_rope_franka_newton( + prim_path: str, + cfg: UsdFileCfg, + translation: tuple[float, float, float] | None = None, + orientation: tuple[float, float, float, float] | None = None, + **kwargs, +) -> Usd.Prim: + """Use analytic finger colliders for stable volumetric rope contacts. + + The boxes occupy the original collision meshes' bounds. Visual meshes and + articulation joints remain those of the shared Franka. + + Args: + prim_path: Prim path or pattern at which to spawn the robot. + cfg: Robot USD spawner configuration. + translation: Translation relative to the parent [m]. + orientation: Rotation quaternion in XYZW order. + **kwargs: Additional arguments forwarded to the USD spawner. + + Returns: + The spawned robot prim. + """ + from pxr import Gf, UsdGeom + + prim = spawn_from_usd(prim_path, cfg, translation, orientation, **kwargs) + for finger in ("panda_leftfinger", "panda_rightfinger"): + collider = prim.GetStage().GetPrimAtPath(f"{prim.GetPath()}/{finger}/collisions/collisions") + assert collider.IsValid() and collider.GetTypeName() == "Mesh", "Expected the legacy Franka finger collider" + points = UsdGeom.Mesh(collider).GetPointsAttr().Get() + lower = Gf.Vec3d(*(min(point[axis] for point in points) for axis in range(3))) + upper = Gf.Vec3d(*(max(point[axis] for point in points) for axis in range(3))) + collider.SetTypeName("Cube") + cube = UsdGeom.Cube(collider) + cube.CreateSizeAttr(1.0) + cube.AddTranslateOp(opSuffix="rope_contact_center").Set((lower + upper) * 0.5) + cube.AddScaleOp(opSuffix="rope_contact_size").Set(upper - lower) + return prim + + +@clone +def spawn_rope_table_newton( + prim_path: str, + cfg: UsdFileCfg, + translation: tuple[float, float, float] | None = None, + orientation: tuple[float, float, float, float] | None = None, + **kwargs, +) -> Usd.Prim: + """Use an analytic box for the Seattle table's existing unit-cube collider. + + Args: + prim_path: Prim path or pattern at which to spawn the table. + cfg: Table USD spawner configuration. + translation: Translation relative to the parent [m]. + orientation: Rotation quaternion in XYZW order. + **kwargs: Additional arguments forwarded to the USD spawner. + + Returns: + The spawned table prim. + """ + from pxr import UsdGeom + + prim = spawn_from_usd(prim_path, cfg, translation, orientation, **kwargs) + collider = prim.GetStage().GetPrimAtPath(f"{prim.GetPath()}/Collisions/Cube") + assert collider.IsValid() and collider.GetTypeName() == "Mesh", "Expected the Seattle table's cube mesh collider" + collider.SetTypeName("Cube") + UsdGeom.Cube(collider).CreateSizeAttr(1.0) + return prim + + +def rope_newton_physics_cfg() -> NewtonCfg: + """Couple the Franka's rigid dynamics to the rope's volumetric VBD solve.""" + return NewtonCfg( + solver_cfg=CouplerProxyCfg( + class_type="autodata_examples.envs.isaac_lab_arena.franka_rope.newton_visuals:NewtonRopeManager", + entries=[ + CouplerEntryCfg( + name="rigid", + solver_cfg=MJWarpSolverCfg(cone="elliptic", ls_iterations=20, integrator="implicitfast"), + bodies=[r"/World/envs/env_[^/]+/Robot"], + include_static_shapes=True, + ), + CouplerEntryCfg( + name="soft", + solver_cfg=VBDSolverCfg( + iterations=10, + rigid_body_particle_contact_buffer_size=4096, + particle_enable_self_contact=True, + particle_self_contact_radius=0.002, + particle_self_contact_margin=0.003, + ), + all_particles=True, + ), + ], + proxies=[ + CouplerProxyMappingCfg( + source="rigid", + destination="soft", + bodies=[ + r"/World/envs/env_[^/]+/Robot/.*panda_hand", + r"/World/envs/env_[^/]+/Robot/.*panda_(left|right)finger", + ], + mode="lagged", + collide_interval=1, + collision_pipeline=NewtonCollisionPipelineCfg(enable_rigid_soft_full_surface_contact=True), + ), + ], + iterations=1, + ), + soft_contact_cfg=NewtonSoftContactCfg( + # Resolve contacts around the narrow rope, with enough friction to carry it. + soft_contact_ke=1.0e5, + soft_contact_kd=10.0, + soft_contact_mu=10.0, + ), + num_substeps=10, + ) diff --git a/autodata_examples/envs/isaac_lab_arena/franka_rope/newton_visuals.py b/autodata_examples/envs/isaac_lab_arena/franka_rope/newton_visuals.py new file mode 100644 index 0000000..c03af54 --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/newton_visuals.py @@ -0,0 +1,117 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Deform the rope's original round visual mesh for Newton's Kit viewport.""" + +import numpy as np + +import warp as wp +from isaaclab.scene_data.deformable_vis_remap import VolumeVisRemap, build_volume_vis_barycentric_remap +from isaaclab.sim.utils import find_matching_prim_paths +from isaaclab.sim.utils.stage import get_current_stage +from isaaclab_contrib.coupling.coupler import NewtonCouplerManager + + +@wp.kernel(enable_backward=False) +def _skin_rope_points( + points: wp.fabricarrayarray(dtype=wp.vec3f), + matrices: wp.fabricarray(dtype=wp.mat44d), + offsets: wp.fabricarray(dtype=wp.uint32), + particles: wp.array(dtype=wp.vec3f), + indices: wp.array2d(dtype=wp.int32), + weights: wp.array2d(dtype=wp.float32), +): + instance, vertex = wp.tid() + offset = int(offsets[instance]) + world = wp.vec3f(0.0) + for corner in range(4): + world += weights[vertex, corner] * particles[offset + indices[vertex, corner]] + world_to_local = wp.inverse(wp.transpose(wp.mat44f(matrices[instance]))) + points[instance][vertex] = wp.transform_point(world_to_local, world) + + +class NewtonRopeManager(NewtonCouplerManager): + """Add visual skinning to the standard rigid/deformable Newton coupler. + + The pinned Lab backend copies simulation nodes into visual vertices by index. + This rope has different simulation and visual topologies. Embed its original + round mesh in the simulation tetrahedra and update every visual vertex before + Kit renders. The solver, collision geometry, and particle state are unchanged. + """ + + _rope_vis_remap: VolumeVisRemap | None = None + _rope_offset_attr = "autodata:rope:particleOffset" + + @classmethod + def _initialize_rope_skin(cls) -> None: + import usdrt + from pxr import Gf, UsdGeom + + assert len(cls._deformable_registry) == 1, "The rope task expects one deformable asset" + entry = cls._deformable_registry[0] + stage = get_current_stage() + paths = find_matching_prim_paths(entry.vis_mesh_prim_path) + assert paths, "Rope visual mesh is missing" + visual = UsdGeom.Mesh(stage.GetPrimAtPath(paths[0])) + simulation = UsdGeom.TetMesh( + stage.GetPrimAtPath(visual.GetPath().GetParentPath().AppendChild("simulation_mesh")) + ) + xforms = UsdGeom.XformCache() + visual_to_sim = ( + xforms.GetLocalToWorldTransform(visual.GetPrim()) + * xforms.GetLocalToWorldTransform(simulation.GetPrim()).GetInverse() + ) + visual_points = np.asarray( + [visual_to_sim.Transform(Gf.Vec3d(point)) for point in visual.GetPointsAttr().Get()], dtype=np.float32 + ) + cls._rope_vis_remap = build_volume_vis_barycentric_remap( + np.asarray(simulation.GetPointsAttr().Get(), dtype=np.float32), + np.asarray(simulation.GetTetVertexIndicesAttr().Get(), dtype=np.int32).reshape(-1), + visual_points, + device=str(cls._device), + ) + assert cls._rope_vis_remap is not None, "Cannot embed the rope visual mesh in its simulation tetrahedra" + for path in paths: + env_index = int(path.split("/envs/env_", 1)[1].split("/", 1)[0]) + prim = cls._usdrt_stage.GetPrimAtPath(path) + usdrt.Rt.Xformable(prim).SetWorldXformFromUsd() + prim.CreateAttribute(cls._rope_offset_attr, usdrt.Sdf.ValueTypeNames.UInt, True).Set( + entry.particle_offsets[env_index] + ) + + @classmethod + def pre_render(cls) -> None: + """Refresh the round skin before Kit renders, including after state restores.""" + super().pre_render() + if cls._usdrt_stage is None or cls._clone_physics_only or cls._state_0 is None: + return + import usdrt + + if cls._rope_vis_remap is None: + cls._initialize_rope_skin() + selection = cls._usdrt_stage.SelectPrims( + require_attrs=[ + (usdrt.Sdf.ValueTypeNames.Point3fArray, "points", usdrt.Usd.Access.ReadWrite), + (usdrt.Sdf.ValueTypeNames.UInt, cls._rope_offset_attr, usdrt.Usd.Access.Read), + (usdrt.Sdf.ValueTypeNames.Matrix4d, "omni:fabric:worldMatrix", usdrt.Usd.Access.Read), + ], + device=str(cls._device), + ) + wp.launch( + _skin_rope_points, + dim=(selection.GetCount(), cls._rope_vis_remap.tet_vertex_indices.shape[0]), + inputs=[ + wp.fabricarrayarray(data=selection, attrib="points", dtype=wp.vec3f), + wp.fabricarray(data=selection, attrib="omni:fabric:worldMatrix"), + wp.fabricarray(data=selection, attrib=cls._rope_offset_attr), + cls._state_0.particle_q, + cls._rope_vis_remap.tet_vertex_indices, + cls._rope_vis_remap.bary_weights, + ], + device=cls._device, + ) + + @classmethod + def _solver_specific_clear(cls) -> None: + cls._rope_vis_remap = None + super()._solver_specific_clear() 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..33ff0c3 --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/rope_asset.py @@ -0,0 +1,61 @@ +# 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" + + +def rope_newton_spawn_cfg() -> UsdFileCfg: + """Load the shared rope topology with Newton material properties.""" + from isaaclab_newton.sim.schemas import NewtonDeformableBodyPropertiesCfg + from isaaclab_newton.sim.spawners.materials import NewtonDeformableBodyMaterialCfg + + # VBD needs a compliant volume material for the gripper to bend this slender + # mesh. Keep the shared topology and density, with backend-specific elasticity. + youngs_modulus, poissons_ratio = 2.0e5, 0.45 + return UsdFileCfg( + usd_path=str(_ROPE_USD_PATH.with_name("RopeNewton.usda")), + deformable_props=NewtonDeformableBodyPropertiesCfg(), + physics_material_path="NewtonPhysicsMaterial", + physics_material=NewtonDeformableBodyMaterialCfg( + density=1000.0, + particle_radius=0.002, + k_mu=youngs_modulus / (2.0 * (1.0 + poissons_ratio)), + k_lambda=youngs_modulus * poissons_ratio / ((1.0 + poissons_ratio) * (1.0 - 2.0 * poissons_ratio)), + k_damp=1000.0, + ), + ) + + +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..c7127f9 --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/task.py @@ -0,0 +1,152 @@ +# 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 + +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.progress_tracking.completion_criteria import CompletionCriteria +from isaaclab_arena.tasks.task_base import TaskBase +from isaaclab_arena.tasks.task_termination_cfg import TaskTerminationCfg + +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"), + }, + ) + + +@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 = TaskTerminationCfg( + timeout_s=self.episode_length_s, + success=[ + CompletionCriteria( + name="rope_ends_close", + predicate_sequence=[mdp.rope_ends_close_tracked], + ), + ], + failures={ + "object_dropping": DoneTerm( + func=mdp.rope_below_minimum, + params={ + "minimum_height": -0.05, + "object_cfg": SceneEntityCfg("object"), + }, + ), + }, + ) + + 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) -> TaskTerminationCfg: + """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..4c4a666 --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/registration.py @@ -0,0 +1,131 @@ +# 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, + presets: str | None = None, +) -> 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. + presets: Optional physics backend preset; None preserves the environment default. + + 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, + presets=presets, + ), + ) + 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..3794901 --- /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: false + 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 1b21c05..5264b9b 100644 --- a/autodata_interfaces/datastream/datastream.py +++ b/autodata_interfaces/datastream/datastream.py @@ -255,6 +255,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]: 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 e799341..1de1485 100644 --- a/autodata_interfaces/env/isaaclab_env_interface.py +++ b/autodata_interfaces/env/isaaclab_env_interface.py @@ -23,7 +23,6 @@ import isaaclab.sim as sim_utils from isaaclab.assets import RigidObjectCfg from isaaclab.envs.mdp.actions.pink_task_space_actions import PinkInverseKinematicsAction -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 @@ -34,6 +33,9 @@ from autodata_interfaces.env.env_profile import EnvironmentProfile, convert_event_params from autodata_interfaces.env.pink_ik_action import PinkInverseKinematicsActionSharedUrdf +from autodata_interfaces.env.recorders import make_action_state_recorder_manager_cfg +from autodata_interfaces.env.reset_request import EnvResetRequest +from autodata_interfaces.env.scene_state import get_scene_state from autodata_interfaces.tasks.generation_policy_spec import GenerationPolicy @@ -220,7 +222,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(record_images=enable_cameras) else: env_cfg.recorders = recorder_cfg env_cfg.recorders.dataset_export_dir_path = output_dir @@ -256,7 +258,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). @@ -269,13 +271,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() @@ -283,25 +289,83 @@ def env_loop( raise exc return False while not env_reset_queue.empty(): - env_id_tensor[0] = env_reset_queue.get_nowait() + request = env_reset_queue.get_nowait() + assert isinstance( + request, EnvResetRequest + ), f"env_reset_queue entries must be EnvResetRequest instances, got {type(request).__name__}" + assert ( + 0 <= request.env_id < env.num_envs + ), f"Reset environment index {request.env_id} is outside [0, {env.num_envs})." + assert ( + request.env_id not in settling_requests + ), f"Environment {request.env_id} already has a reset settling request in progress." + env_id_tensor[0] = request.env_id env.reset(env_ids=env_id_tensor) - env_reset_queue.task_done() + + if reset_settling_steps > 0: + settling_requests[request.env_id] = request + settling_steps_remaining[request.env_id] = reset_settling_steps + else: + request.completion.set_result(None) + env_reset_queue.task_done() + + expected_action_count = env.num_envs - len(settling_requests) + assert env_action_queue.qsize() <= expected_action_count, ( + f"Received {env_action_queue.qsize()} queued actions for only {expected_action_count} " + "non-settling environments." + ) actions = torch.zeros(env.action_space.shape) - # batch-fetch all per-env actions in one gather instead of sequential blocking calls - get_tasks = [env_action_queue.get() for _ in range(env.num_envs)] - results = asyncio_event_loop.run_until_complete(asyncio.gather(*get_tasks)) + # Settling environments use the zero action. Batch-fetch one real action for every + # other environment so their generation trajectories continue without interruption. + action_count = env.num_envs - len(settling_requests) + get_tasks = [env_action_queue.get() for _ in range(action_count)] + results = asyncio_event_loop.run_until_complete(asyncio.gather(*get_tasks)) if get_tasks else [] + action_env_ids: set[int] = set() for env_id, action in results: + assert ( + env_id not in settling_requests + ), f"Environment {env_id} supplied an action while its reset was still settling." + assert env_id not in action_env_ids, f"Received multiple actions for environment {env_id}." + action_env_ids.add(env_id) actions[env_id] = action + expected_action_env_ids = set(range(env.num_envs)) - settling_requests.keys() + assert ( + action_env_ids == expected_action_env_ids + ), f"Expected actions for environments {sorted(expected_action_env_ids)}, got {sorted(action_env_ids)}." # perform action on environment env.step(actions) # mark done so the data generators can continue with the step results - for _ in range(env.num_envs): + for _ in range(action_count): env_action_queue.task_done() + settled_env_ids: list[int] = [] + for env_id in settling_steps_remaining: + settling_steps_remaining[env_id] -= 1 + if settling_steps_remaining[env_id] == 0: + settled_env_ids.append(env_id) + + if settled_env_ids: + settled_env_ids_tensor = torch.tensor(settled_env_ids, dtype=torch.int64, device=env.device) + # Discard temporary settle frames, then capture the settled state as the generated + # episode's initial state before releasing the corresponding generators. + env.recorder_manager.reset(env_ids=settled_env_ids_tensor) + # Settling is not a reset: Arena correctly ignores post-reset callbacks here. + # Snapshot the actual settled state through the public recorder API instead. + env.recorder_manager.add_to_episodes( + "initial_state", + get_scene_state(env.scene, is_relative=True, env_ids=settled_env_ids_tensor), + env_ids=settled_env_ids_tensor, + ) + for env_id in settled_env_ids: + request = settling_requests.pop(env_id) + settling_steps_remaining.pop(env_id) + request.completion.set_result(None) + env_reset_queue.task_done() + if prev_num_attempts != stats["num_attempts"]: prev_num_attempts = stats["num_attempts"] num_success = stats["num_success"] diff --git a/autodata_interfaces/env/recorders.py b/autodata_interfaces/env/recorders.py new file mode 100644 index 0000000..a12bc5c --- /dev/null +++ b/autodata_interfaces/env/recorders.py @@ -0,0 +1,29 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +from isaaclab.envs.mdp.recorders.recorders_cfg import ActionStateRecorderManagerCfg +from isaaclab.managers.recorder_manager import RecorderTerm, RecorderTermCfg + + +class PreStepCameraObservationsRecorder(RecorderTerm): + """Record image observations exposed in a separate camera group.""" + + def record_pre_step(self): + """Return camera observations alongside the recorded policy observations.""" + return "obs", self._env.obs_buf.get("camera_obs", {}) + + +def make_action_state_recorder_manager_cfg(record_images: bool = False) -> ActionStateRecorderManagerCfg: + """Return Isaac Lab's action-state recorder with optional camera observations. + + Args: + record_images: Include the separate camera observation group when present. + """ + + cfg = ActionStateRecorderManagerCfg() + if record_images: + cfg.record_pre_step_camera_observations = RecorderTermCfg(class_type=PreStepCameraObservationsRecorder) + 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..29fe0fa --- /dev/null +++ b/autodata_interfaces/env/scene_state.py @@ -0,0 +1,36 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Scene-state reads with optional environment selection.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + + +def get_scene_state(scene: Any, is_relative: bool = True, env_ids: Sequence[int] | None = None) -> dict: + """Return scene state in world or per-environment coordinates. + + Isaac Lab owns the coordinate conversion for every supported asset type. + Selection is applied afterwards to retain the requested environment order. + + Args: + scene: Isaac Lab interactive scene. + is_relative: Whether positions are relative to each environment origin. + env_ids: Environment rows to include, in the requested order; None selects all. + + Returns: + Nested scene-state mapping containing cloned tensors. + """ + + state = scene.get_state(is_relative=is_relative) + if env_ids is None: + return state + + def select_envs(value): + if isinstance(value, dict): + return {key: select_envs(item) for key, item in value.items()} + return value[env_ids] + + return select_envs(state) diff --git a/autodata_interfaces/tasks/generation_policy_spec.py b/autodata_interfaces/tasks/generation_policy_spec.py index 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..70f5490 100644 --- a/autodata_tests/core/test_pool.py +++ b/autodata_tests/core/test_pool.py @@ -86,6 +86,27 @@ def _skillgen_task() -> TaskDescriptor: }) +def _softmimicgen_task() -> TaskDescriptor: + return TaskDescriptor.from_dict({ + "name": "rope", + "algo": "softmimicgen", + "subtasks": { + "franka": [ + { + "object_ref": "rope", + "subtask_term_signal": "grasp", + "algo_params": {"object_soft": True}, + }, + { + "object_ref": "rope", + "subtask_term_signal": "", + "algo_params": {"object_soft": True}, + }, + ] + }, + }) + + def _step_signal(length: int, edge: int) -> torch.Tensor: """A per-step step-function: False before ``edge``, True from ``edge`` on.""" sig = torch.ones(length, dtype=torch.bool) @@ -157,6 +178,18 @@ def _skillgen_episode( ) +def _softmimicgen_episode(actions_len: int = 10, term_edge: int = 4) -> types.SimpleNamespace: + return _episode( + actions_len=actions_len, + datagen_info={ + "eef_pose": {"franka": torch.zeros(actions_len, 4, 4)}, + "object_nodal_position": {"rope": torch.zeros(actions_len, 12, 3)}, + "target_eef_pose": {"franka": torch.zeros(actions_len, 4, 4)}, + "subtask_term_signals": {"grasp": _step_signal(actions_len, term_edge)}, + }, + ) + + # --------------------------------------------------------------------------------------------------- # __init__: subtask signal-name / offset-range tables built from the task descriptor # --------------------------------------------------------------------------------------------------- @@ -192,6 +225,16 @@ def test_add_episode_populates_datagen_info_and_passthrough(): assert di.passthrough_action["franka"].shape == (8, 1) +def test_add_episode_accepts_deformable_state_without_rigid_object_pose(): + pool = _pool(_softmimicgen_task()) + pool._add_episode(_softmimicgen_episode(actions_len=8, term_edge=3)) + di = pool.datagen_infos[0] + assert di.object_poses is None + assert set(di.object_nodal_positions) == {"rope"} + assert di.object_nodal_positions["rope"].shape == (8, 12, 3) + assert pool.subtask_boundaries["franka"] == [[(0, 4), (4, 8)]] + + def test_add_episode_missing_datagen_info_error(): pool = _pool(_mimicgen_task()) bad = types.SimpleNamespace(data={"actions": torch.zeros(5, 7), "obs": {}}) @@ -266,3 +309,32 @@ def test_add_episode_skillgen_rejects_start_offset_that_empties_subtask(): pool = _pool(task, uses_start_signals=True) with pytest.raises(AssertionError, match="empty in worst case"): pool._add_episode(_skillgen_episode(actions_len=10, start_edges=(1, 6), term_edges=(4, 8))) + + +def test_add_episode_maps_legacy_single_arm_names_without_mutating_source(): + pool = _pool(_softmimicgen_task()) + episode = _softmimicgen_episode() + info = episode.data["obs"]["datagen_info"] + for key in ("eef_pose", "target_eef_pose"): + info[key]["robot0"] = info[key].pop("franka") + + pool._add_episode(episode) + recorded = pool.datagen_infos[0] + for key in ("eef_pose", "target_eef_pose"): + assert set(info[key]) == {"robot0"} + assert set(getattr(recorded, key)) == {"franka"} + torch.testing.assert_close(getattr(recorded, key)["franka"], info[key]["robot0"]) + + +def test_add_episode_rejects_inconsistent_pose_names(): + episode = _softmimicgen_episode() + info = episode.data["obs"]["datagen_info"] + info["eef_pose"]["robot0"] = info["eef_pose"].pop("franka") + with pytest.raises(AssertionError, match="do not match target_eef_pose"): + _pool(_softmimicgen_task())._add_episode(episode) + + +def test_add_episode_rejects_ambiguous_arm_name_mapping(): + episode = _dexmimicgen_episode() + with pytest.raises(AssertionError, match="multi-arm datasets require matching names"): + _pool(_softmimicgen_task())._add_episode(episode) diff --git a/autodata_tests/core/test_selection_strategy.py b/autodata_tests/core/test_selection_strategy.py index 3cfbaaf..d102269 100644 --- a/autodata_tests/core/test_selection_strategy.py +++ b/autodata_tests/core/test_selection_strategy.py @@ -13,6 +13,7 @@ NearestNeighborObjectStrategy, NearestNeighborRobotDistanceStrategy, RandomStrategy, + RegistrationCostStrategy, make_selection_strategy, ) @@ -34,6 +35,7 @@ def test_registry_contents(): "random", "nearest_neighbor_object", "nearest_neighbor_robot_distance", + "registration_cost", } @@ -109,3 +111,67 @@ def test_nearest_neighbor_robot_distance_transforms_source_eef_into_current_obje ) assert int(index) == 0 + + +def test_registration_cost_picks_lowest_cost(monkeypatch): + import autodata_core.deformable_transforms as deformable_transforms + + monkeypatch.setattr( + deformable_transforms, + "nodal_registration_cost", + lambda source, target, **kwargs: float(torch.linalg.vector_norm(source - target)), + ) + current = torch.zeros(6, 3) + infos = [DatagenInfo(object_nodal_positions={"rope": torch.full((1, 6, 3), value)}) for value in (2.0, 0.0, 1.0)] + index = RegistrationCostStrategy().select_source_demo( + None, + None, + infos, + object_nodal_positions=current, + nn_k=1, + ) + assert int(index) == 1 + + +def test_registration_cost_skips_failed_and_nonfinite_candidates(monkeypatch): + import autodata_core.deformable_transforms as deformable_transforms + + def registration_cost(source, target, **kwargs): + del target, kwargs + source_value = source[0, 0].item() + if source_value == 0.0: + raise ValueError("invalid TPS input") + if source_value == 1.0: + return float("nan") + return 1.0 + + monkeypatch.setattr(deformable_transforms, "nodal_registration_cost", registration_cost) + current = torch.zeros(6, 3) + infos = [DatagenInfo(object_nodal_positions={"rope": torch.full((1, 6, 3), value)}) for value in (0.0, 1.0, 2.0)] + + index = RegistrationCostStrategy().select_source_demo( + None, + None, + infos, + object_nodal_positions=current, + nn_k=3, + ) + + assert index == 2 + + +def test_registration_cost_rejects_all_invalid_candidates(monkeypatch): + import autodata_core.deformable_transforms as deformable_transforms + + monkeypatch.setattr(deformable_transforms, "nodal_registration_cost", lambda source, target, **kwargs: float("inf")) + current = torch.zeros(6, 3) + infos = [DatagenInfo(object_nodal_positions={"rope": torch.zeros(1, 6, 3)})] + + with pytest.raises(AssertionError, match="could not compute a finite TPS cost"): + RegistrationCostStrategy().select_source_demo( + None, + None, + infos, + object_nodal_positions=current, + nn_k=1, + ) diff --git a/autodata_tests/e2e/test_annotation.py b/autodata_tests/e2e/test_annotation.py index f5383d3..e6ff381 100644 --- a/autodata_tests/e2e/test_annotation.py +++ b/autodata_tests/e2e/test_annotation.py @@ -16,9 +16,33 @@ @pytest.mark.with_subprocess -def test_automatic_annotation_from_dataset(tmp_path): +@pytest.mark.parametrize( + ("input_dataset", "task_descriptor", "env_name", "signal_names", "device", "require_all_demos"), + [ + ( + "annotated_dataset_franka_stack_mimicgen.hdf5", + "franka_cube_stack.yaml", + "IsaacContrib-Stack-Cube-Franka-IK-Rel", + {"grasp_1", "stack_1", "grasp_2"}, + "cpu", + True, + ), + ( + "annotated_dataset_franka_rope_softmimicgen.hdf5", + "franka_rope_softmimicgen.yaml", + "Isaac-Rope-Franka-Arena-IK-Rel-v0", + {"grasp"}, + "cuda", + False, + ), + ], + ids=["franka_cube_stack", "franka_rope"], +) +def test_automatic_annotation_from_dataset( + tmp_path, input_dataset, task_descriptor, env_name, signal_names, device, require_all_demos +): """Use the stored task name and export usable termination signals during replay.""" - input_file = Path(TestPaths.test_data_dir) / "annotated_dataset_franka_stack_mimicgen.hdf5" + input_file = Path(TestPaths.test_data_dir) / input_dataset output_file = tmp_path / "annotated.hdf5" run_subprocess([ TestPaths.python_path, @@ -27,9 +51,9 @@ def test_automatic_annotation_from_dataset(tmp_path): "--viz", "none", "--device", - "cpu", + device, "--task_descriptor", - str(Path(TestPaths.tasks_dir) / "franka_cube_stack.yaml"), + str(Path(TestPaths.tasks_dir) / task_descriptor), "--embodiment", str(Path(TestPaths.embodiments_dir) / "franka_ik_rel.yaml"), "--input_file", @@ -39,14 +63,25 @@ def test_automatic_annotation_from_dataset(tmp_path): ]) assert_valid_dataset(str(output_file), min_num_demos=1) with h5py.File(input_file, "r") as source, h5py.File(output_file, "r") as annotated: - assert len(annotated["data"]) == len(source["data"]) - assert json.loads(annotated["data"].attrs["env_args"])["env_name"] == "IsaacContrib-Stack-Cube-Franka-IK-Rel" + if require_all_demos: + assert len(annotated["data"]) == len(source["data"]) + else: + # Deformable replay can diverge across simulator versions. Annotation exports only successful replays. + assert len(annotated["data"]) <= len(source["data"]) + assert json.loads(annotated["data"].attrs["env_args"])["env_name"] == env_name + source_info = next(iter(source["data"].values()))["obs/datagen_info"] for episode in annotated["data"].values(): signals = episode["obs/datagen_info/subtask_term_signals"] - assert set(signals) == {"grasp_1", "stack_1", "grasp_2"} + assert set(signals) == signal_names for signal in signals.values(): assert len(signal) == len(episode["actions"]) assert signal[:].any() + if "object_nodal_position" in source_info: + nodes = episode["obs/datagen_info/object_nodal_position/object"] + assert nodes.shape == ( + len(episode["actions"]), + *source_info["object_nodal_position/object"].shape[1:], + ) @pytest.mark.with_subprocess diff --git a/autodata_tests/e2e/test_newton_rope_workflow.py b/autodata_tests/e2e/test_newton_rope_workflow.py new file mode 100644 index 0000000..a37b044 --- /dev/null +++ b/autodata_tests/e2e/test_newton_rope_workflow.py @@ -0,0 +1,184 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Record, annotate, and generate Newton demonstrations through the public tools.""" + +import h5py +import json +import numpy as np +import os +from pathlib import Path + +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 + + +def configure_newton_recording() -> list[str]: + """Feed reproducible SpaceMouse commands while keeping real physics and success checks.""" + import torch + + import __main__ as recorder + from isaaclab.devices import Se3SpaceMouse + from isaaclab.envs import ManagerBasedRLEnv + from isaaclab.utils.datasets import HDF5DatasetFileHandler + + from autodata_examples.envs import register_environments + + assert recorder.args_cli.teleop_device == "spacemouse" + assert not recorder.args_cli.reset_sim_buffer_each_episode + remaining_args = register_environments() + reader = HDF5DatasetFileHandler() + reader.open(str(Path(TestPaths.test_data_dir) / "annotated_dataset_franka_rope_softmimicgen_newton.hdf5")) + # Only the initial state and robot actions serve as test input. No annotations, + # subsequent states, or success labels are copied into the new recording. + episode = reader.load_episode("demo_0", "cpu") + reader.close() + actions = episode.data["actions"] + original_reset = ManagerBasedRLEnv.reset + + def reset_from_source(env, *args, **kwargs): + original_reset(env, *args, **kwargs) + assert type(env.scene["object"]).__module__.startswith("isaaclab_contrib.deformable") + return env.reset_to(episode.get_initial_state(), env_ids=None, is_relative=True) + + def reset_device(device): + device._test_action_index = 0 + + def advance(device): + index = device._test_action_index + device._test_action_index += 1 + assert index < len(actions) + 30, "Scripted recording did not reach the real task success condition" + if index < len(actions): + return actions[index].clone() + hold = torch.zeros_like(actions[-1]) + hold[-1] = actions[-1, -1] + return hold + + ManagerBasedRLEnv.reset = reset_from_source + Se3SpaceMouse._find_device = lambda self: None + Se3SpaceMouse._run_device = lambda self: None + Se3SpaceMouse.reset = reset_device + Se3SpaceMouse.advance = advance + return remaining_args + + +@pytest.mark.with_subprocess +def test_newton_rope_record_annotate_generate(tmp_path): + """Require newly recorded raw data to become usable SoftMimicGen input.""" + raw_file = tmp_path / "recorded.hdf5" + annotated_file = tmp_path / "annotated.hdf5" + generated_file = tmp_path / "generated.hdf5" + lab_record = ( + Path(TestPaths.repo_root) / "submodules/IsaacLab-Arena/submodules/IsaacLab/scripts/tools/record_demos.py" + ) + process_env = {**os.environ, "HEADLESS": "1"} + env_name = "Isaac-Rope-Franka-Arena-IK-Rel-v0" + camera_args = ["--viz", "kit", "--enable_cameras"] + task_args = [ + "--task_descriptor", + str(Path(TestPaths.tasks_dir) / "franka_rope_softmimicgen.yaml"), + "--embodiment", + str(Path(TestPaths.embodiments_dir) / "franka_ik_rel.yaml"), + ] + + run_subprocess( + [ + TestPaths.python_path, + str(lab_record), + "--external_callback", + "autodata_tests.e2e.test_newton_rope_workflow.configure_newton_recording", + "--task", + env_name, + "--presets", + "newton", + "--teleop_device", + "spacemouse", + "--no-reset_sim_buffer_each_episode", + "--dataset_file", + str(raw_file), + "--num_demos", + "2", + "--step_hz", + "20", + *camera_args, + ], + env=process_env, + ) + assert_valid_dataset(str(raw_file), min_num_demos=2) + with h5py.File(raw_file, "r") as raw: + assert len(raw["data"]) == 2 + for episode in raw["data"].values(): + assert episode.attrs["success"] + assert "datagen_info" not in episode["obs"] + + run_subprocess( + [ + TestPaths.python_path, + str(Path(TestPaths.scripts_dir) / "annotate_demos.py"), + "--auto", + "--presets", + "newton", + "--device", + "cuda", + "--input_file", + str(raw_file), + "--output_file", + str(annotated_file), + *task_args, + *camera_args, + ], + env=process_env, + ) + assert_valid_dataset(str(annotated_file), min_num_demos=2) + with h5py.File(raw_file, "r") as raw, h5py.File(annotated_file, "r") as annotated: + assert len(annotated["data"]) == 2 + assert json.loads(annotated["data"].attrs["env_args"])["env_name"] == env_name + for name, episode in annotated["data"].items(): + np.testing.assert_array_equal(episode["actions"][:], raw[f"data/{name}/actions"][:]) + info = episode["obs/datagen_info"] + signal = info["subtask_term_signals/grasp"][:] + assert len(signal) == len(episode["actions"]) + assert not signal[0] and signal[-1] + assert np.all(np.diff(signal.astype(int)) >= 0) + nodes = info["object_nodal_position/object"][:] + assert nodes.shape == (len(signal), 549, 3) + assert np.isfinite(nodes).all() + assert set(info["eef_pose"]) == {"franka"} + for camera in ("agentview_image_rgb", "robot0_eye_in_hand_image_rgb"): + assert episode[f"obs/{camera}"].shape[:3] == (len(signal), 128, 128) + + run_subprocess( + [ + TestPaths.python_path, + TestPaths.generate_dataset_script, + "--env_name", + env_name, + "--presets", + "newton", + "--alg", + "softmimicgen", + "--input_file", + str(annotated_file), + "--output_file", + str(generated_file), + "--generation_num_trials", + "1", + "--num_envs", + "3", + "--device", + "cuda", + *task_args, + *camera_args, + ], + env=process_env, + ) + assert_valid_dataset(str(generated_file), min_num_demos=1) + with h5py.File(generated_file, "r") as generated: + for episode in generated["data"].values(): + assert episode.attrs["success"] + nodes = episode["initial_state/deformable_object/object/nodal_position"][:] + assert np.isfinite(nodes).all() and np.abs(nodes).max() < 2.0 + assert not (tmp_path / "generated_failed.hdf5").exists() diff --git a/autodata_tests/e2e/test_recording.py b/autodata_tests/e2e/test_recording.py index 3de6757..b995fd5 100644 --- a/autodata_tests/e2e/test_recording.py +++ b/autodata_tests/e2e/test_recording.py @@ -14,7 +14,7 @@ from autodata_tests.utils.subprocess import run_subprocess -def configure_recording_smoke() -> None: +def configure_recording_smoke() -> list[str]: """Replace physical HID I/O and export three real simulation steps for the test.""" import torch @@ -22,9 +22,15 @@ def configure_recording_smoke() -> None: from isaaclab.devices import Se3SpaceMouse from isaaclab.envs import ManagerBasedRLEnv + from autodata_examples.envs import register_environments + # This assertion observes the actual script's parser after Python startup, # before its task lookup. No task or environment resolution is mocked. - assert recorder.args_cli.task == "IsaacContrib-Stack-Cube-Franka-IK-Rel" + assert recorder.args_cli.task in { + "IsaacContrib-Stack-Cube-Franka-IK-Rel", + "Isaac-Rope-Franka-Arena-IK-Rel-v0", + } + remaining_args = register_environments() Se3SpaceMouse._find_device = lambda self: None Se3SpaceMouse._run_device = lambda self: None original_step = ManagerBasedRLEnv.step @@ -35,6 +41,11 @@ def step_and_finish(env, actions): result = original_step(env, actions) steps += 1 if steps == 3: + if recorder.args_cli.task == "Isaac-Rope-Franka-Arena-IK-Rel-v0": + camera_obs = env.obs_buf["camera_obs"] + assert len(camera_obs) == 2 + for image in camera_obs.values(): + assert image.shape[:3] == (1, 128, 128) # Synthetic input does not solve the task. Explicitly export this # short test episode to verify recording and metadata, then stop. env.recorder_manager.record_pre_reset([0], force_export_or_skip=False) @@ -46,10 +57,19 @@ def step_and_finish(env, actions): return result ManagerBasedRLEnv.step = step_and_finish + return remaining_args @pytest.mark.with_subprocess -def test_lab_record_demos_task_startup(tmp_path): +@pytest.mark.parametrize( + ("task_name", "camera_args"), + [ + ("IsaacContrib-Stack-Cube-Franka-IK-Rel", []), + ("Isaac-Rope-Franka-Arena-IK-Rel-v0", ["--enable_cameras"]), + ], + ids=["franka_cube_stack", "franka_rope"], +) +def test_lab_record_demos_task_startup(tmp_path, task_name, camera_args): """The direct recorder uses the registered task and records real actions and states.""" output_file = tmp_path / "recording.hdf5" script = Path(TestPaths.repo_root) / "submodules/IsaacLab-Arena/submodules/IsaacLab/scripts/tools/record_demos.py" @@ -60,7 +80,7 @@ def test_lab_record_demos_task_startup(tmp_path): "--viz", "kit", "--task", - "IsaacContrib-Stack-Cube-Franka-IK-Rel", + task_name, "--teleop_device", "spacemouse", "--dataset_file", @@ -69,12 +89,15 @@ def test_lab_record_demos_task_startup(tmp_path): "1", "--external_callback", "autodata_tests.e2e.test_recording.configure_recording_smoke", + *camera_args, ], env={**os.environ, "HEADLESS": "1"}, ) with h5py.File(output_file, "r") as dataset: - assert json.loads(dataset["data"].attrs["env_args"])["env_name"] == "IsaacContrib-Stack-Cube-Franka-IK-Rel" + assert json.loads(dataset["data"].attrs["env_args"])["env_name"] == task_name assert len(dataset["data"]) == 1 episode = next(iter(dataset["data"].values())) assert episode["actions"].shape == (3, 7) assert "states" in episode + if task_name == "Isaac-Rope-Franka-Arena-IK-Rel-v0": + assert episode["states/deformable_object/object/nodal_position"].shape == (3, 549, 3) diff --git a/autodata_tests/e2e/test_rope_environment.py b/autodata_tests/e2e/test_rope_environment.py new file mode 100644 index 0000000..3cf2598 --- /dev/null +++ b/autodata_tests/e2e/test_rope_environment.py @@ -0,0 +1,358 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Simulate the shared rope-folding task with each supported physics backend.""" + +import os + +import pytest +import warp as wp + +from autodata_tests.utils.constants import TestPaths +from autodata_tests.utils.subprocess import run_subprocess + + +@wp.kernel +def _read_rope_visual_points( + points: wp.fabricarrayarray(dtype=wp.vec3f), + matrices: wp.fabricarray(dtype=wp.mat44d), + offsets: wp.fabricarray(dtype=wp.uint32), + result: wp.array2d(dtype=wp.vec3f), + sim_count: int, +): + i, node = wp.tid() + env_index = int(offsets[i]) // sim_count + result[env_index, node] = wp.transform_point(wp.transpose(wp.mat44f(matrices[i])), points[i][node]) + + +def _assert_kit_rope_visual(env): + """Compare Kit's actual GPU mesh points with the simulated rope in each world.""" + import torch + + import usdrt + from isaaclab.sim.utils.stage import get_current_stage + from pxr import UsdGeom + + env.sim.render() + nodes = env.scene["object"].data.nodal_pos_w.torch + stage = get_current_stage() + for env_index in range(env.num_envs): + root = f"/World/envs/env_{env_index}/Object" + visual = UsdGeom.Mesh(stage.GetPrimAtPath(f"{root}/Xform")) + assert visual.GetPrim().IsActive() + assert len(visual.GetPointsAttr().Get()) == 1056 + assert len(visual.GetFaceVertexCountsAttr().Get()) == 1026 + # CPU attribute reads are not valid once Fabric owns the point data on the GPU. + selection = get_current_stage(fabric=True).SelectPrims( + require_attrs=[ + (usdrt.Sdf.ValueTypeNames.Point3fArray, "points", usdrt.Usd.Access.Read), + (usdrt.Sdf.ValueTypeNames.UInt, "autodata:rope:particleOffset", usdrt.Usd.Access.Read), + (usdrt.Sdf.ValueTypeNames.Matrix4d, "omni:fabric:worldMatrix", usdrt.Usd.Access.Read), + ], + device=env.device, + ) + assert selection.GetCount() == env.num_envs, f"Only {selection.GetCount()}/{env.num_envs} rope skins found" + visual_points = wp.empty((env.num_envs, 1056), dtype=wp.vec3f, device=env.device) + wp.launch( + _read_rope_visual_points, + dim=visual_points.shape, + inputs=[ + wp.fabricarrayarray(data=selection, attrib="points", dtype=wp.vec3f), + wp.fabricarray(data=selection, attrib="omni:fabric:worldMatrix"), + wp.fabricarray(data=selection, attrib="autodata:rope:particleOffset"), + visual_points, + nodes.shape[1], + ], + device=env.device, + ) + remap = env.sim.physics_manager._rope_vis_remap + indices = wp.to_torch(remap.tet_vertex_indices).long() + weights = wp.to_torch(remap.bary_weights) + expected = (nodes[:, indices] * weights[None, :, :, None]).sum(2) + torch.testing.assert_close(wp.to_torch(visual_points), expected, atol=1.0e-5, rtol=0) + + +@pytest.mark.with_subprocess +@pytest.mark.parametrize("presets", ["physx", "newton"]) +@pytest.mark.parametrize("num_envs", [1, 3]) +def test_rope_environment(presets, num_envs, tmp_path): + result_file = tmp_path / "passed" + run_subprocess( + [ + TestPaths.python_path, + "-m", + __name__, + "--presets", + presets, + "--num_envs", + str(num_envs), + "--viz", + "kit", + "--result_file", + str(result_file), + "--enable_cameras", + ], + env={**os.environ, "HEADLESS": "1"}, + ) + assert result_file.exists(), "Simulation exited before completing the checks" + + +@pytest.mark.with_subprocess +def test_rope_gripper_table_contacts_newton(tmp_path): + """Require physical table support under downward commands with open and closed fingers.""" + result_file = tmp_path / "table_contacts_passed" + run_subprocess( + [ + TestPaths.python_path, + "-m", + __name__, + "--presets", + "newton", + "--num_envs", + "1", + "--viz", + "kit", + "--enable_cameras", + "--table_contacts", + "--result_file", + str(result_file), + ], + env={**os.environ, "HEADLESS": "1"}, + ) + assert result_file.exists(), "Simulation exited before completing the table-contact checks" + + +def _assert_newton_gripper_table_contacts(env): + import itertools + import math + import torch + + from isaaclab.utils.math import compute_pose_error, quat_apply, quat_apply_inverse + from newton import ShapeFlags + + assert env.num_envs == 1 + model = env.sim.physics_manager.get_model() + shape_names = list(map(str, model.shape_label)) + finger_ids = [i for i, name in enumerate(shape_names) if "finger/collisions" in name] + table_ids = [i for i, name in enumerate(shape_names) if "/table/Collisions/Cube" in name] + assert len(finger_ids) == 2 and len(table_ids) == 1 + for entry in ("rigid", "soft"): + view = env.sim.physics_manager._solver.view(entry) + table_flags = [ + int(view.shape_flags.numpy()[i]) + for i, name in enumerate(view.shape_label) + if "/table/Collisions/Cube" in str(name) + ] + assert table_flags and all(flags & int(ShapeFlags.COLLIDE_SHAPES) for flags in table_flags) + + shape_poses = wp.to_torch(model.shape_transform) + shape_scales = wp.to_torch(model.shape_scale) + shape_bodies = wp.to_torch(model.shape_body).long() + signs = torch.tensor(list(itertools.product((-1.0, 1.0), repeat=3)), device=env.device) + + def box_corners(shape_ids): + poses = shape_poses[shape_ids] + points = quat_apply( + poses[:, None, 3:].expand(-1, 8, -1).reshape(-1, 4), + (shape_scales[shape_ids, None] * signs).reshape(-1, 3), + ).reshape(-1, 8, 3) + points += poses[:, None, :3] + bodies = wp.to_torch(env.sim.physics_manager.get_state_0().body_q) + for index, body in enumerate(shape_bodies[shape_ids].tolist()): + if body >= 0: + points[index] = quat_apply(bodies[body, 3:].expand(8, -1), points[index]) + bodies[body, :3] + return points + + table_top = box_corners(table_ids)[..., 2].max().item() + robot = env.scene["robot"] + action = torch.zeros((1, 7), device=env.device) + orientation = torch.tensor([[math.sqrt(0.5), -math.sqrt(0.5), 0.0, 0.0]], device=env.device) + with torch.inference_mode(): + action[:, -1] = 1.0 + env.step(action) # Refresh the frame sensor after reset. + for gripper in (1.0, -1.0): + for height, count in ((0.13, 60), (-0.04, 75)): + target = torch.tensor([[0.45, -0.35, height]], device=env.device) + max_penetration = 0.0 + for _ in range(count): + frame = env.scene["ee_frame"].data + pos, quat = frame.target_pos_w.torch[:, 0], frame.target_quat_w.torch[:, 0] + _, rotation = compute_pose_error(pos, quat, target, orientation, rot_error_type="axis_angle") + root_quat = robot.data.root_quat_w.torch + action[:, :3] = (quat_apply_inverse(root_quat, target - pos) / 0.5).clamp(-0.04, 0.04) + action[:, 3:6] = (quat_apply_inverse(root_quat, rotation) / 0.5).clamp(-0.2, 0.2) + action[:, -1] = gripper + env.step(action) + max_penetration = max(max_penetration, table_top - box_corners(finger_ids)[..., 2].min().item()) + assert torch.isfinite(env.scene["object"].data.nodal_pos_w.torch).all() + if height < table_top: + pos = env.scene["ee_frame"].data.target_pos_w.torch[0, 0] + assert torch.linalg.vector_norm(pos[:2] - target[0, :2]) < 0.025, "Missed the test contact area" + assert pos[2] < 0.04, "Gripper never reached the tabletop" + assert max_penetration < 0.001, f"Fingers penetrated the table by {max_penetration:.6f} m" + print(f"ROPE table contact: gripper={gripper}, max_penetration={max_penetration:.6f} m", flush=True) + + +def main() -> None: + """Check spawning, stepping, contacts, observations and indexed resets.""" + import argparse + import tempfile + + wp.config.enable_backward = False + from isaaclab.app import AppLauncher + + parser = argparse.ArgumentParser() + parser.add_argument("--presets", choices=("physx", "newton"), required=True) + parser.add_argument("--num_envs", type=int, default=1) + parser.add_argument("--snapshot", type=str, default=None) + parser.add_argument("--result_file", type=str, default=None) + parser.add_argument("--table_contacts", action="store_true") + AppLauncher.add_app_launcher_args(parser) + parser.add_argument("--enable_cameras", action="store_true") + args = parser.parse_args() + enable_cameras = args.enable_cameras + app = AppLauncher(args).app + + import gymnasium as gym + import torch + + from autodata_examples.envs.isaac_lab_arena.franka_rope.environment import FRANKA_ROPE_ARENA_ENV_ID + from autodata_examples.envs.isaac_lab_arena.registration import register_environment_from_cli + + # Use the same backend-selection callback as the public Lab teleop/replay scripts. + register_environment_from_cli() + cfg = gym.spec(FRANKA_ROPE_ARENA_ENV_ID).kwargs["env_cfg_entry_point"] + cfg.seed = 42 + cfg.recorders.dataset_export_dir_path = tempfile.mkdtemp(prefix="rope_env_check_") + if args.table_contacts: + assert args.presets == "newton" + cfg.recorders = {} + cfg.terminations = {} + env = gym.make(FRANKA_ROPE_ARENA_ENV_ID, cfg=cfg).unwrapped + exit_code = 0 + try: + if args.table_contacts: + env.reset() + _assert_newton_gripper_table_contacts(env) + if args.result_file: + from pathlib import Path + + Path(args.result_file).touch() + return + + def camera_observations(): + # Headless Kit renders on demand in the pinned Lab version. Request + # an actual frame before reading the RTX camera annotators. + if args.presets == "physx": + for visualizer in env.sim.visualizers: + if visualizer.pumps_app_update(): + visualizer.render_rgb_array() + for sensor in env.scene.sensors.values(): + if hasattr(sensor.cfg, "renderer_cfg"): + sensor.update(0.0, force_recompute=True) + return env.observation_manager.compute()["camera_obs"] + + env.reset() + if args.presets == "newton": + _assert_kit_rope_visual(env) + rope = env.scene["object"] + robot = env.scene["robot"] + assert rope.data.nodal_pos_w.torch.shape == (args.num_envs, 549, 3) + expected_module = "isaaclab_physx" if args.presets == "physx" else "isaaclab_contrib.deformable" + assert type(rope).__module__.startswith(expected_module) + print(f"ROPE backend={args.presets} implementation={type(rope)}", flush=True) + initial_nodes = rope.data.nodal_pos_w.torch.clone() + actions = torch.zeros((args.num_envs, env.action_manager.total_action_dim), device=env.device) + actions[:, -1] = 1.0 + with torch.inference_mode(): + for step in range(100): + obs, _, terminated, _, _ = env.step(actions) + nodes = rope.data.nodal_pos_w.torch + assert torch.isfinite(nodes).all(), f"Non-finite rope nodes at step {step}" + assert not terminated.any(), ( + f"Unexpected task termination at step {step}: " + f"{env.termination_manager.get_active_iterable_terms(int(terminated.nonzero()[0]))}" + ) + assert torch.isfinite(obs["policy"]["object_nodal_pos"]).all() + if step % 20 == 0: + local = nodes - env.scene.env_origins[:, None, :] + print(f"ROPE step={step} min={local.amin((0, 1))} max={local.amax((0, 1))}", flush=True) + if args.presets == "newton": + _assert_kit_rope_visual(env) + assert (rope.data.nodal_pos_w.torch - initial_nodes).abs().max() > 1.0e-4 + + # Gravity must settle the rope on the tabletop, without penetrating or exploding. + local = rope.data.nodal_pos_w.torch - env.scene.env_origins[:, None, :] + assert local[..., 2].min() > -0.01 + assert local[..., 2].max() < 0.15 + assert (local.amax(1) - local.amin(1)).amax() < 1.5 + + if enable_cameras: + images = camera_observations() + assert len(images) == 2 + for image in images.values(): + assert image.shape[:3] == (args.num_envs, 128, 128) + assert image.float().std() > 1.0 + if args.snapshot: + from PIL import Image + + snapshot = next(value for name, value in images.items() if "agentview" in name) + Image.fromarray(snapshot[0, ..., :3].cpu().numpy()).save(args.snapshot) + + # Relative IK and gripper actions must drive the Newton/PhysX articulation. + before_joints = robot.data.joint_pos.torch.clone() + before_images = {name: value.clone() for name, value in images.items()} if enable_cameras else {} + actions[:, 2] = 0.005 + actions[:, -1] = -1.0 + for _ in range(12): + obs, *_ = env.step(actions) + after_joints = robot.data.joint_pos.torch + assert torch.isfinite(after_joints).all() + assert (after_joints[:, :7] - before_joints[:, :7]).abs().max() > 0.005 + assert after_joints[:, -2:].mean() < before_joints[:, -2:].mean() - 0.005 + if enable_cameras: + images = camera_observations() + assert any( + (images[name].float() - before.float()).abs().mean() > 0.01 for name, before in before_images.items() + ), "Camera images did not change when the robot moved" + + # Reset one environment without changing the other environments' nodes. + before_reset = rope.data.nodal_pos_w.torch.clone() + env.reset(env_ids=torch.tensor([0], device=env.device)) + after_reset = rope.data.nodal_pos_w.torch + assert (after_reset[0] - before_reset[0]).abs().max() > 0.001 + torch.testing.assert_close(after_reset[1:], before_reset[1:], atol=1.0e-5, rtol=0) + if args.presets == "newton": + _assert_kit_rope_visual(env) + + # Restoring a saved state must preserve nodal coordinates, including env offsets. + state = env.scene.get_state(is_relative=True) + saved_nodes = rope.data.nodal_pos_w.torch.clone() + env.reset() + env.reset_to(state, env_ids=None, is_relative=True) + torch.testing.assert_close(rope.data.nodal_pos_w.torch, saved_nodes, atol=1.0e-5, rtol=0) + if args.presets == "newton": + _assert_kit_rope_visual(env) + actions.zero_() + actions[:, -1] = 1.0 + for _ in range(5): + env.step(actions) + assert torch.isfinite(rope.data.nodal_pos_w.torch).all() + print("ROPE smoke passed", flush=True) + if args.result_file: + from pathlib import Path + + Path(args.result_file).touch() + except BaseException: + import traceback + + traceback.print_exc() + exit_code = 1 + finally: + env.close() + # Kit terminates the process directly; preserve failures across shutdown. + app.close(exit_code=exit_code) + + +if __name__ == "__main__": + main() 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..eef0e0f --- /dev/null +++ b/autodata_tests/e2e/test_softmimicgen_data_generation.py @@ -0,0 +1,86 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import h5py +import numpy as np +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, presets: str = "physx") -> 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") + dataset_suffix = "_newton" if presets == "newton" else "" + + args = [ + TestPaths.python_path, + TestPaths.generate_dataset_script, + "--env_name", + "Isaac-Rope-Franka-Arena-IK-Rel-v0", + "--presets", + presets, + "--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, f"annotated_dataset_franka_rope_softmimicgen{dataset_suffix}.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) + with h5py.File(output_file, "r") as dataset: + for episode in dataset["data"].values(): + assert episode.attrs["success"] + initial_nodes = episode["initial_state/deformable_object/object/nodal_position"][:] + assert initial_nodes.shape == (1, 549, 3) + assert np.isfinite(initial_nodes).all() + assert np.abs(initial_nodes).max() < 2.0 + nodes = episode["states/deformable_object/object/nodal_position"][:] + assert nodes.shape == (len(episode["actions"]), 549, 3) + assert np.isfinite(nodes).all() + # Parallel environment origins must not leak into recorded positions. + assert np.abs(nodes).max() < 2.0 + + +@pytest.mark.with_subprocess +@pytest.mark.parametrize("presets", ["physx", "newton"]) +def test_franka_rope_softmimicgen_data_generation_single_env_cuda(presets): + """SoftMimicGen generation for the Franka rope task on a single env on GPU.""" + _run_franka_rope_softmimicgen(num_envs=1, device="cuda", presets=presets) + + +@pytest.mark.with_subprocess +@pytest.mark.parametrize("presets", ["physx", "newton"]) +def test_franka_rope_softmimicgen_data_generation_multi_env_cuda(presets): + """SoftMimicGen generation for the Franka rope task on multiple parallel envs on GPU.""" + _run_franka_rope_softmimicgen(num_envs=10, device="cuda", presets=presets) + + +if __name__ == "__main__": + for presets in ("physx", "newton"): + test_franka_rope_softmimicgen_data_generation_single_env_cuda(presets) + test_franka_rope_softmimicgen_data_generation_multi_env_cuda(presets) diff --git a/autodata_tests/interfaces/datastream/test_datastream.py b/autodata_tests/interfaces/datastream/test_datastream.py index 8a45db1..69a7644 100644 --- a/autodata_tests/interfaces/datastream/test_datastream.py +++ b/autodata_tests/interfaces/datastream/test_datastream.py @@ -57,9 +57,15 @@ def _env(env_origins: torch.Tensor | None = None) -> MockEnv: root_quat_w=torch.tensor(_IDENTITY_QUAT_XYZW), ) ) + rope = MockAsset( + MockArticulationData( + nodal_pos_w=torch.tensor([[[1.0, 1.0, 0.05], [1.2, 1.0, 0.05]]]), + ) + ) scene = MockScene( assets={"robot": robot}, rigid_objects={"cube_1": cube}, + deformable_objects={"rope": rope}, env_origins=torch.tensor([[0.0, 0.0, 0.0]]) if env_origins is None else env_origins, state=_SCENE_STATE_SENTINEL, ) @@ -232,6 +238,23 @@ def test_get_object_poses_subtracts_env_origin(): assert torch.allclose(pose[0, :3, 3], torch.tensor([0.0, 0.0, 0.05])) +def test_get_object_nodal_positions_is_env_relative(): + task = TaskDescriptor.from_dict(_task_dict()) + embodiment_adapter = _embodiment_adapter() + env = _env(env_origins=torch.tensor([[1.0, 1.0, 0.0]])) + pool = DataGenInfoPool( + task_descriptor=task, + embodiment_adapter=embodiment_adapter, + device=env.device, + uses_start_signals=False, + ) + datastream = Datastream(env=env, task_descriptor=task, embodiment_adapter=embodiment_adapter, source_pool=pool) + nodes = datastream.get_object_nodal_positions()["rope"] + assert nodes.shape == (1, 2, 3) + assert torch.allclose(nodes[0, 0], torch.tensor([0.0, 0.0, 0.05])) + assert torch.allclose(nodes[0, 1], torch.tensor([0.2, 0.0, 0.05])) + + def test_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..6e497f4 --- /dev/null +++ b/autodata_tests/interfaces/env/test_env_loop.py @@ -0,0 +1,161 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for reset settling in the synchronous data-generation environment loop.""" + +from __future__ import annotations + +import asyncio +import contextlib +import torch +from types import SimpleNamespace + +import pytest + +from autodata_interfaces.env.isaaclab_env_interface import env_loop +from autodata_interfaces.env.reset_request import EnvResetRequest +from autodata_interfaces.tasks.generation_policy_spec import GenerationPolicy + + +class _MockRecorderManager: + def __init__(self) -> None: + self.reset_env_ids: list[tuple[int, ...]] = [] + self.post_reset_env_ids: list[tuple[int, ...]] = [] + self.initial_states: list[tuple[tuple[int, ...], dict]] = [] + + def reset(self, env_ids: torch.Tensor) -> None: + self.reset_env_ids.append(tuple(env_ids.tolist())) + + def record_post_reset(self, env_ids: torch.Tensor) -> None: + # Arena ignores this callback outside an actual reset. + self.post_reset_env_ids.append(tuple(env_ids.tolist())) + + def add_to_episodes(self, key: str, value: dict, env_ids: torch.Tensor) -> None: + assert key == "initial_state" + self.initial_states.append((tuple(env_ids.tolist()), value)) + + +class _MockSimulation: + def is_stopped(self) -> bool: + return False + + +class _MockScene: + def __init__(self, num_envs: int) -> None: + self.env_origins = torch.arange(num_envs * 3, dtype=torch.float32).reshape(num_envs, 3) + self.nodes = torch.zeros((num_envs, 5, 3)) + + def get_state(self, is_relative: bool) -> dict: + nodes = self.nodes.clone() + if not is_relative: + nodes += self.env_origins[:, None] + return {"deformable_object": {"object": {"nodal_position": nodes}}} + + +class _MockEnv: + def __init__(self, num_envs: int) -> None: + self.num_envs = num_envs + self.device = "cpu" + self.action_space = SimpleNamespace(shape=(num_envs, 1)) + self.recorder_manager = _MockRecorderManager() + self.sim = _MockSimulation() + self.scene = _MockScene(num_envs) + self.reset_env_ids: list[tuple[int, ...]] = [] + self.actions: list[torch.Tensor] = [] + + def reset(self, env_ids: torch.Tensor) -> None: + self.reset_env_ids.append(tuple(env_ids.tolist())) + + def step(self, actions: torch.Tensor) -> None: + self.actions.append(actions.clone()) + self.scene.nodes[:, :, 2] += 0.01 + + +def _close_event_loop(loop: asyncio.AbstractEventLoop, tasks: asyncio.Future) -> None: + if not tasks.done(): + tasks.cancel() + with contextlib.suppress(asyncio.CancelledError): + loop.run_until_complete(tasks) + loop.close() + asyncio.set_event_loop(None) + + +def test_reset_without_settling_completes_without_environment_steps(): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + env = _MockEnv(num_envs=1) + reset_queue: asyncio.Queue = asyncio.Queue() + action_queue: asyncio.Queue = asyncio.Queue() + completion = loop.create_future() + reset_queue.put_nowait(EnvResetRequest(env_id=0, completion=completion)) + + async def wait_for_reset() -> None: + await completion + + tasks = asyncio.gather(wait_for_reset()) + try: + completed = env_loop( + env=env, + env_reset_queue=reset_queue, + env_action_queue=action_queue, + asyncio_event_loop=loop, + generation_policy_params=GenerationPolicy(reset_settling_steps=0), + stats={"num_success": 0, "num_failures": 0, "num_attempts": 0}, + data_gen_tasks=tasks, + ) + assert completed is False + assert env.reset_env_ids == [(0,)] + assert env.actions == [] + assert env.recorder_manager.reset_env_ids == [] + assert env.recorder_manager.post_reset_env_ids == [] + loop.run_until_complete(asyncio.wait_for(reset_queue.join(), timeout=0.1)) + finally: + _close_event_loop(loop, tasks) + + +@pytest.mark.parametrize("settling_env_id", [0, 1]) +def test_settling_env_uses_zero_actions_while_other_env_continues(settling_env_id): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + env = _MockEnv(num_envs=2) + running_env_id = 1 - settling_env_id + reset_queue: asyncio.Queue = asyncio.Queue() + action_queue: asyncio.Queue = asyncio.Queue() + completion = loop.create_future() + reset_queue.put_nowait(EnvResetRequest(env_id=settling_env_id, completion=completion)) + + async def provide_running_env_actions() -> None: + for value in (3.0, 4.0): + await action_queue.put((running_env_id, torch.tensor([value]))) + await action_queue.join() + + async def wait_for_reset() -> None: + await completion + + tasks = asyncio.gather(provide_running_env_actions(), wait_for_reset()) + try: + completed = env_loop( + env=env, + env_reset_queue=reset_queue, + env_action_queue=action_queue, + asyncio_event_loop=loop, + generation_policy_params=GenerationPolicy(reset_settling_steps=2), + stats={"num_success": 0, "num_failures": 0, "num_attempts": 0}, + data_gen_tasks=tasks, + ) + assert completed is False + assert env.reset_env_ids == [(settling_env_id,)] + assert len(env.actions) == 2 + assert [action[settling_env_id, 0].item() for action in env.actions] == [0.0, 0.0] + assert [action[running_env_id, 0].item() for action in env.actions] == [3.0, 4.0] + assert env.recorder_manager.reset_env_ids == [(settling_env_id,)] + assert env.recorder_manager.post_reset_env_ids == [] + assert len(env.recorder_manager.initial_states) == 1 + ids, state = env.recorder_manager.initial_states[0] + assert ids == (settling_env_id,) + nodes = state["deformable_object"]["object"]["nodal_position"] + torch.testing.assert_close(nodes, env.scene.nodes[[settling_env_id]]) + torch.testing.assert_close(nodes[:, :, 2], torch.full((1, 5), 0.02)) + loop.run_until_complete(asyncio.wait_for(reset_queue.join(), timeout=0.1)) + finally: + _close_event_loop(loop, tasks) diff --git a/autodata_tests/interfaces/env/test_franka_rope_arena.py b/autodata_tests/interfaces/env/test_franka_rope_arena.py new file mode 100644 index 0000000..bd2b56b --- /dev/null +++ b/autodata_tests/interfaces/env/test_franka_rope_arena.py @@ -0,0 +1,348 @@ +# 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 + +import pytest +from isaaclab.sim import SimulationCfg +from isaaclab_arena.assets.registries import EnvironmentRegistry +from isaaclab_arena.tasks.task_termination_cfg import TaskTerminationCfg +from isaaclab_arena.utils.physics_backend import PhysicsBackend + +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, + _configure_franka_rope_env, +) +from autodata_examples.envs.isaac_lab_arena.franka_rope.rope_asset import FrankaRopeAsset, rope_newton_spawn_cfg +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 + + +@pytest.mark.parametrize("presets", [None, "physx", "newton"]) +def test_arena_registration_uses_run_configuration_and_preserves_make_kwargs(monkeypatch, presets): + 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.setitem( + sys.modules, + "isaaclab_arena.environments.arena_env_builder", + SimpleNamespace(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 = env_registration.register_environment_for_run( + env_name=FRANKA_ROPE_ARENA_ENV_ID, + enable_cameras=True, + num_envs=8, + device="cuda:1", + seed=23, + presets=presets, + ) + + 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 builder_cfg.presets == (None if presets is None else PhysicsBackend(presets)) + 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_newton_rope_preserves_simulation_topology(): + from pxr import Usd + + physx_stage = Usd.Stage.Open(FrankaRopeAsset().object_cfg.spawn.usd_path) + newton_stage = Usd.Stage.Open(rope_newton_spawn_cfg().usd_path) + for attribute in ("points", "tetVertexIndices", "omniphysics:restShapePoints", "omniphysics:restTetVtxIndices"): + original = physx_stage.GetPrimAtPath("/World/simulation_mesh").GetAttribute(attribute).Get() + adapted = newton_stage.GetPrimAtPath("/World/simulation_mesh").GetAttribute(attribute).Get() + assert original is not None + assert original == adapted + assert len(newton_stage.GetPrimAtPath("/World/simulation_mesh").GetAttribute("points").Get()) == 549 + assert physx_stage.GetPrimAtPath("/World/collision_mesh").IsActive() + assert not newton_stage.GetPrimAtPath("/World/collision_mesh").IsActive() + assert physx_stage.GetPrimAtPath("/World/Xform").IsActive() + assert newton_stage.GetPrimAtPath("/World/Xform").IsActive() + for attribute in ("points", "faceVertexCounts", "faceVertexIndices", "primvars:st"): + original = physx_stage.GetPrimAtPath("/World/Xform").GetAttribute(attribute).Get() + adapted = newton_stage.GetPrimAtPath("/World/Xform").GetAttribute(attribute).Get() + assert adapted == original + assert sum(prim.GetTypeName() == "TetMesh" for prim in newton_stage.Traverse()) == 1 + + +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(episode_length_s=15.0) + assert task.get_observation_cfg().subtask_terms.grasp is not None + terminations = task.get_termination_cfg() + assert isinstance(terminations, TaskTerminationCfg) + assert terminations.timeout_s == 15.0 + assert len(terminations.success) == 1 + assert terminations.success[0].predicate_sequence == [mdp.rope_ends_close_tracked] + assert terminations.failures["object_dropping"].func is mdp.rope_below_minimum + assert task.get_events_cfg().reset_object_position is not None + + +def test_rope_environment_configures_lab_timing(): + env_cfg = SimpleNamespace(sim=SimulationCfg(), scene=SimpleNamespace(replicate_physics=True)) + + assert _configure_franka_rope_env(env_cfg) is env_cfg + assert env_cfg.decimation * env_cfg.sim.dt == 0.05 + assert env_cfg.sim.render_interval == 2 + assert env_cfg.num_rerenders_on_reset == 1 + assert env_cfg.scene.replicate_physics is False + assert set(env_cfg.teleop_devices.devices) == {"keyboard", "spacemouse"} + + +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 + for camera_name in embodiment.camera_config.camera_names(): + camera = getattr(embodiment.camera_config, camera_name) + assert camera.renderer_cfg.renderer_type == "isaac_rtx" + assert camera.renderer_cfg.global_settings.antialiasing_mode == "DLSS" + + +def test_newton_embodiment_compensates_gravity_without_changing_physx(): + physx = FrankaRopeIKEmbodiment() + newton = FrankaRopeIKEmbodiment() + physx.configure_physics_backend(PhysicsBackend.PHYSX) + newton.configure_physics_backend(PhysicsBackend.NEWTON) + assert physx.scene_config.robot.spawn.rigid_props.disable_gravity is True + assert physx.scene_config.robot.spawn.collision_props is None + assert newton.scene_config.robot.spawn.rigid_props.gravcomp == 1.0 + assert newton.scene_config.robot.spawn.make_uninstanceable is True + assert newton.scene_config.robot.spawn.collision_props.sdf_max_resolution == 64 + assert physx.scene_config.robot.actuators["panda_shoulder"].stiffness == 400.0 + assert newton.scene_config.robot.actuators["panda_shoulder"].stiffness == 800.0 + assert physx.action_config.gripper_action.close_command_expr == {"panda_finger_.*": 0.0} + assert newton.action_config.gripper_action.close_command_expr == {"panda_finger_.*": 0.0} + for name in newton.camera_config.camera_names(): + assert getattr(newton.camera_config, name).renderer_cfg.renderer_type == "newton_warp" + assert getattr(physx.camera_config, name).renderer_cfg.renderer_type == "isaac_rtx" + + +def test_newton_rope_reset_preserves_tetrahedra(monkeypatch): + """Reset poses must not invert or excessively stretch the volume mesh.""" + import itertools + import numpy as np + import torch + + from isaaclab.assets.deformable_object.base_deformable_object import BaseDeformableObject + from isaaclab.utils import math as math_utils + from pxr import Gf, Usd, UsdGeom + + from autodata_examples.envs.isaac_lab_arena.franka_rope.mdp import events + + stage = Usd.Stage.Open(FrankaRopeAsset().object_cfg.spawn.usd_path) + mesh = UsdGeom.TetMesh(stage.GetPrimAtPath("/World/simulation_mesh")) + transform = UsdGeom.XformCache().GetLocalToWorldTransform(mesh.GetPrim()) + nodes = torch.tensor( + np.asarray([transform.Transform(Gf.Vec3d(point)) for point in mesh.GetPointsAttr().Get()]), + dtype=torch.float32, + ) + # Match the task's initial orientation: the rope's long axis is world Y. + nodes = math_utils.quat_apply(torch.tensor([0.0, 0.0, 2**-0.5, 2**-0.5]).expand(len(nodes), -1), nodes) + corners = torch.tensor(list(itertools.product((0.0, 0.1), (0.0, 0.1), (-torch.pi / 6, torch.pi / 6)))) + count = len(corners) + state = torch.zeros(count, len(nodes), 6) + state[..., :3] = nodes + written = {} + rope = SimpleNamespace(device="cpu", data=SimpleNamespace(default_nodal_state_w=SimpleNamespace(torch=state))) + rope.transform_nodal_pos = lambda *args: BaseDeformableObject.transform_nodal_pos(rope, *args) + rope.write_nodal_state_to_sim_index = lambda values, env_ids: written.update(state=values, env_ids=env_ids) + translations = torch.zeros(count, 3) + translations[:, :2] = corners[:, :2] + rotations = math_utils.quat_from_euler_xyz(torch.zeros(count), torch.zeros(count), corners[:, 2]) + samples = iter( + ((torch.zeros(count, 3), torch.tensor([0.0, 0.0, 0.0, 1.0]).repeat(count, 1)), (translations, rotations)) + ) + monkeypatch.setattr(events, "_sample_pose_delta", lambda *_: next(samples)) + env_ids = torch.arange(count) + events.reset_rope_nodal_state(SimpleNamespace(scene={"object": rope}), env_ids, {}, {}, transition_length=0.4) + indices = torch.tensor(np.asarray(mesh.GetTetVertexIndicesAttr().Get()), dtype=torch.long) + rest = nodes[indices[:, 1:]] - nodes[indices[:, :1]] + result = written["state"][..., :3] + deformed = result[:, indices[:, 1:]] - result[:, indices[:, :1]] + gradient = torch.linalg.solve(rest.transpose(-1, -2), deformed.transpose(-1, -2), left=False) + assert torch.linalg.det(gradient).min() > 0.25 + assert torch.linalg.svdvals(gradient).max() < 2.0 + torch.testing.assert_close(written["state"][..., 3:], torch.zeros_like(state[..., 3:])) + torch.testing.assert_close(written["env_ids"], env_ids) + torch.testing.assert_close(state[..., :3], nodes.expand(count, -1, -1)) + + +def test_rope_reset_smoothing_is_newton_only(): + from isaaclab_newton.physics import NewtonCfg + + for physics in (None, NewtonCfg()): + task = FrankaRopeTask() + cfg = SimpleNamespace( + sim=SimulationCfg() if physics is None else SimulationCfg(physics=physics), + scene=SimpleNamespace( + object=FrankaRopeAsset().object_cfg, + table=SimpleNamespace(spawn=SimpleNamespace()), + ), + events=task.get_events_cfg(), + ) + _configure_franka_rope_env(cfg) + params = cfg.events.reset_object_position.params + if physics is None: + assert "transition_length" not in params + else: + assert params["transition_length"] == 0.4 + assert "transition_length" not in FrankaRopeTask().get_events_cfg().reset_object_position.params + + +def test_newton_rope_contact_buffer_fits_full_surface(): + from collections import Counter + from itertools import combinations + + from pxr import Usd, UsdGeom + + from autodata_examples.envs.isaac_lab_arena.franka_rope.newton_cfg import rope_newton_physics_cfg + + stage = Usd.Stage.Open(rope_newton_spawn_cfg().usd_path) + mesh = UsdGeom.TetMesh(stage.GetPrimAtPath("/World/simulation_mesh")) + face_counts = Counter( + tuple(sorted(face)) for tet in mesh.GetTetVertexIndicesAttr().Get() for face in combinations(tet, 3) + ) + surface_faces = [face for face, count in face_counts.items() if count == 1] + surface_edges = {edge for face in surface_faces for edge in combinations(face, 2)} + contact_count = len(mesh.GetPointsAttr().Get()) + len(surface_faces) + len(surface_edges) + soft_solver = next( + entry.solver_cfg for entry in rope_newton_physics_cfg().solver_cfg.entries if entry.name == "soft" + ) + assert soft_solver.rigid_body_particle_contact_buffer_size >= contact_count 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/test_data/annotated_dataset_franka_rope_softmimicgen_newton.hdf5 b/autodata_tests/test_data/annotated_dataset_franka_rope_softmimicgen_newton.hdf5 new file mode 100644 index 0000000..54b4f86 --- /dev/null +++ b/autodata_tests/test_data/annotated_dataset_franka_rope_softmimicgen_newton.hdf5 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c56d7d7856b07a7b87255b26ba4a0334dde8e2ce548b88e1215d61798732907a +size 13597900 diff --git a/autodata_tests/utils/test_ci_cache.py b/autodata_tests/utils/test_ci_cache.py new file mode 100644 index 0000000..fe3a21f --- /dev/null +++ b/autodata_tests/utils/test_ci_cache.py @@ -0,0 +1,97 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CI must mount a writable home cache for Arena's composed robot assets.""" + +import json +import os +import pwd +import shutil +import subprocess +import sys +from pathlib import Path + +import pytest + +from autodata_tests.utils.constants import TestPaths + + +@pytest.fixture +def ci_launcher(tmp_path): + """Exercise the real shell entry points while replacing only Docker and X11.""" + root = tmp_path / "checkout with spaces" + for relative in ("scripts/ci/run_tests.sh", "docker/run_docker.sh"): + target = root / relative + target.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(Path(TestPaths.repo_root) / relative, target) + target.chmod(0o755) + + commands = tmp_path / "bin" + commands.mkdir() + docker = commands / "docker" + docker.write_text( + f"#!{sys.executable}\n" + "import json, os, sys\n" + "args = sys.argv[1:]\n" + "if args[:2] == ['images', '-q']:\n" + " print('test-image')\n" + "elif args[:2] == ['container', 'inspect']:\n" + " print('false')\n" + "elif args[:1] == ['run']:\n" + " with open(os.environ['TEST_DOCKER_LOG'], 'w') as output:\n" + " json.dump(args, output)\n" + "elif args[:2] != ['ps', '-a']:\n" + " raise SystemExit(f'Unexpected Docker command: {args}')\n" + ) + docker.chmod(0o755) + xhost = commands / "xhost" + xhost.write_text("#!/bin/sh\nexit 0\n") + xhost.chmod(0o755) + runner_temp = tmp_path / "runner temp" + runner_temp.mkdir() + log = tmp_path / "docker.json" + environment = os.environ.copy() + environment.pop("AUTODATA_DOCKER_CACHE_DIR", None) + environment.update( + PATH=f"{commands}:{environment['PATH']}", + RUNNER_TEMP=str(runner_temp), + TEST_DOCKER_LOG=str(log), + FORCE_REBUILD="false", + AUTODATA_RESULTS_DIR="", + ) + + def run(**variables): + result = subprocess.run( + ["bash", str(root / "scripts/ci/run_tests.sh")], + env=environment | variables, + capture_output=True, + text=True, + timeout=30, + ) + assert result.returncode == 0, result.stdout + result.stderr + args = json.loads(log.read_text()) + target = f"/home/{pwd.getpwuid(os.getuid()).pw_name}/.cache" + sources = [args[i + 1].split(":")[0] for i, arg in enumerate(args[:-1]) if arg == "-v"] + mounts = [source for source in sources if f"{source}:{target}" in args] + assert len(mounts) == 1 + source = Path(mounts[0]) + assert source.is_dir() + assert source.stat().st_uid == os.getuid() + (source / "arena-cache-write-check").write_text("writable") + assert "/isaac-sim/python.sh" in args and "pytest" in args + return source + + return run, runner_temp + + +def test_ci_cache_is_fresh_and_owned_by_runner(ci_launcher): + run, runner_temp = ci_launcher + first, second = run(), run() + assert first.parent == second.parent == runner_temp + assert first != second + + +def test_ci_cache_respects_explicit_mount(ci_launcher, tmp_path): + run, _ = ci_launcher + selected = tmp_path / "selected cache" + assert run(AUTODATA_DOCKER_CACHE_DIR=str(selected)) == selected 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/docker/run_docker.sh b/docker/run_docker.sh index f623012..3946bf4 100755 --- a/docker/run_docker.sh +++ b/docker/run_docker.sh @@ -43,6 +43,8 @@ ISAACLAB_PATH="${WORKDIR}/submodules/IsaacLab-Arena/submodules/IsaacLab" # Optional host datasets directory mounted at /datasets (in addition to the repo's own datasets/). DATASETS_HOST_MOUNT_DIRECTORY="$HOME/datasets" +# CI can provide a writable cache without changing the shared host cache's ownership. +CACHE_HOST_MOUNT_DIRECTORY="${AUTODATA_DOCKER_CACHE_DIR:-$HOME/.cache}" FORCE_REBUILD=false NO_CACHE="" @@ -133,6 +135,9 @@ add_volume_if_it_exists() { [ -d "$1" ] && echo "-v $1:$2" } +# Create it as the host user before Docker can create a root-owned bind source. +mkdir -p "${CACHE_HOST_MOUNT_DIRECTORY}" + # Forward the host SSH agent if available (so git over SSH works in-container). SSH_DOCKER_ARGS=() if [ -n "${SSH_AUTH_SOCK:-}" ] && [ -S "$SSH_AUTH_SOCK" ]; then @@ -151,8 +156,8 @@ DOCKER_RUN_ARGS=( # Live-mount the repo: host edits are reflected in the container's editable installs. "-v" "${REPO_ROOT}:${WORKDIR}" $(add_volume_if_it_exists "$DATASETS_HOST_MOUNT_DIRECTORY" /datasets) - # Share the host Kit/pip cache to speed up shader warmup and reinstalls across runs. - "-v" "$HOME/.cache:/home/$(id -un)/.cache" + # Share the selected cache, including Arena's hard-coded ~/.cache location. + "-v" "${CACHE_HOST_MOUNT_DIRECTORY}:/home/$(id -un)/.cache" # X11 passthrough so "--viz kit" can open a window. "-v" "/tmp/.X11-unix:/tmp/.X11-unix:rw" "${SSH_DOCKER_ARGS[@]}" diff --git a/scripts/annotate_demos.py b/scripts/annotate_demos.py index a334cc4..6685a75 100644 --- a/scripts/annotate_demos.py +++ b/scripts/annotate_demos.py @@ -43,6 +43,15 @@ default=None, help="Environment name. Overrides the env name recorded in the source dataset.", ) +parser.add_argument( + "--presets", choices=("physx", "newton"), default=None, help="Physics backend preset for Arena environments." +) +parser.add_argument( + "--reset_sim_buffer_each_episode", + action=argparse.BooleanOptionalAction, + default=None, + help="Flush simulation buffers before replay. Defaults to enabled for PhysX and disabled for Newton.", +) parser.add_argument( "--task_descriptor", type=str, @@ -97,7 +106,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 sys # noqa: E402 import torch # noqa: E402 import traceback # noqa: E402 @@ -111,9 +122,11 @@ 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 +from autodata_interfaces.env.recorders import PreStepCameraObservationsRecorder # noqa: E402 from autodata_interfaces.tasks.task_descriptor import TaskDescriptor # noqa: E402 is_paused = False @@ -153,6 +166,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), } @@ -228,6 +242,19 @@ 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.record_images, + num_envs=1, + device=args_cli.device, + seed=generation_policy.seed, + presets=args_cli.presets, + ) + # Start signals are required only by SkillGen. annotate_start_signals = generation_policy.use_skillgen @@ -260,10 +287,20 @@ def main() -> int: recorder_cfg=AnnotationRecorderManagerCfg(), enable_cameras=args_cli.record_images, ) + if args_cli.reset_sim_buffer_each_episode is None: + from isaaclab_newton.physics import NewtonCfg + + # Preserve the established PhysX replay behavior. Newton keeps solver buffers + # across episode restores, as in Lab's public replay tool. + args_cli.reset_sim_buffer_each_episode = not isinstance(env_cfg.sim.physics, NewtonCfg) + if args_cli.record_images: + env_cfg.recorders.record_pre_step_camera_observations = RecorderTermCfg( + class_type=PreStepCameraObservationsRecorder + ) # 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) @@ -365,8 +402,9 @@ def replay_episode( initial_state = episode.data["initial_state"] actions = episode.data["actions"] - # Reset the simulation before restoring state (use of env.sim.reset() for deterministic replay) - env.sim.reset() + # A full solver reset is distinct from restoring an episode's scene state. + if args_cli.reset_sim_buffer_each_episode: + env.sim.reset() env.recorder_manager.reset() env.reset_to(initial_state, torch.tensor([0], device=env.device), is_relative=True) diff --git a/scripts/ci/run_tests.sh b/scripts/ci/run_tests.sh index f9701ae..2b91fdc 100755 --- a/scripts/ci/run_tests.sh +++ b/scripts/ci/run_tests.sh @@ -30,13 +30,14 @@ FORCE_REBUILD="${FORCE_REBUILD:-false}" # run_docker.sh gives the container a fresh environment, so this is forwarded # explicitly on the in-container command line rather than exported on the host. SUBPROCESS_TIMEOUT="${AUTODATA_SUBPROCESS_TIMEOUT:-1200}" -# Container-local cache directory. run_docker.sh bind-mounts the host's -# $HOME/.cache into the container; on the CI runner that path is root-owned (or -# auto-created as root), so the recreated non-root container user cannot write -# it -- warp fails to create ~/.cache/warp with a PermissionError. Pointing the -# cache env vars at a writable, container-local /tmp path sidesteps the mounted -# host cache entirely. The tools create these dirs themselves (makedirs). +# Writable container-local cache for Warp and tools that honor XDG_CACHE_HOME. +# Arena's hard-coded home cache is handled separately below. CONTAINER_CACHE_DIR="${CONTAINER_CACHE_DIR:-/tmp/autodata_ci_cache}" +# Arena's composed USD cache ignores XDG_CACHE_HOME and writes under ~/.cache. +# Mount a fresh directory owned by the runner instead of its shared host cache. +# RUNNER_TEMP is cleaned by Actions; local invocations fall back to TMPDIR or /tmp. +AUTODATA_DOCKER_CACHE_DIR="${AUTODATA_DOCKER_CACHE_DIR:-$(mktemp -d "${RUNNER_TEMP:-${TMPDIR:-/tmp}}/autodata-ci-cache.XXXXXX")}" +export AUTODATA_DOCKER_CACHE_DIR # Optional JUnit report dir (repo-relative). The repo is bind-mounted, so a report # written here lands on the host for CI to collect. RESULTS_DIR="${AUTODATA_RESULTS_DIR-}" @@ -56,8 +57,8 @@ fi # container (run_docker.sh passes trailing args through verbatim, and the values # are inherited by the data-generation subprocess pytest spawns). Call # /isaac-sim/python.sh -m pytest explicitly so it does not depend on the -# in-container `pytest` alias expanding. XDG_CACHE_HOME/WARP_CACHE_PATH steer all -# caches away from the mounted host $HOME/.cache (see CONTAINER_CACHE_DIR above). +# in-container `pytest` alias expanding. XDG_CACHE_HOME/WARP_CACHE_PATH keep those +# caches container-local; Arena uses the writable bind mount configured above. PYTEST_ARGS=( env "AUTODATA_SUBPROCESS_TIMEOUT=${SUBPROCESS_TIMEOUT}" diff --git a/scripts/generate_dataset.py b/scripts/generate_dataset.py index 50f96bc..d6a8506 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( @@ -76,6 +77,9 @@ ) parser.add_argument("--generation_num_trials", type=int, default=None, help="Number of demos to generate.") parser.add_argument("--num_envs", type=int, default=1, help="Number of parallel environments.") +parser.add_argument( + "--presets", choices=("physx", "newton"), default=None, help="Physics backend preset for Arena environments." +) parser.add_argument("--input_file", type=str, required=True, help="Source dataset HDF5 file.") parser.add_argument( "--output_file", @@ -125,6 +129,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 @@ -292,6 +297,19 @@ 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.record_images, + num_envs=args_cli.num_envs, + device=args_cli.device, + seed=generation_policy_params.seed, + presets=args_cli.presets, + ) + # 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). @@ -313,11 +331,7 @@ def main() -> None: enable_cameras=args_cli.record_images, ) - 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()