feat(algo): H-Net, diffusion and BC algorithms

[ElmoPA]

algo/hnet, algo/diffusion (with its outer-stage variants) and algo/bc. Builds on
the model zoo in the previous commit.

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

[ElmoPA]

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

• 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 👈 (View in Graphite)
• 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
• main

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

[github-actions]

Claude Code Review

Review of PR #555: H-Net, diffusion and BC algorithms

Summary

Large refactor introducing algo/hnet (with PackedAlgoBase spine), algo/bc (WindowedBC as robomimic BC-RNN-GMM replica), and algo/diffusion (truncated). Also hoists shared embodiment-key resolution / obs-building / default reducer into Algo base class. Adds BATCHFLOW.md describing a stage-based convention (though this PR doesn't fully implement it — it's aspirational).

Key concerns

1. Algo.__init__ breakage of existing subclasses

The new Algo base introduces default implementations of compute_losses and log_info, but the default compute_losses reads predictions[f"{emb_id}_action_loss"] and iterates batch.keys() as if it's {emb_id: sub_batch}. This works for the new packed algos, but if any existing subclass (ACT, HPT, PI) doesn't override and expects the old NotImplementedError contract, silent behavior change or KeyError is possible. The docstring claims "subclasses… still override with their own superset reducer" — please confirm every existing subclass in the repo overrides both, or the default is a footgun.

Also, self.device is referenced in the default compute_losses but is not set by Algo.__init__ — relies on subclasses setting it. This is fine for the new algos but should be documented on the base.

2. _resolve_embodiment_keys — silent missing action key

for key in norm_stats.keys_of_type("action_keys", emb_id):
    if norm_stats.is_key_with_embodiment(key, emb_id) and key == self.ac_keys[emb]:
        self.resolved_ac_keys[emb_id] = key

If self.ac_keys[emb] doesn't match any key returned by norm_stats, resolved_ac_keys[emb_id] is silently missing, and the failure surfaces as a KeyError deep in _unpack_obs_actions. Please raise explicitly with the mismatch shown (embodiment enum + expected vs available keys) — mismatch between config embodiment strings and norm_stats topology is exactly the bug class that eats hours.

3. WindowedBC monkey-patches train_obs_transforms / episode_level_transforms

self.train_obs_transforms: list = []
self.episode_level_transforms: list = []

The comment admits this is a "pre-existing latent bug from the H-Net restructure." That's fine as a fix, but the real issue is that PackedAlgoBase.__init__ isn't being called and the base contract is being papered over. Either:

• Call PackedAlgoBase.__init__ properly and let it initialize these, or
• Move these to PackedAlgoBase class attributes (defaults) rather than instance attributes so subclasses inherit safely.

The current pattern will break again the next time PackedAlgoBase.process_batch_for_training gains a new attribute reference.

4. _cut_windows O(B·T_max) Python loop

pairs = []
for b in range(B):
    L = int(seq_lens[b].item())
    for s in range(L):
        pairs.append((b, s))

For B=32, T_max=500 that's 16k pairs enumerated in Python every training step, plus torch.randperm on a possibly-large list. Not fatal but shows up in profiler. Consider vectorizing (torch.cartesian_prod or masked index math). Low priority.

5. WindowedBCPolicy.step — last_action fallback

if state["queue"]:
    a = state["queue"].pop(0)
else:
    a = state["last_action"]

If state["last_action"] is None (first env step is somehow a non-obs step, e.g. t=0 with obs_stride>1? — actually 0 % anything == 0 so obs step, OK), this crashes with AttributeError on .unsqueeze. Add an explicit assertion that t=0 is always an obs step, or initialize last_action in init_step_state.

6. HNetPolicy._remap_legacy_input_keys contradicts PR description

The PR description references BATCHFLOW.md which says: "Backward ckpt compatibility explicitly NOT required (user 2026-07-10)." Yet the code carries a legacy-key remapping hook. Either delete the remap hook (align with declared policy) or delete the sentence in BATCHFLOW.md. Dead compatibility code compounds.

7. Diff truncated

The diff is cut off mid-HNetPolicy.step. I can't review:

• HNetOuterStage (stage-based H-Net trunk)
• algo/diffusion/ in full
• The PackedAlgoBase class body

Please split this PR. ~1600 lines for hnet/algo.py alone, plus BC, plus diffusion, plus the Algo base refactor — this cannot be reviewed atomically. Suggest:

• PR A: hoist shared helpers into Algo (small, mechanical, testable)
• PR B: algo/hnet package
• PR C: algo/bc package
• PR D: algo/diffusion package
• PR E: BATCHFLOW.md and any actual pipeline stages implementing it

8. BATCHFLOW.md is aspirational documentation of a design not implemented here

The doc claims stage(batch: dict) -> dict is the "one interface" and describes PipelineAlgo / stages_io / stages_hnet / stages_flow. Nothing in this diff shows these modules. Either land the doc with the pipeline it describes, or mark it as a design proposal — otherwise future readers will assume the codebase follows this convention when the algos in this same PR use HNetContext (which the doc explicitly

---

Reviewed by Claude · Review workflow