diff --git a/egomimic/algo/diffusion/__init__.py b/egomimic/algo/diffusion/__init__.py new file mode 100644 index 000000000..40c5b02f9 --- /dev/null +++ b/egomimic/algo/diffusion/__init__.py @@ -0,0 +1,93 @@ +"""DFoT (Diffusion Forcing Transformer) ALGO package (DESIGN.md §2 +``egomimic/algo/diffusion/``). + +Role home for the *algo* half of the diffusion stack — the :class:`DFoT` algo, +its per-variant outer-stages, and the VAE pre-train algo. Relocated here from +``egomimic/algo/dfot`` (+ ``egomimic/algo/vae``) in DESIGN.md step 7 via +``git mv`` (no behaviour change); the *model* half (backbones / diffusion +processes / embeddings / sampling) lives at ``egomimic.models.diffusion``. + + * :class:`DFoT` — the diffusion-forcing algo. + * ``outer_stages/`` — per-variant OuterStage subclasses + (spatial / pixel / obs-action / ...). + * :class:`VAE` + ``load_pretrained_vae`` (``vae_algo``) — image VAE pre-train. + +The legacy ``egomimic.algo.dfot`` / ``egomimic.algo.vae`` facade shims (and the +yaml ``_target_``s that routed through them) were deleted at the DESIGN.md +step-13 final flip; every consumer now imports from ``egomimic.algo.diffusion`` +/ ``egomimic.models.diffusion`` directly. +""" + +from egomimic.algo.diffusion.algo import DFoT +from egomimic.algo.diffusion.outer_stages.outer_stage import ( + DFoTOuterStage, + make_dfot_ctx, +) +from egomimic.algo.diffusion.outer_stages.image_spatial_outer_stage import ( + ImageSpatialDFoTOuterStage, +) +from egomimic.algo.diffusion.outer_stages.obs_action_image_outer_stage import ( + ObsActionImageDFoTOuterStage, +) +from egomimic.algo.diffusion.outer_stages.obs_action_outer_stage import ( + ObsActionDFoTOuterStage, +) +from egomimic.algo.diffusion.outer_stages.pixel_obs_action_outer_stage import ( # noqa: E501 + PixelObsActionDFoTOuterStage, +) +from egomimic.algo.diffusion.outer_stages.pixel_spatial_outer_stage import ( + PixelSpatialDFoTOuterStage, +) +from egomimic.algo.diffusion.outer_stages.pixel_video_outer_stage import ( + PixelVideoDFoTOuterStage, +) +from egomimic.algo.diffusion.outer_stages.spatial_obs_action_policy_outer_stage import ( # noqa: E501 + SpatialObsActionPolicyDFoTOuterStage, +) +from egomimic.algo.diffusion.vae_algo import VAE, load_pretrained_vae + +# Convenience re-exports of the model-half pieces some legacy code imports via +# the algo namespace (kept here so ``from egomimic.algo.diffusion import +# DFoTBackbone`` mirrors the model-half surface). +from egomimic.models.diffusion import ( + ContinuousDiffusion, + DFoTBackbone, + DFoTDiT3DBackbone, + DFoTSpatialBackbone, + DiscreteDiffusion, + causal_ar_schedule, + ddim_sample, + ddpm_sample, + sample, + sample_step, + staircase_ar_schedule, + vanilla_schedule, +) + +__all__ = [ + "DFoT", + "DFoTOuterStage", + "make_dfot_ctx", + "ImageSpatialDFoTOuterStage", + "ObsActionImageDFoTOuterStage", + "ObsActionDFoTOuterStage", + "PixelObsActionDFoTOuterStage", + "PixelSpatialDFoTOuterStage", + "PixelVideoDFoTOuterStage", + "SpatialObsActionPolicyDFoTOuterStage", + "VAE", + "load_pretrained_vae", + # model-half re-exports + "DFoTBackbone", + "DFoTDiT3DBackbone", + "DFoTSpatialBackbone", + "ContinuousDiffusion", + "DiscreteDiffusion", + "sample_step", + "sample", + "vanilla_schedule", + "causal_ar_schedule", + "staircase_ar_schedule", + "ddim_sample", + "ddpm_sample", +] diff --git a/egomimic/algo/diffusion/algo.py b/egomimic/algo/diffusion/algo.py new file mode 100644 index 000000000..61fb37e0b --- /dev/null +++ b/egomimic/algo/diffusion/algo.py @@ -0,0 +1,1082 @@ +""" +DFoT Algo: per-token-noise-level diffusion over action chunks. + +Supports both packed (full variable-length episodes; ``cu_seqlens``-driven +within-episode attention) and padded (fixed-T windows) batches. In packed +mode obs is per-frame (one obs per action timestep, aligned via the +``pushshapes.get_keymap`` per-frame keymap); in padded mode obs may be +single-frame and is broadcast across T at the per-token AdaLN. Loss is the +diffusion (epsilon / v / x0) MSE with the configured weighting strategy. + +Two inference modes (see ``inference_step``): + * "ar": rolling causal-AR staircase. One ``sample_step`` per env tick. + Matches training distribution. Default. + * "chunk": vanilla DDIM over a fixed window with plan-and-execute. Legacy + baseline; doesn't exercise DFoT's per-token-noise capability. + +Teacher-forced offline val viz lives in ``egomimic/eval/eval_dfot_val.py`` +(``DFoTValEval``). +""" + +from collections import OrderedDict +from typing import Optional + +import numpy as np +import torch +import torch.nn as nn +from overrides import override + +from egomimic.algo.algo import Algo +from egomimic.models.diffusion.backbones.backbone import DFoTBackbone +from egomimic.models.diffusion.diffusion.continuous_diffusion import ContinuousDiffusion +from egomimic.models.diffusion.diffusion.discrete_diffusion import DiscreteDiffusion +from egomimic.algo.diffusion.outer_stages.outer_stage import DFoTOuterStage, make_dfot_ctx +from egomimic.models.diffusion.sampling import ddim_sample, ddpm_sample, sample_step +from egomimic.models.stems.cond_encoders import CondEncoderModule +from egomimic.rldb.embodiment.embodiment import get_embodiment, get_embodiment_id + + +class _PackedBackboneWrapper: + """Closure-style wrapper so the diffusion module's 3-arg + ``backbone(x, t, cond)`` call automatically threads cu_seqlens / max_seqlen + into ``DFoTBackbone.forward`` for packed-mode within-episode attention.""" + + def __init__(self, backbone: DFoTBackbone, cu_seqlens, max_seqlen): + self.backbone = backbone + self.cu_seqlens = cu_seqlens + self.max_seqlen = max_seqlen + + def __call__(self, x, noise_levels, external_cond=None): + return self.backbone( + x, + noise_levels, + external_cond=external_cond, + cu_seqlens=self.cu_seqlens, + max_seqlen=self.max_seqlen, + ) + + +class DFoTLoss(nn.Module): + """DFoT epsilon-MSE with sigmoid (SNR-style) weighting. + + Self-contained loss for the ``DFoT`` algo (each model owns its own loss). + Implements the ``forward(batch, ctx) -> scalar`` contract the algo + orchestrator calls. + + Reads: + - ``batch["pred_v"]``: v-prediction emitted by the DFoT backbone + (written by ``DFoTOuterStage.decode``). + - ``ctx.q_state``: dict produced by ``diffusion.q_sample`` during + ``DFoTOuterStage.encode``. Carries ``x_t``, ``noise``, ``alpha_t``, + ``sigma_t``, ``logsnr``. + + The actual math lives in ``diffusion.compute_loss``; this class is the + interface that fits the OuterStage-orchestrated training loop and + reduces the per-token loss to a scalar. + + ``diffusion`` is the same instance held by the outer stage — it has + no learnable params (in continuous mode) so this is a free reference, + not a duplicate submodule. + """ + + def __init__(self, diffusion: nn.Module): + super().__init__() + self.diffusion = diffusion + + def forward(self, batch: dict, ctx) -> torch.Tensor: + # Structured-target outer stages (e.g. the 2D spatial-image + action + # policy) compute their own multi-term loss and stash it here, since a + # single (image, action) target can't flow through the scalar v-MSE + # path below. No-op for every existing 1D/spatial stage. + precomputed = getattr(ctx, "precomputed_loss", None) + if precomputed is not None: + return precomputed + v_pred = batch["pred_v"] + q_state = ctx.q_state + per_token = self.diffusion.compute_loss(v_pred, q_state) + return per_token.mean() + + +class DFoT(Algo): + """Diffusion Forcing Transformer (action-chunk denoising) Algo. + + Args: + action_dim: action feature width. + action_horizon: AR buffer / chunk length T. Also the planning window + for legacy chunk-mode inference. + cond_encoder: ``CondEncoderModule`` for obs -> ``(B, T, cond_dim)`` + (per-frame) or ``(B, cond_dim)`` (single-frame, broadcast). + backbone: ``DFoTBackbone`` (built via Hydra). Owns x / cond / time + projections and the ``Isotropic`` trunk. + norm_stats: ``MultiDataset`` (injected by + ``pl_model._instantiate_model``). + diffusion_type: ``"discrete"`` or ``"continuous"``. + diffusion_kwargs: dict forwarded to the chosen diffusion class. + sampler: ``"ddpm"`` or ``"ddim"`` — used only by chunk-mode inference. + sampler_n_steps: denoising step count for chunk-mode inference. + sampler_eta: DDIM eta (0.0 = deterministic) for chunk-mode inference. + inference_mode: ``"ar"`` (default) for rolling causal-AR staircase or + ``"chunk"`` for legacy plan-and-execute. + ar_inference_chunk_size: tokens committed per env tick in AR mode + (1 = classic causal AR; >1 = chunked staircase rungs). + ``action_horizon`` must be divisible by this. + domains: list of embodiment names (single-element for v1). + ac_keys: dict ``embodiment_name -> action zarr key``. + cond_output_key: key under which the cond encoder exposes its fused + cond. + """ + + def __init__( + self, + outer_stage: DFoTOuterStage, + action_dim: int, + action_horizon: int, + norm_stats, + loss: Optional[nn.Module] = None, + sampler: str = "ddim", + sampler_n_steps: int = 50, + sampler_eta: float = 0.0, + inference_mode: str = "ar", + ar_inference_chunk_size: int = 1, + ar_inference_step_size: int = 1, + cfg_scale: float = 1.0, + sp_n_context: int = 4, + sp_commit: int = 1, + sp_n_samples: int = 1, + domains: Optional[list] = None, + ac_keys: Optional[dict] = None, + device=None, + **kwargs, + ): + """Refactored DFoT algo. + + Args: + outer_stage: ``DFoTOuterStage`` owning the cond_encoder + backbone + + diffusion submodules and implementing the training-path + encode -> q_sample -> backbone -> decode flow. + loss: an ``nn.Module`` (typically ``DFoTLoss``) that consumes + ``batch['pred_v']`` + ``ctx.q_state`` and emits the scalar + training loss. + sampler / sampler_n_steps / sampler_eta / inference_mode / + ar_inference_chunk_size / ar_inference_step_size / cfg_scale: + Inference knobs. Closed-loop AR + chunk-mode inference paths + still live on this algo class and consume the outer_stage's + submodules via the ``cond_encoder`` / ``backbone`` / + ``diffusion`` properties below. + """ + super().__init__() + self.norm_stats = norm_stats + self.domains = list(domains or []) + self.ac_keys = dict(ac_keys or {}) + self.action_dim = int(action_dim) + self.action_horizon = int(action_horizon) + self.cond_output_key = outer_stage.cond_output_key + self.sampler = sampler + self.sampler_n_steps = int(sampler_n_steps) + self.sampler_eta = float(sampler_eta) + if inference_mode not in {"ar", "chunk", "spatial_rh", "spatial_decoupled", "pixel_policy", "pixel_regress", "pixel_decoupled"}: + raise ValueError( + f"inference_mode must be 'ar', 'chunk', 'spatial_rh', or " + f"'spatial_decoupled', got {inference_mode!r}" + ) + self.inference_mode = inference_mode + self.ar_inference_chunk_size = int(ar_inference_chunk_size) + # Number of `sample_step` calls per env tick. step_size>1 advances + # noise levels by 1/(n_rungs*step_size) per sub-step. Mirrors the + # offline staircase_ar_schedule(chunk, step) shape. + self.ar_inference_step_size = int(ar_inference_step_size) + # Classifier-free-guidance scale at inference. 1.0 disables CFG. + self.cfg_scale = float(cfg_scale) + # spatial_rh closed-loop controller: context window + actions committed + # per replan (sp_commit>1 = predict a short chunk open-loop, re-plan + # every sp_commit ticks -> sp_commit-fold fewer diffusion rollouts). + self.sp_n_context = int(sp_n_context) + self.sp_commit = int(sp_commit) + self.sp_n_samples = int(sp_n_samples) + self.device = device or torch.device( + "cuda" if torch.cuda.is_available() else "cpu" + ) + + # outer_stage owns: cond_encoder, inner_stage (backbone), diffusion. + # loss reads ctx.q_state (populated by outer_stage.encode) + batch. + # If no loss is provided, default to DFoTLoss(outer_stage.diffusion). + if loss is None: + loss = DFoTLoss(outer_stage.diffusion) + self.nets = nn.ModuleDict({"outer_stage": outer_stage, "loss": loss}) + self.nets = self.nets.float().to(self.device) + + # Resolve per-embodiment keys via norm_stats (HNet-style). Shared + # base helper (collapse c5): see Algo._resolve_embodiment_keys. + self._resolve_embodiment_keys(norm_stats) + + # ----- Convenience accessors so the inference-path code paths + # (forward_eval, _sample_chunk, _inference_step_ar, _inference_step_chunk) + # can keep referring to ``self.backbone`` / ``self.cond_encoder`` / + # ``self.diffusion`` without going through ``self.nets["outer_stage"]`` + # each call. All three are submodules of the outer_stage. + + @property + def outer_stage(self) -> DFoTOuterStage: + return self.nets["outer_stage"] + + @property + def loss(self) -> nn.Module: + return self.nets["loss"] + + @property + def cond_encoder(self) -> CondEncoderModule: + return self.outer_stage.cond_encoder + + @property + def backbone(self) -> DFoTBackbone: + return self.outer_stage.inner_stage + + @property + def diffusion(self) -> nn.Module: + return self.outer_stage.diffusion + + # Packed-mode metadata that must NOT go through zarr_key_to_keyname + # resolution (these are bookkeeping, not feature tensors). + _PACKED_META_KEYS = ("cu_seqlens", "max_seq_len", "seq_lens") + + # ---- Algo API -------------------------------------------------------- # + + @override + def process_batch_for_training(self, batch): + """Accept both padded ``(B, T, *)`` batches and packed + ``(T_total, *)`` + ``cu_seqlens`` batches.""" + processed = {} + for emb_name, _batch in batch.items(): + emb_id = get_embodiment_id(emb_name) + processed[emb_id] = {} + is_packed = "cu_seqlens" in _batch + + for key, value in _batch.items(): + if is_packed and key in self._PACKED_META_KEYS: + processed[emb_id][key] = value + continue + key_name = self.norm_stats.zarr_key_to_keyname(key, emb_id) + if key_name is not None: + processed[emb_id][key_name] = value + else: + processed[emb_id][key] = value + + processed[emb_id]["_packed"] = is_packed + # Synthesize seq_lens from cu_seqlens for packed batches if the + # collator didn't emit it. Several downstream evaluators + # (``PackedSimEval._infer_n_episodes``, etc.) key off seq_lens + # to find episode boundaries — silently returning 0 episodes + # when it's missing produces no metrics + no videos. + if is_packed and "seq_lens" not in processed[emb_id]: + cu = processed[emb_id].get("cu_seqlens") + if cu is not None and torch.is_tensor(cu): + processed[emb_id]["seq_lens"] = (cu[1:] - cu[:-1]).to(torch.int64) + processed[emb_id] = self.norm_stats.normalize(processed[emb_id], emb_id) + processed[emb_id]["embodiment"] = torch.tensor( + [emb_id], device=self.device, dtype=torch.int64 + ) + for key, value in processed[emb_id].items(): + if isinstance(value, torch.Tensor): + value = value.to(self.device) + if value.is_floating_point(): + value = value.float() + processed[emb_id][key] = value + return processed + + # _build_obs is inherited from Algo (collapse c5 — byte-identical). + + def _encode_cond(self, obs: dict, T: int) -> Optional[torch.Tensor]: + """Encode obs to per-token cond. Honors per-frame obs (no reduction).""" + cond_dict = self.cond_encoder.encode(obs, T) + return cond_dict.get(self.cond_output_key) + + def _encode_cond_packed(self, obs: dict) -> Optional[torch.Tensor]: + """Packed-mode cond. Obs values are (T_total, ...). We fake batch=1 + by ``unsqueeze(0)``-ing each so ``CondEncoderModule.encode`` runs in + its already-per-frame branch (it doesn't broadcast when dim already + matches). Output: (T_total, d_cond) or None.""" + obs_3d = { + k: (v.unsqueeze(0) if torch.is_tensor(v) else v) + for k, v in obs.items() + } + # T_action argument is unused when obs is already per-frame (dim==3 + # for state, dim==5 for image); pass any non-zero placeholder. + cond_dict = self.cond_encoder.encode(obs_3d, T_action=1) + c = cond_dict.get(self.cond_output_key) + if c is None: + return None + if c.dim() == 3 and c.shape[0] == 1: + c = c.squeeze(0) + return c # (T_total, d_cond) + + def _sample_noise_levels(self, shape, device) -> torch.Tensor: + """Per-token random noise level. Discrete -> longs in [0, timesteps); + continuous -> floats in (0, 1).""" + if isinstance(self.diffusion, DiscreteDiffusion): + return torch.randint( + 0, self.diffusion.timesteps, shape, device=device, dtype=torch.long + ) + return torch.rand(shape, device=device).clamp_(1e-5, 1.0 - 1e-5) + + @override + def forward_training(self, batch): + """Refactored training forward: delegates encode/decode + loss to + the outer_stage + loss submodules. + + For each embodiment: + 1. Build a DFoT context with packed/padded mode and obs. + 2. Call outer_stage(batch_emb, ctx) — runs encode -> q_sample, + backbone, decode (writes batch[pred_v]). + 3. Call loss(batch_emb, ctx) — reads pred_v + ctx.q_state, returns + scalar SNR-weighted eps-MSE. + """ + predictions = OrderedDict() + for emb_id, _batch in batch.items(): + ac_key = self.resolved_ac_keys.get(emb_id, "actions") + is_packed = _batch.get("_packed", False) + obs = self._build_obs(_batch, emb_id) + ctx = make_dfot_ctx( + is_packed=is_packed, + action_key=ac_key, + obs=obs, + cu_seqlens=_batch.get("cu_seqlens") if is_packed else None, + max_seqlen=(int(_batch.get("max_seq_len", 0)) or None) if is_packed else None, + ) + self.outer_stage(_batch, ctx) + mse = self.loss(_batch, ctx) + predictions[f"{emb_id}_action_loss"] = mse + return predictions + + @override + def forward_eval(self, batch): + """Returns val-loss + sampled chunks for each embodiment. Sampled + chunks are always single-window (B=1 if packed) at length + ``self.action_horizon`` — packed-mode val skips per-position chunk + sampling for now (rollout drives closed-loop quality via + ``inference_step``).""" + unnorm = {} + backbone = self.backbone + for emb_id, _batch in batch.items(): + ac_key = self.resolved_ac_keys.get(emb_id, "actions") + actions = _batch[ac_key] + is_packed = _batch.get("_packed", False) + obs = self._build_obs(_batch, emb_id) + + # Build the joint diffusion target (= the bundle for obs+action + # outer stages, or just actions for vanilla DFoT). Previously this + # path diffused raw `actions`, which is wrong for obs-action + # variants whose backbone/diffusion are sized to the full bundle. + _eval_ctx = make_dfot_ctx( + is_packed=is_packed, + action_key=ac_key, + obs=obs, + cu_seqlens=_batch.get("cu_seqlens") if is_packed else None, + max_seqlen=(int(_batch.get("max_seq_len", 0)) or None) if is_packed else None, + ) + _build_bundle = getattr(self.outer_stage, "_build_bundle", None) + target = _build_bundle(_batch, _eval_ctx) if _build_bundle is not None else actions + + if is_packed: + T_total = target.shape[0] + cu = _batch["cu_seqlens"] + msl = int(_batch.get("max_seq_len", 0)) or None + cond = self._encode_cond_packed(obs) + k = self._sample_noise_levels((T_total,), target.device) + packed_backbone = _PackedBackboneWrapper(backbone, cu, msl) + _, loss = self.diffusion(packed_backbone, target, k, external_cond=cond) + unnorm[f"emb{emb_id}_loss"] = loss.mean() + # No chunk sampling in packed val — too expensive per episode and + # the closed-loop measure lives in inference_step / sim eval. + else: + B, T, _ = actions.shape + cond = self._encode_cond(obs, T) + k = self._sample_noise_levels((B, T), target.device) + _, loss = self.diffusion(backbone, target, k, external_cond=cond) + unnorm[f"emb{emb_id}_loss"] = loss.mean() + sampled = self._sample_chunk(B, T, cond=cond, device=actions.device) + # _sample_chunk returns the full bundle; pick out the action + # slice so unnormalize sees an action-shaped tensor. + sampled = sampled[..., self.outer_stage.action_slice] + preds = OrderedDict() + preds[ac_key] = sampled + unnorm_actions = self.norm_stats.unnormalize(preds, emb_id) + for key, val in unnorm_actions.items(): + unnorm[f"emb{emb_id}_{key}"] = val + return unnorm + + def _sample_chunk( + self, B: int, T: int, cond: Optional[torch.Tensor], device + ) -> torch.Tensor: + # Build the schedule + run sample() directly so we can plumb + # cfg_scale through. (ddim_sample/ddpm_sample wrappers don't yet + # expose cfg_scale; could be added if needed.) + outer = self.outer_stage + # Prefer ``bundle_shape`` (tuple) when the outer stage has it + # (spatial backbones like ``DFoTSpatialBackbone``); fall back to + # the legacy ``bundle_dim`` (int) for 1D-bundle outer stages. + # ``hasattr(outer, "bundle_shape")`` is always True on the base + # class, so check the tuple-ness explicitly. + b_shape = getattr(outer, "bundle_shape", None) + spatial_mode = isinstance(b_shape, tuple) and len(b_shape) > 1 + if self.sampler not in ("ddpm", "ddim"): + raise ValueError(f"unknown sampler {self.sampler!r}") + discrete_ts = ( + int(self.diffusion.timesteps) + if isinstance(self.diffusion, DiscreteDiffusion) + else None + ) + n_steps = self.sampler_n_steps if self.sampler == "ddim" else ( + discrete_ts or self.sampler_n_steps + ) + from egomimic.models.diffusion.sampling import vanilla_schedule, sample as _sample + sm = vanilla_schedule(n_steps=n_steps, T=T, discrete_timesteps=discrete_ts) + eta = 1.0 if self.sampler == "ddpm" else self.sampler_eta + sample_kwargs = dict( + schedule_matrix=sm, + batch_size=B, + external_cond=cond, + eta=eta, + cfg_scale=self.cfg_scale, + device=device, + ) + if spatial_mode: + sample_kwargs["x_shape"] = b_shape + else: + sample_kwargs["action_dim"] = outer.bundle_dim + return _sample(self.diffusion, self.backbone, **sample_kwargs) + + # compute_losses is inherited from Algo (collapse c7 — the pure + # sum-per-embodiment {emb}_action_loss reducer is now the Algo default). + # log_info is inherited from Algo (collapse c5 — byte-identical default). + + # ---- Sim eval hook ---- # + # + # Two inference modes, selected by ``self.inference_mode`` (config): + # "ar" — DFoT-flavored causal-AR rolling staircase (default). + # Each env step does ONE ``sample_step`` against the AR + # buffer; the front rung commits as the action(s). + # Matches training distribution (per-token random noise). + # "chunk" — legacy chunk-replan: predict ``action_horizon`` actions + # with uniform-per-token DDIM, execute them one at a time, + # replan after the chunk is exhausted. Cheaper but doesn't + # exercise DFoT's per-token-noise capability. + + @torch.no_grad() + def inference_step( + self, obs_zarr: dict, t: int, emb_id: int, T_max=None + ) -> "np.ndarray": + if self.inference_mode == "ar": + return self._inference_step_ar(obs_zarr, t, emb_id) + if self.inference_mode == "chunk": + return self._inference_step_chunk(obs_zarr, t, emb_id) + if self.inference_mode == "spatial_rh": + return self._inference_step_spatial_rh(obs_zarr, t, emb_id) + if self.inference_mode == "spatial_decoupled": + return self._inference_step_spatial_decoupled(obs_zarr, t, emb_id) + if self.inference_mode == "pixel_policy": + return self._inference_step_pixel_policy(obs_zarr, t, emb_id) + if self.inference_mode == "pixel_regress": + return self._inference_step_pixel_regress(obs_zarr, t, emb_id) + if self.inference_mode == "pixel_decoupled": + return self._inference_step_pixel_decoupled(obs_zarr, t, emb_id) + raise ValueError(f"unknown inference_mode {self.inference_mode!r}") + + def _ar_state_init(self, device, external_cond): + """Initialize the AR buffer + staircase geometry AND warm it up so + every slot's actual noise level matches the schedule before any + action is fired. + + Without warmup, the buffer is ``randn`` (all slots at noise 1.0) but + the staircase schedule claims token 0 is at noise 1/K, token 1 at + 2/K, etc. — a mismatch that produces junk for the first K-1 env + ticks. After ``n_rungs - 1`` ``sample_step`` calls, every slot has + been denoised the right number of times and the buffer/schedule + agree. + """ + if self.action_horizon % self.ar_inference_chunk_size != 0: + raise ValueError( + f"action_horizon ({self.action_horizon}) must be divisible by " + f"ar_inference_chunk_size ({self.ar_inference_chunk_size})." + ) + n_rungs = self.action_horizon // self.ar_inference_chunk_size + is_discrete = isinstance(self.diffusion, DiscreteDiffusion) + if is_discrete: + self._ar_unit = max(1, self.diffusion.timesteps // n_rungs) + else: + self._ar_unit = 1.0 / float(n_rungs) + self._ar_discrete = is_discrete + self._ar_buffer = torch.randn( + 1, self.action_horizon, self.outer_stage.bundle_dim, + device=device, dtype=torch.float32, + ) + self._sim_committed_queue = [] + + # ---- Warmup: walk the staircase forward without committing ---- + # We need every slot at the noise level the schedule expects. The + # buffer starts with all slots at noise 1.0 (fully noisy randn). + # The staircase, when first queried, claims slot 0 is at 1/K. We + # need to "demote" each slot's actual noise to match by running + # n_rungs - 1 denoise steps where the schedule entries are shifted + # back by one rung at the top. + # + # Concretely: at warmup step ``w`` (0-indexed, 0..n_rungs-2): + # declared current levels = (rung_idx + (n_rungs-1-w)) / n_rungs, + # clamp(<=1.0); slots beyond the declared front are still 1.0. + # declared next levels = (rung_idx + (n_rungs-2-w)) / n_rungs, + # clamp(<=1.0). + # This walks the schedule from "all slots at 1.0" down to the + # canonical staircase [1/K, 2/K, ..., 1.0] over n_rungs-1 steps. + # On the final canonical step (committing tick 0), the first + # ``inference_step`` call then advances by one more unit, which is + # the correct first-action commit. + if n_rungs > 1: + self._ar_warmup(device, external_cond=external_cond, n_rungs=n_rungs) + + @torch.no_grad() + def _ar_warmup(self, device, external_cond, n_rungs: int): + """Run ``n_rungs - 1`` denoise steps to fill the buffer to the + canonical staircase [1/K, 2/K, ..., 1.0] starting from ``randn``.""" + tok_idx = torch.arange(self.action_horizon, device=device).float() + rung_idx = tok_idx // self.ar_inference_chunk_size # 0..n_rungs-1 + for w in range(n_rungs - 1): + # At warmup step w (0-indexed, 0..n_rungs-2): + # slot i is at rung_idx[i] + (n_rungs - w) entering this step, + # clamped to n_rungs (= level 1.0 ceiling). So at w=0 every + # slot starts at the full-noise ceiling — exactly matching the + # ``randn`` buffer — and slot 0 then denoises by one rung. + # After this step, slot 0 has gone (n_rungs - w)/K -> (n_rungs - w - 1)/K + # (for w=0: 1.0 -> (K-1)/K). + shift_cur = float(n_rungs - w) + shift_nxt = float(n_rungs - 1 - w) + cur_rung = (rung_idx + shift_cur).clamp(max=float(n_rungs)) + nxt_rung = (rung_idx + shift_nxt).clamp(max=float(n_rungs)) + if self._ar_discrete: + cur_levels = (cur_rung * self._ar_unit).long().clamp( + -1, self.diffusion.timesteps - 1 + ) + nxt_levels = (nxt_rung * self._ar_unit).long().clamp( + -1, self.diffusion.timesteps - 1 + ) + else: + cur_levels = (cur_rung * self._ar_unit).clamp(0.0, 1.0) + nxt_levels = (nxt_rung * self._ar_unit).clamp(0.0, 1.0) + self._ar_buffer = sample_step( + self.diffusion, + self.backbone, + x=self._ar_buffer, + current_levels=cur_levels.unsqueeze(0), + next_levels=nxt_levels.unsqueeze(0), + external_cond=external_cond, + cfg_scale=self.cfg_scale, + ) + + def _ar_levels(self, offset: float, device) -> torch.Tensor: + """Per-token noise levels for the staircase at rung-offset ``offset``. + ``offset=0`` = current; ``offset=1`` = after one rung advance. + Tokens within the same chunk share a rung's level.""" + tok_idx = torch.arange(self.action_horizon, device=device).float() + rung_idx = (tok_idx // self.ar_inference_chunk_size) + 1.0 - offset + if self._ar_discrete: + levels = (rung_idx * self._ar_unit).long().clamp( + -1, self.diffusion.timesteps - 1 + ) + else: + levels = (rung_idx * self._ar_unit).clamp(0.0, 1.0) + return levels.unsqueeze(0) # (1, action_horizon) + + @torch.no_grad() + @torch.no_grad() + def _inference_step_pixel_policy(self, obs_zarr, t, emb_id): + """Closed-loop controller for the PIXEL obs+action policy. The action + rides as broadcast channels inside the diffused frame. Pin the last + n_context OBSERVED frames (real RGB + executed-action planes) clean, + denoise the next k frames' [RGB + action] jointly, read the committed + frame's action planes by global-avg-pool -> action. Receding horizon.""" + from egomimic.models.diffusion.sampling import vanilla_schedule + import numpy as _np + + outer = self.outer_stage + diff = self.diffusion + device = next(self.backbone.parameters()).device + ac_key = self.ac_keys[get_embodiment(emb_id).lower()] + n_ctx = max(1, int(getattr(self, "sp_n_context", 1))) + k = max(1, int(getattr(self, "sp_commit", 1))) + n_steps = int(self.sampler_n_steps) + Ci = int(outer._image_channels) + Ca = int(outer._action_channels) + A = int(outer.action_dim) + H = W = int(outer._image_size) + C = Ci + Ca + + if t == 0 or not hasattr(self, "_pp_rgb"): + self._pp_rgb = [] # observed RGB, each (Ci,H,W) in [0,1] + self._pp_act = [] # executed NORMALIZED actions, each (A,) + self._pp_queue = [] # pending committed unnorm actions + + img = obs_zarr[outer.image_key].float().to(device) + if img.max() > 1.5: + img = img / 255.0 + if img.dim() == 4: + img = img[0] + self._pp_rgb.append(img) + + if self._pp_queue: + return self._pp_queue.pop(0) + + def act_plane(a): + return a[:Ca].reshape(Ca, 1, 1).expand(Ca, H, W) + + T = n_ctx + k + nrgb = len(self._pp_rgb) + nact = len(self._pp_act) + + ctx_frames = [] + for j in range(n_ctx): + ri = max(0, nrgb - n_ctx + j) # most-recent n_ctx OBSERVED frames + rgb = self._pp_rgb[ri] + ai = nact - n_ctx + j # the action that produced that obs + if 0 <= ai < nact: + a = self._pp_act[ai] + else: + a = torch.zeros(A, device=device) + try: + _st = obs_zarr["state_agent_obj"] + _st = _st[0] if _st.dim() == 2 else _st + _xy = _st[:A].float().to(device).unsqueeze(0) + a = self.norm_stats.normalize({ac_key: _xy}, emb_id)[ac_key][0][:A] + except Exception: + pass + ctx_frames.append(torch.cat([rgb, act_plane(a)], dim=0)) + ctx_stack = torch.stack(ctx_frames, dim=0).unsqueeze(0) # (1,n_ctx,C,H,W) + + dts = int(diff.timesteps) if isinstance(diff, DiscreteDiffusion) else None + clean = -1 if dts is not None else 0.0 + sched = vanilla_schedule(n_steps, T, discrete_timesteps=dts).to(device).clone() + sched[:, :n_ctx] = clean + + x = torch.randn(1, T, C, H, W, device=device) + x[:, :n_ctx] = ctx_stack + for s in range(sched.shape[0] - 1): + klev = sched[s].clamp_min(0).long().unsqueeze(0) + v = self.backbone(x, klev, external_cond=None) + x = self._struct_ddim_step(diff, x, v, sched[s], sched[s + 1]) + x[:, :n_ctx] = ctx_stack + + pred_planes = x[0, n_ctx:n_ctx + k, Ci:Ci + Ca] # (k,Ca,H,W) + pred_norm = pred_planes.mean(dim=(2, 3))[:, :A] # (k,A) global-avg-pool + for j in range(k): + self._pp_act.append(pred_norm[j].detach()) + unnorm = self.norm_stats.unnormalize({ac_key: pred_norm}, emb_id)[ac_key] + unnorm_np = unnorm.detach().cpu().numpy() + for row in unnorm_np[1:]: + self._pp_queue.append(row.reshape(-1).astype(_np.float32)) + return unnorm_np[0].reshape(-1).astype(_np.float32) + + @torch.no_grad() + def _inference_step_pixel_regress(self, obs_zarr, t, emb_id): + """Closed-loop controller for Design B (regression). Pin the last + n_context observed RGB frames clean, denoise the next k RGB frames, + then read the action off each predicted frame via the outer stage's + conv ``action_head``. Receding horizon.""" + from egomimic.models.diffusion.sampling import vanilla_schedule + import numpy as _np + + outer = self.outer_stage + diff = self.diffusion + device = next(self.backbone.parameters()).device + ac_key = self.ac_keys[get_embodiment(emb_id).lower()] + n_ctx = max(1, int(getattr(self, "sp_n_context", 1))) + k = max(1, int(getattr(self, "sp_commit", 1))) + n_steps = int(self.sampler_n_steps) + Ci = int(outer._image_channels) + H = W = int(outer._image_size) + + if t == 0 or not hasattr(self, "_pr_rgb"): + self._pr_rgb = [] + self._pr_queue = [] + + img = obs_zarr[outer.image_key].float().to(device) + if img.max() > 1.5: + img = img / 255.0 + if img.dim() == 4: + img = img[0] + self._pr_rgb.append(img) + + if self._pr_queue: + return self._pr_queue.pop(0) + + T = n_ctx + k + nrgb = len(self._pr_rgb) + ctx_stack = torch.stack( + [self._pr_rgb[max(0, nrgb - n_ctx + j)] for j in range(n_ctx)], dim=0 + ).unsqueeze(0) + + dts = int(diff.timesteps) if isinstance(diff, DiscreteDiffusion) else None + clean = -1 if dts is not None else 0.0 + sched = vanilla_schedule(n_steps, T, discrete_timesteps=dts).to(device).clone() + sched[:, :n_ctx] = clean + + x = torch.randn(1, T, Ci, H, W, device=device) + x[:, :n_ctx] = ctx_stack + for sidx in range(sched.shape[0] - 1): + klev = sched[sidx].clamp_min(0).long().unsqueeze(0) + v = self.backbone(x, klev, external_cond=None) + x = self._struct_ddim_step(diff, x, v, sched[sidx], sched[sidx + 1]) + x[:, :n_ctx] = ctx_stack + + pred_frames = x[0, n_ctx:n_ctx + k] # (k, Ci, H, W) predicted clean frames + pred_norm = outer.action_head(pred_frames) # (k, A) + unnorm = self.norm_stats.unnormalize({ac_key: pred_norm}, emb_id)[ac_key] + unnorm_np = unnorm.detach().cpu().numpy() + for row in unnorm_np[1:]: + self._pr_queue.append(row.reshape(-1).astype(_np.float32)) + return unnorm_np[0].reshape(-1).astype(_np.float32) + + def _inference_step_ar( + self, obs_zarr: dict, t: int, emb_id: int + ) -> "np.ndarray": + """Causal-AR rolling-staircase inference. One sample_step per env + tick, one (or ``ar_inference_chunk_size``-many) action(s) committed + per call. Buffer carries across calls.""" + embodiment_name = get_embodiment(emb_id).lower() + device = next(self.backbone.parameters()).device + ac_key = self.ac_keys[embodiment_name] + + # Encode current obs into per-call cond. Broadcast across buffer + # tokens inside the backbone (per-token AdaLN). + obs_norm = self.norm_stats.normalize(obs_zarr, emb_id) + cond = self._encode_cond(obs_norm, self.action_horizon) + if cond is not None and cond.dim() == 3: + cond = cond[:, 0] # (1, cond_dim) — backbone broadcasts to T + + # Reset + warm up on episode start. Warmup uses the t=0 obs cond + # for all warmup denoise steps (no future obs available); this is + # the correct online-AR semantics — every future env tick uses its + # own obs. + if t == 0 or not hasattr(self, "_ar_buffer"): + self._ar_state_init(device, external_cond=cond) + + # If we already have committed actions ready, just pop one. + if self._sim_committed_queue: + return self._sim_committed_queue.pop(0) + + # One denoise step on the buffer. + cur_levels = self._ar_levels(offset=0.0, device=device) + nxt_levels = self._ar_levels(offset=1.0, device=device) + self._ar_buffer = sample_step( + self.diffusion, + self.backbone, + x=self._ar_buffer, + current_levels=cur_levels, + next_levels=nxt_levels, + external_cond=cond, + cfg_scale=self.cfg_scale, + ) + + # Commit front rung, slide buffer, push fresh noisy rung at the back. + chunk = self.ar_inference_chunk_size + committed_norm = self._ar_buffer[:, :chunk, :].clone() # (1, chunk, bundle_dim) + new_back = torch.randn( + 1, chunk, self.outer_stage.bundle_dim, + device=device, dtype=torch.float32, + ) + self._ar_buffer = torch.cat( + [self._ar_buffer[:, chunk:, :], new_back], dim=1 + ) + + # For joint obs+action bundles, slice out the action portion before + # returning to the env. For vanilla DFoT this is a no-op (action_slice + # spans the full trailing dim). + committed_actions = committed_norm[..., self.outer_stage.action_slice] + committed_world = self.norm_stats.unnormalize( + {ac_key: committed_actions.squeeze(0)}, emb_id + )[ac_key] + committed_np = committed_world.detach().cpu().numpy() + for row in committed_np[1:]: + self._sim_committed_queue.append(row.reshape(-1).astype(np.float32)) + return committed_np[0].reshape(-1).astype(np.float32) + + @torch.no_grad() + def _inference_step_chunk( + self, obs_zarr: dict, t: int, emb_id: int + ) -> "np.ndarray": + """Legacy chunk-replan inference (uniform per-token noise, plan + + execute action_horizon steps before replanning). Cheaper but does + NOT exercise DFoT's per-token-noise capability.""" + embodiment_name = get_embodiment(emb_id).lower() + device = next(self.backbone.parameters()).device + ac_key = self.ac_keys[embodiment_name] + if t == 0: + self._sim_state = {"chunk": None, "chunk_idx": 0} + state = self._sim_state + + if state["chunk"] is None or state["chunk_idx"] >= self.action_horizon: + obs_norm = self.norm_stats.normalize(obs_zarr, emb_id) + cond = self._encode_cond(obs_norm, self.action_horizon) + sampled = self._sample_chunk( + B=1, T=self.action_horizon, cond=cond, device=device + ) + # Slice action portion out of the (potentially joint) bundle + # before unnormalizing. For vanilla DFoT action_slice is full. + sampled_actions = sampled[..., self.outer_stage.action_slice] + chunk_world = self.norm_stats.unnormalize( + {ac_key: sampled_actions.squeeze(0)}, emb_id + )[ac_key] + state["chunk"] = chunk_world.detach() + state["chunk_idx"] = 0 + + idx = state["chunk_idx"] + action_world = state["chunk"][idx] + state["chunk_idx"] = idx + 1 + return action_world.cpu().numpy().reshape(-1).astype(np.float32) + + # ------------------------------------------------------------------ # + # Closed-loop controller for the 2D spatial obs+action policy. + # ------------------------------------------------------------------ # + def _struct_ddim_step(self, diff, x, v, cur, nxt): + """One eta=0 DDIM step from a v-prediction (discrete diffusion). + ``cur``/``nxt`` are (T,) per-token levels; ``x`` is (1, T, *trailing). + Mirrors DFoTPolicyActionEval._ddim_from_v exactly.""" + T = x.shape[1] + pad = x.dim() - 2 + kBT = cur.clamp_min(0).long().unsqueeze(0) + x0 = diff.predict_start_from_v(x, kBT, v) + eps = diff.predict_noise_from_v(x, kBT, v) + an = diff.alphas_cumprod[nxt.clamp_min(0).long()] + an = torch.where(nxt < 0, torch.ones_like(an), an) + an = an.reshape(1, T, *([1] * pad)) + c = (1.0 - an).clamp_min(0.0).sqrt() + return x0 * an.sqrt() + eps * c + + @torch.no_grad() + def _inference_step_pixel_decoupled(self, obs_zarr, t, emb_id): + """Closed-loop controller for the PIXEL DECOUPLED-action policy, with + optional CHUNKING (sp_commit>1). Pin the current RGB (+ last n_context) + CLEAN; for a chunk, append (sp_commit-1) FUTURE frames whose obs AND + action are denoised (model predicts future obs as a world-model), read + the action at the current frame + the future frames, commit the whole + chunk open-loop. Action never a clean input (no copy), no offset. + sp_n_samples averages K diffusion samples (variance reduction).""" + from egomimic.models.diffusion.sampling import vanilla_schedule + + outer = self.outer_stage + diff = self.diffusion + device = next(self.backbone.parameters()).device + ac_key = self.ac_keys[get_embodiment(emb_id).lower()] + n_ctx = max(1, int(getattr(self, "sp_n_context", 1))) + k = max(1, int(getattr(self, "sp_commit", 1))) + n_steps = int(self.sampler_n_steps) + n_samp = max(1, int(getattr(self, "sp_n_samples", 1))) + A = int(outer.action_dim) + Ci = int(outer._image_channels) + H = W = int(outer._image_size) + + if t == 0 or not hasattr(self, "_pd_rgb"): + self._pd_rgb = [] + self._pd_queue = [] + + img = obs_zarr[outer.image_key].float().to(device) + if img.max() > 1.5: + img = img / 255.0 + if img.dim() == 4: + img = img[0] + self._pd_rgb.append(img) + + if self._pd_queue: + return self._pd_queue.pop(0) + + L = len(self._pd_rgb) + idx = [max(0, L - n_ctx + i) for i in range(n_ctx)] + rgb_ctx = torch.stack([self._pd_rgb[i] for i in idx], dim=0).unsqueeze(0) + + T = n_ctx + (k - 1) + dts = int(diff.timesteps) if isinstance(diff, DiscreteDiffusion) else None + sched = vanilla_schedule(n_steps, T, discrete_timesteps=dts).to(device) + obs_sched = sched.clone() + obs_sched[:, :n_ctx] = 0 # context obs CLEAN (re-pinned each step) + + ctx_b = rgb_ctx.expand(n_samp, -1, -1, -1, -1) + x_obs = torch.randn(n_samp, T, Ci, H, W, device=device) + x_obs[:, :n_ctx] = ctx_b + x_act = torch.randn(n_samp, T, A, device=device) + for st in range(sched.shape[0] - 1): + o_lev = obs_sched[st].unsqueeze(0).expand(n_samp, -1) + a_lev = sched[st].unsqueeze(0).expand(n_samp, -1) + v_img, v_act = self.backbone( + x_obs, o_lev, external_cond=None, action=x_act, action_noise_levels=a_lev, + ) + x_obs = self._struct_ddim_step(diff, x_obs, v_img, obs_sched[st], obs_sched[st + 1]) + x_obs[:, :n_ctx] = ctx_b + x_act = self._struct_ddim_step(diff, x_act, v_act, sched[st], sched[st + 1]) + + chunk = x_act[:, n_ctx - 1:].mean(0) # (k, A): current + (k-1) future actions + unnorm = self.norm_stats.unnormalize({ac_key: chunk}, emb_id)[ac_key] + unnorm_np = unnorm.detach().cpu().numpy() + for row in unnorm_np[1:]: + self._pd_queue.append(row.reshape(-1).astype(np.float32)) + return unnorm_np[0].reshape(-1).astype(np.float32) + + def _inference_step_spatial_rh( + self, obs_zarr: dict, t: int, emb_id: int + ) -> "np.ndarray": + """Reactive receding-horizon controller for the 2D spatial obs+action + policy (``SpatialObsActionPolicyDFoTOuterStage`` + dual-stream DiT3D). + + The backbone uses ONE noise level per frame (shared by the latent and + its action token), so we cannot pin the current image clean while + predicting its action. Instead: pin the last ``n_context`` OBSERVED + frames (real latent + executed action) clean, condition on the REAL + current state via external_cond, and predict the next frame's + (imagined latent, action); commit the action. Real observations enter + as clean context one tick later. This is the closed-loop form of the + validated DFoTPolicyActionEval._rollout. + """ + from egomimic.models.diffusion.sampling import vanilla_schedule + + outer = self.outer_stage + diff = self.diffusion + device = next(self.backbone.parameters()).device + embodiment_name = get_embodiment(emb_id).lower() + ac_key = self.ac_keys[embodiment_name] + n_ctx = int(self.sp_n_context) + k = max(1, int(self.sp_commit)) # actions committed per replan + n_steps = int(self.sampler_n_steps) + + if t == 0 or not hasattr(self, "_sp_lat"): + self._sp_lat = [] # observed clean latents, each (1, C, H, W) + self._sp_state = [] # observed states, each (1, sd) normalized + self._sp_act = [] # executed NORMALIZED actions, each (A,) + self._sp_queue = [] # pending committed unnorm actions (np, A) + + # ---- encode current obs (becomes a clean context frame next tick) ---- + obs_norm = self.norm_stats.normalize(obs_zarr, emb_id) + state = torch.cat( + [obs_norm[kk] for kk in outer.bundle_obs_keys], dim=-1 + ).float().to(device) # (1, sd) + img = obs_zarr[outer.image_key].float().to(device) + if img.max() > 1.5: + img = img / 255.0 + mu, _ = outer.vae.encode(img) # (1, C, H, W) + lat = outer.normalize_latent(mu) # (1, C, H, W) + self._sp_lat.append(lat) + self._sp_state.append(state) + + # Mid-chunk: return the next already-committed action (no replan). + if self._sp_queue: + return self._sp_queue.pop(0) + + # ---- replan: predict the next k (latent, action) frames ---- + C, H, W = outer.bundle_shape + A = outer._action_dim + n_done = len(self._sp_act) # executed actions so far + ctx_idx = [max(0, n_done - n_ctx + i) for i in range(n_ctx)] + T = n_ctx + k + + latent_ctx = torch.cat([self._sp_lat[i] for i in ctx_idx], dim=0) # (n_ctx, C, H, W) + # cond: context-frame states + the REAL current state for each predicted frame + state_seq = [self._sp_state[i] for i in ctx_idx] + [self._sp_state[-1]] * k + cond = outer.state_only_proj(torch.cat(state_seq, dim=0)).unsqueeze(0) # (1, T, proj) + act_ctx = torch.stack( + [self._sp_act[i] if i < len(self._sp_act) else torch.zeros(A, device=device) + for i in ctx_idx], dim=0) # (n_ctx, A) + + dts = int(diff.timesteps) if isinstance(diff, DiscreteDiffusion) else None + clean = -1 if dts is not None else 0.0 + sched = vanilla_schedule(n_steps, T, discrete_timesteps=dts).to(device).clone() + sched[:, :n_ctx] = clean + + x_lat = torch.randn(1, T, C, H, W, device=device) + x_lat[:, :n_ctx] = latent_ctx.unsqueeze(0) + x_act = torch.randn(1, T, A, device=device) + x_act[:, :n_ctx] = act_ctx.unsqueeze(0) + for s in range(sched.shape[0] - 1): + klev = sched[s].clamp_min(0).long().unsqueeze(0) + v_lat, v_act = self.backbone(x_lat, klev, external_cond=cond, action=x_act) + x_lat = self._struct_ddim_step(diff, x_lat, v_lat, sched[s], sched[s + 1]) + x_act = self._struct_ddim_step(diff, x_act, v_act, sched[s], sched[s + 1]) + x_lat[:, :n_ctx] = latent_ctx.unsqueeze(0) + x_act[:, :n_ctx] = act_ctx.unsqueeze(0) + + pred_norm = x_act[0, n_ctx:n_ctx + k] # (k, A) normalized predicted actions + for j in range(k): + self._sp_act.append(pred_norm[j].detach()) # record executed (normalized) + unnorm = self.norm_stats.unnormalize( + {ac_key: pred_norm}, emb_id + )[ac_key] # (k, A) absolute frame + unnorm_np = unnorm.detach().cpu().numpy() + for row in unnorm_np[1:]: + self._sp_queue.append(row.reshape(-1).astype(np.float32)) + return unnorm_np[0].reshape(-1).astype(np.float32) + + @torch.no_grad() + def _inference_step_spatial_decoupled( + self, obs_zarr: dict, t: int, emb_id: int + ) -> "np.ndarray": + """Closed-loop controller for the DECOUPLED-action policy + (``decouple_action_noise=True``). The action token has its own noise + level, so we CAN pin the real current image clean while predicting its + action -> reactive, and the action is never a clean input (no copy). + + Each tick: pin the last ``n_context`` OBSERVED latents clean (incl. the + current frame), keep ALL action tokens noised, denoise the action stream + on its own schedule conditioned on the clean obs + state, and commit the + action predicted at the current (last) frame. Causal: no future frames. + """ + from egomimic.models.diffusion.sampling import vanilla_schedule + + outer = self.outer_stage + diff = self.diffusion + device = next(self.backbone.parameters()).device + embodiment_name = get_embodiment(emb_id).lower() + ac_key = self.ac_keys[embodiment_name] + n_ctx = int(self.sp_n_context) + n_steps = int(self.sampler_n_steps) + + if t == 0 or not hasattr(self, "_sd_lat"): + self._sd_lat = [] # observed clean latents, each (1, C, H, W) + self._sd_state = [] # observed states, each (1, sd) normalized + + obs_norm = self.norm_stats.normalize(obs_zarr, emb_id) + state = torch.cat( + [obs_norm[kk] for kk in outer.bundle_obs_keys], dim=-1 + ).float().to(device) + img = obs_zarr[outer.image_key].float().to(device) + if img.max() > 1.5: + img = img / 255.0 + mu, _ = outer.vae.encode(img) + self._sd_lat.append(outer.normalize_latent(mu)) + self._sd_state.append(state) + + C, H, W = outer.bundle_shape + A = outer._action_dim + L = len(self._sd_lat) + idx = [max(0, L - n_ctx + i) for i in range(n_ctx)] # current frame is last + T = n_ctx + latent_ctx = torch.cat([self._sd_lat[i] for i in idx], dim=0).unsqueeze(0) # (1,T,C,H,W) CLEAN + cond = outer.state_only_proj( + torch.cat([self._sd_state[i] for i in idx], dim=0) + ).unsqueeze(0) # (1,T,proj) + + dts = int(diff.timesteps) if isinstance(diff, DiscreteDiffusion) else None + clean = -1 if dts is not None else 0.0 + # obs given clean for all frames; action denoised on its own schedule. + obs_levels = torch.full( + (1, T), clean, device=device, + dtype=torch.long if dts is not None else torch.float32, + ) + act_sched = vanilla_schedule(n_steps, T, discrete_timesteps=dts).to(device) + + x_lat = latent_ctx # clean obs, never stepped + x_act = torch.randn(1, T, A, device=device) + for s in range(act_sched.shape[0] - 1): + _, v_act = self.backbone( + x_lat, obs_levels, external_cond=cond, action=x_act, + action_noise_levels=act_sched[s].unsqueeze(0), + ) + x_act = self._struct_ddim_step(diff, x_act, v_act, act_sched[s], act_sched[s + 1]) + + pred_norm = x_act[0, -1] # action at the current frame + unnorm = self.norm_stats.unnormalize( + {ac_key: pred_norm.unsqueeze(0)}, emb_id + )[ac_key] + return unnorm.squeeze(0).detach().cpu().numpy().reshape(-1).astype(np.float32) diff --git a/egomimic/algo/diffusion/outer_stages/__init__.py b/egomimic/algo/diffusion/outer_stages/__init__.py new file mode 100644 index 000000000..a129a7914 --- /dev/null +++ b/egomimic/algo/diffusion/outer_stages/__init__.py @@ -0,0 +1,47 @@ +"""DFoT per-variant OUTER-STAGES (DESIGN.md §2 +``egomimic/algo/diffusion/outer_stages/``). + +The :class:`OuterStage` subclasses that wire the diffusion backbone + noise +process into each DFoT training variant (spatial / pixel / obs-action / video / +policy / decoupled). Relocated here from ``egomimic/algo/dfot`` in DESIGN.md +step 7 (``git mv``, no behaviour change); the base :class:`DFoTOuterStage` lives +in ``outer_stage.py`` and the rest subclass it. +""" + +from egomimic.algo.diffusion.outer_stages.outer_stage import ( + DFoTOuterStage, + make_dfot_ctx, +) +from egomimic.algo.diffusion.outer_stages.image_spatial_outer_stage import ( + ImageSpatialDFoTOuterStage, +) +from egomimic.algo.diffusion.outer_stages.obs_action_image_outer_stage import ( + ObsActionImageDFoTOuterStage, +) +from egomimic.algo.diffusion.outer_stages.obs_action_outer_stage import ( + ObsActionDFoTOuterStage, +) +from egomimic.algo.diffusion.outer_stages.pixel_obs_action_outer_stage import ( # noqa: E501 + PixelObsActionDFoTOuterStage, +) +from egomimic.algo.diffusion.outer_stages.pixel_spatial_outer_stage import ( + PixelSpatialDFoTOuterStage, +) +from egomimic.algo.diffusion.outer_stages.pixel_video_outer_stage import ( + PixelVideoDFoTOuterStage, +) +from egomimic.algo.diffusion.outer_stages.spatial_obs_action_policy_outer_stage import ( # noqa: E501 + SpatialObsActionPolicyDFoTOuterStage, +) + +__all__ = [ + "DFoTOuterStage", + "make_dfot_ctx", + "ImageSpatialDFoTOuterStage", + "ObsActionImageDFoTOuterStage", + "ObsActionDFoTOuterStage", + "PixelObsActionDFoTOuterStage", + "PixelSpatialDFoTOuterStage", + "PixelVideoDFoTOuterStage", + "SpatialObsActionPolicyDFoTOuterStage", +] diff --git a/egomimic/algo/diffusion/outer_stages/image_spatial_outer_stage.py b/egomimic/algo/diffusion/outer_stages/image_spatial_outer_stage.py new file mode 100644 index 000000000..28a919ff5 --- /dev/null +++ b/egomimic/algo/diffusion/outer_stages/image_spatial_outer_stage.py @@ -0,0 +1,451 @@ +"""``ImageSpatialDFoTOuterStage`` — DFoT world-model with spatial latent. + +Plan-A variant of the obs+action+image diffusion-forcing pipeline: +**diffuses ONLY the image latent** (kept in its full ``(C, H, W)`` +spatial form, NOT flattened). State + action enter the model as +``external_cond`` via a small MLP, broadcast across patches by the +``DFoTSpatialBackbone``'s AdaLN. + +Sacrifice vs ``ObsActionImageDFoTOuterStage`` (the flat-bundle variant): +this stage CANNOT produce predicted actions — action prediction is +deferred to a follow-up "action-as-extra-token" hybrid. The trade-off +is the spatial inductive bias that ``DiT3D``-style architectures use +to actually make image prediction work (see Optimize.md / the plan-A +write-up for the architectural reasoning). + +Bundle shape: ``(latent_channels, latent_h, latent_w)`` per step. The +algo's sampler / AR code reads ``outer_stage.bundle_shape`` (a tuple) +to allocate tensors; vanilla 1D-bundle stages still use +``outer_stage.bundle_dim`` (an int). Both code paths are supported. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import List, Optional + +import torch +import torch.nn as nn + +from egomimic.algo.diffusion.outer_stages.outer_stage import DFoTOuterStage +from egomimic.algo.diffusion.vae_algo import load_pretrained_vae + + +class ImageSpatialDFoTOuterStage(DFoTOuterStage): + """Spatial-image-only DFoT outer stage. + + Args: + cond_encoder: kept on the class for ABI parity, but typically + unused (pass an empty CondEncoderModule with no obs_specs + and no img_encoders). state + action conditioning is done + internally by ``_state_action_to_cond``. + backbone: a ``DFoTSpatialBackbone``. Its ``cond_dim`` MUST equal + ``state_action_proj_dim`` (this stage's MLP output width). + diffusion: continuous / discrete diffusion module. + bundle_obs_keys: list of obs keys to include in the external_cond + (typically ``["state_agent_obj"]``). + bundle_obs_dims: per-key feature widths matching ``bundle_obs_keys``. + vae_checkpoint_path: path to the frozen ImageVAE Lightning ckpt. + image_key: which obs key carries the image. + action_dim: ACTION feature width (for the MLP input). + state_action_proj_dim: width of the projected ``external_cond`` + fed to the backbone. Must match backbone's ``cond_dim``. + cond_output_key: ABI parity; unused. + """ + + def __init__( + self, + action_dim: int, + cond_encoder, + backbone, + diffusion, + bundle_obs_keys: List[str], + bundle_obs_dims: List[int], + vae_checkpoint_path: str, + image_key: str = "front_img_1", + state_action_proj_dim: int = 128, + cond_output_key: str = "fused_cond", + ): + # We don't call the ObsActionDFoTOuterStage constructor's + # bundle-width sanity check (which assumes 1D bundle). Go + # straight to DFoTOuterStage's __init__. + super().__init__( + action_dim=action_dim, + cond_encoder=cond_encoder, + backbone=backbone, + diffusion=diffusion, + cond_output_key=cond_output_key, + ) + + vae = load_pretrained_vae(vae_checkpoint_path) + if not getattr(vae, "spatial_latent", False): + raise NotImplementedError( + "ImageSpatialDFoTOuterStage requires a VAE with spatial_latent=True." + ) + self.vae = vae + self.image_key = str(image_key) + self._latent_shape = ( + int(vae.latent_channels), + int(vae.bottleneck_size), + int(vae.bottleneck_size), + ) + self.bundle_obs_keys = list(bundle_obs_keys) + self.bundle_obs_dims = [int(d) for d in bundle_obs_dims] + self._state_dim_total = sum(self.bundle_obs_dims) + self._action_dim = int(action_dim) + self._state_action_proj_dim = int(state_action_proj_dim) + + # External-cond MLP: concat(state, action) -> state_action_proj_dim. + self.state_action_proj = nn.Sequential( + nn.Linear(self._state_dim_total + self._action_dim, self._state_action_proj_dim), + nn.SiLU(), + nn.Linear(self._state_action_proj_dim, self._state_action_proj_dim), + ) + + # Backbone sanity-check: its cond_dim should match the projector. + bb_cond_dim = int(getattr(backbone, "cond_dim", -1)) + if bb_cond_dim != self._state_action_proj_dim: + raise ValueError( + f"backbone.cond_dim ({bb_cond_dim}) must equal " + f"state_action_proj_dim ({self._state_action_proj_dim})." + ) + + # Latent normalization: center and scale VAE latents so the + # diffusion prior is standard Gaussian. Without this, the VAE's + # per-channel bias causes sampled latents to decode to near-white. + latent_stats_path = getattr(vae, "_latent_stats_path", None) + if latent_stats_path is None: + # Default: try to load from the same dir as the VAE ckpt. + import os + stats_path = os.path.join( + os.path.dirname(vae_checkpoint_path), "vae_latent_stats.json" + ) + if os.path.exists(stats_path): + latent_stats_path = stats_path + if latent_stats_path is not None: + import json + with open(latent_stats_path) as f: + ls = json.load(f) + self.register_buffer( + "_latent_mean", + torch.tensor(ls["latent_mean"]).float().reshape(1, -1, 1, 1), + ) + self.register_buffer( + "_latent_std", + torch.tensor(ls["latent_std"]).float().reshape(1, -1, 1, 1), + ) + else: + self._latent_mean = None + self._latent_std = None + + # ------------------------------------------------------------------ + # Shape advertised to the sampler / AR allocators. + # ------------------------------------------------------------------ + + @property + def bundle_shape(self) -> tuple: + """Spatial trailing dims of the diffusion target. Sampler builds + ``(B, T, *bundle_shape)``.""" + return self._latent_shape + + @property + def bundle_dim(self) -> int: + """Backwards compatibility: total numel of bundle_shape. The algo + and sampler should prefer ``bundle_shape`` when present, but a + few legacy code paths still read ``bundle_dim``.""" + c, h, w = self._latent_shape + return c * h * w + + @property + def action_slice(self) -> slice: + """No action in this bundle — return an empty slice to signal + downstream code that action extraction from the bundle is N/A. + Callers that want predicted actions from this model should be + upgraded to a different stage (the action-as-extra-token + hybrid follow-up).""" + return slice(0, 0) + + # ------------------------------------------------------------------ + # External cond from state + action. + # ------------------------------------------------------------------ + + def _state_action_to_cond( + self, batch: dict, ctx: SimpleNamespace + ) -> Optional[torch.Tensor]: + """Per-step ``(state || action)`` -> projected ``external_cond``. + + Packed: returns ``(T_total, state_action_proj_dim)``. + Padded: returns ``(B, T, state_action_proj_dim)``. + """ + pieces = [] + for key in self.bundle_obs_keys: + if key not in ctx.obs: + raise KeyError( + f"obs key '{key}' required in external_cond but missing " + f"from ctx.obs (keys: {list(ctx.obs.keys())})." + ) + pieces.append(ctx.obs[key]) + pieces.append(batch[ctx.action_key]) + concat = torch.cat(pieces, dim=-1) + return self.state_action_proj(concat) + + # ------------------------------------------------------------------ + # Build VAE latent target (spatial). + # ------------------------------------------------------------------ + + def normalize_latent(self, z: torch.Tensor) -> torch.Tensor: + if self._latent_mean is not None: + return (z - self._latent_mean.to(z.device)) / self._latent_std.to(z.device) + return z + + def denormalize_latent(self, z: torch.Tensor) -> torch.Tensor: + if self._latent_mean is not None: + return z * self._latent_std.to(z.device) + self._latent_mean.to(z.device) + return z + + # ------------------------------------------------------------------ + # Video-rollout hook (COMBINE A — decode-on-outer-stage). + # + # This stage owns: building per-step (state,action) external_cond, the + # CONDITIONAL chunk/AR sampler (with optional GT-context anchoring), and + # the VAE decode of the predicted spatial latent. Moved byte-for-byte + # from the old ``eval_dfot_spatial_video_rollout.DFoTSpatialVideoRolloutEval`` + # (``_rollout`` + ``_build_cond_seq`` + per-episode body). Side-by-side + # [GT|pred] panel; per-step MSE vs the per-timestep-aligned GT frames. + # ------------------------------------------------------------------ + + video_metric_prefix = "spatial" + video_panel = "sidebyside" + video_has_extra_metrics = False + + @torch.no_grad() + def _rollout_latent( + self, ev, algo, cond_seq: torch.Tensor, device, context_latent=None, + batch_size: int = 1, + ) -> torch.Tensor: + """Returns predicted latent ``(batch_size, T, C, H, W)``. + + If ``context_latent`` (1, n_ctx, C, H, W) is given, the first n_ctx + latents are seeded from GT and pinned clean across all sampling steps + (anchored conditional rollout — predict the continuation given the real + first frame[s] + the action sequence). Else the whole sequence is + generated from noise conditioned only on the actions.""" + from egomimic.models.diffusion.diffusion.discrete_diffusion import DiscreteDiffusion + from egomimic.models.diffusion.sampling import ( + sample as _sample, + staircase_ar_schedule, + vanilla_schedule, + ) + + bundle_shape = self.bundle_shape + T = cond_seq.shape[1] + discrete_ts = ( + int(algo.diffusion.timesteps) + if isinstance(algo.diffusion, DiscreteDiffusion) + else None + ) + if ev.mode == "chunk": + schedule = vanilla_schedule( + n_steps=ev.n_chunk_steps, T=T, discrete_timesteps=discrete_ts, + ).to(device) + else: + schedule = staircase_ar_schedule( + T=T, chunk_size=ev.ar_chunk_size, step_size=ev.ar_step_size, + discrete_timesteps=discrete_ts, + ).to(device) + + if context_latent is None: + return _sample( + algo.diffusion, algo.backbone, schedule_matrix=schedule, + x_shape=bundle_shape, batch_size=batch_size, + external_cond=cond_seq, cfg_scale=ev.cfg_scale, device=device, + ) + + from egomimic.eval.dfot._sampling import anchored_ddim_rollout + return anchored_ddim_rollout( + algo.diffusion, algo.backbone, schedule=schedule, + context=context_latent, total_T=T, trailing_shape=bundle_shape, + device=device, batch_size=batch_size, external_cond=cond_seq, + discrete_ts=discrete_ts, cfg_scale=ev.cfg_scale, + ) + + def _build_cond_seq( + self, algo, _batch: dict, emb_id: int, start: int, T: int + ) -> torch.Tensor: + """``(1, T, state_action_proj_dim)`` projected cond over a slice + of the val batch starting at frame ``start``.""" + ac_key = algo.resolved_ac_keys[emb_id] + pieces = [] + for key in self.bundle_obs_keys: + pieces.append(_batch[key][start : start + T]) + pieces.append(_batch[ac_key][start : start + T]) + concat = torch.cat(pieces, dim=-1) # (T, state_dim + action_dim) + cond = self.state_action_proj(concat) # (T, cond_dim) + return cond.unsqueeze(0) # (1, T, cond_dim) + + @torch.no_grad() + def rollout_video_episode( + self, ev, algo, _batch, emb_id, ep_idx, ep_start, ep_len, device + ) -> SimpleNamespace: + """Per-episode world-model video rollout for the spatial stage.""" + imgs = _batch[ev.image_key] + is_packed = _batch.get("_packed", False) + T_rollout = min(ev.rollout_steps, ep_len) + + if is_packed: + start = ep_start + cond_seq = self._build_cond_seq( + algo, _batch, emb_id, start, T_rollout, + ) + gt_seq = imgs[start : start + T_rollout] + else: + # Padded: take row ep_idx, first T_rollout frames. + ac_key = algo.resolved_ac_keys[emb_id] + pieces = [] + for key in self.bundle_obs_keys: + pieces.append(_batch[key][ep_idx, :T_rollout]) + pieces.append(_batch[ac_key][ep_idx, :T_rollout]) + concat = torch.cat(pieces, dim=-1) + cond_seq = self.state_action_proj(concat).unsqueeze(0) + gt_seq = imgs[ep_idx, :T_rollout] + + cond_seq = cond_seq.to(device).float() + # Optional GT context anchor: VAE-encode the first n_context frames + # to latents and pin them clean -> predict the continuation given + # the real frame(s) + the action sequence (the in-distribution task). + context_latent = None + if ev.n_context_frames > 0: + n_ctx = min(ev.n_context_frames, T_rollout) + ctx_img = gt_seq[:n_ctx].to(device).float() + if ctx_img.max() > 1.5: + ctx_img = ctx_img / 255.0 + mu, _lv = self.vae.encode(ctx_img) + context_latent = self.normalize_latent(mu).unsqueeze(0) + # ---- Sampler returns (1, T, C, H, W) ---- + pred_latent = self._rollout_latent( + ev, algo, cond_seq, device, context_latent=context_latent, + ).squeeze(0) + # ---- VAE decode -> pixel frames (T, 3, H, W) ---- + pred_frames = self.vae.decode(self.denormalize_latent(pred_latent)) + + n_cmp = min(ev.recon_loss_n_frames, pred_frames.shape[0]) + gt_f = gt_seq[:n_cmp].to(device).float() + if gt_f.max() > 1.5: + gt_f = gt_f / 255.0 + + return SimpleNamespace( + pred_frames=pred_frames, + gt_for_mse=gt_f, + # Raw GT slice (NOT pre-normalized) — the unified eval's + # side-by-side panel reproduces the original per-frame + # ``gt_seq[t] / (255 if gt_seq.max()>1.5 else 1)`` exactly. + gt_panel_raw=gt_seq, + extra_metrics={}, + ) + + def _build_latent(self, ctx: SimpleNamespace) -> torch.Tensor: + """``image -> VAE.encode -> spatial latent``. + + Packed mode: image is ``(T_total, C, H, W)``; output is + ``(T_total, latent_c, h, w)``. + + Padded mode: image is ``(B, T, C, H, W)``; output is + ``(B, T, latent_c, h, w)``. + """ + if self.image_key not in ctx.obs: + raise KeyError( + f"image key '{self.image_key}' required but missing from " + f"ctx.obs (keys: {list(ctx.obs.keys())})." + ) + img = ctx.obs[self.image_key] + orig_dim = img.dim() + if orig_dim == 5: + B, T, C, H, W = img.shape + img_flat = img.reshape(B * T, C, H, W) + elif orig_dim == 4: + img_flat = img + else: + raise ValueError( + f"image has unexpected dim {orig_dim}; expected 4 (packed) or 5 (padded)." + ) + with torch.no_grad(): + mu, _logvar = self.vae.encode(img_flat) + # mu: (N, latent_c, h, w). Reshape back to padded form if needed. + if orig_dim == 5: + mu = mu.reshape(B, T, *mu.shape[1:]) + return self.normalize_latent(mu) + + # ------------------------------------------------------------------ + # encode -> q_sample on spatial latent. external_cond computed and + # attached to ctx. + # ------------------------------------------------------------------ + + def forward(self, batch: dict, ctx: SimpleNamespace) -> torch.Tensor: + """Override to handle packed data by running each episode + independently through the backbone (no padding, no mask needed). + """ + x_t = self.encode(batch, ctx) + time_cond = ctx.q_state["time_cond"] + + if not ctx.is_packed: + v_pred = self.inner_stage( + x_t, time_cond, external_cond=ctx.external_cond, + ) + else: + cu = ctx.cu_seqlens + B = cu.shape[0] - 1 + pieces = [] + for i in range(B): + s, e = int(cu[i].item()), int(cu[i + 1].item()) + x_ep = x_t[s:e].unsqueeze(0) + t_ep = time_cond[s:e].unsqueeze(0) + c_ep = ctx.external_cond[s:e].unsqueeze(0) + v_ep = self.inner_stage(x_ep, t_ep, external_cond=c_ep) + pieces.append(v_ep.squeeze(0)) + v_pred = torch.cat(pieces, dim=0) + + self.decode(v_pred, batch, ctx) + return v_pred + + def encode(self, batch: dict, ctx: SimpleNamespace) -> torch.Tensor: + latent = self._build_latent(ctx) + + # Build external_cond from (state, action). + cond = self._state_action_to_cond(batch, ctx) + + if ctx.is_packed: + if latent.dim() != 4: + raise ValueError( + f"packed latent must be (T_total, C, H, W); got " + f"{tuple(latent.shape)}" + ) + T_total = latent.shape[0] + t = self._sample_noise_levels((T_total,), latent.device) + else: + if latent.dim() != 5: + raise ValueError( + f"padded latent must be (B, T, C, H, W); got " + f"{tuple(latent.shape)}" + ) + B, T = latent.shape[:2] + t = self._sample_noise_levels((B, T), latent.device) + + # For discrete diffusion, q_sample needs explicit noise so we can + # store it in q_state for compute_loss. + from egomimic.models.diffusion.diffusion.discrete_diffusion import DiscreteDiffusion as _DD + if isinstance(self.diffusion, _DD): + noise = torch.randn_like(latent).clamp_( + -self.diffusion.clip_noise, self.diffusion.clip_noise + ) + x_t = self.diffusion.q_sample(latent, t, noise=noise) + ctx.q_state = { + "x_t": x_t, "k": t, "time_cond": t, + "noise": noise, "x_start": latent, + } + else: + q = self.diffusion.q_sample(latent, t) + ctx.q_state = q + x_t = q["x_t"] + ctx.external_cond = cond + ctx.latent_clean = latent + return x_t diff --git a/egomimic/algo/diffusion/outer_stages/obs_action_image_outer_stage.py b/egomimic/algo/diffusion/outer_stages/obs_action_image_outer_stage.py new file mode 100644 index 000000000..8e16e640d --- /dev/null +++ b/egomimic/algo/diffusion/outer_stages/obs_action_image_outer_stage.py @@ -0,0 +1,325 @@ +"""``ObsActionImageDFoTOuterStage`` — joint obs+action+image DFoT. + +Extends ``ObsActionDFoTOuterStage`` so the diffusion bundle per step is +``[state, vae_latent_flat, action]``, where ``vae_latent_flat`` is the +flattened spatial latent (e.g. 4x6x6 = 144D) produced by a *frozen* +pretrained ``ImageVAE`` on the input image. + +This is what lets the diffusion model jointly predict actions AND +future video frames: at inference we sample the bundle, slice the +latent portion, and pass it through the VAE decoder to recover pixel +frames. + +The image goes ONLY into the bundle now — the cond_encoder is image- +free (the original obs+action variant used image as AdaLN cond; here +that's redundant since the image is part of the diffused state and we +don't want to double-count). + +Bundle layout (in concat order): + [ bundle_obs_keys ... | vae_latent_flat | action ] + +So ``action_slice`` is the trailing 2D, and ``image_latent_slice`` +indexes the middle ``latent_flat_dim``-wide chunk. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import List + +import torch + +from egomimic.algo.diffusion.outer_stages.obs_action_outer_stage import ObsActionDFoTOuterStage +from egomimic.algo.diffusion.vae_algo import load_pretrained_vae + + +class ObsActionImageDFoTOuterStage(ObsActionDFoTOuterStage): + """DFoT outer stage that diffuses ``[state, image_latent, action]`` + jointly per step, with a frozen pretrained ``ImageVAE`` providing + the image latent. + + Args: + action_dim: ACTION-only width (e.g. 2 for pushshapes). + cond_encoder: ``CondEncoderModule``. Usually image-free for this + variant -- the image is in the bundle, not in the cond path. + Pass an empty-obs encoder if you really want no cond. + backbone: ``DFoTBackbone`` with ``action_dim = bundle_dim`` + (= sum(bundle_obs_dims) + latent_flat_dim + action_dim). + diffusion: continuous / discrete diffusion module. + bundle_obs_keys: obs keys to include in the bundle (e.g. + ``["state_agent_obj"]``). + bundle_obs_dims: per-key feature width for each obs key. + vae_checkpoint_path: path to the frozen ImageVAE Lightning ckpt. + Loaded via ``algo.vae.algo.load_pretrained_vae``. + image_key: which obs key carries the image stream. Default + ``"front_img_1"``. + cond_output_key: unused for this variant since cond_encoder is + typically empty; kept for ABI parity with the base class. + """ + + def __init__( + self, + action_dim: int, + cond_encoder, + backbone, + diffusion, + bundle_obs_keys: List[str], + bundle_obs_dims: List[int], + vae_checkpoint_path: str, + image_key: str = "front_img_1", + action_in_bundle: bool = True, + cond_output_key: str = "fused_cond", + ): + # Load the VAE first so we know its latent_flat_dim — that + # determines the bundle width the backbone has to expect. + vae = load_pretrained_vae(vae_checkpoint_path) + if not getattr(vae, "spatial_latent", False): + raise NotImplementedError( + "ObsActionImageDFoTOuterStage currently only supports " + "VAEs with spatial latent (latent shape = " + "(latent_channels, H', W')). Re-train the VAE with " + "spatial_latent=True (e.g. config vae_pushshapes_v4)." + ) + latent_c = int(vae.latent_channels) + latent_h = int(vae.bottleneck_size) + latent_w = int(vae.bottleneck_size) + latent_flat_dim = latent_c * latent_h * latent_w + + # Treat the image latent as just another "obs" key in the bundle + # construction. We append "_image_latent" to the obs key list + + # latent_flat_dim to the dim list; ``_build_bundle`` (inherited) + # will then concat in order: [bundle_obs_keys..., _image_latent, + # action]. Override ``_build_bundle`` below to compute the + # latent on the fly from the image via the frozen VAE. + extended_keys = list(bundle_obs_keys) + ["_image_latent"] + extended_dims = list(bundle_obs_dims) + [latent_flat_dim] + + super().__init__( + action_dim=action_dim, + cond_encoder=cond_encoder, + backbone=backbone, + diffusion=diffusion, + bundle_obs_keys=extended_keys, + bundle_obs_dims=extended_dims, + action_in_bundle=action_in_bundle, + cond_output_key=cond_output_key, + ) + + # Register the frozen VAE as a submodule so .to(device) moves + # its weights. All its params are already requires_grad=False + # from ``load_pretrained_vae``. + self.vae = vae + self.image_key = str(image_key) + self._latent_shape = (latent_c, latent_h, latent_w) + self._latent_flat_dim = int(latent_flat_dim) + # Save the raw obs-key list (without the synthetic _image_latent) + # so callers can introspect what real keys are in the bundle. + self.real_bundle_obs_keys = list(bundle_obs_keys) + self.real_bundle_obs_dims = list(bundle_obs_dims) + + @property + def image_latent_slice(self) -> slice: + """Slice into the bundle's trailing dim that holds the flat + image latent. Useful for the self-rollout eval to pull out + predicted latents before VAE-decoding them.""" + # bundle layout: [state(dims)... | latent(flat_dim) | action(action_dim)] + start = sum(self.real_bundle_obs_dims) + return slice(start, start + self._latent_flat_dim) + + @property + def latent_shape(self) -> tuple[int, int, int]: + """``(latent_channels, latent_h, latent_w)`` — the un-flattened + spatial shape of the VAE's latent. The eval reshapes the bundle's + latent slice to this before calling ``vae.decode``.""" + return self._latent_shape + + # ------------------------------------------------------------------ + # Video-rollout hook (COMBINE A — decode-on-outer-stage). + # + # The family-agnostic ``DFoTVideoRolloutEval`` calls this per episode. + # This stage owns: the UNCONDITIONAL chunk/AR bundle sampler, the + # latent-slice extraction, and the frozen-VAE decode that turn sampler + # output into pixel frames. Code below is moved byte-for-byte from the + # old ``eval_dfot_video_rollout.DFoTVideoRolloutEval._rollout`` + + # ``_decode_latents_to_frames`` + the per-episode body of its + # ``compute_metrics_and_viz`` (single-panel: t=0 GT prepended). + # ------------------------------------------------------------------ + + #: metric-key infix the unified eval uses (``Valid/embN__recon_mse_*``). + video_metric_prefix = "video" + #: panel layout the unified eval assembles for this family. + video_panel = "single_t0prepend" + #: this family compares preds against the GT episode's first-N frames + #: (unconditioned rollout — no per-step GT alignment guarantee). + video_has_extra_metrics = False + + @torch.no_grad() + def _decode_latents_to_frames( + self, latent_flat: torch.Tensor + ) -> torch.Tensor: + """``(T, latent_flat_dim) -> (T, 3, H, W)``.""" + c, h, w = self.latent_shape + z = latent_flat.view(-1, c, h, w) + # Frozen VAE on the right device. + return self.vae.decode(z) + + @torch.no_grad() + def _rollout_bundle(self, ev, algo, device, cond) -> torch.Tensor: + """Sample a (T, bundle_dim) bundle. Mode = ``"chunk"`` uses + ``algo._sample_chunk`` (uniform per-token noise + DDIM). Mode = + ``"ar"`` builds a staircase schedule matrix and calls + ``sample`` directly so the per-token noise pattern mirrors + DFoT's training AR schedule.""" + from egomimic.models.diffusion.diffusion.discrete_diffusion import DiscreteDiffusion + from egomimic.models.diffusion.sampling import ( + sample as _sample, + staircase_ar_schedule, + ) + + T = ev.rollout_steps + if ev.mode == "chunk": + bundle = algo._sample_chunk(B=1, T=T, cond=cond, device=device) + return bundle.squeeze(0) + + # mode == "ar" + bundle_dim = int(algo.outer_stage.bundle_dim) + discrete_ts = ( + int(algo.diffusion.timesteps) + if isinstance(algo.diffusion, DiscreteDiffusion) + else None + ) + schedule = staircase_ar_schedule( + T=T, + chunk_size=ev.ar_chunk_size, + step_size=ev.ar_step_size, + discrete_timesteps=discrete_ts, + ).to(device) + ec = cond.unsqueeze(0) if cond is not None else None + bundle = _sample( + algo.diffusion, + algo.backbone, + schedule_matrix=schedule, + action_dim=bundle_dim, + batch_size=1, + external_cond=ec, + cfg_scale=ev.cfg_scale, + device=device, + ).squeeze(0) + return bundle + + @torch.no_grad() + def rollout_video_episode( + self, ev, algo, _batch, emb_id, ep_idx, ep_start, ep_len, device + ) -> SimpleNamespace: + """Per-episode video rollout for the joint obs+action+image stage. + + Returns a namespace the unified eval consumes: + * ``pred_frames`` ``(T,3,H,W)`` float in [0,1] — decoded preds. + * ``gt_for_mse`` ``(T,3,H,W)`` float in [0,1] — GT first-N frames + (this family's MSE is against the GT episode's leading frames; + the rollout is unconditioned so alignment is only distributional). + * ``gt_t0_chw`` the raw t=0 GT image to prepend as launch context. + * ``extra_metrics`` empty for this family. + """ + imgs = _batch[ev.image_key] + is_packed = _batch.get("_packed", False) + + # t=0 launch image for this episode. + if is_packed: + img_chw = imgs[ep_start] + else: + img_chw = imgs[ep_idx, 0] + + # The model's _sample_chunk wants a cond tensor that matches the + # backbone's external_cond shape. For the obs+action+image + # variant the cond_encoder is usually empty (image is in the + # bundle, not the cond), so cond should be None. + cond = None + bundle_pred = self._rollout_bundle(ev, algo, device, cond) # (T, bundle_dim) + latent_seq = bundle_pred[:, self.image_latent_slice] + pred_frames = self._decode_latents_to_frames(latent_seq) + + n_cmp = min(ev.recon_loss_n_frames, pred_frames.shape[0]) + if is_packed: + gt_seq = imgs[ep_start : ep_start + n_cmp] + else: + gt_seq = imgs[ep_idx, :n_cmp] + # GT in [0,1] float; predicted also in [0,1] from VAE sigmoid. + gt_f = gt_seq.to(device).float() + if gt_f.max() > 1.5: + gt_f = gt_f / 255.0 + + return SimpleNamespace( + pred_frames=pred_frames, + gt_for_mse=gt_f, + gt_t0_chw=img_chw, + extra_metrics={}, + ) + + # ------------------------------------------------------------------ + # Override the parent's _build_bundle so the "_image_latent" entry + # is computed from the image via the frozen VAE, not pulled from + # ctx.obs directly. + # ------------------------------------------------------------------ + + def _build_bundle(self, batch: dict, ctx: SimpleNamespace) -> torch.Tensor: + actions = batch[ctx.action_key] + pieces: List[torch.Tensor] = [] + + # 1. Real obs keys (state, etc.). + for key, dim in zip(self.real_bundle_obs_keys, self.real_bundle_obs_dims): + if key not in ctx.obs: + raise KeyError( + f"obs key '{key}' required in bundle but missing from " + f"ctx.obs (keys: {list(ctx.obs.keys())})." + ) + v = ctx.obs[key] + if v.shape[-1] != dim: + raise ValueError( + f"obs '{key}' trailing dim {v.shape[-1]} != configured " + f"bundle_obs_dim {dim}." + ) + pieces.append(v) + + # 2. Image latent via frozen VAE encoder. Always no_grad. + if self.image_key not in ctx.obs: + raise KeyError( + f"image key '{self.image_key}' required but missing from " + f"ctx.obs (keys: {list(ctx.obs.keys())})." + ) + img = ctx.obs[self.image_key] + # img can be packed (T_total, C, H, W) or padded (B, T, C, H, W). + # Collapse to (N, C, H, W) for the VAE encoder, then reshape back. + orig_dim = img.dim() + if orig_dim == 5: + B, T, C, H, W = img.shape + img_flat = img.reshape(B * T, C, H, W) + elif orig_dim == 4: + img_flat = img + else: + raise ValueError( + f"image '{self.image_key}' has unexpected dim {orig_dim}; " + f"expected 4 (packed T_total,C,H,W) or 5 (padded B,T,C,H,W)." + ) + with torch.no_grad(): + # Use posterior mean (mu) for the latent — deterministic so + # the diffusion target is well-defined. + mu, _logvar = self.vae.encode(img_flat) + # mu shape: (N, latent_c, h', w'); flatten the spatial dims. + latent_flat = mu.flatten(1) + if orig_dim == 5: + latent_flat = latent_flat.reshape(B, T, -1) + pieces.append(latent_flat) + + # 3. Action (only when it is a diffusion target; in world-model mode + # it is routed through external_cond by the base ``encode``). + if self.action_in_bundle: + pieces.append(actions) + + bundle = torch.cat(pieces, dim=-1) + if bundle.shape[-1] != self._bundle_dim: + raise RuntimeError( + f"bundle width mismatch: built {bundle.shape[-1]}, expected " + f"{self._bundle_dim}." + ) + return bundle diff --git a/egomimic/algo/diffusion/outer_stages/obs_action_outer_stage.py b/egomimic/algo/diffusion/outer_stages/obs_action_outer_stage.py new file mode 100644 index 000000000..bc0ffeeba --- /dev/null +++ b/egomimic/algo/diffusion/outer_stages/obs_action_outer_stage.py @@ -0,0 +1,203 @@ +"""``ObsActionDFoTOuterStage`` — joint obs+action diffusion forcing. + +Mirrors the pattern used in https://github.com/buoyancy99/diffusion-forcing +``df_planning.py``: per step, the diffusion target is a "bundle" tensor +formed by concatenating obs modalities with the action along the last +dim, ``bundle = concat([obs_0, ..., obs_K, action], dim=-1)``. The whole +bundle is noised together at training time with one independent noise +level per (sequence position) token, and the backbone is asked to denoise +it jointly. + +Image obs (or any other high-dimensional modality) is NOT added to the +bundle — it stays in the AdaLN conditioning path via ``cond_encoder``, so +the backbone sees it as un-noised side information. The split between +"bundle modalities" and "cond-only modalities" is decided by the +``bundle_obs_keys`` / ``bundle_obs_dims`` config. + +At inference time, the algo's AR-staircase / chunk samplers run on a +buffer of ``bundle_dim`` width; before the committed slice is sent to the +env, ``DFoT._inference_step_*`` applies ``outer_stage.action_slice`` to +pick the action portion out of the bundle and discards the obs portion. + +Train-time loss is unchanged: ``DFoTLoss`` averages MSE over all bundle +dims, so both obs and action contribute to the gradient — that's the +"joint world-model + policy" objective from the diffusion-forcing paper. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import List + +import torch + +from egomimic.algo.diffusion.outer_stages.outer_stage import DFoTOuterStage + + +class ObsActionDFoTOuterStage(DFoTOuterStage): + """DFoT outer stage that diffuses concat([obs, action]) jointly. + + Args: + action_dim: width of the ACTION portion only (e.g. 2 for pushshapes). + cond_encoder: usually configured with image_encoders only — state + modalities that are in ``bundle_obs_keys`` should NOT also be + in cond_encoder.obs_specs (or they'd appear on both the noised + input and the cond path, which double-counts). + backbone: ``DFoTBackbone`` configured with ``action_dim = obs_dim + + action_dim`` (the full bundle width). The class checks this. + diffusion: ``ContinuousDiffusion`` or ``DiscreteDiffusion``. + bundle_obs_keys: list of obs keys to include in the diffusion + bundle, in the order they are concatenated. Action is always + concatenated LAST (so ``action_slice`` is a single trailing + range). Each key must exist on every training batch. + bundle_obs_dims: per-key feature width matching ``bundle_obs_keys``. + ``sum(bundle_obs_dims) + action_dim`` must equal the backbone's + configured ``action_dim`` (= bundle width). + cond_output_key: same as base. + """ + + def __init__( + self, + action_dim: int, + cond_encoder, + backbone, + diffusion, + bundle_obs_keys: List[str], + bundle_obs_dims: List[int], + action_in_bundle: bool = True, + cond_output_key: str = "fused_cond", + ): + super().__init__( + action_dim=action_dim, + cond_encoder=cond_encoder, + backbone=backbone, + diffusion=diffusion, + cond_output_key=cond_output_key, + ) + if len(bundle_obs_keys) != len(bundle_obs_dims): + raise ValueError( + f"bundle_obs_keys ({len(bundle_obs_keys)}) and " + f"bundle_obs_dims ({len(bundle_obs_dims)}) must align." + ) + # action_in_bundle=True (default): the action is a DIFFUSION TARGET in + # the bundle (joint obs+action prediction = POLICY). False: the action + # is moved to the AdaLN conditioning path (external_cond) and the bundle + # is obs-only (action-conditioned video = WORLD MODEL). In world-model + # mode the backbone must be configured with cond_dim = action_dim. + self.action_in_bundle = bool(action_in_bundle) + self.bundle_obs_keys = list(bundle_obs_keys) + self.bundle_obs_dims = [int(d) for d in bundle_obs_dims] + self._obs_total = sum(self.bundle_obs_dims) + self._bundle_dim = self._obs_total + ( + int(action_dim) if self.action_in_bundle else 0 + ) + + # Backbone width sanity check — its ``action_dim`` must match the + # bundle width since it produces ``v_pred`` of the same shape. + bb_dim = int(getattr(backbone, "action_dim", -1)) + if bb_dim != self._bundle_dim: + raise ValueError( + f"backbone.action_dim ({bb_dim}) must equal bundle width " + f"obs_total({self._obs_total}) + action_dim({action_dim}) " + f"= {self._bundle_dim}." + ) + + # ------------------------------------------------------------------ + # Override base properties so the algo's inference path slices and + # sizes against the bundle correctly. + # ------------------------------------------------------------------ + + @property + def bundle_dim(self) -> int: + return self._bundle_dim + + @property + def action_slice(self) -> slice: + # action is concatenated last; slice picks it out of the bundle. + # World-model mode (action_in_bundle=False) predicts no action, so the + # slice is empty. + if not self.action_in_bundle: + return slice(self._obs_total, self._obs_total) + return slice(self._obs_total, self._bundle_dim) + + # ------------------------------------------------------------------ + # Bundle construction. Reads obs values from ``ctx.obs`` (set by + # algo.forward_training before calling self.outer_stage(...)) and the + # action from ``batch[ctx.action_key]``. Concatenates along the + # trailing dim, then runs ``diffusion.q_sample`` on the whole thing. + # ------------------------------------------------------------------ + + def _build_bundle(self, batch: dict, ctx: SimpleNamespace) -> torch.Tensor: + actions = batch[ctx.action_key] + pieces: List[torch.Tensor] = [] + for key, dim in zip(self.bundle_obs_keys, self.bundle_obs_dims): + if key not in ctx.obs: + raise KeyError( + f"obs key '{key}' required in the bundle but missing from " + f"ctx.obs (keys: {list(ctx.obs.keys())})." + ) + v = ctx.obs[key] + # Allow both packed (T_total, D) and padded (B, T, D); reject + # anything else — image obs (4D) should be in cond_encoder, not + # the bundle. + if v.dim() not in (actions.dim() - 1, actions.dim()): + # Permissive: only require trailing-dim alignment and same + # rank as actions for cat. Anything else is a config bug. + pass + if v.shape[-1] != dim: + raise ValueError( + f"obs '{key}' trailing dim {v.shape[-1]} != configured " + f"bundle_obs_dim {dim}." + ) + pieces.append(v) + if self.action_in_bundle: + pieces.append(actions) + bundle = torch.cat(pieces, dim=-1) + if bundle.shape[-1] != self._bundle_dim: + raise RuntimeError( + f"bundle width mismatch: built {bundle.shape[-1]}, expected " + f"{self._bundle_dim}." + ) + return bundle + + def encode(self, batch: dict, ctx: SimpleNamespace) -> torch.Tensor: + bundle = self._build_bundle(batch, ctx) + # World-model mode: the action is the conditioning signal, fed through + # the backbone's external_cond / AdaLN path instead of being diffused. + action_cond = None if self.action_in_bundle else batch[ctx.action_key] + + if ctx.is_packed: + if bundle.dim() != 2: + raise ValueError( + f"packed bundle must be (T_total, bundle_dim); " + f"got {tuple(bundle.shape)}" + ) + T_total = bundle.shape[0] + cond = ( + action_cond + if action_cond is not None + else self._encode_cond_packed(ctx.obs) + ) + t = self._sample_noise_levels((T_total,), bundle.device) + else: + if bundle.dim() != 3: + raise ValueError( + f"padded bundle must be (B, T, bundle_dim); " + f"got {tuple(bundle.shape)}" + ) + B, T, _ = bundle.shape + cond = ( + action_cond + if action_cond is not None + else self._encode_cond_padded(ctx.obs, T) + ) + t = self._sample_noise_levels((B, T), bundle.device) + + q = self.diffusion.q_sample(bundle, t) + ctx.q_state = q + ctx.external_cond = cond + # Stash the per-step ground-truth bundle so downstream logging / + # diagnostics can compute per-modality losses if desired (the + # default DFoTLoss reduces over all bundle dims uniformly). + ctx.bundle_clean = bundle + return q["x_t"] diff --git a/egomimic/algo/diffusion/outer_stages/outer_stage.py b/egomimic/algo/diffusion/outer_stages/outer_stage.py new file mode 100644 index 000000000..333f10505 --- /dev/null +++ b/egomimic/algo/diffusion/outer_stages/outer_stage.py @@ -0,0 +1,199 @@ +"""``DFoTOuterStage`` — outermost stage for the DFoT algorithm. + +Owns: +- ``cond_encoder``: CondEncoderModule that turns obs into per-token AdaLN cond. +- ``backbone``: DFoTBackbone (Isotropic trunk + per-token AdaLN noise embedding). + Stored as ``self.inner_stage`` per the OuterStage convention. +- ``diffusion``: ContinuousDiffusion (or DiscreteDiffusion) — used to call + ``q_sample`` at training time. Has no learnable params here; lives on the + outer stage so the loss class can call ``compute_loss`` symmetrically. + +The training-time ``forward`` does: + encode -> q_sample -> backbone -> decode (writes batch['pred_v']) + +The DFoTLoss class (in egomimic/algo/diffusion/algo.py) then reads ctx.q_state + +batch['pred_v'] and produces the scalar SNR-weighted epsilon-MSE. + +Inference paths (closed-loop AR sample_step, chunk-mode plan-and-execute) +live on the algo class for now — they use this outer stage's cond_encoder +and backbone directly without going through ``forward``. A later refactor +can fold them in once the training-path refactor is verified. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Optional + +import torch +import torch.nn as nn + +from egomimic.models.diffusion.backbones.backbone import DFoTBackbone +from egomimic.models.diffusion.diffusion.discrete_diffusion import DiscreteDiffusion +from egomimic.models.stems.cond_encoders import CondEncoderModule + + +def make_dfot_ctx( + *, + is_packed: bool, + action_key: str, + obs: dict, + cu_seqlens: Optional[torch.Tensor] = None, + max_seqlen: Optional[int] = None, +) -> SimpleNamespace: + """Build a minimal DFoT context object. Fields are filled in by the + outer stage during ``encode`` (q_state, external_cond) and consumed by + the loss class. ``cu_seqlens``/``max_seqlen`` carry packing info.""" + return SimpleNamespace( + is_packed=is_packed, + action_key=action_key, + obs=obs, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + q_state=None, + external_cond=None, + ) + + +class DFoTOuterStage(nn.Module): + """Training-time outer stage for DFoT. + + Subclass-specific contract: + + * ``encode(batch, ctx)``: reads ``batch[ctx.action_key]`` (clean actions), + encodes obs into per-token cond via ``self.cond_encoder``, samples + per-token noise levels, runs ``diffusion.q_sample`` to get noisy + tokens. Stores the full ``q_state`` and the encoded cond on + ``ctx`` for the loss to read later. Returns the noisy actions + (``x_t``) for the inner_stage (backbone). + + * ``decode(v_pred, batch, ctx)``: writes ``batch['pred_v']`` so the + loss can pick it up. (For inference paths the algo class converts + v -> action separately via the sampler step formulas.) + + * ``forward(batch, ctx)``: overrides the default OuterStage flow to + thread ``cu_seqlens`` / ``max_seqlen`` / ``time_cond`` / ``external_cond`` + into the backbone, which has a richer signature than a plain + ``inner_stage(x, ctx) -> x``. + """ + + def __init__( + self, + action_dim: int, + cond_encoder: CondEncoderModule, + backbone: DFoTBackbone, + diffusion: nn.Module, # ContinuousDiffusion or DiscreteDiffusion + cond_output_key: str = "fused_cond", + ): + super().__init__() + self.inner_stage = backbone + self.action_dim = int(action_dim) + self.cond_encoder = cond_encoder + self.diffusion = diffusion + self.cond_output_key = cond_output_key + + @property + def bundle_dim(self) -> int: + """Width of the diffused tensor. For vanilla DFoT this equals + ``action_dim``; subclasses that diffuse extra modalities (obs+action + joint, etc.) override to return the full bundle width.""" + return self.action_dim + + @property + def action_slice(self) -> slice: + """Slice into the trailing dim of the sampled bundle that + corresponds to actions. Vanilla DFoT is action-only, so this is + the full slice. Subclasses (e.g. obs+action joint) override to + point at just the action portion of their bundle.""" + return slice(0, self.action_dim) + + # ------------------------------------------------------------------- + # Per-mode cond encode helpers (lifted from algo.py to keep modality + # encoding on the outer stage where it belongs). + # ------------------------------------------------------------------- + + def _encode_cond_padded(self, obs: dict, T: int) -> Optional[torch.Tensor]: + cond_dict = self.cond_encoder.encode(obs, T) + return cond_dict.get(self.cond_output_key) + + def _encode_cond_packed(self, obs: dict) -> Optional[torch.Tensor]: + obs_3d = { + k: (v.unsqueeze(0) if torch.is_tensor(v) else v) for k, v in obs.items() + } + cond_dict = self.cond_encoder.encode(obs_3d, T_action=1) + c = cond_dict.get(self.cond_output_key) + if c is None: + return None + if c.dim() == 3 and c.shape[0] == 1: + c = c.squeeze(0) + return c + + # ------------------------------------------------------------------- + # Noise-level sampling. Mirrors DFoT._sample_noise_levels exactly so + # the refactor preserves training-time noise distribution. + # ------------------------------------------------------------------- + + def _sample_noise_levels(self, shape, device) -> torch.Tensor: + if isinstance(self.diffusion, DiscreteDiffusion): + return torch.randint( + 0, self.diffusion.timesteps, shape, device=device, dtype=torch.long + ) + return torch.rand(shape, device=device).clamp_(1e-5, 1.0 - 1e-5) + + # ------------------------------------------------------------------- + # OuterStage API. + # ------------------------------------------------------------------- + + def encode(self, batch: dict, ctx: SimpleNamespace) -> torch.Tensor: + actions = batch[ctx.action_key] + + if ctx.is_packed: + if actions.dim() != 2: + raise ValueError( + f"packed actions must be (T_total, action_dim); " + f"got {tuple(actions.shape)}" + ) + T_total = actions.shape[0] + cond = self._encode_cond_packed(ctx.obs) + t = self._sample_noise_levels((T_total,), actions.device) + else: + if actions.dim() != 3: + raise ValueError( + f"padded actions must be (B, T, action_dim); " + f"got {tuple(actions.shape)}" + ) + B, T, _ = actions.shape + cond = self._encode_cond_padded(ctx.obs, T) + t = self._sample_noise_levels((B, T), actions.device) + + q = self.diffusion.q_sample(actions, t) + ctx.q_state = q + ctx.external_cond = cond + return q["x_t"] + + def decode(self, v_pred: torch.Tensor, batch: dict, ctx: SimpleNamespace) -> None: + # The loss reads ``batch['pred_v']`` and combines with ``ctx.q_state``. + batch["pred_v"] = v_pred + + def forward(self, batch: dict, ctx: SimpleNamespace) -> torch.Tensor: + """Override the default OuterStage.forward to thread the backbone's + extra kwargs (time_cond from q_state, external_cond, cu_seqlens, + max_seqlen for packed) through the inner_stage call.""" + x_t = self.encode(batch, ctx) + time_cond = ctx.q_state["time_cond"] + if ctx.is_packed: + v_pred = self.inner_stage( + x_t, + time_cond, + external_cond=ctx.external_cond, + cu_seqlens=ctx.cu_seqlens, + max_seqlen=ctx.max_seqlen, + ) + else: + v_pred = self.inner_stage( + x_t, + time_cond, + external_cond=ctx.external_cond, + ) + self.decode(v_pred, batch, ctx) + return v_pred diff --git a/egomimic/algo/diffusion/outer_stages/pixel_obs_action_outer_stage.py b/egomimic/algo/diffusion/outer_stages/pixel_obs_action_outer_stage.py new file mode 100644 index 000000000..406824b86 --- /dev/null +++ b/egomimic/algo/diffusion/outer_stages/pixel_obs_action_outer_stage.py @@ -0,0 +1,324 @@ +"""``PixelObsActionDFoTOuterStage`` — UNIFIED pixel-space obs+action DFoT policy. + +This single parameterized stage subsumes the three former near-duplicate +pixel-policy outer stages, selected by the ``pixel_mode`` knob: + + * ``pixel_mode="policy"`` (Design A) — action broadcast into RGB CHANNELS, + jointly diffused by the DiT3D (bundle = 3 RGB + C_a action planes); decode by + global-avg-pool of the predicted action planes. Was + ``PixelObsActionPolicyDFoTOuterStage``. + * ``pixel_mode="regress"`` (Design B) — diffuse ONLY the RGB video + (latent_channels=3); a conv ``action_head`` regresses the action from the + model's predicted clean frame (x0). Action is NOT a diffusion target. Was + ``PixelObsActionRegressPolicyDFoTOuterStage``. + * ``pixel_mode="decoupled"`` (DEC) — action rides as its OWN DiT3D token with an + INDEPENDENT per-frame noise level (backbone ``action_token_dim``); backbone + returns ``(v_image, v_action)``. Was ``PixelObsActionDecoupledDFoTOuterStage``. + +All three subclass the proven no-VAE pixel video model +(``PixelSpatialDFoTOuterStage``). Each mode reproduces the corresponding former +class EXACTLY — identical construction (state_dict), identical forward outputs, +identical duck-typed attribute surface consumed by the algo inference paths +(``_action_channels`` for policy, ``action_head`` for regress, +``decouple_action_noise`` for decoupled, plus the mode-correct ``action_slice``). + +Per-mode kwargs (all accepted; only the selected mode's are used): + * policy: ``action_channels`` (default = ``action_dim``) + * regress: ``action_loss_weight`` (default 1.0), ``head_width`` (default 64) + * decoupled: ``action_loss_weight`` (default 1.0), + ``decouple_action_noise`` (default True) +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import torch +import torch.nn as nn + +from egomimic.algo.diffusion.outer_stages.pixel_spatial_outer_stage import PixelSpatialDFoTOuterStage + + +_PIXEL_MODES = ("policy", "regress", "decoupled") + + +class PixelObsActionDFoTOuterStage(PixelSpatialDFoTOuterStage): + def __init__( + self, + *args, + pixel_mode: str = "policy", + # --- policy (Design A) --- + action_channels: int | None = None, + # --- regress (Design B) --- + head_width: int = 64, + # --- regress + decoupled --- + action_loss_weight: float = 1.0, + # --- decoupled --- + decouple_action_noise: bool = True, + **kwargs, + ): + super().__init__(*args, **kwargs) + if pixel_mode not in _PIXEL_MODES: + raise ValueError( + f"pixel_mode must be one of {_PIXEL_MODES}; got {pixel_mode!r}" + ) + self.pixel_mode = str(pixel_mode) + + if self.pixel_mode == "policy": + Ca = int(action_channels) if action_channels is not None else self.action_dim + if Ca < self.action_dim: + raise ValueError( + f"action_channels ({Ca}) must be >= action_dim ({self.action_dim}); " + f"each action component broadcasts to one plane." + ) + self._action_channels = Ca + # bundle now = image channels + action planes (drives sampler tensor alloc) + self._bundle_shape = ( + self._image_channels + Ca, + self._image_size, + self._image_size, + ) + + elif self.pixel_mode == "regress": + self.action_loss_weight = float(action_loss_weight) + c = self._image_channels + # conv-down: predicted frame (C,H,W) -> action (A,) [your "project down"] + self.action_head = nn.Sequential( + nn.Conv2d(c, head_width, 3, stride=2, padding=1), nn.SiLU(), # H/2 + nn.Conv2d(head_width, head_width, 3, stride=2, padding=1), nn.SiLU(), # H/4 + nn.Conv2d(head_width, head_width, 3, stride=2, padding=1), nn.SiLU(), # H/8 + nn.AdaptiveAvgPool2d(1), nn.Flatten(), + nn.Linear(head_width, self.action_dim), + ) + + else: # decoupled + self.action_loss_weight = float(action_loss_weight) + self.decouple_action_noise = bool(decouple_action_noise) + bb_at = int(getattr(self.inner_stage, "action_token_dim", 0)) + if bb_at != self.action_dim: + raise ValueError( + f"backbone.action_token_dim ({bb_at}) must equal action_dim " + f"({self.action_dim}) for the decoupled pixel policy." + ) + + # ------------------------------------------------------------------ # + @property + def action_slice(self) -> slice: + if self.pixel_mode == "policy": + # action occupies the trailing C_a channels of the per-frame tensor + return slice(self._image_channels, self._image_channels + self._action_channels) + # regress + decoupled: action is a separate output, not sliced from bundle + return slice(0, 0) + + # ------------------------------------------------------------------ # + # policy: action -> broadcast planes + # ------------------------------------------------------------------ # + def _action_to_planes(self, actions: torch.Tensor, h: int, w: int) -> torch.Tensor: + """``(N, A) -> (N, C_a, H, W)`` by broadcasting each action component + across the spatial plane. Continuous + precise; global-avg-pool on + decode recovers the value with prediction error averaged out.""" + n = actions.shape[0] + a = actions[..., : self._action_channels] + return a.reshape(n, self._action_channels, 1, 1).expand( + n, self._action_channels, h, w + ) + + # ------------------------------------------------------------------ # + # joint image/action frame sampling per packed episode: + # ``_sample_windows_packed`` now lives on the parent + # ``PixelSpatialDFoTOuterStage`` (dedup collapse c7) — inherited here. + # ------------------------------------------------------------------ # + + # ------------------------------------------------------------------ # + # decoupled: per-tensor q_state helper + # ------------------------------------------------------------------ # + def _qstate(self, x: torch.Tensor, t: torch.Tensor) -> dict: + noise = torch.randn_like(x).clamp_( + -self.diffusion.clip_noise, self.diffusion.clip_noise + ) + x_t = self.diffusion.q_sample(x, t, noise=noise) + return {"x_t": x_t, "k": t, "time_cond": t, "noise": noise, "x_start": x} + + # ------------------------------------------------------------------ # + # encode — dispatch on pixel_mode + # ------------------------------------------------------------------ # + def encode(self, batch: dict, ctx: SimpleNamespace): + if self.pixel_mode == "policy": + return self._encode_policy(batch, ctx) + if self.pixel_mode == "regress": + return self._encode_regress(batch, ctx) + return self._encode_decoupled(batch, ctx) + + def _encode_policy(self, batch: dict, ctx: SimpleNamespace) -> torch.Tensor: + images = self._extract_images(ctx) # packed (T,3,H,W) | padded (B,T,3,H,W) + ac_key = getattr(ctx, "action_key", None) + if ac_key is None or ac_key not in batch: + raise KeyError( + f"PixelObsActionPolicy needs the action key in the batch " + f"(ctx.action_key={ac_key!r}); keys: {list(batch.keys())}" + ) + actions = batch[ac_key].to(images.device).float() + + if ctx.is_packed: + if self._frame_sampling != "full" and self.training: + images, actions, new_cu = self._sample_windows_packed( + images, actions, ctx.cu_seqlens + ) + ctx.cu_seqlens = new_cu + ctx.max_seqlen = max( + int(new_cu[i + 1] - new_cu[i]) for i in range(new_cu.shape[0] - 1) + ) + h, w = images.shape[-2:] + planes = self._action_to_planes(actions, h, w) # (T, C_a, H, W) + x_start = torch.cat([images, planes], dim=1) # (T, 3+C_a, H, W) + t = self._sample_noise_levels((x_start.shape[0],), images.device) + else: + b, T = images.shape[:2] + h, w = images.shape[-2:] + planes = self._action_to_planes( + actions.reshape(b * T, -1), h, w + ).reshape(b, T, self._action_channels, h, w) + x_start = torch.cat([images, planes], dim=2) # (B,T,3+C_a,H,W) + t = self._sample_noise_levels((b, T), images.device) + + noise = torch.randn_like(x_start).clamp_( + -self.diffusion.clip_noise, self.diffusion.clip_noise + ) + x_t = self.diffusion.q_sample(x_start, t, noise=noise) + ctx.q_state = { + "x_t": x_t, + "k": t, + "time_cond": t, + "noise": noise, + "x_start": x_start, + } + ctx.external_cond = None + ctx.latent_clean = x_start + return x_t + + def _encode_regress(self, batch: dict, ctx: SimpleNamespace) -> torch.Tensor: + images = self._extract_images(ctx) + ac_key = getattr(ctx, "action_key", None) + if ac_key is None or ac_key not in batch: + raise KeyError(f"regress policy needs action key (ctx.action_key={ac_key!r})") + actions = batch[ac_key].to(images.device).float() + + if ctx.is_packed: + if self._frame_sampling != "full" and self.training: + images, actions, new_cu = self._sample_windows_packed(images, actions, ctx.cu_seqlens) + ctx.cu_seqlens = new_cu + ctx.max_seqlen = max(int(new_cu[i + 1] - new_cu[i]) for i in range(new_cu.shape[0] - 1)) + t = self._sample_noise_levels((images.shape[0],), images.device) + else: + b, T = images.shape[:2] + t = self._sample_noise_levels((b, T), images.device) + + noise = torch.randn_like(images).clamp_(-self.diffusion.clip_noise, self.diffusion.clip_noise) + x_t = self.diffusion.q_sample(images, t, noise=noise) + ctx.q_state = {"x_t": x_t, "k": t, "time_cond": t, "noise": noise, "x_start": images} + ctx.external_cond = None + ctx._gt_actions = actions + return x_t + + def _encode_decoupled(self, batch: dict, ctx: SimpleNamespace): + images = self._extract_images(ctx) + ac_key = getattr(ctx, "action_key", None) + if ac_key is None or ac_key not in batch: + raise KeyError( + f"decoupled pixel policy needs the action key in the batch " + f"(ctx.action_key={ac_key!r}); keys: {list(batch.keys())}" + ) + actions = batch[ac_key].to(images.device).float() + + if ctx.is_packed: + t = self._sample_noise_levels((images.shape[0],), images.device) + else: + t = self._sample_noise_levels((images.shape[0], images.shape[1]), images.device) + + ctx.q_state = self._qstate(images, t) + t_a = (self._sample_noise_levels(tuple(t.shape), images.device) + if self.decouple_action_noise else t) + ctx.q_action = self._qstate(actions, t_a) + ctx.external_cond = None + ctx.latent_clean = images + return ctx.q_state["x_t"], ctx.q_action["x_t"] + + # ------------------------------------------------------------------ # + # forward — dispatch on pixel_mode + # ------------------------------------------------------------------ # + def forward(self, batch: dict, ctx: SimpleNamespace) -> torch.Tensor: + if self.pixel_mode == "policy": + # policy reuses the base PixelSpatial forward (encode -> backbone -> decode) + return PixelSpatialDFoTOuterStage.forward(self, batch, ctx) + if self.pixel_mode == "regress": + return self._forward_regress(batch, ctx) + return self._forward_decoupled(batch, ctx) + + def _forward_regress(self, batch: dict, ctx: SimpleNamespace) -> torch.Tensor: + x_t = self.encode(batch, ctx) + time_cond = ctx.q_state["time_cond"] + gt_actions = ctx._gt_actions + + if not ctx.is_packed: + v_pred = self.inner_stage(x_t, time_cond, external_cond=None) + else: + cu = ctx.cu_seqlens + b = cu.shape[0] - 1 + pieces = [] + for i in range(b): + s, e = int(cu[i].item()), int(cu[i + 1].item()) + pieces.append(self.inner_stage(x_t[s:e].unsqueeze(0), time_cond[s:e].unsqueeze(0), + external_cond=None).squeeze(0)) + v_pred = torch.cat(pieces, dim=0) + + video_loss = self.diffusion.compute_loss(v_pred, ctx.q_state).mean() + + # predicted clean frame -> regress action + pred_x0 = self.diffusion.predict_start_from_v(ctx.q_state["x_t"], ctx.q_state["k"], v_pred) + if pred_x0.dim() == 5: + bb, tt = pred_x0.shape[:2] + ap = self.action_head(pred_x0.reshape(bb * tt, *pred_x0.shape[2:])).reshape(bb, tt, -1) + else: + ap = self.action_head(pred_x0) # (T, A) + gta = gt_actions[..., : self.action_dim] + action_loss = ((ap - gta) ** 2).mean() + + ctx.precomputed_loss = video_loss + self.action_loss_weight * action_loss + batch["pred_v"] = v_pred + batch["pred_action"] = ap + return v_pred + + def _forward_decoupled(self, batch: dict, ctx: SimpleNamespace) -> torch.Tensor: + x_t, x_t_a = self.encode(batch, ctx) + time_cond = ctx.q_state["time_cond"] + a_levels = ctx.q_action["time_cond"] if self.decouple_action_noise else None + + if not ctx.is_packed: + v_img, v_act = self.inner_stage( + x_t, time_cond, external_cond=None, + action=x_t_a, action_noise_levels=a_levels, + ) + else: + cu = ctx.cu_seqlens + B = cu.shape[0] - 1 + vi, va = [], [] + for i in range(B): + s, e = int(cu[i].item()), int(cu[i + 1].item()) + vimg, vact = self.inner_stage( + x_t[s:e].unsqueeze(0), time_cond[s:e].unsqueeze(0), + external_cond=None, + action=x_t_a[s:e].unsqueeze(0), + action_noise_levels=( + a_levels[s:e].unsqueeze(0) if a_levels is not None else None + ), + ) + vi.append(vimg.squeeze(0)) + va.append(vact.squeeze(0)) + v_img = torch.cat(vi, dim=0) + v_act = torch.cat(va, dim=0) + + img_loss = self.diffusion.compute_loss(v_img, ctx.q_state).mean() + act_loss = self.diffusion.compute_loss(v_act, ctx.q_action).mean() + ctx.precomputed_loss = img_loss + self.action_loss_weight * act_loss + batch["pred_v"] = v_img + return v_img diff --git a/egomimic/algo/diffusion/outer_stages/pixel_spatial_outer_stage.py b/egomimic/algo/diffusion/outer_stages/pixel_spatial_outer_stage.py new file mode 100644 index 000000000..3883f385f --- /dev/null +++ b/egomimic/algo/diffusion/outer_stages/pixel_spatial_outer_stage.py @@ -0,0 +1,381 @@ +"""``PixelSpatialDFoTOuterStage`` — DFoT on raw pixel images (no VAE). + +Training frame sampling modes (``frame_sampling`` config): + - ``"full"``: use the entire episode as-is. No cropping. + - ``"fixed_window"``: sample a random window of ``sample_n_frames`` + consecutive frames from each episode. Matches the reference repo. + - ``"start_to_end"``: sample a random start index, take everything + from there to the end of the episode. Variable-length sequences. + - ``"random_subsample"``: sample ``sample_n_frames`` frames uniformly + at random (not necessarily consecutive) from the episode, sorted + by time. Preserves temporal coverage but with gaps. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Optional + +import torch +import torch.nn as nn + +from egomimic.algo.diffusion.outer_stages.outer_stage import DFoTOuterStage +from egomimic.models.diffusion.diffusion.discrete_diffusion import DiscreteDiffusion + +try: + from torchmetrics.image import ( # noqa: F401 + PeakSignalNoiseRatio, + StructuralSimilarityIndexMeasure, + ) + from torchmetrics.image.lpip import ( # noqa: F401 + LearnedPerceptualImagePatchSimilarity, + ) + _HAS_METRICS = True +except ImportError: + _HAS_METRICS = False + + +class PixelSpatialDFoTOuterStage(DFoTOuterStage): + + def __init__( + self, + action_dim: int, + cond_encoder, + backbone, + diffusion, + image_key: str = "front_img_1", + image_channels: int = 3, + image_size: int = 96, + frame_sampling: str = "full", + sample_n_frames: int = 9, + cond_output_key: str = "fused_cond", + ): + super().__init__( + action_dim=action_dim, + cond_encoder=cond_encoder, + backbone=backbone, + diffusion=diffusion, + cond_output_key=cond_output_key, + ) + self.image_key = str(image_key) + self._image_channels = int(image_channels) + self._image_size = int(image_size) + self._frame_sampling = str(frame_sampling) + self._sample_n_frames = int(sample_n_frames) + self._bundle_shape = (self._image_channels, self._image_size, self._image_size) + + @property + def bundle_shape(self) -> tuple: + return self._bundle_shape + + @property + def bundle_dim(self) -> int: + c, h, w = self._bundle_shape + return c * h * w + + @property + def action_slice(self) -> slice: + return slice(0, 0) + + def _extract_images(self, ctx: SimpleNamespace) -> torch.Tensor: + if self.image_key not in ctx.obs: + raise KeyError( + f"image key '{self.image_key}' required but missing from " + f"ctx.obs (keys: {list(ctx.obs.keys())})." + ) + img = ctx.obs[self.image_key] + if img.dtype == torch.uint8: + img = img.float() / 255.0 + elif img.max() > 1.5: + img = img.float() / 255.0 + else: + img = img.float() + return img + + def _sample_windows_packed(self, images, actions, cu): + """Frame-sample images (and OPTIONALLY actions) IDENTICALLY per packed + episode, so the action at frame t always lines up with image t after + cropping. + + This is the SUPERSET sampler shared by the image-only base stage and + the image+action subclass (dedup collapse c7). Passing ``actions=None`` + makes it the pure image-only frame sampler — the image cropping + + cu_seqlens are byte-identical whether or not actions are supplied (the + action branch performs no extra RNG draws), proven by + ``tests/test_c7_sampler_reducer_equality``. + + Returns ``(sampled_images, sampled_actions_or_None, new_cu)``. + """ + has_act = actions is not None + b = cu.shape[0] - 1 + n = self._sample_n_frames + mode = self._frame_sampling + img_crops, act_crops = [], [] + for i in range(b): + s, e = int(cu[i].item()), int(cu[i + 1].item()) + L = e - s + if mode == "fixed_window" and L > n: + st = int(torch.randint(0, L - n + 1, (1,)).item()) + sl = slice(s + st, s + st + n) + elif mode == "start_to_end" and L > n: + st = int(torch.randint(0, L - n + 1, (1,)).item()) + sl = slice(s + st, e) + elif mode == "random_subsample" and L > n: + idx = torch.randperm(L)[:n].sort().values + img_crops.append(images[s + idx]) + if has_act: + act_crops.append(actions[s + idx]) + continue + else: + sl = slice(s, e) + img_crops.append(images[sl]) + if has_act: + act_crops.append(actions[sl]) + new_cu = torch.zeros(b + 1, dtype=cu.dtype, device=cu.device) + for i, c in enumerate(img_crops): + new_cu[i + 1] = new_cu[i] + c.shape[0] + sampled_act = torch.cat(act_crops, 0) if has_act else None + return torch.cat(img_crops, 0), sampled_act, new_cu + + def _sample_frames_packed( + self, images: torch.Tensor, cu_seqlens: torch.Tensor + ): + """Image-only frame sampler — thin delegate to the action-aware + superset ``_sample_windows_packed`` with ``actions=None`` (dedup + collapse c7; the standalone duplicated loop was removed after proving + byte-identical output across all sampling modes). + + Returns: + sampled: (T_new, C, H, W) re-packed sampled frames. + new_cu: (B+1,) updated cu_seqlens. + """ + sampled, _, new_cu = self._sample_windows_packed(images, None, cu_seqlens) + return sampled, new_cu + + def encode(self, batch: dict, ctx: SimpleNamespace) -> torch.Tensor: + images = self._extract_images(ctx) + + # Frame sampling for packed data during training. + if (ctx.is_packed and self._frame_sampling != "full" + and self.training): + images, new_cu = self._sample_frames_packed( + images, ctx.cu_seqlens + ) + ctx.cu_seqlens = new_cu + ctx.max_seqlen = max( + int(new_cu[i + 1] - new_cu[i]) + for i in range(new_cu.shape[0] - 1) + ) + # Also crop the action key in the batch so loss shapes match + ac_key = getattr(ctx, 'action_key', None) + if ac_key and ac_key in batch: + actions = batch[ac_key] + crops = [] + old_cu = ctx._original_cu if hasattr(ctx, '_original_cu') else None + # Actions were already packed matching original images, + # but we've now re-packed images. We need to crop actions + # the same way. Store original cu before overwrite. + # Actually, the actions come from the batch which hasn't + # been modified. We need to re-crop them too. + # For simplicity, just truncate actions to match new cu. + # This works because the loss only uses ctx.q_state which + # has the cropped x_start/noise. + + if ctx.is_packed: + T_total = images.shape[0] + t = self._sample_noise_levels((T_total,), images.device) + else: + B, T = images.shape[:2] + t = self._sample_noise_levels((B, T), images.device) + + noise = torch.randn_like(images).clamp_( + -self.diffusion.clip_noise, self.diffusion.clip_noise + ) + x_t = self.diffusion.q_sample(images, t, noise=noise) + ctx.q_state = { + "x_t": x_t, + "k": t, + "time_cond": t, + "noise": noise, + "x_start": images, + } + ctx.external_cond = None + ctx.latent_clean = images + return x_t + + def forward(self, batch: dict, ctx: SimpleNamespace) -> torch.Tensor: + x_t = self.encode(batch, ctx) + time_cond = ctx.q_state["time_cond"] + + if not ctx.is_packed: + v_pred = self.inner_stage( + x_t, time_cond, external_cond=None, + ) + else: + cu = ctx.cu_seqlens + B = cu.shape[0] - 1 + pieces = [] + for i in range(B): + s, e = int(cu[i].item()), int(cu[i + 1].item()) + x_ep = x_t[s:e].unsqueeze(0) + t_ep = time_cond[s:e].unsqueeze(0) + v_ep = self.inner_stage(x_ep, t_ep, external_cond=None) + pieces.append(v_ep.squeeze(0)) + v_pred = torch.cat(pieces, dim=0) + + self.decode(v_pred, batch, ctx) + return v_pred + + # ------------------------------------------------------------------ + # Video-rollout hook (COMBINE A — decode-on-outer-stage). + # + # This stage owns the pixel-family rollout: a sliding-window DDIM rollout + # anchored on the first n_context GT frames (no VAE — output IS pixels), + # plus the PSNR/SSIM/LPIPS perceptual metrics. Moved byte-for-byte from + # ``eval_dfot_pixel_video_rollout.DFoTPixelVideoRolloutEval`` + # (``_rollout_sliding_window`` + per-episode body). Side-by-side + # [GT|pred] panel. + # ------------------------------------------------------------------ + + video_metric_prefix = "pixel" + video_panel = "sidebyside" + video_has_extra_metrics = True + + @torch.no_grad() + def _rollout_sliding_window( + self, ev, algo, total_frames: int, context_frames: torch.Tensor, + window_size: int, device, + ) -> torch.Tensor: + """Sliding window rollout matching the reference DFoT inference. + + Args: + ev: the unified video-rollout eval (carries n_chunk_steps). + algo: the DFoT algo. + total_frames: total number of frames to generate. + context_frames: (n_context, C, H, W) initial context frames. + window_size: number of frames per denoising window. + device: torch device. + + Returns: + (total_frames, C, H, W) generated frames. + """ + from egomimic.models.diffusion.sampling import sample_step, vanilla_schedule + + bundle_shape = self.bundle_shape + discrete_ts = ( + int(algo.diffusion.timesteps) + if isinstance(algo.diffusion, DiscreteDiffusion) + else None + ) + + n_context = context_frames.shape[0] + generated = context_frames.clone() # (n_context, C, H, W) + + while generated.shape[0] < total_frames: + # How many context frames to use (up to window_size - 1) + c = min(generated.shape[0], window_size - 1) + # How many new frames to generate + h = min(total_frames - generated.shape[0], window_size - c) + T_window = c + h + + # Build window: context + noise + context_part = generated[-c:] # (c, C, H, W) + noise_part = torch.randn(h, *bundle_shape, device=device) + x = torch.cat([context_part, noise_part], dim=0).unsqueeze(0) # (1, T_window, C, H, W) + + # Build schedule + schedule = vanilla_schedule( + n_steps=ev.n_chunk_steps, T=T_window, discrete_timesteps=discrete_ts, + ).to(device) + + # Context tokens get noise_level = -1 (clean) throughout + # Modify schedule: context columns stay at -1 + for step_idx in range(schedule.shape[0]): + schedule[step_idx, :c] = -1 if discrete_ts else 0.0 + + # Run DDIM sampling + for s in range(schedule.shape[0] - 1): + x = sample_step( + algo.diffusion, algo.backbone, x=x, + current_levels=schedule[s], + next_levels=schedule[s + 1], + external_cond=None, eta=0.0, + ) + # Revert context frames to clean values + x[0, :c] = context_part + + # Append newly generated frames + new_frames = x[0, c:c + h] + generated = torch.cat([generated, new_frames.clamp(0, 1)], dim=0) + + return generated[:total_frames] + + @torch.no_grad() + def rollout_video_episode( + self, ev, algo, _batch, emb_id, ep_idx, ep_start, ep_len, device + ) -> SimpleNamespace: + """Per-episode pixel-space sliding-window rollout + perceptual metrics.""" + imgs = _batch[ev.image_key] + is_packed = _batch.get("_packed", False) + T_rollout = min(ev.rollout_steps, ep_len) + + if is_packed: + gt_seq = imgs[ep_start : ep_start + T_rollout] + else: + gt_seq = imgs[ep_idx, :T_rollout] + + # GT normalization (also used to seed the context frame(s)). + gt_f = gt_seq[:T_rollout].to(device).float() + if gt_f.max() > 1.5: + gt_f = gt_f / 255.0 + + # Conditional rollout matching the reference DFoT prediction task: + # seed the first n_context GT frames, hold them clean (noise level + # -1) for the whole sampling trajectory, and predict the rest + # conditioned on them. + n_ctx = max(1, min(ev.n_context_frames, T_rollout)) + pred_frames = self._rollout_sliding_window( + ev, + algo, + total_frames=T_rollout, + context_frames=gt_f[:n_ctx], + window_size=min(ev.rollout_window, T_rollout), + device=device, + ) # (T, 3, H, W) + + # Clamp output to [0, 1]. + pred_frames = pred_frames.clamp(0.0, 1.0) + + n_cmp = min(ev.recon_loss_n_frames, pred_frames.shape[0]) + + # PSNR, SSIM, LPIPS per episode (averaged over frames). + extra = {} + if _HAS_METRICS and n_cmp > 0: + from torchmetrics.image import ( + PeakSignalNoiseRatio, + StructuralSimilarityIndexMeasure, + ) + from torchmetrics.image.lpip import ( + LearnedPerceptualImagePatchSimilarity, + ) + pred_cmp = pred_frames[:n_cmp].to(device) + gt_cmp = gt_f[:n_cmp].to(device) + psnr_fn = PeakSignalNoiseRatio(data_range=1.0).to(device) + extra["psnr"] = psnr_fn(pred_cmp, gt_cmp) + ssim_fn = StructuralSimilarityIndexMeasure(data_range=1.0).to(device) + extra["ssim"] = ssim_fn(pred_cmp, gt_cmp) + try: + lpips_fn = LearnedPerceptualImagePatchSimilarity( + net_type="alex", normalize=True).to(device) + extra["lpips"] = lpips_fn(pred_cmp, gt_cmp) + except Exception: + pass + + return SimpleNamespace( + pred_frames=pred_frames, + gt_for_mse=gt_f[:n_cmp], + # Raw GT slice for the side-by-side panel (reproduces the + # original per-frame ``gt_seq[t] / (255 if max>1.5 else 1)``). + gt_panel_raw=gt_seq, + extra_metrics=extra, + ) diff --git a/egomimic/algo/diffusion/outer_stages/pixel_video_outer_stage.py b/egomimic/algo/diffusion/outer_stages/pixel_video_outer_stage.py new file mode 100644 index 000000000..6c6c0ad49 --- /dev/null +++ b/egomimic/algo/diffusion/outer_stages/pixel_video_outer_stage.py @@ -0,0 +1,126 @@ +"""``PixelVideoDFoTOuterStage`` — Clean pixel-space DFoT for video clip batches. + +Designed for ``ZarrVideoClipDataset`` which provides standard batches of +``(B, T, 3, H, W)`` fixed-length video clips. No packing, no per-episode +loops — the backbone processes the whole batch at once, exactly like the +reference DFoT repo. + +This keeps fused_min_snr working correctly since ``k`` stays ``(B, T)``. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Optional + +import torch +import torch.nn as nn + +from egomimic.algo.diffusion.outer_stages.outer_stage import DFoTOuterStage +from egomimic.models.diffusion.diffusion.discrete_diffusion import DiscreteDiffusion + + +class PixelVideoDFoTOuterStage(DFoTOuterStage): + """Clean pixel-space DFoT outer stage for fixed-length video clip batches. + + Expects ``batch[video_key]`` to be ``(B, T, 3, H, W)`` float [0, 1]. + No packed mode support — use ``ZarrVideoClipDataset`` instead of packed datasets. + + Args: + action_dim: unused, kept for interface compat. + cond_encoder: empty CondEncoderModule for ABI parity. + backbone: DFoTDiT3DBackbone configured for pixel space. + diffusion: DiscreteDiffusion instance. + video_key: batch key for video clips. + """ + + def __init__( + self, + action_dim: int, + cond_encoder, + backbone, + diffusion, + video_key: str = "videos", + cond_output_key: str = "fused_cond", + ): + super().__init__( + action_dim=action_dim, + cond_encoder=cond_encoder, + backbone=backbone, + diffusion=diffusion, + cond_output_key=cond_output_key, + ) + self.video_key = str(video_key) + # Infer shape from backbone config + self._channels = int(backbone.latent_channels) + self._size = int(backbone.latent_size) + self._bundle_shape = (self._channels, self._size, self._size) + + @property + def bundle_shape(self) -> tuple: + return self._bundle_shape + + @property + def bundle_dim(self) -> int: + c, h, w = self._bundle_shape + return c * h * w + + @property + def action_slice(self) -> slice: + return slice(0, 0) + + def encode(self, batch: dict, ctx: SimpleNamespace) -> torch.Tensor: + """Extract video clips and apply forward noising. + + Expects ``batch[video_key]``: ``(B, T, C, H, W)`` float [0, 1]. + """ + # Get videos directly from batch (not from ctx.obs) + videos = batch.get(self.video_key) + if videos is None: + raise KeyError( + f"'{self.video_key}' not found in batch. " + f"Keys: {list(batch.keys())}" + ) + + # Ensure float [0, 1] + if videos.dtype == torch.uint8: + videos = videos.float() / 255.0 + elif videos.max() > 1.5: + videos = videos.float() / 255.0 + else: + videos = videos.float() + + B, T = videos.shape[:2] + + # Sample per-token noise levels: (B, T) + k = torch.randint( + 0, self.diffusion.timesteps, (B, T), + device=videos.device, dtype=torch.long, + ) + + # Forward noising + noise = torch.randn_like(videos).clamp_( + -self.diffusion.clip_noise, self.diffusion.clip_noise + ) + x_t = self.diffusion.q_sample(videos, k, noise=noise) + + ctx.q_state = { + "x_t": x_t, + "k": k, + "time_cond": k, + "noise": noise, + "x_start": videos, + } + ctx.external_cond = None + ctx.latent_clean = videos + return x_t + + def forward(self, batch: dict, ctx: SimpleNamespace) -> torch.Tensor: + """Single forward pass on the whole batch — no per-episode loops.""" + x_t = self.encode(batch, ctx) + k = ctx.q_state["k"] # (B, T) — 2D, fused_min_snr works + + v_pred = self.inner_stage(x_t, k, external_cond=None) + + self.decode(v_pred, batch, ctx) + return v_pred diff --git a/egomimic/algo/diffusion/outer_stages/spatial_obs_action_policy_outer_stage.py b/egomimic/algo/diffusion/outer_stages/spatial_obs_action_policy_outer_stage.py new file mode 100644 index 000000000..a0a7f78ee --- /dev/null +++ b/egomimic/algo/diffusion/outer_stages/spatial_obs_action_policy_outer_stage.py @@ -0,0 +1,134 @@ +"""``SpatialObsActionPolicyDFoTOuterStage`` — 2D POLICY variant of obs+action +diffusion forcing. + +Like ``ImageSpatialDFoTOuterStage`` (the 2D world model), the image stays a +SPATIAL VAE latent (NOT flattened) and is the spatial diffusion target. The +difference: the ACTION is ALSO a diffusion target, carried as one per-frame +"action token" appended to the DiT3D token sequence (see +``DFoTDiT3DBackbone(action_token_dim=...)``). State enters as ``external_cond``. + +The backbone returns ``(v_image, v_action)``; this stage computes a structured +loss ``image_v_mse + action_loss_weight * action_v_mse`` and stashes it on +``ctx.precomputed_loss`` (``DFoTLoss`` returns it verbatim). This keeps the +structured (image, action) target out of the scalar 1D v-MSE path. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import torch +import torch.nn as nn + +from egomimic.algo.diffusion.outer_stages.image_spatial_outer_stage import ImageSpatialDFoTOuterStage + + +class SpatialObsActionPolicyDFoTOuterStage(ImageSpatialDFoTOuterStage): + def __init__(self, *args, action_loss_weight: float = 1.0, + decouple_action_noise: bool = False, **kwargs): + super().__init__(*args, **kwargs) + self.action_loss_weight = float(action_loss_weight) + # When True, the action token carries an INDEPENDENT per-frame noise + # level (never reliably clean) so the model must predict a_t from obs + # and cannot copy a_{t-1}. Cut-action-input fix for closed-loop rollout. + self.decouple_action_noise = bool(decouple_action_noise) + + # State-only conditioning MLP. The parent's ``state_action_proj`` mixes + # the action into the cond, but here the action is a TARGET, so we + # condition on state alone. + self.state_only_proj = nn.Sequential( + nn.Linear(self._state_dim_total, self._state_action_proj_dim), + nn.SiLU(), + nn.Linear(self._state_action_proj_dim, self._state_action_proj_dim), + ) + + bb_at = int(getattr(self.inner_stage, "action_token_dim", 0)) + if bb_at != self._action_dim: + raise ValueError( + f"backbone.action_token_dim ({bb_at}) must equal action_dim " + f"({self._action_dim}) for the 2D policy variant." + ) + + @property + def action_slice(self) -> slice: + # Action is predicted as a separate backbone output, not sliced from a + # bundle — empty slice signals "N/A" to bundle-based extractors. + return slice(0, 0) + + def _state_to_cond(self, batch: dict, ctx: SimpleNamespace) -> torch.Tensor: + pieces = [ctx.obs[k] for k in self.bundle_obs_keys] + concat = torch.cat(pieces, dim=-1) + return self.state_only_proj(concat) + + def encode(self, batch: dict, ctx: SimpleNamespace): + latent = self._build_latent(ctx) # (T_total,C,H,W) | (B,T,C,H,W) + action = batch[ctx.action_key] # (T_total,A) | (B,T,A) + state_cond = self._state_to_cond(batch, ctx) + + if ctx.is_packed: + t = self._sample_noise_levels((latent.shape[0],), latent.device) + else: + t = self._sample_noise_levels((latent.shape[0], latent.shape[1]), latent.device) + + ctx.q_state = self._qstate(latent, t) + if self.decouple_action_noise: + t_a = self._sample_noise_levels(tuple(t.shape), latent.device) + else: + t_a = t # legacy: shared per-frame level + ctx.q_action = self._qstate(action, t_a) + ctx.external_cond = state_cond + return ctx.q_state["x_t"], ctx.q_action["x_t"] + + def _qstate(self, x: torch.Tensor, t: torch.Tensor) -> dict: + """Noise ``x`` at level ``t`` and return a q_state dict. Discrete + diffusion's ``q_sample`` returns a bare tensor, so we supply explicit + noise and assemble the dict ourselves (matching ImageSpatial).""" + from egomimic.models.diffusion.diffusion.discrete_diffusion import DiscreteDiffusion as _DD + + if isinstance(self.diffusion, _DD): + noise = torch.randn_like(x).clamp_( + -self.diffusion.clip_noise, self.diffusion.clip_noise + ) + x_t = self.diffusion.q_sample(x, t, noise=noise) + return {"x_t": x_t, "k": t, "time_cond": t, "noise": noise, "x_start": x} + return self.diffusion.q_sample(x, t) + + def forward(self, batch: dict, ctx: SimpleNamespace) -> torch.Tensor: + x_t_latent, x_t_action = self.encode(batch, ctx) + time_cond = ctx.q_state["time_cond"] + + if not ctx.is_packed: + v_latent, v_action = self.inner_stage( + x_t_latent, time_cond, + external_cond=ctx.external_cond, action=x_t_action, + action_noise_levels=( + ctx.q_action["time_cond"] if self.decouple_action_noise else None + ), + ) + else: + # DiT3D has no packed mode -> run each episode as a batch-of-1. + cu = ctx.cu_seqlens + B = cu.shape[0] - 1 + vlat, vact = [], [] + for i in range(B): + s, e = int(cu[i].item()), int(cu[i + 1].item()) + vl, va = self.inner_stage( + x_t_latent[s:e].unsqueeze(0), + time_cond[s:e].unsqueeze(0), + external_cond=ctx.external_cond[s:e].unsqueeze(0), + action=x_t_action[s:e].unsqueeze(0), + action_noise_levels=( + ctx.q_action["time_cond"][s:e].unsqueeze(0) + if self.decouple_action_noise else None + ), + ) + vlat.append(vl.squeeze(0)) + vact.append(va.squeeze(0)) + v_latent = torch.cat(vlat, dim=0) + v_action = torch.cat(vact, dim=0) + + latent_loss = self.diffusion.compute_loss(v_latent, ctx.q_state).mean() + action_loss = self.diffusion.compute_loss(v_action, ctx.q_action).mean() + ctx.precomputed_loss = latent_loss + self.action_loss_weight * action_loss + batch["pred_v"] = v_latent + return v_latent diff --git a/egomimic/algo/diffusion/vae_algo.py b/egomimic/algo/diffusion/vae_algo.py new file mode 100644 index 000000000..63605cbff --- /dev/null +++ b/egomimic/algo/diffusion/vae_algo.py @@ -0,0 +1,315 @@ +"""``VAE`` algo — image autoencoder pre-training stage. + +Wraps ``ImageVAE`` so it plugs into the existing trainHydra / Lightning +pipeline. No norm_stats interactions, no embodiment-specific routing, +no obs/action concept — purely "give me images, I'll reconstruct them". + +Training step: + * Pull the image stream from the packed batch + (``observations.images.`` -> ``self.image_key``). + * For each frame: forward through VAE -> (x_rec, mu, logvar). + * Loss = MSE(x, x_rec) + beta * KL(N(mu, sigma) || N(0, I)). Reported + as ``recon_loss``, ``kl_loss``, ``vae_loss``. + +Eval step: same forward, returns reconstructions for the recon +evaluator (``egomimic/eval/eval_vae_recon.py``) to write a comparison +mp4. + +After training, the saved Lightning checkpoint contains the VAE's +weights under ``nets["vae"]``. ``load_pretrained_vae(path)`` is the +helper used by ObsActionImageDFoTOuterStage to load the frozen +encoder/decoder at policy-training time. +""" + +from __future__ import annotations + +from collections import OrderedDict +from typing import Optional + +import torch +import torch.nn as nn +from overrides import override + +from egomimic.algo.algo import Algo +from egomimic.models.diffusion.image_vae import ImageVAE +from egomimic.rldb.embodiment.embodiment import get_embodiment_id + + +class VAE(Algo): + """Image VAE training algo. + + Args: + vae: ``ImageVAE`` instance (built by Hydra). + image_key: which obs key carries the image stream (default + ``"front_img_1"``). + beta: KL weight in the ELBO. Default 1.0 (standard VAE). Drop + to ~1e-3 for lower-KL "near-AE" behavior if the latent + collapses. + domains: embodiment-name list (for batch routing parity with + the other algos). + device: optional device override. + """ + + def __init__( + self, + vae: ImageVAE, + norm_stats=None, + image_key: str = "front_img_1", + beta: float = 1.0, + lpips_weight: float = 0.0, + lpips_net: str = "alex", + domains: Optional[list] = None, + device=None, + **kwargs, + ): + super().__init__() + self.image_key = str(image_key) + self.beta = float(beta) + self.lpips_weight = float(lpips_weight) + self.lpips_net = str(lpips_net) + self.norm_stats = norm_stats + self.domains = list(domains or []) + self.device = device or torch.device( + "cuda" if torch.cuda.is_available() else "cpu" + ) + # LPIPS perceptual loss. Built lazily and frozen — its job is to + # measure feature-space distance via a small pretrained net + # (alex / vgg / squeeze), which is much more edge-sensitive than + # per-pixel MSE. Inputs MUST be in [-1, 1] before being passed in. + self._lpips = None + if self.lpips_weight > 0: + try: + import lpips + + self._lpips = lpips.LPIPS(net=self.lpips_net).eval() + for p in self._lpips.parameters(): + p.requires_grad = False + except ImportError as e: + raise ImportError( + "LPIPS requested (lpips_weight > 0) but the `lpips` " + "package isn't installed. uv pip install lpips." + ) from e + modules = {"vae": vae} + if self._lpips is not None: + modules["lpips"] = self._lpips + self.nets = nn.ModuleDict(modules).float().to(self.device) + # Mirror the other algos' embodiment routing so PL batches + # arriving as {emb_name: {...}} work without surprise. We only + # care about per-image data here, so the rest is trivial. + self.embodiment_ids = {emb: get_embodiment_id(emb) for emb in self.domains} + + # ---- Algo API ---------------------------------------------------- # + + @override + def process_batch_for_training(self, batch): + """Convert the loader's {emb_name: ...} batch into + {emb_id: {'images': (B, ...)}}. + + Accepts both packed (image: (T_total, C, H, W)) and padded + (image: (B, T, C, H, W)) layouts. Output is always a flat + (N, C, H, W) tensor on device so the VAE forward is one big + batched call. + """ + out = {} + for emb_name, _batch in batch.items(): + emb_id = get_embodiment_id(emb_name) + # Look up the image tensor by zarr key first (raw batch), + # then by friendly key (already remapped batches). + image = None + for cand in (self.image_key, f"observations.images.{self.image_key}"): + if cand in _batch: + image = _batch[cand] + break + if image is None: + continue + if image.dim() == 4: + # Packed: (T_total, C, H, W) -> already flat per-frame. + imgs = image + elif image.dim() == 5: + # Padded: (B, T, C, H, W) -> collapse the (B, T) prefix. + B, T = image.shape[:2] + imgs = image.reshape(B * T, *image.shape[2:]) + else: + raise ValueError( + f"image '{self.image_key}' has unexpected dim {image.dim()}; " + f"expected 4 (packed) or 5 (padded)." + ) + imgs = imgs.to(self.device).float() + # Robomimic image floats are in [0,1] already; if a future + # data path delivers uint8 we'd /255 here. Defensive guard: + if imgs.max() > 1.5: + imgs = imgs / 255.0 + out[emb_id] = {"images": imgs} + return out + + @override + def forward_training(self, batch): + """One forward pass per embodiment -> per-key loss tensors.""" + predictions = OrderedDict() + for emb_id, _batch in batch.items(): + x = _batch["images"] + x_rec, mu, logvar = self.nets["vae"](x) + loss, recon, kl = ImageVAE.vae_loss( + x, + x_rec, + mu, + logvar, + beta=self.beta, + ) + # Add LPIPS perceptual loss if configured. LPIPS expects + # inputs in [-1, 1]; our images are in [0, 1], so rescale. + lpips_term = torch.zeros((), device=x.device, dtype=x.dtype) + if self._lpips is not None: + x_lp = x * 2.0 - 1.0 + xr_lp = x_rec * 2.0 - 1.0 + lpips_term = self.nets["lpips"](x_lp, xr_lp).mean() + loss = loss + self.lpips_weight * lpips_term + predictions[f"{emb_id}_recon_loss"] = recon + predictions[f"{emb_id}_kl_loss"] = kl + predictions[f"{emb_id}_lpips_loss"] = lpips_term.detach() + predictions[f"{emb_id}_vae_loss"] = loss + predictions[f"{emb_id}_x_rec"] = x_rec.detach() + predictions[f"{emb_id}_x_gt"] = x.detach() + return predictions + + @override + def forward_eval(self, batch): + """Same forward as training, but no grads. Returns (x_gt, x_rec) + per embodiment under keys ``emb{id}_x_gt`` / ``emb{id}_x_rec``. + """ + out = {} + with torch.no_grad(): + for emb_id, _batch in batch.items(): + x = _batch["images"] + mu, logvar = self.nets["vae"].encode(x) + # Eval-time: use the posterior mean (no noise) for cleaner + # reconstructions in the recon video. + x_rec = self.nets["vae"].decode(mu) + out[f"emb{emb_id}_x_gt"] = x.detach() + out[f"emb{emb_id}_x_rec"] = x_rec.detach() + return out + + @override + def compute_losses(self, predictions, batch): + total = torch.tensor(0.0, device=self.device) + loss_dict = OrderedDict() + for emb_id in batch.keys(): + r = predictions[f"{emb_id}_recon_loss"] + k = predictions[f"{emb_id}_kl_loss"] + total_emb = predictions[f"{emb_id}_vae_loss"] + loss_dict[f"emb{emb_id}_recon_loss"] = r + loss_dict[f"emb{emb_id}_kl_loss"] = k + loss_dict[f"emb{emb_id}_vae_loss"] = total_emb + if f"{emb_id}_lpips_loss" in predictions: + loss_dict[f"emb{emb_id}_lpips_loss"] = predictions[ + f"{emb_id}_lpips_loss" + ] + total = total + total_emb + # Use "action_loss" key name for back-compat with the trainer's + # progress-bar / logging hook that reads losses["action_loss"] + # as the primary scalar. + loss_dict["action_loss"] = total / max(len(batch), 1) + return loss_dict + + @override + def log_info(self, info): + log = OrderedDict() + log["Loss"] = info["losses"]["action_loss"].item() + for k, v in info["losses"].items(): + log[k] = v.item() + return log + + +def load_pretrained_vae( + checkpoint_path: str, + map_location: str = "cpu", + **vae_kwargs, +) -> ImageVAE: + """Load a frozen VAE from a Lightning checkpoint produced by training + this algo. Returns the inner ``ImageVAE`` with all params + ``requires_grad=False`` and the module in eval mode. + + Path convention: pass the Lightning .ckpt file. The state dict prefix + is ``robomimic_model.nets.vae.*``. + + Architecture detection: shapes are sniffed from the state dict for + common combos (flat vs spatial latent; residual vs plain; channel + widths). You can pass explicit ``vae_kwargs`` to override the + inference if the model was built with a non-default config. + """ + ckpt = torch.load(checkpoint_path, map_location=map_location, weights_only=False) + state = ckpt.get("state_dict", ckpt) + # Different Lightning wrappers stick the VAE under different prefixes + # ("robomimic_model.nets.vae." for the trainHydra+ModelWrapper path, + # or just "nets.vae." for older PL wrappers). Try both in priority + # order; first non-empty match wins. + vae_state = {} + for prefix in ("robomimic_model.nets.vae.", "nets.vae.", "vae."): + candidate = { + k[len(prefix) :]: v for k, v in state.items() if k.startswith(prefix) + } + if candidate: + vae_state = candidate + break + if not vae_state: + raise RuntimeError( + f"No keys with any known VAE prefix in {checkpoint_path}; " + f"sample keys: {list(state.keys())[:5]}" + ) + + # Sniff architecture (overridden by anything passed in vae_kwargs). + spatial_latent = "to_mu.weight" in vae_state + residual = ( + "encoder.0.conv1.weight" in vae_state or "encoder.0.skip.weight" in vae_state + ) + if spatial_latent: + to_mu_w = vae_state["to_mu.weight"] # (latent_c, bottleneck_c, 1, 1) + latent_channels = int(to_mu_w.shape[0]) + bottleneck_c = int(to_mu_w.shape[1]) + latent_dim = 64 # unused in spatial mode + else: + fc_mu_w = vae_state["fc_mu.weight"] + latent_dim = int(fc_mu_w.shape[0]) + flat_dim = int(fc_mu_w.shape[1]) + bottleneck_size_inferred = 6 + bottleneck_c = flat_dim // (bottleneck_size_inferred**2) + latent_channels = 4 # unused in flat mode + + # Encoder channel widths from the state dict. + if residual: + # _ResBlock has conv1 + conv2; encoder.{i}.conv1.weight has shape + # (c_out, c_in, 3, 3). + c_keys = sorted( + [k for k in vae_state if k.startswith("encoder.") and "conv1.weight" in k] + ) + channels = tuple(int(vae_state[k].shape[0]) for k in c_keys) + else: + c_keys = sorted( + [k for k in vae_state if k.startswith("encoder.") and ".conv.weight" in k] + ) + channels = tuple(int(vae_state[k].shape[0]) for k in c_keys) + if not channels: + channels = (32, 64, 128, int(bottleneck_c)) + + # decoder_upsample: only matters for non-residual builds. Detect from + # the presence of a stride-1 stand-alone Conv2d vs ConvTranspose in + # the last decoder block by looking for ``decoder..weight``. + decoder_upsample = "nearest" if residual else "transpose" + + defaults = dict( + in_channels=3, + image_size=96, + latent_dim=latent_dim, + channels=channels, + decoder_upsample=decoder_upsample, + residual=residual, + spatial_latent=spatial_latent, + latent_channels=latent_channels, + ) + defaults.update(vae_kwargs) + vae = ImageVAE(**defaults) + vae.load_state_dict(vae_state, strict=True) + vae.eval() + for p in vae.parameters(): + p.requires_grad = False + return vae diff --git a/egomimic/eval/dfot/__init__.py b/egomimic/eval/dfot/__init__.py new file mode 100644 index 000000000..24a9419bc --- /dev/null +++ b/egomimic/eval/dfot/__init__.py @@ -0,0 +1,39 @@ +"""DFoT closed-loop / rollout evaluators (DESIGN.md §2 ``egomimic/eval/{dfot}``). + +The DFoT self-rollout (world-model) + policy-action + video-rollout evaluators, +curated here in DESIGN.md step 8 (``git mv``, no behaviour change): + + * :class:`DFoTSelfRolloutEval` — joint obs+action world-model rollout. + * :class:`DFoTVideoRolloutEval` — family-agnostic video self-rollout + (COMBINE A): drives the obs+action+image / image-spatial / pixel families + via each outer stage's ``rollout_video_episode`` hook. + :class:`DFoTPixelVideoRolloutEval` / :class:`DFoTSpatialVideoRolloutEval` + are compat aliases of it (the per-family modules were collapsed in + COMBINE A). + * :class:`DFoTPolicyActionEval` / :class:`DFoTPolicyRecedingHorizonEval` — + 2D-policy action prediction (whole-horizon vs receding-horizon). COMBINE B + merged the two into ``eval_dfot_policy`` (they share ``_rollout``). + * :class:`DFoTBundleAnchoredEval` — anchored clean-history bundle rollout. +""" + +from egomimic.eval.dfot.eval_dfot_self_rollout import DFoTSelfRolloutEval +from egomimic.eval.dfot.eval_dfot_video_rollout import ( + DFoTVideoRolloutEval, + DFoTPixelVideoRolloutEval, + DFoTSpatialVideoRolloutEval, +) +from egomimic.eval.dfot.eval_dfot_policy import ( + DFoTPolicyActionEval, + DFoTPolicyRecedingHorizonEval, +) +from egomimic.eval.dfot.eval_dfot_bundle_anchored import DFoTBundleAnchoredEval + +__all__ = [ + "DFoTSelfRolloutEval", + "DFoTVideoRolloutEval", + "DFoTPixelVideoRolloutEval", + "DFoTSpatialVideoRolloutEval", + "DFoTPolicyActionEval", + "DFoTPolicyRecedingHorizonEval", + "DFoTBundleAnchoredEval", +] diff --git a/egomimic/eval/dfot/_base.py b/egomimic/eval/dfot/_base.py new file mode 100644 index 000000000..c07b559e4 --- /dev/null +++ b/egomimic/eval/dfot/_base.py @@ -0,0 +1,53 @@ +"""Shared knob-storage + ``video_dir`` boilerplate for the DFoT evaluators. + +The DFoT rollout/policy/bundle evaluators all subclass :class:`EvalVideo` and +each re-stored the SAME handful of scalar knobs (``embodiment_name``, +``image_key``, ``recon_loss_n_frames``, ``upscale_to``, ``n_chunk_steps``, +``_video_subdir``) and re-defined the SAME 2-line ``video_dir`` override +(``root_dir()/_video_subdir``). COMBINE B hoists that boilerplate here. + +This is a pure storage/path mixin — it adds NO new behaviour and changes NO +resolved attribute value. Each subclass ``__init__`` still declares its own +per-class defaults in its signature and passes them through ``store_dfot_knobs`` +explicitly, so every evaluator lands on byte-identical attribute values (the +configs also pass every knob explicitly). ``limit_val_batches`` / ``max_videos`` +already live on :class:`EvalVideo` and are forwarded via ``super().__init__``. +""" + +from __future__ import annotations + +import os + + +class DFoTVideoEvalMixin: + """Mixin holding the knobs + ``video_dir`` shared by every DFoT evaluator. + + Subclasses call :meth:`store_dfot_knobs` from their ``__init__`` (after + ``super().__init__`` has set up the :class:`EvalVideo` base) with their own + default-bearing arguments. ``upscale_to`` / ``n_chunk_steps`` / + ``recon_loss_n_frames`` accept ``None`` for evaluators that genuinely don't + use a given knob (none currently do, but keeps the mixin honest). + """ + + def store_dfot_knobs( + self, + *, + embodiment_name: str, + image_key: str, + video_subdir: str, + recon_loss_n_frames: int | None = None, + upscale_to: int | None = None, + n_chunk_steps: int | None = None, + ) -> None: + self.embodiment_name = embodiment_name + self.image_key = str(image_key) + self._video_subdir = str(video_subdir) + if recon_loss_n_frames is not None: + self.recon_loss_n_frames = int(recon_loss_n_frames) + if upscale_to is not None: + self.upscale_to = int(upscale_to) + if n_chunk_steps is not None: + self.n_chunk_steps = int(n_chunk_steps) + + def video_dir(self): + return os.path.join(self.root_dir(), self._video_subdir) diff --git a/egomimic/eval/dfot/_sampling.py b/egomimic/eval/dfot/_sampling.py new file mode 100644 index 000000000..5e84293d3 --- /dev/null +++ b/egomimic/eval/dfot/_sampling.py @@ -0,0 +1,86 @@ +"""Shared anchored-DDIM rollout helper for the DFoT evaluators (COMBINE B). + +The "anchored clean-history" rollout pattern is structurally identical across +two single-tensor DFoT samplers: + + * :meth:`DFoTBundleAnchoredEval._rollout` — flat bundle ``(1, T, bundle_dim)``. + * :meth:`ImageSpatialDFoTOuterStage._rollout_latent` (anchored branch) — + spatial latent ``(B, T, C, H, W)``. + +Both clone a per-token noise schedule, PIN the first ``n_ctx`` tokens CLEAN +(``-1`` for discrete diffusion, ``0.0`` for continuous), seed those tokens from +a context tensor, then loop :func:`sample_step` re-pinning the context after +every denoise step. They differ ONLY in the trailing tensor shape, the batch +size, and whether they pass a non-default ``cfg_scale`` — all of which this +helper takes as explicit arguments, so the two call sites stay byte-identical. + +NOTE — the 2D-policy ``DFoTPolicyActionEval._rollout`` is a DIFFERENT sampler: +it co-denoises TWO streams (obs-latent ``x_lat`` + action ``x_act``) through a +dual-output backbone ``(v_lat, v_act)`` with a hand-rolled v-prediction DDIM +step, not the single-tensor ``sample_step`` primitive. It is intentionally NOT +folded into this helper (its loop body has no single-tensor equivalent). +""" + +from __future__ import annotations + +import torch + +from egomimic.models.diffusion.sampling import sample_step + + +@torch.no_grad() +def anchored_ddim_rollout( + diffusion, + backbone, + *, + schedule: torch.Tensor, + context: torch.Tensor, + total_T: int, + trailing_shape, + device, + batch_size: int = 1, + external_cond: torch.Tensor | None = None, + discrete_ts: int | None = None, + cfg_scale: float = 1.0, +) -> torch.Tensor: + """Anchored clean-history DDIM rollout over a single per-token tensor. + + Seeds the first ``n_ctx`` (= ``context.shape[1]``) tokens from ``context``, + pins them CLEAN across every step of ``schedule``, and denoises the rest via + repeated :func:`sample_step` (eta=0). Returns the final ``x`` of shape + ``(batch_size, total_T, *trailing_shape)``. + + Args: + diffusion: ``ContinuousDiffusion`` or ``DiscreteDiffusion``. + backbone: callable passed straight to :func:`sample_step`. + schedule: per-step per-token noise levels ``(n_steps, total_T)``. NOT + mutated — cloned internally before the context columns are pinned. + context: clean prefix ``(batch_size, n_ctx, *trailing_shape)`` seeded + into and re-pinned over the first ``n_ctx`` tokens each step. + total_T: full token length ``T`` of the rollout. + trailing_shape: per-token trailing dims (e.g. ``(bundle_dim,)`` or + ``(C, H, W)``); ``x`` is ``(batch_size, total_T, *trailing_shape)``. + device: device for the noise init. + batch_size: leading batch dim of ``x`` (default 1). + external_cond: optional ``(B, T, cond_dim)`` / ``(B, cond_dim)`` cond. + discrete_ts: ``int`` discrete-timestep count -> clean sentinel ``-1``; + ``None`` (continuous) -> clean sentinel ``0.0``. + cfg_scale: classifier-free-guidance scale forwarded to ``sample_step``. + + Returns: + ``x`` of shape ``(batch_size, total_T, *trailing_shape)``. + """ + n_ctx = context.shape[1] + clean = -1 if discrete_ts is not None else 0.0 + schedule = schedule.clone() + schedule[:, :n_ctx] = clean + x = torch.randn(batch_size, total_T, *trailing_shape, device=device) + x[:, :n_ctx] = context + for s in range(schedule.shape[0] - 1): + x = sample_step( + diffusion, backbone, x=x, + current_levels=schedule[s], next_levels=schedule[s + 1], + external_cond=external_cond, eta=0.0, cfg_scale=cfg_scale, + ) + x[:, :n_ctx] = context + return x diff --git a/egomimic/eval/dfot/eval_dfot_bundle_anchored.py b/egomimic/eval/dfot/eval_dfot_bundle_anchored.py new file mode 100644 index 000000000..fbbc9de0e --- /dev/null +++ b/egomimic/eval/dfot/eval_dfot_bundle_anchored.py @@ -0,0 +1,140 @@ +"""Anchored clean-history rollout eval for the 1D obs-action(+image) DFoT. + +Flat bundle ``[state | vae_mu_flat | action]``. Seeds the first +``n_context_frames`` bundles from GT (state + frozen-VAE mu + executed action), +pins them clean, and predicts the rest with a single-tensor anchored DDIM +rollout. Decodes the latent slice -> ``[GT|pred]`` video; extracts the action +slice -> per-step action MSE on the predicted future (policy variant). For the +world-model variant (``action_in_bundle=False``) the action is fed as +``external_cond`` and only video MSE is reported. +""" + +from __future__ import annotations + +from typing import Dict, List + +import cv2 +import numpy as np +import torch + +from egomimic.models.diffusion.diffusion.discrete_diffusion import DiscreteDiffusion +from egomimic.models.diffusion.sampling import vanilla_schedule +from egomimic.eval.core.eval_video import EvalVideo +from egomimic.eval.core.img_utils import img_chw_to_uint8 +from egomimic.eval.dfot._base import DFoTVideoEvalMixin +from egomimic.eval.dfot._sampling import anchored_ddim_rollout +from egomimic.rldb.embodiment.embodiment import get_embodiment_id + + + +class DFoTBundleAnchoredEval(DFoTVideoEvalMixin, EvalVideo): + def __init__( + self, n_context_frames: int = 4, rollout_steps: int = 32, + n_chunk_steps: int = 50, embodiment_name: str = "pushshapes_sim", + image_key: str = "front_img_1", recon_loss_n_frames: int = 20, + upscale_to: int = 384, limit_val_batches: int = 4, max_videos: int = 2, + video_subdir: str = "videos_bundle_anchored", viz_func=None, + transform_lists=None, + ): + super().__init__(limit_val_batches=limit_val_batches, viz_func=viz_func, + transform_lists=transform_lists, max_videos=max_videos) + self.store_dfot_knobs( + embodiment_name=embodiment_name, image_key=image_key, + video_subdir=video_subdir, recon_loss_n_frames=recon_loss_n_frames, + upscale_to=upscale_to, n_chunk_steps=n_chunk_steps, + ) + self.n_context_frames = int(n_context_frames) + self.rollout_steps = int(rollout_steps) + + @torch.no_grad() + def _rollout(self, algo, ctx_bundle, cond_seq, T, device): + outer, diff = algo.outer_stage, algo.diffusion + dts = int(diff.timesteps) if isinstance(diff, DiscreteDiffusion) else None + sched = vanilla_schedule(self.n_chunk_steps, T, discrete_timesteps=dts).to(device) + return anchored_ddim_rollout( + diff, algo.backbone, schedule=sched, context=ctx_bundle, + total_T=T, trailing_shape=(outer.bundle_dim,), device=device, + batch_size=1, external_cond=cond_seq, discrete_ts=dts, + ) + + def compute_metrics_and_viz(self, batch): + algo = self.model + metrics: Dict[str, torch.Tensor] = {} + images: Dict[int, np.ndarray] = {} + emb_id = get_embodiment_id(self.embodiment_name) + if emb_id not in batch: + return metrics, images + _batch = batch[emb_id] + outer = algo.outer_stage + if not hasattr(outer, "image_latent_slice") or self.image_key not in _batch: + return metrics, images + device = self.trainer.lightning_module.device + ac_key = algo.resolved_ac_keys[emb_id] + in_bundle = bool(getattr(outer, "action_in_bundle", True)) + imgs = _batch[self.image_key] + is_packed = _batch.get("_packed", False) + if is_packed: + cu = _batch["cu_seqlens"].to(imgs.device, dtype=torch.long) + n = min(int(cu.shape[0] - 1), self.max_videos or 99999) + spans = [(int(cu[i].item()), int(cu[i + 1].item())) for i in range(n)] + else: + n = min(imgs.shape[0], self.max_videos or imgs.shape[0]) + spans = [(i, None) for i in range(n)] + + a_sse = np.zeros(self.recon_loss_n_frames); a_n = np.zeros(self.recon_loss_n_frames, dtype=np.int64) + v_sse = np.zeros(self.recon_loss_n_frames); v_n = np.zeros(self.recon_loss_n_frames, dtype=np.int64) + all_frames: List[np.ndarray] = [] + for (s0, s1) in spans: + if is_packed: + Lep = s1 - s0; T = min(self.rollout_steps, Lep); sl = slice(s0, s0 + T) + pick = lambda k: _batch[k][sl] + else: + Lep = imgs.shape[1]; T = min(self.rollout_steps, Lep) + pick = lambda k: _batch[k][s0, :T] + ncx = min(self.n_context_frames, T - 1) + img_seq = pick(self.image_key).to(device).float() + if img_seq.max() > 1.5: img_seq = img_seq / 255.0 + act_seq = pick(ac_key).to(device).float() + state = torch.cat([pick(k).to(device).float() for k in outer.real_bundle_obs_keys], -1) + mu, _ = outer.vae.encode(img_seq) + mu_flat = mu.flatten(1) + + # context bundle for first ncx steps + pieces = [state[:ncx], mu_flat[:ncx]] + if in_bundle: pieces.append(act_seq[:ncx]) + ctx_bundle = torch.cat(pieces, -1).unsqueeze(0) + cond_seq = None if in_bundle else act_seq.unsqueeze(0) + + pred = self._rollout(algo, ctx_bundle, cond_seq, T, device).squeeze(0) # (T, bundle_dim) + lat = pred[:, outer.image_latent_slice].reshape(T, *outer.latent_shape) + pred_frames = outer.vae.decode(lat) + + # video MSE + m = min(self.recon_loss_n_frames, T) + vmse = ((pred_frames[:m] - img_seq[:m]) ** 2).mean(dim=(1, 2, 3)).detach().cpu().numpy() + for t in range(m): v_sse[t] += float(vmse[t]); v_n[t] += 1 + # action MSE on predicted future + if in_bundle: + pa = pred[:, outer.action_slice] + gf = act_seq[ncx:]; pf = pa[ncx:] + ma = min(self.recon_loss_n_frames, gf.shape[0]) + amse = ((pf[:ma] - gf[:ma]) ** 2).mean(dim=-1).detach().cpu().numpy() + for t in range(ma): a_sse[t] += float(amse[t]); a_n[t] += 1 + + for t in range(pred_frames.shape[0]): + g = cv2.resize(img_chw_to_uint8(img_seq[t]), (self.upscale_to,) * 2, interpolation=cv2.INTER_NEAREST) + p = cv2.resize(img_chw_to_uint8(pred_frames[t]), (self.upscale_to,) * 2, interpolation=cv2.INTER_NEAREST) + all_frames.append(np.concatenate([g, p], axis=1)) + + for t in range(self.recon_loss_n_frames): + if v_n[t] > 0: + metrics[f"Valid/emb{emb_id}_video_mse_step_{t:02d}"] = torch.tensor(v_sse[t] / v_n[t], device=device) + if a_n[t] > 0: + metrics[f"Valid/emb{emb_id}_action_mse_step_{t:02d}"] = torch.tensor(a_sse[t] / a_n[t], device=device) + if v_n.sum() > 0: + metrics[f"Valid/emb{emb_id}_video_mse_first{self.recon_loss_n_frames}"] = torch.tensor(v_sse.sum() / v_n.sum(), device=device) + if a_n.sum() > 0: + metrics[f"Valid/emb{emb_id}_action_mse_first{self.recon_loss_n_frames}"] = torch.tensor(a_sse.sum() / a_n.sum(), device=device) + if all_frames: + images[emb_id] = np.stack(all_frames, axis=0) + return metrics, images diff --git a/egomimic/eval/dfot/eval_dfot_policy.py b/egomimic/eval/dfot/eval_dfot_policy.py new file mode 100644 index 000000000..bd98bc1c7 --- /dev/null +++ b/egomimic/eval/dfot/eval_dfot_policy.py @@ -0,0 +1,334 @@ +"""Action-prediction evals for the 2D policy (``SpatialObsActionPolicyDFoTOuterStage``). + +This module hosts the policy-action pair (COMBINE B — merged from the former +``eval_dfot_policy_action`` + ``eval_dfot_policy_receding_horizon`` modules; the +two evaluators share ``_rollout`` / ``_ddim_from_v``, so they live together): + + * :class:`DFoTPolicyActionEval` — whole-horizon: denoise the WHOLE T future + obs-latent + action chunk jointly in one pass. + * :class:`DFoTPolicyRecedingHorizonEval` (subclass) — deployment metric: + slide an anchor across the episode, predict only the next ``k`` actions per + anchor via a tiny ``T = n_context + k`` rollout, score vs GT pooled over all + anchors x ``max_episodes``. + +``DFoTPolicyActionEval`` implements the **clean-history -> predict-next-chunk** +rollout (teacher-forced): the first ``n_context_frames`` obs-latent + action +tokens are pinned CLEAN from GT (the observed history + executed actions), and +the remaining obs-latent AND action tokens are denoised jointly — the DiT3D +backbone returns ``(v_image, v_action)`` and we run a structured DDIM step on +each stream. The predicted FUTURE actions (steps >= n_context) are compared to +GT actions (action MSE), and the predicted latents are VAE-decoded into a +``[GT|pred]`` video. Future obs are predicted (not GT), so the action chunk +conditions only on the clean past + its own predicted future, not on leaked +future frames. + +NOTE: state enters as external_cond and is currently supplied from GT for all +steps (teacher-forced); a fully causal closed-loop policy would predict/withhold +future state. This eval verifies the action-token mechanism end to end. +""" + +from __future__ import annotations + +from typing import Dict, List + +import cv2 +import numpy as np +import torch + +from egomimic.models.diffusion.diffusion.discrete_diffusion import DiscreteDiffusion +from egomimic.models.diffusion.sampling import vanilla_schedule +from egomimic.eval.core.eval_video import EvalVideo +from egomimic.eval.core.img_utils import img_chw_to_uint8 +from egomimic.eval.dfot._base import DFoTVideoEvalMixin +from egomimic.rldb.embodiment.embodiment import get_embodiment_id + + + +class DFoTPolicyActionEval(DFoTVideoEvalMixin, EvalVideo): + def __init__( + self, + n_context_frames: int = 4, + rollout_steps: int = 32, + n_chunk_steps: int = 50, + embodiment_name: str = "pushshapes_sim", + image_key: str = "front_img_1", + recon_loss_n_frames: int = 20, + upscale_to: int = 384, + limit_val_batches: int = 4, + max_videos: int = 2, + video_subdir: str = "videos_policy_action", + viz_func=None, + transform_lists=None, + ): + super().__init__( + limit_val_batches=limit_val_batches, viz_func=viz_func, + transform_lists=transform_lists, max_videos=max_videos, + ) + self.store_dfot_knobs( + embodiment_name=embodiment_name, image_key=image_key, + video_subdir=video_subdir, recon_loss_n_frames=recon_loss_n_frames, + upscale_to=upscale_to, n_chunk_steps=n_chunk_steps, + ) + self.n_context_frames = int(n_context_frames) + self.rollout_steps = int(rollout_steps) + + def _ddim_from_v(self, diff, x, v, cur, nxt): + """One eta=0 DDIM step from a v-prediction. ``cur``/``nxt`` are (T,) + per-token levels; ``x`` is (1, T, *trailing).""" + T = x.shape[1] + pad = x.dim() - 2 + kBT = cur.clamp_min(0).long().unsqueeze(0) # (1, T) + x0 = diff.predict_start_from_v(x, kBT, v) + eps = diff.predict_noise_from_v(x, kBT, v) + an = diff.alphas_cumprod[nxt.clamp_min(0).long()] # (T,) + an = torch.where(nxt < 0, torch.ones_like(an), an) + an = an.reshape(1, T, *([1] * pad)) + c = (1.0 - an).clamp_min(0.0).sqrt() + return x0 * an.sqrt() + eps * c + + @torch.no_grad() + def _rollout(self, algo, latent_ctx, action_ctx, cond_seq, T, device): + """Structured clean-history rollout. latent_ctx (1,n,C,H,W), + action_ctx (1,n,A), cond_seq (1,T,proj). Returns (pred_latent + (1,T,C,H,W), pred_action (1,T,A)).""" + outer, diff = algo.outer_stage, algo.diffusion + n = latent_ctx.shape[1] + C, H, W = outer.bundle_shape + A = outer._action_dim + dts = int(diff.timesteps) if isinstance(diff, DiscreteDiffusion) else None + sched = vanilla_schedule(self.n_chunk_steps, T, discrete_timesteps=dts).to(device) + clean = -1 if dts is not None else 0.0 + sched = sched.clone() + sched[:, :n] = clean + + x_lat = torch.randn(1, T, C, H, W, device=device) + x_lat[:, :n] = latent_ctx + x_act = torch.randn(1, T, A, device=device) + x_act[:, :n] = action_ctx + for s in range(sched.shape[0] - 1): + k = sched[s].clamp_min(0).long().unsqueeze(0) # (1, T) + v_lat, v_act = algo.backbone( + x_lat, k, external_cond=cond_seq, action=x_act + ) + x_lat = self._ddim_from_v(diff, x_lat, v_lat, sched[s], sched[s + 1]) + x_act = self._ddim_from_v(diff, x_act, v_act, sched[s], sched[s + 1]) + x_lat[:, :n] = latent_ctx + x_act[:, :n] = action_ctx + return x_lat, x_act + + def compute_metrics_and_viz(self, batch): + algo = self.model + metrics: Dict[str, torch.Tensor] = {} + images: Dict[int, np.ndarray] = {} + emb_id = get_embodiment_id(self.embodiment_name) + if emb_id not in batch: + return metrics, images + _batch = batch[emb_id] + if self.image_key not in _batch: + return metrics, images + outer = algo.outer_stage + if not hasattr(outer, "_state_to_cond"): + return metrics, images # not the 2D policy + device = self.trainer.lightning_module.device + ac_key = algo.resolved_ac_keys[emb_id] + imgs = _batch[self.image_key] + is_packed = _batch.get("_packed", False) + if is_packed: + cu = _batch["cu_seqlens"].to(imgs.device, dtype=torch.long) + n = min(int(cu.shape[0] - 1), self.max_videos or 99999) + starts = [int(cu[i].item()) for i in range(n)] + lens = [int(cu[i + 1].item() - cu[i].item()) for i in range(n)] + else: + n = min(imgs.shape[0], self.max_videos or imgs.shape[0]) + starts = [None] * n + lens = [imgs.shape[1]] * n + + per_step_sse = np.zeros(self.recon_loss_n_frames) + per_step_n = np.zeros(self.recon_loss_n_frames, dtype=np.int64) + all_frames: List[np.ndarray] = [] + for ep in range(n): + T = min(self.rollout_steps, lens[ep]) + ncx = min(self.n_context_frames, T - 1) + if is_packed: + sl = slice(starts[ep], starts[ep] + T) + img_seq, act_seq = imgs[sl], _batch[ac_key][sl] + state = torch.cat([_batch[k][sl] for k in outer.bundle_obs_keys], -1) + else: + img_seq, act_seq = imgs[ep, :T], _batch[ac_key][ep, :T] + state = torch.cat([_batch[k][ep, :T] for k in outer.bundle_obs_keys], -1) + img_seq = img_seq.to(device).float() + if img_seq.max() > 1.5: + img_seq = img_seq / 255.0 + act_seq = act_seq.to(device).float() + cond = outer.state_only_proj(state.to(device).float()).unsqueeze(0) # (1,T,proj) + + mu, _ = outer.vae.encode(img_seq) + latent_all = outer.normalize_latent(mu) # (T,C,H,W) + latent_ctx = latent_all[:ncx].unsqueeze(0) + action_ctx = act_seq[:ncx].unsqueeze(0) + pred_lat, pred_act = self._rollout(algo, latent_ctx, action_ctx, cond, T, device) + + # ---- action MSE on the PREDICTED future (steps >= ncx) ---- + gt_fut = act_seq[ncx:] + pred_fut = pred_act.squeeze(0)[ncx:] + m = min(self.recon_loss_n_frames, gt_fut.shape[0]) + if m > 0: + amse = ((pred_fut[:m] - gt_fut[:m]) ** 2).mean(dim=-1).detach().cpu().numpy() + for t in range(m): + per_step_sse[t] += float(amse[t]); per_step_n[t] += 1 + + # ---- decode predicted latents -> [GT|pred] frames ---- + pred_frames = outer.vae.decode(outer.denormalize_latent(pred_lat.squeeze(0))) + for t in range(pred_frames.shape[0]): + gt_t = cv2.resize(img_chw_to_uint8(img_seq[t]), (self.upscale_to, self.upscale_to), interpolation=cv2.INTER_NEAREST) + pr_t = cv2.resize(img_chw_to_uint8(pred_frames[t]), (self.upscale_to, self.upscale_to), interpolation=cv2.INTER_NEAREST) + all_frames.append(np.concatenate([gt_t, pr_t], axis=1)) + + for t in range(self.recon_loss_n_frames): + if per_step_n[t] > 0: + metrics[f"Valid/emb{emb_id}_action_mse_step_{t:02d}"] = torch.tensor( + per_step_sse[t] / per_step_n[t], device=device) + if per_step_n.sum() > 0: + metrics[f"Valid/emb{emb_id}_action_mse_first{self.recon_loss_n_frames}"] = torch.tensor( + per_step_sse.sum() / per_step_n.sum(), device=device) + if all_frames: + images[emb_id] = np.stack(all_frames, axis=0) + return metrics, images + + +class DFoTPolicyRecedingHorizonEval(DFoTPolicyActionEval): + """Receding-horizon action eval for the 2D policy. + + The parent ``DFoTPolicyActionEval`` denoises the WHOLE T=32 future + action+latent chunk jointly in one pass — that is the policy's *worst case* + (28 future action tokens co-denoised under a shared per-step noise level), + and its non-monotonic per-step curve is an artifact of that joint long-chunk + coupling plus n=2 sampling noise, NOT the policy's deployable accuracy. + + This eval measures the DEPLOYMENT metric: slide an anchor across the whole + episode; at each anchor feed CLEAN GT obs-latent + action history of length + ``n_context`` and predict ONLY the next ``k`` actions via a tiny + ``T = n_context + k`` rollout (matching how a receding-horizon controller + runs: predict next chunk, execute, re-anchor on the observed history). Score + those k predicted actions against GT in normalized action space, pooled over + ALL anchors and ``max_episodes`` episodes — orders of magnitude more samples + than the old 2-episode single-chunk read. + + Headline: ``rh_k1_action_mse_overall``. If the short-horizon regime truly + holds (old step_00 = 0.046) it lands ~0.02-0.05 across the whole episode — + directly debunking the one-shot-long-chunk blow-up (old step_02 = 2.10). + + Reuses ``_rollout`` and ``_ddim_from_v`` from the parent VERBATIM. + """ + + def __init__( + self, + n_context_frames: int = 4, + n_chunk_steps: int = 50, + k_actions=(1, 2), + max_episodes: int = 8, + anchor_stride: int = 1, + embodiment_name: str = "pushshapes_sim", + image_key: str = "front_img_1", + limit_val_batches: int = 4, + max_videos: int = 2, + rollout_steps: int = 6, # accepted for ABI; T is computed as n_ctx+k + recon_loss_n_frames: int = 20, # accepted for ABI; unused in the RH metric + video_subdir: str = "videos_policy_rh", + viz_func=None, + transform_lists=None, + ): + super().__init__( + n_context_frames=n_context_frames, rollout_steps=rollout_steps, + n_chunk_steps=n_chunk_steps, embodiment_name=embodiment_name, + image_key=image_key, recon_loss_n_frames=recon_loss_n_frames, + limit_val_batches=limit_val_batches, max_videos=max_videos, + video_subdir=video_subdir, viz_func=viz_func, transform_lists=transform_lists, + ) + self.k_actions = tuple(int(k) for k in k_actions) + self.max_episodes = int(max_episodes) + self.anchor_stride = max(1, int(anchor_stride)) + + @torch.no_grad() + def compute_metrics_and_viz(self, batch): + algo = self.model + metrics: Dict[str, torch.Tensor] = {} + emb_id = get_embodiment_id(self.embodiment_name) + if emb_id not in batch: + return metrics, {} + _batch = batch[emb_id] + if self.image_key not in _batch: + return metrics, {} + outer = algo.outer_stage + if not hasattr(outer, "_state_to_cond"): + return metrics, {} # not the 2D policy + device = self.trainer.lightning_module.device + ac_key = algo.resolved_ac_keys[emb_id] + imgs = _batch[self.image_key] + is_packed = _batch.get("_packed", False) + + # ---- per-episode spans, capped by max_episodes (NOT max_videos) ---- + if is_packed: + cu = _batch["cu_seqlens"].to(imgs.device, dtype=torch.long) + ne = min(int(cu.shape[0] - 1), self.max_episodes) + spans = [(int(cu[i].item()), int(cu[i + 1].item() - cu[i].item())) for i in range(ne)] + else: + ne = min(imgs.shape[0], self.max_episodes) + spans = [(ep, imgs.shape[1]) for ep in range(ne)] + + Ks = self.k_actions + n_ctx = self.n_context_frames + sse = {k: np.zeros(k) for k in Ks} + cnt = {k: np.zeros(k, dtype=np.int64) for k in Ks} + n_anchors_total = 0 + + for (start, L) in spans: + if L < n_ctx + max(Ks): + continue + # ---- encode the whole episode ONCE, reuse across all anchors ---- + if is_packed: + sl = slice(start, start + L) + img_seq = imgs[sl] + act_seq = _batch[ac_key][sl] + state = torch.cat([_batch[key][sl] for key in outer.bundle_obs_keys], -1) + else: + img_seq = imgs[start, :L] + act_seq = _batch[ac_key][start, :L] + state = torch.cat([_batch[key][start, :L] for key in outer.bundle_obs_keys], -1) + img_seq = img_seq.to(device).float() + if img_seq.max() > 1.5: + img_seq = img_seq / 255.0 + act_seq = act_seq.to(device).float() # already normalized + mu, _ = outer.vae.encode(img_seq) + latent_all = outer.normalize_latent(mu) # (L, C, H, W) + cond_all = outer.state_only_proj(state.to(device).float()) # (L, proj) + + for k in Ks: + T = n_ctx + k + last_anchor = L - k # inclusive + for a in range(n_ctx, last_anchor + 1, self.anchor_stride): + lo = a - n_ctx + latent_ctx = latent_all[lo:a].unsqueeze(0) # (1, n_ctx, C, H, W) clean + action_ctx = act_seq[lo:a].unsqueeze(0) # (1, n_ctx, A) clean + cond = cond_all[lo:a + k].unsqueeze(0) # (1, T, proj) + _, pred_act = self._rollout(algo, latent_ctx, action_ctx, cond, T, device) + pred_next = pred_act.squeeze(0)[n_ctx:n_ctx + k] # (k, A) predicted future + gt_next = act_seq[a:a + k] # (k, A) + off = ((pred_next - gt_next) ** 2).mean(dim=-1).detach().cpu().numpy() + for o in range(k): + sse[k][o] += float(off[o]) + cnt[k][o] += 1 + if k == Ks[0]: + n_anchors_total += 1 + + for k in Ks: + for o in range(k): + if cnt[k][o] > 0: + metrics[f"Valid/emb{emb_id}_rh_k{k}_action_mse_off{o:02d}"] = torch.tensor( + sse[k][o] / cnt[k][o], device=device) + if cnt[k].sum() > 0: + metrics[f"Valid/emb{emb_id}_rh_k{k}_action_mse_overall"] = torch.tensor( + sse[k].sum() / cnt[k].sum(), device=device) + if n_anchors_total > 0: + metrics[f"Valid/emb{emb_id}_rh_n_anchors"] = torch.tensor(float(n_anchors_total), device=device) + return metrics, {} diff --git a/egomimic/eval/dfot/eval_dfot_self_rollout.py b/egomimic/eval/dfot/eval_dfot_self_rollout.py new file mode 100644 index 000000000..4c0d9a6ba --- /dev/null +++ b/egomimic/eval/dfot/eval_dfot_self_rollout.py @@ -0,0 +1,361 @@ +"""DFoT self-rollout (world-model) evaluator. + +For the joint-obs+action DFoT variant (``ObsActionDFoTOuterStage``), the +model denoises ``concat([state, action])`` jointly per step. The state +portion of the bundle IS the model's predicted future world state. +This evaluator turns that predicted-state trajectory into an actual +RENDERED video frame-by-frame by using the PushShapesEnv as a stateless +renderer: + + 1. Pick the first ``max_videos`` episodes from the val batch. + 2. Take ONLY the t=0 image as the AdaLN cond (no GT state / action + used at inference). Broadcast the cond across ``rollout_steps`` + tokens. + 3. Run chunk-mode DDIM sampling -> predicted bundle + ``(rollout_steps, bundle_dim)``. + 4. Slice the state portion via ``outer_stage`` config: + ``state = bundle[..., :sum(bundle_obs_dims)]`` + and unnormalize per obs key to recover world-frame state. + 5. For each step, ``env.set_state(agent_pos, object_pose)`` and + ``env.render()`` to produce a frame. Stack -> mp4. + +The first 5 dims of the bundle's obs portion are +``(agent_x, agent_y, obj_x, obj_y, obj_theta)`` (the pushshapes +state layout) — same split used by ``eval_sim._state_to_init``. If the +bundle obs is shorter than 5 dims, the env can't be rendered and the +eval falls back to the path-overlay viz from the previous version. + +The starting frame (rendered from the GT initial state) is prepended +to each video so viewers can see the launch condition before the model +takes over. +""" + +from __future__ import annotations + +from typing import Dict, List + +import cv2 +import numpy as np +import torch + +from egomimic.rldb.embodiment.pushshapes_sim import _state_to_init +from egomimic.eval.core.eval_video import EvalVideo +from egomimic.rldb.embodiment.embodiment import get_embodiment_id + + +def _img_chw_to_uint8(img_chw: torch.Tensor) -> np.ndarray: + x = img_chw.detach().cpu().float().numpy() + if x.max() <= 1.5: + x = x * 255.0 + return x.clip(0, 255).astype(np.uint8) + + +def _path_overlay_frame( + img_chw_uint8: np.ndarray, + path_xy: np.ndarray, + cur_idx: int, + world_size: float, + upscale_to: int = 512, +) -> np.ndarray: + """Fallback viz when the env renderer is unavailable: trail overlay + on the starting image. Same as the previous implementation.""" + img_hwc = np.transpose(img_chw_uint8, (1, 2, 0)) + img_hwc = cv2.resize( + img_hwc, (upscale_to, upscale_to), interpolation=cv2.INTER_NEAREST + ) + H, W, _ = img_hwc.shape + out = np.ascontiguousarray(img_hwc.copy()) + sx = W / float(world_size) + sy = H / float(world_size) + pts = [] + for p in path_xy[: cur_idx + 1]: + cx = int(round(float(p[0]) * sx)) + cy = int(round(float(p[1]) * sy)) + if 0 <= cx < W and 0 <= cy < H: + pts.append((cx, cy)) + for i in range(1, len(pts)): + cv2.line(out, pts[i - 1], pts[i], (180, 220, 255), thickness=2) + if pts: + cv2.circle(out, pts[-1], 8, (0, 0, 0), thickness=-1) + cv2.circle(out, pts[-1], 6, (255, 255, 255), thickness=-1) + return out + + +class DFoTSelfRolloutEval(EvalVideo): + """World-model self-rollout for joint obs+action DFoT. + + Args: + rollout_steps: number of bundle tokens to generate from the + single starting image. mp4 length = rollout_steps (+ one + launch-condition frame at the start). + world_size: physics world extent for pixel scaling (PushShapes + uses 512). + image_key: which obs key carries the env image. + agent_state_slice: slice into the predicted state vector that + holds (agent_x, agent_y). Default ``[0, 2]``. + env_kwargs: forwarded to ``PushShapesEnv(**)`` for rendering + (object_shape, pusher_shape, obstacle_level, image_size, ...). + ``render_mode`` defaults to ``"rgb_array"``. + cfg_scale: CFG scale for sampling. 1.0 = off. + upscale_to: per-frame side after upscaling (env's image_size is + usually 96; we upscale for clarity). + max_videos: per-pass cap on rendered episodes. + limit_val_batches: per-pass cap on val batches. + embodiment_name: which embodiment the packed batch carries. + viz_func / transform_lists: forwarded to ``EvalVideo`` base. + """ + + def __init__( + self, + rollout_steps: int = 64, + world_size: float = 512.0, + image_key: str = "front_img_1", + agent_state_slice: slice | tuple = (0, 2), + env_kwargs: dict | None = None, + cfg_scale: float = 1.0, + upscale_to: int = 512, + max_videos: int = 2, + limit_val_batches: int = 4, + embodiment_name: str = "pushshapes_sim", + video_subdir: str = "videos_self_rollout", + viz_func=None, + transform_lists=None, + ): + super().__init__( + limit_val_batches=limit_val_batches, + viz_func=viz_func, + transform_lists=transform_lists, + max_videos=max_videos, + ) + # Per-eval output namespace so DFoTValEval and DFoTSelfRolloutEval + # can be run side-by-side at the top of an EvalList without + # overwriting each other's validation_video_*.mp4 files. + self._video_subdir = str(video_subdir) + self.rollout_steps = int(rollout_steps) + self.world_size = float(world_size) + self.image_key = image_key + if isinstance(agent_state_slice, slice): + self.agent_state_slice = agent_state_slice + else: + start, stop = agent_state_slice + self.agent_state_slice = slice(int(start), int(stop)) + self.env_kwargs = dict(env_kwargs or {}) + self.cfg_scale = float(cfg_scale) + self.upscale_to = int(upscale_to) + self.embodiment_name = embodiment_name + self._env = None # lazy-init on first call + + # Override the base ``EvalVideo.video_dir`` so we don't collide with + # DFoTValEval's mp4 filenames at val time. + def video_dir(self): + import os + + return os.path.join(self.root_dir(), self._video_subdir) + + # ------------------------------------------------------------------ + # Env: lazy PushShapesEnv for state-driven rendering. No physics + # step — we use set_state + render() only. + # ------------------------------------------------------------------ + + def _get_env(self): + if self._env is not None: + return self._env + try: + from Tsimulation.pushshapes import PushShapesEnv + except Exception: + return None + kw = dict(self.env_kwargs) + kw.setdefault("render_mode", "rgb_array") + self._env = PushShapesEnv(**kw) + return self._env + + # ------------------------------------------------------------------ + # Bundle rollout from one starting image. + # ------------------------------------------------------------------ + + @torch.no_grad() + def _rollout_from_image( + self, algo, img_chw_norm: torch.Tensor, emb_id: int + ) -> torch.Tensor: + device = img_chw_norm.device + T = self.rollout_steps + # CondEncoderModule's image branch broadcasts (B, C, H, W) over + # T_action; a 5-D tensor is taken as already-temporal and skips + # the broadcast, which mismatches the backbone at T=rollout_steps. + img_bchw = img_chw_norm.unsqueeze(0) + obs_for_cond = {self.image_key: img_bchw} + cond_dict = algo.cond_encoder.encode(obs_for_cond, T_action=T) + cond = cond_dict.get(algo.cond_output_key) + if cond is None: + raise RuntimeError( + "cond_encoder produced no fused_cond — DFoTSelfRolloutEval " + "requires image conditioning." + ) + bundle = algo._sample_chunk(B=1, T=T, cond=cond, device=device) + return bundle.squeeze(0) # (T, bundle_dim) + + # ------------------------------------------------------------------ + # Frame renderer: per-step set_state -> render -> resize. + # ------------------------------------------------------------------ + + def _render_states_to_frames( + self, + state_world_seq: np.ndarray, # (T, state_dim) — agent_xy + obj_xytheta + prepend_state: np.ndarray | None = None, + ) -> List[np.ndarray] | None: + env = self._get_env() + if env is None: + return None + if not hasattr(env, "set_state"): + return None + + frames: List[np.ndarray] = [] + # reset() must happen once before set_state(); we only need this + # once per episode rollout since set_state overrides everything. + try: + env.reset() + except Exception: + return None + + def _render_state(s_vec: np.ndarray) -> np.ndarray | None: + if s_vec.shape[0] < 5: + return None + agent_pos, obj_pose = _state_to_init(s_vec) + env.set_state(agent_pos=agent_pos, object_pose=obj_pose) + rgb = env.render() + if rgb is None: + return None + if rgb.shape[0] != self.upscale_to: + rgb = cv2.resize( + rgb, + (self.upscale_to, self.upscale_to), + interpolation=cv2.INTER_NEAREST, + ) + return np.ascontiguousarray(rgb) + + if prepend_state is not None: + f0 = _render_state(prepend_state) + if f0 is not None: + frames.append(f0) + + for t in range(state_world_seq.shape[0]): + f = _render_state(state_world_seq[t]) + if f is None: + return None + frames.append(f) + return frames + + # ------------------------------------------------------------------ + # Main entrypoint. + # ------------------------------------------------------------------ + + def compute_metrics_and_viz(self, batch): + algo = self.model + metrics: Dict[str, torch.Tensor] = {} + images_dict: Dict[int, np.ndarray] = {} + + emb_id = get_embodiment_id(self.embodiment_name) + if emb_id not in batch: + return metrics, images_dict + _batch = batch[emb_id] + ac_key = algo.resolved_ac_keys[emb_id] + + if self.image_key not in _batch: + return metrics, images_dict + imgs = _batch[self.image_key] + + outer = algo.outer_stage + action_slice = getattr(outer, "action_slice", slice(0, algo.action_dim)) + obs_end = action_slice.start or 0 + has_obs_in_bundle = obs_end > 0 + obs_keys = list(getattr(outer, "bundle_obs_keys", [])) + obs_dims = list(getattr(outer, "bundle_obs_dims", [])) + + # Pick starting images + (optionally) GT starting states. + if not _batch.get("_packed", False): + B = imgs.shape[0] + n = min(B, self.max_videos or B) + start_imgs = [imgs[i, 0] for i in range(n)] + gt_starts = [] + for i in range(n): + if obs_keys and obs_keys[0] in _batch: + gt_starts.append(_batch[obs_keys[0]][i, 0]) + else: + gt_starts.append(None) + else: + cu = _batch["cu_seqlens"].to(imgs.device, dtype=torch.long) + n = min(int(cu.shape[0] - 1), self.max_videos or 99999) + start_imgs = [imgs[int(cu[i].item())] for i in range(n)] + gt_starts = [] + for i in range(n): + if obs_keys and obs_keys[0] in _batch: + gt_starts.append(_batch[obs_keys[0]][int(cu[i].item())]) + else: + gt_starts.append(None) + + frames_for_emb: List[np.ndarray] = [] + for img_chw, gt_start in zip(start_imgs, gt_starts): + bundle_pred = self._rollout_from_image(algo, img_chw, emb_id) + + if has_obs_in_bundle and obs_keys: + obs_pred = bundle_pred[..., :obs_end] + splits = list(torch.split(obs_pred, obs_dims, dim=-1)) + obs_unnorm_dict = {k: v for k, v in zip(obs_keys, splits)} + obs_unnorm = algo.norm_stats.unnormalize(obs_unnorm_dict, emb_id) + state_world = obs_unnorm[obs_keys[0]] # (T, dim_0) + + # Prepend GT initial state so the video starts from the + # actual launch condition (rendered). + gt_start_world = None + if gt_start is not None: + gt_start_norm = {obs_keys[0]: gt_start.unsqueeze(0)} + gt_start_world_t = algo.norm_stats.unnormalize( + gt_start_norm, emb_id + )[obs_keys[0]].squeeze(0) + gt_start_world = gt_start_world_t.detach().cpu().numpy() + + state_world_np = state_world.detach().cpu().numpy() + rendered = self._render_states_to_frames( + state_world_np, + prepend_state=gt_start_world, + ) + if rendered is not None: + frames_for_emb.extend(rendered) + continue + # Renderer unavailable -> fall back to overlay. + xy = state_world[..., self.agent_state_slice] + xy_np = xy.detach().cpu().numpy() + img_uint8 = _img_chw_to_uint8(img_chw) + for t in range(self.rollout_steps): + frames_for_emb.append( + _path_overlay_frame( + img_uint8, + xy_np, + t, + self.world_size, + self.upscale_to, + ) + ) + else: + # Vanilla DFoT (bundle = action only) — no state to render. + act_pred = bundle_pred[..., action_slice] + act_unnorm = algo.norm_stats.unnormalize({ac_key: act_pred}, emb_id)[ + ac_key + ] + xy = act_unnorm[..., :2] + xy_np = xy.detach().cpu().numpy() + img_uint8 = _img_chw_to_uint8(img_chw) + for t in range(self.rollout_steps): + frames_for_emb.append( + _path_overlay_frame( + img_uint8, + xy_np, + t, + self.world_size, + self.upscale_to, + ) + ) + + if frames_for_emb: + images_dict[emb_id] = np.stack(frames_for_emb, axis=0) + return metrics, images_dict diff --git a/egomimic/eval/dfot/eval_dfot_video_rollout.py b/egomimic/eval/dfot/eval_dfot_video_rollout.py new file mode 100644 index 000000000..fc5fe4a6d --- /dev/null +++ b/egomimic/eval/dfot/eval_dfot_video_rollout.py @@ -0,0 +1,265 @@ +"""Family-agnostic DFoT video self-rollout evaluator (COMBINE A). + +ONE evaluator drives all three DFoT video-rollout families. Each family's +outer stage owns a ``rollout_video_episode`` hook (decode-on-outer-stage) +that turns its sampler output into ``(T, 3, H, W)`` pixel frames: + + * :class:`ObsActionImageDFoTOuterStage` — unconditional chunk/AR bundle + sampler, slice the flat VAE-latent portion, frozen-VAE decode. Single + panel (t=0 GT prepended). Metric prefix ``video``. + * :class:`ImageSpatialDFoTOuterStage` — per-step (state,action) cond, + conditional chunk/AR spatial-latent sampler (optional GT-context + anchor), frozen-VAE decode. Side-by-side [GT|pred]. Prefix ``spatial``. + * :class:`PixelSpatialDFoTOuterStage` — sliding-window pixel-space rollout + anchored on the first GT frame(s), no VAE. Side-by-side [GT|pred]. + Prefix ``pixel``. Also reports PSNR / SSIM / LPIPS. + +This eval owns the family-INVARIANT skeleton: episode indexing (packed / +padded), per-step recon-MSE accumulation, panel assembly, perceptual-metric +averaging, and mp4 emission. The family-VARIANT pieces live on the outer +stage; the eval dispatches on three class attributes it advertises: +``video_metric_prefix`` / ``video_panel`` / ``video_has_extra_metrics``. + +The old per-family classes are kept as compat aliases at the bottom +(:class:`DFoTSpatialVideoRolloutEval`, :class:`DFoTPixelVideoRolloutEval`) — +the unified ``__init__`` is a true superset and every old config now passes +its knobs explicitly, so the aliases are pure ``_target_`` redirects. + +Output: ``//epoch_N// +validation_video_N.mp4``. +""" + +from __future__ import annotations + +from typing import Dict, List + +import cv2 +import numpy as np +import torch + +from egomimic.eval.core.eval_video import EvalVideo +from egomimic.eval.core.img_utils import img_chw_to_uint8 +from egomimic.eval.dfot._base import DFoTVideoEvalMixin +from egomimic.rldb.embodiment.embodiment import get_embodiment_id + + +class DFoTVideoRolloutEval(DFoTVideoEvalMixin, EvalVideo): + """Family-agnostic DFoT video rollout — delegates the family-specific + sampler + decode to ``outer_stage.rollout_video_episode``. + + Args (union of the three families' knobs; per-knob defaults preserved): + rollout_steps: bundle tokens / frames to generate. (video/spatial + default 64; pixel configs pass 9.) + upscale_to: post-render upscale side per panel; default 384. + max_videos: cap on episodes rendered per val pass. + limit_val_batches: cap on val batches per val pass. + embodiment_name: which embodiment the packed val batches carry. + image_key: which obs key carries the image stream. + recon_loss_n_frames: leading steps scored for the scalar recon MSE. + n_context_frames: GT-context anchor length. 0 = unconditional + (video / spatial-default); pixel uses >=1 (its config passes 1). + rollout_window: pixel sliding-window size (pixel-only knob). + mode: ``"chunk"`` (default) or ``"ar"`` (staircase). The video + family's chunk mode = ``algo._sample_chunk`` (uniform per-token + noise + DDIM); spatial/pixel chunk = ``vanilla_schedule``. AR = + staircase per-token schedule (mirrors ``DFoTValEval``). + ar_chunk_size, ar_step_size, n_chunk_steps, cfg_scale: sampler knobs. + video_subdir: output subdir under root_dir (per-family namespaced). + """ + + def __init__( + self, + rollout_steps: int = 64, + upscale_to: int = 384, + max_videos: int = 2, + limit_val_batches: int = 4, + embodiment_name: str = "pushshapes_sim", + image_key: str = "front_img_1", + recon_loss_n_frames: int = 10, + n_context_frames: int = 0, + rollout_window: int = 9, + mode: str = "chunk", + ar_chunk_size: int = 1, + ar_step_size: int = 1, + n_chunk_steps: int = 50, + cfg_scale: float = 1.0, + video_subdir: str = "videos_video_rollout", + viz_func=None, + transform_lists=None, + ): + super().__init__( + limit_val_batches=limit_val_batches, + viz_func=viz_func, + transform_lists=transform_lists, + max_videos=max_videos, + ) + if mode not in {"chunk", "ar"}: + raise ValueError(f"mode must be 'chunk' or 'ar', got {mode!r}") + self.store_dfot_knobs( + embodiment_name=embodiment_name, image_key=image_key, + video_subdir=video_subdir, recon_loss_n_frames=recon_loss_n_frames, + upscale_to=upscale_to, n_chunk_steps=n_chunk_steps, + ) + self.rollout_steps = int(rollout_steps) + self.n_context_frames = int(n_context_frames) + self.rollout_window = int(rollout_window) + self.mode = mode + self.ar_chunk_size = int(ar_chunk_size) + self.ar_step_size = int(ar_step_size) + self.cfg_scale = float(cfg_scale) + + # ------------------------------------------------------------------ + # Panel renderers — one per ``outer_stage.video_panel`` layout. Each + # reproduces its family's original frame construction byte-for-byte. + # ------------------------------------------------------------------ + + def _panel_single_t0prepend(self, res) -> List[np.ndarray]: + """obs+action+image: prepend the GT t=0 launch frame, then preds.""" + frames: List[np.ndarray] = [] + t0_uint = img_chw_to_uint8(res.gt_t0_chw) + t0_uint = cv2.resize( + t0_uint, (self.upscale_to, self.upscale_to), + interpolation=cv2.INTER_NEAREST, + ) + frames.append(np.ascontiguousarray(t0_uint)) + pred_frames = res.pred_frames + for t in range(pred_frames.shape[0]): + f = img_chw_to_uint8(pred_frames[t]) + f = cv2.resize( + f, (self.upscale_to, self.upscale_to), + interpolation=cv2.INTER_NEAREST, + ) + frames.append(np.ascontiguousarray(f)) + return frames + + def _panel_sidebyside(self, res) -> List[np.ndarray]: + """spatial / pixel: [GT | pred] per step.""" + frames: List[np.ndarray] = [] + pred_frames = res.pred_frames + gt_seq = res.gt_panel_raw + gt_scale = 255.0 if gt_seq.max() > 1.5 else 1.0 + for t in range(pred_frames.shape[0]): + gt_t = img_chw_to_uint8(gt_seq[t] / gt_scale) + pr_t = img_chw_to_uint8(pred_frames[t]) + gt_t = cv2.resize( + gt_t, (self.upscale_to, self.upscale_to), + interpolation=cv2.INTER_NEAREST, + ) + pr_t = cv2.resize( + pr_t, (self.upscale_to, self.upscale_to), + interpolation=cv2.INTER_NEAREST, + ) + frames.append(np.concatenate([gt_t, pr_t], axis=1)) + return frames + + def compute_metrics_and_viz(self, batch): + algo = self.model + metrics: Dict[str, torch.Tensor] = {} + images_dict: Dict[int, np.ndarray] = {} + + emb_id = get_embodiment_id(self.embodiment_name) + if emb_id not in batch: + return metrics, images_dict + _batch = batch[emb_id] + if self.image_key not in _batch: + return metrics, images_dict + + outer = algo.outer_stage + # This eval only drives outer stages that advertise the rollout hook. + if not hasattr(outer, "rollout_video_episode"): + return metrics, images_dict + + prefix = outer.video_metric_prefix + panel = outer.video_panel + has_extra = bool(getattr(outer, "video_has_extra_metrics", False)) + + device = self.trainer.lightning_module.device + imgs = _batch[self.image_key] + is_packed = _batch.get("_packed", False) + + # Episode indexing (packed / padded). + if is_packed: + cu = _batch["cu_seqlens"].to(imgs.device, dtype=torch.long) + n = min(int(cu.shape[0] - 1), self.max_videos or 99999) + ep_starts = [int(cu[i].item()) for i in range(n)] + ep_lens = [int(cu[i + 1].item() - cu[i].item()) for i in range(n)] + else: + B = imgs.shape[0] + n = min(B, self.max_videos or B) + ep_starts = [None] * n + ep_lens = [imgs.shape[1]] * n + + per_step_sse = np.zeros(self.recon_loss_n_frames, dtype=np.float64) + per_step_n = np.zeros(self.recon_loss_n_frames, dtype=np.int64) + extra_sums: Dict[str, torch.Tensor] = {} + all_frames: List[np.ndarray] = [] + + for ep_idx in range(n): + res = outer.rollout_video_episode( + self, algo, _batch, emb_id, ep_idx, + ep_starts[ep_idx], ep_lens[ep_idx], device, + ) + + # ---- per-step MSE vs aligned GT ---- + pred = res.pred_frames + gt_f = res.gt_for_mse + n_cmp = min(self.recon_loss_n_frames, pred.shape[0], gt_f.shape[0]) + pred_f = pred[:n_cmp].to(device).float() + gt_cmp = gt_f[:n_cmp].to(device).float() + mse_per_step = ( + (pred_f - gt_cmp) ** 2 + ).mean(dim=(1, 2, 3)).detach().cpu().numpy() + for t in range(n_cmp): + per_step_sse[t] += float(mse_per_step[t]) + per_step_n[t] += 1 + + # ---- per-episode extra (perceptual) metrics, summed ---- + for k, v in res.extra_metrics.items(): + extra_sums.setdefault( + f"Valid/emb{emb_id}_{prefix}_{k}", + torch.tensor(0.0, device=device), + ) + extra_sums[f"Valid/emb{emb_id}_{prefix}_{k}"] += v + + # ---- panel frames ---- + if panel == "single_t0prepend": + all_frames.extend(self._panel_single_t0prepend(res)) + elif panel == "sidebyside": + all_frames.extend(self._panel_sidebyside(res)) + else: + raise ValueError(f"unknown video_panel {panel!r}") + + # ---- emit per-step + scalar recon MSE ---- + for t in range(self.recon_loss_n_frames): + if per_step_n[t] > 0: + metrics[ + f"Valid/emb{emb_id}_{prefix}_recon_mse_step_{t:02d}" + ] = torch.tensor( + per_step_sse[t] / per_step_n[t], device=device + ) + if per_step_n.sum() > 0: + metrics[ + f"Valid/emb{emb_id}_{prefix}_recon_mse_first{self.recon_loss_n_frames}" + ] = torch.tensor( + per_step_sse.sum() / per_step_n.sum(), device=device + ) + + # ---- average perceptual metrics over episodes ---- + if has_extra and n > 0: + for key, total in extra_sums.items(): + metrics[key] = total / n + + if all_frames: + images_dict[emb_id] = np.stack(all_frames, axis=0) + return metrics, images_dict + + +# --------------------------------------------------------------------------- +# Compat aliases (COMBINE A). The unified ``DFoTVideoRolloutEval`` is a true +# superset: every knob each old config passed is still accepted with identical +# semantics, and the family-specific behaviour now lives on the outer stage's +# ``rollout_video_episode`` hook. So the old per-family classes are pure +# redirects — configs that still name them resolve to the unified eval. +# --------------------------------------------------------------------------- +DFoTSpatialVideoRolloutEval = DFoTVideoRolloutEval +DFoTPixelVideoRolloutEval = DFoTVideoRolloutEval diff --git a/egomimic/eval/tf/__init__.py b/egomimic/eval/tf/__init__.py new file mode 100644 index 000000000..d8ab130dd --- /dev/null +++ b/egomimic/eval/tf/__init__.py @@ -0,0 +1,15 @@ +"""TEACHER-FORCED evaluators (DESIGN.md §2 ``egomimic/eval/{tf}``). + +The teacher-forced (GT-conditioned) prediction/overlay evaluators — the cheap, +dense, per-step signal you render + iterate on BEFORE spending on closed-loop +sim. Curated here in DESIGN.md step 8 (``git mv``, no behaviour change): + + * :class:`DFoTValEval` — TF action-prediction val with viz overlay. + * :class:`DFoTControllerTFEval` — TF sanity check for the spatial_rh + closed-loop controller (drives the inference step on GT history). +""" + +from egomimic.eval.tf.eval_dfot_val import DFoTValEval +from egomimic.eval.tf.eval_dfot_controller_tf import DFoTControllerTFEval + +__all__ = ["DFoTValEval", "DFoTControllerTFEval"] diff --git a/egomimic/eval/tf/eval_dfot_controller_tf.py b/egomimic/eval/tf/eval_dfot_controller_tf.py new file mode 100644 index 000000000..1997952d4 --- /dev/null +++ b/egomimic/eval/tf/eval_dfot_controller_tf.py @@ -0,0 +1,121 @@ +"""Teacher-forced sanity check for the spatial_rh closed-loop controller. + +Drives ``algo._inference_step_spatial_rh`` through GT val episodes one frame at +a time, but OVERWRITES the controller's recorded action with the GT action each +step (so its context is clean GT obs + GT actions — no own-action drift). Then +compares the controller's *predicted* action to GT in normalized space. + +Decisive read: + * overall_mse ~= offline RH eval's 0.017 -> controller is faithful given good + context; the closed-loop sim failure is the offline->online DRIFT + (exposure bias), not a controller bug. + * overall_mse >> 0.017 -> the controller code itself is the + bug (context-building / cond / normalization), independent of drift. + +Forces sp_commit=1 (re-plan every tick) so every step is testable. +""" +from __future__ import annotations + +from typing import Dict + +import numpy as np +import torch + +from egomimic.eval.core.eval_video import EvalVideo +from egomimic.rldb.embodiment.embodiment import get_embodiment_id + + +class DFoTControllerTFEval(EvalVideo): + def __init__( + self, n_steps_eval: int = 40, max_episodes: int = 4, + embodiment_name: str = "pushshapes_sim", image_key: str = "front_img_1", + limit_val_batches: int = 4, max_videos: int = 2, + video_subdir: str = "videos_ctrl_tf", viz_func=None, transform_lists=None, + ): + super().__init__(limit_val_batches=limit_val_batches, viz_func=viz_func, + transform_lists=transform_lists, max_videos=max_videos) + self.n_steps_eval = int(n_steps_eval) + self.max_episodes = int(max_episodes) + self.embodiment_name = embodiment_name + self.image_key = str(image_key) + self._video_subdir = str(video_subdir) + + def video_dir(self): + import os + return os.path.join(self.root_dir(), self._video_subdir) + + @torch.no_grad() + def compute_metrics_and_viz(self, batch): + algo = self.model + metrics: Dict[str, torch.Tensor] = {} + emb_id = get_embodiment_id(self.embodiment_name) + if emb_id not in batch: + return metrics, {} + _batch = batch[emb_id] + outer = algo.outer_stage + if self.image_key not in _batch or not hasattr(outer, "_state_to_cond"): + return metrics, {} + device = self.trainer.lightning_module.device + ac_key = algo.resolved_ac_keys[emb_id] + obs_keys = list(outer.bundle_obs_keys) + imgs = _batch[self.image_key] + is_packed = _batch.get("_packed", False) + if is_packed: + cu = _batch["cu_seqlens"].to(imgs.device, dtype=torch.long) + ne = min(int(cu.shape[0] - 1), self.max_episodes) + spans = [(int(cu[i].item()), int(cu[i + 1].item() - cu[i].item())) for i in range(ne)] + else: + ne = min(imgs.shape[0], self.max_episodes) + spans = [(ep, imgs.shape[1]) for ep in range(ne)] + + old_commit = getattr(algo, "sp_commit", 1) + algo.sp_commit = 1 + Tcap = self.n_steps_eval + sse = np.zeros(Tcap); cnt = np.zeros(Tcap, dtype=np.int64) + dbg = [] + for (start, L) in spans: + T = min(Tcap, L) + if T < 2: + continue + for attr in ("_sp_lat", "_sp_state", "_sp_act", "_sp_queue"): + if hasattr(algo, attr): + delattr(algo, attr) + for t in range(T): + if is_packed: + gi = start + t + sl = slice(gi, gi + 1) + img_t = imgs[sl].to(device).float() + gt_act = _batch[ac_key][gi].to(device).float() # normalized (A,) + obs_zarr = {self.image_key: img_t} + for kk in obs_keys: + gv = _batch[kk][sl].to(device).float() # normalized (1, d) + obs_zarr[kk] = algo.norm_stats.unnormalize({kk: gv}, emb_id)[kk] + else: + img_t = imgs[start, t:t + 1].to(device).float() + gt_act = _batch[ac_key][start, t].to(device).float() + obs_zarr = {self.image_key: img_t} + for kk in obs_keys: + gv = _batch[kk][start, t:t + 1].to(device).float() + obs_zarr[kk] = algo.norm_stats.unnormalize({kk: gv}, emb_id)[kk] + + pred_unnorm = algo._inference_step_spatial_rh(obs_zarr, t, emb_id) # (A,) world + pu = torch.tensor(np.asarray(pred_unnorm), device=device).float().unsqueeze(0) + pred_norm = algo.norm_stats.normalize({ac_key: pu}, emb_id)[ac_key].squeeze(0) + e = float(((pred_norm - gt_act) ** 2).mean().item()) + sse[t] += e; cnt[t] += 1 + if len(dbg) < 6: + dbg.append((t, [round(x, 3) for x in pred_norm.detach().cpu().numpy().tolist()], + [round(x, 3) for x in gt_act.detach().cpu().numpy().tolist()])) + # teacher-force: replace recorded predicted action with GT (normalized) + algo._sp_act[-1] = gt_act.clone() + algo.sp_commit = old_commit + + print("[TF_SANITY] first preds (t, pred_norm, gt_norm):") + for d in dbg: + print(" ", d) + if cnt.sum() > 0: + for t in range(Tcap): + if cnt[t] > 0: + metrics[f"Valid/emb{emb_id}_tf_action_mse_step_{t:02d}"] = torch.tensor(sse[t] / cnt[t], device=device) + metrics[f"Valid/emb{emb_id}_tf_action_mse_overall"] = torch.tensor(sse.sum() / cnt.sum(), device=device) + return metrics, {} diff --git a/egomimic/eval/tf/eval_dfot_val.py b/egomimic/eval/tf/eval_dfot_val.py new file mode 100644 index 000000000..6816e087e --- /dev/null +++ b/egomimic/eval/tf/eval_dfot_val.py @@ -0,0 +1,387 @@ +"""DFoT val-data evaluator: teacher-forced action prediction with viz. + +Runs two off-env prediction modes per val episode and renders a single mp4 +overlaying both predictions on top of the GT environment frames: + + * **full_chunk**: one-shot denoising of the entire episode's action + sequence. Single schedule with uniform per-token noise levels driven + from 1.0 -> 0.0 over ``n_chunk_steps`` DDIM steps. Tests how well the + model can recover a full trajectory from a single chunk of obs. + * **ar (staircase)**: a single offline ``sample()`` call over the full + episode driven by ``staircase_ar_schedule(T, ar_chunk_size, + ar_step_size)``. Earlier tokens denoise first; later tokens still at + high noise — the rolling-staircase pattern unrolled across the whole + episode in one pass. Knobs: ``ar_chunk_size`` = tokens per rung + ("width"); ``ar_step_size`` = denoise steps per rung ("height"). + For online (env-tick-by-env-tick) AR see ``DFoT.inference_step`` — + this evaluator is teacher-forced and runs offline. + +For each val episode the evaluator emits one mp4: the GT env image stream +(from ``front_img_1``) overlaid with three coloured action dots per frame +— GT (green), full_chunk pred (blue), AR pred (yellow). Per-mode per- +frame MSE is reported as the validation metric. + +This evaluator does NOT touch a simulator (cf. ``PackedSimEval`` / +``HPTSimEval``). All obs come straight from the packed val batch — no +closed-loop drift between predicted action and next obs. +""" + +from __future__ import annotations + +from typing import Dict, List, Tuple + +import cv2 +import numpy as np +import torch + +from egomimic.models.diffusion.diffusion.discrete_diffusion import DiscreteDiffusion +from egomimic.models.diffusion.sampling import ( + sample, + staircase_ar_schedule, + vanilla_schedule, +) +from egomimic.eval.core.eval_video import EvalVideo + + +def _to_xy_pix(action_world: np.ndarray, img_hw: tuple, world_size: float) -> tuple: + """Map a world-frame xy action into pixel coords for a (H, W) image. + + PushShapesEnv decouples physics world size (``WORLD_SIZE=512``) from + the rendered image_size (e.g. 96). Actions live in + ``[0, world_size]``; we scale to pixel coords using the image dims. + Out-of-bounds returns None. + """ + h, w = img_hw + x, y = float(action_world[0]), float(action_world[1]) + sx = w / float(world_size) + sy = h / float(world_size) + cx, cy = int(round(x * sx)), int(round(y * sy)) + if 0 <= cx < w and 0 <= cy < h: + return cx, cy + return None + + +def _draw_overlay( + img_chw_uint8: np.ndarray, + actions_world: Dict[str, np.ndarray], + palette: Dict[str, Tuple[int, int, int]], + world_size: float, + radius: int = 8, + upscale_to: int | None = 512, +) -> np.ndarray: + """Overlay one dot per mode on a (C, H, W) uint8 image. Returns (H, W, C) + uint8 ready for video stacking. If upscale_to is set, the env image is + nearest-neighbor upscaled before dots are drawn so dot edges stay crisp.""" + if img_chw_uint8.dtype != np.uint8: + img_chw_uint8 = (img_chw_uint8 * 255.0).clip(0, 255).astype(np.uint8) + img_hwc = np.transpose(img_chw_uint8, (1, 2, 0)) + if upscale_to is not None and img_hwc.shape[0] != upscale_to: + img_hwc = cv2.resize( + img_hwc, (upscale_to, upscale_to), interpolation=cv2.INTER_NEAREST + ) + H, W, _ = img_hwc.shape + out = np.ascontiguousarray(img_hwc.copy()) + for name, a in actions_world.items(): + pt = _to_xy_pix(a, (H, W), world_size=world_size) + if pt is None: + continue + cv2.circle(out, pt, radius + 2, (0, 0, 0), thickness=-1) + cv2.circle(out, pt, radius, palette[name], thickness=-1) + return out + + +# ---------------------------------------------------------------------- # +# Eval class +# ---------------------------------------------------------------------- # + + +class DFoTValEval(EvalVideo): + """Teacher-forced val-data evaluator for DFoT. + + Args: + do_full_chunk: run the one-shot full-episode denoising mode. + do_ar: run the causal-AR staircase mode. + n_chunk_steps: number of DDIM steps for the full-chunk mode. + ar_chunk_size: staircase width (tokens per rung). 1 = vanilla + causal-AR. + ar_step_size: staircase height (denoising steps per rung). + limit_val_batches: cap on val batches per validation pass. + max_videos: cap on number of episodes rendered per val pass. + coverage_threshold: unused here (kept for config-shape parity with + sim-eval configs); accepted and ignored. + env_kwargs: unused (no env); accepted for config parity. + embodiment_name: which embodiment to expect in the packed batch + (default ``"pushshapes_sim"``). + viz_func / transform_lists: forwarded to ``EvalVideo`` base. + """ + + def __init__( + self, + do_full_chunk: bool = True, + do_ar: bool = True, + n_chunk_steps: int = 50, + ar_chunk_size: int = 1, + ar_step_size: int = 1, + cfg_scale: float = 1.0, + world_size: float = 512.0, + limit_val_batches: int = 4, + max_videos: int = 2, + coverage_threshold: float = 0.7, + env_kwargs: dict | None = None, + embodiment_name: str = "pushshapes_sim", + viz_func=None, + transform_lists=None, + ): + super().__init__( + limit_val_batches=limit_val_batches, + viz_func=viz_func, + transform_lists=transform_lists, + max_videos=max_videos, + ) + self.do_full_chunk = bool(do_full_chunk) + self.do_ar = bool(do_ar) + self.n_chunk_steps = int(n_chunk_steps) + self.ar_chunk_size = int(ar_chunk_size) + self.ar_step_size = int(ar_step_size) + # Classifier-free-guidance scale at sampling. 1.0 disables CFG. + self.cfg_scale = float(cfg_scale) + # PushShapesEnv physics world extent. Actions live in [0, world_size]. + # Rendered images may be smaller (e.g. 96 px); _to_xy_pix scales. + self.world_size = float(world_size) + self.embodiment_name = embodiment_name + # Kept for config-shape parity with eval_dfot_sim; not used. + _ = coverage_threshold, env_kwargs + + # ---- env-image extraction ---- # + + def _img_uint8(self, img_chw_norm: torch.Tensor) -> np.ndarray: + """Convert one (C, H, W) image tensor (any float range) to uint8 + (C, H, W) numpy. Heuristic: if max <= 1.5 we treat as [0,1] floats.""" + x = img_chw_norm.detach().cpu().float().numpy() + if x.max() <= 1.5: + x = x * 255.0 + return x.clip(0, 255).astype(np.uint8) + + # ---- per-episode prediction passes ---- # + + @torch.no_grad() + def _full_chunk_pred( + self, + algo, + actions_norm: torch.Tensor, # (T, A) + cond_per_frame: torch.Tensor | None, # (T, d_cond) or None + emb_id: int, + ) -> torch.Tensor: + """One-shot DDIM denoise over the full episode. Returns predicted + actions in WORLD frame (un-normalized), shape (T, A). + + Bundle-aware: when the algo's ``outer_stage`` diffuses a joint + ``[obs, action]`` bundle (e.g. ``ObsActionDFoTOuterStage``), the + sampler is allocated at ``bundle_dim`` width and the action + portion is sliced out via ``outer_stage.action_slice`` before + unnormalizing. For vanilla DFoT both reduce to no-ops + (bundle_dim = action_dim, slice = full), so behavior is preserved. + """ + T, _ = actions_norm.shape + device = actions_norm.device + diff = algo.diffusion + backbone = algo.backbone + outer = algo.outer_stage + bundle_dim = int(getattr(outer, "bundle_dim", actions_norm.shape[-1])) + action_slice = getattr(outer, "action_slice", slice(None)) + discrete_ts = ( + int(diff.timesteps) if isinstance(diff, DiscreteDiffusion) else None + ) + schedule = vanilla_schedule( + n_steps=self.n_chunk_steps, T=T, discrete_timesteps=discrete_ts + ).to(device) + ec = cond_per_frame.unsqueeze(0) if cond_per_frame is not None else None + pred_full = sample( + diff, + backbone, + schedule_matrix=schedule, + action_dim=bundle_dim, + batch_size=1, + external_cond=ec, + cfg_scale=self.cfg_scale, + device=device, + ).squeeze(0) # (T, bundle_dim) + pred_act_norm = pred_full[..., action_slice] # (T, A) + ac_key = algo.resolved_ac_keys[emb_id] + pred_world = algo.norm_stats.unnormalize({ac_key: pred_act_norm}, emb_id)[ + ac_key + ] + return pred_world # (T, A) + + @torch.no_grad() + def _ar_pred( + self, + algo, + cond_per_frame: torch.Tensor, # (T, d_cond) + emb_id: int, + T: int, + ) -> torch.Tensor: + """Single-call staircase causal-AR denoising over the whole episode. + + The schedule_matrix is the rolling-staircase geometry produced by + ``staircase_ar_schedule(T, chunk_size=ar_chunk_size, + step_size=ar_step_size)``. ``sample()`` then walks all + ``ceil(T / ar_chunk_size) * ar_step_size`` denoising steps, with + every token at every step taking its noise level from the matrix. + Earlier tokens reach 0 first; later tokens still at high noise + — the canonical DFoT-paper causal-AR pattern. + """ + device = cond_per_frame.device + diff = algo.diffusion + backbone = algo.backbone + outer = algo.outer_stage + bundle_dim = int(getattr(outer, "bundle_dim", algo.action_dim)) + action_slice = getattr(outer, "action_slice", slice(None)) + discrete_ts = ( + int(diff.timesteps) if isinstance(diff, DiscreteDiffusion) else None + ) + ac_key = algo.resolved_ac_keys[emb_id] + schedule = staircase_ar_schedule( + T=T, + chunk_size=int(self.ar_chunk_size), + step_size=int(self.ar_step_size), + discrete_timesteps=discrete_ts, + ).to(device) + ec = cond_per_frame.unsqueeze(0) # (1, T, d_cond) + pred_full = sample( + diff, + backbone, + schedule_matrix=schedule, + action_dim=bundle_dim, + batch_size=1, + external_cond=ec, + cfg_scale=self.cfg_scale, + device=device, + ).squeeze(0) # (T, bundle_dim) + pred_act_norm = pred_full[..., action_slice] # (T, A) + pred_world = algo.norm_stats.unnormalize({ac_key: pred_act_norm}, emb_id)[ + ac_key + ] + return pred_world + + # ---- main eval entrypoint ---- # + + def compute_metrics_and_viz(self, batch): + device = self.trainer.lightning_module.device + algo = self.model + metrics: Dict[str, torch.Tensor] = {} + images_dict: Dict[int, np.ndarray] = {} + + palette = { + "gt": (0, 255, 0), # green + "chunk": (255, 64, 64), # red — full-chunk DDIM + "ar": (255, 255, 0), # yellow — staircase AR + } + + for emb_id, _batch in batch.items(): + if not _batch.get("_packed", False): + # Unsupported in v1; skip silently. + continue + cu = _batch["cu_seqlens"] + if not torch.is_tensor(cu): + continue + B = int(cu.shape[0]) - 1 + if B <= 0: + continue + B_render = min(B, self.max_videos) if self.max_videos is not None else B + + ac_key = algo.resolved_ac_keys[emb_id] + actions_packed = _batch[ac_key] # (T_total, A) + obs = algo._build_obs(_batch, emb_id) + cond_packed = algo._encode_cond_packed(obs) # (T_total, d_cond) or None + img_key = next(iter(algo.camera_keys[emb_id]), None) + imgs_packed = _batch.get(img_key) if img_key is not None else None + + # Accumulators for global mean MSE (sum of squared errors + # across ALL episodes / frames / action dims, divided by the + # TOTAL element count). This weights long episodes + # proportionally rather than treating each episode equally, + # matching the train-loss reduction. + chunk_sse: float = 0.0 + chunk_n: int = 0 + ar_sse: float = 0.0 + ar_n: int = 0 + ep_frames: List[np.ndarray] = [] + + for b in range(B_render): + s = int(cu[b].item()) + e = int(cu[b + 1].item()) + T = e - s + if T <= 1: + continue + + # Slice this episode's tensors. + acts_norm_ep = actions_packed[s:e] # (T, A) + cond_ep = cond_packed[s:e] if cond_packed is not None else None + imgs_ep = imgs_packed[s:e] if imgs_packed is not None else None + + # GT in world frame. + gt_world = algo.norm_stats.unnormalize({ac_key: acts_norm_ep}, emb_id)[ + ac_key + ] # (T, A) + + chunk_world = None + ar_world = None + + if self.do_full_chunk: + chunk_world = self._full_chunk_pred( + algo, acts_norm_ep, cond_ep, emb_id + ) + chunk_sse += float(torch.sum((chunk_world - gt_world) ** 2).item()) + chunk_n += int(chunk_world.numel()) + + if self.do_ar: + if cond_ep is None: + ar_world = None + else: + ar_world = self._ar_pred(algo, cond_ep, emb_id, T=T) + ar_sse += float(torch.sum((ar_world - gt_world) ** 2).item()) + ar_n += int(ar_world.numel()) + + # Build viz video for this episode. + if imgs_ep is not None: + frames_ep: List[np.ndarray] = [] + for t in range(T): + img_chw = self._img_uint8(imgs_ep[t]) + overlay_actions = {"gt": gt_world[t].detach().cpu().numpy()} + if chunk_world is not None: + overlay_actions["chunk"] = ( + chunk_world[t].detach().cpu().numpy() + ) + if ar_world is not None: + overlay_actions["ar"] = ar_world[t].detach().cpu().numpy() + out = _draw_overlay( + img_chw, + overlay_actions, + palette, + world_size=self.world_size, + ) + frames_ep.append(np.ascontiguousarray(out)) + ep_frames.extend(frames_ep) + if b < B_render - 1 and frames_ep: + # 5-frame black separator between consecutive episodes. + H, W, _ = frames_ep[0].shape + sep = np.zeros((5, H, W, 3), dtype=np.uint8) + ep_frames.extend(sep) + + # Aggregate: global mean MSE = sum_squared_errors / N where + # N is total element count across all episodes (frames * + # action_dim). Prefixed with ``Valid/`` so the metric lands + # in the W&B Valid panel alongside sim_coverage etc. + if chunk_n > 0: + metrics[f"Valid/emb{emb_id}_chunk_action_mse"] = torch.tensor( + chunk_sse / chunk_n, device=device + ) + if ar_n > 0: + metrics[f"Valid/emb{emb_id}_ar_action_mse"] = torch.tensor( + ar_sse / ar_n, device=device + ) + if ep_frames: + images_dict[emb_id] = np.stack(ep_frames, axis=0) + + return metrics, images_dict diff --git a/egomimic/hydra_configs/evaluator/dfot/bundle_anchored.yaml b/egomimic/hydra_configs/evaluator/dfot/bundle_anchored.yaml new file mode 100644 index 000000000..d3d115ca8 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/dfot/bundle_anchored.yaml @@ -0,0 +1,11 @@ +_target_: egomimic.eval.core.eval_composite.EvalList +evals: + - _target_: egomimic.eval.dfot.eval_dfot_bundle_anchored.DFoTBundleAnchoredEval + n_context_frames: 4 + rollout_steps: 32 + n_chunk_steps: 50 + embodiment_name: "pushshapes_sim" + image_key: "front_img_1" + recon_loss_n_frames: 20 + max_videos: 2 + limit_val_batches: 4 diff --git a/egomimic/hydra_configs/evaluator/dfot/controller_tf.yaml b/egomimic/hydra_configs/evaluator/dfot/controller_tf.yaml new file mode 100644 index 000000000..bab17bdfd --- /dev/null +++ b/egomimic/hydra_configs/evaluator/dfot/controller_tf.yaml @@ -0,0 +1,11 @@ +# Teacher-forced sanity check for the spatial_rh closed-loop controller. +_target_: egomimic.eval.core.eval_composite.EvalList +evals: + - _target_: egomimic.eval.tf.eval_dfot_controller_tf.DFoTControllerTFEval + n_steps_eval: 40 + max_episodes: 4 + embodiment_name: "pushshapes_sim" + image_key: "front_img_1" + limit_val_batches: 4 + max_videos: 2 + video_subdir: "videos_ctrl_tf" diff --git a/egomimic/hydra_configs/evaluator/dfot/full.yaml b/egomimic/hydra_configs/evaluator/dfot/full.yaml new file mode 100644 index 000000000..f4cf4ab75 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/dfot/full.yaml @@ -0,0 +1,44 @@ +# Composite DFoT evaluator: teacher-forced val viz + closed-loop sim eval. +# 1. eval_dfot_val — full-chunk DDIM + staircase-AR overlaid on GT env +# frames. Tells us "given the right obs, does the model predict the +# right actions?" Independent of env dynamics. +# 2. eval_dfot_sim — closed-loop rollout in PushShapesEnv driven by +# DFoT.inference_step (default: causal-AR staircase). Tells us +# "does the policy actually solve the task?" +# +# Both run every val pass. + +_target_: egomimic.eval.core.eval_composite.EvalList + +evals: + - _target_: egomimic.eval.tf.eval_dfot_val.DFoTValEval + do_full_chunk: true + do_ar: true + n_chunk_steps: 50 + ar_chunk_size: 1 + ar_step_size: 1 + cfg_scale: 1.0 # standard CFG; 1.0 = off + world_size: 512.0 # PushShapesEnv physics extent — for viz overlay scaling + limit_val_batches: 4 + max_videos: 2 + coverage_threshold: 0.7 + embodiment_name: pushshapes_sim + env_kwargs: null + viz_func: null + transform_lists: null + + - _target_: egomimic.eval.core.eval_sim.PackedSimEval + env_kwargs: + object_shape: "T" + pusher_shape: "circle" + obstacle_level: 0 + image_size: 96 + init_mode: "replay" + init_seeds: [0, 1, 2, 3] + max_steps: 1200 + coverage_threshold: 0.7 + video_fps: 30 + limit_val_batches: 4 + max_videos: 2 + viz_func: null + transform_lists: null diff --git a/egomimic/hydra_configs/evaluator/dfot/image_spatial.yaml b/egomimic/hydra_configs/evaluator/dfot/image_spatial.yaml new file mode 100644 index 000000000..dacc6d144 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/dfot/image_spatial.yaml @@ -0,0 +1,26 @@ +# DFoT image-spatial evaluator (plan A: world-model only). +# +# Single sub-eval: the family-agnostic DFoTVideoRolloutEval (COMBINE A), +# dispatched onto ImageSpatialDFoTOuterStage's rollout_video_episode hook -- +# decodes the predicted spatial latent sequence into pixel frames + per-step +# MSE vs GT frames in the val batch (metric prefix "spatial"). +# +# No DFoTValEval / HNetSimEval here because this variant doesn't +# diffuse actions; sim env-rollout requires predicted actions. + +_target_: egomimic.eval.core.eval_composite.EvalList + +evals: + - _target_: egomimic.eval.dfot.eval_dfot_video_rollout.DFoTVideoRolloutEval + rollout_steps: 64 + upscale_to: 384 + max_videos: 2 + limit_val_batches: 4 + embodiment_name: "pushshapes_sim" + image_key: "front_img_1" + recon_loss_n_frames: 10 + n_context_frames: 0 # spatial-default: unconditional rollout (was the class default) + mode: "chunk" + n_chunk_steps: 50 + cfg_scale: 1.0 + video_subdir: "videos_spatial_rollout" diff --git a/egomimic/hydra_configs/evaluator/dfot/image_spatial_policy.yaml b/egomimic/hydra_configs/evaluator/dfot/image_spatial_policy.yaml new file mode 100644 index 000000000..16287b3e3 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/dfot/image_spatial_policy.yaml @@ -0,0 +1,17 @@ +# Action-prediction eval for the 2D policy: clean obs history -> predict next +# action chunk, compare predicted future actions to GT (action MSE) + decode +# predicted latents into a [GT|pred] video. +_target_: egomimic.eval.core.eval_composite.EvalList + +evals: + - _target_: egomimic.eval.dfot.eval_dfot_policy.DFoTPolicyActionEval + n_context_frames: 4 + rollout_steps: 32 + n_chunk_steps: 50 + embodiment_name: "pushshapes_sim" + image_key: "front_img_1" + recon_loss_n_frames: 20 + upscale_to: 384 + max_videos: 2 + limit_val_batches: 4 + video_subdir: "videos_policy_action" diff --git a/egomimic/hydra_configs/evaluator/dfot/image_spatial_policy_rh.yaml b/egomimic/hydra_configs/evaluator/dfot/image_spatial_policy_rh.yaml new file mode 100644 index 000000000..2cf1c67ef --- /dev/null +++ b/egomimic/hydra_configs/evaluator/dfot/image_spatial_policy_rh.yaml @@ -0,0 +1,19 @@ +# Receding-horizon action eval for the 2D policy. Slides an anchor across each +# episode, predicts the next k=1,2 actions from CLEAN GT history (short T=n_ctx+k +# rollout), scores vs GT in normalized action space, pooled over all anchors x +# max_episodes. Headline: rh_k1_action_mse_overall. Decoupled max_episodes (NOT +# max_videos) so statistics aren't limited to 2 episodes. +_target_: egomimic.eval.core.eval_composite.EvalList + +evals: + - _target_: egomimic.eval.dfot.eval_dfot_policy.DFoTPolicyRecedingHorizonEval + n_context_frames: 4 + n_chunk_steps: 50 + k_actions: [1, 2] + max_episodes: 12 + anchor_stride: 2 + embodiment_name: "pushshapes_sim" + image_key: "front_img_1" + limit_val_batches: 4 + max_videos: 2 + video_subdir: "videos_policy_rh" diff --git a/egomimic/hydra_configs/evaluator/dfot/obs_action.yaml b/egomimic/hydra_configs/evaluator/dfot/obs_action.yaml new file mode 100644 index 000000000..073b9c4ad --- /dev/null +++ b/egomimic/hydra_configs/evaluator/dfot/obs_action.yaml @@ -0,0 +1,48 @@ +# DFoT obs+action evaluator. Two independent sub-evals, each writes +# its own mp4 (no composite/side-by-side panel). +# +# 1. DFoTSelfRolloutEval — world-model self-rollout. Renders the +# model's predicted state at each future tick via the pushshapes +# env. Output: ``/videos_self_rollout/epoch_N/ +# PUSHSHAPES_SIM/validation_video_*.mp4`` ("video prediction +# quality" for the obs portion of the bundle). +# +# 2. DFoTValEval — teacher-forced full-chunk + AR overlay. Renders GT +# env frames with GT/chunk/AR action dots. Output: +# ``/videos/epoch_N/PUSHSHAPES_SIM/ +# validation_video_*.mp4``. +# +# Sub-evals namespace their own output via ``video_dir()`` so two mp4s +# land per val pass without filename clashes. + +_target_: egomimic.eval.core.eval_composite.EvalList + +evals: + - _target_: egomimic.eval.dfot.eval_dfot_self_rollout.DFoTSelfRolloutEval + rollout_steps: 64 + world_size: 512.0 + image_key: "front_img_1" + agent_state_slice: [0, 2] + cfg_scale: 1.0 + upscale_to: 512 + max_videos: 2 + limit_val_batches: 4 + embodiment_name: "pushshapes_sim" + video_subdir: "videos_self_rollout" + env_kwargs: + object_shape: T + pusher_shape: circle + obstacle_level: 0 + image_size: 96 + + - _target_: egomimic.eval.tf.eval_dfot_val.DFoTValEval + do_full_chunk: true + do_ar: true + n_chunk_steps: 50 + ar_chunk_size: 1 + ar_step_size: 1 + cfg_scale: 1.0 + world_size: 512.0 + limit_val_batches: 4 + max_videos: 2 + embodiment_name: "pushshapes_sim" diff --git a/egomimic/hydra_configs/evaluator/dfot/obs_action_image.yaml b/egomimic/hydra_configs/evaluator/dfot/obs_action_image.yaml new file mode 100644 index 000000000..ce8e1da95 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/dfot/obs_action_image.yaml @@ -0,0 +1,61 @@ +# DFoT obs+action+image evaluator: video self-rollout only. +# +# DFoTVideoRolloutEval picks the first ``max_videos`` val episodes, +# takes only the t=0 image, runs a chunk-mode rollout, and decodes the +# predicted image latents back to pixels via the frozen VAE. Output is +# one mp4 per episode showing the predicted video. +# +# (DFoTValEval was considered but dropped: with empty cond_encoder its +# sampler runs unconditionally, so the action MSE it reports would be +# "sample from learned prior" quality — not strict teacher-forced +# semantics. A proper teacher-forced action eval needs inpainting the +# GT obs/latent dims during denoising, which is a v2 follow-up.) + +_target_: egomimic.eval.core.eval_composite.EvalList + +evals: + - _target_: egomimic.eval.dfot.eval_dfot_video_rollout.DFoTVideoRolloutEval + rollout_steps: 64 + upscale_to: 384 + max_videos: 2 + limit_val_batches: 4 + embodiment_name: "pushshapes_sim" + image_key: "front_img_1" + video_subdir: "videos_video_rollout_chunk" + mode: "chunk" + n_chunk_steps: 50 + + - _target_: egomimic.eval.dfot.eval_dfot_video_rollout.DFoTVideoRolloutEval + rollout_steps: 64 + upscale_to: 384 + max_videos: 2 + limit_val_batches: 4 + embodiment_name: "pushshapes_sim" + image_key: "front_img_1" + video_subdir: "videos_video_rollout_ar" + mode: "ar" + ar_chunk_size: 1 + ar_step_size: 1 + + # Closed-loop sim rollout. PackedSimEval / HNetSimEval (same class) + # drives the PushShapesEnv with actions from algo.inference_step, + # records frames + coverage. Mp4 lives under videos/epoch_*/... + # + # CAVEAT: this variant's cond_encoder is empty, so DFoT's AR + # inference rollout runs unconditioned -- the action committed at + # each env tick is sampled from the learned joint prior without + # anchoring to the current env obs. Expect random-ish behaviour + # until we add obs/latent anchoring to inference_step. Useful as a + # smoke + visual sanity check. + - _target_: egomimic.eval.core.eval_sim.HNetSimEval + limit_val_batches: 4 + max_videos: 2 + embodiment_name: pushshapes_sim + init_mode: replay + max_steps: 1200 + coverage_threshold: 0.7 + env_kwargs: + object_shape: T + pusher_shape: circle + obstacle_level: 0 + image_size: 96 diff --git a/egomimic/hydra_configs/evaluator/dfot/pixel.yaml b/egomimic/hydra_configs/evaluator/dfot/pixel.yaml new file mode 100644 index 000000000..ae28341c3 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/dfot/pixel.yaml @@ -0,0 +1,25 @@ +# DFoT pixel-space evaluator (no VAE decode needed). +# +# Single sub-eval: the family-agnostic DFoTVideoRolloutEval (COMBINE A), +# dispatched onto PixelSpatialDFoTOuterStage's rollout_video_episode hook -- +# generates video frames directly from the pixel diffusion model (sliding- +# window rollout anchored on the first GT frame) + per-step MSE + PSNR/SSIM/ +# LPIPS vs GT frames (metric prefix "pixel"). + +_target_: egomimic.eval.core.eval_composite.EvalList + +evals: + - _target_: egomimic.eval.dfot.eval_dfot_video_rollout.DFoTVideoRolloutEval + rollout_steps: 9 + upscale_to: 384 + max_videos: 2 + limit_val_batches: 4 + embodiment_name: "pushshapes_sim" + image_key: "front_img_1" + recon_loss_n_frames: 9 + n_context_frames: 1 # pixel-default: anchor on first GT frame (was the class default) + rollout_window: 9 # pixel-default sliding-window size (was the class default) + mode: "chunk" + n_chunk_steps: 50 + cfg_scale: 1.0 + video_subdir: "videos_pixel_rollout" diff --git a/egomimic/hydra_configs/evaluator/dfot/sim.yaml b/egomimic/hydra_configs/evaluator/dfot/sim.yaml new file mode 100644 index 000000000..d953961e3 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/dfot/sim.yaml @@ -0,0 +1,22 @@ +# Closed-loop sim eval for DFoT. Uses the algo-agnostic PackedSimEval +# class with PushShapesEnv; the model side runs ``DFoT.inference_step`` +# (driven by the model config's ``inference_mode`` — "ar" by default, +# i.e. rolling causal-AR staircase). + +_target_: egomimic.eval.core.eval_sim.PackedSimEval + +env_kwargs: + object_shape: "T" + pusher_shape: "circle" + obstacle_level: 0 + image_size: 96 + +init_mode: "replay" +init_seeds: [0, 1, 2, 3] +max_steps: 1200 +coverage_threshold: 0.7 +video_fps: 30 +limit_val_batches: 4 +max_videos: 2 +viz_func: null +transform_lists: null diff --git a/egomimic/hydra_configs/evaluator/dfot/val.yaml b/egomimic/hydra_configs/evaluator/dfot/val.yaml new file mode 100644 index 000000000..5ca3551ba --- /dev/null +++ b/egomimic/hydra_configs/evaluator/dfot/val.yaml @@ -0,0 +1,47 @@ +# DFoT val-data evaluator. Teacher-forced action prediction over val +# episodes with two modes overlaid in a single mp4 per episode: +# * full_chunk: one-shot DDIM denoising over the entire episode action +# sequence. +# * ar (staircase): rolling causal-AR denoising; configurable staircase +# geometry via ar_chunk_size (width) and ar_step_size (height). +# +# Use with packed val data (tsimulation/full.yaml / tsimulation/base.yaml). +# Does NOT touch a simulator — for closed-loop coverage use eval_dfot_sim +# (or the composite eval_dfot_full which runs both). + +_target_: egomimic.eval.tf.eval_dfot_val.DFoTValEval + +# ---- mode toggles ---- +do_full_chunk: true +do_ar: true + +# ---- full-chunk knobs ---- +# DDIM steps over the full episode. Higher = better but slower per val pass. +n_chunk_steps: 50 + +# ---- AR / staircase knobs ---- +# chunk_size = staircase "width" — tokens that share a rung's noise level. +# 1 = vanilla causal-AR (one token per rung) +# k = k tokens denoise together as a rung +# step_size = staircase "height" — denoising steps per rung. +# 1 = each step advances all tokens by one schedule unit +# k = k denoise iterations per rung before the chunk commits +ar_chunk_size: 1 +ar_step_size: 1 + +# Classifier-free guidance scale (1.0 = off; >1 enables CFG, 2x backbone +# passes per denoise step). Same recipe as DFoT paper standard CFG. +cfg_scale: 1.0 + +# PushShapesEnv physics world extent (actions live in [0, world_size]). +# Used by the viz overlay to scale dots onto the smaller rendered image. +world_size: 512.0 + +# ---- shared knobs (match eval_dfot_sim / EvalVideo shape) ---- +limit_val_batches: 4 +max_videos: 2 +coverage_threshold: 0.7 # ignored; present for config-shape parity +embodiment_name: pushshapes_sim +env_kwargs: null # ignored; present for config-shape parity +viz_func: null +transform_lists: null diff --git a/egomimic/hydra_configs/model/dfot/base.yaml b/egomimic/hydra_configs/model/dfot/base.yaml new file mode 100644 index 000000000..9a9d7c86c --- /dev/null +++ b/egomimic/hydra_configs/model/dfot/base.yaml @@ -0,0 +1,138 @@ +_target_: egomimic.pl_utils.pl_model.ModelWrapper +robomimic_model: + _target_: egomimic.algo.diffusion.DFoT + # ``norm_stats`` is injected at runtime by + # ``pl_model.ModelWrapper._instantiate_model``; do not set it here. + + # ---- core dims ---- + # PushShapes circle/basic action dim = 2 (xy). + action_dim: 2 + # AR buffer / chunk length T. Also the planning window for legacy + # chunk-mode inference. Packed-mode training does NOT slice into + # ``action_horizon``-sized windows — the diffusion runs over whole + # variable-length episodes; ``action_horizon`` only governs inference. + action_horizon: 32 + + sampler: "ddim" + sampler_n_steps: 50 + sampler_eta: 0.0 + + # ---- closed-loop inference (sim eval) ---- + # "ar" — causal-AR rolling staircase. One denoise step per env tick; + # buffer carries across steps. Matches DFoT training distribution. + # "chunk" — legacy plan-and-execute (uniform per-token noise, predict + # ``action_horizon`` actions, replan after). Cheaper baseline. + inference_mode: "ar" + ar_inference_chunk_size: 1 # actions committed per env tick (1 = classic AR) + ar_inference_step_size: 1 # sample_step sub-steps per env tick + # Classifier-free-guidance scale. 1.0 = no CFG (single backbone pass). + # >1.0 runs two passes per denoise step (cond + zero-cond) and blends: + # v = v_uncond + cfg_scale * (v_cond - v_uncond) + # Typical action-policy values: 1.2-2.0. Higher = sharper, less variance. + # Requires training with ``cond_dropout_prob > 0`` (we use 0.1). + cfg_scale: 1.0 + + domains: ["pushshapes_sim"] + ac_keys: + pushshapes_sim: "actions" + + # ---- outer_stage: owns the modality boundary + the trunk ---- + # The outer stage holds the cond_encoder (obs -> per-token AdaLN cond), + # the backbone (Isotropic trunk + per-token AdaLN), and the diffusion + # module (q_sample / compute_loss math). DFoTOuterStage.encode samples + # per-token noise at training time, runs q_sample, stores q_state on + # ctx for the loss to read. Its inner_stage IS the backbone. + outer_stage: + _target_: egomimic.algo.diffusion.outer_stages.outer_stage.DFoTOuterStage + action_dim: 2 + cond_output_key: "fused_cond" + + # In packed-mode training, obs is per-frame (one obs aligned with one + # action token), so cond_encoder emits ``(T_total, d_cond)`` and is fed + # per-token into the backbone's AdaLN. At online inference, the same + # encoder runs on the current env obs and is broadcast across the AR + # buffer. + cond_encoder: + _target_: egomimic.models.stems.cond_encoders.CondEncoderModule + d_cond: 512 + output_key: "fused_cond" + cond_proj_widths: [512, 512] + obs_specs: + state_agent_obj: + input_dim: 5 + embed_dim: 256 + widths: [256] + img_encoders: + front_img_1: + _target_: egomimic.models.stems.image_encoders.SimpleConv + in_channels: 3 + channels: [64, 128, 256, 512] + kernel_size: 3 + stride: 2 + embed_dim: 256 + norm_groups: 8 + + # DFoT backbone. Scaled for H200 (141GB VRAM). T12 transformer at + # d_model=512 + d_intermediate=2048 ~= 50M params. Combined with + # packed batch_size=16 this should fill the GPU comfortably while + # leaving headroom for activation memory at long episodes. + backbone: + _target_: egomimic.models.diffusion.DFoTBackbone + action_dim: 2 + action_horizon: 32 + d_model: 512 + d_cond: 512 + cond_dim: 512 # matches cond_encoder.d_cond + time_embed_dim: 512 + cond_emb_dim: 512 + use_fourier_time: true # continuous diffusion -> Fourier (no stochastic dropout) + stochastic_time_p: 0.0 + cond_dropout_prob: 0.1 # CFG-style cond dropout + arch_layout: "T12" + num_heads: 8 + d_intermediate: 2048 + rotary_emb_dim: 64 # head_dim = 512 / 8 = 64 + dropout: 0.1 + resid_dropout: 0.1 + # Causal self-attention. Each action token attends only to tokens at + # equal-or-earlier positions, matching closed-loop control semantics + # (at env-tick t, only obs[0..t] is available). Combined with the + # rolling-staircase schedule at inference, this gives strict AR + # behavior at both the attention level and the noise-level level. + causal: true + + # Continuous-time diffusion (v-prediction, sigmoid loss weighting). + # Used by outer_stage.encode for q_sample at training, and also + # implicitly by the default DFoTLoss (no loss block below -> the + # algo class auto-builds DFoTLoss(outer_stage.diffusion)). + diffusion: + _target_: egomimic.models.diffusion.diffusion.continuous_diffusion.ContinuousDiffusion + action_dim: 2 + precond_scale: 0.25 + sigmoid_bias: -1.0 + clip_noise: 20.0 + logsnr_min: -15.0 + logsnr_max: 15.0 + shift: 1.0 + + # ---- loss: optional ---- + # If omitted, DFoT.__init__ defaults to ``DFoTLoss(outer_stage.diffusion)`` + # (SNR-weighted epsilon-MSE). Override here only when you want a different + # loss policy (e.g. a custom Loss subclass adding obs-reconstruction + # terms in the future fused-obs+action mode). + # loss: + # _target_: egomimic.algo.diffusion.algo.DFoTLoss + +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 1.0e-4 + weight_decay: 0.0001 + +scheduler: + _target_: egomimic.utils.schedulers.warmup_cosine_scheduler + _partial_: true + max_steps: 3200 # matches HPT default budget + warmup_steps: 100 + warmup_start_factor: 0.1 + eta_min: 1.0e-5 diff --git a/egomimic/hydra_configs/model/dfot/image_spatial.yaml b/egomimic/hydra_configs/model/dfot/image_spatial.yaml new file mode 100644 index 000000000..4502ac455 --- /dev/null +++ b/egomimic/hydra_configs/model/dfot/image_spatial.yaml @@ -0,0 +1,90 @@ +_target_: egomimic.pl_utils.pl_model.ModelWrapper + +# DFoT video diffusion with DiT3D backbone (History Guided Video Diffusion). +# +# Diffusion target = VAE latent in spatial form (4, 6, 6) per timestep. +# Actions + state enter via external_cond (broadcast across spatial +# patches by AdaLN-Zero). Model predicts FUTURE IMAGE LATENTS. +# +# Key arch choices matching the reference (kwsong0113/diffusion-forcing-transformer): +# - DiT3D backbone with AdaLN-Zero gating (zero-init residual gates) +# - RoPE 3D positional encoding (temporal + spatial axes) +# - Discrete diffusion, cosine schedule, pred_v, min-SNR loss weighting +# - Non-causal attention (chunk-mode inference) + +robomimic_model: + _target_: egomimic.algo.diffusion.DFoT + action_dim: 2 + action_horizon: 32 + + sampler: "ddim" + sampler_n_steps: 50 + sampler_eta: 0.0 + + inference_mode: "chunk" + cfg_scale: 1.0 + + domains: ["pushshapes_sim"] + ac_keys: + pushshapes_sim: "actions" + + outer_stage: + _target_: egomimic.algo.diffusion.ImageSpatialDFoTOuterStage + action_dim: 2 + bundle_obs_keys: ["state_agent_obj"] + bundle_obs_dims: [5] + image_key: "front_img_1" + state_action_proj_dim: 384 + vae_checkpoint_path: /coc/flash7/paphiwetsa3/projects/EgoVerse-pact/external_ckpts/vae_v5_last.ckpt + cond_output_key: "fused_cond" + + cond_encoder: + _target_: egomimic.models.stems.cond_encoders.CondEncoderModule + d_cond: 384 + output_key: "fused_cond" + cond_proj_widths: [384, 384] + obs_specs: {} + img_encoders: {} + + backbone: + _target_: egomimic.models.diffusion.DFoTDiT3DBackbone + latent_channels: 4 + latent_size: 6 + patch_size: 1 + action_horizon: 32 + d_model: 384 + d_cond: 384 + cond_dim: 384 + depth: 12 + num_heads: 6 + mlp_ratio: 4.0 + time_embed_dim: 256 + cond_emb_dim: 256 + use_fourier_time: false + stochastic_time_p: 0.0 + cond_dropout_prob: 0.1 + causal: false + + diffusion: + _target_: egomimic.models.diffusion.diffusion.discrete_diffusion.DiscreteDiffusion + action_dim: 4 + timesteps: 1000 + beta_schedule: "cosine" + objective: "pred_v" + loss_weighting_strategy: "min_snr" + snr_clip: 5.0 + clip_noise: 20.0 + +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 1.0e-4 + weight_decay: 0.0001 + +scheduler: + _target_: egomimic.utils.schedulers.warmup_cosine_scheduler + _partial_: true + max_steps: 18800 + warmup_steps: 480 + warmup_start_factor: 0.1 + eta_min: 1.0e-5 diff --git a/egomimic/hydra_configs/model/dfot/image_spatial_cont_sigmoid.yaml b/egomimic/hydra_configs/model/dfot/image_spatial_cont_sigmoid.yaml new file mode 100644 index 000000000..a65f2359a --- /dev/null +++ b/egomimic/hydra_configs/model/dfot/image_spatial_cont_sigmoid.yaml @@ -0,0 +1,81 @@ +_target_: egomimic.pl_utils.pl_model.ModelWrapper + +# DFoT video diffusion with DiT3D backbone + continuous diffusion + sigmoid weighting. +# Tests the additive conditioning fusion fix (matching reference) in isolation. + +robomimic_model: + _target_: egomimic.algo.diffusion.DFoT + action_dim: 2 + action_horizon: 32 + + sampler: "ddim" + sampler_n_steps: 100 + sampler_eta: 0.0 + + inference_mode: "chunk" + cfg_scale: 1.0 + + domains: ["pushshapes_sim"] + ac_keys: + pushshapes_sim: "actions" + + outer_stage: + _target_: egomimic.algo.diffusion.ImageSpatialDFoTOuterStage + action_dim: 2 + bundle_obs_keys: ["state_agent_obj"] + bundle_obs_dims: [5] + image_key: "front_img_1" + state_action_proj_dim: 384 + vae_checkpoint_path: /coc/flash7/paphiwetsa3/projects/EgoVerse-pact/external_ckpts/vae_v5_last.ckpt + cond_output_key: "fused_cond" + + cond_encoder: + _target_: egomimic.models.stems.cond_encoders.CondEncoderModule + d_cond: 384 + output_key: "fused_cond" + cond_proj_widths: [384, 384] + obs_specs: {} + img_encoders: {} + + backbone: + _target_: egomimic.models.diffusion.DFoTDiT3DBackbone + latent_channels: 4 + latent_size: 6 + patch_size: 2 + action_horizon: 32 + d_model: 384 + d_cond: 384 + cond_dim: 384 + depth: 12 + num_heads: 6 + mlp_ratio: 4.0 + time_embed_dim: 256 + cond_emb_dim: 256 + use_fourier_time: false + stochastic_time_p: 0.0 + cond_dropout_prob: 0.1 + causal: false + + diffusion: + _target_: egomimic.models.diffusion.diffusion.continuous_diffusion.ContinuousDiffusion + action_dim: 4 + precond_scale: 0.25 + sigmoid_bias: -1.0 + clip_noise: 20.0 + logsnr_min: -15.0 + logsnr_max: 15.0 + shift: 1.0 + +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 1.0e-4 + weight_decay: 0.0001 + +scheduler: + _target_: egomimic.utils.schedulers.warmup_cosine_scheduler + _partial_: true + max_steps: 18800 + warmup_steps: 480 + warmup_start_factor: 0.1 + eta_min: 1.0e-5 diff --git a/egomimic/hydra_configs/model/dfot/image_spatial_continuous.yaml b/egomimic/hydra_configs/model/dfot/image_spatial_continuous.yaml new file mode 100644 index 000000000..a8dfbedf2 --- /dev/null +++ b/egomimic/hydra_configs/model/dfot/image_spatial_continuous.yaml @@ -0,0 +1,81 @@ +_target_: egomimic.pl_utils.pl_model.ModelWrapper + +# DFoT video diffusion with DiT3D backbone + continuous diffusion + uniform weighting. + +robomimic_model: + _target_: egomimic.algo.diffusion.DFoT + action_dim: 2 + action_horizon: 32 + + sampler: "ddim" + sampler_n_steps: 100 + sampler_eta: 0.0 + + inference_mode: "chunk" + cfg_scale: 1.0 + + domains: ["pushshapes_sim"] + ac_keys: + pushshapes_sim: "actions" + + outer_stage: + _target_: egomimic.algo.diffusion.ImageSpatialDFoTOuterStage + action_dim: 2 + bundle_obs_keys: ["state_agent_obj"] + bundle_obs_dims: [5] + image_key: "front_img_1" + state_action_proj_dim: 384 + vae_checkpoint_path: /coc/flash7/paphiwetsa3/projects/EgoVerse-pact/external_ckpts/vae_v5_last.ckpt + cond_output_key: "fused_cond" + + cond_encoder: + _target_: egomimic.models.stems.cond_encoders.CondEncoderModule + d_cond: 384 + output_key: "fused_cond" + cond_proj_widths: [384, 384] + obs_specs: {} + img_encoders: {} + + backbone: + _target_: egomimic.models.diffusion.DFoTDiT3DBackbone + latent_channels: 4 + latent_size: 6 + patch_size: 2 + action_horizon: 32 + d_model: 384 + d_cond: 384 + cond_dim: 384 + depth: 12 + num_heads: 6 + mlp_ratio: 4.0 + time_embed_dim: 256 + cond_emb_dim: 256 + use_fourier_time: false + stochastic_time_p: 0.0 + cond_dropout_prob: 0.1 + causal: false + + diffusion: + _target_: egomimic.models.diffusion.diffusion.continuous_diffusion.ContinuousDiffusion + action_dim: 4 + precond_scale: 0.25 + sigmoid_bias: -1.0 + clip_noise: 20.0 + logsnr_min: -15.0 + logsnr_max: 15.0 + shift: 1.0 + loss_weighting: "uniform" + +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 1.0e-4 + weight_decay: 0.0001 + +scheduler: + _target_: egomimic.utils.schedulers.warmup_cosine_scheduler + _partial_: true + max_steps: 18800 + warmup_steps: 480 + warmup_start_factor: 0.1 + eta_min: 1.0e-5 diff --git a/egomimic/hydra_configs/model/dfot/image_spatial_policy.yaml b/egomimic/hydra_configs/model/dfot/image_spatial_policy.yaml new file mode 100644 index 000000000..0d9e03b6c --- /dev/null +++ b/egomimic/hydra_configs/model/dfot/image_spatial_policy.yaml @@ -0,0 +1,87 @@ +_target_: egomimic.pl_utils.pl_model.ModelWrapper + +# 2D POLICY variant: spatial VAE-latent image target (NOT flattened) PLUS a +# jointly-diffused action, carried as a per-frame action token appended to the +# DiT3D sequence (backbone.action_token_dim > 0). State enters via external_cond. +# Predicts BOTH future image latents AND actions. Contrast: +# - dfot_pushshapes_image_spatial.yaml = 2D world model (action -> cond) +# - dfot_pushshapes_obs_action_image.yaml = 1D policy (flattened bundle) + +robomimic_model: + _target_: egomimic.algo.diffusion.DFoT + action_dim: 2 + action_horizon: 32 + + sampler: "ddim" + sampler_n_steps: 50 + sampler_eta: 0.0 + + inference_mode: "chunk" + cfg_scale: 1.0 + + domains: ["pushshapes_sim"] + ac_keys: + pushshapes_sim: "actions" + + outer_stage: + _target_: egomimic.algo.diffusion.SpatialObsActionPolicyDFoTOuterStage + action_dim: 2 + action_loss_weight: 1.0 + bundle_obs_keys: ["state_agent_obj"] + bundle_obs_dims: [5] + image_key: "front_img_1" + state_action_proj_dim: 384 + vae_checkpoint_path: /coc/flash7/paphiwetsa3/projects/EgoVerse-pact/external_ckpts/vae_v5_last.ckpt + cond_output_key: "fused_cond" + + cond_encoder: + _target_: egomimic.models.stems.cond_encoders.CondEncoderModule + d_cond: 384 + output_key: "fused_cond" + cond_proj_widths: [384, 384] + obs_specs: {} + img_encoders: {} + + backbone: + _target_: egomimic.models.diffusion.DFoTDiT3DBackbone + latent_channels: 4 + latent_size: 6 + patch_size: 1 + action_horizon: 32 + d_model: 384 + d_cond: 384 + cond_dim: 384 # state-only external_cond width + action_token_dim: 2 # <-- jointly diffuse a per-frame action token + depth: 12 + num_heads: 6 + mlp_ratio: 4.0 + time_embed_dim: 256 + cond_emb_dim: 256 + use_fourier_time: false + stochastic_time_p: 0.0 + cond_dropout_prob: 0.1 + causal: false + + diffusion: + _target_: egomimic.models.diffusion.diffusion.discrete_diffusion.DiscreteDiffusion + action_dim: 4 + timesteps: 1000 + beta_schedule: "cosine" + objective: "pred_v" + loss_weighting_strategy: "min_snr" + snr_clip: 5.0 + clip_noise: 20.0 + +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 1.0e-4 + weight_decay: 0.0001 + +scheduler: + _target_: egomimic.utils.schedulers.warmup_cosine_scheduler + _partial_: true + max_steps: 18800 + warmup_steps: 480 + warmup_start_factor: 0.1 + eta_min: 1.0e-5 diff --git a/egomimic/hydra_configs/model/dfot/obs_action.yaml b/egomimic/hydra_configs/model/dfot/obs_action.yaml new file mode 100644 index 000000000..94f234b53 --- /dev/null +++ b/egomimic/hydra_configs/model/dfot/obs_action.yaml @@ -0,0 +1,115 @@ +_target_: egomimic.pl_utils.pl_model.ModelWrapper + +# DFoT with joint obs+action diffusion forcing. +# +# Mirrors the df_planning pattern from buoyancy99/diffusion-forcing: per +# step, the diffused target is concat([state, action]) instead of just +# action. The backbone is a "world model + policy" jointly: it denoises +# the next state alongside the next action, with per-token independent +# noise levels (DFoT semantics preserved). +# +# Bundle layout (PushShapes): +# state_agent_obj (5D) || actions (2D) = 7D per step +# +# Image obs (front_img_1) stays in the AdaLN conditioning path — it's +# observed every tick, not diffused. The cond_encoder below therefore +# has NO obs_specs (state goes into the bundle, not into cond) and +# only img_encoders. + +robomimic_model: + _target_: egomimic.algo.diffusion.DFoT + action_dim: 2 + action_horizon: 32 + + sampler: "ddim" + sampler_n_steps: 50 + sampler_eta: 0.0 + + inference_mode: "ar" + ar_inference_chunk_size: 1 + ar_inference_step_size: 1 + cfg_scale: 1.0 + + domains: ["pushshapes_sim"] + ac_keys: + pushshapes_sim: "actions" + + outer_stage: + _target_: egomimic.algo.diffusion.ObsActionDFoTOuterStage + action_dim: 2 + cond_output_key: "fused_cond" + + # Which obs keys to include in the diffusion bundle (concatenated + # last-dim, in order, with action appended after). Each entry's dim + # below must match the obs tensor's trailing dim coming from the + # data loader. + bundle_obs_keys: ["state_agent_obj"] + bundle_obs_dims: [5] + + # cond_encoder: image-only. state_agent_obj is in the bundle, so we + # do NOT also feed it through the AdaLN cond path (that would be + # double-counting). If you want a cond term that doesn't appear in + # the bundle, add it under obs_specs. + cond_encoder: + _target_: egomimic.models.stems.cond_encoders.CondEncoderModule + d_cond: 512 + output_key: "fused_cond" + cond_proj_widths: [512, 512] + obs_specs: {} + img_encoders: + front_img_1: + _target_: egomimic.models.stems.image_encoders.SimpleConv + in_channels: 3 + channels: [64, 128, 256, 512] + kernel_size: 3 + stride: 2 + embed_dim: 256 + norm_groups: 8 + + # Backbone width = bundle width = obs(5) + action(2) = 7. The output + # projection produces v of the same shape, and DFoTLoss reduces MSE + # over all 7 dims (obs and action contribute jointly to the loss). + backbone: + _target_: egomimic.models.diffusion.DFoTBackbone + action_dim: 7 + action_horizon: 32 + d_model: 512 + d_cond: 512 + cond_dim: 512 + time_embed_dim: 512 + cond_emb_dim: 512 + use_fourier_time: true + stochastic_time_p: 0.0 + cond_dropout_prob: 0.1 + arch_layout: "T12" + num_heads: 8 + d_intermediate: 2048 + rotary_emb_dim: 64 + dropout: 0.1 + resid_dropout: 0.1 + causal: true + + # Continuous diffusion sized to the bundle width too. + diffusion: + _target_: egomimic.models.diffusion.diffusion.continuous_diffusion.ContinuousDiffusion + action_dim: 7 + precond_scale: 0.25 + sigmoid_bias: -1.0 + clip_noise: 20.0 + logsnr_min: -15.0 + logsnr_max: 15.0 + shift: 1.0 + +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 1.0e-4 + weight_decay: 0.0001 + +scheduler: + _target_: egomimic.utils.schedulers.warmup_cosine_scheduler + _partial_: true + max_steps: 3200 + warmup_steps: 100 + warmup_start_factor: 0.1 + eta_min: 1.0e-5 diff --git a/egomimic/hydra_configs/model/dfot/obs_action_image.yaml b/egomimic/hydra_configs/model/dfot/obs_action_image.yaml new file mode 100644 index 000000000..71cbddfd7 --- /dev/null +++ b/egomimic/hydra_configs/model/dfot/obs_action_image.yaml @@ -0,0 +1,98 @@ +_target_: egomimic.pl_utils.pl_model.ModelWrapper + +# Joint obs+action+image diffusion forcing. +# Per step the diffusion target is concat([state, vae_latent_flat, action]). +# A frozen pretrained ImageVAE (spatial latent 4x6x6) supplies the latent. +# +# Bundle layout per step: +# state_agent_obj (5D) || vae_latent_flat (4*6*6 = 144D) || actions (2D) = 151D +# +# cond_encoder is intentionally empty: the image already lives in the +# bundle via its VAE encoding; AdaLN-conditioning on it again would +# double-count. At inference the t=0 image is encoded by the same VAE +# and used as the initial bundle prefix (handled by the eval). + +robomimic_model: + _target_: egomimic.algo.diffusion.DFoT + action_dim: 2 + action_horizon: 32 + + sampler: "ddim" + sampler_n_steps: 50 + sampler_eta: 0.0 + + inference_mode: "ar" + ar_inference_chunk_size: 1 + ar_inference_step_size: 1 + cfg_scale: 1.0 + + domains: ["pushshapes_sim"] + ac_keys: + pushshapes_sim: "actions" + + outer_stage: + _target_: egomimic.algo.diffusion.ObsActionImageDFoTOuterStage + action_dim: 2 + cond_output_key: "fused_cond" + + # State is the only "real" obs key in the bundle. Image latent is + # added automatically (144D) by the outer stage. + bundle_obs_keys: ["state_agent_obj"] + bundle_obs_dims: [5] + image_key: "front_img_1" + vae_checkpoint_path: /coc/flash7/paphiwetsa3/projects/EgoVerse-pact/external_ckpts/vae_v5_last.ckpt + + # Empty cond_encoder -- no AdaLN cond. The image is in the bundle. + cond_encoder: + _target_: egomimic.models.stems.cond_encoders.CondEncoderModule + d_cond: 512 + output_key: "fused_cond" + cond_proj_widths: [512, 512] + obs_specs: {} + img_encoders: {} + + # Bundle width = 5 + 144 + 2 = 151. + backbone: + _target_: egomimic.models.diffusion.DFoTBackbone + action_dim: 151 + action_horizon: 32 + d_model: 512 + d_cond: 512 + # No external cond -- image is in the bundle, not via AdaLN. + cond_dim: 0 + time_embed_dim: 512 + cond_emb_dim: 0 + use_fourier_time: true + stochastic_time_p: 0.0 + cond_dropout_prob: 0.0 # no cond -> CFG meaningless + arch_layout: "T12" + num_heads: 8 + d_intermediate: 2048 + rotary_emb_dim: 64 + dropout: 0.1 + resid_dropout: 0.1 + causal: true + + diffusion: + _target_: egomimic.models.diffusion.diffusion.continuous_diffusion.ContinuousDiffusion + action_dim: 151 + precond_scale: 0.25 + sigmoid_bias: -1.0 + clip_noise: 20.0 + logsnr_min: -15.0 + logsnr_max: 15.0 + shift: 1.0 + +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 1.0e-4 + weight_decay: 0.0001 + +scheduler: + _target_: egomimic.utils.schedulers.warmup_cosine_scheduler + _partial_: true + max_steps: 3200 + warmup_steps: 100 + warmup_start_factor: 0.1 + eta_min: 1.0e-5 diff --git a/egomimic/hydra_configs/model/dfot/obs_action_image_wm.yaml b/egomimic/hydra_configs/model/dfot/obs_action_image_wm.yaml new file mode 100644 index 000000000..adf21cca8 --- /dev/null +++ b/egomimic/hydra_configs/model/dfot/obs_action_image_wm.yaml @@ -0,0 +1,98 @@ +_target_: egomimic.pl_utils.pl_model.ModelWrapper + +# 1D WORLD MODEL variant of obs+action+image diffusion forcing. +# Action is NOT in the bundle: it is fed as AdaLN conditioning (external_cond). +# The diffused bundle is obs-only: concat([state, vae_latent_flat]). +# +# Bundle layout per step: +# state_agent_obj (5D) || vae_latent_flat (4*6*6 = 144D) = 149D (target) +# actions (2D) -> external_cond +# +# This is the action-conditioned video predictor (df_video-style): given the +# action sequence, predict the future obs/frames. Contrast with the POLICY +# variant (dfot_pushshapes_obs_action_image.yaml) where action_in_bundle=True +# so the action is jointly predicted. + +robomimic_model: + _target_: egomimic.algo.diffusion.DFoT + action_dim: 2 + action_horizon: 32 + + sampler: "ddim" + sampler_n_steps: 50 + sampler_eta: 0.0 + + inference_mode: "ar" + ar_inference_chunk_size: 1 + ar_inference_step_size: 1 + cfg_scale: 1.0 + + domains: ["pushshapes_sim"] + ac_keys: + pushshapes_sim: "actions" + + outer_stage: + _target_: egomimic.algo.diffusion.ObsActionImageDFoTOuterStage + action_dim: 2 + cond_output_key: "fused_cond" + action_in_bundle: false # <-- world model: action is conditioning + + bundle_obs_keys: ["state_agent_obj"] + bundle_obs_dims: [5] + image_key: "front_img_1" + vae_checkpoint_path: /coc/flash7/paphiwetsa3/projects/EgoVerse-pact/external_ckpts/vae_v5_last.ckpt + + # cond_encoder stays empty -- the image is in the bundle, and the action + # is supplied directly as external_cond by the outer stage's encode(). + cond_encoder: + _target_: egomimic.models.stems.cond_encoders.CondEncoderModule + d_cond: 512 + output_key: "fused_cond" + cond_proj_widths: [512, 512] + obs_specs: {} + img_encoders: {} + + # Bundle width = 5 + 144 = 149 (no action). Action (2D) -> external_cond. + backbone: + _target_: egomimic.models.diffusion.DFoTBackbone + action_dim: 149 + action_horizon: 32 + d_model: 512 + d_cond: 512 + cond_dim: 2 # <-- action width fed as external_cond + cond_emb_dim: 128 + cond_dropout_prob: 0.0 # set >0 to enable action CFG + time_embed_dim: 512 + use_fourier_time: true + stochastic_time_p: 0.0 + arch_layout: "T12" + num_heads: 8 + d_intermediate: 2048 + rotary_emb_dim: 64 + dropout: 0.1 + resid_dropout: 0.1 + causal: true + + diffusion: + _target_: egomimic.models.diffusion.diffusion.continuous_diffusion.ContinuousDiffusion + action_dim: 149 + precond_scale: 0.25 + sigmoid_bias: -1.0 + clip_noise: 20.0 + logsnr_min: -15.0 + logsnr_max: 15.0 + shift: 1.0 + +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 1.0e-4 + weight_decay: 0.0001 + +scheduler: + _target_: egomimic.utils.schedulers.warmup_cosine_scheduler + _partial_: true + max_steps: 3200 + warmup_steps: 100 + warmup_start_factor: 0.1 + eta_min: 1.0e-5 diff --git a/egomimic/hydra_configs/model/dfot/pixel.yaml b/egomimic/hydra_configs/model/dfot/pixel.yaml new file mode 100644 index 000000000..717ec19ad --- /dev/null +++ b/egomimic/hydra_configs/model/dfot/pixel.yaml @@ -0,0 +1,82 @@ +_target_: egomimic.pl_utils.pl_model.ModelWrapper + +# DFoT pixel-space video diffusion (no VAE). +# Matches reference config exactly (except patch_size=8 for memory). + +robomimic_model: + _target_: egomimic.algo.diffusion.DFoT + action_dim: 2 + action_horizon: 9 + + sampler: "ddim" + sampler_n_steps: 50 + sampler_eta: 0.0 + + inference_mode: "chunk" + cfg_scale: 1.0 + + domains: ["pushshapes_sim"] + ac_keys: + pushshapes_sim: "actions" + + outer_stage: + _target_: egomimic.algo.diffusion.PixelSpatialDFoTOuterStage + action_dim: 2 + image_key: "front_img_1" + image_channels: 3 + image_size: 96 + frame_sampling: "fixed_window" + sample_n_frames: 9 + + cond_encoder: + _target_: egomimic.models.stems.cond_encoders.CondEncoderModule + d_cond: 256 + output_key: "fused_cond" + cond_proj_widths: [256, 256] + obs_specs: {} + img_encoders: {} + + backbone: + _target_: egomimic.models.diffusion.DFoTDiT3DBackbone + latent_channels: 3 + latent_size: 96 + patch_size: 2 + action_horizon: 9 + d_model: 256 + d_cond: 256 + cond_dim: 0 + depth: 6 + num_heads: 4 + mlp_ratio: 4.0 + time_embed_dim: 64 + cond_emb_dim: 0 + use_fourier_time: false + stochastic_time_p: 0.0 + cond_dropout_prob: 0.0 + causal: false + + diffusion: + _target_: egomimic.models.diffusion.diffusion.discrete_diffusion.DiscreteDiffusion + action_dim: 3 + timesteps: 1000 + beta_schedule: "cosine" + objective: "pred_v" + loss_weighting_strategy: "fused_min_snr" + snr_clip: 5.0 + cum_snr_decay: 0.96 + clip_noise: 20.0 + +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 2.0e-4 + weight_decay: 0.0 + betas: [0.9, 0.99] + +scheduler: + _target_: egomimic.utils.schedulers.warmup_cosine_scheduler + _partial_: true + max_steps: 100000 + warmup_steps: 1000 + warmup_start_factor: 0.01 + eta_min: 2.0e-4 diff --git a/egomimic/hydra_configs/model/dfot/pixel_decoupled.yaml b/egomimic/hydra_configs/model/dfot/pixel_decoupled.yaml new file mode 100644 index 000000000..8c51f1d26 --- /dev/null +++ b/egomimic/hydra_configs/model/dfot/pixel_decoupled.yaml @@ -0,0 +1,89 @@ +_target_: egomimic.pl_utils.pl_model.ModelWrapper + +# DFoT pixel-space OBS+ACTION policy (no VAE) — DECOUPLED action token. +# The action is NOT fused into obs channels; it rides as its own DiT3D token +# with an INDEPENDENT per-frame noise level (backbone.action_token_dim=2, +# outer_stage.decouple_action_noise=true). At inference obs_t is pinned clean +# while a_t is denoised in the SAME frame -> predicts a_t from obs_t, NO 1-frame +# offset, and the action is never a clean input (cannot copy). Forces obs->action. + +robomimic_model: + _target_: egomimic.algo.diffusion.DFoT + action_dim: 2 + action_horizon: 9 + + sampler: "ddim" + sampler_n_steps: 50 + sampler_eta: 0.0 + + inference_mode: "pixel_decoupled" + cfg_scale: 1.0 + + domains: ["pushshapes_sim"] + ac_keys: + pushshapes_sim: "actions" + + outer_stage: + _target_: egomimic.algo.diffusion.PixelObsActionDFoTOuterStage + pixel_mode: "decoupled" + action_dim: 2 + action_loss_weight: 1.0 + decouple_action_noise: true + image_key: "front_img_1" + image_channels: 3 + image_size: 96 + frame_sampling: "full" + sample_n_frames: 9 + + cond_encoder: + _target_: egomimic.models.stems.cond_encoders.CondEncoderModule + d_cond: 256 + output_key: "fused_cond" + cond_proj_widths: [256, 256] + obs_specs: {} + img_encoders: {} + + backbone: + _target_: egomimic.models.diffusion.DFoTDiT3DBackbone + latent_channels: 3 # RGB only; action is a separate token + latent_size: 96 + patch_size: 2 + action_horizon: 9 + action_token_dim: 2 # <-- separate decoupled action token + d_model: 256 + d_cond: 256 + cond_dim: 0 + depth: 6 + num_heads: 4 + mlp_ratio: 4.0 + time_embed_dim: 64 + cond_emb_dim: 0 + use_fourier_time: false + stochastic_time_p: 0.0 + cond_dropout_prob: 0.0 + causal: false + + diffusion: + _target_: egomimic.models.diffusion.diffusion.discrete_diffusion.DiscreteDiffusion + action_dim: 3 # obs token channels (RGB) for q_sample/loss + timesteps: 1000 + beta_schedule: "cosine" + objective: "pred_v" + loss_weighting_strategy: "min_snr" + snr_clip: 5.0 + clip_noise: 20.0 + +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 2.0e-4 + weight_decay: 0.0 + betas: [0.9, 0.99] + +scheduler: + _target_: egomimic.utils.schedulers.warmup_cosine_scheduler + _partial_: true + max_steps: 24000 + warmup_steps: 1000 + warmup_start_factor: 0.01 + eta_min: 2.0e-4 diff --git a/egomimic/hydra_configs/model/dfot/pixel_policy.yaml b/egomimic/hydra_configs/model/dfot/pixel_policy.yaml new file mode 100644 index 000000000..f8c990ee2 --- /dev/null +++ b/egomimic/hydra_configs/model/dfot/pixel_policy.yaml @@ -0,0 +1,87 @@ +_target_: egomimic.pl_utils.pl_model.ModelWrapper + +# DFoT pixel-space OBS+ACTION policy (no VAE) — Design A. +# Action carried as broadcast channels appended to the RGB frame and JOINTLY +# diffused by the DiT3D (latent_channels = 3 RGB + action_channels). Decode at +# inference = global-avg-pool of the predicted action planes (a conv decode is a +# config flip / Design B is the regression variant in a sibling config). + +robomimic_model: + _target_: egomimic.algo.diffusion.DFoT + action_dim: 2 + action_horizon: 9 + + sampler: "ddim" + sampler_n_steps: 50 + sampler_eta: 0.0 + + inference_mode: "chunk" + cfg_scale: 1.0 + + domains: ["pushshapes_sim"] + ac_keys: + pushshapes_sim: "actions" + + outer_stage: + _target_: egomimic.algo.diffusion.PixelObsActionDFoTOuterStage + pixel_mode: "policy" + action_dim: 2 + action_channels: 2 + image_key: "front_img_1" + image_channels: 3 + image_size: 96 + frame_sampling: "fixed_window" + sample_n_frames: 9 + + cond_encoder: + _target_: egomimic.models.stems.cond_encoders.CondEncoderModule + d_cond: 256 + output_key: "fused_cond" + cond_proj_widths: [256, 256] + obs_specs: {} + img_encoders: {} + + backbone: + _target_: egomimic.models.diffusion.DFoTDiT3DBackbone + latent_channels: 5 # 3 RGB + 2 broadcast action planes + latent_size: 96 + patch_size: 2 + action_horizon: 9 + d_model: 256 + d_cond: 256 + cond_dim: 0 + depth: 6 + num_heads: 4 + mlp_ratio: 4.0 + time_embed_dim: 64 + cond_emb_dim: 0 + use_fourier_time: false + stochastic_time_p: 0.0 + cond_dropout_prob: 0.0 + causal: false + + diffusion: + _target_: egomimic.models.diffusion.diffusion.discrete_diffusion.DiscreteDiffusion + action_dim: 5 # diffused per-token channel count (3+2) + timesteps: 1000 + beta_schedule: "cosine" + objective: "pred_v" + loss_weighting_strategy: "fused_min_snr" + snr_clip: 5.0 + cum_snr_decay: 0.96 + clip_noise: 20.0 + +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 2.0e-4 + weight_decay: 0.0 + betas: [0.9, 0.99] + +scheduler: + _target_: egomimic.utils.schedulers.warmup_cosine_scheduler + _partial_: true + max_steps: 100000 + warmup_steps: 1000 + warmup_start_factor: 0.01 + eta_min: 2.0e-4 diff --git a/egomimic/hydra_configs/model/dfot/pixel_regress.yaml b/egomimic/hydra_configs/model/dfot/pixel_regress.yaml new file mode 100644 index 000000000..ef52a06df --- /dev/null +++ b/egomimic/hydra_configs/model/dfot/pixel_regress.yaml @@ -0,0 +1,87 @@ +_target_: egomimic.pl_utils.pl_model.ModelWrapper + +# DFoT pixel-space OBS+ACTION policy (no VAE) — Design B (regression). +# Diffuses RGB video (latent_channels=3, the proven foundation); a conv head +# regresses the action from the model's predicted clean frame. Action is NOT a +# diffusion target. Contrast to Design A (broadcast channels, jointly diffused). + +robomimic_model: + _target_: egomimic.algo.diffusion.DFoT + action_dim: 2 + action_horizon: 9 + + sampler: "ddim" + sampler_n_steps: 50 + sampler_eta: 0.0 + + inference_mode: "chunk" + cfg_scale: 1.0 + + domains: ["pushshapes_sim"] + ac_keys: + pushshapes_sim: "actions" + + outer_stage: + _target_: egomimic.algo.diffusion.PixelObsActionDFoTOuterStage + pixel_mode: "regress" + action_dim: 2 + action_loss_weight: 1.0 + head_width: 64 + image_key: "front_img_1" + image_channels: 3 + image_size: 96 + frame_sampling: "fixed_window" + sample_n_frames: 9 + + cond_encoder: + _target_: egomimic.models.stems.cond_encoders.CondEncoderModule + d_cond: 256 + output_key: "fused_cond" + cond_proj_widths: [256, 256] + obs_specs: {} + img_encoders: {} + + backbone: + _target_: egomimic.models.diffusion.DFoTDiT3DBackbone + latent_channels: 3 # RGB video only; action is a regression head + latent_size: 96 + patch_size: 2 + action_horizon: 9 + d_model: 256 + d_cond: 256 + cond_dim: 0 + depth: 6 + num_heads: 4 + mlp_ratio: 4.0 + time_embed_dim: 64 + cond_emb_dim: 0 + use_fourier_time: false + stochastic_time_p: 0.0 + cond_dropout_prob: 0.0 + causal: false + + diffusion: + _target_: egomimic.models.diffusion.diffusion.discrete_diffusion.DiscreteDiffusion + action_dim: 3 # RGB channel count + timesteps: 1000 + beta_schedule: "cosine" + objective: "pred_v" + loss_weighting_strategy: "fused_min_snr" + snr_clip: 5.0 + cum_snr_decay: 0.96 + clip_noise: 20.0 + +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 2.0e-4 + weight_decay: 0.0 + betas: [0.9, 0.99] + +scheduler: + _target_: egomimic.utils.schedulers.warmup_cosine_scheduler + _partial_: true + max_steps: 100000 + warmup_steps: 1000 + warmup_start_factor: 0.01 + eta_min: 2.0e-4 diff --git a/egomimic/hydra_configs/model/dfot/pixel_video.yaml b/egomimic/hydra_configs/model/dfot/pixel_video.yaml new file mode 100644 index 000000000..9a8b640d9 --- /dev/null +++ b/egomimic/hydra_configs/model/dfot/pixel_video.yaml @@ -0,0 +1,79 @@ +_target_: egomimic.pl_utils.pl_model.ModelWrapper + +# DFoT pixel-space video diffusion — clean implementation. +# Uses ZarrVideoClipDataset for (B, 9, 3, 96, 96) batches. +# No packing, no per-episode loops. Matches reference exactly. + +robomimic_model: + _target_: egomimic.algo.diffusion.DFoT + action_dim: 2 + action_horizon: 9 + + sampler: "ddim" + sampler_n_steps: 50 + sampler_eta: 0.0 + + inference_mode: "chunk" + cfg_scale: 1.0 + + domains: ["pushshapes_sim"] + ac_keys: + pushshapes_sim: "actions" + + outer_stage: + _target_: egomimic.algo.diffusion.PixelVideoDFoTOuterStage + action_dim: 2 + video_key: "front_img_1" + + cond_encoder: + _target_: egomimic.models.stems.cond_encoders.CondEncoderModule + d_cond: 256 + output_key: "fused_cond" + cond_proj_widths: [256, 256] + obs_specs: {} + img_encoders: {} + + backbone: + _target_: egomimic.models.diffusion.DFoTDiT3DBackbone + latent_channels: 3 + latent_size: 96 + patch_size: 2 + action_horizon: 9 + d_model: 256 + d_cond: 256 + cond_dim: 0 + depth: 6 + num_heads: 4 + mlp_ratio: 4.0 + time_embed_dim: 64 + cond_emb_dim: 0 + use_fourier_time: false + stochastic_time_p: 0.0 + cond_dropout_prob: 0.0 + causal: false + + diffusion: + _target_: egomimic.models.diffusion.diffusion.discrete_diffusion.DiscreteDiffusion + action_dim: 3 + timesteps: 1000 + beta_schedule: "cosine" + objective: "pred_v" + loss_weighting_strategy: "fused_min_snr" + snr_clip: 5.0 + cum_snr_decay: 0.96 + clip_noise: 20.0 + +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 2.0e-4 + weight_decay: 0.0 + betas: [0.9, 0.99] + +scheduler: + _target_: egomimic.utils.schedulers.warmup_cosine_scheduler + _partial_: true + max_steps: 100000 + warmup_steps: 1000 + warmup_start_factor: 0.01 + eta_min: 2.0e-4 diff --git a/tests/regression/test_dfot_inference.py b/tests/regression/test_dfot_inference.py new file mode 100644 index 000000000..8c7b06eba --- /dev/null +++ b/tests/regression/test_dfot_inference.py @@ -0,0 +1,74 @@ +"""Smoke: refactored DFoT inference paths (closed-loop AR + chunk) still work +via the property accessors.""" +import pytest # noqa: E402 + +pytest.skip( + "manual regression smoke: hardcoded EgoVerse-clone-3 paths + configs " + "removed from this repo + needs GPU/checkpoints; run directly with " + "python, skipped under pytest collection", + allow_module_level=True, +) + +import sys, torch, numpy as np +sys.path.insert(0, "/storage/project/r-dxu345-0/paphiwetsa3/projects/EgoVerse-clone-3") +from omegaconf import OmegaConf +from hydra.utils import instantiate + + +class MockNormStats: + def __init__(self): + self._t = {0: {"action_keys": ["actions"], "proprio_keys": ["state_agent_obj"], + "lang_keys": [], "camera_keys": ["front_img_1"]}} + def keys_of_type(self, kind, emb_id): + return self._t.get(emb_id, {}).get(kind, []) + def is_key_with_embodiment(self, k, eid): return True + def zarr_key_to_keyname(self, k, eid): return k + def normalize(self, b, e): return b + def unnormalize(self, b, e): return b + + +import egomimic.rldb.embodiment.embodiment as _emb +orig_get_id = _emb.get_embodiment_id +orig_get_name = _emb.get_embodiment +_emb.get_embodiment_id = lambda n: 0 if n == "pushshapes_sim" else orig_get_id(n) +_emb.get_embodiment = lambda i: "pushshapes_sim" if i == 0 else orig_get_name(i) + +cfg = OmegaConf.load( + "/storage/project/r-dxu345-0/paphiwetsa3/projects/EgoVerse-clone-3/" + "egomimic/hydra_configs/model/dfot/base.yaml" +) +dfot = instantiate(cfg.robomimic_model, norm_stats=MockNormStats(), _recursive_=True) +dfot.nets = dfot.nets.cpu() +dfot.device = torch.device("cpu") + +# Synthetic single-env-tick obs (normalized). +obs = { + "state_agent_obj": torch.randn(1, 5), + "front_img_1": torch.rand(1, 3, 96, 96), +} + +# ---- AR mode (default) ---- +dfot.inference_mode = "ar" +dfot.ar_inference_chunk_size = 1 +dfot.ar_inference_step_size = 1 +print(f"[ar] inference_mode={dfot.inference_mode}, chunk={dfot.ar_inference_chunk_size}, step={dfot.ar_inference_step_size}") + +action_t0 = dfot.inference_step(obs, t=0, emb_id=0) +print(f"[ar] action @ t=0: {action_t0} shape={action_t0.shape}") +assert action_t0.shape == (2,), f"expected (2,), got {action_t0.shape}" +assert np.isfinite(action_t0).all() + +action_t1 = dfot.inference_step(obs, t=1, emb_id=0) +print(f"[ar] action @ t=1: {action_t1}") +assert action_t1.shape == (2,) + +# ---- Chunk mode ---- +dfot.inference_mode = "chunk" +if hasattr(dfot, "_sim_state"): + del dfot._sim_state +print(f"\n[chunk] inference_mode={dfot.inference_mode}") +action_chunk_t0 = dfot.inference_step(obs, t=0, emb_id=0) +print(f"[chunk] action @ t=0: {action_chunk_t0} shape={action_chunk_t0.shape}") +assert action_chunk_t0.shape == (2,) + +print("\nPASS — refactored DFoT inference_step works in both ar + chunk modes") diff --git a/tests/regression/test_dfot_refactor_e2e.py b/tests/regression/test_dfot_refactor_e2e.py new file mode 100644 index 000000000..23e6e7a25 --- /dev/null +++ b/tests/regression/test_dfot_refactor_e2e.py @@ -0,0 +1,92 @@ +"""End-to-end smoke: instantiate refactored DFoT via hydra config, run +forward_training, check loss is finite and a scalar tensor.""" +import pytest # noqa: E402 + +pytest.skip( + "manual regression smoke: hardcoded EgoVerse-clone-3 paths + configs " + "removed from this repo + needs GPU/checkpoints; run directly with " + "python, skipped under pytest collection", + allow_module_level=True, +) + +import sys, torch +sys.path.insert(0, "/storage/project/r-dxu345-0/paphiwetsa3/projects/EgoVerse-clone-3") +from omegaconf import OmegaConf +from hydra.utils import instantiate + + +class MockNormStats: + """Stub that gives DFoT the bare minimum it needs.""" + def __init__(self): + self._action_keys = {0: ["actions"]} + self._proprio_keys = {0: ["state_agent_obj"]} + self._lang_keys = {0: []} + self._camera_keys = {0: ["front_img_1"]} + + def keys_of_type(self, kind, emb_id): + return { + "action_keys": self._action_keys, + "proprio_keys": self._proprio_keys, + "lang_keys": self._lang_keys, + "camera_keys": self._camera_keys, + }[kind].get(emb_id, []) + + def is_key_with_embodiment(self, key, emb_id): + return True + + def zarr_key_to_keyname(self, key, emb_id): + return key + + def normalize(self, batch, emb_id): + return batch + + def unnormalize(self, batch, emb_id): + return batch + + +# Patch the embodiment lookup to return our pushshapes_sim id. +import egomimic.rldb.embodiment.embodiment as _emb +orig_get_id = _emb.get_embodiment_id +_emb.get_embodiment_id = lambda n: 0 if n == "pushshapes_sim" else orig_get_id(n) + + +cfg = OmegaConf.load( + "/storage/project/r-dxu345-0/paphiwetsa3/projects/EgoVerse-clone-3/" + "egomimic/hydra_configs/model/dfot/base.yaml" +) +model_wrapper_cfg = cfg.robomimic_model + +# Instantiate DFoT via hydra (which builds outer_stage + auto-builds the loss). +norm_stats = MockNormStats() +dfot = instantiate(model_wrapper_cfg, norm_stats=norm_stats, _recursive_=True) +dfot.nets = dfot.nets.cpu() # avoid CUDA assumption for the test +dfot.device = torch.device("cpu") + +print(f"[init] DFoT built. nets keys: {list(dfot.nets.keys())}") +print(f"[init] outer_stage class: {type(dfot.outer_stage).__name__}") +print(f"[init] loss class: {type(dfot.loss).__name__}") +print(f"[init] backbone (via property) class: {type(dfot.backbone).__name__}") +print(f"[init] cond_encoder (via property) class: {type(dfot.cond_encoder).__name__}") +print(f"[init] diffusion (via property) class: {type(dfot.diffusion).__name__}") + +# Build a synthetic processed batch (padded mode). +B, T, A = 2, 8, 2 +batch = { + 0: { # emb_id + "actions": torch.randn(B, T, A), + "state_agent_obj": torch.randn(B, T, 5), + "front_img_1": torch.rand(B, T, 3, 96, 96), + "_packed": False, + }, +} + +torch.manual_seed(42) +predictions = dfot.forward_training(batch) +loss_dict = dfot.compute_losses(predictions, batch) + +print(f"\n[forward_training] predictions keys: {list(predictions.keys())}") +print(f"[forward_training] 0_action_loss: {predictions['0_action_loss'].item():.6f}") +print(f"[compute_losses] action_loss: {loss_dict['action_loss'].item():.6f}") +assert torch.isfinite(predictions["0_action_loss"]).all() +assert predictions["0_action_loss"].ndim == 0, "loss should be scalar" +print("\nPASS — refactored DFoT instantiates from yaml + runs forward_training")