Implement DeMemWM online validation memory

algorithms/dememwm/df_video.py

@@ -17,6 +17,7 @@ from lightning.pytorch.utilities.types import STEP_OUTPUT
 from algorithms.common.metrics import (
     LearnedPerceptualImagePatchSimilarity,
 )
+from datasets.video.memory_selection import select_memory_indices
 from utils.logging_utils import log_video, get_validation_metrics_for_videos
 from .df_base import DiffusionForcingBase
 from .models.vae import VAE_models
@@ -151,6 +152,19 @@ def _preprocess_dememwm_latent_batch(batch):
     }
 
 
+def _gather_online_memory_tensor(source, indices, masks):
+    output_shape = (indices.shape[1], indices.shape[0], *source.shape[2:])
+    if output_shape[0] == 0 or source.shape[0] == 0:
+        return source.new_zeros(output_shape)
+
+    gather_idx = indices.clamp(0, source.shape[0] - 1).T.to(source.device)
+    batch_idx = torch.arange(indices.shape[0], device=source.device).expand_as(gather_idx)
+    gathered = source[gather_idx, batch_idx]
+    mask = masks.T.to(device=source.device, dtype=torch.bool)
+    mask = mask.view(*mask.shape, *((1,) * (gathered.ndim - 2)))
+    return gathered * mask.to(dtype=gathered.dtype)
+
+
 def euler_to_rotation_matrix(pitch, yaw):
     """
     Convert pitch and yaw angles (in radians) to a 3x3 rotation matrix.
@@ -904,76 +918,114 @@ class DeMemWMMinecraft(DiffusionForcingBase):
         memory_condition_length = self.memory_condition_length
         preprocessed = self._preprocess_batch(batch)
         if isinstance(preprocessed, Mapping):
-            xs = preprocessed["latents"]
+            target_tensors = preprocessed["target_tensors"]
+            xs = target_tensors["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()
-            }
+            target_mask = preprocessed["memory_masks"].get("target")
+            target_mask = None if target_mask is None else target_mask.to(device=xs.device)
+            target_poses = target_tensors["poses"].to(device=xs.device, dtype=xs.dtype)
+            target_frame_indices = target_tensors["frame_indices"].to(device=xs.device)
+            stream_lengths = preprocessed["stream_lengths"]
+            memory_selection_cfg = _cfg_get(self.cfg, "memory_selection", {})
+            poses_np = target_poses.detach().cpu().numpy()
 
             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,
-            )
+            n_context_frames = min(self.context_frames // self.frame_stack, target_length)
+            curr_frame = n_context_frames
 
-            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],
+            while curr_frame < target_length:
+                horizon = min(target_length - curr_frame, self.chunk_size) if self.chunk_size > 0 else target_length - curr_frame
+                target_slice = slice(curr_frame, curr_frame + horizon)
+                xs_pred[target_slice] = torch.randn_like(xs_pred[target_slice]).clamp(-self.clip_noise, self.clip_noise)
+
+                stream_indices = {
+                    key: torch.full((batch_size, int(stream_lengths[key])), -1, device=xs.device, dtype=torch.long)
+                    for key in _DEMEMWM_STREAM_KEYS
+                }
+                stream_masks = {key: torch.zeros_like(indices, dtype=torch.bool) for key, indices in stream_indices.items()}
+                target_positions = np.arange(target_slice.start, target_slice.stop, dtype=np.int64)
+                for batch_i in range(batch_size):
+                    selected, selected_masks = select_memory_indices(
+                        poses_np[:, batch_i, :5],
+                        target_positions,
+                        memory_selection_cfg,
+                        split=namespace,
+                    )
+                    for key, indices in stream_indices.items():
+                        count = indices.shape[1]
+                        selected_key = np.asarray(selected.get(key, []), dtype=np.int64)[:count]
+                        mask_key = np.asarray(selected_masks.get(key, []), dtype=bool)[:count]
+                        indices[batch_i, : len(selected_key)] = torch.as_tensor(selected_key, device=xs.device, dtype=torch.long)
+                        stream_masks[key][batch_i, : len(mask_key)] = torch.as_tensor(mask_key, device=xs.device, dtype=torch.bool)
+
+                source_latents = xs_pred[:curr_frame]
+                source_poses = target_poses[:curr_frame]
+                source_frame_indices = target_frame_indices[:curr_frame]
+                frame_memory_masks = {
+                    "target": torch.ones((batch_size, horizon), device=xs.device, dtype=torch.bool)
+                    if target_mask is None
+                    else target_mask[:, target_slice],
+                    **stream_masks,
+                }
+                stream_latents = [_gather_online_memory_tensor(source_latents, stream_indices[key], stream_masks[key]) for key in _DEMEMWM_STREAM_KEYS]
+                stream_poses = [_gather_online_memory_tensor(source_poses, stream_indices[key], stream_masks[key]) for key in _DEMEMWM_STREAM_KEYS]
+                stream_frame_indices = [_gather_online_memory_tensor(source_frame_indices, stream_indices[key], stream_masks[key]) for key in _DEMEMWM_STREAM_KEYS]
+                target_conditions = target_tensors["action_conditions"][target_slice].to(device=xs.device)
+                memory_length = sum(mask.shape[1] for mask in stream_masks.values())
+
+                packed_latents = torch.cat([xs_pred[target_slice], *stream_latents], dim=0)
+                packed_conditions = torch.cat(
+                    [target_conditions, target_conditions.new_zeros((memory_length, batch_size, target_conditions.shape[-1]))],
+                    dim=0,
+                )
+                frame_memory_pose = torch.cat([target_poses[target_slice], *stream_poses], dim=0)
+                frame_indices = torch.cat([target_frame_indices[target_slice], *stream_frame_indices], dim=0)
+                frame_memory_segments = {
+                    "target": int(horizon),
+                    **{key: int(stream_masks[key].shape[1]) for key in _DEMEMWM_STREAM_KEYS},
+                }
+                memory_noise_levels = torch.zeros(
+                    (packed_latents.shape[0] - horizon, batch_size),
                     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]
+                scheduling_matrix = self._generate_scheduling_matrix(horizon)
+
+                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)
+                    sampled_target = self.diffusion_model.sample_step(
+                        packed_latents,
+                        packed_conditions,
+                        None,
+                        torch.cat([from_target, memory_noise_levels], dim=0),
+                        torch.cat([to_target, memory_noise_levels], dim=0),
+                        current_frame=curr_frame,
+                        mode=namespace,
+                        reference_length=0,
+                        frame_idx=frame_indices,
+                        frame_memory_segments=frame_memory_segments,
+                        frame_memory_masks=frame_memory_masks,
+                        frame_memory_pose=frame_memory_pose,
+                        image_hw=image_hw,
+                    )
+                    packed_latents[:horizon] = sampled_target[:horizon]
 
-            latent_loss = F.mse_loss(
-                xs_pred[:target_length],
-                xs[:target_length],
-                reduction="none",
-            )
-            target_mask = frame_memory_masks.get("target")
+                xs_pred[target_slice] = packed_latents[:horizon]
+                curr_frame += horizon
+
+            eval_start = n_context_frames
+            latent_loss = F.mse_loss(xs_pred[eval_start:target_length], xs[eval_start:target_length], reduction="none")
             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:
+                loss_mask = rearrange(target_mask[:, eval_start:target_length].to(dtype=latent_loss.dtype), "b t -> t b")
+                loss_mask = loss_mask.view(*loss_mask.shape, *((1,) * (latent_loss.ndim - 2)))
+                latent_loss = (latent_loss * loss_mask).sum() / loss_mask.expand_as(latent_loss).sum().clamp_min(1.0)
+            elif latent_loss.numel():
                 latent_loss = latent_loss.mean()
+            else:
+                latent_loss = xs_pred.new_tensor(0.0)
 
             self.log(f"{namespace}/latent_mse", latent_loss.detach())
             return latent_loss

tests/test_dememwm_latent_dataset.py

@@ -1,5 +1,6 @@
 import tempfile
 import unittest
+from unittest import mock
 from pathlib import Path
 
 import numpy as np
@@ -300,9 +301,17 @@ 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):
+    def test_validation_step_builds_online_memory_from_committed_latents(self):
+        import algorithms.dememwm.df_video as df_video
         from algorithms.dememwm.df_video import DeMemWMMinecraft, _preprocess_dememwm_latent_batch
 
+        selection_calls = []
+
+        def fake_select_memory_indices(poses, target_positions, cfg=None, split="training", rng=None):
+            start = int(target_positions[0])
+            selection_calls.append((target_positions.tolist(), split, float(cfg.fov_overlap_threshold)))
+            return {"anchor": [0], "dynamic": [max(0, start - 1)], "revisit": [max(0, start - 2)]}, {"anchor": [start > 0], "dynamic": [start > 0], "revisit": [start > 2]}
+
         class FakeDiffusion:
             def __init__(self):
                 self.calls = []
@@ -318,16 +327,15 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
                         "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
+                return torch.full_like(x[: kwargs["frame_memory_segments"]["target"]], float(len(self.calls)))
 
         class Harness:
             def __init__(self):
-                self.memory_condition_length = 3
-                self.clip_noise = 0.0
+                self.memory_condition_length, self.clip_noise = 3, 0.0
+                self.context_frames, self.frame_stack, self.chunk_size = 2, 1, 2
+                self.cfg = OmegaConf.create(
+                    {"memory_selection": {"enabled": True, "max_anchor_frames": 1, "max_dynamic_frames": 1, "max_revisit_frames": 1, "fov_overlap_threshold": 0.75}}
+                )
                 self.diffusion_model = FakeDiffusion()
                 self.logged = []
                 self.horizons = []
@@ -337,10 +345,10 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
 
             def _generate_scheduling_matrix(self, horizon):
                 self.horizons.append(horizon)
-                return np.array([[2, 2], [1, 1], [0, 0]], dtype=np.int64)
+                return np.stack([np.full(horizon, level, dtype=np.int64) for level in (2, 1, 0)])
 
             def _generate_condition_indices(self, *args, **kwargs):
-                raise AssertionError("latent dict validation must not reselect memory")
+                raise AssertionError("latent dict validation must not use the legacy memory_condition_length path")
 
             def encode(self, xs):
                 raise AssertionError("latent dict validation must not VAE-encode inputs")
@@ -352,51 +360,33 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
                 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]]),
-            },
+            "latents": torch.tensor([0, 1, 2, 3, 4, 100, 101, 102], dtype=torch.float32).view(1, 8, 1, 1, 1),
+            "actions": (100 + torch.arange(8, dtype=torch.float32)).view(1, 8, 1),
+            "poses": (200 + torch.arange(8, dtype=torch.float32)).view(1, 8, 1).repeat(1, 1, 5),
+            "frame_indices": torch.tensor([[10, 11, 12, 13, 14, 900, 901, 902]], dtype=torch.long),
+            "memory_segments": {"target": torch.tensor([5]), "anchor": torch.tensor([1]), "dynamic": torch.tensor([1]), "revisit": torch.tensor([1])},
+            "memory_masks": {"target": torch.ones((1, 5), dtype=torch.bool), "anchor": torch.ones((1, 1), dtype=torch.bool), "dynamic": torch.ones((1, 1), dtype=torch.bool), "revisit": torch.ones((1, 1), dtype=torch.bool)},
             "image_hw": torch.tensor([[360, 640]], dtype=torch.long),
         }
         harness = Harness()
 
-        loss = DeMemWMMinecraft.validation_step(harness, batch, 0, namespace="test")
+        with mock.patch.object(df_video, "select_memory_indices", side_effect=fake_select_memory_indices):
+            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"}
 
-        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)
+        self.assertEqual((harness.horizons, selection_calls, len(calls)), ([2, 1], [([2, 3], "test", 0.75), ([4], "test", 0.75)], 4))
+        self.assertTrue(all(required_kwargs <= call["kwargs"].keys() for call in calls))
+        self.assertTrue(all(call["kwargs"]["reference_length"] == 0 for call in calls))
+        self.assertEqual(
+            ([call["kwargs"]["current_frame"] for call in calls[::2]], [call["kwargs"]["frame_memory_segments"]["target"] for call in calls[::2]]),
+            ([2, 4], [2, 1]),
+        )
+        self.assertEqual(calls[0]["x"][:, 0, 0, 0, 0].tolist(), [0.0, 0.0, 0.0, 1.0, 0.0])
+        self.assertEqual(calls[0]["action_cond"][:, 0, 0].tolist(), [102.0, 103.0, 0.0, 0.0, 0.0])
+        self.assertEqual(calls[0]["kwargs"]["frame_idx"][:, 0].tolist(), [12, 13, 10, 11, 0])
+        self.assertEqual(calls[0]["kwargs"]["frame_memory_masks"]["revisit"].tolist(), [[False]])
+        self.assertAlmostEqual(float(loss), 1.0 / 3.0)
 
     def test_dataset_returns_target_anchor_dynamic_revisit_contract(self):
         with tempfile.TemporaryDirectory() as tmp: