Fix DeMemWM memory dropout and inference stabilization

algorithms/dememwm/df_video.py

@@ -645,11 +645,22 @@ def _select_online_event_dynamic_from_cache(
     return _select_dynamic_from_stream(eligible_stream, reference_frames, count)
 
 
+def _stabilized_sampling_levels(reference_levels, count: int):
+    """Schedule level 0 enters sample_step's fixed stabilization path."""
+    levels = reference_levels[:, None] if reference_levels.ndim == 1 else reference_levels
+    return levels.new_zeros((int(count), levels.shape[-1]), dtype=torch.long)
+
+
 def _memory_noise_levels_for_streams(cfg, diffusion_model, query_noise_levels, stream_lengths, mode: str):
     is_training = str(mode) == "training"
     noise_cfg = _cfg_get(cfg, "memory_noise")
     if noise_cfg is None:
-        value = int(getattr(diffusion_model, "stabilization_level", 0)) if is_training else 0
+        if not is_training:
+            return {
+                key: _stabilized_sampling_levels(query_noise_levels, int(stream_lengths[key]))
+                for key in _DEMEMWM_STREAM_KEYS
+            }
+        value = int(getattr(diffusion_model, "stabilization_level", 0))
         return {
             key: torch.full((int(stream_lengths[key]), query_noise_levels.shape[-1]), value, device=query_noise_levels.device, dtype=torch.long)
             for key in _DEMEMWM_STREAM_KEYS
@@ -663,6 +674,11 @@ def _memory_noise_levels_for_streams(cfg, diffusion_model, query_noise_levels, s
     )
     query = query_noise_levels[:, None] if query_noise_levels.ndim == 1 else query_noise_levels
     query_bounds = query.to(dtype=torch.long).clamp_min(0).min(dim=0).values
+    if not noisy_memory:
+        return {
+            key: _stabilized_sampling_levels(query, max(0, int(stream_lengths[key])))
+            for key in _DEMEMWM_STREAM_KEYS
+        }
     if noise_mode == "independent":
         if is_training:
             max_level = int(getattr(diffusion_model, "timesteps", 0) or 0) - 1
@@ -2077,7 +2093,7 @@ class DeMemWMMinecraft(DiffusionForcingBase):
                 for m in range(scheduling_matrix.shape[0] - 1):
                     from_query = torch.as_tensor(scheduling_matrix[m], device=xs.device, dtype=torch.long)[:, None].repeat(1, batch_size)
                     to_query = torch.as_tensor(scheduling_matrix[m + 1], device=xs.device, dtype=torch.long)[:, None].repeat(1, batch_size)
-                    context_levels = from_query.new_zeros((query_offset, batch_size))
+                    context_levels = _stabilized_sampling_levels(from_query, query_offset)
                     from_target = torch.cat([context_levels, from_query], dim=0)
                     to_target = torch.cat([context_levels, to_query], dim=0)
                     memory_noise_levels_by_key = _memory_noise_levels_for_streams(

configurations/algorithm/dememwm_base.yaml

@@ -13,6 +13,8 @@ memory_noise:
   anchor_max_fraction: 0.25
   dynamic_max_fraction: 0.50
   revisit_max_fraction: 0.25
+  # Inference memory follows the WorldMem stabilization path by default.
+  validation_noisy_memory: false
 # Per stream: all = always active, high = active for high query noise,
 # low = active for low query noise.
 noise_route: {anchor: all, dynamic: all, revisit: all}

configurations/dataset/video_minecraft_dememwm_latent.yaml

@@ -26,6 +26,7 @@ memory_selection:
   pose_similarity_threshold: 0.6
   training_use_plucker: true
   training_plucker_weight: 1.0
+  training_dropout: 0.1
   fov_overlap_threshold: 0.6
   min_total_selected_coverage: 0.1
   local_context_exclusion_frames: 8

datasets/video/memory_selection.py

@@ -421,6 +421,27 @@ def _write_segment(indices: MutableMapping[str, np.ndarray], masks: MutableMappi
     masks[key][: len(selected)] = True
 
 
+def _rng_choice(rng, values: np.ndarray, size: int, *, replace: bool) -> np.ndarray:
+    chooser = np.random if rng is None else rng
+    return np.asarray(chooser.choice(values, size=int(size), replace=replace), dtype=np.int64)
+
+
+def _rng_random(rng) -> float:
+    return float(np.random.random() if rng is None else rng.random())
+
+
+def _sample_random_selection(candidates: np.ndarray, count: int, rng=None) -> np.ndarray:
+    count = int(count)
+    if count <= 0:
+        return np.empty((0,), dtype=np.int64)
+
+    candidates = np.asarray(candidates, dtype=np.int64)
+    candidates = candidates[candidates >= 0]
+    if len(candidates) == 0:
+        return np.empty((0,), dtype=np.int64)
+    return _rng_choice(rng, candidates, count, replace=True)
+
+
 def _best_row(rows: torch.Tensor, gains: torch.Tensor, fov_values: torch.Tensor, plucker: torch.Tensor, gaps: torch.Tensor, candidates_t: torch.Tensor) -> int:
     active = rows
     for values in (gains, fov_values, plucker, -gaps.to(dtype=torch.float32), -candidates_t.to(dtype=torch.float32)):
@@ -1319,6 +1340,37 @@ def _select_dynamic_by_policy(
     raise AssertionError(f"unhandled dynamic selection policy: {policy!r}")
 
 
+def _dynamic_random_candidates(
+    num_frames: int,
+    target_start: int,
+    target_positions: np.ndarray,
+    cfg,
+    split: str,
+    min_candidate_frame: int,
+    policy: str,
+    *,
+    excluded: np.ndarray | None = None,
+) -> np.ndarray:
+    if policy == "recent":
+        stop = int(target_start)
+        exclusion = max(0, int(cfg_get(cfg, "local_context_exclusion_frames", 8)))
+        if exclusion > 0:
+            stop = max(int(min_candidate_frame), stop - exclusion)
+        return np.arange(int(min_candidate_frame), stop, dtype=np.int64)
+
+    candidates = _memory_candidate_frames(
+        int(num_frames),
+        target_positions,
+        cfg,
+        split,
+        min_candidate_frame=int(min_candidate_frame),
+    )
+    candidates = _exclude_local_context(candidates, target_positions, cfg)
+    if policy == "multiview" and excluded is not None and len(excluded) > 0 and len(candidates) > 0:
+        candidates = candidates[~np.isin(candidates, np.asarray(excluded, dtype=np.int64))]
+    return candidates.astype(np.int64, copy=False)
+
+
 def _select_revisit(
     poses: np.ndarray,
     target_positions: np.ndarray,
@@ -1433,7 +1485,6 @@ def select_memory_indices(
         min_candidate_frame=min_candidate_frame,
         fov_pool=fov_pool,
     )
-
     if policy == "recent":
         dynamic = _select_dynamic_by_policy(
             target_start,
@@ -1486,6 +1537,30 @@ def select_memory_indices(
     else:
         raise AssertionError(f"unhandled dynamic selection policy: {policy!r}")
 
+    training_dropout = float(cfg_get(cfg, "training_dropout", 0.0)) if split == "training" else 0.0
+    if training_dropout > 0.0 and _rng_random(rng) < training_dropout:
+        revisit_candidates = _memory_candidate_frames(
+            len(poses),
+            target_positions,
+            cfg,
+            split,
+            min_candidate_frame=min_candidate_frame,
+        )
+        revisit_candidates = _exclude_revisit_local_context(revisit_candidates, target_positions, cfg)
+        dynamic_candidates = _dynamic_random_candidates(
+            len(poses),
+            target_start,
+            target_positions,
+            cfg,
+            split,
+            min_candidate_frame,
+            policy,
+            excluded=revisit if policy == "multiview" else None,
+        )
+        anchor = _sample_random_selection(anchor_candidates, counts["anchor"], rng=rng)
+        dynamic = _sample_random_selection(dynamic_candidates, counts["dynamic"], rng=rng)
+        revisit = _sample_random_selection(revisit_candidates, counts["revisit"], rng=rng)
+
     _write_segment(indices, masks, "anchor", anchor)
     _write_segment(indices, masks, "dynamic", dynamic)
     _write_segment(indices, masks, "revisit", revisit)

tests/test_dememwm_latent_dataset.py

@@ -48,6 +48,7 @@ def _selection_cfg(**overrides):
         "pose_similarity_threshold": 0.6,
         "training_use_plucker": False,
         "training_plucker_weight": 0.1,
+        "training_dropout": 0.0,
         "fov_overlap_threshold": 0.0,
         "min_total_selected_coverage": 0.0,
         "local_context_exclusion_frames": 2,
@@ -219,6 +220,7 @@ class MemorySelectionTests(unittest.TestCase):
         self.assertEqual(cfg.memory_selection.dynamic.multiview_selector, "fov_greedy")
         self.assertEqual(cfg.memory_selection.pose_preselect_topk, 32)
         self.assertEqual(cfg.memory_selection.candidate_chunk_size, 0)
+        self.assertEqual(cfg.memory_selection.training_dropout, 0.1)
 
     def test_revisit_uses_sampled_fov_selection(self):
         poses = np.array(
@@ -833,6 +835,32 @@ class MemorySelectionTests(unittest.TestCase):
         self.assertEqual(indices["anchor"].tolist(), [2, 5])
         self.assertEqual(masks["anchor"].tolist(), [True, True])
 
+    def test_training_dropout_randomizes_memory_streams_with_replacement(self):
+        poses = _poses(12)
+        cfg = _selection_cfg(
+            max_anchor_frames=2,
+            max_dynamic_frames=3,
+            max_revisit_frames=2,
+            training_dropout=1.0,
+            dynamic=_dynamic_cfg("recent"),
+        )
+
+        indices, masks = select_memory_indices(
+            poses,
+            np.array([8]),
+            cfg,
+            split="training",
+            rng=np.random.default_rng(0),
+            anchor_candidate_start=2,
+            anchor_candidate_stop=6,
+        )
+
+        self.assertEqual(indices["anchor"].tolist(), [4, 3])
+        self.assertEqual(indices["dynamic"].tolist(), [1, 0, 0])
+        self.assertEqual(indices["revisit"].tolist(), [0, 1])
+        for key in ("anchor", "dynamic", "revisit"):
+            self.assertTrue(masks[key].all())
+
     def test_event_triggered_dynamic_causal_stream_stays_before_target(self):
         poses = _poses(12)
         cfg = _selection_cfg(
@@ -1771,7 +1799,10 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
                 self.memory_condition_length, self.clip_noise = 3, 0.0
                 self.context_frames, self.frame_stack, self.chunk_size, self.n_tokens = 2, 1, 2, 2
                 self.cfg = OmegaConf.create(
-                    {"memory_selection": {"enabled": True, "max_anchor_frames": 1, "max_dynamic_frames": 1, "max_revisit_frames": 1, "fov_overlap_threshold": 0.75}}
+                    {
+                        "memory_selection": {"enabled": True, "max_anchor_frames": 1, "max_dynamic_frames": 1, "max_revisit_frames": 1, "fov_overlap_threshold": 0.75},
+                        "memory_noise": {"enabled": True, "validation_noisy_memory": False},
+                    }
                 )
                 self.diffusion_model = FakeDiffusion()
                 self.logged = []
@@ -1858,6 +1889,8 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         self.assertEqual(calls[0]["action_cond"][:, 0, 0].tolist(), [102.0, 103.0, 0.0, 0.0])
         self.assertEqual(calls[0]["kwargs"]["frame_idx"][:, 0].tolist(), [12, 13, 10, 11])
         self.assertEqual(calls[2]["kwargs"]["frame_idx"][:, 0].tolist(), [13, 14, 10, 12, 12])
+        self.assertEqual(calls[0]["curr_noise_level"][:, 0].tolist(), [2, 2, 0, 0])
+        self.assertEqual(calls[2]["curr_noise_level"][:, 0].tolist(), [0, 2, 0, 0, 0])
         self.assertEqual(calls[0]["kwargs"]["frame_memory_masks"]["revisit"].tolist(), [[]])
         self.assertEqual(calls[2]["kwargs"]["frame_memory_masks"]["revisit"].tolist(), [[True]])
         self.assertAlmostEqual(float(loss), 1.0 / 3.0)