Allow noncausal DeMemWM revisit selection

.exp_artifact/dememwm_future_memory_selection_plan.md

@@ -201,7 +201,7 @@ Add DeMemWM memory causality config
 
 ## Substep 2: Enable Future And Local-Window Revisit Selection
 
-Status: `[ ]`
+Status: `[x]`
 
 Goal:
 

algorithms/dememwm/df_video.py

@@ -1544,20 +1544,14 @@ class DeMemWMMinecraft(DiffusionForcingBase):
                             dynamic_count,
                             memory_selection_cfg,
                         )
-                        revisit_excluded = (
-                            np.concatenate([selected["anchor"], selected["dynamic"]])
-                            if len(selected["anchor"]) or len(selected["dynamic"])
-                            else np.empty((0,), dtype=np.int64)
-                        )
                     else:
                         selected["dynamic"] = np.empty((0,), dtype=np.int64)
-                        revisit_excluded = selected["anchor"]
                     selected["revisit"] = _select_revisit(
                         poses_np[:, batch_i, :5],
                         target_positions,
                         memory_selection_cfg,
                         revisit_count,
-                        revisit_excluded,
+                        np.empty((0,), dtype=np.int64),
                         namespace,
                     )
                     if dynamic_policy == "event_triggered":

datasets/video/memory_selection.py

@@ -785,10 +785,13 @@ def _select_revisit(
     if count <= 0:
         return np.empty((0,), dtype=np.int64)
 
-    target_start = int(target_positions[0])
-    exclusion = max(0, int(cfg_get(cfg, "local_context_exclusion_frames", 8)))
-    causal_stop = max(0, target_start - exclusion)
-    candidates = np.arange(min_candidate_frame, causal_stop, dtype=np.int64)
+    candidates = _memory_candidate_frames(
+        len(poses),
+        target_positions,
+        cfg,
+        split,
+        min_candidate_frame=min_candidate_frame,
+    )
     if len(candidates) == 0:
         return np.empty((0,), dtype=np.int64)
 
@@ -868,15 +871,12 @@ def select_memory_indices(
         anchor_candidates = anchor_candidates[~np.isin(anchor_candidates, dynamic)]
     anchor = _select_anchor(anchor_candidates, counts["anchor"], cfg, poses=poses)
 
-    revisit_excluded = anchor if policy == "event_triggered" else (
-        np.concatenate([anchor, dynamic]) if len(anchor) or len(dynamic) else np.empty((0,), dtype=np.int64)
-    )
     revisit = _select_revisit(
         poses,
         target_positions,
         cfg,
         counts["revisit"],
-        revisit_excluded,
+        np.empty((0,), dtype=np.int64),
         split,
         rng=rng,
         min_candidate_frame=min_candidate_frame,

tests/test_dememwm_latent_dataset.py

@@ -232,17 +232,137 @@ class MemorySelectionTests(unittest.TestCase):
         self.assertEqual(indices["revisit"].tolist(), [0, -1])
         self.assertEqual(masks["revisit"].tolist(), [True, False])
 
+    def test_training_default_causal_revisit_does_not_select_future(self):
+        poses = _poses(8)
+        poses[6] = poses[4]
+        cfg = _selection_cfg(
+            max_anchor_frames=0,
+            max_dynamic_frames=0,
+            max_revisit_frames=1,
+            pose_similarity_threshold=0.0,
+        )
+
+        indices, masks = select_memory_indices(
+            poses,
+            np.array([4]),
+            cfg,
+            split="training",
+            rng=np.random.default_rng(0),
+        )
+
+        selected = indices["revisit"][masks["revisit"]]
+        self.assertTrue(len(selected) > 0)
+        self.assertTrue(np.all(selected < 4))
+
+    def test_training_noncausal_revisit_can_select_future_pose_match(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,
+            max_dynamic_frames=0,
+            max_revisit_frames=2,
+            pose_similarity_threshold=0.95,
+        )
+
+        indices, masks = select_memory_indices(
+            poses,
+            np.array([4]),
+            cfg,
+            split="training",
+            rng=np.random.default_rng(0),
+        )
+
+        selected = indices["revisit"][masks["revisit"]]
+        self.assertIn(6, selected.tolist())
+
+    def test_training_noncausal_revisit_can_select_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)
+        cfg = _selection_cfg(
+            causal=False,
+            max_anchor_frames=0,
+            max_dynamic_frames=0,
+            max_revisit_frames=1,
+            pose_similarity_threshold=0.95,
+        )
+
+        indices, masks = select_memory_indices(
+            poses,
+            np.array([4]),
+            cfg,
+            split="training",
+            rng=np.random.default_rng(0),
+        )
+
+        self.assertEqual(indices["revisit"].tolist(), [4])
+        self.assertEqual(masks["revisit"].tolist(), [True])
+
+    def test_validation_and_test_noncausal_revisit_stays_past_only(self):
+        import datasets.video.memory_selection as memory_selection
+
+        poses = _poses(8)
+        cfg = _selection_cfg(
+            causal=False,
+            max_anchor_frames=0,
+            max_dynamic_frames=0,
+            max_revisit_frames=1,
+        )
+
+        for split in ("validation", "test"):
+            seen_candidates = []
+
+            def fake_point_union(poses_arg, candidates, target_positions, cfg_arg, count, **kwargs):
+                seen_candidates.append(candidates.copy())
+                return candidates[-count:].astype(np.int64, copy=False)
+
+            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(masks["revisit"].tolist(), [True])
+
+    def test_revisit_can_duplicate_anchor_and_dynamic_when_scored_best(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(
+            max_anchor_frames=1,
+            max_dynamic_frames=1,
+            max_revisit_frames=2,
+            pose_similarity_threshold=0.95,
+        )
+
+        indices, masks = select_memory_indices(
+            poses,
+            np.array([6]),
+            cfg,
+            split="training",
+            rng=np.random.default_rng(0),
+        )
+
+        self.assertEqual(indices["anchor"].tolist(), [0])
+        self.assertEqual(indices["dynamic"].tolist(), [5])
+        self.assertEqual(indices["revisit"].tolist(), [0, 5])
+        self.assertTrue(masks["revisit"].all())
+
     def test_dememwm_selection_is_typed_and_causal(self):
         poses = _poses(12)
         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.assertTrue(np.all(indices["revisit"][masks["revisit"]] < 4))
         for key in ("anchor", "dynamic", "revisit"):
             self.assertTrue(np.all(indices[key][masks[key]] < 6))
-        selected = np.concatenate([indices[key][masks[key]] for key in ("anchor", "dynamic", "revisit")])
-        self.assertEqual(len(selected), len(np.unique(selected)))
 
     def test_anchor_selection_is_context_bounded_farthest_point(self):
         poses = _poses(10)
@@ -950,7 +1070,7 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
 
         self.assertEqual(harness.horizons, [2, 1])
         self.assertEqual(anchor_calls, [(2, 1, (2, 5))])
-        self.assertEqual(revisit_calls, [([2, 3], "test", 0.75, [0, 1]), ([4], "test", 0.75, [0, 2])])
+        self.assertEqual(revisit_calls, [([2, 3], "test", 0.75, []), ([4], "test", 0.75, [])])
         self.assertEqual(len(revisit_calls), len(harness.horizons))
         self.assertEqual(len(calls), 4)
         self.assertTrue(all(required_kwargs <= call["kwargs"].keys() for call in calls))
@@ -963,10 +1083,10 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         self.assertTrue(torch.equal(calls[0]["kwargs"]["frame_idx"], calls[1]["kwargs"]["frame_idx"]))
         self.assertTrue(torch.equal(calls[2]["kwargs"]["frame_idx"], calls[3]["kwargs"]["frame_idx"]))
         self.assertEqual(calls[0]["x"][:, 0, 0, 0, 0].tolist(), [0.0, 0.0, 0.0, 1.0, 0.0])
-        self.assertEqual(calls[2]["x"][:, 0, 0, 0, 0].tolist(), [2.0, 0.0, 0.0, 2.0, 1.0])
+        self.assertEqual(calls[2]["x"][:, 0, 0, 0, 0].tolist(), [2.0, 0.0, 0.0, 2.0, 2.0])
         self.assertEqual(calls[0]["action_cond"][:, 0, 0].tolist(), [102.0, 103.0, 0.0, 0.0, 0.0])
         self.assertEqual(calls[0]["kwargs"]["frame_idx"][:, 0].tolist(), [12, 13, 10, 11, 0])
-        self.assertEqual(calls[2]["kwargs"]["frame_idx"][:, 0].tolist(), [13, 14, 10, 12, 11])
+        self.assertEqual(calls[2]["kwargs"]["frame_idx"][:, 0].tolist(), [13, 14, 10, 12, 12])
         self.assertEqual(calls[0]["kwargs"]["frame_memory_masks"]["revisit"].tolist(), [[False]])
         self.assertEqual(calls[2]["kwargs"]["frame_memory_masks"]["revisit"].tolist(), [[True]])
         self.assertAlmostEqual(float(loss), 1.0 / 3.0)
@@ -1209,7 +1329,7 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
                 actions=np.ones((num_frames, 25), dtype=np.float32),
             )
             dataset = MinecraftVideoDeMemWMLatentDataset(
-                _dataset_cfg(root, context_length=6, memory_selection=_selection_cfg()),
+                _dataset_cfg(root, context_length=6, memory_selection=_selection_cfg(training_use_plucker=True)),
                 split="training",
             )
             sample = dataset[0]
@@ -1220,7 +1340,7 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         self.assertEqual(sample["frame_indices"][3:5].tolist(), [100, 101])
         self.assertEqual(sample["frame_indices"][5:7].tolist(), [104, 105])
         revisit_frames = sample["frame_indices"][7:9].numpy()
-        self.assertEqual(revisit_frames.tolist(), [102, 103])
+        self.assertEqual(revisit_frames.tolist(), [104, 105])
         self.assertEqual(len(np.unique(revisit_frames)), 2)
         self.assertTrue(sample["memory_masks"]["target"].all().item())
         self.assertTrue(sample["memory_masks"]["anchor"].all().item())