feat(hpt): stems and trunk as graph stages

[AnikethCheluva]

Decomposes HPT's input side into two stages instead of one opaque node:

observations --HPTStemStage--> hpt/tokens --HPTTrunkStage--> condition

HPTStemStage runs a PolicyStem per modality, each compressing to a fixed latent
count via cross-attention, and concatenates them along the token axis. That
pooling is what lets a 14-D pose and a 512-D image feature meet in one sequence
the trunk can attend over. Stems run in sorted key order so the token layout is
deterministic -- dict ordering would otherwise let DDP ranks diverge.

HPTTrunkStage attends over those tokens and pools to one vector, written as
"condition". That is deliberately the key the DP and flow heads already read,
so an HPT representation feeds any head in the graph without changing it --
which is the whole point of decomposing rather than wrapping.

Per-domain behaviour keeps the runner ignorant of domains: optional
domain_stems and the trunk's domain embedding both resolve their branch from a
batch key via resolve_homogeneous_scalar, the same mechanism
KeyedFeatureProjection uses.

Two adaptations were needed. nn.ModuleDict rejects "." in keys and batch keys
are dotted paths, so modules are stored under sanitized names with the mapping
kept explicit for readable checkpoint keys. And the cross-attention spec is
accepted as either a DictConfig or a plain dict, since hydra gives one and a
hand-built stem the other.

Tests cover the contracts, latent pooling, sorted ordering, dotted keys, width
mismatches, all three pooling modes, position-embedding bounds, domain
selection and rejection, that the position table is a buffer rather than a
trained parameter, and that the two compose into a runnable subgraph whose
condition is differentiable back to the stems.

Co-Authored-By: Claude Opus 5 noreply@anthropic.com

[AnikethCheluva]

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(arc): arc tokenize and detokenize as graph nodes #11 : 4 dependent PRs (#12 , #63 , #68 and 1 other)
• feat(hpt): configured HPT graph for the ABC yam baseline #46
• feat(abc): ABC-130k mcap->zarr pipeline + yam_bimanual embodiment #10
• feat(hpt): Add Qwen stem and annotation prompt stage #53
• fix(data): actually honour bounds_check #45
• feat(flow): flow-matching stages for the HPT head #44
• feat(hpt): stems and trunk as graph stages #43 👈 (View in Graphite)
• feat(hpt): trunk primitives and tensor helpers #42
• feat(val): nested validation dataset groups #9
• refactor(transforms): factor rotation representation out of transform modes #8
• main

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