Read DeMemWM training dynamic stream around revisit

datasets/video/memory_selection.py

@@ -606,6 +606,98 @@ def _nearest_unique_event_anchors(event_anchors: np.ndarray, reference_frames: n
     return np.sort(np.asarray(selected[:count], dtype=np.int64))
 
 
+def _build_dynamic_stream(
+    latents,
+    actions=None,
+    poses=None,
+    cfg=None,
+    min_candidate_frame: int = 0,
+    stop: int | None = None,
+) -> np.ndarray:
+    if latents is None:
+        return np.empty((0,), dtype=np.int64)
+
+    length = len(latents)
+    min_candidate_frame = max(0, int(min_candidate_frame))
+    target_stop = length if stop is None else max(0, min(int(stop), length))
+    if target_stop <= min_candidate_frame:
+        return np.empty((0,), dtype=np.int64)
+
+    event_anchors = _event_triggered_anchor_candidates(
+        target_stop,
+        cfg,
+        poses=poses,
+        latents=latents,
+        actions=actions,
+        min_candidate_frame=min_candidate_frame,
+    )
+    stream = np.concatenate([np.asarray([min_candidate_frame], dtype=np.int64), event_anchors.astype(np.int64, copy=False)])
+    stream = stream[(stream >= min_candidate_frame) & (stream < target_stop)]
+    return np.unique(np.sort(stream.astype(np.int64, copy=False)))
+
+
+def _select_dynamic_from_stream(dynamic_stream, reference_frames, count: int) -> np.ndarray:
+    if count <= 0:
+        return np.empty((0,), dtype=np.int64)
+
+    stream = np.asarray(dynamic_stream, dtype=np.int64)
+    stream = np.unique(np.sort(stream[stream >= 0]))
+    if len(stream) == 0:
+        return np.empty((0,), dtype=np.int64)
+
+    references = np.asarray([] if reference_frames is None else reference_frames, dtype=np.int64)
+    references = np.unique(np.sort(references[references >= 0]))
+    if len(references) == 0:
+        return stream[-count:].astype(np.int64, copy=False)
+
+    selected = []
+    used = set()
+
+    def add_candidate(value) -> bool:
+        value = int(value)
+        if value in used:
+            return False
+        selected.append(value)
+        used.add(value)
+        return True
+
+    base_quota, extra = divmod(int(count), len(references))
+    for ref_idx, ref in enumerate(references):
+        quota = base_quota + (1 if ref_idx < extra else 0)
+        if quota <= 0:
+            continue
+        picked = 0
+        before = stream[stream <= int(ref)]
+        after = stream[stream >= int(ref)]
+        brackets = []
+        if len(before) > 0:
+            brackets.append(int(before[-1]))
+        if len(after) > 0:
+            brackets.append(int(after[0]))
+        for candidate in brackets:
+            picked += int(add_candidate(candidate))
+            if picked >= quota:
+                break
+        if picked >= quota:
+            continue
+
+        order = np.lexsort((stream, np.abs(stream - int(ref))))
+        for row in order:
+            picked += int(add_candidate(stream[row]))
+            if picked >= quota:
+                break
+
+    if len(selected) < count:
+        distance = np.min(np.abs(stream[:, None] - references[None, :]), axis=1)
+        order = np.lexsort((stream, distance))
+        for row in order:
+            add_candidate(stream[row])
+            if len(selected) >= count:
+                break
+
+    return np.sort(np.asarray(selected[:count], dtype=np.int64))
+
+
 def _select_event_triggered_dynamic(
     target_start: int,
     count: int,
@@ -724,6 +816,7 @@ def select_memory_indices(
     actions=None,
     anchor_candidate_start=None,
     anchor_candidate_stop=None,
+    dynamic_stream=None,
 ):
     """Select causal memory indices and masks for [anchor][dynamic][revisit]."""
 
@@ -783,17 +876,24 @@ def select_memory_indices(
     )
 
     if policy == "event_triggered":
-        dynamic = _select_dynamic_by_policy(
-            target_start,
-            counts["dynamic"],
-            cfg,
-            poses=poses,
-            latents=latents,
-            actions=actions,
-            min_candidate_frame=min_candidate_frame,
-            reference_frames=revisit if len(revisit) > 0 else None,
-            excluded=np.concatenate([anchor, revisit]) if len(anchor) or len(revisit) else None,
-        )
+        if dynamic_stream is not None:
+            exclusion = max(0, int(cfg_get(cfg, "local_context_exclusion_frames", 8)))
+            causal_cutoff = max(min_candidate_frame, target_start - exclusion)
+            stream = np.asarray(dynamic_stream, dtype=np.int64)
+            eligible_stream = stream[(stream >= min_candidate_frame) & (stream < causal_cutoff)]
+            dynamic = _select_dynamic_from_stream(eligible_stream, revisit, counts["dynamic"])
+        else:
+            dynamic = _select_dynamic_by_policy(
+                target_start,
+                counts["dynamic"],
+                cfg,
+                poses=poses,
+                latents=latents,
+                actions=actions,
+                min_candidate_frame=min_candidate_frame,
+                reference_frames=revisit if len(revisit) > 0 else None,
+                excluded=np.concatenate([anchor, revisit]) if len(anchor) or len(revisit) else None,
+            )
 
     _write_segment(indices, masks, "anchor", anchor)
     _write_segment(indices, masks, "dynamic", dynamic)

datasets/video/minecraft_video_dememwm_latent_dataset.py

@@ -10,7 +10,7 @@ import numpy as np
 import torch
 from omegaconf import DictConfig
 
-from .memory_selection import SEGMENT_KEYS, cfg_get, memory_segment_lengths, select_memory_indices
+from .memory_selection import SEGMENT_KEYS, _build_dynamic_stream, cfg_get, memory_segment_lengths, select_memory_indices
 
 
 ACTION_KEYS = [
@@ -210,6 +210,7 @@ class MinecraftVideoDeMemWMLatentDataset(torch.utils.data.Dataset):
             actions=actions,
             anchor_candidate_start=anchor_start,
             anchor_candidate_stop=anchor_stop,
+            dynamic_stream=arrays.get("dynamic_stream"),
         )
         ordered_memory = [memory_indices[key] for key in SEGMENT_KEYS]
         source_positions = np.concatenate([target_positions, *ordered_memory]).astype(np.int64)
@@ -292,10 +293,24 @@ class MinecraftVideoDeMemWMLatentDataset(torch.utils.data.Dataset):
             actions = self._normalize_actions(action_path, np.asarray(data["actions"]))
             poses = self._sanitize_poses(action_path, np.asarray(data["poses"], dtype=np.float32), len(actions))
 
+        dynamic_cfg = cfg_get(self.memory_selection, "dynamic", {})
+        dynamic_policy = str(cfg_get(dynamic_cfg, "selection_policy", "recent"))
+        if self.split == "training" and dynamic_policy == "event_triggered":
+            dynamic_stream = _build_dynamic_stream(
+                latents,
+                actions=actions,
+                poses=poses,
+                cfg=self.memory_selection,
+                min_candidate_frame=self.initial_frame_offset,
+            )
+        else:
+            dynamic_stream = np.empty((0,), dtype=np.int64)
+
         return {
             "latents": latents,
             "actions": actions,
             "poses": poses,
+            "dynamic_stream": dynamic_stream,
             "frame_indices": np.arange(len(latents), dtype=np.int64),
             "image_hw": self._feature_image_hw(meta_path),
         }