Fix PI0.5 raw rotation action encoding

[ElmoPA]

Fix PI0.5 raw rotation action encoding

Preserve physical YPR angles when converting normalized Cartesian actions to 6D rotation targets so PI training does not learn rotations from normalized Euler coordinates.

Co-authored-by: Cursor cursoragent@cursor.com

Tighten raw rotation action encoding

Keep the new PI0.5 encoding simple by dropping the version suffix and failing fast when a converter does not implement raw-rotation packing.

Co-authored-by: Cursor cursoragent@cursor.com

Harden raw rotation action handling

Clone normalized actions before unnormalization and assert the decoded bimanual action shape before applying non-rotation unnormalization.

Co-authored-by: Cursor cursoragent@cursor.com

[ElmoPA]

• 6D pi data transform #487 : 3 dependent PRs (#486 , #500 , #501 )
• Fix PI0.5 raw rotation action encoding #483 👈 (View in Graphite)
• feat(zarr): persist camera intrinsics in zarr.json #386 : 1 other dependent PR (#482 )
• main

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

[github-actions]

Claude Code Review

Review: PR #483 — Fix PI0.5 raw rotation action encoding

Summary

Adds a new cartesian_ypr_raw_rot6d action encoding for PI0.5 that packs 6D rotation from raw YPR values (avoiding double-normalization of angles through Euler space), while still normalizing the non-rotation (xyz + gripper) dims. Legacy path is preserved behind a config flag.

Key concerns

1. Silent behavior change for existing PI0.5 Eva runs. The default action_encoding on the class is LEGACY, but pi0.5_bc_eva.yaml is flipped to cartesian_ypr_raw_rot6d. Any in-progress Eva experiment resumed from checkpoint will now decode rotations differently than it was trained. Norm stats aren't versioned against the encoding, so this can silently produce wrong actions at eval. Consider either (a) leaving the config on legacy and opting new runs in explicitly, or (b) saving action_encoding into the checkpoint metadata and asserting on load.

2. _unnormalize_action returns unnormalized values for ALL dims, including rotation. In get_action:

   raw_action = self._unnormalize_action(action, emb_id, ac_key)
   action32 = converter.to32_raw_rotation(raw_action, normalized_actions=action, ...)

   raw_action[..., rot_dims] is the norm_stats-unnormalized YPR — which for quantile/minmax treats angles as bounded scalars and just rescales them. That's fine for the rotation dims since to20_raw_rotation only reads raw rotation from raw_actions. But you're relying on the fact that norm_stats.unnormalize produces "close enough to raw YPR" for the rotation dims — which is only true if the normalization stats on rotation dims are near-identity, or if training also unnormalized them the same way. Please confirm that norm stats for rotation dims during training match what to32_raw_rotation expects, i.e., that YPR was stored raw in the zarr and only wrapped through minmax/quantile scaling. If YPR was stored pre-normalized in some pipeline, this will be wrong.

3. process_batch_for_training was not updated. The PR fixes encoding at inference (get_action) and decoding at validation (train_on_batch? — actually the loop in from32/from32_raw_rotation inside eval), but I don't see a corresponding change to how training targets are constructed. If training still calls converter.to32(...) on already-normalized actions, then the model is still being trained on normalized-YPR-derived 6D rotations, and inference-time raw decoding will be inconsistent. Please point to where the training target now uses to32_raw_rotation, or add it.

4. _unnormalize_action mutates through .clone() but then calls .to(action.device) — the norm_stats path may return a CPU tensor and the .to() is fine, but note that self.norm_stats.unnormalize typically expects a dict including "embodiment"; you pass it, good. However, the returned dict from unnormalize — does it also contain other keys that get discarded? Just a minor concern; verify no side effect on predictions dict elsewhere.

5. Test coverage gap. Tests only cover the converter round-trip and yaw-wrap continuity (nice!). No test for:

   • The PI.get_action path with action_encoding=raw_rot_6d.
   • The training/loss path (the most consequential change).
   • Assertion that legacy config still round-trips unchanged (regression guard).

6. Yaw-wrap test threshold is loose. fixed_rot_distance < 1e-3 — but with eps=1e-4 in raw YPR and a 6D rotation matrix, the actual distance should be ~2e-4. Tighter bound would catch regressions better. Minor.

Suggestions

• Add a training-side test or at least a line-level pointer confirming to32_raw_rotation is used when building training targets. If it isn't, this PR only half-fixes the bug.
• Save action_encoding into the algo checkpoint and assert on resume/eval.
• Consider adding action_encoding to the norm-stats sidecar / model config summary so downstream eval scripts can't accidentally mismatch.
• In _action_stats, catch the case where norm_stats.norm_stats itself is missing (e.g., None) with a clearer error.
• Rename PI05_CARTESIAN_ACTION_ENCODING_LEGACY value from "legacy_normalized_ypr_rot6d" — the string is fine but consider documenting in a docstring why legacy is preserved (bit-for-bit reproducibility of prior checkpoints).
• Add a comment above ROBOT_BIMANUAL_CARTESIAN_ROT_DIMS clarifying these are indices into the 14-D [xyz, ypr, g] x 2 layout, not the 20-D packed layout — easy to misread.
• Tighten the yaw-wrap test tolerance to < 1e-2 at most, ideally < 5e-4.

Verdict: Request Changes

The converter logic and tests look solid, and preserving raw YPR through the 6D encoding is the right fix. But I need confirmation (in the diff or a follow-up comment) that:

1. The training target construction also uses to32_raw_rotation — otherwise train/eval are inconsistent.
2. Existing Eva PI0.5 checkpoints/runs won't silently break given the config default flip.

Once those two are addressed (with tests + checkpoint versioning), happy to approve.

---

Reviewed by Claude · Review workflow