From ef7d3a3fd35e6dacb6f7bf0f2f10a9b9634a992c Mon Sep 17 00:00:00 2001 From: ElmoPA Date: Thu, 20 Aug 2026 23:17:12 -0400 Subject: [PATCH] feat(pipeline): sync current inference graph runtime --- egomimic/eval/inference_graph.py | 134 +++ egomimic/models/diffusion/denoising_nets.py | 33 +- egomimic/models/hnet/multi_stream_trunk.py | 44 +- egomimic/models/stems/input_modules.py | 33 + egomimic/pipeline/algo.py | 568 ++++++++++- egomimic/pipeline/stages_flow.py | 577 ++++++++++- egomimic/pipeline/stages_io.py | 153 ++- egomimic/pipeline/stages_seq.py | 333 ++++++- egomimic/pl_utils/pl_model.py | 140 ++- egomimic/rldb/embodiment/eva.py | 21 +- .../rldb/embodiment/fold_span_transforms.py | 938 +++++++++++++++++- egomimic/rldb/norm_stats.py | 145 +++ tests/test_inference_graph.py | 70 ++ tests/test_norm_stats.py | 48 + tests/test_pipeline_inference_graph.py | 135 +++ tests/test_pipeline_normal_dp_rollout.py | 131 +++ 16 files changed, 3353 insertions(+), 150 deletions(-) create mode 100644 egomimic/eval/inference_graph.py create mode 100644 egomimic/rldb/norm_stats.py create mode 100644 tests/test_inference_graph.py create mode 100644 tests/test_norm_stats.py create mode 100644 tests/test_pipeline_inference_graph.py create mode 100644 tests/test_pipeline_normal_dp_rollout.py diff --git a/egomimic/eval/inference_graph.py b/egomimic/eval/inference_graph.py new file mode 100644 index 000000000..cfd6a18e7 --- /dev/null +++ b/egomimic/eval/inference_graph.py @@ -0,0 +1,134 @@ +"""Small keyed inference graph with episode-scoped action-cache state. + +The evaluator calls one graph. A graph either serves an already committed +action from ``check_cache`` or runs ``preprocess -> model -> update_cache``. +Nodes expose explicit ``in``/``out`` port maps so the same node (or subgraph) +can be reused under different key names. +""" +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Callable, Mapping, MutableMapping, Sequence + + +class GraphContractError(RuntimeError): + """Raised when a node violates its declared keyed interface.""" + + +class KeyedNode: + """Adapt a callable to a shared key-value context. + + ``inputs`` and ``outputs`` use ``{callable_port: graph_key}``. Config + loaders may pass literal YAML keys ``in`` and ``out`` through ``ports``. + """ + + def __init__( + self, + fn: Callable[..., Any], + inputs: Mapping[str, str] | None = None, + outputs: Mapping[str, str] | None = None, + **ports: Any, + ) -> None: + self.fn = fn + self.inputs = dict(ports.pop("in", inputs or {})) + self.outputs = dict(ports.pop("out", outputs or {})) + if ports: + raise TypeError(f"unknown node fields: {sorted(ports)}") + + def __call__(self, context: MutableMapping[str, Any]) -> None: + missing = [key for key in self.inputs.values() if key not in context] + if missing: + raise GraphContractError(f"node input keys missing: {missing}") + result = self.fn(**{ + port: context[key] for port, key in self.inputs.items() + }) + if not self.outputs: + if result is not None: + raise GraphContractError("node with no outputs returned a value") + return + if len(self.outputs) == 1: + port = next(iter(self.outputs)) + # A mapping is commonly the value transported through one port; + # treat it as an output record only when it names that port. + if not isinstance(result, Mapping) or port not in result: + result = {port: result} + if not isinstance(result, Mapping): + raise GraphContractError("multi-output node must return a mapping") + missing_ports = [port for port in self.outputs if port not in result] + if missing_ports: + raise GraphContractError( + f"node output ports missing: {missing_ports}") + for port, key in self.outputs.items(): + context[key] = result[port] + + +class Subgraph(KeyedNode): + """A sequence of keyed nodes presented as one single-entry/exit node.""" + + def __init__(self, nodes: Sequence[KeyedNode], **ports: Any) -> None: + self.nodes = tuple(nodes) + super().__init__(self._run, **ports) + + def _run(self, **values: Any) -> Any: + local: dict[str, Any] = dict(values) + for node in self.nodes: + node(local) + if len(self.outputs) == 1: + port = next(iter(self.outputs)) + if port not in local: + raise GraphContractError(f"subgraph endpoint {port!r} missing") + return local[port] + return {port: local[port] for port in self.outputs} + + +@dataclass +class ActionCacheState: + """Mutable state owned by one policy instance and reset per episode.""" + + actions: list[Any] = field(default_factory=list) + + def reset(self) -> None: + self.actions.clear() + + def replace(self, actions: Sequence[Any]) -> None: + self.actions[:] = list(actions) + + def pop(self) -> Any: + if not self.actions: + raise GraphContractError("attempted to pop an empty action cache") + return self.actions.pop(0) + + def __bool__(self) -> bool: + return bool(self.actions) + + +class InferenceGraph: + """Fixed controller topology with an early terminal action edge.""" + + def __init__( + self, + *, + check_cache: KeyedNode, + inference_preprocess: KeyedNode, + model: KeyedNode, + update_cache: KeyedNode, + terminal_key: str = "policy.action", + ) -> None: + self.check_cache = check_cache + self.inference_preprocess = inference_preprocess + self.model = model + self.update_cache = update_cache + self.terminal_key = terminal_key + + def __call__(self, **inputs: Any) -> Any: + context: dict[str, Any] = dict(inputs) + self.check_cache(context) + if context.get(self.terminal_key) is not None: + return context[self.terminal_key] + self.inference_preprocess(context) + self.model(context) + self.update_cache(context) + if self.terminal_key not in context: + raise GraphContractError( + f"graph completed without terminal key {self.terminal_key!r}") + return context[self.terminal_key] diff --git a/egomimic/models/diffusion/denoising_nets.py b/egomimic/models/diffusion/denoising_nets.py index f14d0e53f..0d08fd46d 100644 --- a/egomimic/models/diffusion/denoising_nets.py +++ b/egomimic/models/diffusion/denoising_nets.py @@ -219,6 +219,7 @@ def __init__( n_groups=8, cond_predict_scale=False, feature_concatenate=False, + dp_exact=False, ): """ local conditioning and global conditioning scheme @@ -228,15 +229,27 @@ def __init__( start_dim = down_dims[0] dsed = diffusion_step_embed_dim - diffusion_step_encoder = nn.Sequential( - SinusoidalPosEmb(dsed // 2), - nn.Linear(dsed // 2, dsed * 2), - nn.Mish(), - nn.Linear(dsed * 2, dsed // 2), - ) - self.proj_cond = nn.Linear(ac_latent_seq * cond_dim, dsed // 2) - - cond_dim = dsed + self.dp_exact = bool(dp_exact) + if self.dp_exact: + # Stock Diffusion Policy: full-width step embedding, conditioning + # concatenated UNPROJECTED, so every residual block sees dsed + G. + diffusion_step_encoder = nn.Sequential( + SinusoidalPosEmb(dsed), + nn.Linear(dsed, dsed * 4), + nn.Mish(), + nn.Linear(dsed * 4, dsed), + ) + self.proj_cond = None + cond_dim = dsed + (cond_dim or 0) + else: + diffusion_step_encoder = nn.Sequential( + SinusoidalPosEmb(dsed // 2), + nn.Linear(dsed // 2, dsed * 2), + nn.Mish(), + nn.Linear(dsed * 2, dsed // 2), + ) + self.proj_cond = nn.Linear(ac_latent_seq * cond_dim, dsed // 2) + cond_dim = dsed in_out = list(zip(all_dims[:-1], all_dims[1:])) @@ -395,7 +408,7 @@ def forward( timesteps = timesteps.expand(sample.shape[0]) global_feature = self.diffusion_step_encoder(timesteps) - cond = self.proj_cond(global_cond) + cond = global_cond if self.proj_cond is None else self.proj_cond(global_cond) if global_cond is not None: global_feature = torch.cat([global_feature, cond], axis=-1) diff --git a/egomimic/models/hnet/multi_stream_trunk.py b/egomimic/models/hnet/multi_stream_trunk.py index 2797c516a..20b419a9a 100644 --- a/egomimic/models/hnet/multi_stream_trunk.py +++ b/egomimic/models/hnet/multi_stream_trunk.py @@ -408,18 +408,32 @@ def allocate_inference_cache(self, batch_size, max_seqlen=None, device=None, dty return None def _init_weights(self, initializer_range: float, parent_residuals: int) -> int: - n_residuals = parent_residuals + self.trunk.height - scaled_std = initializer_range / max(n_residuals, 1) ** 0.5 - for name, m in self.trunk.named_modules(): - if not isinstance(m, nn.Linear): - continue - if getattr(m.weight, "_no_reinit", False): - continue - # residual-out projections (attn .out + SwiGLU down-proj) get scaled - if name.endswith("out") or "w2" in name or "fc2" in name or "down" in name: - nn.init.normal_(m.weight, mean=0.0, std=scaled_std) - else: - nn.init.normal_(m.weight, mean=0.0, std=initializer_range) - if m.bias is not None: - nn.init.zeros_(m.bias) - return n_residuals + return init_multistream_trunk(self.trunk, initializer_range, parent_residuals) + + +def init_multistream_trunk(trunk, initializer_range: float, + parent_residuals: int) -> int: + """Residual-stream-aware Linear init for ONE ``MultiStreamTrunk``. + + Extracted from ``MultiStreamComputeStage._init_weights`` (math UNCHANGED) + so the flat ``pipeline.stages_seq.StreamTrunk`` can reuse it instead of + copying the rule. Returns the cumulative residual depth AFTER this trunk, + so callers can thread it into the next stage. + + ``trunk.height`` is 2 per layer (attention + FFN residual adds). + """ + n_residuals = parent_residuals + trunk.height + scaled_std = initializer_range / max(n_residuals, 1) ** 0.5 + for name, m in trunk.named_modules(): + if not isinstance(m, nn.Linear): + continue + if getattr(m.weight, "_no_reinit", False): + continue + # residual-out projections (attn .out + SwiGLU down-proj) get scaled + if name.endswith("out") or "w2" in name or "fc2" in name or "down" in name: + nn.init.normal_(m.weight, mean=0.0, std=scaled_std) + else: + nn.init.normal_(m.weight, mean=0.0, std=initializer_range) + if m.bias is not None: + nn.init.zeros_(m.bias) + return n_residuals diff --git a/egomimic/models/stems/input_modules.py b/egomimic/models/stems/input_modules.py index 99f84fd16..2cb39a161 100644 --- a/egomimic/models/stems/input_modules.py +++ b/egomimic/models/stems/input_modules.py @@ -166,6 +166,39 @@ def forward_padded(self, *, actions, obs, B, T, device, dtype, embodiment_id=Non c = self._encode(obs, T, embodiment_id) # (B, T, d_model) return c.to(dtype) + + def forward_packed_both(self, *, obs_packed, T_total, embodiment_id=None): + """SINGLE encode -> (fused pooled vector, per-modality tokens). + + Avoids the double-encode: the wrapped CondEncoderModule.encode() with + per_obs_keys=True produces the fused output_key AND every per-key + feature in ONE forward (one ResNet pass per camera). Returns + (fused (T, d_model), tokens (T, K, d) or None). The old + encode_tokens_packed re-ran the encoders; this replaces it so the + token path costs nothing beyond forward_packed. + """ + import torch as _t + enc = self.obs_encoder + if hasattr(enc, "encoders"): # per-embodiment dispatch + sub = (enc.encoders.get(str(embodiment_id)) + if hasattr(enc.encoders, "get") else enc.encoders[str(embodiment_id)]) + if sub is None: + fused = self._encode( + {k: v.unsqueeze(0) for k, v in obs_packed.items()}, + T_total, embodiment_id).squeeze(0) + return fused, None + enc_use = sub + else: + enc_use = enc + out = enc_use.encode({k: v.unsqueeze(0) for k, v in obs_packed.items()}, + T_action=T_total, embodiment_id=embodiment_id) + fused = out[enc_use.output_key].squeeze(0) # (T, d_model) + tokens = None + if getattr(enc_use, "per_obs_keys", False): + keys = [k for k in out.keys() if k != enc_use.output_key] + if keys: + tokens = _t.stack([out[k].squeeze(0) for k in keys], dim=1) # (T,K,d) + return fused, tokens def forward_packed( self, *, actions_packed, obs_packed, cu_seqlens, T_total, device, dtype, embodiment_id=None, diff --git a/egomimic/pipeline/algo.py b/egomimic/pipeline/algo.py index df60c6874..d4543b1e7 100644 --- a/egomimic/pipeline/algo.py +++ b/egomimic/pipeline/algo.py @@ -17,21 +17,38 @@ """ from __future__ import annotations -from collections import OrderedDict -from typing import List, Optional +from collections import OrderedDict, deque +from typing import List import torch import torch.nn as nn from egomimic.algo.algo import Algo +from egomimic.eval.inference_graph import ( + ActionCacheState, + InferenceGraph, + KeyedNode, +) from egomimic.rldb.embodiment.embodiment import get_embodiment_id from egomimic.pipeline.core import Pipeline, Stage, sum_losses +from egomimic.pipeline.stages_hnet import ApexLevel _PACKED_META_KEYS = ( "cu_seqlens", "max_seq_len", "seq_lens", "batch_size", "embodiment", "episode_idx", "chunk_offset", ) +# Explicitly retained after the norm-stats key mapping so evaluation can use +# per-episode camera calibration and invert wrist-frame targets for overlays. +# Do not replace this allowlist with keep_unmapped=True: raw dataset scratch +# keys must not leak into the policy batch. +_VIZ_PASSTHROUGH_KEYS = ( + "front_intrinsics", + "left_camera_extrinsics", + "right_camera_extrinsics", + "viz_current_wrist_poses", +) + class PipelineAlgo(Algo): def __init__( @@ -40,11 +57,20 @@ def __init__( norm_stats, domains: list = None, ac_keys: dict = None, + auxiliary_ac_keys: dict = None, device=None, action_horizon: int = 2560, train_obs_transforms: list | None = None, episode_level_transforms: list | None = None, init_ckpt: str | None = None, + rollout_apex_mode: str = "configured", + rollout_apex_window: int | None = None, + inference_stages: dict | None = None, + # Default OFF: the residual-aware init pass was dead code until + # 2026-08-18, so every existing config trained without it. Opting + # in per-config keeps other experiments' fresh launches unchanged. + init_range: float | None = None, + lr_multipliers: dict | None = None, **kwargs, ): super().__init__() @@ -55,13 +81,41 @@ def __init__( self.domains = list(domains or []) self.domain_by_id = {get_embodiment_id(e): e for e in self.domains} self.ac_keys = dict(ac_keys or {}) + # Evaluators written against the older Algo base read this (5 sites in + # eval/hpt/eval_hpt.py alone, which is the DEFAULT evaluator). Without it + # every PipelineAlgo run raised AttributeError at its first validation. + # Optional ctor arg rather than a hardcoded {} so a config can populate it. + self.auxiliary_ac_keys = dict(auxiliary_ac_keys or {}) self.action_horizon = int(action_horizon) self.device = device or torch.device( "cuda" if torch.cuda.is_available() else "cpu") self.train_obs_transforms = list(train_obs_transforms or []) self.episode_level_transforms = list(episode_level_transforms or []) + self.rollout_apex_mode = str(rollout_apex_mode) + self.inference_stages = inference_stages + self.rollout_apex_window = ( + int(rollout_apex_window) + if rollout_apex_window is not None + else None + ) self._resolve_embodiment_keys(norm_stats) self.nets = nn.ModuleDict({"policy": Pipeline(list(stages))}) + # RESIDUAL-STREAM-AWARE INIT (2026-08-18). The flat pipeline defined + # `_init_weights` on every stage but NOTHING ever called it, so the + # whole seq lineage trained from PyTorch's default kaiming init with no + # `1/sqrt(n_residuals)` damping -- unlike the nested lineage, whose + # container applies it by default (stages_hnet.py:1211,1223, init_range + # 0.02). This walks the flat stage list threading the cumulative + # residual depth, exactly as the nested chain threads it through + # `inner`. Set `init_range: null` to restore the un-initialized + # behaviour of runs launched before this date. + # + # Ordering matters: this runs BEFORE `init_ckpt` / any checkpoint + # restore, so a resume still ends up with the checkpoint's weights. + self.init_range = None if init_range is None else float(init_range) + self.lr_multipliers = dict(lr_multipliers) if lr_multipliers else None + if self.init_range: + self._apply_residual_aware_init(self.init_range) if init_ckpt: _ck = torch.load(init_ckpt, map_location="cpu", weights_only=False) _sd = _ck.get("state_dict", _ck) @@ -84,19 +138,197 @@ def __init__( for s in self.nets["policy"].stages: if type(s).__name__ == "TargetBuilder": self.replan_stride = int(s.stride) + # The evaluator-facing controller is the same keyed three-node graph + # used by DF. Model/history layout remains H-Net-specific and lives in + # the bound node methods below; weights and stages still come from the + # loaded checkpoint pipeline. + self._sim_action_cache = ActionCacheState() + self._sim_action_queue = self._sim_action_cache.actions + self._inference_graph = self._build_inference_graph() # ------------------------------------------------------------------ # + def _apply_residual_aware_init(self, init_range: float) -> None: + """Thread `_init_weights(init_range, parent_residuals)` down the stages. + + A stage returns the cumulative residual depth AFTER itself, so a trunk + adds 2 per layer, a Dechunk adds 1, and a Chunk adds 0 -- the same + accounting the nested chain does via `inner`. Stages without the method + (obs encoders, heads, loss stages) self-init in their own ctors and are + skipped, which is also what the nested container does. + """ + import inspect + + n_residuals = 0 + touched, skipped = [], [] + for stage in self.nets["policy"].stages: + fn = getattr(stage, "_init_weights", None) + if not callable(fn): + continue + # SIGNATURE GUARD. Not every `_init_weights` in this codebase takes + # (range, parent_residuals): `stages_dfot_v3.py:390` defines a + # zero-arg `_init_weights(self)` that its own ctor already called. + # Calling that here would TypeError and kill the job at + # construction -- including a requeue of a live DFoT run. Only + # stages with the 2-arg chain signature participate. + try: + params = [ + p for p in inspect.signature(fn).parameters.values() + if p.kind in (p.POSITIONAL_ONLY, p.POSITIONAL_OR_KEYWORD) + ] + except (TypeError, ValueError): + params = [] + if len(params) != 2: + skipped.append(type(stage).__name__) + continue + out = fn(init_range, n_residuals) + if not isinstance(out, int): + raise TypeError( + f"{type(stage).__name__}._init_weights must return the " + f"cumulative residual count (int), got {out!r}. Returning " + "None here is what silently broke this chain before.") + n_residuals = out + touched.append(f"{type(stage).__name__}->{n_residuals}") + print(f"[init] residual-aware init range={init_range} " + f"final_depth={n_residuals} inited={len(touched)} " + f"skipped={sorted(set(skipped)) or 'none'} :: " + + ", ".join(touched)) + + def parameter_groups(self, base_lr: float): + """Per-stage LR multipliers -> AdamW param groups (opt-in). + + `lr_multipliers` is `{glob-over-parameter-name: multiplier}`, e.g. + `{"policy.stages.10.*": 0.5}`. Matching uses fnmatch over + `named_parameters()` -- the same convention as + `egomimic.pl_utils.param_groups`. Returning None (the default, when no + multipliers are configured) makes `pl_model.configure_optimizers` fall + back to flat `parameters()`, so this is byte-identical for every + existing run and does NOT change the optimizer state layout they resume + into. + """ + if not self.lr_multipliers: + return None + import fnmatch + + buckets: OrderedDict = OrderedDict() + for name, p in self.named_parameters(): + if not p.requires_grad: + continue + mult = 1.0 + for pat, m in self.lr_multipliers.items(): + if fnmatch.fnmatch(name, str(pat)): + mult = float(m) + break + buckets.setdefault(mult, []).append(p) + groups = [{"params": ps, "lr": float(base_lr) * mult} + for mult, ps in buckets.items()] + print("[lr] per-stage multipliers -> %d groups: %s" + % (len(groups), {m: len(ps) for m, ps in buckets.items()})) + return groups + @property def policy(self) -> Pipeline: - return self.nets["policy"] + if hasattr(self, "nets") and "policy" in self.nets: + return self.nets["policy"] + if hasattr(self, "_rollout_policy"): + return self._rollout_policy + raise AttributeError("PipelineAlgo has no policy pipeline") + + @policy.setter + def policy(self, value: Pipeline) -> None: + """Install a pipeline without requiring the full training constructor. + + Small rollout adapters and contract tests build only the deployment + surface. Keeping this setter makes that surface use the same + ``nets['policy']`` ownership as production checkpoints instead of a + second shadow attribute. + """ + if not isinstance(value, Pipeline): + raise TypeError( + f"policy must be a Pipeline, got {type(value).__name__}") + if not hasattr(self, "nets"): + self.nets = nn.ModuleDict() + self.nets["policy"] = value + # Retain the deployment pipeline if a lightweight adapter later + # replaces ``nets`` with a model-only ModuleDict. + self._rollout_policy = value + + def _normal_rollout_adapter(self): + """Return the fixed-history adapter, if this policy owns one. + + Graph-only adapters may override ``step`` and intentionally omit a + training pipeline. In that case there is no normal-history adapter; + cache/preprocess still remain valid and model-specific ``step`` owns + its request layout. + """ + try: + policy = self.policy + except AttributeError: + return None + return next( + (stage for stage in policy.stages + if getattr(stage, "rollout_obs_steps", None) is not None), + None, + ) + + def apex_levels(self) -> list[ApexLevel]: + return [m for m in self.policy.modules() if isinstance(m, ApexLevel)] + + def set_apex_attention_mode( + self, + mode: str, + *, + window: int | None = None, + ) -> None: + apexes = self.apex_levels() + if not apexes: + raise RuntimeError("PipelineAlgo has no ApexLevel") + for apex in apexes: + apex.set_attention_mode(mode, window=window) + + def activate_rollout_apex_attention(self) -> None: + """Apply the CONFIGURED rollout apex regime at episode start. + + This runs from init_step_state on every episode and OVERRIDES any mode + set earlier (e.g. by an eval CLI flag). It is logged once per process + because a silent override here made rollouts look identical across + regimes and cost hours of misdiagnosis (2026-07-31). + """ + if self.rollout_apex_mode.lower() == "configured": + if not getattr(self, "_apex_regime_logged", False): + modes = {a.attention_mode for a in self.apex_levels()} + print("[apex/rollout] regime=CONFIGURED (from the model config); " + "live apex mode(s)=%s" % sorted(modes)) + self._apex_regime_logged = True + return + self.set_apex_attention_mode( + self.rollout_apex_mode, + window=self.rollout_apex_window, + ) + if not getattr(self, "_apex_regime_logged", False): + print("[apex/rollout] regime=%s window=%s APPLIED at episode start " + "(overrides any earlier set_attention_mode)" + % (self.rollout_apex_mode, self.rollout_apex_window)) + self._apex_regime_logged = True def _seed(self, emb_id: int, _batch: dict) -> dict: """Flat batch dict for one embodiment (the ONE carrier).""" obs = self._build_obs(_batch, emb_id) + actions = _batch[self.resolved_ac_keys[emb_id]] + # The NORMAL (per-sample) reader carries no packing metadata: one sample + # per frame, so seed a 1-token-per-sample grid. NormalObsExpand rewrites + # it to the real obs-history grid as the first stage. The packed reader + # supplies both keys and is untouched. + cu = _batch.get("cu_seqlens") + if cu is None: + n = int(actions.shape[0]) + cu = torch.arange(0, n + 1, device=actions.device, dtype=torch.long) + max_seq_len = 1 + else: + max_seq_len = int(_batch["max_seq_len"]) b = { - "actions": _batch[self.resolved_ac_keys[emb_id]], - "cu_seqlens": _batch["cu_seqlens"], - "max_seq_len": int(_batch["max_seq_len"]), + "actions": actions, + "cu_seqlens": cu, + "max_seq_len": max_seq_len, "embodiment": self.domain_by_id.get(emb_id), "aux/chunker": [], } @@ -118,17 +350,11 @@ def process_batch_for_training(self, batch): ) _batch = apply_episode_level_transforms( _batch, self.episode_level_transforms) - out = {} - is_packed = "cu_seqlens" in _batch - for key, value in _batch.items(): - if is_packed and key in _PACKED_META_KEYS: - out[key] = value - continue - key_name = self.norm_stats.zarr_key_to_keyname(key, emb_id) - if key_name is not None: - out[key_name] = value - out["_packed"] = is_packed - out = self.norm_stats.normalize(out, emb_id) + out, _ = self._prepare_loader_batch( + _batch, emb_id, packed_meta_keys=_PACKED_META_KEYS) + for key in _VIZ_PASSTHROUGH_KEYS: + if key in _batch: + out[key] = _batch[key] if self.train_obs_transforms and self.policy.training: for t in self.train_obs_transforms: out = t(out) @@ -190,6 +416,8 @@ def collect_chunkviz(self, batch): # rebinds seed["cu_seqlens"] to the decimated grid inside. raw_cu = seed["cu_seqlens"].cpu() gt_actions = seed["actions"].detach().float().cpu() + seed["aux/trunk_enc"] = [] # opt in to the trunk-level probes + seed["aux/trunk_dec"] = [] with torch.no_grad(): b = self.policy(seed) idxs = sorted({int(k.split("/")[1][1:]) for k in b @@ -202,21 +430,69 @@ def collect_chunkviz(self, batch): mask = b[f"chunk/L{i}/boundary_mask"] levels.append((prob[..., 1].detach().float().cpu(), mask.detach().cpu().to(torch.bool))) + if not levels: + # SEQ-PIPELINE FALLBACK (2026-08-13): stages_seq publishes + # aux/chunker records but its graphs have no flattener stage, + # so no chunk/L* flat keys exist and the scan above finds + # nothing (export then strips the anchor and crashes on an + # empty stack). Read the records directly. Compose contract + # wants INNERMOST FIRST; a record's boundary_mask lives on + # that level's INPUT grid and inner grids are strictly + # shorter, so ascending mask length = innermost..outermost. + for rec in sorted(b.get("aux/chunker") or [], + key=lambda r: int(r["boundary_mask"].shape[0])): + prob = rec["boundary_prob"] + if prob.dim() > 1 and prob.shape[-1] == 2: + prob = prob[..., 1] + levels.append((prob.detach().float().cpu(), + rec["boundary_mask"].detach().cpu().to(torch.bool))) # OUTERMOST pseudo-level: TargetBuilder's fixed stride decimation # exposed on the RAW frame grid so composition can anchor # (mask len == T_total). Kept frames: every `stride`-th per episode. stride = next((int(st.stride) for st in self.policy.stages if hasattr(st, "stride")), 1) + # DP CHUNK GRID (2026-08-14): chunkerless graphs (no learned levels) + # get their FIXED ACTION CHUNK as the display grid -- stride=1 would + # otherwise mean chunk-per-frame. chunk_len 16 -> 16-frame chunks. + if not levels: + stride = max(stride, next( + (int(st.chunk_len) for st in self.policy.stages + if hasattr(st, "chunk_len")), 1)) T_raw = int(raw_cu[-1]) keep = torch.zeros(T_raw, dtype=torch.bool) for _e in range(len(raw_cu) - 1): keep[int(raw_cu[_e]):int(raw_cu[_e + 1]):stride] = True levels.append((keep.float(), keep)) toks = b.get("apex/tokens") + if toks is None: + # CHUNKERLESS FALLBACK (2026-08-13): DP-style graphs have no + # apex; use the per-frame conditioning features so the PCA + # panel still shows the model's latents. Decimated below to the + # anchor chunk grid so PCA rows == chunks (2026-08-14). + toks = b.get("obs_feat") + if toks is not None and toks.shape[0] == keep.shape[0]: + toks = toks[keep] if toks is not None: best = toks.detach().float().cpu().numpy() entry = {"levels": levels, "tokens": best, "anchor": True, "gt_frame": gt_actions.numpy()} + # lowest trunk level (L0) A/S streams, on the decimated token grid + # lowest trunk level = smallest level index. Encoders append + # outermost-first but decoders append innermost-first, so pick by + # the tagged index rather than by list position. + for _key, _recs in (("", b.get("aux/trunk_enc") or []), + ("dec", b.get("aux/trunk_dec") or [])): + if not _recs: + continue + _lvl, _a, _s = min(_recs, key=lambda r: r[0]) + _suf = f"L0{_key}" + # A single-stream level (DualTrunkLevel stream_keys=[A], as in + # the StreamMLP arch) records S as None -- that level HAS no S + # stream. Emit what exists instead of crashing on None.float(). + if _a is not None: + entry[f"{_suf}_A"] = _a.float().cpu().numpy() + if _s is not None: + entry[f"{_suf}_S"] = _s.float().cpu().numpy() # tile pred chunks (kept grid, chunk C) back to the raw frame grid pred = b.get("pred_action") if pred is not None: @@ -245,7 +521,18 @@ def forward_eval(self, batch): predictions = OrderedDict() for emb_id, _batch in batch.items(): b = self.policy(self._seed(emb_id, _batch)) - predictions[self.resolved_ac_keys[emb_id]] = b["pred_action"] + ac = self.resolved_ac_keys[emb_id] + pred = b["pred_action"] + # Embodiment-SCOPED keys. A bare `ac` key is ambiguous in cotrain: + # with a matched action space both embodiments use the same ac_key, + # so the second overwrote the first. It also matched no evaluator -- + # HNetEvalVideo reads emb{id}_{ac} and silently skips when absent, + # which produced empty videos/ and zero Valid/ metrics at exit 0. + predictions[f"emb{emb_id}_{ac}"] = pred + name = self.domain_by_id.get(emb_id) + if name: + predictions[f"{str(name).lower()}_{ac}"] = pred # eval_hpt form + predictions.setdefault(ac, pred) # legacy bare return predictions # ------------------------------------------------------------------ # @@ -254,14 +541,22 @@ def forward_eval(self, batch): # ------------------------------------------------------------------ # def init_step_state(self, batch_size, T_max, device, dtype): assert batch_size == 1, "rollout is batch_size=1 (recompute-over-prefix)" + self.activate_rollout_apex_attention() return {"obs_prefix": [], "device": device, "dtype": dtype, "plan": None} @torch.no_grad() - def step(self, state: dict, obs_norm: dict, t: int, embodiment_id=None): + def step(self, state: dict, obs_norm: dict, t: int, embodiment_id=None, + obs_norm_history=None): state["obs_prefix"].append({k: v for k, v in obs_norm.items()}) - prefix = state["obs_prefix"] + prefix = (list(obs_norm_history) if obs_norm_history is not None + else state["obs_prefix"]) T = len(prefix) - dev, dt = state["device"], state["dtype"] + dev = state["device"] + normal_adapter = next( + (stage for stage in self.policy.stages + if getattr(stage, "rollout_obs_steps", None) is not None), + None, + ) b = { # NO "actions" key: plan() then provably excludes TargetBuilder, # posterior and every loss stage -> DENSE full-rate prefix, exactly @@ -270,51 +565,220 @@ def step(self, state: dict, obs_norm: dict, t: int, embodiment_id=None): "max_seq_len": T, "embodiment": embodiment_id, "aux/chunker": [], + # Set before plan/forward: NormalObsExpand uses this explicit marker + # to distinguish target-free deployment from malformed training. + "rollout_t": t, } - for k in prefix[0].keys(): - vs = [f[k] for f in prefix] - # env->zarr frames carry B=1 (e.g. (1,5), (1,3,H,W)) -> cat - # along dim0 gives the packed (T, ...) layout the encoders expect. - b[f"obs/{k}"] = (torch.cat(vs, 0).to(dev) - if torch.is_tensor(vs[0]) else vs[-1]) + if normal_adapter is not None: + # Standard DP was trained on a fixed n-frame sample axis. Recreate + # that exact shape at rollout: keep the newest n frames and repeat + # the episode's first frame at the left boundary (DP pad_before). + n = int(normal_adapter.rollout_obs_steps) + recent = prefix[-n:] + frames = [prefix[0]] * (n - len(recent)) + recent + for k in prefix[0].keys(): + vs = [f[k] for f in frames] + b[f"obs/{k}"] = ( + torch.cat(vs, 0).unsqueeze(0).to(dev) + if torch.is_tensor(vs[0]) else vs[-1] + ) + else: + for k in prefix[0].keys(): + vs = [f[k] for f in prefix] + # env->zarr frames carry B=1 (e.g. (1,5), (1,3,H,W)) -> cat + # along dim0 gives the packed (T, ...) layout encoders expect. + b[f"obs/{k}"] = (torch.cat(vs, 0).to(dev) + if torch.is_tensor(vs[0]) else vs[-1]) if state["plan"] is None: runnable, excluded = self.policy.plan(list(b.keys())) state["plan"] = runnable if excluded: names = [(type(s).__name__, miss) for s, miss in excluded] print(f"[PipelineAlgo.step] plan excluded (train-only): {names}") - b["rollout_t"] = t # streaming heads (SDPHead) key on this for stage in state["plan"]: b = stage(b) - return b["pred_action"][T - 1] # (C, D) decoded chunk at the last token + # Packed-prefix models produce T rows; NormalObsCollapse produces one. + # In both cases the deployable prediction is the final/current row. + return b["pred_action"][-1] # (C, D) decoded chunk at current token # ------------------------------------------------------------------ # - # Sim-eval entry (PackedSimEval calls this every env frame). + # Keyed inference graph. The control topology is universal; these bound + # nodes own H-Net's packed-prefix layout and checkpoint pipeline call. # ------------------------------------------------------------------ # + def _build_inference_graph(self) -> InferenceGraph: + defaults = { + "check_cache": {"in": {"obs": "obs"}, + "out": {"action": "policy.action"}}, + "inference_preprocess": { + "in": {"obs": "obs"}, + "out": {"request": "model.request"}}, + "model": {"in": {"request": "model.request"}, + "out": {"plan": "model.plan"}}, + "update_cache": { + "in": {"plan": "model.plan", "obs": "obs"}, + "out": {"action": "policy.action"}}, + } + cfg = self.inference_stages or {} + nodes = cfg.get("nodes", cfg) if hasattr(cfg, "get") else {} + + def ports(name): + node = (cfg.get("model", defaults[name]) if name == "model" + else nodes.get(name, defaults[name])) + return { + "in": dict(node.get("in", defaults[name]["in"])), + "out": dict(node.get("out", defaults[name]["out"])), + } + + return InferenceGraph( + check_cache=KeyedNode( + self._graph_check_cache, **ports("check_cache")), + inference_preprocess=KeyedNode( + self._graph_preprocess, **ports("inference_preprocess")), + model=KeyedNode(self._graph_model, **ports("model")), + update_cache=KeyedNode( + self._graph_update_cache, **ports("update_cache")), + terminal_key=str(cfg.get("terminal", "policy.action")), + ) + + def _reset_inference_graph(self, T_max=None) -> None: + param = next(self.nets.parameters()) + self._sim_state = self.init_step_state( + batch_size=1, T_max=int(T_max or self.action_horizon), + device=param.device, dtype=param.dtype) + self._sim_action_cache.reset() + normal = self._normal_rollout_adapter() + self._sim_raw_obs_history = deque( + maxlen=int(normal.rollout_obs_steps) if normal is not None else 1) + # Compatibility view for old diagnostics. Never replace this list: + # the ActionCacheState owns it for the lifetime of the policy. + self._sim_action_queue = self._sim_action_cache.actions + self._sim_prev_chunk = None + self._sim_prev_chunk_t = None + self._sim_ema_a = None + self._sim_blend_announced = False + self._sim_ema_announced = False + + def _graph_check_cache(self, obs: dict): + # Standard DP needs the immediately previous ENV frame at each replan, + # not the previous model-query frame. Keep a tiny raw ring here; cache + # hits still skip transforms, normalization, stacking and the model. + normal = self._normal_rollout_adapter() + if normal is not None: + snap = {} + for key, value in obs.items(): + if torch.is_tensor(value): + snap[key] = value.detach().clone() + elif hasattr(value, "copy"): + snap[key] = value.copy() + else: + snap[key] = value + self._sim_raw_obs_history.append(snap) + if not self._sim_action_cache: + return None + return self._graph_commit(self._sim_action_cache.pop()) + + def _graph_preprocess(self, obs: dict) -> dict: + adapter = getattr(self, "inference_obs_adapter", None) + normal = self._normal_rollout_adapter() + if normal is not None: + raw = list(self._sim_raw_obs_history) + raw = [raw[0]] * (int(normal.rollout_obs_steps) - len(raw)) + raw + history = [] + for frame in raw: + model_obs = adapter(frame) if adapter is not None else frame + history.append(self.norm_stats.normalize( + model_obs, self._sim_emb_id)) + return {"obs_norm": history[-1], + "obs_norm_history": history} + model_obs = adapter(obs) if adapter is not None else obs + return {"obs_norm": self.norm_stats.normalize(model_obs, self._sim_emb_id)} + + def _inference_cache_value(self, key: str, default): + overrides = getattr(self, "inference_cache_overrides", {}) or {} + if key in overrides: + return overrides[key] + cfg = self.inference_stages or {} + cache = cfg.get("cache", {}) if hasattr(cfg, "get") else {} + value = cache.get(key, default) if hasattr(cache, "get") else default + return default if value is None else value + + def _graph_model(self, request: dict) -> dict: + step_kwargs = {} + if request.get("obs_norm_history") is not None: + step_kwargs["obs_norm_history"] = request["obs_norm_history"] + chunk = self.step( + self._sim_state, request["obs_norm"], self._sim_t, + embodiment_id=self.domain_by_id.get(self._sim_emb_id), + **step_kwargs) + if chunk.dim() != 2: # (D,) + return {"actions": [chunk]} + + configured_keep = int(self._inference_cache_value( + "n_keep", self.replan_stride)) + n_keep = max(1, min(configured_keep, chunk.shape[0])) + # Existing eval-only interventions remain model-side and default OFF. + import os as _os + beta_default = float(self._inference_cache_value("blend", 0.0)) + beta = float(_os.environ.get( + "PUSHSHAPES_PLAN_BLEND", str(beta_default)) or 0) + if beta > 0 and chunk.shape[0] >= 2: + n_keep = max(1, chunk.shape[0] // 2) + prev = self._sim_prev_chunk + prev_t = self._sim_prev_chunk_t + if (prev is not None + and prev.shape[0] >= 2 * n_keep + and prev.shape[-1] == chunk.shape[-1] + and self._sim_t == prev_t + n_keep): + blended = chunk.clone() + blended[:n_keep] = ( + beta * prev[n_keep:2 * n_keep] + + (1 - beta) * chunk[:n_keep]) + chunk = blended + if not self._sim_blend_announced: + print(f"[stage4] PLAN_BLEND active beta={beta} " + f"t={self._sim_t} n_keep={n_keep} " + f"C={chunk.shape[0]}", flush=True) + self._sim_blend_announced = True + self._sim_prev_chunk = chunk.detach() + self._sim_prev_chunk_t = self._sim_t + return {"actions": [chunk[j] for j in range(n_keep)]} + + def _graph_update_cache(self, plan: dict, obs: dict): + self._sim_action_cache.replace(plan["actions"]) + return self._graph_commit(self._sim_action_cache.pop()) + + def _graph_commit(self, a_norm): + # PUSHSHAPES_ACTION_EMA=gamma: P6 output low-pass (normalized space). + import os as _os + gamma_default = float(self._inference_cache_value("action_ema", 0.0)) + gamma = float(_os.environ.get( + "PUSHSHAPES_ACTION_EMA", str(gamma_default)) or 0) + if gamma > 0: + prev = self._sim_ema_a + if prev is not None and self._sim_t > 0 and prev.shape == a_norm.shape: + a_norm = gamma * prev + (1 - gamma) * a_norm + if not self._sim_ema_announced: + print(f"[stage4] ACTION_EMA active gamma={gamma} " + f"t={self._sim_t}", flush=True) + self._sim_ema_announced = True + self._sim_ema_a = a_norm.detach() + out = self.norm_stats.unnormalize( + {self._sim_ac_key: a_norm}, self._sim_emb_id)[self._sim_ac_key] + return (out.detach().cpu().numpy().reshape(-1) + .astype("float32")) + + # Sim-eval entry: the evaluator only drives this one public method. def inference_step(self, obs_zarr: dict, t: int, emb_id: int, T_max=None): - import numpy as np from egomimic.rldb.embodiment.embodiment import get_embodiment if t == 0: - device = next(self.nets.parameters()).device - self._sim_state = self.init_step_state( - batch_size=1, T_max=int(T_max or self.action_horizon), - device=device, dtype=next(self.nets.parameters()).dtype) - self._sim_action_queue: list = [] + self._reset_inference_graph(T_max) + elif not hasattr(self, "_sim_state"): + raise RuntimeError("inference_step must begin with t == 0") + self._sim_t = int(t) + self._sim_emb_id = int(emb_id) embodiment_name = get_embodiment(emb_id).lower() - ac_key = (self.ac_keys[embodiment_name] if embodiment_name in self.ac_keys - else self.ac_keys[emb_id]) - if self._sim_action_queue: - a_norm = self._sim_action_queue.pop(0) - else: - obs_norm = self.norm_stats.normalize(obs_zarr, emb_id) - chunk = self.step(self._sim_state, obs_norm, t, - embodiment_id=self.domain_by_id.get(emb_id)) - if chunk.dim() == 2: # (C, D): keep replan_stride actions - n_keep = max(1, min(self.replan_stride, chunk.shape[0])) - self._sim_action_queue = [chunk[j] for j in range(n_keep)] - else: # (D,) - self._sim_action_queue = [chunk] - a_norm = self._sim_action_queue.pop(0) - out = self.norm_stats.unnormalize({ac_key: a_norm}, emb_id)[ac_key] - return out.detach().cpu().numpy().reshape(-1).astype(np.float32) + self._sim_ac_key = ( + self.ac_keys[embodiment_name] + if embodiment_name in self.ac_keys else self.ac_keys[emb_id]) + return self._inference_graph(obs=obs_zarr) diff --git a/egomimic/pipeline/stages_flow.py b/egomimic/pipeline/stages_flow.py index 37f63771d..1b30051c6 100644 --- a/egomimic/pipeline/stages_flow.py +++ b/egomimic/pipeline/stages_flow.py @@ -26,7 +26,11 @@ import torch.nn as nn import torch.nn.functional as F -from egomimic.models.diffusion.denoising_nets import SinusoidalPosEmb +from egomimic.models.hnet.moe_ffn import MoEFFN +from egomimic.models.diffusion.denoising_nets import ( + ConditionalUnet1D, + SinusoidalPosEmb, +) from egomimic.pipeline.core import Stage @@ -203,15 +207,34 @@ def forward(self, batch: dict) -> dict: if not self.training: with torch.no_grad(): T = a_top.shape[0] - x = torch.randn(T, self.C, self.D, device=a_top.device, + # ROLLOUT FAST PATH (2026-08-13). algo.py:375 consumes only + # pred_action[T-1], but this loop denoised ALL T rows at full + # sampler depth -> cost quadratic in episode length (~199x + # wasted on a 397-step episode; one episode cost ~66 min and no + # dn_* sim eval ever completed a single episode). T is a pure + # batch dim in these denoisers -- blocks attend over the chunk + # axis, never across T -- so row -1 is identical whether or not + # the other rows are computed. Mirrors SDPHead's rollout_t + # branch (stages_flow.py:453-455 slice, :495-497 scatter). + _stream = "rollout_t" in batch + _a = a_top[-1:] if _stream else a_top + _s = (s[-1:] if s is not None else None) if _stream else s + Tc = _a.shape[0] + x = torch.randn(Tc, self.C, self.D, device=a_top.device, dtype=a_top.dtype) dt = 1.0 / self.N for i in range(self.N): # t: 1 -> 0 - tt = torch.full((T,), 1.0 - i * dt, device=x.device, + tt = torch.full((Tc,), 1.0 - i * dt, device=x.device, dtype=x.dtype) - v, _, _ = self.net(x, tt, a_top, s, emb) + v, _, _ = self.net(x, tt, _a, _s, emb) x = x - dt * v - batch["pred_action"] = x.clamp(-1.0, 1.0) + x = x.clamp(-1.0, 1.0) + if _stream: + _out = torch.zeros(T, self.C, self.D, + device=a_top.device, dtype=a_top.dtype) + _out[-1] = x[0] + x = _out + batch["pred_action"] = x return batch @@ -510,30 +533,104 @@ class DiffusionHead(Stage): objective; no streaming buffer). Stateless -> TF-val and the rollout step() path are the same computation.""" - reads = ["a_top", "s", "embodiment"] - writes = ["pred_action", "loss/ddpm", "log/ddpm", "log/vA_frac"] + reads = ["a_top", "s", "embodiment"] # narrowed in __init__ when d_s is None + writes = ["pred_action", "loss/ddpm", "log/ddpm", "log/vA_frac", + "loss/moe_lb", "log/*"] def __init__(self, d_a: int, d_s: int, action_dim: int, chunk_len: int, embodiments: Optional[List[str]] = None, num_train_timesteps: int = 100, num_inference_steps: int = 16, d_model_a: int = 256, d_model_s: int = 128, n_layers: int = 4, n_heads: int = 4, ffn_mult: int = 4, - mask_mode: str = "sym"): + mask_mode: str = "sym", + denoiser: str = "dual", + denoiser_arch: str = "adaln", + moe_experts: int = 0, moe_top_k: int = 4, + moe_d_expert: Optional[int] = None, + moe_aux_weight: float = 0.01, + action_dims: Optional[dict] = None, + latent_dim: Optional[int] = None, + enc_hidden=256, enc_layers=3, enc_residual=False, + enc_per_stream: bool = False, + emit_loss: bool = True, + loss_space: str = "eps"): super().__init__() + self.emit_loss = bool(emit_loss) + # loss_space (OBJECTIVE ABLATION, user 2026-08-18): "eps" (default, + # unchanged) scores MSE in NOISE space -- the DDPM eps objective. + # "action" recovers x0 from eps and scores MSE in ACTION space: + # x0 = (x_t - sqrt(1-abar)*eps)/sqrt(abar), clamped to [-1,1] to match + # the sampler's clip_sample and bound the 1/sqrt(abar) blow-up at high + # noise. Published under the SAME aux keys so MaskedActionLoss / + # emit_loss score it unchanged. Isolates the objective (pooled + # conditioning + head unchanged). + if str(loss_space) not in ("eps", "action"): + raise ValueError(f"loss_space must be eps|action, got {loss_space!r}") + self.loss_space = str(loss_space) self.C, self.D = int(chunk_len), int(action_dim) self.N, self.S = int(num_train_timesteps), int(num_inference_steps) self.register_buffer("abar", _cosine_alphas_cumprod(self.N)) self.register_buffer( "inf_levels", torch.linspace(self.N - 1, 0, self.S).round().long()) - self.net = DualStreamDenoiser( - d_a_in=d_a, d_s_in=d_s, action_dim=action_dim, chunk_len=chunk_len, - embodiments=list(embodiments) if embodiments else ["shared"], - d_model_a=d_model_a, d_model_s=d_model_s, n_layers=n_layers, - n_heads=n_heads, ffn_mult=ffn_mult, mask_mode=mask_mode, - n_positions=int(chunk_len)) + if denoiser not in ("dual", "single"): + raise ValueError(f"denoiser must be dual|single, got {denoiser!r}") + self.denoiser_kind = str(denoiser) + # HETERO action dims (robot-human): the SAME construction SDPHead uses + # -- per-emb E_e/D_e MLP codec wrapped around a SHARED denoiser core + # that runs entirely in a common latent, so eva(14) and human(132) can + # cotrain on one core. Configs that pass neither action_dims nor + # latent_dim keep the homogeneous path byte-for-byte. + embs_ = [str(e) for e in embodiments] if embodiments else ["shared"] + self.Dmap = {e: int((action_dims or {}).get(e, action_dim)) for e in embs_} + self.rh = action_dims is not None or latent_dim is not None + if not d_s: + # single-stream on a replica-style obs path: no S key exists + self.reads = ["a_top", "embodiment"] + if self.rh: + L_lat = int(latent_dim if latent_dim is not None + else max(self.Dmap.values())) + self.net = LatentRHDenoiser( + d_a_in=d_a, d_s_in=d_s, action_dims=self.Dmap, latent_dim=L_lat, + chunk_len=chunk_len, embodiments=embs_, + d_model_a=d_model_a, d_model_s=d_model_s, n_layers=n_layers, + n_heads=n_heads, ffn_mult=ffn_mult, mask_mode=mask_mode, + n_positions=int(chunk_len), dual_stream=(denoiser == "dual"), + enc_hidden=enc_hidden, enc_layers=enc_layers, + enc_residual=enc_residual, enc_per_stream=enc_per_stream, + dual_arch="adaln", + moe_experts=moe_experts, moe_top_k=moe_top_k, + moe_d_expert=moe_d_expert, moe_aux_weight=moe_aux_weight) + elif denoiser == "single": + # One denoiser stream; A/S survive only as conditioning projections. + # MoE (when moe_experts > 0) swaps this block's FFN for experts. + self.net = SingleStreamDenoiserV2( + d_a_in=d_a, d_s_in=d_s, action_dim=action_dim, + chunk_len=chunk_len, d_model=d_model_a, n_layers=n_layers, + n_heads=n_heads, ffn_mult=ffn_mult, n_positions=int(chunk_len), + moe_experts=moe_experts, moe_top_k=moe_top_k, + moe_d_expert=moe_d_expert, moe_aux_weight=moe_aux_weight) + else: + # adaLN-Zero by default: the v1 core injects conditioning once, and + # every OTHER cell in this fleet is adaLN. Leaving v1 as the + # homogeneous default made "no latent_dim" silently mean "older + # architecture" and confounded dual-vs-MoE. + _cls = (DualStreamDenoiserV2 if str(denoiser_arch) == "adaln" + else DualStreamDenoiser) + self.net = _cls( + d_a_in=d_a, d_s_in=d_s, action_dim=action_dim, chunk_len=chunk_len, + embodiments=list(embodiments) if embodiments else ["shared"], + d_model_a=d_model_a, d_model_s=d_model_s, n_layers=n_layers, + n_heads=n_heads, ffn_mult=ffn_mult, mask_mode=mask_mode, + n_positions=int(chunk_len)) + + def _D(self, emb) -> int: + """Per-embodiment action dim (== the scalar D when homogeneous).""" + return self.Dmap.get(str(emb), self.D) def forward(self, batch: dict) -> dict: - a_top, s, emb = batch["a_top"], batch["s"], str(batch["embodiment"]) + a_top = batch["a_top"] + s = batch.get("s") + emb = str(batch["embodiment"]) if "target" in batch: x0 = batch["target"] # (T,C,D) T = x0.shape[0] @@ -542,23 +639,55 @@ def forward(self, batch: dict) -> dict: noise = torch.randn_like(x0) x_t = ab.sqrt() * x0 + (1 - ab).sqrt() * noise eps, e_a, e_s = self.net(x_t, t.float(), a_top, s, emb) - loss = F.mse_loss(eps, noise) - batch["loss/ddpm"] = loss - batch["log/ddpm"] = float(loss) - with torch.no_grad(): - na = e_a.norm(dim=-1).mean() - ns = e_s.norm(dim=-1).mean() - batch["log/vA_frac"] = float(na / (na + ns + 1e-8)) + # Publish the eps prediction/target so a downstream loss STAGE can + # score them (e.g. MaskedActionLoss, which drops the gripper dims + # for an embodiment that has no gripper). emit_loss=False hands the + # objective entirely to that stage. + if self.loss_space == "action": + x0_pred = ((x_t - (1 - ab).sqrt() * eps) + / ab.sqrt().clamp_min(1e-4)).clamp(-1.0, 1.0) + batch["aux/eps_pred"] = x0_pred + batch["aux/eps_target"] = x0 + else: + batch["aux/eps_pred"] = eps + batch["aux/eps_target"] = noise + # t is needed to bin the loss by noise level: the uniform-t + # average hides whether the model solves the LOW-t regime, + # which is the one that decides action fidelity. + # (DiffusionDiagnosticEval) + batch["aux/ddpm_t"] = t + if self.emit_loss: + loss = F.mse_loss(batch["aux/eps_pred"], batch["aux/eps_target"]) + batch["loss/ddpm"] = loss + batch["log/ddpm"] = float(loss) + if e_a is not None and e_s is not None: + with torch.no_grad(): + na = e_a.norm(dim=-1).mean() + ns = e_s.norm(dim=-1).mean() + batch["log/vA_frac"] = float(na / (na + ns + 1e-8)) if not self.training: with torch.no_grad(): T = a_top.shape[0] - x = torch.randn(T, self.C, self.D, device=a_top.device, + # ROLLOUT FAST PATH (2026-08-13). algo.py:375 consumes only + # pred_action[T-1], but this loop denoised ALL T rows at full + # sampler depth -> cost quadratic in episode length (~199x + # wasted on a 397-step episode; one episode cost ~66 min and no + # dn_* sim eval ever completed a single episode). T is a pure + # batch dim in these denoisers -- blocks attend over the chunk + # axis, never across T -- so row -1 is identical whether or not + # the other rows are computed. Mirrors SDPHead's rollout_t + # branch (stages_flow.py:453-455 slice, :495-497 scatter). + _stream = "rollout_t" in batch + _a = a_top[-1:] if _stream else a_top + _s = (s[-1:] if s is not None else None) if _stream else s + Tc = _a.shape[0] + x = torch.randn(Tc, self.C, self._D(emb), device=a_top.device, dtype=a_top.dtype) for j in range(self.S): # DDIM eta=0 tl = int(self.inf_levels[j]) - tt = torch.full((T,), float(tl), device=x.device, + tt = torch.full((Tc,), float(tl), device=x.device, dtype=x.dtype) - eps, _, _ = self.net(x, tt, a_top, s, emb) + eps, _, _ = self.net(x, tt, _a, _s, emb) ab_t = self.abar[tl] x0p = ((x - (1 - ab_t).sqrt() * eps) / ab_t.sqrt()) x0p = x0p.clamp(-1.0, 1.0) @@ -567,13 +696,159 @@ def forward(self, batch: dict) -> dict: x = ab_n.sqrt() * x0p + (1 - ab_n).sqrt() * eps else: x = x0p - batch["pred_action"] = x.clamp(-1.0, 1.0) + x = x.clamp(-1.0, 1.0) + if _stream: + _out = torch.zeros(T, self.C, self._D(emb), + device=a_top.device, dtype=a_top.dtype) + _out[-1] = x[0] + x = _out + batch["pred_action"] = x + moes = [m for m in self.net.modules() if isinstance(m, MoEFFN)] + aux = [m.last_aux_loss for m in moes if m.last_aux_loss is not None] + if aux: + batch["loss/moe_lb"] = torch.stack(aux).sum() + e = str(batch["embodiment"]) + f = torch.stack([m.last_expert_frac for m in moes + if m.last_expert_frac is not None]).mean(0) + for i_ in range(f.numel()): + batch[f"log/moe_expert_frac_{e}_e{i_}"] = f[i_] + batch[f"log/moe_gate_entropy_{e}"] = torch.stack( + [m.last_gate_entropy for m in moes + if m.last_gate_entropy is not None]).mean() return batch # --------------------------------------------------------------------------- # # Cursor/prev-action proprio (copycat-vs-smoothness experiment, user 2026-07-18) # --------------------------------------------------------------------------- # + +def _dp_alphas_cumprod(N: int, max_beta: float = 0.999, s: float = 0.008): + """diffusers' squaredcos_cap_v2 alphas_cumprod, EXACTLY. + + betas_for_alpha_bar caps each beta at 0.999 before the cumprod; forming + abar = f(t)/f(0) directly (what _cosine_alphas_cumprod does) skips that cap + and bottoms out at the clamp instead. Identical in the middle, ~24x apart at + the last level, which matters because x0_hat divides by sqrt(abar). + """ + ab = lambda u: math.cos((u + s) / (1 + s) * math.pi / 2) ** 2 + betas = [min(1.0 - ab((i + 1) / N) / ab(i / N), max_beta) for i in range(N)] + alphas = 1.0 - torch.tensor(betas, dtype=torch.float64) + return torch.cumprod(alphas, dim=0).float() + + +class DPUNetHead(Stage): + """Diffusion Policy's denoiser as a pipeline head. + + Wraps :class:`ConditionalUnet1D` in ``dp_exact`` mode -- proven bit-identical + to stock Diffusion Policy's UNet by ``unet_equiv.py`` (same param count, same + state_dict, 0.0 output difference). DDPM eps-prediction on the cosine + (``squaredcos_cap_v2``) ladder, which matches diffusers' schedule to 2.3e-07. + + Conditioning is whatever keys ``cond_keys`` names, concatenated: + + ``[a_top]`` -- the DP replica: one global conditioning vector + ``[a_top, s]`` -- dual-stream DP: the specific stream conditions too + + so the dual-stream and MoE variants are a one-line config change rather than + a different head. One conditioning vector per token, and the UNet denoises + that token's ``chunk_len`` action window, so batch = tokens. + """ + + writes = ["pred_action", "loss/dp", "log/dp"] + + def __init__(self, cond_keys: List[str], cond_dims: List[int], + action_dim: int, chunk_len: int, + num_train_timesteps: int = 100, num_inference_steps: int = 100, + down_dims: Optional[List[int]] = None, kernel_size: int = 5, + n_groups: int = 8, diffusion_step_embed_dim: int = 128, + cond_predict_scale: bool = True, clip_sample: bool = True, + sampler: str = "ddpm", + embodiments: Optional[List[str]] = None, + emit_loss: bool = True): + super().__init__() + self.emit_loss = bool(emit_loss) + if len(cond_keys) != len(cond_dims): + raise ValueError("DPUNetHead: one cond_dim per cond_key.") + self.cond_keys = [str(k) for k in cond_keys] + self.reads = list(self.cond_keys) + ["embodiment"] + self.C, self.D = int(chunk_len), int(action_dim) + self.N, self.S = int(num_train_timesteps), int(num_inference_steps) + self.clip_sample = bool(clip_sample) + if sampler not in ("ddpm", "ddim"): + raise ValueError(f"sampler must be ddpm|ddim, got {sampler!r}") + self.sampler = str(sampler) + G = int(sum(cond_dims)) + self.net = ConditionalUnet1D( + input_dim=self.D, cond_dim=G, + diffusion_step_embed_dim=int(diffusion_step_embed_dim), + down_dims=list(down_dims or [512, 1024, 2048]), + kernel_size=int(kernel_size), n_groups=int(n_groups), + cond_predict_scale=bool(cond_predict_scale), + dp_exact=True) + self.register_buffer("abar", _dp_alphas_cumprod(self.N)) + self.register_buffer( + "inf_levels", torch.linspace(self.N - 1, 0, self.S).round().long()) + + def _cond(self, batch: dict) -> torch.Tensor: + return torch.cat([batch[k] for k in self.cond_keys], dim=-1) # (T, G) + + def forward(self, batch: dict) -> dict: + g = self._cond(batch) + T, dev = g.shape[0], g.device + + if "target" in batch: + x0 = batch["target"] # (T,C,D) + t = torch.randint(0, self.N, (T,), device=dev) + ab = self.abar[t][:, None, None] + noise = torch.randn_like(x0) + x_t = ab.sqrt() * x0 + (1 - ab).sqrt() * noise + eps = self.net(x_t, t, global_cond=g) + batch["aux/eps_pred"] = eps + batch["aux/eps_target"] = noise + # t is needed to bin the loss by noise level: the uniform-t + # average hides whether the model solves the LOW-t regime, + # which is the one that decides action fidelity. + # (DiffusionDiagnosticEval) + batch["aux/ddpm_t"] = t + if self.emit_loss: + loss = F.mse_loss(eps, noise) + batch["loss/dp"] = loss + batch["log/dp"] = loss.detach() + + # Sampling is for ROLLOUT/EVAL only. Without this guard the loop ran on + # every training step -- 100 extra UNet forwards, all discarded (2.16s + # of a 2.27s step). DiffusionHead has had this guard all along. + if self.training: + return batch + with torch.no_grad(): + x = torch.randn(T, self.C, self.D, device=dev, dtype=g.dtype) + for i in range(self.S): + lv = self.inf_levels[i] + tt = torch.full((T,), int(lv), device=dev, dtype=torch.long) + eps = self.net(x, tt, global_cond=g) + ab = self.abar[lv] + x0_hat = (x - (1 - ab).sqrt() * eps) / ab.sqrt() + if self.clip_sample: # DP sets clip_sample=True + x0_hat = x0_hat.clamp(-1.0, 1.0) + if i + 1 >= self.S: + x = x0_hat + elif self.sampler == "ddim": + ab_prev = self.abar[self.inf_levels[i + 1]] + x = ab_prev.sqrt() * x0_hat + (1 - ab_prev).sqrt() * eps + else: + # DDPM ancestral, variance_type="fixed_small" -- what DP + # actually samples with. posterior mean from (x0_hat, x_t) + # plus noise scaled by the posterior variance. + ab_prev = self.abar[self.inf_levels[i + 1]] + beta = 1.0 - ab / ab_prev + coef_x0 = ab_prev.sqrt() * beta / (1.0 - ab) + coef_xt = (ab / ab_prev).sqrt() * (1.0 - ab_prev) / (1.0 - ab) + mean = coef_x0 * x0_hat + coef_xt * x + var = beta * (1.0 - ab_prev) / (1.0 - ab) + x = mean + var.clamp_min(1e-20).sqrt() * torch.randn_like(x) + batch["pred_action"] = x + return batch + class _DualStreamBlockAdaLN(_DualStreamBlock): """AdaLN-Zero variant of _DualStreamBlock: identical joint masked attention and per-stream FFNs, but each branch input is modulated per-position by @@ -730,7 +1005,10 @@ class _SingleStreamBlockAdaLN(nn.Module): conditioning c (shift, scale, gate) x2. Zero-init modulation => identity at init.""" - def __init__(self, d: int, n_heads: int, ffn_mult: int = 4): + def __init__(self, d: int, n_heads: int, ffn_mult: int = 4, + moe_experts: int = 0, moe_top_k: int = 4, + moe_d_expert: Optional[int] = None, + moe_aux_weight: float = 0.01): super().__init__() self.h, self.d = int(n_heads), int(d) assert self.d % self.h == 0 @@ -738,8 +1016,17 @@ def __init__(self, d: int, n_heads: int, ffn_mult: int = 4): self.qkv = nn.Linear(d, 3 * d) self.out = nn.Linear(d, d) self.norm2 = nn.LayerNorm(d) - self.ffn = nn.Sequential(nn.Linear(d, ffn_mult * d), nn.GELU(), - nn.Linear(ffn_mult * d, d)) + # MoE swaps ONLY the FFN -- attention, adaLN modulation and the + # residual structure are untouched, which is the standard way experts + # enter a transformer block. + if moe_experts: + self.ffn = MoEFFN(d, int(moe_d_expert or ffn_mult * d), + num_experts=int(moe_experts), + top_k=int(moe_top_k), + aux_weight=float(moe_aux_weight)) + else: + self.ffn = nn.Sequential(nn.Linear(d, ffn_mult * d), nn.GELU(), + nn.Linear(ffn_mult * d, d)) self.mod = nn.Sequential(nn.SiLU(), nn.Linear(d, 6 * d)) nn.init.zeros_(self.mod[1].weight) nn.init.zeros_(self.mod[1].bias) @@ -765,9 +1052,13 @@ class SingleStreamDenoiserV2(nn.Module): denoiser stream and no per-emb in/out (E_e/D_e own the per-emb mapping). Returns (v, None, None): no v_a/v_s partition => no vA_frac probe.""" - def __init__(self, d_a_in: int, d_s_in: int, action_dim: int, chunk_len: int, + def __init__(self, d_a_in: int, d_s_in: Optional[int], action_dim: int, + chunk_len: int, d_model: int = 256, n_layers: int = 4, n_heads: int = 4, - ffn_mult: int = 4, n_positions: Optional[int] = None): + ffn_mult: int = 4, n_positions: Optional[int] = None, + moe_experts: int = 0, moe_top_k: int = 4, + moe_d_expert: Optional[int] = None, + moe_aux_weight: float = 0.01): super().__init__() C, D = int(chunk_len), int(action_dim) L = int(n_positions) if n_positions else C @@ -775,14 +1066,21 @@ def __init__(self, d_a_in: int, d_s_in: int, action_dim: int, chunk_len: int, self.C, self.D, self.L = C, D, L self.in_x = nn.Linear(D, d) self.cond_a = nn.Linear(int(d_a_in), d) - self.cond_s = nn.Linear(int(d_s_in), d) + # d_s_in None -> no S conditioning at all (the DP-replica obs path + # produces a_top only). Otherwise A/S enter as two summed projections. + self.cond_s = nn.Linear(int(d_s_in), d) if d_s_in else None self.temb = nn.Sequential(SinusoidalPosEmb(d), nn.Linear(d, d), nn.GELU(), nn.Linear(d, d)) self.pos = nn.Parameter(torch.zeros(L, d)) nn.init.trunc_normal_(self.pos, std=0.02) self.vout = nn.Linear(d, D) self.blocks = nn.ModuleList( - _SingleStreamBlockAdaLN(d, n_heads, ffn_mult) for _ in range(int(n_layers))) + _SingleStreamBlockAdaLN(d, n_heads, ffn_mult, + moe_experts=moe_experts, + moe_top_k=moe_top_k, + moe_d_expert=moe_d_expert, + moe_aux_weight=moe_aux_weight) + for _ in range(int(n_layers))) self.norm_f = nn.LayerNorm(d) self.fmod = nn.Sequential(nn.SiLU(), nn.Linear(d, 2 * d)) nn.init.zeros_(self.fmod[1].weight) @@ -795,8 +1093,9 @@ def _temb(self, t, T, L): def forward(self, x_t, t, a_top, s, emb: Optional[str] = None): T, L, _ = x_t.shape - c = (self.cond_a(a_top)[:, None, :] + self.cond_s(s)[:, None, :] - + self._temb(t, T, L)) + c = self.cond_a(a_top)[:, None, :] + self._temb(t, T, L) + if self.cond_s is not None and s is not None: + c = c + self.cond_s(s)[:, None, :] X = self.in_x(x_t) + self.pos[None, :L] for blk in self.blocks: X = blk(X, c) @@ -832,7 +1131,10 @@ def __init__(self, d_a_in: int, d_s_in: int, action_dims: dict, latent_dim: int, mask_mode: str = "sym", n_positions: Optional[int] = None, dual_stream: bool = True, enc_hidden=256, dual_arch: str = "adaln", enc_layers=3, - enc_residual=False, enc_per_stream: bool = False): + enc_residual=False, enc_per_stream: bool = False, + moe_experts: int = 0, moe_top_k: int = 4, + moe_d_expert: Optional[int] = None, + moe_aux_weight: float = 0.01): super().__init__() self.dual_stream = bool(dual_stream) self.enc_per_stream = bool(enc_per_stream) @@ -888,7 +1190,9 @@ def _mlp(emb, d_in, d_out): self.core = SingleStreamDenoiserV2( d_a_in=d_a_in, d_s_in=d_s_in, action_dim=L, chunk_len=chunk_len, d_model=d_model_a, n_layers=n_layers, n_heads=n_heads, - ffn_mult=ffn_mult, n_positions=n_positions) + ffn_mult=ffn_mult, n_positions=n_positions, + moe_experts=moe_experts, moe_top_k=moe_top_k, + moe_d_expert=moe_d_expert, moe_aux_weight=moe_aux_weight) def forward(self, x_t, t, a_top, s, emb: str): e = emb if emb in self.enc else next(iter(self.enc)) @@ -904,3 +1208,202 @@ def forward(self, x_t, t, a_top, s, emb: str): v_lat, v_a, v_s = self.core(z, t, a_top, s, emb) v = self.dec[e](v_lat) # (T, L_pos, action_dim_e) return v, v_a, v_s + + +class MaskedActionLoss(Stage): + """DDPM eps-MSE with per-embodiment action dims EXCLUDED from the score. + + Reads ``aux/eps_pred`` / ``aux/eps_target`` (published by DiffusionHead when + a loss stage owns the objective) and writes ``loss/``. + + ``exclude_dims`` maps an embodiment to the action-dim indices that must not + be scored. The motivating case: with the gripper included the action is + 20-D per the layout [L xyz(3) rot6d(6) grip(1), R xyz(3) rot6d(6) grip(1)], + so the gripper slots are dims 9 and 19. eva has both; Aria has neither, and + scoring its unpopulated gripper columns would train the model to regress a + constant and pollute the reported MSE. The action space stays MATCHED at 20 + -- this is a masked loss, not a hetero head. + + Excluding a dim for EVERY embodiment would silently make it untrained while + still being emitted at rollout, so that is refused. + """ + + reads = ["aux/eps_pred", "aux/eps_target", "embodiment"] + + def __init__(self, exclude_dims: Optional[dict] = None, + name: str = "ddpm", embodiments: Optional[List[str]] = None, + weights: Optional[dict] = None): + super().__init__() + self.name = str(name) + # Per-embodiment loss weight. Scales that stream's gradient into the + # SHARED core as well as its own encoders -- which per-param-group LRs + # cannot do, since the shared trunk is one tensor set. Default 1.0 + # leaves every existing run byte-identical. + self.weights = {str(k): float(v) for k, v in (weights or {}).items()} + self.writes = [f"loss/{self.name}", f"log/{self.name}"] + self.exclude = {str(k): [int(i) for i in v] + for k, v in (exclude_dims or {}).items()} + embs = [str(e) for e in (embodiments or [])] + if embs and self.exclude: + common = set.intersection(*[set(self.exclude.get(e, [])) for e in embs]) \ + if all(e in self.exclude for e in embs) else set() + if common: + raise ValueError( + f"MaskedActionLoss: dims {sorted(common)} are excluded for " + f"EVERY embodiment {embs} -- they would never be trained yet " + f"still be emitted at rollout. Drop them from the action " + f"space instead.") + + def forward(self, batch: dict) -> dict: + pred = batch.get("aux/eps_pred") + if pred is None: # rollout / eval: nothing to score + return batch + tgt = batch["aux/eps_target"] + emb = str(batch["embodiment"]) + drop = self.exclude.get(emb) + if drop: + D = pred.shape[-1] + bad = [d for d in drop if d >= D] + if bad: + raise IndexError( + f"MaskedActionLoss: exclude_dims {bad} out of range for " + f"embodiment {emb!r} with action_dim {D}. The gripper dims " + f"are only valid on the 20-D cartesian layout.") + keep = torch.ones(D, dtype=torch.bool, device=pred.device) + keep[torch.tensor(drop, device=pred.device)] = False + pred, tgt = pred[..., keep], tgt[..., keep] + loss = F.mse_loss(pred, tgt) + w = self.weights.get(emb, 1.0) + batch[f"loss/{self.name}"] = loss * w + # LOG THE UNWEIGHTED loss on purpose: the weight changes the gradient, + # not the quantity we compare across cells. Logging loss*w would make + # a down-weighted run look better than an unweighted one for free. + batch[f"log/{self.name}"] = float(loss) + batch[f"log/{self.name}_weight"] = float(w) + batch[f"log/{self.name}_dims_scored"] = float(pred.shape[-1]) + return batch + + +# ------------------------------------------------------------------------- # +# TOKEN-CONDITIONING ABLATION (user, 2026-08-18) +# The DP failure survives every head (UNet/single/dual/MoE) and every encoder +# (VisualCore/HPTVisualEncoder), and cond_gain ~= 0.3%: the model ignores its +# observation. Common factor = the obs is POOLED into one a_top/s vector and +# injected via AdaLN. HPT instead keeps per-modality TOKENS and ATTENDS to +# them. This head is the single-variable test: same DDPM eps head, but the +# action chunk CROSS-ATTENDS to obs_tokens (from ObsEncoders.expose_tokens) +# instead of AdaLN on the pooled vector. +# ------------------------------------------------------------------------- # +class _XAttnBlock(nn.Module): + """DiT-ish block: self-attn over the chunk positions + cross-attn to obs + tokens + FFN. Self-contained MHA so shapes are unambiguous.""" + def __init__(self, d: int, n_heads: int, ffn_mult: int = 4, dropout: float = 0.1): + super().__init__() + assert d % n_heads == 0 + self.h, self.d = int(n_heads), int(d) + self.n1 = nn.LayerNorm(d); self.sq = nn.Linear(d, d); self.sk = nn.Linear(d, d) + self.sv = nn.Linear(d, d); self.so = nn.Linear(d, d) + self.nc = nn.LayerNorm(d); self.cq = nn.Linear(d, d); self.ck = nn.Linear(d, d) + self.cv = nn.Linear(d, d); self.co = nn.Linear(d, d) + self.n2 = nn.LayerNorm(d) + self.ff = nn.Sequential(nn.Linear(d, d * ffn_mult), nn.GELU(), + nn.Linear(d * ffn_mult, d)) + self.drop = nn.Dropout(dropout) + + def _mha(self, q, k, v): # (T,Nq,d),(T,Nk,d),(T,Nk,d) + T, Nq, d = q.shape; h = self.h; hd = d // h + q = q.view(T, Nq, h, hd).transpose(1, 2) + k = k.view(T, -1, h, hd).transpose(1, 2) + v = v.view(T, -1, h, hd).transpose(1, 2) + a = (q @ k.transpose(-1, -2)) * (hd ** -0.5) + a = a.softmax(-1) + o = (a @ v).transpose(1, 2).reshape(T, Nq, d) + return o + + def forward(self, x, ctx): # x (T,C,d) ctx (T,K,d) + s = self.n1(x); x = x + self.so(self._mha(self.sq(s), self.sk(s), self.sv(s))) + q = self.nc(x); x = x + self.co(self._mha(self.cq(q), self.ck(ctx), self.cv(ctx))) + x = x + self.ff(self.n2(x)) + return x + + +class XAttnDiffusionHead(Stage): + """DDPM eps head whose action chunk cross-attends to obs_tokens. + + reads obs_tokens (T,K,d_token) + target (train). No a_top/s AdaLN -- the + ONLY conditioning is cross-attention to the per-modality tokens. This is + the pooled-vs-token ablation; everything else (DDPM ladder, chunk, DDIM + sampling) mirrors DiffusionHead. + """ + reads = ["obs_tokens", "embodiment"] + writes = ["pred_action", "aux/eps_pred", "aux/eps_target", "loss/ddpm"] + + def __init__(self, action_dim: int, chunk_len: int, d_token: int, + embodiments=None, d_model: int = 384, n_layers: int = 6, + n_heads: int = 6, ffn_mult: int = 4, dropout: float = 0.1, + num_train_timesteps: int = 100, num_inference_steps: int = 16, + emit_loss: bool = True): + super().__init__() + self.C, self.D = int(chunk_len), int(action_dim) + self.N, self.S = int(num_train_timesteps), int(num_inference_steps) + self.emit_loss = bool(emit_loss) + self.register_buffer("abar", _cosine_alphas_cumprod(self.N)) + self.register_buffer("inf_levels", + torch.linspace(self.N - 1, 0, self.S).round().long()) + self.in_proj = nn.Linear(self.D, d_model) + self.ctx_proj = nn.Linear(int(d_token), d_model) + self.pos = nn.Parameter(torch.zeros(self.C, d_model)) + nn.init.trunc_normal_(self.pos, std=0.02) + self.temb = nn.Sequential(SinusoidalPosEmb(d_model), nn.Linear(d_model, d_model), + nn.SiLU(), nn.Linear(d_model, d_model)) + self.blocks = nn.ModuleList( + _XAttnBlock(d_model, n_heads, ffn_mult, dropout) for _ in range(int(n_layers))) + self.norm_f = nn.LayerNorm(d_model) + self.out = nn.Linear(d_model, self.D) + + def _denoise(self, x_t, t, ctx): # (T,C,D),(T,),(T,K,d) + h = self.in_proj(x_t) + self.pos[None] + h = h + self.temb(t.float())[:, None, :] + for blk in self.blocks: + h = blk(h, ctx) + return self.out(self.norm_f(h)) + + def forward(self, batch: dict) -> dict: + ot = batch.get("obs_tokens") + if ot is None: + raise KeyError("XAttnDiffusionHead needs batch['obs_tokens'] -- set " + "ObsEncoders.expose_tokens=True and the specific " + "obs_encoder per_obs_keys=True.") + ctx = self.ctx_proj(ot) # (T,K,d_model) + if "target" in batch: + x0 = batch["target"] # (T,C,D) + T = x0.shape[0] + t = torch.randint(0, self.N, (T,), device=x0.device) + ab = self.abar[t][:, None, None] + noise = torch.randn_like(x0) + x_t = ab.sqrt() * x0 + (1 - ab).sqrt() * noise + eps = self._denoise(x_t, t, ctx) + batch["aux/eps_pred"] = eps + batch["aux/eps_target"] = noise + batch["aux/ddpm_t"] = t + if self.emit_loss: + loss = F.mse_loss(eps, noise) + batch["loss/ddpm"] = loss + batch["log/ddpm"] = float(loss) + if not self.training: + with torch.no_grad(): + T = ctx.shape[0] + x = torch.randn(T, self.C, self.D, device=ctx.device, dtype=ctx.dtype) + for j in range(self.S): + tl = int(self.inf_levels[j]) + tt = torch.full((T,), float(tl), device=x.device, dtype=x.dtype) + eps = self._denoise(x, tt, ctx) + ab_t = self.abar[tl] + x0p = ((x - (1 - ab_t).sqrt() * eps) / ab_t.sqrt()).clamp(-1.0, 1.0) + if j + 1 < self.S: + ab_n = self.abar[int(self.inf_levels[j + 1])] + x = ab_n.sqrt() * x0p + (1 - ab_n).sqrt() * eps + else: + x = x0p + batch["pred_action"] = x.clamp(-1.0, 1.0) + return batch diff --git a/egomimic/pipeline/stages_io.py b/egomimic/pipeline/stages_io.py index 36c12475f..6d4b73980 100644 --- a/egomimic/pipeline/stages_io.py +++ b/egomimic/pipeline/stages_io.py @@ -34,12 +34,24 @@ class ObsEncoders(Stage): Writes A, S, time_pos.""" reads = ["obs/*", "cu_seqlens", "embodiment"] # actions NOT required (rollout) - writes = ["A", "S", "time_pos"] + writes = ["A", "S", "time_pos"] # (+ "obs_tokens" when expose_tokens) - def __init__(self, agnostic: nn.Module, specific: List[nn.Module]): + def __init__(self, agnostic: nn.Module, specific: List[nn.Module], + expose_tokens: bool = False): super().__init__() self.agnostic = agnostic self.specific = nn.ModuleList(specific) + # expose_tokens (token-conditioning ablation, 2026-08-18): when True, + # ALSO write batch["obs_tokens"] = per-modality feature tokens + # (T, K, d) from the specific encoder, so a cross-attention head can + # attend to modalities instead of the pooled S vector. Default False => + # byte-identical to before (A/S only). + self.expose_tokens = bool(expose_tokens) + if self.expose_tokens: + # Declare the extra key dynamically so the routing graph + # (tools/config_graph.py) sees the ObsEncoders -> head edge; + # a written-but-undeclared key reads as a dangling input. + self.writes = list(type(self).writes) + ["obs_tokens"] def forward(self, batch: dict) -> dict: obs_packed = {k.split("/", 1)[1]: v for k, v in batch.items() @@ -63,6 +75,28 @@ def forward(self, batch: dict) -> dict: kw = dict(actions_packed=actions, obs_packed=obs_packed, cu_seqlens=cu, T_total=T, device=dev, dtype=actions.dtype, embodiment_id=emb) + if self.expose_tokens: + # SINGLE-ENCODE path: one encode per stream yields fused (A/S) AND + # tokens, so the ResNets run ONCE (not twice). Mirrors HPT's + # encode-once cost. + token_sets = [] + a_fused, a_tok = self.agnostic.forward_packed_both( + obs_packed=obs_packed, T_total=T, embodiment_id=emb) + batch["A"] = a_fused.to(actions.dtype) + if a_tok is not None: + token_sets.append(a_tok) + S = None + for mod in self.specific: + s_fused, s_tok = mod.forward_packed_both( + obs_packed=obs_packed, T_total=T, embodiment_id=emb) + S = s_fused if S is None else S + s_fused + if s_tok is not None: + token_sets.append(s_tok) + batch["S"] = S.to(actions.dtype) + batch["time_pos"] = packed.frame_idx(cu) + if token_sets: + batch["obs_tokens"] = torch.cat(token_sets, dim=1).to(actions.dtype) + return batch S = None for mod in self.specific: c = mod.forward_packed(**kw) @@ -70,6 +104,24 @@ def forward(self, batch: dict) -> dict: batch["A"] = self.agnostic.forward_packed(**kw) batch["S"] = S batch["time_pos"] = packed.frame_idx(cu) + if False: + # collect per-modality tokens from BOTH streams so no encoder is + # left unused (the agnostic front encoder would otherwise get no + # gradient -> DDP unused-param crash). All per-key features share + # the encoder feature_dimension, so they concat along the token + # axis into (T, K_total, d). + token_sets = [] + at = self.agnostic.encode_tokens_packed( + obs_packed=obs_packed, T_total=T, embodiment_id=emb) + if at is not None: + token_sets.append(at) + for mod in self.specific: + stk = mod.encode_tokens_packed( + obs_packed=obs_packed, T_total=T, embodiment_id=emb) + if stk is not None: + token_sets.append(stk) + if token_sets: + batch["obs_tokens"] = torch.cat(token_sets, dim=1) # (T, K, d) return batch @@ -1011,3 +1063,100 @@ def __init__(self, rec): self.boundary_mask = rec["boundary_mask_s"] self.boundary_prob = rec["boundary_prob_s"] self.selected_probs = rec["selected_probs_s"] + + +# --------------------------------------------------------------------------- # +# Normal (per-sample) dataloader adapters +# +# The standard MultiDataset reader returns ONE SAMPLE PER FRAME: obs keys carry +# an (B, N_OBS, ...) history and the action chunk is already fixed-size +# (B, C, D). Everything downstream in this graph speaks the packed token +# contract (T flat tokens + cu_seqlens). These two stages bridge the two, so no +# other stage needs a second code path: +# +# NormalObsExpand (B,N,...) -> T=B*N tokens, one N-frame "episode" per +# sample. ObsEncoders and ObsStack's within-episode +# lookback clamp then work verbatim. +# NormalObsCollapse keep the LAST frame of each episode (the current one), +# restore a 1-token-per-sample grid, install the dataset's +# action chunk as "target". +# +# Put NormalObsExpand first in the stage list and NormalObsCollapse immediately +# after ObsStack. TargetBuilder is NOT used on this path -- the chunk comes from +# the dataset, which is the whole point of using the standard reader. +# --------------------------------------------------------------------------- # +class NormalObsExpand(Stage): + """Flatten the per-sample obs history into packed tokens. + + During training ``actions`` is required and preserved as the target chunk. + Rollout sends the same ``(1, n_obs, ...)`` observation shape with a + ``rollout_t`` marker and intentionally has no target. Supporting both + modes here keeps the encoder/stack/head path identical at train and deploy + time; only the dataset target is absent. + """ + + # ``actions`` is deliberately not a declared read: an observation-only + # rollout must be able to select this stage. forward() still requires it + # unless the explicit rollout marker is present. + reads = ["obs/*"] + writes = ["cu_seqlens", "max_seq_len", "_normal_target"] + + def __init__(self, n_obs_steps: int = 2): + super().__init__() + self.n = int(n_obs_steps) + self.rollout_obs_steps = self.n + + def forward(self, batch: dict) -> dict: + n = self.n + obs_keys = [k for k in batch + if k.startswith("obs/") and torch.is_tensor(batch[k])] + if not obs_keys: + raise ValueError("NormalObsExpand: no obs/* tensors in batch.") + ref = batch[obs_keys[0]] + B = ref.shape[0] + for k in obs_keys: + v = batch[k] + if v.ndim < 2 or v.shape[1] != n: + raise ValueError( + f"NormalObsExpand: {k} has shape {tuple(v.shape)}; expected " + f"(B, {n}, ...). Every obs key must be fetched with " + f"horizon=n_obs_steps={n} in the keymap.") + batch[k] = v.reshape(B * n, *v.shape[2:]) + dev = ref.device + target = batch.pop("actions", None) + if target is None and "rollout_t" not in batch: + raise ValueError( + "NormalObsExpand: batch has no 'actions' chunk outside rollout.") + # Keep the key present so the stage contract remains literal. Collapse + # consumes it and installs ``target`` only for a real training chunk. + batch["_normal_target"] = target # (B, C, D) | None + batch["cu_seqlens"] = torch.arange(0, B * n + 1, n, + device=dev, dtype=torch.long) + batch["max_seq_len"] = n + return batch + + +class NormalObsCollapse(Stage): + """Keep the current frame per sample; restore a 1-token-per-sample grid.""" + + writes = ["target", "cu_seqlens", "max_seq_len", "frame_idx", "time_pos"] + + def __init__(self, keys: List[str]): + super().__init__() + self.keys = [str(k) for k in keys] + self.reads = list(self.keys) + ["cu_seqlens", "_normal_target"] + + def forward(self, batch: dict) -> dict: + cu = batch["cu_seqlens"] + last = (cu[1:] - 1).to(dtype=torch.long) # current frame + for k in self.keys: + batch[k] = batch[k].index_select(0, last.to(batch[k].device)) + B, dev = last.numel(), last.device + batch["cu_seqlens"] = torch.arange(0, B + 1, device=dev, dtype=torch.long) + batch["max_seq_len"] = 1 + batch["frame_idx"] = torch.zeros(B, dtype=torch.long, device=dev) + batch["time_pos"] = torch.zeros(B, dtype=torch.long, device=dev) + target = batch.pop("_normal_target") + if target is not None: + batch["target"] = target + return batch diff --git a/egomimic/pipeline/stages_seq.py b/egomimic/pipeline/stages_seq.py index e0446daae..573ca9976 100644 --- a/egomimic/pipeline/stages_seq.py +++ b/egomimic/pipeline/stages_seq.py @@ -46,8 +46,12 @@ import torch.nn as nn from egomimic.models.hnet.isotropic_builder import build_isotropic -from egomimic.models.hnet.multi_stream_trunk import MultiStreamTrunk -from egomimic.models.hnet.per_emb import per_emb, pick +from egomimic.models.hnet.moe_ffn import MoEFFN +from egomimic.models.hnet.multi_stream_trunk import ( + MultiStreamTrunk, + init_multistream_trunk, +) +from egomimic.models.hnet.per_emb import PerEmb, per_emb, pick from egomimic.models.hnet.residual_mixer import build_residual_mixer from egomimic.models.hnet.routing import ChunkLayer, DeChunkLayer from egomimic.models.hnet.stages import _init_isotropic_linears @@ -103,6 +107,50 @@ def per_emb_module(factory_or_mod, use_per_emb: bool, embodiments): return per_emb(factory_or_mod, embodiments) # module-level helper, NOT the flag +def _flat_modules(obj): + """Yield the real sub-modules behind a shared module / PerEmb / container. + + Same idea as ``DualStreamChunkerStage._init_weights._members`` + (dual_stream_chunker.py:462-469): a per-emb slot is a ``PerEmb`` holding one + copy per embodiment, and every copy needs initializing -- ``pick()`` only + returns one. Used by the ``_init_weights`` pass, never in forward. + """ + if obj is None: + return + if isinstance(obj, PerEmb): + for m in obj.table.values(): + yield from _flat_modules(m) + elif isinstance(obj, nn.ModuleDict): + for m in obj.values(): + yield from _flat_modules(m) + elif isinstance(obj, nn.ModuleList): + for m in obj: + yield from _flat_modules(m) + elif isinstance(obj, nn.Module): + yield obj + + +def _init_plain_linears(obj, rng: float): + """``normal_(0, rng)`` every Linear under ``obj``, honouring ``_no_reinit``. + + For the NON-residual-stream params of a stage (down/up projections, fusion + heads): they do not feed a residual add, so they take the unscaled range -- + the same treatment ``_init_isotropic_linears`` gives a non-``out_proj`` + Linear, and the same rule ``DualStreamChunkerStage._init_weights`` applies + to its combine/proj list. + """ + if not isinstance(obj, nn.Module): + return + for m in obj.modules(): + if not isinstance(m, nn.Linear): + continue + if getattr(m.weight, "_no_reinit", False): + continue + nn.init.normal_(m.weight, mean=0.0, std=rng) + if m.bias is not None: + nn.init.zeros_(m.bias) + + class StreamTrunk(Stage): """MultiStreamTrunk over one or more named stream keys, at one resolution. @@ -168,7 +216,13 @@ def forward(self, batch: dict) -> dict: return batch def _init_weights(self, rng: float, parent_residuals: int) -> int: - return _init_isotropic_linears(self.trunk, rng, parent_residuals) + # REUSES the nested path's rule verbatim (MultiStreamComputeStage): + # this trunk ADDS 2 residual adds per layer, and its own residual-out + # projections are scaled by 1/sqrt(cumulative depth). The old body + # called `_init_isotropic_linears`, which (a) never added the trunk's + # height so the depth scale was frozen at the parent's value, and + # (b) returned None, breaking the chain for every later stage. + return init_multistream_trunk(self.trunk, rng, parent_residuals) class Framewise(Stage): @@ -226,6 +280,108 @@ def forward(self, batch: dict) -> dict: _write_grid(batch, ok, batch[ck], batch[mk], batch[tk]) return batch + def _init_weights(self, rng: float, parent_residuals: int) -> int: + # `_MLPBlocks` does `x = x + fc2(act(fc1(norm(x))))` per block, so each + # block is ONE residual add and `fc2` is the residual-out projection + # (`_init_isotropic_linears` matches it by name). Streams are parallel, + # not stacked, so the depth added is one stream's block count. + depth = max((len(m.blocks) for m in _flat_modules(self.blocks)), default=0) + n = parent_residuals + depth + return _init_isotropic_linears(self, rng, n) + + +class FramewiseMoE(Stage): + """Framewise, but the two streams are FUSED and routed through MoE experts. + + ``Framewise`` gives A one shared MLP and S one MLP per embodiment: the + per-embodiment capacity is HAND-WIRED. This stage concatenates the streams + and runs the fused token through ``n_layers`` pre-norm residual blocks whose + FFN is a :class:`MoEFFN` (top-k over ``num_experts`` SwiGLU experts), then + projects back to the original per-stream widths. Expert choice is learned + from token CONTENT ONLY -- the gate never sees the embodiment id -- so the + question this arm answers is whether per-embodiment specialisation is + DISCOVERED rather than declared. ``log/moe_expert_frac_`` is the readout. + + Fusing A with S means ``A_p`` is no longer embodiment-agnostic. That is not a + new leak in the arms this replaces: their final trunk already runs + ``mask_mode="sym"`` with ``allow_agnostic_cross=True`` over [A_mix, S_out], + so A already attends to S. This only moves the mixing earlier. + + Sizing note: match on ACTIVE params, not total. Active per token is + ``n_layers * top_k * 3 * d_total * d_expert`` (SwiGLU is 3 matrices), so the + dense Framewise it replaces sets ``d_expert``, while ``num_experts`` buys + total capacity for free at fixed compute. + """ + + def __init__(self, in_keys: List[str], dims: List[int], + out_keys: Optional[List[str]] = None, + n_layers: int = 4, num_experts: int = 8, top_k: int = 2, + d_expert: int = 288, aux_weight: float = 0.01, + embodiments: Optional[List[str]] = None, + final_norm: bool = True): + super().__init__() + self.in_keys = [str(k) for k in in_keys] + self.out_keys = [str(k) for k in (out_keys or in_keys)] + if len(self.in_keys) != len(dims): + raise ValueError("FramewiseMoE: one dim per in_key.") + if len(self.out_keys) != len(self.in_keys): + raise ValueError("FramewiseMoE: out_keys must match in_keys in length.") + self.dims = [int(d) for d in dims] + self.d_total = sum(self.dims) + self.reads = list(self.in_keys) + ["embodiment"] + self.writes = list(self.out_keys) + ["loss/moe_lb", "log/*"] + + self.norms = nn.ModuleList( + [nn.LayerNorm(self.d_total) for _ in range(int(n_layers))]) + self.moes = nn.ModuleList([ + MoEFFN(self.d_total, int(d_expert), num_experts=int(num_experts), + top_k=int(top_k), aux_weight=float(aux_weight)) + for _ in range(int(n_layers))]) + self.final = nn.LayerNorm(self.d_total) if final_norm else None + # One head per output stream: the fused token carries both, and the + # downstream Mix/trunk still expect the original per-stream widths. + self.heads = nn.ModuleList( + [nn.Linear(self.d_total, d) for d in self.dims]) + + def forward(self, batch: dict) -> dict: + x = torch.cat([batch[k] for k in self.in_keys], dim=-1) # (T, d_total) + for norm, moe in zip(self.norms, self.moes): + x = x + moe(norm(x)) + if self.final is not None: + x = self.final(x) + + for ik, ok, head in zip(self.in_keys, self.out_keys, self.heads): + batch[ok] = head(x) + if ok != ik: + ck, mk, tk = _grid_keys(ik) + if ck in batch: + _write_grid(batch, ok, batch[ck], batch[mk], batch[tk]) + + # Load balance: without it the router collapses onto one expert and the + # arm silently degenerates into a (smaller) dense MLP. + aux = [m.last_aux_loss for m in self.moes if m.last_aux_loss is not None] + if aux: + batch["loss/moe_lb"] = torch.stack(aux).sum() + + emb = str(batch["embodiment"]) + fracs = [m.last_expert_frac for m in self.moes + if m.last_expert_frac is not None] + if fracs: + # Per-embodiment, because the whole point is whether a content-only + # gate partitions experts BY embodiment on its own. + f = torch.stack(fracs).mean(0) + for e in range(f.numel()): + batch[f"log/moe_expert_frac_{emb}_e{e}"] = f[e] + batch[f"log/moe_gate_entropy_{emb}"] = torch.stack( + [m.last_gate_entropy for m in self.moes + if m.last_gate_entropy is not None]).mean() + return batch + + def _init_weights(self, rng: float, parent_residuals: int) -> int: + # `x = x + moe(norm(x))` once per layer -> n_layers residual adds. + n = parent_residuals + len(self.moes) + return _init_isotropic_linears(self, rng, n) + class Chunk(Stage): """Router + compress, over a LIST of streams sharing one boundary. @@ -260,7 +416,12 @@ def __init__(self, in_keys: List[str], out_keys: List[str], route_key: str, router_mixer_n_layers: int = 4, router_hidden_mult: float = 4.0, router_temp: float = 1.0, router_sample: bool = False, - res_scale: Optional[float] = None): + router_pre_layout: Optional[str] = None, + router_pre_detach: bool = False, + res_scale: Optional[float] = None, + # "default" = pre-2026-08-18 behaviour. Opt in with "zero" per config, + # so a config that never enabled init_range is untouched. + residual_proj_init: str = "default"): super().__init__() self.in_keys = [str(k) for k in in_keys] self.out_keys = [str(k) for k in out_keys] @@ -279,22 +440,32 @@ def __init__(self, in_keys: List[str], out_keys: List[str], route_key: str, # target_ratio by position, so stamping the level makes per-level # targets/schedules unambiguous instead of order-dependent. self.level = None if level is None else int(level) + # Stable full-resolution confidence diagnostic. ``selected_probs`` is + # exactly c_t = max(p_t, 1-p_t), but previously it was only reachable + # through the internal route / aux records. Publishing it gives + # checkpoint diagnostics and histogram evaluators a durable batch key. + self.ct_key = f"diag/L{self.level}_c_t" if n == 1 and self.route_on != "a": raise ValueError( "Chunk: a single-stream in_keys has no S to fuse; " "set route_on='a'.") self.reads = list(self.in_keys) + ["embodiment", "cu_seqlens", "max_seq_len", "time_pos"] self.writes = (list(self.out_keys) + [f"{k}@res" for k in self.in_keys] - + [self.route_key, "aux/chunker"]) + + [self.route_key, "aux/chunker", self.ct_key]) self.target_compression_ratio = float(target_compression_ratio) self.ratio_loss_weight = float(ratio_loss_weight) # MATCHES DualChunkerLevel (stages_hnet.py:471): an unset `_s` weight - # falls back to the A weight. With one shared boundary that applies the - # ratio term TWICE on the same b_t -- i.e. the effective weight is 2x - # the configured value. That looks accidental upstream, but every - # existing baseline trained under it, so reproducing it is required for - # the sequential baseline to be comparable. Pass ratio_loss_weight_s=0 - # explicitly to opt out. + # falls back to the A weight. + # + # NOTE (verified 2026-08-18, audit): on THIS stage that fallback is + # INERT -- there is no 2x double-count. The value is stamped into the + # aux record as `ratio_weight_s` below, but this stage always emits + # `separate_boundaries: False`, and `stages_io.RatioLoss` builds its S + # term ONLY from records with `separate_boundaries` True + # (stages_io.py:716-717). `ratio_weight_s` has no other consumer, so + # `loss/ratio_s` is never written and the effective ratio weight is + # exactly the configured `ratio_loss_weight`, applied once per level. + # (The 2x DOES apply on the nested path when separate_boundaries is on.) self.ratio_loss_weight_s = (float(ratio_loss_weight_s) if ratio_loss_weight_s else float(ratio_loss_weight)) @@ -311,7 +482,10 @@ def __init__(self, in_keys: List[str], out_keys: List[str], route_key: str, n_layers=int(router_mixer_n_layers), fusion=str(router_fusion), hidden_mult=float(router_hidden_mult), router_temp=float(router_temp), - router_sample=bool(router_sample), route_on=self.route_on) + router_sample=bool(router_sample), + router_pre_layout=router_pre_layout, + router_pre_detach=bool(router_pre_detach), + route_on=self.route_on) self._router = per_emb(mk_router, self._embs) if router_per_emb else mk_router() self.chunk_layer = ChunkLayer() @@ -319,8 +493,41 @@ def __init__(self, in_keys: List[str], out_keys: List[str], route_key: str, def _mk(fac, i): return per_emb(fac, self._embs) if self._per_emb[i] else fac() + # U-NET SKIP PROJECTION. Upstream H-Net and BOTH nested ports here + # ZERO-init this and mark it _no_reinit (models/hnet/stages.py + # `_mk_res_proj` equivalent at :1004-1007, dual_stream_chunker.py + # :305-310): the skip must start CLOSED so the model has to earn the + # bypass around the compressed path. The flat rewrite silently left it + # at PyTorch's default kaiming init, i.e. a full-rank random dense + # bypass live from step 0 -- exactly the trunk-bypass loophole + # residual_mixer.py says the mixer exists to close. + # ``residual_proj_init="default"`` restores the old (pre-2026-08-18) + # behaviour for reproducing runs launched before this fix. + if str(residual_proj_init) not in ("zero", "default"): + raise ValueError( + "Chunk: residual_proj_init must be 'zero' (paper/nested " + f"parity, default) or 'default', got {residual_proj_init!r}") + self.residual_proj_init = str(residual_proj_init) + + def _mk_res_proj(d, a): + # Reference parity (dual_stream_chunker.py:305-310, stages.py:1004-1007): + # a ZERO Linear is built UNCONDITIONALLY so the U-net skip starts + # CLOSED. The old `nn.Identity()` shortcut fired whenever + # apex_dims[i] == dims[i] and left the skip FULLY OPEN -- an + # accidental identity init, and the trunk-bypass loophole the + # residual mixer exists to close. + if self.residual_proj_init == "zero": + m = nn.Linear(d, d) + nn.init.zeros_(m.weight) + nn.init.zeros_(m.bias) + m.weight._no_reinit = True + return m + # "default" == pre-2026-08-18 behaviour, kept so existing runs + # resume with an unchanged state_dict. + return nn.Identity() if a == d else nn.Linear(d, d) + self.residual_proj = nn.ModuleList([ - _mk((lambda d=d: nn.Identity() if a == d else nn.Linear(d, d)), i) + _mk((lambda d=d, a=a: _mk_res_proj(d, a)), i) for i, (d, a) in enumerate(zip(self.dims, self.apex_dims))]) if self.grab_prev_end: self.prev_end_combine = nn.ModuleList([ @@ -377,6 +584,9 @@ def forward(self, batch: dict) -> dict: cu_seqlens_s=cu_s) bmask, bprob, sprob = (bpred.boundary_mask, bpred.boundary_prob, bpred.selected_probs) + # Keep gradients intact: consumers may use c_t diagnostically during + # training, while plotting code can detach at its own boundary. + batch[self.ct_key] = sprob residuals, next_cu, next_max = [], None, None for i, (ik, ok, d) in enumerate(zip(self.in_keys, self.out_keys, self.dims)): @@ -473,6 +683,22 @@ def forward(self, batch: dict) -> dict: return batch def _init_weights(self, rng: float, parent_residuals: int) -> int: + # Mirrors DualStreamChunkerStage._init_weights (dual_stream_chunker.py + # :461-500) exactly: + # * the combine / proj Linears take the UNSCALED range, + # * the ROUTER is left alone -- it self-inits in its ctor (proj_a + # identity, proj_s zero, `fuse` normal(0,0.02)); re-initing it here + # would destroy the warm start, + # * `residual_proj` is left alone -- zero-init IS its init, + # * a router pre-net is an isotropic stack scaled by its OWN depth, + # * the chunker adds NO residual-stream depth. + for attr in ("prev_end_combine", "proj_in"): + for m in _flat_modules(getattr(self, attr, None)): + _init_plain_linears(m, rng) + for r in _flat_modules(self._router): + rp = getattr(r, "router_pre", None) + if rp is not None: + _init_isotropic_linears(rp, rng, rp.height) return parent_residuals @@ -547,6 +773,12 @@ def forward(self, batch: dict) -> dict: return batch def _init_weights(self, rng: float, parent_residuals: int) -> int: + # `proj_out` is a plain up/down projection -> unscaled range. The + # `mixers` are NOT touched: their output layer is zero-init (that is + # what makes Dechunk start as `up + residual`) and now carries + # `_no_reinit`. Dechunk contributes ONE residual add. + for m in _flat_modules(self.proj_out): + _init_plain_linears(m, rng) return parent_residuals + 1 @@ -573,7 +805,12 @@ def forward(self, batch: dict) -> dict: return batch def _init_weights(self, rng: float, parent_residuals: int) -> int: - return _init_isotropic_linears(self.net, rng, parent_residuals) + # `Isotropic.height` counts 2 residual adds per transformer layer + # (blocks.py:1034,1066) -- T10 -> 20. Adding it here is what makes the + # apex's out_proj/fc2 scale by 1/sqrt(depth-through-the-apex) instead + # of the parent's (much smaller) count. + n = parent_residuals + int(getattr(self.net, "height", 0)) + return _init_isotropic_linears(self.net, rng, n) class Mix(Stage): @@ -854,6 +1091,74 @@ def forward(self, batch: dict) -> dict: return batch + +class ObsStack(Stage): + """Concatenate the last ``n_obs_steps`` frames per token (DP's obs history). + + Diffusion Policy conditions on a short observation window, not a single + frame: ``global_cond`` is the per-frame feature repeated over + ``n_obs_steps`` and concatenated, so ``global_cond_dim = feat * n_obs_steps``. + + Lookback is clamped WITHIN the episode, so the first frames of an episode + repeat themselves rather than reaching into the previous episode's tail -- + the same thing DP's ``pad_before`` does at a buffer boundary. Getting this + wrong is silent: it leaks across episode seams and only shows up as slightly + better-than-real validation. + """ + + def __init__(self, in_keys: List[str], out_keys: List[str], + n_obs_steps: int = 2): + super().__init__() + self.in_keys = [str(k) for k in in_keys] + self.out_keys = [str(k) for k in out_keys] + if len(self.in_keys) != len(self.out_keys): + raise ValueError("ObsStack: one out_key per in_key.") + self.n = int(n_obs_steps) + self.reads = list(self.in_keys) + ["cu_seqlens"] + self.writes = list(self.out_keys) + + def forward(self, batch: dict) -> dict: + ref = batch[self.in_keys[0]] + T, dev = ref.shape[0], ref.device + cu = batch["cu_seqlens"].to(device=dev, dtype=torch.long) + pos = _within_episode_time_pos(cu, T, dev) # (T,) index in episode + idx = torch.arange(T, device=dev) + # frame t-k, but never before this episode's first frame + gathers = [idx - torch.minimum(torch.full_like(pos, k), pos) + for k in range(self.n - 1, -1, -1)] + for ik, ok in zip(self.in_keys, self.out_keys): + x = batch[ik] + batch[ok] = torch.cat([x.index_select(0, g) for g in gathers], dim=-1) + ck, mk, tk = _grid_keys(ik) + if ok != ik and ck in batch: + _write_grid(batch, ok, batch[ck], batch[mk], batch[tk]) + return batch + + +class Concat(Stage): + """Plain concatenation along the feature axis -- no projection, no MLP. + + ``Fuse`` puts a learned MLP between the encoders and the trunk. Diffusion + Policy does not: its obs feature is the spatial-softmax keypoints + concatenated with proprio (32*2 + 2 = 66 for PushShapes), fed straight to + the UNet. This exists so a DP replica can be exact rather than + approximately-DP. + """ + + def __init__(self, in_keys: List[str], out_key: str): + super().__init__() + self.in_keys = [str(k) for k in in_keys] + self.out_key = str(out_key) + self.reads = list(self.in_keys) + self.writes = [self.out_key] + + def forward(self, batch: dict) -> dict: + batch[self.out_key] = torch.cat([batch[k] for k in self.in_keys], dim=-1) + ck, mk, tk = _grid_keys(self.in_keys[0]) + if ck in batch: + _write_grid(batch, self.out_key, batch[ck], batch[mk], batch[tk]) + return batch + class TimePos(Stage): """Publish ``time_pos`` -- the within-episode frame index from cu_seqlens. diff --git a/egomimic/pl_utils/pl_model.py b/egomimic/pl_utils/pl_model.py index 39ded4953..e100ea52b 100644 --- a/egomimic/pl_utils/pl_model.py +++ b/egomimic/pl_utils/pl_model.py @@ -10,7 +10,7 @@ from omegaconf import DictConfig, OmegaConf import egomimic.vendored.robomimic_tensor_utils as TensorUtils -from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset +from egomimic.rldb.norm_stats import NormStats class ModelWrapper(LightningModule): @@ -92,6 +92,14 @@ def __init__( cfg, "model.log_per_layer_grad_norms", default=False ) ) + # per_emb_grad_scale (2026-08-02): compensate the cotrain loss-average + # diluting per-emb branches by 1/len(embs). 1.0 = off (default). + self.per_emb_grad_scale = 1.0 + if cfg is not None and OmegaConf.select(cfg, "model") is not None: + self.per_emb_grad_scale = float( + OmegaConf.select(cfg, "model.per_emb_grad_scale", default=1.0) + ) + self._per_emb_param_ids = None # lazily cached id set self.epoch_memory_stats = [] # Store memory stats per epoch self.evaluator = evaluator @@ -106,9 +114,17 @@ def _as_config(cfg): def _instantiate_model(self, config_tree, norm_stats_state): cfg = self._as_config(config_tree) - norm_stats = MultiDataset.from_state(norm_stats_state) + norm_stats = NormStats(norm_stats_state) + model_cfg = OmegaConf.create(OmegaConf.to_container( + cfg.model.robomimic_model, resolve=False)) + inference_cfg = OmegaConf.select( + cfg, "inference_stages", default=None) + if (inference_cfg is not None and + "inference_stages" not in model_cfg): + model_cfg.inference_stages = OmegaConf.create( + OmegaConf.to_container(inference_cfg, resolve=False)) return hydra.utils.instantiate( - cfg.model.robomimic_model, + model_cfg, norm_stats=norm_stats, ) @@ -170,7 +186,7 @@ def training_step(self, batch, batch_idx): info = {} info["losses"] = TensorUtils.detach(losses) for k, v in self.model.log_info(info).items(): - self.log("Train/" + k, v, sync_dist=True, on_step=False, on_epoch=True) + self.log("Train/" + k, v, sync_dist=True, on_step=True, on_epoch=True) return losses["action_loss"] @@ -224,7 +240,36 @@ def _log_per_layer_grad_norms(self): sync_dist=True, ) + def _collect_per_emb_param_ids(self): + ids = set() + n = 0 + for _, m in self.named_modules(): + cn = type(m).__name__ + if cn == "PerEmb": + src = m + elif cn == "MultiEmbodimentCondEncoder" and hasattr(m, "encoders"): + src = m.encoders + else: + continue + for p in src.parameters(): + if id(p) not in ids: + ids.add(id(p)) + n += p.numel() + print(f"[per_emb_grad_scale] x{self.per_emb_grad_scale} on " + f"{len(ids)} per-emb tensors ({n/1e6:.1f}M params)") + return ids + def on_after_backward(self): + # per_emb_grad_scale: restore per-datapoint rate parity for per-emb + # branches (the cotrain loss-average divides their grads by len(embs) + # while they only receive their own embodiment's term). Applied BEFORE + # any clipping so downstream norms see the corrected gradients. + if self.per_emb_grad_scale != 1.0: + if self._per_emb_param_ids is None: + self._per_emb_param_ids = self._collect_per_emb_param_ids() + for p in self.parameters(): + if p.grad is not None and id(p) in self._per_emb_param_ids: + p.grad.mul_(self.per_emb_grad_scale) # Per-layer raw grad norms are gated by their own flag and taken here # (after backward, before any clipping) independent of the global # grad-norm/MAD machinery below. @@ -344,9 +389,23 @@ def configure_optimizers(self) -> Dict[str, Any]: # Method exists but doesn't accept base_lr — skip. groups = None - params_arg = ( - groups if groups is not None else self.trainer.model.parameters() + optimizer_target = str( + OmegaConf.select(cfg, "model.optimizer._target_", default="") ) + if optimizer_target == ( + "egomimic.utils.muon_factory." + "build_single_device_muon_with_aux_adam" + ): + if groups is not None: + raise ValueError( + "The Muon factory performs its own named-parameter routing " + "and cannot be combined with model parameter_groups()." + ) + params_arg = self.trainer.model.named_parameters() + else: + params_arg = ( + groups if groups is not None else self.trainer.model.parameters() + ) # When ``params_arg`` is a ``list[dict]`` of param groups, passing # it through ``hydra.utils.instantiate`` (a kwarg to a _partial_ # target) wraps the dicts as OmegaConf ``DictConfig`` objects, @@ -399,6 +458,75 @@ def on_fit_start(self): flush=True, ) + def on_train_start(self): + # CosineAnnealingLR.state_dict() INCLUDES eta_min, so on any resume + # scheduler.load_state_dict() silently clobbers a config change to + # eta_min (2026-08 dfot_v3 attempt 3: the intended change never took + # effect). Re-assert eta_min from the LIVE config after Lightning has + # restored scheduler state, and log when the restored value differed. + # Recurses into SequentialLR children via ``_schedulers``. + want = None + cfg = self._as_config(getattr(self.hparams, "config_tree", None)) + if cfg is not None: + want = OmegaConf.select(cfg, "model.scheduler.eta_min", default=None) + if want is None: + part = getattr(self.hparams, "scheduler", None) + keywords = getattr(part, "keywords", None) + if keywords: + want = keywords.get("eta_min") + if want is not None: + want = float(want) + for lrs_cfg in self.trainer.lr_scheduler_configs: + stack = [lrs_cfg.scheduler] + while stack: + sch = stack.pop() + stack.extend(list(getattr(sch, "_schedulers", []) or [])) + if hasattr(sch, "eta_min"): + if sch.eta_min != want: + print( + f"[ETA_MIN_REASSERT] {type(sch).__name__}" + f".eta_min {sch.eta_min} -> {want} " + f"(checkpoint state had clobbered the live " + f"config value)", + flush=True, + ) + sch.eta_min = want + + # Same clobber applies to T_max: CosineAnnealingLR.state_dict() + # carries it, so EXTENDING a run (larger model.scheduler.max_steps on + # a resume) is silently inert without this. T_max spans the + # post-warmup window, matching ``warmup_cosine_scheduler``. + want_max, want_warm = None, None + if cfg is not None: + want_max = OmegaConf.select( + cfg, "model.scheduler.max_steps", default=None) + want_warm = OmegaConf.select( + cfg, "model.scheduler.warmup_steps", default=None) + if want_max is None: + part = getattr(self.hparams, "scheduler", None) + keywords = getattr(part, "keywords", None) or {} + want_max = keywords.get("max_steps") + want_warm = keywords.get("warmup_steps", want_warm) + if want_max is not None: + want_warm = int(want_warm) if want_warm is not None else 0 + want_t = max(1, int(want_max) - want_warm) + for lrs_cfg in self.trainer.lr_scheduler_configs: + stack = [lrs_cfg.scheduler] + while stack: + sch = stack.pop() + stack.extend(list(getattr(sch, "_schedulers", []) or [])) + if hasattr(sch, "T_max"): + if sch.T_max != want_t: + print( + f"[T_MAX_REASSERT] {type(sch).__name__}" + f".T_max {sch.T_max} -> {want_t} " + f"(checkpoint state had clobbered the live " + f"config value)", + flush=True, + ) + sch.T_max = want_t + return super().on_train_start() + def on_train_epoch_start(self): for i, param_group in enumerate(self.optimizers().param_groups): self.log( diff --git a/egomimic/rldb/embodiment/eva.py b/egomimic/rldb/embodiment/eva.py index 4c9cdbb87..53377eb65 100644 --- a/egomimic/rldb/embodiment/eva.py +++ b/egomimic/rldb/embodiment/eva.py @@ -5,7 +5,6 @@ import numpy as np from egomimic.rldb.embodiment.embodiment import Embodiment -from egomimic.rldb.embodiment.human import ARIA_INTRINSICS from egomimic.rldb.zarr.action_chunk_transforms import ( ActionChunkCoordinateFrameTransform, BatchQuaternionPoseToYPR, @@ -24,6 +23,13 @@ _matrix_to_xyzwxyz, ) +ARIA_INTRINSICS = np.array( + [[266.50860444, 0.0, 320.0, 0.0], + [0.0, 266.50860444, 240.0, 0.0], + [0.0, 0.0, 1.0, 0.0]], + dtype=np.float64, +) + class Eva(Embodiment): INTRINSICS = ARIA_INTRINSICS @@ -151,6 +157,19 @@ def dinov3_keymap(cls): } +def build_fold_cartesian_wristframe_revert_transform_list( + *, + action_key: str = "actions_cartesian", + state_key: str = "state_ee_pose", +) -> list[Transform]: + """Decode a fold EVA wrist-frame chunk through the shared transforms.""" + from egomimic.rldb.embodiment.fold_span_transforms import ( + build_bimanual_rot6d_wrist_revert_transforms, + ) + + return build_bimanual_rot6d_wrist_revert_transforms(action_key, state_key) + + def _build_eva_bimanual_revert_eef_frame_transform_list( *, action_key: str = "actions_cartesian", diff --git a/egomimic/rldb/embodiment/fold_span_transforms.py b/egomimic/rldb/embodiment/fold_span_transforms.py index 85565f62d..45572540d 100644 --- a/egomimic/rldb/embodiment/fold_span_transforms.py +++ b/egomimic/rldb/embodiment/fold_span_transforms.py @@ -12,21 +12,25 @@ These minimal transforms instead produce PER-FRAME actions + proprio over the whole span by concatenating the raw per-frame arrays (frame conventions are the -raw quaternion / head-frame-relative-nothing forms — good enough to verify the +raw quaternion / head-frame-relative-nothing forms — good enough to verify the ARCH runs + loss descends, which is the smoke's only goal): eva : actions_cartesian (14) = [left.cmd_ee_pose(7), right.cmd_ee_pose(7)] observations.state.ee_pose (14) = [left.obs_ee_pose(7), right.obs_ee_pose(7)] - human: actions_keypoints (138) = [left wrist_ypr(6), left kp(63), - right wrist_ypr(6), right kp(63)] - observations.keypoints (138) = same as the action (teacher-forced) + human: actions_keypoints (132) = [left wrist_xyz(3), left kp(63), + right wrist_xyz(3), right kp(63)] + observations.keypoints (132) = same as the action (teacher-forced) + (was 138 with wrist ypr; dropped -- see HeadFrameWristPos) obs_head_pose (7) = raw head pose """ +import os + import numpy as np from scipy.spatial.transform import Rotation as R from egomimic.rldb.zarr.action_chunk_transforms import ( + ActionChunkCoordinateFrameTransform, ConcatKeys, NumpyToTensor, Transform, @@ -51,6 +55,17 @@ def _wxyz_to_matrix(q): """(N,4) wxyz quats -> (N,3,3) rotation matrices (scipy wants xyzw).""" xyzw = np.concatenate([q[:, 1:4], q[:, 0:1]], axis=-1) + # Aria arrays are zero-padded past attrs["total_frames"], and an action + # window near the end of an episode gets clamped short then padded back to + # the horizon -- either way some rows arrive as [0,0,0,0]. A zero-norm + # quaternion has NO rotation: scipy raises, and any silent normalize turns + # numerical dust into a garbage frame. Substitute identity there so the pad + # is a no-op rotation instead of corrupting the whole chunk. + n = np.linalg.norm(xyzw, axis=-1) + bad = n < 1e-8 + if bad.any(): + xyzw = np.asarray(xyzw, dtype=np.float64).copy() + xyzw[bad] = np.array([0.0, 0.0, 0.0, 1.0]) # identity, xyzw order return R.from_quat(xyzw).as_matrix() @@ -117,6 +132,38 @@ def transform(self, batch: dict) -> dict: return batch +class HeadFrameWristPos(Transform): + """Head-frame wrist POSITION only -- xyz, no orientation (option B). + + ``wrist`` (T,7) [or (7,)] xyz+quat(wxyz) -> (T,3) [or (3,)]: + t_wh = R_head^T @ (t_wrist - t_head) + + Deliberately drops the ZYX-euler orientation that HeadFrameWristYPR emits: + those 3 dims/wrist are discontinuous (wrap at +-pi, gimbal lock at pitch + +-pi/2) and a diffusion model cannot represent the branch cut, which showed + up as ~80%% of the human action error concentrated in 12 of 138 dims. + Orientation remains recoverable from the 21 head-frame keypoints. + """ + + def __init__(self, head_key: str, wrist_key: str, out_key: str): + self.head_key = head_key + self.wrist_key = wrist_key + self.out_key = out_key + + def transform(self, batch: dict) -> dict: + head = np.asarray(batch[self.head_key], dtype=np.float64) + wr = np.asarray(batch[self.wrist_key], dtype=np.float64) + single = head.ndim == 1 + if single: + head, wr = head[None, :], wr[None, :] + t_head = head[:, :3] + R_head = _wxyz_to_matrix(head[:, 3:7]) + R_headT = np.transpose(R_head, (0, 2, 1)) + t_wh = np.einsum("tij,tj->ti", R_headT, wr[:, :3] - t_head) # (T,3) + batch[self.out_key] = t_wh[0] if single else t_wh + return batch + + # --------------------------------------------------------------------------- # # eva_bimanual # --------------------------------------------------------------------------- # @@ -126,7 +173,9 @@ def _drop_camera_keys(km): return { k: v for k, v in km.items() - if v.get("key_type") not in ("camera_keys", "annotation_keys") + if v.get("key_type") not in ( + "camera_keys", "annotation_keys", "metadata_keys" + ) } @@ -145,7 +194,38 @@ def eva_span_keymap(norm_mode: bool = False, annotation_key=None): return _drop_camera_keys(km) if norm_mode else km -def eva_span_transforms(): +class CanonicalizeQuatSign(Transform): + """Force the quaternion hemisphere to w >= 0 on an (..., 7) xyz+quat(wxyz). + + q and -q are the SAME rotation but OPPOSITE regression targets. eva's raw + cmd_ee_pose/obs_ee_pose quats are bimodal -- 53% of frames sit at w<0 with + 443 mid-episode sign flips per 20 episodes -- so a denoiser fitting the + conditional mean blends the two hemispheres and carries an irreducible + error. Human poses were already canonicalised inside HeadFramePose; this + is the eva-side mirror. Stateless (per frame), so it also holds for the + single-frame reads norm_stats issues. + """ + + def __init__(self, keys, quat_slice=(3, 7)): + self.keys = list(keys) + self.qs = tuple(quat_slice) + + def transform(self, batch: dict) -> dict: + a, b = self.qs + for k in self.keys: + if k not in batch: + continue + v = np.array(batch[k], dtype=np.float64, copy=True) + single = v.ndim == 1 + if single: + v = v[None, :] + w = v[:, a:a + 1] + v[:, a:b] = np.where(w < 0.0, -v[:, a:b], v[:, a:b]) + batch[k] = (v[0] if single else v).astype(np.float32) + return batch + + +def eva_span_transforms_quat14(): return [ ConcatKeys( key_list=["left.cmd_ee_pose", "right.cmd_ee_pose"], @@ -176,30 +256,862 @@ def human_span_keymap(norm_mode: bool = False, annotation_key=None): return _drop_camera_keys(km) if norm_mode else km +_WRIST_MODE = os.environ.get("RH_WRIST_MODE", "pos").lower() +#: "ypr" -> 138 (legacy, wrist xyz+ypr) | "pos" -> 132 (wrist xyz) | +#: "none" -> 126 (KEYPOINTS ONLY, no wrist pose) +_WRIST = (HeadFrameWristYPR if _WRIST_MODE == "ypr" + else (None if _WRIST_MODE == "none" else HeadFrameWristPos)) + + def human_span_transforms(): # SPAN-SAFE HEAD-FRAME (fold2 FIX 2): keypoints + wrist expressed in each # frame's head frame (relative to obs_head_pose), matching the canonical # Aria "keypoints_headframe_ypr" representation but computed PER-FRAME so it # works on variable-length span reads AND the single-frame windowed probe. - # action_keypoints (138) = [Lwrist_ypr(6), Lkp_hf(63), Rwrist_ypr(6), Rkp_hf(63)] + # action_keypoints (132) = [Lwrist_xyz(3), Lkp_hf(63), Rwrist_xyz(3), Rkp_hf(63)] # (order matches _build_aria_keypoints_bimanual_transform_list). The action # is the teacher-forced per-frame head-frame state (obs == action here). + # RH_WRIST_MODE: ypr -> 138 (legacy wrist xyz+ypr) | pos -> 132 (wrist xyz) + # | none -> 126 (KEYPOINTS ONLY). 138 is needed to load pre-2026-07-30 + # checkpoints, whose norm stats and per-emb codecs are keyed to it. + _wrist_ops = ([] if _WRIST is None else [ + _WRIST("obs_head_pose", "left.obs_wrist_pose", "L_wrist_hf"), + _WRIST("obs_head_pose", "right.obs_wrist_pose", "R_wrist_hf"), + ]) + _keys = (["L_kp_hf", "R_kp_hf"] if _WRIST is None + else ["L_wrist_hf", "L_kp_hf", "R_wrist_hf", "R_kp_hf"]) return [ - HeadFrameWristYPR("obs_head_pose", "left.obs_wrist_pose", "L_wrist_hf"), - HeadFrameWristYPR("obs_head_pose", "right.obs_wrist_pose", "R_wrist_hf"), + *_wrist_ops, HeadFrameKeypoints("obs_head_pose", "left.obs_keypoints", "L_kp_hf"), HeadFrameKeypoints("obs_head_pose", "right.obs_keypoints", "R_kp_hf"), - # action target (138): head-frame wrist-ypr + head-frame keypoints per hand. + # action target (132): head-frame wrist-xyz + head-frame keypoints per hand. ConcatKeys( - key_list=["L_wrist_hf", "L_kp_hf", "R_wrist_hf", "R_kp_hf"], + key_list=_keys, new_key_name="actions_keypoints", delete_old_keys=False, ), - # proprio obs (138): same head-frame layout; consumes the components. + # proprio obs (132): same head-frame layout; consumes the components. ConcatKeys( - key_list=["L_wrist_hf", "L_kp_hf", "R_wrist_hf", "R_kp_hf"], + key_list=_keys, new_key_name="state_keypoints", delete_old_keys=True, ), NumpyToTensor(keys=["actions_keypoints", "state_keypoints", "obs_head_pose"]), ] + + +# --------------------------------------------------------------------------- # +# human_bimanual -- CARTESIAN (action space MATCHED to eva) +# +# "Same action space" cotrain (user, 2026-08-12): both embodiments emit a 14-D +# cartesian end-effector action [left(7), right(7)], each 7 = xyz(3) + quat +# wxyz(4). Because the two spaces are IDENTICAL in dim AND semantics, a +# STANDARD Diffusion Policy works unchanged -- one scalar action_dim, one +# shared head, no per-embodiment codec, no per-emb obs encoder. +# +# eva : actions_cartesian(14) = [left.cmd_ee_pose(7), right.cmd_ee_pose(7)] +# state_ee_pose(14) = [left.obs_ee_pose(7), right.obs_ee_pose(7)] +# human: actions_cartesian(14) = [L_ee_hf(7), R_ee_hf(7)] +# state_ee_pose(14) = the same tensor (teacher-forced: the human +# has no separately commanded pose, exactly as +# the keypoints feed already uses obs as its +# own target) +# +# FRAME -- a representation choice, stated rather than buried: the human's +# left/right.obs_ee_pose live in the Aria WORLD frame, whose origin+yaw are +# arbitrary PER RECORDING, so raw world coordinates are not a learnable action +# space (the same fold would sit at a different absolute xyz in every episode). +# They are converted PER FRAME into that frame's HEAD frame -- the convention +# every other human transform in this file already uses, and the closest +# analogue to eva's fixed robot-base frame. +# --------------------------------------------------------------------------- # +class HeadFramePose(Transform): + """Per-frame head-frame conversion of a full 7-D pose. + + ``pose`` (T, 7) [or (7,) single frame] as xyz + quat(wxyz) -> same shape, + expressed in that frame's head frame:: + + t_hf[t] = R_head[t]^T @ (t_pose[t] - t_head[t]) + R_hf[t] = R_head[t]^T @ R_pose[t] + + The output quaternion is sign-canonicalised to ``w >= 0``. q and -q are the + same rotation but NOT the same regression target: without this the sign + flips arbitrarily between frames and the normalizer sees a bimodal + distribution straddling zero. Canonicalising is stateless, so it also holds + for the single-frame ``(7,)`` reads that ``norm_stats.populate_from_datasets`` + issues (a previous-frame continuity fix could not). + """ + + def __init__(self, head_key: str, pose_key: str, out_key: str): + self.head_key = head_key + self.pose_key = pose_key + self.out_key = out_key + + def transform(self, batch: dict) -> dict: + head = np.asarray(batch[self.head_key], dtype=np.float64) + pose = np.asarray(batch[self.pose_key], dtype=np.float64) + single = pose.ndim == 1 # decided by POSE, not head + if single: + pose = pose[None, :] + if head.ndim == 1: + # one observation frame's head pose applies to every action in that + # frame's chunk -> broadcast, do not unsqueeze in lockstep. Keying + # "single" off head.ndim gave head (1,7) vs pose (1,H,7). + head = np.broadcast_to(head[None, :], (pose.shape[0], head.shape[-1])) + R_head = _wxyz_to_matrix(head[:, 3:7]) # (T,3,3) + R_pose = _wxyz_to_matrix(pose[:, 3:7]) + t_hf = np.einsum("tij,tj->ti", R_head.transpose(0, 2, 1), + pose[:, 0:3] - head[:, 0:3]) # (T,3) + R_hf = np.einsum("tij,tjk->tik", R_head.transpose(0, 2, 1), R_pose) + xyzw = R.from_matrix(R_hf).as_quat() # scipy: xyzw + q = np.concatenate([xyzw[:, 3:4], xyzw[:, 0:3]], axis=-1) # -> wxyz + q = np.where(q[:, 0:1] < 0.0, -q, q) # canonical w>=0 + out = np.concatenate([t_hf, q], axis=-1).astype(np.float32) # (T,7) + batch[self.out_key] = out[0] if single else out + return batch + + +def human_span_cart_keymap(norm_mode: bool = False, annotation_key=None): + """Cartesian counterpart of :func:`human_span_keymap`: no keypoints, just + the per-hand end-effector pose + the head pose that defines the frame.""" + km = { + "front_img_1": {"key_type": "camera_keys", "zarr_key": "images.front_1"}, + "left.obs_ee_pose": {"key_type": "proprio_keys", "zarr_key": "left.obs_ee_pose"}, + "right.obs_ee_pose": {"key_type": "proprio_keys", "zarr_key": "right.obs_ee_pose"}, + "obs_head_pose": {"key_type": "proprio_keys", "zarr_key": "obs_head_pose"}, + } + return _drop_camera_keys(km) if norm_mode else km + + +def human_span_cart_transforms_quat14(): + """14-D head-frame cartesian action+state, key names IDENTICAL to eva's so + the model config needs no per-embodiment branching at all. + + NOTE: unaffected by ``RH_WRIST_MODE`` -- that env var only sizes the + keypoints feed (126/132/138). This feed is always 14. + """ + return [ + HeadFramePose("obs_head_pose", "left.obs_ee_pose", "L_ee_hf"), + HeadFramePose("obs_head_pose", "right.obs_ee_pose", "R_ee_hf"), + ConcatKeys( + key_list=["L_ee_hf", "R_ee_hf"], + new_key_name="actions_cartesian", + delete_old_keys=False, + ), + ConcatKeys( + key_list=["L_ee_hf", "R_ee_hf"], + new_key_name="state_ee_pose", + delete_old_keys=True, + ), + NumpyToTensor(keys=["actions_cartesian", "state_ee_pose", "obs_head_pose"]), + ] + +# --------------------------------------------------------------------------- # +# 6D ROTATION ACTION SPACE (user, 2026-08-15) +# +# Quaternions double-cover SO(3): q and -q are the same rotation but opposite +# regression targets. Measured on eva cmd_ee_pose: 53% of frames at w<0 with +# 443 mid-episode sign flips per 20 episodes -- so a mean-seeking denoiser +# blends hemispheres and carries an irreducible error floor. Canonicalising to +# w>=0 only MOVES the seam; the 6D representation (first two columns of R, +# Zhou et al. 2019 "On the Continuity of Rotation Representations") removes it, +# being continuous on SO(3). +# +# per hand xyz(3) + rot6d(6) = 9 bimanual = 18 (quat version: 14) +# +# Both embodiments use it, so the matched action space is preserved. Recover R +# with Gram-Schmidt on the two columns; the decoder side is +# rot6d_to_matrix below. +# --------------------------------------------------------------------------- # +def _rot6d_from_matrix(R_): + """(N,3,3) -> (N,6): first TWO COLUMNS, flattened. Column order matters -- + the Gram-Schmidt inverse must read them back in the same order.""" + return np.concatenate([R_[:, :, 0], R_[:, :, 1]], axis=-1) + + +def rot6d_to_matrix(d6): + """(N,6) -> (N,3,3) via Gram-Schmidt. Inverse of _rot6d_from_matrix.""" + a1, a2 = d6[:, 0:3], d6[:, 3:6] + b1 = a1 / (np.linalg.norm(a1, axis=-1, keepdims=True) + 1e-8) + a2p = a2 - (b1 * a2).sum(-1, keepdims=True) * b1 + b2 = a2p / (np.linalg.norm(a2p, axis=-1, keepdims=True) + 1e-8) + b3 = np.cross(b1, b2) + return np.stack([b1, b2, b3], axis=-1) + + +class PoseToRot6D(Transform): + """xyz + quat(wxyz) (..., 7) -> xyz + rot6d (..., 9). Stateless per frame, + so it also holds for the single-frame reads norm_stats issues.""" + + def __init__(self, in_key: str, out_key: str = None): + self.in_key = in_key + self.out_key = out_key or in_key + + def transform(self, batch: dict) -> dict: + v = np.asarray(batch[self.in_key], dtype=np.float64) + single = v.ndim == 1 + if single: + v = v[None, :] + R_ = _wxyz_to_matrix(v[:, 3:7]) + out = np.concatenate([v[:, 0:3], _rot6d_from_matrix(R_)], axis=-1) + out = out.astype(np.float32) + batch[self.out_key] = out[0] if single else out + return batch + + +def eva_span_transforms(): + """eva: actions_cartesian(18) / state_ee_pose(18) in xyz+rot6d.""" + return [ + PoseToRot6D("left.cmd_ee_pose"), PoseToRot6D("right.cmd_ee_pose"), + PoseToRot6D("left.obs_ee_pose"), PoseToRot6D("right.obs_ee_pose"), + ConcatKeys(key_list=["left.cmd_ee_pose", "right.cmd_ee_pose"], + new_key_name="actions_cartesian", delete_old_keys=True), + ConcatKeys(key_list=["left.obs_ee_pose", "right.obs_ee_pose"], + new_key_name="state_ee_pose", delete_old_keys=True), + NumpyToTensor(keys=["actions_cartesian", "state_ee_pose"]), + ] + + +def eva_dfot_keymap(norm_mode: bool = False, annotation_key=None): + """Robot-only DFoT feed: front RGB plus the commanded 20-D action.""" + km = { + "front_img_1": { + "key_type": "camera_keys", "zarr_key": "images.front_1"}, + "left.cmd_ee_pose": { + "key_type": "action_keys", "zarr_key": "left.cmd_ee_pose"}, + "left.cmd_gripper": { + "key_type": "action_keys", "zarr_key": "left.cmd_gripper"}, + "right.cmd_ee_pose": { + "key_type": "action_keys", "zarr_key": "right.cmd_ee_pose"}, + "right.cmd_gripper": { + "key_type": "action_keys", "zarr_key": "right.cmd_gripper"}, + } + return _drop_camera_keys(km) if norm_mode else km + + +def eva_dfot_transforms(): + """Per-frame ``[L xyz+rot6d+grip, R xyz+rot6d+grip]`` action.""" + return [ + PoseToRot6D("left.cmd_ee_pose"), + PoseToRot6D("right.cmd_ee_pose"), + ConcatKeys( + key_list=["left.cmd_ee_pose", "left.cmd_gripper", + "right.cmd_ee_pose", "right.cmd_gripper"], + new_key_name="actions_cartesian", delete_old_keys=True, + ), + NumpyToTensor(keys=["actions_cartesian"]), + ] + + +def human_span_cart_transforms(): + """human: same 18-D xyz+rot6d, in the per-frame HEAD frame (HeadFramePose + first, then 6D) -- so eva and human remain a MATCHED action space.""" + return [ + HeadFramePose("obs_head_pose", "left.obs_ee_pose", "L_ee_hf"), + HeadFramePose("obs_head_pose", "right.obs_ee_pose", "R_ee_hf"), + PoseToRot6D("L_ee_hf"), PoseToRot6D("R_ee_hf"), + ConcatKeys(key_list=["L_ee_hf", "R_ee_hf"], + new_key_name="actions_cartesian", delete_old_keys=False), + ConcatKeys(key_list=["L_ee_hf", "R_ee_hf"], + new_key_name="state_ee_pose", delete_old_keys=True), + NumpyToTensor(keys=["actions_cartesian", "state_ee_pose", + "obs_head_pose"]), + ] + + +# Back-compat aliases: the 6D builders keep their explicit names too, so a +# config can ask for either representation by name rather than by default. +eva_span_rot6d_transforms = eva_span_transforms +human_span_cart_rot6d_transforms = human_span_cart_transforms + +# --------------------------------------------------------------------------- # +# NORMAL-DATALOADER (MultiDataset) keymaps + transforms for fold. +# +# This is the STANDARD reader (one sample per FRAME, annotation_collate), which +# is what stock Diffusion Policy trains on. It decodes only the frames it reads +# -- no img_decode_stride, no zero-filled placeholders -- so the black-image +# class of bug the packed h264 path had cannot occur here. +# +# Obs keys carry horizon=N_OBS_STEPS so DP's observation history is REAL data +# rather than the current frame duplicated. Action keys are fetched one frame +# longer (ACTION_HORIZON + N_OBS_STEPS - 1) and then sliced, so the chunk starts +# at the LAST obs frame: +# +# obs frames t ........ t+N-1 (N_OBS_STEPS of them) +# action chunk t+N-1 ........ t+N-2+ACTION_HORIZON +# +# i.e. the policy sees N frames up to and including the current one and predicts +# ACTION_HORIZON actions starting at the current one. Fetching without the +N-1 +# and skipping the slice would silently shift the chunk one frame into the past. +# +# Both embodiments emit the SAME 18-D action (xyz + rot6d per hand), so the +# matched action space carries over from the packed configs unchanged. +# +# Tail frames are repeat-last padded to exactly the requested horizon by +# ZarrDataset._pad_sequences, which is what keeps the chunk fixed-size. +# --------------------------------------------------------------------------- # +ACTION_HORIZON = 100 +N_OBS_STEPS = 2 +_ACT_FETCH = ACTION_HORIZON + N_OBS_STEPS - 1 + + +class ZeroOut(Transform): + """DIAGNOSTIC ONLY: overwrite a key with zeros, preserving shape/dtype. + Used to hand the model a trivially-learnable target: with x0==0 the DDPM + forward gives x_t = sqrt(1-abar)*eps, so eps is EXACTLY recoverable from + x_t alone. A model that cannot fit this has a wiring fault, not a data one.""" + + def __init__(self, key): + self.key = str(key) + + def transform(self, batch): + v = batch[self.key] + batch[self.key] = np.zeros_like(v) + return batch + + +class SliceFrames(Transform): + """Keep ``[start:stop]`` along the leading (time) axis of ``key``. + + Used to drop the lead-in frames of an over-fetched action chunk so the + chunk starts at the current obs frame. + """ + + def __init__(self, key, start=0, stop=None, new_key=None): + self.key = key + self.start = start + self.stop = stop + self.new_key = new_key or key + + def transform(self, batch): + v = np.asarray(batch[self.key]) + batch[self.new_key] = v[self.start:self.stop] + return batch + + +class SelectFrame(Transform): + """Copy one frame from a time-major array into a standalone pose key.""" + + def __init__(self, key, index=-1, new_key=None): + self.key = key + self.index = int(index) + self.new_key = new_key or key + + def transform(self, batch): + batch[self.new_key] = np.asarray(batch[self.key])[self.index] + return batch + + +class ReshapePoints(Transform): + """Switch between flattened ``(..., N*3)`` and ``(..., N, 3)`` points.""" + + def __init__(self, key, n_points=21, flatten=False, new_key=None): + self.key = key + self.n_points = int(n_points) + self.flatten = bool(flatten) + self.new_key = new_key or key + + def transform(self, batch): + value = np.asarray(batch[self.key]) + if self.flatten: + value = value.reshape(*value.shape[:-2], self.n_points * 3) + else: + value = value.reshape(*value.shape[:-1], self.n_points, 3) + batch[self.new_key] = value + return batch + + +class PoseXYZ(Transform): + """Keep only xyz from an ``(..., 7)`` pose, preserving leading axes.""" + + def __init__(self, key, new_key=None): + self.key = key + self.new_key = new_key or key + + def transform(self, batch): + batch[self.new_key] = np.asarray(batch[self.key])[..., :3] + return batch + + +class ZerosLike(Transform): + """(T, W) zeros with the leading axis of ``ref_key``. + + Used to pad an embodiment that lacks a channel the shared action layout + reserves -- Aria has no gripper. The pad is NEVER scored: MaskedActionLoss + excludes those dims for that embodiment. + """ + + def __init__(self, ref_key, new_key, width=1): + self.ref_key, self.new_key, self.width = ref_key, new_key, int(width) + + def transform(self, batch): + ref = np.asarray(batch[self.ref_key]) + batch[self.new_key] = np.zeros((*ref.shape[:-1], self.width), + dtype=np.float32) + return batch + + +class DropKeys(Transform): + """Remove scratch keys so they never reach the collate / norm-stat layer.""" + + def __init__(self, keys): + self.keys = list(keys) + + def transform(self, batch): + for k in self.keys: + batch.pop(k, None) + return batch + + +def eva_normal_keymap(norm_mode: bool = False, annotation_key=None): + km = { + "front_img_1": {"key_type": "camera_keys", "zarr_key": "images.front_1", + "horizon": N_OBS_STEPS}, + "front_intrinsics": {"key_type": "metadata_keys", + "zarr_key": "intrinsics.front_1"}, + "left_camera_extrinsics": {"key_type": "metadata_keys", + "zarr_key": "extrinsics.left"}, + "right_camera_extrinsics": {"key_type": "metadata_keys", + "zarr_key": "extrinsics.right"}, + # eva HAS a commanded stream -> use it as the action target + "left.cmd_ee_pose": {"key_type": "action_keys", "zarr_key": "left.cmd_ee_pose", + "horizon": _ACT_FETCH}, + "right.cmd_ee_pose": {"key_type": "action_keys", "zarr_key": "right.cmd_ee_pose", + "horizon": _ACT_FETCH}, + "left.obs_ee_pose": {"key_type": "proprio_keys", "zarr_key": "left.obs_ee_pose", + "horizon": N_OBS_STEPS}, + "right.obs_ee_pose": {"key_type": "proprio_keys", "zarr_key": "right.obs_ee_pose", + "horizon": N_OBS_STEPS}, + # GRIPPER (2026-08-16): the fold task is grasp/release, so a policy + # without this predicts arm motion it can never act on. + "left.cmd_gripper": {"key_type": "action_keys", "zarr_key": "left.cmd_gripper", + "horizon": _ACT_FETCH}, + "right.cmd_gripper": {"key_type": "action_keys", "zarr_key": "right.cmd_gripper", + "horizon": _ACT_FETCH}, + "left.obs_gripper": {"key_type": "proprio_keys", "zarr_key": "left.obs_gripper", + "horizon": N_OBS_STEPS}, + "right.obs_gripper": {"key_type": "proprio_keys", "zarr_key": "right.obs_gripper", + "horizon": N_OBS_STEPS}, + # WRIST CAMERAS -- eva only; Aria has no wrist views, so these belong in + # the per-embodiment (specific) obs branch, never the agnostic one. + "left_wrist_img": {"key_type": "camera_keys", "zarr_key": "images.left_wrist", + "horizon": N_OBS_STEPS}, + "right_wrist_img": {"key_type": "camera_keys", "zarr_key": "images.right_wrist", + "horizon": N_OBS_STEPS}, + } + return _drop_camera_keys(km) if norm_mode else km + + +def human_normal_keymap(norm_mode: bool = False, annotation_key=None): + km = { + "front_img_1": {"key_type": "camera_keys", "zarr_key": "images.front_1", + "horizon": N_OBS_STEPS}, + "front_intrinsics": {"key_type": "metadata_keys", + "zarr_key": "intrinsics.front_1"}, + # human has NO commanded stream; obs IS the target (teacher-forced), + # exactly as the packed human feed already does. + "left.act_ee_pose": {"key_type": "action_keys", "zarr_key": "left.obs_ee_pose", + "horizon": _ACT_FETCH}, + "right.act_ee_pose": {"key_type": "action_keys", "zarr_key": "right.obs_ee_pose", + "horizon": _ACT_FETCH}, + "left.obs_ee_pose": {"key_type": "proprio_keys", "zarr_key": "left.obs_ee_pose", + "horizon": N_OBS_STEPS}, + "right.obs_ee_pose": {"key_type": "proprio_keys", "zarr_key": "right.obs_ee_pose", + "horizon": N_OBS_STEPS}, + "obs_head_pose": {"key_type": "proprio_keys", "zarr_key": "obs_head_pose", + "horizon": N_OBS_STEPS}, + } + return _drop_camera_keys(km) if norm_mode else km + + +def eva_normal_transforms(): + """Wrist-relative actions + base-frame proprio, both in rot6d layout. + + Each arm's future command is expressed relative to that arm's CURRENT + observed EEF pose. This is the same action-frame contract used by the + human transforms below; only the observation frames remain embodiment + native (robot base vs human head camera). + """ + return [ + SelectFrame("left.obs_ee_pose", -1, "left.current_ee_pose"), + SelectFrame("right.obs_ee_pose", -1, "right.current_ee_pose"), + SliceFrames("left.cmd_ee_pose", start=N_OBS_STEPS - 1), + SliceFrames("right.cmd_ee_pose", start=N_OBS_STEPS - 1), + SliceFrames("left.cmd_gripper", start=N_OBS_STEPS - 1), + SliceFrames("right.cmd_gripper", start=N_OBS_STEPS - 1), + ActionChunkCoordinateFrameTransform( + "left.current_ee_pose", "left.cmd_ee_pose", "left.cmd_ee_wrist", + mode="xyzwxyz", + ), + ActionChunkCoordinateFrameTransform( + "right.current_ee_pose", "right.cmd_ee_pose", "right.cmd_ee_wrist", + mode="xyzwxyz", + ), + PoseToRot6D("left.cmd_ee_wrist"), PoseToRot6D("right.cmd_ee_wrist"), + PoseToRot6D("left.obs_ee_pose"), PoseToRot6D("right.obs_ee_pose"), + # 20-D: [L xyz3 rot6d6 grip1, R xyz3 rot6d6 grip1] -> gripper at 9 and 19 + ConcatKeys(key_list=["left.cmd_ee_wrist", "left.cmd_gripper", + "right.cmd_ee_wrist", "right.cmd_gripper"], + new_key_name="actions_cartesian", delete_old_keys=True), + ConcatKeys(key_list=["left.obs_ee_pose", "left.obs_gripper", + "right.obs_ee_pose", "right.obs_gripper"], + new_key_name="state_ee_pose", delete_old_keys=True), + DropKeys(["left.cmd_ee_pose", "right.cmd_ee_pose", + "left.current_ee_pose", "right.current_ee_pose"]), + NumpyToTensor(keys=["actions_cartesian", "state_ee_pose"]), + ] + + +def human_normal_transforms(): + """Wrist-relative actions + head-frame proprio in the shared 20-D layout.""" + return [ + SelectFrame("left.obs_ee_pose", -1, "left.current_ee_pose"), + SelectFrame("right.obs_ee_pose", -1, "right.current_ee_pose"), + SliceFrames("left.act_ee_pose", start=N_OBS_STEPS - 1), + SliceFrames("right.act_ee_pose", start=N_OBS_STEPS - 1), + ActionChunkCoordinateFrameTransform( + "left.current_ee_pose", "left.act_ee_pose", "L_act_wrist", + mode="xyzwxyz", + ), + ActionChunkCoordinateFrameTransform( + "right.current_ee_pose", "right.act_ee_pose", "R_act_wrist", + mode="xyzwxyz", + ), + HeadFramePose("obs_head_pose", "left.obs_ee_pose", "L_obs_hf"), + HeadFramePose("obs_head_pose", "right.obs_ee_pose", "R_obs_hf"), + PoseToRot6D("L_act_wrist"), PoseToRot6D("R_act_wrist"), + PoseToRot6D("L_obs_hf"), PoseToRot6D("R_obs_hf"), + # Match eva's 20-D layout. Aria has no gripper, so those two slots are + # zero pads -- MaskedActionLoss excludes dims 9 and 19 for this + # embodiment, so they are never scored and never learned. + ZerosLike("L_act_wrist", "L_grip_pad", 1), ZerosLike("R_act_wrist", "R_grip_pad", 1), + ConcatKeys(key_list=["L_act_wrist", "L_grip_pad", "R_act_wrist", "R_grip_pad"], + new_key_name="actions_cartesian", delete_old_keys=True), + ZerosLike("L_obs_hf", "L_grip_pad_o", 1), ZerosLike("R_obs_hf", "R_grip_pad_o", 1), + ConcatKeys(key_list=["L_obs_hf", "L_grip_pad_o", "R_obs_hf", "R_grip_pad_o"], + new_key_name="state_ee_pose", delete_old_keys=True), + # HeadFramePose writes NEW keys, so the raw streams survive the + # concats (unlike eva, where ConcatKeys consumes them directly). + # Drop them: nothing downstream reads them, but they would be + # collated, normalized and shipped to the GPU every step. + DropKeys(["left.act_ee_pose", "right.act_ee_pose", + "left.current_ee_pose", "right.current_ee_pose", + "left.obs_ee_pose", "right.obs_ee_pose"]), + NumpyToTensor(keys=["actions_cartesian", "state_ee_pose", "obs_head_pose"]), + ] + + +# --------------------------------------------------------------------------- # +# NORMAL-DATALOADER keypoint variant (human hetero action space). +# +# eva is UNCHANGED -- it reuses eva_normal_{keymap,transforms} (cartesian 18). +# A robot has no body keypoints, so only the human side switches: +# actions_keypoints (ACTION_HORIZON, 132) state_keypoints (N_OBS, 132) +# +# The human action is teacher-forced (obs IS the action), so each keypoint / +# wrist array is fetched TWICE from the same zarr key at two horizons: once at +# _ACT_FETCH for the action chunk, once at N_OBS_STEPS for the observation. Each +# copy is paired with a head pose of matching length, because HeadFrameKeypoints +# / HeadFrameWristPos convert per-frame and need equal leading lengths. +# --------------------------------------------------------------------------- # +def human_normal_transforms_synth(): + """human_normal_transforms + the action target ZEROED. Diagnostic: separates + 'the model cannot fit this target' from 'the model cannot fit anything here'.""" + tl = list(human_normal_transforms()) + out = [] + for t in tl: + out.append(t) + if isinstance(t, ConcatKeys) and getattr(t, "new_key_name", None) == "actions_cartesian": + out.append(ZeroOut("actions_cartesian")) + return out + + +def human_normal_keymap_kp(norm_mode: bool = False, annotation_key=None): + P = "proprio_keys" + km = { + "front_img_1": {"key_type": "camera_keys", "zarr_key": "images.front_1", + "horizon": N_OBS_STEPS}, + "front_intrinsics": {"key_type": "metadata_keys", + "zarr_key": "intrinsics.front_1"}, + # ---- action copies (teacher-forced), fetched at the action horizon ---- + "left.act_keypoints": {"key_type": P, "zarr_key": "left.obs_keypoints", + "horizon": _ACT_FETCH}, + "right.act_keypoints": {"key_type": P, "zarr_key": "right.obs_keypoints", + "horizon": _ACT_FETCH}, + "left.act_wrist_pose": {"key_type": P, "zarr_key": "left.obs_wrist_pose", + "horizon": _ACT_FETCH}, + "right.act_wrist_pose": {"key_type": P, "zarr_key": "right.obs_wrist_pose", + "horizon": _ACT_FETCH}, + # ---- observation copies, fetched at the obs-history horizon ---- + "left.obs_keypoints": {"key_type": P, "zarr_key": "left.obs_keypoints", + "horizon": N_OBS_STEPS}, + "right.obs_keypoints": {"key_type": P, "zarr_key": "right.obs_keypoints", + "horizon": N_OBS_STEPS}, + "left.obs_wrist_pose": {"key_type": P, "zarr_key": "left.obs_wrist_pose", + "horizon": N_OBS_STEPS}, + "right.obs_wrist_pose": {"key_type": P, "zarr_key": "right.obs_wrist_pose", + "horizon": N_OBS_STEPS}, + "obs_head_pose": {"key_type": P, "zarr_key": "obs_head_pose", + "horizon": N_OBS_STEPS}, + } + return _drop_camera_keys(km) if norm_mode else km + + +def human_normal_transforms_kp(): + """132 = [Lwrist_xyz(3), Lkp_hf(63), Rwrist_xyz(3), Rkp_hf(63)] per frame, + same layout and same RH_WRIST_MODE sizing as the packed human_span_transforms + -- so kp-vs-cart is not confounded by a different human representation.""" + _wrist_act = ([] if _WRIST is None else [ + _WRIST("act_head_pose", "left.act_wrist_pose", "L_wrist_act"), + _WRIST("act_head_pose", "right.act_wrist_pose", "R_wrist_act"), + ]) + _wrist_obs = ([] if _WRIST is None else [ + _WRIST("obs_head_pose", "left.obs_wrist_pose", "L_wrist_obs"), + _WRIST("obs_head_pose", "right.obs_wrist_pose", "R_wrist_obs"), + ]) + # The action contract is fixed at 132-D regardless of the legacy + # RH_WRIST_MODE switch: per arm = relative wrist xyz3 + relative kp63. + act_keys = ["L_wrist_act", "L_kp_act", "R_wrist_act", "R_kp_act"] + obs_keys = (["L_kp_obs", "R_kp_obs"] if _WRIST is None + else ["L_wrist_obs", "L_kp_obs", "R_wrist_obs", "R_kp_obs"]) + return [ + SelectFrame("left.obs_wrist_pose", -1, "left.current_wrist_pose"), + SelectFrame("right.obs_wrist_pose", -1, "right.current_wrist_pose"), + SliceFrames("left.act_wrist_pose", start=N_OBS_STEPS - 1), + SliceFrames("right.act_wrist_pose", start=N_OBS_STEPS - 1), + SliceFrames("left.act_keypoints", start=N_OBS_STEPS - 1), + SliceFrames("right.act_keypoints", start=N_OBS_STEPS - 1), + ReshapePoints("left.act_keypoints"), + ReshapePoints("right.act_keypoints"), + ActionChunkCoordinateFrameTransform( + "left.current_wrist_pose", "left.act_wrist_pose", "L_wrist_act", + mode="xyzwxyz", + ), + ActionChunkCoordinateFrameTransform( + "right.current_wrist_pose", "right.act_wrist_pose", "R_wrist_act", + mode="xyzwxyz", + ), + ActionChunkCoordinateFrameTransform( + "left.current_wrist_pose", "left.act_keypoints", "L_kp_act", + mode="xyz", + ), + ActionChunkCoordinateFrameTransform( + "right.current_wrist_pose", "right.act_keypoints", "R_kp_act", + mode="xyz", + ), + PoseXYZ("L_wrist_act"), PoseXYZ("R_wrist_act"), + ReshapePoints("L_kp_act", flatten=True), + ReshapePoints("R_kp_act", flatten=True), + ConcatKeys(key_list=act_keys, new_key_name="actions_keypoints", + delete_old_keys=True), + *_wrist_obs, + HeadFramePose("obs_head_pose", "left.obs_wrist_pose", "L_wrist_obs_pose_hf"), + HeadFramePose("obs_head_pose", "right.obs_wrist_pose", "R_wrist_obs_pose_hf"), + SelectFrame("L_wrist_obs_pose_hf", -1, "L_current_wrist_hf"), + SelectFrame("R_wrist_obs_pose_hf", -1, "R_current_wrist_hf"), + ConcatKeys( + key_list=["L_current_wrist_hf", "R_current_wrist_hf"], + new_key_name="viz_current_wrist_poses", + delete_old_keys=True, + ), + HeadFrameKeypoints("obs_head_pose", "left.obs_keypoints", "L_kp_obs"), + HeadFrameKeypoints("obs_head_pose", "right.obs_keypoints", "R_kp_obs"), + ConcatKeys(key_list=obs_keys, new_key_name="state_keypoints", + delete_old_keys=True), + DropKeys(["left.current_wrist_pose", "right.current_wrist_pose", + "left.act_keypoints", "right.act_keypoints", + "left.act_wrist_pose", "right.act_wrist_pose", + "left.obs_keypoints", "right.obs_keypoints", + "left.obs_wrist_pose", "right.obs_wrist_pose"]), + NumpyToTensor(keys=["actions_keypoints", "state_keypoints", "obs_head_pose", + "viz_current_wrist_poses"]), + ] + + +# --------------------------------------------------------------------------- # +# ROLLOUT transform lists (2026-08-16). +# +# Deploy must preprocess obs EXACTLY as training did, and must undo the action +# encoding exactly, or the policy is fed/read in a different space than it was +# fitted in -- the classic silent train/deploy skew. These live beside the +# forward lists on purpose: if eva_normal_transforms changes, the mismatch is +# visible in the same file rather than in robot code nobody re-reads. +# --------------------------------------------------------------------------- # +class SplitConcat(Transform): + """Inverse of ConcatKeys: split one array into named parts along the last + axis. ``parts`` is [(name, width), ...] and must sum to the array width.""" + + def __init__(self, in_key, parts, delete_old_key=True): + self.in_key = in_key + self.parts = [(str(n), int(w)) for n, w in parts] + self.delete_old_key = bool(delete_old_key) + + def transform(self, batch): + v = np.asarray(batch[self.in_key]) + total = sum(w for _, w in self.parts) + if v.shape[-1] != total: + raise ValueError( + f"SplitConcat: {self.in_key!r} is {v.shape[-1]}-D but parts " + f"{self.parts} sum to {total}. The action space changed -- " + f"update the split rather than letting it mis-slice.") + off = 0 + for name, w in self.parts: + batch[name] = v[..., off:off + w] + off += w + if self.delete_old_key: + batch.pop(self.in_key, None) + return batch + + +class Rot6DToPoseYPR(Transform): + """xyz + rot6d (..., 9) -> xyz + ypr (..., 6). + + Inverts PoseToRot6D (rot6d_to_matrix is its documented inverse) and then + converts to the ZYX euler the robot interface speaks -- NOT back to the + quaternion the raw zarr held, because rollout.py's + cam_frame_to_base_frame / rot_ee_frame_to_ee_pose_batch both read + pose[..., 3:6] as ypr. + """ + + def __init__(self, in_key, out_key=None): + self.in_key = in_key + self.out_key = out_key or in_key + + def transform(self, batch): + v = np.asarray(batch[self.in_key], dtype=np.float64) + single = v.ndim == 1 + if single: + v = v[None, :] + if v.shape[-1] != 9: + raise ValueError( + f"Rot6DToPoseYPR: {self.in_key!r} is {v.shape[-1]}-D, expected 9 " + f"(xyz3 + rot6d6).") + ypr = R.from_matrix(rot6d_to_matrix(v[:, 3:9])).as_euler("ZYX") + out = np.concatenate([v[:, 0:3], ypr], axis=-1).astype(np.float32) + batch[self.out_key] = out[0] if single else out + return batch + + +def build_bimanual_rot6d_wrist_revert_transforms( + action_key="actions_cartesian", state_key="state_ee_pose" +): + """20-D fold wrist chunk -> 14-D parent-frame pose chunk. + + This is only composition: rotation conversion and SE(3) application stay in + the existing ``Rot6DToPoseYPR`` and ``ActionChunkCoordinateFrameTransform`` + utilities used by the training/rollout pipelines. + """ + return [ + SplitConcat( + action_key, + [("L_action_pose", 9), ("L_action_grip", 1), + ("R_action_pose", 9), ("R_action_grip", 1)], + ), + SplitConcat( + state_key, + [("L_state_pose", 9), ("L_state_grip", 1), + ("R_state_pose", 9), ("R_state_grip", 1)], + delete_old_key=False, + ), + SelectFrame("L_state_pose", -1, "L_current_pose"), + SelectFrame("R_state_pose", -1, "R_current_pose"), + Rot6DToPoseYPR("L_action_pose"), + Rot6DToPoseYPR("R_action_pose"), + Rot6DToPoseYPR("L_current_pose"), + Rot6DToPoseYPR("R_current_pose"), + ActionChunkCoordinateFrameTransform( + "L_current_pose", "L_action_pose", "L_action_parent", + mode="xyzypr", inverse=False, + ), + ActionChunkCoordinateFrameTransform( + "R_current_pose", "R_action_pose", "R_action_parent", + mode="xyzypr", inverse=False, + ), + ConcatKeys( + ["L_action_parent", "L_action_grip", + "R_action_parent", "R_action_grip"], + action_key, + delete_old_keys=True, + ), + ] + + +def build_bimanual_keypoint_wrist_revert_transforms( + action_key="actions_keypoints", wrist_pose_key="viz_current_wrist_poses" +): + """132-D fold wrist/keypoint chunk -> head-camera coordinates.""" + return [ + SplitConcat( + action_key, + [("L_wrist_xyz", 3), ("L_keypoints", 63), + ("R_wrist_xyz", 3), ("R_keypoints", 63)], + ), + SplitConcat( + wrist_pose_key, + [("L_current_wrist", 7), ("R_current_wrist", 7)], + delete_old_key=False, + ), + ReshapePoints("L_keypoints"), + ReshapePoints("R_keypoints"), + ActionChunkCoordinateFrameTransform( + "L_current_wrist", "L_wrist_xyz", "L_wrist_head", + mode="xyz", inverse=False, + ), + ActionChunkCoordinateFrameTransform( + "R_current_wrist", "R_wrist_xyz", "R_wrist_head", + mode="xyz", inverse=False, + ), + ActionChunkCoordinateFrameTransform( + "L_current_wrist", "L_keypoints", "L_keypoints_head", + mode="xyz", inverse=False, + ), + ActionChunkCoordinateFrameTransform( + "R_current_wrist", "R_keypoints", "R_keypoints_head", + mode="xyz", inverse=False, + ), + ReshapePoints("L_keypoints_head", flatten=True), + ReshapePoints("R_keypoints_head", flatten=True), + ConcatKeys( + ["L_wrist_head", "L_keypoints_head", + "R_wrist_head", "R_keypoints_head"], + action_key, + delete_old_keys=True, + ), + ] + + +def eva_rollout_obs_transforms(): + """OBS-ONLY subset of eva_normal_transforms. + + Identical ops on the same keys, minus everything touching cmd_* -- those + keys do not exist at rollout time. state_ee_pose comes out 20-D, matching + what the encoder was trained on. + """ + return [ + PoseToRot6D("left.obs_ee_pose"), PoseToRot6D("right.obs_ee_pose"), + ConcatKeys(key_list=["left.obs_ee_pose", "left.obs_gripper", + "right.obs_ee_pose", "right.obs_gripper"], + new_key_name="state_ee_pose", delete_old_keys=True), + NumpyToTensor(keys=["state_ee_pose"]), + ] + + +def eva_action_revert_transforms(in_key="actions_cartesian", + out_key="robot_action"): + """MODEL action (..., 20) -> ROBOT action (..., 14). + + model : [L xyz3 rot6d6 grip1 | R xyz3 rot6d6 grip1] = 20 + robot : [L xyz3 ypr3 grip1 | R xyz3 ypr3 grip1] = 14 + + The exact inverse of the action path in eva_normal_transforms, re-expressed + in the robot's ypr convention. Frame transforms (cam->base, rot-ee->ee) are + NOT here: they need the per-arm extrinsics and live in the rollout node, so + this list stays pure and testable. + """ + return [ + SplitConcat(in_key, parts=[("L_pose", 9), ("L_grip", 1), + ("R_pose", 9), ("R_grip", 1)]), + Rot6DToPoseYPR("L_pose"), Rot6DToPoseYPR("R_pose"), + ConcatKeys(key_list=["L_pose", "L_grip", "R_pose", "R_grip"], + new_key_name=out_key, delete_old_keys=True), + ] diff --git a/egomimic/rldb/norm_stats.py b/egomimic/rldb/norm_stats.py new file mode 100644 index 000000000..2ceeb8c7e --- /dev/null +++ b/egomimic/rldb/norm_stats.py @@ -0,0 +1,145 @@ +"""Dataset-free normalization state used by training and deployment models.""" + +from __future__ import annotations + +import copy + +import numpy as np +import torch + + +class NormStats: + """Checkpoint normalization metadata without importing a dataset backend.""" + + NORMALIZE_KEY_TYPES = ("proprio_keys", "action_keys") + + def __init__(self, state: dict): + if state is None: + raise ValueError("normalization state is required") + self.norm_mode = state.get("norm_mode", "zscore") + self.embodiments = set(state.get("embodiments", [])) + self.key_types = copy.deepcopy(state.get("key_types", {})) + self.zarr_keys = copy.deepcopy(state.get("zarr_keys", {})) + self.shapes = copy.deepcopy(state.get("shapes", {})) + self.norm_stats = self._clone_norm_stats(state.get("norm_stats", {})) + for embodiment_id in self.embodiments: + self.key_types.setdefault(embodiment_id, {}) + self.zarr_keys.setdefault(embodiment_id, {}) + self.shapes.setdefault(embodiment_id, {}) + self.norm_stats.setdefault(embodiment_id, {}) + + @staticmethod + def _clone_norm_stats(norm_stats): + return { + embodiment_id: { + key: { + name: ( + value.detach().cpu().clone() + if torch.is_tensor(value) + else copy.deepcopy(value) + ) + for name, value in stats.items() + } + for key, stats in per_embodiment.items() + } + for embodiment_id, per_embodiment in (norm_stats or {}).items() + } + + def keys_of_type(self, key_type: str, embodiment_id: int) -> list[str]: + return [ + key + for key, actual_type in self.key_types.get(embodiment_id, {}).items() + if actual_type == key_type + ] + + def is_key_with_embodiment(self, key_name: str, embodiment_id: int) -> bool: + return key_name in self.key_types.get(embodiment_id, {}) + + def keyname_to_zarr_key(self, key_name: str, embodiment_id: int) -> str | None: + return self.zarr_keys.get(embodiment_id, {}).get(key_name) + + def zarr_key_to_keyname(self, zarr_key: str, embodiment_id: int) -> str | None: + for key_name, candidate in self.zarr_keys.get(embodiment_id, {}).items(): + if candidate == zarr_key: + return key_name + return None + + def key_shape(self, key_name: str, embodiment_id: int) -> tuple: + try: + return self.shapes[embodiment_id][key_name] + except KeyError as exc: + raise ValueError( + f"Shape for key {key_name!r} on embodiment {embodiment_id} " + "is unavailable." + ) from exc + + def _apply_norm_one(self, tensor, stats): + if self.norm_mode == "zscore": + mean = torch.as_tensor(stats["mean"], device=tensor.device).float() + std = torch.as_tensor(stats["std"], device=tensor.device).float() + return (tensor - mean) / (std + 1e-6) + if self.norm_mode == "minmax": + minimum = torch.as_tensor(stats["min"], device=tensor.device).float() + maximum = torch.as_tensor(stats["max"], device=tensor.device).float() + return 2.0 * ((tensor - minimum) / (maximum - minimum + 1e-6)) - 1.0 + if self.norm_mode == "quantile": + q1 = torch.as_tensor(stats["quantile_1"], device=tensor.device).float() + q99 = torch.as_tensor(stats["quantile_99"], device=tensor.device).float() + return 2.0 * ((tensor - q1) / (q99 - q1 + 1e-6)) - 1.0 + raise ValueError(f"Invalid normalization mode: {self.norm_mode}") + + def _apply_unnorm_one(self, tensor, stats): + if self.norm_mode == "zscore": + mean = torch.as_tensor(stats["mean"], device=tensor.device).float() + std = torch.as_tensor(stats["std"], device=tensor.device).float() + return tensor * (std + 1e-6) + mean + if self.norm_mode == "minmax": + minimum = torch.as_tensor(stats["min"], device=tensor.device).float() + maximum = torch.as_tensor(stats["max"], device=tensor.device).float() + return (tensor + 1) * 0.5 * (maximum - minimum + 1e-6) + minimum + if self.norm_mode == "quantile": + q1 = torch.as_tensor(stats["quantile_1"], device=tensor.device).float() + q99 = torch.as_tensor(stats["quantile_99"], device=tensor.device).float() + return (tensor + 1) * 0.5 * (q99 - q1 + 1e-6) + q1 + raise ValueError(f"Invalid normalization mode: {self.norm_mode}") + + def normalize(self, data: dict, embodiment_id: int) -> dict: + if not self.norm_stats.get(embodiment_id): + return data + out = dict(data) + for key_name, key_type in self.key_types.get(embodiment_id, {}).items(): + if key_type not in self.NORMALIZE_KEY_TYPES: + continue + stats = self.norm_stats[embodiment_id].get(key_name) + zarr_key = self.zarr_keys[embodiment_id].get(key_name) + if stats is None or zarr_key not in out: + continue + value = out[zarr_key] + if isinstance(value, np.ndarray): + value = torch.from_numpy(value).float() + if torch.is_tensor(value): + out[zarr_key] = self._apply_norm_one(value, stats) + return out + + def unnormalize(self, data: dict, embodiment_id: int) -> dict: + if not self.norm_stats.get(embodiment_id): + return data + out = dict(data) + zarr_to_name = { + value: key + for key, value in self.zarr_keys.get(embodiment_id, {}).items() + } + for data_key, value in data.items(): + key_name = ( + data_key + if data_key in self.norm_stats[embodiment_id] + else zarr_to_name.get(data_key) + ) + stats = self.norm_stats[embodiment_id].get(key_name) + if stats is None: + continue + if isinstance(value, np.ndarray): + value = torch.from_numpy(value).float() + if torch.is_tensor(value): + out[data_key] = self._apply_unnorm_one(value, stats) + return out diff --git a/tests/test_inference_graph.py b/tests/test_inference_graph.py new file mode 100644 index 000000000..2e65fb5b6 --- /dev/null +++ b/tests/test_inference_graph.py @@ -0,0 +1,70 @@ +from egomimic.eval.inference_graph import ( + ActionCacheState, + InferenceGraph, + KeyedNode, + Subgraph, +) + + +def _node(fn, inputs, outputs): + return KeyedNode(fn, **{"in": inputs, "out": outputs}) + + +def test_cache_hit_exits_before_preprocess_and_model(): + calls = [] + graph = InferenceGraph( + check_cache=_node(lambda obs: calls.append("check") or 7, + {"obs": "obs"}, {"action": "policy.action"}), + inference_preprocess=_node( + lambda obs: calls.append("pre") or obs, + {"obs": "obs"}, {"request": "request"}), + model=_node(lambda request: calls.append("model") or request, + {"request": "request"}, {"plan": "plan"}), + update_cache=_node(lambda plan: calls.append("update") or plan, + {"plan": "plan"}, + {"action": "policy.action"}), + ) + assert graph(obs=3) == 7 + assert calls == ["check"] + + +def test_cache_miss_runs_preprocess_model_update_in_order(): + calls = [] + graph = InferenceGraph( + check_cache=_node(lambda obs: calls.append("check"), + {"obs": "rollout.obs"}, + {"action": "policy.action"}), + inference_preprocess=_node( + lambda obs: calls.append("pre") or obs + 1, + {"obs": "rollout.obs"}, {"request": "model.request"}), + model=_node(lambda request: calls.append("model") or request * 2, + {"request": "model.request"}, {"plan": "model.plan"}), + update_cache=_node( + lambda plan: calls.append("update") or plan + 3, + {"plan": "model.plan"}, {"action": "policy.action"}), + ) + assert graph(**{"rollout.obs": 4}) == 13 + assert calls == ["check", "pre", "model", "update"] + + +def test_subgraph_has_single_entry_and_endpoint(): + subgraph = Subgraph( + [ + _node(lambda x: x + 2, {"x": "x"}, {"y": "y"}), + _node(lambda y: y * 3, {"y": "y"}, {"result": "result"}), + ], + **{"in": {"x": "outer.x"}, "out": {"result": "outer.y"}}, + ) + context = {"outer.x": 5} + subgraph(context) + assert context["outer.y"] == 21 + + +def test_action_cache_is_instance_scoped_and_resettable(): + left, right = ActionCacheState(), ActionCacheState() + left.replace([1, 2]) + right.replace([9]) + assert left.pop() == 1 + assert right.pop() == 9 + left.reset() + assert not left and not right diff --git a/tests/test_norm_stats.py b/tests/test_norm_stats.py new file mode 100644 index 000000000..3fd5f1db6 --- /dev/null +++ b/tests/test_norm_stats.py @@ -0,0 +1,48 @@ +import numpy as np +import torch + +from egomimic.rldb.norm_stats import NormStats + + +def _state(mode): + return { + "norm_mode": mode, + "embodiments": [8], + "key_types": {8: {"action": "action_keys", "image": "camera_keys"}}, + "zarr_keys": {8: {"action": "actions_cartesian", "image": "front"}}, + "shapes": {8: {"action": (2,)}}, + "norm_stats": { + 8: { + "action": { + "mean": np.array([2.0, 4.0]), + "std": np.array([2.0, 4.0]), + "min": np.array([0.0, 0.0]), + "max": np.array([4.0, 8.0]), + "quantile_1": np.array([0.0, 0.0]), + "quantile_99": np.array([4.0, 8.0]), + } + } + }, + } + + +def test_checkpoint_norm_stats_round_trip_all_modes(): + action = torch.tensor([[1.0, 6.0]]) + for mode in ("zscore", "minmax", "quantile"): + stats = NormStats(_state(mode)) + normalized = stats.normalize({"actions_cartesian": action}, 8) + restored = stats.unnormalize(normalized, 8) + torch.testing.assert_close(restored["actions_cartesian"], action) + + +def test_checkpoint_norm_stats_key_interface_and_numpy_input(): + stats = NormStats(_state("zscore")) + normalized = stats.normalize( + {"actions_cartesian": np.array([2.0, 4.0], dtype=np.float32)}, 8 + ) + + torch.testing.assert_close(normalized["actions_cartesian"], torch.zeros(2)) + assert stats.keys_of_type("action_keys", 8) == ["action"] + assert stats.keyname_to_zarr_key("action", 8) == "actions_cartesian" + assert stats.zarr_key_to_keyname("actions_cartesian", 8) == "action" + assert stats.key_shape("action", 8) == (2,) diff --git a/tests/test_pipeline_inference_graph.py b/tests/test_pipeline_inference_graph.py new file mode 100644 index 000000000..28fd048ef --- /dev/null +++ b/tests/test_pipeline_inference_graph.py @@ -0,0 +1,135 @@ +import numpy as np +import torch + +from egomimic.eval.inference_graph import ActionCacheState +from egomimic.pipeline.algo import PipelineAlgo + + +class _NormStats: + def __init__(self): + self.normalize_calls = 0 + + def normalize(self, obs, emb_id): + self.normalize_calls += 1 + return dict(obs) + + def unnormalize(self, values, emb_id): + return values + + +class _TinyHNet(PipelineAlgo): + """Minimal rollout surface; no training pipeline construction required.""" + + def __init__(self, inference_stages=None): + self.nets = torch.nn.ModuleDict({"anchor": torch.nn.Linear(1, 1)}) + self.norm_stats = _NormStats() + self.action_horizon = 32 + self.replan_stride = 2 + self.domain_by_id = {15: "pushshapes_sim"} + self.ac_keys = {"pushshapes_sim": "actions"} + self.inference_stages = inference_stages + self._sim_action_cache = ActionCacheState() + self._sim_action_queue = self._sim_action_cache.actions + self._inference_graph = self._build_inference_graph() + self.model_calls = 0 + self.reset_calls = 0 + + def init_step_state(self, batch_size, T_max, device, dtype): + self.reset_calls += 1 + return {"T_max": T_max} + + def step(self, state, obs_norm, t, embodiment_id=None): + self.model_calls += 1 + base = 10 * t + return torch.tensor( + [[base + 1.0], [base + 2.0], [base + 3.0]], + device=next(self.nets.parameters()).device, + ) + + +def _act(policy, t): + return policy.inference_step({"x": torch.tensor([float(t)])}, t, 15) + + +def test_hnet_graph_cache_hit_skips_preprocess_and_model(): + policy = _TinyHNet() + + np.testing.assert_array_equal(_act(policy, 0), np.array([1], np.float32)) + assert policy.model_calls == 1 + assert policy.norm_stats.normalize_calls == 1 + assert len(policy._sim_action_cache.actions) == 1 + + np.testing.assert_array_equal(_act(policy, 1), np.array([2], np.float32)) + assert policy.model_calls == 1 + assert policy.norm_stats.normalize_calls == 1 + assert not policy._sim_action_cache + + np.testing.assert_array_equal(_act(policy, 2), np.array([21], np.float32)) + assert policy.model_calls == 2 + assert policy.norm_stats.normalize_calls == 2 + + +def test_runtime_obs_adapter_runs_only_after_cache_miss(): + policy = _TinyHNet() + calls = [] + + def adapt(obs): + calls.append(float(obs["x"][0])) + return {"x": obs["x"] + 100} + + policy.inference_obs_adapter = adapt + _act(policy, 0) # miss: preprocess and model + _act(policy, 1) # hit: graph exits before environment/model adaptation + _act(policy, 2) # miss again + + assert calls == [0.0, 2.0] + + +def test_hnet_graph_t0_resets_history_and_cache_per_episode(): + policy = _TinyHNet() + _act(policy, 0) + assert policy._sim_action_cache + _act(policy, 0) + assert policy.reset_calls == 2 + assert policy.model_calls == 2 + assert len(policy._sim_action_cache.actions) == 1 + + +def test_hnet_graph_state_is_policy_instance_scoped(): + first, second = _TinyHNet(), _TinyHNet() + _act(first, 0) + assert first._sim_action_cache + assert not second._sim_action_cache + _act(second, 0) + assert first._sim_action_cache.actions is not second._sim_action_cache.actions + + +def test_hnet_graph_accepts_literal_key_remapping(): + cfg = { + "terminal": "hnet.action", + "nodes": { + "check_cache": { + "in": {"obs": "raw.obs"}, + "out": {"action": "hnet.action"}, + }, + "inference_preprocess": { + "in": {"obs": "raw.obs"}, + "out": {"request": "hnet.request"}, + }, + "update_cache": { + "in": {"plan": "hnet.plan", "obs": "raw.obs"}, + "out": {"action": "hnet.action"}, + }, + }, + "model": { + "in": {"request": "hnet.request"}, + "out": {"plan": "hnet.plan"}, + }, + } + policy = _TinyHNet(cfg) + policy._reset_inference_graph() + policy._sim_t = 0 + policy._sim_emb_id = 15 + policy._sim_ac_key = "actions" + out = policy._inference_graph(**{"raw.obs": {"x": torch.tensor([0.0])}}) + np.testing.assert_array_equal(out, np.array([1], np.float32)) diff --git a/tests/test_pipeline_normal_dp_rollout.py b/tests/test_pipeline_normal_dp_rollout.py new file mode 100644 index 000000000..ef5b14ad8 --- /dev/null +++ b/tests/test_pipeline_normal_dp_rollout.py @@ -0,0 +1,131 @@ +import torch + +from egomimic.eval.inference_graph import ActionCacheState +from egomimic.pipeline.algo import PipelineAlgo +from egomimic.pipeline.core import Pipeline, Stage +from egomimic.pipeline.stages_io import NormalObsCollapse, NormalObsExpand +from egomimic.pipeline.stages_seq import ObsStack + + +class _Encode(Stage): + reads = ["obs/value"] + writes = ["a_top", "s"] + + def forward(self, batch): + batch["a_top"] = batch["obs/value"].float() + batch["s"] = batch["obs/value"].float() + return batch + + +class _Head(Stage): + reads = ["a_top", "s", "embodiment"] + writes = ["pred_action"] + + def forward(self, batch): + # One action row whose two values expose [previous, current]. + batch["pred_action"] = batch["a_top"].unsqueeze(1).repeat(1, 2, 1) + return batch + + +class _NormalDP(PipelineAlgo): + def __init__(self): + # Minimal rollout-only surface; avoid constructing the training runner. + self.policy = Pipeline([ + NormalObsExpand(n_obs_steps=2), + _Encode(), + ObsStack( + in_keys=["a_top", "s"], out_keys=["a_top", "s"], + n_obs_steps=2, + ), + NormalObsCollapse(keys=["a_top", "s"]), + _Head(), + ]) + + def activate_rollout_apex_attention(self): + return None + + +def test_normal_dp_rollout_reuses_training_adapters_without_target(): + policy = _NormalDP() + state = policy.init_step_state(1, 20, torch.device("cpu"), torch.float32) + + first = policy.step( + state, {"value": torch.tensor([[1.0]])}, 0, "eva_bimanual") + second = policy.step( + state, {"value": torch.tensor([[2.0]])}, 1, "eva_bimanual") + + torch.testing.assert_close(first[0], torch.tensor([1.0, 1.0])) + torch.testing.assert_close(second[0], torch.tensor([1.0, 2.0])) + assert [type(stage).__name__ for stage in state["plan"]] == [ + "NormalObsExpand", "_Encode", "ObsStack", "NormalObsCollapse", "_Head" + ] + + +def test_normal_obs_expand_still_requires_actions_outside_rollout(): + stage = NormalObsExpand(n_obs_steps=2) + try: + stage({"obs/value": torch.zeros(1, 2, 1)}) + except ValueError as error: + assert "outside rollout" in str(error) + else: + raise AssertionError("target-free training batch was accepted") + + +def test_normal_obs_training_target_contract_is_unchanged(): + expand = NormalObsExpand(n_obs_steps=2) + collapse = NormalObsCollapse(keys=["a_top", "s"]) + actions = torch.arange(6, dtype=torch.float32).reshape(1, 3, 2) + batch = expand({ + "obs/value": torch.tensor([[[1.0], [2.0]]]), + "actions": actions.clone(), + }) + batch["a_top"] = batch["obs/value"] + batch["s"] = batch["obs/value"] + batch = collapse(batch) + + torch.testing.assert_close(batch["target"], actions) + torch.testing.assert_close(batch["a_top"], torch.tensor([[2.0]])) + + +class _IdentityNorm: + def normalize(self, obs, emb_id): + return obs + + def unnormalize(self, values, emb_id): + return values + + +class _GraphNormalDP(_NormalDP): + def __init__(self): + super().__init__() + self.nets = torch.nn.ModuleDict({"anchor": torch.nn.Linear(1, 1)}) + self.norm_stats = _IdentityNorm() + self.action_horizon = 20 + self.replan_stride = 2 + self.domain_by_id = {15: "eva_bimanual"} + self.ac_keys = {"eva_bimanual": "actions"} + self.inference_stages = None + self._sim_action_cache = ActionCacheState() + self._inference_graph = self._build_inference_graph() + + +def test_dp_cache_keeps_immediately_previous_environment_frame(): + policy = _GraphNormalDP() + policy._reset_inference_graph() + policy._sim_emb_id = 15 + policy._sim_ac_key = "actions" + + outputs = [] + for t, value in enumerate((1.0, 2.0, 3.0)): + policy._sim_t = t + outputs.append(policy._inference_graph( + obs={"value": torch.tensor([[value]])})) + + # t1 is a cache hit. At t2 the two-frame DP input must be [t1,t2], not + # [previous model query t0, current t2]. + torch.testing.assert_close(torch.from_numpy(outputs[0]), + torch.tensor([1.0, 1.0])) + torch.testing.assert_close(torch.from_numpy(outputs[1]), + torch.tensor([1.0, 1.0])) + torch.testing.assert_close(torch.from_numpy(outputs[2]), + torch.tensor([2.0, 3.0]))