feat(pipeline): sync current inference graph runtime

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]

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)
 

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

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,