Fix PI0.5 raw rotation action encoding

[aidang3019]

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
(cherry picked from commit 349e6c804bea6b7606376fad366abacaf0797a98)

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
(cherry picked from commit b2223eaa6788a762598ff4561f0fa7b25e5915d9)

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
(cherry picked from commit 309e5f99385d2a90c9962ca1191d08568d7cf948)

6D pi data transform

(cherry picked from commit 6494762f2d22c534565422aaa1c6f990676c9312)

6D proprio, rotation-free bounds check, constant LR, mecka 6D configs

• eva/human cartesian_6d + wristframe_6d modes 6D-encode
  observations.state.ee_pose (revert lists convert it back before the
  eef-frame revert) so proprio normalization avoids the ±π yaw/roll wrap
• bounds checking skips rotation channels of actions_cartesian / ee_pose
  via BIMANUAL_CARTESIAN_LAYOUTS (ported from the wristframe-6d stack)
• pi0.5_base: constant-with-warmup LR schedule (cosine decayed to 0 @ 60k)
• new pi0.5_bc_mecka_6d + mecka_all_pi_6d configs: all SQL mecka episodes,
  bs 64, cartesian_6d stride 1, proprio in prompt, no embodiment label

Co-Authored-By: Claude Fable 5 noreply@anthropic.com
(cherry picked from commit 19c59e50f0e3eb920288c140719dee72e3d81546)

mecka wristframe-6d: launch infra, precomputed norm stats, left-wrist Rz(180) fix

Squash of the two mecka-all pipeline commits that used to sit inside
aidan/eval_gating, so the whole wristframe 6D pipeline lives on this
branch.

mecka-all launch infra (was a8ae76b):

• pi0.5_base default_prompt "" (was "This is a bad action."; remote parity
  — applies to frames outside any annotation span)
• mecka_all_pi_6d: cartesian_wristframe_6d stride 1 — 18D wrist-relative
  actions (world frame AND moving head frame cancel out of the targets),
  head-frame 6D proprio; shared PACE mirror folder_path; DatasetFilter
  lambda (dict filters raise in the zarr resolver path)
• get_embodiment_id: vendor-tagged human names (MECKA_/SCALE_/ARIA_/
  LIGHTWHEEL_) collapse to HUMAN_ — mirror episodes carry MECKA_BIMANUAL
  metadata from the vendor-split registry
• norm collection: float32 + 2M-sample cap (a 0.1-frac (N,100,18) float64
  stack plus np.percentile's sort copy is ~250GB)
• precompute_norm_stats.py: CPU port of the remote's script against the
  local trainHydra norm loop; sbatch variants (24c CPU / 1xa100 32c /
  2xa100 64c) + smoke (viz gate proof) + 8xH200 train scripts, paired
  with evaluator=eval_pi_wristframe_6d

fix mecka left wrist frame (was 68ef3aa):

• mecka_to_zarr.compute_hand_pose_xyzquat double-mirrored the LEFT hand
  onto the right hand's spatial convention (palm normal from
  cross(thumb, pinky) already mirrors chirality; rot_left flipped x/y
  again). Converter fixed; RotateLocalFrame + the fix_mecka_left_wrist
  data-config flag right-multiply the stored left pose keys by Rz(180°)
  before any frame math, exactly equivalent to reconverting processed_v3.
• viz_mecka_episode / viz_eva_episode debug renderers.

Co-Authored-By: Claude Fable 5 noreply@anthropic.com
(cherry picked from commit 96b439d15fcd9e786d7ed54173cec0a2ca02737c)

6D pipeline hardening: loud config mismatches, bounds tolerance, norm-stats provenance, independent round-trip tests

Adversarial review of the wristframe 6D pipeline (2026-08-28) found no
math bug: the forward transform, quantile normalization, 32D pack/unpack
and the evaluator revert all round-trip to float32 precision against
independent SE(3) math, for both embodiments and both left-wrist
conventions. It did find three silent-failure modes:

• SplitKeys sliced an 18/20-dim 6D batch as "xyz + ypr" without complaint
  whenever an evaluator's ypr revert list was paired with a *_6d data
  config (eval_pi.yaml vs cartesian_wristframe_6d). It now raises with the
  widths and points at transform_lists.
• MultiDataset._check_bounds compared strictly against quantile bounds
  that collapse to [0, 0] at t=0 of wrist-frame chunks (the reference IS
  the obs pose). Those cells are exactly 0.0 today, so nothing was
  rejected, but any roundoff (other BLAS / dtype) would have resampled
  every frame. Add a 1e-6 absolute tolerance, and warn once instead of
  silently skipping the key when the stats do not broadcast.
• norm_stats.json carried no provenance: a cached file from another
  norm_mode / keymap / transform mode with the same dims was applied
  as-is at train AND eval. cache_stats now writes {norm_mode, stat_shapes};
  _load_precomputed_stats refuses a norm_mode or key-set mismatch (files
  written before provenance still load, with a warning).

test_wrist6d_roundtrip.py runs the REAL Human/Eva transform lists on
synthetic world-frame poses through normalize -> to32/from32 ->
unnormalize -> revert and compares every stage with plain numpy/scipy
SE(3) math that shares no code with the pipeline (both
fix_mecka_left_wrist settings, both eva extrinsics), plus torch-vs-scipy
Euler-convention and Gram-Schmidt cross-checks and tests for the three
guards above. The existing tests only compared the pipeline with itself.

Co-Authored-By: Claude Fable 5 noreply@anthropic.com
(cherry picked from commit 5820bea75e88665fefb66bf53f544abfa2a3f1c7)

split val metrics/viz

(cherry picked from commit 96ce8873e93c358b8a42ea5c9e9f42f80a4c94da)

val every 20 epochs, viz every 100; align viz gate with Lightning val epochs

• trainer/ddp_pi: check_val_every_n_epoch 10 -> 20 (80 val batches inherited
  from trainer/default; valid loader shuffle pinned false in mecka_all_pi_6d)
• evaluator/eval_pi: viz_every_n_epochs 200 -> 100
• _should_viz now gates on (current_epoch + 1) % viz_every_n_epochs, matching
  Lightning's (current_epoch + 1) % check_val_every_n_epoch val condition —
  the plain current_epoch % n never coincided with a validation epoch unless
  check_val_every_n_epoch == 1, so viz never fired

Co-Authored-By: Claude Fable 5 noreply@anthropic.com
(cherry picked from commit 51e979c656499f13cc6789714e56edffb689b9f0)

sampler: widen bad-sample retries to global index space after 25 attempts

One wholly-out-of-bounds episode (692ea23dc621d7f4aac3aaa1) exhausted the
within-episode retry pool, killed rank 3, and NCCL-timeouted the other 7
ranks 8h into the mecka-all run. Retries now leave the failing episode
after 25 attempts; 1000 consecutive failures still raise loudly.
Also 6->8 dataloader workers; dedupe two accidentally duplicated tests.

Co-Authored-By: Claude Fable 5 noreply@anthropic.com
(cherry picked from commit 954e47063ca8ceb4d186fac7a72b44a74ac3e0f2)

wrap-aware cam-frame MSE: wrap YPR errors to ±pi, keep *_nowrap companions

Euler wrap at ±pi inflates cam-frame MSE: +pi-eps vs -pi+eps scores ~(2pi)^2
per dim. Wrap rotation-dim errors before squaring (12/14-wide YPR layouts
only; 6D widths untouched). Unwrapped values stay logged as *_nowrap to
quantify the inflation.

Co-Authored-By: Claude Fable 5 noreply@anthropic.com
(cherry picked from commit a3787eba4f9f852890c3935e5c9c635f173b8ffe)

val videos: scale playback fps by 1/world_size (DistributedSampler stride)

Distributed val deals unshuffled indices round-robin, so rank 0 renders
every world_size-th source frame; writing at the raw 30fps produced a
world_size-times timelapse (8x on the current run).

Co-Authored-By: Claude Fable 5 noreply@anthropic.com
(cherry picked from commit 47505c63332506b1c40d3a9624e76618c49abdc4)

port better-val-metrics + train viz from ryanco/better-val-metrics

Val metrics (eval_pi): xyz/ypr MSE splits (native + cam), Frechet Gaussian
over time (avg/min/max), reverse-KL + best/mean/worst-of-M + sample
diversity from M stochastic rollouts (sample_action_chunks, gated on
config.reverse_kl_samples; mecka config enables M=4). Kept our wrap-aware
cam MSE + *_nowrap companions.

Train viz second head: TrainVizEvalVideo wraps the wristframe evaluator,
runs against a train-split loader (dataloader_idx=1), metrics prefixed
train_viz/, videos under videos_train_viz/, M-sample metrics forced off.
val_at_start=true gives a pre-fit baseline val on fresh launches (skipped
on requeue/ckpt-resume, where it would score un-restored weights).

Co-Authored-By: Claude Fable 5 noreply@anthropic.com
(cherry picked from commit 70bc694e69fd8741591194ac303f6d69db18a746)

val metrics: batched DTW distance (native frame, per-step normalized)

Wavefront DP over anti-diagonals, vectorized across the batch (2T-1 steps
for a TxT grid). Tolerates temporal misalignment that paired MSE penalizes;
logged as Valid/_dtw_avg on every val (train_viz inherits). Verified
against a brute-force reference in tests.

Co-Authored-By: Claude Fable 5 noreply@anthropic.com
(cherry picked from commit 44175ef79b9b6e2615dae3abab36820801fd62cb)

pi trainer: validate under no_grad, not inference_mode (cudagraph capture crash)

The compiled sampler's cudagraph recording does in-place setup on static
buffers - illegal on inference tensors. Hit by the M-sample val metrics'
repeated sample_actions calls (fold-debug 11112709, val_at_start).

Co-Authored-By: Claude Fable 5 noreply@anthropic.com
(cherry picked from commit 4c13207d1b0512f69fc310897482121683ac96a0)

val_at_start off by default

Co-Authored-By: Claude Fable 5 noreply@anthropic.com
(cherry picked from commit 39e39b30acb9d2d7d6d5c83fc2e98ab31aeb5a0a)

eval_pi: stateless MSE helpers, geodesic rotation-error metric

• _split_mse re-instantiated two torchmetrics MeanSquaredError objects per
  call (twice per embodiment per val batch) and the shared mse was a
  stateful accumulator; both replaced by a stateless _paired_mse.
• New Valid/__rot_geodesic_avg (native output, after the same
  Gram-Schmidt the decode uses) and _cam_rot_geodesic_avg (cam frame):
  arccos((tr(Rp^T Rg) - 1) / 2), computed in float64. The wrap-aware ypr
  MSE is still degenerate at pitch ~ +-pi/2, where yaw and roll trade off
  and a 0.004 rad orientation error scores > 0.1; the geodesic error is
  immune to both the +-pi wrap and that degeneracy. Tests cover both.

Co-Authored-By: Claude Fable 5 noreply@anthropic.com
(cherry picked from commit 9e206dceee07f4fdbb726af6cad61be71c7fc047)

pi0.5 fold-freeform fit tests on the eval_gating recipe: cam6d + wrist6d data configs, cam-frame train_viz evaluator, CPU norm-stats + 2xH200 train launchers (flat 5e-5)

(cherry picked from commit bb4128f)

pi0.5 keypoint action path + held-out-operator split for the fold-freeform fit tests: keypoints_wristframe_pi transform (138-D wrist-first keypoints + cartesian ee_pose prompt proprio), HumanBimanualKeypoints converter, PI resizes openpi action projections for action_dim!=32, opsplit data configs (cam6d/wrist6d/kp), keypoint evaluator/viz, norm + train launchers, CPU smoke

(cherry picked from commit 294f2ec)

smoke: CPU keypoint smoke needs 1-process DDP (find_unused_parameters) — pl_model on_fit_start barriers; add working smoke_keypoints_cpu_ddp launcher

(cherry picked from commit 38712fd)

opsplit: hold out a full 10% of frames by operator (was 8%)

Co-Authored-By: Claude Fable 5.1 noreply@anthropic.com
Claude-Session: https://claude.ai/code/session_012AaLp3ZC4xVapZZvE5Pd4Z
(cherry picked from commit af90992)

opsplit launchers: validate every 20 epochs instead of 50 (viz still every 100, ckpt every 50)

Co-Authored-By: Claude Fable 5.1 noreply@anthropic.com
Claude-Session: https://claude.ai/code/session_012AaLp3ZC4xVapZZvE5Pd4Z
(cherry picked from commit 94f7187)

opsplit launchers: 72h wall, exclude the two thermally throttled H200 nodes

Co-Authored-By: Claude Fable 5.1 noreply@anthropic.com
Claude-Session: https://claude.ai/code/session_012AaLp3ZC4xVapZZvE5Pd4Z
(cherry picked from commit e2a677f)

eval: palm-origin cam-frame metric for keypoint runs (PIEvalKeypointPalm) + eval-only launcher for the epoch-49 checkpoint

Co-Authored-By: Claude Fable 5.1 noreply@anthropic.com
Claude-Session: https://claude.ai/code/session_012AaLp3ZC4xVapZZvE5Pd4Z
(cherry picked from commit 8dd5d5a)

eval launcher: typed h200 gres (untyped gres over a partition list landed on a V100)

Co-Authored-By: Claude Fable 5.1 noreply@anthropic.com
Claude-Session: https://claude.ai/code/session_012AaLp3ZC4xVapZZvE5Pd4Z
(cherry picked from commit 05982e3)

eval: offline launcher for cartesian opsplit checkpoints (same path as the keypoint palm eval)

Co-Authored-By: Claude Fable 5.1 noreply@anthropic.com
Claude-Session: https://claude.ai/code/session_012AaLp3ZC4xVapZZvE5Pd4Z
(cherry picked from commit 1671bf0)

eval launcher: per-job inductor/triton cache dirs (same-node autotune race)

Co-Authored-By: Claude Fable 5.1 noreply@anthropic.com
Claude-Session: https://claude.ai/code/session_012AaLp3ZC4xVapZZvE5Pd4Z
(cherry picked from commit 53101ea)

eval: viz_max_batches render cap on EvalVideo (partial pick of d0eda867)

Only the evaluator part of d0eda867 (combined mecka+abc pretrain): the
viz_max_batches kwarg on EvalVideo / TrainVizEvalVideo and the
add_dataloader_idx=False log_dict. eval_pi_keypoints_wristframe.yaml and
train_viz_pi_keypoints_wristframe.yaml pass viz_max_batches, so without
this the keypoint evaluator fails to instantiate. The abc/eva extrinsics
registry and combined-pretrain configs from that commit are left out.

(cherry picked from commit d0eda8673398e884d57e0bb6292e2999acce2d04, partial)

launchers: run from the kp-wristframe-port worktree (WT env override)

The cherry-picked fold-freeform / opsplit / eval launchers hard-coded the
fold-freeform-kp and fold-freeform-frames worktrees of the original stack.
They now default to .claude/worktrees/kp-wristframe-port (this branch, the
clean port onto main) and accept WT= to run from any checkout.

HPT twin of the keypoints-wristframe recipe: ~300M HPT + Qwen text stem on the 138-D wrist-frame keypoint action

Same data, split, transform and normalization as the pi0.5 keypoint fit
test (mecka_fold_freeform_opsplit_pi_keypoints / pi0.5_bc_mecka_keypoints),
for an HPT model:

• data/mecka_fold_freeform_opsplit_hpt_keypoints.yaml: byte-identical
  31-operator held-out lists and keypoints_wristframe_pi transform
  (fix_mecka_left_wrist, pad_proprio_gripper, stride 1); only the keymap
  changes to keypoints so the image lands on front_img_1. Because the
  norm-mode keymap strips camera keys, the pi twin's precomputed quantile
  stats (fold_freeform_opsplit_kp) carry exactly this config's key set and
  load through the provenance check unchanged (verified on the login node).
• model/hpt_bc_keypoints_wrist_300M.yaml: HPT flow BC, act_dim 138, stems
  = ResNet front image + 138-D state_keypoints + 20-D state_ee_pose + frozen
  Qwen3-Embedding-0.6B per-token stem on the sampled annotation. Widths from
  the lab's 300M reference (840 bus / 10 heads / 18x10x140 cross-attn / 320
  head); 23 trunk blocks -> ~303.7M trainable (24 blocks measured 312.2M:
  trunk 8.48M/block, heads 74.3M, stems 22.7M, encoder 11.6M; Qwen 595.8M
  frozen). Constant lr 5e-5, no scheduler.
• evaluator/eval_hpt_keypoints_wrist.yaml + train_viz_hpt_keypoints_wrist:
  HPTEvalVideo with the keypoint wrist->head revert (138 -> 126-D verified
  on a real sample) and the keypoints_wrist overlay; viz every 100 epochs,
  10 batches.
• train_zarr_mecka_kp_wrist_hpt.yaml: composed recipe.
• scripts: train_hpt_keypoints_fold_freeform_opsplit.sbatch (L40S per the
  repo's non-pi rule; H200 via sbatch -p/--gres override), a norm-stats job
  for a separate copy of the stats, and a 2-step GPU smoke.

Qwen/Qwen3-Embedding-0.6B must be in the HF cache before a compute-node run
(fetched on the login node 2026-09-09).

[aidang3019]

Clean port of the wrist-frame 6D + wrist-frame keypoint recipe onto main, for pi0.5 and HPT. 27 cherry-picks (Aseem raw-rot6d #483, aidan/pi-6d #487, a 10-commit subset of eval_gating #501, fold-freeform #602, keypoints/opsplit #603) + the viz_max_batches cap + launchers repointed + an HPT twin (same data/split/transform/norm stats; ~300M HPT + frozen Qwen3-Embedding stem, 138-D wrist-frame keypoint action, eval_hpt_keypoints_wrist). 48 pipeline tests pass; all four recipes compose. Known: fix_mecka_left_wrist rotates ee_pose only, the keypoint wrist frame (obs_wrist_pose) keeps the old left-hand convention.

[aidang3019]

• Fix PI0.5 raw rotation action encoding #616 👈 (View in Graphite)
• main

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