Remove DeMemWM raw training fallback

algorithms/dememwm/df_video.py

@@ -719,6 +719,12 @@ class DeMemWMMinecraft(DiffusionForcingBase):
         Returns:
             dict: A dictionary containing the training loss.
         """
+        if not isinstance(batch, Mapping):
+            raise TypeError(
+                "DeMemWM training requires the latent dict batch contract "
+                "from video_minecraft_dememwm_latent; raw tuple batches are unsupported."
+            )
+
         preprocessed = self._preprocess_batch(batch)
         if isinstance(preprocessed, Mapping):
             xs = preprocessed["latents"]
@@ -778,64 +784,10 @@ class DeMemWMMinecraft(DiffusionForcingBase):
                 self.log("training/loss", loss.cpu())
 
             return {"loss": loss}
-        xs, conditions, pose_conditions, c2w_mat, frame_idx = preprocessed
-
-        if self.use_plucker:
-            if self.relative_embedding:
-                input_pose_condition = []
-                frame_idx_list = []
-                for i in range(self.n_frames):
-                    input_pose_condition.append(
-                        convert_to_plucker(
-                            torch.cat([c2w_mat[i:i + 1], c2w_mat[-self.memory_condition_length:]]).clone(),
-                            0,
-                            focal_length=self.focal_length,
-                            image_height=xs.shape[-2],image_width=xs.shape[-1]
-                        ).to(xs.dtype)
-                    )
-                    frame_idx_list.append(
-                        torch.cat([
-                            frame_idx[i:i + 1] - frame_idx[i:i + 1],
-                            frame_idx[-self.memory_condition_length:] - frame_idx[i:i + 1]
-                        ]).clone()
-                    )
-                input_pose_condition = torch.cat(input_pose_condition)
-                frame_idx_list = torch.cat(frame_idx_list)
-            else:
-                input_pose_condition = convert_to_plucker(
-                    c2w_mat, 0, focal_length=self.focal_length
-                ).to(xs.dtype)
-                frame_idx_list = frame_idx
-        else:
-            input_pose_condition = pose_conditions.to(xs.dtype)
-            frame_idx_list = None
-
-        xs = self.encode(xs)
-
-        noise_levels = self._generate_noise_levels(xs)
-
-        if self.memory_condition_length:
-            noise_levels[-self.memory_condition_length:] = self.diffusion_model.stabilization_level
-            conditions[-self.memory_condition_length:] *= 0
-
-        _, loss = self.diffusion_model(
-            xs,
-            conditions,
-            input_pose_condition,
-            noise_levels=noise_levels,
-            reference_length=self.memory_condition_length,
-            frame_idx=frame_idx_list
+        raise TypeError(
+            "DeMemWM _preprocess_batch must return the latent dict contract; "
+            "legacy raw training tuples are unsupported."
         )
-
-        if self.memory_condition_length:
-            loss = loss[:-self.memory_condition_length]
-
-        loss = self.reweight_loss(loss, None)
-
-        if batch_idx % 20 == 0:
-            self.log("training/loss", loss.cpu())
-
-        return {"loss": loss}
     
     def on_validation_epoch_start(self) -> None:
         self._reset_metric_accumulators()

tests/test_dememwm_latent_dataset.py

@@ -302,6 +302,16 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         self.assertEqual(harness.logged[0][0], "training/loss")
         self.assertEqual(float(harness.logged[0][1]), 2.0)
 
+    def test_training_step_rejects_raw_tuple_batches(self):
+        from algorithms.dememwm.df_video import DeMemWMMinecraft
+
+        class Harness:
+            def _preprocess_batch(self, batch):
+                raise AssertionError("raw batches must be rejected before preprocessing")
+
+        with self.assertRaisesRegex(TypeError, "latent dict batch contract"):
+            DeMemWMMinecraft.training_step(Harness(), (None, None, None, None), 0)
+
     def test_memory_noise_helpers_route_shift_and_keep_validation_clean(self):
         from algorithms.dememwm.df_video import _apply_memory_route_masks, _memory_noise_levels_for_streams