rollout: config-driven viz, wristframe revert, HPT annotation support

[rl2aloha]

Read transform_list mode from .hydra/config.yaml instead of hardcoding
cartesian_wristframe_ypr. The rollout now matches whatever frame the
model was trained on (e.g. mode: cartesian for QWEN-HPT,
mode: cartesian_wristframe_ypr for pi05_eva_aria_obj_gen_lang_wristframe).

Add config-driven prediction visualization. build_viz_func_from_config
reads evaluator.viz_func from .hydra/config.yaml and uses hydra.utils.instantiate
to bind image_key / action_key / mode. A runtime 'v' toggle in the
intervention menu enables/disables saving viz.png to debug/
on every inference. The live camera image is resized to 640x480
inside _save_viz so the projection matches ARIA_INTRINSICS (the
Aria camera publishes at 960x720 per configs.yaml, but the intrinsics
in INTRINSICS['base'] are calibrated for 640x480). cv2.imwrite gets
BGR conversion to avoid the inverted-color save bug.

Add wrist-frame revert. _build_revert_transform_from_config reads
evaluator.transform_lists. from config and instantiates
the revert transform that converts wrist-frame predictions back to
cam frame using the current obs ee_pose as the reference. rollout_step
applies this to preds and the GT batch BEFORE both the viz call and
the cam_frame_to_base_frame post-processing. Wrist-frame models now
execute correctly on the robot AND project to the right pixels.
Cam-frame models (revert is None) skip the revert as a no-op, so
QWEN-HPT and old PI cartesian checkpoints continue to work unchanged.

Add HPT (QWEN) annotation support. _apply_annotation_to_algo now
sets annotation_sampling_mode='first' in addition to PI's
sampling_mode='first', so --annotation-path works for both algos
without subclass-specific code paths.

Co-Authored-By: Claude Opus 4.7 noreply@anthropic.com

[rl2aloha]

• rollout: config-driven viz, wristframe revert, HPT annotation support #482 👈 (View in Graphite)
• feat(zarr): persist camera intrinsics in zarr.json #386 : 1 other dependent PR (#483 )
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[github-actions]

Claude Code Review

Review of PR #482

Summary

This PR bundles three logically distinct changes: (1) config-driven rollout viz + wristframe revert (per the PR title), (2) HPT annotation support, and (3) a large hierarchical subtask-prediction feature (PI0Subtask, new configs, embodiment keymap changes, 6D rotation transforms, proprio normalization in collate). The scope mismatch between title/description and diff is a serious concern.

Key concerns

1. Scope creep — PR description does not match the diff

The description covers only the rollout changes, but the diff adds:

• PI0SubtaskPytorch model (~370 lines)
• New pi0.5_subtask / sort_pi_subtask configs
• New 6D rotation transforms (CartesianYPRToRot6D, CartesianRot6DToYPR)
• New action_encoding field on PI with three modes
• Changes to Embodiment.get_keymap signature (adds high_annotation_key)
• Changes to ZarrDataset._load_annotations signature and cache structure (dict-keyed)
• Non-trivial changes to build_tokenized_collate (now requires data_schematic, hard-fails without stats)
• New ROLLOUT_INSTALL.md documenting hand-editing an external submodule

This should be split into at least 3 PRs. The rollout fix is small and safe; the subtask work is a large research feature that needs its own review. Landing them together makes bisecting regressions painful.

2. Breaking API change without callers audited

ZarrDataset._annotations changed from list[dict] | None to dict[str, list[dict]] | None. Any external code (or __setstate__, checkpoint restores, tests) touching self._annotations directly will silently misbehave. Also _annotation_text_for_frame now returns list[str] instead of str (per the docstring/signature) — this is a return-type change; all callers need to be verified.

3. build_tokenized_collate hard-fails on missing schematic/stats

The new code raises ValueError if data_schematic is None or norm_stats is None when proprio=True. This is correct behavior (silent skip would corrupt bins), but existing callers that construct the collate without passing data_schematic will now crash. Please confirm all call sites in pl_data_utils.py / trainHydra.py / any tests are updated, and that ordering (infer_norm_from_dataset before first batch) is guaranteed under DDP with multiple workers.

4. Rollout viz resize assumption

_save_viz resizes the live camera to 640×480 to match ARIA_INTRINSICS. This is Aria-specific — will it silently mis-project for Eva / Mecka rollouts that have different intrinsics? The revert transform is embodiment-aware (evaluator.transform_lists.<embodiment>), but the intrinsics used for viz appear hardcoded. Please confirm or gate on embodiment.

5. Subtask decoding correctness

In sample_actions, next_tok is a [B] tensor and zeros_col[:, 0] is used as the "inactive" branch of torch.where. Verify dtype/device match — this looks fine but is worth a unit test. More importantly: when the high-only prefix arrives right-padded and tokens are inserted at insert_idx = cur_mask.sum(dim=1), this writes inside the original padding region only if the prefix was left-packed and padding is on the right. That's the assumption in _tokenize_highonly (ids[i, :len(pre)] = ...), so it holds — but please add a comment or assertion because the correctness of return_subtask gathering depends on it.

6. Non-idempotent openpi submodule setup

ROLLOUT_INSTALL.md documents hand-editing files in external/openpi at a detached HEAD with uncommitted changes and an untracked script. This is not reproducible and will cause silent drift. Please push a pi-rollout-host branch to the fork and reference a commit SHA, as the doc itself suggests.

7. No tests

Zero tests added for: 6D↔ypr round-trip (CartesianYPRToRot6D / CartesianRot6DToYPR), subtask CE loss shape, hierarchical decoding, _load_annotations filtering by level, or the wristframe revert transform. For a change this large touching training loss computation and data layout, at minimum:

• Round-trip test for the 6D transforms
• Snapshot test that pi0.5_cotrain_eva_aria_6d produces identical batch tensors (modulo action rep) to the ypr config
• Unit test for _annotation_text_for_frame with level="high" / "low" split

8. Norm stats regression risk

action_encoding=cartesian_normalized_rot6d uses norm stats computed over the 6D representation. If any existing checkpoint's norm stats were computed over ypr, loading it under a 6D config will silently produce garbage actions. Is there a version/compatibility check on the norm_stats file, or at least a config-level assertion that transform_list mode and action_encoding are consistent?

Suggestions

1. Split the PR. Land the rollout fix + HPT annotation support as a small PR first (matches the description, low risk). Land subtask + 6D as a separate feature PR.
2. Add a config-level validator: action_encoding=cartesian_normalized_rot6d ⇒ transform mode must be cartesian_6d or cartesian_wristframe_6d. Fail fast in PI.__init__.
3. Add a norm_stats fingerprint (embodiment × ac_key × action_encoding) and check it at load time.
4. Push a stable openpi fork branch and reference a SHA in ROLLOUT_INSTALL.md.
5. Grep for

---

Reviewed by Claude · Review workflow