refactoring aciton transforms

[AnikethCheluva]

No description provided.

[AnikethCheluva]

• Aniketh/arc #573
• feat(abc): ABC-130k mcap to zarr pipeline + yam_bimanual embodiment #606
• feat(val): nested validation dataset groups #605
• refactoring aciton transforms #604 👈 (View in Graphite)
• main

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

[github-actions]

Claude Code Review

Review of PR #604: Refactoring action transforms

Summary

Refactors get_transform_list on both Human and Eva embodiments from a single opaque mode string into three orthogonal parameters (action_mode, coord_frame, rotation_mode), and adds a new XYZWXYZ_to_XYZRot6D transform + transforms_for_rotation_mode dispatcher to support Zhou 6D rotations. All Hydra data configs and downstream callers are updated in lockstep.

Key concerns

1. Config renaming is a breaking change with no back-compat shim. Every existing training config that pins transform_list.mode: <something> in a checkpoint dir or external launch script will break with a Hydra key error. Given this is a research codebase with active experiments, please confirm no in-flight runs / launched sweeps reference the old mode: key, or add a deprecation path (accept mode and translate).

2. Silent semantic change in _build_human_keypoints_bimanual_transform_list. In the old code, if not is_quat: <append rotation converters> ran regardless of concat_keys. In the new code, transforms_for_rotation_mode(...) is called inside if concat_keys:. That means when _build_human_keypoints_eef_frame_transform_list calls the bimanual builder with concat_keys=False (which it does), the ypr conversion no longer happens inside the inner builder — it's re-done later by the outer function. This is probably fine because the eef-frame builder overrides with rotation_mode="quat" for the inner call, but it's brittle and the coupling is invisible. Please add a comment, or better, move the rotation dispatch back outside the concat_keys guard and let the outer builder pass rotation_mode="quat" explicitly (already done — good). Just add an assertion or comment inside the inner builder documenting "callers with concat_keys=False must pass rotation_mode='quat'".

3. cartesian_padded → cartesian_gripper_padded rename + new padding of obs. Old behavior padded only actions_cartesian (12→14). New _pad_human_cartesian_gripper also pads observations.state.ee_pose. This changes proprio dim seen by the policy — any existing checkpoints trained with cartesian_padded will have mismatched proprio input dims on resume/eval. Please confirm this is intentional and note it in the description; also update any norm-stats caches keyed on old dims.

4. PadGripperZeros semantic change. The pad slot location changed from "fixed indices 6 and 13" to "after each pose_dim block". For pose_dim=6 this is arithmetically identical, so 14D behavior is preserved. Good — but worth a quick unit test asserting pad(2*6) == old_behavior and a new test for pose_dim=9 → 20D.

5. No new tests. XYZWXYZ_to_XYZRot6D, transforms_for_rotation_mode, and the generalized PadGripperZeros have no tests in this diff. _matrix_to_xyzrot6d is imported but I can't see it validated end-to-end. Please add:

   • transforms_for_rotation_mode(["k"], "6D") on a known quat produces the expected 9D output
   • PadGripperZeros(pose_dim=9) on 18D → 20D with zeros at indices 9 and 19
   • Round-trip smoke test that all three rotation_modes produce the right last-dim for one Eva and one Human config

6. Eva.get_transform_list accepts only action_mode="cartesian" — fine, but the Literal type says only "cartesian" and then the code still raises on anything else. Consider dropping the redundant check or expanding the Literal for future-proofing (matching Human's shape).

Suggestions

• Add a mode→(action_mode, coord_frame, rotation_mode) compatibility shim (dict lookup) that emits a DeprecationWarning. Removes the "silent break" risk across configs you don't own.
• Docstring on Human.get_transform_list: clarify that cartesian_gripper_padded now also pads observations (call this out explicitly — this is the kind of thing that will bite someone).
• Add an assert on the rotation-mode dispatcher output dim per key so mis-wiring surfaces at data-load rather than at loss-backward.
• The notebook re-execution churn (execution_count: null) is fine but noisy — consider nbstripout in pre-commit.
• Minor: _pad_human_cartesian_gripper hard-codes actions_key="actions_cartesian" and obs_key="observations.state.ee_pose". These are the canonical names, but if a future config renames them the pad will silently no-op or KeyError. Not blocking.

Verdict: Request Changes

Blocking items are (3) the proprio-padding semantic change (needs explicit confirmation / release note / norm-stat invalidation), (5) missing tests for the new 6D path and generalized padding, and ideally (1) a compat shim for the config rename. The refactor itself is a clear improvement — orthogonal knobs beat a combinatoric enum — and the code reads well.

---

Reviewed by Claude · Review workflow