diff --git a/autodata_core/algorithms.py b/autodata_core/algorithms.py index 3f6b8b6..65d6daf 100644 --- a/autodata_core/algorithms.py +++ b/autodata_core/algorithms.py @@ -156,10 +156,10 @@ class MimicGen(GenerationAlgorithm): class DexMimicGen(GenerationAlgorithm): - """Two-arm MimicGen with subtask coordination constraints.""" + """Two-arm MimicGen with subtask coordination and an optional navigation channel.""" name = "dexmimicgen" - expected_eef_count = 2 + expected_eef_count = (2, 3) requires_motion_planner = False uses_subtask_start_signals = False supports_coordination = True diff --git a/autodata_core/data_generator.py b/autodata_core/data_generator.py index dc91e69..59481cb 100644 --- a/autodata_core/data_generator.py +++ b/autodata_core/data_generator.py @@ -540,6 +540,9 @@ async def generate( randomized_subtask_boundaries: dict[str, np.ndarray] | None = None prev_pool_size = 0 + was_navigating = False + navigation_hold_steps = 0 + generation_policy = self.datastream.get_generation_policy() while True: await asyncio.sleep(0) @@ -586,6 +589,15 @@ async def generate( env_action_queue=env_action_queue, ) buffers.success = buffers.success or exec_success + if generation_policy.use_navigation_controller: + hold, was_navigating = self._synchronize_navigation(eef_states, env_id, was_navigating) + navigation_hold_steps = navigation_hold_steps + 1 if hold else 0 + if hold: + if navigation_hold_steps >= generation_policy.max_navigation_hold_steps: + print(f"Navigation hold limit reached for env {env_id} after {navigation_hold_steps} steps.") + buffers.success = False + break + continue self._advance_subtask_progress( eef_states=eef_states, runtime_constraints=runtime_constraints, @@ -614,6 +626,33 @@ async def generate( # Generation loop helpers # ------------------------------------------------------------------ + def _synchronize_navigation( + self, eef_states: dict[str, _EEFGenerationState], env_id: int, was_navigating: bool + ) -> tuple[bool, bool]: + """Hold playback for an unfinished navigation goal or skip its remaining recorded idle steps.""" + navigation = self.datastream.get_navigation_state(env_id) + assert navigation is not None, "Navigation generation requires an embodiment navigation configuration." + channel, is_navigating, goal_reached = navigation + state = eef_states[channel] + assert state.subtask_step_index is not None + if state.subtasks_done or state.is_currently_paused: + return False, is_navigating + if is_navigating and not goal_reached and state.subtask_step_index + 1 >= len(state.current_trajectory) - 1: + # Keep all channels synchronized while the controller finishes a segment. + return True, is_navigating + if was_navigating and not is_navigating and goal_reached: + steps_to_skip = max(0, len(state.current_trajectory) - state.subtask_step_index - 2) + for other_state in eef_states.values(): + if ( + other_state.subtask_step_index is not None + and not other_state.is_currently_paused + and not other_state.subtasks_done + ): + index = other_state.subtask_step_index + last_index = max(index, len(other_state.current_trajectory) - 2) + other_state.subtask_step_index = min(index + steps_to_skip, last_index) + return False, is_navigating + async def _reset_environment_for_generation( self, env_id: int, diff --git a/autodata_examples/embodiments/g1_wbc_pink.yaml b/autodata_examples/embodiments/g1_wbc_pink.yaml new file mode 100644 index 0000000..3ed539f --- /dev/null +++ b/autodata_examples/embodiments/g1_wbc_pink.yaml @@ -0,0 +1,25 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +type: absolute_pose +name: g1_wbc_pink +description: Bimanual wrist pose control in the pelvis frame, scalar grippers, and locomotion commands. +action_dim: 23 +reference_frame: {asset_name: robot, body_name: pelvis} +eefs: + left: + pose_obs_keys: {pos: left_eef_pos, quat: left_eef_quat} + pose_action_slice: [2, 9] + passthrough_action_slice: [0, 1] + right: + pose_obs_keys: {pos: right_eef_pos, quat: right_eef_quat} + pose_action_slice: [9, 16] + passthrough_action_slice: [1, 2] + body: + pose_obs_keys: {pos: body_eef_pos, quat: body_eef_quat} + passthrough_action_slice: [16, 23] +navigation: + eef_name: body + observation_group: wbc + active_key: is_navigating + goal_reached_key: navigation_goal_reached diff --git a/autodata_examples/env_profiles/g1_locomanip.yaml b/autodata_examples/env_profiles/g1_locomanip.yaml new file mode 100644 index 0000000..51dbf2e --- /dev/null +++ b/autodata_examples/env_profiles/g1_locomanip.yaml @@ -0,0 +1,14 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +name: g1_locomanip_navigation +base_env: galileo_g1_locomanip_pick_and_place +description: Closed-loop navigation for the recorded table-to-bin manipulation route. +actions: + g1_action: + use_p_control: true + navigation_subgoals: + - [[0.18, 0.18, 0.0], false] + - [[0.18, 0.18, -1.78], true] + - [[-0.0955, -1.1070, -1.78], false] + - [[-0.0955, -1.1070, -1.78], false] diff --git a/autodata_examples/tasks/g1_locomanip.yaml b/autodata_examples/tasks/g1_locomanip.yaml new file mode 100644 index 0000000..872c170 --- /dev/null +++ b/autodata_examples/tasks/g1_locomanip.yaml @@ -0,0 +1,44 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +name: g1_locomanip +description: Navigate to the box, grasp with both hands, carry it to the bin, and release it. +algo: dexmimicgen +generation_policy: + name: g1_locomanip + num_trials: 100 + guarantee_success: false + use_navigation_controller: true +subtasks: + right: + - &manipulation + object_ref: brown_box + subtask_term_signal: idle_right + selection_strategy: nearest_neighbor_object + selection_strategy_kwargs: {nn_k: 3} + action_noise: 0.003 + num_interpolation_steps: 0 + num_fixed_steps: 0 + subtask_term_offset_range: [0, 0] + - <<: *manipulation + subtask_term_signal: grasp_and_idle_right + - <<: *manipulation + subtask_term_signal: null + left: + - <<: *manipulation + subtask_term_signal: idle_left + - <<: *manipulation + subtask_term_signal: grasp_and_idle_left + - <<: *manipulation + subtask_term_signal: null + body: + - &navigation + <<: *manipulation + subtask_term_signal: navigate_to_table + action_noise: 0.0 + - <<: *navigation + subtask_term_signal: navigate_turn_inplace + - <<: *navigation + subtask_term_signal: navigate_to_bin + - <<: *navigation + subtask_term_signal: null diff --git a/autodata_examples/tasks/gr1_open_microwave.yaml b/autodata_examples/tasks/gr1_open_microwave.yaml new file mode 100644 index 0000000..153d177 --- /dev/null +++ b/autodata_examples/tasks/gr1_open_microwave.yaml @@ -0,0 +1,34 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +name: gr1_open_microwave +description: Reach the microwave handle with the right hand and open its door. +algo: dexmimicgen +generation_policy: + name: gr1_open_microwave + num_trials: 50 + guarantee_success: false + select_src_per_subtask: true +subtasks: + right: + - &door + object_ref: microwave + description: Reach the door handle. + subtask_term_signal: move_to_door + subtask_term_offset_range: [10, 20] + selection_strategy: nearest_neighbor_object + selection_strategy_kwargs: {nn_k: 3} + action_noise: 0.005 + num_interpolation_steps: 5 + num_fixed_steps: 0 + apply_noise_during_interpolation: false + - <<: *door + description: Open the door. + subtask_term_signal: null + subtask_term_offset_range: [0, 0] + left: + - <<: *door + description: Hold the supporting arm steady. + subtask_term_signal: null + subtask_term_offset_range: [0, 0] + num_interpolation_steps: 0 diff --git a/autodata_examples/tasks/gr1_put_and_close_door.yaml b/autodata_examples/tasks/gr1_put_and_close_door.yaml new file mode 100644 index 0000000..cb0b98c --- /dev/null +++ b/autodata_examples/tasks/gr1_put_and_close_door.yaml @@ -0,0 +1,46 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +name: gr1_put_and_close_door +description: Place the ranch bottle inside the refrigerator and close its door. +algo: dexmimicgen +generation_policy: + name: gr1_put_and_close_door + num_trials: 100 + guarantee_success: false + select_src_per_subtask: true +subtasks: + right: + - &grasp + object_ref: ranch_dressing_hope_robolab + description: Grasp the bottle. + subtask_term_signal: subtask_0_right_grasp_1 + subtask_term_offset_range: [0, 0] + selection_strategy: nearest_neighbor_object + selection_strategy_kwargs: {nn_k: 3} + action_noise: 0.003 + num_interpolation_steps: 5 + num_fixed_steps: 0 + apply_noise_during_interpolation: false + - <<: *grasp + object_ref: refrigerator + description: Place the bottle on the refrigerator shelf. + subtask_term_signal: subtask_0_right_last_mimic_subtask + - <<: *grasp + object_ref: refrigerator + description: Reach the refrigerator door. + subtask_term_signal: subtask_1_right_move_to_door + - <<: *grasp + object_ref: refrigerator + description: Close the door. + subtask_term_signal: null + left: + - <<: *grasp + description: Support the bottle placement. + subtask_term_signal: subtask_0_left_last_mimic_subtask + num_interpolation_steps: 0 + - <<: *grasp + object_ref: refrigerator + description: Hold the supporting arm steady while the door closes. + subtask_term_signal: null + num_interpolation_steps: 0 diff --git a/autodata_interfaces/datastream/datastream.py b/autodata_interfaces/datastream/datastream.py index 1b21c05..9733fa7 100644 --- a/autodata_interfaces/datastream/datastream.py +++ b/autodata_interfaces/datastream/datastream.py @@ -80,13 +80,7 @@ def __init__( "the embodiment adapter env unset before creating the Datastream." ) - # Ensure that task descriptor and embodiment adapter declare the same EEFs. - task_eefs = set(task_descriptor.get_eef_names()) - adapter_eefs = set(embodiment_adapter.get_eef_names()) - assert task_eefs == adapter_eefs, ( - f"EEF mismatch: task descriptor declares {sorted(task_eefs)}, " - f"embodiment adapter exposes {sorted(adapter_eefs)}" - ) + task_descriptor.validate_embodiment(embodiment_adapter) # Initialize the source-demo pool. assert (source_pool is None) != ( @@ -242,17 +236,39 @@ def actions_to_passthrough_actions(self, actions: torch.Tensor) -> dict[str, tor # Other runtime queries # ------------------------------------------------------------------ + def get_navigation_state(self, env_id: int) -> tuple[str, bool, bool] | None: + """Return the navigation channel and its controller's active and goal-reached flags. + + Args: + env_id: Environment index to query. + + Returns: + Channel name, active flag, and goal-reached flag, or None when unavailable. + """ + return self.embodiment_adapter.get_navigation_state(env_id) + + def _reference_frame_inverse(self, env_ids: Sequence[int] | None) -> torch.Tensor | None: + reference = self.embodiment_adapter.get_reference_frame_pose(env_ids) + return None if reference is None else torch.linalg.inv(reference) + def get_object_poses(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor]: - """Get all rigid object poses from the environment.""" + """Get non-robot rigid and articulated object poses in the embodiment's reference frame.""" 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_pose_matrix: dict[str, torch.Tensor] = {} - for obj_name, obj in scene.rigid_objects.items(): + assets = {**scene.rigid_objects, **getattr(scene, "articulations", {})} + reference_inverse = self._reference_frame_inverse(env_ids) + robot_name = getattr(self.embodiment_adapter, "robot_asset_name", "robot") + for obj_name, obj in assets.items(): + if obj_name == robot_name: + continue pos_rel = as_torch(obj.data.root_pos_w)[index] - env_origins quat = as_torch(obj.data.root_quat_w)[index] object_pose_matrix[obj_name] = pose_math.make_pose(pos_rel, pose_math.matrix_from_quat(quat)) + if reference_inverse is not None: + object_pose_matrix[obj_name] = reference_inverse @ object_pose_matrix[obj_name] return object_pose_matrix def get_subtask_term_signals( @@ -298,12 +314,12 @@ def get_subtask_term_signals( return signals def get_robot_root_pose(self, env_ids: Sequence[int] | None = None) -> torch.Tensor: - """Get the robot articulation root pose in the env-relative frame. + """Get the robot articulation root pose in the embodiment's reference frame. Returns a tensor of shape ``(len(env_ids), 4, 4)`` using the same frame convention as - :meth:`get_object_poses` (per-env origin subtracted). Motion planners read this to - reconcile the env-relative frame they exchange with callers and the robot base frame - their solver plans in. + :meth:`get_object_poses`. Without a configured controller frame, both queries return + environment-relative poses. Motion planners use this pose to transform between the + embodiment frame and the robot base frame their solver plans in. """ index: slice | Sequence[int] = slice(None) if env_ids is None else env_ids @@ -312,7 +328,9 @@ def get_robot_root_pose(self, env_ids: Sequence[int] | None = None) -> torch.Ten robot = scene[robot_name] pos_rel = as_torch(robot.data.root_pos_w)[index] - scene.env_origins[index] quat = as_torch(robot.data.root_quat_w)[index] - return pose_math.make_pose(pos_rel, pose_math.matrix_from_quat(quat)) + pose = pose_math.make_pose(pos_rel, pose_math.matrix_from_quat(quat)) + reference_inverse = self._reference_frame_inverse(env_ids) + return reference_inverse @ pose if reference_inverse is not None else pose def get_scene_state(self, is_relative: bool = True) -> dict: """Return the raw scene-state snapshot from the underlying env. @@ -330,7 +348,7 @@ def get_scene_state(self, is_relative: bool = True) -> dict: # Motion planners build their collision world from the live scene. The Datastream is the # single owner of that simulator-facing access so the planners stay backend-agnostic: a # cuRobo planner feeds the stage + prim scoping below into its own USD parser, and per-step - # obstacle pose sync goes through :meth:`get_object_poses` (env-relative frame). + # obstacle pose sync goes through :meth:`get_object_poses` (embodiment reference frame). def get_usd_stage(self) -> Any: """Return the live USD stage backing the scene. diff --git a/autodata_interfaces/embodiments/__init__.py b/autodata_interfaces/embodiments/__init__.py index 6dc5bef..39cddd4 100644 --- a/autodata_interfaces/embodiments/__init__.py +++ b/autodata_interfaces/embodiments/__init__.py @@ -10,6 +10,7 @@ IK control). YAML construction goes through :mod:`factory`. """ +from autodata_interfaces.embodiments.absolute_pose_embodiment_adapter import AbsolutePoseEmbodimentAdapter from autodata_interfaces.embodiments.bimanual_embodiment_adapter import ( AbsolutePoseWholeBodyBimanualAdapter, BimanualEefConfig, @@ -27,6 +28,7 @@ ) __all__ = [ + "AbsolutePoseEmbodimentAdapter", "AbsolutePoseWholeBodyBimanualAdapter", "BimanualEefConfig", "BimanualEmbodimentAdapter", diff --git a/autodata_interfaces/embodiments/absolute_pose_embodiment_adapter.py b/autodata_interfaces/embodiments/absolute_pose_embodiment_adapter.py new file mode 100644 index 0000000..c1d6c09 --- /dev/null +++ b/autodata_interfaces/embodiments/absolute_pose_embodiment_adapter.py @@ -0,0 +1,167 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Configurable absolute-pose control with independently laid out action channels.""" + +from __future__ import annotations + +import torch +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +from autodata_interfaces.embodiments.embodiment_adapter import EmbodimentAdapter +from autodata_utils import pose_math +from autodata_utils.tensor_utils import as_torch + + +@dataclass(kw_only=True) +class AbsolutePoseEmbodimentAdapter(EmbodimentAdapter): + """Encode arbitrary absolute-pose and passthrough channels from a descriptor. + + Each entry in ``eefs`` declares ``pose_obs_keys``, an optional seven-component + ``pose_action_slice`` (position [m] and quaternion xyzw), and a ``passthrough_action_slice``. + Tracking-only channels omit the pose slice; their actions contain only passthrough commands. + + Args: + name: Embodiment name. + eefs: Per-channel observation keys and action slices. + action_dim: Width of the complete action vector. + description: Human-readable description. + obs_group: Observation group containing the tracked poses. + reference_frame: Optional controller frame, specified by asset name and optional body name. + navigation: Optional navigation channel and observation keys used to synchronize playback. + canonicalize_quat: Whether encoded quaternions have a nonnegative scalar component. + """ + + name: str + eefs: dict[str, dict[str, Any]] + action_dim: int + description: str = "" + obs_group: str = "policy" + reference_frame: dict[str, str] | None = None + navigation: dict[str, str] | None = None + canonicalize_quat: bool = True + + def __post_init__(self) -> None: + assert self.name and self.eefs and self.action_dim > 0 + occupied: set[int] = set() + for name, spec in self.eefs.items(): + assert set(spec) <= { + "pose_obs_keys", + "pose_action_slice", + "passthrough_action_slice", + }, f"Unknown channel parameters for {name!r}: {set(spec)}." + assert set(spec["pose_obs_keys"]) == {"pos", "quat"} + for key in ("pose_action_slice", "passthrough_action_slice"): + indices = spec.get(key) + if indices is None: + continue + assert len(indices) == 2 and 0 <= indices[0] < indices[1] <= self.action_dim + if key == "pose_action_slice": + assert indices[1] - indices[0] == 7 + positions = set(range(*indices)) + assert not positions & occupied, f"Overlapping action slices for {name!r}." + occupied.update(positions) + assert "passthrough_action_slice" in spec, f"Missing passthrough action slice for {name!r}." + assert occupied == set(range(self.action_dim)), "Every action component must belong to exactly one channel." + if self.reference_frame is not None: + assert "asset_name" in self.reference_frame + assert set(self.reference_frame) <= {"asset_name", "body_name"} + if self.navigation is not None: + assert set(self.navigation) == {"eef_name", "observation_group", "active_key", "goal_reached_key"} + assert self.navigation["eef_name"] in self.eefs + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> AbsolutePoseEmbodimentAdapter: + """Build the adapter from a validated descriptor.""" + return cls(**data) + + def get_eef_names(self) -> tuple[str, ...]: + return tuple(self.eefs) + + def get_eef_poses(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor]: + assert self.env is not None, "Call bind_env(env) before reading state." + index = slice(None) if env_ids is None else env_ids + obs = self.env.obs_buf[self.obs_group] + return { + name: pose_math.make_pose( + as_torch(obs[spec["pose_obs_keys"]["pos"]])[index], + pose_math.matrix_from_quat(as_torch(obs[spec["pose_obs_keys"]["quat"]])[index]), + ) + for name, spec in self.eefs.items() + } + + def get_reference_frame_pose(self, env_ids: Sequence[int] | None = None) -> torch.Tensor | None: + if self.reference_frame is None: + return None + assert self.env is not None, "Call bind_env(env) before reading state." + index = slice(None) if env_ids is None else env_ids + asset = self.env.scene[self.reference_frame["asset_name"]] + body_name = self.reference_frame.get("body_name") + if body_name is None: + position = as_torch(asset.data.root_pos_w)[index] + quat = as_torch(asset.data.root_quat_w)[index] + else: + body_index = asset.data.body_names.index(body_name) + state = as_torch(asset.data.body_link_state_w)[index, body_index] + position, quat = state[..., :3], state[..., 3:7] + return pose_math.make_pose(position - self.env.scene.env_origins[index], pose_math.matrix_from_quat(quat)) + + def get_navigation_eef_name(self) -> str | None: + """Return the navigation channel declared by the descriptor, if any.""" + return self.navigation["eef_name"] if self.navigation is not None else None + + def get_navigation_state(self, env_id: int) -> tuple[str, bool, bool] | None: + """Return the navigation channel, active flag, and goal flag from standard observations.""" + if self.navigation is None: + return None + assert self.env is not None, "Call bind_env(env) before reading state." + obs = self.env.obs_buf[self.navigation["observation_group"]] + return ( + self.navigation["eef_name"], + bool(obs[self.navigation["active_key"]][env_id]), + bool(obs[self.navigation["goal_reached_key"]][env_id]), + ) + + def action_to_target_eef_pose(self, action: torch.Tensor) -> dict[str, torch.Tensor]: + assert action.ndim == 2 and action.shape[-1] == self.action_dim + result = {} + for name, spec in self.eefs.items(): + indices = spec.get("pose_action_slice") + if indices is None: + # A tracking-only channel has no commanded target pose. Its pose does not + # participate in encoding; use a valid neutral transform for the source pool. + result[name] = torch.eye(4, dtype=action.dtype, device=action.device).repeat(len(action), 1, 1) + else: + values = action[:, slice(*indices)] + result[name] = pose_math.make_pose(values[:, :3], pose_math.matrix_from_quat(values[:, 3:7])) + return result + + def actions_to_passthrough_actions(self, actions: torch.Tensor) -> dict[str, torch.Tensor]: + assert actions.shape[-1] == self.action_dim + return {name: actions[..., slice(*spec["passthrough_action_slice"])] for name, spec in self.eefs.items()} + + def target_eef_pose_to_action( + self, + target_eef_pose_dict: dict[str, torch.Tensor], + passthrough_action_dict: dict[str, torch.Tensor], + action_noise_dict: dict[str, float] | None = None, + env_id: int = 0, + ) -> torch.Tensor: + assert set(target_eef_pose_dict) == set(self.eefs) + assert set(passthrough_action_dict) == set(self.eefs) + first_action = next(iter(passthrough_action_dict.values())) + action = first_action.new_zeros(self.action_dim) + for name, spec in self.eefs.items(): + action[slice(*spec["passthrough_action_slice"])] = passthrough_action_dict[name].reshape(-1) + indices = spec.get("pose_action_slice") + if indices is not None: + position, rotation = pose_math.unmake_pose(target_eef_pose_dict[name]) + quat = pose_math.quat_from_matrix(rotation) + if self.canonicalize_quat: + quat = pose_math.quat_unique(quat) + values = torch.cat((position, quat)) + noise = (action_noise_dict or {}).get(name, 0.0) + action[slice(*indices)] = values + noise * torch.randn_like(values) if noise else values + return action diff --git a/autodata_interfaces/embodiments/embodiment_adapter.py b/autodata_interfaces/embodiments/embodiment_adapter.py index 6082c58..4f63034 100644 --- a/autodata_interfaces/embodiments/embodiment_adapter.py +++ b/autodata_interfaces/embodiments/embodiment_adapter.py @@ -38,6 +38,19 @@ class EmbodimentAdapter(ABC): env: Any = None robot_asset_name: str = "robot" + def get_navigation_eef_name(self) -> str | None: + """Return the configured navigation channel without requiring a live environment.""" + + def get_navigation_state(self, env_id: int) -> tuple[str, bool, bool] | None: + """Return a navigation channel and its active/goal flags, or None when unavailable.""" + + def get_reference_frame_pose(self, env_ids: Sequence[int] | None = None) -> torch.Tensor | None: + """Return the controller frame in environment coordinates, or None for environment coordinates. + + Args: + env_ids: Environment indices to query. None selects all environments. + """ + def bind_env(self, env: Any) -> None: """Attach the env after construction. diff --git a/autodata_interfaces/embodiments/factory.py b/autodata_interfaces/embodiments/factory.py index 5baa081..f21aea5 100644 --- a/autodata_interfaces/embodiments/factory.py +++ b/autodata_interfaces/embodiments/factory.py @@ -17,11 +17,13 @@ from pathlib import Path from typing import Any +from autodata_interfaces.embodiments.absolute_pose_embodiment_adapter import AbsolutePoseEmbodimentAdapter from autodata_interfaces.embodiments.bimanual_embodiment_adapter import AbsolutePoseWholeBodyBimanualAdapter from autodata_interfaces.embodiments.embodiment_adapter import EmbodimentAdapter from autodata_interfaces.embodiments.single_arm_embodiment_adapter import DeltaPoseIKSingleArmAdapter EMBODIMENT_TYPE_REGISTRY: dict[str, type[EmbodimentAdapter]] = { + "absolute_pose": AbsolutePoseEmbodimentAdapter, "delta_pose_ik_single_arm": DeltaPoseIKSingleArmAdapter, "absolute_pose_whole_body_bimanual": AbsolutePoseWholeBodyBimanualAdapter, } diff --git a/autodata_interfaces/env/env_profile.py b/autodata_interfaces/env/env_profile.py index 921f353..937f358 100644 --- a/autodata_interfaces/env/env_profile.py +++ b/autodata_interfaces/env/env_profile.py @@ -124,6 +124,7 @@ class EnvironmentProfile: base_env: Env id the profile is written against; must equal the id used at apply time. scene: Scene changes (:class:`SceneSpec`). events: Event-term changes (:class:`EventsSpec`). + actions: Parameter overrides for existing action terms. planner: Name of the motion-planner profile tuned for this scene (key into :meth:`CuroboPlannerCfg.from_profile`), or None to fall back to task-name matching. """ @@ -134,6 +135,8 @@ class EnvironmentProfile: scene: SceneSpec = field(default_factory=SceneSpec) events: EventsSpec = field(default_factory=EventsSpec) planner: str | None = None + actions: dict[str, dict[str, Any]] = field(default_factory=dict) + """Action-term parameter overrides, applied before the action manager is built.""" @classmethod def from_yaml(cls, path: str) -> EnvironmentProfile: @@ -148,6 +151,8 @@ def from_yaml(cls, path: str) -> EnvironmentProfile: description: # optional base_env: planner: # optional + actions: # optional + : {: } scene: # optional rigid_objects: add: @@ -211,6 +216,7 @@ def from_dict(cls, data: dict[str, Any]) -> EnvironmentProfile: scene=scene, events=events, planner=data.get("planner"), + actions=data.get("actions", {}), ) def referenced_asset_names(self) -> set[str]: @@ -231,7 +237,7 @@ def validate_profile_dict(data: dict[str, Any]) -> None: assert isinstance(data, dict), f"Expected top-level dict, got {type(data).__name__}" required_keys = {"name", "base_env"} - optional_keys = {"description", "planner", "scene", "events"} + optional_keys = {"description", "planner", "scene", "events", "actions"} keys = set(data) missing = required_keys - keys assert not missing, f"Missing required top-level keys: {sorted(missing)}" @@ -248,6 +254,12 @@ def validate_profile_dict(data: dict[str, Any]) -> None: _validate_scene_dict(data.get("scene", {})) _validate_events_dict(data.get("events", {})) + actions = data.get("actions", {}) + assert isinstance(actions, dict), "'actions' must map action term names to parameter overrides." + assert all( + isinstance(name, str) and isinstance(params, dict) and all(isinstance(key, str) for key in params) + for name, params in actions.items() + ), "Each action override must map a term name to a parameter dict with string keys." def _validate_scene_dict(scene: dict[str, Any]) -> None: diff --git a/autodata_interfaces/env/isaaclab_env_interface.py b/autodata_interfaces/env/isaaclab_env_interface.py index e799341..3a8f916 100644 --- a/autodata_interfaces/env/isaaclab_env_interface.py +++ b/autodata_interfaces/env/isaaclab_env_interface.py @@ -83,6 +83,14 @@ def apply_env_profile(env_cfg: Any, profile: EnvironmentProfile, env_name: str) f"but is being applied to {env_name!r}." ) + # Actions: override controller settings before the manager constructs its terms. + for term_name, overrides in profile.actions.items(): + term_cfg = getattr(env_cfg.actions, term_name, None) + assert term_cfg is not None, f"Cannot override unknown action term {term_name!r}." + for parameter, value in overrides.items(): + assert hasattr(term_cfg, parameter), f"Action term {term_name!r} has no parameter {parameter!r}." + setattr(term_cfg, parameter, value) + # Scene: add rigid objects. for asset_name, add_spec in profile.scene.rigid_objects.add.items(): assert ( diff --git a/autodata_interfaces/env/registration.py b/autodata_interfaces/env/registration.py new file mode 100644 index 0000000..e9444aa --- /dev/null +++ b/autodata_interfaces/env/registration.py @@ -0,0 +1,53 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Load externally registered environments without coupling generation to their provider.""" + +import argparse +import sys +from collections.abc import Callable +from typing import Any + +from isaaclab.utils.string import string_to_callable + +from autodata_utils.cli import callable_path + + +def load_callable(path: str) -> Callable[..., Any]: + """Resolve a callable written as ``module:function`` or ``module.function``.""" + path = callable_path(path) + return string_to_callable(path, separator=":" if ":" in path else ".") + + +def register_external_environment(args: argparse.Namespace, extra_args: list[str]) -> None: + """Invoke an Isaac Lab registration callback after the simulation app starts. + + The callback receives environment arguments through ``sys.argv``, as in Isaac Lab's + native tools. AutoData's descriptor arguments are excluded to avoid collisions with + provider-specific options. The original arguments are restored even if registration fails. + Options after ``--`` are forwarded verbatim, allowing provider options to reuse AutoData names. + + Args: + args: AutoData CLI arguments, including ``external_callback`` and ``env_name``. + extra_args: Environment-provider arguments left by ``parse_known_args``. + """ + if args.external_callback is None: + assert not extra_args, f"Unrecognized arguments: {extra_args}. Use --external_callback for environment options." + return + assert args.env_name, "--env_name is required when using --external_callback." + callback_args = ["--task", args.env_name, "--device", args.device] + if hasattr(args, "num_envs"): + callback_args += ["--num_envs", str(args.num_envs)] + if args.record_images: + callback_args.append("--enable_cameras") + if "--" in extra_args: + separator = extra_args.index("--") + extra_args = extra_args[:separator] + extra_args[separator + 1 :] + callback_args.extend(extra_args) + original_argv = sys.argv + try: + sys.argv = [original_argv[0], *callback_args] + remaining = load_callable(args.external_callback)() + assert not remaining, f"Environment callback did not consume arguments: {remaining}" + finally: + sys.argv = original_argv diff --git a/autodata_interfaces/tasks/generation_policy_spec.py b/autodata_interfaces/tasks/generation_policy_spec.py index 67cf8c2..0883ed3 100644 --- a/autodata_interfaces/tasks/generation_policy_spec.py +++ b/autodata_interfaces/tasks/generation_policy_spec.py @@ -23,6 +23,8 @@ class GenerationPolicy: use_skillgen: Whether SkillGen is used to generate motion trajectories. use_navigation_controller: Whether a navigation controller generates loco-manipulation trajectories. + max_navigation_hold_steps: Maximum consecutive simulation steps spent waiting for a + navigation goal at the end of its recorded segment. Reaching the limit fails the attempt. Segment stitching parameters: select_src_per_subtask: If True, re-select a source demo for every subtask. If False, the @@ -47,9 +49,13 @@ class GenerationPolicy: task_name: str | None = None use_skillgen: bool = False use_navigation_controller: bool = False + max_navigation_hold_steps: int = 1000 # --- segment stitching --- select_src_per_subtask: bool = False select_src_per_arm: bool = False transform_first_robot_pose: bool = False interpolate_from_last_target_pose: bool = True + + def __post_init__(self) -> None: + assert self.max_navigation_hold_steps > 0, "max_navigation_hold_steps must be positive." diff --git a/autodata_interfaces/tasks/task_descriptor.py b/autodata_interfaces/tasks/task_descriptor.py index 32b9f25..23c2a07 100644 --- a/autodata_interfaces/tasks/task_descriptor.py +++ b/autodata_interfaces/tasks/task_descriptor.py @@ -6,13 +6,16 @@ import yaml from dataclasses import MISSING, dataclass, field from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any from autodata_interfaces.tasks.generation_policy_spec import GenerationPolicy from autodata_interfaces.tasks.subtask_constraint_spec import SubtaskConstraint from autodata_interfaces.tasks.subtask_spec import ALGO_PARAMS_REGISTRY, Subtask, SubtaskAlgoParams from autodata_interfaces.tasks.task_descriptor_utils import build_constraint, build_subtask, validate_task_dict +if TYPE_CHECKING: + from autodata_interfaces.embodiments.embodiment_adapter import EmbodimentAdapter + @dataclass class TaskDescriptor: @@ -43,6 +46,23 @@ def bind_env(self, env: Any) -> None: assert self.env is None, "env already bound" self.env = env + def validate_embodiment(self, embodiment_adapter: EmbodimentAdapter) -> None: + """Validate task/embodiment compatibility without starting the simulator. + + Args: + embodiment_adapter: Adapter whose channels and navigation configuration must match this task. + """ + task_eefs = set(self.get_eef_names()) + adapter_eefs = set(embodiment_adapter.get_eef_names()) + assert task_eefs == adapter_eefs, ( + f"EEF mismatch: task descriptor declares {sorted(task_eefs)}, " + f"embodiment adapter exposes {sorted(adapter_eefs)}" + ) + has_navigation = embodiment_adapter.get_navigation_eef_name() is not None + assert ( + self.generation_policy.use_navigation_controller == has_navigation + ), "Set generation_policy.use_navigation_controller exactly when the embodiment declares navigation." + def get_eef_names(self) -> list[str]: """Return the end-effector names declared by this task.""" diff --git a/autodata_tests/core/test_algorithms.py b/autodata_tests/core/test_algorithms.py index d37e319..b3b0867 100644 --- a/autodata_tests/core/test_algorithms.py +++ b/autodata_tests/core/test_algorithms.py @@ -56,7 +56,7 @@ def test_mimicgen_attributes(): def test_dexmimicgen_attributes(): algo = DexMimicGen() assert algo.name == "dexmimicgen" - assert algo.expected_eef_count == 2 + assert algo.expected_eef_count == (2, 3) assert algo.supports_coordination is True assert algo.requires_motion_planner is False diff --git a/autodata_tests/core/test_data_generator.py b/autodata_tests/core/test_data_generator.py index 51566da..806c0e9 100644 --- a/autodata_tests/core/test_data_generator.py +++ b/autodata_tests/core/test_data_generator.py @@ -3,16 +3,22 @@ """Unit tests for :class:`autodata_core.data_generator.DataGenerator`. -Only the construction-time validation is covered here. The async generation loop is -exercised by the e2e suite. +Construction-time validation and navigation playback decisions are covered here. +Simulation integration is covered by the e2e suite. """ +import asyncio import dataclasses +import torch +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock import pytest from autodata_core.algorithms import DexMimicGen, MimicGen, SkillGen -from autodata_core.data_generator import DataGenerator, GenerationResult +from autodata_core.data_generator import DataGenerator, GenerationResult, _EEFGenerationState +from autodata_core.waypoint import MultiWaypoint, Waypoint +from autodata_interfaces.tasks.generation_policy_spec import GenerationPolicy from autodata_interfaces.tasks.subtask_spec import MimicGenSubtaskAlgoParams, Subtask @@ -25,12 +31,22 @@ def _subtask(term_signal: str = "s", term_offset: tuple[int, int] = (0, 0)) -> S class _MockDatastream: - """Minimal stand-in exposing only what ``DataGenerator.__init__`` reads.""" + """Minimal datastream for construction and navigation playback tests.""" def __init__(self, subtasks_by_eef: dict, constraints: list | None = None) -> None: - self.source_pool = object() + self.source_pool = SimpleNamespace( + asyncio_lock=None, + datagen_infos=[object()], + subtask_boundaries={ + name: [[(2 * i, 2 * i + 2) for i in range(len(subtasks))]] for name, subtasks in subtasks_by_eef.items() + }, + ) self._subtasks = subtasks_by_eef self._constraints = constraints or [] + self.navigation_state = None + self.generation_policy = GenerationPolicy() + self.device = "cpu" + self.env = None def get_eef_names(self): return list(self._subtasks) @@ -41,6 +57,15 @@ def get_subtasks(self, eef_name): def get_task_constraints(self): return self._constraints + def get_navigation_state(self, env_id): + return self.navigation_state + + def get_generation_policy(self): + return self.generation_policy + + def get_env(self): + return self.env + def _assert_constructed(generator: DataGenerator, datastream: _MockDatastream, algorithm) -> None: assert generator.algorithm is algorithm @@ -114,3 +139,87 @@ def test_generation_result_fields_and_frozen(): assert result.success is True with pytest.raises(dataclasses.FrozenInstanceError): result.success = False # type: ignore[misc] + + +def _navigation_generator(active: bool, reached: bool) -> DataGenerator: + datastream = _MockDatastream({name: [_subtask("")] for name in ("left", "right", "base")}) + datastream.navigation_state = ("base", active, reached) + datastream.generation_policy.use_navigation_controller = True + return DataGenerator(datastream=datastream, algorithm=DexMimicGen()) + + +def _playback_state(length: int, index: int, *, paused: bool = False, done: bool = False) -> _EEFGenerationState: + waypoint = Waypoint(pose=torch.eye(4), passthrough_action={}) + return _EEFGenerationState( + current_trajectory=[waypoint] * length, + subtask_step_index=index, + is_currently_paused=paused, + subtasks_done=done, + ) + + +def test_navigation_holds_playback_until_goal_reached(): + states = {"base": _playback_state(10, 8), "left": _playback_state(20, 8), "right": _playback_state(20, 8)} + generator = _navigation_generator(active=True, reached=False) + assert generator._synchronize_navigation(states, 0, True) == (True, True) + assert all(state.subtask_step_index == 8 for state in states.values()) + + generator.datastream.navigation_state = ("base", False, True) + assert generator._synchronize_navigation(states, 0, True) == (False, False) + + +@pytest.mark.parametrize( + "length, paused, done", + [(20, True, False), (20, False, True), (5, False, False)], + ids=["paused", "completed", "last_waypoint"], +) +def test_navigation_completion_skips_idle_steps_without_rewinding_other_channels(length, paused, done): + states = { + "base": _playback_state(10, 4), + "left": _playback_state(20, 4), + "right": _playback_state(length, 4, paused=paused, done=done), + } + generator = _navigation_generator(active=False, reached=True) + assert generator._synchronize_navigation(states, 0, True) == (False, False) + assert {name: state.subtask_step_index for name, state in states.items()} == {"base": 8, "left": 8, "right": 4} + + +@pytest.mark.parametrize("paused, reached", [(False, False), (True, True)], ids=["cancelled", "paused"]) +def test_navigation_does_not_skip_playback_after_cancellation_or_while_paused(paused, reached): + states = { + "base": _playback_state(10, 4, paused=paused), + "left": _playback_state(20, 4), + "right": _playback_state(20, 4), + } + generator = _navigation_generator(active=False, reached=reached) + assert generator._synchronize_navigation(states, 0, True) == (False, False) + assert all(state.subtask_step_index == 4 for state in states.values()) + + +def test_unresponsive_navigation_fails_and_finalizes_the_attempt(monkeypatch): + generator = _navigation_generator(active=True, reached=False) + generator.datastream.generation_policy.max_navigation_hold_steps = 3 + recorder = Mock() + generator.datastream.env = SimpleNamespace(recorder_manager=recorder) + waypoint = Waypoint(pose=torch.eye(4), passthrough_action={}) + monkeypatch.setattr(generator.algorithm, "plan_subtask_trajectory", lambda **kwargs: ([waypoint, waypoint], False)) + monkeypatch.setattr(generator, "_reset_environment_for_generation", AsyncMock(return_value=(torch.tensor([0]), {}))) + execute = AsyncMock(return_value=False) + monkeypatch.setattr(MultiWaypoint, "execute", execute) + + result = asyncio.run( + asyncio.wait_for( + generator.generate( + env_id=0, + success_term=None, + env_reset_queue=asyncio.Queue(), + env_action_queue=asyncio.Queue(), + ), + timeout=1.0, + ) + ) + + assert result.success is False + assert execute.await_count == 3 + assert recorder.set_success_to_episodes.call_args.args[1].item() is False + recorder.export_episodes.assert_called_once() diff --git a/autodata_tests/e2e/test_arena_data_generation.py b/autodata_tests/e2e/test_arena_data_generation.py new file mode 100644 index 0000000..a051d3f --- /dev/null +++ b/autodata_tests/e2e/test_arena_data_generation.py @@ -0,0 +1,99 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Generate one successful demonstration for each supported Arena workflow.""" + +import yaml +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, read_generation_result + + +@pytest.mark.with_subprocess +@pytest.mark.parametrize( + "env_name, task_descriptor, embodiment, dataset_file, env_profile", + [ + pytest.param( + "gr1_open_microwave", + "gr1_open_microwave.yaml", + "gr1_ik_abs.yaml", + "arena_gr1_manipulation_dataset_annotated.hdf5", + None, + id="gr1_open_microwave", + ), + pytest.param( + "put_item_in_fridge_and_close_door", + "gr1_put_and_close_door.yaml", + "gr1_ik_abs.yaml", + "ranch_bottle_into_fridge_annotated.hdf5", + None, + id="gr1_put_and_close_door", + ), + pytest.param( + "galileo_g1_locomanip_pick_and_place", + "g1_locomanip.yaml", + "g1_wbc_pink.yaml", + "arena_g1_loco_manipulation_dataset_annotated.hdf5", + "g1_locomanip.yaml", + id="g1_locomanip", + ), + ], +) +def test_arena_data_generation( + env_name: str, + task_descriptor: str, + embodiment: str, + dataset_file: str, + env_profile: str | None, + tmp_path: Path, +) -> None: + """Run one CPU environment until one successful episode is generated and exported.""" + with (Path(TestPaths.tasks_dir) / task_descriptor).open(encoding="utf-8") as descriptor_file: + descriptor = yaml.safe_load(descriptor_file) + # These examples normally count attempts; the test requests one actual success. + descriptor["generation_policy"]["guarantee_success"] = True + test_descriptor = tmp_path / "task_descriptor.yaml" + test_descriptor.write_text(yaml.safe_dump(descriptor), encoding="utf-8") + + output_file = tmp_path / "generated.hdf5" + result_file = tmp_path / "generation_result.json" + args = [ + TestPaths.python_path, + TestPaths.generate_dataset_script, + "--external_callback", + "isaaclab_arena.environments.isaaclab_interop.environment_registration_callback", + "--env_name", + env_name, + "--alg", + "dexmimicgen", + "--task_descriptor", + str(test_descriptor), + "--embodiment", + str(Path(TestPaths.embodiments_dir) / embodiment), + "--input_file", + str(Path(TestPaths.test_data_dir) / "arena_data_generation" / dataset_file), + "--output_file", + str(output_file), + "--result_file", + str(result_file), + "--generation_num_trials", + "1", + "--num_envs", + "1", + "--device", + "cpu", + "--viz", + "none", + ] + if env_profile is not None: + args.extend(["--env_profile", str(Path(TestPaths.env_profiles_dir) / env_profile)]) + + run_subprocess(args) + + num_success, _ = read_generation_result(str(result_file)) + assert num_success == 1, f"Expected one successful episode for {env_name}, got {num_success}." + assert_valid_dataset(str(output_file), min_num_demos=1) diff --git a/autodata_tests/interfaces/datastream/test_datastream.py b/autodata_tests/interfaces/datastream/test_datastream.py index 8a45db1..56486fa 100644 --- a/autodata_tests/interfaces/datastream/test_datastream.py +++ b/autodata_tests/interfaces/datastream/test_datastream.py @@ -232,6 +232,48 @@ 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_articulated_object_poses_use_controller_frame(): + """A moving controller frame must transform articulated roots just like rigid objects.""" + datastream, _, adapter, env = _datastream() + env.scene.articulations = { + "robot": env.scene["robot"], + "door": MockAsset( + MockArticulationData( + root_pos_w=torch.tensor([[2.0, 3.0, 0.5]]), + root_quat_w=torch.tensor([[0.0, 0.0, 0.0, 1.0]]), + ) + ), + } + reference = torch.tensor([ + [0.0, -1.0, 0.0, 1.0], + [1.0, 0.0, 0.0, 2.0], + [0.0, 0.0, 1.0, 0.5], + [0.0, 0.0, 0.0, 1.0], + ]).unsqueeze(0) + adapter.get_reference_frame_pose = lambda env_ids: reference + poses = datastream.get_object_poses() + assert set(poses) == {"cube_1", "door"} + assert torch.allclose( + poses["door"][0], + torch.tensor([ + [0.0, 1.0, 0.0, 1.0], + [-1.0, 0.0, 0.0, -1.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 1.0], + ]), + ) + assert torch.allclose(poses["cube_1"][0, :3, 3], torch.tensor([-1.0, 0.0, -0.45])) + assert torch.allclose( + datastream.get_robot_root_pose(env_ids=[0])[0], + torch.tensor([ + [0.0, 1.0, 0.0, -2.0], + [-1.0, 0.0, 0.0, 0.5], + [0.0, 0.0, 1.0, -0.4], + [0.0, 0.0, 0.0, 1.0], + ]), + ) + + def test_get_robot_root_pose(): datastream, _, _, _ = _datastream() pose = datastream.get_robot_root_pose(env_ids=[0]) diff --git a/autodata_tests/interfaces/embodiments/test_absolute_pose_embodiment_adapter.py b/autodata_tests/interfaces/embodiments/test_absolute_pose_embodiment_adapter.py new file mode 100644 index 0000000..b58fff6 --- /dev/null +++ b/autodata_tests/interfaces/embodiments/test_absolute_pose_embodiment_adapter.py @@ -0,0 +1,76 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Regression checks for absolute-pose embodiments with independently arranged action channels.""" + +import torch +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from autodata_interfaces.embodiments.factory import embodiment_adapter_from_dict, embodiment_adapter_from_yaml +from autodata_tests.interfaces.mocks import MockAsset, MockEnv, MockScene + + +def _adapter(): + return embodiment_adapter_from_yaml(Path(__file__).parents[3] / "autodata_examples/embodiments/g1_wbc_pink.yaml") + + +def test_pose_conversion_preserves_scalar_grippers_and_body_commands(): + adapter = _adapter() + actions = torch.arange(46, dtype=torch.float32).reshape(2, 23) + diagonal = 2.0**-0.5 + actions[:, 5:9] = torch.tensor([0.0, 0.0, diagonal, diagonal]) + actions[:, 12:16] = torch.tensor([diagonal, 0.0, 0.0, diagonal]) + poses = adapter.action_to_target_eef_pose(actions) + channels = adapter.actions_to_passthrough_actions(actions) + for i in range(2): + reconstructed = adapter.target_eef_pose_to_action( + {name: pose[i] for name, pose in poses.items()}, + {name: command[i] for name, command in channels.items()}, + ) + assert torch.allclose(reconstructed, actions[i]) + noisy_action = adapter.target_eef_pose_to_action( + {name: pose[i] for name, pose in poses.items()}, + {name: command[i] for name, command in channels.items()}, + action_noise_dict={name: 0.01 for name in poses}, + ) + noisy_channels = adapter.actions_to_passthrough_actions(noisy_action) + assert all(torch.equal(noisy_channels[name], command[i]) for name, command in channels.items()) + + +def test_controller_state_queries_select_requested_environment(): + adapter = _adapter() + state = torch.zeros(2, 2, 13) + state[..., 6] = 1.0 + state[1, 1, :3] = torch.tensor([3.0, 4.0, 0.8]) + robot = MockAsset(SimpleNamespace(body_names=["other", "pelvis"], body_link_state_w=state)) + scene = MockScene(assets={"robot": robot}, env_origins=torch.tensor([[0.0, 0.0, 0.0], [1.0, 2.0, 0.0]])) + adapter.bind_env( + MockEnv( + scene=scene, + obs_buf={ + "wbc": { + "is_navigating": torch.tensor([False, True]), + "navigation_goal_reached": torch.tensor([True, False]), + } + }, + ) + ) + frame = adapter.get_reference_frame_pose([1]) + assert frame.shape == (1, 4, 4) + assert torch.allclose(frame[0, :3, 3], torch.tensor([2.0, 2.0, 0.8])) + assert adapter.get_navigation_state(1) == ("body", True, False) + + +def test_bad_action_layout_rejected(): + adapter = _adapter() + adapter.eefs["right"]["pose_action_slice"] = [8, 15] + with pytest.raises(AssertionError, match="Overlapping"): + embodiment_adapter_from_dict({ + "type": "absolute_pose", + "name": "overlapping_layout", + "action_dim": adapter.action_dim, + "eefs": adapter.eefs, + }) diff --git a/autodata_tests/interfaces/embodiments/test_factory.py b/autodata_tests/interfaces/embodiments/test_factory.py index fab09e5..4b48a1a 100644 --- a/autodata_tests/interfaces/embodiments/test_factory.py +++ b/autodata_tests/interfaces/embodiments/test_factory.py @@ -14,6 +14,7 @@ embodiment_adapter_from_dict, embodiment_adapter_from_yaml, ) +from autodata_interfaces.embodiments.absolute_pose_embodiment_adapter import AbsolutePoseEmbodimentAdapter from autodata_tests.interfaces.mocks import MockEnv from autodata_tests.utils.constants import TestPaths @@ -28,6 +29,7 @@ def test_registry_contents(): assert EMBODIMENT_TYPE_REGISTRY == { + "absolute_pose": AbsolutePoseEmbodimentAdapter, "delta_pose_ik_single_arm": DeltaPoseIKSingleArmAdapter, "absolute_pose_whole_body_bimanual": AbsolutePoseWholeBodyBimanualAdapter, } diff --git a/autodata_tests/interfaces/env/test_env_profile.py b/autodata_tests/interfaces/env/test_env_profile.py index c319ac1..dd2f5cd 100644 --- a/autodata_tests/interfaces/env/test_env_profile.py +++ b/autodata_tests/interfaces/env/test_env_profile.py @@ -26,6 +26,7 @@ def _full_valid() -> dict: "description": "d", "base_env": "IsaacContrib-Stack-Cube-Franka-IK-Rel", "planner": "franka_stack_cube_bin", + "actions": {"controller": {"gain": 2.0}}, "scene": { "rigid_objects": { "add": { @@ -66,6 +67,13 @@ def test_validate_full_valid_passes(): validate_profile_dict(_full_valid()) # must not raise +def test_action_override_parameter_names_must_be_strings(): + data = _minimal_valid() + data["actions"] = {"controller": {1: 2.0}} + with pytest.raises(AssertionError, match="string keys"): + EnvironmentProfile.from_dict(data) + + def test_validate_missing_required_keys_fails(): for key in ("name", "base_env"): data = _minimal_valid() diff --git a/autodata_tests/interfaces/env/test_env_profile_apply.py b/autodata_tests/interfaces/env/test_env_profile_apply.py index dc21b69..f5bb068 100644 --- a/autodata_tests/interfaces/env/test_env_profile_apply.py +++ b/autodata_tests/interfaces/env/test_env_profile_apply.py @@ -9,6 +9,7 @@ import copy import os +from types import SimpleNamespace import pytest @@ -23,6 +24,24 @@ BASE_ENV = "IsaacContrib-Stack-Cube-Franka-IK-Rel" +def test_profile_action_overrides_are_applied_before_construction(): + profile = EnvironmentProfile.from_dict( + {"name": "controller", "base_env": "external-task", "actions": {"controller": {"gain": 2.0}}} + ) + cfg = SimpleNamespace(actions=SimpleNamespace(controller=SimpleNamespace(gain=1.0)), events=SimpleNamespace()) + apply_env_profile(cfg, profile, "external-task") + assert cfg.actions.controller.gain == 2.0 + + +def test_profile_rejects_unknown_action_parameters(): + profile = EnvironmentProfile.from_dict( + {"name": "controller", "base_env": "external-task", "actions": {"controller": {"unknown": 2.0}}} + ) + cfg = SimpleNamespace(actions=SimpleNamespace(controller=SimpleNamespace(gain=1.0))) + with pytest.raises(AssertionError, match="has no parameter"): + apply_env_profile(cfg, profile, "external-task") + + def _parse_base_cfg(): return parse_env_cfg(BASE_ENV, device="cpu", num_envs=1) diff --git a/autodata_tests/interfaces/env/test_registration.py b/autodata_tests/interfaces/env/test_registration.py new file mode 100644 index 0000000..61d7e0b --- /dev/null +++ b/autodata_tests/interfaces/env/test_registration.py @@ -0,0 +1,64 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Registration callbacks consume provider options without seeing AutoData descriptors.""" + +import argparse +import sys +from types import SimpleNamespace + +import pytest + +from autodata_interfaces.env.registration import register_external_environment + + +def test_callback_receives_only_environment_options(monkeypatch): + original = ["annotate_demos.py", "--embodiment", "robot.yaml", "--task_descriptor", "task.yaml"] + monkeypatch.setattr(sys, "argv", original) + captured = [] + + def register(): + captured.extend(sys.argv[1:]) + return [] + + monkeypatch.setitem(sys.modules, "external_provider", SimpleNamespace(register=register)) + args = argparse.Namespace( + external_callback="external_provider.register", + env_name="external-task", + device="cpu", + record_images=True, + num_envs=2, + ) + register_external_environment(args, ["--object", "box", "--", "--embodiment", "provider_robot"]) + assert captured == [ + "--task", + "external-task", + "--device", + "cpu", + "--num_envs", + "2", + "--enable_cameras", + "--object", + "box", + "--embodiment", + "provider_robot", + ] + assert sys.argv is original + + +def test_callback_failure_restores_process_arguments(monkeypatch): + original = sys.argv + + def register(): + raise RuntimeError("registration failed") + + monkeypatch.setitem(sys.modules, "external_provider", SimpleNamespace(register=register)) + args = argparse.Namespace( + external_callback="external_provider:register", + env_name="external-task", + device="cpu", + record_images=False, + ) + with pytest.raises(RuntimeError, match="registration failed"): + register_external_environment(args, []) + assert sys.argv is original diff --git a/autodata_tests/interfaces/tasks/test_generation_policy_spec.py b/autodata_tests/interfaces/tasks/test_generation_policy_spec.py index 2cfcb4a..8c78c7e 100644 --- a/autodata_tests/interfaces/tasks/test_generation_policy_spec.py +++ b/autodata_tests/interfaces/tasks/test_generation_policy_spec.py @@ -3,6 +3,8 @@ """Unit tests for :mod:`autodata_interfaces.tasks.generation_policy_spec`.""" +import pytest + from autodata_interfaces.tasks.generation_policy_spec import GenerationPolicy @@ -50,3 +52,9 @@ def test_generation_policy_instances_are_independent(): b = GenerationPolicy() a.num_trials = 999 assert b.num_trials == 10 + + +@pytest.mark.parametrize("limit", [0, -1]) +def test_navigation_hold_limit_must_be_positive(limit): + with pytest.raises(AssertionError, match="max_navigation_hold_steps must be positive"): + GenerationPolicy(max_navigation_hold_steps=limit) diff --git a/autodata_tests/interfaces/tasks/test_task_descriptor.py b/autodata_tests/interfaces/tasks/test_task_descriptor.py index 567cfce..ae22d65 100644 --- a/autodata_tests/interfaces/tasks/test_task_descriptor.py +++ b/autodata_tests/interfaces/tasks/test_task_descriptor.py @@ -7,6 +7,7 @@ import pytest +from autodata_interfaces.embodiments import embodiment_adapter_from_yaml from autodata_interfaces.tasks.generation_policy_spec import GenerationPolicy from autodata_interfaces.tasks.subtask_constraint_spec import SubtaskConstraint, SubTaskConstraintType from autodata_interfaces.tasks.subtask_spec import MimicGenSubtaskAlgoParams, SkillGenSubtaskAlgoParams @@ -207,6 +208,27 @@ def test_constraints_parsed_and_exposed(): assert set(runtime) == {("left", 0), ("right", 0)} +# --------------------------------------------------------------------------------------------------- +# Task/embodiment compatibility before simulation startup +# --------------------------------------------------------------------------------------------------- +@pytest.mark.parametrize( + "uses_navigation, has_navigation", + [(False, False), (True, True), (True, False), (False, True)], +) +def test_navigation_setup_requires_matching_task_and_embodiment(uses_navigation, has_navigation): + task = TaskDescriptor.from_yaml(os.path.join(TestPaths.tasks_dir, "g1_locomanip.yaml")) + adapter = embodiment_adapter_from_yaml(os.path.join(TestPaths.embodiments_dir, "g1_wbc_pink.yaml")) + task.generation_policy.use_navigation_controller = uses_navigation + if not has_navigation: + adapter.navigation = None + + if uses_navigation == has_navigation: + task.validate_embodiment(adapter) + else: + with pytest.raises(AssertionError, match="use_navigation_controller"): + task.validate_embodiment(adapter) + + # --------------------------------------------------------------------------------------------------- # bind_env # --------------------------------------------------------------------------------------------------- diff --git a/autodata_tests/test_data/arena_data_generation/arena_g1_loco_manipulation_dataset_annotated.hdf5 b/autodata_tests/test_data/arena_data_generation/arena_g1_loco_manipulation_dataset_annotated.hdf5 new file mode 100644 index 0000000..46ddc24 --- /dev/null +++ b/autodata_tests/test_data/arena_data_generation/arena_g1_loco_manipulation_dataset_annotated.hdf5 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:429ce048e98af7e4fe7706fb5ef3d0d28908a019eaa0e11c1cb08d904855f41d +size 1461939 diff --git a/autodata_tests/test_data/arena_data_generation/arena_gr1_manipulation_dataset_annotated.hdf5 b/autodata_tests/test_data/arena_data_generation/arena_gr1_manipulation_dataset_annotated.hdf5 new file mode 100644 index 0000000..7101398 --- /dev/null +++ b/autodata_tests/test_data/arena_data_generation/arena_gr1_manipulation_dataset_annotated.hdf5 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cec97dbf285e89ec212a2994e9a85ce4b7198654d3172e9dbfba8e91c35aa674 +size 4546487 diff --git a/autodata_tests/test_data/arena_data_generation/ranch_bottle_into_fridge_annotated.hdf5 b/autodata_tests/test_data/arena_data_generation/ranch_bottle_into_fridge_annotated.hdf5 new file mode 100644 index 0000000..d43babc --- /dev/null +++ b/autodata_tests/test_data/arena_data_generation/ranch_bottle_into_fridge_annotated.hdf5 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2dd0ded921638ed78651b6bdda0e0cec6aec8f0ab942109a725e882495d6cb1c +size 5729085 diff --git a/autodata_utils/cli.py b/autodata_utils/cli.py new file mode 100644 index 0000000..bf4ca1c --- /dev/null +++ b/autodata_utils/cli.py @@ -0,0 +1,24 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Lightweight CLI validation that can run before simulation startup.""" + +import argparse + + +def callable_path(value: str) -> str: + """Validate a callable path without importing its module. + + Args: + value: Callable written as module.function or module:function. + + Returns: + The validated path. + """ + separator = ":" if ":" in value else "." + module_name, _, name = value.rpartition(separator) + if not name.isidentifier() or not all(part.isidentifier() for part in module_name.split(".")): + raise argparse.ArgumentTypeError( + f"Expected a full callable path (module.function or module:function), got {value!r}." + ) + return value diff --git a/scripts/annotate_demos.py b/scripts/annotate_demos.py index a334cc4..b0d6b2b 100644 --- a/scripts/annotate_demos.py +++ b/scripts/annotate_demos.py @@ -33,6 +33,8 @@ from isaaclab.app import AppLauncher +from autodata_utils.cli import callable_path + parser = argparse.ArgumentParser( description="Manually annotate AutoData source demonstrations with subtask signals.", formatter_class=argparse.RawDescriptionHelpFormatter, @@ -87,7 +89,10 @@ help="Render and record the task's image observations.", ) AppLauncher.add_app_launcher_args(parser) -args_cli = parser.parse_args() +parser.add_argument("--external_callback", type=callable_path, help="Callable registering environments after startup.") +args_cli, environment_args = parser.parse_known_args() +if environment_args and args_cli.external_callback is None: + parser.error(f"unrecognized arguments: {' '.join(environment_args)}") app_launcher = AppLauncher(args_cli, enable_cameras=args_cli.record_images) simulation_app = app_launcher.app @@ -114,6 +119,7 @@ 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.registration import register_external_environment # noqa: E402 from autodata_interfaces.tasks.task_descriptor import TaskDescriptor # noqa: E402 is_paused = False @@ -214,6 +220,7 @@ def main() -> int: """Annotate the source dataset and export the annotated copy.""" global _datastream, is_paused, skip_episode, marked_subtask_action_indices + register_external_environment(args_cli, environment_args) assert args_cli.auto or app_launcher.has_window, ( "annotate_demos.py performs manual keyboard annotation and cannot run headless. " "Re-run with a window (--viz kit), or pass --auto." diff --git a/scripts/generate_dataset.py b/scripts/generate_dataset.py index d168d7d..d6f5167 100755 --- a/scripts/generate_dataset.py +++ b/scripts/generate_dataset.py @@ -36,7 +36,11 @@ from isaaclab.app import AppLauncher +from autodata_interfaces.embodiments import embodiment_adapter_from_yaml +from autodata_interfaces.embodiments.embodiment_adapter import EmbodimentAdapter +from autodata_interfaces.tasks.task_descriptor import TaskDescriptor from autodata_interfaces.tasks.task_descriptor_utils import resolve_task_algorithm +from autodata_utils.cli import callable_path # Hardcoded to keep argparse importable without pulling in the heavy core package. # Add new algorithms here when registering them in autodata_core.algorithms. @@ -113,8 +117,17 @@ help="Render and record the task's image observations.", ) AppLauncher.add_app_launcher_args(parser) -args_cli = parser.parse_args() +parser.add_argument("--external_callback", type=callable_path, help="Callable registering environments after startup.") +parser.add_argument( + "--recorder_cfg", type=callable_path, help="Callable returning an Isaac Lab recorder configuration." +) +args_cli, environment_args = parser.parse_known_args() +if environment_args and args_cli.external_callback is None: + parser.error(f"unrecognized arguments: {' '.join(environment_args)}") args_cli.alg = resolve_task_algorithm(args_cli.alg, args_cli.task_descriptor) +task_descriptor = TaskDescriptor.from_yaml(args_cli.task_descriptor) +embodiment_adapter = embodiment_adapter_from_yaml(args_cli.embodiment) +task_descriptor.validate_embodiment(embodiment_adapter) app_launcher = AppLauncher(args_cli, enable_cameras=args_cli.record_images) simulation_app = app_launcher.app @@ -133,7 +146,6 @@ from autodata_core import DataGenerator, get_algorithm # noqa: E402 from autodata_core.algorithms import REGISTERED_ALGORITHMS # 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 EnvironmentProfile, env_loop, @@ -141,7 +153,7 @@ setup_env_config, setup_output_paths, ) -from autodata_interfaces.tasks.task_descriptor import TaskDescriptor # noqa: E402 +from autodata_interfaces.env.registration import load_callable, register_external_environment # noqa: E402 from autodata_utils.generation_result import write_generation_result # noqa: E402 @@ -181,13 +193,12 @@ def build_datastream( env: Any, input_file: str, task_descriptor: TaskDescriptor, - embodiment_yaml: str, + embodiment_adapter: EmbodimentAdapter, ) -> Datastream: """Compose the read interface (task + embodiment + source pool) over the live env. Built before the motion planners so the planners can read all world state through it. """ - embodiment_adapter = embodiment_adapter_from_yaml(embodiment_yaml) datastream = Datastream( env=env, task_descriptor=task_descriptor, @@ -287,12 +298,12 @@ def _close_motion_planners(planners: dict | None) -> None: def main() -> None: + register_external_environment(args_cli, environment_args) output_dir, output_file_name = setup_output_paths(args_cli.output_file) # The env name (CLI override) falls back to the name recorded in the source dataset. env_name = args_cli.env_name.split(":")[-1] if args_cli.env_name else get_env_name_from_dataset(args_cli.input_file) # The task descriptor's GenerationPolicy is the source for generation policy parameters. - task_descriptor = TaskDescriptor.from_yaml(args_cli.task_descriptor) generation_policy_params = task_descriptor.get_generation_policy() # The CLI --generation_num_trials, when given, overrides the descriptor's num_trials. @@ -316,6 +327,7 @@ def main() -> None: num_envs=args_cli.num_envs, device=args_cli.device, generation_policy_params=generation_policy_params, + recorder_cfg=load_callable(args_cli.recorder_cfg)() if args_cli.recorder_cfg else None, env_profile=env_profile, enable_cameras=args_cli.record_images, ) @@ -333,7 +345,7 @@ def main() -> None: env=env, input_file=args_cli.input_file, task_descriptor=task_descriptor, - embodiment_yaml=args_cli.embodiment, + embodiment_adapter=embodiment_adapter, ) # SkillGen needs one curobo planner per env. Mimic/DexMimic take no kwargs.