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 : 2 dependent PRs (#486 , #500 )
• Fix PI0.5 raw rotation action encoding #483 👈 (View in Graphite)
• main

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

[github-actions]

Claude Code Review

PR Review: Fix PI0.5 raw rotation action encoding

Summary

Adds a new cartesian_ypr_raw_rot6d action encoding mode for PI0.5 that constructs the 6D rotation representation from raw (unnormalized) YPR angles, fixing a bug where rotations were being learned from normalized Euler coordinates that don't respect SO(3) topology. Legacy behavior is preserved behind an action_encoding config flag.

Key concerns

1. Eva config opts into the new encoding by default — potential training regression

egomimic/hydra_configs/model/pi0.5_bc_eva.yaml is switched to cartesian_ypr_raw_rot6d. Any in-progress runs/checkpoints trained with the legacy encoding will be incompatible at inference time without manually setting action_encoding: "legacy_normalized_ypr_rot6d". Please confirm:

• No active eva PI0.5 experiments will silently break.
• Release notes / commit message flag this for collaborators.
• Existing checkpoints have clear migration guidance (since the field defaults to legacy, restoring an old checkpoint with the new yaml will mismatch).

2. serialize path computes raw_action via _unnormalize_action, then re-passes normalized_actions=action

In to32_raw_rotation, when both normalized_actions and stats are provided, only normalized_actions is used to overwrite non-rotation dims of the cloned raw_actions. This is fine, but it means the unnormalize→re-normalize round-trip happens twice (once in _unnormalize_action, once implicitly via the normalized dims you pass through). Not a correctness bug, but consider whether passing only raw_action + stats (letting the converter normalize) is simpler and avoids drift from eps=1e-6 round-trip error accumulation across two transforms.

3. _apply_unnorm_one for minmax/quantile uses (mx - mn + 1e-6) symmetrically with _apply_norm_one

The forward uses / (mx - mn + 1e-6) and the inverse uses * (mx - mn + 1e-6). These are not exact inverses (off by the eps term), which is why the round-trip test passes only at 1e-5 tolerance. This matches the existing convention in norm_stats, but please double-check it actually does — if norm_stats.unnormalize uses a slightly different eps or no eps, then the _unnormalize_action path (real) and _apply_unnorm_one path (local) will disagree, producing subtle bias in non-rotation dims. Worth a unit test asserting parity with NormStats.unnormalize for at least one mode.

4. Dead _action_stats call path in inference

_action_stats is fetched and passed to from32_raw_rotation in the prediction path, but from20_raw_rotation only uses stats when unnormalize_non_rotation=True (which is the case here). Fine — just flagging that error messages should mention the embodiment name as well as id for easier debugging in multi-embodiment runs.

5. Test file naming

egomimic/utils/test_pi05_action_encoding.py lives in the source tree rather than under a tests/ directory. Confirm this matches the repo's test discovery convention (pytest will find it via test_*.py, but most modules in egomimic/utils/ aren't test files).

Suggestions

• Add a parity test between _apply_unnorm_one and NormStats.unnormalize to lock in agreement; otherwise these two implementations can drift silently.
• Test the from32_raw_rotation decode path explicitly with quantile norm_mode at training-realistic ranges, since that's the production code path through _get_action_trajectory.
• Add an assertion in _get_action_trajectory for the raw-rot branch confirming pred_actions_orig.shape[-1] == 14 before slicing to [:, :T, :D] — the PR description mentions this guard but I don't see it in the diff. If it landed in a later commit, ignore.
• Document the encoding semantics in a docstring on RobotBimanualCartesianEuler so future contributors understand that rotation dims (3,4,5,10,11,12) are intentionally not normalized when this encoding is active.
• Consider exposing the encoding constant in the eva yaml as a Hydra reference rather than a string literal, to avoid typos silently falling through to the ValueError branch only at runtime.
• Minor: _stat_tensor's dtype conversion logic (ref.dtype if ref.is_floating_point() else torch.float32) is duplicative — torch.as_tensor(..., dtype=torch.float32) is already set, then immediately recast. Simplify.

Verdict: Comment

The core fix is well-motivated and the round-trip test demonstrates the bug it solves (yaw wrap discontinuity). I'd like to see (a) confirmation that _apply_unnorm_one matches NormStats.unnormalize exactly, and (b) explicit owner sign-off that the eva yaml flip won't disrupt active runs, before merging. Not blocking, but please address before this goes into a training run we care about.

---

Reviewed by Claude · Review workflow