Add reproducible EasyCom causal head pipeline

README.md

@@ -32,9 +32,10 @@ classification head.
 | `smoke_test.py` | Verified on one NVIDIA L20 | Text and 12-second EasyCom audio-video inference |
 | `easycom_manifest.py` | Implemented, awaiting dataset-side run | Discover one-minute clips, preserve session timecodes, and check 5 Hz targets |
 | `qwen_omni_features.py` | Implemented, not executed here | Load only the 3B Thinker and extract fused states from one AV window |
-| `streaming_turn_model.py` | Prototype implemented, not executed here | Produce one causal Omni state every 200 ms, with reference and KV-cache modes |
+| `streaming_turn_model.py` | Implemented | Produce one causal Omni state every 200 ms, with causal-prefix and open-stream modes |
 | `TurnPredictionHead` | Defined, untrained | Map `[2048]` to `[10, 3]` logits |
-| Head training and EgoSpeak evaluation script | Not yet implemented | Align saved states with `target_perframe`, train the head, and reproduce Table 2 mAP |
+| `train_turn_head.py` | Implemented | Train the frozen-Omni anticipation head with the official EasyCom split |
+| `infer_turn_head.py` | Implemented | Save EgoSpeak-format predictions and reproduce per-horizon AP/mAP |
 
 The official full model, Thinker input path, embedded audio, video decoding,
 and FlashAttention-2 have passed a 12-second EasyCom smoke test on one NVIDIA
@@ -284,13 +285,17 @@ Session_1.json         clip order, exact T, crop and sampling provenance
 Session_1.ffconcat     reproducible dense concatenation order
 ```
 
-The script fails rather than silently padding or truncating when
+The target length must equal the number of complete 5 Hz intervals,
-`sum(source_frames) / 4 != target.shape[0]`.
+`floor(sum(source_frames) / 4)`. If a session ends with one to three residual
+20 Hz frames, the target-derived duration discards only that incomplete final
+interval; the discarded frame count is recorded in the sidecar.
 
 Each clip's decoded video and audio are trimmed to the exact duration implied
 by its video frame count before concatenation. This prevents AAC priming/tail
 padding from accumulating at the 24 clip boundaries and causing progressive
-audio/video drift.
+audio/video drift. The session-global visual grid is padded defensively at EOF
+and then trimmed by frame index, so it contains exactly
+`ceil(target_steps / 5 * video_fps)` frames on every FFmpeg version.
 
 On the remote L20 machine, use its own dataset paths. For example:
 
@@ -336,7 +341,7 @@ That early state can contain audio from 0.2--2.0 seconds. At prediction time
 
 ## Causal 5 Hz extraction modes
 
-`streaming_turn_model.py` implements three modes.
+`streaming_turn_model.py` implements two modes.
 
 ### 1. `causal-prefix`: correctness and semantic baseline
 
@@ -368,57 +373,7 @@ python streaming_turn_model.py \
   --output features/easycom_20s_prefix.pt
 ```
 
-### 2. `commit-replay`: closed-turn compatibility baseline
+### 2. `open-stream-replay`: recommended efficient mode
-
-The extractor retains Thinker KV for completed two-second media blocks. During
-the unfinished block it branches from the same committed cache and evaluates
-only the available prefix:
-
-```text
-committed historical KV + current block [0.0, 0.2] -> q_t
-committed historical KV + current block [0.0, 0.4] -> q_t
-...
-committed historical KV + current block [0.0, 2.0] -> q_t
-                                                    -> commit media KV
-```
-
-The temporary readout/assistant query is never committed. Only the completed
-media block enters long-term cache. This avoids recomputing the full history,
-but still re-encodes at most the current two-second prefix.
-
-Example with numerical checking:
-
-```bash
-python streaming_turn_model.py \
-  --video /path/to/easycom_session.mp4 \
-  --use-audio-in-video \
-  --mode commit-replay \
-  --start 0 \
-  --end 20 \
-  --memory-seconds 20 \
-  --verify-every 5 \
-  --verify-continuous-every 5 \
-  --output features/easycom_20s_stream.pt
-```
-
-`--verify-every 5` compares cached processing with a from-scratch replay of the
-same sequence of two-second media turns. The saved cosine similarity should be
-extremely close to 1. If it is not, the TMRoPE/cache path is invalid and must
-not be used for a paper result.
-
-`--verify-continuous-every 5` is a different check: it compares the cached
-sequence of two-second media turns with a single continuous causal window.
-This comparison need not be numerically identical, but it measures whether the
-input organization materially changes the representation. Save and report the
-two check types separately (`cache_replay` and `continuous_window_semantics`).
-
-There is an important experimental distinction: `commit-replay` represents
-history as a sequence of two-second media turns. Its numerical check proves
-cache equivalence to that same turn sequence, not equivalence to one continuous
-20-second video placeholder. Compare it against `causal-prefix` on identical
-clips and report the representation if it changes mAP.
-
-### 3. `open-stream-replay`: recommended efficient mode
 
 This mode keeps one user/media turn open across the whole bounded-memory
 segment. The first block contributes the chat and AV opening markers once;
@@ -468,20 +423,6 @@ Twenty seconds is a practical memory/context setting for the 3090, not the
 maximum video duration. The complete session is processed sequentially and a
 prediction is emitted every 0.2 seconds.
 
-The legacy `commit-replay` prototype resets Thinker KV at each non-overlapping
-20-second boundary:
-
-```text
-session:  0 -------- 20 -------- 40 -------- 60 ...
-cache:   [segment 1] [segment 2] [segment 3] ...
-output:   continuous predictions across the entire session
-```
-
-This preserves complete output coverage but limits available history after a
-reset. Report the memory length and additionally evaluate the first two seconds
-after each reset versus all other timesteps. Rolling KV eviction is a possible
-future optimization, but requires correct multimodal position rebasing.
-
 `open-stream-replay` instead uses blockwise rolling memory. With a 20-second
 budget it rebuilds the open cache from the latest nine completed blocks
 (18 seconds) when a new block starts, then adds the current 0--2 second tail.
@@ -515,6 +456,10 @@ uses only frames with timestamps `<= t`, and duplicates the most recent past
 frame when the visual tower needs an even frame count. Audio is loaded once at
 16 kHz and sliced exactly at each causal boundary. The processor still receives
 `use_audio_in_video=True`, so Qwen's audio/video token interleaving is retained.
+Long runs use `tqdm` and update the progress bar after every completed state,
+showing percentage, media time, elapsed time, ETA, and throughput. Pass
+`--no-progress` to disable it. The Thinker-only path also suppresses Qwen's
+irrelevant speech-output system prompt warning.
 
 Before a full session, run only the first four seconds and inspect both check
 types:
@@ -536,7 +481,7 @@ python streaming_turn_model.py \
 
 Do not launch all sessions unless `cache_replay` cosine is nearly 1. The
 `continuous_window_semantics` cosine answers a different question and may be
-lower because media-turn boundaries change the representation.
+lower because blockwise media preprocessing changes the representation.
 
 Video can be sampled at 2 FPS while predictions remain at 5 Hz. Audio and the
 causal boundary update every 200 ms; consecutive prediction times may use the
@@ -552,7 +497,7 @@ Each `.pt` file contains:
 - `label_indices`: `[0, ..., N-1]`, direct `target_perframe` row alignment;
 - `valid_10_horizon_steps`: `N-10`, because the last ten current rows do not
   have all future labels;
-- `mode`: `causal-prefix`, `commit-replay`, or `open-stream-replay`;
+- `mode`: `causal-prefix` or `open-stream-replay`;
 - `stream_representation`: explicit input-organization metadata;
 - `config`: step, native block, memory, FPS, and pixel settings;
 - `verification`: cached-versus-exact numerical checks;
@@ -575,7 +520,8 @@ LayerNorm(2048)
 -> reshape [10, 3]
 ```
 
-`streaming_turn_model.py` exposes:
+`easycom_head.py` defines the head and loss; they remain re-exported by
+`streaming_turn_model.py` for compatibility:
 
 ```python
 from streaming_turn_model import (
@@ -584,16 +530,37 @@ from streaming_turn_model import (
 )
 ```
 
-Recommended initial procedure:
+The supplied LSTR config uses Adam, 50 epochs, learning rate `7e-5`, weight
+decay `5e-5`, dropout `0.2`, and 10 warmup epochs followed by cosine decay.
+`train_turn_head.py` uses those defaults. Its batch size is 512 causal states,
+not LSTR's 16 temporal windows, because the two model inputs have different
+shapes.
+
+Train once on the official nine sessions with a fixed 50-epoch protocol:
+
+```bash
+python train_turn_head.py \
+  --feature-root features/causal_prefix_2fps \
+  --target-root target_frame \
+  --output checkpoints/causal_prefix_2fps_head.pth \
+  --device cuda:0
+```
+
+The default loss is `lstr-mce`, exactly matching LSTR's MCE treatment of a
+multi-positive `[0,1,1]` row: it sums both active negative-log-softmax terms
+without normalizing the target. The previous L2-normalized objective belongs
+to MiniROAD and is available only as the explicit
+`--loss miniroad-normalized` ablation.
+
+Recommended experimental procedure:
 
 1. Freeze Qwen and extract all training-session states offline.
 2. Verify state/label time offset and truncate only the final ten invalid
    current positions.
 3. Keep EgoSpeak targets as float multi-positive tensors `[N, 10, 3]`; never
    convert them with `argmax`.
-4. Train with `egospeak_multihorizon_loss`. It exactly mirrors MiniROAD's
+4. For comparison with the supplied LSTR shell, train with the default
-   anticipation objective: L2-normalize each three-class target, then compute
+   `egospeak_multihorizon_loss`/`lstr-mce` objective.
-   its weighted negative log-softmax.
 5. Select hyperparameters using only training sessions. For EasyCom, reserve a
    training session for development or use session-level cross-validation;
    never tune on Sessions 1--3.
@@ -605,7 +572,22 @@ ablations. They should not replace the frozen-Omni linear/MLP-head baseline.
 
 ## EgoSpeak-compatible evaluation
 
-For every valid current timestep, produce probabilities shaped:
+Run inference on the official test sessions (`Session_3`, `Session_2`, and
+`Session_1`):
+
+```bash
+python infer_turn_head.py \
+  --feature-root features/causal_prefix_2fps \
+  --target-root target_frame \
+  --checkpoint checkpoints/causal_prefix_2fps_head.pth \
+  --output-dir results/causal_prefix_2fps_head \
+  --device cuda:0
+```
+
+This writes `pred/Session_X.npy` with the exact LSTR/EgoSpeak layout
+`[target_time, class, future_horizon]`, i.e. `[T,3,10]`, plus `metrics.json`.
+For every valid current timestep, the head itself produces probabilities
+shaped:
 
 ```text
 [number of timesteps, 10 future horizons, 3 classes]
@@ -621,8 +603,22 @@ Then reproduce EgoSpeak's Table 2 procedure:
 5. average the ten horizon mAP values for the final score.
 
 With EgoSpeak's `ignore_index=-100`, the released evaluator does not remove
-class index zero, so background is included. Reuse the same implementation or
+class index zero, so background is included. The script reports two protocols:
-verify a reimplementation numerically.
+
+- `egospeak_full_timeline`: the same full `[T,3,10]` array and evaluator as
+  LSTR, with zeros where a prediction would require pre-session context;
+- `causal_valid_only`: excludes only those unavailable first `h` rows for
+  horizon `h`. This is the primary strictly causal score.
+
+The difference involves at most the first ten rows of each session, but both
+are saved so the comparison is explicit.
+
+The convenience wrapper mirrors EgoSpeak's train/inference shells:
+
+```bash
+chmod +x run_easycom_head_experiment.sh
+CUDA_VISIBLE_DEVICES=0 ./run_easycom_head_experiment.sh all
+```
 
 Reference Table 2 targets from the EgoSpeak paper:
 
@@ -701,7 +697,7 @@ pip install -r requirements.txt
 ```
 
 The implementation is pinned to `transformers==4.52.3` because the
-commit/replay path uses version-specific Thinker internals. On Linux,
+open-stream replay path uses version-specific Thinker internals. On Linux,
 `qwen-omni-utils[decord]` is recommended for faster video decoding. On a
 compatible NVIDIA GPU, optionally install `flash-attn` and pass
 `--flash-attention-2`.
@@ -732,8 +728,9 @@ Do not launch full EasyCom/Ego4D extraction until all of these pass:
   time `t`;
 - `states` and `target_perframe` lengths/time origins are reconciled and the
   chosen label offset is saved;
-- commit/replay cosine equivalence is near 1 on short clips;
+- open-stream cache replay is exactly equivalent to its from-scratch replay on
-- commit/replay and causal-prefix predictions are compared on identical clips;
+  short clips;
+- open-stream and causal-prefix predictions are compared on identical clips;
 - a 20-second run fits the 3090 with the selected pixel budget and attention
   implementation;
 - saved outputs contain exactly ten future horizons in the order 0.2--2.0 s;

easycom_head.py

@@ -0,0 +1,299 @@
+"""Shared head, labels, and metrics for the EasyCom EgoSpeak experiment."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Literal, Sequence
+
+import numpy as np
+import torch
+from torch import nn
+
+
+OFFICIAL_TRAIN_SESSIONS = (
+    "Session_7",
+    "Session_11",
+    "Session_10",
+    "Session_6",
+    "Session_9",
+    "Session_4",
+    "Session_8",
+    "Session_5",
+    "Session_12",
+)
+OFFICIAL_TEST_SESSIONS = ("Session_3", "Session_2", "Session_1")
+CLASS_NAMES = (
+    "background",
+    "camera_wearer_speaking",
+    "other_people_speaking",
+)
+STEP_SECONDS = 0.2
+NUM_HORIZONS = 10
+NUM_CLASSES = 3
+
+
+class TurnPredictionHead(nn.Module):
+    """Map one Omni state to EgoSpeak-compatible ``[10, 3]`` logits."""
+
+    def __init__(
+        self,
+        hidden_size: int = 2048,
+        num_horizons: int = NUM_HORIZONS,
+        num_classes: int = NUM_CLASSES,
+        dropout: float = 0.1,
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_horizons = num_horizons
+        self.num_classes = num_classes
+        self.dropout = dropout
+        self.net = nn.Sequential(
+            nn.LayerNorm(hidden_size),
+            nn.Linear(hidden_size, hidden_size // 2),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_size // 2, num_horizons * num_classes),
+        )
+
+    def forward(self, states: torch.Tensor) -> torch.Tensor:
+        logits = self.net(states)
+        return logits.view(
+            *states.shape[:-1], self.num_horizons, self.num_classes
+        )
+
+
+def _validate_loss_inputs(logits: torch.Tensor, targets: torch.Tensor) -> None:
+    if logits.shape != targets.shape:
+        raise ValueError(
+            f"logits/targets shape mismatch: {tuple(logits.shape)} vs "
+            f"{tuple(targets.shape)}"
+        )
+    if logits.shape[-1] != NUM_CLASSES:
+        raise ValueError(
+            f"Expected {NUM_CLASSES} classes, got {logits.shape[-1]}"
+        )
+
+
+def egospeak_multihorizon_loss(
+    logits: torch.Tensor,
+    targets: torch.Tensor,
+    *,
+    reduction: Literal["mean", "sum", "none"] = "mean",
+) -> torch.Tensor:
+    """Reproduce the LSTR ``MCE`` loss used by the supplied shell/config.
+
+    EasyCom targets may contain both active speaker classes as ``[0, 1, 1]``.
+    LSTR preserves those weights directly and sums their negative log-softmax
+    terms. It does not normalize the target row.
+    """
+
+    _validate_loss_inputs(logits, targets)
+    targets = targets.to(device=logits.device, dtype=logits.dtype)
+    loss = -(targets * nn.functional.log_softmax(logits, dim=-1)).sum(dim=-1)
+    if reduction == "mean":
+        return loss.mean()
+    if reduction == "sum":
+        return loss.sum()
+    if reduction == "none":
+        return loss
+    raise ValueError(f"Unknown reduction: {reduction}")
+
+
+def miniroad_multihorizon_loss(
+    logits: torch.Tensor,
+    targets: torch.Tensor,
+    *,
+    reduction: Literal["mean", "sum", "none"] = "mean",
+) -> torch.Tensor:
+    """Reproduce MiniROAD's L2-normalized multi-positive anticipation loss."""
+
+    _validate_loss_inputs(logits, targets)
+    targets = nn.functional.normalize(
+        targets.to(device=logits.device, dtype=logits.dtype),
+        p=2,
+        dim=-1,
+    )
+    loss = -(targets * nn.functional.log_softmax(logits, dim=-1)).sum(dim=-1)
+    if reduction == "mean":
+        return loss.mean()
+    if reduction == "sum":
+        return loss.sum()
+    if reduction == "none":
+        return loss
+    raise ValueError(f"Unknown reduction: {reduction}")
+
+
+def multihorizon_cross_entropy(
+    logits: torch.Tensor,
+    targets: torch.Tensor,
+) -> torch.Tensor:
+    """Backward-compatible alias for the LSTR/EgoSpeak loss."""
+
+    return egospeak_multihorizon_loss(logits, targets)
+
+
+def normalize_session_name(value: str | int) -> str:
+    text = str(value)
+    if text.startswith("Session_"):
+        suffix = text.removeprefix("Session_")
+    else:
+        suffix = text
+    if not suffix.isdigit():
+        raise ValueError(f"Invalid EasyCom session: {value}")
+    return f"Session_{int(suffix)}"
+
+
+@dataclass
+class SessionData:
+    name: str
+    states: torch.Tensor
+    targets: torch.Tensor
+    label_indices: torch.LongTensor
+    feature_metadata: dict[str, Any]
+
+    @property
+    def hidden_size(self) -> int:
+        return int(self.states.shape[-1])
+
+    def future_examples(
+        self,
+        num_horizons: int = NUM_HORIZONS,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.LongTensor]:
+        """Return states, future labels, and target indices for valid rows."""
+
+        offsets = torch.arange(1, num_horizons + 1, dtype=torch.long)
+        future_indices = self.label_indices[:, None] + offsets[None, :]
+        valid = (future_indices >= 0).all(dim=1) & (
+            future_indices < len(self.targets)
+        ).all(dim=1)
+        rows = future_indices[valid]
+        return self.states[valid], self.targets[rows], rows
+
+
+def load_session_data(
+    feature_root: str | Path,
+    target_root: str | Path,
+    session: str | int,
+    *,
+    expected_mode: str | None = "causal-prefix",
+) -> SessionData:
+    """Load and validate one saved Omni feature payload and target array."""
+
+    name = normalize_session_name(session)
+    feature_path = Path(feature_root) / f"{name}.pt"
+    target_path = Path(target_root) / f"{name}.npy"
+    if not feature_path.is_file():
+        raise FileNotFoundError(feature_path)
+    if not target_path.is_file():
+        raise FileNotFoundError(target_path)
+
+    payload = torch.load(feature_path, map_location="cpu", weights_only=False)
+    if not isinstance(payload, dict) or "states" not in payload:
+        raise ValueError(f"Invalid feature payload: {feature_path}")
+    states = payload["states"]
+    if not isinstance(states, torch.Tensor) or states.ndim != 2:
+        raise ValueError(f"states must be [T,D] in {feature_path}")
+    if not torch.is_floating_point(states):
+        raise ValueError(f"states must be floating point in {feature_path}")
+    if not torch.isfinite(states).all():
+        raise ValueError(f"states contain NaN/Inf in {feature_path}")
+
+    mode = payload.get("mode")
+    if expected_mode is not None and mode != expected_mode:
+        raise ValueError(
+            f"{feature_path} mode is {mode!r}, expected {expected_mode!r}"
+        )
+    label_indices = payload.get("label_indices")
+    if label_indices is None:
+        raise ValueError(f"Missing label_indices in {feature_path}")
+    label_indices = torch.as_tensor(label_indices, dtype=torch.long)
+    if label_indices.ndim != 1 or len(label_indices) != len(states):
+        raise ValueError(f"Invalid label_indices in {feature_path}")
+    if len(label_indices) and not torch.equal(
+        label_indices,
+        torch.arange(len(label_indices), dtype=torch.long),
+    ):
+        raise ValueError(
+            f"{feature_path} does not use the verified zero-based dense alignment"
+        )
+
+    target_array = np.load(target_path, allow_pickle=False)
+    if target_array.ndim != 2 or target_array.shape[1] != NUM_CLASSES:
+        raise ValueError(
+            f"Expected target [T,{NUM_CLASSES}], got {target_array.shape}: "
+            f"{target_path}"
+        )
+    targets = torch.from_numpy(np.asarray(target_array, dtype=np.float32))
+    if len(label_indices) and int(label_indices[-1]) >= len(targets):
+        raise ValueError(
+            f"Last label index {int(label_indices[-1])} exceeds target length "
+            f"{len(targets)} in {name}"
+        )
+    if len(states) != len(targets):
+        raise ValueError(
+            f"{name}: feature/target length mismatch: {len(states)} vs "
+            f"{len(targets)}"
+        )
+
+    metadata = {
+        key: payload.get(key)
+        for key in (
+            "model_id",
+            "mode",
+            "config",
+            "start",
+            "end",
+            "num_steps",
+            "video_sampling_grid",
+        )
+    }
+    return SessionData(name, states, targets, label_indices, metadata)
+
+
+def average_precision_metrics(
+    scores: np.ndarray,
+    targets: np.ndarray,
+) -> dict[str, Any]:
+    """Compute EgoSpeak frame AP for all three classes, including background."""
+
+    from sklearn.metrics import average_precision_score
+
+    scores = np.asarray(scores, dtype=np.float64)
+    targets = np.asarray(targets, dtype=np.float64)
+    if scores.shape != targets.shape or scores.ndim != 2:
+        raise ValueError(
+            f"scores/targets must have equal [N,C] shape, got "
+            f"{scores.shape} and {targets.shape}"
+        )
+    if scores.shape[1] != NUM_CLASSES:
+        raise ValueError(f"Expected {NUM_CLASSES} classes, got {scores.shape[1]}")
+
+    aps = average_precision_score(targets, scores, average=None)
+    return {
+        "num_steps": int(len(scores)),
+        "class_ap": {
+            name: float(value) for name, value in zip(CLASS_NAMES, aps)
+        },
+        "map": float(np.mean(aps)),
+    }
+
+
+def checkpoint_head_config(payload: dict[str, Any]) -> dict[str, Any]:
+    """Return a validated head config, including legacy-checkpoint defaults."""
+
+    config = dict(payload.get("head_config", {}))
+    config.setdefault("hidden_size", 2048)
+    config.setdefault("num_horizons", NUM_HORIZONS)
+    config.setdefault("num_classes", NUM_CLASSES)
+    config.setdefault("dropout", 0.1)
+    return config
+
+
+def validate_disjoint_sessions(
+    first: Sequence[str],
+    second: Sequence[str],
+) -> None:
+    overlap = set(first) & set(second)
+    if overlap:
+        raise ValueError(f"Session sets overlap: {sorted(overlap)}")

infer_turn_head.py

@@ -0,0 +1,289 @@
+"""Run the Omni anticipation head and evaluate it with EgoSpeak's EasyCom metric."""
+
+from __future__ import annotations
+
+import argparse
+import json
+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_CLASSES,
+    NUM_HORIZONS,
+    OFFICIAL_TEST_SESSIONS,
+    STEP_SECONDS,
+    TurnPredictionHead,
+    average_precision_metrics,
+    checkpoint_head_config,
+    load_session_data,
+    normalize_session_name,
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--feature-root", type=Path, required=True)
+    parser.add_argument("--target-root", type=Path, required=True)
+    parser.add_argument("--checkpoint", type=Path, required=True)
+    parser.add_argument("--output-dir", type=Path, required=True)
+    parser.add_argument(
+        "--sessions",
+        nargs="+",
+        default=list(OFFICIAL_TEST_SESSIONS),
+        help="Evaluation sessions; default is the official Session 3/2/1 test split.",
+    )
+    parser.add_argument("--batch-size", type=int, default=1024)
+    parser.add_argument("--device", default="cuda:0")
+    parser.add_argument(
+        "--expected-mode",
+        default="causal-prefix",
+        help="Expected feature mode, or 'any' to disable the check.",
+    )
+    parser.add_argument("--no-progress", action="store_true")
+    return parser.parse_args()
+
+
+def _resolve_device(name: str) -> torch.device:
+    device = torch.device(name)
+    if device.type == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError(f"CUDA requested but unavailable: {name}")
+    return device
+
+
+def _load_head(
+    checkpoint: Path,
+    device: torch.device,
+) -> tuple[TurnPredictionHead, dict[str, Any]]:
+    payload = torch.load(checkpoint, map_location="cpu", weights_only=False)
+    if not isinstance(payload, dict):
+        raise ValueError(f"Invalid checkpoint: {checkpoint}")
+    state_dict = payload.get("head")
+    if not isinstance(state_dict, dict):
+        raise ValueError(f"Checkpoint has no head state_dict: {checkpoint}")
+    config = checkpoint_head_config(payload)
+    if int(config["num_horizons"]) != NUM_HORIZONS:
+        raise ValueError(
+            f"Checkpoint has {config['num_horizons']} horizons, expected {NUM_HORIZONS}"
+        )
+    if int(config["num_classes"]) != NUM_CLASSES:
+        raise ValueError(
+            f"Checkpoint has {config['num_classes']} classes, expected {NUM_CLASSES}"
+        )
+    head = TurnPredictionHead(**config)
+    head.load_state_dict(state_dict)
+    return head.to(device).eval(), payload
+
+
+@torch.inference_mode()
+def _predict_states(
+    head: TurnPredictionHead,
+    states: torch.Tensor,
+    device: torch.device,
+    batch_size: int,
+    *,
+    description: str,
+    show_progress: bool,
+) -> np.ndarray:
+    chunks: list[torch.Tensor] = []
+    starts = range(0, len(states), batch_size)
+    progress = tqdm(
+        starts,
+        total=(len(states) + batch_size - 1) // batch_size,
+        desc=description,
+        unit="batch",
+        disable=not show_progress,
+    )
+    for start in progress:
+        batch = states[start : start + batch_size].to(
+            device=device,
+            dtype=torch.float32,
+            non_blocking=True,
+        )
+        chunks.append(head(batch).softmax(dim=-1).cpu())
+    return torch.cat(chunks, dim=0).numpy()
+
+
+def _align_to_egospeak_timeline(
+    probabilities: np.ndarray,
+    label_indices: torch.LongTensor,
+    num_target_steps: int,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Return EgoSpeak ``[T,3,10]`` scores and their causal-valid mask."""
+
+    predictions = np.zeros(
+        (num_target_steps, NUM_CLASSES, NUM_HORIZONS),
+        dtype=np.float32,
+    )
+    valid = np.zeros((num_target_steps, NUM_HORIZONS), dtype=bool)
+    current = label_indices.numpy()
+    for horizon in range(NUM_HORIZONS):
+        target_indices = current + horizon + 1
+        usable = (target_indices >= 0) & (target_indices < num_target_steps)
+        rows = target_indices[usable]
+        predictions[rows, :, horizon] = probabilities[usable, horizon, :]
+        valid[rows, horizon] = True
+    return predictions, valid
+
+
+def _evaluate(
+    session_predictions: dict[str, np.ndarray],
+    session_valid: dict[str, np.ndarray],
+    session_targets: dict[str, np.ndarray],
+) -> dict[str, Any]:
+    full_results = []
+    valid_results = []
+    for horizon in range(NUM_HORIZONS):
+        full_scores = np.concatenate(
+            [scores[:, :, horizon] for scores in session_predictions.values()],
+            axis=0,
+        )
+        full_targets = np.concatenate(list(session_targets.values()), axis=0)
+        full = average_precision_metrics(full_scores, full_targets)
+        full["horizon"] = horizon + 1
+        full["seconds"] = (horizon + 1) * STEP_SECONDS
+        full_results.append(full)
+
+        valid_scores = np.concatenate(
+            [
+                session_predictions[name][session_valid[name][:, horizon], :, horizon]
+                for name in session_predictions
+            ],
+            axis=0,
+        )
+        valid_targets = np.concatenate(
+            [
+                session_targets[name][session_valid[name][:, horizon]]
+                for name in session_predictions
+            ],
+            axis=0,
+        )
+        causal = average_precision_metrics(valid_scores, valid_targets)
+        causal["horizon"] = horizon + 1
+        causal["seconds"] = (horizon + 1) * STEP_SECONDS
+        valid_results.append(causal)
+
+    return {
+        "egospeak_full_timeline": {
+            "description": (
+                "Exact EgoSpeak array/evaluator layout. Unavailable pre-session "
+                "causal predictions remain zero in the preallocated arrays."
+            ),
+            "horizons": full_results,
+            "mean_map": float(np.mean([item["map"] for item in full_results])),
+        },
+        "causal_valid_only": {
+            "description": (
+                "Primary causal metric; excludes only the first h rows of each "
+                "session for horizon h because no pre-session context exists."
+            ),
+            "horizons": valid_results,
+            "mean_map": float(np.mean([item["map"] for item in valid_results])),
+        },
+    }
+
+
+def _print_metrics(metrics: dict[str, Any]) -> None:
+    for protocol_name in ("egospeak_full_timeline", "causal_valid_only"):
+        protocol = metrics[protocol_name]
+        print(f"\n{protocol_name}:")
+        for result in protocol["horizons"]:
+            print(
+                f"Action anticipation ({result['seconds']:.2f}s) "
+                f"perframe mAP: {result['map'] * 100:.5f}"
+            )
+            class_text = ", ".join(
+                f"{name}={value * 100:.5f}"
+                for name, value in result["class_ap"].items()
+            )
+            print(f"  AP: {class_text}; n={result['num_steps']}")
+        print(
+            f"Action anticipation (mean) perframe mAP: "
+            f"{protocol['mean_map'] * 100:.5f}"
+        )
+
+
+def main() -> None:
+    args = parse_args()
+    if args.batch_size <= 0:
+        raise ValueError("--batch-size must be positive")
+    device = _resolve_device(args.device)
+    expected_mode = None if args.expected_mode == "any" else args.expected_mode
+    sessions = [normalize_session_name(item) for item in args.sessions]
+    if len(sessions) != len(set(sessions)):
+        raise ValueError("Duplicate sessions were supplied")
+    head, checkpoint_payload = _load_head(args.checkpoint, device)
+    head_config = checkpoint_head_config(checkpoint_payload)
+
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+    prediction_dir = args.output_dir / "pred"
+    valid_dir = args.output_dir / "valid"
+    prediction_dir.mkdir(parents=True, exist_ok=True)
+    valid_dir.mkdir(parents=True, exist_ok=True)
+    session_predictions: dict[str, np.ndarray] = {}
+    session_valid: dict[str, np.ndarray] = {}
+    session_targets: dict[str, np.ndarray] = {}
+
+    print(f"device: {device}")
+    print(f"sessions: {', '.join(sessions)}")
+    for name in sessions:
+        data = load_session_data(
+            args.feature_root,
+            args.target_root,
+            name,
+            expected_mode=expected_mode,
+        )
+        if data.hidden_size != int(head_config["hidden_size"]):
+            raise ValueError(
+                f"{name} feature dim is {data.hidden_size}, checkpoint expects "
+                f"{head_config['hidden_size']}"
+            )
+        probabilities = _predict_states(
+            head,
+            data.states,
+            device,
+            args.batch_size,
+            description=f"infer {name}",
+            show_progress=not args.no_progress,
+        )
+        predictions, valid = _align_to_egospeak_timeline(
+            probabilities,
+            data.label_indices,
+            len(data.targets),
+        )
+        target_array = data.targets.numpy()
+        session_predictions[name] = predictions
+        session_valid[name] = valid
+        session_targets[name] = target_array
+        np.save(prediction_dir / f"{name}.npy", predictions)
+        np.save(valid_dir / f"{name}.npy", valid)
+        print(
+            f"{name}: saved pred {predictions.shape}, "
+            f"causal-valid counts={valid.sum(axis=0).tolist()}"
+        )
+
+    metrics = _evaluate(session_predictions, session_valid, session_targets)
+    metrics.update(
+        {
+            "sessions": sessions,
+            "classes": list(CLASS_NAMES),
+            "prediction_layout": "[target_time, class, future_horizon]",
+            "checkpoint": str(args.checkpoint.resolve()),
+            "checkpoint_epoch": checkpoint_payload.get("epoch"),
+            "checkpoint_train_sessions": checkpoint_payload.get("train_sessions"),
+            "checkpoint_dev_sessions": checkpoint_payload.get("dev_sessions"),
+        }
+    )
+    metrics_path = args.output_dir / "metrics.json"
+    metrics_path.write_text(json.dumps(metrics, indent=2) + "\n")
+    _print_metrics(metrics)
+    print(f"\nmetrics: {metrics_path}")
+    print(f"EgoSpeak-compatible predictions: {prediction_dir}")
+
+
+if __name__ == "__main__":
+    main()

prepare_easycom_session.py

@@ -42,10 +42,22 @@ class ClipInfo:
 
 
 def _run_json(command: list[str]) -> dict:
+    try:
         completed = subprocess.run(
             command, check=True, capture_output=True, text=True
         )
+    except subprocess.CalledProcessError as exc:
+        stderr = (exc.stderr or "").strip()
+        detail = f"\n{stderr}" if stderr else ""
+        raise RuntimeError(
+            f"Command failed ({exc.returncode}): {shlex.join(command)}{detail}"
+        ) from exc
+    try:
         return json.loads(completed.stdout)
+    except json.JSONDecodeError as exc:
+        raise RuntimeError(
+            f"Command returned invalid JSON: {shlex.join(command)}"
+        ) from exc
 
 
 def probe_clip(path: Path) -> ClipInfo:
@@ -108,16 +120,22 @@ def validate_inventory(clips: list[ClipInfo], target_steps: int | None) -> int:
                 f"{clip.width}x{clip.height}: {clip.path}"
             )
     total_frames = sum(clip.frames for clip in clips)
-    if total_frames % (SOURCE_FPS // TASK_FPS):
+    frames_per_step = SOURCE_FPS // TASK_FPS
-        raise ValueError(
+    media_steps = total_frames // frames_per_step
-            f"Total source frames {total_frames} are not divisible into 5 Hz steps"
-        )
-    media_steps = total_frames // (SOURCE_FPS // TASK_FPS)
     if target_steps is not None and media_steps != target_steps:
         raise ValueError(
-            f"Dense media gives {media_steps} steps but target has {target_steps}"
+            f"Dense media gives {media_steps} complete 5 Hz steps from "
+            f"{total_frames} frames, but target has {target_steps}"
         )
-    return media_steps
+    # A session may end with 1--3 source frames, which cannot form another
+    # complete 5 Hz target interval. The target length is authoritative; the
+    # final video/audio trim below discards only this sub-step tail.
+    return target_steps if target_steps is not None else media_steps
+
+
+def fixed_grid_frame_count(*, fps: float, dense_duration: float) -> int:
+    """Number of samples on the grid ``n / fps < dense_duration``."""
+    return int(math.ceil(dense_duration * fps - 1e-7))
 
 
 def build_ffmpeg_command(
@@ -127,6 +145,9 @@ def build_ffmpeg_command(
     fps: float,
     dense_duration: float,
 ) -> list[str]:
+    expected_frames = fixed_grid_frame_count(
+        fps=fps, dense_duration=dense_duration
+    )
     # Trim every decoded stream independently before concatenation. Released
     # AAC clips can contain roughly one codec-padding frame beyond their video
     # duration; concatenating them first would accumulate AV drift.
@@ -158,11 +179,19 @@ def build_ffmpeg_command(
         "".join(concat_inputs)
         + f"concat=n={len(clips)}:v=1:a=1[vcat][acat]"
     )
-    # fps=start_time=0 establishes one global visual lattice after the clips
+    # Establish one global visual lattice after the clips have become a dense
-    # have become a dense timeline. The final atrim is a defensive guard
+    # timeline. Some FFmpeg versions omit the last grid point when the input
-    # against decoder rounding at the last sample.
+    # ends exactly on an integer-second boundary. Add two disposable clones,
+    # then trim by frame index, so EOF timestamp rounding cannot change the
+    # number of samples. The clones are retained only if a real final grid
+    # sample was missing.
     filters.append(
-        f"[vcat]fps=fps={fps}:start_time=0:round=down,format=yuv420p[vout]"
+        f"[vcat]trim=start=0:end={dense_duration:.9f},"
+        f"setpts=PTS-STARTPTS,"
+        f"fps=fps={fps}:start_time=0:round=down,"
+        f"tpad=stop_mode=clone:stop=2,"
+        f"trim=start_frame=0:end_frame={expected_frames},"
+        f"setpts=N/({fps}*TB),format=yuv420p[vout]"
     )
     filters.append(
         f"[acat]atrim=start=0:end={dense_duration:.9f},"
@@ -220,20 +249,15 @@ def validate_output(path: Path, *, fps: float, dense_duration: float) -> dict:
         raise RuntimeError(f"Prepared file must contain video and audio: {path}")
     actual_fps = float(Fraction(video["avg_frame_rate"]))
     actual_frames = int(video["nb_frames"])
-    # FFmpeg may either retain or discard the final partial visual interval,
+    expected_frames = fixed_grid_frame_count(
-    # depending on timestamps delivered by the concat filter. Both are valid
+        fps=fps, dense_duration=dense_duration
-    # fixed grids; for 1388.4 s at 2 FPS this permits 2776 or 2777 frames.
+    )
-    frame_product = dense_duration * fps
-    expected_frames = {
-        int(math.floor(frame_product + 1e-7)),
-        int(math.ceil(frame_product - 1e-7)),
-    }
     if not math.isclose(actual_fps, fps, rel_tol=0.0, abs_tol=1e-6):
         raise RuntimeError(f"Prepared FPS is {actual_fps}, expected {fps}: {path}")
-    if actual_frames not in expected_frames:
+    if actual_frames != expected_frames:
         raise RuntimeError(
-            f"Prepared frame count is {actual_frames}, expected one of "
+            f"Prepared frame count is {actual_frames}, expected "
-            f"{sorted(expected_frames)}: {path}"
+            f"{expected_frames}: {path}"
         )
     audio_duration = float(audio["duration"])
     # AAC frame padding can shift the container duration by a few milliseconds.
@@ -245,7 +269,7 @@ def validate_output(path: Path, *, fps: float, dense_duration: float) -> dict:
     return {
         "video_fps": actual_fps,
         "video_frames": actual_frames,
-        "allowed_video_frames": sorted(expected_frames),
+        "expected_video_frames": expected_frames,
         "video_duration_seconds": float(video["duration"]),
         "audio_duration_seconds": audio_duration,
         "audio_sample_rate": int(audio["sample_rate"]),
@@ -322,12 +346,16 @@ def main() -> None:
                 output, fps=fps, dense_duration=duration
             )
 
+    source_frames = sum(clip.frames for clip in clips)
+    used_source_frames = num_steps * (SOURCE_FPS // TASK_FPS)
     metadata = {
         "session": session,
         "source_session_dir": str(args.session_dir.resolve()),
         "source_clip_order": [asdict(clip) for clip in clips],
         "source_num_clips": len(clips),
-        "source_frames_20hz": sum(clip.frames for clip in clips),
+        "source_frames_20hz": source_frames,
+        "used_source_frames_20hz": used_source_frames,
+        "discarded_tail_frames_20hz": source_frames - used_source_frames,
         "task_fps": TASK_FPS,
         "num_steps": num_steps,
         "dense_duration_seconds": duration,

requirements.txt

@@ -5,5 +5,7 @@ accelerate
 qwen-omni-utils[decord]
 librosa
 Pillow
+tqdm
+scikit-learn
 
 # Optional on compatible NVIDIA GPUs: flash-attn

run_easycom_head_experiment.sh

@@ -0,0 +1,35 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Frozen-Qwen head equivalent of EgoSpeak's run_LSTR_{train,inference}.sh.
+# Override any path or device by setting the corresponding environment value.
+ACTION="${1:-all}"
+if [[ "${ACTION}" != "train" && "${ACTION}" != "infer" && "${ACTION}" != "all" ]]; then
+  echo "Usage: $0 [train|infer|all]" >&2
+  exit 2
+fi
+PROJECT_ROOT="${PROJECT_ROOT:-/data1/yfzx/wangx/Code/test_qwen_omni}"
+FEATURE_ROOT="${FEATURE_ROOT:-${PROJECT_ROOT}/features/causal_prefix_2fps}"
+TARGET_ROOT="${TARGET_ROOT:-${PROJECT_ROOT}/target_frame}"
+CHECKPOINT="${CHECKPOINT:-${PROJECT_ROOT}/checkpoints/causal_prefix_2fps_head.pth}"
+RESULT_ROOT="${RESULT_ROOT:-${PROJECT_ROOT}/results/causal_prefix_2fps_head}"
+DEVICE="${DEVICE:-cuda:0}"
+
+cd "${PROJECT_ROOT}"
+
+if [[ "${ACTION}" == "train" || "${ACTION}" == "all" ]]; then
+  python train_turn_head.py \
+    --feature-root "${FEATURE_ROOT}" \
+    --target-root "${TARGET_ROOT}" \
+    --output "${CHECKPOINT}" \
+    --device "${DEVICE}"
+fi
+
+if [[ "${ACTION}" == "infer" || "${ACTION}" == "all" ]]; then
+  python infer_turn_head.py \
+    --feature-root "${FEATURE_ROOT}" \
+    --target-root "${TARGET_ROOT}" \
+    --checkpoint "${CHECKPOINT}" \
+    --output-dir "${RESULT_ROOT}" \
+    --device "${DEVICE}"
+fi

streaming_turn_model.py

@@ -4,24 +4,22 @@ This module implements two inference paths for Qwen2.5-Omni-3B:
 
 1. ``causal-prefix``: recompute a causal AV window ending at every 200 ms.
    This is deliberately slow, but is the simplest correctness baseline.
-2. ``commit-replay``: keep the Thinker KV cache for completed 2-second
+2. ``open-stream-replay``: cache completed AV blocks inside one
-   multimodal blocks, and replay only the causal prefix of the current block.
-3. ``open-stream-replay``: cache the same completed AV blocks inside one
    still-open user/media turn; closure and readout tokens remain temporary.
 
 Neither path uses an RNN.  A small head maps one Thinker state to the same
 ``10 future horizons x 3 classes`` output used by EgoSpeak.
 
-The commit/replay path touches version-specific Hugging Face internals.  It is
+The open-stream replay path touches version-specific Hugging Face internals.
-therefore pinned to Transformers 4.52.3 and includes an optional exact-replay
+It is therefore pinned to Transformers 4.52.3 and includes an optional
-check.  Run that check before extracting a full dataset.
+exact-replay check. Run that check before extracting a full dataset.
 
 Example (no command is run merely by importing this file):
 
     python streaming_turn_model.py \
         --video /path/to/session.mp4 \
         --use-audio-in-video \
-        --mode commit-replay \
+        --mode open-stream-replay \
         --memory-seconds 20 \
         --verify-every 25 \
         --output features/session.pt
@@ -32,6 +30,7 @@ from __future__ import annotations
 import argparse
 import copy
 import json
+import logging
 import math
 import subprocess
 from dataclasses import asdict, dataclass
@@ -39,16 +38,58 @@ from pathlib import Path
 from typing import Any, Literal
 
 import torch
+from easycom_head import (
+    TurnPredictionHead,
+    egospeak_multihorizon_loss,
+    miniroad_multihorizon_loss,
+    multihorizon_cross_entropy,
+)
 from qwen_omni_utils import process_mm_info
 from torch import nn
+from tqdm.auto import tqdm
 from transformers import (
     Qwen2_5OmniProcessor,
     Qwen2_5OmniThinkerForConditionalGeneration,
 )
 from transformers.cache_utils import DynamicCache
 
-Mode = Literal["causal-prefix", "commit-replay", "open-stream-replay"]
+Mode = Literal["causal-prefix", "open-stream-replay"]
 MODEL_ID = "Qwen/Qwen2.5-Omni-3B"
+_KNOWN_QWEN_WARNING_TEXT = (
+    "System prompt modified, audio output may not work",
+    "Unused or unrecognized kwargs: images.",
+    "Unused or unrecognized kwargs: return_tensors, images.",
+    "Unused or unrecognized kwargs: images, return_tensors.",
+)
+
+
+class _KnownQwenWarningFilter(logging.Filter):
+    """Hide Qwen's speech-output warning in Thinker-only extraction."""
+
+    def filter(self, record: logging.LogRecord) -> bool:
+        message = record.getMessage()
+        return not any(text in message for text in _KNOWN_QWEN_WARNING_TEXT)
+
+
+def _configure_qwen_logging() -> None:
+    # Qwen2.5-Omni 4.52.3 emits this through Python's root logger whenever the
+    # conversation does not use its speech-generation system prompt. This
+    # extractor uses only the Thinker and never requests audio output.
+    def add_filter(target: logging.Filterer) -> None:
+        if not any(
+            isinstance(item, _KnownQwenWarningFilter)
+            for item in target.filters
+        ):
+            target.addFilter(_KnownQwenWarningFilter())
+
+    # The system-prompt warning is logged directly on root. The unused-kwargs
+    # warning is emitted by a Transformers child logger and handled by the
+    # library's top-level StreamHandler, so cover both paths.
+    for logger_name in (None, "transformers"):
+        logger = logging.getLogger(logger_name)
+        add_filter(logger)
+        for handler in logger.handlers:
+            add_filter(handler)
 
 
 def _choose_dtype(
@@ -132,6 +173,7 @@ class OpenMediaPieces:
     payload_embeddings: torch.Tensor
     payload_position_ids: torch.LongTensor
     closure_ids: torch.LongTensor
+    closure_position_ids: torch.LongTensor
 
 
 class FixedGridAV:
@@ -167,19 +209,22 @@ class FixedGridAV:
         )
 
     def frames(self, start: float, end: float) -> list[Any]:
-        """Return fixed-grid frames whose timestamps lie in [start, end]."""
+        """Return fixed-grid frames whose timestamps lie in [start, end)."""
 
         import numpy as np
         from PIL import Image
 
         first = max(0, int(math.ceil(start * self.fps - 1e-7)))
-        last = min(
+        last_exclusive = min(
-            self.num_frames - 1,
+            self.num_frames,
-            int(math.floor(end * self.fps + 1e-7)),
+            int(math.ceil(end * self.fps - 1e-7)),
         )
+        last = last_exclusive - 1
         if first > last:
-            # This can occur only when a caller starts between global visual
+            # A short interval can fall entirely between global visual ticks.
-            # ticks. Reuse the most recent past frame, never ceil into future.
+            # Reuse the latest frame strictly before ``end``, never a future
+            # frame. Native 2-second blocks start on-grid and do not take this
+            # fallback path.
             first = last = max(0, min(self.num_frames - 1, last))
         indices = list(range(first, last + 1))
         arrays = self.reader.get_batch(indices).asnumpy()
@@ -198,76 +243,6 @@ class FixedGridAV:
         return self.audio[first:last]
 
 
-class TurnPredictionHead(nn.Module):
-    """Map one Omni state to EgoSpeak-compatible [10, 3] logits."""
-
-    def __init__(
-        self,
-        hidden_size: int = 2048,
-        num_horizons: int = 10,
-        num_classes: int = 3,
-        dropout: float = 0.1,
-    ) -> None:
-        super().__init__()
-        self.num_horizons = num_horizons
-        self.num_classes = num_classes
-        self.net = nn.Sequential(
-            nn.LayerNorm(hidden_size),
-            nn.Linear(hidden_size, hidden_size // 2),
-            nn.GELU(),
-            nn.Dropout(dropout),
-            nn.Linear(hidden_size // 2, num_horizons * num_classes),
-        )
-
-    def forward(self, states: torch.Tensor) -> torch.Tensor:
-        logits = self.net(states)
-        return logits.view(*states.shape[:-1], self.num_horizons, self.num_classes)
-
-
-def egospeak_multihorizon_loss(
-    logits: torch.Tensor,
-    targets: torch.Tensor,
-    *,
-    reduction: Literal["mean", "sum", "none"] = "mean",
-) -> torch.Tensor:
-    """Reproduce EgoSpeak's anticipation loss for multi-positive targets.
-
-    ``logits`` and ``targets`` both have shape ``[..., 10, 3]``. EasyCom can
-    label camera-wearer and other speech simultaneously as ``[0, 1, 1]``.
-    MiniROAD L2-normalizes that row and applies a log-softmax loss.
-    """
-
-    if logits.shape != targets.shape:
-        raise ValueError(
-            f"logits/targets shape mismatch: {tuple(logits.shape)} vs "
-            f"{tuple(targets.shape)}"
-        )
-    if logits.shape[-1] != 3:
-        raise ValueError(f"Expected three classes, got {logits.shape[-1]}")
-    targets = nn.functional.normalize(
-        targets.to(device=logits.device, dtype=logits.dtype),
-        p=2,
-        dim=-1,
-    )
-    loss = -(targets * nn.functional.log_softmax(logits, dim=-1)).sum(dim=-1)
-    if reduction == "mean":
-        return loss.mean()
-    if reduction == "sum":
-        return loss.sum()
-    if reduction == "none":
-        return loss
-    raise ValueError(f"Unknown reduction: {reduction}")
-
-
-def multihorizon_cross_entropy(
-    logits: torch.Tensor,
-    targets: torch.Tensor,
-) -> torch.Tensor:
-    """Backward-compatible name for :func:`egospeak_multihorizon_loss`."""
-
-    return egospeak_multihorizon_loss(logits, targets)
-
-
 def _require_integer_ratio(numerator: float, denominator: float, name: str) -> int:
     ratio = numerator / denominator
     rounded = round(ratio)
@@ -437,16 +412,23 @@ def _prepare_conversation(
         # Qwen2.5-Omni uses this value to compute video_second_per_grid for
         # TMRoPE. This is essential for the 5 FPS ablation (default is 2 FPS).
         processor_kwargs["fps"] = video_fps
-    batch = processor(
+    processor_inputs: dict[str, Any] = {
-        text=[text],
+        "text": [text],
-        audio=audios,
+        "audio": audios,
-        images=images,
+        "videos": videos,
-        videos=videos,
+        "padding": True,
-        padding=True,
+        "use_audio_in_video": use_audio_in_video,
-        return_tensors="pt",
-        use_audio_in_video=use_audio_in_video,
         **processor_kwargs,
-    )
+    }
+    # Passing images=None and return_tensors through Qwen's 4.52.3 processor
+    # produces a warning on every streaming step. Omit the absent modality and
+    # perform the same BatchFeature tensor conversion explicitly afterward.
+    if images is not None and (
+        not isinstance(images, (list, tuple)) or len(images) > 0
+    ):
+        processor_inputs["images"] = images
+    batch = processor(**processor_inputs)
+    batch = batch.convert_to_tensors("pt")
     # BatchFeature.to(dtype) changes floating tensors only; IDs stay integral.
     batch = batch.to(model.device).to(model.dtype)
     return dict(batch)
@@ -662,18 +644,10 @@ def _open_media_pieces(
         payload_embeddings=embeddings[:, first : last + 1, :],
         payload_position_ids=positions[:, :, first : last + 1],
         closure_ids=closure_ids,
+        closure_position_ids=positions[:, :, last + 1 :],
     )
 
 
-def _positions_at_offset(
-    position_ids: torch.LongTensor,
-    offset: int,
-) -> torch.LongTensor:
-    """Preserve local 3D geometry while moving its minimum to ``offset``."""
-
-    return position_ids - position_ids.min() + offset
-
-
 def _tmrope_positions_per_chunk(
     model: Qwen2_5OmniThinkerForConditionalGeneration,
 ) -> int:
@@ -823,13 +797,22 @@ def _read_open_state(
     model: Qwen2_5OmniThinkerForConditionalGeneration,
     open_media_cache: Any,
     *,
-    next_rope_offset: int,
     closure_ids: torch.LongTensor,
+    closure_position_ids: torch.LongTensor,
     query_ids: torch.LongTensor,
 ) -> torch.Tensor:
     """Temporarily close an open media turn, read q_t, and discard the branch."""
 
     branch_cache = clone_cache(open_media_cache)
+    if closure_position_ids.shape != (
+        3,
+        closure_ids.shape[0],
+        closure_ids.shape[1],
+    ):
+        raise ValueError(
+            "closure position/ID mismatch: "
+            f"{tuple(closure_position_ids.shape)} vs {tuple(closure_ids.shape)}"
+        )
     readout_ids = torch.cat(
         [closure_ids.to(query_ids.device), query_ids],
         dim=1,
@@ -837,13 +820,24 @@ def _read_open_state(
     past_length = _cache_length(branch_cache)
     readout_length = readout_ids.shape[1]
     embeddings = model.get_input_embeddings()(readout_ids)
-    scalar_positions = torch.arange(
+
-        next_rope_offset,
+    # Qwen's AV TMRoPE gives the paired video/audio closing tokens the same
-        next_rope_offset + readout_length,
+    # position. Preserve those model-derived positions instead of replacing
+    # the whole closure with a scalar arange. Only the artificial assistant
+    # query follows the closure's maximum position.
+    closure_position_ids = closure_position_ids.to(query_ids.device)
+    query_start = int(closure_position_ids.max().item()) + 1
+    query_scalar_positions = torch.arange(
+        query_start,
+        query_start + query_ids.shape[1],
         device=readout_ids.device,
     )
-    position_ids = scalar_positions.view(1, 1, -1).expand(
+    query_position_ids = query_scalar_positions.view(1, 1, -1).expand(
-        3, readout_ids.shape[0], -1
+        3, query_ids.shape[0], -1
+    )
+    position_ids = torch.cat(
+        [closure_position_ids, query_position_ids],
+        dim=2,
     )
     cache_position = torch.arange(
         past_length,
@@ -876,32 +870,41 @@ def _append_open_pieces(
     past_key_values: Any | None,
     rope_offset: int,
     include_prefix: bool,
-) -> tuple[Any, int]:
+) -> tuple[Any, int, torch.LongTensor]:
-    """Append one AV payload, adding the chat/media prefix exactly once."""
+    """Append one AV payload and return its shifted closure positions."""
 
     if include_prefix:
         if past_key_values is not None:
             raise ValueError("An open-media prefix can only start an empty cache")
+        payload_base = int(pieces.payload_position_ids.min().item())
+        if rope_offset != payload_base:
+            raise ValueError(
+                "The first open payload must retain its local TMRoPE base: "
+                f"got {rope_offset}, expected {payload_base}"
+            )
         embeddings = torch.cat(
             [pieces.prefix_embeddings, pieces.payload_embeddings], dim=1
         )
         position_ids = torch.cat(
             [pieces.prefix_position_ids, pieces.payload_position_ids], dim=2
         )
+        closure_position_ids = pieces.closure_position_ids
     else:
         if past_key_values is None:
             raise ValueError("Payload-only append requires an existing open cache")
         embeddings = pieces.payload_embeddings
-        position_ids = _positions_at_offset(
+        position_shift = rope_offset - int(
-            pieces.payload_position_ids,
+            pieces.payload_position_ids.min().item()
-            rope_offset,
         )
-    return _append_embeddings(
+        position_ids = pieces.payload_position_ids + position_shift
+        closure_position_ids = pieces.closure_position_ids + position_shift
+    next_cache, next_rope_offset = _append_embeddings(
         model,
         embeddings,
         position_ids,
         past_key_values=past_key_values,
     )
+    return next_cache, next_rope_offset, closure_position_ids
 
 
 class CausalPrefixExtractor:
@@ -982,11 +985,23 @@ class CausalPrefixExtractor:
         start: float,
         end: float,
         use_audio_in_video: bool,
+        show_progress: bool = True,
     ) -> ExtractionResult:
         times = _time_grid(start, end, self.config.step_seconds)
         states = []
+        with tqdm(
+            total=len(times),
+            desc="causal-prefix",
+            unit="state",
+            dynamic_ncols=True,
+            mininterval=0.0,
+            miniters=1,
+            disable=not show_progress,
+        ) as progress:
             for current_time in times:
-            context_start = max(start, current_time - self.config.memory_seconds)
+                context_start = max(
+                    start, current_time - self.config.memory_seconds
+                )
                 state = self.state_for_intervals(
                     video_path=video_path,
                     audio_path=audio_path,
@@ -994,6 +1009,10 @@ class CausalPrefixExtractor:
                     use_audio_in_video=use_audio_in_video,
                 )
                 states.append(state.cpu())
+                progress.set_postfix_str(
+                    f"media={current_time:.1f}s", refresh=False
+                )
+                progress.update(1)
         return ExtractionResult(
             times=torch.tensor(times, dtype=torch.float64),
             states=torch.cat(states, dim=0),
@@ -1001,202 +1020,6 @@ class CausalPrefixExtractor:
         )
 
 
-class CommitReplayExtractor:
-    """Cache complete 2 s blocks and replay only the current causal prefix."""
-
-    def __init__(
-        self,
-        model: Qwen2_5OmniThinkerForConditionalGeneration,
-        processor: Qwen2_5OmniProcessor,
-        config: StreamConfig,
-        fixed_grid_av: FixedGridAV | None = None,
-    ) -> None:
-        config.validate()
-        self.model = model
-        self.processor = processor
-        self.config = config
-        self.fixed_grid_av = fixed_grid_av
-        self.query_ids = _readout_query_ids(processor, model.device)
-        self.reference = CausalPrefixExtractor(
-            model, processor, config, fixed_grid_av=fixed_grid_av
-        )
-
-    def _prepare_one(
-        self,
-        *,
-        video_path: str | Path | None,
-        audio_path: str | Path | None,
-        interval: tuple[float, float],
-        use_audio_in_video: bool,
-        turn_index: int,
-    ) -> dict[str, torch.Tensor]:
-        conversation, manual_audios = build_chunk_conversation(
-            video_path=video_path,
-            audio_path=audio_path,
-            intervals=[interval],
-            use_audio_in_video=use_audio_in_video,
-            config=self.config,
-            fixed_grid_av=self.fixed_grid_av,
-        )
-        text_override = _media_turn_template_suffix(
-            self.processor,
-            conversation[0],
-            turn_index=turn_index,
-        )
-        return _prepare_conversation(
-            self.processor,
-            conversation,
-            use_audio_in_video=use_audio_in_video,
-            model=self.model,
-            manual_audios=manual_audios,
-            video_fps=self.config.video_fps,
-            text_override=text_override,
-        )
-
-    def _extract_segment(
-        self,
-        *,
-        video_path: str | Path | None,
-        audio_path: str | Path | None,
-        segment_start: float,
-        segment_end: float,
-        use_audio_in_video: bool,
-        verify_every: int,
-        verify_continuous_every: int,
-        global_step_offset: int,
-    ) -> tuple[list[float], list[torch.Tensor], list[dict[str, Any]]]:
-        times = _time_grid(segment_start, segment_end, self.config.step_seconds)
-        committed_cache = None
-        committed_rope_offset = 0
-        block_start = segment_start
-        states: list[torch.Tensor] = []
-        checks: list[dict[str, Any]] = []
-
-        for local_step, current_time in enumerate(times, start=1):
-            while current_time > block_start + self.config.native_block_seconds + 1e-6:
-                block_start += self.config.native_block_seconds
-
-            # DynamicCache is mutated by forward.  Every unfinished prefix must
-            # branch from the identical committed history.
-            branch_cache = clone_cache(committed_cache)
-            inputs = self._prepare_one(
-                video_path=video_path,
-                audio_path=audio_path,
-                interval=(block_start, current_time),
-                use_audio_in_video=use_audio_in_video,
-                turn_index=round(
-                    (block_start - segment_start)
-                    / self.config.native_block_seconds
-                ),
-            )
-            media_cache, next_rope_offset = _append_media_turn(
-                self.model,
-                inputs,
-                past_key_values=branch_cache,
-                rope_offset=committed_rope_offset,
-                use_audio_in_video=use_audio_in_video,
-            )
-            state = _read_state(
-                self.model,
-                media_cache,
-                next_rope_offset=next_rope_offset,
-                query_ids=self.query_ids,
-            )
-            states.append(state.cpu())
-
-            global_step = global_step_offset + local_step
-            if verify_every > 0 and global_step % verify_every == 0:
-                intervals = _block_intervals(
-                    segment_start,
-                    current_time,
-                    self.config.native_block_seconds,
-                )
-                exact = self.reference.state_for_intervals(
-                    video_path=video_path,
-                    audio_path=audio_path,
-                    intervals=intervals,
-                    use_audio_in_video=use_audio_in_video,
-                )
-                checks.append(
-                    _compare_states(current_time, state, exact, kind="cache_replay")
-                )
-            if (
-                verify_continuous_every > 0
-                and global_step % verify_continuous_every == 0
-            ):
-                continuous = self.reference.state_for_intervals(
-                    video_path=video_path,
-                    audio_path=audio_path,
-                    intervals=[(segment_start, current_time)],
-                    use_audio_in_video=use_audio_in_video,
-                )
-                checks.append(
-                    _compare_states(
-                        current_time,
-                        state,
-                        continuous,
-                        kind="continuous_window_semantics",
-                    )
-                )
-
-            block_end = min(
-                block_start + self.config.native_block_seconds,
-                segment_end,
-            )
-            if math.isclose(current_time, block_end, rel_tol=0.0, abs_tol=1e-6):
-                # media_cache excludes the readout query, so it is safe to keep.
-                committed_cache = media_cache
-                committed_rope_offset = next_rope_offset
-                block_start = block_end
-
-        return times, states, checks
-
-    def extract(
-        self,
-        *,
-        video_path: str | Path | None,
-        audio_path: str | Path | None,
-        start: float,
-        end: float,
-        use_audio_in_video: bool,
-        verify_every: int = 0,
-        verify_continuous_every: int = 0,
-    ) -> ExtractionResult:
-        all_times: list[float] = []
-        all_states: list[torch.Tensor] = []
-        all_checks: list[dict[str, Any]] = []
-        segment_start = start
-        global_steps = 0
-
-        # A bounded cache is implemented as non-overlapping memory segments.
-        # Predictions are still emitted for the entire session, but history is
-        # reset at each boundary.  Rolling KV eviction is intentionally not
-        # attempted because it also requires rebasing multimodal RoPE.
-        while segment_start < end - 1e-7:
-            segment_end = min(segment_start + self.config.memory_seconds, end)
-            times, states, checks = self._extract_segment(
-                video_path=video_path,
-                audio_path=audio_path,
-                segment_start=segment_start,
-                segment_end=segment_end,
-                use_audio_in_video=use_audio_in_video,
-                verify_every=verify_every,
-                verify_continuous_every=verify_continuous_every,
-                global_step_offset=global_steps,
-            )
-            all_times.extend(times)
-            all_states.extend(states)
-            all_checks.extend(checks)
-            global_steps += len(times)
-            segment_start = segment_end
-
-        return ExtractionResult(
-            times=torch.tensor(all_times, dtype=torch.float64),
-            states=torch.cat(all_states, dim=0),
-            verification=all_checks,
-        )
-
-
 class OpenStreamReplayExtractor:
     """Cache completed AV payloads inside one still-open multimodal turn."""
 
@@ -1271,7 +1094,13 @@ class OpenStreamReplayExtractor:
     ) -> torch.Tensor:
         """Replay the same open stream from an empty cache for verification."""
 
-        media_cache, next_offset, _, closure_ids = self._cache_from_intervals(
+        (
+            media_cache,
+            _,
+            _,
+            closure_ids,
+            closure_position_ids,
+        ) = self._cache_from_intervals(
             video_path=video_path,
             audio_path=audio_path,
             intervals=intervals,
@@ -1280,8 +1109,8 @@ class OpenStreamReplayExtractor:
         return _read_open_state(
             self.model,
             media_cache,
-            next_rope_offset=next_offset,
             closure_ids=closure_ids,
+            closure_position_ids=closure_position_ids,
             query_ids=self.query_ids,
         )
 
@@ -1292,7 +1121,7 @@ class OpenStreamReplayExtractor:
         audio_path: str | Path | None,
         intervals: list[tuple[float, float]],
         use_audio_in_video: bool,
-    ) -> tuple[Any, int, int, torch.LongTensor]:
+    ) -> tuple[Any, int, int, torch.LongTensor, torch.LongTensor]:
         """Build an open payload cache and expose its local TMRoPE base."""
 
         if not intervals:
@@ -1300,6 +1129,7 @@ class OpenStreamReplayExtractor:
         media_cache = None
         next_offset = 0
         closure_ids = None
+        closure_position_ids = None
         payload_base = None
         positions_per_chunk = _tmrope_positions_per_chunk(self.model)
         for index, interval in enumerate(intervals):
@@ -1312,7 +1142,11 @@ class OpenStreamReplayExtractor:
             if payload_base is None:
                 payload_base = int(pieces.payload_position_ids.min().item())
             payload_offset = payload_base + index * positions_per_chunk
-            media_cache, next_offset = _append_open_pieces(
+            (
+                media_cache,
+                next_offset,
+                closure_position_ids,
+            ) = _append_open_pieces(
                 self.model,
                 pieces,
                 past_key_values=media_cache,
@@ -1322,9 +1156,17 @@ class OpenStreamReplayExtractor:
             closure_ids = pieces.closure_ids
         if closure_ids is None:
             raise RuntimeError("Open-stream replay did not produce closure tokens")
+        if closure_position_ids is None:
+            raise RuntimeError("Open-stream replay did not produce closure positions")
         if payload_base is None:
             raise RuntimeError("Open-stream replay did not establish a payload base")
-        return media_cache, next_offset, payload_base, closure_ids
+        return (
+            media_cache,
+            next_offset,
+            payload_base,
+            closure_ids,
+            closure_position_ids,
+        )
 
     def _extract_segment(
         self,
@@ -1337,6 +1179,7 @@ class OpenStreamReplayExtractor:
         verify_every: int,
         verify_continuous_every: int,
         global_step_offset: int,
+        progress: Any,
     ) -> tuple[list[float], list[torch.Tensor], list[dict[str, Any]]]:
         times = _time_grid(segment_start, segment_end, self.config.step_seconds)
         committed_cache = None
@@ -1373,6 +1216,7 @@ class OpenStreamReplayExtractor:
                     _,
                     payload_base,
                     _,
+                    _,
                 ) = self._cache_from_intervals(
                     video_path=video_path,
                     audio_path=audio_path,
@@ -1394,7 +1238,11 @@ class OpenStreamReplayExtractor:
             )
             payload_offset = payload_base + block_index * positions_per_chunk
             branch_cache = clone_cache(committed_cache)
-            media_cache, next_offset = _append_open_pieces(
+            (
+                media_cache,
+                _,
+                closure_position_ids,
+            ) = _append_open_pieces(
                 self.model,
                 pieces,
                 past_key_values=branch_cache,
@@ -1404,8 +1252,8 @@ class OpenStreamReplayExtractor:
             state = _read_open_state(
                 self.model,
                 media_cache,
-                next_rope_offset=next_offset,
                 closure_ids=pieces.closure_ids,
+                closure_position_ids=closure_position_ids,
                 query_ids=self.query_ids,
             )
             states.append(state.cpu())
@@ -1454,6 +1302,11 @@ class OpenStreamReplayExtractor:
                 committed_cache = media_cache
                 block_start = block_end
 
+            progress.set_postfix_str(
+                f"media={current_time:.1f}s", refresh=False
+            )
+            progress.update(1)
+
         return times, states, checks
 
     def extract(
@@ -1466,7 +1319,18 @@ class OpenStreamReplayExtractor:
         use_audio_in_video: bool,
         verify_every: int = 0,
         verify_continuous_every: int = 0,
+        show_progress: bool = True,
     ) -> ExtractionResult:
+        total_steps = len(_time_grid(start, end, self.config.step_seconds))
+        with tqdm(
+            total=total_steps,
+            desc="open-stream-replay",
+            unit="state",
+            dynamic_ncols=True,
+            mininterval=0.0,
+            miniters=1,
+            disable=not show_progress,
+        ) as progress:
             all_times, all_states, all_checks = self._extract_segment(
                 video_path=video_path,
                 audio_path=audio_path,
@@ -1476,6 +1340,7 @@ class OpenStreamReplayExtractor:
                 verify_every=verify_every,
                 verify_continuous_every=verify_continuous_every,
                 global_step_offset=0,
+                progress=progress,
             )
         return ExtractionResult(
             times=torch.tensor(all_times, dtype=torch.float64),
@@ -1485,7 +1350,7 @@ class OpenStreamReplayExtractor:
 
 
 def _block_intervals(start: float, end: float, block_seconds: float) -> list[tuple[float, float]]:
-    """Exact media-turn sequence represented by a commit/replay cache."""
+    """Exact AV payload intervals represented by an open-stream cache."""
 
     intervals = []
     cursor = start
@@ -1637,7 +1502,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--output", type=Path, required=True)
     parser.add_argument(
         "--mode",
-        choices=("causal-prefix", "commit-replay", "open-stream-replay"),
+        choices=("causal-prefix", "open-stream-replay"),
         default="open-stream-replay",
     )
     parser.add_argument("--model-id", default=MODEL_ID)
@@ -1661,6 +1526,11 @@ def parse_args() -> argparse.Namespace:
         default=0,
         help="Also compare block replay with one continuous causal media window.",
     )
+    parser.add_argument(
+        "--no-progress",
+        action="store_true",
+        help="Disable the per-state tqdm progress bar.",
+    )
     parser.add_argument("--head-checkpoint", type=Path)
     parser.add_argument(
         "--dtype",
@@ -1674,6 +1544,7 @@ def parse_args() -> argparse.Namespace:
 
 def main() -> None:
     args = parse_args()
+    _configure_qwen_logging()
     session_metadata = _apply_session_metadata(args)
     if args.video is None and args.audio is None:
         raise SystemExit("Pass --video, --audio, or both")
@@ -1683,7 +1554,6 @@ def main() -> None:
         raise SystemExit("--verify-every must be >= 0")
     if args.verify_continuous_every < 0:
         raise SystemExit("--verify-continuous-every must be >= 0")
-
     if args.num_steps is not None and args.num_steps <= 0:
         raise SystemExit("--num-steps must be positive")
     if args.target is not None:
@@ -1758,19 +1628,7 @@ def main() -> None:
             start=args.start,
             end=end,
             use_audio_in_video=args.use_audio_in_video,
-        )
+            show_progress=not args.no_progress,
-    elif args.mode == "commit-replay":
-        extractor = CommitReplayExtractor(
-            model, processor, config, fixed_grid_av=fixed_grid_av
-        )
-        result = extractor.extract(
-            video_path=args.video,
-            audio_path=args.audio,
-            start=args.start,
-            end=end,
-            use_audio_in_video=args.use_audio_in_video,
-            verify_every=args.verify_every,
-            verify_continuous_every=args.verify_continuous_every,
         )
     else:
         extractor = OpenStreamReplayExtractor(
@@ -1784,6 +1642,7 @@ def main() -> None:
             use_audio_in_video=args.use_audio_in_video,
             verify_every=args.verify_every,
             verify_continuous_every=args.verify_continuous_every,
+            show_progress=not args.no_progress,
         )
 
     payload: dict[str, Any] = {
@@ -1792,11 +1651,7 @@ def main() -> None:
         "stream_representation": (
             "single_causal_window"
             if args.mode == "causal-prefix"
-            else (
-                "sequence_of_2s_media_turns"
-                if args.mode == "commit-replay"
             else "single_open_media_turn_with_2s_payload_commits"
-            )
         ),
         "config": asdict(config),
         "start": args.start,
@@ -1808,7 +1663,7 @@ def main() -> None:
         "session_metadata": session_metadata,
         "input_video_fps": input_video_fps,
         "video_sampling_grid": (
-            "explicit_fixed_past_frames_with_duplication"
+            "explicit_fixed_half_open_past_frames_with_tail_duplication"
             if fixed_grid_av is not None
             else "qwen_utility_adaptive_sampling"
         ),

train_turn_head.py

@@ -0,0 +1,476 @@
+"""Train the frozen-Omni EasyCom anticipation head.
+
+The defaults follow EgoSpeak's EasyCom LSTR optimization config where the
+architectures are comparable: 50 epochs, Adam, 7e-5 learning rate, 5e-5 weight
+decay, linear warmup for 10 epochs, cosine decay, dropout 0.2, and MCE loss.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import random
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from torch.optim.lr_scheduler import LambdaLR
+from torch.utils.data import DataLoader, TensorDataset
+from tqdm.auto import tqdm
+
+from easycom_head import (
+    CLASS_NAMES,
+    NUM_CLASSES,
+    NUM_HORIZONS,
+    OFFICIAL_TRAIN_SESSIONS,
+    STEP_SECONDS,
+    TurnPredictionHead,
+    average_precision_metrics,
+    egospeak_multihorizon_loss,
+    load_session_data,
+    miniroad_multihorizon_loss,
+    normalize_session_name,
+    validate_disjoint_sessions,
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--feature-root", type=Path, required=True)
+    parser.add_argument("--target-root", type=Path, required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument(
+        "--sessions",
+        nargs="+",
+        default=list(OFFICIAL_TRAIN_SESSIONS),
+        help="Training sessions; default is the official nine-session split.",
+    )
+    parser.add_argument(
+        "--dev-sessions",
+        nargs="*",
+        default=[],
+        help=(
+            "Optional held-out training sessions used for checkpoint selection. "
+            "They are automatically removed from --sessions. Never use 1/2/3."
+        ),
+    )
+    parser.add_argument("--epochs", type=int, default=50)
+    parser.add_argument("--batch-size", type=int, default=512)
+    parser.add_argument("--learning-rate", type=float, default=7e-5)
+    parser.add_argument("--weight-decay", type=float, default=5e-5)
+    parser.add_argument("--warmup-epochs", type=float, default=10.0)
+    parser.add_argument("--warmup-factor", type=float, default=0.3)
+    parser.add_argument("--dropout", type=float, default=0.2)
+    parser.add_argument(
+        "--loss",
+        choices=("lstr-mce", "miniroad-normalized"),
+        default="lstr-mce",
+    )
+    parser.add_argument("--device", default="cuda:0")
+    parser.add_argument("--num-workers", type=int, default=0)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument(
+        "--expected-mode",
+        default="causal-prefix",
+        help="Expected feature mode, or 'any' to disable the check.",
+    )
+    parser.add_argument("--no-progress", action="store_true")
+    return parser.parse_args()
+
+
+def _validate_args(args: argparse.Namespace) -> None:
+    if args.epochs <= 0:
+        raise ValueError("--epochs must be positive")
+    if args.batch_size <= 0:
+        raise ValueError("--batch-size must be positive")
+    if args.learning_rate <= 0:
+        raise ValueError("--learning-rate must be positive")
+    if args.weight_decay < 0:
+        raise ValueError("--weight-decay cannot be negative")
+    if not 0 <= args.warmup_epochs <= args.epochs:
+        raise ValueError("--warmup-epochs must be in [0, epochs]")
+    if not 0 < args.warmup_factor <= 1:
+        raise ValueError("--warmup-factor must be in (0, 1]")
+    if not 0 <= args.dropout < 1:
+        raise ValueError("--dropout must be in [0, 1)")
+    if args.num_workers < 0:
+        raise ValueError("--num-workers cannot be negative")
+
+
+def _resolve_device(name: str) -> torch.device:
+    device = torch.device(name)
+    if device.type == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError(f"CUDA requested but unavailable: {name}")
+    return device
+
+
+def _seed_everything(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)
+
+
+def _load_examples(
+    feature_root: Path,
+    target_root: Path,
+    sessions: list[str],
+    expected_mode: str | None,
+) -> tuple[torch.Tensor, torch.Tensor, int, list[dict[str, Any]]]:
+    all_states: list[torch.Tensor] = []
+    all_targets: list[torch.Tensor] = []
+    metadata: list[dict[str, Any]] = []
+    hidden_size: int | None = None
+    for name in sessions:
+        data = load_session_data(
+            feature_root,
+            target_root,
+            name,
+            expected_mode=expected_mode,
+        )
+        states, targets, _ = data.future_examples(NUM_HORIZONS)
+        if hidden_size is None:
+            hidden_size = data.hidden_size
+        elif data.hidden_size != hidden_size:
+            raise ValueError(
+                f"Feature dimensions differ: {name} has {data.hidden_size}, "
+                f"expected {hidden_size}"
+            )
+        all_states.append(states)
+        all_targets.append(targets)
+        metadata.append(
+            {
+                "session": name,
+                "num_states": len(data.states),
+                "num_training_examples": len(states),
+                **data.feature_metadata,
+            }
+        )
+        print(
+            f"{name}: states={len(data.states)}, "
+            f"valid_10_horizon_examples={len(states)}"
+        )
+    if hidden_size is None:
+        raise ValueError("No sessions selected")
+    return (
+        torch.cat(all_states, dim=0),
+        torch.cat(all_targets, dim=0),
+        hidden_size,
+        metadata,
+    )
+
+
+def _scheduler(
+    optimizer: torch.optim.Optimizer,
+    *,
+    steps_per_epoch: int,
+    epochs: int,
+    warmup_epochs: float,
+    warmup_factor: float,
+) -> LambdaLR:
+    total_steps = max(1, steps_per_epoch * epochs)
+    warmup_steps = steps_per_epoch * warmup_epochs
+
+    def multiplier(step: int) -> float:
+        cosine = 0.5 * (1.0 + math.cos(math.pi * step / total_steps))
+        if warmup_steps > 0 and step < warmup_steps:
+            alpha = step / warmup_steps
+            warmup = warmup_factor * (1.0 - alpha) + alpha
+        else:
+            warmup = 1.0
+        return warmup * cosine
+
+    return LambdaLR(optimizer, multiplier)
+
+
+@torch.inference_mode()
+def _evaluate_valid_horizons(
+    head: TurnPredictionHead,
+    states: torch.Tensor,
+    targets: torch.Tensor,
+    device: torch.device,
+    batch_size: int,
+) -> dict[str, Any]:
+    head.eval()
+    chunks: list[torch.Tensor] = []
+    for start in range(0, len(states), batch_size):
+        x = states[start : start + batch_size].to(
+            device=device,
+            dtype=torch.float32,
+            non_blocking=True,
+        )
+        chunks.append(head(x).softmax(dim=-1).cpu())
+    probabilities = torch.cat(chunks, dim=0).numpy()
+    target_array = targets.numpy()
+
+    results = []
+    for horizon in range(NUM_HORIZONS):
+        result = average_precision_metrics(
+            probabilities[:, horizon],
+            target_array[:, horizon],
+        )
+        result["horizon"] = horizon + 1
+        result["seconds"] = (horizon + 1) * STEP_SECONDS
+        results.append(result)
+    return {
+        "horizons": results,
+        "mean_map": float(np.mean([item["map"] for item in results])),
+    }
+
+
+def _checkpoint(
+    head: TurnPredictionHead,
+    *,
+    epoch: int,
+    best_epoch: int,
+    best_dev_map: float | None,
+    train_loss: float,
+    args: argparse.Namespace,
+    train_sessions: list[str],
+    dev_sessions: list[str],
+    hidden_size: int,
+    feature_metadata: list[dict[str, Any]],
+    history: list[dict[str, Any]],
+) -> dict[str, Any]:
+    return {
+        "format_version": 1,
+        "head": head.state_dict(),
+        "head_config": {
+            "hidden_size": hidden_size,
+            "num_horizons": NUM_HORIZONS,
+            "num_classes": NUM_CLASSES,
+            "dropout": args.dropout,
+        },
+        "alignment": {
+            "state_label_offset": 0,
+            "future_offsets": list(range(1, NUM_HORIZONS + 1)),
+            "future_seconds": [
+                (index + 1) * STEP_SECONDS for index in range(NUM_HORIZONS)
+            ],
+        },
+        "classes": list(CLASS_NAMES),
+        "epoch": epoch,
+        "best_epoch": best_epoch,
+        "best_dev_map": best_dev_map,
+        "train_loss": train_loss,
+        "train_sessions": train_sessions,
+        "dev_sessions": dev_sessions,
+        "feature_metadata": feature_metadata,
+        "training_config": {
+            "epochs": args.epochs,
+            "batch_size": args.batch_size,
+            "learning_rate": args.learning_rate,
+            "weight_decay": args.weight_decay,
+            "warmup_epochs": args.warmup_epochs,
+            "warmup_factor": args.warmup_factor,
+            "loss": args.loss,
+            "optimizer": "Adam",
+            "scheduler": "linear_warmup_times_cosine",
+            "seed": args.seed,
+            "expected_mode": args.expected_mode,
+        },
+        "history": history,
+    }
+
+
+def main() -> None:
+    args = parse_args()
+    _validate_args(args)
+    _seed_everything(args.seed)
+    device = _resolve_device(args.device)
+    expected_mode = None if args.expected_mode == "any" else args.expected_mode
+
+    requested_train = [normalize_session_name(item) for item in args.sessions]
+    dev_sessions = [normalize_session_name(item) for item in args.dev_sessions]
+    train_sessions = [item for item in requested_train if item not in dev_sessions]
+    validate_disjoint_sessions(train_sessions, dev_sessions)
+    if not train_sessions:
+        raise ValueError("No training sessions remain after removing dev sessions")
+    forbidden_test = {"Session_1", "Session_2", "Session_3"}
+    leaked = forbidden_test & (set(train_sessions) | set(dev_sessions))
+    if leaked:
+        raise ValueError(
+            f"Official test sessions cannot be used for training/selection: "
+            f"{sorted(leaked)}"
+        )
+
+    print(f"device: {device}")
+    print(f"train sessions: {', '.join(train_sessions)}")
+    print(f"dev sessions: {', '.join(dev_sessions) if dev_sessions else 'none'}")
+    states, targets, hidden_size, feature_metadata = _load_examples(
+        args.feature_root,
+        args.target_root,
+        train_sessions,
+        expected_mode,
+    )
+    print(f"training examples: {len(states)}, feature_dim={hidden_size}")
+    dev_states: torch.Tensor | None = None
+    dev_targets: torch.Tensor | None = None
+    if dev_sessions:
+        dev_states, dev_targets, dev_hidden_size, _ = _load_examples(
+            args.feature_root,
+            args.target_root,
+            dev_sessions,
+            expected_mode,
+        )
+        if dev_hidden_size != hidden_size:
+            raise ValueError(
+                f"Dev feature dimension {dev_hidden_size} differs from "
+                f"training dimension {hidden_size}"
+            )
+        print(f"development examples: {len(dev_states)}")
+
+    dataset = TensorDataset(states, targets)
+    generator = torch.Generator().manual_seed(args.seed)
+    loader = DataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=args.num_workers,
+        pin_memory=device.type == "cuda",
+        persistent_workers=args.num_workers > 0,
+        generator=generator,
+    )
+    head = TurnPredictionHead(
+        hidden_size=hidden_size,
+        num_horizons=NUM_HORIZONS,
+        num_classes=NUM_CLASSES,
+        dropout=args.dropout,
+    ).to(device)
+    optimizer = torch.optim.Adam(
+        head.parameters(),
+        lr=args.learning_rate,
+        weight_decay=args.weight_decay,
+    )
+    scheduler = _scheduler(
+        optimizer,
+        steps_per_epoch=len(loader),
+        epochs=args.epochs,
+        warmup_epochs=args.warmup_epochs,
+        warmup_factor=args.warmup_factor,
+    )
+    loss_function = (
+        egospeak_multihorizon_loss
+        if args.loss == "lstr-mce"
+        else miniroad_multihorizon_loss
+    )
+
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    history: list[dict[str, Any]] = []
+    best_dev_map = -math.inf if dev_sessions else None
+    best_epoch = args.epochs if not dev_sessions else 0
+    for epoch in range(1, args.epochs + 1):
+        head.train()
+        loss_sum = 0.0
+        example_count = 0
+        progress = tqdm(
+            loader,
+            desc=f"head epoch {epoch}/{args.epochs}",
+            unit="batch",
+            disable=args.no_progress,
+        )
+        for state_batch, target_batch in progress:
+            state_batch = state_batch.to(
+                device=device,
+                dtype=torch.float32,
+                non_blocking=True,
+            )
+            target_batch = target_batch.to(
+                device=device,
+                dtype=torch.float32,
+                non_blocking=True,
+            )
+            optimizer.zero_grad(set_to_none=True)
+            logits = head(state_batch)
+            loss = loss_function(logits, target_batch)
+            loss.backward()
+            optimizer.step()
+            scheduler.step()
+            batch_count = len(state_batch)
+            loss_sum += float(loss.detach()) * batch_count
+            example_count += batch_count
+            progress.set_postfix(
+                loss=f"{float(loss.detach()):.5f}",
+                lr=f"{scheduler.get_last_lr()[0]:.3g}",
+            )
+        train_loss = loss_sum / example_count
+        record: dict[str, Any] = {
+            "epoch": epoch,
+            "train_loss": train_loss,
+            "learning_rate": scheduler.get_last_lr()[0],
+        }
+
+        if dev_sessions:
+            assert dev_states is not None and dev_targets is not None
+            dev_metrics = _evaluate_valid_horizons(
+                head,
+                dev_states,
+                dev_targets,
+                device,
+                args.batch_size,
+            )
+            dev_map = float(dev_metrics["mean_map"])
+            record["dev_valid_mean_map"] = dev_map
+            print(
+                f"epoch {epoch:02d}: train_loss={train_loss:.6f}, "
+                f"dev_valid_mean_mAP={dev_map * 100:.3f}"
+            )
+            if dev_map > float(best_dev_map):
+                best_dev_map = dev_map
+                best_epoch = epoch
+                torch.save(
+                    _checkpoint(
+                        head,
+                        epoch=epoch,
+                        best_epoch=best_epoch,
+                        best_dev_map=best_dev_map,
+                        train_loss=train_loss,
+                        args=args,
+                        train_sessions=train_sessions,
+                        dev_sessions=dev_sessions,
+                        hidden_size=hidden_size,
+                        feature_metadata=feature_metadata,
+                        history=history + [record],
+                    ),
+                    args.output,
+                )
+                print(f"saved new best checkpoint: {args.output}")
+        else:
+            print(f"epoch {epoch:02d}: train_loss={train_loss:.6f}")
+        history.append(record)
+
+    if not dev_sessions:
+        torch.save(
+            _checkpoint(
+                head,
+                epoch=args.epochs,
+                best_epoch=best_epoch,
+                best_dev_map=None,
+                train_loss=history[-1]["train_loss"],
+                args=args,
+                train_sessions=train_sessions,
+                dev_sessions=dev_sessions,
+                hidden_size=hidden_size,
+                feature_metadata=feature_metadata,
+                history=history,
+            ),
+            args.output,
+        )
+    print(f"checkpoint: {args.output}")
+    if dev_sessions:
+        print(
+            f"best epoch: {best_epoch}, "
+            f"dev valid-only mean mAP: {float(best_dev_map) * 100:.3f}"
+        )
+    summary_path = args.output.with_suffix(args.output.suffix + ".json")
+    payload = torch.load(args.output, map_location="cpu", weights_only=False)
+    summary = {key: value for key, value in payload.items() if key != "head"}
+    summary_path.write_text(json.dumps(summary, indent=2) + "\n")
+    print(f"training summary: {summary_path}")
+
+
+if __name__ == "__main__":
+    main()