Fix DeMemWM nested split filtering

.exp_artifact/dememwm_remaining_fix_plan.md

@@ -0,0 +1,368 @@
+# DeMemWM Remaining Fix Plan
+
+Repo: `/share_1/users/bonan_ding/WorldMem`
+
+Branch: `dememwm`
+
+Purpose: fix the remaining DeMemWM bugs one at a time. Each substep must be implemented, reviewed, fixed if needed, validated, committed, and marked complete before moving to the next substep.
+
+## 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_remaining_fix_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.
+```
+
+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, inspect 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 add config flags, new abstractions, compatibility fallbacks, or broad refactors unless the substep explicitly requires them.
+- 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
+
+## Substep 1: Nested Validation/Test Split Filter Must Not Fall Back To Unfiltered Files
+
+Status: `[x]`
+
+Bug:
+
+The DeMemWM latent dataset filters validation/test files by split type, such as `w_updown` or `wo_updown`. The direct-file path filtering was fixed, but when the direct filtered candidates are empty, the nested fallback can still add nested `.npz` files without applying the same split filter. That silently loads the wrong validation/test subset.
+
+Primary files:
+
+- `datasets/video/minecraft_video_dememwm_latent_dataset.py`
+- `tests/test_dememwm_latent_dataset.py`
+
+Known code area to inspect:
+
+- `get_data_paths()` in `datasets/video/minecraft_video_dememwm_latent_dataset.py`
+- Validation/test filtering around the existing `w_updown` and `wo_updown` logic
+- Nested `.npz` fallback collection after the direct candidates are empty
+
+Required behavior:
+
+- Apply the same validation/test split filter to nested fallback files.
+- Never re-add an unfiltered nested file list after validation/test filtering removed all candidates.
+- If validation/test filtering leaves no files, emit a warning and then raise `FileNotFoundError`.
+- The warning and error must make the reason clear: the selected split is empty because the split filter, such as `w_updown` or `wo_updown`, removed all candidates.
+- Do not add a compatibility fallback that silently loads mismatched files.
+
+Implementation constraints:
+
+- Keep the dataset change local to path collection/filtering.
+- Do not change sample loading semantics for valid datasets.
+- Use standard `warnings.warn(...)` unless the local file already uses another warning pattern.
+- Keep tests small and filesystem-local with temporary directories.
+
+Validation for this substep:
+
+```bash
+python -m py_compile datasets/video/minecraft_video_dememwm_latent_dataset.py tests/test_dememwm_latent_dataset.py
+python -B -m unittest tests.test_dememwm_latent_dataset
+git diff --check
+```
+
+Review must specifically confirm:
+
+- Mismatched nested validation/test files are rejected.
+- A warning is emitted before `FileNotFoundError`.
+- Positive filtering behavior still works.
+- Train split behavior is not accidentally changed.
+
+Suggested commit message:
+
+```text
+Fix DeMemWM nested split filtering
+```
+
+## Substep 2: Remove Stale Raw-Video Helper Path From DeMemWM Algorithm
+
+Status: `[ ]`
+
+Bug:
+
+The public DeMemWM train, validation, test, and interactive paths now use typed latent/memory dictionary batches, but stale legacy raw-video helper code remains in `algorithms/dememwm/df_video.py`. This dead path makes the implementation harder to audit and can accidentally preserve old tuple/raw-video assumptions.
+
+Primary file:
+
+- `algorithms/dememwm/df_video.py`
+
+Known stale symbols to inspect:
+
+- `_preprocess_batch`
+- `_generate_condition_indices`
+- `_prepare_conditions`
+- `_prepare_noise_levels`
+
+Required behavior:
+
+- Make `_preprocess_batch()` strict for the current mapping/dictionary batch format.
+- Remove unused legacy raw-video helper methods and dead code.
+- Before deleting each helper, use `rg` to confirm it has no remaining references.
+- Keep active encode/decode helpers only if they are still used by current DeMemWM code.
+- Do not rewrite the training loop.
+- Do not change current public train/validation/test/interactive behavior except removing stale unsupported paths.
+
+Implementation constraints:
+
+- Avoid broad refactors.
+- Keep the diff centered on dead-code removal and the strict current-batch boundary.
+- If a helper is still used, do not delete it. Instead, report why it must remain.
+
+Validation for this substep:
+
+```bash
+rg -n "_generate_condition_indices|_prepare_conditions|_prepare_noise_levels" algorithms/dememwm tests
+python -m py_compile algorithms/dememwm/df_video.py
+python -B -m unittest tests.test_checkpoint_loading tests.test_resume_checkpoint_logic
+git diff --check
+```
+
+Review must specifically confirm:
+
+- No stale raw-video/tuple helper path remains in active DeMemWM preprocessing.
+- No live references were broken.
+- No unrelated training/eval behavior was changed.
+- The code is cleaner, not merely rearranged.
+
+Suggested commit message:
+
+```text
+Remove stale DeMemWM raw-video helpers
+```
+
+## Substep 3: Reject Stale Experiment Flags For DeMemWM
+
+Status: `[ ]`
+
+Bug:
+
+The generic experiment layer still accepts old experiment flags such as `zero_init_gate` and `only_tune_memory`. For `algorithm._name == "dememwm_base"`, these flags can override or corrupt DeMemWM trainability assumptions. The reference DeMemWM path rejects these stale flags instead of applying them.
+
+Primary files:
+
+- `experiments/exp_base.py`
+- The smallest relevant existing test file, likely one of:
+  - `tests/test_checkpoint_loading.py`
+  - `tests/test_resume_checkpoint_logic.py`
+
+Known code area to inspect:
+
+- Experiment setup logic in `experiments/exp_base.py`
+- Existing handling for `zero_init_gate`
+- Existing handling for `only_tune_memory`
+- How tests instantiate experiment/config objects
+
+Required behavior:
+
+- Add an early, surgical validation for DeMemWM experiments.
+- If `algorithm._name == "dememwm_base"` and `zero_init_gate` is enabled, raise a clear `ValueError`.
+- If `algorithm._name == "dememwm_base"` and `only_tune_memory` is enabled, raise a clear `ValueError`.
+- Do not affect non-DeMemWM algorithms.
+- Do not refactor experiment setup.
+
+Implementation constraints:
+
+- Keep the validation near existing flag handling or experiment initialization, whichever is cleanest in the current file.
+- Do not add new config keys.
+- Do not change default trainability behavior beyond rejecting these stale flags for DeMemWM.
+
+Validation for this substep:
+
+```bash
+python -m py_compile experiments/exp_base.py tests/test_checkpoint_loading.py tests/test_resume_checkpoint_logic.py
+python -B -m unittest tests.test_checkpoint_loading tests.test_resume_checkpoint_logic
+git diff --check
+```
+
+Review must specifically confirm:
+
+- DeMemWM rejects both stale flags with clear errors.
+- Non-DeMemWM behavior is unchanged.
+- The check is not duplicated in multiple messy locations.
+- The patch does not introduce a broad config-validation layer.
+
+Suggested commit message:
+
+```text
+Reject stale DeMemWM experiment flags
+```
+
+## Substep 4: Add Minimal DeMemWM DiT Boundary Checks
+
+Status: `[ ]`
+
+Bug:
+
+The DeMemWM DiT frame/memory splitting and stream-mask helpers are weaker than the reference implementation. Bad segment lengths or bad mask shapes can produce confusing downstream tensor failures. The fix should add only targeted checks at the model boundary, not broad internal guard clutter.
+
+Primary files:
+
+- `algorithms/dememwm/models/dit.py`
+- `tests/test_dememwm_temporal_attention.py`
+
+Known code area to inspect:
+
+- `_split_frame_memory()`
+- `_frame_memory_stream_mask()`
+
+Required behavior:
+
+- In `_split_frame_memory()`, validate that frame/memory segment lengths are nonnegative and sum to `x.shape[1]`.
+- In `_frame_memory_stream_mask()`, validate that the stream mask shape matches the expected stream length.
+- Error messages should be concise and useful for debugging a bad config or caller.
+- Keep checks surgical. Do not add repetitive guards for every tensor property.
+
+Implementation constraints:
+
+- Follow the `codex-ml-cv-clean-implementation` guard policy.
+- Add checks only where they catch this known realistic failure.
+- Prefer clear local code over helper abstractions.
+- Keep CPU-safe tests.
+
+Validation for this substep:
+
+```bash
+python -m py_compile algorithms/dememwm/models/dit.py tests/test_dememwm_temporal_attention.py
+python -B -m unittest tests.test_dememwm_temporal_attention
+git diff --check
+```
+
+Review must specifically confirm:
+
+- Bad segment sums fail early with a clear error.
+- Bad stream-mask shapes fail early with a clear error.
+- The checks do not sprawl into broad defensive programming.
+- Existing temporal attention behavior is preserved.
+
+Suggested commit message:
+
+```text
+Validate DeMemWM DiT stream boundaries
+```
+
+## Substep 5: Compute Plucker Geometry In fp32 Under Mixed Precision
+
+Status: `[ ]`
+
+Bug:
+
+The DeMemWM DiT Plucker/ray geometry path casts camera pose/ray math to the model activation dtype. With mixed precision training, for example `16-mixed`, this can run geometry in fp16/bf16 and lose precision. Geometry should be computed in fp32, then cast back to the activation dtype before use.
+
+Primary files:
+
+- `algorithms/dememwm/models/dit.py`
+- `tests/test_dememwm_temporal_attention.py` or another existing focused model test file
+
+Known code area to inspect:
+
+- Camera pose handling in `algorithms/dememwm/models/dit.py`
+- Ray or Plucker feature construction
+- Any `dtype=x.dtype` or activation-dtype casting around pose/ray geometry
+
+Required behavior:
+
+- Compute camera/ray/Plucker geometry in fp32.
+- Disable autocast only around the geometry math block if needed.
+- Cast the final geometry tensor back to the activation/model dtype before concatenation or use.
+- Keep the change local to the geometry path.
+- Do not add a config flag.
+
+Implementation constraints:
+
+- Preserve device placement.
+- Do not force CPU.
+- Do not disrupt AMP for unrelated model operations.
+- Avoid broad dtype guards; the important invariant is fp32 geometry math and final cast back to the caller's dtype.
+
+Validation for this substep:
+
+```bash
+python -m py_compile algorithms/dememwm/models/dit.py tests/test_dememwm_temporal_attention.py
+python -B -m unittest tests.test_dememwm_temporal_attention
+git diff --check
+```
+
+Review must specifically confirm:
+
+- Geometry math runs in fp32 under autocast/mixed precision.
+- Final geometry output returns to the expected activation dtype before model use.
+- The patch does not disable autocast for unrelated transformer/model compute.
+- The test covers finite output and dtype behavior as far as the local environment allows.
+
+Suggested commit message:
+
+```text
+Use fp32 for DeMemWM Plucker geometry
+```
+
+## Final Audit After All Substeps
+
+Status: `[ ]`
+
+After all implementation substeps are complete and committed, run a final read-only audit.
+
+Required commands:
+
+```bash
+git status --short --branch
+python -m py_compile datasets/video/minecraft_video_dememwm_latent_dataset.py algorithms/dememwm/df_video.py experiments/exp_base.py algorithms/dememwm/models/dit.py
+python -B -m unittest tests.test_dememwm_latent_dataset tests.test_dememwm_temporal_attention tests.test_checkpoint_loading tests.test_resume_checkpoint_logic
+git diff --check
+```
+
+Final audit review must check:
+
+- All substeps are marked `[x]`.
+- Each substep has a separate commit.
+- No unrelated files were changed.
+- No stale raw-video helper path remains.
+- No unfiltered nested validation/test fallback remains.
+- DeMemWM rejects stale experiment flags.
+- DeMemWM DiT boundary checks are concise and targeted.
+- Plucker geometry is fp32 for math and returns to activation dtype for model use.
+- There are no remaining uncommitted changes except intentional final plan status updates.
+
+If the final audit passes, mark this final audit status `[x]` and commit that status update only if it is not already included in the last substep commit.

datasets/video/minecraft_video_dememwm_latent_dataset.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import warnings
 from pathlib import Path
 import numpy as np
 import torch
@@ -85,16 +86,28 @@ class MinecraftVideoDeMemWMLatentDataset(torch.utils.data.Dataset):
         if not split_dir.exists():
             return []
         paths = sorted((path for path in split_dir.glob("**/*.npz") if path.is_file()), key=lambda path: path.name)
+        num_candidates = len(paths)
 
-        if self.wo_updown:
-            paths = [path for path in paths if "w_updown" not in str(path)]
-
-        if split in {"validation", "test"} and self.wo_updown:
+        split_filter = None
+        if split in {"validation", "test"}:
+            split_filter = "wo_updown" if self.wo_updown else "w_updown"
+            if self.wo_updown:
+                paths = [path for path in paths if "w_updown" not in str(path)]
+            else:
+                paths = [path for path in paths if "w_updown" in str(path)]
+        elif self.wo_updown:
             paths = [path for path in paths if "w_updown" not in str(path)]
-        elif split in {"validation", "test"}:
-            paths = [path for path in paths if "w_updown" in str(path)]
 
         if not paths:
+            if split_filter is not None and num_candidates:
+                message = (
+                    f"Selected split '{split}' under {split_dir} is empty because the "
+                    f"{split_filter} split filter removed all {num_candidates} candidate .npz file(s)."
+                )
+                warnings.warn(message, stacklevel=2)
+                raise FileNotFoundError(message)
+            if split_filter is not None:
+                return paths
             for sub_path in sorted(split_dir.iterdir(), key=lambda path: path.name):
                 paths += sorted((path for path in sub_path.glob("**/*.npz") if path.is_file()), key=lambda path: path.name)
         return paths

tests/test_dememwm_latent_dataset.py

@@ -197,6 +197,33 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         self.assertEqual([path.name for path in dataset.data_paths], ["clip_w_updown.npz"])
         self.assertEqual([path.name for path in dataset_wo.data_paths], ["plain.npz"])
 
+    def test_dataset_filters_nested_validation_files_by_selected_split(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp)
+            _write_latent_npz(root / "validation" / "nested" / "plain.npz", 104)
+            _write_latent_npz(root / "validation" / "nested" / "clip_w_updown.npz", 104)
+
+            dataset = MinecraftVideoDeMemWMLatentDataset(_dataset_cfg(root), split="validation")
+            dataset_wo = MinecraftVideoDeMemWMLatentDataset(_dataset_cfg(root, wo_updown=True), split="validation")
+
+        self.assertEqual([path.name for path in dataset.data_paths], ["clip_w_updown.npz"])
+        self.assertEqual([path.name for path in dataset_wo.data_paths], ["plain.npz"])
+
+    def test_dataset_rejects_nested_validation_test_files_removed_by_split_filter(self):
+        cases = [
+            ("validation", False, "plain.npz", "w_updown"),
+            ("test", True, "clip_w_updown.npz", "wo_updown"),
+        ]
+        for split, wo_updown, filename, split_filter in cases:
+            with self.subTest(split=split, split_filter=split_filter), tempfile.TemporaryDirectory() as tmp:
+                root = Path(tmp)
+                _write_latent_npz(root / split / "nested" / filename, 104)
+                message = f"Selected split '{split}'.*{split_filter} split filter removed all"
+
+                with self.assertWarnsRegex(UserWarning, message):
+                    with self.assertRaisesRegex(FileNotFoundError, message):
+                        MinecraftVideoDeMemWMLatentDataset(_dataset_cfg(root, wo_updown=wo_updown), split=split)
+
     def test_dataset_rejects_files_shorter_than_original_frame_skip_window(self):
         with tempfile.TemporaryDirectory() as tmp:
             root = Path(tmp)
@@ -524,7 +551,7 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         class Harness:
             def __init__(self):
                 self.memory_condition_length, self.clip_noise = 3, 0.0
-                self.context_frames, self.frame_stack, self.chunk_size = 2, 1, 2
+                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}}
                 )
@@ -592,26 +619,26 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         self.assertTrue(all(call["kwargs"]["reference_length"] == 0 for call in calls))
         self.assertEqual(
             ([call["kwargs"]["current_frame"] for call in calls[::2]], [call["kwargs"]["frame_memory_segments"]["target"] for call in calls[::2]]),
-            ([2, 4], [2, 1]),
+            ([2, 4], [2, 2]),
         )
         self.assertEqual([call["kwargs"]["current_frame"] for call in calls], [2, 2, 4, 4])
         self.assertTrue(torch.equal(calls[0]["kwargs"]["frame_idx"], calls[1]["kwargs"]["frame_idx"]))
         self.assertTrue(torch.equal(calls[2]["kwargs"]["frame_idx"], calls[3]["kwargs"]["frame_idx"]))
         self.assertEqual(calls[0]["x"][:, 0, 0, 0, 0].tolist(), [0.0, 0.0, 0.0, 1.0, 0.0])
-        self.assertEqual(calls[2]["x"][:, 0, 0, 0, 0].tolist(), [0.0, 0.0, 2.0, 2.0])
+        self.assertEqual(calls[2]["x"][:, 0, 0, 0, 0].tolist(), [2.0, 0.0, 0.0, 2.0, 2.0])
         self.assertEqual(calls[0]["action_cond"][:, 0, 0].tolist(), [102.0, 103.0, 0.0, 0.0, 0.0])
         self.assertEqual(calls[0]["kwargs"]["frame_idx"][:, 0].tolist(), [12, 13, 10, 11, 0])
-        self.assertEqual(calls[2]["kwargs"]["frame_idx"][:, 0].tolist(), [14, 10, 13, 12])
+        self.assertEqual(calls[2]["kwargs"]["frame_idx"][:, 0].tolist(), [13, 14, 10, 12, 12])
         self.assertEqual(calls[0]["kwargs"]["frame_memory_masks"]["revisit"].tolist(), [[False]])
         self.assertEqual(calls[2]["kwargs"]["frame_memory_masks"]["revisit"].tolist(), [[True]])
         self.assertAlmostEqual(float(loss), 1.0 / 3.0)
         self.assertEqual([name for name, _ in harness.logged], ["test/latent_mse"])
         self.assertEqual(len(harness.decoded_inputs), 2)
-        self.assertEqual(harness.decoded_inputs[0][:, 0, 0, 0, 0].tolist(), [2.0, 2.0, 4.0])
+        self.assertEqual(harness.decoded_inputs[0][:, 0, 0, 0, 0].tolist(), [2.0, 4.0, 4.0])
         self.assertEqual(harness.decoded_inputs[1][:, 0, 0, 0, 0].tolist(), [2.0, 3.0, 4.0])
         self.assertEqual(len(harness.metric_updates), 1)
         metric_pred, metric_gt = harness.metric_updates[0]
-        self.assertEqual(metric_pred[:, 0, 0, 0, 0].tolist(), [12.0, 12.0, 14.0])
+        self.assertEqual(metric_pred[:, 0, 0, 0, 0].tolist(), [12.0, 14.0, 14.0])
         self.assertEqual(metric_gt[:, 0, 0, 0, 0].tolist(), [12.0, 13.0, 14.0])
         self.assertEqual(len(video_calls), 1)
         self.assertEqual(video_calls[0][1]["namespace"], "test_vis")