Make DeMemWM FOV candidates deterministic

.exp_artifact/dememwm_dynamic_multiview_memory_selection_plan.md

@@ -342,7 +342,7 @@ Add DeMemWM multiview dynamic policy
 
 ## Substep 2: Make FOV Candidate Construction Deterministic And Reusable
 
-Status: `[ ]`
+Status: `[x]`
 
 Goal:
 

datasets/video/memory_selection.py

@@ -17,6 +17,7 @@ _FOV_HALF_H = 105.0 / 2.0
 _FOV_HALF_V = 75.0 / 2.0
 _POSE_DISTANCE_SCALE = 30.0
 _ANGLE_DISTANCE_SCALE = 180.0
+_DETERMINISTIC_UNIT_BALL_CACHE: Dict[Tuple[int, str, torch.dtype], torch.Tensor] = {}
 
 
 def cfg_get(cfg, key: str, default=None):
@@ -80,25 +81,47 @@ def _pose_directions(poses: torch.Tensor) -> torch.Tensor:
     return directions / torch.linalg.vector_norm(directions, dim=-1, keepdim=True).clamp_min(1e-6)
 
 
-def _sample_points_in_sphere(center: torch.Tensor) -> torch.Tensor:
-    device = center.device
-    dtype = center.dtype
+def _deterministic_points_in_unit_ball(num_points: int, device, dtype) -> torch.Tensor:
+    """Return deterministic points in the unit ball."""
+
+    num_points = int(num_points)
+    device = torch.device(device)
+    if num_points <= 0:
+        return torch.zeros((0, 3), device=device, dtype=dtype)
 
-    samples_r = torch.rand(_FOV_NUM_POINTS, device=device, dtype=dtype)
-    samples_phi = torch.rand(_FOV_NUM_POINTS, device=device, dtype=dtype)
-    samples_u = torch.rand(_FOV_NUM_POINTS, device=device, dtype=dtype)
-    r = _FOV_RADIUS * torch.pow(samples_r, 1.0 / 3.0)
-    phi = 2.0 * torch.pi * samples_phi
-    theta = torch.acos(1.0 - 2.0 * samples_u)
+    cache_key = (num_points, str(device), dtype)
+    cached = _DETERMINISTIC_UNIT_BALL_CACHE.get(cache_key)
+    if cached is not None:
+        return cached
 
-    offsets = torch.stack(
+    compute_dtype = torch.float64 if dtype == torch.float64 else torch.float32
+    idx = torch.arange(num_points, device=device, dtype=compute_dtype)
+    inv_n = 1.0 / float(num_points)
+    z = 1.0 - 2.0 * ((idx + 0.5) * inv_n)
+    xy = torch.sqrt((1.0 - z * z).clamp_min(0.0))
+    theta = idx * 2.399963229728653
+    directions = torch.stack(
         [
-            r * torch.sin(theta) * torch.cos(phi),
-            r * torch.sin(theta) * torch.sin(phi),
-            r * torch.cos(theta),
+            xy * torch.cos(theta),
+            xy * torch.sin(theta),
+            z,
         ],
         dim=-1,
     )
+    radii = torch.pow((idx + 0.5) * inv_n, 1.0 / 3.0)
+    points = (directions * radii[:, None]).to(dtype=dtype)
+
+    if len(_DETERMINISTIC_UNIT_BALL_CACHE) >= 8:
+        _DETERMINISTIC_UNIT_BALL_CACHE.clear()
+    _DETERMINISTIC_UNIT_BALL_CACHE[cache_key] = points
+    return points
+
+
+def _sample_points_in_sphere(center: torch.Tensor) -> torch.Tensor:
+    device = center.device
+    dtype = center.dtype
+
+    offsets = _deterministic_points_in_unit_ball(_FOV_NUM_POINTS, device, dtype) * _FOV_RADIUS
     return center[None, :3] + offsets
 
 
@@ -194,6 +217,57 @@ def _plucker_scores_tensor(candidate_poses: torch.Tensor, target_poses: torch.Te
     return (direction_sim * moment_sim).max(dim=1).values.to(dtype=torch.float32)
 
 
+def _empty_fov_candidate_pool(device) -> dict[str, torch.Tensor]:
+    return {
+        "candidates_t": torch.empty((0,), device=device, dtype=torch.long),
+        "candidate_poses": torch.empty((0, 5), device=device, dtype=torch.float32),
+        "inside": torch.zeros((0, 0), device=device, dtype=torch.bool),
+        "fov_values": torch.empty((0,), device=device, dtype=torch.float32),
+        "plucker": torch.empty((0,), device=device, dtype=torch.float32),
+        "gaps": torch.empty((0,), device=device, dtype=torch.long),
+        "positive_fov": torch.empty((0,), device=device, dtype=torch.bool),
+    }
+
+
+def _build_fov_candidate_pool(
+    poses: np.ndarray,
+    candidates: np.ndarray,
+    target_positions: np.ndarray,
+    cfg,
+    *,
+    use_plucker: bool,
+) -> dict[str, torch.Tensor]:
+    """Build deterministic FOV/Plucker/coverage data for candidate frames."""
+
+    poses_t = _as_pose_tensor(poses)
+    device = poses_t.device
+    candidates = np.asarray(candidates, dtype=np.int64)
+    target_positions = np.asarray(target_positions, dtype=np.int64)
+    if len(candidates) == 0 or len(target_positions) == 0:
+        return _empty_fov_candidate_pool(device)
+
+    candidates = _pose_preselect(candidates, poses_t, target_positions, cfg)
+    inside, fov_values = _candidate_fov_masks(poses_t, candidates, target_positions, cfg)
+    candidates_t = torch.as_tensor(candidates, device=device, dtype=torch.long)
+    candidate_poses = poses_t.index_select(0, candidates_t)
+    target_t = torch.as_tensor(target_positions, device=device, dtype=torch.long)
+    if use_plucker:
+        plucker = _plucker_scores_tensor(candidate_poses, poses_t.index_select(0, target_t), cfg)
+    else:
+        plucker = torch.zeros((candidates_t.shape[0],), device=device, dtype=torch.float32)
+
+    first_target = torch.full_like(candidates_t, int(target_positions[0]))
+    return {
+        "candidates_t": candidates_t,
+        "candidate_poses": candidate_poses,
+        "inside": inside,
+        "fov_values": fov_values,
+        "plucker": plucker,
+        "gaps": first_target - candidates_t,
+        "positive_fov": inside.any(dim=1),
+    }
+
+
 def _empty_selection(counts: Mapping[str, int]) -> Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray]]:
     indices = {}
     masks = {}
@@ -288,37 +362,59 @@ def _select_by_point_union(
     fov_threshold: float | None = None,
     min_total_coverage: float | None = None,
     use_plucker: bool = False,
+    fov_pool: dict[str, torch.Tensor] | None = None,
 ) -> np.ndarray:
-    if count <= 0 or len(candidates) == 0 or len(target_positions) == 0:
+    if count <= 0:
         return np.empty((0,), dtype=np.int64)
 
-    poses_t = _as_pose_tensor(poses)
-    candidates = _pose_preselect(candidates.astype(np.int64), poses_t, target_positions, cfg)
-    inside, fov_values = _candidate_fov_masks(poses_t, candidates, target_positions, cfg)
-    if inside.shape[0] == 0:
+    if fov_pool is None:
+        if len(candidates) == 0 or len(target_positions) == 0:
+            return np.empty((0,), dtype=np.int64)
+        fov_pool = _build_fov_candidate_pool(poses, candidates, target_positions, cfg, use_plucker=use_plucker)
+
+    return _select_by_point_union_from_pool(
+        fov_pool,
+        count,
+        fov_threshold=fov_threshold,
+        min_total_coverage=min_total_coverage,
+    )
+
+
+def _select_by_point_union_from_pool(
+    pool: dict[str, torch.Tensor],
+    count: int,
+    *,
+    fov_threshold: float | None = None,
+    min_total_coverage: float | None = None,
+) -> np.ndarray:
+    if count <= 0:
+        return np.empty((0,), dtype=np.int64)
+
+    candidates_t = pool["candidates_t"]
+    inside = pool["inside"]
+    if candidates_t.numel() == 0 or inside.shape[0] == 0 or inside.shape[1] == 0:
         return np.empty((0,), dtype=np.int64)
 
-    valid = torch.ones((len(candidates),), device=poses_t.device, dtype=torch.bool)
+    device = candidates_t.device
+    fov_values = pool["fov_values"]
+    plucker = pool["plucker"]
+    gaps = pool["gaps"]
+    valid = torch.ones((candidates_t.shape[0],), device=device, dtype=torch.bool)
     if fov_threshold is not None:
         valid &= fov_values >= float(fov_threshold)
     if not bool(valid.any()):
         return np.empty((0,), dtype=np.int64)
 
     valid_idx = torch.nonzero(valid, as_tuple=False).flatten()
-    candidates_t = torch.as_tensor(candidates, device=poses_t.device, dtype=torch.long).index_select(0, valid_idx)
+    candidates_t = candidates_t.index_select(0, valid_idx)
     inside = inside.index_select(0, valid_idx)
     fov_values = fov_values.index_select(0, valid_idx)
-    target_t = torch.as_tensor(target_positions, device=poses_t.device, dtype=torch.long)
-    if use_plucker:
-        plucker = _plucker_scores_tensor(poses_t.index_select(0, candidates_t), poses_t.index_select(0, target_t), cfg)
-    else:
-        plucker = torch.zeros((candidates_t.shape[0],), device=poses_t.device, dtype=torch.float32)
+    plucker = plucker.index_select(0, valid_idx)
+    gaps = gaps.index_select(0, valid_idx)
 
-    remaining = torch.ones((candidates_t.shape[0],), device=poses_t.device, dtype=torch.bool)
-    covered = torch.zeros((inside.shape[1],), device=poses_t.device, dtype=torch.bool)
+    remaining = torch.ones((candidates_t.shape[0],), device=device, dtype=torch.bool)
+    covered = torch.zeros((inside.shape[1],), device=device, dtype=torch.bool)
     selected_rows = []
-    first_target = int(target_positions[0])
-    gaps = first_target - candidates_t
 
     for _ in range(count):
         gains = (inside & ~covered[None, :]).float().mean(dim=1)
@@ -338,7 +434,7 @@ def _select_by_point_union(
     coverage = float(covered.float().mean().item()) if covered.numel() else 0.0
     if min_total_coverage is not None and coverage < float(min_total_coverage):
         return np.empty((0,), dtype=np.int64)
-    selected = candidates_t.index_select(0, torch.as_tensor(selected_rows, device=poses_t.device, dtype=torch.long))
+    selected = candidates_t.index_select(0, torch.as_tensor(selected_rows, device=device, dtype=torch.long))
     return np.sort(selected.cpu().numpy().astype(np.int64))
 
 

tests/test_dememwm_latent_dataset.py

@@ -206,6 +206,84 @@ class MemorySelectionTests(unittest.TestCase):
         self.assertEqual(indices["revisit"].tolist(), [0])
         self.assertEqual(masks["revisit"].tolist(), [True])
 
+    def test_sample_points_in_sphere_is_deterministic_without_torch_rand(self):
+        import datasets.video.memory_selection as memory_selection
+
+        center = torch.tensor([1.0, 2.0, 3.0, 0.0, 0.0], dtype=torch.float32)
+        with mock.patch.object(torch, "rand", side_effect=AssertionError("unexpected random FOV sampling")):
+            first = memory_selection._sample_points_in_sphere(center)
+            second = memory_selection._sample_points_in_sphere(center)
+
+        self.assertEqual(tuple(first.shape), (memory_selection._FOV_NUM_POINTS, 3))
+        self.assertTrue(torch.equal(first, second))
+
+    def test_fov_revisit_selection_is_deterministic_across_calls(self):
+        poses = np.array(
+            [
+                [0, 0, 0, 0, 0],
+                [0, 0, 0, 0, 180],
+                [0, 0, 0, 0, 180],
+                [0, 0, 0, 0, 180],
+                [0, 0, 0, 0, 180],
+                [0, 0, 0, 0, 0],
+            ],
+            dtype=np.float32,
+        )
+        cfg = _selection_cfg(
+            max_anchor_frames=0,
+            max_dynamic_frames=0,
+            max_revisit_frames=1,
+            fov_overlap_threshold=0.5,
+            local_context_exclusion_frames=1,
+        )
+
+        first_indices, first_masks = select_memory_indices(poses, np.array([5]), cfg, split="validation")
+        second_indices, second_masks = select_memory_indices(poses, np.array([5]), cfg, split="validation")
+
+        self.assertEqual(first_indices["revisit"].tolist(), [0])
+        self.assertEqual(second_indices["revisit"].tolist(), first_indices["revisit"].tolist())
+        self.assertEqual(second_masks["revisit"].tolist(), first_masks["revisit"].tolist())
+
+    def test_empty_fov_candidate_pool_returns_empty_selection(self):
+        import datasets.video.memory_selection as memory_selection
+
+        pool = memory_selection._build_fov_candidate_pool(
+            _poses(4),
+            np.empty((0,), dtype=np.int64),
+            np.array([3], dtype=np.int64),
+            _selection_cfg(),
+            use_plucker=True,
+        )
+        selected = memory_selection._select_by_point_union_from_pool(pool, count=1)
+
+        self.assertEqual(tuple(pool["candidates_t"].shape), (0,))
+        self.assertEqual(tuple(pool["candidate_poses"].shape), (0, 5))
+        self.assertEqual(tuple(pool["inside"].shape), (0, 0))
+        self.assertEqual(pool["positive_fov"].dtype, torch.bool)
+        self.assertEqual(selected.tolist(), [])
+
+    def test_validation_fov_revisit_selection_stays_causal(self):
+        poses = np.zeros((8, 5), dtype=np.float32)
+        poses[:, 0] = 100.0
+        poses[:, 4] = 180.0
+        poses[0] = np.asarray([0, 0, 0, 0, 0], dtype=np.float32)
+        poses[5] = np.asarray([0, 0, 0, 0, 0], dtype=np.float32)
+        poses[6] = np.asarray([0, 0, 0, 0, 0], dtype=np.float32)
+        cfg = _selection_cfg(
+            causal=False,
+            max_anchor_frames=0,
+            max_dynamic_frames=0,
+            max_revisit_frames=2,
+            fov_overlap_threshold=0.5,
+            local_context_exclusion_frames=1,
+        )
+
+        indices, masks = select_memory_indices(poses, np.array([5]), cfg, split="validation")
+
+        selected = indices["revisit"][masks["revisit"]]
+        self.assertEqual(selected.tolist(), [0])
+        self.assertTrue(np.all(selected < 5))
+
     def test_training_revisit_uses_pose_similarity_threshold(self):
         poses = np.array(
             [