Restore revisit local context exclusion

datasets/video/memory_selection.py

@@ -32,6 +32,17 @@ 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:
+    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
+
+    target_start = int(target_positions[0])
+    target_stop = int(target_positions[-1]) + 1
+    local_start = max(0, target_start - exclusion)
+    return candidates[(candidates < local_start) | (candidates >= target_stop)]
+
+
 def _as_pose_array(poses) -> np.ndarray:
     poses = np.asarray(poses, dtype=np.float32)
     if poses.ndim != 2 or poses.shape[1] < 5:
@@ -795,6 +806,10 @@ def _select_revisit(
     if len(candidates) == 0:
         return np.empty((0,), dtype=np.int64)
 
+    candidates = _exclude_revisit_local_context(candidates, target_positions, cfg)
+    if len(candidates) == 0:
+        return np.empty((0,), dtype=np.int64)
+
     if len(excluded) > 0:
         candidates = candidates[~np.isin(candidates, excluded)]
     if len(candidates) == 0:
@@ -828,7 +843,7 @@ def select_memory_indices(
     anchor_candidate_stop=None,
     dynamic_stream=None,
 ):
-    """Select causal memory indices and masks for [anchor][dynamic][revisit]."""
+    """Select memory indices and masks for [anchor][dynamic][revisit]."""
 
     poses = _as_pose_array(poses)
     target_positions = np.asarray(target_positions, dtype=np.int64)

tests/test_dememwm_latent_dataset.py

@@ -141,6 +141,30 @@ class MemorySelectionTests(unittest.TestCase):
 
                 self.assertEqual(candidates.tolist(), [1, 2])
 
+    def test_revisit_excludes_local_context_and_target_window_but_keeps_future(self):
+        import datasets.video.memory_selection as memory_selection
+
+        poses = _poses(12)
+        cfg = _selection_cfg(
+            causal=False,
+            max_anchor_frames=0,
+            max_dynamic_frames=0,
+            max_revisit_frames=2,
+            local_context_exclusion_frames=2,
+        )
+        seen_candidates = []
+
+        def fake_pose_similarity(poses_arg, candidates, target_positions, cfg_arg, count, rng=None):
+            seen_candidates.append(candidates.copy())
+            return np.asarray([1, 8], dtype=np.int64)
+
+        with mock.patch.object(memory_selection, "_select_by_pose_similarity", side_effect=fake_pose_similarity):
+            indices, masks = select_memory_indices(poses, np.array([4, 5, 6, 7]), cfg, split="training")
+
+        self.assertEqual([candidates.tolist() for candidates in seen_candidates], [[0, 1, 8, 9, 10, 11]])
+        self.assertEqual(indices["revisit"].tolist(), [1, 8])
+        self.assertEqual(masks["revisit"].tolist(), [True, True])
+
     def test_dememwm_latent_config_has_dataset_side_memory_causal_default(self):
         config_path = Path(__file__).resolve().parents[1] / "configurations" / "dataset" / "video_minecraft_dememwm_latent.yaml"
 
@@ -279,11 +303,12 @@ class MemorySelectionTests(unittest.TestCase):
         selected = indices["revisit"][masks["revisit"]]
         self.assertIn(6, selected.tolist())
 
-    def test_training_noncausal_revisit_can_select_target_window_frame(self):
+    def test_training_noncausal_revisit_excludes_target_window_frame(self):
         poses = np.zeros((8, 5), dtype=np.float32)
         poses[:, 0] = 100.0
         poses[:, 4] = 180.0
         poses[4] = 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,
@@ -300,7 +325,7 @@ class MemorySelectionTests(unittest.TestCase):
             rng=np.random.default_rng(0),
         )
 
-        self.assertEqual(indices["revisit"].tolist(), [4])
+        self.assertEqual(indices["revisit"].tolist(), [6])
         self.assertEqual(masks["revisit"].tolist(), [True])
 
     def test_validation_and_test_noncausal_revisit_stays_past_only(self):
@@ -324,8 +349,8 @@ class MemorySelectionTests(unittest.TestCase):
             with self.subTest(split=split), mock.patch.object(memory_selection, "_select_by_point_union", side_effect=fake_point_union):
                 indices, masks = select_memory_indices(poses, np.array([4]), cfg, split=split)
 
-            self.assertEqual([candidates.tolist() for candidates in seen_candidates], [[0, 1, 2, 3]])
-            self.assertEqual(indices["revisit"].tolist(), [3])
+            self.assertEqual([candidates.tolist() for candidates in seen_candidates], [[0, 1]])
+            self.assertEqual(indices["revisit"].tolist(), [1])
             self.assertEqual(masks["revisit"].tolist(), [True])
 
     def test_revisit_can_duplicate_anchor_and_dynamic_when_scored_best(self):
@@ -340,6 +365,7 @@ class MemorySelectionTests(unittest.TestCase):
             max_dynamic_frames=1,
             max_revisit_frames=2,
             pose_similarity_threshold=0.95,
+            local_context_exclusion_frames=0,
         )
 
         indices, masks = select_memory_indices(