Add DeMemWM latent validation path

algorithms/dememwm/df_video.py

@@ -880,9 +880,79 @@ class WorldMemMinecraft(DiffusionForcingBase):
         memory_condition_length = self.memory_condition_length
         preprocessed = self._preprocess_batch(batch)
         if isinstance(preprocessed, Mapping):
-            raise NotImplementedError(
-                "DeMemWM latent validation needs Step 5.6 target-only sampling before evaluation."
+            xs = preprocessed["latents"]
+            target_length = preprocessed["target_length"]
+            batch_size = xs.shape[1]
+            conditions = preprocessed["action_conditions"].to(device=xs.device)
+            frame_indices = preprocessed["frame_indices"].to(device=xs.device)
+            frame_memory_pose = preprocessed["frame_memory_pose"].to(device=xs.device, dtype=xs.dtype)
+            image_hw = preprocessed["image_hw"].to(device=xs.device)
+            frame_memory_masks = {
+                key: mask.to(device=xs.device)
+                for key, mask in preprocessed["memory_masks"].items()
+            }
+
+            xs_pred = xs.clone()
+            # Single-window latent validation: only target frames are sampled;
+            # appended memory latents stay exactly as provided by the dataset.
+            xs_pred[:target_length] = torch.randn_like(xs_pred[:target_length]).clamp(
+                -self.clip_noise,
+                self.clip_noise,
+            )
+            scheduling_matrix = self._generate_scheduling_matrix(target_length)
+            memory_length = xs.shape[0] - target_length
+            memory_noise_levels = torch.zeros(
+                (memory_length, batch_size),
+                device=xs.device,
+                dtype=torch.long,
             )
+
+            for m in range(scheduling_matrix.shape[0] - 1):
+                from_target = torch.as_tensor(
+                    scheduling_matrix[m],
+                    device=xs.device,
+                    dtype=torch.long,
+                )[:, None].repeat(1, batch_size)
+                to_target = torch.as_tensor(
+                    scheduling_matrix[m + 1],
+                    device=xs.device,
+                    dtype=torch.long,
+                )[:, None].repeat(1, batch_size)
+                from_noise_levels = torch.cat([from_target, memory_noise_levels], dim=0)
+                to_noise_levels = torch.cat([to_target, memory_noise_levels], dim=0)
+
+                sampled_target = self.diffusion_model.sample_step(
+                    xs_pred,
+                    conditions,
+                    None,
+                    from_noise_levels,
+                    to_noise_levels,
+                    current_frame=0,
+                    mode=namespace,
+                    reference_length=0,
+                    frame_idx=frame_indices,
+                    frame_memory_segments=preprocessed["memory_segments"],
+                    frame_memory_masks=frame_memory_masks,
+                    frame_memory_pose=frame_memory_pose,
+                    image_hw=image_hw,
+                )
+                xs_pred[:target_length] = sampled_target[:target_length]
+
+            latent_loss = F.mse_loss(
+                xs_pred[:target_length],
+                xs[:target_length],
+                reduction="none",
+            )
+            target_mask = frame_memory_masks.get("target")
+            if target_mask is not None:
+                target_mask = rearrange(target_mask.to(dtype=latent_loss.dtype), "b t -> t b")
+                target_mask = target_mask.view(*target_mask.shape, *((1,) * (latent_loss.ndim - 2)))
+                latent_loss = (latent_loss * target_mask).sum() / target_mask.expand_as(latent_loss).sum().clamp_min(1.0)
+            else:
+                latent_loss = latent_loss.mean()
+
+            self.log(f"{namespace}/latent_mse", latent_loss.detach())
+            return latent_loss
         xs_raw, conditions, pose_conditions, c2w_mat, frame_idx = preprocessed
 
 

tests/test_dememwm_latent_dataset.py

@@ -293,6 +293,104 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         self.assertEqual(harness.logged[0][0], "training/loss")
         self.assertEqual(float(harness.logged[0][1]), 2.0)
 
+    def test_validation_step_uses_latent_dict_target_sampling_and_metric(self):
+        from algorithms.dememwm.df_video import WorldMemMinecraft, _preprocess_dememwm_latent_batch
+
+        class FakeDiffusion:
+            def __init__(self):
+                self.calls = []
+
+            def sample_step(self, x, action_cond, pose_cond, curr_noise_level, next_noise_level, **kwargs):
+                self.calls.append(
+                    {
+                        "x": x.detach().clone(),
+                        "action_cond": action_cond.detach().clone(),
+                        "pose_cond": pose_cond,
+                        "curr_noise_level": curr_noise_level.detach().clone(),
+                        "next_noise_level": next_noise_level.detach().clone(),
+                        "kwargs": kwargs,
+                    }
+                )
+                target_length = kwargs["frame_memory_segments"]["target"]
+                sampled = x.clone()
+                sampled[:target_length] = float(len(self.calls))
+                sampled[target_length:] = 999.0
+                return sampled
+
+        class Harness:
+            def __init__(self):
+                self.memory_condition_length = 3
+                self.clip_noise = 0.0
+                self.diffusion_model = FakeDiffusion()
+                self.logged = []
+                self.horizons = []
+
+            def _preprocess_batch(self, batch):
+                return _preprocess_dememwm_latent_batch(batch)
+
+            def _generate_scheduling_matrix(self, horizon):
+                self.horizons.append(horizon)
+                return np.array([[2, 2], [1, 1], [0, 0]], dtype=np.int64)
+
+            def _generate_condition_indices(self, *args, **kwargs):
+                raise AssertionError("latent dict validation must not reselect memory")
+
+            def encode(self, xs):
+                raise AssertionError("latent dict validation must not VAE-encode inputs")
+
+            def decode(self, xs):
+                raise AssertionError("latent dict validation must not decode for the latent metric")
+
+            def log(self, name, value):
+                self.logged.append((name, value))
+
+        batch = {
+            "latents": torch.arange(5, dtype=torch.float32).view(1, 5, 1, 1, 1),
+            "actions": (100 + torch.arange(5, dtype=torch.float32)).view(1, 5, 1),
+            "poses": (200 + torch.arange(5, dtype=torch.float32)).view(1, 5, 1).repeat(1, 1, 5),
+            "frame_indices": (10 + torch.arange(5, dtype=torch.long)).view(1, 5),
+            "memory_segments": {
+                "target": torch.tensor([2]),
+                "anchor": torch.tensor([1]),
+                "dynamic": torch.tensor([1]),
+                "revisit": torch.tensor([1]),
+            },
+            "memory_masks": {
+                "target": torch.tensor([[True, False]]),
+                "anchor": torch.tensor([[True]]),
+                "dynamic": torch.tensor([[False]]),
+                "revisit": torch.tensor([[True]]),
+            },
+            "image_hw": torch.tensor([[360, 640]], dtype=torch.long),
+        }
+        harness = Harness()
+
+        loss = WorldMemMinecraft.validation_step(harness, batch, 0, namespace="test")
+        calls = harness.diffusion_model.calls
+
+        self.assertEqual(harness.horizons, [2])
+        self.assertEqual(len(calls), 2)
+        for call in calls:
+            self.assertEqual(tuple(call["x"].shape), (5, 1, 1, 1, 1))
+            self.assertEqual(call["x"][2:, 0, 0, 0, 0].tolist(), [2.0, 3.0, 4.0])
+            self.assertEqual(call["action_cond"][:, 0, 0].tolist(), [0.0, 101.0, 0.0, 0.0, 0.0])
+            self.assertIsNone(call["pose_cond"])
+            self.assertEqual(call["kwargs"]["reference_length"], 0)
+            self.assertEqual(call["kwargs"]["current_frame"], 0)
+            self.assertEqual(call["kwargs"]["mode"], "test")
+            self.assertEqual(call["kwargs"]["frame_idx"][:, 0].tolist(), [10, 11, 12, 13, 14])
+            self.assertEqual(call["kwargs"]["frame_memory_segments"], {"target": 2, "anchor": 1, "dynamic": 1, "revisit": 1})
+            self.assertEqual(call["kwargs"]["frame_memory_pose"][:, 0, 0].tolist(), [200.0, 201.0, 202.0, 203.0, 204.0])
+            self.assertEqual(call["kwargs"]["frame_memory_masks"]["dynamic"].tolist(), [[False]])
+            self.assertEqual(call["kwargs"]["image_hw"].tolist(), [[360, 640]])
+        self.assertEqual(calls[0]["curr_noise_level"][:, 0].tolist(), [2, 2, 0, 0, 0])
+        self.assertEqual(calls[0]["next_noise_level"][:, 0].tolist(), [1, 1, 0, 0, 0])
+        self.assertEqual(calls[1]["curr_noise_level"][:, 0].tolist(), [1, 1, 0, 0, 0])
+        self.assertEqual(calls[1]["next_noise_level"][:, 0].tolist(), [0, 0, 0, 0, 0])
+        self.assertEqual(float(loss), 4.0)
+        self.assertEqual(harness.logged[0][0], "test/latent_mse")
+        self.assertEqual(float(harness.logged[0][1]), 4.0)
+
     def test_dataset_returns_target_anchor_dynamic_revisit_contract(self):
         with tempfile.TemporaryDirectory() as tmp:
             root = Path(tmp)