Add causal person track GRU residual experiment

.gitignore

@@ -11,6 +11,7 @@ __pycache__/
 /manifests/
 /outputs/
 /target_frame/
+/third_party/EgoSpeak/
 
 # Local editors and environments.
 .idea/

PROJECT_CONTEXT.md

@@ -43,17 +43,27 @@ result.
 - Ego4D batch and 8-GPU local-feature launchers, with strict frame audits.
 - Timing-derived wearer floor-state and event diagnostics.
 
+### Measured screening result
+
+- F1 Frame-set Attention uses fused G+A as one query over the latest half-open
+  causal person set. On the existing three-seed, four-epoch test screen it
+  obtained `81.77745 +/- 0.15638` 0.2-second mAP and
+  `64.95099 +/- 0.09451` ten-horizon mean mAP. This is stable, but is not a
+  matched improvement over the historical no-dev pooled result; correct/wrong
+  visual controls remain missing.
+
 ### Implemented, awaiting remote runtime verification
 
-- F1 Frame-set Attention loader/head/train/infer path. It uses fused G+A as
-  one query over the latest half-open causal person set, retains direct A in
-  the head, and dynamically pads people per batch. Do not report it as a
-  measured method until the L20 smoke and full controls pass.
 - Two current-sidecar screening controls: `audio-person-frame-set` adds a
   zero-gated audio-to-person read before the F1 query, and
   `gated-person-state-adapter` adds a zero-gated person residual after the
   frozen Thinker readout. They are explicitly recorded as F5/F6 screening,
   not full history-conditioned F5 or an internal Thinker F6 adapter.
+- `track-gru-residual` builds strictly causal histories from anonymous provided
+  track IDs, applies one shared GRU per track, lets G+A query current track
+  states, and adds a zero-gated correction to an exactly initialized pooled
+  `G+A+V` checkpoint. See
+  [`docs/easycom_track_gru_residual.md`](docs/easycom_track_gru_residual.md).
 
 ### Current blocker
 
@@ -69,12 +79,17 @@ semantic event ground truth.
 
 ### Next experiment gates
 
+The user has temporarily chosen to run the dense-AP model screen while P0 is
+unfinished. This does not relax the paper claim: event numbers remain proxy
+diagnostics until competing-risk censoring is fixed.
+
 1. **P0:** fix/freeze event attribution, splits and scenario definitions.
 2. **P1:** probe aggregate activity and per-person active-speaker grounding;
    compare correct, wrong-person and time-shifted visual inputs.
 3. **P2:** re-evaluate `G`, `G+A`, `G+A+pooled-V` under the frozen event protocol.
-4. **P3:** begin the progressive person-fusion route with Frame-set Attention,
-   then add spatial or track-temporal structure only after its controls pass.
+4. **P3:** run the matched pooled/anonymous-track-GRU residual screen, then add
+   broken-ID, wrong-time and remove-person controls before claiming that track
+   continuity matters.
 5. **P4:** attempt person-memory adapter/Thinker LoRA only if P1/P3 pass.
 
 The detailed architecture, controls and stopping gates for P3 are frozen in
@@ -227,6 +242,7 @@ Do not reorganize them physically without compatibility wrappers.
 | `infer_turn_head.py` | Save EgoSpeak-compatible predictions and AP/mAP |
 | `evaluate_floor_state_changes.py` | CPU floor-state/event proxy evaluation |
 | `summarize_local_baselines.py` | Aggregate local baseline `metrics.json` files |
+| `run_easycom_track_gru_residual.sh` | Matched pooled and anonymous track-GRU training/inference launcher |
 | `train_omni_turn_lora.py` | Thinker LoRA training from cached AV |
 | `eval_omni_turn_lora.py` | Thinker LoRA evaluation |
 

README.md

@@ -1899,6 +1899,31 @@ acceptance rule using Sessions 1--3.  LLM-adjudicated candidates are deliberatel
 kept in a separate JSONL until pilot precision is accepted; they are not yet
 silently merged into the training arrays.
 
+## Anonymous person-track GRU residual
+
+The first person-temporal experiment is implemented as
+`--local-fusion track-gru-residual`. It preserves an exactly initialized
+`G+A+pooled-V` head and learns a zero-gated correction from shared causal GRU
+states over temporary person tracks. IDs route history only; they are not model
+inputs. Source-clip boundaries, gaps over the configured maximum and unknown
+IDs reset history. Existing local Qwen sidecars can be reused without
+re-extraction.
+
+Run the matched pooled/track development protocol with:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 bash run_easycom_track_gru_residual.sh all 0
+```
+
+Use separate GPUs for additional seeds, then summarize the common result root
+with `summarize_local_baselines.py --paired-deltas`. After method selection is
+frozen, `bash run_easycom_track_gru_residual.sh test 0 1 2` evaluates the saved
+checkpoints once on Sessions 1--3. The exact causal contract,
+EgoSpeak code comparison, L20 smoke command and stopping rule are in
+[`docs/easycom_track_gru_residual.md`](docs/easycom_track_gru_residual.md).
+This is a frozen-backbone side branch, not an internal Thinker modification or
+an EgoSpeak MiniROAD reproduction.
+
 ## Candidate improvement: explicit working and long memory
 
 Motivated by EgoSpeak's separation of long-term and working memory, test an

docs/README.md

@@ -23,6 +23,7 @@ tests.
 | [`easycom_encoder_only.md`](easycom_encoder_only.md) | Comparing EasyCom Qwen encoder-only A/V/AV representations |
 | [`ego4d_encoder_only.md`](ego4d_encoder_only.md) | Preparing and running the Ego4D encoder-only comparison |
 | [`ego4d_local_baselines.md`](ego4d_local_baselines.md) | Previewing Ego4D boxes, extracting local features and running 8-GPU jobs |
+| [`easycom_track_gru_residual.md`](easycom_track_gru_residual.md) | Matched pooled-to-track GRU experiment, causal identity contract and L20 commands |
 
 ## Design history and implementation options
 

docs/easycom_track_gru_residual.md

@@ -0,0 +1,158 @@
+# EasyCom anonymous track-GRU residual
+
+## Decision
+
+The first track-temporal experiment keeps the strongest pooled local baseline
+as an exact, optionally frozen base and learns only a causal person-track
+correction:
+
+```text
+pooled base:  G_t + A_t + pooled-V_t -> concat-MLP -> base logits
+
+per visible person i:
+  past crop observations on the same temporary track
+      -> shared projection -> shared one-layer GRU -> H_t,i
+
+G_t + A_t -> masked attention over {NULL, H_t,1, ..., H_t,N}
+          -> zero-gated residual logits
+
+final logits = base logits + tanh(alpha) * person_valid * residual logits
+alpha starts at exactly zero
+```
+
+This is the smallest experiment that asks whether anonymous visual continuity
+adds information beyond `G+A+pooled-V`. It is not yet person-conditioned audio
+grounding and does not modify the Omni Thinker.
+
+## What was taken from EgoSpeak
+
+The official EgoSpeak repository is cloned for inspection at
+`third_party/EgoSpeak` and excluded from Git. The inspected revision is
+`b5bee35e796d8f2d6aa7da2bf1b07836270f8bfb`.
+
+Its `MiniROADA` model projects concatenated RGB/audio features, applies a
+single-layer GRU, expands the recurrent state into ten anticipation states,
+and uses a shared classifier. The EasyCom configuration uses a 128-step
+window, stride 4, 2048D embedding, 1024D GRU and ten future steps.
+
+We reuse the defensible ideas—causal recurrence, shared weights and a ten-step
+anticipation output—but not the implementation. Our GRU is shared across
+anonymous person tracks, reads only crop observations at or before the current
+cutoff, and corrects the already measured pooled Omni baseline. Consequently,
+this method must be described as an EgoSpeak-inspired track-temporal adapter,
+not an EgoSpeak GRU reproduction.
+
+## Causal and identity contract
+
+- A 5 Hz row with cutoff `(t+1)/5` reads the latest local video frame strictly
+  before that cutoff.
+- Track history contains at most `--person-history-frames` real 2 FPS crop
+  observations and is ordered oldest to newest.
+- Temporary IDs only link observations. They are never embedded or passed to
+  the model, so global ID renaming cannot change a prediction.
+- A source-clip boundary always resets a track. A time gap larger than
+  `--person-track-max-gap-seconds` also resets it.
+- Missing, `-1`, `None`, `null` and `unknown` IDs are isolated observations.
+- Duplicate known IDs in one frame fail instead of being silently merged.
+- Empty person sets produce exactly zero visual correction.
+
+The existing EasyCom local sidecars already contain crop features, frame times,
+source-clip indices and person IDs; this experiment does not require feature
+re-extraction.
+
+## Fair initialization
+
+`track-gru-residual` requires `--initialize-pooled-checkpoint`. The loader
+checks dataset, train/dev split, loss, layer fusion, dimensions and local
+feature contract before copying the pooled path. By default those copied
+parameters are frozen. The residual gate starts at zero, so before optimization
+the track model must reproduce the pooled logits exactly, element for element.
+
+The matched screening protocol is:
+
+- EasyCom Session 4--11: training;
+- Session 12: development/model selection;
+- Session 1--3: final test, run only after history/gap and the method are frozen;
+- 4 epochs, three seeds, identical optimizer/loss for pooled and track runs;
+- 16 observations at 2 FPS (up to 8 seconds), 1-second maximum link gap.
+
+This protocol intentionally re-runs the pooled base. Historical no-dev results
+cannot be used as the matched control.
+
+## Run
+
+After pulling on the L20 machine, run the focused CPU regression tests first:
+
+```bash
+python -m unittest \
+  test_local_turn_features.LocalTests.test_track_loader_builds_causal_histories_and_resets_tracks \
+  test_local_turn_features.LocalTests.test_track_gru_residual_starts_from_exact_pooled_logits \
+  test_local_turn_features.RoundTripTests.test_train_infer_easycom_and_ego4d
+```
+
+One-seed runtime smoke on an L20:
+
+```bash
+EPOCHS=1 \
+SESSIONS="4 5" \
+DEV_SESSIONS=5 \
+SCREEN_SESSIONS=5 \
+CHECKPOINT_ROOT=checkpoints/easycom_track_gru_smoke \
+RESULT_ROOT=results/easycom_track_gru_smoke \
+CUDA_VISIBLE_DEVICES=0 \
+bash run_easycom_track_gru_residual.sh all 0
+```
+
+Then use three GPUs in three foreground terminals. These runs report Session
+12 development scores, not the official test set:
+
+```bash
+SUMMARIZE=0 CUDA_VISIBLE_DEVICES=0 bash run_easycom_track_gru_residual.sh all 0
+SUMMARIZE=0 CUDA_VISIBLE_DEVICES=1 bash run_easycom_track_gru_residual.sh all 1
+SUMMARIZE=0 CUDA_VISIBLE_DEVICES=2 bash run_easycom_track_gru_residual.sh all 2
+```
+
+After all development jobs finish:
+
+```bash
+python summarize_local_baselines.py \
+  --root results/easycom_track_gru_residual \
+  --dataset easycom \
+  --protocol causal_valid_only \
+  --paired-deltas
+```
+
+Only after choosing and freezing the method, evaluate the saved checkpoints
+once on Sessions 1--3:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 bash run_easycom_track_gru_residual.sh test 0 1 2
+```
+
+This writes and summarizes `results/easycom_track_gru_residual_test` without
+retraining. A single GPU is sufficient because head inference is inexpensive.
+
+Paths and hyperparameters can be overridden with the environment variables at
+the top of `run_easycom_track_gru_residual.sh`. The launcher runs in the
+foreground and refuses to overwrite an existing checkpoint or result.
+
+## Interpretation and stopping rule
+
+The initial screen answers only whether track-temporal visual structure can
+improve the dense EgoSpeak-compatible metric. Advance this branch only if
+same-seed `track_gru - pooled` deltas are directionally consistent on Session
+12 without sacrificing short-horizon wearer AP. The official test is a final
+confirmation, not another model-selection stage.
+Because prior F1 variation was roughly a few tenths of a percentage point,
+tiny single-seed changes are not evidence.
+
+If the track branch wins, next run correct-track versus broken-track,
+past-shifted and no-ID temporal controls, followed by person-conditioned audio
+grounding. If it does not win, first compare 8/16/32 observations and optionally
+unfreeze the pooled base using Session 12 only; do not add deeper cross-attention
+merely to rescue a null result.
+
+The paper-level claim still requires the frozen competing-risk event evaluator
+and matched-recall false-entry analysis. Dense AP/mAP improvement alone is a
+representation result, not proof that the model understands multi-person floor
+behavior.

docs/progressive_feature_fusion_route_2026-09-02.md

@@ -21,11 +21,14 @@
 融合实验从 F1 Frame-set Attention 开始。后续模块只有在前一级通过预先定义的 gate 后
 才实施。
 
-实现状态（2026-09-02）：F1 的 loader、head、train/infer 路径和回归测试已经加入代码，
-但尚未在远程 PyTorch/L20 环境完成 runtime smoke，不能写成已验证结果。另加入两个仅用
+实现状态（更新于 2026-09-07）：F1 已完成三 seed、四 epoch 的远程筛选，得到
+`81.77745 +/- 0.15638` 的 0.2 秒 mAP 和 `64.95099 +/- 0.09451` 的十 horizon mean
+mAP。它表现稳定，但与旧 pooled 结果的 split/dev 协议并不匹配，而且尚无 correct/wrong
+视觉控制，因此不能写成注意力优于 pooling。另加入两个仅用
 当前 sidecar 的 F5/F6 **screening** 模式：`audio-person-frame-set` 和
 `gated-person-state-adapter`。前者没有 F4 人物历史，后者位于 Thinker readout 之后；二者
-都不能写成完整 F5/F6。F2--F4、真正的 F5/F6 和 D1 仍然只是计划。
+都不能写成完整 F5/F6。第一版 F3 `track-gru-residual` 已实现但尚待远程运行；F2、F4、
+真正的 F5/F6 和 D1 仍然只是计划。
 
 ## 2. 不可改变的实验条件
 
@@ -193,6 +196,14 @@ floor-action 结果。
 对照包括 frame-set、无 ID temporal pool、正确 track、打断连续性的 track 以及 past-shift。
 正确 track 必须优于无 ID 和打断 track，才能说明身份连续性有价值。
 
+第一版实现为 `track-gru-residual`。它没有直接替换 pooled 头，而是从同 seed、同 split、
+同 local contract 的 `G+A+pooled-V` checkpoint 精确复制并默认冻结 base path，再增加共享
+track GRU、G+A 对 track states 的 masked attention 以及零初始化 scalar gate。这样训练
+开始前 logits 与 pooled baseline 逐元素相同，空人物时 correction 恒为零，实验首先回答
+“track continuity 是否提供增量信息”，而不会把随机重训 base 的波动误当作 GRU 收益。
+实现、契约和命令见
+[`easycom_track_gru_residual.md`](easycom_track_gru_residual.md)。
+
 ## 8. F4：Person-conditioned audio/Omni history
 
 F4 让每个匿名人物分别读取已有的 5 Hz 因果历史：
@@ -268,11 +279,11 @@ README 必须写明 split、seed、checkpoint、输入特征版本、评测协
 
 ## 12. 最近执行顺序
 
-1. 完成并测试 P0 competing-risk censoring，人工打印至少一条 response 和 yield 时间轴。
-2. 扩展 local sidecar loader，使其在不破坏现有 pooled 加载路径的情况下返回逐帧人物
-   feature、offset 和 mask；测试成功、失败、空人物和越界路径。
-3. 实现 F1 最小 Frame-set Attention，不加入 bbox 或 track temporal。
-4. 单 session/短片 smoke：检查 shape、cutoff、chosen frame time、mask、空人物和排列
-   不变性。
-5. 在固定 train-internal dev 上运行 B1、B3、F1-a、F1-c、F1-d、F1-e 和 F1-g。
-6. F1 通过后才进入 F2；若未通过，优先审计人物 crop 和 AV 对齐，不继续堆模型。
+1. 用户当前选择先做 dense-AP model screen；运行 matched pooled 与
+   `track-gru-residual`，固定 Session 12 dev，不用测试集选历史长度。
+2. 若 track GRU 有稳定增益，补 broken-track、past-shift、no-ID 和 remove-person 控制；若
+   无增益，只在 dev 上比较 8/16/32 observations 后停止该分支。
+3. correct track 通过控制后，再实现 F4 person-conditioned audio/Omni history；否则回到
+   对齐审计，不直接堆更深 interaction。
+4. 在形成论文事件结论前，仍须完成 P0 competing-risk censoring，并人工打印 response 与
+   yield 时间轴。当前事件结果只能称为 timing-derived proxy。

easycom_head.py

@@ -61,17 +61,23 @@ LAYER_FUSION_MODES = (
     "scalar-mixture",
     "concat-mlp",
 )
-PERSON_SET_FUSION_MODES = (
+CURRENT_FRAME_PERSON_SET_FUSION_MODES = (
     "frame-set-attention",
     "audio-person-frame-set",
     "gated-person-state-adapter",
 )
+TRACK_PERSON_FUSION_MODES = ("track-gru-residual",)
+PERSON_SET_FUSION_MODES = (
+    *CURRENT_FRAME_PERSON_SET_FUSION_MODES,
+    *TRACK_PERSON_FUSION_MODES,
+)
 LOCAL_FUSION_MODES = ("pooled", *PERSON_SET_FUSION_MODES)
 LOCAL_FUSION_SCOPES = {
     "pooled": "pooled-baseline",
     "frame-set-attention": "f1-current-frame-person-set",
     "audio-person-frame-set": "f5-screening-current-frame-no-person-history",
     "gated-person-state-adapter": "f6-screening-post-thinker-not-internal",
+    "track-gru-residual": "causal-anonymous-track-gru-pooled-baseline-residual",
 }
 DatasetName = Literal["easycom", "ego4d"]
 
@@ -95,6 +101,8 @@ class TurnPredictionHead(nn.Module):
         local_fusion: str = "pooled",
         person_attention_dim: int = 256,
         person_attention_heads: int = 4,
+        person_history_frames: int = 16,
+        person_track_max_gap_seconds: float = 1.0,
     ) -> None:
         super().__init__()
         if num_input_layers <= 0:
@@ -143,6 +151,10 @@ class TurnPredictionHead(nn.Module):
             raise ValueError("Person attention dimension/heads must be positive")
         if person_attention_dim % person_attention_heads:
             raise ValueError("Person attention dimension must be divisible by heads")
+        if person_history_frames <= 0:
+            raise ValueError("person_history_frames must be positive")
+        if person_track_max_gap_seconds <= 0:
+            raise ValueError("person_track_max_gap_seconds must be positive")
         if local_fusion in PERSON_SET_FUSION_MODES and (
             local_visual_dim <= 0 or local_audio_dim <= 0
         ):
@@ -152,6 +164,8 @@ class TurnPredictionHead(nn.Module):
         self.local_fusion = local_fusion
         self.person_attention_dim = person_attention_dim
         self.person_attention_heads = person_attention_heads
+        self.person_history_frames = person_history_frames
+        self.person_track_max_gap_seconds = person_track_max_gap_seconds
         if self.learnable_layer_mixture:
             # Zero initialization exactly reproduces a uniform layer mean at
             # step zero while allowing training to select useful query layers.
@@ -171,7 +185,7 @@ class TurnPredictionHead(nn.Module):
         fused_size = hidden_size
         projected_modalities = (
             (("audio", local_audio_dim),)
-            if local_fusion in PERSON_SET_FUSION_MODES
+            if local_fusion in CURRENT_FRAME_PERSON_SET_FUSION_MODES
             else (("visual", local_visual_dim), ("audio", local_audio_dim))
         )
         for name, dimension in projected_modalities:
@@ -195,6 +209,9 @@ class TurnPredictionHead(nn.Module):
         self.audio_person_gate: nn.Parameter | None = None
         self.person_state_adapter: nn.Module | None = None
         self.person_state_adapter_gate: nn.Parameter | None = None
+        self.person_track_gru: nn.GRU | None = None
+        self.person_track_delta: nn.Module | None = None
+        self.person_track_residual_gate: nn.Parameter | None = None
         if local_fusion in PERSON_SET_FUSION_MODES:
             self.person_input_projection = nn.Sequential(
                 nn.LayerNorm(local_visual_dim),
@@ -244,6 +261,25 @@ class TurnPredictionHead(nn.Module):
                     nn.Linear(person_attention_dim, hidden_size),
                 )
                 self.person_state_adapter_gate = nn.Parameter(torch.zeros(()))
+            elif local_fusion == "track-gru-residual":
+                # Apply one shared causal GRU independently to every anonymous
+                # face track. The regular pooled G+A+V head remains the exact
+                # base path; this branch can only add a zero-gated correction.
+                self.person_track_gru = nn.GRU(
+                    person_attention_dim,
+                    person_attention_dim,
+                    num_layers=1,
+                    batch_first=True,
+                )
+                delta_input = hidden_size // 2 + person_attention_dim + 1
+                self.person_track_delta = nn.Sequential(
+                    nn.LayerNorm(delta_input),
+                    nn.Linear(delta_input, hidden_size // 2),
+                    nn.GELU(),
+                    nn.Dropout(dropout),
+                    nn.Linear(hidden_size // 2, num_horizons * num_classes),
+                )
+                self.person_track_residual_gate = nn.Parameter(torch.zeros(()))
             else:
                 fused_size += person_attention_dim + 1
         self.net = nn.Sequential(
@@ -291,7 +327,7 @@ class TurnPredictionHead(nn.Module):
 
     @property
     def local_input_size(self) -> int:
-        if self.local_fusion in PERSON_SET_FUSION_MODES:
+        if self.local_fusion in CURRENT_FRAME_PERSON_SET_FUSION_MODES:
             return self.local_audio_dim + 1
         return sum(d + 1 for d in (self.local_visual_dim, self.local_audio_dim) if d)
 
@@ -369,6 +405,95 @@ class TurnPredictionHead(nn.Module):
         person_valid = person_mask.any(dim=1, keepdim=True).to(states.dtype)
         return audio_projected, audio_valid, context[:, 0, :], person_valid
 
+    def _track_gru_context(
+        self,
+        states: torch.Tensor,
+        local: torch.Tensor,
+        person_histories: torch.Tensor | None,
+        history_mask: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Encode each causal track with a shared GRU, then query the set."""
+
+        if states.ndim != 2:
+            raise ValueError("Track GRU fusion expects batched [B,D] states")
+        if local.shape != (len(states), self.local_input_size):
+            raise ValueError("Track GRU fusion requires pooled [visual,audio,valid] input")
+        if (
+            person_histories is None
+            or person_histories.ndim != 4
+            or person_histories.shape[0] != len(states)
+            or person_histories.shape[2] != self.person_history_frames
+            or person_histories.shape[3] != self.local_visual_dim
+        ):
+            raise ValueError(
+                "Track GRU fusion requires person histories "
+                "[B,P,person_history_frames,visual_dim]"
+            )
+        if (
+            history_mask is None
+            or history_mask.dtype != torch.bool
+            or history_mask.shape != person_histories.shape[:3]
+        ):
+            raise ValueError("Track GRU fusion requires a boolean history mask [B,P,L]")
+        if not torch.isfinite(person_histories).all():
+            raise ValueError("Person track histories contain NaN/Inf")
+        if bool((history_mask[:, :, 1:] & ~history_mask[:, :, :-1]).any()):
+            raise ValueError("Person track histories must be left aligned")
+        assert self.person_input_projection is not None
+        assert self.person_state_query is not None
+        assert self.person_audio_query is not None
+        assert self.person_attention is not None
+        assert self.null_person is not None
+        assert self.person_track_gru is not None
+
+        batch, people, length, _ = person_histories.shape
+        person_valid = history_mask.any(dim=-1)
+        if people:
+            projected = self.person_input_projection(person_histories)
+            flat = projected.reshape(batch * people, length, -1)
+            lengths = history_mask.sum(dim=-1).reshape(-1)
+            safe_lengths = lengths.clamp_min(1).to(device="cpu")
+            packed = nn.utils.rnn.pack_padded_sequence(
+                flat,
+                safe_lengths,
+                batch_first=True,
+                enforce_sorted=False,
+            )
+            _, hidden = self.person_track_gru(packed)
+            track_states = hidden[-1].reshape(batch, people, -1)
+            track_states = track_states * person_valid[..., None]
+        else:
+            track_states = states.new_zeros(
+                (batch, 0, self.person_attention_dim)
+            )
+
+        visual_width = self.local_visual_dim + 1
+        audio_values = local[:, visual_width : visual_width + self.local_audio_dim]
+        audio_valid = local[:, visual_width + self.local_audio_dim :]
+        if audio_valid.shape != (batch, 1):
+            raise AssertionError("Unexpected pooled AV local feature layout")
+        query = self.person_state_query(states)
+        query = query + self.person_audio_query(audio_values) * audio_valid
+        memory = torch.cat(
+            (self.null_person.expand(batch, -1, -1), track_states), dim=1
+        )
+        memory_valid = torch.cat(
+            (
+                torch.ones((batch, 1), dtype=torch.bool, device=states.device),
+                person_valid,
+            ),
+            dim=1,
+        )
+        context, _ = self.person_attention(
+            query[:, None, :],
+            memory,
+            memory,
+            key_padding_mask=~memory_valid,
+            need_weights=False,
+        )
+        any_person = person_valid.any(dim=1, keepdim=True).to(states.dtype)
+        return context[:, 0, :] * any_person, any_person
+
     def _shared_features(
         self,
         states: torch.Tensor,
@@ -379,7 +504,7 @@ class TurnPredictionHead(nn.Module):
         """Return the activity MLP representation shared by both readouts."""
 
         states = self._fuse_layers(states)
-        if self.local_fusion in PERSON_SET_FUSION_MODES:
+        if self.local_fusion in CURRENT_FRAME_PERSON_SET_FUSION_MODES:
             if local is None:
                 raise ValueError("Person-set fusion requires local audio features")
             audio, audio_valid, person_context, person_valid = self._frame_set_context(
@@ -428,6 +553,42 @@ class TurnPredictionHead(nn.Module):
             raise ValueError("Person sets supplied to a non-frame-set head")
         return self.net[:-1](states)
 
+    def _activity_logits(
+        self,
+        states: torch.Tensor,
+        local: torch.Tensor | None,
+        person_features: torch.Tensor | None,
+        person_mask: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        features = self._shared_features(
+            states,
+            local,
+            person_features,
+            person_mask,
+        )
+        flat_logits = self.net[-1](features)
+        if self.local_fusion == "track-gru-residual":
+            if local is None:
+                raise ValueError("Track GRU residual requires pooled local AV")
+            assert self.person_track_delta is not None
+            assert self.person_track_residual_gate is not None
+            fused_states = self._fuse_layers(states)
+            context, person_valid = self._track_gru_context(
+                fused_states,
+                local,
+                person_features,
+                person_mask,
+            )
+            delta = self.person_track_delta(
+                torch.cat((features, context, person_valid), dim=-1)
+            )
+            flat_logits = flat_logits + (
+                torch.tanh(self.person_track_residual_gate)
+                * delta
+                * person_valid
+            )
+        return features, flat_logits
+
     def forward(
         self,
         states: torch.Tensor,
@@ -435,13 +596,12 @@ class TurnPredictionHead(nn.Module):
         person_features: torch.Tensor | None = None,
         person_mask: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        features = self._shared_features(
+        features, logits = self._activity_logits(
             states,
             local,
             person_features,
             person_mask,
         )
-        logits = self.net[-1](features)
         return logits.view(
             *features.shape[:-1], self.num_horizons, self.num_classes
         )
@@ -459,13 +619,13 @@ class TurnPredictionHead(nn.Module):
             raise RuntimeError(
                 "This head was created without transition_auxiliary=True"
             )
-        features = self._shared_features(
+        features, flat_activity_logits = self._activity_logits(
             states,
             local,
             person_features,
             person_mask,
         )
-        activity_logits = self.net[-1](features).view(
+        activity_logits = flat_activity_logits.view(
             *features.shape[:-1],
             self.num_horizons,
             self.num_classes,
@@ -492,6 +652,10 @@ class TurnPredictionHead(nn.Module):
             result["person_state_adapter"] = float(
                 torch.tanh(self.person_state_adapter_gate.detach()).cpu()
             )
+        if self.person_track_residual_gate is not None:
+            result["person_track_residual"] = float(
+                torch.tanh(self.person_track_residual_gate.detach()).cpu()
+            )
         return result
 
 
@@ -1004,6 +1168,7 @@ def load_session_data(
     local_modalities: str = "av",
     local_contract: dict[str, Any] | None = None,
     local_fusion: str = "pooled",
+    person_track_max_gap_seconds: float = 1.0,
 ) -> SessionData:
     """Load and validate one saved Omni feature payload and target array."""
 
@@ -1110,7 +1275,22 @@ def load_session_data(
             name,
             dataset=dataset,
         )
-        if local_fusion in PERSON_SET_FUSION_MODES:
+        if local_fusion in TRACK_PERSON_FUSION_MODES:
+            if local_modalities != "av":
+                raise ValueError(
+                    "Track person fusion requires --local-modalities av"
+                )
+            from local_turn_features import load_track_features
+
+            local, local_people, contract = load_track_features(
+                local_path,
+                payload,
+                sample_id=name,
+                dataset=dataset,
+                track_max_gap_seconds=person_track_max_gap_seconds,
+                expected_contract=local_contract,
+            )
+        elif local_fusion in CURRENT_FRAME_PERSON_SET_FUSION_MODES:
             if local_modalities != "av":
                 raise ValueError(
                     "Person-set fusion requires --local-modalities av"
@@ -1196,6 +1376,8 @@ def checkpoint_head_config(payload: dict[str, Any]) -> dict[str, Any]:
     config.setdefault("local_fusion", "pooled")
     config.setdefault("person_attention_dim", 256)
     config.setdefault("person_attention_heads", 4)
+    config.setdefault("person_history_frames", 16)
+    config.setdefault("person_track_max_gap_seconds", 1.0)
     config.setdefault(
         "layer_fusion",
         (

infer_turn_head.py

@@ -149,7 +149,13 @@ def _predict_states(
                 start,
                 min(start + batch_size, len(states)),
             )
-            person_batch, person_mask = local_person_features.padded(rows)
+            if head.local_fusion == "track-gru-residual":
+                person_batch, person_mask = local_person_features.padded_history(
+                    rows,
+                    max_observations=head.person_history_frames,
+                )
+            else:
+                person_batch, person_mask = local_person_features.padded(rows)
             person_batch = person_batch.to(
                 device=device,
                 dtype=torch.float32,
@@ -440,6 +446,7 @@ def main() -> None:
             local_modalities=local_contract["modalities"] if local_contract else "av",
             local_contract=local_contract,
             local_fusion=head.local_fusion,
+            person_track_max_gap_seconds=head.person_track_max_gap_seconds,
         )
         if data.hidden_size != int(head_config["hidden_size"]):
             raise ValueError(
@@ -519,12 +526,18 @@ def main() -> None:
             "prediction_layout": "[target_time, class, future_horizon]",
             "checkpoint": str(args.checkpoint.resolve()),
             "checkpoint_epoch": checkpoint_payload.get("epoch"),
+            "checkpoint_training_epochs": checkpoint_payload.get(
+                "training_config",
+                {},
+            ).get("epochs"),
             "checkpoint_train_sessions": checkpoint_payload.get("train_sessions"),
             "checkpoint_dev_sessions": checkpoint_payload.get("dev_sessions"),
             "layer_fusion": head.layer_fusion,
             "local_fusion": head.local_fusion,
             "local_fusion_scope": LOCAL_FUSION_SCOPES[head.local_fusion],
             "person_interaction_gates": head.person_interaction_gates(),
+            "person_history_frames": head.person_history_frames,
+            "person_track_max_gap_seconds": head.person_track_max_gap_seconds,
             "local_contract": local_contract,
             "transition_auxiliary": head.transition_auxiliary,
             "transition_objective": checkpoint_payload.get(

literature_notes/02_egospeak.md

@@ -3,6 +3,8 @@
 - 会议：Findings of NAACL 2025
 - 阅读深度：深入阅读
 - 官方链接：https://aclanthology.org/2025.findings-naacl.163/
+- 官方代码：https://github.com/jun297/EgoSpeak
+- 本项目核对的代码 revision：`b5bee35e796d8f2d6aa7da2bf1b07836270f8bfb`
 
 ## 摘要（中文转述）
 
@@ -12,3 +14,14 @@ EgoSpeak 研究第一人称智能体在自然多人对话中“何时开口”
 
 它定义了本项目复现的核心协议：5 Hz 目标、多个未来预测时距、EasyCom/Ego4D 划分及 AP/mAP 评测。我们发现 EasyCom 中 wearer 与 others 可以同时为真，因此三路 softmax 的互斥输出与真实标签结构不完全一致；四状态 softmax 和独立 BCE 的消融虽更匹配标签结构，但当前四个 epoch 的收益很小，说明主要瓶颈可能不是输出空间。
 
+## GRU 实现核对
+
+官方 `MiniROADA` 将同一时刻的 RGB 与音频 feature 拼接，经
+`Linear -> LayerNorm -> ReLU -> Dropout` 后送入单层 GRU，再把每个 recurrent state
+映射成十个 anticipation hidden states 并共享分类器。EasyCom 配置使用 128 个 5 Hz
+step（25.6 秒）、stride 4、2048D embedding、1024D hidden 和 10 个 future steps。
+
+本项目的 `track-gru-residual` 没有复制这段实现，也不是该基线的复现：它把 GRU 参数在
+匿名人物轨迹之间共享，并用零门控残差修正 matched pooled Omni head。官方实现仍应作为
+单独的 sequence-GRU 强基线；逐人 track continuity 是否有效，需要 broken-ID/no-ID 控制
+而不能由架构名称推断。

local_turn_features.py

@@ -25,6 +25,7 @@ class FramePersonFeatures:
     frame_offsets: torch.LongTensor
     row_to_frame: torch.LongTensor
     chosen_frame_times: torch.Tensor
+    previous_observation: torch.LongTensor | None = None
 
     def __post_init__(self) -> None:
         if (
@@ -65,6 +66,19 @@ class FramePersonFeatures:
             raise ValueError("Only rows without an available frame may have NaN frame time")
         if not torch.isfinite(self.chosen_frame_times[~missing]).all():
             raise ValueError("Chosen person frame times must be finite")
+        if self.previous_observation is not None:
+            previous = self.previous_observation
+            if (
+                not isinstance(previous, torch.Tensor)
+                or previous.dtype != torch.long
+                or previous.shape != (len(self.features),)
+            ):
+                raise ValueError("previous_observation must be int64 [N]")
+            indices = torch.arange(len(previous), dtype=torch.long)
+            if bool(((previous < -1) | (previous >= indices)).any()):
+                raise ValueError(
+                    "Each previous observation must be -1 or point strictly backward"
+                )
 
     def __len__(self) -> int:
         return len(self.row_to_frame)
@@ -89,13 +103,14 @@ class FramePersonFeatures:
             self.frame_offsets,
             self.row_to_frame[indices],
             self.chosen_frame_times[indices],
+            self.previous_observation,
         )
 
-    def padded(
+    def _padded_observation_indices(
         self,
         row_indices: torch.Tensor | None = None,
-    ) -> tuple[torch.Tensor, torch.BoolTensor]:
-        """Materialize only one batch as [B,P,D] plus a True-for-real mask."""
+    ) -> tuple[torch.LongTensor, torch.BoolTensor]:
+        """Return current-frame observation indices and a real-person mask."""
 
         if row_indices is None:
             frames = self.row_to_frame
@@ -121,12 +136,77 @@ class FramePersonFeatures:
         width = int(counts.max()) if len(counts) else 0
         mask = torch.arange(width)[None, :] < counts[:, None]
         if width == 0:
-            return self.features.new_zeros((len(frames), 0, self.dimension)), mask
-        person_indices = starts[:, None] + torch.arange(width)[None, :]
-        safe_indices = person_indices.masked_fill(~mask, 0)
+            return torch.zeros((len(frames), 0), dtype=torch.long), mask
+        indices = starts[:, None] + torch.arange(width)[None, :]
+        return indices.masked_fill(~mask, 0), mask
+
+    def padded(
+        self,
+        row_indices: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.BoolTensor]:
+        """Materialize only one batch as [B,P,D] plus a True-for-real mask."""
+
+        safe_indices, mask = self._padded_observation_indices(row_indices)
+        if safe_indices.shape[1] == 0:
+            return self.features.new_zeros((len(safe_indices), 0, self.dimension)), mask
         values = self.features[safe_indices]
         return values.masked_fill(~mask[..., None], 0), mask
 
+    def padded_history(
+        self,
+        row_indices: torch.Tensor | None = None,
+        *,
+        max_observations: int,
+    ) -> tuple[torch.Tensor, torch.BoolTensor]:
+        """Return causal per-current-person histories as ``[B,P,L,D]``.
+
+        Histories are chronological and left aligned. Track links are built
+        from provided IDs with an explicit causal gap/reset policy when the
+        sidecar is loaded; no identity value is exposed to the model.
+        """
+
+        if max_observations <= 0:
+            raise ValueError("max_observations must be positive")
+        if self.previous_observation is None:
+            raise ValueError("Track histories were not constructed for this person set")
+        current, current_mask = self._padded_observation_indices(row_indices)
+        batch, people = current.shape
+        if people == 0:
+            return (
+                self.features.new_zeros(
+                    (batch, 0, max_observations, self.dimension)
+                ),
+                torch.zeros(
+                    (batch, 0, max_observations), dtype=torch.bool
+                ),
+            )
+
+        reverse_indices = torch.zeros(
+            (batch, people, max_observations), dtype=torch.long
+        )
+        reverse_mask = torch.zeros(
+            (batch, people, max_observations), dtype=torch.bool
+        )
+        cursor = current.clone()
+        cursor_valid = current_mask.clone()
+        for step in range(max_observations):
+            reverse_indices[:, :, step] = cursor.masked_fill(~cursor_valid, 0)
+            reverse_mask[:, :, step] = cursor_valid
+            safe = cursor.masked_fill(~cursor_valid, 0)
+            cursor = self.previous_observation[safe]
+            cursor_valid &= cursor >= 0
+
+        lengths = reverse_mask.sum(dim=-1)
+        output_positions = torch.arange(max_observations)[None, None, :]
+        source_positions = lengths[:, :, None] - 1 - output_positions
+        history_mask = output_positions < lengths[:, :, None]
+        source_positions = source_positions.clamp_min(0)
+        chronological = reverse_indices.gather(dim=-1, index=source_positions)
+        chronological = chronological.masked_fill(~history_mask, 0)
+        histories = self.features[chronological]
+        histories = histories.masked_fill(~history_mask[..., None], 0)
+        return histories, history_mask
+
     @classmethod
     def concatenate(
         cls,
@@ -143,6 +223,10 @@ class FramePersonFeatures:
         time_parts = []
         feature_base = 0
         frame_base = 0
+        previous_parts = []
+        has_track_links = [part.previous_observation is not None for part in parts]
+        if any(has_track_links) and not all(has_track_links):
+            raise ValueError("Track history links are missing for some sessions")
         for part in parts:
             feature_parts.append(part.features)
             offset_parts.append(part.frame_offsets[1:] + feature_base)
@@ -151,6 +235,11 @@ class FramePersonFeatures:
             rows[valid] += frame_base
             row_parts.append(rows)
             time_parts.append(part.chosen_frame_times)
+            if part.previous_observation is not None:
+                previous = part.previous_observation.clone()
+                valid_previous = previous >= 0
+                previous[valid_previous] += feature_base
+                previous_parts.append(previous)
             feature_base += len(part.features)
             frame_base += len(part.frame_offsets) - 1
         return cls(
@@ -158,6 +247,7 @@ class FramePersonFeatures:
             torch.cat(offset_parts, dim=0),
             torch.cat(row_parts, dim=0),
             torch.cat(time_parts, dim=0),
+            torch.cat(previous_parts, dim=0) if previous_parts else None,
         )
 
 
@@ -257,6 +347,7 @@ def load_frame_set_features(
     sample_id: str,
     dataset: str,
     expected_contract: dict[str, Any] | None = None,
+    track_max_gap_seconds: float | None = None,
 ) -> tuple[torch.Tensor, FramePersonFeatures, dict[str, Any]]:
     """Load direct audio and the latest causal per-frame person set per 5 Hz row."""
 
@@ -340,11 +431,23 @@ def load_frame_set_features(
     chosen[has_frame] = frame_times[row_to_frame[has_frame]]
     if bool((chosen[has_frame] >= times[has_frame]).any()):
         raise ValueError(f"Person frame mapping selected future visual evidence: {path}")
+    previous_observation = None
+    if track_max_gap_seconds is not None:
+        if not math.isfinite(track_max_gap_seconds) or track_max_gap_seconds <= 0:
+            raise ValueError("track_max_gap_seconds must be finite and positive")
+        previous_observation = _build_previous_track_links(
+            records,
+            frame_times,
+            frame_offsets,
+            max_gap_seconds=track_max_gap_seconds,
+            path=path,
+        )
     frame_sets = FramePersonFeatures(
         features,
         frame_offsets,
         row_to_frame.to(dtype=torch.long),
         chosen,
+        previous_observation,
     )
     direct_audio = torch.cat(
         (audio.float(), audio_valid[:, None].float()),
@@ -353,6 +456,96 @@ def load_frame_set_features(
     return direct_audio, frame_sets, contract
 
 
+def _build_previous_track_links(
+    records: list[Any],
+    frame_times: torch.Tensor,
+    frame_offsets: torch.LongTensor,
+    *,
+    max_gap_seconds: float,
+    path: Path,
+) -> torch.LongTensor:
+    """Link observations causally by provided track ID without exposing IDs.
+
+    A source-clip boundary always resets a track. Missing/unknown IDs are
+    isolated observations instead of being incorrectly merged into one track.
+    Duplicate known IDs within one frame fail closed.
+    """
+
+    previous = torch.full((len(records),), -1, dtype=torch.long)
+    last_by_track: dict[tuple[int, str], tuple[int, float]] = {}
+    unknown_values = {"", "-1", "none", "null", "unknown"}
+    for frame_index, time_value in enumerate(frame_times.tolist()):
+        start = int(frame_offsets[frame_index])
+        end = int(frame_offsets[frame_index + 1])
+        seen: set[tuple[int, str]] = set()
+        for observation_index in range(start, end):
+            record = records[observation_index]
+            if not isinstance(record, dict):
+                raise ValueError(f"Invalid person record at {observation_index}: {path}")
+            raw_id = record.get("id")
+            text_id = "" if raw_id is None else str(raw_id).strip()
+            source_clip_index = record.get("source_clip_index", 0)
+            if isinstance(source_clip_index, bool) or not isinstance(
+                source_clip_index, int
+            ):
+                raise ValueError(
+                    f"Invalid source_clip_index at person record "
+                    f"{observation_index}: {path}"
+                )
+            if text_id.lower() in unknown_values:
+                # Never join multiple anonymous detections under a shared -1.
+                continue
+            key = (source_clip_index, text_id)
+            if key in seen:
+                raise ValueError(
+                    f"Duplicate person/track ID {text_id!r} in frame "
+                    f"{frame_index}: {path}"
+                )
+            seen.add(key)
+            old = last_by_track.get(key)
+            if old is not None:
+                old_index, old_time = old
+                gap = float(time_value) - old_time
+                if gap <= 0:
+                    raise ValueError(f"Non-causal person track order: {path}")
+                if gap <= max_gap_seconds + 1e-8:
+                    previous[observation_index] = old_index
+            last_by_track[key] = (observation_index, float(time_value))
+    return previous
+
+
+def load_track_features(
+    path: Path,
+    reference: dict[str, Any],
+    *,
+    sample_id: str,
+    dataset: str,
+    track_max_gap_seconds: float,
+    expected_contract: dict[str, Any] | None = None,
+) -> tuple[torch.Tensor, FramePersonFeatures, dict[str, Any]]:
+    """Load the exact pooled AV baseline plus causal anonymous track links."""
+
+    pooled, contract = load_local_features(
+        path,
+        reference,
+        sample_id=sample_id,
+        dataset=dataset,
+        modalities="av",
+        expected_contract=expected_contract,
+    )
+    _, people, second_contract = load_frame_set_features(
+        path,
+        reference,
+        sample_id=sample_id,
+        dataset=dataset,
+        expected_contract=contract,
+        track_max_gap_seconds=track_max_gap_seconds,
+    )
+    if second_contract != contract:
+        raise AssertionError("Track and pooled local contracts unexpectedly differ")
+    return pooled, people, contract
+
+
 def local_head_dimensions(contract: dict[str, Any] | None) -> dict[str, int]:
     if contract is None:
         return {"local_visual_dim": 0, "local_audio_dim": 0}

run_easycom_track_gru_residual.sh

@@ -0,0 +1,177 @@
+#!/usr/bin/env bash
+# Foreground matched EasyCom pooled AV -> causal anonymous-track GRU residual.
+# Usage: bash run_easycom_track_gru_residual.sh [all|pooled|track|test] [seed ...]
+set -euo pipefail
+cd "$(dirname "$0")"
+
+FEATURE_ROOT=${FEATURE_ROOT:-/data1/yfzx/wangx/Datasets/easycom/omni_features/open_stream_2fps_8s_last4_stack}
+TARGET_ROOT=${TARGET_ROOT:-/data1/yfzx/wangx/Datasets/easycom/easycom_target_frame}
+LOCAL_ROOT=${LOCAL_ROOT:-/data1/yfzx/wangx/Datasets/easycom/omni_features/local_qwen_1s_encoded_tf4576}
+CHECKPOINT_ROOT=${CHECKPOINT_ROOT:-checkpoints/easycom_track_gru_residual}
+RESULT_ROOT=${RESULT_ROOT:-results/easycom_track_gru_residual}
+TEST_RESULT_ROOT=${TEST_RESULT_ROOT:-results/easycom_track_gru_residual_test}
+EPOCHS=${EPOCHS:-4}
+WARMUP_EPOCHS=${WARMUP_EPOCHS:-1.6}
+LOCAL_PROJECTION_DIM=${LOCAL_PROJECTION_DIM:-256}
+PERSON_ATTENTION_DIM=${PERSON_ATTENTION_DIM:-256}
+PERSON_ATTENTION_HEADS=${PERSON_ATTENTION_HEADS:-4}
+PERSON_HISTORY_FRAMES=${PERSON_HISTORY_FRAMES:-16}
+PERSON_TRACK_MAX_GAP_SECONDS=${PERSON_TRACK_MAX_GAP_SECONDS:-1.0}
+SUMMARIZE=${SUMMARIZE:-1}
+read -r -a SESSION_ARGS <<< "${SESSIONS:-4 5 6 7 8 9 10 11 12}"
+read -r -a DEV_SESSION_ARGS <<< "${DEV_SESSIONS:-12}"
+read -r -a SCREEN_SESSION_ARGS <<< "${SCREEN_SESSIONS:-12}"
+export OMP_NUM_THREADS=${OMP_NUM_THREADS:-4}
+export MKL_NUM_THREADS=${MKL_NUM_THREADS:-4}
+export OPENBLAS_NUM_THREADS=${OPENBLAS_NUM_THREADS:-4}
+
+variant=${1:-all}
+if (( $# )); then shift; fi
+if (( $# )); then seeds=("$@"); else seeds=(0 1 2); fi
+case "$variant" in
+  all) variants=(pooled track) ;;
+  pooled|track) variants=("$variant") ;;
+  test) variants=(pooled track) ;;
+  *) echo "Expected all, pooled, track, or test" >&2; exit 2 ;;
+esac
+
+mkdir -p "$CHECKPOINT_ROOT" "$RESULT_ROOT"
+
+infer_checkpoint() {
+  local checkpoint=$1 result=$2
+  shift 2
+  if [[ -e "$result/metrics.json" ]]; then
+    echo "Refusing to overwrite $result/metrics.json" >&2
+    return 1
+  fi
+  python infer_turn_head.py \
+    --feature-root "$FEATURE_ROOT" \
+    --target-root "$TARGET_ROOT" \
+    --local-feature-root "$LOCAL_ROOT" \
+    --checkpoint "$checkpoint" \
+    --output-dir "$result" \
+    --sessions "$@" \
+    --expected-mode open-stream-replay \
+    --device cuda:0
+}
+
+train_and_infer_pooled() {
+  local seed=$1 checkpoint result
+  checkpoint="$CHECKPOINT_ROOT/pooled_seed${seed}.pth"
+  result="$RESULT_ROOT/pooled_seed${seed}"
+  if [[ -e "$checkpoint" || -e "$result/metrics.json" ]]; then
+    echo "Refusing to overwrite $checkpoint or $result" >&2
+    return 1
+  fi
+  python train_turn_head.py \
+    --feature-root "$FEATURE_ROOT" \
+    --target-root "$TARGET_ROOT" \
+    --local-feature-root "$LOCAL_ROOT" \
+    --local-modalities av \
+    --local-fusion pooled \
+    --local-projection-dim "$LOCAL_PROJECTION_DIM" \
+    --output "$checkpoint" \
+    --sessions "${SESSION_ARGS[@]}" \
+    --dev-sessions "${DEV_SESSION_ARGS[@]}" \
+    --epochs "$EPOCHS" \
+    --batch-size 512 \
+    --learning-rate 7e-5 \
+    --weight-decay 5e-5 \
+    --warmup-epochs "$WARMUP_EPOCHS" \
+    --warmup-factor 0.3 \
+    --dropout 0.2 \
+    --loss lstr-mce \
+    --transition-objective none \
+    --layer-fusion concat-mlp \
+    --expected-mode open-stream-replay \
+    --seed "$seed" \
+    --device cuda:0
+  # Screen architecture/hyperparameters only on the held-out training session.
+  infer_checkpoint "$checkpoint" "$result" "${SCREEN_SESSION_ARGS[@]}"
+}
+
+train_and_infer_track() {
+  local seed=$1 base checkpoint result
+  base="$CHECKPOINT_ROOT/pooled_seed${seed}.pth"
+  checkpoint="$CHECKPOINT_ROOT/track_gru_seed${seed}.pth"
+  result="$RESULT_ROOT/track_gru_seed${seed}"
+  if [[ ! -s "$base" ]]; then
+    echo "Missing matched pooled checkpoint: $base; run variant pooled or all first" >&2
+    return 1
+  fi
+  if [[ -e "$checkpoint" || -e "$result/metrics.json" ]]; then
+    echo "Refusing to overwrite $checkpoint or $result" >&2
+    return 1
+  fi
+  python train_turn_head.py \
+    --feature-root "$FEATURE_ROOT" \
+    --target-root "$TARGET_ROOT" \
+    --local-feature-root "$LOCAL_ROOT" \
+    --local-modalities av \
+    --local-fusion track-gru-residual \
+    --local-projection-dim "$LOCAL_PROJECTION_DIM" \
+    --person-attention-dim "$PERSON_ATTENTION_DIM" \
+    --person-attention-heads "$PERSON_ATTENTION_HEADS" \
+    --person-history-frames "$PERSON_HISTORY_FRAMES" \
+    --person-track-max-gap-seconds "$PERSON_TRACK_MAX_GAP_SECONDS" \
+    --initialize-pooled-checkpoint "$base" \
+    --output "$checkpoint" \
+    --sessions "${SESSION_ARGS[@]}" \
+    --dev-sessions "${DEV_SESSION_ARGS[@]}" \
+    --epochs "$EPOCHS" \
+    --batch-size 512 \
+    --learning-rate 7e-5 \
+    --weight-decay 5e-5 \
+    --warmup-epochs "$WARMUP_EPOCHS" \
+    --warmup-factor 0.3 \
+    --dropout 0.2 \
+    --loss lstr-mce \
+    --transition-objective none \
+    --layer-fusion concat-mlp \
+    --expected-mode open-stream-replay \
+    --seed "$seed" \
+    --device cuda:0
+  infer_checkpoint "$checkpoint" "$result" "${SCREEN_SESSION_ARGS[@]}"
+}
+
+if [[ "$variant" == test ]]; then
+  mkdir -p "$TEST_RESULT_ROOT"
+  for selected in "${variants[@]}"; do
+    for seed in "${seeds[@]}"; do
+      checkpoint="$CHECKPOINT_ROOT/${selected}_seed${seed}.pth"
+      if [[ "$selected" == track ]]; then
+        checkpoint="$CHECKPOINT_ROOT/track_gru_seed${seed}.pth"
+        result="$TEST_RESULT_ROOT/track_gru_seed${seed}"
+      else
+        result="$TEST_RESULT_ROOT/pooled_seed${seed}"
+      fi
+      if [[ ! -s "$checkpoint" ]]; then
+        echo "Missing checkpoint: $checkpoint" >&2
+        exit 1
+      fi
+      echo "$(date): final test ${selected}, seed=${seed}"
+      infer_checkpoint "$checkpoint" "$result" 3 2 1
+    done
+  done
+  summary_root=$TEST_RESULT_ROOT
+else
+  for selected in "${variants[@]}"; do
+    for seed in "${seeds[@]}"; do
+      echo "$(date): ${selected}, seed=${seed}"
+      if [[ "$selected" == pooled ]]; then
+        train_and_infer_pooled "$seed"
+      else
+        train_and_infer_track "$seed"
+      fi
+    done
+  done
+  summary_root=$RESULT_ROOT
+fi
+
+if [[ "$SUMMARIZE" == 1 ]]; then
+  python summarize_local_baselines.py \
+    --root "$summary_root" \
+    --dataset easycom \
+    --protocol causal_valid_only \
+    --paired-deltas
+fi

summarize_local_baselines.py

@@ -26,11 +26,17 @@ def main():
         if args.dataset is not None and metrics.get("dataset") != args.dataset:
             raise ValueError(f"Dataset mismatch in {path}: {metrics.get('dataset')}")
         if args.paired_deltas:
-            signature = (metrics.get("dataset"), metrics.get("sessions"), metrics.get("checkpoint_epoch"),
+            # Candidates may select different best epochs on the same dev set.
+            # Match the maximum training budget, not the selected epoch.
+            training_epochs = metrics.get(
+                "checkpoint_training_epochs",
+                metrics.get("checkpoint_epoch"),
+            )
+            signature = (metrics.get("dataset"), metrics.get("sessions"), training_epochs,
                          metrics.get("checkpoint_train_sessions"), metrics.get("checkpoint_dev_sessions"),
                          metrics.get("loss"), metrics.get("layer_fusion"))
             if comparison_signature is not None and signature != comparison_signature:
-                raise ValueError(f"Paired runs have different splits/epochs/loss/fusion: {path}")
+                raise ValueError(f"Paired runs have different splits/budget/loss/fusion: {path}")
             comparison_signature = signature
         score = metrics[args.protocol]
         first = score["horizons"][0]
@@ -54,7 +60,8 @@ def main():
     if args.paired_deltas:
         print("\nPaired deltas (percentage points; same column order; matching seeds only)")
         for treatment, control in (("audio", "baseline"), ("visual", "baseline"),
-                                   ("av", "baseline"), ("av", "audio")):
+                                   ("av", "baseline"), ("av", "audio"),
+                                   ("track_gru", "pooled")):
             left, right = by_seed.get(treatment, {}), by_seed.get(control, {})
             seeds = sorted(left.keys() & right.keys())
             if not seeds:

test_local_turn_features.py

@@ -33,7 +33,8 @@ from extract_turn_local_features import (
 )
 from local_turn_features import (
     LOCAL_FORMAT, FramePersonFeatures, causal_frame_ticks, crop_box, easycom_boxes,
-    load_frame_set_features, load_local_features, local_head_dimensions,
+    load_frame_set_features, load_local_features, load_track_features,
+    local_head_dimensions,
     pooled_recent_visual, reference_alignment,
 )
 from train_turn_head import _load_examples
@@ -352,6 +353,72 @@ class LocalTests(unittest.TestCase):
                         dataset="easycom",
                     )
 
+    def test_track_loader_builds_causal_histories_and_resets_tracks(self):
+        reference, side, _ = fixture(n=12)
+        frame_times = torch.tensor([0., .5, 1., 1.5, 2.], dtype=torch.float64)
+        side["persons"] = {
+            "features": torch.arange(1, 9, dtype=torch.float32)[:, None].repeat(1, 6),
+            "frame_times": frame_times,
+            "frame_offsets": torch.tensor([0, 2, 4, 6, 7, 8]),
+            "records": [
+                {"id": "A", "source_clip_index": 0},
+                {"id": "B", "source_clip_index": 0},
+                {"id": "B", "source_clip_index": 0},
+                {"id": "A", "source_clip_index": 0},
+                {"id": "A", "source_clip_index": 0},
+                {"id": -1, "source_clip_index": 0},
+                # The same ID in another retained source clip is a new track.
+                {"id": "A", "source_clip_index": 1},
+                {"id": "A", "source_clip_index": 1},
+            ],
+        }
+        with tempfile.TemporaryDirectory() as directory:
+            path = Path(directory) / "Session_4.pt"
+            torch.save(side, path)
+            pooled, people, contract = load_track_features(
+                path,
+                reference,
+                sample_id="Session_4",
+                dataset="easycom",
+                track_max_gap_seconds=.75,
+            )
+            self.assertEqual(pooled.shape, (12, 13))
+            self.assertEqual(contract["modalities"], "av")
+            self.assertEqual(
+                people.previous_observation.tolist(),
+                [-1, -1, 1, 0, 3, -1, -1, 6],
+            )
+            # Cutoff 1.2s selects frame 1.0s. A has observations 1,4,5;
+            # anonymous -1 remains a one-observation track.
+            history, mask = people.padded_history(
+                torch.tensor([5]),
+                max_observations=4,
+            )
+            self.assertEqual(mask.tolist(), [[[True, True, True, False],
+                                              [True, False, False, False]]])
+            self.assertEqual(history[0, 0, :, 0].tolist(), [1., 4., 5., 0.])
+            self.assertEqual(history[0, 1, :, 0].tolist(), [6., 0., 0., 0.])
+            # Cutoff 1.8s selects frame 1.5s; source clip reset prevents A's
+            # earlier observations from crossing the boundary.
+            history, mask = people.padded_history(
+                torch.tensor([8]),
+                max_observations=4,
+            )
+            self.assertEqual(mask.tolist(), [[[True, False, False, False]]])
+            self.assertEqual(history[0, 0, :, 0].tolist(), [7., 0., 0., 0.])
+
+            broken = copy.deepcopy(side)
+            broken["persons"]["records"][1]["id"] = "A"
+            torch.save(broken, path)
+            with self.assertRaisesRegex(ValueError, "Duplicate person/track ID"):
+                load_track_features(
+                    path,
+                    reference,
+                    sample_id="Session_4",
+                    dataset="easycom",
+                    track_max_gap_seconds=.75,
+                )
+
     def test_frame_person_concatenation_selection_and_empty_padding(self):
         features = torch.tensor([[1., 2.], [3., 4.]])
         a = FramePersonFeatures(
@@ -567,6 +634,86 @@ class LocalTests(unittest.TestCase):
         restored.eval()
         self.assertTrue(torch.equal(restored(states, local, people, mask), logits))
 
+    def test_track_gru_residual_starts_from_exact_pooled_logits(self):
+        torch.manual_seed(17)
+        common = {
+            "hidden_size": 8,
+            "num_input_layers": 4,
+            "layer_fusion": "concat-mlp",
+            "local_visual_dim": 6,
+            "local_audio_dim": 5,
+            "local_projection_dim": 3,
+            "dropout": 0,
+        }
+        pooled = TurnPredictionHead(**common, local_fusion="pooled").eval()
+        track = TurnPredictionHead(
+            **common,
+            local_fusion="track-gru-residual",
+            person_attention_dim=4,
+            person_attention_heads=2,
+            person_history_frames=3,
+        ).eval()
+        track_state = track.state_dict()
+        for name, value in pooled.state_dict().items():
+            self.assertIn(name, track_state)
+            self.assertEqual(track_state[name].shape, value.shape)
+            track_state[name] = value
+        track.load_state_dict(track_state)
+
+        states = torch.randn(3, 4, 8)
+        local = torch.randn(3, 13)
+        local[:, 6] = 1
+        local[:, 12] = 1
+        histories = torch.randn(3, 2, 3, 6)
+        mask = torch.tensor(
+            [
+                [[True, True, True], [True, True, False]],
+                [[True, False, False], [False, False, False]],
+                [[False, False, False], [False, False, False]],
+            ]
+        )
+        histories[~mask] = 0
+        baseline = pooled(states, local)
+        self.assertTrue(torch.equal(track(states, local, histories, mask), baseline))
+        self.assertEqual(track.person_interaction_gates(), {"person_track_residual": 0.0})
+
+        with torch.no_grad():
+            assert track.person_track_residual_gate is not None
+            track.person_track_residual_gate.fill_(.25)
+        logits = track(states, local, histories, mask)
+        self.assertFalse(torch.equal(logits[:2], baseline[:2]))
+        # With no current person the correction must be exactly zero.
+        self.assertTrue(torch.equal(logits[2], baseline[2]))
+        permutation = torch.tensor([1, 0])
+        self.assertTrue(
+            torch.allclose(
+                logits,
+                track(states, local, histories[:, permutation], mask[:, permutation]),
+                atol=1e-6,
+                rtol=1e-6,
+            )
+        )
+        changed_padding = histories.clone()
+        changed_padding[~mask] = 1000
+        self.assertTrue(torch.equal(logits, track(states, local, changed_padding, mask)))
+        logits.square().mean().backward()
+        self.assertIsNotNone(track.person_track_gru.weight_ih_l0.grad)
+        self.assertIsNotNone(track.person_track_delta[-1].weight.grad)
+        self.assertIsNotNone(track.person_track_residual_gate.grad)
+
+        config = {
+            **common,
+            "local_fusion": "track-gru-residual",
+            "person_attention_dim": 4,
+            "person_attention_heads": 2,
+            "person_history_frames": 3,
+        }
+        restored = TurnPredictionHead(
+            **checkpoint_head_config({"head_config": config})
+        ).eval()
+        restored.load_state_dict(track.state_dict())
+        self.assertTrue(torch.equal(restored(states, local, histories, mask), logits))
+
     def test_head_shapes_gradients_missing_and_roundtrip(self):
         states = torch.randn(3, 4, 8)
         for v, a in ((0, 0), (6, 0), (0, 5), (6, 5)):
@@ -855,6 +1002,7 @@ class RoundTripTests(unittest.TestCase):
                 "audio",
                 "visual",
                 "av",
+                "track-gru",
                 "frame-set",
                 "f5-screen",
                 "f6-screen",
@@ -877,6 +1025,7 @@ class RoundTripTests(unittest.TestCase):
                 if mode == "av":
                     train += ["--transition-objective", "r1-supervised"]
                 person_mode = {
+                    "track-gru": "track-gru-residual",
                     "frame-set": "frame-set-attention",
                     "f5-screen": "audio-person-frame-set",
                     "f6-screen": "gated-person-state-adapter",
@@ -887,10 +1036,16 @@ class RoundTripTests(unittest.TestCase):
                         "--person-attention-dim", "4",
                         "--person-attention-heads", "2",
                     ]
+                if mode == "track-gru":
+                    train += [
+                        "--person-history-frames", "3",
+                        "--person-track-max-gap-seconds", "1",
+                        "--initialize-pooled-checkpoint", str(root / "av.pth"),
+                    ]
                 if mode != "baseline":
                     modality = (
                         "av"
-                        if mode in ("ego4d", "frame-set", "f5-screen", "f6-screen")
+                        if mode in ("ego4d", "track-gru", "frame-set", "f5-screen", "f6-screen")
                         else mode
                     )
                     train += ["--local-feature-root", str(root / "local"), "--local-modalities", modality]

third_party/README.md

@@ -0,0 +1,15 @@
+# Third-party reference code
+
+Reference repositories are inspected here but are not vendored into this
+project. They are not runtime dependencies.
+
+The track-GRU design was checked against the official EgoSpeak repository at
+revision `b5bee35e796d8f2d6aa7da2bf1b07836270f8bfb`:
+
+```bash
+git clone https://github.com/jun297/EgoSpeak.git third_party/EgoSpeak
+git -C third_party/EgoSpeak checkout b5bee35e796d8f2d6aa7da2bf1b07836270f8bfb
+```
+
+`third_party/EgoSpeak/` is intentionally ignored. No EgoSpeak source file is
+copied or imported by the project implementation.

train_turn_head.py

@@ -33,6 +33,8 @@ from easycom_head import (
     TRANSITION_OBJECTIVE_NAMES,
     TurnPredictionHead,
     average_precision_metrics,
+    checkpoint_head_config,
+    checkpoint_loss_name,
     egospeak_multihorizon_loss,
     four_state_multihorizon_loss,
     head_output_names_for_loss,
@@ -67,6 +69,7 @@ def parse_args() -> argparse.Namespace:
         default="pooled",
         help=(
             "pooled keeps the existing baseline; frame-set-attention is F1; "
+            "track-gru-residual adds a zero-gated causal track correction; "
             "audio-person-frame-set and gated-person-state-adapter are "
             "current-sidecar F5/F6 screening variants, not full F5/F6."
         ),
@@ -77,6 +80,31 @@ def parse_args() -> argparse.Namespace:
     )
     parser.add_argument("--person-attention-dim", type=int, default=256)
     parser.add_argument("--person-attention-heads", type=int, default=4)
+    parser.add_argument(
+        "--person-history-frames",
+        type=int,
+        default=16,
+        help="Maximum causal observations per current face track (16 at 2 FPS = 8s).",
+    )
+    parser.add_argument(
+        "--person-track-max-gap-seconds",
+        type=float,
+        default=1.0,
+        help="Reset a provided track when consecutive observations are farther apart.",
+    )
+    parser.add_argument(
+        "--initialize-pooled-checkpoint",
+        type=Path,
+        help=(
+            "Required by track-gru-residual: copy the exact pooled G+A+V base "
+            "before training the zero-gated track correction."
+        ),
+    )
+    parser.add_argument(
+        "--train-pooled-base",
+        action="store_true",
+        help="Also update the initialized pooled base; default keeps it frozen.",
+    )
     parser.add_argument(
         "--dataset",
         choices=DATASET_NAMES,
@@ -184,6 +212,10 @@ def _validate_args(args: argparse.Namespace) -> None:
         raise ValueError("Person attention dimension/heads must be positive")
     if args.person_attention_dim % args.person_attention_heads:
         raise ValueError("--person-attention-dim must be divisible by --person-attention-heads")
+    if args.person_history_frames <= 0:
+        raise ValueError("--person-history-frames must be positive")
+    if args.person_track_max_gap_seconds <= 0:
+        raise ValueError("--person-track-max-gap-seconds must be positive")
     if args.local_fusion != "pooled" and (
         args.local_feature_root is None or args.local_modalities != "av"
     ):
@@ -191,6 +223,16 @@ def _validate_args(args: argparse.Namespace) -> None:
             "Person-set --local-fusion modes require --local-feature-root and "
             "--local-modalities av"
         )
+    if args.local_fusion == "track-gru-residual":
+        if args.initialize_pooled_checkpoint is None:
+            raise ValueError(
+                "track-gru-residual requires --initialize-pooled-checkpoint"
+            )
+    elif args.initialize_pooled_checkpoint is not None or args.train_pooled_base:
+        raise ValueError(
+            "--initialize-pooled-checkpoint/--train-pooled-base require "
+            "--local-fusion track-gru-residual"
+        )
     if args.epochs <= 0:
         raise ValueError("--epochs must be positive")
     if args.batch_size <= 0:
@@ -252,6 +294,7 @@ def _load_examples(
     local_contract: dict[str, Any] | None = None,
     compute_transition_targets: bool = True,
     local_fusion: str = "pooled",
+    person_track_max_gap_seconds: float = 1.0,
 ) -> tuple[
     torch.Tensor,
     torch.Tensor,
@@ -281,6 +324,7 @@ def _load_examples(
             local_modalities=local_modalities,
             local_contract=local_contract,
             local_fusion=local_fusion,
+            person_track_max_gap_seconds=person_track_max_gap_seconds,
         )
         local_contract = data.feature_metadata.get("local_contract")
         states, targets, future_indices = data.future_examples(NUM_HORIZONS)
@@ -467,7 +511,13 @@ def _evaluate_valid_horizons(
         person_mask = None
         if people is not None:
             rows = torch.arange(start, min(start + batch_size, len(states)))
-            person_batch, person_mask = people.padded(rows)
+            if head.local_fusion == "track-gru-residual":
+                person_batch, person_mask = people.padded_history(
+                    rows,
+                    max_observations=head.person_history_frames,
+                )
+            else:
+                person_batch, person_mask = people.padded(rows)
             person_batch = person_batch.to(device=device, dtype=torch.float32)
             person_mask = person_mask.to(device=device)
         chunks.append(
@@ -510,7 +560,7 @@ def _checkpoint(
     history: list[dict[str, Any]],
 ) -> dict[str, Any]:
     return {
-        "format_version": 6,
+        "format_version": 7,
         "head": head.state_dict(),
         "head_config": {
             "hidden_size": hidden_size,
@@ -527,6 +577,8 @@ def _checkpoint(
             "local_fusion": head.local_fusion,
             "person_attention_dim": head.person_attention_dim,
             "person_attention_heads": head.person_attention_heads,
+            "person_history_frames": head.person_history_frames,
+            "person_track_max_gap_seconds": head.person_track_max_gap_seconds,
         },
         "local_contract": feature_metadata[0].get("local_contract"),
         "alignment": {
@@ -576,6 +628,14 @@ def _checkpoint(
             "local_projection_dim": args.local_projection_dim,
             "person_attention_dim": args.person_attention_dim,
             "person_attention_heads": args.person_attention_heads,
+            "person_history_frames": args.person_history_frames,
+            "person_track_max_gap_seconds": args.person_track_max_gap_seconds,
+            "initialize_pooled_checkpoint": (
+                str(args.initialize_pooled_checkpoint.expanduser().resolve())
+                if args.initialize_pooled_checkpoint is not None
+                else None
+            ),
+            "train_pooled_base": args.train_pooled_base,
             "learnable_layer_mixture": args.layer_fusion == "scalar-mixture",
             "supervision_policy": (
                 "transition-sparse-manifest"
@@ -591,6 +651,94 @@ def _checkpoint(
     }
 
 
+def _initialize_track_residual_from_pooled(
+    head: TurnPredictionHead,
+    checkpoint: Path,
+    *,
+    args: argparse.Namespace,
+    train_sessions: list[str],
+    dev_sessions: list[str],
+    local_contract: dict[str, Any] | None,
+) -> None:
+    """Load and optionally freeze the exact pooled base inside a track model."""
+
+    path = checkpoint.expanduser().resolve()
+    if not path.is_file():
+        raise FileNotFoundError(path)
+    payload = torch.load(path, map_location="cpu", weights_only=False)
+    if not isinstance(payload, dict) or not isinstance(payload.get("head"), dict):
+        raise ValueError(f"Invalid pooled checkpoint: {path}")
+    config = checkpoint_head_config(payload)
+    if config["local_fusion"] != "pooled":
+        raise ValueError(
+            "Track residual initialization requires a pooled checkpoint, got "
+            f"{config['local_fusion']!r}: {path}"
+        )
+    if checkpoint_loss_name(payload) != args.loss:
+        raise ValueError("Pooled checkpoint loss differs from the track experiment")
+    if payload.get("dataset", args.dataset) != args.dataset:
+        raise ValueError("Pooled checkpoint dataset differs from the track experiment")
+    if payload.get("train_sessions") != train_sessions:
+        raise ValueError("Pooled checkpoint train sessions differ from this run")
+    if payload.get("dev_sessions", []) != dev_sessions:
+        raise ValueError("Pooled checkpoint dev sessions differ from this run")
+    if payload.get("local_contract") != local_contract:
+        raise ValueError("Pooled checkpoint local feature contract differs from this run")
+
+    expected = {
+        "hidden_size": head.hidden_size,
+        "num_horizons": head.num_horizons,
+        "num_classes": head.num_classes,
+        "num_input_layers": head.num_input_layers,
+        "layer_fusion": head.layer_fusion,
+        "local_visual_dim": head.local_visual_dim,
+        "local_audio_dim": head.local_audio_dim,
+        "local_projection_dim": head.local_projection_dim,
+    }
+    mismatches = {
+        key: (config.get(key), value)
+        for key, value in expected.items()
+        if config.get(key) != value
+    }
+    if mismatches:
+        raise ValueError(
+            f"Pooled checkpoint architecture differs from track base: {mismatches}"
+        )
+
+    def is_base_parameter(name: str) -> bool:
+        return name == "layer_logits" or name.startswith(
+            ("layer_fusion_net.", "local_projections.", "net.")
+        )
+
+    source = payload["head"]
+    current = head.state_dict()
+    required = [name for name in current if is_base_parameter(name)]
+    missing = [name for name in required if name not in source]
+    bad_shapes = [
+        name
+        for name in required
+        if name in source and source[name].shape != current[name].shape
+    ]
+    if missing or bad_shapes:
+        raise ValueError(
+            "Pooled checkpoint cannot initialize the exact base; "
+            f"missing={missing[:8]}, shape_mismatch={bad_shapes[:8]}"
+        )
+    for name in required:
+        current[name] = source[name]
+    head.load_state_dict(current)
+
+    if not args.train_pooled_base:
+        for name, parameter in head.named_parameters():
+            if is_base_parameter(name):
+                parameter.requires_grad_(False)
+    print(
+        f"initialized pooled base: {path}; "
+        f"frozen={not args.train_pooled_base}",
+        flush=True,
+    )
+
+
 def main() -> None:
     args = parse_args()
     _validate_args(args)
@@ -679,6 +827,7 @@ def main() -> None:
         args.local_modalities,
         compute_transition_targets=args.transition_objective != "none",
         local_fusion=args.local_fusion,
+        person_track_max_gap_seconds=args.person_track_max_gap_seconds,
     )
     local_contract = feature_metadata[0].get("local_contract")
     if not len(states):
@@ -732,6 +881,7 @@ def main() -> None:
             local_contract,
             compute_transition_targets=args.transition_objective != "none",
             local_fusion=args.local_fusion,
+            person_track_max_gap_seconds=args.person_track_max_gap_seconds,
         )
         if dev_hidden_size != hidden_size:
             raise ValueError(
@@ -778,8 +928,20 @@ def main() -> None:
         local_fusion=args.local_fusion,
         person_attention_dim=args.person_attention_dim,
         person_attention_heads=args.person_attention_heads,
+        person_history_frames=args.person_history_frames,
+        person_track_max_gap_seconds=args.person_track_max_gap_seconds,
         **local_head_dimensions(local_contract),
     ).to(device)
+    if args.local_fusion == "track-gru-residual":
+        assert args.initialize_pooled_checkpoint is not None
+        _initialize_track_residual_from_pooled(
+            head,
+            args.initialize_pooled_checkpoint,
+            args=args,
+            train_sessions=train_sessions,
+            dev_sessions=dev_sessions,
+            local_contract=local_contract,
+        )
     if local_contract is not None:
         print(
             f"local fusion: {args.local_fusion}/{args.local_modalities}; "
@@ -861,6 +1023,13 @@ def main() -> None:
     best_epoch = args.epochs if not dev_sessions else 0
     for epoch in range(1, args.epochs + 1):
         head.train()
+        if args.local_fusion == "track-gru-residual" and not args.train_pooled_base:
+            # Keep dropout/norm behavior in the frozen pooled checkpoint fixed
+            # while the new track branch learns a residual correction.
+            if head.layer_fusion_net is not None:
+                head.layer_fusion_net.eval()
+            head.local_projections.eval()
+            head.net.eval()
         loss_sum = 0.0
         activity_loss_sum = 0.0
         transition_loss_sum = 0.0
@@ -900,9 +1069,15 @@ def main() -> None:
             person_batch = None
             person_mask = None
             if person_features is not None:
-                person_batch, person_mask = person_features.padded(
-                    example_index_batch
-                )
+                if head.local_fusion == "track-gru-residual":
+                    person_batch, person_mask = person_features.padded_history(
+                        example_index_batch,
+                        max_observations=head.person_history_frames,
+                    )
+                else:
+                    person_batch, person_mask = person_features.padded(
+                        example_index_batch
+                    )
                 person_batch = person_batch.to(
                     device=device,
                     dtype=torch.float32,