Implement DeMemWM batch contract

algorithms/dememwm/df_video.py

@@ -1,3 +1,4 @@
+from collections.abc import Mapping
 import os
 import random
 import math
@@ -24,6 +25,53 @@ from .models.pose_prediction import PosePredictionNet
 import glob
 
 # Utility Functions
+_DEMEMWM_SEGMENT_KEYS = ("target", "anchor", "dynamic", "revisit")
+
+
+def _segment_value_to_int(value, key):
+    if torch.is_tensor(value):
+        flat = value.reshape(-1)
+        if flat.numel() == 0:
+            raise ValueError(f"memory_segments['{key}'] is empty")
+        first = flat[0]
+        if flat.numel() > 1 and not bool(torch.all(flat == first).item()):
+            raise ValueError(f"memory_segments['{key}'] must be identical across the batch")
+        return int(first.item())
+    if isinstance(value, (list, tuple)):
+        values = [_segment_value_to_int(item, key) for item in value]
+        if not values:
+            raise ValueError(f"memory_segments['{key}'] is empty")
+        if any(item != values[0] for item in values[1:]):
+            raise ValueError(f"memory_segments['{key}'] must be identical across the batch")
+        return values[0]
+    return int(value)
+
+
+def _preprocess_dememwm_latent_batch(batch):
+    memory_segments = {
+        key: _segment_value_to_int(batch["memory_segments"][key], key)
+        for key in _DEMEMWM_SEGMENT_KEYS
+    }
+    memory_masks = {key: batch["memory_masks"][key] for key in _DEMEMWM_SEGMENT_KEYS}
+
+    # Latent dataset batches are already VAE-encoded: B x T_all x C x H_lat x W_lat.
+    latents = rearrange(batch["latents"], "b t c ... -> t b c ...").contiguous()
+    actions = rearrange(batch["actions"], "b t d -> t b d").contiguous()
+    poses = rearrange(batch["poses"], "b t d -> t b d").contiguous()
+    frame_indices = rearrange(batch["frame_indices"], "b t -> t b").contiguous()
+
+    # Keep original image H/W for later ray geometry; latent H/W is not a substitute.
+    return {
+        "latents": latents,
+        "actions": actions,
+        "poses": poses,
+        "frame_indices": frame_indices,
+        "memory_segments": memory_segments,
+        "memory_masks": memory_masks,
+        "image_hw": batch["image_hw"],
+    }
+
+
 def euler_to_rotation_matrix(pitch, yaw):
     """
     Convert pitch and yaw angles (in radians) to a 3x3 rotation matrix.
@@ -417,7 +465,12 @@ class WorldMemMinecraft(DiffusionForcingBase):
         Returns:
             dict: A dictionary containing the training loss.
         """
-        xs, conditions, pose_conditions, c2w_mat, frame_idx = self._preprocess_batch(batch)
+        preprocessed = self._preprocess_batch(batch)
+        if isinstance(preprocessed, Mapping):
+            raise NotImplementedError(
+                "DeMemWM latent training_step needs Step 5.2/5.5 target/memory split before diffusion forward."
+            )
+        xs, conditions, pose_conditions, c2w_mat, frame_idx = preprocessed
 
         if self.use_plucker:
             if self.relative_embedding:
@@ -551,6 +604,8 @@ class WorldMemMinecraft(DiffusionForcingBase):
         self.validation_step_outputs.clear()
 
     def _preprocess_batch(self, batch):
+        if isinstance(batch, Mapping):
+            return _preprocess_dememwm_latent_batch(batch)
 
         xs, conditions, pose_conditions, frame_index = batch
 
@@ -715,7 +770,12 @@ class WorldMemMinecraft(DiffusionForcingBase):
         """
         # Preprocess the input batch
         memory_condition_length = self.memory_condition_length
-        xs_raw, conditions, pose_conditions, c2w_mat, frame_idx = self._preprocess_batch(batch)
+        preprocessed = self._preprocess_batch(batch)
+        if isinstance(preprocessed, Mapping):
+            raise NotImplementedError(
+                "DeMemWM latent validation needs Step 5.6 target-only sampling before evaluation."
+            )
+        xs_raw, conditions, pose_conditions, c2w_mat, frame_idx = preprocessed
 
 
         # Encode frames in chunks if necessary

tests/test_dememwm_latent_dataset.py

@@ -182,6 +182,22 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         self.assertTrue(sample["memory_masks"]["revisit"].all().item())
         self.assertEqual(sample["image_hw"].tolist(), [360, 640])
 
+        from torch.utils.data import default_collate
+        from algorithms.dememwm.df_video import _preprocess_dememwm_latent_batch
+
+        preprocessed = _preprocess_dememwm_latent_batch(default_collate([sample]))
+        self.assertEqual(tuple(preprocessed["latents"].shape), (9, 1, 1, 1, 2))
+        self.assertEqual(tuple(preprocessed["actions"].shape), (9, 1, 25))
+        self.assertEqual(tuple(preprocessed["poses"].shape), (9, 1, 5))
+        self.assertEqual(tuple(preprocessed["frame_indices"].shape), (9, 1))
+        self.assertEqual(preprocessed["memory_segments"], sample["memory_segments"])
+        self.assertTrue(all(isinstance(value, int) for value in preprocessed["memory_segments"].values()))
+        self.assertEqual(tuple(preprocessed["memory_masks"]["target"].shape), (1, 3))
+        self.assertEqual(tuple(preprocessed["memory_masks"]["anchor"].shape), (1, 2))
+        self.assertEqual(tuple(preprocessed["memory_masks"]["dynamic"].shape), (1, 2))
+        self.assertEqual(tuple(preprocessed["memory_masks"]["revisit"].shape), (1, 2))
+        self.assertEqual(preprocessed["image_hw"].tolist(), [[360, 640]])
+
 
 if __name__ == "__main__":
     unittest.main()