Add DeMemWM image-space validation metrics

algorithms/dememwm/df_video.py

@@ -1156,6 +1156,34 @@ class DeMemWMMinecraft(DiffusionForcingBase):
                 latent_loss = xs_pred.new_tensor(0.0)
 
             self.log(f"{namespace}/latent_mse", latent_loss.detach())
+            if eval_start < target_length:
+                xs_pred_decode = self.decode(xs_pred[eval_start:target_length].to(target_poses.device))
+                xs_decode = self.decode(xs[eval_start:target_length].to(target_poses.device))
+
+                if self.logger and self.log_video:
+                    log_video(
+                        xs_pred_decode,
+                        xs_decode,
+                        step=batch_idx,
+                        namespace=namespace + "_vis",
+                        context_frames=self.context_frames,
+                        logger=self.logger.experiment,
+                        save_local=self.save_local,
+                        local_save_dir=self.local_save_dir,
+                    )
+
+                if target_mask is None:
+                    self._update_metric_accumulators(xs_pred_decode, xs_decode)
+                else:
+                    valid_mask = rearrange(
+                        target_mask[:, eval_start:target_length].to(device=xs_pred_decode.device),
+                        "b t -> t b",
+                    )
+                    if valid_mask.any():
+                        self._update_metric_accumulators(
+                            xs_pred_decode[valid_mask].unsqueeze(1),
+                            xs_decode[valid_mask].unsqueeze(1),
+                        )
             return latent_loss
 
     def test_step(self, batch, batch_idx) -> STEP_OUTPUT:

tests/test_dememwm_latent_dataset.py

@@ -538,6 +538,12 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
                 self.diffusion_model = FakeDiffusion()
                 self.logged = []
                 self.horizons = []
+                self.decoded_inputs = []
+                self.metric_updates = []
+                self.logger = type("Logger", (), {"experiment": object()})()
+                self.log_video = True
+                self.save_local = False
+                self.local_save_dir = None
 
             def _preprocess_batch(self, batch):
                 return _preprocess_dememwm_latent_batch(batch)
@@ -553,7 +559,11 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
                 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")
+                self.decoded_inputs.append(xs.detach().clone())
+                return xs + 10.0
+
+            def _update_metric_accumulators(self, xs_pred, xs_gt):
+                self.metric_updates.append((xs_pred.detach().clone(), xs_gt.detach().clone()))
 
             def log(self, name, value):
                 self.logged.append((name, value))
@@ -569,7 +579,14 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         }
         harness = Harness()
 
-        with mock.patch.object(df_video, "select_memory_indices", side_effect=fake_select_memory_indices):
+        video_calls = []
+
+        def fake_log_video(*args, **kwargs):
+            video_calls.append((args, kwargs))
+
+        with mock.patch.object(
+            df_video, "select_memory_indices", side_effect=fake_select_memory_indices
+        ), mock.patch.object(df_video, "log_video", side_effect=fake_log_video):
             loss = DeMemWMMinecraft.validation_step(harness, batch, 0, namespace="test")
         calls = harness.diffusion_model.calls
         required_kwargs = {"reference_length", "frame_memory_segments", "frame_memory_masks", "frame_memory_pose", "frame_idx", "image_hw"}
@@ -595,6 +612,17 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         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[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_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")
+        self.assertEqual(video_calls[0][1]["step"], 0)
 
     def test_interactive_generation_fails_until_packed_memory_path_exists(self):
         from algorithms.dememwm.df_video import DeMemWMMinecraft