pi train human new

[RyanPCo]

No description provided.

[RyanPCo]

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

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

[github-actions]

Claude Code Review

PR Review: pi train human new

Summary

Adds Mecka support to the Pi 0.5 training pipeline (keymap/transforms/configs), tightens optimizer/scheduler defaults, broadens MultiDataset retry to be cross-leaf, and adds split xyz/ypr eval metrics. Removes the mecka+scale cotrain config.

Key concerns

1. default_prompt: "" change in pi0.5_base.yaml — switching from "This is a bad action." to empty string is a meaningful training-behavior change. This silently affects every Pi 0.5 run, not just Mecka-human. Empty prompts may interact poorly with tokenizer / PaliGemma prompt conditioning (some tokenizers produce only BOS/EOS, some warn). Was this validated against existing Aria/Scale runs, or is this intended to become the new default for all experiments?

2. sample_frac: 1.0 → 0.2 in train_zarr_cartesian.yaml — this is a global default for the Cartesian training config. A 5× reduction in normalization sample fraction will change norm stats for every run using this config, not just this PR's experiments. If you have cached precomputed_norm_path references from prior runs, those become inconsistent with fresh recomputes. Consider overriding per-run instead of changing the base.

3. dataset_dir path change in paths/default.yaml — switching the default from /coc/flash7/scratch/... (Skynet) to /storage/project/r-dxu345-0/... (PACE) will break anyone training on the old cluster who relies on the default. Consider a cluster-specific override file rather than mutating the shared default.

4. Filter semantics change in mecka_pi.yaml — the new DatasetFilter with filter_lambdas uses task == 'fold_clothes', but the old (and deleted cotrain) config used task == 'folding_clothes' for mecka. Confirm the actual task string in the SQL DB — one of these is wrong, and if it's the new one, training silently gets zero episodes (or vice versa). The deleted cotrain file used fold_clothes, so this is probably the correct one, but worth double-checking against the DB.

5. Mecka keymap front_key branching — for cartesian_pi, the front cam is renamed base_0_rgb; for cartesian it stays VIZ_IMAGE_KEY. The viz configs (pi_cartesian.yaml) hardcode image_key: base_0_rgb for mecka, which is consistent with cartesian_pi mode. But get_transform_list(mode="cartesian") calls _build_aria_cartesian_bimanual_transform_list — does that transform expect base_0_rgb or the Mecka VIZ_IMAGE_KEY? If the transform list expects an Aria-style image key but the keymap emits base_0_rgb, you'll get a silent key-miss in the transform. Please verify.

6. MultiDataset global retry — broadening to range(len(self.index_map)) is the right call to escape a bad leaf, but it changes sampling distribution: a single bad leaf can now redirect its samples to other leaves, biasing the effective mixture ratio away from the configured train_dataloader_params. For small-scale debugging this is fine; for production cotraining this could matter. A warning log of cumulative retry counts per leaf during training would help catch silent rebalancing.

7. PI_DEBUG_PROMPTS debug logging — fine as gated debug, but self._debug_prompt_count is per-rank under DDP, so under 8 GPUs you'll see up to 24 log lines, not 3. Minor.

Suggestions

• Move default_prompt: "" and sample_frac: 0.2 into a Mecka-specific override (e.g., model/pi0.5_mecka.yaml, train_zarr_cartesian_mecka.yaml) instead of mutating shared base configs.
• Don't change paths/default.yaml; use a cluster-specific paths override.
• Verify the fold_clothes vs folding_clothes task string against the active SQL DB and add a test/assertion that the resolver returns >0 episodes.
• Verify get_transform_list(mode="cartesian") for Mecka uses the same image key (base_0_rgb) that the keymap emits in cartesian_pi mode — and ideally pass mode consistently.
• The _check_bounds short-circuit assumption in test_multi_retry.py is fragile to refactors; consider adding an embodiment attribute to _DummyLeaf and a real bounds path test, or comment-mark the test as intentionally bypassing bounds.
• Rename egomimic/rldb/zarr/test_multi_retry.py → tests/... or test_zarr_dataset_multi.py to match repo test conventions (and avoid pytest picking up a test file inside the importable package).
• The _split_mse helper in eval_pi.py is hardcoded to the 12-D bimanual layout — add a comment or assertion mapping it to HumanBimanualCartesianEuler.from32, and consider moving it to a util so unimanual / other layouts can later opt in.

Verdict: Request Changes

Main blockers are (a) mutating shared base configs (default_prompt, sample_frac, dataset_dir) that affect non-Mecka experiments, and (b) confirming the task-string filter actually matches DB rows. Once those are scoped to Mecka-specific overrides or confirmed safe, the rest is good — the retry fix and Mecka keymap extension are nice additions.

---

Reviewed by Claude · Review workflow