Add DeMemWM dynamic selector diagnostics

.exp_artifact/dememwm_revisit_dynamic_selector_diagnostic.py

@@ -0,0 +1,363 @@
+#!/usr/bin/env python3
+"""Generate selector diagnostics for the DeMemWM dynamic-stream fix."""
+
+from __future__ import annotations
+
+import csv
+import sys
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+import numpy as np
+from PIL import Image, ImageDraw
+
+from datasets.video.memory_selection import (
+    SEGMENT_KEYS,
+    _build_dynamic_stream,
+    _select_dynamic_from_stream,
+    select_memory_indices,
+)
+from datasets.video.minecraft_video_dememwm_latent_dataset import (
+    MinecraftVideoDeMemWMLatentDataset,
+    _ACTION_DIM,
+    convert_action_space,
+)
+
+DATA_ROOT = Path("/share_1/users/bonan_ding/worldmem_data/minecraft")
+INPUT_CSV = REPO_ROOT / ".exp_artifact" / "worldmem_vs_dememwm_memory_samples.csv"
+OUTPUT_DIR = REPO_ROOT / ".exp_artifact" / "dememwm_revisit_dynamic_selector_diagnostic"
+
+INITIAL_FRAME_OFFSET = 100
+CONTEXT_LENGTH = 100
+N_FRAMES = 8
+
+MEMORY_SELECTION = {
+    "enabled": True,
+    "max_anchor_frames": 2,
+    "max_dynamic_frames": 4,
+    "max_revisit_frames": 2,
+    "pose_similarity_threshold": 0.6,
+    "training_use_plucker": True,
+    "training_plucker_weight": 1.0,
+    "fov_overlap_threshold": 0.6,
+    "min_total_selected_coverage": 0.1,
+    "local_context_exclusion_frames": 8,
+    "plucker_moment_radius": 30.0,
+    "anchor_diverse_selection": True,
+    "pose_preselect_topk": 64,
+    "candidate_chunk_size": 64,
+    "dynamic": {
+        "selection_policy": "event_triggered",
+        "scene_threshold": 2.5,
+        "state_threshold": 2.5,
+        "stable_threshold": 1.0,
+        "stable_frames": 3,
+        "min_event_gap": 8,
+        "min_anchor_score": 2.0,
+        "max_event_anchors": None,
+        "b_pose": 0.3,
+        "b_action": 0.2,
+    },
+}
+
+
+def _feature_path(video_path: Path) -> Path:
+    relative = video_path.relative_to(DATA_ROOT)
+    return DATA_ROOT / "vae_features" / relative.parent / f"{relative.stem}_vae_feature.npy"
+
+
+def _load_video_arrays(relative_video: str):
+    video_path = DATA_ROOT / relative_video
+    feature_path = _feature_path(video_path)
+    action_path = video_path.with_suffix(".npz")
+    latents = np.load(feature_path, mmap_mode="r")
+    with np.load(action_path, allow_pickle=False) as data:
+        raw_actions = np.asarray(data["actions"])
+        if raw_actions.ndim == 2 and raw_actions.shape[1] == _ACTION_DIM:
+            actions = raw_actions.astype(np.float32, copy=False)
+        else:
+            actions = convert_action_space(raw_actions).numpy()
+        poses = MinecraftVideoDeMemWMLatentDataset._sanitize_poses(
+            action_path,
+            np.asarray(data["poses"], dtype=np.float32),
+            len(actions),
+        )
+    return video_path, latents, actions, poses
+
+
+def _valid(indices, masks, key: str) -> list[int]:
+    return [int(x) for x in np.asarray(indices[key], dtype=np.int64)[np.asarray(masks[key], dtype=bool)]]
+
+
+def _rel(values: list[int], target: int) -> list[int]:
+    return [int(value) - int(target) for value in values]
+
+
+def _nearest_revisit_delta(frame: int, revisit: list[int]) -> int | None:
+    if not revisit:
+        return None
+    revisit_np = np.asarray(revisit, dtype=np.int64)
+    nearest = int(revisit_np[np.argmin(np.abs(revisit_np - int(frame)))])
+    return int(frame) - nearest
+
+
+def _stream_ranks(stream: np.ndarray, frames: list[int]) -> list[int]:
+    ranks = []
+    for frame in frames:
+        hits = np.nonzero(stream == int(frame))[0]
+        ranks.append(int(hits[0]) if len(hits) else -1)
+    return ranks
+
+
+def _stream_window(eligible_stream: np.ndarray, revisit: list[int], count: int = 8) -> list[int]:
+    if len(eligible_stream) == 0:
+        return []
+    if not revisit:
+        return [int(x) for x in eligible_stream[-count:]]
+    refs = np.asarray(revisit, dtype=np.int64)
+    distance = np.min(np.abs(eligible_stream[:, None] - refs[None, :]), axis=1)
+    order = np.lexsort((eligible_stream, distance))[:count]
+    return [int(x) for x in np.sort(eligible_stream[order])]
+
+
+def _list_text(values) -> str:
+    return "[" + ", ".join("" if value is None else str(value) for value in values) + "]"
+
+
+def _read_samples() -> list[dict]:
+    with INPUT_CSV.open("r", newline="", encoding="utf-8") as handle:
+        return list(csv.DictReader(handle))
+
+
+
+def _analyze_sample(row: dict) -> tuple[dict, dict]:
+    sample = int(row["sample"])
+    relative_video = row["video"]
+    target_raw = int(row["worldmem_at_demem_target_start"])
+    clip_offset = int(row["clip_offset"])
+    video_path, latents, actions, poses = _load_video_arrays(relative_video)
+
+    dynamic_stream = _build_dynamic_stream(
+        latents,
+        actions=actions,
+        poses=poses,
+        cfg=MEMORY_SELECTION,
+        min_candidate_frame=INITIAL_FRAME_OFFSET,
+    )
+    target_positions = target_raw + np.arange(N_FRAMES, dtype=np.int64)
+    anchor_start = max(INITIAL_FRAME_OFFSET, target_raw - CONTEXT_LENGTH)
+    indices, masks = select_memory_indices(
+        poses,
+        target_positions,
+        MEMORY_SELECTION,
+        split="training",
+        min_candidate_frame=INITIAL_FRAME_OFFSET,
+        latents=latents,
+        actions=actions,
+        anchor_candidate_start=anchor_start,
+        anchor_candidate_stop=target_raw,
+        dynamic_stream=dynamic_stream,
+    )
+
+    selected = {key: _valid(indices, masks, key) for key in SEGMENT_KEYS}
+    exclusion = int(MEMORY_SELECTION["local_context_exclusion_frames"])
+    causal_cutoff = target_raw - exclusion
+    eligible_stream = dynamic_stream[(dynamic_stream >= INITIAL_FRAME_OFFSET) & (dynamic_stream < causal_cutoff)]
+    stream_window = _stream_window(eligible_stream, selected["revisit"], count=8)
+    dynamic_revisit_rel = [_nearest_revisit_delta(frame, selected["revisit"]) for frame in selected["dynamic"]]
+    dynamic_stream_rank = _stream_ranks(dynamic_stream, selected["dynamic"])
+    stream_flags = {
+        key: {
+            "causal": bool(all(frame < target_raw for frame in selected[key])),
+            "outside": bool(all(frame < causal_cutoff for frame in selected[key])),
+        }
+        for key in SEGMENT_KEYS
+    }
+
+    record = {
+        "sample": sample,
+        "video": relative_video,
+        "clip_offset": clip_offset,
+        "target_raw": target_raw,
+        "anchor_candidate_window": f"{anchor_start}..{target_raw - 1}",
+        "causal_cutoff": causal_cutoff,
+        "anchor_raw": _list_text(selected["anchor"]),
+        "anchor_rel": _list_text(_rel(selected["anchor"], target_raw)),
+        "revisit_raw": _list_text(selected["revisit"]),
+        "revisit_rel": _list_text(_rel(selected["revisit"], target_raw)),
+        "eligible_stream_count": int(len(eligible_stream)),
+        "stream_window_raw": _list_text(stream_window),
+        "dynamic_raw": _list_text(selected["dynamic"]),
+        "dynamic_rel": _list_text(_rel(selected["dynamic"], target_raw)),
+        "dynamic_rel_to_nearest_revisit": _list_text(dynamic_revisit_rel),
+        "dynamic_stream_rank": _list_text(dynamic_stream_rank),
+        "anchor_causal": stream_flags["anchor"]["causal"],
+        "anchor_outside_local_exclusion": stream_flags["anchor"]["outside"],
+        "revisit_causal": stream_flags["revisit"]["causal"],
+        "revisit_outside_local_exclusion": stream_flags["revisit"]["outside"],
+        "dynamic_causal": stream_flags["dynamic"]["causal"],
+        "dynamic_outside_local_exclusion": stream_flags["dynamic"]["outside"],
+        "dynamic_causal_outside_local_exclusion": stream_flags["dynamic"]["causal"] and stream_flags["dynamic"]["outside"],
+        "old_dynamic_raw": row["demem_dynamic"],
+        "old_revisit_raw": row["demem_revisit"],
+    }
+    visual = {
+        "sample": sample,
+        "video_path": video_path,
+        "target_raw": target_raw,
+        "selected": selected,
+        "dynamic_revisit_rel": dynamic_revisit_rel,
+        "dynamic_stream_rank": dynamic_stream_rank,
+    }
+    return record, visual
+
+
+def _read_frame(video_path: Path, frame_idx: int):
+    try:
+        import cv2
+    except ImportError:
+        return None
+    cap = cv2.VideoCapture(str(video_path))
+    if not cap.isOpened():
+        return None
+    cap.set(cv2.CAP_PROP_POS_FRAMES, int(frame_idx))
+    ok, frame = cap.read()
+    cap.release()
+    if not ok:
+        return None
+    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+    return Image.fromarray(frame)
+
+
+def _tile(video_path: Path, frame_idx: int | None, label: str, size=(160, 90), label_h=52) -> Image.Image:
+    image = Image.new("RGB", (size[0], size[1] + label_h), (230, 230, 230))
+    if frame_idx is not None and frame_idx >= 0:
+        frame = _read_frame(video_path, frame_idx)
+        if frame is not None:
+            image.paste(frame.resize(size, Image.BILINEAR), (0, label_h))
+    draw = ImageDraw.Draw(image)
+    draw.rectangle((0, 0, size[0], label_h), fill=(20, 20, 20))
+    for line_idx, line in enumerate(label.split("\n")[:3]):
+        draw.text((4, 4 + line_idx * 15), line[:28], fill=(255, 255, 255))
+    return image
+
+
+def _make_contact_sheet(visual_rows: list[dict], output_path: Path) -> bool:
+    try:
+        import cv2  # noqa: F401
+    except ImportError:
+        return False
+
+    columns = ["target", "anchor0", "anchor1", "revisit0", "revisit1", "dynamic0", "dynamic1", "dynamic2", "dynamic3"]
+    tile_w, tile_h = 160, 142
+    gap = 6
+    sheet = Image.new(
+        "RGB",
+        (len(columns) * tile_w + (len(columns) - 1) * gap, len(visual_rows) * tile_h + (len(visual_rows) - 1) * gap),
+        (245, 245, 245),
+    )
+    for row_idx, row in enumerate(visual_rows):
+        target = int(row["target_raw"])
+        selected = row["selected"]
+        video_path = row["video_path"]
+        cells: list[tuple[int | None, str]] = [
+            (target, f"target raw{target}\nrelT 0\nsample {row['sample']}"),
+        ]
+        for slot in range(2):
+            frame = selected["anchor"][slot] if slot < len(selected["anchor"]) else None
+            label = f"anchor{slot}" if frame is None else f"anchor{slot} raw{frame}\nrelT {frame - target}\n"
+            cells.append((frame, label))
+        for slot in range(2):
+            frame = selected["revisit"][slot] if slot < len(selected["revisit"]) else None
+            label = f"revisit{slot}" if frame is None else f"revisit{slot} raw{frame}\nrelT {frame - target}\n"
+            cells.append((frame, label))
+        for slot in range(4):
+            frame = selected["dynamic"][slot] if slot < len(selected["dynamic"]) else None
+            if frame is None:
+                label = f"dynamic{slot}"
+            else:
+                rel_r = row["dynamic_revisit_rel"][slot] if slot < len(row["dynamic_revisit_rel"]) else None
+                rank = row["dynamic_stream_rank"][slot] if slot < len(row["dynamic_stream_rank"]) else -1
+                label = f"dynamic{slot} raw{frame}\nrelT {frame - target} relR {rel_r}\nrank {rank}"
+            cells.append((frame, label))
+        y = row_idx * (tile_h + gap)
+        for col_idx, (frame, label) in enumerate(cells):
+            x = col_idx * (tile_w + gap)
+            sheet.paste(_tile(video_path, frame, label), (x, y))
+    sheet.save(output_path)
+    return True
+
+
+def _write_outputs(records: list[dict], visual_rows: list[dict]) -> None:
+    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+    csv_path = OUTPUT_DIR / "selector_diagnostic.csv"
+    with csv_path.open("w", newline="", encoding="utf-8") as handle:
+        writer = csv.DictWriter(handle, fieldnames=list(records[0].keys()))
+        writer.writeheader()
+        writer.writerows(records)
+
+    sheet_written = _make_contact_sheet(visual_rows, OUTPUT_DIR / "selector_contact_sheet.png")
+    md_path = OUTPUT_DIR / "selector_diagnostic.md"
+    with md_path.open("w", encoding="utf-8") as handle:
+        handle.write("# DeMemWM Revisit-Conditioned Dynamic Selector Diagnostic\n\n")
+        handle.write(f"- samples: {len(records)} from `{INPUT_CSV}`\n")
+        handle.write(f"- data root: `{DATA_ROOT}`\n")
+        handle.write("- selector: anchor prefix + wrapped pose deltas + cached dynamic stream around revisit\n")
+        handle.write(f"- contact sheet: `selector_contact_sheet.png` ({'written' if sheet_written else 'skipped: cv2 unavailable'})\n\n")
+        header = [
+            "sample",
+            "target",
+            "anchor",
+            "revisit",
+            "stream window",
+            "dynamic",
+            "dyn rel target",
+            "dyn rel revisit",
+            "dyn rank",
+            "stream causal/outside",
+        ]
+        handle.write("| " + " | ".join(header) + " |\n")
+        handle.write("| " + " | ".join(["---"] * len(header)) + " |\n")
+        for row in records:
+            causal = (
+                f"a {row['anchor_causal']}/{row['anchor_outside_local_exclusion']}; "
+                f"r {row['revisit_causal']}/{row['revisit_outside_local_exclusion']}; "
+                f"d {row['dynamic_causal']}/{row['dynamic_outside_local_exclusion']}"
+            )
+            values = [
+                str(row["sample"]),
+                str(row["target_raw"]),
+                row["anchor_raw"],
+                row["revisit_raw"],
+                row["stream_window_raw"],
+                row["dynamic_raw"],
+                row["dynamic_rel"],
+                row["dynamic_rel_to_nearest_revisit"],
+                row["dynamic_stream_rank"],
+                causal,
+            ]
+            handle.write("| " + " | ".join(values) + " |\n")
+
+
+def main() -> None:
+    records = []
+    visual_rows = []
+    for row in _read_samples():
+        record, visual = _analyze_sample(row)
+        records.append(record)
+        visual_rows.append(visual)
+    _write_outputs(records, visual_rows)
+    for row in records:
+        print(
+            f"sample={row['sample']} target={row['target_raw']} "
+            f"anchor={row['anchor_raw']} revisit={row['revisit_raw']} "
+            f"dynamic={row['dynamic_raw']} outside={row['dynamic_causal_outside_local_exclusion']}"
+        )
+    print(f"wrote {OUTPUT_DIR}")
+
+
+if __name__ == "__main__":
+    main()

.exp_artifact/dememwm_revisit_dynamic_selector_diagnostic/selector_diagnostic.csv

@@ -0,0 +1,7 @@
+sample,video,clip_offset,target_raw,anchor_candidate_window,causal_cutoff,anchor_raw,anchor_rel,revisit_raw,revisit_rel,eligible_stream_count,stream_window_raw,dynamic_raw,dynamic_rel,dynamic_rel_to_nearest_revisit,dynamic_stream_rank,anchor_causal,anchor_outside_local_exclusion,revisit_causal,revisit_outside_local_exclusion,dynamic_causal,dynamic_outside_local_exclusion,dynamic_causal_outside_local_exclusion,old_dynamic_raw,old_revisit_raw
+0,training/place_item_4/000011.mp4,783,983,883..982,975,"[883, 923]","[-100, -60]","[973, 974]","[-10, -9]",52,"[808, 822, 834, 851, 863, 880, 915, 943]","[863, 880, 915, 943]","[-120, -103, -68, -40]","[-110, -93, -58, -30]","[48, 49, 50, 51]",True,True,True,True,True,True,True,"[880, 915, 932, 943]","[973, 974]"
+1,training/savanna_2/000056.mp4,670,870,770..869,862,"[770, 869]","[-100, -1]","[542, 543]","[-328, -327]",41,"[481, 493, 512, 532, 553, 577, 591, 607]","[512, 532, 553, 577]","[-358, -338, -317, -293]","[-30, -10, 10, 34]","[21, 22, 23, 24]",True,False,True,True,True,True,True,"[512, 532, 553, 577]","[542, 543]"
+2,training/sunflower_0/000070.mp4,150,350,250..349,342,"[250, 301]","[-100, -49]","[150, 151]","[-200, -199]",15,"[100, 117, 147, 158, 177, 188, 207, 218]","[117, 147, 158, 177]","[-233, -203, -192, -173]","[-33, -3, 7, 26]","[1, 2, 3, 4]",True,True,True,True,True,True,True,"[117, 147, 177, 193]","[150, 151]"
+3,training/ice_plains_w_updown_0/000071.mp4,686,886,786..885,878,"[786, 866]","[-100, -20]","[425, 427]","[-461, -459]",38,"[333, 352, 397, 417, 433, 447, 482, 517]","[397, 417, 433, 447]","[-489, -469, -453, -439]","[-28, -8, 6, 20]","[15, 16, 17, 18]",True,True,True,True,True,True,True,"[397, 417, 433, 447]","[425, 427]"
+4,training/place_item_w_updown_3/000169.mp4,481,681,581..680,673,"[581, 680]","[-100, -1]","[453, 454]","[-228, -227]",34,"[391, 410, 421, 438, 460, 471, 498, 515]","[421, 438, 460, 471]","[-260, -243, -221, -210]","[-32, -15, 6, 17]","[20, 21, 22, 23]",True,False,True,True,True,True,True,"[410, 438, 460, 496]","[453, 454]"
+5,training/place_item_w_updown_7/001332.mp4,1137,1337,1237..1336,1329,"[1237, 1332]","[-100, -5]","[998, 1004]","[-339, -333]",70,"[923, 935, 969, 1008, 1020, 1039, 1050, 1081]","[969, 1008, 1020, 1039]","[-368, -329, -317, -298]","[-29, 4, 16, 35]","[51, 52, 53, 54]",True,False,True,True,True,True,True,"[969, 1008, 1020, 1039]","[998, 1004]"

.exp_artifact/dememwm_revisit_dynamic_selector_diagnostic/selector_diagnostic.md

@@ -0,0 +1,15 @@
+# DeMemWM Revisit-Conditioned Dynamic Selector Diagnostic
+
+- samples: 6 from `/share_1/users/bonan_ding/WorldMem/.exp_artifact/worldmem_vs_dememwm_memory_samples.csv`
+- data root: `/share_1/users/bonan_ding/worldmem_data/minecraft`
+- selector: anchor prefix + wrapped pose deltas + cached dynamic stream around revisit
+- contact sheet: `selector_contact_sheet.png` (written)
+
+| sample | target | anchor | revisit | stream window | dynamic | dyn rel target | dyn rel revisit | dyn rank | stream causal/outside |
+| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
+| 0 | 983 | [883, 923] | [973, 974] | [808, 822, 834, 851, 863, 880, 915, 943] | [863, 880, 915, 943] | [-120, -103, -68, -40] | [-110, -93, -58, -30] | [48, 49, 50, 51] | a True/True; r True/True; d True/True |
+| 1 | 870 | [770, 869] | [542, 543] | [481, 493, 512, 532, 553, 577, 591, 607] | [512, 532, 553, 577] | [-358, -338, -317, -293] | [-30, -10, 10, 34] | [21, 22, 23, 24] | a True/False; r True/True; d True/True |
+| 2 | 350 | [250, 301] | [150, 151] | [100, 117, 147, 158, 177, 188, 207, 218] | [117, 147, 158, 177] | [-233, -203, -192, -173] | [-33, -3, 7, 26] | [1, 2, 3, 4] | a True/True; r True/True; d True/True |
+| 3 | 886 | [786, 866] | [425, 427] | [333, 352, 397, 417, 433, 447, 482, 517] | [397, 417, 433, 447] | [-489, -469, -453, -439] | [-28, -8, 6, 20] | [15, 16, 17, 18] | a True/True; r True/True; d True/True |
+| 4 | 681 | [581, 680] | [453, 454] | [391, 410, 421, 438, 460, 471, 498, 515] | [421, 438, 460, 471] | [-260, -243, -221, -210] | [-32, -15, 6, 17] | [20, 21, 22, 23] | a True/False; r True/True; d True/True |
+| 5 | 1337 | [1237, 1332] | [998, 1004] | [923, 935, 969, 1008, 1020, 1039, 1050, 1081] | [969, 1008, 1020, 1039] | [-368, -329, -317, -298] | [-29, 4, 16, 35] | [51, 52, 53, 54] | a True/False; r True/True; d True/True |