feat: add causal frame set attention fusion

PROJECT_CONTEXT.md

@@ -43,6 +43,13 @@ result.
 - Ego4D batch and 8-GPU local-feature launchers, with strict frame audits.
 - Timing-derived wearer floor-state and event diagnostics.
 
+### 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.
+
 ### Current blocker
 
 `evaluate_floor_state_changes.py` counts a wearer transition anywhere in a

docs/local_feature_baseline.md

@@ -174,6 +174,27 @@ warmup_factor=0.3、dropout=0.2、LSTR-MCE、不加 transition objective。
 Loader 会检查 sample、时间、label_indices、维度、valid mask 和 encoder/pooling 配置。
 训练与评测必须用同一种 AP 协议作比较；不要混用 full-timeline 与 causal-valid 的表格。
 
+## 4. F1 Frame-set Attention
+
+已有完整局部 sidecar 的 persons 字段保存了逐帧、逐人的原生 2048D 特征，因此 F1
+不需要重新提取图片或音频。新增训练参数为：
+
+    --local-feature-root "$LOCAL_ROOT" \
+    --local-modalities av \
+    --local-fusion frame-set-attention \
+    --local-projection-dim 256 \
+    --person-attention-dim 256 \
+    --person-attention-heads 4
+
+F1 用融合后的 G 和直接局部音频 A 构造一个 query，读取当前最新半开因果帧中的可变人物
+集合。人物集合只在 batch 内动态 padding；NULL person 处理没有有效框的时刻。A 仍直接
+进入最终 head。第一版不使用 bbox、track ID、跨帧 GRU 或双向 AV attention。
+
+训练输出会打印 source frame 数、person observation 数、max_people 和 attention 配置。
+推理只需继续提供 local feature root；融合方式和 attention 参数从 checkpoint 恢复。
+正式比较需要重新训练同 split、seed 和优化设置的 pooled AV 对照，不能直接把不同开发
+协议的旧 checkpoint 当作唯一基线。
+
 ## Ego4D 接口
 
 完整命令见 [Ego4D 四组局部增强实验](ego4d_local_baselines.md)：新增 manifest 批量提取、

docs/progressive_feature_fusion_route_2026-09-02.md

@@ -21,14 +21,17 @@
 融合实验从 F1 Frame-set Attention 开始。后续模块只有在前一级通过预先定义的 gate 后
 才实施。
 
-这是一份实验设计，不代表 F1--F6 或 D1 已经实现或验证。
+实现状态（2026-09-02）：F1 的 loader、head、train/infer 路径和回归测试已经加入代码，
+但尚未在远程 PyTorch/L20 环境完成 runtime smoke，不能写成已验证结果。F2--F6 和 D1
+仍然只是计划。
 
 ## 2. 不可改变的实验条件
 
 所有阶段共同遵守：
 
 - 目标与全局 state 均为 5 Hz；第 t 行只能使用截止到 (t+1)/5 秒的媒体。
-- 视频人物特征若为 2 FPS，在 5 Hz 预测点选择时间不晚于 cutoff 的最新人物集合。
+- 视频人物特征若为 2 FPS，在 5 Hz 预测点选择严格早于 cutoff 的最新人物集合，与已有
+  half-open local audio/visual 契约一致。
 - 不允许用未来帧插值；同一 2 FPS 人物集合可以供多个 5 Hz query 读取，但不能伪装成
   新观测去重复更新时序状态。
 - EasyCom Session 4--12 为官方训练集，Session 1--3 为最终测试集。模型和阈值选择使用
@@ -128,11 +131,11 @@ F1 暂不加入 bbox、track GRU、人物自注意力或 audio-to-person 反向
 
 对每个 5 Hz cutoff c_t：
 
-1. 找到 frame_time 小于或等于 c_t 的最新 local frame；
+1. 找到 frame_time 严格小于 c_t 的最新 local frame；
 2. 读取该帧 frame_offsets 范围内全部逐人 feature；
 3. 如果尚无帧，使用仅含 NULL person 的集合；
 4. 不允许 nearest-neighbor 选到 c_t 之后的帧；
-5. 记录 chosen_frame_time，测试所有 chosen_frame_time 小于或等于 cutoff。
+5. 记录 chosen_frame_time，测试所有 chosen_frame_time 严格小于 cutoff。
 
 ### 5.5 F1 必跑对照
 

easycom_head.py

@@ -5,13 +5,16 @@ from __future__ import annotations
 import csv
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Literal, Sequence
+from typing import TYPE_CHECKING, Any, Literal, Sequence
 from uuid import UUID
 
 import numpy as np
 import torch
 from torch import nn
 
+if TYPE_CHECKING:
+    from local_turn_features import FramePersonFeatures
+
 
 OFFICIAL_TRAIN_SESSIONS = (
     "Session_7",
@@ -58,6 +61,7 @@ LAYER_FUSION_MODES = (
     "scalar-mixture",
     "concat-mlp",
 )
+LOCAL_FUSION_MODES = ("pooled", "frame-set-attention")
 DatasetName = Literal["easycom", "ego4d"]
 
 
@@ -77,6 +81,9 @@ class TurnPredictionHead(nn.Module):
         local_visual_dim: int = 0,
         local_audio_dim: int = 0,
         local_projection_dim: int = 256,
+        local_fusion: str = "pooled",
+        person_attention_dim: int = 256,
+        person_attention_heads: int = 4,
     ) -> None:
         super().__init__()
         if num_input_layers <= 0:
@@ -116,6 +123,24 @@ class TurnPredictionHead(nn.Module):
         self.local_visual_dim = local_visual_dim
         self.local_audio_dim = local_audio_dim
         self.local_projection_dim = local_projection_dim
+        if local_fusion not in LOCAL_FUSION_MODES:
+            raise ValueError(
+                f"Unknown local_fusion={local_fusion!r}; expected one of "
+                f"{LOCAL_FUSION_MODES}"
+            )
+        if person_attention_dim <= 0 or person_attention_heads <= 0:
+            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 local_fusion == "frame-set-attention" and (
+            local_visual_dim <= 0 or local_audio_dim <= 0
+        ):
+            raise ValueError(
+                "Frame-set Attention requires both local visual people and audio"
+            )
+        self.local_fusion = local_fusion
+        self.person_attention_dim = person_attention_dim
+        self.person_attention_heads = person_attention_heads
         if self.learnable_layer_mixture:
             # Zero initialization exactly reproduces a uniform layer mean at
             # step zero while allowing training to select useful query layers.
@@ -133,7 +158,12 @@ class TurnPredictionHead(nn.Module):
             self.layer_fusion_net = None
         self.local_projections = nn.ModuleDict()
         fused_size = hidden_size
-        for name, dimension in (("visual", local_visual_dim), ("audio", local_audio_dim)):
+        projected_modalities = (
+            (("audio", local_audio_dim),)
+            if local_fusion == "frame-set-attention"
+            else (("visual", local_visual_dim), ("audio", local_audio_dim))
+        )
+        for name, dimension in projected_modalities:
             if dimension:
                 self.local_projections[name] = (
                     nn.Sequential(
@@ -145,6 +175,36 @@ class TurnPredictionHead(nn.Module):
                     else nn.Identity()
                 )
                 fused_size += (local_projection_dim or dimension) + 1
+        self.person_input_projection: nn.Module | None = None
+        self.person_state_query: nn.Module | None = None
+        self.person_audio_query: nn.Module | None = None
+        self.person_attention: nn.MultiheadAttention | None = None
+        self.null_person: nn.Parameter | None = None
+        if local_fusion == "frame-set-attention":
+            self.person_input_projection = nn.Sequential(
+                nn.LayerNorm(local_visual_dim),
+                nn.Linear(local_visual_dim, person_attention_dim),
+                nn.GELU(),
+            )
+            self.person_state_query = nn.Sequential(
+                nn.LayerNorm(hidden_size),
+                nn.Linear(hidden_size, person_attention_dim),
+            )
+            self.person_audio_query = nn.Sequential(
+                nn.LayerNorm(local_audio_dim),
+                nn.Linear(local_audio_dim, person_attention_dim),
+            )
+            self.person_attention = nn.MultiheadAttention(
+                person_attention_dim,
+                person_attention_heads,
+                dropout=dropout,
+                batch_first=True,
+            )
+            self.null_person = nn.Parameter(
+                torch.zeros(1, 1, person_attention_dim)
+            )
+            nn.init.normal_(self.null_person, std=0.02)
+            fused_size += person_attention_dim + 1
         self.net = nn.Sequential(
             nn.LayerNorm(fused_size),
             nn.Linear(fused_size, hidden_size // 2),
@@ -190,17 +250,99 @@ class TurnPredictionHead(nn.Module):
 
     @property
     def local_input_size(self) -> int:
+        if self.local_fusion == "frame-set-attention":
+            return self.local_audio_dim + 1
         return sum(d + 1 for d in (self.local_visual_dim, self.local_audio_dim) if d)
 
+    def _frame_set_context(
+        self,
+        states: torch.Tensor,
+        local: torch.Tensor,
+        person_features: torch.Tensor | None,
+        person_mask: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if states.ndim != 2:
+            raise ValueError("Frame-set Attention expects batched [B,D] states")
+        if local.shape != (len(states), self.local_audio_dim + 1):
+            raise ValueError(
+                "Frame-set Attention requires aligned [audio,valid] local features"
+            )
+        if (
+            person_features is None
+            or person_features.ndim != 3
+            or person_features.shape[0] != len(states)
+            or person_features.shape[2] != self.local_visual_dim
+        ):
+            raise ValueError(
+                "Frame-set Attention requires person features [B,P,visual_dim]"
+            )
+        if (
+            person_mask is None
+            or person_mask.dtype != torch.bool
+            or person_mask.shape != person_features.shape[:2]
+        ):
+            raise ValueError("Frame-set Attention requires a boolean person mask [B,P]")
+        if not torch.isfinite(person_features).all():
+            raise ValueError("Person features contain NaN/Inf")
+        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
+
+        audio_values = local[:, : self.local_audio_dim]
+        audio_valid = local[:, self.local_audio_dim : self.local_audio_dim + 1]
+        audio_projected = self.local_projections["audio"](audio_values) * audio_valid
+        query = (
+            self.person_state_query(states)
+            + self.person_audio_query(audio_values) * audio_valid
+        )
+        people = self.person_input_projection(person_features)
+        null = self.null_person.expand(len(states), -1, -1)
+        memory = torch.cat((null, people), dim=1)
+        null_mask = torch.ones(
+            (len(states), 1),
+            dtype=torch.bool,
+            device=person_mask.device,
+        )
+        memory_valid = torch.cat((null_mask, person_mask), dim=1)
+        context, _ = self.person_attention(
+            query[:, None, :],
+            memory,
+            memory,
+            key_padding_mask=~memory_valid,
+            need_weights=False,
+        )
+        person_valid = person_mask.any(dim=1, keepdim=True).to(states.dtype)
+        return audio_projected, audio_valid, torch.cat(
+            (context[:, 0, :], person_valid),
+            dim=-1,
+        )
+
     def _shared_features(
         self,
         states: torch.Tensor,
         local: torch.Tensor | None = None,
+        person_features: torch.Tensor | None = None,
+        person_mask: torch.Tensor | None = None,
     ) -> torch.Tensor:
         """Return the activity MLP representation shared by both readouts."""
 
         states = self._fuse_layers(states)
-        if self.local_input_size:
+        if self.local_fusion == "frame-set-attention":
+            if local is None:
+                raise ValueError("Frame-set Attention requires local audio features")
+            audio, audio_valid, person_context = self._frame_set_context(
+                states,
+                local,
+                person_features,
+                person_mask,
+            )
+            states = torch.cat(
+                (states, audio, audio_valid, person_context),
+                dim=-1,
+            )
+        elif self.local_input_size:
             if local is None or local.shape != (*states.shape[:-1], self.local_input_size):
                 raise ValueError(
                     f"Head requires aligned local features of width {self.local_input_size}"
@@ -220,14 +362,25 @@ class TurnPredictionHead(nn.Module):
             states = torch.cat(pieces, dim=-1)
         elif local is not None and local.shape[-1] != 0:
             raise ValueError("Local features supplied to a baseline head")
+        if self.local_fusion != "frame-set-attention" and (
+            person_features is not None or person_mask is not None
+        ):
+            raise ValueError("Person sets supplied to a non-frame-set head")
         return self.net[:-1](states)
 
     def forward(
         self,
         states: torch.Tensor,
         local: torch.Tensor | None = None,
+        person_features: torch.Tensor | None = None,
+        person_mask: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        features = self._shared_features(states, local)
+        features = self._shared_features(
+            states,
+            local,
+            person_features,
+            person_mask,
+        )
         logits = self.net[-1](features)
         return logits.view(
             *features.shape[:-1], self.num_horizons, self.num_classes
@@ -237,6 +390,8 @@ class TurnPredictionHead(nn.Module):
         self,
         states: torch.Tensor,
         local: torch.Tensor | None = None,
+        person_features: torch.Tensor | None = None,
+        person_mask: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Return activity logits and binary transition logits."""
 
@@ -244,7 +399,12 @@ class TurnPredictionHead(nn.Module):
             raise RuntimeError(
                 "This head was created without transition_auxiliary=True"
             )
-        features = self._shared_features(states, local)
+        features = self._shared_features(
+            states,
+            local,
+            person_features,
+            person_mask,
+        )
         activity_logits = self.net[-1](features).view(
             *features.shape[:-1],
             self.num_horizons,
@@ -734,6 +894,7 @@ class SessionData:
     label_indices: torch.LongTensor
     feature_metadata: dict[str, Any]
     local_features: torch.Tensor | None = None
+    local_person_features: "FramePersonFeatures | None" = None
 
     @property
     def hidden_size(self) -> int:
@@ -768,6 +929,7 @@ def load_session_data(
     local_feature_root: str | Path | None = None,
     local_modalities: str = "av",
     local_contract: dict[str, Any] | None = None,
+    local_fusion: str = "pooled",
 ) -> SessionData:
     """Load and validate one saved Omni feature payload and target array."""
 
@@ -864,22 +1026,55 @@ def load_session_data(
         "".join(str(int(value)) for value in row.tolist()): int(count)
         for row, count in zip(unique_rows, row_counts)
     }
+    if local_fusion not in LOCAL_FUSION_MODES:
+        raise ValueError(f"Unknown local fusion: {local_fusion}")
     local = None
+    local_people = None
     if local_feature_root is not None:
-        from local_turn_features import load_local_features
+        local_path = resolve_feature_path(
-
+            local_feature_root,
-        local, contract = load_local_features(
+            name,
-            resolve_feature_path(local_feature_root, name, dataset=dataset),
-            payload,
-            sample_id=name,
             dataset=dataset,
-            modalities=local_modalities,
-            expected_contract=local_contract,
         )
+        if local_fusion == "frame-set-attention":
+            if local_modalities != "av":
+                raise ValueError(
+                    "Frame-set Attention requires --local-modalities av"
+                )
+            from local_turn_features import load_frame_set_features
+
+            local, local_people, contract = load_frame_set_features(
+                local_path,
+                payload,
+                sample_id=name,
+                dataset=dataset,
+                expected_contract=local_contract,
+            )
+        else:
+            from local_turn_features import load_local_features
+
+            local, contract = load_local_features(
+                local_path,
+                payload,
+                sample_id=name,
+                dataset=dataset,
+                modalities=local_modalities,
+                expected_contract=local_contract,
+            )
         metadata["local_contract"] = contract
     elif local_contract is not None:
         raise ValueError("--local-feature-root is required for this checkpoint")
-    return SessionData(name, states, targets, label_indices, metadata, local)
+    elif local_fusion != "pooled":
+        raise ValueError("Frame-set Attention requires --local-feature-root")
+    return SessionData(
+        name,
+        states,
+        targets,
+        label_indices,
+        metadata,
+        local,
+        local_people,
+    )
 
 
 def average_precision_metrics(
@@ -924,6 +1119,9 @@ def checkpoint_head_config(payload: dict[str, Any]) -> dict[str, Any]:
     config.setdefault("local_visual_dim", 0)
     config.setdefault("local_audio_dim", 0)
     config.setdefault("local_projection_dim", 256)
+    config.setdefault("local_fusion", "pooled")
+    config.setdefault("person_attention_dim", 256)
+    config.setdefault("person_attention_heads", 4)
     config.setdefault(
         "layer_fusion",
         (

infer_turn_head.py

@@ -30,6 +30,7 @@ from easycom_head import (
     output_space_for_loss,
     transition_targets_from_activity,
 )
+from local_turn_features import FramePersonFeatures
 
 
 def parse_args() -> argparse.Namespace:
@@ -118,6 +119,7 @@ def _predict_states(
     show_progress: bool,
     loss_name: str,
     local_features: torch.Tensor | None = None,
+    local_person_features: FramePersonFeatures | None = None,
 ) -> tuple[np.ndarray, np.ndarray | None]:
     chunks: list[torch.Tensor] = []
     transition_chunks: list[torch.Tensor] = []
@@ -139,11 +141,35 @@ def _predict_states(
             local_features[start : start + batch_size].to(device=device, dtype=torch.float32)
             if local_features is not None else None
         )
+        person_batch = None
+        person_mask = None
+        if local_person_features is not None:
+            rows = torch.arange(
+                start,
+                min(start + batch_size, len(states)),
+            )
+            person_batch, person_mask = local_person_features.padded(rows)
+            person_batch = person_batch.to(
+                device=device,
+                dtype=torch.float32,
+                non_blocking=True,
+            )
+            person_mask = person_mask.to(device=device, non_blocking=True)
         if head.transition_auxiliary:
-            activity_logits, transition_logits = head.forward_with_transition(batch, side)
+            activity_logits, transition_logits = head.forward_with_transition(
+                batch,
+                side,
+                person_batch,
+                person_mask,
+            )
             transition_chunks.append(transition_logits.sigmoid().cpu())
         else:
-            activity_logits = head(batch, side)
+            activity_logits = head(
+                batch,
+                side,
+                person_batch,
+                person_mask,
+            )
         chunks.append(
             marginal_probabilities_from_logits(activity_logits, loss_name).cpu()
         )
@@ -403,6 +429,7 @@ def main() -> None:
             local_feature_root=args.local_feature_root,
             local_modalities=local_contract["modalities"] if local_contract else "av",
             local_contract=local_contract,
+            local_fusion=head.local_fusion,
         )
         if data.hidden_size != int(head_config["hidden_size"]):
             raise ValueError(
@@ -423,6 +450,7 @@ def main() -> None:
             show_progress=not args.no_progress,
             loss_name=loss_name,
             local_features=data.local_features,
+            local_person_features=data.local_person_features,
         )
         predictions, valid = _align_to_egospeak_timeline(
             probabilities,
@@ -484,6 +512,7 @@ def main() -> None:
             "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_contract": local_contract,
             "transition_auxiliary": head.transition_auxiliary,
             "transition_objective": checkpoint_payload.get(

local_turn_features.py

@@ -1,4 +1,4 @@
-"""Validated, causal local-feature sidecars for the pooled B1 baseline.
+"""Validated, causal pooled and frame-set local-feature sidecars.
 
 No model/decoder imports: these utilities are also used by CPU head training.
 Sidecars store encoder-native vectors, never a randomly initialized projection.
@@ -6,6 +6,7 @@ Sidecars store encoder-native vectors, never a randomly initialized projection.
 
 from __future__ import annotations
 
+from dataclasses import dataclass
 import math
 from pathlib import Path
 from typing import Any
@@ -16,6 +17,150 @@ LOCAL_FORMAT = "qwen-local-pooled-v1"
 MODALITIES = {"visual": ("visual",), "audio": ("audio",), "av": ("visual", "audio")}
 
 
+@dataclass(frozen=True)
+class FramePersonFeatures:
+    """Ragged person sets indexed by aligned 5 Hz rows without frame duplication."""
+
+    features: torch.Tensor
+    frame_offsets: torch.LongTensor
+    row_to_frame: torch.LongTensor
+    chosen_frame_times: torch.Tensor
+
+    def __post_init__(self) -> None:
+        if (
+            not isinstance(self.features, torch.Tensor)
+            or self.features.ndim != 2
+            or not self.features.is_floating_point()
+            or not torch.isfinite(self.features).all()
+            or self.features.shape[1] <= 0
+        ):
+            raise ValueError("Person features must be finite floating [N,D]")
+        if (
+            not isinstance(self.frame_offsets, torch.Tensor)
+            or self.frame_offsets.dtype != torch.long
+            or self.frame_offsets.ndim != 1
+            or not len(self.frame_offsets)
+            or int(self.frame_offsets[0]) != 0
+            or int(self.frame_offsets[-1]) != len(self.features)
+            or bool((self.frame_offsets[1:] < self.frame_offsets[:-1]).any())
+        ):
+            raise ValueError("Invalid person frame_offsets")
+        if (
+            not isinstance(self.row_to_frame, torch.Tensor)
+            or self.row_to_frame.dtype != torch.long
+            or self.row_to_frame.ndim != 1
+        ):
+            raise ValueError("row_to_frame must be int64 [T]")
+        num_frames = len(self.frame_offsets) - 1
+        if bool(((self.row_to_frame < -1) | (self.row_to_frame >= num_frames)).any()):
+            raise ValueError("row_to_frame contains an invalid frame index")
+        if (
+            not isinstance(self.chosen_frame_times, torch.Tensor)
+            or self.chosen_frame_times.dtype != torch.float64
+            or self.chosen_frame_times.shape != self.row_to_frame.shape
+        ):
+            raise ValueError("chosen_frame_times must be float64 [T]")
+        missing = self.row_to_frame < 0
+        if not torch.equal(torch.isnan(self.chosen_frame_times), missing):
+            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")
+
+    def __len__(self) -> int:
+        return len(self.row_to_frame)
+
+    @property
+    def dimension(self) -> int:
+        return int(self.features.shape[1])
+
+    @property
+    def max_people(self) -> int:
+        counts = self.frame_offsets[1:] - self.frame_offsets[:-1]
+        return int(counts.max()) if len(counts) else 0
+
+    def select_rows(self, indices: torch.Tensor) -> "FramePersonFeatures":
+        indices = torch.as_tensor(indices)
+        if indices.dtype != torch.long or indices.ndim != 1:
+            raise ValueError("Person row selection must be int64 [N]")
+        if len(indices) and (int(indices.min()) < 0 or int(indices.max()) >= len(self)):
+            raise IndexError("Person row selection is out of range")
+        return FramePersonFeatures(
+            self.features,
+            self.frame_offsets,
+            self.row_to_frame[indices],
+            self.chosen_frame_times[indices],
+        )
+
+    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."""
+
+        if row_indices is None:
+            frames = self.row_to_frame
+        else:
+            row_indices = torch.as_tensor(row_indices)
+            if row_indices.dtype != torch.long or row_indices.ndim != 1:
+                raise ValueError("Batch row indices must be int64 [B]")
+            if len(row_indices) and (
+                int(row_indices.min()) < 0 or int(row_indices.max()) >= len(self)
+            ):
+                raise IndexError("Batch person row index is out of range")
+            frames = self.row_to_frame[row_indices]
+        valid_frames = frames >= 0
+        safe_frames = frames.clamp_min(0)
+        starts = torch.zeros_like(frames)
+        counts = torch.zeros_like(frames)
+        if bool(valid_frames.any()):
+            starts[valid_frames] = self.frame_offsets[safe_frames[valid_frames]]
+            counts[valid_frames] = (
+                self.frame_offsets[safe_frames[valid_frames] + 1]
+                - starts[valid_frames]
+            )
+        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)
+        values = self.features[safe_indices]
+        return values.masked_fill(~mask[..., None], 0), mask
+
+    @classmethod
+    def concatenate(
+        cls,
+        parts: list["FramePersonFeatures"],
+    ) -> "FramePersonFeatures":
+        if not parts:
+            raise ValueError("Cannot concatenate no person feature sets")
+        dimension = parts[0].dimension
+        if any(part.dimension != dimension for part in parts):
+            raise ValueError("Person feature dimensions differ across sessions")
+        feature_parts = []
+        offset_parts = [torch.zeros(1, dtype=torch.long)]
+        row_parts = []
+        time_parts = []
+        feature_base = 0
+        frame_base = 0
+        for part in parts:
+            feature_parts.append(part.features)
+            offset_parts.append(part.frame_offsets[1:] + feature_base)
+            rows = part.row_to_frame.clone()
+            valid = rows >= 0
+            rows[valid] += frame_base
+            row_parts.append(rows)
+            time_parts.append(part.chosen_frame_times)
+            feature_base += len(part.features)
+            frame_base += len(part.frame_offsets) - 1
+        return cls(
+            torch.cat(feature_parts, dim=0),
+            torch.cat(offset_parts, dim=0),
+            torch.cat(row_parts, dim=0),
+            torch.cat(time_parts, dim=0),
+        )
+
+
 def reference_alignment(payload: dict[str, Any]) -> tuple[torch.Tensor, torch.Tensor]:
     """Require the same full, zero-origin 5 Hz timeline used by head training."""
     states = payload.get("states")
@@ -34,16 +179,15 @@ def reference_alignment(payload: dict[str, Any]) -> tuple[torch.Tensor, torch.Te
     return times, indices
 
 
-def load_local_features(
+def _load_validated_payload(
     path: Path,
     reference: dict[str, Any],
     *,
     sample_id: str,
     dataset: str,
     modalities: str,
-    expected_contract: dict[str, Any] | None = None,
+    expected_contract: dict[str, Any] | None,
-) -> tuple[torch.Tensor, dict[str, Any]]:
+) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any], dict[str, Any], dict[str, Any]]:
-    """Pack each requested modality as [raw vector, validity bit]. Fail closed."""
     if modalities not in MODALITIES:
         raise ValueError(f"Unknown local modalities: {modalities}")
     times, indices = reference_alignment(reference)
@@ -66,6 +210,27 @@ def load_local_features(
     contract = {"format": LOCAL_FORMAT, "modalities": modalities, "config": config}
     if expected_contract is not None and contract != expected_contract:
         raise ValueError(f"Local encoder/pooling contract differs across sessions/checkpoint: {path}")
+    return times, indices, payload, config, contract
+
+
+def load_local_features(
+    path: Path,
+    reference: dict[str, Any],
+    *,
+    sample_id: str,
+    dataset: str,
+    modalities: str,
+    expected_contract: dict[str, Any] | None = None,
+) -> tuple[torch.Tensor, dict[str, Any]]:
+    """Pack each requested modality as [raw vector, validity bit]. Fail closed."""
+    times, _, payload, config, contract = _load_validated_payload(
+        path,
+        reference,
+        sample_id=sample_id,
+        dataset=dataset,
+        modalities=modalities,
+        expected_contract=expected_contract,
+    )
     chunks = []
     for name in MODALITIES[modalities]:
         values = payload.get(name)
@@ -85,6 +250,109 @@ def load_local_features(
     return torch.cat(chunks, dim=-1), contract
 
 
+def load_frame_set_features(
+    path: Path,
+    reference: dict[str, Any],
+    *,
+    sample_id: str,
+    dataset: str,
+    expected_contract: dict[str, Any] | None = None,
+) -> tuple[torch.Tensor, FramePersonFeatures, dict[str, Any]]:
+    """Load direct audio and the latest causal per-frame person set per 5 Hz row."""
+
+    times, _, payload, config, contract = _load_validated_payload(
+        path,
+        reference,
+        sample_id=sample_id,
+        dataset=dataset,
+        modalities="av",
+        expected_contract=expected_contract,
+    )
+    audio = payload.get("audio")
+    audio_valid = payload.get("audio_valid")
+    audio_dim = config.get("audio_dim")
+    if not isinstance(audio_dim, int) or audio_dim <= 0:
+        raise ValueError(f"Missing audio_dim: {path}")
+    if not isinstance(audio, torch.Tensor) or audio.shape != (len(times), audio_dim):
+        raise ValueError(f"Invalid audio shape: {path}")
+    if not audio.is_floating_point() or not torch.isfinite(audio).all():
+        raise ValueError(f"Non-finite/non-floating audio features: {path}")
+    if (
+        not isinstance(audio_valid, torch.Tensor)
+        or audio_valid.dtype != torch.bool
+        or audio_valid.shape != (len(times),)
+    ):
+        raise ValueError(f"Invalid audio_valid mask: {path}")
+    if torch.count_nonzero(audio[~audio_valid]):
+        raise ValueError(f"Invalid audio rows must be stored as zeros: {path}")
+
+    visual_dim = config.get("visual_dim")
+    persons = payload.get("persons")
+    if not isinstance(visual_dim, int) or visual_dim <= 0:
+        raise ValueError(f"Missing visual_dim: {path}")
+    if not isinstance(persons, dict):
+        raise ValueError(f"Missing ragged person features: {path}")
+    features = persons.get("features")
+    frame_times = persons.get("frame_times")
+    frame_offsets = persons.get("frame_offsets")
+    records = persons.get("records")
+    if (
+        not isinstance(features, torch.Tensor)
+        or features.ndim != 2
+        or features.shape[1] != visual_dim
+        or not features.is_floating_point()
+        or not torch.isfinite(features).all()
+    ):
+        raise ValueError(f"Invalid person feature shape/content: {path}")
+    if (
+        not isinstance(frame_times, torch.Tensor)
+        or frame_times.ndim != 1
+        or not frame_times.is_floating_point()
+    ):
+        raise ValueError(f"Invalid person frame_times: {path}")
+    frame_times = frame_times.to(dtype=torch.float64)
+    if (
+        not torch.isfinite(frame_times).all()
+        or bool((frame_times < 0).any())
+        or bool((frame_times[1:] <= frame_times[:-1]).any())
+    ):
+        raise ValueError(f"Person frame_times must be finite, nonnegative and increasing: {path}")
+    if (
+        not isinstance(frame_offsets, torch.Tensor)
+        or frame_offsets.dtype != torch.long
+        or frame_offsets.shape != (len(frame_times) + 1,)
+        or not len(frame_offsets)
+        or int(frame_offsets[0]) != 0
+        or int(frame_offsets[-1]) != len(features)
+        or bool((frame_offsets[1:] < frame_offsets[:-1]).any())
+    ):
+        raise ValueError(f"Invalid person frame_offsets: {path}")
+    if not isinstance(records, list) or len(records) != len(features):
+        raise ValueError(f"Person records/features length mismatch: {path}")
+    if len(frame_times) and float(frame_times[-1]) > float(times[-1]) + 1e-8:
+        raise ValueError(f"Person frame timeline extends beyond the reference: {path}")
+
+    # Match the sidecar's half-open policy: a frame stamped exactly at the
+    # state cutoff is not yet part of that state's observed prefix.
+    row_to_frame = torch.searchsorted(frame_times, times, right=False) - 1
+    chosen = torch.full(times.shape, float("nan"), dtype=torch.float64)
+    has_frame = row_to_frame >= 0
+    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}")
+    frame_sets = FramePersonFeatures(
+        features,
+        frame_offsets,
+        row_to_frame.to(dtype=torch.long),
+        chosen,
+    )
+    direct_audio = torch.cat(
+        (audio.float(), audio_valid[:, None].float()),
+        dim=-1,
+    )
+    return direct_audio, frame_sets, 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}

test_local_turn_features.py

@@ -28,8 +28,9 @@ from extract_turn_local_features import (
     main as extract_main,
 )
 from local_turn_features import (
-    LOCAL_FORMAT, causal_frame_ticks, crop_box, easycom_boxes, load_local_features,
+    LOCAL_FORMAT, FramePersonFeatures, causal_frame_ticks, crop_box, easycom_boxes,
-    local_head_dimensions, pooled_recent_visual, reference_alignment,
+    load_frame_set_features, load_local_features, local_head_dimensions,
+    pooled_recent_visual, reference_alignment,
 )
 from train_turn_head import _load_examples
 
@@ -47,6 +48,17 @@ def fixture(n=20, name="Session_4", dataset="easycom"):
                        "causal_policy": "half-open-trailing-window", "uses_speech_labels": False}}
     side["visual_valid"][0] = False
     side["visual"][0] = 0
+    frame_times = torch.arange(max(1, (n + 2) // 3), dtype=torch.float64) * .5
+    counts = [2 if index % 3 == 0 else (0 if index % 3 == 1 else 1)
+              for index in range(len(frame_times))]
+    offsets = torch.tensor([0, *np.cumsum(counts).tolist()], dtype=torch.long)
+    person_features = torch.randn(int(offsets[-1]), 6)
+    side["persons"] = {
+        "features": person_features,
+        "frame_times": frame_times,
+        "frame_offsets": offsets,
+        "records": [{"id": str(index)} for index in range(len(person_features))],
+    }
     labels = np.tile(np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1], [0, 1, 1]], dtype=np.float32), (n // 4 + 1, 1))[:n]
     return reference, side, labels
 
@@ -280,6 +292,148 @@ class LocalTests(unittest.TestCase):
             with self.assertRaisesRegex(ValueError, "times"):
                 reference_alignment({**reference, "times": reference["times"] + .2})
 
+    def test_frame_set_loader_is_causal_ragged_and_fail_closed(self):
+        reference, side, _ = fixture()
+        with tempfile.TemporaryDirectory() as directory:
+            path = Path(directory) / "Session_4.pt"
+            torch.save(side, path)
+            audio, people, contract = load_frame_set_features(
+                path,
+                reference,
+                sample_id="Session_4",
+                dataset="easycom",
+            )
+            self.assertEqual(audio.shape, (20, 6))
+            self.assertEqual(contract["modalities"], "av")
+            self.assertEqual(people.dimension, 6)
+            self.assertTrue(
+                torch.all(
+                    people.chosen_frame_times[
+                        torch.isfinite(people.chosen_frame_times)
+                    ]
+                    < reference["times"][
+                        torch.isfinite(people.chosen_frame_times)
+                    ]
+                )
+            )
+            # At 0.2/0.4s, both rows reuse the 0.0s frame with two people.
+            values, mask = people.padded(torch.tensor([0, 1, 2]))
+            self.assertEqual(values.shape, (3, 2, 6))
+            self.assertEqual(mask.tolist(), [[True, True], [True, True], [False, False]])
+            self.assertTrue(torch.equal(values[0], values[1]))
+            self.assertEqual(torch.count_nonzero(values[2]), 0)
+            # The frame stamped exactly 1.0s is excluded at the 1.0s cutoff.
+            self.assertEqual(people.chosen_frame_times[4].item(), .5)
+
+            broken_cases = []
+            broken = copy.deepcopy(side)
+            broken["persons"]["frame_offsets"][-1] -= 1
+            broken_cases.append(broken)
+            broken = copy.deepcopy(side)
+            broken["persons"]["frame_times"][1] = broken["persons"]["frame_times"][0]
+            broken_cases.append(broken)
+            broken = copy.deepcopy(side)
+            broken["persons"]["records"] = []
+            broken_cases.append(broken)
+            broken = copy.deepcopy(side)
+            broken["persons"]["features"][0, 0] = float("nan")
+            broken_cases.append(broken)
+            for index, broken in enumerate(broken_cases):
+                torch.save(broken, path)
+                with self.subTest(index=index), self.assertRaises(ValueError):
+                    load_frame_set_features(
+                        path,
+                        reference,
+                        sample_id="Session_4",
+                        dataset="easycom",
+                    )
+
+    def test_frame_person_concatenation_selection_and_empty_padding(self):
+        features = torch.tensor([[1., 2.], [3., 4.]])
+        a = FramePersonFeatures(
+            features,
+            torch.tensor([0, 2], dtype=torch.long),
+            torch.tensor([0, -1], dtype=torch.long),
+            torch.tensor([0., float("nan")], dtype=torch.float64),
+        )
+        b = FramePersonFeatures(
+            torch.empty(0, 2),
+            torch.tensor([0, 0], dtype=torch.long),
+            torch.tensor([0], dtype=torch.long),
+            torch.tensor([.5], dtype=torch.float64),
+        )
+        combined = FramePersonFeatures.concatenate([a, b])
+        self.assertEqual(len(combined), 3)
+        self.assertEqual(combined.row_to_frame.tolist(), [0, -1, 1])
+        selected = combined.select_rows(torch.tensor([2, 0]))
+        values, mask = selected.padded()
+        self.assertEqual(mask.tolist(), [[False, False], [True, True]])
+        self.assertEqual(torch.count_nonzero(values[0]), 0)
+        self.assertTrue(torch.equal(values[1], features))
+
+    def test_frame_set_head_is_permutation_invariant_and_masks_padding(self):
+        torch.manual_seed(7)
+        head = TurnPredictionHead(
+            hidden_size=8,
+            num_input_layers=4,
+            layer_fusion="concat-mlp",
+            local_visual_dim=6,
+            local_audio_dim=5,
+            local_projection_dim=3,
+            local_fusion="frame-set-attention",
+            person_attention_dim=4,
+            person_attention_heads=2,
+            transition_auxiliary=True,
+            dropout=0,
+        ).eval()
+        states = torch.randn(3, 4, 8)
+        local = torch.randn(3, 6)
+        local[:, -1] = 1
+        people = torch.randn(3, 3, 6)
+        mask = torch.tensor(
+            [[True, True, False], [True, False, False], [False, False, False]]
+        )
+        people[~mask] = 0
+        logits, transitions = head.forward_with_transition(
+            states,
+            local,
+            people,
+            mask,
+        )
+        self.assertEqual(logits.shape, (3, 10, 3))
+        self.assertEqual(transitions.shape, (3, 10))
+        (logits.square().mean() + transitions.square().mean()).backward()
+        self.assertIsNotNone(head.person_attention.in_proj_weight.grad)
+
+        permutation = torch.tensor([1, 0, 2])
+        permuted = head(states, local, people[:, permutation], mask[:, permutation])
+        self.assertTrue(torch.allclose(logits, permuted, atol=1e-6, rtol=1e-6))
+        changed = people.clone()
+        changed[~mask] = 1000
+        self.assertTrue(torch.allclose(logits, head(states, local, changed, mask)))
+        with self.assertRaisesRegex(ValueError, "person features"):
+            head(states, local)
+
+        config = {
+            "hidden_size": 8,
+            "num_input_layers": 4,
+            "layer_fusion": "concat-mlp",
+            "local_visual_dim": 6,
+            "local_audio_dim": 5,
+            "local_projection_dim": 3,
+            "local_fusion": "frame-set-attention",
+            "person_attention_dim": 4,
+            "person_attention_heads": 2,
+            "transition_auxiliary": True,
+            "dropout": 0,
+        }
+        restored = TurnPredictionHead(
+            **checkpoint_head_config({"head_config": config})
+        )
+        restored.load_state_dict(head.state_dict())
+        restored.eval()
+        self.assertTrue(torch.equal(restored(states, local, people, 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)):
@@ -338,7 +492,8 @@ class LocalTests(unittest.TestCase):
             np.save(root / "targets/Session_4.npy", targets)
             data = _load_examples(root / "features", root / "targets", ["Session_4"], "open-stream-replay",
                                   "easycom", {"Session_4": {1, 4}}, root / "local", "audio")
-            self.assertTrue(torch.equal(data[-1][:, :5], side["audio"][[1, 4]]))
+            self.assertTrue(torch.equal(data[-2][:, :5], side["audio"][[1, 4]]))
+            self.assertIsNone(data[-1])
             self.assertTrue(torch.equal(data[0], reference["states"][[1, 4]]))
             self.assertTrue(torch.equal(data[1][:, 0], torch.from_numpy(targets[[2, 5]])))
 
@@ -562,7 +717,7 @@ class RoundTripTests(unittest.TestCase):
                 np.save(root / "targets" / f"{name}.npy", targets)
             common = ["--feature-root", str(root / "features"), "--target-root", str(root / "targets"),
                       "--device", "cpu", "--no-progress", "--expected-mode", "open-stream-replay"]
-            for mode in ("baseline", "audio", "visual", "av", "ego4d"):
+            for mode in ("baseline", "audio", "visual", "av", "frame-set", "ego4d"):
                 checkpoint = root / f"{mode}.pth"
                 train = [sys.executable, str(repo / "train_turn_head.py"), *common,
                          "--epochs", "1", "--warmup-epochs", "0", "--batch-size", "4",
@@ -579,8 +734,15 @@ class RoundTripTests(unittest.TestCase):
                     train += ["--dev-sessions", "5", "--local-projection-dim", "0"]
                 if mode == "av":
                     train += ["--transition-objective", "r1-supervised"]
+                if mode == "frame-set":
+                    train += [
+                        "--local-fusion", "frame-set-attention",
+                        "--person-attention-dim", "4",
+                        "--person-attention-heads", "2",
+                    ]
                 if mode != "baseline":
-                    train += ["--local-feature-root", str(root / "local"), "--local-modalities", "av" if mode == "ego4d" else mode]
+                    modality = "av" if mode in ("ego4d", "frame-set") else mode
+                    train += ["--local-feature-root", str(root / "local"), "--local-modalities", modality]
                     infer += ["--local-feature-root", str(root / "local")]
                 for command in (train, infer):
                     completed = subprocess.run(command, cwd=repo, env=env, capture_output=True, text=True, timeout=60)
@@ -590,6 +752,8 @@ class RoundTripTests(unittest.TestCase):
                 self.assertEqual(bool(metrics["local_contract"]), mode != "baseline")
                 if mode == "av":
                     self.assertTrue(metrics["transition_auxiliary"])
+                if mode == "frame-set":
+                    self.assertEqual(metrics["local_fusion"], "frame-set-attention")
             # The reporting command uses the same metric keys as inference.
             completed = subprocess.run(
                 [sys.executable, str(repo / "summarize_local_baselines.py"), "--root", str(root)],

train_turn_head.py

@@ -24,6 +24,7 @@ from easycom_head import (
     CLASS_NAMES,
     DATASET_NAMES,
     LAYER_FUSION_MODES,
+    LOCAL_FUSION_MODES,
     LOSS_NAMES,
     NUM_HORIZONS,
     OFFICIAL_TRAIN_SESSIONS,
@@ -46,6 +47,7 @@ from easycom_head import (
     transition_targets_from_activity,
     validate_disjoint_sessions,
 )
+from local_turn_features import FramePersonFeatures
 
 
 def parse_args() -> argparse.Namespace:
@@ -58,10 +60,18 @@ def parse_args() -> argparse.Namespace:
         help="Optional frozen local-feature sidecars; omit for the original baseline.",
     )
     parser.add_argument("--local-modalities", choices=("visual", "audio", "av"), default="av")
+    parser.add_argument(
+        "--local-fusion",
+        choices=LOCAL_FUSION_MODES,
+        default="pooled",
+        help="pooled keeps the existing baseline; frame-set-attention queries the latest causal person set.",
+    )
     parser.add_argument(
         "--local-projection-dim", type=int, default=256,
         help="Trainable projection per local modality; 0 retains encoder dimensions.",
     )
+    parser.add_argument("--person-attention-dim", type=int, default=256)
+    parser.add_argument("--person-attention-heads", type=int, default=4)
     parser.add_argument(
         "--dataset",
         choices=DATASET_NAMES,
@@ -165,6 +175,17 @@ def parse_args() -> argparse.Namespace:
 def _validate_args(args: argparse.Namespace) -> None:
     if args.local_projection_dim < 0:
         raise ValueError("--local-projection-dim must be non-negative")
+    if args.person_attention_dim <= 0 or args.person_attention_heads <= 0:
+        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.local_fusion == "frame-set-attention" and (
+        args.local_feature_root is None or args.local_modalities != "av"
+    ):
+        raise ValueError(
+            "--local-fusion frame-set-attention requires --local-feature-root "
+            "and --local-modalities av"
+        )
     if args.epochs <= 0:
         raise ValueError("--epochs must be positive")
     if args.batch_size <= 0:
@@ -225,6 +246,7 @@ def _load_examples(
     local_modalities: str = "av",
     local_contract: dict[str, Any] | None = None,
     compute_transition_targets: bool = True,
+    local_fusion: str = "pooled",
 ) -> tuple[
     torch.Tensor,
     torch.Tensor,
@@ -233,11 +255,13 @@ def _load_examples(
     int,
     list[dict[str, Any]],
     torch.Tensor,
+    FramePersonFeatures | None,
 ]:
     all_states: list[torch.Tensor] = []
     all_targets: list[torch.Tensor] = []
     all_transition_targets: list[torch.Tensor] = []
     all_local: list[torch.Tensor] = []
+    all_people: list[FramePersonFeatures] = []
     metadata: list[dict[str, Any]] = []
     hidden_size: int | None = None
     num_input_layers: int | None = None
@@ -251,6 +275,7 @@ def _load_examples(
             local_feature_root=local_feature_root,
             local_modalities=local_modalities,
             local_contract=local_contract,
+            local_fusion=local_fusion,
         )
         local_contract = data.feature_metadata.get("local_contract")
         states, targets, future_indices = data.future_examples(NUM_HORIZONS)
@@ -260,6 +285,11 @@ def _load_examples(
             if data.local_features is not None
             else torch.empty(len(states), 0)
         )
+        people = (
+            data.local_person_features.select_rows(current_indices)
+            if data.local_person_features is not None
+            else None
+        )
         if compute_transition_targets:
             activity_sequence = torch.cat(
                 (data.targets[current_indices, None, :], targets),
@@ -286,6 +316,8 @@ def _load_examples(
             targets = targets[keep]
             transition_targets = transition_targets[keep]
             local = local[keep]
+            if people is not None:
+                people = people.select_rows(torch.nonzero(keep).flatten())
             if not len(states):
                 raise ValueError(f"Sparse anchor selection kept no examples for {name}")
         if hidden_size is None:
@@ -306,6 +338,8 @@ def _load_examples(
         all_targets.append(targets)
         all_transition_targets.append(transition_targets)
         all_local.append(local)
+        if people is not None:
+            all_people.append(people)
         metadata.append(
             {
                 "session": name,
@@ -332,6 +366,8 @@ def _load_examples(
         )
     if hidden_size is None or num_input_layers is None:
         raise ValueError("No sessions selected")
+    if all_people and len(all_people) != len(all_states):
+        raise ValueError("Person sets are missing for some sessions")
     return (
         torch.cat(all_states, dim=0),
         torch.cat(all_targets, dim=0),
@@ -340,6 +376,7 @@ def _load_examples(
         num_input_layers,
         metadata,
         torch.cat(all_local, dim=0),
+        FramePersonFeatures.concatenate(all_people) if all_people else None,
     )
 
 
@@ -407,6 +444,7 @@ def _evaluate_valid_horizons(
     batch_size: int,
     loss_name: str,
     local: torch.Tensor | None = None,
+    people: FramePersonFeatures | None = None,
 ) -> dict[str, Any]:
     head.eval()
     chunks: list[torch.Tensor] = []
@@ -420,7 +458,19 @@ def _evaluate_valid_horizons(
             local[start : start + batch_size].to(device=device, dtype=torch.float32)
             if local is not None else None
         )
-        chunks.append(marginal_probabilities_from_logits(head(x, side), loss_name).cpu())
+        person_batch = None
+        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)
+            person_batch = person_batch.to(device=device, dtype=torch.float32)
+            person_mask = person_mask.to(device=device)
+        chunks.append(
+            marginal_probabilities_from_logits(
+                head(x, side, person_batch, person_mask),
+                loss_name,
+            ).cpu()
+        )
     probabilities = torch.cat(chunks, dim=0).numpy()
     target_array = targets.numpy()
 
@@ -455,7 +505,7 @@ def _checkpoint(
     history: list[dict[str, Any]],
 ) -> dict[str, Any]:
     return {
-        "format_version": 4,
+        "format_version": 5,
         "head": head.state_dict(),
         "head_config": {
             "hidden_size": hidden_size,
@@ -469,6 +519,9 @@ def _checkpoint(
             "local_visual_dim": head.local_visual_dim,
             "local_audio_dim": head.local_audio_dim,
             "local_projection_dim": head.local_projection_dim,
+            "local_fusion": head.local_fusion,
+            "person_attention_dim": head.person_attention_dim,
+            "person_attention_heads": head.person_attention_heads,
         },
         "local_contract": feature_metadata[0].get("local_contract"),
         "alignment": {
@@ -513,7 +566,10 @@ def _checkpoint(
             "layer_fusion": args.layer_fusion,
             "local_feature_root": str(args.local_feature_root) if args.local_feature_root else None,
             "local_modalities": args.local_modalities if args.local_feature_root else None,
+            "local_fusion": args.local_fusion,
             "local_projection_dim": args.local_projection_dim,
+            "person_attention_dim": args.person_attention_dim,
+            "person_attention_heads": args.person_attention_heads,
             "learnable_layer_mixture": args.layer_fusion == "scalar-mixture",
             "supervision_policy": (
                 "transition-sparse-manifest"
@@ -605,6 +661,7 @@ def main() -> None:
         num_input_layers,
         feature_metadata,
         local_features,
+        person_features,
     ) = _load_examples(
         args.feature_root,
         args.target_root,
@@ -615,6 +672,7 @@ def main() -> None:
         args.local_feature_root,
         args.local_modalities,
         compute_transition_targets=args.transition_objective != "none",
+        local_fusion=args.local_fusion,
     )
     local_contract = feature_metadata[0].get("local_contract")
     if not len(states):
@@ -645,6 +703,7 @@ def main() -> None:
     dev_states: torch.Tensor | None = None
     dev_targets: torch.Tensor | None = None
     dev_local: torch.Tensor | None = None
+    dev_people: FramePersonFeatures | None = None
     if dev_sessions:
         (
             dev_states,
@@ -654,6 +713,7 @@ def main() -> None:
             dev_num_input_layers,
             _,
             dev_local,
+            dev_people,
         ) = _load_examples(
             args.feature_root,
             args.target_root,
@@ -665,6 +725,7 @@ def main() -> None:
             args.local_modalities,
             local_contract,
             compute_transition_targets=args.transition_objective != "none",
+            local_fusion=args.local_fusion,
         )
         if dev_hidden_size != hidden_size:
             raise ValueError(
@@ -678,7 +739,13 @@ def main() -> None:
             )
         print(f"development examples: {len(dev_states)}")
 
-    dataset = TensorDataset(states, targets, transition_targets, local_features)
+    dataset = TensorDataset(
+        states,
+        targets,
+        transition_targets,
+        local_features,
+        torch.arange(len(states), dtype=torch.long),
+    )
     generator = torch.Generator().manual_seed(args.seed)
     loader = DataLoader(
         dataset,
@@ -702,13 +769,26 @@ def main() -> None:
         layer_fusion=args.layer_fusion,
         transition_auxiliary=args.transition_objective != "none",
         local_projection_dim=args.local_projection_dim,
+        local_fusion=args.local_fusion,
+        person_attention_dim=args.person_attention_dim,
+        person_attention_heads=args.person_attention_heads,
         **local_head_dimensions(local_contract),
     ).to(device)
     if local_contract is not None:
         print(
-            f"local fusion: {args.local_modalities}; raw V={head.local_visual_dim}, "
+            f"local fusion: {args.local_fusion}/{args.local_modalities}; "
-            f"A={head.local_audio_dim}; projection={args.local_projection_dim or 'identity'}"
+            f"raw V={head.local_visual_dim}, A={head.local_audio_dim}; "
+            f"projection={args.local_projection_dim or 'identity'}"
         )
+        if person_features is not None:
+            print(
+                "frame sets: "
+                f"{len(person_features.frame_offsets) - 1} source frames, "
+                f"{len(person_features.features)} person observations, "
+                f"max_people={person_features.max_people}, "
+                f"attention={args.person_attention_dim}D/"
+                f"{args.person_attention_heads} heads"
+            )
     trainable_parameters = sum(
         parameter.numel() for parameter in head.parameters() if parameter.requires_grad
     )
@@ -782,7 +862,13 @@ def main() -> None:
             unit="batch",
             disable=args.no_progress,
         )
-        for state_batch, target_batch, transition_target_batch, local_batch in progress:
+        for (
+            state_batch,
+            target_batch,
+            transition_target_batch,
+            local_batch,
+            example_index_batch,
+        ) in progress:
             local_batch = local_batch.to(
                 device=device, dtype=torch.float32, non_blocking=True,
             )
@@ -801,15 +887,37 @@ def main() -> None:
                 dtype=torch.float32,
                 non_blocking=True,
             )
+            person_batch = None
+            person_mask = None
+            if person_features is not None:
+                person_batch, person_mask = person_features.padded(
+                    example_index_batch
+                )
+                person_batch = person_batch.to(
+                    device=device,
+                    dtype=torch.float32,
+                    non_blocking=True,
+                )
+                person_mask = person_mask.to(device=device, non_blocking=True)
             optimizer.zero_grad(set_to_none=True)
             if args.transition_objective == "none":
-                logits = head(state_batch, local_batch)
+                logits = head(
+                    state_batch,
+                    local_batch,
+                    person_batch,
+                    person_mask,
+                )
                 activity_loss = loss_function(logits, target_batch)
                 loss = activity_loss
                 transition_loss = activity_loss.detach().new_zeros(())
                 expected_reward = activity_loss.detach().new_zeros(())
             else:
-                logits, transition_logits = head.forward_with_transition(state_batch, local_batch)
+                logits, transition_logits = head.forward_with_transition(
+                    state_batch,
+                    local_batch,
+                    person_batch,
+                    person_mask,
+                )
                 activity_loss = loss_function(logits, target_batch)
                 components = transition_aware_objective(
                     activity_loss,
@@ -890,6 +998,7 @@ def main() -> None:
                 args.batch_size,
                 args.loss,
                 dev_local,
+                dev_people,
             )
             dev_map = float(dev_metrics["mean_map"])
             record["dev_valid_mean_map"] = dev_map