Enable Arena workflows to migrate annotation and generation in AutoData

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

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,

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

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]

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

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

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

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.

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",