From 8c8f835a77a6535a5e156f7d6741fb730f2191ea Mon Sep 17 00:00:00 2001 From: Xin Wang Date: Wed, 16 Sep 2026 15:39:31 +0800 Subject: [PATCH] Add frozen EgoSpeak Omni residual fusion stage --- PROJECT_CONTEXT.md | 9 + docs/README.md | 1 + docs/easycom_egospeak_omni_residual.md | 170 +++++++ egospeak_omni_residual.py | 618 ++++++++++++++++++++++++ test_egospeak_omni_residual.py | 169 +++++++ train_egospeak_omni_residual.py | 631 +++++++++++++++++++++++++ 6 files changed, 1598 insertions(+) create mode 100644 docs/easycom_egospeak_omni_residual.md create mode 100644 egospeak_omni_residual.py create mode 100644 test_egospeak_omni_residual.py create mode 100755 train_egospeak_omni_residual.py diff --git a/PROJECT_CONTEXT.md b/PROJECT_CONTEXT.md index ed3e0d1..e757df3 100644 --- a/PROJECT_CONTEXT.md +++ b/PROJECT_CONTEXT.md @@ -135,6 +135,15 @@ result. described as upstream Stage-3 code because that entry point remains unpublished. See the [`runbook`](docs/easycom_muvap_stage3_baseline.md). +- `egospeak-omni-zero-residual` is implemented but unmeasured. It strictly + loads and freezes the released MiniROADA A+V/GRU/anticipation path, reuses + existing causal 5 Hz Omni readouts, adds an explicit ten-bin partial-tail + phase embedding, and learns only a zero-output residual over the original + logits. The launcher verifies exact initial recovery and exports fused and + matched frozen-EgoSpeak predictions together. Existing Omni states already + combine completed history with the current partial tail; separating those + states remains a future extraction ablation. See the + [`runbook`](docs/easycom_egospeak_omni_residual.md). ### Measured oracle-information screen diff --git a/docs/README.md b/docs/README.md index cf3535c..6a14b52 100644 --- a/docs/README.md +++ b/docs/README.md @@ -28,6 +28,7 @@ tests. | [`easycom_track_gru_residual.md`](easycom_track_gru_residual.md) | Matched pooled-to-track GRU experiment, causal identity contract and L20 commands | | [`easycom_causal_person_adapter_v2.md`](easycom_causal_person_adapter_v2.md) | Dynamic-gated, zero-output causal person-track residual and matched L20 screen | | [`easycom_egospeak_style_person_gru.md`](easycom_egospeak_style_person_gru.md) | Controlled V2 temporal-stem ablation using EgoSpeak MiniROADA's layer order | +| [`easycom_egospeak_omni_residual.md`](easycom_egospeak_omni_residual.md) | Stage-one frozen released MiniROADA plus zero-initialized tail-aware Omni residual fusion | | [`easycom_highrate_person_av.md`](easycom_highrate_person_av.md) | Frozen 25 Hz causal ASD person evidence fused with 5 Hz global Omni context | | [`easycom_muvap_stage3_baseline.md`](easycom_muvap_stage3_baseline.md) | MuVAP-only frozen VAP+ASD Stage-3 paper reconstruction on EgoSpeak/EasyCom targets | | [`easycom_oracle_information.md`](easycom_oracle_information.md) | Causal oracle per-person speaking-history upper bound and Session-12 screening commands | diff --git a/docs/easycom_egospeak_omni_residual.md b/docs/easycom_egospeak_omni_residual.md new file mode 100644 index 0000000..82dc846 --- /dev/null +++ b/docs/easycom_egospeak_omni_residual.md @@ -0,0 +1,170 @@ +# Frozen EgoSpeak + Omni tail-aware residual fusion + +## Scope + +This is the first-stage dual-stream screen. It preserves the released +EgoSpeak MiniROADA RGB, wav2vec2-audio, GRU, and ten-horizon anticipation path +as a frozen base. Existing causal 5 Hz Qwen2.5-Omni states provide a second +branch. Only a tail-phase-aware residual fusion module is trained: + +```text +released EgoSpeak RGB + audio -> frozen MiniROADA -> h_ego -> z_ego ---------+ +existing causal Omni g_full -> projection -----------------------> fusion -> Δz + | + z_fused = z_ego + Δz +``` + +The final residual projection is initialized to exactly zero. Before any +optimization, the launcher verifies bit-exact equality between fused and base +logits on a real training example. Unlike the earlier scalar-zero-gate track +screen, the zero is placed only on the final residual output layer; after its +first update, gradients can reach the fusion branch. + +This is **baseline-recoverable initialization**, not a guarantee that the +trained fusion will be non-inferior on unseen Sessions 1--3. The launcher +therefore exports matched frozen-base and fused predictions from the same +forward pass. + +## Inputs + +No Qwen re-extraction is needed for this stage. The Omni root must contain +complete `Session_X.pt` files with: + +- `mode=open-stream-replay` by default; +- dense zero-origin `label_indices`; +- one state per 5 Hz target row; +- `step_seconds=0.2` and `native_block_seconds=2.0`; +- a full-session `times` grid and no `partial=True` marker. + +The saved state is already the readout over rolling completed blocks plus the +current temporary partial block. The fusion branch derives the ten-value tail +phase (`0.2, ..., 2.0` seconds) from the validated grid. It does not claim to +separate a completed-history state from a partial-tail state. Such an +ablation would require a new Omni extraction field. + +The EgoSpeak feature root must contain the released directories: + +```text +rgb_kinetics_resnet50/Session_X.npy [T,2048] +audio_wav2vec2/Session_X.npy [T,5120] +target_perframe/Session_X.npy [T,3] +``` + +The base checkpoint must be the released raw MiniROADA state dict. It is +loaded strictly, and every base parameter remains frozen and in evaluation +mode so its dropout cannot reactivate during residual training. + +## Audit the inputs + +Set machine-specific paths rather than committing `/data1` paths: + +```bash +EGO_ROOT=third_party/EgoSpeak/MiniROAD/data/EasyCom +OMNI_ROOT=/data1/yfzx/wangx/Datasets/easycom/omni_features/open_stream_2fps_8s_last4_stack +BASE_CKPT=third_party/EgoSpeak/MiniROAD/checkpoints/miniroad_easycom_AV.pth +``` + +The default `--omni-layer-fusion mean` uses a fixed mean for the existing +last-four stack and is also valid for a single `[T,D]` state. Use `last` only +for an explicit layer ablation. + +Audit the released EgoSpeak arrays: + +```bash +python3 audit_egospeak_features.py \ + --feature-root "$EGO_ROOT" +``` + +The training launcher independently validates every selected Omni and +EgoSpeak session and aborts on a length, time-grid, mode, dimension, or +zero-origin mismatch. + +## Session 4 smoke test + +Use a new output directory for every invocation; the launcher refuses to +overwrite an existing nonempty result. + +```bash +CUDA_VISIBLE_DEVICES=0 python3 -u train_egospeak_omni_residual.py \ + --egospeak-feature-root "$EGO_ROOT" \ + --omni-feature-root "$OMNI_ROOT" \ + --base-checkpoint "$BASE_CKPT" \ + --train-sessions Session_4 \ + --eval-sessions Session_4 \ + --epochs 1 \ + --max-train-examples 32 \ + --batch-size 8 \ + --num-workers 0 \ + --output-dir results/easycom_egospeak_omni_residual_smoke +``` + +The required early output is: + +```text +Initialization audit: fused logits == frozen EgoSpeak logits exactly +``` + +If a raw released Session 4 prediction is available, add: + +```bash + --reference-prediction-root \ + third_party/EgoSpeak/MiniROAD/output_cfg/miniroad_aa_easycom_session4_AV/eval +``` + +This is a stronger implementation audit: the launcher checks that its frozen +base reproduces the upstream raw `[source,horizon,class]` probabilities. + +## Fixed four-epoch Sessions 4--12 run + +After the smoke test passes, omit session lists to use the official EasyCom +train/test split. Evaluation happens once after epoch 4; the test split is not +used for checkpoint selection. + +```bash +CUDA_VISIBLE_DEVICES=0 python3 -u train_egospeak_omni_residual.py \ + --egospeak-feature-root "$EGO_ROOT" \ + --omni-feature-root "$OMNI_ROOT" \ + --base-checkpoint "$BASE_CKPT" \ + --epochs 4 \ + --batch-size 64 \ + --num-workers 4 \ + --seed 0 \ + --output-dir results/easycom_egospeak_omni_residual_seed0 +``` + +Do not add `--amp` to the first reproducibility run. Once the float32 base has +been compared against released predictions, AMP can be screened separately. +The default `--eval-chunk-size 0` processes each evaluation session in one GRU +pass, matching the released evaluator. Set a positive value only if inference +memory requires chunked recurrent execution, and repeat the reference audit. + +The output contains: + +```text +checkpoint_last.pt self-contained model/checkpoint provenance +metrics.json fused dense metrics plus matched frozen base +raw_eval/ fused MiniROADA source-indexed probabilities +pred/ and valid/ fused evaluator-ready aligned arrays +frozen_egospeak_raw_eval/ matched frozen-base raw probabilities +frozen_egospeak_pred/ matched frozen-base aligned arrays +``` + +`metrics.json` keeps the fused `causal_valid_only` protocol at the top level, +so it remains readable by `summarize_local_baselines.py`. + +## Interpretation and next gate + +This first stage uses only the released dense future-activity objective. It +tests whether Omni adds useful context without discarding the task-specific +EgoSpeak path. It does not train semantic ENTRY/YIELD heads, and Sessions +1--3 LLM-assisted reviews must never be used for optimization or model +selection. + +Advance to selective MiniROADA unfreezing or semantic opportunity heads only +if the frozen residual passes these audits: + +1. exact zero-residual initialization and upstream base reproduction; +2. no causal/time-grid mismatch; +3. stable dense and boundary changes over seeds; +4. opportunity-conditioned ENTRY/YIELD reported separately, without treating + dense wearer probability as semantic intent. diff --git a/egospeak_omni_residual.py b/egospeak_omni_residual.py new file mode 100644 index 0000000..83ddfc6 --- /dev/null +++ b/egospeak_omni_residual.py @@ -0,0 +1,618 @@ +"""Frozen EgoSpeak MiniROADA plus a trainable Omni residual branch. + +This module deliberately does not modify the vendored EgoSpeak checkout. The +``MiniROADABase`` parameter names match the released implementation so its raw +state dict can be loaded strictly. Stage-one training freezes that complete +base and learns only a zero-initialized residual over the ten anticipation +logits. +""" + +from __future__ import annotations + +import hashlib +from collections.abc import Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal + +import numpy as np +import torch +from torch import nn +from torch.nn import functional as F +from torch.utils.data import Dataset + +from easycom_head import NUM_CLASSES, NUM_HORIZONS, STEP_SECONDS + +OmniLayerFusion = Literal["single", "last", "mean"] + + +def load_torch_payload(path: Path, *, weights_only: bool = False) -> Any: + """Load a torch payload while retaining compatibility with older PyTorch.""" + + try: + return torch.load(path, map_location="cpu", weights_only=weights_only) + except TypeError: # pragma: no cover - only exercised by older PyTorch. + return torch.load(path, map_location="cpu") + + +def sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def normalize_session_name(value: str | int) -> str: + text = str(value).strip() + suffix = text.removeprefix("Session_") + if not suffix.isdigit(): + raise ValueError(f"Invalid EasyCom session: {value!r}") + return f"Session_{int(suffix)}" + + +class MiniROADABase(nn.Module): + """Clean-room parameter-compatible implementation of EgoSpeak MiniROADA.""" + + def __init__( + self, + *, + rgb_dim: int = 2048, + audio_dim: int = 5120, + embedding_dim: int = 2048, + hidden_dim: int = 1024, + num_layers: int = 1, + num_horizons: int = NUM_HORIZONS, + num_classes: int = NUM_CLASSES, + dropout: float = 0.2, + ) -> None: + super().__init__() + self.rgb_dim = int(rgb_dim) + self.audio_dim = int(audio_dim) + self.embedding_dim = int(embedding_dim) + self.hidden_dim = int(hidden_dim) + self.num_layers = int(num_layers) + self.num_horizons = int(num_horizons) + self.num_classes = int(num_classes) + + self.layer1 = nn.Sequential( + nn.Linear(self.rgb_dim + self.audio_dim, self.embedding_dim), + nn.LayerNorm(self.embedding_dim), + nn.ReLU(), + nn.Dropout(p=dropout), + ) + self.relu = nn.ReLU() + self.gru = nn.GRU( + self.embedding_dim, + self.hidden_dim, + self.num_layers, + batch_first=True, + ) + self.f_classification = nn.Sequential( + nn.Linear(self.hidden_dim, self.num_classes) + ) + self.anticipation_layer = nn.Sequential( + nn.Linear( + self.hidden_dim, + self.num_horizons * self.hidden_dim, + ) + ) + + def encode( + self, + rgb: torch.Tensor, + audio: torch.Tensor, + recurrent_state: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if rgb.ndim != 3 or rgb.shape[-1] != self.rgb_dim: + raise ValueError( + f"Expected RGB [B,S,{self.rgb_dim}], got {tuple(rgb.shape)}" + ) + if ( + audio.ndim != 3 + or audio.shape[:-1] != rgb.shape[:-1] + or audio.shape[-1] != self.audio_dim + ): + raise ValueError( + f"Expected audio [B,S,{self.audio_dim}] aligned with RGB, got " + f"{tuple(audio.shape)}" + ) + embedded = self.layer1(torch.cat((rgb, audio), dim=-1)) + hidden, next_state = self.gru(embedded, recurrent_state) + return self.relu(hidden), next_state + + def anticipation_logits(self, hidden: torch.Tensor) -> torch.Tensor: + shape = hidden.shape[:-1] + anticipated = self.anticipation_layer(self.relu(hidden)).view( + *shape, + self.num_horizons, + self.hidden_dim, + ) + return self.f_classification(self.relu(anticipated)) + + def current_logits(self, hidden: torch.Tensor) -> torch.Tensor: + return self.f_classification(self.relu(hidden)) + + def forward( + self, rgb: torch.Tensor, audio: torch.Tensor + ) -> dict[str, torch.Tensor]: + hidden, _ = self.encode(rgb, audio) + return { + "logits": self.current_logits(hidden), + "anticipation_logits": self.anticipation_logits(hidden), + } + + +def _unwrap_state_dict(payload: Any) -> dict[str, torch.Tensor]: + if not isinstance(payload, dict): + raise TypeError("EgoSpeak checkpoint must contain a state dict") + for key in ("state_dict", "model_state_dict"): + nested = payload.get(key) + if isinstance(nested, dict): + payload = nested + break + if not payload or not all( + isinstance(value, torch.Tensor) for value in payload.values() + ): + raise ValueError("EgoSpeak checkpoint does not contain tensor parameters") + state = dict(payload) + if all(key.startswith("module.") for key in state): + state = {key.removeprefix("module."): value for key, value in state.items()} + return state + + +def load_frozen_miniroada(checkpoint: Path, **model_kwargs: Any) -> MiniROADABase: + checkpoint = checkpoint.expanduser().resolve() + if not checkpoint.is_file(): + raise FileNotFoundError(checkpoint) + base = MiniROADABase(**model_kwargs) + state = _unwrap_state_dict(load_torch_payload(checkpoint, weights_only=True)) + base.load_state_dict(state, strict=True) + base.requires_grad_(False) + base.eval() + return base + + +class EgoSpeakOmniResidual(nn.Module): + """Feature fusion whose zero residual exactly recovers frozen MiniROADA.""" + + def __init__( + self, + base: MiniROADABase, + *, + omni_dim: int, + fusion_dim: int = 256, + tail_embedding_dim: int = 32, + num_tail_bins: int = NUM_HORIZONS, + dropout: float = 0.2, + ) -> None: + super().__init__() + if fusion_dim <= 0 or tail_embedding_dim <= 0 or num_tail_bins <= 0: + raise ValueError("Fusion and tail dimensions must be positive") + self.base = base + self.omni_dim = int(omni_dim) + self.fusion_dim = int(fusion_dim) + self.tail_embedding_dim = int(tail_embedding_dim) + self.num_tail_bins = int(num_tail_bins) + + self.ego_projection = nn.Sequential( + nn.LayerNorm(base.hidden_dim), + nn.Linear(base.hidden_dim, fusion_dim), + nn.GELU(), + ) + self.omni_projection = nn.Sequential( + nn.LayerNorm(omni_dim), + nn.Linear(omni_dim, fusion_dim), + nn.GELU(), + ) + self.tail_embedding = nn.Embedding(num_tail_bins, tail_embedding_dim) + self.fusion = nn.Sequential( + nn.LayerNorm(3 * fusion_dim + tail_embedding_dim), + nn.Linear(3 * fusion_dim + tail_embedding_dim, fusion_dim), + nn.GELU(), + nn.Dropout(dropout), + ) + self.residual_output = nn.Linear( + fusion_dim, + base.num_horizons * base.num_classes, + ) + nn.init.zeros_(self.residual_output.weight) + nn.init.zeros_(self.residual_output.bias) + + self.base.requires_grad_(False) + self.base.eval() + + def train(self, mode: bool = True) -> EgoSpeakOmniResidual: + super().train(mode) + # The released base remains deterministic: in particular, its dropout + # must never reactivate during stage-one residual training. + self.base.eval() + return self + + def trainable_parameters(self) -> list[nn.Parameter]: + return [parameter for parameter in self.parameters() if parameter.requires_grad] + + def _residual_logits( + self, + ego_hidden: torch.Tensor, + omni: torch.Tensor, + tail_bin: torch.Tensor, + ) -> torch.Tensor: + if omni.shape != (*ego_hidden.shape[:-1], self.omni_dim): + raise ValueError( + f"Omni/hidden shape mismatch: {tuple(omni.shape)} versus " + f"{tuple(ego_hidden.shape)}" + ) + if tail_bin.shape != ego_hidden.shape[:-1]: + raise ValueError( + f"Tail-bin shape mismatch: {tuple(tail_bin.shape)} versus " + f"{tuple(ego_hidden.shape[:-1])}" + ) + if tail_bin.dtype != torch.long: + raise ValueError("tail_bin must be int64") + if bool(((tail_bin < 0) | (tail_bin >= self.num_tail_bins)).any()): + raise ValueError("tail_bin lies outside the configured block phase") + + ego = self.ego_projection(ego_hidden) + global_state = self.omni_projection(omni) + fused = torch.cat( + ( + ego, + global_state, + ego * global_state, + self.tail_embedding(tail_bin), + ), + dim=-1, + ) + residual = self.residual_output(self.fusion(fused)) + return residual.view( + *ego_hidden.shape[:-1], + self.base.num_horizons, + self.base.num_classes, + ) + + def forward_last( + self, + rgb: torch.Tensor, + audio: torch.Tensor, + omni: torch.Tensor, + tail_bin: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + with torch.no_grad(): + hidden, _ = self.base.encode(rgb, audio) + ego_hidden = hidden[:, -1, :] + base_logits = self.base.anticipation_logits(ego_hidden) + residual = self._residual_logits(ego_hidden.detach(), omni, tail_bin) + return base_logits + residual, base_logits + + def forward_sequence_chunk( + self, + rgb: torch.Tensor, + audio: torch.Tensor, + omni: torch.Tensor, + tail_bin: torch.Tensor, + recurrent_state: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + with torch.no_grad(): + hidden, next_state = self.base.encode(rgb, audio, recurrent_state) + base_logits = self.base.anticipation_logits(hidden) + residual = self._residual_logits(hidden.detach(), omni, tail_bin) + return base_logits + residual, base_logits, next_state + + +def collapse_omni_states( + states: torch.Tensor, + mode: OmniLayerFusion, +) -> torch.Tensor: + if not isinstance(states, torch.Tensor) or states.ndim not in (2, 3): + raise ValueError("Omni states must be [T,D] or [T,L,D]") + if not torch.is_floating_point(states) or not torch.isfinite(states).all(): + raise ValueError("Omni states must be finite floating-point values") + if states.ndim == 2: + if mode not in ("single", "last", "mean"): + raise ValueError(f"Unsupported Omni layer fusion: {mode}") + return states.float().contiguous() + if mode == "single": + raise ValueError("Stacked Omni states require --omni-layer-fusion last or mean") + if mode == "last": + return states[:, -1, :].float().contiguous() + if mode == "mean": + return states.float().mean(dim=1).contiguous() + raise ValueError(f"Unsupported Omni layer fusion: {mode}") + + +@dataclass +class EasyComSessionArrays: + name: str + rgb: np.ndarray + audio: np.ndarray + targets: np.ndarray + omni: torch.Tensor + tail_bins: torch.LongTensor + omni_metadata: dict[str, Any] + + @property + def num_steps(self) -> int: + return len(self.targets) + + +def _require_feature_array(path: Path, expected_dim: int) -> np.ndarray: + if not path.is_file(): + raise FileNotFoundError(path) + array = np.load(path, mmap_mode="r", allow_pickle=False) + if array.ndim != 2 or array.shape[1] != expected_dim: + raise ValueError(f"Expected [T,{expected_dim}], got {array.shape}: {path}") + if not np.issubdtype(array.dtype, np.floating): + raise ValueError(f"Expected floating-point features: {path}") + return array + + +def _require_target_array(path: Path, expected_dim: int) -> np.ndarray: + if not path.is_file(): + raise FileNotFoundError(path) + array = np.load(path, mmap_mode="r", allow_pickle=False) + if array.ndim != 2 or array.shape[1] != expected_dim: + raise ValueError(f"Expected [T,{expected_dim}], got {array.shape}: {path}") + if not ( + np.issubdtype(array.dtype, np.number) or np.issubdtype(array.dtype, np.bool_) + ): + raise ValueError(f"Expected numeric activity targets: {path}") + return array + + +def load_easycom_session_arrays( + egospeak_feature_root: Path, + omni_feature_root: Path, + session: str | int, + *, + expected_omni_mode: str | None = "open-stream-replay", + omni_layer_fusion: OmniLayerFusion = "single", + rgb_dim: int = 2048, + audio_dim: int = 5120, + num_classes: int = NUM_CLASSES, +) -> EasyComSessionArrays: + name = normalize_session_name(session) + ego_root = egospeak_feature_root.expanduser().resolve() + omni_root = omni_feature_root.expanduser().resolve() + rgb = _require_feature_array( + ego_root / "rgb_kinetics_resnet50" / f"{name}.npy", + rgb_dim, + ) + audio = _require_feature_array( + ego_root / "audio_wav2vec2" / f"{name}.npy", + audio_dim, + ) + targets = _require_target_array( + ego_root / "target_perframe" / f"{name}.npy", + num_classes, + ) + lengths = {len(rgb), len(audio), len(targets)} + if len(lengths) != 1: + raise ValueError( + f"{name}: unequal EgoSpeak RGB/audio/target lengths " + f"{len(rgb)}/{len(audio)}/{len(targets)}" + ) + unique_targets = np.unique(np.asarray(targets)) + if not np.isin(unique_targets, (0, 1)).all(): + raise ValueError( + f"{name}: target_perframe contains values outside binary activity: " + f"{unique_targets.tolist()}" + ) + + omni_path = omni_root / f"{name}.pt" + if not omni_path.is_file(): + raise FileNotFoundError(omni_path) + payload = load_torch_payload(omni_path) + if not isinstance(payload, dict): + raise TypeError(f"Invalid Omni payload: {omni_path}") + mode = payload.get("mode") + if expected_omni_mode is not None and mode != expected_omni_mode: + raise ValueError( + f"{name}: Omni mode is {mode!r}, expected {expected_omni_mode!r}" + ) + omni = collapse_omni_states(payload.get("states"), omni_layer_fusion) + if len(omni) != len(targets): + raise ValueError( + f"{name}: Omni/target length mismatch {len(omni)} vs {len(targets)}" + ) + if payload.get("label_indices") is None: + raise ValueError(f"{name}: Omni payload lacks label_indices") + label_indices = torch.as_tensor(payload["label_indices"], dtype=torch.long) + expected_indices = torch.arange(len(targets), dtype=torch.long) + if label_indices.ndim != 1 or not torch.equal(label_indices, expected_indices): + raise ValueError(f"{name}: Omni label_indices are not dense zero-origin rows") + + config = payload.get("config") + if not isinstance(config, dict): + raise TypeError(f"{name}: Omni payload lacks extraction config") + step_seconds = float(config.get("step_seconds", float("nan"))) + native_block_seconds = float(config.get("native_block_seconds", float("nan"))) + if not np.isclose(step_seconds, STEP_SECONDS, rtol=0.0, atol=1e-9): + raise ValueError( + f"{name}: Omni step_seconds={step_seconds}, expected {STEP_SECONDS}" + ) + block_steps_float = native_block_seconds / step_seconds + block_steps = round(block_steps_float) + if block_steps <= 0 or not np.isclose( + block_steps_float, block_steps, rtol=0.0, atol=1e-9 + ): + raise ValueError( + f"{name}: native block {native_block_seconds}s is not an integer " + f"multiple of {step_seconds}s" + ) + if block_steps != NUM_HORIZONS: + raise ValueError( + f"{name}: expected ten 0.2-s tail bins in a 2-s block, got {block_steps}" + ) + + times = torch.as_tensor(payload.get("times"), dtype=torch.float64) + start = float(payload.get("start", 0.0)) + expected_times = ( + start + (torch.arange(len(targets), dtype=torch.float64) + 1) * step_seconds + ) + if times.ndim != 1 or not torch.allclose( + times, expected_times, rtol=0.0, atol=1e-6 + ): + raise ValueError(f"{name}: Omni times do not match the dense 5 Hz grid") + if bool(payload.get("partial", False)): + raise ValueError(f"{name}: partial Omni smoke feature cannot train the model") + + tail_bins = torch.remainder(expected_indices, block_steps) + metadata = { + "path": str(omni_path), + "mode": mode, + "model_id": payload.get("model_id"), + "config": config, + "video_sampling_grid": payload.get("video_sampling_grid"), + "media_interval_policy": payload.get("media_interval_policy"), + "open_stream_policy": payload.get("open_stream_policy"), + "state_shape_before_layer_fusion": list(payload["states"].shape), + "state_shape": list(omni.shape), + } + return EasyComSessionArrays( + name=name, + rgb=rgb, + audio=audio, + targets=targets, + omni=omni, + tail_bins=tail_bins, + omni_metadata=metadata, + ) + + +class EgoSpeakOmniWindowDataset(Dataset[tuple[torch.Tensor, ...]]): + """MiniROADA windows plus the aligned current Omni state.""" + + def __init__( + self, + sessions: Sequence[EasyComSessionArrays], + *, + window_size: int = 128, + stride: int = 4, + num_horizons: int = NUM_HORIZONS, + seed: int = 0, + max_examples: int | None = None, + ) -> None: + if not sessions: + raise ValueError("At least one training session is required") + if window_size <= 0 or stride <= 0 or num_horizons <= 0: + raise ValueError("Window, stride, and horizon counts must be positive") + if max_examples is not None and max_examples <= 0: + raise ValueError("max_examples must be positive") + self.sessions = list(sessions) + self.window_size = int(window_size) + self.stride = int(stride) + self.num_horizons = int(num_horizons) + self.seed = int(seed) + self.max_examples = max_examples + self.indices: list[tuple[int, int]] = [] + self.resample(0) + + def resample(self, epoch: int) -> None: + rng = np.random.default_rng(self.seed + int(epoch)) + indices: list[tuple[int, int]] = [] + for session_index, session in enumerate(self.sessions): + last_source = session.num_steps - self.num_horizons - 1 + if last_source < 0: + raise ValueError( + f"{session.name}: only {session.num_steps} rows for " + f"{self.num_horizons} horizons" + ) + offset = int(rng.integers(self.stride)) + indices.extend( + (session_index, source_index) + for source_index in range(offset, last_source + 1, self.stride) + ) + if self.max_examples is not None: + rng.shuffle(indices) + indices = indices[: self.max_examples] + if not indices: + raise ValueError("Training resampling produced no examples") + self.indices = indices + + def __len__(self) -> int: + return len(self.indices) + + def __getitem__(self, index: int) -> tuple[torch.Tensor, ...]: + session_index, source_index = self.indices[index] + session = self.sessions[session_index] + first = max(0, source_index - self.window_size + 1) + length = source_index - first + 1 + rgb = torch.zeros( + self.window_size, + session.rgb.shape[1], + dtype=torch.float32, + ) + audio = torch.zeros( + self.window_size, + session.audio.shape[1], + dtype=torch.float32, + ) + rgb[-length:] = torch.from_numpy( + np.asarray(session.rgb[first : source_index + 1], dtype=np.float32).copy() + ) + audio[-length:] = torch.from_numpy( + np.asarray(session.audio[first : source_index + 1], dtype=np.float32).copy() + ) + future = torch.from_numpy( + np.asarray( + session.targets[ + source_index + 1 : source_index + self.num_horizons + 1 + ], + dtype=np.float32, + ).copy() + ) + return ( + rgb, + audio, + session.omni[source_index].float(), + session.tail_bins[source_index], + future, + torch.tensor(session_index, dtype=torch.long), + torch.tensor(source_index, dtype=torch.long), + ) + + +def miniroad_normalized_loss( + logits: torch.Tensor, + targets: torch.Tensor, +) -> torch.Tensor: + """Released MiniROAD multi-positive anticipation objective.""" + + if logits.shape != targets.shape or logits.shape[-1] != NUM_CLASSES: + raise ValueError( + f"Expected equal [...,{NUM_CLASSES}] logits/targets, got " + f"{tuple(logits.shape)} and {tuple(targets.shape)}" + ) + normalized_targets = F.normalize(targets, dim=-1) + return torch.sum( + -normalized_targets * F.log_softmax(logits, dim=-1), + dim=-1, + ).mean() + + +def align_source_predictions( + raw_prediction: np.ndarray, + target_steps: int, +) -> tuple[np.ndarray, np.ndarray]: + """Align MiniROADA [source,horizon,class] forecasts to target rows.""" + + raw_prediction = np.asarray(raw_prediction, dtype=np.float32) + expected = (target_steps - NUM_HORIZONS, NUM_HORIZONS, NUM_CLASSES) + if raw_prediction.shape != expected: + raise ValueError( + f"Expected raw prediction {expected}, got {raw_prediction.shape}" + ) + if not np.isfinite(raw_prediction).all(): + raise ValueError("Raw prediction contains NaN/Inf") + predictions = np.zeros( + (target_steps, NUM_CLASSES, NUM_HORIZONS), + dtype=np.float32, + ) + valid = np.zeros((target_steps, NUM_HORIZONS), dtype=bool) + source_rows = np.arange(len(raw_prediction), dtype=np.int64) + for horizon in range(NUM_HORIZONS): + target_rows = source_rows + horizon + 1 + predictions[target_rows, :, horizon] = raw_prediction[:, horizon, :] + valid[target_rows, horizon] = True + return predictions, valid diff --git a/test_egospeak_omni_residual.py b/test_egospeak_omni_residual.py new file mode 100644 index 0000000..3d81da1 --- /dev/null +++ b/test_egospeak_omni_residual.py @@ -0,0 +1,169 @@ +import tempfile +import unittest +from pathlib import Path + +import numpy as np +import torch + +from egospeak_omni_residual import ( + EgoSpeakOmniResidual, + EgoSpeakOmniWindowDataset, + MiniROADABase, + align_source_predictions, + load_easycom_session_arrays, + load_frozen_miniroada, + miniroad_normalized_loss, +) + + +class EgoSpeakOmniResidualTest(unittest.TestCase): + def test_raw_miniroada_state_dict_loads_strictly_and_freezes(self) -> None: + with tempfile.TemporaryDirectory() as directory: + checkpoint = Path(directory) / "base.pth" + source = MiniROADABase( + rgb_dim=4, + audio_dim=6, + embedding_dim=8, + hidden_dim=5, + num_horizons=10, + ) + torch.save(source.state_dict(), checkpoint) + loaded = load_frozen_miniroada( + checkpoint, + rgb_dim=4, + audio_dim=6, + embedding_dim=8, + hidden_dim=5, + num_horizons=10, + ) + self.assertFalse(loaded.training) + self.assertTrue( + all(not parameter.requires_grad for parameter in loaded.parameters()) + ) + for key, value in source.state_dict().items(): + torch.testing.assert_close(loaded.state_dict()[key], value) + + def test_zero_residual_exactly_recovers_frozen_base(self) -> None: + torch.manual_seed(7) + base = MiniROADABase( + rgb_dim=4, + audio_dim=6, + embedding_dim=8, + hidden_dim=5, + num_horizons=10, + dropout=0.4, + ) + model = EgoSpeakOmniResidual( + base, + omni_dim=7, + fusion_dim=4, + tail_embedding_dim=3, + dropout=0.3, + ) + model.train() + self.assertFalse(model.base.training) + self.assertTrue( + all(not parameter.requires_grad for parameter in model.base.parameters()) + ) + + rgb = torch.randn(2, 9, 4) + audio = torch.randn(2, 9, 6) + omni = torch.randn(2, 7) + tail = torch.tensor([0, 9], dtype=torch.long) + fused, frozen = model.forward_last(rgb, audio, omni, tail) + torch.testing.assert_close(fused, frozen, rtol=0.0, atol=0.0) + + targets = torch.zeros_like(fused) + targets[..., 1] = 1 + optimizer = torch.optim.SGD(model.trainable_parameters(), lr=0.1) + loss = miniroad_normalized_loss(fused, targets) + loss.backward() + self.assertGreater(float(model.residual_output.weight.grad.abs().sum()), 0.0) + optimizer.step() + updated, frozen_again = model.forward_last(rgb, audio, omni, tail) + self.assertGreater( + float((updated - frozen_again).abs().max().detach()), + 0.0, + ) + + def test_feature_contract_and_window_alignment(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + ego = root / "ego" + omni = root / "omni" + for subdir in ( + "rgb_kinetics_resnet50", + "audio_wav2vec2", + "target_perframe", + ): + (ego / subdir).mkdir(parents=True, exist_ok=True) + omni.mkdir() + steps = 15 + rgb = np.arange(steps * 4, dtype=np.float32).reshape(steps, 4) + audio = np.arange(steps * 6, dtype=np.float32).reshape(steps, 6) + target = np.zeros((steps, 3), dtype=np.float32) + target[:, 0] = 1 + np.save(ego / "rgb_kinetics_resnet50" / "Session_4.npy", rgb) + np.save(ego / "audio_wav2vec2" / "Session_4.npy", audio) + np.save(ego / "target_perframe" / "Session_4.npy", target) + torch.save( + { + "mode": "open-stream-replay", + "states": torch.arange(steps * 7, dtype=torch.float32).view( + steps, 7 + ), + "times": (torch.arange(steps, dtype=torch.float64) + 1) / 5, + "label_indices": torch.arange(steps, dtype=torch.long), + "start": 0.0, + "config": { + "step_seconds": 0.2, + "native_block_seconds": 2.0, + }, + }, + omni / "Session_4.pt", + ) + session = load_easycom_session_arrays( + ego, + omni, + "Session_4", + rgb_dim=4, + audio_dim=6, + ) + self.assertEqual( + session.tail_bins.tolist(), list(range(10)) + list(range(5)) + ) + dataset = EgoSpeakOmniWindowDataset( + [session], + window_size=4, + stride=1, + num_horizons=3, + seed=0, + ) + first_index = dataset.indices.index((0, 0)) + rgb_window, audio_window, state, tail, future, _, source = dataset[ + first_index + ] + self.assertEqual(int(source), 0) + self.assertEqual(int(tail), 0) + torch.testing.assert_close(rgb_window[-1], torch.from_numpy(rgb[0])) + self.assertEqual(int(torch.count_nonzero(rgb_window[:-1])), 0) + torch.testing.assert_close(audio_window[-1], torch.from_numpy(audio[0])) + torch.testing.assert_close(state, session.omni[0]) + torch.testing.assert_close(future, torch.from_numpy(target[1:4])) + + def test_source_prediction_alignment(self) -> None: + target_steps = 14 + raw = np.arange( + (target_steps - 10) * 10 * 3, + dtype=np.float32, + ).reshape(target_steps - 10, 10, 3) + pred, valid = align_source_predictions(raw, target_steps) + np.testing.assert_array_equal(pred[1, :, 0], raw[0, 0, :]) + np.testing.assert_array_equal(pred[13, :, 9], raw[3, 9, :]) + self.assertTrue(valid[1, 0]) + self.assertFalse(valid[0, 0]) + self.assertTrue(valid[13, 9]) + + +if __name__ == "__main__": + unittest.main() diff --git a/train_egospeak_omni_residual.py b/train_egospeak_omni_residual.py new file mode 100755 index 0000000..d175703 --- /dev/null +++ b/train_egospeak_omni_residual.py @@ -0,0 +1,631 @@ +#!/usr/bin/env python3 +"""Train stage-one frozen-EgoSpeak plus frozen-Omni residual fusion. + +The released MiniROADA A+V checkpoint is loaded strictly and remains in eval +mode with gradients disabled. Existing causal 5 Hz Omni states are fused with +the MiniROADA GRU state, and only a zero-initialized residual anticipation head +is optimized. Evaluation happens once after the requested fixed epoch budget; +the script saves both fused and exactly matched frozen-base predictions. +""" + +from __future__ import annotations + +import argparse +import json +import math +import random +from collections.abc import Sequence +from contextlib import nullcontext +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from torch.utils.data import DataLoader +from tqdm.auto import tqdm + +from easycom_head import ( + CLASS_NAMES, + NUM_HORIZONS, + OFFICIAL_TEST_SESSIONS, + OFFICIAL_TRAIN_SESSIONS, + STEP_SECONDS, + average_precision_metrics, +) +from egospeak_omni_residual import ( + EasyComSessionArrays, + EgoSpeakOmniResidual, + EgoSpeakOmniWindowDataset, + align_source_predictions, + load_easycom_session_arrays, + load_frozen_miniroada, + miniroad_normalized_loss, + normalize_session_name, + sha256_file, +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--egospeak-feature-root", + type=Path, + required=True, + help=( + "EasyCom root containing rgb_kinetics_resnet50/, " + "audio_wav2vec2/, and target_perframe/." + ), + ) + parser.add_argument( + "--omni-feature-root", + type=Path, + required=True, + help="Existing directory containing causal Session_X.pt Omni states.", + ) + parser.add_argument( + "--base-checkpoint", + type=Path, + required=True, + help="Released EgoSpeak miniroad_easycom_AV.pth state dict.", + ) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--train-sessions", nargs="+", default=None) + parser.add_argument("--eval-sessions", nargs="+", default=None) + parser.add_argument( + "--expected-omni-mode", + default="open-stream-replay", + help="Expected saved extraction mode; pass 'any' only for an audit run.", + ) + parser.add_argument( + "--omni-layer-fusion", + choices=("single", "last", "mean"), + default="mean", + help="How to collapse saved [T,L,D] Thinker states.", + ) + parser.add_argument("--epochs", type=int, default=4) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument( + "--eval-chunk-size", + type=int, + default=0, + help=( + "GRU inference chunk size; 0 processes each session in one pass " + "to match released MiniROADA evaluation." + ), + ) + parser.add_argument("--num-workers", type=int, default=4) + parser.add_argument("--window-size", type=int, default=128) + parser.add_argument("--stride", type=int, default=4) + parser.add_argument("--fusion-dim", type=int, default=256) + parser.add_argument("--tail-embedding-dim", type=int, default=32) + parser.add_argument("--dropout", type=float, default=0.2) + parser.add_argument("--lr", type=float, default=1e-4) + parser.add_argument("--weight-decay", type=float, default=0.05) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--device", default="cuda") + parser.add_argument("--amp", action="store_true") + parser.add_argument( + "--max-train-examples", + type=int, + help="Smoke-test cap; omit for a complete run.", + ) + parser.add_argument( + "--reference-prediction-root", + type=Path, + help=( + "Optional released MiniROADA raw eval directory. The frozen base " + "is compared against Session_X.npy files before results are accepted." + ), + ) + parser.add_argument("--reference-tolerance", type=float, default=1e-5) + parser.add_argument("--no-progress", action="store_true") + return parser.parse_args() + + +def set_seed(seed: int) -> None: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.benchmark = False + torch.backends.cudnn.deterministic = True + + +def resolve_device(value: str) -> torch.device: + device = torch.device(value) + if device.type == "cuda" and not torch.cuda.is_available(): + raise RuntimeError("CUDA was requested but is unavailable") + return device + + +def normalize_sessions( + values: Sequence[str] | None, + default: Sequence[str], +) -> list[str]: + sessions = [normalize_session_name(value) for value in (values or default)] + if len(sessions) != len(set(sessions)): + raise ValueError("Duplicate sessions were supplied") + return sessions + + +def load_sessions( + args: argparse.Namespace, + sessions: Sequence[str], +) -> list[EasyComSessionArrays]: + expected_mode = ( + None if args.expected_omni_mode == "any" else args.expected_omni_mode + ) + loaded = [] + for session in sessions: + item = load_easycom_session_arrays( + args.egospeak_feature_root, + args.omni_feature_root, + session, + expected_omni_mode=expected_mode, + omni_layer_fusion=args.omni_layer_fusion, + ) + loaded.append(item) + print( + f"{item.name}: steps={item.num_steps}, " + f"omni={tuple(item.omni.shape)}, " + f"mode={item.omni_metadata['mode']}" + ) + dims = {int(item.omni.shape[-1]) for item in loaded} + if len(dims) != 1: + raise ValueError(f"Omni dimensions differ across sessions: {sorted(dims)}") + signatures = { + ( + item.omni_metadata["model_id"], + item.omni_metadata["mode"], + item.omni_metadata["config"].get("memory_seconds"), + item.omni_metadata["config"].get("video_fps"), + item.omni_metadata["config"].get("readout_last_n_layers"), + item.omni_metadata["config"].get("readout_layer_fusion"), + tuple(item.omni_metadata["state_shape_before_layer_fusion"][1:]), + ) + for item in loaded + } + if len(signatures) != 1: + raise ValueError( + "Selected sessions mix incompatible Omni extraction contracts: " + f"{sorted(map(str, signatures))}" + ) + return loaded + + +def omni_contract_signature(item: EasyComSessionArrays) -> tuple[Any, ...]: + config = item.omni_metadata["config"] + return ( + item.omni_metadata["model_id"], + item.omni_metadata["mode"], + config.get("memory_seconds"), + config.get("video_fps"), + config.get("readout_last_n_layers"), + config.get("readout_layer_fusion"), + tuple(item.omni_metadata["state_shape_before_layer_fusion"][1:]), + ) + + +def autocast_context(device: torch.device, enabled: bool): + if not enabled: + return nullcontext() + if device.type != "cuda": + raise ValueError("--amp currently requires CUDA") + return torch.autocast(device_type="cuda", dtype=torch.float16) + + +def make_loader( + dataset: EgoSpeakOmniWindowDataset, + args: argparse.Namespace, + epoch: int, +) -> DataLoader: + generator = torch.Generator() + generator.manual_seed(args.seed + epoch) + return DataLoader( + dataset, + batch_size=args.batch_size, + shuffle=True, + num_workers=args.num_workers, + pin_memory=torch.cuda.is_available(), + persistent_workers=False, + generator=generator, + ) + + +@torch.inference_mode() +def verify_zero_initialization( + model: EgoSpeakOmniResidual, + dataset: EgoSpeakOmniWindowDataset, + device: torch.device, +) -> float: + model.eval() + rgb, audio, omni, tail_bin, _, _, _ = dataset[0] + fused, base = model.forward_last( + rgb[None].to(device), + audio[None].to(device), + omni[None].to(device), + tail_bin[None].to(device), + ) + maximum = float((fused - base).abs().max().cpu()) + if maximum != 0.0: + raise RuntimeError( + "Zero-initialized residual does not exactly recover MiniROADA: " + f"max_abs_difference={maximum}" + ) + print("Initialization audit: fused logits == frozen EgoSpeak logits exactly") + return maximum + + +def train_one_epoch( + model: EgoSpeakOmniResidual, + loader: DataLoader, + optimizer: torch.optim.Optimizer, + scaler: torch.cuda.amp.GradScaler, + device: torch.device, + *, + amp: bool, + show_progress: bool, +) -> dict[str, float]: + model.train() + total_loss = 0.0 + total_examples = 0 + iterator = tqdm(loader, desc="train", disable=not show_progress, dynamic_ncols=True) + for rgb, audio, omni, tail_bin, targets, _, _ in iterator: + rgb = rgb.to(device=device, dtype=torch.float32, non_blocking=True) + audio = audio.to(device=device, dtype=torch.float32, non_blocking=True) + omni = omni.to(device=device, dtype=torch.float32, non_blocking=True) + tail_bin = tail_bin.to(device=device, dtype=torch.long, non_blocking=True) + targets = targets.to(device=device, dtype=torch.float32, non_blocking=True) + + optimizer.zero_grad(set_to_none=True) + with autocast_context(device, amp): + logits, _ = model.forward_last(rgb, audio, omni, tail_bin) + loss = miniroad_normalized_loss(logits.float(), targets.float()) + scaler.scale(loss).backward() + scaler.step(optimizer) + scaler.update() + + batch_size = len(rgb) + total_examples += batch_size + total_loss += float(loss.detach().cpu()) * batch_size + iterator.set_postfix_str(f"loss={total_loss / total_examples:.5f}") + return { + "loss": total_loss / total_examples, + "num_examples": total_examples, + } + + +@torch.inference_mode() +def predict_session( + model: EgoSpeakOmniResidual, + session: EasyComSessionArrays, + device: torch.device, + *, + chunk_size: int, + amp: bool, + show_progress: bool, +) -> tuple[np.ndarray, np.ndarray]: + model.eval() + num_sources = session.num_steps - NUM_HORIZONS + if num_sources <= 0: + raise ValueError(f"{session.name}: no valid ten-horizon source rows") + if chunk_size == 0: + chunk_size = num_sources + fused_chunks: list[torch.Tensor] = [] + base_chunks: list[torch.Tensor] = [] + recurrent_state = None + starts = range(0, num_sources, chunk_size) + for start in tqdm( + starts, + total=math.ceil(num_sources / chunk_size), + desc=f"infer {session.name}", + disable=not show_progress, + dynamic_ncols=True, + ): + end = min(start + chunk_size, num_sources) + rgb = torch.from_numpy( + np.asarray(session.rgb[start:end], dtype=np.float32).copy() + )[None].to(device) + audio = torch.from_numpy( + np.asarray(session.audio[start:end], dtype=np.float32).copy() + )[None].to(device) + omni = session.omni[start:end][None].to(device=device, dtype=torch.float32) + tail_bin = session.tail_bins[start:end][None].to(device) + with autocast_context(device, amp): + fused, base, recurrent_state = model.forward_sequence_chunk( + rgb, + audio, + omni, + tail_bin, + recurrent_state, + ) + recurrent_state = recurrent_state.detach() + fused_chunks.append(torch.softmax(fused.float(), dim=-1)[0].cpu()) + base_chunks.append(torch.softmax(base.float(), dim=-1)[0].cpu()) + return torch.cat(fused_chunks).numpy(), torch.cat(base_chunks).numpy() + + +def evaluate_predictions( + predictions: dict[str, np.ndarray], + valid: dict[str, np.ndarray], + targets: dict[str, np.ndarray], +) -> dict[str, Any]: + full_results = [] + causal_results = [] + for horizon in range(NUM_HORIZONS): + full_scores = np.concatenate( + [item[:, :, horizon] for item in predictions.values()], axis=0 + ) + full_targets = np.concatenate(list(targets.values()), axis=0) + full = average_precision_metrics(full_scores, full_targets) + full.update(horizon=horizon + 1, seconds=(horizon + 1) * STEP_SECONDS) + full_results.append(full) + + causal_scores = np.concatenate( + [ + predictions[name][valid[name][:, horizon], :, horizon] + for name in predictions + ], + axis=0, + ) + causal_targets = np.concatenate( + [targets[name][valid[name][:, horizon]] for name in predictions], + axis=0, + ) + causal = average_precision_metrics(causal_scores, causal_targets) + causal.update(horizon=horizon + 1, seconds=(horizon + 1) * STEP_SECONDS) + causal_results.append(causal) + return { + "egospeak_full_timeline": { + "horizons": full_results, + "mean_map": float(np.mean([item["map"] for item in full_results])), + }, + "causal_valid_only": { + "horizons": causal_results, + "mean_map": float(np.mean([item["map"] for item in causal_results])), + }, + } + + +def print_metrics(label: str, metrics: dict[str, Any]) -> None: + print(f"\n=== {label} ===") + for protocol_name in ("egospeak_full_timeline", "causal_valid_only"): + protocol = metrics[protocol_name] + print(f"{protocol_name}:") + for item in protocol["horizons"]: + class_text = ", ".join( + f"{name}={value * 100:.3f}" for name, value in item["class_ap"].items() + ) + print(f" {item['seconds']:.1f}s mAP={item['map'] * 100:.3f}; {class_text}") + print(f" mean mAP={protocol['mean_map'] * 100:.3f}") + + +def prepare_output(path: Path) -> Path: + path = path.expanduser().resolve() + if path.exists() and any(path.iterdir()): + raise FileExistsError(f"Refusing to overwrite nonempty output: {path}") + path.mkdir(parents=True, exist_ok=True) + for name in ( + "raw_eval", + "pred", + "valid", + "frozen_egospeak_raw_eval", + "frozen_egospeak_pred", + ): + (path / name).mkdir() + return path + + +def main() -> None: + args = parse_args() + if args.epochs <= 0 or args.batch_size <= 0 or args.eval_chunk_size < 0: + raise ValueError( + "Epoch and batch sizes must be positive; evaluation chunk size " + "must be nonnegative" + ) + if args.reference_tolerance < 0: + raise ValueError("--reference-tolerance must be nonnegative") + set_seed(args.seed) + device = resolve_device(args.device) + + train_names = normalize_sessions(args.train_sessions, OFFICIAL_TRAIN_SESSIONS) + eval_names = normalize_sessions(args.eval_sessions, OFFICIAL_TEST_SESSIONS) + train_sessions = load_sessions(args, train_names) + eval_sessions = load_sessions(args, eval_names) + omni_dims = {int(item.omni.shape[-1]) for item in (*train_sessions, *eval_sessions)} + if len(omni_dims) != 1: + raise ValueError(f"Omni dimensions differ: {sorted(omni_dims)}") + contracts = { + omni_contract_signature(item) for item in (*train_sessions, *eval_sessions) + } + if len(contracts) != 1: + raise ValueError( + "Train/evaluation sessions use incompatible Omni extraction " + f"contracts: {sorted(map(str, contracts))}" + ) + omni_dim = omni_dims.pop() + + base = load_frozen_miniroada(args.base_checkpoint) + model = EgoSpeakOmniResidual( + base, + omni_dim=omni_dim, + fusion_dim=args.fusion_dim, + tail_embedding_dim=args.tail_embedding_dim, + dropout=args.dropout, + ).to(device) + trainable = model.trainable_parameters() + print(f"device: {device}") + print(f"train sessions: {', '.join(train_names)}") + print(f"evaluation sessions: {', '.join(eval_names)}") + print( + f"parameters: total={sum(p.numel() for p in model.parameters()):,}, " + f"trainable={sum(p.numel() for p in trainable):,}, " + "MiniROADA frozen=True" + ) + + dataset = EgoSpeakOmniWindowDataset( + train_sessions, + window_size=args.window_size, + stride=args.stride, + seed=args.seed, + max_examples=args.max_train_examples, + ) + initialization_difference = verify_zero_initialization(model, dataset, device) + output_dir = prepare_output(args.output_dir) + optimizer = torch.optim.AdamW( + trainable, + lr=args.lr, + weight_decay=args.weight_decay, + ) + scaler = torch.cuda.amp.GradScaler(enabled=args.amp) + history = [] + for epoch in range(1, args.epochs + 1): + dataset.resample(epoch - 1) + loader = make_loader(dataset, args, epoch) + result = train_one_epoch( + model, + loader, + optimizer, + scaler, + device, + amp=args.amp, + show_progress=not args.no_progress, + ) + result["epoch"] = epoch + history.append(result) + print( + f"epoch {epoch}/{args.epochs}: loss={result['loss']:.6f}, " + f"examples={result['num_examples']}" + ) + + checkpoint_path = output_dir / "checkpoint_last.pt" + checkpoint_payload = { + "format": "egospeak-omni-zero-residual-v1", + "model_state_dict": model.state_dict(), + "model_config": { + "omni_dim": omni_dim, + "fusion_dim": args.fusion_dim, + "tail_embedding_dim": args.tail_embedding_dim, + "num_tail_bins": NUM_HORIZONS, + "dropout": args.dropout, + }, + "base_checkpoint": str(args.base_checkpoint.expanduser().resolve()), + "base_checkpoint_sha256": sha256_file( + args.base_checkpoint.expanduser().resolve() + ), + "train_sessions": train_names, + "eval_sessions": eval_names, + "training_config": { + "epochs": args.epochs, + "batch_size": args.batch_size, + "window_size": args.window_size, + "stride": args.stride, + "lr": args.lr, + "weight_decay": args.weight_decay, + "seed": args.seed, + "amp": args.amp, + "omni_layer_fusion": args.omni_layer_fusion, + "expected_omni_mode": args.expected_omni_mode, + "max_train_examples": args.max_train_examples, + }, + "initialization_max_abs_logit_difference": initialization_difference, + "history": history, + } + torch.save(checkpoint_payload, checkpoint_path) + print(f"Saved checkpoint: {checkpoint_path}") + + fused_predictions: dict[str, np.ndarray] = {} + base_predictions: dict[str, np.ndarray] = {} + valid: dict[str, np.ndarray] = {} + targets: dict[str, np.ndarray] = {} + reference_audit: dict[str, Any] = {} + for session in eval_sessions: + fused_raw, base_raw = predict_session( + model, + session, + device, + chunk_size=args.eval_chunk_size, + amp=args.amp, + show_progress=not args.no_progress, + ) + fused_aligned, session_valid = align_source_predictions( + fused_raw, session.num_steps + ) + base_aligned, base_valid = align_source_predictions(base_raw, session.num_steps) + if not np.array_equal(session_valid, base_valid): + raise RuntimeError("Fused/base causal-valid masks unexpectedly differ") + np.save(output_dir / "raw_eval" / f"{session.name}.npy", fused_raw) + np.save(output_dir / "pred" / f"{session.name}.npy", fused_aligned) + np.save(output_dir / "valid" / f"{session.name}.npy", session_valid) + np.save( + output_dir / "frozen_egospeak_raw_eval" / f"{session.name}.npy", + base_raw, + ) + np.save( + output_dir / "frozen_egospeak_pred" / f"{session.name}.npy", + base_aligned, + ) + fused_predictions[session.name] = fused_aligned + base_predictions[session.name] = base_aligned + valid[session.name] = session_valid + targets[session.name] = np.asarray(session.targets, dtype=np.float32) + + if args.reference_prediction_root is not None: + reference_path = ( + args.reference_prediction_root.expanduser().resolve() + / f"{session.name}.npy" + ) + if not reference_path.is_file(): + raise FileNotFoundError(reference_path) + reference = np.load(reference_path, allow_pickle=False) + if reference.shape != base_raw.shape: + raise ValueError( + f"{session.name}: reference shape {reference.shape} != " + f"{base_raw.shape}" + ) + difference = np.abs(reference.astype(np.float32) - base_raw) + maximum = float(difference.max()) + reference_audit[session.name] = { + "path": str(reference_path), + "max_abs_difference": maximum, + "mean_abs_difference": float(difference.mean()), + } + print( + f"{session.name}: frozen-base/reference max abs " + f"difference={maximum:.8g}" + ) + if maximum > args.reference_tolerance: + raise RuntimeError( + f"{session.name}: frozen MiniROADA did not reproduce the " + f"reference within {args.reference_tolerance}: {maximum}" + ) + + fused_metrics = evaluate_predictions(fused_predictions, valid, targets) + base_metrics = evaluate_predictions(base_predictions, valid, targets) + metrics = { + **fused_metrics, + "format": "egospeak-omni-zero-residual-eval-v1", + "classes": list(CLASS_NAMES), + "sessions": eval_names, + "prediction_layout": "[target_time, class, future_horizon]", + "raw_prediction_layout": "[source_time, future_horizon, class]", + "checkpoint": str(checkpoint_path), + "train_sessions": train_names, + "frozen_egospeak_baseline": base_metrics, + "reference_prediction_audit": reference_audit, + "initialization_max_abs_logit_difference": initialization_difference, + "omni_sources": { + item.name: item.omni_metadata for item in (*train_sessions, *eval_sessions) + }, + } + metrics_path = output_dir / "metrics.json" + metrics_path.write_text(json.dumps(metrics, indent=2) + "\n") + print_metrics("fused", fused_metrics) + print_metrics("frozen EgoSpeak", base_metrics) + print(f"\nSaved metrics: {metrics_path}") + print(f"Fused evaluator input: {output_dir}") + + +if __name__ == "__main__": + main()