Add DeMemWM latent training step

algorithms/dememwm/df_video.py

@@ -528,9 +528,56 @@ class WorldMemMinecraft(DiffusionForcingBase):
         """
         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 = preprocessed["latents"]
+            target_length = preprocessed["target_length"]
+            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()
+            }
+
+            target_noise_levels = self._generate_noise_levels(xs[:target_length])
+            memory_noise_levels = [
+                torch.full(
+                    (preprocessed["stream_lengths"][key], xs.shape[1]),
+                    self.diffusion_model.stabilization_level,
+                    device=xs.device,
+                    dtype=target_noise_levels.dtype,
+                )
+                for key in _DEMEMWM_STREAM_KEYS
+            ]
+            # Appended memory stays at a fixed stabilization level for the first
+            # latent smoke path; routed/SNR-shifted policies belong in a later step.
+            noise_levels = torch.cat([target_noise_levels, *memory_noise_levels], dim=0)
+
+            _, loss = self.diffusion_model(
+                xs,
+                conditions,
+                None,
+                noise_levels=noise_levels,
+                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,
             )
+            loss = loss[:target_length]
+            target_mask = frame_memory_masks.get("target")
+            if target_mask is not None:
+                target_mask = rearrange(target_mask.to(dtype=loss.dtype), "b t -> t b")
+                target_mask = target_mask.view(*target_mask.shape, *((1,) * (loss.ndim - 2)))
+                loss = (loss * target_mask).sum() / target_mask.expand_as(loss).sum().clamp_min(1.0)
+            else:
+                loss = self.reweight_loss(loss, None)
+
+            if batch_idx % 20 == 0:
+                self.log("training/loss", loss.cpu())
+
+            return {"loss": loss}
         xs, conditions, pose_conditions, c2w_mat, frame_idx = preprocessed
 
         if self.use_plucker:

tests/test_dememwm_latent_dataset.py

@@ -218,6 +218,81 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         self.assertEqual(tuple(preprocessed["image_hw"].shape), (2, 2))
         self.assertEqual(preprocessed["image_hw"].tolist(), [[720, 1280], [720, 1280]])
 
+    def test_training_step_uses_latent_dict_target_loss_and_memory_metadata(self):
+        from algorithms.dememwm.df_video import WorldMemMinecraft, _preprocess_dememwm_latent_batch
+
+        class FakeDiffusion:
+            stabilization_level = 7
+
+            def __call__(self, x, action_cond, pose_cond, **kwargs):
+                self.x = x
+                self.action_cond = action_cond
+                self.pose_cond = pose_cond
+                self.kwargs = kwargs
+                loss = (torch.arange(x.shape[0], device=x.device, dtype=x.dtype) + 2).view(-1, 1, 1, 1, 1)
+                return x, loss.expand_as(x).clone()
+
+        class Harness:
+            def __init__(self):
+                self.diffusion_model = FakeDiffusion()
+                self.logged = []
+
+            def _preprocess_batch(self, batch):
+                return _preprocess_dememwm_latent_batch(batch)
+
+            def _generate_noise_levels(self, xs):
+                self.noise_input_shape = tuple(xs.shape)
+                return torch.arange(xs.shape[0], device=xs.device, dtype=torch.long).view(-1, 1)
+
+            def encode(self, xs):
+                raise AssertionError("latent dict training must not VAE-encode inputs")
+
+            def reweight_loss(self, loss, weight=None):
+                raise AssertionError("target mask path should reduce the latent dict loss directly")
+
+            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()
+
+        output = WorldMemMinecraft.training_step(harness, batch, 0)
+        call = harness.diffusion_model
+
+        self.assertEqual(harness.noise_input_shape, (2, 1, 1, 1, 1))
+        self.assertTrue(torch.equal(call.x[:, 0, 0, 0, 0], torch.arange(5, dtype=torch.float32)))
+        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["frame_memory_segments"], {"target": 2, "anchor": 1, "dynamic": 1, "revisit": 1})
+        self.assertEqual(call.kwargs["noise_levels"][:, 0].tolist(), [0, 1, 7, 7, 7])
+        self.assertEqual(call.kwargs["frame_idx"][:, 0].tolist(), [10, 11, 12, 13, 14])
+        self.assertEqual(call.kwargs["frame_memory_pose"][:, 0, 0].tolist(), [200.0, 201.0, 202.0, 203.0, 204.0])
+        self.assertEqual(call.kwargs["image_hw"].tolist(), [[360, 640]])
+        self.assertEqual(call.kwargs["frame_memory_masks"]["target"].device, call.x.device)
+        self.assertEqual(call.kwargs["frame_memory_masks"]["dynamic"].tolist(), [[False]])
+        self.assertEqual(float(output["loss"]), 2.0)
+        self.assertEqual(harness.logged[0][0], "training/loss")
+        self.assertEqual(float(harness.logged[0][1]), 2.0)
+
     def test_dataset_returns_target_anchor_dynamic_revisit_contract(self):
         with tempfile.TemporaryDirectory() as tmp:
             root = Path(tmp)