Optimize DeMemWM dynamic selection defaults

algorithms/dememwm/df_video.py

@@ -2005,9 +2005,14 @@ class DeMemWMMinecraft(DiffusionForcingBase):
                             curr_frame,
                             memory_selection_cfg,
                         )
+                        dynamic_stop = max(
+                            0,
+                            int(target_positions[0])
+                            - max(0, int(_cfg_get(memory_selection_cfg, "local_context_exclusion_frames", self.n_tokens))),
+                        )
                         selected["dynamic"] = _select_online_event_dynamic_from_cache(
                             cache,
-                            start_frame,
+                            dynamic_stop,
                             dynamic_count,
                             memory_selection_cfg,
                             reference_frames=selected["revisit"] if len(selected["revisit"]) > 0 else None,

configurations/dataset/video_minecraft_dememwm_latent.yaml

@@ -31,8 +31,8 @@ memory_selection:
   local_context_exclusion_frames: 8
   plucker_moment_radius: 30.0
   anchor_diverse_selection: true
-  pose_preselect_topk: 64
-  candidate_chunk_size: 64
+  pose_preselect_topk: 32
+  candidate_chunk_size: 0
   dynamic:
     selection_policy: multiview
     multiview_selector: fov_greedy

datasets/video/memory_selection.py

@@ -15,6 +15,8 @@ _FOV_NUM_POINTS = 10000
 _FOV_RADIUS = 30.0
 _FOV_HALF_H = 105.0 / 2.0
 _FOV_HALF_V = 75.0 / 2.0
+_FOV_COS_HALF_H = float(np.cos(np.deg2rad(_FOV_HALF_H)))
+_FOV_COS_HALF_V = float(np.cos(np.deg2rad(_FOV_HALF_V)))
 _POSE_DISTANCE_SCALE = 30.0
 _ANGLE_DISTANCE_SCALE = 180.0
 _DETERMINISTIC_UNIT_BALL_CACHE: Dict[Tuple[int, str, torch.dtype], torch.Tensor] = {}
@@ -35,7 +37,7 @@ def _memory_candidate_frames(num_frames, target_positions, cfg, split, min_candi
     return np.arange(min_candidate_frame, stop, dtype=np.int64)
 
 
-def _exclude_revisit_local_context(candidates: np.ndarray, target_positions: np.ndarray, cfg) -> np.ndarray:
+def _exclude_local_context(candidates: np.ndarray, target_positions: np.ndarray, cfg) -> np.ndarray:
     exclusion = max(0, int(cfg_get(cfg, "local_context_exclusion_frames", 8)))
     if exclusion <= 0 or len(candidates) == 0 or len(target_positions) == 0:
         return candidates
@@ -46,6 +48,10 @@ def _exclude_revisit_local_context(candidates: np.ndarray, target_positions: np.
     return candidates[(candidates < local_start) | (candidates >= target_stop)]
 
 
+def _exclude_revisit_local_context(candidates: np.ndarray, target_positions: np.ndarray, cfg) -> np.ndarray:
+    return _exclude_local_context(candidates, target_positions, cfg)
+
+
 def _as_pose_array(poses) -> np.ndarray:
     poses = np.asarray(poses, dtype=np.float32)
     if poses.ndim != 2 or poses.shape[1] < 5:
@@ -135,11 +141,17 @@ def _inside_fov_points(points: torch.Tensor, poses: torch.Tensor) -> torch.Tenso
     x = vectors[..., 0]
     y = vectors[..., 1]
     z = vectors[..., 2]
-    azimuth = torch.rad2deg(torch.atan2(x, z))
-    elevation = torch.rad2deg(torch.atan2(y, torch.sqrt(x * x + z * z)))
-    yaw_diff = _wrap_degrees(azimuth - poses[:, None, 4])
-    pitch_diff = _wrap_degrees(elevation - poses[:, None, 3])
-    return (yaw_diff < _FOV_HALF_H) & (pitch_diff < _FOV_HALF_V)
+    horizontal = torch.sqrt((x * x + z * z).clamp_min(0.0))
+    ray_norm = torch.sqrt((horizontal * horizontal + y * y).clamp_min(0.0))
+
+    yaw = torch.deg2rad(poses[:, 4])[:, None]
+    pitch = torch.deg2rad(poses[:, 3])[:, None]
+    yaw_aligned = x * torch.sin(yaw) + z * torch.cos(yaw)
+    pitch_aligned = horizontal * torch.cos(pitch) + y * torch.sin(pitch)
+
+    # Same yaw/elevation half-angle convention as the previous atan2 path,
+    # but without per-point inverse trig or degree wrapping.
+    return (yaw_aligned > horizontal * _FOV_COS_HALF_H) & (pitch_aligned > ray_norm * _FOV_COS_HALF_V)
 
 
 def _target_fov_points(target_poses: torch.Tensor) -> torch.Tensor:
@@ -162,7 +174,7 @@ def _pose_preselect(
     *,
     extra_topk: int = 0,
 ) -> np.ndarray:
-    topk = cfg_get(cfg, "pose_preselect_topk", 64)
+    topk = cfg_get(cfg, "pose_preselect_topk", 32)
     if topk is None:
         return candidates
     topk = int(topk)
@@ -196,7 +208,7 @@ def _candidate_fov_masks(poses_t: torch.Tensor, candidates: np.ndarray, target_p
         )
 
     candidates_t = torch.as_tensor(candidates, device=poses_t.device, dtype=torch.long)
-    chunk_size = int(cfg_get(cfg, "candidate_chunk_size", 64))
+    chunk_size = int(cfg_get(cfg, "candidate_chunk_size", 0))
     chunk_size = len(candidates) if chunk_size <= 0 else chunk_size
     inside_parts = []
     score_parts = []
@@ -358,10 +370,11 @@ def _build_shared_fov_candidate_pool(
     poses_t = _as_pose_tensor(poses)
     dynamic_candidates = None
     if int(dynamic_count) > 0:
+        dynamic_base_candidates = _exclude_local_context(base_candidates, target_positions, cfg)
         # Revisit frames are excluded after selection; include replacement rows
         # that can enter dynamic's pose-preselected top-k after that exclusion.
         dynamic_candidates = _pose_preselect(
-            base_candidates,
+            dynamic_base_candidates,
             poses_t,
             target_positions,
             cfg,
@@ -556,7 +569,7 @@ def _select_dynamic_multiview_pose_plucker_fps(
     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)
+    topk = cfg_get(cfg, "pose_preselect_topk", 32)
     if topk is not None:
         topk = int(topk)
         if topk > 0:
@@ -846,6 +859,7 @@ def _select_dynamic_multiview(
         split,
         min_candidate_frame=min_candidate_frame,
     )
+    candidates = _exclude_local_context(candidates, target_positions, cfg)
     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:
@@ -1240,6 +1254,7 @@ def _select_event_triggered_dynamic(
         actions=actions,
         min_candidate_frame=min_candidate_frame,
     )
+    event_anchors = _exclude_local_context(event_anchors, np.asarray([target_start], dtype=np.int64), cfg)
     if excluded is not None and len(excluded) > 0 and len(event_anchors) > 0:
         event_anchors = event_anchors[~np.isin(event_anchors, np.asarray(excluded, dtype=np.int64))]
     if len(event_anchors) == 0:
@@ -1268,7 +1283,11 @@ def _select_dynamic_by_policy(
 
     policy = _dynamic_policy(cfg)
     if policy == "recent":
-        return _select_dynamic(target_start, count, min_candidate_frame)
+        dynamic_stop = int(target_start)
+        if target_positions is not None:
+            exclusion = max(0, int(cfg_get(cfg, "local_context_exclusion_frames", 8)))
+            dynamic_stop = max(min_candidate_frame, dynamic_stop - exclusion) if exclusion > 0 else dynamic_stop
+        return _select_dynamic(dynamic_stop, count, min_candidate_frame)
     if policy == "event_triggered":
         return _select_event_triggered_dynamic(
             target_start,
@@ -1421,6 +1440,7 @@ def select_memory_indices(
             counts["dynamic"],
             cfg,
             min_candidate_frame=min_candidate_frame,
+            target_positions=target_positions,
         )
     elif policy == "event_triggered":
         dynamic_references = revisit if len(revisit) > 0 else target_positions
@@ -1432,6 +1452,7 @@ def select_memory_indices(
                 split,
                 min_candidate_frame=min_candidate_frame,
             )
+            candidates = _exclude_local_context(candidates, target_positions, cfg)
             stream = np.asarray(dynamic_stream, dtype=np.int64)
             if len(candidates) > 0:
                 eligible_stream = stream[np.isin(stream, candidates)]

tests/test_dememwm_latent_dataset.py

@@ -53,8 +53,8 @@ def _selection_cfg(**overrides):
         "local_context_exclusion_frames": 2,
         "plucker_moment_radius": 30.0,
         "anchor_diverse_selection": False,
-        "pose_preselect_topk": 64,
-        "candidate_chunk_size": 64,
+        "pose_preselect_topk": 32,
+        "candidate_chunk_size": 0,
         "dynamic": _dynamic_cfg(),
     }
     cfg.update(overrides)
@@ -217,8 +217,8 @@ class MemorySelectionTests(unittest.TestCase):
         self.assertIs(cfg.memory_selection.causal, True)
         self.assertEqual(cfg.memory_selection.dynamic.selection_policy, "multiview")
         self.assertEqual(cfg.memory_selection.dynamic.multiview_selector, "fov_greedy")
-        self.assertEqual(cfg.memory_selection.pose_preselect_topk, 64)
-        self.assertEqual(cfg.memory_selection.candidate_chunk_size, 64)
+        self.assertEqual(cfg.memory_selection.pose_preselect_topk, 32)
+        self.assertEqual(cfg.memory_selection.candidate_chunk_size, 0)
 
     def test_revisit_uses_sampled_fov_selection(self):
         poses = np.array(
@@ -551,7 +551,7 @@ class MemorySelectionTests(unittest.TestCase):
     def test_dynamic_multiview_provided_fov_pool_keeps_excluded_reference_coverage(self):
         import datasets.video.memory_selection as memory_selection
 
-        poses = _poses(6)
+        poses = _poses(8)
         cfg = _selection_cfg(dynamic=_dynamic_cfg("multiview", multiview_selector="fov_greedy"))
         pool = _fov_pool(
             [0, 2, 3],
@@ -569,7 +569,7 @@ class MemorySelectionTests(unittest.TestCase):
         ):
             selected = memory_selection._select_dynamic_multiview(
                 poses,
-                np.asarray([5], dtype=np.int64),
+                np.asarray([7], dtype=np.int64),
                 cfg,
                 count=1,
                 excluded=np.asarray([0], dtype=np.int64),
@@ -587,10 +587,10 @@ class MemorySelectionTests(unittest.TestCase):
         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]),
+            (default_cfg, "training", np.asarray([5], dtype=np.int64), [1, 2], [1, 2]),
+            (noncausal_cfg, "training", np.asarray([3], dtype=np.int64), [4, 5, 6, 7], [6, 7]),
+            (noncausal_cfg, "validation", np.asarray([5], dtype=np.int64), [1, 2], [1, 2]),
+            (noncausal_cfg, "test", np.asarray([5], dtype=np.int64), [1, 2], [1, 2]),
         ]
 
         def fake_select(pool, count, **kwargs):
@@ -619,7 +619,7 @@ class MemorySelectionTests(unittest.TestCase):
             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(first.tolist(), [2])
         self.assertEqual(second.tolist(), first.tolist())
         self.assertTrue(set(first.tolist()).isdisjoint({0, 1}))
 
@@ -785,11 +785,11 @@ class MemorySelectionTests(unittest.TestCase):
         indices, masks = select_memory_indices(poses, np.array([6, 7, 8]), _selection_cfg())
 
         self.assertEqual(indices["anchor"].tolist(), [0, 1])
-        self.assertEqual(indices["dynamic"].tolist(), [4, 5])
+        self.assertEqual(indices["dynamic"].tolist(), [2, 3])
         for key in ("anchor", "dynamic", "revisit"):
             self.assertTrue(np.all(indices[key][masks[key]] < 6))
 
-    def test_anchor_can_duplicate_recent_dynamic_when_selected(self):
+    def test_recent_dynamic_excludes_anchor_inside_local_context(self):
         poses = _poses(8)
         cfg = _selection_cfg(
             max_anchor_frames=1,
@@ -808,7 +808,7 @@ class MemorySelectionTests(unittest.TestCase):
         )
 
         self.assertEqual(indices["anchor"].tolist(), [5])
-        self.assertEqual(indices["dynamic"].tolist(), [5])
+        self.assertEqual(indices["dynamic"].tolist(), [3])
         self.assertEqual(masks["anchor"].tolist(), [True])
         self.assertEqual(masks["dynamic"].tolist(), [True])
         self.assertTrue(np.all(indices["dynamic"][masks["dynamic"]] < 6))
@@ -851,7 +851,7 @@ class MemorySelectionTests(unittest.TestCase):
         )
 
         selected = indices["dynamic"][masks["dynamic"]]
-        self.assertEqual(selected.tolist(), [1, 4])
+        self.assertEqual(selected.tolist(), [1])
         self.assertTrue(np.all(selected < 5))
 
     def test_event_triggered_dynamic_noncausal_stream_can_select_future(self):
@@ -873,7 +873,7 @@ class MemorySelectionTests(unittest.TestCase):
         )
 
         selected = indices["dynamic"][masks["dynamic"]]
-        self.assertEqual(selected.tolist(), [4, 5, 7])
+        self.assertEqual(selected.tolist(), [1, 7, 9])
         self.assertTrue(np.any(selected > 5))
 
     def test_event_triggered_dynamic_eval_stream_stays_past_when_noncausal(self):
@@ -897,7 +897,7 @@ class MemorySelectionTests(unittest.TestCase):
                 )
 
                 selected = indices["dynamic"][masks["dynamic"]]
-                self.assertEqual(selected.tolist(), [1, 4])
+                self.assertEqual(selected.tolist(), [1])
                 self.assertTrue(np.all(selected < 5))
 
     def test_event_triggered_dynamic_can_duplicate_anchor_and_revisit(self):
@@ -923,8 +923,8 @@ class MemorySelectionTests(unittest.TestCase):
 
         self.assertEqual(indices["anchor"].tolist(), [0])
         self.assertEqual(indices["revisit"].tolist(), [0, 5])
-        self.assertEqual(indices["dynamic"].tolist(), [0, 5])
-        self.assertTrue(masks["dynamic"].all())
+        self.assertEqual(indices["dynamic"].tolist(), [0, -1])
+        self.assertEqual(masks["dynamic"].tolist(), [True, False])
 
     def test_recent_dynamic_remains_past_only_when_memory_selection_is_noncausal(self):
         poses = _poses(10)
@@ -944,7 +944,7 @@ class MemorySelectionTests(unittest.TestCase):
             dynamic_stream=np.asarray([6, 7, 8, 9], dtype=np.int64),
         )
 
-        self.assertEqual(indices["dynamic"].tolist(), [4, 5])
+        self.assertEqual(indices["dynamic"].tolist(), [2, 3])
         self.assertEqual(masks["dynamic"].tolist(), [True, True])
         self.assertTrue(np.all(indices["dynamic"][masks["dynamic"]] < 6))
 
@@ -2194,7 +2194,7 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
 
         self.assertEqual(dataset.context_length, 4)
         self.assertEqual(len(dataset), 112 - (100 + 4 + 3) + 1)
-        self.assertEqual(sample["frame_indices"].tolist(), [106, 107, 108, 102, 105, 104, 105])
+        self.assertEqual(sample["frame_indices"].tolist(), [106, 107, 108, 102, 105, 102, 103])
         self.assertTrue(sample["memory_masks"]["anchor"].all().item())
         self.assertTrue(sample["memory_masks"]["dynamic"].all().item())
         self.assertEqual(tuple(sample["memory_masks"]["revisit"].shape), (0,))
@@ -2302,7 +2302,7 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         self.assertEqual(tuple(sample["latents"].shape), (9, 1, 1, 2))
         self.assertEqual(sample["frame_indices"][:3].tolist(), [106, 107, 108])
         self.assertEqual(sample["frame_indices"][3:5].tolist(), [100, 101])
-        self.assertEqual(sample["frame_indices"][5:7].tolist(), [104, 105])
+        self.assertEqual(sample["frame_indices"][5:7].tolist(), [102, 103])
         revisit_frames = sample["frame_indices"][7:9].numpy()
         self.assertEqual(revisit_frames.tolist(), [102, 103])
         self.assertEqual(len(np.unique(revisit_frames)), 2)
@@ -2326,7 +2326,7 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         self.assertTrue(torch.equal(preprocessed["actions"][5:], torch.ones_like(preprocessed["actions"][5:])))
         self.assertEqual(tuple(preprocessed["poses"].shape), (9, 1, 5))
         self.assertIs(preprocessed["frame_memory_pose"], preprocessed["poses"])
-        self.assertEqual(preprocessed["frame_memory_pose"][:7, 0, 0].tolist(), [0.0, 1.0, 2.0, -6.0, -5.0, -2.0, -1.0])
+        self.assertEqual(preprocessed["frame_memory_pose"][:7, 0, 0].tolist(), [0.0, 1.0, 2.0, -6.0, -5.0, -4.0, -3.0])
         self.assertEqual(tuple(preprocessed["frame_indices"].shape), (9, 1))
         self.assertEqual(preprocessed["memory_segments"], sample["memory_segments"])
         self.assertEqual(preprocessed["segment_lengths"], sample["memory_segments"])
@@ -2344,7 +2344,7 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         )
         self.assertEqual(preprocessed["target_tensors"]["frame_indices"][:, 0].tolist(), [106, 107, 108])
         self.assertEqual(preprocessed["stream_tensors"]["anchor"]["frame_indices"][:, 0].tolist(), [100, 101])
-        self.assertEqual(preprocessed["stream_tensors"]["dynamic"]["frame_indices"][:, 0].tolist(), [104, 105])
+        self.assertEqual(preprocessed["stream_tensors"]["dynamic"]["frame_indices"][:, 0].tolist(), [102, 103])
         self.assertEqual(preprocessed["stream_tensors"]["revisit"]["frame_indices"].shape[0], 2)
         self.assertEqual(tuple(preprocessed["memory_masks"]["target"].shape), (1, 3))
         self.assertEqual(tuple(preprocessed["memory_masks"]["anchor"].shape), (1, 2))