Add DeMemWM memory causality config

.exp_artifact/dememwm_future_memory_selection_plan.md

@@ -0,0 +1,489 @@
+# DeMemWM Future Memory Selection Plan
+
+Repo: `/share_1/users/bonan_ding/WorldMem`
+
+Branch: `dememwm`
+
+Purpose: allow the DeMemWM dataset-side memory selector to choose future frames
+during training when explicitly configured, while keeping the packed DeMemWM
+temporal-attention memory streams diagonal/self-only.
+
+## Required Orchestration Flow
+
+Use this plan with a new session prompt such as:
+
+```text
+/goal Complete /share_1/users/bonan_ding/WorldMem/.exp_artifact/dememwm_future_memory_selection_plan.md one substep at a time. For every substep, use a fresh implementation subagent, a fresh review subagent, and if needed fresh fix subagents until review is clean. Validate, inspect diff, commit the substep, and update this plan before moving to the next substep.
+```
+
+The main session is the orchestra leader. It owns sequencing, diff inspection,
+validation, commits, and plan-status updates. Subagents do bounded implementation,
+review, or fix work only for the current substep.
+
+For every substep below:
+
+1. Start a fresh implementation subagent.
+2. The implementation subagent must read:
+   - `AGENTS.md`
+   - `.codex/AGENT.md`
+   - `.codex/memory.md`
+   - `/share_1/users/bonan_ding/.codex/skills/codex-ml-cv-clean-implementation/SKILL.md`
+3. The implementation subagent must use:
+   - `[$codex-ml-cv-clean-implementation](/share_1/users/bonan_ding/.codex/skills/codex-ml-cv-clean-implementation/SKILL.md)`
+4. The implementation subagent must implement only the current substep.
+5. After implementation, the main session inspects the diff before review.
+6. Start a fresh review subagent.
+7. The review subagent must read the same files and use the same skill.
+8. The review subagent must check:
+   - correctness
+   - bugs or regressions
+   - improper or unnecessary changes
+   - code cleanliness
+   - adherence to this plan
+   - adherence to repository conventions
+   - targeted test adequacy
+9. If review finds issues, start a fresh fix subagent.
+10. The fix subagent must read the same files, use the same skill, and fix only
+    the review findings.
+11. Repeat review/fix/re-review until the review subagent reports no remaining
+    issues.
+12. Run the validation commands listed for the substep.
+13. Run `git diff --check`.
+14. Inspect `git diff`.
+15. Commit only the completed substep with a clear commit message.
+16. Mark the substep complete in this plan file and include that plan-status
+    update in the same substep commit if possible.
+17. Do not proceed to the next substep until the current one is reviewed clean,
+    validated, committed, and marked complete.
+
+Global rules:
+
+- Do not skip substeps.
+- Do not combine unrelated substeps into one commit.
+- Keep changes surgical.
+- Preserve unrelated user changes.
+- Do not train or submit Slurm jobs.
+- Do not modify `algorithms/worldmem`; use it only as a reference.
+- Do not couple `dataset.memory_selection.causal` to `algorithm.causal`.
+- Do not change DeMemWM local target temporal-attention causality in this plan.
+- Keep anchor selection semantics unchanged: same anchor candidate window and
+  same `_select_anchor` scoring/diversity logic.
+- Allow cross-stream duplicate memory frames: `anchor`, `dynamic`, and `revisit`
+  may select the same raw frame, including frames in the target/local-context
+  window when training is configured non-causal.
+- If validation fails, fix through the same fresh fix-subagent and
+  review-subagent loop, then rerun validation.
+
+Status legend:
+
+- `[ ]` not started
+- `[~]` in progress
+- `[x]` complete and committed
+
+## Current Code Facts
+
+- DeMemWM training samples are built by
+  `datasets/video/minecraft_video_dememwm_latent_dataset.py`, which calls
+  `select_memory_indices()` in `datasets/video/memory_selection.py`.
+- `anchor` currently uses a clip-local context prefix through
+  `anchor_candidate_start` and `anchor_candidate_stop`. This must stay
+  unchanged.
+- `_select_revisit()` currently builds a past-only pool using
+  `local_context_exclusion_frames`; this blocks future frames and local-window
+  frames.
+- `select_memory_indices()` currently enforces some cross-stream uniqueness:
+  recent `dynamic` can be removed from anchor candidates, and `revisit` can be
+  told to exclude anchor/dynamic frames.
+- DeMemWM validation/test rollout also calls `_select_revisit()` from
+  `algorithms/dememwm/df_video.py`; if duplicate stream semantics are changed
+  in the dataset selector, the online path should be checked for the same
+  cross-stream exclusion issue.
+- In `algorithms/dememwm/models/attention.py`,
+  `_frame_memory_attn_bias()` currently makes `anchor` and `revisit`
+  diagonal/self-only, but gives `dynamic` full intra-stream attention.
+- WorldMem memory frames do not need memory-stream self-attention: memory is
+  consumed as keys/values by target queries. In DeMemWM, this corresponds to
+  `FrameMemoryReferenceAttention`, not temporal mixing among memory frames.
+
+## Non-Goals
+
+- Do not change the frame-memory segment order:
+
+  ```text
+  [target][anchor][dynamic][revisit]
+  ```
+
+- Do not change the DeMemWM reference-attention architecture.
+- Do not change `algorithm.causal` behavior.
+- Do not add a new selector framework or backend.
+- Do not use `local_context_exclusion_frames` to block training memory
+  candidates in the new non-causal path.
+- Do not add compatibility aliases such as `casual`; use `causal`.
+- Do not run training or submit cluster jobs.
+
+## Substep 1: Add Dataset-Side Memory Causal Config And Candidate Helper
+
+Status: `[x]`
+
+Goal:
+
+Add the explicit dataset-side switch that controls whether training memory
+selection may use future frames.
+
+Primary files:
+
+- `configurations/dataset/video_minecraft_dememwm_latent.yaml`
+- `datasets/video/memory_selection.py`
+- `tests/test_dememwm_latent_dataset.py`
+
+Required behavior:
+
+- Add:
+
+  ```yaml
+  memory_selection:
+    causal: true
+  ```
+
+- The runtime override for future-memory training is:
+
+  ```bash
+  dataset.memory_selection.causal=false
+  ```
+
+- Do not set `memory_selection.causal: ${algorithm.causal}`.
+- Add a small helper in `memory_selection.py` for memory candidate frames:
+
+  ```python
+  def _memory_candidate_frames(num_frames, target_positions, cfg, split, min_candidate_frame=0):
+      target_start = int(target_positions[0])
+      causal = split != "training" or bool(cfg_get(cfg, "causal", True))
+      stop = target_start if causal else num_frames
+      return np.arange(min_candidate_frame, stop, dtype=np.int64)
+  ```
+
+- This helper intentionally does not apply `local_context_exclusion_frames`.
+  For training with `causal=false`, the candidate pool may include future
+  frames, target frames, and local-window frames, matching the desired
+  WorldMem-style training behavior.
+- For validation/test, the helper remains causal regardless of the config value
+  because `split != "training"` forces `causal=True`.
+
+Implementation constraints:
+
+- Keep the helper local to `memory_selection.py`; do not add a new module.
+- Do not yet rewrite all selectors in this substep unless it stays cleaner to
+  include one direct use in `_select_revisit()`.
+- Keep tests tiny and deterministic.
+
+Validation for this substep:
+
+```bash
+python -m py_compile datasets/video/memory_selection.py tests/test_dememwm_latent_dataset.py
+python -B -m unittest tests.test_dememwm_latent_dataset
+git diff --check
+```
+
+Review must specifically confirm:
+
+- The config key is dataset-side only.
+- `algorithm.causal` remains unrelated.
+- Training with `causal=false` can produce a candidate pool reaching the end of
+  the video.
+- Validation/test remain no-future.
+- No local-context exclusion was introduced in the new helper.
+
+Suggested commit message:
+
+```text
+Add DeMemWM memory causality config
+```
+
+## Substep 2: Enable Future And Local-Window Revisit Selection
+
+Status: `[ ]`
+
+Goal:
+
+Use the new candidate helper for `revisit` selection so training can choose
+future frames when `dataset.memory_selection.causal=false`.
+
+Primary files:
+
+- `datasets/video/memory_selection.py`
+- `algorithms/dememwm/df_video.py`
+- `tests/test_dememwm_latent_dataset.py`
+
+Required behavior:
+
+- Replace `_select_revisit()` candidate construction with the helper from
+  Substep 1.
+- Do not apply `local_context_exclusion_frames` in `_select_revisit()` candidate
+  construction for this path.
+- Preserve the existing scoring:
+  - training uses pose-similarity selection;
+  - validation/test uses point-union/FOV selection.
+- Preserve `_select_revisit()`'s `excluded` argument for any direct diagnostic
+  caller, but top-level DeMemWM selection should stop using cross-stream
+  exclusions to enforce uniqueness.
+- In `select_memory_indices()`, pass an empty exclusion set to `_select_revisit()`
+  so `revisit` may reuse frames selected by `anchor` or `dynamic`.
+- In `algorithms/dememwm/df_video.py`, check the online validation/test call to
+  `_select_revisit()`. Remove only the cross-stream exclusion that prevents
+  `revisit` from reusing online `anchor`/`dynamic` frames. Do not change online
+  causal history limits.
+
+Implementation constraints:
+
+- Keep anchor candidate window and `_select_anchor()` logic unchanged.
+- Do not change dynamic selection in this substep except where needed to stop
+  blocking revisit duplicates.
+- Do not modify WorldMem.
+
+Required tests:
+
+- Training with default `causal=true` does not select future revisit frames.
+- Training with `causal=false` can select a future pose-matching revisit frame.
+- Training with `causal=false` can select a target/local-window frame when it is
+  the best pose match.
+- Validation/test with `causal=false` still does not select future revisit
+  frames.
+- Revisit can duplicate an anchor or dynamic frame when scoring selects it.
+
+Validation for this substep:
+
+```bash
+python -m py_compile datasets/video/memory_selection.py algorithms/dememwm/df_video.py tests/test_dememwm_latent_dataset.py
+python -B -m unittest tests.test_dememwm_latent_dataset
+git diff --check
+```
+
+Review must specifically confirm:
+
+- Future memory is gated by `dataset.memory_selection.causal=false`.
+- Validation/test are still no-future.
+- Local-window candidates are not filtered out.
+- Revisit duplicates across streams are allowed by the top-level DeMemWM path.
+- Anchor selection itself was not changed.
+
+Suggested commit message:
+
+```text
+Allow noncausal DeMemWM revisit selection
+```
+
+## Substep 3: Enable Future Event-Triggered Dynamic Selection
+
+Status: `[ ]`
+
+Goal:
+
+Make event-triggered `dynamic` memory use the same training causal/non-causal
+candidate pool as `revisit`, while still keeping `recent` dynamic semantics
+unchanged.
+
+Primary files:
+
+- `datasets/video/memory_selection.py`
+- `tests/test_dememwm_latent_dataset.py`
+
+Required behavior:
+
+- For `dynamic.selection_policy=event_triggered`, filter `dynamic_stream` using
+  `_memory_candidate_frames()`.
+- For training with `causal=false`, eligible dynamic stream frames may include
+  future, target, and local-window frames.
+- For validation/test, eligible dynamic stream frames remain no-future.
+- Do not pass `anchor` or `revisit` as exclusions when selecting dynamic memory;
+  dynamic may reuse frames from either stream.
+- Use `revisit` as reference frames when available; otherwise use
+  `target_positions` so non-causal selection stays query-local instead of
+  drifting to unrelated late-video events.
+- Leave `dynamic.selection_policy=recent` past-only. "Recent" means immediate
+  history; future-frame memory should come from `revisit` and event-triggered
+  `dynamic`.
+
+Implementation constraints:
+
+- Do not change `_build_dynamic_stream()` event scoring.
+- Do not change anchor selection.
+- Do not add a new dynamic policy.
+
+Required tests:
+
+- Event-triggered dynamic with `causal=true` stays before the target.
+- Event-triggered dynamic with `causal=false` can select future stream frames.
+- Event-triggered dynamic can duplicate a revisit or anchor frame.
+- Recent dynamic behavior remains unchanged by this substep.
+
+Validation for this substep:
+
+```bash
+python -m py_compile datasets/video/memory_selection.py tests/test_dememwm_latent_dataset.py
+python -B -m unittest tests.test_dememwm_latent_dataset
+git diff --check
+```
+
+Review must specifically confirm:
+
+- The change applies only to event-triggered dynamic selection.
+- Recent dynamic remains past-only.
+- Cross-stream duplicate frames are allowed.
+- No local-context exclusion was reintroduced.
+
+Suggested commit message:
+
+```text
+Allow noncausal DeMemWM dynamic events
+```
+
+## Substep 4: Make All Packed Memory Streams Diagonal/Self-Only
+
+Status: `[ ]`
+
+Goal:
+
+Make `anchor`, `dynamic`, and `revisit` all use diagonal/self-only temporal
+attention rows in the packed DeMemWM `TemporalAxialAttention` path.
+
+Primary files:
+
+- `algorithms/dememwm/models/attention.py`
+- `tests/test_dememwm_temporal_attention.py`
+
+Required behavior:
+
+- In `_frame_memory_attn_bias()`, remove the special full intra-stream attention
+  for `dynamic`.
+- Each memory stream should only allow each memory frame to attend itself in the
+  temporal self-attention block:
+
+  ```python
+  idx = torch.arange(cursor, cursor + length, device=device)
+  allow[:, idx, idx] = True
+  ```
+
+- Keep target-target causal lower-triangular attention unchanged.
+- Keep mask handling unchanged.
+- Keep the final diagonal restore for padded/all-invalid rows unchanged.
+- Do not change `FrameMemoryReferenceAttention`; target frames should still
+  consume stream memory as keys/values there.
+
+Implementation constraints:
+
+- This is an attention-mask cleanup only.
+- Do not change model dimensions, segment order, or reference-attention logic.
+- Do not modify WorldMem.
+
+Required tests:
+
+- Packed target segment remains causal.
+- Anchor stream is diagonal/self-only.
+- Dynamic stream is diagonal/self-only.
+- Revisit stream is diagonal/self-only.
+- Padded invalid memory rows remain finite through the diagonal restore.
+
+Validation for this substep:
+
+```bash
+python -m py_compile algorithms/dememwm/models/attention.py tests/test_dememwm_temporal_attention.py
+python -B -m unittest tests.test_dememwm_temporal_attention
+git diff --check
+```
+
+Review must specifically confirm:
+
+- `dynamic` no longer has full intra-stream temporal self-attention.
+- Target local temporal causal masking was not changed.
+- Memory consumption through reference attention remains untouched.
+
+Suggested commit message:
+
+```text
+Make DeMemWM memory streams self-only
+```
+
+## Substep 5: Final Integration Validation
+
+Status: `[ ]`
+
+Goal:
+
+Run the complete cheap validation set for this plan after all code substeps are
+committed.
+
+Primary files:
+
+- No code changes expected unless validation or review finds a bug.
+- This plan file should be updated to mark the final substep complete.
+
+Validation:
+
+```bash
+python -m py_compile \
+  datasets/video/memory_selection.py \
+  algorithms/dememwm/df_video.py \
+  algorithms/dememwm/models/attention.py \
+  tests/test_dememwm_latent_dataset.py \
+  tests/test_dememwm_temporal_attention.py
+
+python -B -m unittest \
+  tests.test_dememwm_latent_dataset \
+  tests.test_dememwm_temporal_attention
+
+git diff --check
+git status --short
+```
+
+Review must specifically confirm:
+
+- The worktree contains only expected committed changes plus this plan status
+  update.
+- There are no unrelated refactors.
+- The final behavior can be enabled with:
+
+  ```bash
+  dataset.memory_selection.causal=false
+  ```
+
+- No training or Slurm job was run.
+
+Suggested commit message:
+
+```text
+Validate DeMemWM future memory selection
+```
+
+## New-Session Goal Prompt
+
+Use this prompt to start the implementation session:
+
+```text
+Use $codex-ml-cv-clean-implementation.
+Repo: /share_1/users/bonan_ding/WorldMem
+Goal: Complete .exp_artifact/dememwm_future_memory_selection_plan.md step by step.
+
+Act as the main orchestra leader. For each substep in the plan:
+1. Use a fresh implementation subagent for only that substep.
+2. Inspect the diff yourself.
+3. Use a fresh review subagent to review correctness, bugs, design, cleanliness, and tests.
+4. If review finds any issue, use a fresh fix subagent to fix only those findings.
+5. Repeat review/fix with fresh subagents until review is clean.
+6. Run the substep validation commands and git diff --check.
+7. Inspect git diff.
+8. Commit only that substep with the suggested or similarly clear commit message.
+9. Mark the substep complete in the plan file and include the status update in the same commit if possible.
+10. Move to the next substep only after the current substep is reviewed clean, validated, committed, and marked complete.
+
+Hard constraints:
+- Do not train or submit Slurm jobs.
+- Do not modify algorithms/worldmem.
+- Do not couple dataset.memory_selection.causal to algorithm.causal.
+- Keep anchor selection semantics unchanged.
+- Allow anchor, dynamic, and revisit to reuse the same raw frame.
+- Do not filter training memory candidates with local_context_exclusion_frames.
+- Make all packed DeMemWM memory streams diagonal/self-only in TemporalAxialAttention.
+- Preserve unrelated user changes.
+
+Final report must include changed files, commits made, validation commands and results, and remaining risks.
+```

configurations/dataset/video_minecraft_dememwm_latent.yaml

@@ -19,6 +19,7 @@ shuffle_clips: true
 
 memory_selection:
   enabled: true
+  causal: true
   max_anchor_frames: 2
   max_dynamic_frames: 4
   max_revisit_frames: 2

datasets/video/memory_selection.py

@@ -25,6 +25,13 @@ def cfg_get(cfg, key: str, default=None):
     return getattr(cfg, key, default)
 
 
+def _memory_candidate_frames(num_frames, target_positions, cfg, split, min_candidate_frame=0):
+    target_start = int(target_positions[0])
+    causal = split != "training" or bool(cfg_get(cfg, "causal", True))
+    stop = target_start if causal else num_frames
+    return np.arange(min_candidate_frame, stop, dtype=np.int64)
+
+
 def _as_pose_array(poses) -> np.ndarray:
     poses = np.asarray(poses, dtype=np.float32)
     if poses.ndim != 2 or poses.shape[1] < 5:

tests/test_dememwm_latent_dataset.py

@@ -9,7 +9,7 @@ import torch
 from torch import nn
 from omegaconf import OmegaConf
 
-from datasets.video.memory_selection import select_memory_indices
+from datasets.video.memory_selection import _memory_candidate_frames, select_memory_indices
 from datasets.video.minecraft_video_dememwm_latent_dataset import MinecraftVideoDeMemWMLatentDataset
 
 
@@ -33,6 +33,7 @@ def _dynamic_cfg(selection_policy="recent", **overrides):
 def _selection_cfg(**overrides):
     cfg = {
         "enabled": True,
+        "causal": True,
         "max_anchor_frames": 2,
         "max_dynamic_frames": 2,
         "max_revisit_frames": 2,
@@ -117,6 +118,36 @@ def _write_vae_feature_clip(root, feature_root, split="training", subdir="scene"
 
 
 class MemorySelectionTests(unittest.TestCase):
+    def test_memory_candidate_frames_default_training_is_causal_without_local_exclusion(self):
+        cfg = _selection_cfg(local_context_exclusion_frames=99)
+
+        candidates = _memory_candidate_frames(10, np.array([5]), cfg, "training", min_candidate_frame=1)
+
+        self.assertEqual(candidates.tolist(), [1, 2, 3, 4])
+
+    def test_memory_candidate_frames_training_can_be_noncausal(self):
+        cfg = _selection_cfg(causal=False, local_context_exclusion_frames=99)
+
+        candidates = _memory_candidate_frames(8, np.array([3]), cfg, "training", min_candidate_frame=1)
+
+        self.assertEqual(candidates.tolist(), [1, 2, 3, 4, 5, 6, 7])
+
+    def test_memory_candidate_frames_validation_and_test_stay_causal(self):
+        cfg = _selection_cfg(causal=False)
+
+        for split in ("validation", "test"):
+            with self.subTest(split=split):
+                candidates = _memory_candidate_frames(8, np.array([3]), cfg, split, min_candidate_frame=1)
+
+                self.assertEqual(candidates.tolist(), [1, 2])
+
+    def test_dememwm_latent_config_has_dataset_side_memory_causal_default(self):
+        config_path = Path(__file__).resolve().parents[1] / "configurations" / "dataset" / "video_minecraft_dememwm_latent.yaml"
+
+        cfg = OmegaConf.load(config_path)
+
+        self.assertIs(cfg.memory_selection.causal, True)
+
     def test_revisit_uses_sampled_fov_selection(self):
         poses = np.array(
             [
@@ -433,8 +464,9 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
             _write_vae_feature_clip(root, root / "vae_features", stem="sample", poses=poses)
             dataset = MinecraftVideoDeMemWMLatentDataset(_dataset_cfg(root), split="training")
 
-            with self.assertRaisesRegex(ValueError, "Pose height variation"):
+            with self.assertRaisesRegex(RuntimeError, "Pose height variation") as captured:
                 dataset[0]
+            self.assertIsInstance(captured.exception.__cause__, ValueError)
 
     def test_dataset_converts_raw_actions_to_original_25d_contract(self):
         with tempfile.TemporaryDirectory() as tmp:
@@ -857,6 +889,7 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
                 self.metric_updates = []
                 self.logger = type("Logger", (), {"experiment": object()})()
                 self.log_video = True
+                self.log_memory_selection_sheet = False
                 self.save_local = False
                 self.local_save_dir = None
 
@@ -877,6 +910,12 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
             def _update_metric_accumulators(self, xs_pred, xs_gt):
                 self.metric_updates.append((xs_pred.detach().clone(), xs_gt.detach().clone()))
 
+            def _update_per_frame_metric_accumulators(self, xs_pred, xs_gt, mask, eval_start):
+                pass
+
+            def _update_revisit_metric_accumulators(self, xs_pred, xs_gt, poses, frame_indices, mask, eval_start, batch_idx, namespace):
+                pass
+
             def log(self, name, value):
                 self.logged.append((name, value))
 
@@ -980,6 +1019,7 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
                 self.logged = []
                 self.logger = None
                 self.log_video = False
+                self.log_memory_selection_sheet = False
                 self.save_local = False
                 self.local_save_dir = None
 
@@ -995,6 +1035,12 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
             def _update_metric_accumulators(self, xs_pred, xs_gt):
                 pass
 
+            def _update_per_frame_metric_accumulators(self, xs_pred, xs_gt, mask, eval_start):
+                pass
+
+            def _update_revisit_metric_accumulators(self, xs_pred, xs_gt, poses, frame_indices, mask, eval_start, batch_idx, namespace):
+                pass
+
             def log(self, name, value):
                 self.logged.append((name, value))
 
@@ -1025,9 +1071,9 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
 
         self.assertEqual(len(harness.diffusion_model.calls), 1)
         call = harness.diffusion_model.calls[0]
-        self.assertEqual(call["kwargs"]["frame_idx"][:, 0].tolist(), [13, 14, 0])
-        self.assertEqual(call["kwargs"]["frame_memory_masks"]["dynamic"].tolist(), [[False]])
-        self.assertEqual(call["x"][:, 0, 0, 0, 0].tolist(), [0.0, 0.0, 0.0])
+        self.assertEqual(call["kwargs"]["frame_idx"][:, 0].tolist(), [13, 14, 10])
+        self.assertEqual(call["kwargs"]["frame_memory_masks"]["dynamic"].tolist(), [[True]])
+        self.assertEqual(call["x"][:, 0, 0, 0, 0].tolist(), [0.0, 0.0, 1.0])
 
     def test_interactive_generation_fails_until_packed_memory_path_exists(self):
         from algorithms.dememwm.df_video import DeMemWMMinecraft
@@ -1096,7 +1142,7 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
 
         self.assertEqual(dataset.context_length, 4)
         self.assertEqual(len(dataset), 112 - (100 + 4 + 3) + 1)
-        self.assertEqual(sample["frame_indices"].tolist(), [106, 107, 108, 100, 103, 104, 105])
+        self.assertEqual(sample["frame_indices"].tolist(), [106, 107, 108, 102, 103, 104, 105])
         self.assertTrue(sample["memory_masks"]["anchor"].all().item())
         self.assertTrue(sample["memory_masks"]["dynamic"].all().item())
         self.assertEqual(tuple(sample["memory_masks"]["revisit"].shape), (0,))
@@ -1131,7 +1177,7 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
                 stem="sample",
                 num_frames=num_frames,
                 actions=np.zeros((num_frames, 25), dtype=np.float32),
-                latents=_event_latents(num_frames, event_frame=103),
+                latents=_event_latents(num_frames, event_frame=102),
             )
             dataset = MinecraftVideoDeMemWMLatentDataset(
                 _dataset_cfg(
@@ -1148,8 +1194,8 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
             )
             sample = dataset[0]
 
-        self.assertEqual(sample["frame_indices"].tolist(), [106, 107, 108, 104, -1])
-        self.assertEqual(sample["memory_masks"]["dynamic"].tolist(), [True, False])
+        self.assertEqual(sample["frame_indices"].tolist(), [106, 107, 108, 100, 103])
+        self.assertEqual(sample["memory_masks"]["dynamic"].tolist(), [True, True])
 
     def test_dataset_returns_target_anchor_dynamic_revisit_contract(self):
         with tempfile.TemporaryDirectory() as tmp: