Add pooled local AV baselines and research notes

easycom_head.py

@@ -74,6 +74,9 @@ class TurnPredictionHead(nn.Module):
         learnable_layer_mixture: bool = False,
         layer_fusion: str | None = None,
         transition_auxiliary: bool = False,
+        local_visual_dim: int = 0,
+        local_audio_dim: int = 0,
+        local_projection_dim: int = 256,
     ) -> None:
         super().__init__()
         if num_input_layers <= 0:
@@ -108,6 +111,11 @@ class TurnPredictionHead(nn.Module):
         self.layer_fusion = layer_fusion
         self.learnable_layer_mixture = layer_fusion == "scalar-mixture"
         self.transition_auxiliary = transition_auxiliary
+        if min(local_visual_dim, local_audio_dim, local_projection_dim) < 0:
+            raise ValueError("Local dimensions must be non-negative")
+        self.local_visual_dim = local_visual_dim
+        self.local_audio_dim = local_audio_dim
+        self.local_projection_dim = local_projection_dim
         if self.learnable_layer_mixture:
             # Zero initialization exactly reproduces a uniform layer mean at
             # step zero while allowing training to select useful query layers.
@@ -123,9 +131,23 @@ class TurnPredictionHead(nn.Module):
             )
         else:
             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)):
+            if dimension:
+                self.local_projections[name] = (
+                    nn.Sequential(
+                        nn.LayerNorm(dimension),
+                        nn.Linear(dimension, local_projection_dim),
+                        nn.GELU(),
+                    )
+                    if local_projection_dim
+                    else nn.Identity()
+                )
+                fused_size += (local_projection_dim or dimension) + 1
         self.net = nn.Sequential(
-            nn.LayerNorm(hidden_size),
-            nn.Linear(hidden_size, hidden_size // 2),
+            nn.LayerNorm(fused_size),
+            nn.Linear(fused_size, hidden_size // 2),
             nn.GELU(),
             nn.Dropout(dropout),
             nn.Linear(hidden_size // 2, num_horizons * num_classes),
@@ -166,14 +188,46 @@ class TurnPredictionHead(nn.Module):
             states = self.layer_fusion_net(states)
         return states
 
-    def _shared_features(self, states: torch.Tensor) -> torch.Tensor:
+    @property
+    def local_input_size(self) -> int:
+        return sum(d + 1 for d in (self.local_visual_dim, self.local_audio_dim) if d)
+
+    def _shared_features(
+        self,
+        states: torch.Tensor,
+        local: 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 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}"
+                )
+            pieces = [states]
+            offset = 0
+            for name, dimension in (
+                ("visual", self.local_visual_dim),
+                ("audio", self.local_audio_dim),
+            ):
+                if dimension:
+                    values = local[..., offset : offset + dimension]
+                    valid = local[..., offset + dimension : offset + dimension + 1]
+                    # Mask AFTER projection too: a learned bias is not an observation.
+                    pieces.extend((self.local_projections[name](values) * valid, valid))
+                    offset += dimension + 1
+            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")
         return self.net[:-1](states)
 
-    def forward(self, states: torch.Tensor) -> torch.Tensor:
-        features = self._shared_features(states)
+    def forward(
+        self,
+        states: torch.Tensor,
+        local: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        features = self._shared_features(states, local)
         logits = self.net[-1](features)
         return logits.view(
             *features.shape[:-1], self.num_horizons, self.num_classes
@@ -182,6 +236,7 @@ class TurnPredictionHead(nn.Module):
     def forward_with_transition(
         self,
         states: torch.Tensor,
+        local: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Return activity logits and binary transition logits."""
 
@@ -189,7 +244,7 @@ class TurnPredictionHead(nn.Module):
             raise RuntimeError(
                 "This head was created without transition_auxiliary=True"
             )
-        features = self._shared_features(states)
+        features = self._shared_features(states, local)
         activity_logits = self.net[-1](features).view(
             *features.shape[:-1],
             self.num_horizons,
@@ -631,6 +686,7 @@ class SessionData:
     targets: torch.Tensor
     label_indices: torch.LongTensor
     feature_metadata: dict[str, Any]
+    local_features: torch.Tensor | None = None
 
     @property
     def hidden_size(self) -> int:
@@ -662,6 +718,9 @@ def load_session_data(
     *,
     expected_mode: str | None = "causal-prefix",
     dataset: DatasetName = "easycom",
+    local_feature_root: str | Path | None = None,
+    local_modalities: str = "av",
+    local_contract: dict[str, Any] | None = None,
 ) -> SessionData:
     """Load and validate one saved Omni feature payload and target array."""
 
@@ -738,7 +797,22 @@ def load_session_data(
             "video_sampling_grid",
         )
     }
-    return SessionData(name, states, targets, label_indices, metadata)
+    local = None
+    if local_feature_root is not None:
+        from local_turn_features import load_local_features
+
+        local, contract = load_local_features(
+            Path(local_feature_root) / f"{name}.pt",
+            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)
 
 
 def average_precision_metrics(
@@ -780,6 +854,9 @@ def checkpoint_head_config(payload: dict[str, Any]) -> dict[str, Any]:
     config.setdefault("num_input_layers", 1)
     config.setdefault("learnable_layer_mixture", False)
     config.setdefault("transition_auxiliary", False)
+    config.setdefault("local_visual_dim", 0)
+    config.setdefault("local_audio_dim", 0)
+    config.setdefault("local_projection_dim", 256)
     config.setdefault(
         "layer_fusion",
         (