feat: port the batchflow lineage onto current main

[ElmoPA]

elmo/batchflow-core forked from main on 2026-05-18 and never restacked: 170
commits ahead, 73 behind, carrying code main had since deleted. Rebasing it would
replay 170 commits to move files that mostly cannot conflict, so this ports the
net result onto current main instead. 88% of the change is purely additive --
whole new subsystems (models/hnet, models/diffusion, eval/core, the pipeline, 223
model configs) with no counterpart on main to fight with.

Composition: 775 added, 25 renamed (batchflow's file reorganisation), 13 deleted,
14 modified.

WHERE THE LINEAGES DISAGREED, MAIN WINS. Almost every conflict traces to one main
commit, 0bde7a7 "mandatory zarr.json intrinsics + collapse human embodiments",
which deleted CameraTransforms and renumbered the embodiments. batchflow's
changes to those files exist only to serve an API that is gone, so they are
dropped rather than reconciled:

• embodiment.py takes main's collapsed enum (HUMAN_* 1-3, EVA_* 4-6) and
  re-adds PUSHSHAPES_SIM 15 / _STICK 16 / _SMALL_CIRCLE 17. Those IDs are
  pinned because trained checkpoints and collected datasets encode them --
  renumbering would silently re-route every existing pushshapes result.
• data_schematic/default.yaml and hpt.yaml collapse aria_/mecka_/scale_bimanual
  into one human_bimanual block. Verified non-destructive: aria was a strict
  superset -- it alone carried state_keypoints and actions_keypoints -- and
  nothing was unique to mecka or scale.
• the CameraTransforms-era hpt/act/pi0.5 configs, viz_utils, rollout and
  egomimicUtils keep main's per-episode-intrinsics versions.

Three files needed a union rather than a choice, done as a 3-way merge against
the fork point so disjoint additions from both sides survive:

• zarr_dataset_multi.py keeps main's SafeS3EpisodeResolver, EvenStrideDataset,
  _evenly_spaced_indices, _jpeg_probe_failed and the intrinsics property, plus
  batchflow's _read_span, _annotations_for_span and
  LocalEpisodeResolverWithEmbodimentOverride. Ten ported files depend on that
  resolver and zarr_dataset_packed is written around _read_span. Both genuine
  overlaps resolved to main. ZarrEpisode._get_store is kept: the merged
  init carries its _pid state and read() calls it, and zarr v3 uses asyncio
  internally so the store is not fork-safe under DataLoader workers.
• action_chunk_transforms.py takes batchflow's DeltaAction alongside main's
  PadGripperZeros.
• trainHydra.py takes wiring from both; no overlap.

Two of git's 27 detected renames were SKIPPED as spurious -- it matched empty
init.py files across unrelated trees, e.g. egomimic/rldb/scripts/init.py
to Tsimulation/init.py, and applying them would have deleted live packages.
Likewise 33 of the 46 net "deletions" are files main added after the fork rather
than batchflow deletions; only the 13 batchflow actually removed are applied.

Verified by importing every module in the tree on a compute node: 209 import
clean. The 8 that do not are pre-existing or environmental, none introduced here:
missing optional deps (streamlit, oculus_reader, ppadb, robot_utils), a hardcoded
PACE path, a test importing build_human_bimanual_transform_list which exists on
neither branch since transform lists moved onto the embodiment classes, and three
modules that each register a global OmegaConf "eval" resolver without
replace=True.

Co-Authored-By: Claude Opus 5 (1M context) noreply@anthropic.com
Claude-Session: https://claude.ai/code/session_012V58H37tmcvgDthELMd5Xk

[ElmoPA]

• feat: port the batchflow lineage onto current main #548 👈 (View in Graphite)
• main

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

[github-actions]

Claude Code Review

Review

Summary

Massive port of the elmo/batchflow-core fork (170 commits ahead / 73 behind) onto current main, adopting main's decisions on every conflict point (collapsed embodiment enum, per-episode intrinsics, deleted CameraTransforms) while adding batchflow's new subsystems (hnet, diffusion, eval/core, pipeline, 223 model configs). 775 added / 25 renamed / 13 deleted / 14 modified.

Key concerns

1. Cannot actually review the code diff. The visible diff is 100% new documentation files (AUDIT.md, BATCHFLOW.md, CLAUDE.md, DESIGN.md) — the actual 775 file additions/renames and 14 modifications to zarr_dataset_multi.py, action_chunk_transforms.py, trainHydra.py, embodiment.py, data_schematic/*.yaml, etc. are all past the 80k-char truncation. The high-risk parts of this PR are unreviewed. For a port of this size, this needs to land as a stack of smaller PRs or with the critical files posted inline.

2. Embodiment enum re-adds (15/16/17) — verify no collision. PR claims PUSHSHAPES_SIM=15/STICK=16/SMALL_CIRCLE=17 are re-added on top of main's collapsed HUMAN_* 1-3, EVA_* 4-6. Confirm main hasn't since allocated 15-17 for anything else and that no zarr.json intrinsics/norm-stats keyed on the old batchflow numbering leak into any config under the 223 new model configs. The DESIGN.md even flags the "embodiment label double-resolution" risk (zarr.json=15 vs dict-key=16, silent merge, norm-stats last-writer-wins) — this must be explicitly verified before merge.

3. ZarrEpisode._get_store retention rationale needs a test. Keeping per-PID store state to avoid asyncio fork-safety issues under DataLoader workers is correct in spirit, but there's no mention of a test that exercises num_workers>1 end-to-end on the merged path. Zarr v3 fork-safety bugs typically manifest as intermittent decode failures, not clean crashes — silent data corruption risk.

4. Documented known bug not fixed in the port. AUDIT.md flags _Timesteps defaults (flip_sin_to_cos, downscale_freq_shift) as diverging from the reference DFoT. The PR ports the buggy defaults forward. Either fix now, or file an issue and reference it — leaving a ❌ BUG marked in a landed audit doc is a landmine for the next person.

5. CameraTransforms drop is irreversible for batchflow-era configs. 223 new model configs are added; confirm none of them reference the deleted CameraTransforms API. "Import clean on 209 modules" doesn't cover config resolution.

6. 8 import failures dismissed as "pre-existing or environmental." The three modules registering global OmegaConf eval resolvers without replace=True will crash any test that imports two of them in the same process. This is not "pre-existing" if the port increases the chance of two being imported together. Please add replace=True — one-line fix.

7. data_schematic collapse to human_bimanual. Claim of "aria was strict superset" needs a per-key diff posted, not asserted. If any existing checkpoint's norm-stats were computed under the separate aria_/mecka_/scale_bimanual blocks, this silently re-buckets them.

8. No tests listed in the diff summary. For a port touching training entry point, data loading, and action chunk transforms, "verified by importing every module" is not sufficient. At minimum need: (a) fixed-seed forward parity check on one existing training config, (b) norm-stats byte-equality check before/after on one existing dataset, (c) pytest tests/ green.

9. DeltaAction + PadGripperZeros union in action_chunk_transforms.py. Verify ordering when both are configured in the same transform list — delta computation vs gripper padding is order-sensitive.

Suggestions

• Split this PR. At minimum: (1) additive new subsystems (hnet/diffusion/pipeline/configs) as one PR — low risk, big diff; (2) the 3-way merges on zarr_dataset_multi.py, action_chunk_transforms.py, trainHydra.py as a second PR with explicit before/after tests; (3) embodiment enum + data_schematic changes as a third PR with a norm-stats parity check.
• Post the actual diff for the 14 modified files inline in the PR description, since they're where all conflict resolution lives.
• Add a parity test that runs one epoch of an existing training config pre/post-port and asserts identical loss to N decimal places. This is the only way to know the port didn't silently regress training.
• Fix the _Timesteps defaults while you're touching this code — cost is one line, benefit is checkpoint interoperability.
• Fix the OmegaConf eval resolver duplicate-registration — replace=True on all three sites.
• Delete or gitignore the AUDIT/BATCHFLOW/CLAUDE/DESIGN.md files or move them under docs/. Landing 4 top-level markdown files that mix stale port notes, TODO lists, and internal Claude session context into main is going to age badly. AUDIT.md in particular documents a known unfixed bug at the repo root.
• Explicitly list the 13 files being deleted in the PR description so reviewers can sanity-check none are still imported.

Verdict: Request Changes

The port strategy (main wins on conflicts, 3-way merge on genuine unions, skip spurious renames) is sound and the description is unusually thorough. But: the actual code diff is not visible in this review, the PR is far too large to land as a single unit, there's no

---

Reviewed by Claude · Review workflow