Fix disabled DeMemWM memory selection slots

datasets/video/memory_selection.py

@@ -392,15 +392,15 @@ def select_memory_indices(poses, target_positions, cfg=None, split: str = "train
     poses = _as_pose_array(poses)
     target_positions = np.asarray(target_positions, dtype=np.int64)
     target_positions = target_positions[(target_positions >= 0) & (target_positions < len(poses))]
+    enabled = bool(cfg_get(cfg, "enabled", True))
 
     counts = {
-        "anchor": int(cfg_get(cfg, "max_anchor_frames", 0)),
-        "dynamic": int(cfg_get(cfg, "max_dynamic_frames", 0)),
-        "revisit": int(cfg_get(cfg, "max_revisit_frames", 0)),
+        "anchor": int(cfg_get(cfg, "max_anchor_frames", 0)) if enabled else 0,
+        "dynamic": int(cfg_get(cfg, "max_dynamic_frames", 0)) if enabled else 0,
+        "revisit": int(cfg_get(cfg, "max_revisit_frames", 0)) if enabled else 0,
     }
     indices, masks = _empty_selection(counts)
 
-    enabled = bool(cfg_get(cfg, "enabled", True))
     backend = cfg_get(cfg, "backend", "dememwm")
     backend = "dememwm" if backend is None else str(backend).lower()
     if enabled and backend != "dememwm":
@@ -429,9 +429,10 @@ def select_memory_indices(poses, target_positions, cfg=None, split: str = "train
 
 
 def memory_segment_lengths(target_length: int, cfg=None) -> Dict[str, int]:
+    enabled = bool(cfg_get(cfg, "enabled", True))
     return {
         "target": int(target_length),
-        "anchor": int(cfg_get(cfg, "max_anchor_frames", 0)),
-        "dynamic": int(cfg_get(cfg, "max_dynamic_frames", 0)),
-        "revisit": int(cfg_get(cfg, "max_revisit_frames", 0)),
+        "anchor": int(cfg_get(cfg, "max_anchor_frames", 0)) if enabled else 0,
+        "dynamic": int(cfg_get(cfg, "max_dynamic_frames", 0)) if enabled else 0,
+        "revisit": int(cfg_get(cfg, "max_revisit_frames", 0)) if enabled else 0,
     }

tests/test_dememwm_latent_dataset.py

@@ -630,6 +630,24 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         with self.assertRaisesRegex(NotImplementedError, "packed .* frame-memory API"):
             DeMemWMMinecraft.interactive(object(), None, None, None, None, None, None, None, None, None)
 
+    def test_disabled_memory_selection_does_not_allocate_memory_slots(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp)
+            _write_latent_npz(root / "training" / "sample.npz", 104)
+            dataset = MinecraftVideoDeMemWMLatentDataset(
+                _dataset_cfg(root, memory_selection=_selection_cfg(enabled=False)),
+                split="training",
+            )
+            sample = dataset[0]
+
+        self.assertEqual(dataset.memory_condition_length, 0)
+        self.assertEqual(sample["memory_segments"], {"target": 3, "anchor": 0, "dynamic": 0, "revisit": 0})
+        self.assertEqual(tuple(sample["latents"].shape), (3, 1, 1, 2))
+        self.assertEqual(sample["frame_indices"].tolist(), [100, 101, 102])
+        self.assertTrue(sample["memory_masks"]["target"].all().item())
+        for key in ("anchor", "dynamic", "revisit"):
+            self.assertEqual(tuple(sample["memory_masks"][key].shape), (0,))
+
     def test_dataset_returns_target_anchor_dynamic_revisit_contract(self):
         with tempfile.TemporaryDirectory() as tmp:
             root = Path(tmp)