Add selective EgoSpeak Omni stage-two finetuning

PROJECT_CONTEXT.md

@@ -154,6 +154,19 @@ result.
   states remains a future extraction ablation. See the
   [`runbook`](docs/easycom_egospeak_omni_residual.md).
 
+### Unmeasured dual-stream continuation
+
+Stage-two `egospeak-omni-selective-finetune` is implemented but unmeasured.
+It strictly continues the trained stage-one checkpoint, selectively adapts
+MiniROADA's GRU/forecast heads with a lower learning rate, and leaves original
+feature extractors, input projection and Omni frozen. It exports immutable
+stage-one/original references and automatic threshold-free raw endpoint AP.
+The user requested this exploratory continuation while semantic annotations
+are audited; P0 remains unfinished. Four additional epochs means eight total
+adaptation epochs, requiring the fusion-only continuation control for a
+budget-matched unfreezing comparison. See the
+[`stage-two runbook`](docs/easycom_egospeak_omni_finetune.md).
+
 ### Measured oracle-information screen
 
 - A seed-0, four-epoch Session-12 development screen supplied transcript-derived

docs/README.md

@@ -29,6 +29,7 @@ tests.
 | [`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_egospeak_omni_finetune.md`](easycom_egospeak_omni_finetune.md) | Stage-two selective MiniROADA GRU/head continuation with immutable stage-one/base references and raw endpoint AP |
 | [`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 |

docs/easycom_egospeak_omni_finetune.md

@@ -0,0 +1,188 @@
+# Stage two: selective MiniROADA + Omni residual fine-tuning
+
+## Decision and claim boundary
+
+The user requested this continuation after the
+[`stage-one screen`](../results/easycom_egospeak_omni_residual_2026-09-16/README.md):
+causal-valid dense mean mAP improved by +1.521 points, raw onset endpoint AP
+by +0.600 and offset endpoint AP by +2.580. These are single-seed reported
+results, not statistically significant semantic floor-control evidence.
+Provisional ENTRY/YIELD evaluation is paused pending an annotation audit;
+historical provisional results remain archived. P0 is still unfinished.
+
+This stage means selective MiniROADA temporal/forecast adaptation, **not**
+Omni LoRA, a person-memory adapter, new feature extraction or a semantic head.
+It is an exploratory implementation of the earlier selective-unfreezing route.
+Do not use Sessions 1--3 to choose learning rates, windows, loss weights,
+unfreezing policy or the stopping epoch. A comparison observed on the exposed
+test split remains exploratory; choose any later settings on training-side
+development data and report the test exposure honestly.
+
+## What changes
+
+Load `checkpoint_last.pt` from stage one strictly. Preserve its learned
+nonzero residual and verify exact initial recovery of its fused/base logits
+on a training example. The original released checkpoint is not needed again:
+the first-stage checkpoint already contains its full frozen base parameters.
+
+| Component | Stage-two default |
+|---|---|
+| Released RGB and wav2vec2 feature extractors | Frozen; reuse existing arrays |
+| MiniROADA RGB+audio input projection (`layer1`) | Frozen, dropout remains off |
+| MiniROADA GRU | Trainable, learning rate 1e-5 |
+| Anticipation projection and shared classification head | Trainable, learning rate 1e-5 |
+| Ego/Omni projections, tail embedding and fusion residual | Continue training, learning rate 1e-4 |
+| Qwen2.5-Omni backbone | Frozen; reuse exactly the first-stage sidecars |
+
+The GRU state contributes to both the original anticipation logits and the
+feature-fusion residual. Unlike stage one, neither path is enclosed in
+`no_grad` or detached when selected for training. The default trainable
+MiniROADA modules contain 19,942,403 parameters, in addition to the existing
+residual branch. The RGB/audio encoders themselves are not loaded or trained;
+this is partial task-model fine-tuning with frozen feature extractors, not
+raw-media end-to-end training.
+
+Loss, window/stride, grid, class order and ten horizons stay unchanged.
+Window size and stride inherit the stage-one checkpoint and mismatched
+overrides fail. Readout mode/layer fusion/dimension must also match; when the
+stage-one sibling `metrics.json` is available, its full saved Omni extraction
+signature is checked. A relocated checkpoint without that provenance prints
+a warning, so the same original Omni root must be reused explicitly.
+
+Default optimizer: AdamW, weight decay 0.05, gradient clipping at norm 1.0.
+These defaults are specified before the new run, not tuned on test results.
+The first training batch logs separate fusion/base gradient norms. Final
+evaluation occurs only after the fixed additional epoch budget, with no
+test-side early stopping or best-checkpoint selection.
+
+## Session 4 smoke test
+
+From the repository root, reuse the same two feature roots as stage one:
+
+```bash
+EGO_ROOT=third_party/EgoSpeak/MiniROAD/data/EasyCom
+OMNI_ROOT=/path/to/the/same/first_stage/omni_features
+STAGE1=results/easycom_egospeak_omni_residual_seed0/checkpoint_last.pt
+
+CUDA_VISIBLE_DEVICES=0 python3 -u train_egospeak_omni_finetune.py \
+  --egospeak-feature-root "$EGO_ROOT" \
+  --omni-feature-root "$OMNI_ROOT" \
+  --stage-one-checkpoint "$STAGE1" \
+  --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_finetune_smoke_2026-09-16
+```
+
+Check exact initialization, nonzero fusion/base gradients, finite loss,
+`[T,3,10]` predictions, `[T,10]` common valid masks and raw endpoint counts.
+The printed cutoff is `(source_row+1)/5` s; input stops at that row, and
+future labels are rows source+1 through source+10. Smoke predictions are not
+the final model's calibration predictions. Start the formal run from the
+original stage-one checkpoint, not the smoke checkpoint.
+
+## Fixed-budget continuation
+
+After the GPU smoke audit passes:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python3 -u train_egospeak_omni_finetune.py \
+  --egospeak-feature-root "$EGO_ROOT" \
+  --omni-feature-root "$OMNI_ROOT" \
+  --stage-one-checkpoint "$STAGE1" \
+  --base-trainable gru-heads \
+  --fusion-lr 1e-4 \
+  --base-lr 1e-5 \
+  --epochs 4 \
+  --batch-size 64 \
+  --num-workers 4 \
+  --seed 0 \
+  --output-dir results/easycom_egospeak_omni_finetune_seed0_2026-09-16
+```
+
+Defaults use Sessions 4--12 for training and Sessions 1--3 for final
+evaluation. Training or a first-stage checkpoint trained on Sessions 1--3 is
+rejected. Float32 is the first reproducibility run; do not enable AMP until
+its output/causality audit passes. Full-session GRU inference is the default;
+`--eval-chunk-size N` carries recurrent state across chunks if memory requires.
+Nonempty output directories are never overwritten.
+
+An additional four epochs means **4+4=8 total adaptation epochs**. It cannot
+isolate unfreezing against the original four-epoch first stage. The required
+matched-budget continuation control changes only these arguments:
+
+```bash
+  --base-trainable none \
+  --output-dir results/easycom_egospeak_omni_fusion_continue_seed0_2026-09-16
+```
+
+Keep the same stage-one checkpoint, seed, training split, four additional
+epochs, sampling, optimizer and fusion learning rate. `--base-trainable heads`
+is a separate predefined ablation that freezes the GRU and adapts only the
+forecast heads plus residual; do not select among policies on test scores.
+
+## Outputs and reading AP
+
+```text
+RESULT/
+  README.md
+  checkpoint_last.pt              # full adapted model, config, parent hash, budget
+  metrics.json                    # adapted dense + raw endpoint AP
+  raw_endpoint_metrics.csv        # three methods, ten horizons, AP and counts
+  pred/, valid/, raw_eval/         # second-stage predictions
+  stage_one_reference/
+    metrics.json, pred/, valid/, raw_eval/
+  frozen_egospeak_reference/
+    metrics.json, pred/, valid/, raw_eval/
+```
+
+Both references are immutable copies of the first-stage checkpoint: one
+includes its learned fusion, the other uses its original frozen MiniROADA
+path. The adapted base branch is never mislabeled as frozen EgoSpeak.
+Every method uses the same valid source rows. Reference directories are
+directly compatible with the existing evaluator. A loader for the new full
+model is `egospeak_omni_finetune.load_stage_two_checkpoint`; stage-one and
+stage-two checkpoint formats are intentionally distinct.
+
+Dense metrics retain both full-timeline and causal-valid protocols. Raw onset
+uses current wearer OFF, future ON versus future OFF with score p_wearer;
+offset uses current wearer ON, future OFF versus future ON with score
+1-p_wearer. AP pools eligible pairs over evaluated sessions separately for
+each horizon; the headline averages ten APs. This is endpoint anticipation,
+not timestamp detection or semantic ENTRY/YIELD. No Session 4 threshold
+calibration or annotation file is needed for these APs. An undefined horizon
+is preserved as null, not silently removed from the mean.
+
+The CSV/JSON store AP **fractions**, while console output uses percentages.
+To print the raw results:
+
+```bash
+python3 - <<'PY'
+import pandas as pd
+p = "results/easycom_egospeak_omni_finetune_seed0_2026-09-16/raw_endpoint_metrics.csv"
+x = pd.read_csv(p)
+cols = ["onset_endpoint_ap", "offset_endpoint_ap"]
+x[cols] *= 100
+print(x[["run", "nominal_seconds", *cols]].round(3).to_string(index=False))
+print(x.groupby("run")[cols].mean().round(3).to_string())
+PY
+```
+
+If any horizon is undefined, use the JSON headline (which stays null) instead
+of pandas' default mean that skips NaN. No semantic opportunity evaluation is
+run by this launcher.
+
+## Local validation
+
+```bash
+python3 -m unittest test_egospeak_omni_finetune test_egospeak_omni_residual -v
+```
+
+Tests cover trained residual preservation, selected-module gradients,
+reference immutability, policy masks, causal perturbation, carried-state
+chunk equivalence, checkpoint failure/success paths, raw AP and a tiny CPU
+training/export/reload smoke. They do not substitute for the real Session 4
+GPU smoke or establish second-stage performance; that run remains unmeasured.

docs/easycom_egospeak_omni_residual.md

@@ -154,6 +154,13 @@ so it remains readable by `summarize_local_baselines.py`.
 
 ## Interpretation and next gate
 
+The first seed-0 dense/raw screen is now
+[`recorded`](../results/easycom_egospeak_omni_residual_2026-09-16/README.md).
+At the user's request, provisional semantic evaluation is paused for an
+annotation audit. An exploratory selective GRU/head continuation is implemented
+in the [`stage-two runbook`](easycom_egospeak_omni_finetune.md); this does not
+imply that the semantic P0 gate below has passed.
+
 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

egospeak_omni_finetune.py

@@ -0,0 +1,195 @@
+"""Stage-two selective MiniROADA adaptation from a trained Omni residual.
+
+The released input features and Omni states remain frozen. This module only
+enables gradients in selected MiniROADA temporal/forecast modules and the
+existing residual. Stage-one checkpoints are loaded strictly, without resetting
+the learned residual or requiring the original checkpoint to be available.
+"""
+
+from __future__ import annotations
+
+from copy import deepcopy
+from pathlib import Path
+from typing import Any, Literal
+
+import torch
+
+from easycom_head import NUM_HORIZONS, OFFICIAL_TRAIN_SESSIONS
+from egospeak_omni_residual import (
+    EgoSpeakOmniResidual,
+    MiniROADABase,
+    load_torch_payload,
+    normalize_session_name,
+)
+
+BaseTrainable = Literal["none", "heads", "gru-heads"]
+STAGE_ONE_FORMAT = "egospeak-omni-zero-residual-v1"
+STAGE_TWO_FORMAT = "egospeak-omni-selective-finetune-v1"
+
+
+def validate_training_sessions(values: list[str]) -> None:
+    normalized = [normalize_session_name(value) for value in values]
+    if not normalized or len(normalized) != len(set(normalized)):
+        raise ValueError("Training sessions must be nonempty and unique")
+    if not set(normalized).issubset(OFFICIAL_TRAIN_SESSIONS):
+        raise ValueError("Training may only use EasyCom Sessions 4--12")
+
+
+def load_stage_one_checkpoint(
+    path: Path,
+) -> tuple[EgoSpeakOmniResidual, dict[str, Any]]:
+    path = path.expanduser().resolve()
+    if not path.is_file():
+        raise FileNotFoundError(path)
+    payload = load_torch_payload(path)
+    if not isinstance(payload, dict) or payload.get("format") != STAGE_ONE_FORMAT:
+        raise ValueError(
+            "--stage-one-checkpoint must be a stage-one residual checkpoint"
+        )
+    config = payload.get("model_config")
+    if not isinstance(config, dict):
+        raise TypeError("Stage-one checkpoint lacks model_config")
+    required = {
+        "omni_dim",
+        "fusion_dim",
+        "tail_embedding_dim",
+        "num_tail_bins",
+        "dropout",
+    }
+    if set(config) != required or config["num_tail_bins"] != NUM_HORIZONS:
+        raise ValueError("Incompatible stage-one residual model_config")
+    training = payload.get("training_config")
+    if not isinstance(training, dict) or int(training.get("epochs", 0)) <= 0:
+        raise ValueError("Stage-one checkpoint lacks a positive training epoch budget")
+    validate_training_sessions(payload.get("train_sessions", []))
+    base_config = payload.get("base_model_config", {})
+    if not isinstance(base_config, dict):
+        raise TypeError("Invalid base_model_config")
+    # Released v1 checkpoints omit base_model_config and use default MiniROADA.
+    model = EgoSpeakOmniResidual(MiniROADABase(**base_config), **config)
+    model.load_state_dict(payload["model_state_dict"], strict=True)
+    if any(not torch.isfinite(value).all() for value in model.state_dict().values()):
+        raise ValueError("Stage-one checkpoint contains nonfinite parameters")
+    model.requires_grad_(False)
+    model.eval()
+    return model, payload
+
+
+class EgoSpeakOmniFinetune(EgoSpeakOmniResidual):
+    """Continuation with explicit base policy and intact stage-one logits."""
+
+    def __init__(
+        self, *args: Any, base_trainable: BaseTrainable = "gru-heads", **kwargs: Any
+    ):
+        super().__init__(*args, **kwargs)
+        if base_trainable not in ("none", "heads", "gru-heads"):
+            raise ValueError(f"Invalid base-trainable policy: {base_trainable}")
+        self.base_trainable = base_trainable
+        self.requires_grad_(True)
+        self.base.requires_grad_(False)
+        if base_trainable in ("heads", "gru-heads"):
+            self.base.anticipation_layer.requires_grad_(True)
+            self.base.f_classification.requires_grad_(True)
+        if base_trainable == "gru-heads":
+            self.base.gru.requires_grad_(True)
+        self.train(False)
+
+    @classmethod
+    def from_stage_one(
+        cls,
+        reference: EgoSpeakOmniResidual,
+        config: dict[str, Any],
+        *,
+        base_trainable: BaseTrainable = "gru-heads",
+    ) -> EgoSpeakOmniFinetune:
+        model = cls(deepcopy(reference.base), **config, base_trainable=base_trainable)
+        # Copy rather than share weights with the immutable first-stage reference.
+        model.load_state_dict(reference.state_dict(), strict=True)
+        return model
+
+    def train(self, mode: bool = True) -> EgoSpeakOmniFinetune:
+        super().train(mode)
+        # Keep frozen layer1's dropout off. Enabling GRU/head train mode does
+        # not turn the original frozen RGB+audio projection back on.
+        if self.base_trainable in ("heads", "gru-heads"):
+            self.base.anticipation_layer.train(mode)
+            self.base.f_classification.train(mode)
+        if self.base_trainable == "gru-heads":
+            self.base.gru.train(mode)
+        return self
+
+    def optimizer_groups(
+        self, *, fusion_lr: float, base_lr: float
+    ) -> list[dict[str, Any]]:
+        if fusion_lr <= 0 or base_lr <= 0:
+            raise ValueError("Both learning rates must be positive")
+        fusion = [
+            p
+            for name, p in self.named_parameters()
+            if p.requires_grad and not name.startswith("base.")
+        ]
+        base = [
+            p
+            for name, p in self.named_parameters()
+            if p.requires_grad and name.startswith("base.")
+        ]
+        groups = [{"name": "fusion", "params": fusion, "lr": fusion_lr}]
+        if base:
+            groups.append({"name": "miniroada", "params": base, "lr": base_lr})
+        return groups
+
+    def forward_last(
+        self,
+        rgb: torch.Tensor,
+        audio: torch.Tensor,
+        omni: torch.Tensor,
+        tail_bin: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        hidden, _ = self.base.encode(rgb, audio)
+        ego_hidden = hidden[:, -1, :]
+        base_logits = self.base.anticipation_logits(ego_hidden)
+        # No no_grad/detach: gradients must reach the selected MiniROADA GRU.
+        residual = self._residual_logits(ego_hidden, 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]:
+        hidden, next_state = self.base.encode(rgb, audio, recurrent_state)
+        base_logits = self.base.anticipation_logits(hidden)
+        residual = self._residual_logits(hidden, omni, tail_bin)
+        return base_logits + residual, base_logits, next_state
+
+
+def base_model_config(base: MiniROADABase) -> dict[str, Any]:
+    return {
+        "rgb_dim": base.rgb_dim,
+        "audio_dim": base.audio_dim,
+        "embedding_dim": base.embedding_dim,
+        "hidden_dim": base.hidden_dim,
+        "num_layers": base.num_layers,
+        "num_horizons": base.num_horizons,
+        "num_classes": base.num_classes,
+        "dropout": base.layer1[-1].p,
+    }
+
+
+def load_stage_two_checkpoint(
+    path: Path,
+) -> tuple[EgoSpeakOmniFinetune, dict[str, Any]]:
+    payload = load_torch_payload(path)
+    if not isinstance(payload, dict) or payload.get("format") != STAGE_TWO_FORMAT:
+        raise ValueError("Not a stage-two fine-tuning checkpoint")
+    model = EgoSpeakOmniFinetune(
+        MiniROADABase(**payload["base_model_config"]),
+        **payload["model_config"],
+        base_trainable=payload["base_trainable"],
+    )
+    model.load_state_dict(payload["model_state_dict"], strict=True)
+    model.eval()
+    return model, payload

test_egospeak_omni_finetune.py

@@ -0,0 +1,374 @@
+import csv
+import io
+import json
+import tempfile
+import unittest
+from contextlib import redirect_stdout
+from pathlib import Path
+from unittest.mock import patch
+
+import numpy as np
+import torch
+
+from egospeak_omni_finetune import (
+    STAGE_ONE_FORMAT,
+    STAGE_TWO_FORMAT,
+    EgoSpeakOmniFinetune,
+    base_model_config,
+    load_stage_one_checkpoint,
+    load_stage_two_checkpoint,
+    validate_training_sessions,
+)
+from egospeak_omni_residual import (
+    EgoSpeakOmniResidual,
+    MiniROADABase,
+    load_easycom_session_arrays,
+    miniroad_normalized_loss,
+)
+from train_egospeak_omni_finetune import (
+    main,
+    parse_args,
+    prepare_output,
+    raw_endpoint_metrics,
+    validate_checkpoint_contract,
+)
+
+
+def reference_model():
+    torch.manual_seed(7)
+    base = MiniROADABase(
+        rgb_dim=4, audio_dim=6, embedding_dim=8, hidden_dim=5, dropout=0.4
+    )
+    config = {
+        "omni_dim": 7,
+        "fusion_dim": 4,
+        "tail_embedding_dim": 3,
+        "num_tail_bins": 10,
+        "dropout": 0.3,
+    }
+    model = EgoSpeakOmniResidual(base, **config)
+    # Simulate a trained, nonzero first-stage residual that must not be reset.
+    with torch.no_grad():
+        model.residual_output.weight.normal_(std=0.1)
+        model.residual_output.bias.normal_(std=0.1)
+    model.requires_grad_(False)
+    model.eval()
+    return model, config
+
+
+def stage_one_payload(reference, config):
+    return {
+        "format": STAGE_ONE_FORMAT,
+        "model_config": config,
+        "base_model_config": base_model_config(reference.base),
+        "model_state_dict": reference.state_dict(),
+        "train_sessions": ["Session_4"],
+        "training_config": {
+            "epochs": 4,
+            "window_size": 8,
+            "stride": 1,
+            "expected_omni_mode": "open-stream-replay",
+            "omni_layer_fusion": "mean",
+        },
+    }
+
+
+class EgoSpeakOmniFinetuneTest(unittest.TestCase):
+    def setUp(self):
+        # Small synthetic tests should not spin up many CPU BLAS workers.
+        self.previous_threads = torch.get_num_threads()
+        torch.set_num_threads(1)
+
+    def tearDown(self):
+        torch.set_num_threads(self.previous_threads)
+
+    def test_nonzero_stage_one_recovery_gradients_and_reference_immutability(self):
+        reference, config = reference_model()
+        model = EgoSpeakOmniFinetune.from_stage_one(reference, config)
+        inputs = (
+            torch.randn(2, 9, 4),
+            torch.randn(2, 9, 6),
+            torch.randn(2, 7),
+            torch.tensor([0, 9]),
+        )
+        expected, _ = reference.forward_last(*inputs)
+        actual, _ = model.forward_last(*inputs)
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+        self.assertGreater(float(reference.residual_output.weight.abs().sum()), 0)
+        original = {k: v.clone() for k, v in reference.state_dict().items()}
+        model.train()
+        self.assertFalse(model.base.layer1.training)
+        self.assertTrue(model.base.gru.training)
+        self.assertTrue(model.base.anticipation_layer.training)
+        logits, _ = model.forward_last(*inputs)
+        targets = torch.zeros_like(logits)
+        targets[..., 1] = 1
+        loss = miniroad_normalized_loss(logits, targets)
+        loss.backward()
+        for module in (
+            model.base.gru,
+            model.base.anticipation_layer,
+            model.base.f_classification,
+            model.residual_output,
+            model.omni_projection,
+        ):
+            norm = sum(
+                float(p.grad.abs().sum())
+                for p in module.parameters()
+                if p.grad is not None
+            )
+            self.assertGreater(norm, 0)
+        self.assertTrue(all(p.grad is None for p in model.base.layer1.parameters()))
+        self.assertTrue(all(p.grad is None for p in reference.parameters()))
+        groups = model.optimizer_groups(fusion_lr=1e-3, base_lr=1e-4)
+        ids = [id(p) for group in groups for p in group["params"]]
+        self.assertEqual(len(ids), len(set(ids)))
+        self.assertEqual(set(ids), {id(p) for p in model.trainable_parameters()})
+        before = model.base.gru.weight_ih_l0.detach().clone()
+        torch.optim.AdamW(groups).step()
+        self.assertGreater(
+            float((before - model.base.gru.weight_ih_l0).abs().max().detach()), 0
+        )
+        for key, value in reference.state_dict().items():
+            torch.testing.assert_close(value, original[key], rtol=0, atol=0)
+
+    def test_freezing_policies(self):
+        reference, config = reference_model()
+        for policy in ("none", "heads", "gru-heads"):
+            with self.subTest(policy=policy):
+                model = EgoSpeakOmniFinetune.from_stage_one(
+                    reference, config, base_trainable=policy
+                )
+                model.train()
+                for name, parameter in model.base.named_parameters():
+                    expected = (name.startswith("gru.") and policy == "gru-heads") or (
+                        name.startswith(("anticipation_layer.", "f_classification."))
+                        and policy != "none"
+                    )
+                    self.assertEqual(parameter.requires_grad, expected, name)
+                self.assertFalse(model.base.layer1[-1].training)
+                self.assertEqual(
+                    len(model.optimizer_groups(fusion_lr=1e-4, base_lr=1e-5)),
+                    1 if policy == "none" else 2,
+                )
+        with self.assertRaises(ValueError):
+            EgoSpeakOmniFinetune.from_stage_one(reference, config, base_trainable="all")
+
+    def test_continuation_inherits_and_locks_window_and_readout(self):
+        reference, config = reference_model()
+        payload = stage_one_payload(reference, config)
+        required = [
+            "--egospeak-feature-root",
+            "/tmp/ego",
+            "--omni-feature-root",
+            "/tmp/omni",
+            "--stage-one-checkpoint",
+            "/tmp/stage1.pt",
+            "--output-dir",
+            "/tmp/stage2",
+        ]
+        args = parse_args(required)
+        validate_checkpoint_contract(args, payload)
+        self.assertEqual((args.window_size, args.stride), (8, 1))
+        for options in (
+            ["--window-size", "16"],
+            ["--stride", "4"],
+            ["--omni-layer-fusion", "last"],
+            ["--expected-omni-mode", "any"],
+        ):
+            with self.subTest(options=options), self.assertRaises(ValueError):
+                validate_checkpoint_contract(parse_args(required + options), payload)
+
+    def test_sequence_is_causal_and_chunked_recurrent_inference_matches(self):
+        reference, config = reference_model()
+        model = EgoSpeakOmniFinetune.from_stage_one(reference, config)
+        rgb, audio = torch.randn(1, 15, 4), torch.randn(1, 15, 6)
+        omni, tail = torch.randn(1, 15, 7), torch.arange(15)[None] % 10
+        with torch.no_grad():
+            full, _, _ = model.forward_sequence_chunk(rgb, audio, omni, tail, None)
+            prefix, _, state = model.forward_sequence_chunk(
+                rgb[:, :6], audio[:, :6], omni[:, :6], tail[:, :6], None
+            )
+            suffix, _, _ = model.forward_sequence_chunk(
+                rgb[:, 6:], audio[:, 6:], omni[:, 6:], tail[:, 6:], state
+            )
+            torch.testing.assert_close(
+                full, torch.cat((prefix, suffix), dim=1), rtol=1e-5, atol=1e-6
+            )
+            rgb[:, 6:] += 100
+            audio[:, 6:] -= 100
+            omni[:, 6:] += 100
+            changed, _, _ = model.forward_sequence_chunk(rgb, audio, omni, tail, None)
+            torch.testing.assert_close(full[:, :6], changed[:, :6], rtol=0, atol=0)
+
+    def test_checkpoint_loaders_reject_wrong_formats_test_training_and_nonfinite(self):
+        reference, config = reference_model()
+        with tempfile.TemporaryDirectory() as directory:
+            path = Path(directory) / "checkpoint.pt"
+            payload = stage_one_payload(reference, config)
+            torch.save(payload, path)
+            loaded, _ = load_stage_one_checkpoint(path)
+            torch.testing.assert_close(
+                loaded.residual_output.weight, reference.residual_output.weight
+            )
+            for change in (
+                {"format": STAGE_TWO_FORMAT},
+                {"train_sessions": ["Session_1"]},
+                {"model_config": {**config, "num_tail_bins": 9}},
+            ):
+                torch.save({**payload, **change}, path)
+                with self.assertRaises(ValueError):
+                    load_stage_one_checkpoint(path)
+            with torch.no_grad():
+                payload["model_state_dict"]["residual_output.bias"][0] = float("nan")
+            torch.save(payload, path)
+            with self.assertRaisesRegex(ValueError, "nonfinite"):
+                load_stage_one_checkpoint(path)
+        with self.assertRaises(FileNotFoundError):
+            load_stage_one_checkpoint(Path("/does-not-exist/checkpoint.pt"))
+        for sessions in (["Session_1"], ["Session_4", "Session_4"], []):
+            with self.assertRaises(ValueError):
+                validate_training_sessions(sessions)
+
+    def test_raw_ap_uses_endpoint_labels_and_has_no_calibration_requirement(self):
+        targets = np.zeros((14, 3), dtype=np.float32)
+        targets[:, 0] = 1
+        wearer = np.array([0, 0, 1, 1] * 3 + [0, 0], dtype=bool)
+        targets[wearer, 0] = 0
+        targets[wearer, 1] = 1
+        predictions = np.zeros((14, 3, 10), dtype=np.float32)
+        predictions[:, 1, :] = targets[:, 1, None]
+        valid = np.ones((14, 10), dtype=bool)
+        metrics = raw_endpoint_metrics(
+            {"Session_1": predictions}, {"Session_1": valid}, {"Session_1": targets}
+        )
+        self.assertEqual(metrics["horizons"][0]["onset"]["average_precision"], 1)
+        self.assertEqual(metrics["horizons"][0]["offset"]["average_precision"], 1)
+        # Four-step horizons have no positives, so headline stays undefined.
+        self.assertIsNone(metrics["mean_onset_ap"])
+
+    def test_cpu_end_to_end_smoke_export_reload_and_fixed_reference_predictions(self):
+        reference, config = reference_model()
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            checkpoint = root / "stage1" / "checkpoint_last.pt"
+            checkpoint.parent.mkdir()
+            torch.save(stage_one_payload(reference, config), checkpoint)
+            ego, omni = root / "ego", root / "omni"
+            for sub in ("rgb_kinetics_resnet50", "audio_wav2vec2", "target_perframe"):
+                (ego / sub).mkdir(parents=True)
+            omni.mkdir()
+            steps = 36
+            rng = np.random.default_rng(3)
+            np.save(
+                ego / "rgb_kinetics_resnet50" / "Session_4.npy",
+                rng.standard_normal((steps, 4)).astype(np.float32),
+            )
+            np.save(
+                ego / "audio_wav2vec2" / "Session_4.npy",
+                rng.standard_normal((steps, 6)).astype(np.float32),
+            )
+            target = np.eye(3, dtype=np.float32)[np.arange(steps) % 3]
+            np.save(ego / "target_perframe" / "Session_4.npy", target)
+            torch.save(
+                {
+                    "mode": "open-stream-replay",
+                    "states": torch.randn(steps, 2, 7),
+                    "times": (torch.arange(steps, dtype=torch.float64) + 1) / 5,
+                    "label_indices": torch.arange(steps),
+                    "start": 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, omni_layer_fusion="mean"
+            )
+            output = root / "stage2"
+            arguments = [
+                "--egospeak-feature-root",
+                str(ego),
+                "--omni-feature-root",
+                str(omni),
+                "--stage-one-checkpoint",
+                str(checkpoint),
+                "--output-dir",
+                str(output),
+                "--train-sessions",
+                "Session_4",
+                "--eval-sessions",
+                "Session_4",
+                "--device",
+                "cpu",
+                "--epochs",
+                "1",
+                "--batch-size",
+                "2",
+                "--num-workers",
+                "0",
+                "--max-train-examples",
+                "4",
+                "--no-progress",
+            ]
+            with (
+                patch(
+                    "train_egospeak_omni_finetune.load_sessions", return_value=[session]
+                ),
+                redirect_stdout(io.StringIO()),
+            ):
+                main(arguments)
+            loaded, payload = load_stage_two_checkpoint(output / "checkpoint_last.pt")
+            self.assertEqual(payload["budgets"]["total_adaptation_epochs"], 5)
+            self.assertGreater(
+                payload["history"][0]["first_batch_gradient_norms"]["miniroada"], 0
+            )
+            self.assertEqual(loaded.base_trainable, "gru-heads")
+            self.assertTrue((output / "README.md").is_file())
+            for method in (
+                output,
+                output / "stage_one_reference",
+                output / "frozen_egospeak_reference",
+            ):
+                pred = np.load(method / "pred" / "Session_4.npy")
+                valid = np.load(method / "valid" / "Session_4.npy")
+                self.assertEqual(pred.shape, (steps, 3, 10))
+                self.assertEqual(valid.shape, (steps, 10))
+                self.assertEqual(valid.sum(), (steps - 10) * 10)
+                self.assertTrue(np.isfinite(pred).all())
+                np.testing.assert_array_equal(
+                    valid, np.load(output / "valid" / "Session_4.npy")
+                )
+            with torch.no_grad():
+                expected_stage1, expected_original, _ = (
+                    reference.forward_sequence_chunk(
+                        torch.from_numpy(np.asarray(session.rgb).copy())[
+                            None, : steps - 10
+                        ],
+                        torch.from_numpy(np.asarray(session.audio).copy())[
+                            None, : steps - 10
+                        ],
+                        session.omni[None, : steps - 10],
+                        session.tail_bins[None, : steps - 10],
+                        None,
+                    )
+                )
+            for name, expected in (
+                ("stage_one_reference", expected_stage1),
+                ("frozen_egospeak_reference", expected_original),
+            ):
+                np.testing.assert_allclose(
+                    np.load(output / name / "raw_eval" / "Session_4.npy"),
+                    expected.softmax(-1)[0].numpy(),
+                    rtol=0,
+                    atol=0,
+                )
+            metrics = json.loads((output / "metrics.json").read_text())
+            self.assertIn("raw_endpoint_anticipation", metrics)
+            with (output / "raw_endpoint_metrics.csv").open() as handle:
+                self.assertEqual(len(list(csv.DictReader(handle))), 30)
+            with self.assertRaises(FileExistsError):
+                prepare_output(output)
+
+
+if __name__ == "__main__":
+    unittest.main()

train_egospeak_omni_finetune.py

@@ -0,0 +1,520 @@
+#!/usr/bin/env python3
+"""Continue stage-one fusion with selective MiniROADA GRU/head fine-tuning.
+
+Frozen RGB/audio features and causal Omni states are reused. Final evaluation
+exports stage-two, immutable stage-one and original EgoSpeak predictions under
+identical masks, plus dense and threshold-free raw endpoint AP. No semantic
+annotations, dev checkpoint selection or test-side tuning are used.
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import math
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from tqdm.auto import tqdm
+
+from easycom_head import (
+    CLASS_NAMES,
+    NUM_HORIZONS,
+    OFFICIAL_TEST_SESSIONS,
+    OFFICIAL_TRAIN_SESSIONS,
+)
+from egospeak_omni_finetune import (
+    STAGE_TWO_FORMAT,
+    EgoSpeakOmniFinetune,
+    base_model_config,
+    load_stage_one_checkpoint,
+    validate_training_sessions,
+)
+from egospeak_omni_residual import (
+    EgoSpeakOmniResidual,
+    EgoSpeakOmniWindowDataset,
+    align_source_predictions,
+    miniroad_normalized_loss,
+    sha256_file,
+)
+from evaluate_floor_state_changes import evaluate_horizon
+from train_egospeak_omni_residual import (
+    autocast_context,
+    evaluate_predictions,
+    load_sessions,
+    make_loader,
+    normalize_sessions,
+    omni_contract_signature,
+    predict_session,
+    print_metrics,
+    resolve_device,
+    set_seed,
+)
+
+
+def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--egospeak-feature-root", type=Path, required=True)
+    parser.add_argument("--omni-feature-root", type=Path, required=True)
+    parser.add_argument("--stage-one-checkpoint", type=Path, required=True)
+    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")
+    parser.add_argument(
+        "--omni-layer-fusion", choices=("single", "last", "mean"), default="mean"
+    )
+    parser.add_argument(
+        "--base-trainable",
+        choices=("none", "heads", "gru-heads"),
+        default="gru-heads",
+        help="none is a matched extra-epoch fusion-only control; layer1 always stays frozen.",
+    )
+    parser.add_argument(
+        "--epochs", type=int, default=4, help="Additional stage-two epochs."
+    )
+    parser.add_argument("--batch-size", type=int, default=64)
+    parser.add_argument("--num-workers", type=int, default=4)
+    parser.add_argument(
+        "--window-size", type=int, default=None, help="Inherit stage-one window."
+    )
+    parser.add_argument(
+        "--stride", type=int, default=None, help="Inherit stage-one sampling stride."
+    )
+    parser.add_argument("--fusion-lr", type=float, default=1e-4)
+    parser.add_argument("--base-lr", type=float, default=1e-5)
+    parser.add_argument("--weight-decay", type=float, default=0.05)
+    parser.add_argument("--grad-clip", type=float, default=1.0)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--amp", action="store_true")
+    parser.add_argument("--eval-chunk-size", type=int, default=0)
+    parser.add_argument("--max-train-examples", type=int, default=None)
+    parser.add_argument("--no-progress", action="store_true")
+    return parser.parse_args(argv)
+
+
+def validate_checkpoint_contract(
+    args: argparse.Namespace, payload: dict[str, Any]
+) -> None:
+    training = payload["training_config"]
+    for key in ("omni_layer_fusion", "expected_omni_mode"):
+        if key not in training or getattr(args, key) != training[key]:
+            raise ValueError(f"Stage-two {key} must match stage-one extraction/readout")
+    if args.window_size is None:
+        args.window_size = int(training["window_size"])
+    if args.stride is None:
+        args.stride = int(training["stride"])
+    if args.window_size != training["window_size"] or args.stride != training["stride"]:
+        raise ValueError(
+            "Keep stage-one window_size/stride fixed for this continuation"
+        )
+
+
+@torch.inference_mode()
+def verify_stage_one_recovery(
+    model: EgoSpeakOmniFinetune,
+    reference: EgoSpeakOmniResidual,
+    dataset: EgoSpeakOmniWindowDataset,
+    device: torch.device,
+) -> float:
+    model.eval()
+    reference.eval()
+    rgb, audio, omni, tail, _, _, source = dataset[0]
+    inputs = (
+        rgb[None].to(device),
+        audio[None].to(device),
+        omni[None].to(device),
+        tail[None].to(device),
+    )
+    actual, actual_base = model.forward_last(*inputs)
+    expected, expected_base = reference.forward_last(*inputs)
+    maximum = max(
+        float((actual - expected).abs().max()),
+        float((actual_base - expected_base).abs().max()),
+    )
+    if maximum != 0.0:
+        raise RuntimeError(
+            f"Stage-two initialization changed stage-one logits: {maximum}"
+        )
+    print("Initialization audit: stage-two logits == trained stage-one logits exactly")
+    print(
+        f"Causal sample: source_row={int(source)}, cutoff={(int(source) + 1) / 5:.2f}s; "
+        "input ends at source row, labels are source+1,...,source+10"
+    )
+    return maximum
+
+
+def train_epoch(model, loader, optimizer, scaler, device, args) -> dict[str, Any]:
+    model.train()
+    total_loss = 0.0
+    num_examples = 0
+    first_batch_gradient_norms = {}
+    iterator = tqdm(
+        loader, desc="train stage2", disable=args.no_progress, dynamic_ncols=True
+    )
+    for rgb, audio, omni, tail, targets, _, _ in iterator:
+        rgb = rgb.to(device=device, dtype=torch.float32)
+        audio = audio.to(device=device, dtype=torch.float32)
+        omni = omni.to(device=device, dtype=torch.float32)
+        tail = tail.to(device)
+        targets = targets.to(device=device, dtype=torch.float32)
+        optimizer.zero_grad(set_to_none=True)
+        with autocast_context(device, args.amp):
+            logits, _ = model.forward_last(rgb, audio, omni, tail)
+            loss = miniroad_normalized_loss(logits.float(), targets)
+        if not torch.isfinite(loss):
+            raise RuntimeError("Nonfinite stage-two training loss")
+        scaler.scale(loss).backward()
+        scaler.unscale_(optimizer)
+        if not num_examples:
+            for group in optimizer.param_groups:
+                norms = [
+                    p.grad.detach().norm()
+                    for p in group["params"]
+                    if p.grad is not None
+                ]
+                norm = float(torch.stack(norms).norm()) if norms else 0.0
+                first_batch_gradient_norms[group["name"]] = norm
+            print(
+                f"First-batch gradient norms (before clipping): {first_batch_gradient_norms}"
+            )
+        torch.nn.utils.clip_grad_norm_(
+            model.trainable_parameters(),
+            args.grad_clip,
+            error_if_nonfinite=True,
+        )
+        scaler.step(optimizer)
+        scaler.update()
+        num_examples += len(rgb)
+        total_loss += float(loss.detach()) * len(rgb)
+        iterator.set_postfix_str(f"loss={total_loss / num_examples:.5f}")
+    return {
+        "loss": total_loss / num_examples,
+        "num_examples": num_examples,
+        "first_batch_gradient_norms": first_batch_gradient_norms,
+    }
+
+
+def raw_endpoint_metrics(predictions, valid, targets) -> dict[str, Any]:
+    names = list(predictions)
+    target = np.concatenate([targets[name] for name in names])
+    prediction = np.concatenate([predictions[name] for name in names])
+    mask = np.concatenate([valid[name] for name in names])
+    session_ids = np.concatenate(
+        [np.full(len(targets[name]), i) for i, name in enumerate(names)]
+    )
+    horizons = []
+    for horizon in range(NUM_HORIZONS):
+        item = evaluate_horizon(target, prediction, mask, session_ids, horizon, [])
+        horizons.append(
+            {
+                "seconds": item["nominal_seconds"],
+                "onset": item["wearer_off_to_on_endpoint"],
+                "offset": item["wearer_on_to_off_endpoint"],
+            }
+        )
+    means = {}
+    for family in ("onset", "offset"):
+        values = [h[family]["average_precision"] for h in horizons]
+        # Undefined horizons stay undefined rather than being silently dropped.
+        means[f"mean_{family}_ap"] = (
+            None if any(v is None for v in values) else float(np.mean(values))
+        )
+    return {
+        "protocol": "causal_valid_endpoint_state_change_ap",
+        "horizons": horizons,
+        **means,
+    }
+
+
+def prepare_output(path: Path) -> Path:
+    path = path.expanduser().resolve()
+    if path.exists() and (not path.is_dir() or any(path.iterdir())):
+        raise FileExistsError(f"Refusing to overwrite output: {path}")
+    path.mkdir(parents=True, exist_ok=True)
+    for method in (
+        path,
+        path / "stage_one_reference",
+        path / "frozen_egospeak_reference",
+    ):
+        for name in ("raw_eval", "pred", "valid"):
+            (method / name).mkdir(parents=True)
+    return path
+
+
+def main(argv: list[str] | None = None) -> None:
+    args = parse_args(argv)
+    if (
+        args.epochs <= 0
+        or args.batch_size <= 0
+        or args.eval_chunk_size < 0
+        or args.num_workers < 0
+        or args.grad_clip <= 0
+        or args.weight_decay < 0
+        or not all(
+            math.isfinite(value)
+            for value in (
+                args.fusion_lr,
+                args.base_lr,
+                args.grad_clip,
+                args.weight_decay,
+            )
+        )
+    ):
+        raise ValueError("Invalid stage-two epoch/batch/chunk/worker/gradient settings")
+    set_seed(args.seed)
+    device = resolve_device(args.device)
+    if args.amp and device.type != "cuda":
+        raise ValueError("--amp requires CUDA")
+    reference, parent = load_stage_one_checkpoint(args.stage_one_checkpoint)
+    validate_checkpoint_contract(args, parent)
+    train_names = normalize_sessions(args.train_sessions, OFFICIAL_TRAIN_SESSIONS)
+    eval_names = normalize_sessions(args.eval_sessions, OFFICIAL_TEST_SESSIONS)
+    validate_training_sessions(train_names)
+    train_sessions = load_sessions(args, train_names)
+    eval_sessions = load_sessions(args, eval_names)
+    sessions = (*train_sessions, *eval_sessions)
+    if {int(s.omni.shape[-1]) for s in sessions} != {reference.omni_dim}:
+        raise ValueError("Omni feature dimensions differ from stage-one checkpoint")
+    signatures = {omni_contract_signature(s) for s in sessions}
+    if len(signatures) != 1:
+        raise ValueError("Train/evaluation Omni extraction contracts differ")
+    # v1 saves readout/mode but not the entire extraction contract in its
+    # checkpoint. Verify sibling metrics provenance when it is available.
+    parent_metrics = (
+        args.stage_one_checkpoint.expanduser().resolve().parent / "metrics.json"
+    )
+    contract_verified = False
+    if parent_metrics.is_file():
+        metadata = json.loads(parent_metrics.read_text()).get("omni_sources", {})
+        for source in metadata.values():
+            config = source["config"]
+            signature = (
+                source["model_id"],
+                source["mode"],
+                config.get("memory_seconds"),
+                config.get("video_fps"),
+                config.get("readout_last_n_layers"),
+                config.get("readout_layer_fusion"),
+                tuple(source["state_shape_before_layer_fusion"][1:]),
+            )
+            if signature not in signatures:
+                raise ValueError(
+                    "Current Omni extraction differs from stage-one metrics provenance"
+                )
+        contract_verified = bool(metadata)
+    if not contract_verified:
+        print(
+            "WARNING: stage-one full extraction provenance unavailable; reuse its exact Omni root."
+        )
+
+    model = EgoSpeakOmniFinetune.from_stage_one(
+        reference,
+        parent["model_config"],
+        base_trainable=args.base_trainable,
+    ).to(device)
+    reference = reference.to(device)
+    groups = model.optimizer_groups(fusion_lr=args.fusion_lr, base_lr=args.base_lr)
+    for group in groups:
+        print(
+            f"parameters {group['name']}: {sum(p.numel() for p in group['params']):,}; lr={group['lr']:g}"
+        )
+    print(f"base_trainable={args.base_trainable}; layer1 frozen=True; Omni frozen=True")
+    print(f"train={train_names}; final evaluation={eval_names}")
+    dataset = EgoSpeakOmniWindowDataset(
+        train_sessions,
+        window_size=args.window_size,
+        stride=args.stride,
+        seed=args.seed,
+        max_examples=args.max_train_examples,
+    )
+    difference = verify_stage_one_recovery(model, reference, dataset, device)
+    output = prepare_output(args.output_dir)
+    optimizer = torch.optim.AdamW(groups, weight_decay=args.weight_decay)
+    try:
+        scaler = torch.amp.GradScaler("cuda", enabled=args.amp)
+    except (AttributeError, TypeError):  # Older PyTorch compatibility.
+        scaler = torch.cuda.amp.GradScaler(enabled=args.amp)
+    history = []
+    for epoch in range(1, args.epochs + 1):
+        dataset.resample(epoch - 1)
+        result = train_epoch(
+            model, make_loader(dataset, args, epoch), optimizer, scaler, device, args
+        )
+        result["epoch"] = epoch
+        history.append(result)
+        print(
+            f"stage2 epoch {epoch}/{args.epochs}: loss={result['loss']:.6f}; examples={result['num_examples']}"
+        )
+
+    training_config = vars(args).copy()
+    training_config = {
+        k: str(v) if isinstance(v, Path) else v for k, v in training_config.items()
+    }
+    budgets = {
+        "stage_one_epochs": int(parent["training_config"]["epochs"]),
+        "stage_two_epochs": args.epochs,
+        "total_adaptation_epochs": int(parent["training_config"]["epochs"])
+        + args.epochs,
+    }
+    checkpoint = output / "checkpoint_last.pt"
+    provenance = {
+        "stage_one_checkpoint": str(args.stage_one_checkpoint.expanduser().resolve()),
+        "stage_one_checkpoint_sha256": sha256_file(
+            args.stage_one_checkpoint.expanduser().resolve()
+        ),
+        "original_base_checkpoint": parent.get("base_checkpoint"),
+        "original_base_checkpoint_sha256": parent.get("base_checkpoint_sha256"),
+        "stage_one_full_extraction_contract_verified": contract_verified,
+    }
+    (output / "README.md").write_text(
+        "# Selective EgoSpeak + Omni fine-tuning\n\n"
+        f"Train: {train_names}. Final evaluation: {eval_names}. Seed: {args.seed}.\n\n"
+        f"Checkpoint: `checkpoint_last.pt`. Base policy: `{args.base_trainable}`. "
+        "RGB/audio feature extractors, input projection and Omni remain frozen.\n\n"
+        f"Epoch budgets: {budgets}. This continuation is not budget-matched "
+        "against stage one; use the `--base-trainable none` continuation control.\n\n"
+        "`metrics.json` reports dense AP/mAP and threshold-free raw endpoint AP "
+        "as fractions. `raw_endpoint_metrics.csv` has per-horizon AP and counts. "
+        "Raw onset uses OFF->ON versus OFF->OFF, scored by wearer probability; "
+        "offset uses ON->OFF versus ON->ON, scored by 1-wearer probability.\n\n"
+        "Immutable reference predictions are in `stage_one_reference/` and "
+        "`frozen_egospeak_reference/`, each with pred/valid/raw_eval and metrics. "
+        "No provisional semantic annotations are used.\n\n"
+        f"Run arguments:\n```json\n{json.dumps(training_config, indent=2)}\n```\n"
+    )
+    torch.save(
+        {
+            "format": STAGE_TWO_FORMAT,
+            "model_state_dict": model.state_dict(),
+            "base_model_config": base_model_config(model.base),
+            "model_config": parent["model_config"],
+            "base_trainable": args.base_trainable,
+            "training_config": training_config,
+            "train_sessions": train_names,
+            "eval_sessions": eval_names,
+            "stage_one_train_sessions": parent["train_sessions"],
+            "budgets": budgets,
+            "provenance": provenance,
+            "omni_sources": {item.name: item.omni_metadata for item in sessions},
+            "history": history,
+            "initialization_max_abs_stage_one_logit_difference": difference,
+        },
+        checkpoint,
+    )
+
+    directories = {
+        "stage_two": output,
+        "stage_one": output / "stage_one_reference",
+        "egospeak": output / "frozen_egospeak_reference",
+    }
+    predictions = {name: {} for name in directories}
+    valid = {}
+    targets = {}
+    for session in eval_sessions:
+        kwargs = {
+            "chunk_size": args.eval_chunk_size,
+            "amp": args.amp,
+            "show_progress": not args.no_progress,
+        }
+        fused, _ = predict_session(model, session, device, **kwargs)
+        first_stage, original = predict_session(reference, session, device, **kwargs)
+        for method, raw in (
+            ("stage_two", fused),
+            ("stage_one", first_stage),
+            ("egospeak", original),
+        ):
+            aligned, mask = align_source_predictions(raw, session.num_steps)
+            if session.name in valid and not np.array_equal(valid[session.name], mask):
+                raise RuntimeError("Reference and adapted valid masks differ")
+            valid[session.name] = mask
+            predictions[method][session.name] = aligned
+            directory = directories[method]
+            np.save(directory / "raw_eval" / f"{session.name}.npy", raw)
+            np.save(directory / "pred" / f"{session.name}.npy", aligned)
+            np.save(directory / "valid" / f"{session.name}.npy", mask)
+        targets[session.name] = np.asarray(session.targets, dtype=np.float32)
+
+    dense = {
+        name: evaluate_predictions(pred, valid, targets)
+        for name, pred in predictions.items()
+    }
+    endpoints = {
+        name: raw_endpoint_metrics(pred, valid, targets)
+        for name, pred in predictions.items()
+    }
+    for name, directory in directories.items():
+        metrics = {
+            **dense[name],
+            "format": "egospeak-omni-stage-two-eval-v1",
+            "classes": list(CLASS_NAMES),
+            "sessions": eval_names,
+            "train_sessions": train_names,
+            "checkpoint": str(checkpoint)
+            if name == "stage_two"
+            else provenance["stage_one_checkpoint"],
+            "method": name,
+            "seed": args.seed,
+            "prediction_layout": "[target_time, class, future_horizon]",
+            "raw_prediction_layout": "[source_time, future_horizon, class]",
+            "raw_endpoint_anticipation": endpoints[name],
+        }
+        if name == "stage_two":
+            metrics.update(
+                provenance=provenance,
+                budgets=budgets,
+                base_trainable=args.base_trainable,
+                training_config=training_config,
+                omni_sources={item.name: item.omni_metadata for item in sessions},
+                stage_one_reference=dense["stage_one"],
+                frozen_egospeak_baseline=dense["egospeak"],
+                semantic_evaluation="not run: provisional annotation audit pending",
+            )
+        (directory / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n")
+        print_metrics(name, dense[name])
+    csv_path = output / "raw_endpoint_metrics.csv"
+    with csv_path.open("w", newline="") as handle:
+        fields = [
+            "run",
+            "nominal_seconds",
+            "onset_endpoint_ap",
+            "offset_endpoint_ap",
+            "onset_positive",
+            "onset_negative",
+            "offset_positive",
+            "offset_negative",
+        ]
+        writer = csv.DictWriter(handle, fieldnames=fields)
+        writer.writeheader()
+        for name, metrics in endpoints.items():
+            for item in metrics["horizons"]:
+                writer.writerow(
+                    {
+                        "run": name,
+                        "nominal_seconds": item["seconds"],
+                        "onset_endpoint_ap": item["onset"]["average_precision"],
+                        "offset_endpoint_ap": item["offset"]["average_precision"],
+                        "onset_positive": item["onset"]["num_positive"],
+                        "onset_negative": item["onset"]["num_negative"],
+                        "offset_positive": item["offset"]["num_positive"],
+                        "offset_negative": item["offset"]["num_negative"],
+                    }
+                )
+    print("\nRaw endpoint ten-horizon means (percent; threshold-free):")
+    for name, metrics in endpoints.items():
+        print(
+            name,
+            {
+                k: None if metrics[k] is None else round(metrics[k] * 100, 3)
+                for k in ("mean_onset_ap", "mean_offset_ap")
+            },
+        )
+    print(f"Saved {checkpoint}; raw endpoint CSV (fractions)={csv_path}")
+
+
+if __name__ == "__main__":
+    main()