Add offline sparse supervision control

train_turn_head.py

@@ -72,6 +72,18 @@ def parse_args() -> argparse.Namespace:
             "They are automatically removed from --sessions. Never use 1/2/3."
         ),
     )
+    parser.add_argument(
+        "--sparse-anchor-cache-dir",
+        type=Path,
+        action="append",
+        default=[],
+        help=(
+            "Optional transition-sparse AV cache directory containing "
+            "manifest.json. Repeat once per training session. Only its anchor "
+            "target_index values select examples; states still come from "
+            "--feature-root. Development examples always remain dense."
+        ),
+    )
     parser.add_argument("--epochs", type=int, default=50)
     parser.add_argument("--batch-size", type=int, default=512)
     parser.add_argument("--learning-rate", type=float, default=7e-5)
@@ -138,6 +150,7 @@ def _load_examples(
     sessions: list[str],
     expected_mode: str | None,
     dataset: str,
+    selected_indices: dict[str, set[int]] | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor, int, list[dict[str, Any]]]:
     all_states: list[torch.Tensor] = []
     all_targets: list[torch.Tensor] = []
@@ -151,7 +164,21 @@ def _load_examples(
             expected_mode=expected_mode,
             dataset=dataset,
         )
-        states, targets, _ = data.future_examples(NUM_HORIZONS)
+        states, targets, future_indices = data.future_examples(NUM_HORIZONS)
+        dense_example_count = len(states)
+        if selected_indices is not None:
+            if name not in selected_indices:
+                raise ValueError(f"No sparse anchor manifest was supplied for {name}")
+            wanted = selected_indices[name]
+            current_indices = future_indices[:, 0] - 1
+            keep = torch.tensor(
+                [int(index) in wanted for index in current_indices.tolist()],
+                dtype=torch.bool,
+            )
+            states = states[keep]
+            targets = targets[keep]
+            if not len(states):
+                raise ValueError(f"Sparse anchor selection kept no examples for {name}")
         if hidden_size is None:
             hidden_size = data.hidden_size
         elif data.hidden_size != hidden_size:
@@ -166,12 +193,19 @@ def _load_examples(
                 "session": name,
                 "num_states": len(data.states),
                 "num_training_examples": len(states),
+                "num_dense_training_examples": dense_example_count,
+                "supervision_policy": (
+                    "transition-sparse-manifest"
+                    if selected_indices is not None
+                    else "dense-5hz"
+                ),
                 **data.feature_metadata,
             }
         )
         print(
             f"{name}: states={len(data.states)}, "
-            f"valid_10_horizon_examples={len(states)}"
+            f"valid_10_horizon_examples={dense_example_count}, "
+            f"selected_examples={len(states)}"
         )
     if hidden_size is None:
         raise ValueError("No sessions selected")
@@ -183,6 +217,38 @@ def _load_examples(
     )
 
 
+def _load_sparse_anchor_indices(
+    cache_dirs: list[Path],
+) -> dict[str, set[int]] | None:
+    if not cache_dirs:
+        return None
+    result: dict[str, set[int]] = {}
+    for cache_dir in cache_dirs:
+        manifest_path = cache_dir.expanduser().resolve() / "manifest.json"
+        if not manifest_path.is_file():
+            raise FileNotFoundError(manifest_path)
+        manifest = json.loads(manifest_path.read_text())
+        session = normalize_sample_name(str(manifest.get("session", "")), "easycom")
+        policy = manifest.get("anchor_selection", {}).get("policy")
+        if policy != "transition-sparse":
+            raise ValueError(
+                f"Expected transition-sparse manifest, got {policy!r}: "
+                f"{manifest_path}"
+            )
+        if session in result:
+            raise ValueError(f"Duplicate sparse anchor manifest for {session}")
+        indices = {
+            int(anchor["target_index"])
+            for anchor in manifest.get("anchors", [])
+            if bool(anchor.get("has_future_target", True))
+        }
+        if not indices:
+            raise ValueError(f"No valid sparse anchors in {manifest_path}")
+        result[session] = indices
+        print(f"sparse anchors {session}: {len(indices)} from {manifest_path}")
+    return result
+
+
 def _scheduler(
     optimizer: torch.optim.Optimizer,
     *,
@@ -297,6 +363,15 @@ def _checkpoint(
             "scheduler": "linear_warmup_times_cosine",
             "seed": args.seed,
             "expected_mode": args.expected_mode,
+            "supervision_policy": (
+                "transition-sparse-manifest"
+                if args.sparse_anchor_cache_dir
+                else "dense-5hz"
+            ),
+            "sparse_anchor_cache_dirs": [
+                str(path.expanduser().resolve())
+                for path in args.sparse_anchor_cache_dir
+            ],
         },
         "history": history,
     }
@@ -308,6 +383,11 @@ def main() -> None:
     _seed_everything(args.seed)
     device = _resolve_device(args.device)
     expected_mode = None if args.expected_mode == "any" else args.expected_mode
+    sparse_anchor_indices = _load_sparse_anchor_indices(
+        args.sparse_anchor_cache_dir
+    )
+    if sparse_anchor_indices is not None and args.dataset != "easycom":
+        raise ValueError("--sparse-anchor-cache-dir currently supports EasyCom only")
 
     if args.dataset == "ego4d":
         assert args.manifest is not None
@@ -342,6 +422,14 @@ def main() -> None:
     validate_disjoint_sessions(train_sessions, dev_sessions)
     if not train_sessions:
         raise ValueError("No training sessions remain after removing dev sessions")
+    if sparse_anchor_indices is not None:
+        missing = set(train_sessions) - set(sparse_anchor_indices)
+        extra = set(sparse_anchor_indices) - set(train_sessions)
+        if missing or extra:
+            raise ValueError(
+                "Sparse anchor manifests must match training sessions exactly; "
+                f"missing={sorted(missing)}, extra={sorted(extra)}"
+            )
     leaked = official_test & (set(train_sessions) | set(dev_sessions))
     if leaked:
         raise ValueError(
@@ -358,6 +446,7 @@ def main() -> None:
         train_sessions,
         expected_mode,
         args.dataset,
+        sparse_anchor_indices,
     )
     print(f"training examples: {len(states)}, feature_dim={hidden_size}")
     dev_states: torch.Tensor | None = None
@@ -369,6 +458,7 @@ def main() -> None:
             dev_sessions,
             expected_mode,
             args.dataset,
+            None,
         )
         if dev_hidden_size != hidden_size:
             raise ValueError(
@@ -527,4 +617,4 @@ def main() -> None:
 
 
 if __name__ == "__main__":
-    main()
+    main()