remove rotation bounds checks for norm stats

[RyanPCo]

No description provided.

[RyanPCo]

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

• remove rotation bounds checks for norm stats #480 : 2 dependent PRs (#481 , #490 ) 👈 (View in Graphite)
• optional viz videos on train dataset #476
• optional flag to generate validation videos for each task #475
• pi train human new #473
• main

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

[github-actions]

Claude Code Review

Review

Summary

Adds a reject_outliers flag to MultiDataset bounds checking and excludes Cartesian action rotation columns (YPR, indices 3-5, 9-11) from quantile checks while always retaining NaN/Inf rejection. Also adds a new Mecka pi fold-clothes config and bumps validation frequency.

Key concerns

1. Hardcoded YPR layout assumption. CARTESIAN_ACTION_YPR_INDICES = (3,4,5,9,10,11) assumes a fixed 12-dim bimanual [xyz, ypr, xyz, ypr] layout, gated only by arr.shape[-1] == 12. If any embodiment uses [xyz, xyz, ypr, ypr] or a different rotation representation (e.g., 6D rot, axis-angle) but happens to have 12 dims, rotation outliers will silently bypass checks AND xyz indices will be wrong. Worth asserting/documenting the layout contract, or keying off embodiment/transform metadata rather than shape.

2. Single-arm Cartesian (6-dim) not covered. If any embodiment emits a 6-dim actions_cartesian (single arm), the shape[-1] == 12 gate means rotation columns will be quantile-checked there. Probably fine if no such embodiment exists, but should be confirmed — Aria/Eva single-arm variants come to mind.

3. check_val_every_n_epoch: 100 → 150 in ddp_pi.yaml is unrelated to the PR title and affects every pi training run. Should be called out in the description, or split into its own PR. Easy to miss in a "norm stats" PR.

4. New config mecka_pi_fold_clothes_freeform.yaml — also unrelated to the PR title. Two minor things:

   • No reject_outliers override, so it inherits the new default True. Fine, but if the intent of this PR is to enable training freeform fold-clothes data that was previously being filtered out by bounds violations, you may actually want reject_outliers=false here. Worth confirming the intended interaction.
   • Filter lambda uses row['task'] == 'folding_clothes' — verify this matches the SQL task string exactly (the config name says fold_clothes).

5. NaN/Inf check moved before the reject_outliers short-circuit — good, this preserves the "NaN/Inf always rejected" invariant claimed in the docstring. ✅

Suggestions

• Add an assertion or comment in _check_bounds documenting the [xyz(3), ypr(3), xyz(3), ypr(3)] bimanual layout, ideally referencing where it's produced (Mecka/Eva transform).
• Consider sourcing rotation column indices from the embodiment's transform metadata rather than a class constant, so future rotation representations don't silently break this check.
• Split the check_val_every_n_epoch and new data config changes into a separate PR, or update the PR description.
• Add a test for the 6-dim single-arm case to lock in current behavior (quantile-checked end-to-end).
• The new tests reach into mds.key_types / mds.zarr_keys / mds.norm_stats directly — fine for unit testing, but consider a small helper to construct a stats-only MultiDataset to reduce coupling to internal field names.

Verdict

Request Changes — primarily to (a) confirm/document the bimanual Cartesian layout assumption and (b) separate or justify the unrelated ddp_pi.yaml and new Mecka config changes. The core norm-stats logic looks correct and well-tested.

---

Reviewed by Claude · Review workflow