Fix DeMemWM training anchor prefix window

datasets/video/minecraft_video_dememwm_latent_dataset.py

@@ -198,7 +198,8 @@ class MinecraftVideoDeMemWMLatentDataset(torch.utils.data.Dataset):
 
         # === 2. Select target and memory indices ===
         target_positions = frame_idx + self._target_offsets
-        anchor_start = self.initial_frame_offset
+        anchor_start = max(self.initial_frame_offset, frame_idx - self.context_length)
+        anchor_stop = frame_idx
         memory_indices, memory_masks = select_memory_indices(
             poses_pool,
             target_positions,
@@ -208,7 +209,7 @@ class MinecraftVideoDeMemWMLatentDataset(torch.utils.data.Dataset):
             latents=latents,
             actions=actions,
             anchor_candidate_start=anchor_start,
-            anchor_candidate_stop=self.initial_frame_offset + self.context_length,
+            anchor_candidate_stop=anchor_stop,
         )
         ordered_memory = [memory_indices[key] for key in SEGMENT_KEYS]
         source_positions = np.concatenate([target_positions, *ordered_memory]).astype(np.int64)