#!/usr/bin/env python3 """Motion-stratified and correspondence-aware cache analysis for the three AR4 models. The generation jobs deliberately remain project-native. This file is an offline consumer of their saved feature snapshots and RGB anchor frames. It keeps the feature comparison at token locations (rather than comparing only a pooled vector), estimates RAFT flow with a forward/backward consistency mask, and reports raw, global-transform, homography, dense-flow, and (for WorldPlay) camera-action rotation alignment. """ from __future__ import annotations import argparse import csv import gc import json import math import os from collections import defaultdict from pathlib import Path from typing import Any, Iterable import cv2 import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import numpy as np import torch import torch.nn.functional as F IMG_W, IMG_H = 416, 240 ROLE_ORDER = ("early", "middle", "late", "final") def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--self_root", type=Path, required=True) parser.add_argument("--causal_root", type=Path, required=True) parser.add_argument("--hy_root", type=Path, required=True) parser.add_argument("--hy_right_root", type=Path, default=None) parser.add_argument("--output_root", type=Path, required=True) parser.add_argument("--flow_device", default="auto") parser.add_argument("--flow_backend", choices=("raft", "farneback"), default="raft") parser.add_argument("--corr_max_tokens", type=int, default=512) parser.add_argument("--overwrite", action="store_true") return parser.parse_args() def finite(value: Any) -> bool: try: return math.isfinite(float(value)) except (TypeError, ValueError): return False def mean_or_nan(values: Iterable[float]) -> float: values = [float(v) for v in values if finite(v)] return float(np.mean(values)) if values else float("nan") def write_csv(path: Path, rows: list[dict[str, Any]]) -> None: if not rows: return path.parent.mkdir(parents=True, exist_ok=True) fields: list[str] = [] for row in rows: for key in row: if key not in fields: fields.append(key) with path.open("w", newline="", encoding="utf-8") as handle: writer = csv.DictWriter(handle, fieldnames=fields, extrasaction="ignore") writer.writeheader() writer.writerows(rows) def json_dump(path: Path, value: Any) -> None: path.parent.mkdir(parents=True, exist_ok=True) path.write_text(json.dumps(value, indent=2, ensure_ascii=False) + "\n", encoding="utf-8") def _as_numpy(value: Any) -> np.ndarray: if isinstance(value, torch.Tensor): return value.detach().cpu().numpy() return np.asarray(value) def _key_parts(key: str) -> tuple[int, int]: left, right = str(key).split(":")[:2] return int(left), int(right) class FeatureRun: """A normalized view of one prompt/action's saved features.""" def __init__( self, model: str, action: str, prompt_id: int, path: Path, anchors: np.ndarray, chunks: int, chunk_size: int, stages: dict[tuple[str, int, int], tuple[np.ndarray, np.ndarray]], camera: tuple[np.ndarray, np.ndarray] | None = None, dense_path: Path | None = None, ): self.model = model self.action = action self.prompt_id = int(prompt_id) self.path = path self.anchors = anchors self.chunk_size = int(chunk_size) self.chunks = int(chunks) self.stages = stages self.camera = camera self.dense_path = dense_path self._dense_features = None self._dense_coords = None self._dense_stages = None self._dense_chunks = None self._dense_steps = None self.grid_h = 30 self.grid_w = 52 for _, (_, coords) in stages.items(): if len(coords): self.grid_h = max(self.grid_h, int(np.max(coords[:, 1])) + 1) self.grid_w = max(self.grid_w, int(np.max(coords[:, 2])) + 1) @property def stage_names(self) -> list[str]: return sorted({key[0] for key in self.stages}, key=stage_sort) def get(self, stage: str, chunk: int, step: int): # HY's block_53 dense snapshots are large compressed archives. Keep # them lazy so that one run is decompressed only while it is being # analyzed, instead of retaining ~1.2 GB per prompt/action in RAM. if stage == "block_53" and self.dense_path is not None and self.dense_path.exists(): if self._dense_features is None: data = np.load(self.dense_path, allow_pickle=False) # NpzFile lazily decompresses an array on every indexing # operation. Materialize each member once per run so the # 6.2-GB float16 feature tensor is not decompressed once per # chunk/timestep query. try: self._dense_features = np.asarray(data["features"]) self._dense_coords = np.asarray(data["coords"], dtype=np.int32) self._dense_stages = np.asarray(data["stages"]).astype(str) self._dense_chunks = np.asarray(data["chunks"], dtype=np.int32) self._dense_steps = np.asarray(data["steps"], dtype=np.int32) finally: data.close() stages = self._dense_stages chunks = self._dense_chunks steps = self._dense_steps matches = np.flatnonzero( (stages == stage) & (chunks == int(chunk)) & (steps == int(step)) ) if len(matches): index = int(matches[0]) return self._dense_features[index].astype(np.float32), self._dense_coords return self.stages.get((stage, int(chunk), int(step))) def release(self): if self._dense_features is not None: self._dense_features = None self._dense_coords = None self._dense_stages = None self._dense_chunks = None self._dense_steps = None gc.collect() def stage_sort(stage: str): if stage == "projected": return (10000, stage) try: return (int(stage.split("_")[1]), stage) except (IndexError, ValueError): return (9000, stage) def stage_role(run: FeatureRun, stage: str) -> str: if stage == "projected": return "final" block_stages = [name for name in run.stage_names if name.startswith("block_")] if stage in block_stages: index = block_stages.index(stage) return ROLE_ORDER[min(index, len(ROLE_ORDER) - 1)] return stage def _anchor(path: Path) -> np.ndarray: data = np.load(path, allow_pickle=False) value = data["frames"] if value.ndim != 4: raise ValueError(f"Expected [T,H,W,3] anchors at {path}, got {value.shape}") return value.astype(np.uint8) def load_self(root: Path) -> list[FeatureRun]: result: list[FeatureRun] = [] for path in sorted((root / "runs").glob("prompt_*.pt")): state = torch.load(path, map_location="cpu", weights_only=False) anchor_path = path.with_suffix(".anchors.npz") if not anchor_path.exists(): raise FileNotFoundError(anchor_path) coords_hidden = _as_numpy(state["sample_coords"]["hidden"]).astype(np.int32) stages: dict[tuple[str, int, int], tuple[np.ndarray, np.ndarray]] = {} for stage, values in state["records"].items(): if not stage.endswith("_hidden"): continue for key, value in values.items(): chunk, step = _key_parts(key) stages[(stage, chunk, step)] = (_as_numpy(value).astype(np.float32), coords_hidden) # The native recorder has a compact full-grid random projection for its # last block. Keep it as an additional final-stage representation. full_coords = np.stack( np.meshgrid(np.arange(3), np.arange(30), np.arange(52), indexing="ij"), axis=-1 ).reshape(-1, 3).astype(np.int32) for key, value in state.get("projected", {}).items(): chunk, step = _key_parts(key) array = _as_numpy(value).astype(np.float32).reshape(-1, _as_numpy(value).shape[-1]) stages[("projected", chunk, step)] = (array, full_coords) result.append( FeatureRun( "self_forcing", "none", int(state.get("run_index", len(result))), path, _anchor(anchor_path), int(state["num_frames"]) // int(state["num_frame_per_block"]), int(state["num_frame_per_block"]), stages, ) ) del state if len(result) != 10: print(f"[warn] Self-Forcing runs found {len(result)} prompts, expected 10") return result def _flat_indices_to_coords(indices: np.ndarray, frame_count: int = 3) -> np.ndarray: indices = indices.astype(np.int64).reshape(-1) frame_size = 30 * 52 return np.stack( [indices // frame_size, (indices % frame_size) // 52, indices % 52], axis=1 ).astype(np.int32) def load_causal(root: Path) -> list[FeatureRun]: result: list[FeatureRun] = [] for run_dir in sorted((root / "runs").glob("prompt_*")): path = run_dir / "feature_snapshots.pt" if not path.exists(): continue state = torch.load(path, map_location="cpu", weights_only=False) anchor_path = run_dir / "rgb_anchor_frames.npz" if not anchor_path.exists(): raise FileNotFoundError(anchor_path) stages: dict[tuple[str, int, int], tuple[np.ndarray, np.ndarray]] = {} index_dict = state.get("feature_indices", {}) for key, value in state["features"].items(): layer, chunk, step = (int(v) for v in str(key).split(":")) indices = index_dict.get(key) if indices is None: token_count = 3 * 30 * 52 indices = np.linspace(0, token_count - 1, int(state["max_tokens"])).round().astype(np.int64) coords = _flat_indices_to_coords(_as_numpy(indices)) stages[(f"block_{layer:02d}", chunk, step)] = (_as_numpy(value).astype(np.float32), coords) prompt_id = int(state.get("prompt_id", run_dir.name.split("_")[-1])) result.append( FeatureRun( "causal_forcing", "none", prompt_id, path, _anchor(anchor_path), int(state["num_chunks"]), 3, stages, ) ) del state if len(result) != 10: print(f"[warn] Causal-Forcing runs found {len(result)} prompts, expected 10") return result def _load_hy_action(root: Path, action: str) -> list[FeatureRun]: result: list[FeatureRun] = [] for case_dir in sorted((root / "runs").glob("prompt_*")): run_dir = case_dir / action if not run_dir.exists(): continue anchor_path = run_dir / "rgb_anchor_frames.npz" dense_path = run_dir / "dense_selected_snapshots.npz" sampled_path = run_dir / "final_hidden_snapshots.npz" # The sampled archive contains all requested layers and is sufficient # for the ordinary token metrics. Keep the expensive dense archive # only as a lazy source for HY's final block (block_53). snapshot_path = sampled_path if sampled_path.exists() else dense_path if not anchor_path.exists() or not snapshot_path.exists(): print(f"[warn] skip incomplete HY run {run_dir}") continue data = np.load(snapshot_path, allow_pickle=False) features = data["features"].astype(np.float32) chunks = data["chunks"].astype(int) steps = data["steps"].astype(int) stages_np = data["stages"].astype(str) coords = data["coords"].astype(np.int32) grid_shape = tuple(int(v) for v in data["grid_shape"]) stages: dict[tuple[str, int, int], tuple[np.ndarray, np.ndarray]] = {} for index in range(features.shape[0]): stages[(stages_np[index], int(chunks[index]), int(steps[index]))] = ( features[index], coords, ) camera = None camera_path = run_dir / "camera_trajectory.npz" if camera_path.exists(): camera_data = np.load(camera_path, allow_pickle=False) camera = (camera_data["viewmats"].astype(np.float64), camera_data["intrinsics"].astype(np.float64)) metadata = run_dir / "run_metadata.json" prompt_id = int(case_dir.name.split("_")[-1]) chunks_count = int(json.loads(metadata.read_text()).get("video_length", 45) if metadata.exists() else 45) latent_count = (chunks_count - 1) // 4 + 1 # Prefer the actual snapshot chunk count when available. chunks_count = max(1, int(np.max(chunks)) + 1) result.append( FeatureRun( "hy_worldplay", action, prompt_id, run_dir, _anchor(anchor_path), chunks_count, 4, stages, camera, # Keep the cross-model comparison at the same 240-token # sampling budget. The dense 6,240-token archive is retained # as an optional artifact, but using it for all-pairs matching # would make HY's final layer incomparable and unnecessarily # expensive (6,240^2 distances per query). dense_path=None, ) ) return result class FlowEstimator: def __init__(self, backend: str, device: str): self.backend = backend self.device = torch.device("cuda" if device == "auto" and torch.cuda.is_available() else (device if device != "auto" else "cpu")) self.model = None self.transforms = None if backend == "raft": try: from torchvision.models.optical_flow import Raft_Small_Weights, raft_small weights = Raft_Small_Weights.DEFAULT self.model = raft_small(weights=weights, progress=True).eval().to(self.device) self.transforms = weights.transforms() print(f"[flow] RAFT-small on {self.device}") except Exception as error: print(f"[flow] RAFT unavailable ({error}); using Farneback") self.backend = "farneback" @staticmethod def _image(frame: np.ndarray) -> np.ndarray: if frame.shape[:2] != (IMG_H, IMG_W): return cv2.resize(frame, (IMG_W, IMG_H), interpolation=cv2.INTER_AREA) return frame def _farneback(self, first: np.ndarray, second: np.ndarray) -> np.ndarray: a = cv2.cvtColor(self._image(first), cv2.COLOR_RGB2GRAY) b = cv2.cvtColor(self._image(second), cv2.COLOR_RGB2GRAY) return cv2.calcOpticalFlowFarneback(a, b, None, 0.5, 3, 21, 5, 7, 1.5, 0) @torch.inference_mode() def pair(self, first: np.ndarray, second: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """Return first->second flow, reverse flow, and a mask on first pixels.""" first = self._image(first) second = self._image(second) if self.backend != "raft" or self.model is None: forward = self._farneback(first, second) backward = self._farneback(second, first) else: x = torch.from_numpy(first).permute(2, 0, 1).float().div_(255).unsqueeze(0).to(self.device) y = torch.from_numpy(second).permute(2, 0, 1).float().div_(255).unsqueeze(0).to(self.device) x, y = self.transforms(x, y) forward = self.model(x, y)[-1][0].permute(1, 2, 0).float().cpu().numpy() backward = self.model(y, x)[-1][0].permute(1, 2, 0).float().cpu().numpy() # The feature query is made at pixels in ``first`` (the current/target # frame), so the consistency mask must also be indexed in that frame. # A previous implementation used the reverse flow and a mask indexed # in ``second``; that silently inverted the alignment direction. h, w = forward.shape[:2] yy, xx = np.mgrid[0:h, 0:w].astype(np.float32) sx = xx + forward[..., 0] sy = yy + forward[..., 1] sampled_backward = cv2.remap(backward, sx, sy, cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT) fb = np.linalg.norm(forward + sampled_backward, axis=-1) magnitude = np.linalg.norm(forward, axis=-1) in_bounds = (sx >= 0) & (sx < w) & (sy >= 0) & (sy < h) # The threshold scales mildly with motion, avoiding an overly strict # rejection of the fast camera actions while retaining occlusion masks. valid = in_bounds & (fb <= 1.5 + 0.05 * magnitude) return forward, backward, valid def flow_at(flow: np.ndarray, xy: np.ndarray) -> np.ndarray: h, w = flow.shape[:2] x = np.clip(xy[:, 0], 0, w - 1).astype(np.float32) y = np.clip(xy[:, 1], 0, h - 1).astype(np.float32) return np.stack( [cv2.remap(flow[..., dim], x, y, cv2.INTER_LINEAR).reshape(-1) for dim in range(2)], axis=1 ) def mask_at(mask: np.ndarray, xy: np.ndarray) -> np.ndarray: h, w = mask.shape[:2] x = np.clip(xy[:, 0], 0, w - 1).astype(np.float32) y = np.clip(xy[:, 1], 0, h - 1).astype(np.float32) value = cv2.remap(mask.astype(np.uint8), x, y, cv2.INTER_NEAREST).reshape(-1) return value.astype(bool) def fit_transforms(flow: np.ndarray, valid: np.ndarray): h, w = flow.shape[:2] yy, xx = np.mgrid[0:h:8, 0:w:8].astype(np.float32) points = np.stack([xx.reshape(-1), yy.reshape(-1)], axis=1) selected = valid[::8, ::8].reshape(-1) points = points[selected] if len(points) < 6: fallback_y, fallback_x = np.mgrid[0:h:4, 0:w:4] points = np.stack([fallback_x.reshape(-1), fallback_y.reshape(-1)], axis=1).astype(np.float32) selected_flow = flow[::4, ::4].reshape(-1, 2) else: selected_flow = flow[::8, ::8].reshape(-1, 2)[selected] destinations = points + selected_flow affine = None homography = None if len(points) >= 3: affine, _ = cv2.estimateAffine2D(points, destinations, method=cv2.RANSAC, ransacReprojThreshold=3.0) if len(points) >= 4: homography, _ = cv2.findHomography(points, destinations, cv2.RANSAC, 4.0) median = np.median(flow[valid] if bool(np.any(valid)) else flow.reshape(-1, 2), axis=0) translation = np.asarray([[1.0, 0.0, median[0]], [0.0, 1.0, median[1]]], dtype=np.float32) return translation, affine, homography def apply_transform(points: np.ndarray, transform: np.ndarray | None) -> np.ndarray | None: if transform is None or len(points) == 0: return None if transform.shape == (2, 3): return points @ transform[:, :2].T + transform[:, 2] homogeneous = np.concatenate([points, np.ones((len(points), 1), dtype=np.float32)], axis=1) mapped = homogeneous @ transform.T return mapped[:, :2] / np.clip(mapped[:, 2:3], 1e-6, None) def nearest_features(source_xy: np.ndarray, source_features: np.ndarray, query_xy: np.ndarray): if len(source_xy) == 0 or len(query_xy) == 0: return np.empty((0, source_features.shape[-1]), dtype=np.float32), np.zeros(len(query_xy), bool), np.empty(len(query_xy)) distances = ((query_xy[:, None, :] - source_xy[None, :, :]) ** 2).sum(axis=-1) indices = np.argmin(distances, axis=1) return source_features[indices], np.ones(len(query_xy), bool), np.sqrt(distances[np.arange(len(indices)), indices]) def cosine_mean(left: np.ndarray, right: np.ndarray, mask: np.ndarray | None = None) -> float: if len(left) == 0 or len(right) == 0: return float("nan") value = np.sum(left * right, axis=-1) / ( np.linalg.norm(left, axis=-1) * np.linalg.norm(right, axis=-1) + 1e-8 ) if mask is not None: value = value[mask] return mean_or_nan(value) def select_frame(features: np.ndarray, coords: np.ndarray, frame: int): selected = coords[:, 0] == int(frame) return features[selected], coords[selected, 1:3][:, ::-1].astype(np.float32) def token_correspondence(target: np.ndarray, target_xy: np.ndarray, source: np.ndarray, source_xy: np.ndarray, max_tokens: int): if len(target) == 0 or len(source) == 0: return {"top1_cosine": float("nan"), "top5_cosine": float("nan"), "mnn_fraction": float("nan"), "match_distance": float("nan")} def thin(array, xy): if len(array) <= max_tokens: return array, xy idx = np.linspace(0, len(array) - 1, max_tokens).round().astype(int) return array[idx], xy[idx] target, target_xy = thin(target, target_xy) source, source_xy = thin(source, source_xy) target_n = target / (np.linalg.norm(target, axis=-1, keepdims=True) + 1e-8) source_n = source / (np.linalg.norm(source, axis=-1, keepdims=True) + 1e-8) similarity = target_n @ source_n.T topk = np.sort(similarity, axis=1)[:, -min(5, similarity.shape[1]):] best_source = np.argmax(similarity, axis=1) best_target = np.argmax(similarity, axis=0) target_ids = np.arange(len(best_source)) mnn = best_target[best_source] == target_ids distance = np.linalg.norm(target_xy - source_xy[best_source], axis=-1) return { "top1_cosine": float(np.mean(topk[:, -1])), "top5_cosine": float(np.mean(topk)), "mnn_fraction": float(np.mean(mnn)), "match_distance": float(np.mean(distance)), } def camera_rotation_homography(camera, source_index: int, target_index: int, image_shape=(IMG_H, IMG_W)): if camera is None: return None viewmats, intrinsics = camera if source_index >= len(viewmats) or target_index >= len(viewmats): return None source_k = intrinsics[source_index].copy() target_k = intrinsics[target_index].copy() h, w = image_shape for k in (source_k, target_k): k[0, 0] *= w k[0, 2] *= w k[1, 1] *= h k[1, 2] *= h source_r = viewmats[source_index][:3, :3] target_r = viewmats[target_index][:3, :3] try: return source_k @ source_r @ target_r.T @ np.linalg.inv(target_k) except np.linalg.LinAlgError: return None def action_motion(camera, source_index: int, target_index: int): if camera is None: return {"action_translation": float("nan"), "action_rotation_deg": float("nan")} viewmats, _ = camera if source_index >= len(viewmats) or target_index >= len(viewmats): return {"action_translation": float("nan"), "action_rotation_deg": float("nan")} source_to_world = np.linalg.inv(viewmats[source_index]) target_to_world = np.linalg.inv(viewmats[target_index]) relative_rotation = source_to_world[:3, :3].T @ target_to_world[:3, :3] trace = np.clip((np.trace(relative_rotation) - 1.0) / 2.0, -1.0, 1.0) return { "action_translation": float(np.linalg.norm(target_to_world[:3, 3] - source_to_world[:3, 3])), "action_rotation_deg": float(np.degrees(np.arccos(trace))), } def token_points(run: FeatureRun, coords: np.ndarray) -> np.ndarray: # coords are (frame,y,x) in the latent feature grid; convert to pixel # centers in the resized RGB canvas, represented as (x,y). return np.stack( [ (coords[:, 2].astype(np.float32) + 0.5) * IMG_W / run.grid_w, (coords[:, 1].astype(np.float32) + 0.5) * IMG_H / run.grid_h, ], axis=1, ) def alignment_metrics(run: FeatureRun, stage: str, target_pair, source_pair, target_slot: int, flow: np.ndarray, valid_flow: np.ndarray, transforms, action_h, corr_max_tokens: int): target_features, target_coords = target_pair source_features, source_coords = source_pair target_features, target_xy = select_frame(target_features, target_coords, target_slot) source_boundary, source_boundary_xy = select_frame(source_features, source_coords, run.chunk_size - 1) source_same, source_same_xy = select_frame(source_features, source_coords, min(target_slot, run.chunk_size - 1)) if len(target_features) == 0 or len(source_boundary) == 0: return None target_pixels = token_points(run, np.concatenate([np.full((len(target_xy), 1), target_slot), target_xy[:, ::-1]], axis=1).astype(np.int32)) source_boundary_pixels = token_points(run, np.concatenate([np.full((len(source_boundary_xy), 1), run.chunk_size - 1), source_boundary_xy[:, ::-1]], axis=1).astype(np.int32)) source_same_pixels = token_points(run, np.concatenate([np.full((len(source_same_xy), 1), min(target_slot, run.chunk_size - 1)), source_same_xy[:, ::-1]], axis=1).astype(np.int32)) raw_features, _, raw_dist = nearest_features(source_boundary_pixels, source_boundary, target_pixels) same_features, _, _ = nearest_features(source_same_pixels, source_same, target_pixels) target_flow = flow_at(flow, target_pixels) flow_query = target_pixels + target_flow flow_mask = mask_at(valid_flow, target_pixels) flow_features, _, flow_dist = nearest_features(source_boundary_pixels, source_boundary, flow_query) translation, affine, homography = transforms def transformed_metric(transform): query = apply_transform(target_pixels, transform) if query is None: return float("nan"), float("nan") values, _, distances = nearest_features(source_boundary_pixels, source_boundary, query) in_bounds = (query[:, 0] >= 0) & (query[:, 0] < IMG_W) & (query[:, 1] >= 0) & (query[:, 1] < IMG_H) return cosine_mean(target_features, values, in_bounds), mean_or_nan(distances[in_bounds]) trans_cos, trans_dist = transformed_metric(translation) affine_cos, affine_dist = transformed_metric(affine) homo_cos, homo_dist = transformed_metric(homography) action_cos, action_dist = transformed_metric(action_h) correspondence = token_correspondence(target_features, target_pixels, source_boundary, source_boundary_pixels, corr_max_tokens) return { "same_slot_cosine": cosine_mean(target_features, same_features), "boundary_raw_cosine": cosine_mean(target_features, raw_features), "translation_aligned_cosine": trans_cos, "affine_aligned_cosine": affine_cos, "homography_aligned_cosine": homo_cos, "flow_aligned_cosine": cosine_mean(target_features, flow_features, flow_mask), "action_rotation_cosine": action_cos, "raw_match_distance": mean_or_nan(raw_dist), "flow_match_distance": mean_or_nan(flow_dist[flow_mask]), "translation_match_distance": trans_dist, "affine_match_distance": affine_dist, "homography_match_distance": homo_dist, "action_match_distance": action_dist, "valid_flow_ratio": float(np.mean(flow_mask)), "occlusion_ratio": float(1.0 - np.mean(flow_mask)), **correspondence, } def motion_rows(runs: list[FeatureRun], estimator: FlowEstimator, corr_max_tokens: int) -> tuple[list[dict[str, Any]], list[dict[str, Any]], dict[str, Any]]: rows: list[dict[str, Any]] = [] correspondence_rows: list[dict[str, Any]] = [] flow_cache: dict[tuple[str, int, int, int], tuple[np.ndarray, np.ndarray, np.ndarray]] = {} for run in runs: for chunk in range(1, run.chunks): source_index = (chunk - 1) * run.chunk_size + run.chunk_size - 1 if source_index >= len(run.anchors): continue for slot in range(run.chunk_size): target_index = chunk * run.chunk_size + slot if target_index >= len(run.anchors): continue cache_key = (str(run.path), chunk, source_index, target_index) if cache_key not in flow_cache: # The first return is target->source and is therefore the # displacement used to query the previous chunk's tokens. forward, _, valid = estimator.pair(run.anchors[target_index], run.anchors[source_index]) flow_cache[cache_key] = (forward, valid, np.asarray([])) flow, valid, _ = flow_cache[cache_key] flow_mag = np.linalg.norm(flow, axis=-1) median = np.median(flow[valid] if bool(np.any(valid)) else flow.reshape(-1, 2), axis=0) residual = flow - median[None, None] source_img = cv2.resize(run.anchors[source_index], (IMG_W, IMG_H), interpolation=cv2.INTER_AREA) target_img = cv2.resize(run.anchors[target_index], (IMG_W, IMG_H), interpolation=cv2.INTER_AREA) scene_cut = float(np.mean(np.abs(source_img.astype(np.float32) - target_img.astype(np.float32))) / 255.0) action_motion_values = action_motion(run.camera, source_index, target_index) action_h = camera_rotation_homography(run.camera, source_index, target_index) transforms = fit_transforms(flow, valid) stage_names = [stage for stage in run.stage_names if stage.startswith("block_") or stage == "projected"] for stage in stage_names: common_steps = sorted( set(step for st, ch, step in run.stages if st == stage and ch == chunk) & set(step for st, ch, step in run.stages if st == stage and ch == chunk - 1) ) for step in common_steps: target_pair = run.get(stage, chunk, step) source_pair = run.get(stage, chunk - 1, step) metrics = alignment_metrics(run, stage, target_pair, source_pair, slot, flow, valid, transforms, action_h, corr_max_tokens) if metrics is None: continue row = { "model": run.model, "action": run.action, "prompt_id": run.prompt_id, "chunk": chunk, "target_slot": slot, "step": step, "stage": stage, "role": stage_role(run, stage), "source_frame": source_index, "target_frame": target_index, "total_motion": float(np.mean(flow_mag)), "camera_motion": float(np.linalg.norm(median)), "object_motion": float(np.mean(np.linalg.norm(residual, axis=-1))), "scene_cut": scene_cut, **action_motion_values, **metrics, } rows.append(row) if row["role"] == "final": correspondence_rows.append({ key: row[key] for key in ("model", "action", "prompt_id", "chunk", "target_slot", "step", "stage", "role", "total_motion") } | {key: metrics[key] for key in ("top1_cosine", "top5_cosine", "mnn_fraction", "match_distance")}) run.release() values = np.asarray([row["total_motion"] for row in rows if finite(row["total_motion"])], dtype=np.float64) if len(values) >= 3: low, high = np.quantile(values, [1 / 3, 2 / 3]) for row in rows: row["motion_bin"] = "low" if row["total_motion"] <= low else "high" if row["total_motion"] > high else "medium" for row in correspondence_rows: matches = [item for item in rows if all(item[k] == row[k] for k in ("model", "action", "prompt_id", "chunk", "target_slot", "step", "stage", "role"))] row["motion_bin"] = matches[0].get("motion_bin", "unknown") if matches else "unknown" metadata = { "flow_backend": estimator.backend, "flow_device": str(estimator.device), "flow_resolution": [IMG_W, IMG_H], "forward_backward_consistency": "valid = in-bounds and FB error <= 1.5 + 0.05*|flow|", "motion_bin_edges": [float(low), float(high)] if len(values) >= 3 else [], } return rows, correspondence_rows, metadata def group_rows(rows: list[dict[str, Any]], keys: list[str], metrics: list[str]) -> list[dict[str, Any]]: groups: dict[tuple[Any, ...], list[dict[str, Any]]] = defaultdict(list) for row in rows: groups[tuple(row.get(key) for key in keys)].append(row) output: list[dict[str, Any]] = [] for group, values in sorted(groups.items(), key=lambda item: tuple(map(str, item[0]))): item = {key: value for key, value in zip(keys, group)} item["count"] = len(values) for metric in metrics: item[metric] = mean_or_nan([row.get(metric) for row in values]) output.append(item) return output def plot_outputs(rows: list[dict[str, Any]], correspondence: list[dict[str, Any]], output: Path) -> None: output.mkdir(parents=True, exist_ok=True) if not rows: return models = sorted({str(row["model"]) for row in rows}) colors = {model: plt.cm.tab10(index) for index, model in enumerate(models)} fig, axes = plt.subplots(2, 2, figsize=(14, 10)) for model in models: selected = [row for row in rows if row["model"] == model] axes[0, 0].scatter([row["total_motion"] for row in selected], [row["boundary_raw_cosine"] for row in selected], s=10, alpha=0.45, label=model, color=colors[model]) axes[0, 0].set(xlabel="flow magnitude", ylabel="raw boundary cosine", title="Motion vs raw cross-chunk similarity") axes[0, 0].legend(fontsize=8) methods = ["boundary_raw_cosine", "translation_aligned_cosine", "affine_aligned_cosine", "homography_aligned_cosine", "flow_aligned_cosine", "action_rotation_cosine"] labels = ["raw", "translation", "affine", "homography", "dense flow", "action rotation"] positions = np.arange(len(methods)) for index, model in enumerate(models): values = [mean_or_nan([row.get(method) for row in rows if row["model"] == model]) for method in methods] axes[0, 1].plot(positions, values, marker="o", label=model, color=colors[model]) axes[0, 1].set_xticks(positions, labels, rotation=25, ha="right") axes[0, 1].set_ylim(-1, 1) axes[0, 1].set_title("Alignment methods") axes[0, 1].legend(fontsize=8) for index, model in enumerate(models): selected = [row for row in rows if row["model"] == model] bins = ["low", "medium", "high"] raw = [mean_or_nan([row["boundary_raw_cosine"] for row in selected if row.get("motion_bin") == name]) for name in bins] flow = [mean_or_nan([row["flow_aligned_cosine"] for row in selected if row.get("motion_bin") == name]) for name in bins] axes[1, 0].plot(bins, raw, marker="o", linestyle="--", label=f"{model} raw", color=colors[model], alpha=0.55) axes[1, 0].plot(bins, flow, marker="o", label=f"{model} flow", color=colors[model]) axes[1, 0].set_ylim(-1, 1) axes[1, 0].set_title("Motion bins") axes[1, 0].set_ylabel("cosine") axes[1, 0].legend(fontsize=7) if correspondence: for model in models: selected = [row for row in correspondence if row["model"] == model] axes[1, 1].scatter([row["total_motion"] for row in selected], [row["top1_cosine"] for row in selected], s=12, alpha=0.5, label=model, color=colors[model]) axes[1, 1].set(xlabel="flow magnitude", ylabel="top-1 feature match cosine", title="Token correspondence") axes[1, 1].legend(fontsize=8) fig.tight_layout() fig.savefig(output / "motion_alignment_overview.png", dpi=180) plt.close(fig) if correspondence: fig, axes = plt.subplots(1, 3, figsize=(15, 4.5)) for model in models: selected = [row for row in correspondence if row["model"] == model] for axis, metric, title in zip(axes, ("top1_cosine", "mnn_fraction", "match_distance"), ("Top-1 cosine", "MNN fraction", "Match distance")): axis.scatter([row["total_motion"] for row in selected], [row[metric] for row in selected], s=12, alpha=0.5, label=model, color=colors[model]) axis.set_title(title) axis.set_xlabel("flow magnitude") axes[0].set_ylabel("value") axes[0].legend(fontsize=8) fig.tight_layout() fig.savefig(output / "token_correspondence_summary.png", dpi=180) plt.close(fig) def markdown_table(rows: list[dict[str, Any]], columns: list[str]) -> str: if not rows: return "_No rows._" lines = ["| " + " | ".join(columns) + " |", "|" + "|".join("---" for _ in columns) + "|"] for row in rows: values = [] for column in columns: value = row.get(column, "") values.append(f"{float(value):.4f}" if isinstance(value, (float, np.floating)) and finite(value) else str(value)) lines.append("| " + " | ".join(values) + " |") return "\n".join(lines) def write_report(output: Path, rows: list[dict[str, Any]], correspondence: list[dict[str, Any]], metadata: dict[str, Any], runs: list[FeatureRun]) -> None: stage_summary = group_rows(rows, ["model", "action", "role", "stage"], ["boundary_raw_cosine", "translation_aligned_cosine", "affine_aligned_cosine", "homography_aligned_cosine", "flow_aligned_cosine", "action_rotation_cosine", "total_motion", "valid_flow_ratio"]) bin_summary = group_rows(rows, ["model", "action", "motion_bin"], ["total_motion", "boundary_raw_cosine", "flow_aligned_cosine", "homography_aligned_cosine", "object_motion", "occlusion_ratio"]) json_dump(output / "motion_alignment_summary.json", {"metadata": metadata, "stage_summary": stage_summary, "motion_bin_summary": bin_summary, "run_count": len(runs), "row_count": len(rows), "correspondence_count": len(correspondence)}) report = [ "# Experiment 3: motion stratification and explicit alignment", "", f"Runs: {len(runs)} prompt/action runs; metric rows: {len(rows)}; correspondence rows: {len(correspondence)}.", "", "The raw baseline compares the current chunk token with the previous chunk's boundary frame at the same spatial coordinate. `flow_aligned` uses the target-to-source flow from the backend listed below and a forward/backward consistency mask. Affine and homography are global-transform fits to the valid flow. `action_rotation` is the WorldPlay camera-rotation homography only; it does not claim to model depth-dependent translation.", "", "## By model/action/stage", "", markdown_table(stage_summary, ["model", "action", "role", "stage", "count", "boundary_raw_cosine", "translation_aligned_cosine", "affine_aligned_cosine", "homography_aligned_cosine", "flow_aligned_cosine", "action_rotation_cosine"]), "", "## Motion bins", "", markdown_table(bin_summary, ["model", "action", "motion_bin", "count", "total_motion", "boundary_raw_cosine", "homography_aligned_cosine", "flow_aligned_cosine", "object_motion", "occlusion_ratio"]), "", "## Interpretation guardrails", "", "- A positive dense-flow gain with lower raw cosine supports spatial migration rather than loss of content information.", "- Action rotation is a physically informed control for HY; forward translation requires depth and is therefore reported separately rather than treated as an exact warp.", "- Correlation and alignment rows are prompt-level repeated measurements; use prompt bootstrap or a mixed-effects model for significance claims.", "", "## Run settings", "", "```json", json.dumps(metadata, indent=2, ensure_ascii=False), "```", ] (output / "REPORT.md").write_text("\n".join(report) + "\n", encoding="utf-8") write_csv(output / "motion_alignment_metrics.csv", rows) write_csv(output / "motion_alignment_summary.csv", stage_summary) write_csv(output / "motion_alignment_motion_bins.csv", bin_summary) write_csv(output / "token_correspondence.csv", correspondence) def main() -> None: args = parse_args() output = args.output_root.resolve() output.mkdir(parents=True, exist_ok=True) if not args.overwrite and (output / "motion_alignment_metrics.csv").exists(): print(f"[skip] existing analysis at {output}; pass --overwrite to recompute") return runs: list[FeatureRun] = [] runs.extend(load_self(args.self_root.resolve())) runs.extend(load_causal(args.causal_root.resolve())) runs.extend(_load_hy_action(args.hy_root.resolve(), "forward")) runs.extend(_load_hy_action(args.hy_root.resolve(), "static")) if args.hy_right_root is not None: runs.extend(_load_hy_action(args.hy_right_root.resolve(), "right")) if not runs: raise RuntimeError("No normalized runs found") estimator = FlowEstimator(args.flow_backend, args.flow_device) rows, correspondence, metadata = motion_rows(runs, estimator, args.corr_max_tokens) metadata.update({ "models": sorted({run.model for run in runs}), "actions": sorted({run.action for run in runs}), "prompt_ids": sorted({run.prompt_id for run in runs}), "run_descriptions": [ {"model": run.model, "action": run.action, "prompt_id": run.prompt_id, "chunk_size": run.chunk_size, "chunks": run.chunks, "anchor_frames": len(run.anchors), "stages": run.stage_names} for run in runs ], }) plot_outputs(rows, correspondence, output) write_report(output, rows, correspondence, metadata, runs) json_dump(output / "analysis_config.json", metadata) print(f"[complete] {output}: {len(runs)} runs, {len(rows)} motion rows, {len(correspondence)} correspondence rows") if __name__ == "__main__": main()