| """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() |
|
|
