"""Linear probes for FlowMo latent state and drift context.""" from __future__ import annotations import argparse import importlib import json from dataclasses import dataclass from pathlib import Path import numpy as np import torch from torch.utils.data import DataLoader, Dataset from experiments.shared.src.vision.clean_renderer import render_clean_boat_history_tensor from experiments.train_image_world_models import autocast_context, configure_training_runtime, selected_history_indices @dataclass class ProbeSet: z: np.ndarray c: np.ndarray momentum: np.ndarray local_flow: np.ndarray episode_drift: np.ndarray class FlowMoProbeDataset(Dataset): def __init__( self, source_npz: str, history_len: int, episodes: int, max_windows: int, seed: int, ): src = np.load(source_npz, allow_pickle=False) self.states = src["states"][:episodes].astype(np.float32) self.actions = src["actions"][:episodes].astype(np.float32) self.true_flow = src["true_flow"][:episodes].astype(np.float32) self.boat_ids = src["boat_ids"][:episodes].astype(np.int64) self.history_len = int(history_len) steps = self.actions.shape[1] all_indices = [(ep, t) for ep in range(episodes) for t in range(self.history_len - 1, steps)] rng = np.random.default_rng(seed) selected = rng.choice(len(all_indices), size=min(max_windows, len(all_indices)), replace=False) self.indices = [all_indices[int(i)] for i in selected] self.episode_drift = self.true_flow.mean(axis=1).astype(np.float32) def __len__(self) -> int: return len(self.indices) def __getitem__(self, index: int): ep, t = self.indices[index] state_hist = self.states[ep, t - self.history_len + 1 : t + 1, :6].copy() padded_actions = np.zeros((self.actions.shape[1] + 1, self.actions.shape[2]), dtype=np.float32) padded_actions[1:] = self.actions[ep] action_hist = padded_actions[t - self.history_len + 1 : t + 1].copy() return ( torch.from_numpy(state_hist), torch.from_numpy(action_hist), torch.tensor(self.boat_ids[ep], dtype=torch.long), torch.from_numpy(self.states[ep, t, 3:6].copy()), torch.from_numpy(self.true_flow[ep, t].copy()), torch.from_numpy(self.episode_drift[ep].copy()), ) def build_flowmo(): config_module = importlib.import_module("experiments.flowmo.src.config") model_module = importlib.import_module("experiments.flowmo.src.model") cfg = config_module.default_config() return cfg, model_module.build_model(cfg) def loader_kwargs(num_workers: int) -> dict: if num_workers <= 0: return {} return {"multiprocessing_context": "spawn", "persistent_workers": True, "prefetch_factor": 4} @torch.no_grad() def collect_probe_set(model, source: str, episodes: int, max_windows: int, args, seed: int) -> ProbeSet: dataset = FlowMoProbeDataset(source, args.history_len, episodes, max_windows, seed) loader = DataLoader( dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=torch.device(args.device).type == "cuda", **loader_kwargs(args.num_workers), ) device = torch.device(args.device) history_indices = selected_history_indices(model, args.history_len) zs, cs, momentum, local_flow, episode_drift = [], [], [], [], [] model.eval() for states, actions, boat_ids, batch_momentum, batch_flow, batch_drift in loader: states = states[:, history_indices].to(device, non_blocking=True) actions = actions[:, history_indices].to(device, non_blocking=True) boat_ids = boat_ids.to(device, non_blocking=True) images = render_clean_boat_history_tensor( states, boat_ids, image_size=args.image_size, visual_scale=args.visual_scale, ) with autocast_context(device, args.precision): z, c = model.encode(images, actions) zs.append(z.float().cpu().numpy()) cs.append(c.float().cpu().numpy()) momentum.append(batch_momentum.numpy()) local_flow.append(batch_flow.numpy()) episode_drift.append(batch_drift.numpy()) return ProbeSet( z=np.concatenate(zs, axis=0), c=np.concatenate(cs, axis=0), momentum=np.concatenate(momentum, axis=0), local_flow=np.concatenate(local_flow, axis=0), episode_drift=np.concatenate(episode_drift, axis=0), ) def fit_ridge(features: np.ndarray, targets: np.ndarray, alpha: float) -> dict[str, np.ndarray]: x_mean = features.mean(axis=0, keepdims=True) x_std = features.std(axis=0, keepdims=True) + 1.0e-6 x = (features - x_mean) / x_std x = np.concatenate([x, np.ones((x.shape[0], 1), dtype=x.dtype)], axis=1) penalty = np.eye(x.shape[1], dtype=np.float64) * float(alpha) penalty[-1, -1] = 0.0 weights = np.linalg.solve(x.T @ x + penalty, x.T @ targets) return {"x_mean": x_mean, "x_std": x_std, "weights": weights} def predict_ridge(model: dict[str, np.ndarray], features: np.ndarray) -> np.ndarray: x = (features - model["x_mean"]) / model["x_std"] x = np.concatenate([x, np.ones((x.shape[0], 1), dtype=x.dtype)], axis=1) return x @ model["weights"] def regression_metrics(prediction: np.ndarray, target: np.ndarray, dims: list[str]) -> dict: error = prediction - target mse_per_dim = np.mean(error * error, axis=0) rmse_per_dim = np.sqrt(mse_per_dim) baseline_mse = np.mean((target - target.mean(axis=0, keepdims=True)) ** 2, axis=0) + 1.0e-12 r2_per_dim = 1.0 - mse_per_dim / baseline_mse return { "rmse": float(np.sqrt(np.mean(error * error))), "r2_mean": float(np.mean(r2_per_dim)), "rmse_by_dim": {dim: float(value) for dim, value in zip(dims, rmse_per_dim)}, "r2_by_dim": {dim: float(value) for dim, value in zip(dims, r2_per_dim)}, } def feature_matrix(probe_set: ProbeSet, feature_name: str) -> np.ndarray: if feature_name == "z": return probe_set.z if feature_name == "c": return probe_set.c if feature_name == "z_c": return np.concatenate([probe_set.z, probe_set.c], axis=1) raise ValueError(f"unknown feature set: {feature_name}") def target_matrix(probe_set: ProbeSet, target_name: str) -> tuple[np.ndarray, list[str]]: if target_name == "momentum": return probe_set.momentum, ["vx", "vy", "omega"] if target_name == "local_flow": return probe_set.local_flow, ["flow_x", "flow_y"] if target_name == "episode_drift": return probe_set.episode_drift, ["drift_x", "drift_y"] raise ValueError(f"unknown target: {target_name}") def train_probe_bank(train_set: ProbeSet, alpha: float) -> dict: bank = {} for target_name in ["momentum", "local_flow", "episode_drift"]: targets, _dims = target_matrix(train_set, target_name) bank[target_name] = {} for feature_name in ["z", "c", "z_c"]: bank[target_name][feature_name] = fit_ridge(feature_matrix(train_set, feature_name), targets, alpha) return bank def evaluate_probe_bank(bank: dict, probe_set: ProbeSet) -> dict: payload = {} for target_name in ["momentum", "local_flow", "episode_drift"]: targets, dims = target_matrix(probe_set, target_name) payload[target_name] = {} for feature_name in ["z", "c", "z_c"]: prediction = predict_ridge(bank[target_name][feature_name], feature_matrix(probe_set, feature_name)) payload[target_name][feature_name] = regression_metrics(prediction, targets, dims) return payload def parse_eval_split(text: str) -> tuple[str, str, int]: name, source, episodes = text.split(":", maxsplit=2) return name, source, int(episodes) def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--train-source", default="data/paper/train.npz") parser.add_argument("--train-episodes", type=int, default=2400) parser.add_argument("--train-windows", type=int, default=32768) parser.add_argument("--eval-splits", nargs="+", default=[ "test:data/paper/test.npz:480", ]) parser.add_argument("--eval-windows", type=int, default=8192) parser.add_argument("--history-len", type=int, default=32) parser.add_argument("--batch-size", type=int, default=256) parser.add_argument("--ridge-alpha", type=float, default=1.0e-3) parser.add_argument("--checkpoint-name", default="paper.pt") parser.add_argument("--image-size", type=int, default=160) parser.add_argument("--visual-scale", type=float, default=2.5) parser.add_argument("--num-workers", type=int, default=4) parser.add_argument("--seed", type=int, default=909) parser.add_argument("--device", default="cuda") parser.add_argument("--precision", choices=["fp32", "bf16", "fp16"], default="bf16") parser.add_argument("--out", default="experiments/reports/paper_flowmo_latent_probes.json") args = parser.parse_args() device = torch.device(args.device) configure_training_runtime(device) _cfg, model = build_flowmo() state = torch.load(Path("experiments/flowmo/checkpoint") / args.checkpoint_name, map_location="cpu") model.load_state_dict(state) model.to(device) if device.type == "cuda": model.to(memory_format=torch.channels_last) train_set = collect_probe_set(model, args.train_source, args.train_episodes, args.train_windows, args, args.seed) bank = train_probe_bank(train_set, args.ridge_alpha) split_results = {} for index, split in enumerate(args.eval_splits): name, source, episodes = parse_eval_split(split) eval_set = collect_probe_set(model, source, episodes, args.eval_windows, args, args.seed + 100 + index) split_results[name] = evaluate_probe_bank(bank, eval_set) payload = { "method": "flowmo", "checkpoint": args.checkpoint_name, "train_source": args.train_source, "train_windows": args.train_windows, "eval_windows": args.eval_windows, "ridge_alpha": args.ridge_alpha, "feature_sets": { "z": "short-history object-motion latent", "c": "long-history ambient-drift context", "z_c": "concatenated state and context", }, "targets": { "momentum": ["vx", "vy", "omega"], "local_flow": ["flow_x", "flow_y"], "episode_drift": ["mean_flow_x", "mean_flow_y"], }, "splits": split_results, } out = Path(args.out) out.parent.mkdir(parents=True, exist_ok=True) out.write_text(json.dumps(payload, indent=2)) print(json.dumps(payload, indent=2)) if __name__ == "__main__": main()