Add DeMemWM multiview selector backends

.exp_artifact/dememwm_dynamic_multiview_memory_selection_plan.md

@@ -463,7 +463,7 @@ Make DeMemWM FOV candidates deterministic
 
 ## Substep 3: Add Deterministic Multiview Selector Backends
 
-Status: `[ ]`
+Status: `[x]`
 
 Goal:
 

datasets/video/memory_selection.py

@@ -298,25 +298,38 @@ def _best_row(rows: torch.Tensor, gains: torch.Tensor, fov_values: torch.Tensor,
     return int(active[0].item())
 
 
-def _select_by_pose_similarity(
+def _deterministic_score_order(scores: torch.Tensor, candidates_t: torch.Tensor) -> torch.Tensor:
+    order = np.lexsort((candidates_t.detach().cpu().numpy(), -scores.detach().cpu().numpy()))
+    return torch.as_tensor(order, device=candidates_t.device, dtype=torch.long)
+
+
+def _pose_distance_matrix(left: torch.Tensor, right: torch.Tensor) -> torch.Tensor:
+    if left.shape[0] == 0 or right.shape[0] == 0:
+        return torch.zeros((left.shape[0], right.shape[0]), device=left.device, dtype=torch.float32)
+    right = right.to(device=left.device, dtype=left.dtype)
+    translation = torch.linalg.vector_norm(left[:, None, :3] - right[None, :, :3], dim=-1) / _POSE_DISTANCE_SCALE
+    angular = torch.linalg.vector_norm(_wrap_degrees(left[:, None, 3:5] - right[None, :, 3:5]), dim=-1) / _ANGLE_DISTANCE_SCALE
+    return translation + angular
+
+
+def _rank_pose_plucker_candidates(
     poses: np.ndarray,
     candidates: np.ndarray,
     target_positions: np.ndarray,
     cfg,
-    count: int,
-    rng=None,
-) -> np.ndarray:
-    if count <= 0 or len(candidates) == 0 or len(target_positions) == 0:
-        return np.empty((0,), dtype=np.int64)
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Return candidate ids, candidate poses, and deterministic relevance scores."""
 
     poses_t = _as_pose_tensor(poses)
-    candidates_t = torch.as_tensor(candidates.astype(np.int64), device=poses_t.device, dtype=torch.long)
-    targets_t = torch.as_tensor(target_positions.astype(np.int64), device=poses_t.device, dtype=torch.long)
+    candidates_t = torch.as_tensor(np.asarray(candidates, dtype=np.int64), device=poses_t.device, dtype=torch.long)
+    targets_t = torch.as_tensor(np.asarray(target_positions, dtype=np.int64), device=poses_t.device, dtype=torch.long)
+    if candidates_t.numel() == 0 or targets_t.numel() == 0:
+        empty_ids = torch.empty((0,), device=poses_t.device, dtype=torch.long)
+        empty_poses = torch.empty((0, 5), device=poses_t.device, dtype=torch.float32)
+        return empty_ids, empty_poses, torch.empty((0,), device=poses_t.device)
+
     candidate_poses = poses_t.index_select(0, candidates_t)
     target_poses = poses_t.index_select(0, targets_t)
-
-    # Training preselects with cheap pose similarity; Plucker ranking keeps it
-    # closer to the inference FOV selector without running the full 10k-point path.
     spatial_scale = max(float(cfg_get(cfg, "pose_similarity_radius", _FOV_RADIUS)), 1e-6)
     angle_scale = max(float(cfg_get(cfg, "pose_similarity_angle_scale", 180.0)), 1e-6)
     spatial_delta = torch.linalg.vector_norm(candidate_poses[:, None, :3] - target_poses[None, :, :3], dim=-1)
@@ -328,30 +341,210 @@ def _select_by_pose_similarity(
     threshold = float(cfg_get(cfg, "pose_similarity_threshold", 0.6))
     valid = pose_scores >= threshold
     if not bool(valid.any()):
-        return np.empty((0,), dtype=np.int64)
+        empty_ids = torch.empty((0,), device=poses_t.device, dtype=torch.long)
+        empty_poses = torch.empty((0, 5), device=poses_t.device, dtype=torch.float32)
+        return empty_ids, empty_poses, torch.empty((0,), device=poses_t.device)
 
     valid_idx = torch.nonzero(valid, as_tuple=False).flatten()
     valid_candidates = candidates_t.index_select(0, valid_idx)
-    valid_scores = pose_scores.index_select(0, valid_idx)
-
+    valid_poses = candidate_poses.index_select(0, valid_idx)
+    rank_scores = pose_scores.index_select(0, valid_idx)
     if bool(cfg_get(cfg, "training_use_plucker", True)):
-        plucker = _plucker_scores_tensor(
-            poses_t.index_select(0, valid_candidates),
-            target_poses,
-            cfg,
-        )
-        rank_scores = valid_scores + float(cfg_get(cfg, "training_plucker_weight", 1.0)) * plucker
-        order = torch.argsort(rank_scores, descending=True)[:count]
-        selected = valid_candidates.index_select(0, order).cpu().numpy()
-    else:
-        valid_np = valid_candidates.cpu().numpy()
-        size = min(count, len(valid_np))
-        choice = np.random.choice if rng is None else rng.choice
-        selected = choice(valid_np, size=size, replace=False)
+        plucker = _plucker_scores_tensor(valid_poses, target_poses, cfg)
+        rank_scores = rank_scores + float(cfg_get(cfg, "training_plucker_weight", 1.0)) * plucker
+
+    order = _deterministic_score_order(rank_scores, valid_candidates)
+    return valid_candidates.index_select(0, order), valid_poses.index_select(0, order), rank_scores.index_select(0, order)
+
+
+def _select_by_pose_similarity(
+    poses: np.ndarray,
+    candidates: np.ndarray,
+    target_positions: np.ndarray,
+    cfg,
+    count: int,
+    rng=None,
+) -> np.ndarray:
+    if count <= 0 or len(candidates) == 0 or len(target_positions) == 0:
+        return np.empty((0,), dtype=np.int64)
+
+    ranked_candidates, _, _ = _rank_pose_plucker_candidates(poses, candidates, target_positions, cfg)
+    if ranked_candidates.numel() == 0:
+        return np.empty((0,), dtype=np.int64)
 
+    selected = ranked_candidates[:count].cpu().numpy()
     return np.sort(np.asarray(selected, dtype=np.int64))
 
 
+def _select_pose_fps(
+    candidate_ids: torch.Tensor,
+    candidate_poses: torch.Tensor,
+    count: int,
+    *,
+    seed_ids: np.ndarray | None = None,
+    all_poses: np.ndarray | torch.Tensor | None = None,
+) -> np.ndarray:
+    """Select diverse camera poses from a relevance-filtered candidate set."""
+
+    candidate_ids = torch.as_tensor(candidate_ids, dtype=torch.long)
+    candidate_poses = _as_pose_tensor(candidate_poses).to(device=candidate_ids.device)
+    if count <= 0 or candidate_ids.numel() == 0:
+        return np.empty((0,), dtype=np.int64)
+
+    available = torch.ones((candidate_ids.shape[0],), device=candidate_ids.device, dtype=torch.bool)
+    selected_rows = []
+    seed_poses = torch.empty((0, 5), device=candidate_ids.device, dtype=torch.float32)
+    if seed_ids is not None and all_poses is not None:
+        all_poses_t = _as_pose_tensor(all_poses).to(device=candidate_ids.device)
+        seed_t = torch.as_tensor(np.asarray(seed_ids, dtype=np.int64), device=candidate_ids.device, dtype=torch.long)
+        seed_t = seed_t[(seed_t >= 0) & (seed_t < all_poses_t.shape[0])]
+        if seed_t.numel() > 0:
+            seed_poses = all_poses_t.index_select(0, torch.unique(seed_t, sorted=True))
+
+    min_dist = None
+    if seed_poses.shape[0] > 0:
+        min_dist = _pose_distance_matrix(candidate_poses, seed_poses).min(dim=1).values
+
+    for step in range(min(count, int(candidate_ids.numel()))):
+        rows = torch.nonzero(available, as_tuple=False).flatten()
+        if rows.numel() == 0:
+            break
+        if step == 0 and min_dist is None:
+            row = int(rows[0].item())
+        else:
+            vals = min_dist.index_select(0, rows)
+            tied = rows.index_select(0, torch.nonzero(vals == vals.max(), as_tuple=False).flatten())
+            row = int(tied[torch.argmin(candidate_ids.index_select(0, tied))].item())
+        selected_rows.append(row)
+        available[row] = False
+
+        dists = _pose_distance_matrix(candidate_poses, candidate_poses[row : row + 1]).flatten()
+        min_dist = dists if min_dist is None else torch.minimum(min_dist, dists)
+        min_dist = min_dist.masked_fill(~available, float("-inf"))
+
+    selected = candidate_ids.index_select(0, torch.as_tensor(selected_rows, device=candidate_ids.device, dtype=torch.long))
+    return selected.cpu().numpy().astype(np.int64)
+
+
+def _select_dynamic_multiview_pose_plucker_fps(
+    poses: np.ndarray,
+    candidates: np.ndarray,
+    target_positions: np.ndarray,
+    cfg,
+    count: int,
+    *,
+    seed_ids: np.ndarray | None = None,
+) -> np.ndarray:
+    ranked_ids, ranked_poses, _ = _rank_pose_plucker_candidates(poses, candidates, target_positions, cfg)
+    topk = cfg_get(cfg, "pose_preselect_topk", 64)
+    if topk is not None:
+        topk = int(topk)
+        if topk > 0:
+            ranked_ids = ranked_ids[:topk]
+            ranked_poses = ranked_poses[:topk]
+    return _select_pose_fps(ranked_ids, ranked_poses, count, seed_ids=seed_ids, all_poses=poses)
+
+
+def _select_dynamic_multiview_from_fov_pool(
+    pool: dict[str, torch.Tensor],
+    count: int,
+    *,
+    excluded: np.ndarray | None = None,
+    reference_frames: np.ndarray | None = None,
+    all_poses: np.ndarray | torch.Tensor | None = None,
+) -> np.ndarray:
+    """Select dynamic multiview frames by residual FOV coverage."""
+
+    if count <= 0:
+        return np.empty((0,), dtype=np.int64)
+
+    candidates_t = pool["candidates_t"]
+    candidate_poses = pool["candidate_poses"]
+    if candidates_t.numel() == 0:
+        return np.empty((0,), dtype=np.int64)
+
+    valid = torch.ones((candidates_t.shape[0],), device=candidates_t.device, dtype=torch.bool)
+    if excluded is not None and len(excluded) > 0:
+        excluded_t = torch.as_tensor(np.asarray(excluded, dtype=np.int64), device=candidates_t.device, dtype=torch.long)
+        valid &= ~(candidates_t[:, None] == excluded_t[None, :]).any(dim=1)
+    if not bool(valid.any()):
+        return np.empty((0,), dtype=np.int64)
+
+    positive = pool.get("positive_fov", torch.zeros_like(valid)).to(device=candidates_t.device, dtype=torch.bool) & valid
+    if not bool(positive.any()):
+        rows = torch.nonzero(valid, as_tuple=False).flatten()
+        return _select_pose_fps(
+            candidates_t.index_select(0, rows),
+            candidate_poses.index_select(0, rows),
+            count,
+            seed_ids=reference_frames,
+            all_poses=all_poses,
+        )
+
+    rows = torch.nonzero(positive, as_tuple=False).flatten()
+    candidates_t = candidates_t.index_select(0, rows)
+    candidate_poses = candidate_poses.index_select(0, rows)
+    inside = pool["inside"].index_select(0, rows)
+    fov_values = pool["fov_values"].index_select(0, rows)
+    plucker = pool["plucker"].index_select(0, rows)
+    gaps = pool["gaps"].index_select(0, rows)
+    temporal_gap = gaps.abs()
+
+    covered = torch.zeros((inside.shape[1],), device=candidates_t.device, dtype=torch.bool)
+    reference_poses = []
+    if reference_frames is not None and len(reference_frames) > 0:
+        ref_t = torch.as_tensor(np.asarray(reference_frames, dtype=np.int64), device=candidates_t.device, dtype=torch.long)
+        ref_rows = torch.nonzero((pool["candidates_t"][:, None] == ref_t[None, :]).any(dim=1), as_tuple=False).flatten()
+        if ref_rows.numel() > 0:
+            covered |= pool["inside"].index_select(0, ref_rows).any(dim=0)
+            reference_poses.append(pool["candidate_poses"].index_select(0, ref_rows))
+        if all_poses is not None:
+            all_poses_t = _as_pose_tensor(all_poses).to(device=candidates_t.device)
+            ref_t = ref_t[(ref_t >= 0) & (ref_t < all_poses_t.shape[0])]
+            if ref_t.numel() > 0:
+                reference_poses.append(all_poses_t.index_select(0, torch.unique(ref_t, sorted=True)))
+
+    selected_rows = []
+    remaining = torch.ones((candidates_t.shape[0],), device=candidates_t.device, dtype=torch.bool)
+    for _ in range(min(count, int(candidates_t.numel()))):
+        active = torch.nonzero(remaining, as_tuple=False).flatten()
+        if active.numel() == 0:
+            break
+
+        if selected_rows:
+            selected_t = torch.as_tensor(selected_rows, device=candidates_t.device, dtype=torch.long)
+            refs = reference_poses + [candidate_poses.index_select(0, selected_t)]
+        else:
+            refs = reference_poses
+        if refs:
+            reference_pose_t = torch.cat(refs, dim=0)
+            pose_dist = _pose_distance_matrix(candidate_poses, reference_pose_t).min(dim=1).values
+        else:
+            pose_dist = torch.zeros((candidate_poses.shape[0],), device=candidates_t.device, dtype=torch.float32)
+        gains = (inside & ~covered[None, :]).float().mean(dim=1)
+
+        row = active
+        for values in (
+            gains,
+            fov_values,
+            plucker,
+            pose_dist,
+            -temporal_gap.to(dtype=torch.float32),
+            -candidates_t.to(dtype=torch.float32),
+        ):
+            vals = values.index_select(0, row)
+            row = row.index_select(0, torch.nonzero(vals == vals.max(), as_tuple=False).flatten())
+            if row.numel() == 1:
+                break
+        selected_row = int(row[0].item())
+        selected_rows.append(selected_row)
+        covered |= inside[selected_row]
+        remaining[selected_row] = False
+
+    selected = candidates_t.index_select(0, torch.as_tensor(selected_rows, device=candidates_t.device, dtype=torch.long))
+    return selected.cpu().numpy().astype(np.int64)
+
+
 def _select_by_point_union(
     poses: np.ndarray,
     candidates: np.ndarray,

tests/test_dememwm_latent_dataset.py

@@ -66,6 +66,38 @@ def _poses(num_frames):
     return poses
 
 
+def _fov_pool(candidate_ids, inside, *, fov_values=None, plucker=None, gaps=None, poses=None):
+    candidate_ids = np.asarray(candidate_ids, dtype=np.int64)
+    candidates_t = torch.as_tensor(candidate_ids, dtype=torch.long)
+    inside_t = torch.as_tensor(inside, dtype=torch.bool)
+    if fov_values is None:
+        fov_values_t = inside_t.float().mean(dim=1) if inside_t.shape[1] > 0 else torch.zeros((len(candidate_ids),), dtype=torch.float32)
+    else:
+        fov_values_t = torch.as_tensor(fov_values, dtype=torch.float32)
+    if plucker is None:
+        plucker_t = torch.zeros((len(candidate_ids),), dtype=torch.float32)
+    else:
+        plucker_t = torch.as_tensor(plucker, dtype=torch.float32)
+    if gaps is None:
+        gaps_t = torch.arange(len(candidate_ids), 0, -1, dtype=torch.long)
+    else:
+        gaps_t = torch.as_tensor(gaps, dtype=torch.long)
+    if poses is None:
+        candidate_poses = torch.zeros((len(candidate_ids), 5), dtype=torch.float32)
+        candidate_poses[:, 0] = candidates_t.to(dtype=torch.float32)
+    else:
+        candidate_poses = torch.as_tensor(poses[candidate_ids], dtype=torch.float32)
+    return {
+        "candidates_t": candidates_t,
+        "candidate_poses": candidate_poses,
+        "inside": inside_t,
+        "fov_values": fov_values_t,
+        "plucker": plucker_t,
+        "gaps": gaps_t,
+        "positive_fov": inside_t.any(dim=1),
+    }
+
+
 def _dataset_cfg(root, **overrides):
     cfg = {
         "save_dir": str(root),
@@ -344,6 +376,172 @@ class MemorySelectionTests(unittest.TestCase):
         self.assertEqual(indices["revisit"].tolist(), [0, -1])
         self.assertEqual(masks["revisit"].tolist(), [True, False])
 
+    def test_pose_similarity_refactor_preserves_deterministic_ranked_result(self):
+        import datasets.video.memory_selection as memory_selection
+
+        poses = np.array(
+            [
+                [0.1, 0, 0, 0, 1],
+                [2.0, 0, 0, 0, 4],
+                [40.0, 0, 0, 0, 180],
+                [0.0, 0, 0, 0, 0],
+            ],
+            dtype=np.float32,
+        )
+        cfg = _selection_cfg(pose_similarity_threshold=0.0, training_use_plucker=True, training_plucker_weight=0.0)
+
+        selected = memory_selection._select_by_pose_similarity(poses, np.asarray([0, 1, 2]), np.asarray([3]), cfg, count=2)
+
+        self.assertEqual(selected.tolist(), [0, 1])
+
+    def test_pose_fps_returns_at_most_count_frames(self):
+        import datasets.video.memory_selection as memory_selection
+
+        ids = torch.as_tensor([0, 1], dtype=torch.long)
+        poses = torch.as_tensor([[0, 0, 0, 0, 0], [10, 0, 0, 0, 0]], dtype=torch.float32)
+
+        selected = memory_selection._select_pose_fps(ids, poses, count=5)
+
+        self.assertEqual(selected.tolist(), [0, 1])
+
+    def test_pose_fps_is_deterministic(self):
+        import datasets.video.memory_selection as memory_selection
+
+        ids = torch.as_tensor([0, 1, 2], dtype=torch.long)
+        poses = torch.as_tensor([[0, 0, 0, 0, 0], [1, 0, 0, 0, 0], [60, 0, 0, 0, 0]], dtype=torch.float32)
+
+        first = memory_selection._select_pose_fps(ids, poses, count=2)
+        second = memory_selection._select_pose_fps(ids, poses, count=2)
+
+        self.assertEqual(first.tolist(), second.tolist())
+
+    def test_pose_fps_caps_pose_preselect_topk_before_diversity(self):
+        import datasets.video.memory_selection as memory_selection
+
+        poses = np.zeros((6, 5), dtype=np.float32)
+        poses[1, 0] = 1.0
+        poses[2, 0] = 30.0
+        poses[3, 0] = 60.0
+        cfg = _selection_cfg(
+            pose_similarity_threshold=0.0,
+            training_use_plucker=True,
+            training_plucker_weight=0.0,
+            pose_preselect_topk=2,
+        )
+
+        selected = memory_selection._select_dynamic_multiview_pose_plucker_fps(
+            poses,
+            np.asarray([0, 1, 2, 3], dtype=np.int64),
+            np.asarray([5], dtype=np.int64),
+            cfg,
+            count=4,
+        )
+
+        self.assertLessEqual(len(selected), 2)
+        self.assertTrue(set(selected.tolist()) <= {0, 1})
+
+    def test_pose_fps_pose_preselect_topk_zero_is_uncapped(self):
+        import datasets.video.memory_selection as memory_selection
+
+        poses = np.zeros((6, 5), dtype=np.float32)
+        poses[1, 0] = 1.0
+        poses[2, 0] = 30.0
+        poses[3, 0] = 60.0
+        cfg = _selection_cfg(
+            pose_similarity_threshold=0.0,
+            training_use_plucker=True,
+            training_plucker_weight=0.0,
+            pose_preselect_topk=0,
+        )
+
+        selected = memory_selection._select_dynamic_multiview_pose_plucker_fps(
+            poses,
+            np.asarray([0, 1, 2, 3], dtype=np.int64),
+            np.asarray([5], dtype=np.int64),
+            cfg,
+            count=4,
+        )
+
+        self.assertEqual(len(selected), 4)
+        self.assertEqual(set(selected.tolist()), {0, 1, 2, 3})
+
+    def test_pose_fps_prefers_diverse_candidate_poses(self):
+        import datasets.video.memory_selection as memory_selection
+
+        ids = torch.as_tensor([0, 1, 2], dtype=torch.long)
+        poses = torch.as_tensor([[0, 0, 0, 0, 0], [1, 0, 0, 0, 0], [60, 0, 0, 0, 0]], dtype=torch.float32)
+
+        selected = memory_selection._select_pose_fps(ids, poses, count=2)
+
+        self.assertEqual(selected.tolist(), [0, 2])
+
+    def test_pose_fps_seed_views_make_first_selection_far_when_possible(self):
+        import datasets.video.memory_selection as memory_selection
+
+        all_poses = np.zeros((3, 5), dtype=np.float32)
+        all_poses[1, 0] = 1.0
+        all_poses[2, 0] = 60.0
+        ids = torch.as_tensor([1, 2], dtype=torch.long)
+        candidate_poses = torch.as_tensor(all_poses[ids.numpy()], dtype=torch.float32)
+
+        selected = memory_selection._select_pose_fps(ids, candidate_poses, count=1, seed_ids=np.asarray([0]), all_poses=all_poses)
+
+        self.assertEqual(selected.tolist(), [2])
+
+    def test_fov_greedy_dynamic_selection_avoids_excluded_revisit_frames(self):
+        import datasets.video.memory_selection as memory_selection
+
+        pool = _fov_pool([0, 1, 2], [[True, True], [True, False], [False, True]], gaps=[3, 2, 1])
+
+        selected = memory_selection._select_dynamic_multiview_from_fov_pool(pool, count=1, excluded=np.asarray([0]))
+
+        self.assertNotIn(0, selected.tolist())
+        self.assertEqual(len(selected), 1)
+
+    def test_fov_greedy_dynamic_selection_is_deterministic(self):
+        import datasets.video.memory_selection as memory_selection
+
+        pool = _fov_pool([0, 1, 2], [[True, False, False], [False, True, False], [False, False, True]], gaps=[3, 2, 1])
+
+        first = memory_selection._select_dynamic_multiview_from_fov_pool(pool, count=2)
+        second = memory_selection._select_dynamic_multiview_from_fov_pool(pool, count=2)
+
+        self.assertEqual(first.tolist(), second.tolist())
+
+    def test_fov_greedy_dynamic_selection_uses_absolute_temporal_gap_tie_break(self):
+        import datasets.video.memory_selection as memory_selection
+
+        pool = _fov_pool(
+            [4, 10],
+            [[True], [True]],
+            fov_values=[1.0, 1.0],
+            plucker=[0.0, 0.0],
+            gaps=[1, -5],
+        )
+
+        selected = memory_selection._select_dynamic_multiview_from_fov_pool(pool, count=1)
+
+        self.assertEqual(selected.tolist(), [4])
+
+    def test_fov_greedy_dynamic_selection_prefers_residual_coverage(self):
+        import datasets.video.memory_selection as memory_selection
+
+        pool = _fov_pool(
+            [0, 1, 2],
+            [[True, False], [True, False], [False, True]],
+            plucker=[0.0, 10.0, 0.0],
+            gaps=[3, 1, 2],
+        )
+
+        selected = memory_selection._select_dynamic_multiview_from_fov_pool(
+            pool,
+            count=1,
+            excluded=np.asarray([0]),
+            reference_frames=np.asarray([0]),
+        )
+
+        self.assertEqual(selected.tolist(), [2])
+
     def test_training_default_causal_revisit_does_not_select_future(self):
         poses = _poses(8)
         poses[6] = poses[4]