feat(infra): batchflow pipeline core, rldb unions, collapsed embodiment enum

[ElmoPA]

The reconciliation layer. Every file where main and the batchflow lineage
genuinely disagreed lands here, so everything above it is bulk-new code.

embodiment.py takes main's collapsed enum (HUMAN_* 1-3, EVA_* 4-6) and re-adds
PUSHSHAPES_SIM 15 / _STICK 16 / _SMALL_CIRCLE 17 -- pinned because trained
checkpoints and collected datasets encode those IDs.

zarr_dataset_multi.py is a 3-way merge against the fork point: main's
SafeS3EpisodeResolver, EvenStrideDataset, _evenly_spaced_indices and intrinsics
property, plus batchflow's _read_span, _annotations_for_span and
LocalEpisodeResolverWithEmbodimentOverride. action_chunk_transforms.py keeps
batchflow's DeltaAction alongside main's PadGripperZeros.

Co-Authored-By: Claude Opus 5 (1M context) noreply@anthropic.com

[ElmoPA]

• feat(configs): hydra config groups for the batchflow stack #558 : 2 dependent PRs (#560 , #564 )
• test: batchflow test suite and retained tooling #557
• feat(eval): batchflow evaluation stack #556
• feat(algo): H-Net, diffusion and BC algorithms #555
• feat(models): H-Net and diffusion model zoo #554
• refactor(sim): give agents their own interface so the env stops branching on them #559
• sim: Tsimulation with sim_v1 and sim_v2 #553
• feat(infra): batchflow pipeline core, rldb unions, collapsed embodiment enum #552 👈 (View in Graphite)
• main

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

[github-actions]

Claude Code Review

Review of PR #552

Summary

A large "reconciliation" PR merging batchflow and main branches, adding packed-episode dataloading, dual-stream H-Net probes, EMA/warm-start callbacks, span-safe transforms, and re-adding pushshapes sim embodiment enums. Scope is enormous and the diff was truncated at 80k chars, so this review covers what's visible.

Key concerns

1. Embodiment enum — ID pinning is correct, but verify (embodiment.py)

Re-adding PUSHSHAPES_SIM = 15, PUSHSHAPES_SIM_STICK = 16, PUSHSHAPES_SIM_SMALL_CIRCLE = 17 at pinned IDs is the right call for checkpoint/dataset compatibility. However:

• The enum strings need to exactly match what's stored on-disk. Confirm the lowercase forms (pushshapes_sim, pushshapes_sim_stick, pushshapes_sim_small_circle) are what existing zarr episodes actually carry.
• Add a test asserting EMBODIMENT["PUSHSHAPES_SIM"].value == 15 etc., so a future refactor can't silently renumber.

2. Coordinate frames in fold_span_transforms.py — smoke-quality math in a shared module

The docstring explicitly says "frame conventions are the raw quaternion / head-frame-relative-nothing forms — good enough to verify the ARCH runs + loss descends." This is fine for a smoke, but:

• The file lives under egomimic/rldb/embodiment/ (a production path), not tests/ or a smoke_utils/ module. Someone will import it in a real training run.
• The Aria action for humans is keypoints echoed as both action and proprio ("teacher-forced" in the docstring). If a real experiment picks this up, it will silently train on a trivialized task.
• Recommend: move to tests/regression/fixtures/ or gate with a loud runtime warning; at minimum rename to fold_span_transforms_SMOKE.py and add a module-level warnings.warn.

3. pl_data_utils.py — packed/annotation collate dispatch

_collate_fn_for only checks isinstance(dataset, ZarrEpisodePackedDataset). If the annotation-level packed dataset (ZarrAnnotationSpanPackedDataset, referenced in the docstring) already exists in the truncated portion, it won't hit pack_collate. Verify the dispatch covers all packed subclasses — probably safer as isinstance(dataset, _PackedBase) or a duck-typed getattr(dataset, "is_packed", False).

4. pl_model.py — torch.compile gate is fine, but silent

The EGO_COMPILE=1 env-gate replaces self.model.nets["outer_stage"] in-place. Two nits:

• If a downstream callback holds a reference to the pre-compile module, they'll diverge. The comment "measure before trusting any speedup" is honest but should probably be a log-warning at rank 0 only (currently the print fires on every rank).
• The compile happens before optimizer construction. Confirm parameter_groups still resolves post-compile (parameter names get an _orig_mod. prefix under compile — this can break the param_groups.py fnmatch specs and the dual_stream_probe name-based classification).

5. param_groups.py — fnmatch on named_parameters() is brittle under torch.compile

_leaf_patterns matches patterns like "*.stages.2.levels.1.*". Once compile wraps a module, names become _orig_mod.stages.2.... The globs will still match (leading *) but the dual-stream probe callback uses exact segment matches like "agnostic_input" in segs and segs.index("blocks") — those are fine with the prefix, but confirm across all probes. Add a test that runs the probe against a compiled module.

6. structure_ema_callback.py — sequencing race on resume

The comment acknowledges the ordering: setup() runs before on_load_checkpoint(), so _pending is loaded after freeze already happened. The on_train_epoch_start reconciliation looks correct, but:

• If a checkpoint is loaded mid-epoch (rare but possible with --ckpt_path), the shadow won't be restored until the next epoch starts. The batch-level EMA update in on_train_batch_end guards with if not self._active, so this is safe but subtle. Consider a on_train_start hook that also reconciles.
• No test coverage visible for the resume path.

7. warm_start_callback.py — resume guard uses env var + trainer.ckpt_path

The WARM_START_RESUME env var is a footgun: it must be set on every requeue. Slurm's automatic requeue won't set it. The trainer.ckpt_path check is the real guard; the env var is redundant and confusing — recommend dropping it or documenting it explicitly in AGENTS.md.

8. AGENTS.md — massive expansion (550+ lines added)

This is fine, but a lot of it reads like personal dev-log ("bit the s1c1/hireg smokes 2026-06-10"). Consider splitting: keep the operational rules (SLURM, kernels install, invocation) in AGENTS.md and move the H-Net design notes / debug histories to docs/hnet.md so AGENTS stays scannable.

9. utils/utils.py → pl_utils/utils.py move

This is a rename with an import path change. Grep for any external imports of egomimic.utils.utils — any downstream consumer (eval scripts, notebooks) will break. Same concern for timing_callback, instantiators, logging_utils, rich_utils moves.

10. Test coverage

---

Reviewed by Claude · Review workflow