Align DeMemWM latent dataset loading

datasets/video/minecraft_video_dememwm_latent_dataset.py

@@ -10,6 +10,26 @@ 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
+    return vec_25
+
+
 class MinecraftVideoDeMemWMLatentDataset(torch.utils.data.Dataset):
     """Load precomputed VAE latents and append causal memory frames.
 
@@ -23,6 +43,8 @@ class MinecraftVideoDeMemWMLatentDataset(torch.utils.data.Dataset):
         self.cfg = cfg
         self.split = split
         self.save_dir = Path(cfg.save_dir)
+        self.wo_updown = bool(cfg_get(cfg, "wo_updown", False))
+        self.initial_frame_offset = _INITIAL_FRAME_OFFSET
         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)
@@ -32,12 +54,23 @@ class MinecraftVideoDeMemWMLatentDataset(torch.utils.data.Dataset):
         self._target_offsets = np.arange(self.n_frames, dtype=np.int64) * self.frame_skip
         self._memory_segments = memory_segment_lengths(self.n_frames, self.memory_selection)
         self.memory_condition_length = sum(self._memory_segments[key] for key in SEGMENT_KEYS)
+        self._cached_arrays_path = None
+        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}")
+            raise FileNotFoundError(f"No latent .npz 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)
-        clips = np.clip(self.lengths - self.target_span + 1, a_min=1, a_max=None).astype(np.int64)
+        required_length = self.initial_frame_offset + self.target_span
+        too_short = self.lengths < required_length
+        if bool(too_short.any()):
+            bad_idx = int(np.nonzero(too_short)[0][0])
+            raise ValueError(
+                f"Latent file {self.data_paths[bad_idx]} has {int(self.lengths[bad_idx])} frames, "
+                f"but DeMemWM requires at least {required_length} for split '{split}'"
+            )
+
+        clips = (self.lengths - required_length + 1).astype(np.int64)
         if self.single_eval_clip:
             clips = np.ones_like(clips)
         self.clips_per_file = clips
@@ -49,8 +82,22 @@ class MinecraftVideoDeMemWMLatentDataset(torch.utils.data.Dataset):
 
     def get_data_paths(self, split: str):
         split_dir = self.save_dir / split
-        data_dir = split_dir if split_dir.exists() else self.save_dir
-        return sorted(path for path in data_dir.glob("**/*.npz") if path.is_file())
+        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)
+
+        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:
+            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:
+            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
 
     def __len__(self):
         return int(self.cum_clips_per_file[-1])
@@ -66,6 +113,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
         arrays = self._load_arrays(self.data_paths[file_idx])
         latents = arrays["latents"]
         actions = arrays["actions"]
@@ -116,10 +164,13 @@ class MinecraftVideoDeMemWMLatentDataset(torch.utils.data.Dataset):
             return int(data["latents"].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 = np.asarray(data["actions"])
-            poses = np.asarray(data["poses"], dtype=np.float32)
+            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:
@@ -131,13 +182,40 @@ class MinecraftVideoDeMemWMLatentDataset(torch.utils.data.Dataset):
                 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)
-        return {
+        arrays = {
             "latents": latents,
             "actions": actions,
             "poses": poses,
             "frame_indices": frame_indices,
             "image_hw": image_hw,
         }
+        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:
+            raise ValueError(f"actions in {path} must have shape [num_frames, action_dim]")
+        if actions.shape[1] == _ACTION_DIM:
+            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)
+
+    @staticmethod
+    def _sanitize_poses(path: Path, poses: np.ndarray, action_length: int) -> np.ndarray:
+        if poses.ndim != 2 or poses.shape[1] < 5:
+            raise ValueError(f"poses in {path} must have shape [num_frames, >=5]")
+
+        poses = poses.astype(np.float32, copy=True)
+        if len(poses) > 1:
+            poses[0, 1] = poses[1, 1]
+        height_range = float(np.ptp(poses[:, 1])) if len(poses) else 0.0
+        if height_range >= 2:
+            raise ValueError(f"Pose height variation too large for {path}: {height_range}")
+        if len(poses) < action_length:
+            poses = np.pad(poses, ((1, 0), (0, 0)))
+        return poses
 
     def _validate_arrays(self, path: Path, latents, actions, poses, frame_indices) -> None:
         length = len(latents)

tests/test_dememwm_latent_dataset.py

@@ -42,6 +42,41 @@ def _poses(num_frames):
     return poses
 
 
+def _dataset_cfg(root, **overrides):
+    cfg = {
+        "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(max_anchor_frames=0, max_dynamic_frames=0, max_revisit_frames=0),
+    }
+    cfg.update(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)
+    if actions is None:
+        actions = np.ones((num_frames, 25), dtype=np.float32)
+    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),
+    )
+
+
 class MemorySelectionTests(unittest.TestCase):
     def test_worldmem_backend_is_explicitly_rejected_until_implemented(self):
         poses = _poses(8)
@@ -149,6 +184,75 @@ class MemorySelectionTests(unittest.TestCase):
 
 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)
+
+            with self.assertRaisesRegex(FileNotFoundError, "validation"):
+                MinecraftVideoDeMemWMLatentDataset(_dataset_cfg(root), split="validation")
+
+    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)
+
+            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_files_shorter_than_original_frame_skip_window(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp)
+            _write_latent_npz(root / "training" / "short.npz", 102)
+
+            with self.assertRaisesRegex(ValueError, "requires at least 103"):
+                MinecraftVideoDeMemWMLatentDataset(_dataset_cfg(root), split="training")
+
+    def test_dataset_applies_original_pose_height_sanity_check(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            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)
+            dataset = MinecraftVideoDeMemWMLatentDataset(_dataset_cfg(root), split="training")
+
+            with self.assertRaisesRegex(ValueError, "Pose height variation"):
+                dataset[0]
+
+    def test_dataset_converts_raw_actions_to_original_25d_contract(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp)
+            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_latent_npz(root / "training" / "sample.npz", 104, actions=raw_actions)
+            sample = MinecraftVideoDeMemWMLatentDataset(_dataset_cfg(root), split="training")[0]
+
+        self.assertEqual(tuple(sample["actions"].shape), (3, 25))
+        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_reuses_cached_npz_arrays_for_same_file(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp)
+            _write_latent_npz(root / "training" / "sample.npz", 104)
+            dataset = MinecraftVideoDeMemWMLatentDataset(_dataset_cfg(root), split="training")
+            first = dataset[0]
+
+            with mock.patch("datasets.video.minecraft_video_dememwm_latent_dataset.np.load", side_effect=AssertionError("cache miss")):
+                second = dataset[1]
+
+        self.assertEqual(first["frame_indices"][:3].tolist(), [100, 101, 102])
+        self.assertEqual(second["frame_indices"][:3].tolist(), [101, 102, 103])
+
     def test_preprocess_splits_sequence_tensors_from_memory_segments(self):
         from algorithms.dememwm.df_video import _preprocess_dememwm_latent_batch
 
@@ -497,13 +601,13 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
             root = Path(tmp)
             split_dir = root / "training"
             split_dir.mkdir()
-            num_frames = 12
+            num_frames = 112
             np.savez(
                 split_dir / "sample.npz",
                 latents=np.arange(num_frames * 2, dtype=np.float32).reshape(num_frames, 1, 1, 2),
                 actions=np.ones((num_frames, 25), dtype=np.float32),
                 poses=_poses(num_frames),
-                frame_indices=np.arange(100, 100 + num_frames, dtype=np.int64),
+                frame_indices=np.arange(num_frames, dtype=np.int64),
                 image_hw=np.array([360, 640], dtype=np.int64),
             )
             cfg = OmegaConf.create(
@@ -526,7 +630,7 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         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(), [100, 101])
+        self.assertEqual(sample["frame_indices"][3:5].tolist(), [0, 1])
         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)
@@ -549,7 +653,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, -6.0, -5.0, -2.0, -1.0])
+        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(tuple(preprocessed["frame_indices"].shape), (9, 1))
         self.assertEqual(preprocessed["memory_segments"], sample["memory_segments"])
         self.assertEqual(preprocessed["segment_lengths"], sample["memory_segments"])
@@ -566,7 +670,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(), [100, 101])
+        self.assertEqual(preprocessed["stream_tensors"]["anchor"]["frame_indices"][:, 0].tolist(), [0, 1])
         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))