Select DeMemWM dynamic memory by multiview retrieval

.exp_artifact/dememwm_dynamic_multiview_memory_selection_plan.md

@@ -604,7 +604,7 @@ Add DeMemWM multiview selector backends
 
 ## Substep 4: Implement Dynamic-As-Multiview Selection
 
-Status: `[ ]`
+Status: `[x]`
 
 Goal:
 

datasets/video/memory_selection.py

@@ -691,6 +691,83 @@ def _dynamic_multiview_selector(cfg) -> str:
     return selector
 
 
+def _select_dynamic_multiview(
+    poses: np.ndarray,
+    target_positions: np.ndarray,
+    cfg,
+    count: int,
+    *,
+    excluded: np.ndarray | None = None,
+    reference_frames: np.ndarray | None = None,
+    split: str = "training",
+    min_candidate_frame: int = 0,
+    fov_pool: dict[str, torch.Tensor] | None = None,
+) -> np.ndarray:
+    """Select dynamic-as-multiview memory frames."""
+
+    if count <= 0:
+        return np.empty((0,), dtype=np.int64)
+
+    poses = _as_pose_array(poses)
+    target_positions = np.asarray(target_positions, dtype=np.int64)
+    if len(target_positions) == 0:
+        return np.empty((0,), dtype=np.int64)
+
+    candidates = _memory_candidate_frames(
+        len(poses),
+        target_positions,
+        cfg,
+        split,
+        min_candidate_frame=min_candidate_frame,
+    )
+    if excluded is not None and len(excluded) > 0 and len(candidates) > 0:
+        candidates = candidates[~np.isin(candidates, np.asarray(excluded, dtype=np.int64))]
+    if len(candidates) == 0:
+        return np.empty((0,), dtype=np.int64)
+
+    selector = _dynamic_multiview_selector(cfg)
+    if selector == "fov_greedy":
+        pool = fov_pool
+        if pool is None:
+            pool = _build_fov_candidate_pool(poses, candidates, target_positions, cfg, use_plucker=True)
+        else:
+            pool_candidates_t = pool["candidates_t"]
+            allowed_t = torch.as_tensor(candidates, device=pool_candidates_t.device, dtype=torch.long)
+            keep = (pool_candidates_t[:, None] == allowed_t[None, :]).any(dim=1)
+            if reference_frames is not None and len(reference_frames) > 0:
+                ref_t = torch.as_tensor(np.asarray(reference_frames, dtype=np.int64), device=pool_candidates_t.device, dtype=torch.long)
+                # Reference rows seed residual coverage below even when excluded from selection.
+                keep |= (pool_candidates_t[:, None] == ref_t[None, :]).any(dim=1)
+            if not bool(keep.all()):
+                rows = torch.nonzero(keep, as_tuple=False).flatten()
+                pool = {
+                    key: value.index_select(0, rows)
+                    if torch.is_tensor(value) and value.shape[:1] == pool_candidates_t.shape[:1]
+                    else value
+                    for key, value in pool.items()
+                }
+        selected = _select_dynamic_multiview_from_fov_pool(
+            pool,
+            count,
+            excluded=excluded,
+            reference_frames=reference_frames,
+            all_poses=poses,
+        )
+    elif selector == "pose_plucker_fps":
+        selected = _select_dynamic_multiview_pose_plucker_fps(
+            poses,
+            candidates,
+            target_positions,
+            cfg,
+            count,
+            seed_ids=reference_frames,
+        )
+    else:
+        raise AssertionError(f"unhandled dynamic multiview selector: {selector!r}")
+
+    return np.sort(np.asarray(selected, dtype=np.int64))
+
+
 def _latent_frame_vectors(latents, start: int, stop: int):
     if latents is None or stop <= start:
         return None

tests/test_dememwm_latent_dataset.py

@@ -542,6 +542,113 @@ class MemorySelectionTests(unittest.TestCase):
 
         self.assertEqual(selected.tolist(), [2])
 
+    def test_dynamic_multiview_provided_fov_pool_keeps_excluded_reference_coverage(self):
+        import datasets.video.memory_selection as memory_selection
+
+        poses = _poses(6)
+        cfg = _selection_cfg(dynamic=_dynamic_cfg("multiview", multiview_selector="fov_greedy"))
+        pool = _fov_pool(
+            [0, 2, 3],
+            [[True, False], [True, False], [False, True]],
+            fov_values=[0.5, 1.0, 0.5],
+            plucker=[0.0, 0.0, 0.0],
+            gaps=[5, 3, 2],
+            poses=poses,
+        )
+
+        with mock.patch.object(
+            memory_selection,
+            "_build_fov_candidate_pool",
+            side_effect=AssertionError("unexpected FOV pool build"),
+        ):
+            selected = memory_selection._select_dynamic_multiview(
+                poses,
+                np.asarray([5], dtype=np.int64),
+                cfg,
+                count=1,
+                excluded=np.asarray([0], dtype=np.int64),
+                reference_frames=np.asarray([0], dtype=np.int64),
+                split="training",
+                fov_pool=pool,
+            )
+
+        self.assertEqual(selected.tolist(), [3])
+
+    def test_dynamic_multiview_fov_greedy_candidates_pool_exclusion_and_determinism(self):
+        import datasets.video.memory_selection as memory_selection
+
+        poses = _poses(8)
+        default_cfg = _selection_cfg(dynamic=_dynamic_cfg("multiview", multiview_selector="fov_greedy"))
+        noncausal_cfg = _selection_cfg(causal=False, dynamic=_dynamic_cfg("multiview", multiview_selector="fov_greedy"))
+        cases = [
+            (default_cfg, "training", np.asarray([5], dtype=np.int64), [1, 2, 3, 4], [3, 4]),
+            (noncausal_cfg, "training", np.asarray([3], dtype=np.int64), [1, 2, 3, 4, 5, 6, 7], [6, 7]),
+            (noncausal_cfg, "validation", np.asarray([3], dtype=np.int64), [1, 2], [1, 2]),
+            (noncausal_cfg, "test", np.asarray([3], dtype=np.int64), [1, 2], [1, 2]),
+        ]
+
+        def fake_select(pool, count, **kwargs):
+            return pool["candidates_t"].cpu().numpy()[-count:]
+
+        for cfg, split, targets, expected_candidates, expected_selected in cases:
+            seen_candidates = []
+
+            def fake_build(poses_arg, candidates, target_positions, cfg_arg, *, use_plucker):
+                seen_candidates.append((candidates.copy(), use_plucker))
+                return _fov_pool(candidates, np.ones((len(candidates), 1), dtype=bool), poses=poses_arg)
+
+            with self.subTest(split=split), mock.patch.object(memory_selection, "_build_fov_candidate_pool", side_effect=fake_build), mock.patch.object(
+                memory_selection, "_select_dynamic_multiview_from_fov_pool", side_effect=fake_select
+            ):
+                selected = memory_selection._select_dynamic_multiview(poses, targets, cfg, count=2, split=split, min_candidate_frame=1)
+
+            self.assertEqual(seen_candidates[0][0].tolist(), expected_candidates)
+            self.assertTrue(seen_candidates[0][1])
+            self.assertEqual(selected.tolist(), expected_selected)
+
+        provided_poses = _poses(6)
+        pool = _fov_pool([0, 1, 2, 3], [[True, False], [False, True], [True, False], [False, True]], poses=provided_poses)
+        kwargs = dict(excluded=np.asarray([0, 1], dtype=np.int64), reference_frames=np.asarray([0, 1], dtype=np.int64), split="training", fov_pool=pool)
+        with mock.patch.object(memory_selection, "_build_fov_candidate_pool", side_effect=AssertionError("unexpected FOV pool build")):
+            first = memory_selection._select_dynamic_multiview(provided_poses, np.asarray([5], dtype=np.int64), default_cfg, count=2, **kwargs)
+            second = memory_selection._select_dynamic_multiview(provided_poses, np.asarray([5], dtype=np.int64), default_cfg, count=2, **kwargs)
+
+        self.assertEqual(first.tolist(), [2, 3])
+        self.assertEqual(second.tolist(), first.tolist())
+        self.assertTrue(set(first.tolist()).isdisjoint({0, 1}))
+
+    def test_dynamic_multiview_pose_plucker_fps_uses_pose_path_without_fov_masks(self):
+        import datasets.video.memory_selection as memory_selection
+
+        poses = _poses(8)
+        cfg = _selection_cfg(
+            pose_similarity_threshold=0.0,
+            training_use_plucker=False,
+            pose_preselect_topk=3,
+            dynamic=_dynamic_cfg("multiview", multiview_selector="pose_plucker_fps"),
+        )
+        reference = np.asarray([0], dtype=np.int64)
+
+        with mock.patch.object(memory_selection, "_candidate_fov_masks", side_effect=AssertionError("unexpected exact FOV masks")), mock.patch.object(
+            memory_selection, "_build_fov_candidate_pool", side_effect=AssertionError("unexpected FOV pool build")
+        ), mock.patch.object(memory_selection, "_rank_pose_plucker_candidates", wraps=memory_selection._rank_pose_plucker_candidates) as rank_mock, mock.patch.object(
+            memory_selection, "_select_pose_fps", wraps=memory_selection._select_pose_fps
+        ) as fps_mock:
+            selected = memory_selection._select_dynamic_multiview(
+                poses,
+                np.asarray([6], dtype=np.int64),
+                cfg,
+                count=2,
+                reference_frames=reference,
+                split="training",
+            )
+
+        self.assertTrue(rank_mock.called)
+        self.assertTrue(fps_mock.called)
+        self.assertEqual(fps_mock.call_args.kwargs["seed_ids"].tolist(), [0])
+        self.assertEqual(selected.tolist(), sorted(selected.tolist()))
+        self.assertLessEqual(len(selected), 2)
+
     def test_training_default_causal_revisit_does_not_select_future(self):
         poses = _poses(8)
         poses[6] = poses[4]