Fix DeMemWM validation rollout packing

algorithms/dememwm/df_video.py

@@ -212,6 +212,33 @@ def _gather_online_memory_tensor(source, indices, masks):
     return gathered * mask.to(dtype=gathered.dtype)
 
 
+def _select_online_anchor_indices(curr_frame: int, n_context_frames: int, count: int, cfg) -> np.ndarray:
+    if count <= 0:
+        return np.empty((0,), dtype=np.int64)
+    candidates = np.arange(0, min(int(curr_frame), int(n_context_frames)), dtype=np.int64)
+    if len(candidates) == 0:
+        return candidates
+    if not bool(_cfg_get(cfg, "anchor_diverse_selection", True)) or len(candidates) <= count:
+        return candidates[:count].astype(np.int64)
+
+    positions = np.linspace(0, len(candidates) - 1, num=count, dtype=np.int64)
+    selected = list(dict.fromkeys(candidates[positions].tolist()))
+    if len(selected) < count:
+        for idx in candidates:
+            if int(idx) not in selected:
+                selected.append(int(idx))
+            if len(selected) == count:
+                break
+    return np.asarray(selected[:count], dtype=np.int64)
+
+
+def _select_online_dynamic_indices(start_frame: int, count: int) -> np.ndarray:
+    if count <= 0 or start_frame <= 0:
+        return np.empty((0,), dtype=np.int64)
+    start = max(0, int(start_frame) - int(count))
+    return np.arange(start, int(start_frame), dtype=np.int64)[-count:]
+
+
 def _shift_memory_noise_levels(levels, shift: float, diffusion_model, bounds):
     if levels.numel() == 0 or float(shift) == 0.0 or not hasattr(diffusion_model, "snr"):
         return levels
@@ -1057,6 +1084,10 @@ class DeMemWMMinecraft(DiffusionForcingBase):
                 horizon = min(target_length - curr_frame, self.chunk_size) if self.chunk_size > 0 else target_length - curr_frame
                 target_slice = slice(curr_frame, curr_frame + horizon)
                 xs_pred[target_slice] = torch.randn_like(xs_pred[target_slice]).clamp(-self.clip_noise, self.clip_noise)
+                start_frame = max(0, curr_frame + horizon - self.n_tokens)
+                local_slice = slice(start_frame, curr_frame + horizon)
+                target_length_step = local_slice.stop - local_slice.start
+                query_offset = curr_frame - start_frame
 
                 stream_indices = {
                     key: torch.full((batch_size, int(stream_lengths[key])), -1, device=xs.device, dtype=torch.long)
@@ -1071,6 +1102,17 @@ class DeMemWMMinecraft(DiffusionForcingBase):
                         memory_selection_cfg,
                         split=namespace,
                     )
+                    anchor_count = int(stream_indices["anchor"].shape[1])
+                    dynamic_count = int(stream_indices["dynamic"].shape[1])
+                    selected["anchor"] = _select_online_anchor_indices(
+                        curr_frame,
+                        n_context_frames,
+                        anchor_count,
+                        memory_selection_cfg,
+                    )
+                    selected_masks["anchor"] = np.ones(len(selected["anchor"]), dtype=bool)
+                    selected["dynamic"] = _select_online_dynamic_indices(start_frame, dynamic_count)
+                    selected_masks["dynamic"] = np.ones(len(selected["dynamic"]), dtype=bool)
                     for key, indices in stream_indices.items():
                         count = indices.shape[1]
                         selected_key = np.asarray(selected.get(key, []), dtype=np.int64)[:count]
@@ -1082,44 +1124,47 @@ class DeMemWMMinecraft(DiffusionForcingBase):
                 source_poses = target_poses[:curr_frame]
                 source_frame_indices = target_frame_indices[:curr_frame]
                 frame_memory_masks = {
-                    "target": torch.ones((batch_size, horizon), device=xs.device, dtype=torch.bool)
+                    "target": torch.ones((batch_size, target_length_step), device=xs.device, dtype=torch.bool)
                     if target_mask is None
-                    else target_mask[:, target_slice],
+                    else target_mask[:, local_slice],
                     **stream_masks,
                 }
                 stream_latents = [_gather_online_memory_tensor(source_latents, stream_indices[key], stream_masks[key]) for key in _DEMEMWM_STREAM_KEYS]
                 stream_poses = [_gather_online_memory_tensor(source_poses, stream_indices[key], stream_masks[key]) for key in _DEMEMWM_STREAM_KEYS]
                 stream_frame_indices = [_gather_online_memory_tensor(source_frame_indices, stream_indices[key], stream_masks[key]) for key in _DEMEMWM_STREAM_KEYS]
-                target_conditions = target_tensors["action_conditions"][target_slice].to(device=xs.device)
+                target_conditions = target_tensors["action_conditions"][local_slice].to(device=xs.device)
                 memory_length = sum(mask.shape[1] for mask in stream_masks.values())
 
-                packed_latents = torch.cat([xs_pred[target_slice], *stream_latents], dim=0)
+                packed_latents = torch.cat([xs_pred[local_slice], *stream_latents], dim=0)
                 packed_conditions = torch.cat(
                     [target_conditions, target_conditions.new_zeros((memory_length, batch_size, target_conditions.shape[-1]))],
                     dim=0,
                 )
-                frame_memory_pose = torch.cat([target_poses[target_slice], *stream_poses], dim=0)
-                frame_indices = torch.cat([target_frame_indices[target_slice], *stream_frame_indices], dim=0)
+                frame_memory_pose = torch.cat([target_poses[local_slice], *stream_poses], dim=0)
+                frame_indices = torch.cat([target_frame_indices[local_slice], *stream_frame_indices], dim=0)
                 frame_memory_segments = {
-                    "target": int(horizon),
+                    "target": int(target_length_step),
                     **{key: int(stream_masks[key].shape[1]) for key in _DEMEMWM_STREAM_KEYS},
                 }
                 scheduling_matrix = self._generate_scheduling_matrix(horizon)
 
                 for m in range(scheduling_matrix.shape[0] - 1):
-                    from_target = torch.as_tensor(scheduling_matrix[m], device=xs.device, dtype=torch.long)[:, None].repeat(1, batch_size)
-                    to_target = torch.as_tensor(scheduling_matrix[m + 1], device=xs.device, dtype=torch.long)[:, None].repeat(1, batch_size)
+                    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))
+                    from_target = torch.cat([context_levels, from_query], dim=0)
+                    to_target = torch.cat([context_levels, to_query], dim=0)
                     memory_noise_levels = _memory_noise_levels_for_streams(
                         getattr(self, "cfg", None),
                         self.diffusion_model,
-                        from_target,
+                        from_query,
                         frame_memory_segments,
                         mode=namespace,
                     )
                     memory_noise_levels = torch.cat([memory_noise_levels[key] for key in _DEMEMWM_STREAM_KEYS], dim=0)
                     routed_frame_memory_masks = _apply_memory_route_masks(
                         frame_memory_masks,
-                        from_target,
+                        from_query,
                         getattr(self, "cfg", None),
                         self.diffusion_model,
                         mode=namespace,
@@ -1139,9 +1184,9 @@ class DeMemWMMinecraft(DiffusionForcingBase):
                         frame_memory_pose=frame_memory_pose,
                         image_hw=image_hw,
                     )
-                    packed_latents[:horizon] = sampled_target[:horizon]
+                    packed_latents[:target_length_step] = sampled_target[:target_length_step]
 
-                xs_pred[target_slice] = packed_latents[:horizon]
+                xs_pred[local_slice] = packed_latents[:target_length_step]
                 curr_frame += horizon
 
             eval_start = n_context_frames