feat(pipeline): port arc-length action tokens

[ElmoPA]

No description provided.

[ElmoPA]

Warning

This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
Learn more

• feat(sim): add PushShapes Sim V2 with rigid parallel-gripper physics #586
• feat(pipeline): port arc-length action tokens #585 👈 (View in Graphite)
• feat(pipeline): consolidate Fold, PushT, overlays, and MoE #579 : 2 other dependent PRs (#580 , #583 )
• main

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

[github-actions]

Claude Code Review

Review of PR #585: feat(pipeline): port arc-length action tokens

Summary

Ports the arc-length action tokenizer from a GR00T branch, wires it through Eva/Human embodiments and the Pipeline sampler, and adds a U-socket rotvec variant. Substantial change touching data transforms, model configs, rollout, and normalization statistics.

Key concerns

1. get_embodiment_id semantic change is risky. The function used to be case-insensitive via .upper() on EMBODIMENT[name]. The new implementation lowercases first, then applies a legacy alias table, then uppercases. Two issues:

   • The alias map keys (aria_bimanual, etc.) are already lowercase, but if any caller passes "ARIA_BIMANUAL" (previously worked via .upper()), it now goes: "aria_bimanual" → aliased → "HUMAN_BIMANUAL" — actually OK. But the legacy string collapse silently rewrites embodiment identity at read time. This is a behavioral change beyond the PR title: any downstream code that inspected embodiment for aria_bimanual will now see human_bimanual. Please confirm all filters (row['embodiment'] == 'aria_bimanual') in existing configs still work — from the SQL side these are filtered pre-alias, but from cached rows post-alias they won't. This deserves its own PR + tests.

2. Velocity dt derivations look fragile. The Eva formula (ARC_TOK_ACTION_HORIZON - 1) / (30.0 * (ARC_TOK_INTERPOLATED_HORIZON - 1)) and the Human formula involving ((600-1)//stride)*stride / (30*(100-1)) are load-bearing physical scale factors buried inside get_transform_list. If InterpolatePose changes semantics or stride shifts, the velocity token silently mis-scales and norm stats drift. Please:

   • Add a docstring/assertion that couples these to the actual InterpolatePose config.
   • Add a unit test that round-trips a synthetic trajectory: tokenize → detokenize → check reconstructed timing against source dt.

3. _unnormalize_prediction slotwise stats change may break existing eval runs. Existing checkpoints with per-dim (D,) stats broadcast fine, but any old code path that relied on the flatten trick (unlikely, but check EvalVideo subclasses) may regress. Low risk, but worth grepping.

4. Arc-length norm stats — where do trans_vel_left/right get injected? arc_length_stats.py computes them, but I don't see the trainer/norm-stats pipeline wiring that calls compute_trans_vel_stats and merges into the action-key stats. If this isn't hooked into trainHydra.py / NormStats, the trans_vel channels will be normalized with wrong (or zero) statistics on first run. Please confirm the wire-up exists or add it.

5. ArcLengthRolloutAdapter._current_ypr returns zeros silently when context or state_key is missing. During rollout that would produce a chunk with identity orientation — a subtle failure mode. Prefer raising, since arc tokens explicitly cannot carry rotation and correct behavior requires the current wrist pose.

6. _module_key dotted-key encoding — checkpoint compatibility. The "__dotted_" + hex scheme preserves undotted keys unchanged (good), but any existing checkpoint that was trained with an obs key containing a dot (were there any? observations.images.front_img_1 looks like it may have been used) will now fail to load. Please confirm no released checkpoint has dotted keys in obs_encoders/img_encoders state dict.

7. PUSHSHAPES_SIM_U_SOCKET = 19 — comment says ID 18 belonged to HUMAN_BIMANUAL and must not be reused. Good. But is 19 already used by any prior dataset under a different name? A quick grep of the SQL/zarr metadata is warranted.

Suggestions

• Add tests for:
  • TokenizeBimanualArcLength round-trip on a synthetic bimanual chunk (including one stationary arm).
  • ArcLengthRolloutAdapter.decode shape and orientation-hold behavior with a real 14-D state_ee_pose.
  • ThetaToRotVec ∘ RotVecToTheta identity on random angles.
  • get_embodiment_id legacy alias mapping.
• Split the embodiment alias change out — it's an independent semantic change and shouldn't be buried in a tokenizer PR.
• Docstring on Eva.ARC_TOK_ACTION_HORIZON = 200 / Human.ARC_TOK_ACTION_HORIZON = 600 explaining the physical span assumption and coupling to InterpolatePose.
• _run_rollout_policy — the excluded[0] == ["<train-only>"] sentinel string comparison is brittle. Prefer an explicit exclusion reason enum.
• PR description is empty. For a change this size touching training-critical paths, please add a description explaining scope, active experiment context, and which configs are known to be exercised.

Verdict: Request Changes

Primary blockers: (a) the get_embodiment_id legacy-alias behavior change should be separated and tested; (b) confirm/wire the trans_vel norm stats into the training pipeline; (c) add at least round-trip and rollout-shape tests for the new tokenizer paths. The core port looks careful (respect for invalid-pose sentinel, per-arm arc length, ypr via SLERP), but for a data-format change touching norm stats and rollout, we need explicit test coverage before merging.

---

Reviewed by Claude · Review workflow