Align DeMemWM latent memory frame offsets

configurations/dataset/video_minecraft_dememwm_latent.yaml

@@ -1,7 +1,8 @@
 defaults:
   - base_dataset
 
-save_dir: data/minecraft_simple_backforward_latents
+save_dir: data/minecraft_simple_backforward
+precomputed_feature_dir: ${dataset.save_dir}/vae_features
 metadata: null
 n_frames: 8
 n_frames_valid: 500

datasets/video/memory_selection.py

@@ -336,10 +336,10 @@ def _select_anchor(candidates: np.ndarray, count: int, cfg) -> np.ndarray:
     return np.asarray(selected[:count], dtype=np.int64)
 
 
-def _select_dynamic(target_start: int, count: int) -> np.ndarray:
-    if count <= 0 or target_start <= 0:
+def _select_dynamic(target_start: int, count: int, min_candidate_frame: int = 0) -> np.ndarray:
+    if count <= 0 or target_start <= min_candidate_frame:
         return np.empty((0,), dtype=np.int64)
-    start = max(0, target_start - count)
+    start = max(min_candidate_frame, target_start - count)
     return np.arange(start, target_start, dtype=np.int64)[-count:]
 
 
@@ -351,6 +351,7 @@ def _select_revisit(
     excluded: np.ndarray,
     split: str,
     rng=None,
+    min_candidate_frame: int = 0,
 ) -> np.ndarray:
     if count <= 0:
         return np.empty((0,), dtype=np.int64)
@@ -358,7 +359,7 @@ def _select_revisit(
     target_start = int(target_positions[0])
     exclusion = max(0, int(cfg_get(cfg, "local_context_exclusion_frames", 8)))
     causal_stop = max(0, target_start - exclusion)
-    candidates = np.arange(0, causal_stop, dtype=np.int64)
+    candidates = np.arange(min_candidate_frame, causal_stop, dtype=np.int64)
     if len(candidates) == 0:
         return np.empty((0,), dtype=np.int64)
 
@@ -382,13 +383,14 @@ def _select_revisit(
     )
 
 
-def select_memory_indices(poses, target_positions, cfg=None, split: str = "training", rng=None):
+def select_memory_indices(poses, target_positions, cfg=None, split: str = "training", rng=None, min_candidate_frame: int = 0):
     """Select causal memory indices and masks for [anchor][dynamic][revisit]."""
 
     poses = _as_pose_array(poses)
     target_positions = np.asarray(target_positions, dtype=np.int64)
     target_positions = target_positions[(target_positions >= 0) & (target_positions < len(poses))]
     enabled = bool(cfg_get(cfg, "enabled", True))
+    min_candidate_frame = max(0, int(min_candidate_frame))
 
     counts = {
         "anchor": int(cfg_get(cfg, "max_anchor_frames", 0)) if enabled else 0,
@@ -402,12 +404,21 @@ def select_memory_indices(poses, target_positions, cfg=None, split: str = "train
 
     target_start = int(target_positions[0])
 
-    dynamic = _select_dynamic(target_start, counts["dynamic"])
+    dynamic = _select_dynamic(target_start, counts["dynamic"], min_candidate_frame)
     anchor_stop = max(0, target_start - counts["dynamic"])
-    anchor_candidates = np.arange(0, anchor_stop, dtype=np.int64)
+    anchor_candidates = np.arange(min_candidate_frame, anchor_stop, dtype=np.int64)
     anchor = _select_anchor(anchor_candidates, counts["anchor"], cfg)
     excluded = np.concatenate([anchor, dynamic]) if len(anchor) or len(dynamic) else np.empty((0,), dtype=np.int64)
-    revisit = _select_revisit(poses, target_positions, cfg, counts["revisit"], excluded, split, rng=rng)
+    revisit = _select_revisit(
+        poses,
+        target_positions,
+        cfg,
+        counts["revisit"],
+        excluded,
+        split,
+        rng=rng,
+        min_candidate_frame=min_candidate_frame,
+    )
 
     _write_segment(indices, masks, "anchor", anchor)
     _write_segment(indices, masks, "dynamic", dynamic)

datasets/video/minecraft_video_dememwm_latent_dataset.py

@@ -2,7 +2,9 @@
 
 from __future__ import annotations
 
-import warnings
+import json
+import os
+import random
 from pathlib import Path
 import numpy as np
 import torch
@@ -11,32 +13,60 @@ from omegaconf import DictConfig
 from .memory_selection import SEGMENT_KEYS, cfg_get, memory_segment_lengths, select_memory_indices
 
 
-_INITIAL_FRAME_OFFSET = 100
-_ACTION_DIM = 25
-
-
-def _convert_action_space(actions: np.ndarray) -> np.ndarray:
-    vec_25 = np.zeros((len(actions), _ACTION_DIM), dtype=np.float32)
-    vec_25[actions[:, 0] == 1, 11] = 1.0
-    vec_25[actions[:, 0] == 2, 12] = 1.0
-    vec_25[actions[:, 4] == 11, 16] = -1.0
-    vec_25[actions[:, 4] == 13, 16] = 1.0
-    vec_25[actions[:, 3] == 11, 15] = -1.0
-    vec_25[actions[:, 3] == 13, 15] = 1.0
-    vec_25[actions[:, 5] == 6, 24] = 1.0
-    vec_25[actions[:, 5] == 1, 24] = 1.0
-    vec_25[actions[:, 1] == 1, 13] = 1.0
-    vec_25[actions[:, 1] == 2, 14] = 1.0
-    vec_25[actions[:, 7] == 1, 2] = 1.0
+ACTION_KEYS = [
+    "inventory",
+    "ESC",
+    "hotbar.1",
+    "hotbar.2",
+    "hotbar.3",
+    "hotbar.4",
+    "hotbar.5",
+    "hotbar.6",
+    "hotbar.7",
+    "hotbar.8",
+    "hotbar.9",
+    "forward",
+    "back",
+    "left",
+    "right",
+    "cameraY",
+    "cameraX",
+    "jump",
+    "sneak",
+    "sprint",
+    "swapHands",
+    "attack",
+    "use",
+    "pickItem",
+    "drop",
+]
+
+def convert_action_space(actions):
+    vec_25 = torch.zeros(len(actions), len(ACTION_KEYS))
+    vec_25[actions[:,0]==1, 11] = 1
+    vec_25[actions[:,0]==2, 12] = 1
+    vec_25[actions[:,4]==11, 16] = -1
+    vec_25[actions[:,4]==13, 16] = 1
+    vec_25[actions[:,3]==11, 15] = -1
+    vec_25[actions[:,3]==13, 15] = 1
+    vec_25[actions[:,5]==6, 24] = 1
+    vec_25[actions[:,5]==1, 24] = 1
+    vec_25[actions[:,1]==1, 13] = 1
+    vec_25[actions[:,1]==2, 14] = 1
+    vec_25[actions[:,7]==1, 2] = 1
     return vec_25
 
 
+_INITIAL_FRAME_OFFSET = 100
+_ACTION_DIM = len(ACTION_KEYS)
+
+
 class MinecraftVideoDeMemWMLatentDataset(torch.utils.data.Dataset):
     """Load precomputed VAE latents and append causal memory frames.
 
-    Expected files are ``.npz`` archives under ``save_dir/<split>`` with
-    ``latents``, ``actions``, and ``poses`` arrays. Optional ``frame_indices``
-    and ``image_hw`` arrays are propagated when present.
+    The index comes from the raw Minecraft ``.mp4`` tree, matching
+    ``MinecraftVideoDataset``. Latents are read from the original DeMemWM
+    ``vae_features`` cache at the same relative path.
     """
 
     def __init__(self, cfg: DictConfig, split: str = "training"):
@@ -44,8 +74,13 @@ class MinecraftVideoDeMemWMLatentDataset(torch.utils.data.Dataset):
         self.cfg = cfg
         self.split = split
         self.save_dir = Path(cfg.save_dir)
+        precomputed_feature_dir = cfg_get(cfg, "precomputed_feature_dir")
+        self.feature_root = Path(precomputed_feature_dir) if precomputed_feature_dir is not None else self.save_dir / "vae_features"
         self.wo_updown = bool(cfg_get(cfg, "wo_updown", False))
         self.initial_frame_offset = _INITIAL_FRAME_OFFSET
+        context_window = cfg_get(cfg, "context_window", cfg_get(cfg, "context_length", 0))
+        self.context_window = max(0, int(context_window)) if split == "training" else 0
+        self.target_start_offset = self.initial_frame_offset + self.context_window
         self.n_frames = int(cfg_get(cfg, "n_frames_valid", cfg.n_frames)) if split in {"validation", "test"} else int(cfg.n_frames)
         self.frame_skip = int(cfg_get(cfg, "frame_skip", 1))
         self.target_span = max(1, (self.n_frames - 1) * self.frame_skip + 1)
@@ -59,10 +94,10 @@ class MinecraftVideoDeMemWMLatentDataset(torch.utils.data.Dataset):
         self._cached_arrays = None
         self.data_paths = self.get_data_paths(split)
         if not self.data_paths:
-            raise FileNotFoundError(f"No latent .npz files found for split '{split}' under {self.save_dir / split}")
+            raise FileNotFoundError(f"No raw .mp4 files found for split '{split}' under {self.save_dir / split}")
 
         self.lengths = np.asarray([self._file_length(path) for path in self.data_paths], dtype=np.int64)
-        required_length = self.initial_frame_offset + self.target_span
+        required_length = self.target_start_offset + self.target_span
         too_short = self.lengths < required_length
         if bool(too_short.any()):
             bad_idx = int(np.nonzero(too_short)[0][0])
@@ -72,44 +107,47 @@ class MinecraftVideoDeMemWMLatentDataset(torch.utils.data.Dataset):
             )
 
         clips = (self.lengths - required_length + 1).astype(np.int64)
+        if split == "training":
+            # Keep targets inside the original 1300-frame training window after
+            # reserving the eval-style context prefix before each target clip.
+            original_training_clips = max(1, 1300 - self.context_window - self.target_span + 1)
+            clips = np.minimum(clips, original_training_clips)
         if self.single_eval_clip:
             clips = np.ones_like(clips)
         self.clips_per_file = clips
         self.cum_clips_per_file = np.cumsum(self.clips_per_file)
         self.idx_remap = list(range(self.__len__()))
-        if split == "training" and bool(cfg_get(cfg, "shuffle_clips", True)):
-            rng = np.random.default_rng(int(cfg_get(cfg, "seed", 0)))
+        if bool(cfg_get(cfg, "shuffle_clips", True)):
+            rng = random.Random(int(cfg_get(cfg, "seed", 0)))
             rng.shuffle(self.idx_remap)
 
     def get_data_paths(self, split: str):
-        split_dir = self.save_dir / split
-        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)
-
-        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)]
+        """
+        Retrieve all video file paths for the given split.
+
+        Args:
+            split (str): Dataset split ("training" or "validation").
+
+        Returns:
+            List[Path]: List of video file paths.
+        """
+        data_dir = self.save_dir / split
+        paths = sorted(list(data_dir.glob("**/*.mp4")), key=lambda x: x.name)
+
+        if self.wo_updown:
+            # Filter out paths containing "w_updown"
+            paths = [p for p in paths if "w_updown" not in str(p)]
+
+        if (split == "validation" or split == "test") and self.wo_updown:
+            paths = [p for p in paths if "w_updown" not in str(p)]
+        elif split == "validation" or split == "test":
+            paths = [p for p in paths if "w_updown" in str(p)]
 
         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)
+            sub_dirs = os.listdir(data_dir)
+            for sub_dir in sub_dirs:
+                sub_path = data_dir / sub_dir
+                paths += sorted(list(sub_path.glob("**/*.mp4")), key=lambda x: x.name)
         return paths
 
     def __len__(self):
@@ -126,7 +164,7 @@ class MinecraftVideoDeMemWMLatentDataset(torch.utils.data.Dataset):
     def load_data(self, idx: int):
         # === 1. Resolve clip and load arrays ===
         file_idx, frame_idx = self.split_idx(int(idx))
-        frame_idx += self.initial_frame_offset
+        frame_idx += self.target_start_offset
         arrays = self._load_arrays(self.data_paths[file_idx])
         latents = arrays["latents"]
         actions = arrays["actions"]
@@ -140,6 +178,7 @@ class MinecraftVideoDeMemWMLatentDataset(torch.utils.data.Dataset):
             target_positions,
             self.memory_selection,
             split=self.split,
+            min_candidate_frame=self.initial_frame_offset,
         )
         ordered_memory = [memory_indices[key] for key in SEGMENT_KEYS]
         source_positions = np.concatenate([target_positions, *ordered_memory]).astype(np.int64)
@@ -172,39 +211,63 @@ class MinecraftVideoDeMemWMLatentDataset(torch.utils.data.Dataset):
             "image_hw": torch.as_tensor(arrays["image_hw"], dtype=torch.long),
         }
 
+    def _feature_cache_paths(self, video_path: Path) -> tuple[Path, Path]:
+        relative_video_path = video_path.relative_to(self.save_dir)
+        feature_dir = self.feature_root / relative_video_path.parent
+        video_index = relative_video_path.stem
+        return (
+            feature_dir / f"{video_index}_vae_feature.npy",
+            feature_dir / f"{video_index}_vae_feature_meta.json",
+        )
+
+    def _default_image_hw(self) -> np.ndarray:
+        return np.asarray(
+            cfg_get(self.cfg, "image_hw", [cfg_get(self.cfg, "resolution", 0), cfg_get(self.cfg, "resolution", 0)]),
+            dtype=np.int64,
+        ).reshape(-1)[:2]
+
+    def _feature_image_hw(self, meta_path: Path) -> np.ndarray:
+        if not meta_path.exists():
+            return self._default_image_hw()
+        with open(meta_path, "r", encoding="utf-8") as handle:
+            meta = json.load(handle)
+        return np.asarray([int(meta["image_height"]), int(meta["image_width"])], dtype=np.int64)
+
     def _file_length(self, path: Path) -> int:
-        with np.load(path, allow_pickle=False) as data:
-            return int(data["latents"].shape[0])
+        feature_path, _ = self._feature_cache_paths(path)
+        if not feature_path.exists():
+            raise FileNotFoundError(f"Missing precomputed VAE features: {feature_path}")
+        return int(np.load(feature_path, mmap_mode="r").shape[0])
 
     def _load_arrays(self, path: Path):
         if self._cached_arrays_path == path:
             return self._cached_arrays
 
-        with np.load(path, allow_pickle=False) as data:
-            latents = np.asarray(data["latents"])
-            actions = self._normalize_actions(path, np.asarray(data["actions"]))
-            poses = self._sanitize_poses(path, np.asarray(data["poses"], dtype=np.float32), len(actions))
-            if "frame_indices" in data:
-                frame_indices = np.asarray(data["frame_indices"], dtype=np.int64)
-            else:
-                frame_indices = np.arange(len(latents), dtype=np.int64)
-
-            if "image_hw" in data:
-                image_hw = np.asarray(data["image_hw"], dtype=np.int64).reshape(-1)[:2]
-            else:
-                image_hw = np.asarray(cfg_get(self.cfg, "image_hw", [cfg_get(self.cfg, "resolution", 0), cfg_get(self.cfg, "resolution", 0)]), dtype=np.int64)
-
-        self._validate_arrays(path, latents, actions, poses, frame_indices)
-        arrays = {
+        arrays = self._load_feature_arrays(path)
+
+        self._validate_arrays(path, arrays["latents"], arrays["actions"], arrays["poses"], arrays["frame_indices"])
+        self._cached_arrays_path = path
+        self._cached_arrays = arrays
+        return arrays
+
+    def _load_feature_arrays(self, video_path: Path):
+        feature_path, meta_path = self._feature_cache_paths(video_path)
+        if not feature_path.exists():
+            raise FileNotFoundError(f"Missing precomputed VAE features: {feature_path}")
+
+        action_path = video_path.with_suffix(".npz")
+        with np.load(action_path, allow_pickle=False) as data:
+            latents = np.load(feature_path, mmap_mode="r")
+            actions = self._normalize_actions(action_path, np.asarray(data["actions"]))
+            poses = self._sanitize_poses(action_path, np.asarray(data["poses"], dtype=np.float32), len(actions))
+
+        return {
             "latents": latents,
             "actions": actions,
             "poses": poses,
-            "frame_indices": frame_indices,
-            "image_hw": image_hw,
+            "frame_indices": np.arange(len(latents), dtype=np.int64),
+            "image_hw": self._feature_image_hw(meta_path),
         }
-        self._cached_arrays_path = path
-        self._cached_arrays = arrays
-        return arrays
 
     def _normalize_actions(self, path: Path, actions: np.ndarray) -> np.ndarray:
         if actions.ndim != 2:
@@ -213,7 +276,7 @@ class MinecraftVideoDeMemWMLatentDataset(torch.utils.data.Dataset):
             return actions.astype(np.float32, copy=False)
         if actions.shape[1] < 8:
             raise ValueError(f"raw actions in {path} must have at least 8 columns to convert to 25-D actions")
-        return _convert_action_space(actions)
+        return convert_action_space(actions).numpy()
 
     @staticmethod
     def _sanitize_poses(path: Path, poses: np.ndarray, action_length: int) -> np.ndarray:

tests/test_dememwm_latent_dataset.py

@@ -1,3 +1,4 @@
+import json
 import tempfile
 import unittest
 from unittest import mock
@@ -43,6 +44,7 @@ def _poses(num_frames):
 def _dataset_cfg(root, **overrides):
     cfg = {
         "save_dir": str(root),
+        "precomputed_feature_dir": str(Path(root) / "vae_features"),
         "n_frames": 3,
         "n_frames_valid": 3,
         "frame_skip": 1,
@@ -57,22 +59,21 @@ def _dataset_cfg(root, **overrides):
     return OmegaConf.create(cfg)
 
 
-def _write_latent_npz(path, num_frames, actions=None, poses=None, frame_indices=None):
-    path.parent.mkdir(parents=True, exist_ok=True)
+def _write_vae_feature_clip(root, feature_root, split="training", subdir="scene", stem="000001", num_frames=104, actions=None, poses=None):
+    video_path = Path(root) / split / subdir / f"{stem}.mp4"
+    video_path.parent.mkdir(parents=True, exist_ok=True)
+    video_path.write_bytes(b"")
     if actions is None:
-        actions = np.ones((num_frames, 25), dtype=np.float32)
+        actions = np.ones((num_frames, 8), dtype=np.int32)
     if poses is None:
         poses = _poses(num_frames)
-    if frame_indices is None:
-        frame_indices = np.arange(num_frames, dtype=np.int64)
-    np.savez(
-        path,
-        latents=np.arange(num_frames * 2, dtype=np.float32).reshape(num_frames, 1, 1, 2),
-        actions=actions,
-        poses=poses,
-        frame_indices=frame_indices,
-        image_hw=np.array([360, 640], dtype=np.int64),
-    )
+    np.savez(video_path.with_suffix(".npz"), actions=actions, poses=poses)
+
+    feature_dir = Path(feature_root) / split / subdir
+    feature_dir.mkdir(parents=True, exist_ok=True)
+    np.save(feature_dir / f"{stem}_vae_feature.npy", np.arange(num_frames * 2, dtype=np.float16).reshape(num_frames, 1, 1, 2))
+    with open(feature_dir / f"{stem}_vae_feature_meta.json", "w", encoding="utf-8") as handle:
+        json.dump({"image_height": 360, "image_width": 640}, handle)
 
 
 class MemorySelectionTests(unittest.TestCase):
@@ -178,7 +179,7 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
     def test_dataset_requires_split_directory_without_root_fallback(self):
         with tempfile.TemporaryDirectory() as tmp:
             root = Path(tmp)
-            _write_latent_npz(root / "sample.npz", 104)
+            (root / "sample.mp4").write_bytes(b"")
 
             with self.assertRaisesRegex(FileNotFoundError, "validation"):
                 MinecraftVideoDeMemWMLatentDataset(_dataset_cfg(root), split="validation")
@@ -186,46 +187,44 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
     def test_dataset_matches_original_w_updown_split_filtering(self):
         with tempfile.TemporaryDirectory() as tmp:
             root = Path(tmp)
-            _write_latent_npz(root / "validation" / "plain.npz", 104)
-            _write_latent_npz(root / "validation" / "clip_w_updown.npz", 104)
+            _write_vae_feature_clip(root, root / "vae_features", split="validation", stem="plain")
+            _write_vae_feature_clip(root, root / "vae_features", split="validation", stem="clip_w_updown")
 
             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"])
+        self.assertEqual([path.name for path in dataset.data_paths], ["clip_w_updown.mp4"])
+        self.assertEqual([path.name for path in dataset_wo.data_paths], ["plain.mp4"])
 
     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)
+            _write_vae_feature_clip(root, root / "vae_features", split="validation", subdir="nested", stem="plain")
+            _write_vae_feature_clip(root, root / "vae_features", split="validation", subdir="nested", stem="clip_w_updown")
 
             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"])
+        self.assertEqual([path.name for path in dataset.data_paths], ["clip_w_updown.mp4"])
+        self.assertEqual([path.name for path in dataset_wo.data_paths], ["plain.mp4"])
 
-    def test_dataset_rejects_nested_validation_test_files_removed_by_split_filter(self):
+    def test_dataset_matches_original_fallback_when_split_filter_has_no_matches(self):
         cases = [
-            ("validation", False, "plain.npz", "w_updown"),
-            ("test", True, "clip_w_updown.npz", "wo_updown"),
+            ("validation", False, "plain"),
+            ("test", True, "clip_w_updown"),
         ]
-        for split, wo_updown, filename, split_filter in cases:
-            with self.subTest(split=split, split_filter=split_filter), tempfile.TemporaryDirectory() as tmp:
+        for split, wo_updown, stem in cases:
+            with self.subTest(split=split, wo_updown=wo_updown), 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"
+                _write_vae_feature_clip(root, root / "vae_features", split=split, subdir="nested", stem=stem)
+                dataset = MinecraftVideoDeMemWMLatentDataset(_dataset_cfg(root, wo_updown=wo_updown), split=split)
 
-                with self.assertWarnsRegex(UserWarning, message):
-                    with self.assertRaisesRegex(FileNotFoundError, message):
-                        MinecraftVideoDeMemWMLatentDataset(_dataset_cfg(root, wo_updown=wo_updown), split=split)
+                self.assertEqual([path.name for path in dataset.data_paths], [f"{stem}.mp4"])
 
     def test_dataset_rejects_files_shorter_than_original_frame_skip_window(self):
         with tempfile.TemporaryDirectory() as tmp:
             root = Path(tmp)
-            _write_latent_npz(root / "training" / "short.npz", 102)
+            _write_vae_feature_clip(root, root / "vae_features", stem="short", num_frames=102)
 
             with self.assertRaisesRegex(ValueError, "requires at least 103"):
                 MinecraftVideoDeMemWMLatentDataset(_dataset_cfg(root), split="training")
@@ -235,7 +234,7 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
             root = Path(tmp)
             poses = _poses(104)
             poses[:, 1] = np.linspace(0.0, 3.0, 104, dtype=np.float32)
-            _write_latent_npz(root / "training" / "sample.npz", 104, poses=poses)
+            _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"):
@@ -249,7 +248,7 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
             raw_actions[:, 3] = 11
             raw_actions[:, 4] = 13
             raw_actions[:, 7] = 1
-            _write_latent_npz(root / "training" / "sample.npz", 104, actions=raw_actions)
+            _write_vae_feature_clip(root, root / "vae_features", stem="sample", actions=raw_actions)
             sample = MinecraftVideoDeMemWMLatentDataset(_dataset_cfg(root), split="training")[0]
 
         self.assertEqual(tuple(sample["actions"].shape), (3, 25))
@@ -258,10 +257,40 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         self.assertTrue(torch.equal(sample["actions"][:, 15], -torch.ones(3)))
         self.assertTrue(torch.equal(sample["actions"][:, 2], torch.ones(3)))
 
-    def test_dataset_reuses_cached_npz_arrays_for_same_file(self):
+    def test_dataset_reads_dememwm_vae_feature_layout(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp) / "minecraft"
+            feature_root = root / "vae_features"
+            raw_actions = np.zeros((104, 8), dtype=np.int64)
+            raw_actions[:, 0] = 1
+            raw_actions[:, 3] = 11
+            raw_actions[:, 4] = 13
+            raw_actions[:, 7] = 1
+            _write_vae_feature_clip(root, feature_root, actions=raw_actions)
+            dataset = MinecraftVideoDeMemWMLatentDataset(_dataset_cfg(root), split="training")
+            sample = dataset[0]
+
+        self.assertEqual([path.suffix for path in dataset.data_paths], [".mp4"])
+        self.assertEqual(tuple(sample["latents"].shape), (3, 1, 1, 2))
+        self.assertEqual(tuple(sample["actions"].shape), (3, 25))
+        self.assertEqual(sample["image_hw"].tolist(), [360, 640])
+        self.assertTrue(torch.equal(sample["actions"][:, 11], torch.ones(3)))
+        self.assertTrue(torch.equal(sample["actions"][:, 16], torch.ones(3)))
+        self.assertTrue(torch.equal(sample["actions"][:, 15], -torch.ones(3)))
+        self.assertTrue(torch.equal(sample["actions"][:, 2], torch.ones(3)))
+
+    def test_dataset_uses_original_training_clip_window(self):
         with tempfile.TemporaryDirectory() as tmp:
             root = Path(tmp)
-            _write_latent_npz(root / "training" / "sample.npz", 104)
+            _write_vae_feature_clip(root, root / "vae_features", stem="sample", num_frames=1500)
+            dataset = MinecraftVideoDeMemWMLatentDataset(_dataset_cfg(root, n_frames=8), split="training")
+
+        self.assertEqual(len(dataset), 1300 - 8 + 1)
+
+    def test_dataset_reuses_cached_feature_arrays_for_same_file(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp)
+            _write_vae_feature_clip(root, root / "vae_features", stem="sample")
             dataset = MinecraftVideoDeMemWMLatentDataset(_dataset_cfg(root), split="training")
             first = dataset[0]
 
@@ -714,7 +743,7 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
     def test_disabled_memory_selection_does_not_allocate_memory_slots(self):
         with tempfile.TemporaryDirectory() as tmp:
             root = Path(tmp)
-            _write_latent_npz(root / "training" / "sample.npz", 104)
+            _write_vae_feature_clip(root, root / "vae_features", stem="sample")
             dataset = MinecraftVideoDeMemWMLatentDataset(
                 _dataset_cfg(root, memory_selection=_selection_cfg(enabled=False)),
                 split="training",
@@ -729,44 +758,97 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         for key in ("anchor", "dynamic", "revisit"):
             self.assertEqual(tuple(sample["memory_masks"][key].shape), (0,))
 
+    def test_dataset_excludes_initial_offset_from_memory_candidates(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp)
+            _write_vae_feature_clip(
+                root,
+                root / "vae_features",
+                stem="sample",
+                num_frames=112,
+                actions=np.ones((112, 25), dtype=np.float32),
+            )
+            dataset = MinecraftVideoDeMemWMLatentDataset(
+                _dataset_cfg(root, memory_selection=_selection_cfg()),
+                split="training",
+            )
+            sample = dataset[0]
+
+        self.assertEqual(sample["frame_indices"][:3].tolist(), [100, 101, 102])
+        self.assertEqual(sample["frame_indices"][3:].tolist(), [-1, -1, -1, -1, -1, -1])
+        for key in ("anchor", "dynamic", "revisit"):
+            self.assertFalse(sample["memory_masks"][key].any().item())
+
+    def test_training_context_window_shifts_target_start_but_not_memory_floor(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp)
+            _write_vae_feature_clip(
+                root,
+                root / "vae_features",
+                stem="sample",
+                num_frames=112,
+                actions=np.ones((112, 25), dtype=np.float32),
+            )
+            dataset = MinecraftVideoDeMemWMLatentDataset(
+                _dataset_cfg(
+                    root,
+                    context_length=4,
+                    memory_selection=_selection_cfg(max_anchor_frames=2, max_dynamic_frames=2, max_revisit_frames=0),
+                ),
+                split="training",
+            )
+            sample = dataset[0]
+
+        self.assertEqual(dataset.context_window, 4)
+        self.assertEqual(len(dataset), 112 - (100 + 4 + 3) + 1)
+        self.assertEqual(sample["frame_indices"].tolist(), [104, 105, 106, 100, 101, 102, 103])
+        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,))
+
+    def test_validation_ignores_training_context_window_shift(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp)
+            _write_vae_feature_clip(
+                root,
+                root / "vae_features",
+                split="validation",
+                stem="clip_w_updown",
+                num_frames=112,
+                actions=np.ones((112, 25), dtype=np.float32),
+            )
+            dataset = MinecraftVideoDeMemWMLatentDataset(
+                _dataset_cfg(root, context_length=4),
+                split="validation",
+            )
+            sample = dataset[0]
+
+        self.assertEqual(dataset.context_window, 0)
+        self.assertEqual(sample["frame_indices"].tolist(), [100, 101, 102])
+
     def test_dataset_returns_target_anchor_dynamic_revisit_contract(self):
         with tempfile.TemporaryDirectory() as tmp:
             root = Path(tmp)
-            split_dir = root / "training"
-            split_dir.mkdir()
             num_frames = 112
-            np.savez(
-                split_dir / "sample.npz",
-                latents=np.arange(num_frames * 2, dtype=np.float32).reshape(num_frames, 1, 1, 2),
+            _write_vae_feature_clip(
+                root,
+                root / "vae_features",
+                stem="sample",
+                num_frames=num_frames,
                 actions=np.ones((num_frames, 25), dtype=np.float32),
-                poses=_poses(num_frames),
-                frame_indices=np.arange(num_frames, dtype=np.int64),
-                image_hw=np.array([360, 640], dtype=np.int64),
             )
-            cfg = OmegaConf.create(
-                {
-                    "save_dir": str(root),
-                    "n_frames": 3,
-                    "n_frames_valid": 3,
-                    "frame_skip": 1,
-                    "action_cond_dim": 25,
-                    "resolution": 128,
-                    "image_hw": [360, 640],
-                    "shuffle_clips": False,
-                    "single_eval_clip": True,
-                    "memory_selection": _selection_cfg(),
-                }
+            dataset = MinecraftVideoDeMemWMLatentDataset(
+                _dataset_cfg(root, memory_selection=_selection_cfg()),
+                split="training",
             )
-            dataset = MinecraftVideoDeMemWMLatentDataset(cfg, split="training")
             sample = dataset[6]
 
         self.assertEqual(sample["memory_segments"], {"target": 3, "anchor": 2, "dynamic": 2, "revisit": 2})
         self.assertEqual(tuple(sample["latents"].shape), (9, 1, 1, 2))
         self.assertEqual(sample["frame_indices"][:3].tolist(), [106, 107, 108])
-        self.assertEqual(sample["frame_indices"][3:5].tolist(), [0, 1])
+        self.assertEqual(sample["frame_indices"][3:5].tolist(), [100, 101])
         self.assertEqual(sample["frame_indices"][5:7].tolist(), [104, 105])
-        self.assertTrue(all(idx < 104 for idx in sample["frame_indices"][7:9].tolist() if idx >= 0))
-        self.assertEqual(len(sample["frame_indices"][7:9]), 2)
+        self.assertEqual(sample["frame_indices"][7:9].tolist(), [102, 103])
         self.assertTrue(sample["memory_masks"]["target"].all().item())
         self.assertTrue(sample["memory_masks"]["anchor"].all().item())
         self.assertTrue(sample["memory_masks"]["dynamic"].all().item())
@@ -786,7 +868,7 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         self.assertTrue(torch.equal(preprocessed["actions"][3:], torch.ones_like(preprocessed["actions"][3:])))
         self.assertEqual(tuple(preprocessed["poses"].shape), (9, 1, 5))
         self.assertIs(preprocessed["frame_memory_pose"], preprocessed["poses"])
-        self.assertEqual(preprocessed["frame_memory_pose"][:7, 0, 0].tolist(), [0.0, 1.0, 2.0, -106.0, -105.0, -2.0, -1.0])
+        self.assertEqual(preprocessed["frame_memory_pose"][:7, 0, 0].tolist(), [0.0, 1.0, 2.0, -6.0, -5.0, -2.0, -1.0])
         self.assertEqual(tuple(preprocessed["frame_indices"].shape), (9, 1))
         self.assertEqual(preprocessed["memory_segments"], sample["memory_segments"])
         self.assertEqual(preprocessed["segment_lengths"], sample["memory_segments"])
@@ -803,7 +885,7 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
             {"anchor": (3, 5), "dynamic": (5, 7), "revisit": (7, 9)},
         )
         self.assertEqual(preprocessed["target_tensors"]["frame_indices"][:, 0].tolist(), [106, 107, 108])
-        self.assertEqual(preprocessed["stream_tensors"]["anchor"]["frame_indices"][:, 0].tolist(), [0, 1])
+        self.assertEqual(preprocessed["stream_tensors"]["anchor"]["frame_indices"][:, 0].tolist(), [100, 101])
         self.assertEqual(preprocessed["stream_tensors"]["dynamic"]["frame_indices"][:, 0].tolist(), [104, 105])
         self.assertEqual(preprocessed["stream_tensors"]["revisit"]["frame_indices"].shape[0], 2)
         self.assertEqual(tuple(preprocessed["memory_masks"]["target"].shape), (1, 3))