diff --git a/autodata_examples/embodiments/g1_ik_abs.yaml b/autodata_examples/embodiments/g1_ik_abs.yaml index 5c17b06..0da6d5c 100644 --- a/autodata_examples/embodiments/g1_ik_abs.yaml +++ b/autodata_examples/embodiments/g1_ik_abs.yaml @@ -1,4 +1,6 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026, The Isaac AutoData Project Developers. +# All rights reserved. +# # SPDX-License-Identifier: Apache-2.0 # Embodiment descriptor for the Unitree G1 (whole-body IK, Absolute-pose). @@ -27,6 +29,12 @@ type: absolute_pose_whole_body_bimanual name: g1 description: Unitree G1 humanoid with whole-body IK and three-finger hands, absolute-pose control. +# Each EEF's `offset` is the SE(3) offset from the tracked control-link frame (wrist_yaw_link) to a +# canonical grasp frame. Rotate-then-translate: rotate by axis_angle [rad], then translate [m] in that +# grasp frame. Retargeting extracts the trajectory in the canonical frame (source_pose @ offset) and +# re-anchors to the target's control link (@ inv(target offset)), so different grippers align at the +# grasp point (handles gripper length + wrist convention). Seeded from the control-link -> grasped- +# object transform at grasp; translation is shared L/R (same gripper), rotation is per-hand. eefs: left: pose_obs_keys: @@ -34,6 +42,9 @@ eefs: quat: left_eef_quat # 7 left-hand joints (index/middle/thumb proximal + distal + thumb tip) gripper_action_indices: [0, 1, 2, 6, 7, 8, 12] + offset: + axis_angle: [1.2092, -1.2092, 1.2092] + translation: [0.0, -0.105, 0.03] right: pose_obs_keys: @@ -41,6 +52,9 @@ eefs: quat: right_eef_quat # 7 right-hand joints (index/middle/thumb proximal + distal + thumb tip) gripper_action_indices: [3, 4, 5, 9, 10, 11, 13] + offset: + axis_angle: [-1.2092, -1.2092, -1.2092] + translation: [0.0, 0.105, 0.03] action_layout: left_pose_slice: [0, 7] @@ -50,3 +64,13 @@ action_layout: # Non-eef passthrough channels: name -> contiguous action slice. Add more here freely. passthrough_channels: body: [28, 32] # [vel_x, vel_y, rot_vel_z, hip_height] + +# Per-EEF open and closed hand joint configurations (7 hand joints each, in gripper_action_indices +# order). Used by retargeting's 'binary' / 'interpolation' hand policies to score the source hand's +# closedness. Measured from the G1 demos. +hand_open: + left: [0.0, 0.0, 0.0, 0.0, 0.0, 0.72, 0.0] + right: [0.0, 0.0, 0.0, 0.0, 0.0, -0.72, 0.0] +hand_close: + left: [-1.57, -1.57, 0.0, -1.74, -1.74, 1.04, 1.74] + right: [1.57, 1.57, .0, 1.74, 1.74, -1.04, -1.74] diff --git a/autodata_examples/embodiments/gr1_ik_abs.yaml b/autodata_examples/embodiments/gr1_ik_abs.yaml index e1707ec..314564b 100644 --- a/autodata_examples/embodiments/gr1_ik_abs.yaml +++ b/autodata_examples/embodiments/gr1_ik_abs.yaml @@ -1,4 +1,6 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026, The Isaac AutoData Project Developers. +# All rights reserved. +# # SPDX-License-Identifier: Apache-2.0 # Embodiment descriptor for the GR1 (IK, Absolute-pose). @@ -18,6 +20,10 @@ type: absolute_pose_whole_body_bimanual name: gr1t2 description: GR1 humanoid with whole-body IK and dexterous hands, absolute-pose control. +# Each EEF's `offset` is the SE(3) offset from the tracked control-link frame (hand_pitch_link) to a +# canonical grasp frame. See g1_ik_abs.yaml for the convention (rotate then translate; translation +# shared L/R, rotation per-hand). Seeded from the control-link -> grasped-object transform from the +# GR1 demos. eefs: left: pose_obs_keys: @@ -25,6 +31,9 @@ eefs: quat: left_eef_quat # 11 left-hand joints: 5 proximal + 5 intermediate + 1 distal gripper_action_indices: [0, 1, 2, 3, 4, 10, 11, 12, 13, 14, 20] + offset: + axis_angle: [1.5708, 0.0, 0.0] + translation: [0.01, -0.08, 0.035] right: pose_obs_keys: @@ -32,9 +41,22 @@ eefs: quat: right_eef_quat # 11 right-hand joints: 5 proximal + 5 intermediate + 1 distal gripper_action_indices: [5, 6, 7, 8, 9, 15, 16, 17, 18, 19, 21] + offset: + axis_angle: [-1.5708, -0.0, -0.0] + translation: [0.01, 0.08, 0.035] action_layout: left_pose_slice: [0, 7] right_pose_slice: [7, 14] hand_joints_slice: [14, 36] canonicalize_quat: true + +# Per-EEF open and closed hand joint configurations (11 hand joints each, in gripper_action_indices +# order). Used by retargeting's 'binary' / 'interpolation' hand policies to score the source hand's +# closedness and drive the target hand between these configs. Measured from the GR1 demos. +hand_open: # fully extended (robot default) + left: [0.0, 0.0, 0.0, 0.0, -1.57, 0.0, 0.0, 0.0, 0.0, 0.5, -0.238] + right: [0.0, 0.0, 0.0, 0.0, -1.57, 0.0, 0.0, 0.0, 0.0, 0.5, -0.238] +hand_close: # full grip (limits, matching what the demos actually reach) + left: [-1.57, -1.57, -1.57, -1.57, -1.74, 0.35, 0.35, 0.35, 0.35, 1.22, -0.238] + right: [-1.57, -1.57, -1.57, -1.57, -1.74, 0.35, 0.35, 0.35, 0.35, 1.22, -0.238] diff --git a/autodata_examples/retarget/g1_to_gr1_pick_place.yaml b/autodata_examples/retarget/g1_to_gr1_pick_place.yaml new file mode 100644 index 0000000..e107433 --- /dev/null +++ b/autodata_examples/retarget/g1_to_gr1_pick_place.yaml @@ -0,0 +1,65 @@ +# Copyright (c) 2026, The Isaac AutoData Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +# Retarget descriptor: Unitree G1 -> Fourier GR1T2 pick-and-place (bimanual humanoids). +# +# The right arm picks the object and transports it; the left arm does a support motion. Self-contained: +# the per-EEF `subtasks` below segment the source demo and track the carried object during transport +# (replacing the old auto carry-detection). + +name: g1_to_gr1_pick_place +description: G1 -> GR1T2 humanoid picks an object with the right arm and transfers it. +write_datagen_info: true # write full datagen_info -> output is a drop-in generate_dataset source + +source_embodiment: ../embodiments/g1_ik_abs.yaml +target_embodiment: ../embodiments/gr1_ik_abs.yaml +target_env_name: Isaac-PickPlace-GR1T2-Abs-v0 + +hand_policy: binary +init_robot_from_ik: false +replay_speed: 2.0 +reference_pose: eef_pose +eef_reference_link: controlled +retarget_frame: world +scene_translation: [0.0, 0.0, 0.30] + +segment_settle_steps: 0 + +default_object_tracking: + interpolation_step_start: 15 + interpolation_step_after: 5 + +subtasks: + right: + - object_ref: world + description: Approach and grasp the object with the right hand. + subtask_end: + method: gripper_close + offset: 10 + + - frame_ref: world + description: Place the object in the bin. + subtask_end: + method: gripper_open + object_tracking: object + + - object_ref: world + description: Go back to iddle + + left: + - object_ref: object + description: Approach and grasp the object with the left hand. + subtask_end: + method: gripper_close + offset: 10 + + - frame_ref: world + description: Move the object next to the right end + subtask_end: + method: gripper_open + object_tracking: object + + - object_ref: world + description: Go back to iddle \ No newline at end of file diff --git a/autodata_examples/retarget/gr1_to_g1_pick_place.yaml b/autodata_examples/retarget/gr1_to_g1_pick_place.yaml new file mode 100644 index 0000000..31b1b5b --- /dev/null +++ b/autodata_examples/retarget/gr1_to_g1_pick_place.yaml @@ -0,0 +1,69 @@ +# Copyright (c) 2026, The Isaac AutoData Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +# Retarget descriptor: Fourier GR1T2 -> Unitree G1 pick-and-place (bimanual humanoids). +# +# The right arm picks the object and transports it; the left arm does a support motion. Self-contained: +# the per-EEF `subtasks` below segment the source demo and track the carried object during transport +# (replacing the old auto carry-detection). + +name: gr1_to_g1_pick_place +description: GR1T2 -> G1 humanoid picks an object with the right arm and transfers it. +write_datagen_info: true # write full datagen_info -> output is a drop-in generate_dataset source + +source_embodiment: ../embodiments/gr1_ik_abs.yaml +target_embodiment: ../embodiments/g1_ik_abs.yaml +target_env_name: Isaac-PickPlace-Locomanipulation-G1-Abs-v0 + +hand_policy: binary +# The G1 target's locomanipulation action has a body channel; hold it at this default during replay. +target_channel_defaults: + body: [0.0, 0.0, 0.0, 0.72] +init_robot_from_ik: false +replay_speed: 0.5 +reference_pose: target_eef_pose +eef_reference_link: controlled +retarget_frame: world +scene_translation: [0.0, 0.0, -0.30] + +segment_settle_steps: 0 + +default_object_tracking: + interpolation_step_start: 15 + interpolation_step_after: 5 + +subtasks: + right: + - object_ref: world + description: Approach and grasp the object with the right hand. + subtask_end: + method: gripper_close + offset: 10 + + - frame_ref: world + description: Place the object in the bin. + subtask_end: + method: gripper_open + object_tracking: object + + - object_ref: world + description: Go back to iddle + + left: + - object_ref: object + description: Approach and grasp the object with the left hand. + subtask_end: + method: gripper_close + offset: 10 + + - frame_ref: world + description: Move the object next to the right end + subtask_end: + method: gripper_open + object_tracking: object + + - object_ref: world + description: Go back to iddle + diff --git a/autodata_retargeting/README.md b/autodata_retargeting/README.md new file mode 100644 index 0000000..2041b41 --- /dev/null +++ b/autodata_retargeting/README.md @@ -0,0 +1,209 @@ +# Cross-embodiment dataset retargeting + +Replay a demonstration dataset recorded on one robot (the **source** embodiment) onto a different +robot (the **target** embodiment), one output demo per input demo. The source end-effector (EEF) +trajectory is transferred verbatim onto the target robot and the resulting rollout is recorded — no +object-centric regeneration happens, so the motion is copied 1:1 rather than re-planned. + +Entry point: [`scripts/retarget_dataset.py`](../scripts/retarget_dataset.py). A run is driven by a +single **retarget descriptor** YAML (`--retarget_config`) that references two **embodiment +descriptor** YAMLs and the target env. Working examples for the Unitree **G1 ⇄** Fourier **GR1T2** +humanoid pair live in [`autodata_examples/`](../autodata_examples). + +> All commands must run inside the project Docker container (`./docker/run_docker.sh`). + +--- + +## 1. What it does + +Given a source dataset and a source→target embodiment pair, retargeting: + +1. reads each source episode's recorded **absolute EEF pose trajectory** + (`obs/datagen_info/target_eef_pose`), which is embodiment-independent (env-relative SE(3)); +2. resets the target env so the **task scene** (object poses) matches the source episode's initial + state, while the **target robot** starts from its own default configuration; +3. re-encodes each step's EEF pose (plus gripper/hand commands) into the target embodiment's action + and steps it through the target env's IK controller; +4. records the target rollout (actions + states, and optionally full `datagen_info`) as one output + episode, marked success/failure by the target env's success term. + +Because the trajectory is copied rather than regenerated, the source dataset **must already carry +`datagen_info`** — i.e. it was produced by `generate_dataset.py` or passed through +`annotate_demos.py`. + +## 2. How it works + +Per source episode, per EEF: + +- **Reference trajectory.** The commanded/observed EEF path is extracted in a canonical *grasp + frame*: `source_pose @ source_offset`, then re-anchored to the target's control link via + `@ inv(target_offset)`. The per-EEF `offset` in each embodiment descriptor is the SE(3) transform + from the tracked control link to that canonical grasp frame, so different wrists/grippers line up + at the grasp point (handles gripper length and wrist convention). Which source signal drives the + replay is set by `reference_pose` (observed `eef_pose` vs commanded `target_eef_pose`). + +- **Segmentation into subtasks.** The descriptor's per-EEF `subtasks` list splits each demo into + segments. Each subtask declares how it ends (`subtask_end`) from an event detected on the *source* + demo — the gripper crossing open/closed, a subtask-term signal edge, or a fixed frame count. + Segmentation drives object tracking, offsets, boundary settling, and (optionally) cross-arm + synchronization. + +- **Hand / gripper policy.** Gripper and other non-EEF channels (e.g. a locomotion `body` channel) + are extracted from the source actions and remapped to the target's layout. `hand_policy` chooses + how the source hand drives the target hand: `passthrough` (copy raw), `binary` (open/close by a + closedness threshold), `interpolation` (scale between the target's `hand_open`/`hand_close` + configs), or `joint_mapping` (wire target joints to source joints). + +- **Object tracking.** When a subtask sets `object_tracking`, the EEF tracks the named object's pose + (a rigid `eef_T_object`) over that segment instead of following the source EEF path — this keeps a + carried object on its recorded world path even if the target's grasp differs slightly. The command + eases between the EEF path and the object-centric pose over `interpolation_step_start` / + `interpolation_step_after` steps. + +- **Per-subtask offsets.** An optional `offset` on a subtask applies an SE(3) nudge to that EEF's + commanded pose across the subtask's span, in a chosen frame (world / eef / the tracked object). + Offsets form a continuous timeline: adjacent subtask offsets blend into each other, and the last + subtask **holds** its offset (it is not decayed away), so an offset applied through the final + placement stays applied through the success settle. Use it to correct a systematic grasp/place bias + between the two embodiments. + +- **Boundary settling.** At each subtask boundary the replayer holds the pose until the robot's + joints (arm and gripper) stop moving, up to `segment_settle_steps` sim steps — motion-aware, so a + fast gripper exits in ~1–2 steps while a slow linkage gets time to finish. A per-boundary + `settle_steps` on `subtask_end` overrides the cap. + +- **Success & recording.** After the trajectory completes, the pose is held for + `success_settle_steps` while the target env's success term is re-checked. Episodes that satisfy it + are exported; failures are dropped unless `--keep_failed` is set. + +## 3. Configuration + +A run is defined by two kinds of YAML file, both under +[`autodata_examples/`](../autodata_examples). + +### 3.1 Embodiment descriptor (`autodata_examples/embodiments/*.yaml`) + +Describes one robot's action layout so retargeting can read/write its EEF poses, hands, and +passthrough channels. Key sections (see [`g1_ik_abs.yaml`](../autodata_examples/embodiments/g1_ik_abs.yaml) +and [`gr1_ik_abs.yaml`](../autodata_examples/embodiments/gr1_ik_abs.yaml)): + +| Field | Meaning | +| --- | --- | +| `type`, `name`, `description` | Adapter kind (e.g. `absolute_pose_whole_body_bimanual`) and labels. | +| `eefs..pose_obs_keys` | Obs keys for this EEF's position / quaternion. | +| `eefs..gripper_action_indices` | Which hand-joint action indices this EEF owns. | +| `eefs..offset` | SE(3) transform (`axis_angle` [rad], `translation` [m]) from the tracked control link to the canonical grasp frame — how this robot's wrist aligns to the grasp point. | +| `action_layout` | Pose slices per EEF, hand-joint slice, `canonicalize_quat`, and any `passthrough_channels` (name → contiguous action slice, e.g. a `body` locomotion command). | +| `hand_open` / `hand_close` | Per-EEF fully-open / fully-closed hand joint configs, used to score source closedness and drive `binary` / `interpolation` hand policies. | + +Both the **single-arm** (e.g. Franka) and **bimanual/whole-body** (G1, GR1) adapters are supported; +bimanual descriptors simply list two `eefs`. + +### 3.2 Retarget descriptor (`autodata_examples/retarget/*.yaml`) + +Bundles everything task/pair-specific into one file. Required references plus the most common knobs +(full schema in [`config.py`](config.py) — `RetargetConfig`): + +| Field | Meaning | +| --- | --- | +| `source_embodiment` / `target_embodiment` | Paths (relative to this YAML) to the embodiment descriptors. | +| `target_env_name` | Gym env id to instantiate for the target robot. | +| `subtasks` | Per-EEF list of segments (see below). | +| `hand_policy` | `passthrough` / `binary` / `interpolation` / `joint_mapping`. | +| `reference_pose` | `eef_pose` (source *achieved* path, default) or `target_eef_pose` (commanded). | +| `eef_reference_link` | `controlled` rebuilds the reference at the IK-controlled link (needed when a robot observes the EEF a joint short of the link it drives, e.g. GR1). | +| `replay_speed` | Retime the source trajectory (>1 faster, <1 slower). | +| `retarget_frame`, `scene_translation` | Frame the trajectory is anchored in, and a world offset applied to the scene. | +| `target_channel_defaults` | Hold a target-only passthrough channel (e.g. G1's `body`) at a fixed value during replay. | +| `default_object_tracking` | Default `interpolation_step_start` / `interpolation_step_after` for object tracking. | +| `segment_settle_steps`, `success_settle_steps` | Boundary settle cap and post-success hold. | +| `write_datagen_info` | Write full `obs/datagen_info` so the output is a drop-in `generate_dataset.py` source. | +| `synchronization` | Barrier groups of subtask `name`s across EEFs that must conclude together (bimanual joins). | + +**Subtasks.** Each EEF under `subtasks` is a list of segments in execution order. A segment has: + +- `object_ref` / `frame_ref` — the object or frame it is planned relative to (also the default frame + for its `offset`); +- `subtask_end` — how it ends: `method` is `gripper_open`/`gripper_close` + (or `gripper_opening`/`gripper_closing` for the leading edge), `signal_on`/`signal_off`, or + `fixed_length` (with `length`); `offset` shifts the boundary ± frames; `settle_steps` overrides the + settle cap here. The **last** subtask of an EEF needs no `subtask_end` (it runs to the end); +- `object_tracking` — bare object name or a section, to track that object over the segment; +- `offset` — optional per-subtask SE(3) nudge (`frame`, `translation`, `axis_angle`, + `interpolation_start`/`interpolation_end`); +- `name` — optional stable id (required only if referenced in `synchronization`). + +Minimal example (right arm picks, then places, then returns): + +```yaml +subtasks: + right: + - object_ref: world + description: Approach and grasp the object with the right hand. + subtask_end: { method: gripper_close, offset: 10 } + - frame_ref: world + description: Place the object in the bin. + subtask_end: { method: gripper_open } + object_tracking: object # track the carried object over the transport + - object_ref: world + description: Return to idle. # last subtask: no subtask_end +``` + +Adding a per-subtask offset (nudge the place 2 cm higher in world frame, eased in over 15 steps): + +```yaml + - frame_ref: world + subtask_end: { method: gripper_open } + object_tracking: object + offset: + frame: world + translation: [0.0, 0.0, 0.02] + interpolation_start: 15 +``` + +## 4. Running + +Config-based invocation (all task/pair parameters come from the descriptor): + +```bash +python scripts/retarget_dataset.py \ + --retarget_config autodata_examples/retarget/g1_to_gr1_pick_place.yaml \ + --input_file ./datasets/.hdf5 \ + --output_file ./datasets/retargeted_g1_to_gr1.hdf5 \ + --headless +``` + +The reverse direction just swaps the descriptor: + +```bash +python scripts/retarget_dataset.py \ + --retarget_config autodata_examples/retarget/gr1_to_g1_pick_place.yaml \ + --input_file ./datasets/.hdf5 \ + --output_file ./datasets/retargeted_gr1_to_g1.hdf5 \ + --headless +``` + +### Runtime / IO flags (CLI only) + +| Flag | Default | Meaning | +| --- | --- | --- | +| `--retarget_config` | — | Retarget descriptor YAML (recommended path). | +| `--input_file` | *(required)* | Source dataset HDF5 (must carry `datagen_info`). | +| `--output_file` | `./datasets/retargeted_dataset.hdf5` | Destination HDF5. | +| `--num_envs` | `1` | `>1` replays episodes across that many parallel envs on one sim for a large speedup (the per-step tracking report is single-env only). | +| `--select_episodes IDX ...` | all | Retarget only these source-episode indices. | +| `--target_successes N` | — | Stop once `N` replays succeed. | +| `--target_runs N` | — | Stop once `N` replays are attempted. | +| `--keep_failed` | off | Also export episodes that did not satisfy the success term. | +| `--device`, `--headless`, … | — | Standard Isaac Lab `AppLauncher` flags. | + +Without `--retarget_config`, only `--source_embodiment` / `--target_embodiment` / +`--target_env_name` are read and every other knob (including `subtasks`) takes its default, so a +descriptor is needed for any real run. + +## 5. Output + +The retargeted HDF5 records the target robot's actions and states. When the descriptor sets +`write_datagen_info: true`, the output also carries the full `obs/datagen_info` (observed `eef_pose`, +commanded `target_eef_pose`, per-object `object_pose`) and the forwarded `subtask_term_signals`, +making it a drop-in source for `generate_dataset.py` without a separate `annotate_demos.py` pass. diff --git a/autodata_retargeting/__init__.py b/autodata_retargeting/__init__.py new file mode 100644 index 0000000..034e33f --- /dev/null +++ b/autodata_retargeting/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2026, The Isaac AutoData Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Cross-embodiment dataset retargeting. Import submodules after launching the Isaac app.""" diff --git a/autodata_retargeting/config.py b/autodata_retargeting/config.py new file mode 100644 index 0000000..c03705c --- /dev/null +++ b/autodata_retargeting/config.py @@ -0,0 +1,453 @@ +# Copyright (c) 2026, The Isaac AutoData Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Task/pair retargeting configuration (the retarget-descriptor YAML schema). + +The descriptor is now self-contained: it carries its own task structure (a per-EEF list of +:class:`Subtask`) instead of referencing an external task descriptor. Each subtask declares how it ends +(:class:`SubtaskEnd` -- a gripper or signal event, offsettable) and, optionally, which object the EEF +tracks over it (:class:`SubtaskObjectTracking`), so segmentation and object-centric planning come +straight from this file. Retargeting replays each recorded source trajectory 1:1 onto the target robot. +""" + +import yaml +from dataclasses import dataclass, field, fields +from pathlib import Path +from typing import Any + + +@dataclass +class DefaultObjectTracking: + """Default object-tracking interpolation parameters, overridable per :class:`SubtaskObjectTracking`. + + When a subtask tracks an object, the commanded pose eases between the source EEF path and the + object-centric pose: over ``interpolation_step_start`` steps at the subtask's start (EEF path -> + object) and ``interpolation_step_after`` steps past its end (object -> back onto the EEF path). + 0 = switch immediately. + """ + + interpolation_step_start: int = 0 + interpolation_step_after: int = 0 + + @classmethod + def parse(cls, value: "dict | DefaultObjectTracking | None") -> "DefaultObjectTracking": + """Normalize object-tracking defaults from a descriptor value. + + Args: + value: Existing defaults, a mapping of field values, or ``None`` for defaults. + + Returns: + Parsed object-tracking defaults. + """ + if value is None: + return cls() + if isinstance(value, cls): + return value + assert isinstance(value, dict), f"default_object_tracking must be a section (dict), got {type(value).__name__}." + allowed = {f.name for f in fields(cls)} + unknown = set(value) - allowed + assert not unknown, f"unknown default_object_tracking keys {sorted(unknown)}; allowed: {sorted(allowed)}." + return cls(**value) + + +@dataclass +class SubtaskEnd: + """When a subtask ends: an event detected on the source demo, optionally offset. + + ``method``: + * ``gripper_close`` / ``gripper_open`` -- the step ``eef``'s gripper crosses closed / open. + * ``gripper_closing`` / ``gripper_opening`` -- same crossing, but backed up to the frame the gripper + *begins* to move (the leading edge). For a handoff, end the releaser at ``gripper_opening`` so it + holds the object until the barrier releases, instead of the fragile ``gripper_open`` + ``offset: -1``. + * ``signal_on`` / ``signal_off`` -- the rising / falling edge of subtask-term signal ``signal``. + * ``fixed_length`` -- a fixed ``length`` frames after the subtask's start (no source event). + + ``eef`` selects which gripper the event watches (default: the EEF this subtask is listed under). + ``length`` (``fixed_length`` only) is the subtask duration in source frames. ``offset`` shifts the + boundary: negative ends the subtask that many steps *before* the trigger, positive prolongs it that + many steps *after*. ``settle_steps`` overrides the run-wide ``segment_settle_steps`` cap for the + motion-aware hold at *this* boundary (None = use the config default; 0 = no hold here). + """ + + method: str + signal: str | None = None + eef: str | None = None + offset: int = 0 + settle_steps: int | None = None + length: int | None = None + + _SIGNAL_METHODS = ("signal_on", "signal_off") + _GRIPPER_METHODS = ("gripper_open", "gripper_close", "gripper_opening", "gripper_closing") + _LENGTH_METHODS = ("fixed_length",) + + def __post_init__(self) -> None: + """Validate fields that depend on the selected end method.""" + allowed = self._SIGNAL_METHODS + self._GRIPPER_METHODS + self._LENGTH_METHODS + assert self.method in allowed, f"subtask_end method must be one of {allowed}, got {self.method!r}." + if self.method in self._SIGNAL_METHODS: + assert self.signal, f"subtask_end method {self.method!r} needs a 'signal' name." + if self.method in self._LENGTH_METHODS: + assert ( + self.length is not None and self.length > 0 + ), f"subtask_end method {self.method!r} needs a positive 'length' (frames), got {self.length!r}." + + @classmethod + def parse(cls, value: "str | dict | SubtaskEnd") -> "SubtaskEnd": + """Normalize a subtask-end descriptor. + + Args: + value: Existing descriptor, method name, or mapping of field values. + + Returns: + Parsed subtask-end descriptor. + """ + if isinstance(value, cls): + return value + if isinstance(value, str): + return cls(method=value) + assert isinstance(value, dict), f"subtask_end must be a string or a section (dict), got {type(value).__name__}." + allowed = {f.name for f in fields(cls)} + unknown = set(value) - allowed + assert not unknown, f"unknown subtask_end keys {sorted(unknown)}; allowed: {sorted(allowed)}." + return cls(**value) + + +@dataclass +class SubtaskObjectTracking: + """Which object an EEF tracks over a subtask (rigid ``eef_T_object``), with optional interp overrides. + + Written as the bare object name (``object_tracking: cube_2``) or a section + (``{object: cube_2, interpolation_step_start: 10}``). ``interpolation_step_*`` default to + :class:`DefaultObjectTracking` when unset. + """ + + object: str + interpolation_step_start: int | None = None + interpolation_step_after: int | None = None + + @classmethod + def parse(cls, value: "str | dict | SubtaskObjectTracking | None") -> "SubtaskObjectTracking | None": + """Normalize an optional object-tracking descriptor. + + Args: + value: Existing descriptor, object name, mapping of field values, or ``None``. + + Returns: + Parsed object-tracking descriptor, or ``None`` when tracking is disabled. + """ + if value is None: + return None + if isinstance(value, cls): + return value + if isinstance(value, str): + return cls(object=value) + assert isinstance( + value, dict + ), f"object_tracking must be a string or a section (dict), got {type(value).__name__}." + allowed = {f.name for f in fields(cls)} + unknown = set(value) - allowed + assert not unknown, f"unknown object_tracking keys {sorted(unknown)}; allowed: {sorted(allowed)}." + assert "object" in value, "object_tracking section needs an 'object' field." + return cls(**value) + + +@dataclass +class Offset: + """A per-subtask SE(3) offset applied to the commanded EEF pose over the subtask's span. + + ``frame`` selects the frame the offset acts in; it defaults to the subtask's own reference frame + (``object_ref`` -> that tracked object's frame, else ``frame_ref``; ``world``/``eef`` are always + available -- ``eef`` nudges in the gripper's own frame). ``translation`` [m] and ``axis_angle`` + [rad] give the offset transform (either may be omitted -> that part is identity). + ``interpolation_start`` / ``interpolation_end`` ramp the offset in linearly over the first N steps of + the span and out over the last N steps, so it can be eased in/out (0 = applied fully across the span). + """ + + frame: str | None = None + translation: list | None = None + axis_angle: list | None = None + interpolation_start: int = 0 + interpolation_end: int = 0 + + def __post_init__(self) -> None: + """Normalize and validate the offset vectors and interpolation lengths.""" + for name in ("translation", "axis_angle"): + value = getattr(self, name) + if value is not None: + assert ( + isinstance(value, list | tuple) and len(value) == 3 + ), f"offset.{name} must be a 3-vector [x, y, z], got {value!r}." + setattr(self, name, [float(v) for v in value]) + self.interpolation_start = int(self.interpolation_start) + self.interpolation_end = int(self.interpolation_end) + assert ( + self.interpolation_start >= 0 and self.interpolation_end >= 0 + ), f"offset interpolation_start/end must be >= 0, got {self.interpolation_start}/{self.interpolation_end}." + + @classmethod + def parse(cls, value: Any) -> "Offset | None": + """Normalize an optional offset descriptor. + + Args: + value: Mapping of offset field values, or ``None``. + + Returns: + Parsed offset descriptor, or ``None`` when no offset is configured. + """ + if value is None: + return None + assert isinstance(value, dict), f"offset must be a section (dict), got {type(value).__name__}." + allowed = {f.name for f in fields(cls)} + unknown = set(value) - allowed + assert not unknown, f"unknown offset keys {sorted(unknown)}; allowed: {sorted(allowed)}." + return cls(**value) + + +@dataclass +class Subtask: + """One subtask in an EEF's sequence. + + ``name`` (optional) is a stable identifier for the subtask; when given it must be unique across all + subtasks of all EEFs. ``object_ref`` / ``frame_ref`` name the object or frame the subtask is planned + relative to; they default the reference frame a per-subtask :class:`Offset` acts in. ``subtask_end`` is + required except on the last subtask of an EEF (which runs to the end of the trajectory). + ``object_tracking`` (optional) makes the EEF track an object over the whole subtask instead of + following the source EEF path. + """ + + name: str | None = None + object_ref: str | None = None + frame_ref: str | None = None + description: str = "" + subtask_end: SubtaskEnd | None = None + object_tracking: SubtaskObjectTracking | None = None + # Per-subtask SE(3) pose offset applied to this EEF's commanded trajectory over the subtask's span + # (see :class:`Offset`). None = no offset. + offset: "Offset | None" = None + + @classmethod + def parse(cls, value: dict) -> "Subtask": + """Parse one subtask and its nested descriptors. + + Args: + value: Mapping of subtask field values. + + Returns: + Parsed subtask descriptor. + """ + assert isinstance(value, dict), f"each subtask must be a section (dict), got {type(value).__name__}." + data = dict(value) + allowed = {f.name for f in fields(cls)} + unknown = set(data) - allowed + assert not unknown, f"unknown subtask keys {sorted(unknown)}; allowed: {sorted(allowed)}." + if data.get("subtask_end") is not None: + data["subtask_end"] = SubtaskEnd.parse(data["subtask_end"]) + data["object_tracking"] = SubtaskObjectTracking.parse(data.get("object_tracking")) + data["offset"] = Offset.parse(data.get("offset")) + return cls(**data) + + +@dataclass +class RetargetConfig: + """Self-contained task/pair retargeting parameters, loaded from a retarget-descriptor YAML. + + Bundles what defines a source->target retarget for one task (both embodiments, the target env, the + per-EEF ``subtasks``) plus every retargeting knob, so a run is one ``--retarget_config`` file instead + of a long flag list. Runtime/IO (input/output files, keep_failed, diagnostics, device) stay on the CLI. + """ + + source_embodiment: str + target_embodiment: str + target_env_name: str + name: str = "retarget" + description: str = "" + hand_policy: str = "passthrough" + hand_interp_norm: str = "l1" + hand_binary_close_threshold: float = 0.5 + # For hand_policy "interpolation": remap the source closedness fraction so this ``[lo, hi]`` band + # spans the target's full open->close travel, clamping outside it. Lets a source grasp with a dead + # band near its own open/close still drive the target hand fully. None disables (identity). + hand_interp_band: tuple[float, float] | None = None + # For hand_policy "joint_mapping": ``{target_joint_idx: [source_joint_idx, ...]}`` wiring each target + # hand joint to the average of a group of source hand joints (indices into each embodiment's per-EEF + # hand vector, i.e. gripper_action_indices order), mapped through open/close. Target joints not listed + # stay at hand_open. + joint_mapping: dict[int, list[int]] | None = None + num_interpolation_steps: int = 0 + init_robot_from_ik: bool = False + replay_speed: float = 1.0 + # Cap the commanded EEF speed by subdividing the trajectory wherever a step moves + # faster than the limit (extra waypoints are interpolated in). Keeps a reactive controller (e.g. Galbot's + # RmpFlow) from being handed a command it cannot physically follow -- a differential-IK robot snaps to + # each waypoint in one step and is unaffected, so leave these ``None`` there. Measured per replayed + # waypoint on the grasp-frame path: ``max_eef_linear_velocity`` [m/step], ``max_eef_rotation_speed`` + # [deg/step]. ``None`` disables that axis' cap. + max_eef_linear_velocity: float | None = None + max_eef_rotation_speed: float | None = None + # Early-abort a doomed replay to save sim time. When ``stop_early_on_failure`` is set, the replay stops + # the moment the monitored achieved pose misses its target beyond ``max_translation_error`` [m] or + # ``max_rotation_error`` [deg] -- the tracked OBJECT vs its source path during a carry, else the EEF vs + # its commanded target -- and the run is recorded as a failure (success never fired). ``None`` thresholds + # are not checked. Speeds up generation by not replaying past an unrecoverable miss. + stop_early_on_failure: bool = False + max_translation_error: float | None = None + max_rotation_error: float | None = None + # At each subtask boundary the replayer holds the pose until the robot's joints (arm AND gripper) stop + # moving, up to this many extra sim steps -- so a slow gripper (e.g. the Robotiq 2F-85 linkage) finishes + # closing/opening before the arm moves on. Motion-aware: a fast gripper exits in ~1-2 steps. 0 = no hold. + # A per-subtask ``SubtaskEnd.settle_steps`` overrides this cap for that boundary. + segment_settle_steps: int = 30 + # Tolerances for the motion-aware boundary hold -- it exits early once motion falls below all of these. + # ``settle_joint_tol`` (primary): max per-step joint position change [rad or m]; robust to the Robotiq + # linkage's saturated velocity reading. ``settle_pos_tol_m`` / ``settle_rot_tol_deg``: max per-step EEF + # pose change [m] / [deg] (the arm has reached the waypoint). + settle_joint_tol: float = 0.003 + settle_pos_tol_m: float = 0.01 + settle_rot_tol_deg: float = 2.5 + # Hold the final pose this many steps once the success term fires (re-checking success). From the old + # task descriptor's generation policy; re-homed here now that the descriptor is self-contained. + success_settle_steps: int = 0 + # Seed for the reproducible scene resets (was the task descriptor's generation-policy seed). + seed: int = 1 + retarget_frame: str = "world" + scene_translation: tuple[float, float, float] = (0.0, 0.0, 0.0) + eef_offsets: dict[str, dict] | None = None + target_channel_defaults: dict[str, list[float]] | None = None + eef_name_map: dict[str, str] | None = None + # Which source-demo trajectory drives the replay (and is scored against): + # "eef_pose" -- the source robot's *achieved* (executed, task-successful) path (default). + # "target_eef_pose" -- the *commanded* controller targets (the ideal the source may never have + # reached; a more-capable target robot can overshoot toward it). + reference_pose: str = "eef_pose" + # Default object-tracking interpolation, overridable per subtask (see DefaultObjectTracking). + default_object_tracking: "dict | DefaultObjectTracking" = field(default_factory=DefaultObjectTracking) + # Per-EEF subtask lists (keyed by *source* EEF name; eef_name_map renames onto the target). Each entry + # is parsed into a Subtask. Segmentation and object tracking come from these. + subtasks: dict[str, list[Subtask]] = field(default_factory=dict) + # Bimanual/multi-arm synchronization: a list of barrier groups, each a list of subtask *names* that must + # conclude together. Every EEF that reaches its named segment holds at that pose until all segments in + # the group have concluded (a join). Rules: a group's names belong to different EEFs; each subtask is in + # at most one group; and groups are declared in execution order -- each EEF's participating segments + # appear in its temporal subtask order (a simple check that also guarantees deadlock-freedom). Empty for + # single-arm / uncoordinated runs. + synchronization: list[list[str]] = field(default_factory=list) + # Write the full ``obs/datagen_info`` into the retargeted dataset so the output is a drop-in source + # for ``generate_dataset.py`` (no separate ``annotate_demos`` pass): observed ``eef_pose``, commanded + # ``target_eef_pose`` and per-object ``object_pose`` (read live each step, matching annotate_demos' + # recorder), plus the forwarded ``subtask_term_signals`` (resampled to the replay trajectory and held + # across each waypoint's frames -- lead-in, settle, replay_speed retiming -- so boundaries line up with + # the recorded rollout). Default off. + write_datagen_info: bool = False + # Reconstruct the "eef_pose" reference at each EEF's IK-controlled link instead of the source env's + # *observed* link. Some embodiments observe the EEF a joint short of the link the IK drives (e.g. GR1 + # observes ``hand_roll_link`` but controls ``hand_pitch_link``), so the observed ``eef_pose`` drops + # that joint's rotation. When set, the reference (and the object-tracking grasp frame) is rebuilt + # at the controlled link from the recorded per-step link states. Values: ``None`` (off, default), + # ``"controlled"`` to find the controlled link by matching the source's own ``target_eef_pose`` + # (needs no body names -> works cross-embodiment, e.g. GR1->G1), or a ``{eef_name: link_name}`` map as + # a same-embodiment override. Only affects ``reference_pose="eef_pose"``. + eef_reference_link: dict[str, str] | str | None = None + + def __post_init__(self) -> None: + """Parse nested descriptors and validate synchronization barriers.""" + self.default_object_tracking = DefaultObjectTracking.parse(self.default_object_tracking) + self.subtasks = { + eef: [st if isinstance(st, Subtask) else Subtask.parse(st) for st in entries] + for eef, entries in (self.subtasks or {}).items() + } + names = [st.name for entries in self.subtasks.values() for st in entries if st.name is not None] + duplicates = sorted({name for name in names if names.count(name) > 1}) + assert not duplicates, f"subtask names must be unique across all subtasks; duplicates: {duplicates}" + + by_name = { + st.name: (eef, i) + for eef, entries in self.subtasks.items() + for i, st in enumerate(entries) + if st.name is not None + } + # Synchronization barriers: names must exist; a group's members belong to different EEFs; each name + # is in at most one group; and the declaration order is a valid schedule -- for each EEF, the groups + # it joins (in declaration order) reference its subtasks in temporal order. That last check both + # keeps authoring intuitive (list groups as they happen) and guarantees the barriers are acyclic. + group_of: dict[str, int] = {} + eef_join_order: dict[str, list[tuple[int, int]]] = {} # eef -> [(group_index, subtask_index)] + for g_idx, group in enumerate(self.synchronization): + assert len(group) >= 2, f"synchronization group {g_idx} needs >= 2 subtask names, got {group}." + group_eefs: set[str] = set() + for nm in group: + assert nm in by_name, f"synchronization group {g_idx} references unknown subtask name {nm!r}." + assert nm not in group_of, ( + f"subtask {nm!r} is in synchronization groups {group_of[nm]} and {g_idx}; " + "each subtask may be in at most one group." + ) + group_of[nm] = g_idx + eef, s_i = by_name[nm] + assert eef not in group_eefs, ( + f"synchronization group {g_idx} lists two subtasks of EEF {eef!r}; a group's members " + "must belong to different EEFs." + ) + group_eefs.add(eef) + eef_join_order.setdefault(eef, []).append((g_idx, s_i)) + for eef, joins in eef_join_order.items(): + s_indices = [s_i for _, s_i in sorted(joins)] # subtask indices in group-declaration order + assert s_indices == sorted(s_indices) and len(set(s_indices)) == len(s_indices), ( + f"EEF {eef!r} joins synchronization groups out of temporal order; declare groups in the " + "order their subtasks execute." + ) + + @classmethod + def from_yaml(cls, path: str) -> "RetargetConfig": + """Load a retarget descriptor; the embodiment paths resolve relative to the YAML's dir.""" + config_dir = Path(path).resolve().parent + with open(path) as f: + data = yaml.safe_load(f) or {} + + def _resolve(ref: str) -> str: + ref_path = Path(ref) + return str(ref_path if ref_path.is_absolute() else (config_dir / ref_path).resolve()) + + return cls( + source_embodiment=_resolve(data["source_embodiment"]), + target_embodiment=_resolve(data["target_embodiment"]), + target_env_name=data["target_env_name"], + name=data.get("name", "retarget"), + description=data.get("description", ""), + hand_policy=data.get("hand_policy", "passthrough"), + hand_interp_norm=data.get("hand_interp_norm", "l1"), + hand_binary_close_threshold=float(data.get("hand_binary_close_threshold", 0.5)), + hand_interp_band=(tuple(data["hand_interp_band"]) if data.get("hand_interp_band") is not None else None), + joint_mapping=data.get("joint_mapping"), + num_interpolation_steps=int(data.get("num_interpolation_steps", 0)), + init_robot_from_ik=bool(data.get("init_robot_from_ik", False)), + replay_speed=float(data.get("replay_speed", 1.0)), + max_eef_linear_velocity=( + None if data.get("max_eef_linear_velocity") is None else float(data["max_eef_linear_velocity"]) + ), + max_eef_rotation_speed=( + None if data.get("max_eef_rotation_speed") is None else float(data["max_eef_rotation_speed"]) + ), + stop_early_on_failure=bool(data.get("stop_early_on_failure", False)), + max_translation_error=( + None if data.get("max_translation_error") is None else float(data["max_translation_error"]) + ), + max_rotation_error=(None if data.get("max_rotation_error") is None else float(data["max_rotation_error"])), + segment_settle_steps=int(data.get("segment_settle_steps", 30)), + settle_joint_tol=float(data.get("settle_joint_tol", 0.003)), + settle_pos_tol_m=float(data.get("settle_pos_tol_m", 0.01)), + settle_rot_tol_deg=float(data.get("settle_rot_tol_deg", 2.5)), + success_settle_steps=int(data.get("success_settle_steps", 0)), + seed=int(data.get("seed", 1)), + retarget_frame=data.get("retarget_frame", "world"), + scene_translation=tuple(data.get("scene_translation", (0.0, 0.0, 0.0))), + eef_offsets=data.get("eef_offsets"), + target_channel_defaults=data.get("target_channel_defaults"), + eef_name_map=data.get("eef_name_map"), + reference_pose=data.get("reference_pose", "eef_pose"), + default_object_tracking=data.get("default_object_tracking"), + subtasks=data.get("subtasks", {}), + synchronization=[list(group) for group in data.get("synchronization", [])], + write_datagen_info=bool(data.get("write_datagen_info", False)), + eef_reference_link=data.get("eef_reference_link"), + ) diff --git a/autodata_retargeting/eef_offset.py b/autodata_retargeting/eef_offset.py new file mode 100644 index 0000000..e69c6cd --- /dev/null +++ b/autodata_retargeting/eef_offset.py @@ -0,0 +1,85 @@ +# Copyright (c) 2026, The Isaac AutoData Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Per-EEF SE(3) offset (control link -> canonical grasp frame): build, load, compose.""" + +import torch +import yaml + +import isaaclab.utils.math as math_utils + + +def build_eef_offsets( + offset_data: dict[str, dict], eef_names: list[str], device: torch.device +) -> dict[str, torch.Tensor]: + """Build per-EEF ``(4, 4)`` SE(3) offsets from ``{eef: {axis_angle, translation}}``. + + Rotate-then-translate convention: the frame is first rotated by ``axis_angle`` (a compact + axis-angle vector, axis * angle [rad]) then translated by ``translation`` [m] *in that rotated + (canonical) frame* — so the translation reads as forward/lateral/up at the grasp, consistent + across embodiments. Missing entries default to identity. Post-multiplied onto the source pose. + """ + offsets: dict[str, torch.Tensor] = {} + for eef_name in eef_names: + entry = offset_data.get(eef_name, {}) or {} + rotation = torch.eye(3, dtype=torch.float32, device=device) + axis_angle = torch.tensor(entry.get("axis_angle", [0.0, 0.0, 0.0]), dtype=torch.float32, device=device) + angle = torch.linalg.norm(axis_angle) + if float(angle) > 1e-8: + axis = (axis_angle / angle).unsqueeze(0) + rotation = math_utils.matrix_from_quat(math_utils.quat_from_angle_axis(angle.unsqueeze(0), axis))[0] + translation = torch.tensor(entry.get("translation", [0.0, 0.0, 0.0]), dtype=torch.float32, device=device) + offset = torch.eye(4, dtype=torch.float32, device=device) + offset[:3, :3] = rotation + offset[:3, 3] = rotation @ translation # translation applied in the rotated (canonical) frame + offsets[eef_name] = offset + return offsets + + +def load_embodiment_eef_offset( + embodiment_yaml: str, eef_names: list[str], device: torch.device +) -> dict[str, torch.Tensor] | None: + """Load an embodiment's per-EEF ``offset`` (control link → canonical grasp frame) as ``(4, 4)``. + + Supports both embodiment schemas (see ``calibrate_offset.py``): + + * **Bimanual**: each EEF's ``offset`` (``axis_angle`` + ``translation``) lives under ``eefs.``. + * **Single-arm**: a top-level ``eef_offset`` (translation [m]) and ``eef_rotation`` (axis-angle + [rad]). A single-arm embodiment has exactly one EEF, so that one offset is applied to whichever + EEF name is requested (``eef_names`` comes from the target adapter, e.g. ``["ur10"]``). + + Returns None if the embodiment declares no offset (so composition falls back to identity). + """ + with open(embodiment_yaml) as f: + data = yaml.safe_load(f) or {} + eefs = data.get("eefs") or {} + if eefs: # bimanual: per-EEF offset under eefs..offset + offset_data = {name: cfg["offset"] for name, cfg in eefs.items() if isinstance(cfg, dict) and "offset" in cfg} + elif "eef_offset" in data or "eef_rotation" in data: # single-arm: top-level eef_offset/eef_rotation + entry = { + "translation": data.get("eef_offset", [0.0, 0.0, 0.0]), + "axis_angle": data.get("eef_rotation", [0.0, 0.0, 0.0]), + } + # One-EEF embodiment: apply its single offset to whichever EEF name(s) the retarget requests. + offset_data = {name: entry for name in eef_names} + else: + offset_data = {} + return build_eef_offsets(offset_data, eef_names, device) if offset_data else None + + +def compose_retarget_eef_offsets( + source_embodiment: str, target_embodiment: str, eef_names: list[str], device: torch.device +) -> dict[str, torch.Tensor] | None: + """Per-EEF ``T_src @ inv(T_tgt)`` from the two embodiments' ``eef_offset`` declarations, or None. + + Expresses the source trajectory in the canonical grasp frame (``@ T_src``) and re-anchors to the + target's control link (``@ inv(T_tgt)``), so grippers with different geometry/convention align at + the grasp point. Post-multiplied onto the source pose, exactly like a manual ``eef_offsets``. + """ + source_offset = load_embodiment_eef_offset(source_embodiment, eef_names, device) + target_offset = load_embodiment_eef_offset(target_embodiment, eef_names, device) + if source_offset is None or target_offset is None: + return None + return {eef: source_offset[eef] @ torch.linalg.inv(target_offset[eef]) for eef in eef_names} diff --git a/autodata_retargeting/gripper_retargeting.py b/autodata_retargeting/gripper_retargeting.py new file mode 100644 index 0000000..979ead9 --- /dev/null +++ b/autodata_retargeting/gripper_retargeting.py @@ -0,0 +1,299 @@ +# Copyright (c) 2026, The Isaac AutoData Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Gripper / hand passthrough remapping (binary and interpolation hand policies).""" + +import torch +import yaml +from collections.abc import Callable + +from autodata_interfaces.embodiments.embodiment_adapter import EmbodimentAdapter + +# Passthrough channel layout: maps each passthrough channel name to its per-step width. +PassthroughLayout = dict[str, int] +# A remapper turns the source passthrough dict ({channel: (num_steps, width)}) into the target's. +PassthroughRemapper = Callable[[dict[str, torch.Tensor]], dict[str, torch.Tensor]] + + +def _passthrough_layout(adapter: EmbodimentAdapter) -> PassthroughLayout: + """Return the ``{channel_name: width}`` passthrough layout an adapter produces/consumes. + + Probes the adapter with a zero action of its own ``action_dim`` so the layout is derived the + same way for every morphology (single-arm gripper, bimanual hands + extra channels). + """ + assert hasattr(adapter, "action_dim"), f"{type(adapter).__name__} does not expose action_dim" + probe = torch.zeros(1, adapter.action_dim) + channels = adapter.actions_to_passthrough_actions(probe) + return {name: tensor.shape[-1] for name, tensor in channels.items()} + + +def _hand_close_fraction( + source_hand: torch.Tensor, + src_open: torch.Tensor, + src_close: torch.Tensor, + norm: str = "l1", +) -> torch.Tensor: + """Per-step closedness fraction in ``[0, 1]`` for a source hand trajectory ``(T, dim)``. + + Scores how far the hand sits from its open toward its closed posture as the distance ratio + ``d_open / (d_open + d_close)`` under the L1 (``norm="l1"``) or L2 (``norm="l2"``) norm: 0 when the + hand matches ``src_open``, 1 when it matches ``src_close``. This turns a gradual source grasp into a + smooth percentage (unlike a binary open/closed classification), so the target hand can interpolate + instead of snapping shut. + """ + p = 1 if norm == "l1" else 2 + d_open = (source_hand - src_open).norm(p=p, dim=-1) + d_close = (source_hand - src_close).norm(p=p, dim=-1) + return (d_open / (d_open + d_close + 1e-6)).clamp(0.0, 1.0) + + +def _apply_interp_band(alpha: torch.Tensor, band: tuple[float, float] | None) -> torch.Tensor: + """Linearly remap a closedness fraction so ``[lo, hi]`` spans ``[0, 1]``, clamped outside the band. + + With ``band=(0.1, 0.9)``: ``alpha <= 0.1 -> 0`` (fully open), ``alpha >= 0.9 -> 1`` (fully closed), + and the middle stretches linearly across the target's full open->close travel. Lets a source grasp + that never quite reaches its own extremes (dead band near open/close) still drive the target hand all + the way. ``None`` is the identity. + """ + if band is None: + return alpha + lo, hi = band + return ((alpha - lo) / (hi - lo)).clamp(0.0, 1.0) + + +def _apply_joint_mapping( + source_hand: torch.Tensor, + src_open: torch.Tensor, + src_close: torch.Tensor, + tgt_open: torch.Tensor, + tgt_close: torch.Tensor, + joint_mapping: dict[int, list[int]], +) -> torch.Tensor: + """Drive each target hand joint from a chosen group of source joints, via a linear open->close map. + + For target joint ``X`` with sources ``[Y, Z, ...]`` (``joint_mapping[X]``): take the average source + position over those joints, score its closedness against the averaged source ``open``/``close`` + (``0`` = open, ``1`` = closed, clamped), and lerp target joint ``X`` between its own ``open`` and + ``close`` by that fraction. Calibrating through each side's open/close makes it robust to differing + joint ranges and to the mirrored left/right hands (a raw copy would not). Target joints absent from + ``joint_mapping`` are left at ``tgt_open``. + """ + num_steps = source_hand.shape[0] + out = tgt_open.unsqueeze(0).repeat(num_steps, 1).clone() # (T, tgt_width); unmapped joints stay open + for target_index, source_indices in joint_mapping.items(): + src_idx = torch.tensor(source_indices, dtype=torch.long, device=source_hand.device) + avg = source_hand[:, src_idx].mean(dim=1) # (T,) averaged source position + avg_open = src_open[src_idx].mean() + avg_close = src_close[src_idx].mean() + alpha = ((avg - avg_open) / (avg_close - avg_open + 1e-6)).clamp(0.0, 1.0) # (T,) closedness + out[:, target_index] = tgt_open[target_index] + alpha * (tgt_close[target_index] - tgt_open[target_index]) + return out + + +def load_hand_postures(embodiment_yaml: str, eef_names: list[str]) -> dict[str, dict[str, list[float]]] | None: + """Read per-EEF ``{eef: {open, close}}`` hand joint configs from an embodiment YAML, or None. + + Reads the ``hand_open`` / ``hand_close`` sections (each ``{eef: [joint values]}``). Returns None + if either is absent, so retargeting can require them only for the ``binary``/``interpolation`` + hand policies. + + Bimanual YAMLs (``eefs:`` block) key the postures by their own EEF names (``left``/``right``, shared + across the pair). A single-arm YAML has exactly one gripper, so its single posture is applied to + whichever EEF name(s) the retarget requests (``eef_names`` comes from the target adapter) — this + lets e.g. a Franka source (``eef_name: franka``) map onto a UR10 target (``eef_name: ur10``). + """ + with open(embodiment_yaml) as f: + data = yaml.safe_load(f) or {} + hand_open = data.get("hand_open") + hand_close = data.get("hand_close") + if not hand_open or not hand_close: + return None + if "eefs" not in data: # single-arm: one posture, applied to the requested eef name(s) + (open_vals,) = hand_open.values() + (close_vals,) = hand_close.values() + return {eef: {"open": open_vals, "close": close_vals} for eef in eef_names} + return {eef: {"open": hand_open[eef], "close": hand_close[eef]} for eef in eef_names} + + +def _fill_non_eef_channel( + channel_name: str, + width: int, + source_passthrough: dict[str, torch.Tensor], + target_channel_defaults: dict[str, list[float]] | None, +) -> torch.Tensor: + """Value for a non-EEF target channel: the source's when present, else a configured default. + + Handles a target-only channel (e.g. G1's locomotion ``body`` command absent from a fixed-base + GR1 source) by holding the ``target_channel_defaults`` value across the demo. Errors if the + source lacks it and no default is set (zero-filling a locomotion command could drop the robot). + """ + if channel_name in source_passthrough and source_passthrough[channel_name].shape[-1] == width: + return source_passthrough[channel_name] + assert target_channel_defaults and channel_name in target_channel_defaults, ( + f"target passthrough channel {channel_name!r} is not provided by the source embodiment; " + "set its value under target_channel_defaults in the retarget config." + ) + reference = next(iter(source_passthrough.values())) + value = torch.tensor(target_channel_defaults[channel_name], dtype=reference.dtype, device=reference.device) + assert len(value) == width, f"target_channel_defaults[{channel_name!r}] must have {width} values" + return value.unsqueeze(0).expand(reference.shape[0], width).contiguous() + + +def build_passthrough_remapper( + source_adapter: EmbodimentAdapter, + target_adapter: EmbodimentAdapter, + hand_policy: str = "passthrough", + hand_interp_norm: str = "l1", + hand_binary_close_threshold: float = 0.5, + hand_interp_band: tuple[float, float] | None = None, + joint_mapping: dict[int, list[int]] | None = None, + source_hand_postures: dict[str, dict[str, list[float]]] | None = None, + target_hand_postures: dict[str, dict[str, list[float]]] | None = None, + target_channel_defaults: dict[str, list[float]] | None = None, + eef_name_map: dict[str, str] | None = None, +) -> PassthroughRemapper: + """Build a function that maps source passthrough (gripper/hand) actions to the target's layout. + + Policies: + + * ``"passthrough"`` — the two embodiments must share a passthrough layout (matching per-channel + widths); the source channels are copied verbatim. Channel *names* may differ if bridged by + ``eef_name_map`` (e.g. a Franka 1-D gripper ``franka`` → a UR10 1-D gripper ``ur10`` both copy + through under the renamed channel). Differing *widths* (e.g. a Franka parallel gripper → a GR1 + dexterous hand) are rejected here — use ``binary``/``interpolation`` or implement the mapping. + * ``"binary"`` — per step, score the source hand's closedness fraction (0 = its ``hand_open``, + 1 = its ``hand_close``; same distance ratio as ``interpolation``, under ``hand_interp_norm``) and + snap the target hand to ``hand_close`` once that fraction reaches ``hand_binary_close_threshold``, + else ``hand_open``. The threshold sets *how far* the source must close before the target grips + (e.g. ``0.25`` = switch after 25% closed); ``0.5`` reproduces the nearest-posture classification. + * ``"interpolation"`` — per step, score how far the source hand is along its ``hand_open`` → + ``hand_close`` span as a distance ratio (see :func:`_hand_close_fraction`, L1/L2 via + ``hand_interp_norm``) and linearly interpolate the target hand between its own ``hand_open`` and + ``hand_close`` by that fraction, so a gradual grasp stays gradual. ``hand_interp_band`` optionally + remaps that fraction first (see :func:`_apply_interp_band`) so a source that never fully opens/closes + still drives the target across its whole range. + * ``"joint_mapping"`` — per **target joint**, average a chosen group of **source joints** and map that + through open/close (see :func:`_apply_joint_mapping`); the ``joint_mapping`` dict + ``{target_idx: [source_idx, ...]}`` (indices into each embodiment's per-EEF hand vector) defines the + wiring. Finer-grained than ``interpolation``'s single whole-hand fraction — use it to route specific + source fingers to specific target fingers. Target joints not listed stay at ``hand_open``. + + ``"binary"``, ``"interpolation"``, and ``"joint_mapping"`` all need ``hand_open``/``hand_close`` in + both embodiment YAMLs, and drop source-only channels; a non-EEF channel the target requires but the + source lacks must be supplied via ``target_channel_defaults``. + + Args: + source_adapter: Adapter of the embodiment the input dataset was recorded on. + target_adapter: Adapter of the embodiment to retarget onto. + hand_policy: ``"passthrough"``, ``"binary"``, or ``"interpolation"`` (see above). + hand_interp_norm: Distance norm (``"l1"``/``"l2"``) used to score closedness under + ``"binary"`` and ``"interpolation"``. + hand_binary_close_threshold: For ``"binary"``, the source closedness fraction in ``[0, 1]`` at + or above which the target snaps to ``hand_close`` (default ``0.5`` = nearest posture). + hand_interp_band: For ``"interpolation"``, ``[lo, hi]`` sub-range of the source closedness + fraction remapped onto the target's full ``[0, 1]`` open->close travel (values outside clamp); + ``None`` disables it. + joint_mapping: For ``"joint_mapping"``, ``{target_idx: [source_idx, ...]}`` wiring each target + hand joint to the average of a group of source hand joints (indices into the per-EEF vector). + source_hand_postures: ``{eef: {open, close}}`` for the source embodiment (binary/interpolation). + target_hand_postures: ``{eef: {open, close}}`` for the target embodiment (binary/interpolation). + eef_name_map: ``{source_eef: target_eef}`` renaming source channels onto the target names for + ``"passthrough"`` (so a same-width gripper with a different EEF name passes through). + + Returns: + A remapper ``{channel: (num_steps, width)} -> {channel: (num_steps, width)}``. + """ + source_layout = _passthrough_layout(source_adapter) + target_layout = _passthrough_layout(target_adapter) + target_eef_names = set(target_adapter.get_eef_names()) + + if hand_policy in ("binary", "interpolation", "joint_mapping"): + assert ( + source_hand_postures and target_hand_postures + ), f"--hand_policy {hand_policy} needs hand_open/hand_close in both the source and target embodiment YAMLs." + assert hand_interp_norm in ("l1", "l2"), f"hand_interp_norm must be 'l1' or 'l2', got {hand_interp_norm!r}" + if hand_interp_band is not None: + lo, hi = hand_interp_band + assert ( + 0.0 <= lo < hi <= 1.0 + ), f"hand_interp_band must be [lo, hi] with 0 <= lo < hi <= 1, got {hand_interp_band}" + assert hand_policy == "interpolation", "hand_interp_band only applies to --hand_policy interpolation." + for eef_name in target_eef_names: + assert ( + len(target_hand_postures[eef_name]["open"]) == target_layout[eef_name] + ), f"target hand_open[{eef_name!r}] must have {target_layout[eef_name]} values" + if hand_policy == "joint_mapping": + assert ( + joint_mapping + ), "--hand_policy joint_mapping needs a 'joint_mapping' {target_idx: [source_idx, ...]} in the descriptor." + joint_mapping = {int(target): [int(i) for i in sources] for target, sources in joint_mapping.items()} + for eef_name in target_eef_names: + tgt_width, src_width = target_layout[eef_name], source_layout[eef_name] + bad_t = sorted(x for x in joint_mapping if not 0 <= x < tgt_width) + assert ( + not bad_t + ), f"joint_mapping target indices {bad_t} out of range for eef {eef_name!r} (0..{tgt_width - 1})." + bad_s = sorted({i for sources in joint_mapping.values() for i in sources if not 0 <= i < src_width}) + assert ( + not bad_s + ), f"joint_mapping source indices {bad_s} out of range for eef {eef_name!r} (0..{src_width - 1})." + assert all(joint_mapping.values()), "each joint_mapping entry needs at least one source joint index." + + def map_hand(source_hand: torch.Tensor, eef_name: str) -> torch.Tensor: + kwargs = {"dtype": source_hand.dtype, "device": source_hand.device} + src_open = torch.tensor(source_hand_postures[eef_name]["open"], **kwargs) + src_close = torch.tensor(source_hand_postures[eef_name]["close"], **kwargs) + tgt_open = torch.tensor(target_hand_postures[eef_name]["open"], **kwargs) + tgt_close = torch.tensor(target_hand_postures[eef_name]["close"], **kwargs) + if hand_policy == "joint_mapping": + return _apply_joint_mapping(source_hand, src_open, src_close, tgt_open, tgt_close, joint_mapping) + if hand_policy == "binary": + # Snap to close once the source is >= threshold of the way closed, else open. The 1e-4 + # tolerance keeps the boundary inclusive despite the epsilon in the distance ratio, so a + # threshold of 0.25 fires exactly at the "25% closed" point (e.g. Franka 0.03). + alpha = _hand_close_fraction(source_hand, src_open, src_close, hand_interp_norm) + is_closed = alpha >= hand_binary_close_threshold - 1e-4 + return torch.where(is_closed.unsqueeze(1), tgt_close.unsqueeze(0), tgt_open.unsqueeze(0)) + # interpolation: lerp the target open->close by the source's closedness fraction, after the + # optional band remap stretches [lo, hi] of that fraction onto the target's full open->close. + alpha = _apply_interp_band( + _hand_close_fraction(source_hand, src_open, src_close, hand_interp_norm), hand_interp_band + ).unsqueeze(1) + return tgt_open.unsqueeze(0) + alpha * (tgt_close - tgt_open).unsqueeze(0) + + def remap_hands(source_passthrough: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + remapped: dict[str, torch.Tensor] = {} + for channel_name, width in target_layout.items(): + if channel_name in target_eef_names: + remapped[channel_name] = map_hand(source_passthrough[channel_name], channel_name) + else: + remapped[channel_name] = _fill_non_eef_channel( + channel_name, width, source_passthrough, target_channel_defaults + ) + return remapped + + return remap_hands + + assert hand_policy == "passthrough", f"unknown hand_policy {hand_policy!r}" + # Rename the source channels onto the target names (eef_name_map: {source: target}) before comparing + # layouts, so a same-width gripper with a different EEF name (e.g. franka -> ur10) passes through by + # copying the value under the renamed channel. Unlisted channels are kept as-is; the rename is + # idempotent, so pre-renamed input (the generation pool / replay path already rename) is unaffected. + name_map = eef_name_map or {} + renamed_source_layout = {name_map.get(channel, channel): width for channel, width in source_layout.items()} + if renamed_source_layout == target_layout: + + def passthrough(source_passthrough: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + return {name_map.get(channel, channel): value for channel, value in source_passthrough.items()} + + return passthrough + + raise NotImplementedError( + "Source and target embodiments have different passthrough (gripper/hand) layouts " + f"(source={renamed_source_layout} after eef_name_map, target={target_layout}). If only the " + "channel name differs, add an eef_name_map entry; if the widths differ (different grippers), " + "pass --hand_policy binary or interpolation to map the hands via each embodiment's " + "hand_open/hand_close, or add the mapping in build_passthrough_remapper()." + ) diff --git a/autodata_retargeting/object_tracking.py b/autodata_retargeting/object_tracking.py new file mode 100644 index 0000000..7c92fb0 --- /dev/null +++ b/autodata_retargeting/object_tracking.py @@ -0,0 +1,216 @@ +# Copyright (c) 2026, The Isaac AutoData Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Build object-centric carry segments and subtask boundaries from the descriptor's subtasks. + +Each EEF's subtask list (from the retarget descriptor) partitions the source trajectory: a subtask ends +on an event -- a gripper open/close, or a subtask-signal edge -- and a subtask may declare an object it +tracks over its whole span. This module turns those declarations into :class:`CarrySegment`s (where the +replay plans for the *object* instead of the EEF, so the grasped object reproduces its source path +regardless of gripper geometry) plus the boundary steps where the replayer holds a settle. The live +per-step object-centric override lives in the replay loop (it needs the *target* env's grasp transform). +""" + +import torch +import warnings +from dataclasses import dataclass +from typing import Any + + +@dataclass +class CarrySegment: + """A span in which ``eef`` carries ``obj`` (object-centric planning applies on ``[start, end]``). + + ``grasp_transform`` (target ``eef_T_object``) is re-measured live each replay step (kept here mainly + for the debug readout). + """ + + eef: str + obj: str + start: int + end: int + # Object-centric transition smoothing over this segment (see ObjectTracking interpolation params): + # ramp EEF-path -> object over ``interp_start`` steps at the start, object -> EEF-path over + # ``interp_after`` steps past the end. Per-segment so different subtasks can ease differently. + interp_start: int = 0 + interp_after: int = 0 + grasp_transform: torch.Tensor | None = None + last_target: torch.Tensor | None = None # last commanded pose in-segment; the start of the end-ramp + object_start_real: torch.Tensor | None = None # object's real pose at the carry start (for the ease-in) + + +def _event_step(mask: torch.Tensor, start: int, rising: bool) -> int | None: + """First step at or after ``start`` where boolean ``mask`` transitions (rising: F->T, else T->F).""" + values = mask.flatten().bool() + for t in range(max(start, 1), values.numel()): + prev, curr = bool(values[t - 1]), bool(values[t]) + if (rising and curr and not prev) or (not rising and prev and not curr): + return t + return None + + +def _motion_start(frac: torch.Tensor, crossing: int, rising: bool, lower_bound: int) -> int: + """Back up from a gripper crossing to the last frame *before* the motion began (its leading edge). + + ``frac`` is the per-step closedness (0 open -> 1 closed). From ``crossing`` (the first ``>=`` / ``<`` + close-fraction step), step back through the still-monotonic transition -- rising for a close, falling + for an open -- and return the frame just before it started, clamped to ``lower_bound``. Speed-agnostic: + a binary gripper backs up one frame, a slow ramped one backs up to where the ramp left its plateau. + """ + f = frac.flatten() + eps = 1e-4 + t = crossing + while t > lower_bound and ((rising and f[t - 1] < f[t] - eps) or (not rising and f[t - 1] > f[t] + eps)): + t -= 1 + return t + + +def _resolve_subtask_end( + end: Any, + eef_key: str, + start: int, + num_steps: int, + gripper_closed: dict[str, torch.Tensor], + signals: dict[str, torch.Tensor], + name_map: dict[str, str], + close_fraction: float, +) -> int | None: + """Trajectory step where subtask-end event ``end`` fires (searching from ``start``), offset applied. + + ``end`` is a :class:`~.config.SubtaskEnd` (duck-typed: ``.method``, ``.signal``, ``.eef``, ``.offset``, + ``.length``). ``fixed_length`` ends the subtask ``.length`` frames after ``start`` (no source event). + Gripper events read ``gripper_closed[target_eef]`` (a per-step closedness fraction, keyed by the + *target* EEF -- the trigger EEF name is mapped through ``name_map``); ``gripper_closing`` / + ``gripper_opening`` back the crossing up to the motion's leading edge (:func:`_motion_start`). Signal + events read ``signals[name]``. Returns None (with a warning) if the referenced channel is missing or the + event never fires, so a partially-specified demo still runs. + """ + method, offset = end.method, end.offset + if method == "fixed_length": + step = start + end.length # a fixed duration from the subtask's start (no source event) + elif method in ("gripper_close", "gripper_open", "gripper_closing", "gripper_opening"): + trigger = end.eef or eef_key + frac = gripper_closed.get(name_map.get(trigger, trigger)) + if frac is None: + warnings.warn( + f"subtask_end {method!r} references EEF {trigger!r} with no source gripper signal; " + "the subtask boundary is set to the trajectory end. Provide source hand postures.", + stacklevel=2, + ) + return None + f = frac.flatten() + rising = method in ("gripper_close", "gripper_closing") + step = _event_step(f >= close_fraction, start, rising=rising) + if step is not None and method in ("gripper_closing", "gripper_opening"): + step = _motion_start(f, step, rising, start) # back up to the leading edge of the motion + else: # signal_on / signal_off + sig = signals.get(end.signal) + if sig is None: + warnings.warn( + f"subtask_end signal {end.signal!r} is not in the source demo's subtask_term_signals " + f"({sorted(signals)}); the subtask boundary is set to the trajectory end.", + stacklevel=2, + ) + return None + step = _event_step(sig.flatten() > 0.5, start, rising=(method == "signal_on")) + if step is None: + warnings.warn(f"subtask_end {method!r} never fired after step {start}; boundary set to trajectory end.", 2) + return None + return max(0, min(num_steps - 1, step + offset)) + + +def iter_subtask_spans( + subtasks: dict[str, list], + gripper_closed: dict[str, torch.Tensor], + signals: dict[str, torch.Tensor], + num_steps: int, + eef_name_map: dict[str, str] | None = None, + close_fraction: float = 0.5, +): + """Yield ``(eef_key, target_eef, index, subtask, start, end, boundary)`` for every subtask. + + Each EEF's subtasks partition ``[0, num_steps)``: subtask ``i`` runs from the previous subtask's end + (0 for the first) to its own ``subtask_end`` event step; the last subtask (which omits ``subtask_end``) + runs to the trajectory end. ``boundary`` is the resolved end-event step (equal to ``end`` when the + event fired) or ``None`` for the last subtask / when the event never fired. ``eef_key`` is the *source* + EEF name the subtask is listed under; ``target_eef`` is it renamed through ``eef_name_map``. + """ + name_map = eef_name_map or {} + for eef_key, entries in subtasks.items(): + target_eef = name_map.get(eef_key, eef_key) + start = 0 + for i, st in enumerate(entries): + if i == len(entries) - 1: + end, boundary = num_steps - 1, None + else: + assert ( + st.subtask_end is not None + ), f"subtask {i} of EEF {eef_key!r} needs a 'subtask_end' (only the last subtask may omit it)." + boundary = _resolve_subtask_end( + st.subtask_end, eef_key, start, num_steps, gripper_closed, signals, name_map, close_fraction + ) + end = num_steps - 1 if boundary is None else boundary + end = max(end, start) + yield eef_key, target_eef, i, st, start, end, boundary + start = end + + +def carry_segments_and_boundaries_from_subtasks( + subtasks: dict[str, list], + gripper_closed: dict[str, torch.Tensor], + signals: dict[str, torch.Tensor], + num_steps: int, + eef_name_map: dict[str, str] | None = None, + default_interp_start: int = 0, + default_interp_after: int = 0, + default_settle_steps: int = 0, + close_fraction: float = 0.5, +) -> tuple[list[CarrySegment], dict[str, dict[int, int]], dict[str, tuple[str, int]]]: + """Turn the descriptor's per-EEF subtask lists into carry segments + subtask-boundary steps. + + Each EEF's subtasks partition the trajectory: subtask ``i`` runs from the previous subtask's end (0 + for the first) to its own ``subtask_end`` event step; the last subtask (which omits ``subtask_end``) + runs to the trajectory end. A subtask that declares ``object_tracking`` becomes a + :class:`CarrySegment` over its whole span (its per-subtask interp overrides, else the defaults). The + non-final subtask ends are returned as the boundary steps where the replayer holds its settle, each + mapped to its settle-hold cap (``SubtaskEnd.settle_steps`` if set, else ``default_settle_steps``). + + ``subtasks`` values are :class:`~.config.Subtask` objects (duck-typed). Keyed by *source* EEF name; + ``eef_name_map`` renames each onto the target EEF the segment commands. ``gripper_closed`` is keyed by + *target* EEF name (per-step closedness fraction); ``signals`` by signal name. + + Returns: + ``(carry_segments sorted by start, {target_eef: {boundary_step: settle_cap}}, {name: (target_eef, + end_step)})``. Each subtask boundary is registered under its **own** EEF (per-EEF settle), so each + arm settles at its own subtask ends and the executor schedules the arms independently; declare a + ``synchronization`` group to re-couple them (the third return maps each named subtask to its + ``(target_eef, end_step)`` for that barrier). + """ + segments: list[CarrySegment] = [] + boundaries: dict[str, dict[int, int]] = {} + name_end_step: dict[str, tuple[str, int]] = {} + for _eef_key, target_eef, _i, st, start, end, boundary in iter_subtask_spans( + subtasks, gripper_closed, signals, num_steps, eef_name_map, close_fraction + ): + eef_boundaries = boundaries.setdefault(target_eef, {}) + if boundary is not None: + override = st.subtask_end.settle_steps + cap = default_settle_steps if override is None else override + eef_boundaries[boundary] = max(eef_boundaries.get(boundary, 0), cap) + if st.name is not None: + name_end_step[st.name] = (target_eef, end) + if st.object_tracking is not None: + iss, isa = st.object_tracking.interpolation_step_start, st.object_tracking.interpolation_step_after + segments.append( + CarrySegment( + eef=target_eef, + obj=st.object_tracking.object, + start=start, + end=end, + interp_start=default_interp_start if iss is None else iss, + interp_after=default_interp_after if isa is None else isa, + ) + ) + return sorted(segments, key=lambda seg: seg.start), boundaries, name_end_step diff --git a/autodata_retargeting/parallel_replay.py b/autodata_retargeting/parallel_replay.py new file mode 100644 index 0000000..0aaa87d --- /dev/null +++ b/autodata_retargeting/parallel_replay.py @@ -0,0 +1,413 @@ +# Copyright (c) 2026, The Isaac AutoData Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Parallel (multi-env) retargeting replay. + +Replays many source demos at once on a single vectorized env, reusing the proven async-worker + +``env_loop`` machinery that generation runs. One async worker per env pulls episodes from a shared +feed, resets *its* env to the source scene (``prepare_episode`` with ``reset_sim=False``), and drives +the trajectory by putting one action per env-step on the shared action queue; ``env_loop`` batches all +envs into a single ``env.step``. When the episodes run out, a worker holds (zero action) so the loop +keeps stepping until every episode is done. + +Not supported in parallel mode (use ``--num_envs 1``): ``init_robot_from_ik`` (env-0 PinkIK) and the +per-step tracking-error report -- both single-env only. +""" + +import asyncio +import contextlib +import sys +import torch +import traceback +from copy import deepcopy +from dataclasses import dataclass +from typing import Any + +from autodata_interfaces.env import env_loop +from autodata_utils.tensor_utils import as_torch + +from .config import DefaultObjectTracking +from .provider import PlanProvider, ReplayResult +from .replay import ( + _apply_object_centric_override, + _build_controlled_link_pose_reader, + _monitored_miss, + _record_datagen_poses, + _record_signal_frame, + prepare_episode, + read_achieved_eef_poses, +) +from .util import pose_tracking_error + + +@dataclass +class ReplayParams: + """Bundle of per-run retargeting knobs shared by every parallel worker (see ``RetargetConfig``).""" + + robot_asset_name: str + reference_pose: str + eef_name_map: dict[str, str] | None + replay_speed: float + retarget_frame: str + scene_translation: tuple[float, float, float] + eef_offsets: dict[str, torch.Tensor] | None + num_interpolation_steps: int + init_robot_from_ik: bool + subtasks: dict + default_object_tracking: DefaultObjectTracking + write_datagen_info: bool + eef_reference_link: dict[str, str] | str | None + source_hand_postures: dict[str, dict[str, list[float]]] | None + success_settle_steps: int + segment_settle_steps: int + settle_pos_tol_m: float + settle_rot_tol_deg: float + settle_joint_tol: float + config: Any = None # RetargetConfig (early-abort thresholds, synchronization, speed caps) + + +async def _async_step( + env, env_id, action_queue, target_adapter, target_eef_pose_dict, passthrough_action_dict +) -> None: + """Encode one action for ``env_id`` and hand it to ``env_loop`` (which batches all envs and steps).""" + action = target_adapter.target_eef_pose_to_action( + target_eef_pose_dict=target_eef_pose_dict, + passthrough_action_dict=passthrough_action_dict, + env_id=env_id, + ) + if action.dim() > 1: + action = action[0] + await action_queue.put((env_id, action.to(device=env.device))) + await action_queue.join() + + +async def _replay_one_episode( + env, env_id, plan, action_queue, target_default_state, adapters, params, success_term +) -> bool: + """Reset ``env_id`` to the source scene and drive one plan's trajectory async; return whether it + succeeded.""" + target_adapter, source_adapter, remap_passthrough = adapters + prep = prepare_episode( + env, + env_id, + plan.episode, + target_adapter, + source_adapter, + remap_passthrough, + target_default_state, + params.robot_asset_name, + params.reference_pose, + params.eef_name_map, + params.replay_speed, + params.retarget_frame, + params.scene_translation, + params.eef_offsets, + params.num_interpolation_steps, + init_robot_from_ik=params.init_robot_from_ik, + subtasks=params.subtasks, + default_object_tracking=params.default_object_tracking, + source_hand_postures=params.source_hand_postures, + need_source_objects=any( + st.object_tracking is not None for entries in params.subtasks.values() for st in entries + ), + reset_sim=False, + eef_reference_link=params.eef_reference_link, + write_datagen_info=params.write_datagen_info, + segment_settle_steps=params.segment_settle_steps, + max_eef_linear_velocity=params.config.max_eef_linear_velocity, + max_eef_rotation_speed=params.config.max_eef_rotation_speed, + synchronization=params.config.synchronization, + ) + eef_names = prep["eef_names"] + commanded_poses, commanded_passthrough = prep["commanded_poses"], prep["commanded_passthrough"] + num_interpolation_steps, num_steps = prep["num_interpolation_steps"], prep["num_steps"] + carry_segments, source_objects = prep["carry_segments"], prep["source_objects"] + source_signals = prep["source_signals"] # {name: (num_traj_steps, 1)} to forward, or {} + segment_ends = prep["segment_ends"] # {eef: {trajectory step: settle-hold cap}} for the motion-aware settle + sync_of = prep.get("sync_of", {}) # {eef: {end step: group id}} -- a cross-EEF rendezvous barrier + group_members = prep.get("group_members", {}) # {group id: [(eef, end step)]} + # Measure the object-centric grasp transform (and any debug tracking) at the same IK-controlled link + # the command drives when the reference is reconstructed there, else the observed EEF frame. + controlled_reader = ( + _build_controlled_link_pose_reader(env, target_adapter) if params.eef_reference_link is not None else None + ) + + # Per-EEF single-step scheduler (mirrors replay_episode_on_target, one ``_async_step`` per tick for THIS + # env): each EEF advances its own pointer through its (possibly unequal-length) commanded sequence; at a + # subtask boundary it HOLDS until its motion-aware settle AND its sync group's rendezvous are satisfied. + # This is per-env local state, so the barrier joins the two arms of one env -- bimanual + hand-off run in + # parallel. Pose / passthrough / override are (re)computed only for EEFs that just advanced (a holding + # carrier keeps its pose so the override does not remeasure the grasp mid-release). + task_succeeded = success_term is None + lengths = {eef_name: commanded_poses[eef_name].shape[0] for eef_name in eef_names} + ptr = {eef_name: 0 for eef_name in eef_names} + hold = {eef_name: 0 for eef_name in eef_names} + prev_pose: dict[str, torch.Tensor | None] = {eef_name: None for eef_name in eef_names} + prev_ptr = {eef_name: -1 for eef_name in eef_names} + prev_qpos = None + channel_eef = {name: name if name in eef_names else eef_names[0] for name in commanded_passthrough} + target_eef_pose_dict: dict[str, torch.Tensor] = {} + passthrough_action_dict: dict[str, torch.Tensor] = {} + record_signals = None + # Full datagen_info (poses) per step when write_datagen_info (copy) -> the output is generate-ready. + object_names = list(env.scene.rigid_objects.keys()) if params.write_datagen_info else [] + tick = 0 + early_failure = False + monitor_early = params.config.stop_early_on_failure and ( + params.config.max_translation_error is not None or params.config.max_rotation_error is not None + ) + while any(ptr[eef_name] < lengths[eef_name] for eef_name in eef_names): + idx = {eef_name: min(ptr[eef_name], lengths[eef_name] - 1) for eef_name in eef_names} + traj_step = {eef_name: idx[eef_name] - num_interpolation_steps for eef_name in eef_names} + advanced = {eef_name for eef_name in eef_names if ptr[eef_name] != prev_ptr[eef_name]} + for eef_name in advanced: + target_eef_pose_dict[eef_name] = commanded_poses[eef_name][idx[eef_name]] + prev_ptr[eef_name] = ptr[eef_name] + for name, tensor in commanded_passthrough.items(): + if channel_eef[name] in advanced or name not in passthrough_action_dict: + passthrough_action_dict[name] = tensor[idx[channel_eef[name]]] + if carry_segments and advanced: + _apply_object_centric_override( + carry_segments, + traj_step, + target_eef_pose_dict, + source_objects, + target_adapter, + env, + env_id, + controlled_reader=controlled_reader, + commanded_poses=commanded_poses, + num_interpolation_steps=num_interpolation_steps, + eefs=advanced, + ) + ref_ts = max(traj_step.values()) # shared clock for signals only + signal_frame = {name: sig[min(max(ref_ts, 0), sig.shape[0] - 1)] for name, sig in source_signals.items()} + record_signals = (lambda sf=signal_frame: _record_signal_frame(env, env_id, sf)) if signal_frame else None + + await _async_step( + env, + env_id, + action_queue, + target_adapter, + target_eef_pose_dict, + passthrough_action_dict, + ) + if record_signals is not None: + record_signals() + if params.write_datagen_info: + _record_datagen_poses(env, env_id, target_adapter, target_eef_pose_dict, object_names) + if success_term is not None and bool(success_term.func(env, **success_term.params)[env_id]): + task_succeeded = True + + if monitor_early: + achieved_poses = read_achieved_eef_poses(target_adapter, controlled_reader, env_id) + if _monitored_miss( + eef_names, + traj_step, + ptr, + lengths, + carry_segments, + source_objects, + target_eef_pose_dict, + achieved_poses, + env, + env_id, + params.config.max_translation_error, + params.config.max_rotation_error, + tick, + ): + early_failure = True + break + + curr_poses = target_adapter.get_eef_poses(env_ids=[env_id]) + curr_qpos = ( + as_torch(env.scene[params.robot_asset_name].data.joint_pos)[env_id] if params.robot_asset_name else None + ) + joint_moved = ( + prev_qpos is not None + and curr_qpos is not None + and float(torch.max(torch.abs(curr_qpos - prev_qpos))) > params.settle_joint_tol + ) + for eef_name in eef_names: + if ptr[eef_name] >= lengths[eef_name]: + continue + cap = segment_ends.get(eef_name, {}).get(traj_step[eef_name], 0) + gid = sync_of.get(eef_name, {}).get(traj_step[eef_name]) + if cap > 0 or gid is not None: + hold[eef_name] += 1 + settle_ok = True + if cap > 0: + moved = joint_moved + if prev_pose[eef_name] is not None: + dpos, drot = pose_tracking_error(prev_pose[eef_name], curr_poses[eef_name][0]) + moved = moved or dpos > params.settle_pos_tol_m or drot > params.settle_rot_tol_deg + settle_ok = (not moved) or hold[eef_name] >= cap + sync_ok = gid is None or all( + ptr[m_eef] - num_interpolation_steps >= m_step for m_eef, m_step in group_members[gid] + ) + if settle_ok and sync_ok: + ptr[eef_name] += 1 + hold[eef_name] = 0 + else: + ptr[eef_name] += 1 + prev_pose[eef_name] = curr_poses[eef_name][0] + prev_qpos = curr_qpos.clone() if curr_qpos is not None else None + tick += 1 + + # Final success settle: hold the last pose (final gripper release) until success or the cap. + if ( + not task_succeeded + and not early_failure + and params.success_settle_steps > 0 + and num_steps > 0 + and success_term is not None + ): + for _ in range(params.success_settle_steps): + await _async_step(env, env_id, action_queue, target_adapter, target_eef_pose_dict, passthrough_action_dict) + if record_signals is not None: # hold the last waypoint's signal across the success settle + record_signals() + if params.write_datagen_info: + _record_datagen_poses(env, env_id, target_adapter, target_eef_pose_dict, object_names) + if bool(success_term.func(env, **success_term.params)[env_id]): + task_succeeded = True + break + return task_succeeded, passthrough_action_dict + + +def _zeros_passthrough(env, target_adapter) -> dict[str, torch.Tensor]: + """A zero-filled passthrough dict matching the adapter's channel layout (for never-replayed workers).""" + zeros = torch.zeros(1, int(env.action_space.shape[-1]), device=env.device) + return {name: value[0] for name, value in target_adapter.actions_to_passthrough_actions(zeros).items()} + + +def _hold_action(env, env_id, target_adapter, passthrough) -> torch.Tensor: + """A valid 'stay here' action for any controller: command the EEF's *current* pose. + + For a delta-pose env this is a zero delta; for an absolute-pose PinkIK env it is the current pose -- + both hold the robot in place, unlike a zero action, whose zero/identity target pose is unreachable + and makes the IK solver spam "could not find a solution / NaN" every step for each finished env. + """ + current = {eef: pose[0] for eef, pose in target_adapter.get_eef_poses(env_ids=[env_id]).items()} + action = target_adapter.target_eef_pose_to_action(current, passthrough, env_id=env_id) + return action[0] if action.dim() > 1 else action + + +async def _replay_worker( + env, + env_id, + provider: PlanProvider, + action_queue, + target_default_state, + adapters, + params, + success_term, + results, + stats, +) -> None: + """One env's worker: pull plans from ``provider`` and replay them. + + When the provider hands out nothing: if it is fully drained (nothing in flight on any env) the worker + *returns* so ``env_loop`` terminates via its task-completion check (this covers an unreachable success + target on a finite source); otherwise it *holds* at the current pose so the batched step keeps + advancing the still-running workers. + """ + target_adapter = adapters[0] + env_ids = torch.tensor([env_id], device=env.device) + hold_passthrough = None # last episode's passthrough, reused to hold the gripper at a valid command + while True: + # ``next_for_env`` is env-agnostic for the copy provider (its plans do not depend on which env runs + # them). Synchronous, so each worker claims a distinct source episode + index. + plan = provider.next_for_env(env_id) + if plan is None: + if provider.in_flight() == 0: + return + if hold_passthrough is None: + hold_passthrough = _zeros_passthrough(env, target_adapter) + await action_queue.put((env_id, _hold_action(env, env_id, target_adapter, hold_passthrough))) + await action_queue.join() + continue + try: + success, hold_passthrough = await _replay_one_episode( + env, env_id, plan, action_queue, target_default_state, adapters, params, success_term + ) + except Exception: + sys.stderr.write(f"[parallel-replay] env {env_id} failed on {plan.name}:\n{traceback.format_exc()}") + sys.stderr.flush() + success = False + env.recorder_manager.set_success_to_episodes( + env_ids, torch.tensor([[success]], dtype=torch.bool, device=env.device) + ) + env.recorder_manager.export_episodes(env_ids) + provider.observe(plan, ReplayResult(success=bool(success))) + results.append((plan.name, bool(success))) + stats["num_attempts"] += 1 + stats["num_success"] += int(success) + + +def run_parallel_replay( + env, + provider: PlanProvider, + num_envs: int, + adapters: tuple, + target_default_state: dict, + success_term, + params: ReplayParams, + generation_policy, +) -> list[tuple[str, bool]]: + """Replay plans from ``provider`` across ``num_envs`` parallel workers; return ``[(name, success), ...]``. + + Reuses generation's ``env_loop``: it batches one action per env each step and terminates on the + provider's stop condition -- ``provider.loop_policy()`` gives ``(guarantee_success, num_trials)`` so + the loop exits on enough successes or attempts, and the workers exit once the provider is drained (so + an unreachable target on a finite source still terminates). Workers reset their env with + ``env.reset_to`` directly (not the reset queue) since each needs its plan's scene. + """ + event_loop = asyncio.get_event_loop() + action_queue: asyncio.Queue = asyncio.Queue() + reset_queue: asyncio.Queue = asyncio.Queue() # unused here; env_loop expects the handle + results: list[tuple[str, bool]] = [] + stats = {"num_success": 0, "num_failures": 0, "num_attempts": 0} + + tasks = [ + event_loop.create_task( + _replay_worker( + env, + env_id, + provider, + action_queue, + target_default_state, + adapters, + params, + success_term, + results, + stats, + ) + ) + for env_id in range(num_envs) + ] + data_gen_tasks = asyncio.ensure_future(asyncio.gather(*tasks)) + + # Terminate the step loop on the provider's stop condition (successes or attempts). + guarantee_success, num_trials = provider.loop_policy() + loop_policy = deepcopy(generation_policy) + loop_policy.num_trials = num_trials + loop_policy.guarantee_success = guarantee_success + try: + env_loop( + env, + reset_queue, + action_queue, + event_loop, + generation_policy_params=loop_policy, + stats=stats, + data_gen_tasks=data_gen_tasks, + ) + finally: + # Cancelling the workers makes the gathered future raise CancelledError (a BaseException, so it + # is not caught by ``suppress(Exception)``) — list it explicitly so results are returned cleanly. + data_gen_tasks.cancel() + with contextlib.suppress(asyncio.CancelledError, Exception): + event_loop.run_until_complete(data_gen_tasks) + return results diff --git a/autodata_retargeting/provider.py b/autodata_retargeting/provider.py new file mode 100644 index 0000000..d04174a --- /dev/null +++ b/autodata_retargeting/provider.py @@ -0,0 +1,147 @@ +# Copyright (c) 2026, The Isaac AutoData Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Plan providers: the *source* of retargeting examples, decoupled from the replayer. + +A :class:`PlanProvider` hands out :class:`Plan` objects (one retargeting example each) and owns the +**stop condition**. The replayer just pulls plans until ``next()`` returns ``None`` and reports each +outcome back via ``observe()``. This separates "which example / how many" (provider) from "execute it +on the target robot" (replayer). + +:class:`DatasetReplayProvider` is the current behaviour: replay every source episode once (optionally a +selected subset). Stop when the source is exhausted, or early once a target number of successes or of +total runs is reached. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Any + + +@dataclass +class Plan: + """One retargeting example to replay. + + For :class:`DatasetReplayProvider` this wraps a loaded source episode. ``index`` / ``name`` are for + logging + output. + """ + + index: int + name: str + episode: Any = None # isaaclab EpisodeData (or an equivalent trajectory the replayer can consume) + + +@dataclass +class ReplayResult: + """Outcome of replaying one :class:`Plan`, fed back to the provider via :meth:`PlanProvider.observe`.""" + + success: bool + eef_errors: dict[str, dict[str, Any]] | None = None + + +class PlanProvider(ABC): + """Hands out retargeting examples and owns the stop condition. + + Contract: call :meth:`next` to get the next :class:`Plan` or ``None`` when done (target reached or + source exhausted); after replaying it, report the outcome with :meth:`observe` so success/run targets + can stop early. Providers are pulled cooperatively (single-env loop, or the async parallel workers), + so implementations need no locks. + """ + + @abstractmethod + def next(self) -> Plan | None: + """Return the next example, or ``None`` when the provider is done.""" + + def next_for_env(self, env_id: int) -> Plan | None: + """Next plan for a specific parallel worker's ``env_id``. + + Defaults to :meth:`next` (env-agnostic -- the copy provider's plans do not depend on which env + runs them). + """ + return self.next() + + def observe(self, plan: Plan, result: ReplayResult) -> None: + """Record one replay outcome (default: no-op). Overridden to drive success/run targets.""" + + def in_flight(self) -> int: + """Plans handed out but not yet observed (used by the parallel drain to know when to exit).""" + return 0 + + @abstractmethod + def loop_policy(self) -> tuple[bool, int]: + """``(guarantee_success, num_trials)`` mirroring the stop condition, for the parallel ``env_loop``. + + ``guarantee_success`` True -> the loop terminates on that many *successes*; False -> on that many + *attempts*. The async workers also exit once the provider is drained, so an unreachable success + target on a finite source still terminates (via ``env_loop``'s task-completion check). + """ + + +class DatasetReplayProvider(PlanProvider): + """Replay each source episode once (optionally a selected subset), i.e. today's full-replay behaviour. + + Args: + dataset_handler: An opened HDF5 dataset handler exposing ``load_episode(name, device)``. + episode_names: Ordered source-episode names to replay. + device: Torch device to load episodes onto. + target_successes: Stop once this many replays have succeeded (``None`` = no success target). + target_runs: Stop once this many replays have been attempted (``None`` = no run target). + + With neither target set, it stops when the source is exhausted. The two targets and exhaustion + compose: whichever is reached first ends the run. + """ + + def __init__( + self, + dataset_handler: Any, + episode_names: list[str], + device: Any, + target_successes: int | None = None, + target_runs: int | None = None, + ) -> None: + assert target_successes is None or target_successes > 0, "target_successes must be positive" + assert target_runs is None or target_runs > 0, "target_runs must be positive" + self._handler = dataset_handler + self._names = episode_names + self._device = device + self._target_successes = target_successes + self._target_runs = target_runs + self._cursor = 0 # next source index to hand out + self._handed_out = 0 + self._completed = 0 + self._num_success = 0 + + def _target_reached(self) -> bool: + if self._target_successes is not None and self._num_success >= self._target_successes: + return True + if self._target_runs is not None and self._completed >= self._target_runs: + return True + return False + + def next(self) -> Plan | None: + """Load and return the next source episode unless a stop condition was reached.""" + if self._target_reached() or self._cursor >= len(self._names): + return None + name = self._names[self._cursor] + index = self._cursor + self._cursor += 1 + self._handed_out += 1 + return Plan(index=index, name=name, episode=self._handler.load_episode(name, self._device)) + + def observe(self, plan: Plan, result: ReplayResult) -> None: + self._completed += 1 + self._num_success += int(result.success) + + def in_flight(self) -> int: + return self._handed_out - self._completed + + def loop_policy(self) -> tuple[bool, int]: + if self._target_successes is not None: + return True, self._target_successes + if self._target_runs is not None: + return False, min(self._target_runs, len(self._names)) + return False, len(self._names) diff --git a/autodata_retargeting/replay.py b/autodata_retargeting/replay.py new file mode 100644 index 0000000..7b51b8e --- /dev/null +++ b/autodata_retargeting/replay.py @@ -0,0 +1,1302 @@ +# Copyright (c) 2026, The Isaac AutoData Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""The retargeting replay engine: reset state, IK warm-start, closed-loop replay, error report.""" + +import torch +from copy import deepcopy +from typing import Any + +import isaaclab.utils.math as math_utils +from isaaclab.managers import TerminationTermCfg +from isaaclab.utils.datasets import EpisodeData + +from autodata_interfaces.embodiments.embodiment_adapter import EmbodimentAdapter +from autodata_utils.tensor_utils import as_torch + +from .config import DefaultObjectTracking +from .gripper_retargeting import PassthroughRemapper, _hand_close_fraction +from .object_tracking import CarrySegment, carry_segments_and_boundaries_from_subtasks, iter_subtask_spans +from .trajectory import ( + cap_shared_timeline_speed, + resample_trajectory, + source_datagen_poses, + source_eef_poses_at_link, + source_object_poses, + source_subtask_signals, +) +from .util import ( + pose_tracking_error, + poses_from_root_pose, + quat_slerp_batch, + read_target_base_pose, + reanchor_to_target_base, + se3_inverse, +) + +# A commanded EEF pose is counted as "reached" (IK effectively tracked it) when the achieved pose +# is within these tolerances. Used only for the IK-performance report, not for success. +_IK_POS_TOL_M = 0.05 +_IK_ROT_TOL_DEG = 15.0 + +# Per-step EEF motion below these tolerances is treated as "settled" (motion has died down), used to +# end a segment settle-hold early. Compared as a delta between consecutive steps, so the pitch/roll +# frame offset that inflates absolute tracking error does not affect it. +_SETTLE_POS_TOL_M = 0.01 +_SETTLE_ROT_TOL_DEG = 2.5 +# At a subtask boundary (grasp / place) the replayer holds the pose until the robot's joints -- arm AND +# gripper -- stop moving, up to ``config.segment_settle_steps`` extra sim steps (a run knob), then exits +# early, so a slow gripper (e.g. the under-actuated Robotiq 2F-85 linkage) finishes closing/opening before +# the arm advances. Non-boundary waypoints still replay at source speed (one step). "Stopped" is measured +# by per-step joint *position* change, NOT joint_vel: the Robotiq linkage reports a velocity-saturated +# ~1 rad/s joint forever even at rest, whereas the position delta decays cleanly to ~0 once it has closed. +# A fast gripper settles in a few steps and exits well before the cap. +_JOINT_SETTLE_TOL = 0.003 # [rad or m] max per-step joint position change below which the robot is settled + +_IK_SOLVE_ITERATIONS = 80 +"""Iterations of the differential IK solver used to converge the first-pose joint configuration.""" + +# A source EEF's gripper is treated as "closed" (grasping) for carry detection once its closedness +# fraction (0 open -> 1 closed, vs the source hand_open/hand_close postures) reaches this. +_CARRY_GRIPPER_CLOSED_FRACTION = 0.5 + + +def validate_eef_agreement( + source_adapter: EmbodimentAdapter, + target_adapter: EmbodimentAdapter, + eef_name_map: dict[str, str] | None = None, +) -> None: + """Assert both embodiments declare the same EEF names (after ``eef_name_map``). + + The retarget replays the source EEF pose trajectory on the target, so both embodiments must speak of + the same EEFs. ``eef_name_map`` (``{source_eef: target_eef}``) renames the source EEF names onto the + target names before the check, so a cross-named single-arm pair (e.g. ``franka`` -> ``ur10``) agrees. + """ + name_map = eef_name_map or {} + source_eefs = {name_map.get(eef, eef) for eef in source_adapter.get_eef_names()} + target_eefs = set(target_adapter.get_eef_names()) + assert ( + source_eefs == target_eefs + ), f"EEF name mismatch across embodiments: source(mapped)={sorted(source_eefs)}, target={sorted(target_eefs)}" + ( + f" (eef_name_map={name_map})" if name_map else "" + ) + + +def retargeted_initial_state( + source_state: dict, + target_default_state: dict, + robot_asset_name: str, + object_translation: tuple[float, float, float] = (0.0, 0.0, 0.0), +) -> dict: + """Build a reset state placing the task scene from the source demo on the target robot's env. + + The result uses the *target* env's exact scene schema (from ``target_default_state``) so + ``reset_to`` never sees an unexpected/missing key, and overrides every non-robot entity's + matching fields with the source demo's values (object poses, etc.). The robot is kept at the + target default here; ``init_robot_from_ik`` (in the replay) later overwrites it with the IK + solution of the first trajectory pose. + + Args: + source_state: The source episode's recorded ``initial_state`` (scene state dict). + target_default_state: A freshly-reset target scene state (``scene.get_state``). + robot_asset_name: Scene key of the robot articulation. + object_translation: Constant world translation [m] added to every non-robot entity's root + position, shifting the task objects with the trajectory. + + Returns: + A scene-state dict compatible with the target env's ``reset_to``. + """ + merged = deepcopy(target_default_state) + for category, entities in merged.items(): + source_category = source_state.get(category, {}) + for name, fields in entities.items(): + if name == robot_asset_name: + continue # Keep the target robot's default configuration (IK warm-start overrides it). + source_entity = source_category.get(name) + if source_entity is None: + continue + for field, target_value in fields.items(): + source_value = source_entity.get(field) + if source_value is not None and tuple(source_value.shape) == tuple(target_value.shape): + fields[field] = source_value + + # Shift the (non-robot) task objects by the same constant translation applied to the trajectory. + if any(object_translation): + for category, entities in merged.items(): + if category == "articulation": + continue + for name, fields in entities.items(): + root_pose = fields.get("root_pose") + if name == robot_asset_name or root_pose is None: + continue + translation = torch.tensor(object_translation, dtype=root_pose.dtype, device=root_pose.device) + shifted = root_pose.clone() + shifted[..., :3] = shifted[..., :3] + translation + fields["root_pose"] = shifted + return merged + + +def _read_env_object_pose(env: Any, object_name: str, env_id: int = 0) -> torch.Tensor: + """Read a scene object's current pose as an env-relative ``(4, 4)`` SE(3) transform for ``env_id``. + + The origin is subtracted so it matches the adapter's env-relative EEF frame, letting the two be + composed into the grasp transform for object-centric planning. + """ + obj = env.scene[object_name] + origin = as_torch(env.scene.env_origins)[env_id] + pos = as_torch(obj.data.root_pos_w)[env_id] - origin + quat = as_torch(obj.data.root_quat_w)[env_id] + pose = torch.eye(4, dtype=pos.dtype, device=pos.device) + pose[:3, :3] = math_utils.matrix_from_quat(quat.unsqueeze(0))[0] + pose[:3, 3] = pos + return pose + + +def _blend_pose(pose_a: torch.Tensor, pose_b: torch.Tensor, alpha: float) -> torch.Tensor: + """Blend two ``(4, 4)`` SE(3) poses: position lerp, orientation slerp. ``alpha=0`` -> a, ``1`` -> b.""" + out = torch.eye(4, dtype=pose_a.dtype, device=pose_a.device) + out[:3, 3] = (1.0 - alpha) * pose_a[:3, 3] + alpha * pose_b[:3, 3] + quat_a = math_utils.quat_from_matrix(pose_a[:3, :3].unsqueeze(0))[:, [0, 1, 2, 3]] + quat_b = math_utils.quat_from_matrix(pose_b[:3, :3].unsqueeze(0))[:, [0, 1, 2, 3]] + frac = torch.full((1, 1), float(alpha), dtype=pose_a.dtype, device=pose_a.device) + blended = quat_slerp_batch(quat_a, quat_b, frac)[:, [0, 1, 2, 3]] + out[:3, :3] = math_utils.matrix_from_quat(blended)[0] + return out + + +def _apply_object_centric_override( + carry_segments: list[CarrySegment], + trajectory_step_by_eef: dict[str, int], + target_eef_pose_dict: dict[str, torch.Tensor], + source_objects: dict[str, torch.Tensor], + target_adapter: EmbodimentAdapter, + env: Any, + env_id: int = 0, + controlled_reader=None, + commanded_poses: dict[str, torch.Tensor] | None = None, + num_interpolation_steps: int = 0, + eefs: set[str] | None = None, +) -> None: + """During a carry segment, replace the carrying EEF's commanded pose with the object-centric one. + + Each step the grasp transform is **re-measured** live in the target env + (``target_eef_T_object = inv(target_eef) @ target_object``) and the EEF is commanded to + ``src_object_pose @ inv(grasp_T)`` -- a best-effort, closed-loop move that brings the object from + wherever it currently sits in the grip to the source object pose this step. Re-measuring every step + (rather than freezing ``grasp_T`` at the segment start) corrects any object slip relative to the EEF, + so a non-rigid dexterous grasp still tracks. Mutates ``target_eef_pose_dict``; a no-op outside segments. + + ``grasp_T`` must be measured in the **same frame the command drives** -- the IK-controlled link, via + ``controlled_reader`` (:func:`read_achieved_eef_poses`), when the reference is reconstructed there. + Measuring it at the observed link while commanding the controlled link would place the object off by + the fixed observed->controlled offset (e.g. GR1's ~9 deg wrist_pitch). Defaults to the observed frame. + + Transitions are eased (slerp) rather than snapped: over the first ``interp_start`` steps the command + ramps from the EEF path to the object-centric pose (step i weights the object by ``(i+1)/interp_start``); + for ``interp_after`` steps past the segment it ramps from the last tracked pose toward the EEF target + ``interp_after`` steps ahead (needs ``commanded_poses`` + ``num_interpolation_steps`` to index it), + landing back on the EEF path. + """ + for seg in carry_segments: + if eefs is not None and seg.eef not in eefs: # skip a holding eef -- do not remeasure its grasp + continue + ts = trajectory_step_by_eef[seg.eef] # this carrying EEF's own trajectory step + obj_len = source_objects[seg.obj].shape[0] + in_segment = seg.start <= ts <= seg.end and ts < obj_len + in_after = 0 < seg.interp_after and seg.end < ts <= seg.end + seg.interp_after + if not (in_segment or in_after): + continue + eef_path_pose = target_eef_pose_dict[seg.eef] # the source EEF-path pose (this step, pre-override) + if in_segment: + eef_now = read_achieved_eef_poses(target_adapter, controlled_reader, env_id)[seg.eef] + # Live grasp transform (eef_T_object), re-measured every step so slip in the grip self-corrects. + seg.grasp_transform = se3_inverse(eef_now) @ _read_env_object_pose(env, seg.obj, env_id) + step_in = ts - seg.start + if step_in == 0: # capture where the object really is as the carry begins (the ease-in start) + seg.object_start_real = _read_env_object_pose(env, seg.obj, env_id) + if 0 < seg.interp_start and step_in < seg.interp_start: + # Ease in *object* space: move the object from where it really is (grasp may have nudged it) + # to the trajectory target ``interp_start`` steps in, then apply the live grasp to get the EEF. + target_index = min(seg.start + seg.interp_start, seg.end, obj_len - 1) + object_target = _blend_pose( + seg.object_start_real, source_objects[seg.obj][target_index], (step_in + 1) / seg.interp_start + ) + else: + object_target = source_objects[seg.obj][ts] + object_pose = object_target @ se3_inverse(seg.grasp_transform) + seg.last_target = object_pose + target_eef_pose_dict[seg.eef] = object_pose + else: # end-ramp: ease the last tracked pose -> the EEF target interp_after steps ahead + start_pose = seg.last_target if seg.last_target is not None else eef_path_pose + endpoint = eef_path_pose + if commanded_poses is not None and seg.eef in commanded_poses: + traj = commanded_poses[seg.eef] + endpoint = traj[min(seg.end + seg.interp_after + num_interpolation_steps, traj.shape[0] - 1)] + step_after = ts - seg.end - 1 + target_eef_pose_dict[seg.eef] = _blend_pose(start_pose, endpoint, (step_after + 1) / seg.interp_after) + + +def _record_signal_frame(env: Any, env_id: int, signal_frame: dict[str, torch.Tensor]) -> None: + """Append this frame's forwarded subtask signals to the recorded episode (once per ``env.step``). + + ``signal_frame`` maps each signal name to its ``(width,)`` value for the current waypoint (held across + the waypoint's settle/lead-in frames), recorded under ``obs/datagen_info/subtask_term_signals`` so the + output dataset's signals stay frame-aligned with the recorded rollout. + """ + if not signal_frame: + return + value = {name: signal.reshape(1, -1) for name, signal in signal_frame.items()} + env.recorder_manager.add_to_episodes("obs/datagen_info/subtask_term_signals", value, env_ids=[env_id]) + + +def _record_datagen_poses( + env: Any, + env_id: int, + target_adapter: EmbodimentAdapter, + target_eef_pose_dict: dict[str, torch.Tensor], + object_names: list[str], +) -> None: + """Record this frame's ``obs/datagen_info`` poses (once per ``env.step``). + + Writes the observed ``eef_pose`` (``target_adapter.get_eef_poses``), the commanded ``target_eef_pose``, + and per-object ``object_pose`` -- the same quantities ``annotate_demos``' ``PreStepDatagenInfoRecorder`` + writes -- so a ``copy`` output carries a complete ``datagen_info`` and is a drop-in source for + ``generate_dataset.py`` (no separate annotation pass). All env-relative ``(1, 4, 4)`` SE(3), per key. + """ + recorder = env.recorder_manager + achieved = target_adapter.get_eef_poses(env_ids=[env_id]) # {eef: (1, 4, 4)} observed EEF + recorder.add_to_episodes("obs/datagen_info/eef_pose", achieved, env_ids=[env_id]) + recorder.add_to_episodes( + "obs/datagen_info/target_eef_pose", + {eef: target_eef_pose_dict[eef].reshape(1, 4, 4) for eef in achieved}, + env_ids=[env_id], + ) + if object_names: + recorder.add_to_episodes( + "obs/datagen_info/object_pose", + {name: _read_env_object_pose(env, name, env_id).reshape(1, 4, 4) for name in object_names}, + env_ids=[env_id], + ) + + +def _build_commanded_sequences( + target_adapter: EmbodimentAdapter, + target_eef_poses: dict[str, torch.Tensor], + target_passthrough: dict[str, torch.Tensor], + eef_names: list[str], + num_interpolation_steps: int, + env_id: int = 0, +) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: + """Build the per-step commanded pose + passthrough sequences, with an optional lead-in ramp. + + When ``num_interpolation_steps > 0``, each EEF is linearly interpolated from the target robot's + *current* pose (read from the freshly-reset env) to the first source-trajectory pose, and the + passthrough channels hold their first value across the ramp. Both returned dicts share the same + length ``num_interpolation_steps + T``. + """ + commanded_poses: dict[str, torch.Tensor] = {} + if num_interpolation_steps > 0: + current_poses = target_adapter.get_eef_poses(env_ids=[env_id]) + for eef_name in eef_names: + trajectory = target_eef_poses[eef_name] + if num_interpolation_steps > 0: + # interpolate_poses returns [start, ...N interp..., end]; keep only the N interior poses + # (the trajectory itself supplies its own first pose). + lead_in, _ = math_utils.interpolate_poses( + current_poses[eef_name][0], trajectory[0], num_steps=num_interpolation_steps + ) + commanded_poses[eef_name] = torch.cat([lead_in[1:-1], trajectory], dim=0) + else: + commanded_poses[eef_name] = trajectory + + commanded_passthrough: dict[str, torch.Tensor] = {} + for name, tensor in target_passthrough.items(): + if num_interpolation_steps > 0: + held = tensor[0:1].expand(num_interpolation_steps, *tensor.shape[1:]) + commanded_passthrough[name] = torch.cat([held, tensor], dim=0) + else: + commanded_passthrough[name] = tensor + return commanded_poses, commanded_passthrough + + +def _build_controlled_link_pose_reader(env: Any, target_adapter: EmbodimentAdapter): + """Return a fn giving each EEF's IK-controlled-link pose ``{eef: (4,4)}``, or None if unavailable. + + Reads the world (env-relative) pose of the link each EEF's commanded pose actually drives (the + PinkIK ``target_eef_link_names``), mapped to the adapter's EEF names by order. This is the frame + the retarget commands, so measuring here removes the target's observed-vs-controlled frame offset. + """ + eef_names = list(target_adapter.get_eef_names()) + manager = getattr(env, "action_manager", None) + if manager is None: + return None + for term_name in getattr(manager, "active_terms", None) or []: + term = manager.get_term(term_name) + link_map = getattr(getattr(term, "cfg", None), "target_eef_link_names", None) + if not link_map or len(link_map) != len(eef_names): + continue + robot = env.scene[getattr(target_adapter, "robot_asset_name", "robot")] + body_names = list(robot.data.body_names) + idx_by_eef = {eef: body_names.index(link) for eef, link in zip(eef_names, link_map.values())} + + def _torch(x): + # Body poses come back as warp arrays (need .torch); joint/origin data may already be torch. + return x.torch if hasattr(x, "torch") else as_torch(x) + + def read_controlled_poses(env_id: int = 0) -> dict[str, torch.Tensor]: + pos_w = _torch(robot.data.body_pos_w) + quat_w = _torch(robot.data.body_quat_w) + origin = _torch(env.scene.env_origins)[env_id] + poses: dict[str, torch.Tensor] = {} + for eef, idx in idx_by_eef.items(): + pose = torch.eye(4, dtype=pos_w.dtype, device=pos_w.device) + pose[:3, :3] = math_utils.matrix_from_quat(quat_w[env_id, idx : idx + 1])[0] + pose[:3, 3] = pos_w[env_id, idx] - origin + poses[eef] = pose + return poses + + return read_controlled_poses + return None + + +def read_achieved_eef_poses( + target_adapter: EmbodimentAdapter, controlled_reader, env_id: int = 0 +) -> dict[str, torch.Tensor]: + """``{eef: (4,4)}`` achieved EEF poses at the controlled link (via ``controlled_reader``) or observed. + + The single frame both the tracking metric and the object-centric grasp transform must use: when the + reference is reconstructed at the IK-controlled link, ``controlled_reader`` reads there; otherwise + fall back to the adapter's observed EEF frame. Squeezes the adapter's ``(1,4,4)`` to ``(4,4)``. + """ + if controlled_reader is not None: + return controlled_reader(env_id) + return {eef: pose[0] for eef, pose in target_adapter.get_eef_poses(env_ids=[env_id]).items()} + + +def resolve_eef_reference_links( + target_adapter: EmbodimentAdapter, eef_reference_link: dict[str, str] | str | None +) -> dict[str, str] | str | None: + """Validate/normalize ``eef_reference_link`` (see :func:`.trajectory.source_eef_poses_at_link`). + + ``None`` -> feature off. ``"controlled"``/``"auto"`` -> the universal, cross-embodiment sentinel + (the controlled link is found per-episode by matching the source's own ``target_eef_pose``, so no + env introspection is needed). A dict -> an explicit ``{eef: link}`` override (same-embodiment); it + must cover every EEF. Resolved once, before the replay loop. + """ + if eef_reference_link is None or eef_reference_link in ("controlled", "auto"): + return eef_reference_link + assert isinstance( + eef_reference_link, dict + ), f"eef_reference_link must be a {{eef: link}} map or 'controlled'/'auto', got {eef_reference_link!r}." + eef_names = list(target_adapter.get_eef_names()) + missing = [eef for eef in eef_names if eef not in eef_reference_link] + assert not missing, f"eef_reference_link is missing entries for EEFs {missing}." + return {eef: eef_reference_link[eef] for eef in eef_names} + + +def _find_pink_ik_term(env: Any): + """Return the env's PinkIK action term (holds the per-env IK controllers), or None.""" + manager = getattr(env, "action_manager", None) + for term_name in (getattr(manager, "active_terms", None) or []) if manager is not None else []: + term = manager.get_term(term_name) + if getattr(term, "_ik_controllers", None) and getattr(term, "_isaaclab_controlled_joint_ids", None) is not None: + return term + return None + + +def _ik_solved_robot_state( + env: Any, + target_adapter: EmbodimentAdapter, + robot_asset_name: str, + base_state: dict, + first_pose_dict: dict[str, torch.Tensor], + first_passthrough_dict: dict[str, torch.Tensor], + env_id: int = 0, +) -> dict: + """Solve the IK for the first commanded pose and return a reset state starting the robot there. + + The env's PinkIK is a *differential* solver (it returns joint velocities), so we iterate its + ``compute`` in configuration space until the controlled joints converge — pure kinematics, with no + physics, gravity, or settling: the robot never actually moves. The hand/passthrough joints are not + IK'd, so they are taken straight from the first action. Returns ``base_state`` with the robot's + ``joint_position`` overwritten by the solved configuration, so ``reset_to`` begins the recorded demo + exactly at the IK solution of the first trajectory pose (works for any target embodiment). + + Solves for ``env_id`` only: the first-pose target is written into that env's row of the term action + (others zeroed) and only ``env_id``'s controller/joints are read, so this works per-env in the + parallel replay. The scratch task targets set on the other envs' controllers are overwritten by the + next real ``env.step``, so they are harmless. + """ + term = _find_pink_ik_term(env) + assert term is not None, "init_robot_from_ik requires a PinkIK action term to solve the first pose." + first_action = ( + target_adapter.target_eef_pose_to_action( + target_eef_pose_dict=first_pose_dict, passthrough_action_dict=first_passthrough_dict, env_id=env_id + ) + .reshape(-1) + .to(device=env.device) + ) + # The PinkIK term consumes only its own leading slice of the env action (EEF poses + hand joints); + # any other action terms (e.g. G1's 4-D body/locomotion channel) own the trailing columns. Write the + # first-pose slice into this env's row of a full (num_envs, term_dim) action so process_actions sets + # the right task target for env_id's controller. + term_dim = term._raw_actions.shape[-1] + pink_row = first_action[:term_dim] + full_action = torch.zeros((env.num_envs, term_dim), device=env.device, dtype=pink_row.dtype) + full_action[env_id] = pink_row + term.process_actions(full_action) # set the IK task targets to the first pose (base-relative) + + robot = env.scene[robot_asset_name] + controller = term._ik_controllers[env_id] + controlled_ids = list(term._isaaclab_controlled_joint_ids) + dt = getattr(term, "_sim_dt", None) or env.sim.get_physics_dt() + joints = as_torch(robot.data.joint_pos)[env_id].clone() # full joint vector, starting at target default + for _ in range(_IK_SOLVE_ITERATIONS): + solved = controller.compute(joints.detach().cpu().numpy(), dt) # differential step toward the target + joints[controlled_ids] = solved.to(device=joints.device, dtype=joints.dtype) + hand_ids = list(getattr(term, "_hand_joint_ids", []) or []) + if hand_ids: # hands are passthrough, not IK'd: take them from the term's trailing hand-joint slice + joints[hand_ids] = pink_row[-len(hand_ids) :].to(device=joints.device, dtype=joints.dtype) + + warm = deepcopy(base_state) + fields = warm["articulation"][robot_asset_name] + joint_pos = joints.unsqueeze(0) + if "joint_position" in fields and tuple(fields["joint_position"].shape) == tuple(joint_pos.shape): + fields["joint_position"] = joint_pos + if "joint_velocity" in fields: + fields["joint_velocity"] = torch.zeros_like(fields["joint_velocity"]) + return warm + + +def _slice_scene_state(state: dict, env_id: int) -> dict: + """Return a scene-state dict with every leaf tensor sliced to ``env_id``'s single row ((1, ...)). + + A freshly-read ``scene.get_state`` is shaped ``(num_envs, ...)``; ``reset_to`` (and the source + episode's ``(1, ...)`` overrides) need one env's row. For a single-env run (``env_id == 0``) this is + the identity slice. + + Most categories nest ``category -> name -> {field: tensor}`` (3 levels), but surface grippers nest + ``gripper -> name -> tensor`` (the entity value is a bare tensor, not a field dict), so slicing + recurses over dicts and slices any tensor leaf regardless of depth. + """ + + def _slice(node): + if isinstance(node, dict): + return {key: _slice(value) for key, value in node.items()} + if isinstance(node, torch.Tensor): + return node[env_id : env_id + 1] + return node + + return _slice(state) + + +def _build_offset_timeline(segs: list, num_steps: int) -> tuple[list, list, list]: + """Per-step ``(translation, axis_angle, frame)`` for one EEF from its ordered subtask offsets. + + Each subtask *holds* its offset over its core ``[start + interpolation_start, end - interpolation_end]``; + between cores the offset interpolates **linearly between the two adjacent subtasks' offsets** (not + through zero), so a 5 cm subtask followed by a 10 cm subtask blends 5 -> 10 across the boundary. A + leading ramp eases identity -> the first offset over the first ``interpolation_start``. The **last + subtask HOLDS its offset through the end** (and thus through any success-settle hold) -- it is not + decayed to identity, since there is no next offset to blend toward. A subtask with no ``offset`` is an + identity key, so a neighbour eases to/from zero into it. ``segs`` are dicts sorted by ``start``. + """ + zero = [0.0, 0.0, 0.0] + keyframes: list[tuple[int, list, list]] = [] + if segs[0]["i_start"] > 0: # ease in from identity to the first offset + keyframes.append((segs[0]["start"], zero, zero)) + for s in segs: + hstart = s["start"] + s["i_start"] + # The last subtask holds to its end (no ramp-down); interior subtasks reach their own hold end. + hend = s["end"] if s is segs[-1] else max(hstart, s["end"] - s["i_end"]) + hend = max(hstart, hend) + keyframes.append((hstart, s["trans"], s["aa"])) + if hend != hstart: + keyframes.append((hend, s["trans"], s["aa"])) + keyframes.sort(key=lambda k: k[0]) + + frame_by_step = [segs[-1]["frame"]] * num_steps + for s in segs: + for t in range(s["start"], min(s["end"], num_steps)): + frame_by_step[t] = s["frame"] + + trans_by_step: list = [zero] * num_steps + aa_by_step: list = [zero] * num_steps + j = 0 + for t in range(num_steps): + while j + 1 < len(keyframes) and keyframes[j + 1][0] <= t: + j += 1 + step0, tr0, aa0 = keyframes[j] + step1, tr1, aa1 = keyframes[j + 1] if j + 1 < len(keyframes) else keyframes[j] + alpha = 0.0 if (t <= step0 or step1 <= step0) else (1.0 if t >= step1 else (t - step0) / (step1 - step0)) + trans_by_step[t] = [tr0[k] + alpha * (tr1[k] - tr0[k]) for k in range(3)] + aa_by_step[t] = [aa0[k] + alpha * (aa1[k] - aa0[k]) for k in range(3)] + return trans_by_step, aa_by_step, frame_by_step + + +def _offset_se3(translation, axis_angle, scale: float, ref: torch.Tensor) -> torch.Tensor: + """Scaled SE(3) offset ``(4, 4)`` from ``translation`` [m] and ``axis_angle`` [rad] (either optional), + both linearly scaled by ``scale`` (angle scaled, so the rotation eases in). ``ref`` sets dtype/device.""" + mat = torch.eye(4, dtype=ref.dtype, device=ref.device) + if axis_angle is not None: + aa = torch.tensor(axis_angle, dtype=ref.dtype, device=ref.device) * scale + theta = torch.linalg.vector_norm(aa) + if float(theta) > 1e-8: + k = aa / theta # unit axis + skew = torch.zeros((3, 3), dtype=ref.dtype, device=ref.device) + skew[0, 1], skew[0, 2], skew[1, 0] = -k[2], k[1], k[2] + skew[1, 2], skew[2, 0], skew[2, 1] = -k[0], -k[1], k[0] + eye3 = torch.eye(3, dtype=ref.dtype, device=ref.device) + mat[:3, :3] = eye3 + torch.sin(theta) * skew + (1 - torch.cos(theta)) * (skew @ skew) + if translation is not None: + mat[:3, 3] = torch.tensor(translation, dtype=ref.dtype, device=ref.device) * scale + return mat + + +def apply_subtask_offsets( + target_eef_poses: dict[str, torch.Tensor], + subtasks: dict, + gripper_closed: dict[str, torch.Tensor], + signals: dict[str, torch.Tensor], + num_steps: int, + eef_name_map: dict[str, str] | None, + close_fraction: float, + source_objects: dict[str, torch.Tensor], +) -> dict[str, torch.Tensor]: + """Apply each subtask's ``offset`` to that EEF's commanded trajectory over the subtask's span. + + The offset ``O`` acts in a frame ``F`` as ``P' = F @ O @ inv(F) @ P``, which specializes to + ``world`` (``F = I`` -> ``O @ P``), ``eef``/``controlled`` (``F = P`` -> ``P @ O``, a nudge in the + gripper's own frame), and a tracked object name (``F`` = that object's live pose). The frame defaults + to the subtask's own reference (``object_ref`` else ``frame_ref`` else ``eef``). The offset is a + continuous per-EEF curve across subtasks (:func:`_build_offset_timeline`): boundaries interpolate + between adjacent subtask offsets (5 cm -> 10 cm blends 5 -> 10, not 5 -> 0 -> 10), eased by each + subtask's ``interpolation_start`` / ``interpolation_end``. + """ + per_eef: dict[str, list] = {} + for _eef_key, eef, _index, st, start, end, _boundary in iter_subtask_spans( + subtasks, gripper_closed, signals, num_steps, eef_name_map, close_fraction + ): + off = getattr(st, "offset", None) + per_eef.setdefault(eef, []).append( + { + "start": start, + "end": end, + "trans": list(off.translation) if (off and off.translation) else [0.0, 0.0, 0.0], + "aa": list(off.axis_angle) if (off and off.axis_angle) else [0.0, 0.0, 0.0], + "frame": (off.frame if (off and off.frame) else None) or st.object_ref or st.frame_ref or "eef", + "i_start": off.interpolation_start if off else 0, + "i_end": off.interpolation_end if off else 0, + "has_offset": off is not None, + } + ) + for eef, segs in per_eef.items(): + if eef not in target_eef_poses or not any(s["has_offset"] for s in segs): + continue + segs.sort(key=lambda s: s["start"]) + poses = target_eef_poses[eef] = target_eef_poses[eef].clone() # copy-on-write + num = poses.shape[0] + trans_by_step, aa_by_step, frame_by_step = _build_offset_timeline(segs, num) + eye4 = torch.eye(4, dtype=poses.dtype, device=poses.device) + for step in range(num): + trans, axis_angle = trans_by_step[step], aa_by_step[step] + if not any(trans) and not any(axis_angle): # identity offset -> leave pose untouched + continue + o_mat = _offset_se3(trans, axis_angle, 1.0, poses) + pose = poses[step] + frame = frame_by_step[step] + if frame in ("eef", "controlled"): + frame_pose = pose + elif source_objects and frame in source_objects: + frame_pose = source_objects[frame][step] + else: # "world" (or an unrecognized frame_ref -> world) + frame_pose = eye4 + poses[step] = frame_pose @ o_mat @ torch.linalg.inv(frame_pose) @ pose + return target_eef_poses + + +def prepare_episode( + env: Any, + env_id: int, + episode: EpisodeData, + target_adapter: EmbodimentAdapter, + source_adapter: EmbodimentAdapter, + remap_passthrough: PassthroughRemapper, + target_default_state: dict, + robot_asset_name: str, + reference_pose: str, + eef_name_map: dict[str, str] | None, + replay_speed: float, + retarget_frame: str, + scene_translation: tuple[float, float, float], + eef_offsets: dict[str, torch.Tensor] | None, + num_interpolation_steps: int, + init_robot_from_ik: bool, + subtasks: dict, + default_object_tracking: DefaultObjectTracking, + source_hand_postures: dict[str, dict[str, list[float]]] | None, + need_source_objects: bool, + reset_sim: bool, + eef_reference_link: dict[str, str] | str | None = None, + write_datagen_info: bool = False, + segment_settle_steps: int = 0, + max_eef_linear_velocity: float | None = None, + max_eef_rotation_speed: float | None = None, + synchronization: list | None = None, +) -> dict: + """Reset ``env_id`` to the source episode's scene and build its commanded trajectory. + + Shared by the sequential replay and the parallel workers. Everything here is env-side (reset) plus + pure trajectory math, keyed by ``env_id``. ``reset_sim`` triggers a global ``env.sim.reset()`` (only + valid single-env; parallel passes False and relies on per-env ``reset_to``). Returns a dict with the + commanded pose/passthrough sequences, step count, segment boundaries, carry segments, and (when + requested) the source object trajectories -- ready for the step loop. + """ + eef_names = list(target_adapter.get_eef_names()) + + # When eef_reference_link is set, rebuild the achieved ("eef_pose") reference at each EEF's + # IK-controlled link (e.g. GR1 hand_pitch) instead of the source's observed link (hand_roll), which + # can be a joint short of the gripper. Only "eef_pose" is a link-observed quantity; "target_eef_pose" + # is already the commanded controlled-link pose, so it needs no reconstruction. + def _achieved_eef_poses() -> dict[str, torch.Tensor]: + if eef_reference_link is None: + return source_datagen_poses(episode, eef_names, "eef_pose", eef_name_map) + body_names = list(env.scene[robot_asset_name].data.body_names) + return source_eef_poses_at_link(episode, eef_names, eef_reference_link, body_names, eef_name_map) + + if reference_pose == "eef_pose": + target_eef_poses = _achieved_eef_poses() + else: + target_eef_poses = source_datagen_poses(episode, eef_names, reference_pose, eef_name_map) + + source_actions = episode.data["actions"] + source_passthrough = source_adapter.actions_to_passthrough_actions(source_actions) + if eef_name_map: + source_passthrough = {eef_name_map.get(name, name): value for name, value in source_passthrough.items()} + target_passthrough = remap_passthrough(source_passthrough) + + source_objects = source_object_poses(episode) if need_source_objects else {} + # Signals are read when a subtask ends on a signal event and/or to forward them into the output. + needs_signals = write_datagen_info or any( + st.subtask_end is not None and st.subtask_end.method in ("signal_on", "signal_off") + for entries in subtasks.values() + for st in entries + ) + source_signals = source_subtask_signals(episode) if needs_signals else {} + # Per-EEF source gripper closedness (0 open -> 1 closed), keyed by target EEF name (source_passthrough + # was already renamed via eef_name_map); used to detect gripper_open/close subtask boundaries. + source_gripper_closed: dict[str, torch.Tensor] = {} + if source_hand_postures: + for eef in eef_names: + posture, hand = source_hand_postures.get(eef), source_passthrough.get(eef) + if posture is None or hand is None: + continue + src_open = torch.tensor(posture["open"], dtype=hand.dtype, device=hand.device) + src_close = torch.tensor(posture["close"], dtype=hand.dtype, device=hand.device) + source_gripper_closed[eef] = _hand_close_fraction(hand, src_open, src_close).unsqueeze(1) # (T, 1) + + if replay_speed != 1.0: + orig_num_steps = target_eef_poses[eef_names[0]].shape[0] + (target_eef_poses,), (target_passthrough,), _ = resample_trajectory( + [target_eef_poses], + [target_passthrough], + num_steps=orig_num_steps, + segment_ends=[], + replay_speed=replay_speed, + ) + # Per-step value dicts resampled nearest-neighbor (binary/gripper/signals stay sharp). + extras = [d for d in (source_objects,) if d] + value_dicts = [d for d in (source_gripper_closed, source_signals) if d] + if extras or value_dicts: + resampled, resampled_values, _ = resample_trajectory( + extras, value_dicts, num_steps=orig_num_steps, segment_ends=[], replay_speed=replay_speed + ) + resampled_iter, resampled_values_iter = iter(resampled), iter(resampled_values) + if source_objects: + source_objects = next(resampled_iter) + if source_gripper_closed: + source_gripper_closed = next(resampled_values_iter) + if source_signals: + source_signals = next(resampled_values_iter) + + # Reset this env to the source demo's task scene; the target robot starts at its default config. + env_ids = torch.tensor([env_id], device=env.device) + if reset_sim: + env.sim.reset() + env.recorder_manager.reset(env_ids=env_ids) + # ``reset_to`` expects a state shaped (len(env_ids), ...) == (1, ...); slice this env's row out of the + # (num_envs, ...) default so the source object poses (also (1, ...)) shape-match and get applied. + reset_state = retargeted_initial_state( + episode.data["initial_state"], + _slice_scene_state(target_default_state, env_id), + robot_asset_name, + object_translation=scene_translation, + ) + env.reset_to(reset_state, env_ids, is_relative=True) + + if retarget_frame == "robot_base": + source_root = as_torch(episode.data["states"]["articulation"][robot_asset_name]["root_pose"]).to(env.device) + source_base = poses_from_root_pose(source_root) + target_base = read_target_base_pose(env, robot_asset_name, env_id) + target_eef_poses = { + eef: reanchor_to_target_base(pose, source_base, target_base) for eef, pose in target_eef_poses.items() + } + + if any(scene_translation): + translation = torch.tensor( + scene_translation, dtype=target_eef_poses[eef_names[0]].dtype, device=target_eef_poses[eef_names[0]].device + ) + for eef in eef_names: + target_eef_poses[eef] = target_eef_poses[eef].clone() + target_eef_poses[eef][:, :3, 3] += translation + for name in source_objects: + source_objects[name] = source_objects[name].clone() + source_objects[name][:, :3, 3] += translation + + if eef_offsets is not None: + for eef in eef_names: + target_eef_poses[eef] = target_eef_poses[eef] @ eef_offsets[eef].to(target_eef_poses[eef].dtype) + + # Cap the commanded EEF speed for a reactive controller (e.g. Galbot RmpFlow): subdivide the shared + # timeline wherever any EEF moves faster than the caps, resampling every EEF pose, the passthrough, the + # tracked-object paths, gripper closedness, and signals together so they stay index-aligned. Measured on + # the final (post-offset) commanded poses. A differential-IK robot snaps to each waypoint in one step, + # so leave the caps ``None`` there. Boundaries below are then computed on this capped timeline. + if max_eef_linear_velocity is not None or max_eef_rotation_speed is not None: + target_eef_poses, (target_passthrough, source_gripper_closed, source_signals), source_objects = ( + cap_shared_timeline_speed( + target_eef_poses, + [target_passthrough, source_gripper_closed, source_signals], + source_objects, + max_eef_linear_velocity, + max_eef_rotation_speed, + ) + ) + + # Segmentation + object tracking come straight from the descriptor's subtasks, on the final (possibly + # speed-capped) timeline: each non-final subtask's end marks a per-EEF boundary (motion-aware settle), + # a subtask with object_tracking a carry segment, and a named subtask's end feeds the cross-EEF + # ``synchronization`` barriers so a per-EEF settle can't leave the arms desynchronized. + num_traj_steps = target_eef_poses[eef_names[0]].shape[0] + # Per-subtask SE(3) offsets: nudge each EEF's commanded pose over its subtask span (in the subtask's + # frame, eased in/out) before segmentation, so carry/settle are computed on the offset trajectory. + target_eef_poses = apply_subtask_offsets( + target_eef_poses, + subtasks, + source_gripper_closed, + source_signals, + num_traj_steps, + eef_name_map, + _CARRY_GRIPPER_CLOSED_FRACTION, + source_objects, + ) + carry_segments, segment_ends, name_end_step = carry_segments_and_boundaries_from_subtasks( + subtasks, + source_gripper_closed, + source_signals, + num_traj_steps, + eef_name_map=eef_name_map, + default_interp_start=default_object_tracking.interpolation_step_start, + default_interp_after=default_object_tracking.interpolation_step_after, + default_settle_steps=segment_settle_steps, + close_fraction=_CARRY_GRIPPER_CLOSED_FRACTION, + ) + sync_of, group_members = _build_sync_barriers(synchronization or [], name_end_step) + + if init_robot_from_ik: + warm_state = _ik_solved_robot_state( + env, + target_adapter, + robot_asset_name, + reset_state, + first_pose_dict={eef: target_eef_poses[eef][0] for eef in eef_names}, + first_passthrough_dict={name: tensor[0] for name, tensor in target_passthrough.items()}, + env_id=env_id, + ) + env.recorder_manager.reset(env_ids=env_ids) + env.reset_to(warm_state, env_ids, is_relative=True) + num_interpolation_steps = 0 + + commanded_poses, commanded_passthrough = _build_commanded_sequences( + target_adapter, target_eef_poses, target_passthrough, eef_names, num_interpolation_steps, env_id + ) + return { + "eef_names": eef_names, + "commanded_poses": commanded_poses, + "commanded_passthrough": commanded_passthrough, + "num_interpolation_steps": num_interpolation_steps, + "num_steps": commanded_poses[eef_names[0]].shape[0], + "carry_segments": carry_segments, + "source_objects": source_objects, + "source_signals": source_signals, # resampled to trajectory length; forwarded when write_datagen_info + "segment_ends": segment_ends, # {eef: {trajectory step: settle cap}} -- per-EEF motion-aware settle + "sync_of": sync_of, # {eef: {end step: group id}} -- cross-EEF rendezvous from ``synchronization`` + "group_members": group_members, # {group id: [(eef, end step)]} + } + + +def _monitored_error( + eef_name: str, + ts: int, + carry_segments: list, + source_objects: dict[str, torch.Tensor], + target_eef_pose_dict: dict[str, torch.Tensor], + achieved_poses: dict[str, torch.Tensor], + env: Any, + env_id: int, +) -> tuple[float, float, bool]: + """The tracking error to score for an EEF this step: the tracked **object** vs its source-path pose + during a carry (so a slipped grasp shows even though the EEF still follows the live grasp), else the + **EEF** vs its commanded target. Returns ``(pos [m], rot [deg], is_object)``. + """ + seg = next( + ( + s + for s in carry_segments + if s.eef == eef_name and s.start <= ts <= s.end and ts < source_objects[s.obj].shape[0] + ), + None, + ) + if seg is not None: + pos_err, rot_err = pose_tracking_error(source_objects[seg.obj][ts], _read_env_object_pose(env, seg.obj, env_id)) + return pos_err, rot_err, True + pos_err, rot_err = pose_tracking_error(target_eef_pose_dict[eef_name], achieved_poses[eef_name]) + return pos_err, rot_err, False + + +def _early_stop_hit( + pos_err: float, + rot_err: float, + is_object: bool, + max_translation_error: float | None, + max_rotation_error: float | None, + env_id: int, + tick: int, + eef_name: str, +) -> bool: + """True (and prints the reason) if the monitored error exceeds a threshold (``stop_early_on_failure``).""" + if (max_translation_error is not None and pos_err > max_translation_error) or ( + max_rotation_error is not None and rot_err > max_rotation_error + ): + tag = f"[retarget env{env_id}]" if env_id else "[retarget]" + print( + f"{tag} early stop at tick {tick + 1}: {eef_name} {'object' if is_object else 'eef'} miss " + f"{pos_err * 100:.1f}cm / {rot_err:.1f}deg exceeds threshold -> failure.", + flush=True, + ) + return True + return False + + +def _monitored_miss( + eef_names: list[str], + traj_step: dict[str, int], + ptr: dict[str, int], + lengths: dict[str, int], + carry_segments: list, + source_objects: dict[str, torch.Tensor], + target_eef_pose_dict: dict[str, torch.Tensor], + achieved_poses: dict[str, torch.Tensor], + env: Any, + env_id: int, + max_translation_error: float | None, + max_rotation_error: float | None, + tick: int, +) -> bool: + """True if any active EEF's monitored error (see :func:`_monitored_error`) exceeds the thresholds. + + Prints the reason. Used to early-abort a doomed replay (``stop_early_on_failure``) on the parallel path + (the single-env path scores + checks together in :func:`_score_step`). + """ + for eef_name in eef_names: + ts = traj_step[eef_name] + if ts < 0 or ptr[eef_name] >= lengths[eef_name]: + continue + pos_err, rot_err, is_object = _monitored_error( + eef_name, ts, carry_segments, source_objects, target_eef_pose_dict, achieved_poses, env, env_id + ) + if _early_stop_hit( + pos_err, rot_err, is_object, max_translation_error, max_rotation_error, env_id, tick, eef_name + ): + return True + return False + + +def _score_step( + eef_names: list[str], + traj_step: dict[str, int], + ptr: dict[str, int], + lengths: dict[str, int], + target_eef_pose_dict: dict[str, torch.Tensor], + achieved_poses: dict[str, torch.Tensor], + pos_errors: dict[str, list[float]], + rot_errors: dict[str, list[float]], + monitor_early: bool, + carry_segments: list, + source_objects: dict[str, torch.Tensor], + env: Any, + env_id: int, + max_translation_error: float | None, + max_rotation_error: float | None, + tick: int, +) -> bool: + """Record the per-EEF tracking error and test the early-abort thresholds (real, unfinished steps only). + + The recorded error is the **monitored** quantity (:func:`_monitored_error`): the tracked object vs its + source path during a carry, else the EEF vs its command -- so the per-episode report reflects how well + the *object* was tracked, not just how well the arm followed the (object-centric) command. Returns True + if the replay should early-abort as a failure (``stop_early_on_failure``). + """ + early_failure = False + for eef_name in eef_names: + if traj_step[eef_name] < 0 or ptr[eef_name] >= lengths[eef_name]: + continue + pos_err, rot_err, is_object = _monitored_error( + eef_name, + traj_step[eef_name], + carry_segments, + source_objects, + target_eef_pose_dict, + achieved_poses, + env, + env_id, + ) + pos_errors[eef_name].append(pos_err) + rot_errors[eef_name].append(rot_err) + if monitor_early and not early_failure: + early_failure = _early_stop_hit( + pos_err, rot_err, is_object, max_translation_error, max_rotation_error, env_id, tick, eef_name + ) + return early_failure + + +def replay_episode_on_target( + env: Any, + episode: EpisodeData | None, + source_adapter: EmbodimentAdapter | None, + target_adapter: EmbodimentAdapter, + remap_passthrough: PassthroughRemapper | None, + success_term: TerminationTermCfg | None, + robot_asset_name: str, + target_default_state: dict | None, + prep: dict | None = None, + num_interpolation_steps: int = 0, + init_robot_from_ik: bool = False, + replay_speed: float = 1.0, + success_settle_steps: int = 0, + segment_settle_steps: int = 0, + max_eef_linear_velocity: float | None = None, + max_eef_rotation_speed: float | None = None, + synchronization: list | None = None, + stop_early_on_failure: bool = False, + max_translation_error: float | None = None, + max_rotation_error: float | None = None, + settle_pos_tol_m: float = _SETTLE_POS_TOL_M, + settle_rot_tol_deg: float = _SETTLE_ROT_TOL_DEG, + settle_joint_tol: float = _JOINT_SETTLE_TOL, + retarget_frame: str = "world", + scene_translation: tuple[float, float, float] = (0.0, 0.0, 0.0), + eef_offsets: dict[str, torch.Tensor] | None = None, + eef_name_map: dict[str, str] | None = None, + reference_pose: str = "eef_pose", + subtasks: dict | None = None, + default_object_tracking: DefaultObjectTracking | None = None, + write_datagen_info: bool = False, + eef_reference_link: dict[str, str] | str | None = None, + source_hand_postures: dict[str, dict[str, list[float]]] | None = None, +) -> tuple[bool, dict[str, dict[str, torch.Tensor]], torch.Tensor | None]: + """Replay one source episode on the target embodiment and record the target rollout. + + Args: + env: The target env (with a recorder manager) to replay in. + episode: The source episode to retarget. + source_adapter: Adapter used to extract the source passthrough (gripper) actions. + target_adapter: Adapter used to re-encode each step into the target's action space. + remap_passthrough: Maps source passthrough actions to the target passthrough layout. + success_term: Optional success termination term evaluated each step. + robot_asset_name: Scene key of the target robot articulation. + target_default_state: Snapshot of the target env's home scene state (captured once from a + clean reset); the target robot starts each replay from its ``robot_asset_name`` entry. + num_interpolation_steps: Lead-in steps ramping each EEF from the target's start pose to the + first trajectory pose (0 disables). Ignored when ``init_robot_from_ik`` is set, since the + robot then already starts at the first pose. + replay_speed: Retime the extracted trajectory by ``1 / replay_speed`` waypoints per segment + (1.0 = unchanged; 0.5 = twice as many waypoints, replayed slower so the IK tracks a longer + source path more closely). Segment-boundary waypoints are preserved exactly. + init_robot_from_ik: Start the recorded demo at the IK solution of the first trajectory pose + (solved kinematically, no physics), so there is no startup transient. Works for any target + embodiment. + + Returns: + ``(task_succeeded, eef_errors)``. ``eef_errors`` maps each EEF to ``{"pos": tensor, "rot": + tensor}`` of per-step tracking errors [m], [deg] between the target robot's achieved EEF pose and + the (ideal) commanded pose it was driven to (measured at the controlled link when + ``eef_reference_link`` is set; lead-in steps excluded). ``task_succeeded`` is True if the success + condition held on any step (or no success term). + """ + subtasks = subtasks or {} + default_object_tracking = default_object_tracking or DefaultObjectTracking() + need_source_objects = any( + st.object_tracking is not None for entries in subtasks.values() for st in entries + ) + if prep is None: + prep = prepare_episode( + env, + 0, + episode, + target_adapter, + source_adapter, + remap_passthrough, + target_default_state, + robot_asset_name, + reference_pose, + eef_name_map, + replay_speed, + retarget_frame, + scene_translation, + eef_offsets, + num_interpolation_steps, + init_robot_from_ik, + subtasks, + default_object_tracking, + source_hand_postures=source_hand_postures, + need_source_objects=need_source_objects, + reset_sim=True, + eef_reference_link=eef_reference_link, + write_datagen_info=write_datagen_info, + segment_settle_steps=segment_settle_steps, + max_eef_linear_velocity=max_eef_linear_velocity, + max_eef_rotation_speed=max_eef_rotation_speed, + synchronization=synchronization, + ) + eef_names = prep["eef_names"] + commanded_poses = prep["commanded_poses"] + commanded_passthrough = prep["commanded_passthrough"] + num_interpolation_steps = prep["num_interpolation_steps"] + num_steps = prep["num_steps"] + carry_segments = prep["carry_segments"] + source_objects = prep["source_objects"] + source_signals = prep["source_signals"] # {name: (num_traj_steps, 1)} to forward, or {} + segment_ends = prep["segment_ends"] # {eef: {trajectory step: settle-hold cap}} for the motion-aware settle + sync_of = prep.get("sync_of", {}) # {eef: {end step: group id}} -- a cross-EEF rendezvous barrier + group_members = prep.get("group_members", {}) # {group id: [(eef, end step)]} + + # Where the achieved pose is read for the tracking error. When the reference is reconstructed at the + # IK-controlled link (eef_reference_link), the command drives that link, so the achieved must be read + # there too; reading the observed EEF frame would add the fixed observed->controlled offset (e.g. + # GR1's ~20 deg wrist_pitch: the angle between the hand and the forearm) and inflate the error even + # when the hand is aligned. Otherwise (no reconstruction) the observed frame is the reference frame. + controlled_reader = ( + _build_controlled_link_pose_reader(env, target_adapter) if eef_reference_link is not None else None + ) + if eef_reference_link is not None and controlled_reader is None: + print("\t (warning: eef_reference_link set but no PinkIK controlled link found; using observed frame)") + + def read_achieved() -> dict[str, torch.Tensor]: + return read_achieved_eef_poses(target_adapter, controlled_reader, 0) + + task_succeeded = success_term is None + pos_errors: dict[str, list[float]] = {eef_name: [] for eef_name in eef_names} + rot_errors: dict[str, list[float]] = {eef_name: [] for eef_name in eef_names} + + # Single-step per-EEF scheduler (replaces the nested ``_settle_at_pose`` loop): exactly one ``env.step`` + # per tick, each EEF advancing its OWN pointer through its (possibly unequal-length) commanded sequence. + # A subtask boundary is an inline HOLD -- the EEF repeats its pose and each tick checks its motion-aware + # settle (arm+gripper below tolerance), advancing once settled or at the cap; cross-EEF sync extends the + # same hold (M3). Pose / passthrough / object override are (re)computed only for EEFs that just arrived at + # a new waypoint, so a HOLDING EEF keeps its pose (re-running the override mid-release would remeasure the + # grasp and corrupt it). Aligned equal-length arms move in lockstep -> reproduces the pre-refactor path. + lengths = {eef_name: commanded_poses[eef_name].shape[0] for eef_name in eef_names} + ptr = {eef_name: 0 for eef_name in eef_names} + hold = {eef_name: 0 for eef_name in eef_names} # frames held so far at the current segment end + prev_pose: dict[str, torch.Tensor | None] = {eef_name: None for eef_name in eef_names} + prev_ptr = {eef_name: -1 for eef_name in eef_names} + prev_qpos = None + # Passthrough channels named after an EEF (the de-interleaved hands) advance with that EEF; any extra + # non-EEF channel (e.g. a mobile base) follows the first EEF's clock. + channel_eef = {name: name if name in eef_names else eef_names[0] for name in commanded_passthrough} + target_eef_pose_dict: dict[str, torch.Tensor] = {} + passthrough_action_dict: dict[str, torch.Tensor] = {} + record_signals = None + # When ``write_datagen_info`` (copy), record the full ``obs/datagen_info`` (poses + signals) each step so + # the output is a drop-in ``generate_dataset.py`` source. The scene's rigid objects (for ``object_pose``). + object_names = list(env.scene.rigid_objects.keys()) if write_datagen_info else [] + tick = 0 + early_failure = False + monitor_early = stop_early_on_failure and (max_translation_error is not None or max_rotation_error is not None) + while any(ptr[eef_name] < lengths[eef_name] for eef_name in eef_names): + idx = {eef_name: min(ptr[eef_name], lengths[eef_name] - 1) for eef_name in eef_names} + traj_step = {eef_name: idx[eef_name] - num_interpolation_steps for eef_name in eef_names} + advanced = {eef_name for eef_name in eef_names if ptr[eef_name] != prev_ptr[eef_name]} + + for eef_name in advanced: + target_eef_pose_dict[eef_name] = commanded_poses[eef_name][idx[eef_name]] + prev_ptr[eef_name] = ptr[eef_name] + for name, tensor in commanded_passthrough.items(): + if channel_eef[name] in advanced or name not in passthrough_action_dict: + passthrough_action_dict[name] = tensor[idx[channel_eef[name]]] + # Object-centric override for the EEFs that just advanced (a holding carrier keeps its pose). + if carry_segments and advanced: + _apply_object_centric_override( + carry_segments, + traj_step, + target_eef_pose_dict, + source_objects, + target_adapter, + env, + controlled_reader=controlled_reader, + commanded_poses=commanded_poses, + num_interpolation_steps=num_interpolation_steps, + eefs=advanced, + ) + ref_ts = max(traj_step.values()) # shared clock for signals only + signal_frame = {name: sig[min(max(ref_ts, 0), sig.shape[0] - 1)] for name, sig in source_signals.items()} + record_signals = (lambda sf=signal_frame: _record_signal_frame(env, 0, sf)) if signal_frame else None + + action = target_adapter.target_eef_pose_to_action( + target_eef_pose_dict=target_eef_pose_dict, + passthrough_action_dict=passthrough_action_dict, + env_id=0, + ) + env.step(action.reshape(1, -1).to(device=env.device)) + if record_signals is not None: + record_signals() + if write_datagen_info: + _record_datagen_poses(env, 0, target_adapter, target_eef_pose_dict, object_names) + if success_term is not None and bool(success_term.func(env, **success_term.params)[0]): + task_succeeded = True + + # Record per-EEF tracking error (real, unfinished steps only) and test the early-abort thresholds. + achieved_poses = read_achieved() + early_failure = _score_step( + eef_names, + traj_step, + ptr, + lengths, + target_eef_pose_dict, + achieved_poses, + pos_errors, + rot_errors, + monitor_early, + carry_segments, + source_objects, + env, + 0, + max_translation_error, + max_rotation_error, + tick, + ) + if early_failure: + break + + # Advance each EEF: a mid-segment step advances by one; at a segment end (a per-EEF boundary cap) + # HOLD until the arm+gripper settle or the cap is hit. + curr_poses = target_adapter.get_eef_poses(env_ids=[0]) + curr_qpos = as_torch(env.scene[robot_asset_name].data.joint_pos)[0] if robot_asset_name else None + joint_moved = ( + prev_qpos is not None + and curr_qpos is not None + and float(torch.max(torch.abs(curr_qpos - prev_qpos))) > settle_joint_tol + ) + for eef_name in eef_names: + if ptr[eef_name] >= lengths[eef_name]: + continue + cap = segment_ends.get(eef_name, {}).get(traj_step[eef_name], 0) + gid = sync_of.get(eef_name, {}).get(traj_step[eef_name]) # a sync barrier at this step, or None + if cap > 0 or gid is not None: # a hold point: motion-aware settle and/or a cross-EEF rendezvous + hold[eef_name] += 1 + settle_ok = True + if cap > 0: + moved = joint_moved + if prev_pose[eef_name] is not None: + dpos, drot = pose_tracking_error(prev_pose[eef_name], curr_poses[eef_name][0]) + moved = moved or dpos > settle_pos_tol_m or drot > settle_rot_tol_deg + settle_ok = (not moved) or hold[eef_name] >= cap + # Rendezvous: hold until every group member has reached its own barrier (concluded). ``ptr`` + # is in commanded-index space (lead-in included); the barrier step is a trajectory step. + sync_ok = gid is None or all( + ptr[m_eef] - num_interpolation_steps >= m_step for m_eef, m_step in group_members[gid] + ) + if settle_ok and sync_ok: + ptr[eef_name] += 1 + hold[eef_name] = 0 + else: + ptr[eef_name] += 1 + prev_pose[eef_name] = curr_poses[eef_name][0] + prev_qpos = curr_qpos.clone() if curr_qpos is not None else None + tick += 1 + + # Final success settle: after the last waypoint, hold the final commanded pose (which carries the + # final gripper *release*) for up to success_settle_steps, re-checking success each step. Unlike the + # per-waypoint/segment settle this deliberately does NOT stop on arm-settle: a suction cup needs + # several steps to reach fully-open (state -1) after release, and the cubes may still be settling, so + # cubes_stacked can only turn True a few steps after the arm has already stopped. Without this a + # perfectly stacked demo reads as failure because the gripper-open half of the success term is not + # yet satisfied when success is sampled. + if ( + not task_succeeded + and not early_failure + and success_settle_steps > 0 + and num_steps > 0 + and success_term is not None + ): + for _ in range(success_settle_steps): + action = target_adapter.target_eef_pose_to_action( + target_eef_pose_dict=target_eef_pose_dict, + passthrough_action_dict=passthrough_action_dict, + env_id=0, + ) + env.step(action.reshape(1, -1).to(device=env.device)) + if record_signals is not None: # hold the last waypoint's signal across the success settle + record_signals() + if write_datagen_info: + _record_datagen_poses(env, 0, target_adapter, target_eef_pose_dict, object_names) + if bool(success_term.func(env, **success_term.params)[0]): + task_succeeded = True + break + + eef_errors = { + eef_name: {"pos": torch.tensor(pos_errors[eef_name]), "rot": torch.tensor(rot_errors[eef_name])} + for eef_name in eef_names + } + return task_succeeded, eef_errors + + +def _build_sync_barriers( + synchronization: list[list[str]], name_end_step: dict[str, tuple[str, int]] +) -> tuple[dict[str, dict[int, int]], dict[int, list[tuple[str, int]]]]: + """Turn the descriptor's sync groups into the executor's per-EEF barrier lookups. + + ``name_end_step`` maps each named subtask to ``(eef, end_trajectory_step)`` -- the step where that + subtask concludes. Returns ``(sync_of, group_members)`` where ``sync_of[eef][end_step]`` is the group + id an EEF must rendezvous at when it reaches that step, and ``group_members[gid]`` lists every + ``(eef, end_step)`` in the group. At runtime an EEF holds at ``end_step`` until every member has reached + its own ``end_step`` (all ``ptr[m_eef] >= m_step``), then all advance -- the hand-off join. + """ + sync_of: dict[str, dict[int, int]] = {} + group_members: dict[int, list[tuple[str, int]]] = {} + for gid, group in enumerate(synchronization): + for name in group: + if name not in name_end_step: # a named subtask that produced no segment this run -- skip + continue + eef, end_step = name_end_step[name] + sync_of.setdefault(eef, {})[end_step] = gid + group_members.setdefault(gid, []).append((eef, end_step)) + return sync_of, group_members + + +def format_eef_errors(eef_errors: dict[str, dict[str, torch.Tensor]]) -> str: + """One line per EEF summarizing position/orientation tracking error and the matched rate.""" + lines: list[str] = [] + for eef_name, errors in eef_errors.items(): + pos = errors["pos"] + rot = errors["rot"] + matched = ((pos <= _IK_POS_TOL_M) & (rot <= _IK_ROT_TOL_DEG)).float().mean().item() * 100.0 + lines.append( + f"\t {eef_name}: pos err mean {pos.mean() * 100:.1f} / max {pos.max() * 100:.1f} cm | " + f"rot err mean {rot.mean():.1f} / max {rot.max():.1f} deg | " + f"matched (<{int(_IK_POS_TOL_M * 100)}cm,<{int(_IK_ROT_TOL_DEG)}deg) {matched:.0f}%" + ) + return "\n".join(lines) diff --git a/autodata_retargeting/replayer.py b/autodata_retargeting/replayer.py new file mode 100644 index 0000000..2783a2e --- /dev/null +++ b/autodata_retargeting/replayer.py @@ -0,0 +1,182 @@ +# Copyright (c) 2026, The Isaac AutoData Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""The Replayer: executes retargeting :class:`~.provider.Plan` objects on the target robot. + +Holds the execution context (target env + adapters + resolved retargeting artifacts) and pulls plans +from a :class:`~.provider.PlanProvider` until it is done, driving each on the target robot's IK and +recording the rollout. Knows nothing about *where* plans come from -- full-replay of a dataset today, a +scene generator tomorrow -- so the two concerns stay separate. Supports single-env (with the per-step +tracking-error report) and parallel multi-env replay. +""" + +from __future__ import annotations + +import torch +from dataclasses import dataclass +from typing import Any + +from .parallel_replay import ReplayParams, run_parallel_replay +from .provider import PlanProvider, ReplayResult +from .replay import format_eef_errors, replay_episode_on_target + + +@dataclass +class Replayer: + """Executes plans on the target embodiment. Built once per run from the resolved run context.""" + + env: Any + source_adapter: Any + target_adapter: Any + remap_passthrough: Any + config: Any # RetargetConfig + generation_policy: Any + success_term: Any + robot_asset_name: str + target_default_state: dict + eef_offsets: dict[str, torch.Tensor] | None + eef_reference_link: dict[str, str] | str | None + source_hand_postures: dict[str, dict[str, list[float]]] | None + output_file: str + + def _ref_label(self) -> str: + return ( + "source executed path (eef_pose)" + if self.config.reference_pose == "eef_pose" + else "commanded ideal (target_eef_pose)" + ) + + def replay_plan(self, plan) -> ReplayResult: + """Replay one plan on the target env (single-env), record it, and return its outcome.""" + succeeded, eef_errors = replay_episode_on_target( + env=self.env, + episode=plan.episode, + source_adapter=self.source_adapter, + target_adapter=self.target_adapter, + remap_passthrough=self.remap_passthrough, + success_term=self.success_term, + robot_asset_name=self.robot_asset_name, + target_default_state=self.target_default_state, + num_interpolation_steps=self.config.num_interpolation_steps, + init_robot_from_ik=self.config.init_robot_from_ik, + replay_speed=self.config.replay_speed, + success_settle_steps=self.config.success_settle_steps, + segment_settle_steps=self.config.segment_settle_steps, + max_eef_linear_velocity=self.config.max_eef_linear_velocity, + max_eef_rotation_speed=self.config.max_eef_rotation_speed, + synchronization=self.config.synchronization, + settle_pos_tol_m=self.config.settle_pos_tol_m, + settle_rot_tol_deg=self.config.settle_rot_tol_deg, + settle_joint_tol=self.config.settle_joint_tol, + retarget_frame=self.config.retarget_frame, + scene_translation=self.config.scene_translation, + eef_offsets=self.eef_offsets, + eef_name_map=self.config.eef_name_map, + reference_pose=self.config.reference_pose, + subtasks=self.config.subtasks, + default_object_tracking=self.config.default_object_tracking, + write_datagen_info=self.config.write_datagen_info, + eef_reference_link=self.eef_reference_link, + source_hand_postures=self.source_hand_postures, + stop_early_on_failure=self.config.stop_early_on_failure, + max_translation_error=self.config.max_translation_error, + max_rotation_error=self.config.max_rotation_error, + ) + # The recorder's export mode (from the generation policy's keep_failed flag) decides whether a + # failed replay is written; always report the outcome to it. + env_ids = torch.tensor([0], device=self.env.device) + self.env.recorder_manager.set_success_to_episodes( + env_ids, torch.tensor([[succeeded]], dtype=torch.bool, device=self.env.device) + ) + self.env.recorder_manager.export_episodes(env_ids) + return ReplayResult(success=succeeded, eef_errors=eef_errors) + + def _run_single(self, provider: PlanProvider, simulation_app: Any) -> int: + """Pull plans from ``provider`` and replay them one at a time until it is done.""" + ref_label = self._ref_label() + num_success = 0 + num_processed = 0 + all_eef_errors: list[dict[str, dict[str, torch.Tensor]]] = [] + with torch.inference_mode(): + while True: + if not simulation_app.is_running() or simulation_app.is_exiting(): + break + plan = provider.next() + if plan is None: + break + print(f"\nRetargeting example #{plan.index} ({plan.name})") + result = self.replay_plan(plan) + provider.observe(plan, result) + num_processed += 1 + num_success += int(result.success) + all_eef_errors.append(result.eef_errors) + print( + f"\t{'Succeeded' if result.success else 'Failed'}. Tracking error " + f"(object while carried, else EEF vs {ref_label}):" + ) + print(format_eef_errors(result.eef_errors)) + + rate = 100 * num_success / num_processed if num_processed > 0 else 0.0 + print( + f"\nRetargeted {num_processed} example{'s' if num_processed != 1 else ''} " + f"({num_success} successful, {rate:.1f}%) to {self.output_file}." + ) + if all_eef_errors: + eef_names = list(all_eef_errors[0].keys()) + overall = { + eef_name: {key: torch.cat([demo[eef_name][key] for demo in all_eef_errors]) for key in ("pos", "rot")} + for eef_name in eef_names + } + print(f"Overall tracking error (object while carried, else EEF vs {ref_label}, all examples):") + print(format_eef_errors(overall)) + return num_success + + def _run_parallel(self, provider: PlanProvider, num_envs: int) -> int: + """Replay plans across ``num_envs`` workers pulling from ``provider``.""" + params = ReplayParams( + robot_asset_name=self.robot_asset_name, + reference_pose=self.config.reference_pose, + eef_name_map=self.config.eef_name_map, + replay_speed=self.config.replay_speed, + retarget_frame=self.config.retarget_frame, + scene_translation=self.config.scene_translation, + eef_offsets=self.eef_offsets, + num_interpolation_steps=self.config.num_interpolation_steps, + init_robot_from_ik=self.config.init_robot_from_ik, + subtasks=self.config.subtasks, + default_object_tracking=self.config.default_object_tracking, + write_datagen_info=self.config.write_datagen_info, + eef_reference_link=self.eef_reference_link, + source_hand_postures=self.source_hand_postures, + success_settle_steps=self.config.success_settle_steps, + segment_settle_steps=self.config.segment_settle_steps, + settle_pos_tol_m=self.config.settle_pos_tol_m, + settle_rot_tol_deg=self.config.settle_rot_tol_deg, + settle_joint_tol=self.config.settle_joint_tol, + config=self.config, + ) + results = run_parallel_replay( + self.env, + provider, + num_envs, + (self.target_adapter, self.source_adapter, self.remap_passthrough), + self.target_default_state, + self.success_term, + params, + self.generation_policy, + ) + n_success = sum(1 for _, ok in results if ok) + rate = 100 * n_success / len(results) if results else 0.0 + print( + f"\nRetargeted {len(results)} example{'s' if len(results) != 1 else ''} " + f"({n_success} successful, {rate:.1f}%) to {self.output_file} across {num_envs} envs." + ) + return n_success + + def run(self, provider: PlanProvider, num_envs: int, simulation_app: Any) -> int: + """Drive ``provider`` to completion on the target env; return the number of successful replays.""" + if num_envs > 1: + return self._run_parallel(provider, num_envs) + return self._run_single(provider, simulation_app) diff --git a/autodata_retargeting/runner.py b/autodata_retargeting/runner.py new file mode 100644 index 0000000..f05f8b1 --- /dev/null +++ b/autodata_retargeting/runner.py @@ -0,0 +1,182 @@ +# Copyright (c) 2026, The Isaac AutoData Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Top-level orchestration: resolve config, build the env, replay every episode.""" + +import gymnasium as gym +import numpy as np +import random +import torch +from copy import deepcopy + +from isaaclab.utils.datasets import HDF5DatasetFileHandler + +from autodata_interfaces.embodiments import embodiment_adapter_from_yaml +from autodata_interfaces.env import setup_env_config, setup_output_paths +from autodata_interfaces.tasks.generation_policy_spec import GenerationPolicy + +from .config import RetargetConfig +from .eef_offset import build_eef_offsets, compose_retarget_eef_offsets +from .gripper_retargeting import build_passthrough_remapper, load_hand_postures +from .provider import DatasetReplayProvider +from .replay import resolve_eef_reference_links, validate_eef_agreement +from .replayer import Replayer + + +def _resolve_retarget_params(args) -> RetargetConfig: + """Resolve the task/pair retargeting parameters from ``--retarget_config`` or the individual flags.""" + if args.retarget_config: + return RetargetConfig.from_yaml(args.retarget_config) + missing = [ + name for name in ("source_embodiment", "target_embodiment", "target_env_name") if getattr(args, name) is None + ] + assert not missing, f"Without --retarget_config, these flags are required: {missing}" + # The retargeting knobs (hand policy, replay_speed, frame reconciliation, reference_pose, + # eef_reference_link, subtasks/object tracking, write_datagen_info, ...) are descriptor-only; a flag-only + # run takes their RetargetConfig defaults (empty subtasks -> no segmentation / object tracking). + return RetargetConfig( + source_embodiment=args.source_embodiment, + target_embodiment=args.target_embodiment, + target_env_name=args.target_env_name, + ) + + +def run(args, simulation_app) -> int: + """Retarget every source episode onto the target embodiment and export the results.""" + output_dir, output_file_name = setup_output_paths(args.output_file) + config = _resolve_retarget_params(args) + + # Generation policy is now built straight from the (self-contained) retarget config -- the retargeting + # path only reads seed and keep_failed (success_settle_steps is read directly from the config); + # num_trials/guarantee_success are superseded by the provider's stop condition + # (--target_successes / --target_runs). + generation_policy = GenerationPolicy( + name=config.name, + seed=config.seed, + keep_failed=args.keep_failed, + ) + + # Adapters. The source adapter only decodes recorded actions (no env needed); the target adapter + # encodes actions against the live target env and is bound after the env is created. + source_adapter = embodiment_adapter_from_yaml(config.source_embodiment) + target_adapter = embodiment_adapter_from_yaml(config.target_embodiment) + validate_eef_agreement(source_adapter, target_adapter, config.eef_name_map) + eef_names = list(target_adapter.get_eef_names()) + # Source hand_open/hand_close postures (keyed by the target EEF names), reused for both the hand + # policy and, under object_tracking="auto", the gripper-closed gate in carry detection. + source_hand_postures = load_hand_postures(config.source_embodiment, eef_names) + remap_passthrough = build_passthrough_remapper( + source_adapter, + target_adapter, + hand_policy=config.hand_policy, + hand_interp_norm=config.hand_interp_norm, + hand_binary_close_threshold=config.hand_binary_close_threshold, + hand_interp_band=config.hand_interp_band, + joint_mapping=config.joint_mapping, + source_hand_postures=source_hand_postures, + target_hand_postures=load_hand_postures(config.target_embodiment, eef_names), + target_channel_defaults=config.target_channel_defaults, + eef_name_map=config.eef_name_map, + ) + device = torch.device(args.device) + # EEF frame reconciliation: the descriptor's inline ``eef_offsets`` if given, else composed from each + # embodiment's own eef_offset (the general per-robot design). + if config.eef_offsets: + eef_offsets = build_eef_offsets(config.eef_offsets, eef_names, device) + else: + eef_offsets = compose_retarget_eef_offsets( + config.source_embodiment, config.target_embodiment, eef_names, device + ) + + num_envs = getattr(args, "num_envs", 1) + env_cfg, success_term = setup_env_config( + env_name=config.target_env_name, + output_dir=output_dir, + output_file_name=output_file_name, + num_envs=num_envs, + device=args.device, + generation_policy_params=generation_policy, + ) + + # Compensate the target env's (auto-detected) arm-action scale *in the encoded action* rather than + # changing the env: the adapter emits (delta / scale) so the env re-applying its own (unchanged) scale + # realizes the full commanded delta each step. The source per-step delta was physically achieved, so it + # is reachable in one step; without this the Franka IK-Rel scale=0.5 halves it and the robot lags. The + # compensation lives in the action, so the recorded action stays correct for the standard env. Only + # single-arm delta-pose adapters carry the ``action_scale`` field (bimanual/whole-body left untouched). + arm_action = getattr(env_cfg.actions, "arm_action", None) + action_scale = getattr(arm_action, "scale", None) + if isinstance(action_scale, (int, float)) and hasattr(target_adapter, "action_scale"): + target_adapter.action_scale = float(action_scale) + print(f"[retarget] compensating target arm-action scale={float(action_scale)} in the encoded action") + + env = gym.make(config.target_env_name, cfg=env_cfg).unwrapped + + random.seed(generation_policy.seed) + np.random.seed(generation_policy.seed) + torch.manual_seed(generation_policy.seed) + + try: + target_adapter.bind_env(env) + env.reset() + # Snapshot the target's home scene state once; each replay restores its robot from here + # (deepcopy so later sim steps cannot mutate the cached tensors). + target_default_state = deepcopy(env.scene.get_state(is_relative=True)) + + # Resolve the controlled-link reference reconstruction once (needs the env's action term to + # infer links when set to "controlled"); None (default) leaves the observed eef_pose in place. + eef_reference_link = resolve_eef_reference_links(target_adapter, config.eef_reference_link) + if eef_reference_link is not None: + where = "matched target_eef_pose per episode" if isinstance(eef_reference_link, str) else eef_reference_link + print( + f"[retarget] reconstructing eef_pose reference at controlled links ({where}); " + "tracking error is measured there too." + ) + + dataset_file_handler = HDF5DatasetFileHandler() + dataset_file_handler.open(args.input_file) + episode_names = list(dataset_file_handler.get_episode_names()) + if len(episode_names) == 0: + print("No episodes found in the source dataset.") + return 0 + + # Optionally retarget only a chosen subset of source episodes (by index into the dataset). + select_episodes = getattr(args, "select_episodes", None) or [] + if select_episodes: + out_of_range = [i for i in select_episodes if not 0 <= i < len(episode_names)] + assert ( + not out_of_range + ), f"--select_episodes {out_of_range} out of range; the source dataset has {len(episode_names)} episodes." + episode_names = [episode_names[i] for i in select_episodes] + print(f"[retarget] replaying {len(episode_names)} selected episode(s): {select_episodes}") + + # Provider: hands out examples and owns the stop condition (target successes / runs / exhaust). + # Copy retargeting replays each recorded source demo once on the target robot. + provider = DatasetReplayProvider( + dataset_file_handler, + episode_names, + env.device, + target_successes=getattr(args, "target_successes", None), + target_runs=getattr(args, "target_runs", None), + ) + # Replayer: executes each plan on the target robot (single-env or parallel). + replayer = Replayer( + env=env, + source_adapter=source_adapter, + target_adapter=target_adapter, + remap_passthrough=remap_passthrough, + config=config, + generation_policy=generation_policy, + success_term=success_term, + robot_asset_name=getattr(target_adapter, "robot_asset_name", "robot"), + target_default_state=target_default_state, + eef_offsets=eef_offsets, + eef_reference_link=eef_reference_link, + source_hand_postures=source_hand_postures, + output_file=args.output_file, + ) + return replayer.run(provider, num_envs, simulation_app) + finally: + env.close() diff --git a/autodata_retargeting/trajectory.py b/autodata_retargeting/trajectory.py new file mode 100644 index 0000000..262a7b2 --- /dev/null +++ b/autodata_retargeting/trajectory.py @@ -0,0 +1,406 @@ +# Copyright (c) 2026, The Isaac AutoData Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Source-trajectory extraction, subtask-segment detection, and replay-speed resampling.""" + +import itertools +import torch + +import isaaclab.utils.math as math_utils +from isaaclab.utils.datasets import EpisodeData + +from .util import quat_slerp_batch, se3_inverse + +# Component order that maps a recorded ``robot_links_state`` quaternion to Isaac Lab's ``(w, x, y, z)``. +# ``robot_links_state`` comes from ``body_link_pose_w`` and stores its quaternion in an order that this +# permutation reorders to ``(w, x, y, z)`` -- verified to reconstruct the datagen ``eef_pose`` exactly. +# :func:`source_eef_poses_at_link` re-validates it per dataset (matching the observed link), so a future +# format change fails loudly rather than silently producing a wrong reference. +_LINK_STATE_QUAT_ORDER = (1, 2, 3, 0) + + +def source_datagen_poses( + episode: EpisodeData, + eef_names: list[str], + key: str, + eef_name_map: dict[str, str] | None = None, +) -> dict[str, torch.Tensor]: + """Read a per-step, per-EEF ``(T, 4, 4)`` pose trajectory from a source episode's datagen info. + + ``key`` selects the datagen-info field: ``"target_eef_pose"`` (the commanded controller targets, + which drive the replay) or ``"eef_pose"`` (the source robot's *achieved* EEF poses, used as the + consistent-frame reference for the reproduction-error report). + + The source demo keys its poses by the *source* embodiment's EEF names; ``eef_name_map`` + (``{source_eef: target_eef}``) renames them onto the requested target/task ``eef_names`` (identity + when the pair shares names, e.g. bimanual left/right). Returned dict is keyed by ``eef_names``. + """ + obs = episode.data.get("obs", {}) + datagen_info = obs.get("datagen_info") if isinstance(obs, dict) else None + assert datagen_info is not None and key in datagen_info, ( + f"Source episode is missing 'obs/datagen_info/{key}'. Retargeting needs an annotated " + "dataset — run annotate_demos.py (or use a generate_dataset.py output) first." + ) + poses = datagen_info[key] + inverse = {tgt: src for src, tgt in (eef_name_map or {}).items()} + result: dict[str, torch.Tensor] = {} + for eef_name in eef_names: + source_key = inverse.get(eef_name, eef_name) + assert source_key in poses, f"datagen_info/{key} has no entry for eef '{source_key}'" + result[eef_name] = poses[source_key] + return result + + +def _link_state_pose(links_state: torch.Tensor, body_index: int) -> torch.Tensor: + """``(T, 4, 4)`` world pose of one body from a recorded ``robot_links_state`` ``(T, num_bodies, 13)``. + + Row layout is ``[pos(3), quat(4), lin_vel(3), ang_vel(3)]``; the quaternion is reordered to + ``(w, x, y, z)`` via :data:`_LINK_STATE_QUAT_ORDER` before building the rotation. + """ + row = links_state[:, body_index] + num_steps = row.shape[0] + pose = torch.eye(4, dtype=row.dtype, device=row.device).repeat(num_steps, 1, 1) + pose[:, :3, :3] = math_utils.matrix_from_quat(row[:, 3:7][:, list(_LINK_STATE_QUAT_ORDER)]) + pose[:, :3, 3] = row[:, :3] + return pose + + +def _match_body(links_state: torch.Tensor, target: torch.Tensor) -> tuple[int, float]: + """Body index whose recorded orientation best matches ``target``, plus that mean error [deg]. + + Matches on orientation only (translation-free), so it is unaffected by the world-vs-env-origin + offset between the world-frame link states and the env-relative datagen ``target``. This locates a + datagen pose (observed ``eef_pose`` or commanded ``target_eef_pose``) among the robot's bodies + without needing body names, and the match residual doubles as a check of the link-state quaternion + convention (:data:`_LINK_STATE_QUAT_ORDER`). + """ + target_quat = math_utils.quat_from_matrix(target[:, :3, :3]) + best_index, best_err = 0, float("inf") + for body_index in range(links_state.shape[1]): + quat = links_state[:, body_index, 3:7][:, list(_LINK_STATE_QUAT_ORDER)] + err = float(torch.rad2deg(math_utils.quat_error_magnitude(target_quat, quat)).mean()) + if err < best_err: + best_index, best_err = body_index, err + return best_index, best_err + + +def source_eef_poses_at_link( + episode: EpisodeData, + eef_names: list[str], + controlled_spec: dict[str, str] | str, + robot_body_names: list[str] | None = None, + eef_name_map: dict[str, str] | None = None, + observed_tol_deg: float = 5.0, + controlled_tol_deg: float = 20.0, +) -> dict[str, torch.Tensor]: + """Reconstruct the source *achieved* EEF poses at each EEF's IK-controlled link, from link states. + + The datagen ``eef_pose`` is recorded at whatever link the source env *observed* (e.g. GR1's + ``hand_roll_link``), which can be a joint short of the link the IK actually *controls* + (``hand_pitch_link``) -- dropping that joint's rotation from the reference. This rebuilds the + reference at the controlled link from the recorded per-step link poses:: + + achieved_controlled = datagen_eef_pose_observed @ inv(links[observed]) @ links[controlled] + + The datagen pose is the correct-frame anchor; the ``inv(observed) @ controlled`` term is the + intra-robot observed->controlled transform (the missing joint), and being relative it cancels the + world/env-origin offset of the link states. Both links are located in ``robot_links_state`` by + orientation match (:func:`_match_body`), which also validates the link-state format. + + ``controlled_spec`` selects how the controlled link is found: + + * ``"controlled"``/``"auto"`` -- match each EEF's datagen ``target_eef_pose`` (the commanded + controller target, i.e. the controlled frame by definition). Needs no body names, so it works + **cross-embodiment** (the source's controlled link is found purely from the source's own datagen). + * ``{eef_name: link_name}`` -- an explicit override; the named link is located by ``robot_body_names`` + order, so it requires **same-embodiment** retargeting (``robot_body_names`` must be the source + robot's, taken from the target env). + + Args: + controlled_spec: ``"controlled"``/``"auto"`` (match ``target_eef_pose``) or a ``{eef: link}`` map. + robot_body_names: Source robot body names in ``robot_links_state`` order (only for the map form). + observed_tol_deg: Max orientation error [deg] tolerated when locating the observed link. + controlled_tol_deg: Max error [deg] when matching the controlled link to ``target_eef_pose`` + (looser than ``observed_tol_deg``: the achieved link lags its command by the tracking error). + """ + obs = episode.data.get("obs", {}) + links_state = obs.get("robot_links_state") if isinstance(obs, dict) else None + assert links_state is not None, ( + "eef_reference_link reconstruction needs 'obs/robot_links_state' (per-step link poses) in the " + "source demo; this env's dataset does not record it. Use reference_pose without eef_reference_link." + ) + match_controlled = isinstance(controlled_spec, str) + datagen_observed = source_datagen_poses(episode, eef_names, "eef_pose", eef_name_map) + datagen_commanded = ( + source_datagen_poses(episode, eef_names, "target_eef_pose", eef_name_map) if match_controlled else {} + ) + if not match_controlled: + assert robot_body_names is not None and links_state.shape[1] == len(robot_body_names), ( + "an explicit {eef: link} eef_reference_link needs the source robot's body names in link-state " + f"order (got {None if robot_body_names is None else len(robot_body_names)} for " + f"{links_state.shape[1]} bodies); this requires same-embodiment retargeting." + ) + + result: dict[str, torch.Tensor] = {} + for eef_name in eef_names: + anchor = datagen_observed[eef_name] + observed_index, err = _match_body(links_state, anchor) + assert err <= observed_tol_deg, ( + f"could not locate eef {eef_name!r}'s observed link in the recorded link states (nearest body " + f"off by {err:.1f} deg > {observed_tol_deg}); the link-state quaternion convention may differ " + "for this dataset." + ) + if match_controlled: + controlled_index, cerr = _match_body(links_state, datagen_commanded[eef_name]) + assert cerr <= controlled_tol_deg, ( + f"could not locate eef {eef_name!r}'s controlled link (matching target_eef_pose): nearest " + f"body off by {cerr:.1f} deg > {controlled_tol_deg}." + ) + else: + controlled_link = controlled_spec[eef_name] + assert controlled_link in robot_body_names, ( + f"controlled link {controlled_link!r} for eef {eef_name!r} is not a robot body " + f"(e.g. {robot_body_names[:3]}...)." + ) + controlled_index = robot_body_names.index(controlled_link) + relative = se3_inverse(_link_state_pose(links_state, observed_index)) @ _link_state_pose( + links_state, controlled_index + ) + result[eef_name] = anchor @ relative + return result + + +def source_object_poses(episode: EpisodeData) -> dict[str, torch.Tensor]: + """Read the source demo's recorded per-object pose trajectories (env-relative ``(T, 4, 4)`` SE(3)). + + Returns ``{object_name: (T, 4, 4)}`` from ``obs/datagen_info/object_pose`` (e.g. ``cube_1/2/3`` for + cube-stack), or ``{}`` if the annotated source carries no object poses. Used only by the debug + overlay that draws where each object *was* in the source demo, to see whether the retargeted + object follows that path or deviates. + """ + obs = episode.data.get("obs", {}) + datagen_info = obs.get("datagen_info") if isinstance(obs, dict) else None + object_pose = datagen_info.get("object_pose") if isinstance(datagen_info, dict) else None + if not isinstance(object_pose, dict): + return {} + return {name: traj for name, traj in object_pose.items()} + + +def source_subtask_signals(episode: EpisodeData) -> dict[str, torch.Tensor]: + """Read the source demo's per-step subtask-term signals as ``{name: (T, 1)}`` float tensors. + + From ``obs/datagen_info/subtask_term_signals`` (each stored ``(1, T)``); reshaped to ``(T, 1)`` so it + resamples alongside the other per-step values. Empty when the source carries no signals. Used by the + signal-driven object tracking (``object_tracking="signals"``) to mark when an object is attached. + """ + obs = episode.data.get("obs", {}) + datagen_info = obs.get("datagen_info") if isinstance(obs, dict) else None + signals = datagen_info.get("subtask_term_signals") if isinstance(datagen_info, dict) else None + if not isinstance(signals, dict): + return {} + return {name: value.reshape(-1, 1).to(torch.float32) for name, value in signals.items()} + + +def _sample_indices(num_src: int, new_len: int, device: torch.device, dtype: torch.dtype): + """Fractional resampling indices over ``num_src`` waypoints: bracketing index ``i0`` and lerp ``frac``.""" + t = torch.linspace(0.0, num_src - 1, new_len, device=device, dtype=dtype) + i0 = t.floor().long().clamp(max=num_src - 2) + return i0, (t - i0.to(dtype)).unsqueeze(-1) + + +def _resample_pose_trajectory(poses: torch.Tensor, new_len: int) -> torch.Tensor: + """Resample an SE(3) pose trajectory ``(T, 4, 4)`` to ``new_len`` waypoints along its polyline. + + Position is lerped and rotation slerped between bracketing source waypoints; the first and last + poses are preserved exactly. + """ + num_src = poses.shape[0] + if num_src < 2 or new_len == num_src: + return poses + i0, frac = _sample_indices(num_src, new_len, poses.device, poses.dtype) + p0, p1 = poses[i0], poses[i0 + 1] + out = torch.eye(4, device=poses.device, dtype=poses.dtype).repeat(new_len, 1, 1) + out[:, :3, 3] = torch.lerp(p0[:, :3, 3], p1[:, :3, 3], frac) + q0 = math_utils.quat_from_matrix(p0[:, :3, :3]) + q1 = math_utils.quat_from_matrix(p1[:, :3, :3]) + out[:, :3, :3] = math_utils.matrix_from_quat(quat_slerp_batch(q0, q1, frac)) + return out + + +def _resample_values(values: torch.Tensor, new_len: int) -> torch.Tensor: + """Nearest-neighbor (step-hold) resample of a per-step value trajectory ``(T, W)`` to ``new_len``. + + These are gripper / passthrough commands, so a discrete open<->close must stay **sharp**. Linear + interpolation would ramp the transition (a binary ``-1<->+1`` through 0, a parallel target + ``0.04<->0.0`` through half-open); at ``replay_speed < 1`` that ramps the close over the extra + waypoints, so the gripper shuts gradually instead of snapping and fails to secure the object while + the arm moves on. Nearest-neighbor keeps each command at its discrete value. Endpoints are kept. + """ + num_src = values.shape[0] + if num_src < 2 or new_len == num_src: + return values + t = torch.linspace(0.0, num_src - 1, new_len, device=values.device, dtype=values.dtype) + nearest = t.round().long().clamp(max=num_src - 1) + return values[nearest] + + +def _cap_trajectory_speed( + poses: torch.Tensor, + values: dict[str, torch.Tensor], + max_lin: float | None, + max_rot_deg: float | None, + max_expand: float = 10.0, +) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """Subdivide a pose trajectory so no step exceeds ``max_lin`` [m] / ``max_rot_deg`` [deg] of motion. + + Each interval ``i -> i+1`` is split into ``k_i = ceil(max(dpos/max_lin, dang/max_rot_deg))`` equal + sub-steps (position lerp, orientation slerp), so every replayed step stays under the per-waypoint speed + caps -- a reactive controller then never receives a command faster than it can physically track. Only + the axes with a non-``None`` cap constrain ``k``. ``values`` (gripper/passthrough) are nearest-neighbor + resampled onto the new timeline so a discrete open<->close stays sharp. Total length is bounded to + ``max_expand * T`` (uniform-resample fallback) so a tiny cap cannot explode the trajectory. Returns the + inputs unchanged when both caps are ``None``, ``T < 2``, or the trajectory is already within the caps. + """ + num = poses.shape[0] + if num < 2 or (max_lin is None and max_rot_deg is None): + return poses, values + pos = poses[:, :3, 3] + quat = math_utils.quat_from_matrix(poses[:, :3, :3]) + dpos = torch.linalg.vector_norm(pos[1:] - pos[:-1], dim=1) # (num-1,) per-step translation [m] + dots = (quat[1:] * quat[:-1]).sum(dim=1).abs().clamp(max=1.0) + dang_deg = torch.rad2deg(2.0 * torch.acos(dots)) # (num-1,) per-step rotation [deg] + k = torch.ones(num - 1, dtype=torch.long, device=poses.device) + if max_lin is not None: + k = torch.maximum(k, torch.ceil(dpos / max(max_lin, 1e-9)).long()) + if max_rot_deg is not None: + k = torch.maximum(k, torch.ceil(dang_deg / max(max_rot_deg, 1e-9)).long()) + total = int(k.sum()) + 1 + if total <= num: + return poses, values # already within the caps -- nothing to subdivide + max_len = int(max_expand * num) + if total > max_len: # tiny-cap guard: fall back to a uniform resample at the length ceiling + return _resample_pose_trajectory(poses, max_len), {c: _resample_values(v, max_len) for c, v in values.items()} + # Sampling indices: ``k_i`` evenly-spaced fractions per interval, then the final pose exactly. + i0_list: list[int] = [] + frac_list: list[float] = [] + for i in range(num - 1): + ki = int(k[i]) + for j in range(ki): + i0_list.append(i) + frac_list.append(j / ki) + i0_list.append(num - 2) + frac_list.append(1.0) + i0 = torch.tensor(i0_list, device=poses.device) + frac = torch.tensor(frac_list, device=poses.device, dtype=poses.dtype).unsqueeze(-1) # (out, 1) + out = torch.eye(4, device=poses.device, dtype=poses.dtype).repeat(len(i0_list), 1, 1) + out[:, :3, 3] = torch.lerp(pos[i0], pos[i0 + 1], frac) + out[:, :3, :3] = math_utils.matrix_from_quat(quat_slerp_batch(quat[i0], quat[i0 + 1], frac)) + nearest = (i0 + frac.squeeze(-1).round().long()).clamp(max=num - 1) # step-hold gripper onto new timeline + new_values = {channel: value[nearest] for channel, value in values.items()} + return out, new_values + + +def cap_shared_timeline_speed( + poses: dict[str, torch.Tensor], + value_dicts: list[dict[str, torch.Tensor]], + objects: dict[str, torch.Tensor], + max_lin: float | None, + max_rot_deg: float | None, + max_expand: float = 10.0, +) -> tuple[dict[str, torch.Tensor], list[dict[str, torch.Tensor]], dict[str, torch.Tensor]]: + """Cap the commanded EEF speed on a SHARED multi-EEF timeline (the copy path). + + Unlike :func:`_cap_trajectory_speed` (one trajectory), copy replays every EEF on one shared timeline, + so the subdivision is computed once as the per-interval **maximum over every EEF's** motion and applied + to all EEF poses, their passthrough ``value_dicts`` (nearest-neighbor, keeping a discrete open<->close + sharp), and the tracked ``objects`` (SE(3)) together -- so they stay index-aligned and no step exceeds + ``max_lin`` [m] / ``max_rot_deg`` [deg]. Total length is bounded to ``max_expand * T``. Returns the + inputs unchanged when both caps are ``None``, ``T < 2``, or the timeline is already within the caps. + """ + if (max_lin is None and max_rot_deg is None) or not poses: + return poses, value_dicts, objects + any_traj = next(iter(poses.values())) + num = any_traj.shape[0] + if num < 2: + return poses, value_dicts, objects + # Per-interval subdivision = the max over every EEF (the fastest EEF sets the shared timeline). + k = torch.ones(num - 1, dtype=torch.long, device=any_traj.device) + for traj in poses.values(): + pos = traj[:, :3, 3] + quat = math_utils.quat_from_matrix(traj[:, :3, :3]) + dpos = torch.linalg.vector_norm(pos[1:] - pos[:-1], dim=1) + dots = (quat[1:] * quat[:-1]).sum(dim=1).abs().clamp(max=1.0) + dang_deg = torch.rad2deg(2.0 * torch.acos(dots)) + if max_lin is not None: + k = torch.maximum(k, torch.ceil(dpos / max(max_lin, 1e-9)).long()) + if max_rot_deg is not None: + k = torch.maximum(k, torch.ceil(dang_deg / max(max_rot_deg, 1e-9)).long()) + total = int(k.sum()) + 1 + if total <= num: + return poses, value_dicts, objects + max_len = int(max_expand * num) + if total > max_len: # tiny-cap guard: uniform resample at the length ceiling + return ( + {e: _resample_pose_trajectory(t, max_len) for e, t in poses.items()}, + [{n: _resample_values(v, max_len) for n, v in d.items()} for d in value_dicts], + {n: _resample_pose_trajectory(t, max_len) for n, t in objects.items()}, + ) + i0_list: list[int] = [] + frac_list: list[float] = [] + for i in range(num - 1): + ki = int(k[i]) + for j in range(ki): + i0_list.append(i) + frac_list.append(j / ki) + i0_list.append(num - 2) + frac_list.append(1.0) + i0 = torch.tensor(i0_list, device=any_traj.device) + frac = torch.tensor(frac_list, device=any_traj.device, dtype=any_traj.dtype).unsqueeze(-1) # (out, 1) + nearest = (i0 + frac.squeeze(-1).round().long()).clamp(max=num - 1) + + def _interp(traj: torch.Tensor) -> torch.Tensor: + pos, quat = traj[:, :3, 3], math_utils.quat_from_matrix(traj[:, :3, :3]) + out = torch.eye(4, device=traj.device, dtype=traj.dtype).repeat(len(i0_list), 1, 1) + out[:, :3, 3] = torch.lerp(pos[i0], pos[i0 + 1], frac) + out[:, :3, :3] = math_utils.matrix_from_quat(quat_slerp_batch(quat[i0], quat[i0 + 1], frac)) + return out + + return ( + {e: _interp(t) for e, t in poses.items()}, + [{n: v[nearest] for n, v in d.items()} for d in value_dicts], + {n: _interp(t) for n, t in objects.items()}, + ) + + +def resample_trajectory( + pose_dicts: list[dict[str, torch.Tensor]], + value_dicts: list[dict[str, torch.Tensor]], + num_steps: int, + segment_ends: list[int], + replay_speed: float, +) -> tuple[list[dict[str, torch.Tensor]], list[dict[str, torch.Tensor]], list[int]]: + """Retime the trajectory by ``1 / replay_speed`` waypoints, resampling each segment independently. + + Splits ``[0, num_steps - 1]`` at the ``segment_ends`` (subtask boundaries) and resamples each chunk + so its span scales by ``1 / replay_speed`` (e.g. 0.5 -> twice as many waypoints, replayed slower), + keeping every segment-boundary waypoint exact. ``pose_dicts`` are resampled as SE(3), ``value_dicts`` + (passthrough) linearly. Returns the resampled dicts plus the boundaries mapped to new indices. + """ + breaks = sorted({0, num_steps - 1} | {b for b in segment_ends if 0 < b < num_steps - 1}) + chunks = list(zip(breaks[:-1], breaks[1:])) + new_spans = [max(1, round((end - start) / replay_speed)) for start, end in chunks] + new_bounds = list(itertools.accumulate(new_spans)) + new_segment_ends = new_bounds[:-1] # internal boundaries only (drop the final trajectory end) + + def resample(traj: torch.Tensor, fn) -> torch.Tensor: + pieces = [] + for index, (start, end) in enumerate(chunks): + piece = fn(traj[start : end + 1], new_spans[index] + 1) + pieces.append(piece if index == 0 else piece[1:]) # boundary waypoint is shared; keep it once + return torch.cat(pieces, dim=0) + + out_poses = [{eef: resample(traj, _resample_pose_trajectory) for eef, traj in d.items()} for d in pose_dicts] + out_values = [{name: resample(traj, _resample_values) for name, traj in d.items()} for d in value_dicts] + return out_poses, out_values, new_segment_ends diff --git a/autodata_retargeting/util.py b/autodata_retargeting/util.py new file mode 100644 index 0000000..d90ef49 --- /dev/null +++ b/autodata_retargeting/util.py @@ -0,0 +1,79 @@ +# Copyright (c) 2026, The Isaac AutoData Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Generic pose / quaternion math helpers shared across the retargeting pipeline.""" + +import torch +from typing import Any + +import isaaclab.utils.math as math_utils + +from autodata_utils.tensor_utils import as_torch + + +def se3_inverse(pose: torch.Tensor) -> torch.Tensor: + """Inverse of a batched SE(3) homogeneous transform ``(..., 4, 4)`` (``R^T``, ``-R^T t``).""" + rot = pose[..., :3, :3] + trans = pose[..., :3, 3:] + rot_t = rot.transpose(-1, -2) + out = torch.zeros_like(pose) + out[..., :3, :3] = rot_t + out[..., :3, 3:] = -rot_t @ trans + out[..., 3, 3] = 1.0 + return out + + +def pose_tracking_error(commanded: torch.Tensor, achieved: torch.Tensor) -> tuple[float, float]: + """Return ``(position error [m], orientation error [deg])`` between two ``(4, 4)`` poses.""" + pos_err = float(torch.linalg.norm(commanded[:3, 3] - achieved[:3, 3]).item()) + quat_cmd = math_utils.quat_from_matrix(commanded[:3, :3].unsqueeze(0)) + quat_ach = math_utils.quat_from_matrix(achieved[:3, :3].unsqueeze(0)) + rot_err_deg = float(torch.rad2deg(math_utils.quat_error_magnitude(quat_cmd, quat_ach))[0].item()) + return pos_err, rot_err_deg + + +def quat_slerp_batch(q0: torch.Tensor, q1: torch.Tensor, frac: torch.Tensor) -> torch.Tensor: + """Batched SLERP between quaternion rows ``q0``/``q1`` (N, 4) by ``frac`` (N, 1), in (x,y,z,w).""" + dot = (q0 * q1).sum(dim=-1, keepdim=True) + q1 = torch.where(dot < 0, -q1, q1) # take the shorter arc + dot = dot.abs().clamp(max=1.0) + angle = torch.acos(dot) + sin = torch.sin(angle) + small = sin.abs() < 1e-6 # nearly parallel -> fall back to (normalized) linear + w0 = torch.where(small, 1.0 - frac, torch.sin((1.0 - frac) * angle) / sin) + w1 = torch.where(small, frac, torch.sin(frac * angle) / sin) + return math_utils.normalize(w0 * q0 + w1 * q1) + + +def poses_from_root_pose(root_pose: torch.Tensor) -> torch.Tensor: + """Convert a recorded ``root_pose`` ``(T, 7)`` = [pos(3), quat wxyz(4)] into ``(T, 4, 4)`` poses.""" + num_steps = root_pose.shape[0] + poses = torch.eye(4, dtype=root_pose.dtype, device=root_pose.device).repeat(num_steps, 1, 1) + poses[:, :3, :3] = math_utils.matrix_from_quat(root_pose[:, 3:7]) + poses[:, :3, 3] = root_pose[:, :3] + return poses + + +def read_target_base_pose(env: Any, robot_asset_name: str, env_id: int = 0) -> torch.Tensor: + """Read the target robot base (root) pose in the env-relative frame as a ``(4, 4)`` tensor.""" + robot = env.scene[robot_asset_name] + pos_w = robot.data.root_pos_w + quat_w = robot.data.root_quat_w + pos_w = pos_w.torch if hasattr(pos_w, "torch") else as_torch(pos_w) + quat_w = quat_w.torch if hasattr(quat_w, "torch") else as_torch(quat_w) + origin = env.scene.env_origins + origin = origin.torch if hasattr(origin, "torch") else as_torch(origin) + pose = torch.eye(4, dtype=pos_w.dtype, device=pos_w.device) + pose[:3, :3] = math_utils.matrix_from_quat(quat_w[env_id : env_id + 1])[0] + pose[:3, 3] = pos_w[env_id] - origin[env_id] + return pose + + +def reanchor_to_target_base(poses: torch.Tensor, source_base: torch.Tensor, target_base: torch.Tensor) -> torch.Tensor: + """Re-express world ``poses`` relative to ``source_base`` (per step), then anchor to ``target_base``. + + ``poses`` and ``source_base`` are ``(T, 4, 4)``; ``target_base`` is ``(4, 4)``. Returns ``(T, 4, 4)``. + """ + return target_base.unsqueeze(0) @ torch.linalg.inv(source_base) @ poses diff --git a/scripts/retarget_dataset.py b/scripts/retarget_dataset.py new file mode 100755 index 0000000..2f199eb --- /dev/null +++ b/scripts/retarget_dataset.py @@ -0,0 +1,156 @@ +#!/usr/bin/env python +# Copyright (c) 2026, The Isaac AutoData Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Cross-embodiment dataset retargeting entrypoint (replay-based). + +Usage (config-based — the task/pair parameters live in a retarget descriptor YAML):: + + python scripts/retarget_dataset.py \\ + --retarget_config \\ + --input_file \\ + --output_file \\ + [--keep_failed] + +The retarget descriptor (see ``autodata_examples/retarget/``) is self-contained: it references the +source/target embodiments and target env, declares its own per-EEF ``subtasks`` (segmentation + object +tracking), and carries every retargeting knob (hand_policy / hand_interp_* / joint_mapping, +reference_pose, eef_reference_link, write_datagen_info, init_robot_from_ik, replay_speed, retarget_frame, +scene_translation, eef_offsets, ...) -- all descriptor-only. Without ``--retarget_config`` only the three +references (``--source_embodiment``/``--target_embodiment``/``--target_env_name``) build a config and +every other knob (including subtasks) takes its default, so a descriptor is needed to set them. + +Retargeting transfers each demonstration recorded on a *source* embodiment onto a *target* +embodiment, one output demo per input demo (replay 1:1). Unlike +:mod:`scripts.generate_dataset`, no object-centric regeneration happens: the source +end-effector trajectory is replayed on the target robot and the resulting rollout is recorded. + +How it works, per source episode: + +* The recorded per-step target EEF pose trajectory (``obs/datagen_info/target_eef_pose``) is an + absolute, env-relative SE(3) trajectory and is embodiment-independent. The source dataset must + therefore already carry ``datagen_info`` (i.e. it went through ``annotate_demos.py`` or was + produced by ``generate_dataset.py``). +* The gripper/hand "passthrough" channels are extracted from the source actions with the *source* + embodiment adapter and remapped to the *target* embodiment's passthrough layout (see + :func:`build_passthrough_remapper`). v1 supports embodiments whose passthrough layouts match; + differing grippers (e.g. parallel gripper → dexterous hand) are a documented extension seam. +* The target env is reset so the *task scene* (objects) matches the source episode's initial state + while the *target robot* starts from its own default configuration, then each step's target EEF + pose + passthrough actions are re-encoded into the target embodiment's action via + :meth:`EmbodimentAdapter.target_eef_pose_to_action` and stepped through the target env. + +The task descriptor is shared with the source dataset (same subtasks / EEF names) and is used here +for the generation policy (export mode, seed) and to validate that both embodiments agree on the +task's EEF names. + +Note: the output records target actions + states. Re-run ``annotate_demos.py`` on it before using +it as a source for ``generate_dataset.py`` (the subtask signals are not carried over). +""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse + +from isaaclab.app import AppLauncher + +parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) +parser.add_argument( + "--retarget_config", + type=str, + default=None, + help=( + "Path to a retarget descriptor YAML bundling the task/pair-specific parameters (task " + "descriptor, source/target embodiments, target env, and every retargeting knob: hand_policy, " + "num_interpolation_steps, init_robot_from_ik, replay_speed, retarget_frame, scene_translation, " + "object_tracking, eef_offsets, ...). Almost all knobs are descriptor-only; runtime/IO flags stay " + "on the CLI." + ), +) +parser.add_argument( + "--source_embodiment", + type=str, + default=None, + help="Path to the source embodiment YAML (the embodiment the input dataset was recorded on).", +) +parser.add_argument( + "--target_embodiment", + type=str, + default=None, + help="Path to the target embodiment YAML (the embodiment to retarget the dataset onto).", +) +parser.add_argument( + "--target_env_name", + type=str, + default=None, + help="Environment id to instantiate for the target embodiment (the target robot's env).", +) +parser.add_argument( + "--num_envs", + type=int, + default=1, + help=( + "Number of parallel environments. 1 (default) replays episodes sequentially. >1 replays them " + "across that many envs on one sim (async workers, one episode per env at a time) for a large " + "speedup. The per-step tracking report is single-env only and unavailable in parallel mode." + ), +) +parser.add_argument("--input_file", type=str, required=True, help="Source dataset HDF5 file (must carry datagen_info).") +parser.add_argument( + "--output_file", + type=str, + default="./datasets/retargeted_dataset.hdf5", + help="Destination HDF5 for the retargeted episodes.", +) +parser.add_argument( + "--select_episodes", + type=int, + nargs="+", + default=[], + metavar="IDX", + help="Retarget only these source-episode indices (0-based, into the dataset); empty (default) = all.", +) +parser.add_argument( + "--target_successes", + type=int, + default=None, + metavar="N", + help=( + "Stop once N replays have succeeded (early-stop). Unset (default) = run every source example. " + "On a finite source that can't reach N, it stops when the source is exhausted." + ), +) +parser.add_argument( + "--target_runs", + type=int, + default=None, + metavar="N", + help="Stop once N replays have been attempted (successful or not); capped at the source size.", +) +parser.add_argument( + "--keep_failed", + action="store_true", + help=( + "Also export episodes whose replay did not satisfy the success term (overrides the task " + "descriptor's keep_failed). Useful for inspecting the retargeted motion when success is " + "expected to be low." + ), +) +AppLauncher.add_app_launcher_args(parser) +args_cli = parser.parse_args() + +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +from autodata_retargeting.runner import run # noqa: E402 + +if __name__ == "__main__": + try: + run(args_cli, simulation_app) + except KeyboardInterrupt: + print("\nInterrupted; exiting.") + simulation_app.close()