Allow noncausal DeMemWM dynamic events

.exp_artifact/dememwm_future_memory_selection_plan.md

@@ -274,7 +274,7 @@ Allow noncausal DeMemWM revisit selection
 
 ## Substep 3: Enable Future Event-Triggered Dynamic Selection
 
-Status: `[ ]`
+Status: `[x]`
 
 Goal:
 

datasets/video/memory_selection.py

@@ -883,12 +883,21 @@ def select_memory_indices(
     )
 
     if policy == "event_triggered":
+        dynamic_references = revisit if len(revisit) > 0 else target_positions
         if dynamic_stream is not None:
-            exclusion = max(0, int(cfg_get(cfg, "local_context_exclusion_frames", 8)))
-            causal_cutoff = max(min_candidate_frame, target_start - exclusion)
+            candidates = _memory_candidate_frames(
+                len(poses),
+                target_positions,
+                cfg,
+                split,
+                min_candidate_frame=min_candidate_frame,
+            )
             stream = np.asarray(dynamic_stream, dtype=np.int64)
-            eligible_stream = stream[(stream >= min_candidate_frame) & (stream < causal_cutoff)]
-            dynamic = _select_dynamic_from_stream(eligible_stream, revisit, counts["dynamic"])
+            if len(candidates) > 0:
+                eligible_stream = stream[np.isin(stream, candidates)]
+            else:
+                eligible_stream = np.empty((0,), dtype=np.int64)
+            dynamic = _select_dynamic_from_stream(eligible_stream, dynamic_references, counts["dynamic"])
         else:
             dynamic = _select_dynamic_by_policy(
                 target_start,
@@ -898,8 +907,7 @@ def select_memory_indices(
                 latents=latents,
                 actions=actions,
                 min_candidate_frame=min_candidate_frame,
-                reference_frames=revisit if len(revisit) > 0 else None,
-                excluded=np.concatenate([anchor, revisit]) if len(anchor) or len(revisit) else None,
+                reference_frames=dynamic_references,
             )
 
     _write_segment(indices, masks, "anchor", anchor)

tests/test_dememwm_latent_dataset.py

@@ -384,6 +384,120 @@ class MemorySelectionTests(unittest.TestCase):
         self.assertEqual(indices["anchor"].tolist(), [2, 5])
         self.assertEqual(masks["anchor"].tolist(), [True, True])
 
+    def test_event_triggered_dynamic_causal_stream_stays_before_target(self):
+        poses = _poses(12)
+        cfg = _selection_cfg(
+            max_anchor_frames=0,
+            max_dynamic_frames=3,
+            max_revisit_frames=0,
+            dynamic=_dynamic_cfg("event_triggered"),
+        )
+
+        indices, masks = select_memory_indices(
+            poses,
+            np.array([5]),
+            cfg,
+            split="training",
+            dynamic_stream=np.asarray([1, 4, 5, 7, 9], dtype=np.int64),
+        )
+
+        selected = indices["dynamic"][masks["dynamic"]]
+        self.assertEqual(selected.tolist(), [1, 4])
+        self.assertTrue(np.all(selected < 5))
+
+    def test_event_triggered_dynamic_noncausal_stream_can_select_future(self):
+        poses = _poses(12)
+        cfg = _selection_cfg(
+            causal=False,
+            max_anchor_frames=0,
+            max_dynamic_frames=3,
+            max_revisit_frames=0,
+            dynamic=_dynamic_cfg("event_triggered"),
+        )
+
+        indices, masks = select_memory_indices(
+            poses,
+            np.array([5]),
+            cfg,
+            split="training",
+            dynamic_stream=np.asarray([1, 4, 5, 7, 9], dtype=np.int64),
+        )
+
+        selected = indices["dynamic"][masks["dynamic"]]
+        self.assertEqual(selected.tolist(), [4, 5, 7])
+        self.assertTrue(np.any(selected > 5))
+
+    def test_event_triggered_dynamic_eval_stream_stays_past_when_noncausal(self):
+        poses = _poses(12)
+        cfg = _selection_cfg(
+            causal=False,
+            max_anchor_frames=0,
+            max_dynamic_frames=3,
+            max_revisit_frames=0,
+            dynamic=_dynamic_cfg("event_triggered"),
+        )
+
+        for split in ("validation", "test"):
+            with self.subTest(split=split):
+                indices, masks = select_memory_indices(
+                    poses,
+                    np.array([5]),
+                    cfg,
+                    split=split,
+                    dynamic_stream=np.asarray([1, 4, 5, 7, 9], dtype=np.int64),
+                )
+
+                selected = indices["dynamic"][masks["dynamic"]]
+                self.assertEqual(selected.tolist(), [1, 4])
+                self.assertTrue(np.all(selected < 5))
+
+    def test_event_triggered_dynamic_can_duplicate_anchor_and_revisit(self):
+        import datasets.video.memory_selection as memory_selection
+
+        poses = _poses(8)
+        cfg = _selection_cfg(
+            max_anchor_frames=1,
+            max_dynamic_frames=2,
+            max_revisit_frames=2,
+            anchor_diverse_selection=False,
+            dynamic=_dynamic_cfg("event_triggered"),
+        )
+
+        with mock.patch.object(memory_selection, "_select_revisit", return_value=np.asarray([0, 5], dtype=np.int64)):
+            indices, masks = select_memory_indices(
+                poses,
+                np.array([6]),
+                cfg,
+                split="training",
+                dynamic_stream=np.asarray([0, 5], dtype=np.int64),
+            )
+
+        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())
+
+    def test_recent_dynamic_remains_past_only_when_memory_selection_is_noncausal(self):
+        poses = _poses(10)
+        cfg = _selection_cfg(
+            causal=False,
+            max_anchor_frames=0,
+            max_dynamic_frames=2,
+            max_revisit_frames=0,
+            dynamic=_dynamic_cfg("recent"),
+        )
+
+        indices, masks = select_memory_indices(
+            poses,
+            np.array([6]),
+            cfg,
+            split="training",
+            dynamic_stream=np.asarray([6, 7, 8, 9], dtype=np.int64),
+        )
+
+        self.assertEqual(indices["dynamic"].tolist(), [4, 5])
+        self.assertEqual(masks["dynamic"].tolist(), [True, True])
+        self.assertTrue(np.all(indices["dynamic"][masks["dynamic"]] < 6))
 
     def test_event_triggered_dynamic_selects_first_stable_post_event_frame(self):
         poses = np.zeros((10, 5), dtype=np.float32)