Add causal person track GRU residual experiment

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}