[RyanPCo]

Summary

Two test tiers that prove HPT and Pi0.5 build and take optimizer steps for every vendor / robot that has configs (eva, aria, mecka, scale), plus the repo fixes they needed.

CPU tier (every PR, ~46 s on a compute node) — tests/unit/test_train_step.py: synthetic zarr episodes written through the real ZarrWriter (tests/fixtures/synthetic_episodes.py), the real top-level + per-vendor data + model configs composed with only the resolver swapped to LocalEpisodeResolver (tests/fixtures/recipes.py), the real trainHydra.train() for 2 steps. Asserts a finite loss and that the policy saw the canonical batch keys. Pi's openpi network is a stub that asserts on every observation field the EgoVerse wrapper must fill (state dim, image masks per camera, 32-wide action chunk).

GPU tier (opt-in) — tests/integration/test_train_step_gpu.py, pytest --integration -m gpu on a GPU node: real configs through the S3 resolver + SQL filters, real Pi weights, 5 bf16 steps, checkpoint written and reloaded.

Also — tests/unit/test_data_configs_compose.py: every data/*.yaml composes, resolves, and its key_map / transform_list nodes instantiate; known-broken configs are strict xfails with a reason.

Repo fixes the tests required

• pl_model.py: the two torch.distributed.barrier() calls crashed every single-process run; now guarded (and the guard is tested for the DDP path too).
• 14 data configs used an absolute ${train_datasets.X} for valid_datasets that cannot resolve under the data: mount (only data/human.yaml composed). Converted to the relative form eva_pi/cotrain_pi_base already use on ryanco/infra-fixes, line for line.
• has_head_pose: false was stale: scale exports do write obs_head_pose (external/scale/scripts/sfs_to_egoverse_zarr.py), and the cartesian transform pivots on it, so scale cartesian had never been trainable. Dropped from scale, scale_pi, mecka_scale_cotrain_pi, industry_eva_pi; docstring + CONTRIBUTING_DATA.md corrected.
• cotrain_pi_base.yaml human domain: keymap_mode: cartesian_pi (Pi reads base_0_rgb).
• Test harness: hermetic_env set HF_HUB_OFFLINE with monkeypatch, but huggingface_hub/transformers read it once at import, so the Pi tests actually ran online and 401'd in CI probing the gated PaliGemma repo for chat_template.jinja. It now patches the imported constants too, and the CI unit-test step sets HF_HUB_OFFLINE=1.
• GPU tier: last.ckpt is opened with torch.load(mmap=True) and its keys/shapes checked against the trained module, instead of reading the ~21 GB Pi checkpoint back just to assert a non-empty state_dict.
• MultiDataset: norm-stat samples with NaN/Inf are dropped (with a warning) before np.percentile. Found by the GPU tier: eva episode 2026-03-28-19-49-52-786000 has an all-NaN right.cmd_gripper; its NaN went into quantile_1/99 for one action dim and every normalized sample of the run became NaN.

Findings not fixed here (maintainer calls)

• Shipped pins are stale: data/scale.yaml pins 2026-03-16-01-22-26-448000, which app.episodes records as eva_bimanual; data/mecka.yaml pins 69199812208123403bbdb24f, which is not in the table. The GPU tier pins the eva/aria episodes their configs use and falls back to lab + debug=2 for mecka/scale.
• aria_pi, mecka_pi, scale_pi, industry_eva_pi, mecka_scale_cotrain_pi call get_keymap without keymap_mode (TypeError at instantiate); cotrain_pi_latent does not compose. All listed as strict xfails.
• A fully corrupt episode still exhausts the in-episode retry loop and crashes training (Entire dataset bad); left as-is so bad data surfaces rather than being skipped silently.
• human.py / zarr_dataset_multi.py carry ruff-format churn because pre-commit formats every touched file; the real changes are one docstring sentence and _drop_nonfinite_rows.

Test plan

Run (Slurm)	Result
pytest tests/unit on a CPU node	91 passed, 7 xfailed
pytest tests/unit/test_train_step.py	8 passed (4 HPT + 4 Pi stub) in 46 s
pytest tests/unit in a CI-equivalent venv (openpi --no-deps, tokenizer-only HF cache, no token, HF_HUB_OFFLINE=1)	264 passed, 5 xfailed
GPU tier, eva/aria/scale x {hpt, pi}, A40	6 passed
GPU tier, mecka x {hpt, pi}, A40 (post-review code)	2 passed
Final full GPU run on the final branch code, A40 (hoffman-lab)	8 passed in 29 min 30 s
GPU tier after the HF-offline fix, A40 (rl2-lab)	8 passed in 51 min (full checkpoint reload, basetemp on NFS)
GPU tier with the mmap checkpoint check, A40 (rl2-lab)	8 passed in 35 min

Spike that chose the Pi stub over openpi's tiny variant: stub 15-28 s / 1.3 GB vs dummy-gemma 71 s / 13 GB (and openpi's dummy variant needs a projector patch to run at all); both caught the same wiring bugs.

🤖 Generated with Claude Code

https://claude.ai/code/session_01N254Pmk9beYNooGou6T29R

[RyanPCo]

• import cloud packages lazily to support local training #631 : 2 dependent PRs (#632 , #633 )
• cache dataset resolver for faster training initialization #630
• validate episode hashes in data config before training #628
• modify pi camera map to match data #629
• add tests to ensure training works between every model and data vendor #626 👈 (View in Graphite)
• run CI tests on every PR #625
• add env variables for user-specific paths #624
• build: set pyproject and uv as the dependency managers #623
• fix(configs): fix paths for data hydra yaml configs #622
• refactor(algo): delete ACT #627 : 1 other dependent PR (#621 )
• tests: create dedicated test folder with unit and integration tests #620
• add pyright type checker #619
• main

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