Match DeMemWM online dynamic stream readout

algorithms/dememwm/df_video.py

@@ -20,8 +20,8 @@ from algorithms.common.metrics import (
 )
 from datasets.video.memory_selection import (
     _event_triggered_anchor_candidates_from_deltas,
-    _nearest_unique_event_anchors,
     _pose_delta_values,
+    _select_dynamic_from_stream,
     _select_anchor,
     _select_dynamic_by_policy,
     _select_revisit,
@@ -331,6 +331,7 @@ def _new_online_event_cache(capacity=None):
     return {
         "stop": 0,
         "capacity": None if capacity is None else max(0, int(capacity)),
+        "anchors": np.empty((0,), dtype=np.int64),
         "vectors": None,
         "d_vis": None,
         "d_pose": None,
@@ -403,7 +404,7 @@ def _append_online_event_values(cached, values, old_stop: int, stop: int, capaci
     return torch.cat([cached, values.to(device=cached.device)], dim=0)
 
 
-def _extend_online_event_cache(cache, latents, actions, poses, stop: int) -> None:
+def _extend_online_event_cache(cache, latents, actions, poses, stop: int, cfg=None) -> None:
     if latents is None:
         return
     old_stop = int(cache["stop"])
@@ -451,6 +452,18 @@ def _extend_online_event_cache(cache, latents, actions, poses, stop: int) -> Non
     cache["d_act"] = _append_online_event_values(cache["d_act"], d_act, old_stop, target_stop, capacity)
     cache["stop"] = target_stop
 
+    frames = np.arange(0, target_stop, dtype=np.int64)
+    event_anchors = _event_triggered_anchor_candidates_from_deltas(
+        frames,
+        cache["d_vis"][:target_stop],
+        cache["d_pose"][:target_stop],
+        cache["d_act"][:target_stop],
+        cfg,
+    )
+    stream = np.concatenate([np.asarray([0], dtype=np.int64), event_anchors.astype(np.int64, copy=False)])
+    stream = stream[(stream >= 0) & (stream < target_stop)]
+    cache["anchors"] = np.unique(np.sort(stream.astype(np.int64, copy=False)))
+
 
 def _select_online_event_dynamic_from_cache(
     cache,
@@ -458,7 +471,6 @@ def _select_online_event_dynamic_from_cache(
     count: int,
     cfg,
     reference_frames=None,
-    excluded=None,
 ) -> np.ndarray:
     if count <= 0:
         return np.empty((0,), dtype=np.int64)
@@ -471,24 +483,9 @@ def _select_online_event_dynamic_from_cache(
         return np.empty((0,), dtype=np.int64)
 
     stop = min(max(0, int(start_frame)), int(cache["stop"]))
-    if stop <= 0 or cache["d_vis"] is None:
-        return np.empty((0,), dtype=np.int64)
-
-    frames = np.arange(0, stop, dtype=np.int64)
-    event_anchors = _event_triggered_anchor_candidates_from_deltas(
-        frames,
-        cache["d_vis"][:stop],
-        cache["d_pose"][:stop],
-        cache["d_act"][:stop],
-        cfg,
-    )
-    if excluded is not None and len(excluded) > 0 and len(event_anchors) > 0:
-        event_anchors = event_anchors[~np.isin(event_anchors, np.asarray(excluded, dtype=np.int64))]
-    if len(event_anchors) == 0:
-        return np.empty((0,), dtype=np.int64)
-
-    references = np.asarray([int(start_frame)] if reference_frames is None or len(reference_frames) == 0 else reference_frames, dtype=np.int64)
-    return _nearest_unique_event_anchors(event_anchors, references, count)
+    anchors = np.asarray(cache.get("anchors", np.empty((0,), dtype=np.int64)), dtype=np.int64)
+    eligible_stream = anchors[(anchors >= 0) & (anchors < stop)]
+    return _select_dynamic_from_stream(eligible_stream, reference_frames, count)
 
 
 def _memory_noise_levels_for_streams(cfg, diffusion_model, query_noise_levels, stream_lengths, mode: str):
@@ -1570,7 +1567,8 @@ class DeMemWMMinecraft(DiffusionForcingBase):
                             source_latents[:, batch_i],
                             source_actions[:, batch_i],
                             source_poses[:, batch_i, :5],
-                            start_frame,
+                            curr_frame,
+                            memory_selection_cfg,
                         )
                         selected["dynamic"] = _select_online_event_dynamic_from_cache(
                             cache,
@@ -1578,11 +1576,6 @@ class DeMemWMMinecraft(DiffusionForcingBase):
                             dynamic_count,
                             memory_selection_cfg,
                             reference_frames=selected["revisit"] if len(selected["revisit"]) > 0 else None,
-                            excluded=(
-                                np.concatenate([selected["anchor"], selected["revisit"]])
-                                if len(selected["anchor"]) or len(selected["revisit"])
-                                else None
-                            ),
                         )
                     if memory_sheet_records is not None and int(target_positions[0]) in memory_sheet_targets:
                         memory_sheet_records[batch_i].append(