Initial FlowMo-WM public code release

README.md

@@ -1,3 +1,41 @@
----
-license: mit
----
+# FlowMo: Flow-Momentum World Model
+
+FlowMo is a clean-image world-model benchmark for surface vehicles under hidden water drift. The proposed model separates short-history endogenous state and momentum from long-history exogenous drift context, then evaluates whether that factorization improves rollout prediction and closed-loop planning.
+
+## Paper Pipeline
+
+Run the complete paper-facing experiment:
+
+```bash
+python -m experiments.run_paper_image_pipeline
+```
+
+The default command trains all learned world models, evaluates prediction, runs FlowMo latent probes, evaluates planning on all configured tasks and boat morphologies, generates GIFs, and writes:
+
+```text
+experiments/reports/paper_prediction_seen_flow_diagnostic.json
+experiments/reports/paper_prediction_unseen_flow.json
+experiments/reports/paper_prediction_unseen_boat_params.json
+experiments/reports/paper_flowmo_latent_probes.json
+experiments/reports/paper_planning/
+experiments/reports/paper_report.md
+```
+
+Images are rendered online from simulator states. Model inputs are clean top-down RGB frames with no flow arrows, no goal markers, no velocity vectors, and no trajectory overlays.
+
+## Compared Methods
+
+- `flowmo`: proposed Flow-Momentum World Model.
+- `leworldmodel`: LeWorldModel-style JEPA latent predictor.
+- `planet`: PlaNet-style RSSM world model.
+- `tdmpc2`: TD-MPC2-style latent dynamics world model.
+- `pid_los_controller`, `physics_mpc_no_flow`, `current_estimator_mpc`, `oracle_flow_mpc`: traditional planning/control baselines.
+
+Baseline fidelity and naming rules are documented in `experiments/BASELINES.md`.
+The complete paper experiment matrix is documented in `experiments/EXPERIMENT_MATRIX.md`.
+
+## Tests
+
+```bash
+python -m pytest -q
+```

data/paper/dataset_card.md

@@ -0,0 +1,47 @@
+# FlowMo Paper Dataset
+
+This directory contains the canonical datasets used by the paper-facing
+experiments. File names are stable and intentionally do not include version
+suffixes; when the dataset is regenerated, these files are replaced in place.
+
+## Splits
+
+| File | Role |
+| --- | --- |
+| `train.npz` | Training split shared by all learned world models. |
+| `test_unseen_flow.npz` | Primary split with unseen flow families. |
+| `test_unseen_boat_params.npz` | Primary split with unseen boat dynamics. |
+| `diagnostic_seen_flow.npz` | Seen-flow-family diagnostic split used only for optimization sanity checks. |
+
+## Sizes
+
+| File | Episodes | Steps per episode |
+| --- | ---: | ---: |
+| `train.npz` | 2400 | 300 |
+| `diagnostic_seen_flow.npz` | 480 | 300 |
+| `test_unseen_flow.npz` | 480 | 300 |
+| `test_unseen_boat_params.npz` | 480 | 300 |
+
+## Stored Arrays
+
+Each `.npz` stores low-dimensional simulator state and metadata. Image-input
+models receive clean rendered images generated online from the same states.
+The image observation contains only the boat and clean workspace; flow vectors,
+velocity arrows, and visualization overlays are not part of the model input.
+
+All learned world models use the same split files, the same window sampling
+rules, the same image renderer, and the same train/evaluation budgets.
+
+## Flow Families
+
+The training split and seen-flow-family diagnostic split use `noflow`, `uniform`,
+`slowly_varying`, `vortex_center`, `gradient`, and `turbulent_patch` flows.
+The unseen-flow split uses `noflow`, `shear`, `moving_vortex`, and
+`random_fourier` flows. The unseen-boat-dynamics split uses the training flow
+families with held-out boat mass, drag, inertia, and actuator-delay ranges.
+
+All splits use the fixed paper flow-strength constants in `driftwm/sim/flow.py`,
+both boat morphologies, and clean image observations without flow overlays.
+`flow_pool_size=80` means that each nonzero flow family is represented by 80
+hidden flow conditions before trajectories are sampled; it is a dataset
+diversity constant, not a model input or experiment mode.

data/paper/diagnostic_seen_flow.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e8366dd25e98b81b3b159a600e4e158b8b719e0ad41a0e3fbc915b6af34ce924
+size 6601980

data/paper/generation_config.json

@@ -0,0 +1,49 @@
+{
+  "boats": [
+    "twin",
+    "triangle"
+  ],
+  "train_flow_types": [
+    "noflow",
+    "uniform",
+    "slowly_varying",
+    "vortex_center",
+    "gradient",
+    "turbulent_patch"
+  ],
+  "trajectory_types": [
+    "noflow_random_action",
+    "noflow_action_then_zero",
+    "flow_zero_action",
+    "flow_active_control",
+    "flow_waypoint_control"
+  ],
+  "episodes": {
+    "train": 2400,
+    "diagnostic_seen_flow": 480,
+    "test_unseen_flow": 480,
+    "test_unseen_boat_params": 480
+  },
+  "steps": 300,
+  "flow_pool_size": 80,
+  "boundary": "terminate",
+  "seeds": {
+    "train": 4301,
+    "diagnostic_seen_flow": 4302,
+    "test_unseen_flow": 4303,
+    "test_unseen_boat_params": 4304
+  },
+  "image_size": 160,
+  "visual_scale": 2.5,
+  "workspace": [
+    0.0,
+    10.0,
+    0.0,
+    10.0
+  ],
+  "unseen_flow_types": [
+    "shear",
+    "moving_vortex",
+    "random_fourier"
+  ]
+}

data/paper/test_unseen_boat_params.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fa4161b8917935a7a159f3e1604d4551c455f55160d8dd83b234bd726fc154a4
+size 6400836

data/paper/test_unseen_flow.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3beab06303edd2a81c802972d1055fed3ceffa54ac90f4686b11b08479682b87
+size 6493086

data/paper/train.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3632528ec850fd526e028f49bf1641d8f9f9399dc6c07e495b9351df29f76f2a
+size 32273806

driftwm/__init__.py

@@ -0,0 +1,3 @@
+"""Passive-drift disentangled world model package."""
+
+__version__ = "0.1.0"

driftwm/data/__init__.py

@@ -0,0 +1 @@
+"""Dataset generation and loading utilities."""

driftwm/data/generate.py

@@ -0,0 +1,260 @@
+from __future__ import annotations
+
+import argparse
+import copy
+import json
+from pathlib import Path
+
+import numpy as np
+
+from driftwm.sim.env import SurfaceBoatEnv
+from driftwm.sim.flow import sample_flow
+from driftwm.utils import ensure_dir, pad_action
+
+
+BOAT_TO_ID = {"twin": 0, "triangle": 1}
+ID_TO_BOAT = {v: k for k, v in BOAT_TO_ID.items()}
+FLOW_TO_ID = {
+    "noflow": 0,
+    "uniform": 1,
+    "slowly_varying": 2,
+    "vortex": 3,
+    "vortex_center": 4,
+    "gradient": 5,
+    "turbulent_patch": 6,
+    "shear": 7,
+    "moving_vortex": 8,
+    "random_fourier": 9,
+}
+ID_TO_FLOW = {v: k for k, v in FLOW_TO_ID.items()}
+TRAJ_TO_ID = {
+    "noflow_random_action": 0,
+    "noflow_action_then_zero": 1,
+    "flow_zero_action": 2,
+    "flow_active_control": 3,
+    "flow_waypoint_control": 4,
+}
+ID_TO_TRAJ = {v: k for k, v in TRAJ_TO_ID.items()}
+PAPER_FLOW_POOL_SIZE = 80
+
+
+def smooth_random_actions(rng: np.random.Generator, steps: int, action_dim: int, scale: float = 1.0) -> np.ndarray:
+    actions = np.zeros((steps, action_dim), dtype=np.float32)
+    current = rng.uniform(-0.2, 0.2, size=action_dim).astype(np.float32)
+    for t in range(steps):
+        if t % 12 == 0:
+            target = rng.uniform(-1.0, 1.0, size=action_dim).astype(np.float32)
+        current = 0.82 * current + 0.18 * target
+        actions[t] = np.clip(scale * current, -1.0, 1.0)
+    return actions
+
+
+def simple_goal_policy(env: SurfaceBoatEnv, goal: np.ndarray, rng: np.random.Generator) -> np.ndarray:
+    pos = env.state[:2]
+    theta = float(env.state[2])
+    delta = goal - pos
+    target_angle = float(np.arctan2(delta[1], delta[0]))
+    err = float(np.arctan2(np.sin(target_angle - theta), np.cos(target_angle - theta)))
+    if env.spec.name == "twin":
+        forward = np.clip(0.35 + 0.55 * np.cos(err), -0.4, 0.9)
+        turn = np.clip(-0.75 * np.sin(err), -0.8, 0.8)
+        action = np.array([forward + turn, forward - turn], dtype=np.float32)
+    else:
+        base = np.array([np.sin(err), -0.5 * np.sin(err) + 0.35, -0.5 * np.sin(err) - 0.35], dtype=np.float32)
+        action = 0.7 * base
+    action += rng.normal(0.0, 0.08, size=env.action_dim).astype(np.float32)
+    return np.clip(action, -1.0, 1.0)
+
+
+def choose_trajectory_type(rng: np.random.Generator) -> str:
+    keys = list(TRAJ_TO_ID)
+    weights = np.array([0.14, 0.14, 0.14, 0.28, 0.30], dtype=np.float64)
+    weights = weights[: len(keys)]
+    weights = weights / weights.sum()
+    return str(rng.choice(keys, p=weights))
+
+
+def make_flow_pool(
+    flow_types: list[str],
+    rng: np.random.Generator,
+    workspace: tuple[float, float, float, float],
+    pool_size: int = PAPER_FLOW_POOL_SIZE,
+) -> dict[str, list]:
+    pool = {ft: [] for ft in flow_types}
+    fid = 1
+    for ft in flow_types:
+        if ft == "noflow":
+            pool[ft] = [sample_flow("noflow", rng, flow_id=0, workspace=workspace)]
+            continue
+        for _ in range(pool_size):
+            pool[ft].append(sample_flow(ft, rng, flow_id=fid, workspace=workspace))
+            fid += 1
+    return pool
+
+
+def generate_dataset(
+    boats: list[str],
+    flow_types: list[str],
+    episodes: int,
+    steps: int,
+    out: str | Path,
+    seed: int = 0,
+    workspace: tuple[float, float, float, float] = (0.0, 10.0, 0.0, 10.0),
+    boundary: str = "terminate",
+    randomize_params: bool = True,
+    flow_pool_size: int = PAPER_FLOW_POOL_SIZE,
+    unseen_boat_params: bool = False,
+) -> None:
+    rng = np.random.default_rng(seed)
+    out = Path(out)
+    ensure_dir(out.parent)
+    flow_pool = make_flow_pool(flow_types, rng, workspace, flow_pool_size)
+
+    obs = np.zeros((episodes, steps + 1, 4), dtype=np.float32)
+    actions = np.zeros((episodes, steps, 3), dtype=np.float32)
+    states = np.zeros((episodes, steps + 1, 9), dtype=np.float32)
+    true_flow = np.zeros((episodes, steps + 1, 2), dtype=np.float32)
+    boat_ids = np.zeros((episodes,), dtype=np.int64)
+    action_dims = np.zeros((episodes,), dtype=np.int64)
+    flow_type_ids = np.zeros((episodes,), dtype=np.int64)
+    flow_ids = np.zeros((episodes,), dtype=np.int64)
+    traj_type_ids = np.zeros((episodes,), dtype=np.int64)
+
+    env = SurfaceBoatEnv(
+        seed=seed,
+        episode_steps=steps,
+        workspace=workspace,
+        boundary=boundary,
+        randomize_params=randomize_params,
+    )
+
+    for ep in range(episodes):
+        traj_type = choose_trajectory_type(rng)
+        boat = boats[int(rng.integers(0, len(boats)))]
+        if traj_type.startswith("noflow"):
+            flow_type = "noflow"
+        else:
+            available = [ft for ft in flow_types if ft != "noflow"] or ["uniform"]
+            flow_type = available[int(rng.integers(0, len(available)))]
+        flow_template = flow_pool[flow_type][int(rng.integers(0, len(flow_pool[flow_type])))]
+        flow = copy.deepcopy(flow_template)
+        random_velocity = traj_type != "flow_zero_action"
+        env.reset(
+            boat=boat,
+            flow_type=flow_type,
+            flow=flow,
+            random_velocity=random_velocity,
+            randomize_params=randomize_params,
+        )
+        if unseen_boat_params:
+            for key in list(env.params):
+                if key in {"mass", "inertia", "actuator_tau"}:
+                    factor = rng.choice([rng.uniform(0.55, 0.72), rng.uniform(1.45, 1.85)])
+                else:
+                    factor = rng.choice([rng.uniform(0.45, 0.70), rng.uniform(1.45, 1.90)])
+                env.params[key] = float(env.params[key] * factor)
+
+        if traj_type == "noflow_random_action":
+            planned_actions = smooth_random_actions(rng, steps, env.action_dim, scale=1.0)
+        elif traj_type == "noflow_action_then_zero":
+            planned_actions = np.zeros((steps, env.action_dim), dtype=np.float32)
+            push_len = max(12, min(50, steps // 4))
+            planned_actions[:push_len] = smooth_random_actions(rng, push_len, env.action_dim, scale=0.95)
+        elif traj_type == "flow_zero_action":
+            planned_actions = np.zeros((steps, env.action_dim), dtype=np.float32)
+            env.state[3:6] = 0.0
+        else:
+            planned_actions = smooth_random_actions(rng, steps, env.action_dim, scale=0.7)
+            goal = rng.uniform([1.5, 1.5], [8.5, 8.5]).astype(np.float32)
+            waypoints = rng.uniform([1.5, 1.5], [8.5, 8.5], size=(4, 2)).astype(np.float32)
+            waypoint_idx = 0
+
+        boat_ids[ep] = BOAT_TO_ID[boat]
+        action_dims[ep] = env.action_dim
+        flow_type_ids[ep] = FLOW_TO_ID[flow_type]
+        flow_ids[ep] = int(env.flow.flow_id)
+        traj_type_ids[ep] = TRAJ_TO_ID[traj_type]
+        obs[ep, 0] = env.observation()
+        states[ep, 0, : 6 + env.action_dim] = env.full_state()
+        true_flow[ep, 0] = env.flow_at(env.state[:2])
+
+        done = False
+        for t in range(steps):
+            if traj_type == "flow_waypoint_control":
+                goal = waypoints[min(waypoint_idx, len(waypoints) - 1)]
+                if np.linalg.norm(env.state[:2] - goal) < 0.75 and waypoint_idx < len(waypoints) - 1:
+                    waypoint_idx += 1
+                    goal = waypoints[waypoint_idx]
+                if rng.random() < 0.88:
+                    action = simple_goal_policy(env, goal, rng)
+                else:
+                    action = planned_actions[t]
+            elif traj_type == "flow_active_control" and rng.random() < 0.75:
+                action = simple_goal_policy(env, goal, rng)
+            else:
+                action = planned_actions[t]
+            actions[ep, t] = pad_action(action)
+            ob, _, done, _ = env.step(action)
+            obs[ep, t + 1] = ob
+            states[ep, t + 1, : 6 + env.action_dim] = env.full_state()
+            true_flow[ep, t + 1] = env.flow_at(env.state[:2])
+            if done:
+                obs[ep, t + 2 :] = obs[ep, t + 1]
+                states[ep, t + 2 :] = states[ep, t + 1]
+                true_flow[ep, t + 2 :] = true_flow[ep, t + 1]
+                break
+
+    metadata = {
+        "boats": BOAT_TO_ID,
+        "flows": FLOW_TO_ID,
+        "trajectories": TRAJ_TO_ID,
+        "steps": steps,
+        "workspace": list(workspace),
+        "boundary": boundary,
+        "seed": seed,
+        "max_action_dim": 3,
+        "flow_pool_size": flow_pool_size,
+        "unseen_boat_params": bool(unseen_boat_params),
+    }
+    np.savez_compressed(
+        out,
+        obs=obs,
+        actions=actions,
+        states=states,
+        true_flow=true_flow,
+        boat_ids=boat_ids,
+        action_dims=action_dims,
+        flow_type_ids=flow_type_ids,
+        flow_ids=flow_ids,
+        traj_type_ids=traj_type_ids,
+        metadata=json.dumps(metadata),
+    )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--boats", nargs="+", choices=list(BOAT_TO_ID), default=["twin"])
+    parser.add_argument("--flow-types", nargs="+", choices=list(FLOW_TO_ID), default=["noflow", "uniform"])
+    parser.add_argument("--episodes", type=int, default=2000)
+    parser.add_argument("--steps", type=int, default=200)
+    parser.add_argument("--out", required=True)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--boundary", choices=["terminate", "bounce", "clip"], default="terminate")
+    parser.add_argument("--no-randomize-params", action="store_true")
+    parser.add_argument("--unseen-boat-params", action="store_true")
+    args = parser.parse_args()
+    generate_dataset(
+        boats=args.boats,
+        flow_types=args.flow_types,
+        episodes=args.episodes,
+        steps=args.steps,
+        out=args.out,
+        seed=args.seed,
+        boundary=args.boundary,
+        randomize_params=not args.no_randomize_params,
+        unseen_boat_params=args.unseen_boat_params,
+    )
+
+
+if __name__ == "__main__":
+    main()

driftwm/sim/__init__.py

@@ -0,0 +1 @@
+"""Simulation components for 2D surface vehicles."""

driftwm/sim/boat.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+
+
+@dataclass(frozen=True)
+class BoatSpec:
+    name: str
+    action_dim: int
+    thruster_positions: np.ndarray
+    thruster_dirs: np.ndarray
+    hull_vertices: np.ndarray
+
+
+def get_boat_spec(name: str) -> BoatSpec:
+    name = name.lower()
+    if name == "twin":
+        length = 0.62
+        beam = 0.34
+        positions = np.array(
+            [[-length / 2.0, beam / 2.0], [-length / 2.0, -beam / 2.0]],
+            dtype=np.float32,
+        )
+        dirs = np.array([[1.0, 0.0], [1.0, 0.0]], dtype=np.float32)
+        hull = np.array(
+            [
+                [length / 2.0, 0.0],
+                [0.12, beam / 2.0],
+                [-length / 2.0, beam / 2.0],
+                [-length / 2.0, -beam / 2.0],
+                [0.12, -beam / 2.0],
+            ],
+            dtype=np.float32,
+        )
+        return BoatSpec("twin", 2, positions, dirs, hull)
+    if name == "triangle":
+        radius = 0.34
+        phis = np.array([0.0, 2.0 * np.pi / 3.0, 4.0 * np.pi / 3.0], dtype=np.float32)
+        positions = np.stack([radius * np.cos(phis), radius * np.sin(phis)], axis=-1)
+        dirs = np.stack([-np.sin(phis), np.cos(phis)], axis=-1).astype(np.float32)
+        hull = np.stack([0.44 * np.cos(phis), 0.44 * np.sin(phis)], axis=-1).astype(np.float32)
+        return BoatSpec("triangle", 3, positions.astype(np.float32), dirs, hull)
+    raise ValueError(f"unknown boat type: {name}")
+
+
+def default_boat_params(name: str) -> dict[str, float]:
+    if name == "triangle":
+        return {
+            "mass": 1.15,
+            "inertia": 0.18,
+            "t_max": 1.25,
+            "drag_linear_x": 0.62,
+            "drag_linear_y": 0.72,
+            "drag_quad_x": 0.18,
+            "drag_quad_y": 0.22,
+            "drag_angular": 0.32,
+            "drag_angular_quad": 0.08,
+            "actuator_tau": 0.22,
+        }
+    return {
+        "mass": 1.0,
+        "inertia": 0.13,
+        "t_max": 1.1,
+        "drag_linear_x": 0.55,
+        "drag_linear_y": 0.88,
+        "drag_quad_x": 0.16,
+        "drag_quad_y": 0.25,
+        "drag_angular": 0.28,
+        "drag_angular_quad": 0.06,
+        "actuator_tau": 0.18,
+    }
+
+
+def sample_boat_params(name: str, rng: np.random.Generator, randomize: bool = True) -> dict[str, float]:
+    params = default_boat_params(name)
+    if not randomize:
+        return params
+    sampled: dict[str, float] = {}
+    for key, value in params.items():
+        if key == "actuator_tau":
+            factor = rng.uniform(0.75, 1.35)
+        elif key in {"mass", "inertia"}:
+            factor = rng.uniform(0.85, 1.25)
+        else:
+            factor = rng.uniform(0.75, 1.30)
+        sampled[key] = float(value * factor)
+    return sampled
+
+
+def thruster_wrench_body(
+    spec: BoatSpec,
+    actuator_state: np.ndarray,
+    params: dict[str, float],
+) -> tuple[np.ndarray, float]:
+    u = np.asarray(actuator_state, dtype=np.float32)
+    forces = params["t_max"] * u[:, None] * spec.thruster_dirs
+    force_body = forces.sum(axis=0)
+    torques = spec.thruster_positions[:, 0] * forces[:, 1] - spec.thruster_positions[:, 1] * forces[:, 0]
+    return force_body.astype(np.float32), float(torques.sum())
+
+
+def boat_metadata(name: str, params: dict[str, float]) -> dict[str, Any]:
+    return {
+        "boat_type": name,
+        "action_dim": get_boat_spec(name).action_dim,
+        "params": {k: float(v) for k, v in params.items()},
+    }

driftwm/sim/dynamics.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+import numpy as np
+
+from driftwm.sim.boat import BoatSpec, thruster_wrench_body
+from driftwm.utils import wrap_angle
+
+
+def rot_body_to_world(theta: float) -> np.ndarray:
+    c = np.cos(theta)
+    s = np.sin(theta)
+    return np.array([[c, -s], [s, c]], dtype=np.float32)
+
+
+def step_dynamics(
+    state: np.ndarray,
+    action: np.ndarray,
+    spec: BoatSpec,
+    params: dict[str, float],
+    flow_velocity: np.ndarray,
+    dt: float,
+    workspace: tuple[float, float, float, float],
+    boundary: str = "terminate",
+) -> tuple[np.ndarray, bool]:
+    action = np.clip(np.asarray(action, dtype=np.float32), -1.0, 1.0)
+    state = np.asarray(state, dtype=np.float32).copy()
+    x, y, theta, vx, vy, omega = state[:6]
+    u = state[6 : 6 + spec.action_dim]
+
+    tau_a = max(params["actuator_tau"], 1e-3)
+    alpha = min(1.0, dt / tau_a)
+    u_next = u + alpha * (action - u)
+
+    rot = rot_body_to_world(float(theta))
+    vel_world = np.array([vx, vy], dtype=np.float32)
+    rel_world = vel_world - np.asarray(flow_velocity, dtype=np.float32)
+    rel_body = rot.T @ rel_world
+
+    d1 = np.array([params["drag_linear_x"], params["drag_linear_y"]], dtype=np.float32)
+    d2 = np.array([params["drag_quad_x"], params["drag_quad_y"]], dtype=np.float32)
+    drag_body = -d1 * rel_body - d2 * np.abs(rel_body) * rel_body
+
+    thrust_body, tau_thr = thruster_wrench_body(spec, u_next, params)
+    force_world = rot @ (thrust_body + drag_body)
+
+    tau_drag = -params["drag_angular"] * omega - params["drag_angular_quad"] * abs(float(omega)) * omega
+    tau_total = tau_thr + tau_drag
+
+    vel_next = vel_world + dt * force_world / params["mass"]
+    pos_next = np.array([x, y], dtype=np.float32) + dt * vel_next
+    omega_next = omega + dt * tau_total / params["inertia"]
+    theta_next = wrap_angle(theta + dt * omega_next)
+
+    xmin, xmax, ymin, ymax = workspace
+    done = bool(pos_next[0] < xmin or pos_next[0] > xmax or pos_next[1] < ymin or pos_next[1] > ymax)
+    if done and boundary == "bounce":
+        restitution = 0.45
+        if pos_next[0] < xmin or pos_next[0] > xmax:
+            vel_next[0] *= -restitution
+        if pos_next[1] < ymin or pos_next[1] > ymax:
+            vel_next[1] *= -restitution
+        pos_next[0] = np.clip(pos_next[0], xmin, xmax)
+        pos_next[1] = np.clip(pos_next[1], ymin, ymax)
+        done = False
+    elif done and boundary == "clip":
+        pos_next[0] = np.clip(pos_next[0], xmin, xmax)
+        pos_next[1] = np.clip(pos_next[1], ymin, ymax)
+        done = False
+
+    next_state = state.copy()
+    next_state[:6] = np.array([pos_next[0], pos_next[1], theta_next, vel_next[0], vel_next[1], omega_next], dtype=np.float32)
+    next_state[6 : 6 + spec.action_dim] = u_next
+    return next_state.astype(np.float32), done

driftwm/sim/env.py

@@ -0,0 +1,142 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+
+from driftwm.sim.boat import BoatSpec, get_boat_spec, sample_boat_params
+from driftwm.sim.dynamics import step_dynamics
+from driftwm.sim.flow import Flow, sample_flow
+from driftwm.utils import obs_from_state
+
+
+@dataclass
+class EnvConfig:
+    boat: str = "twin"
+    flow_type: str = "noflow"
+    dt: float = 0.05
+    episode_steps: int = 200
+    workspace: tuple[float, float, float, float] = (0.0, 10.0, 0.0, 10.0)
+    boundary: str = "terminate"
+    randomize_params: bool = True
+
+
+class SurfaceBoatEnv:
+    def __init__(
+        self,
+        boat: str = "twin",
+        flow_type: str = "noflow",
+        dt: float = 0.05,
+        episode_steps: int = 200,
+        workspace: tuple[float, float, float, float] = (0.0, 10.0, 0.0, 10.0),
+        boundary: str = "terminate",
+        randomize_params: bool = True,
+        seed: int | None = None,
+    ):
+        self.config = EnvConfig(boat, flow_type, dt, episode_steps, workspace, boundary, randomize_params)
+        self.rng = np.random.default_rng(seed)
+        self.spec: BoatSpec = get_boat_spec(boat)
+        self.params: dict[str, float] = sample_boat_params(boat, self.rng, randomize_params)
+        self.flow: Flow = sample_flow(flow_type, self.rng, flow_id=1, workspace=workspace)
+        self.state = np.zeros(6 + self.spec.action_dim, dtype=np.float32)
+        self.t = 0
+        self.time = 0.0
+        self.last_flow_velocity = np.zeros(2, dtype=np.float32)
+
+    @property
+    def action_dim(self) -> int:
+        return self.spec.action_dim
+
+    @property
+    def workspace(self) -> tuple[float, float, float, float]:
+        return self.config.workspace
+
+    def reset(
+        self,
+        *,
+        boat: str | None = None,
+        flow_type: str | None = None,
+        flow: Flow | None = None,
+        flow_id: int | None = None,
+        random_velocity: bool = True,
+        initial_state: np.ndarray | None = None,
+        randomize_params: bool | None = None,
+    ) -> tuple[np.ndarray, dict[str, Any]]:
+        if boat is not None:
+            self.config.boat = boat
+        if flow_type is not None:
+            self.config.flow_type = flow_type
+        if randomize_params is not None:
+            self.config.randomize_params = randomize_params
+
+        self.spec = get_boat_spec(self.config.boat)
+        self.params = sample_boat_params(self.config.boat, self.rng, self.config.randomize_params)
+        if flow is not None:
+            self.flow = flow
+        else:
+            fid = int(flow_id if flow_id is not None else self.rng.integers(1, 2_000_000))
+            self.flow = sample_flow(self.config.flow_type, self.rng, fid, self.config.workspace)
+
+        if initial_state is not None:
+            self.state = np.asarray(initial_state, dtype=np.float32).copy()
+        else:
+            xmin, xmax, ymin, ymax = self.config.workspace
+            margin = 1.0
+            pos = np.array(
+                [self.rng.uniform(xmin + margin, xmax - margin), self.rng.uniform(ymin + margin, ymax - margin)],
+                dtype=np.float32,
+            )
+            theta = self.rng.uniform(-np.pi, np.pi)
+            vel = self.rng.uniform(-0.12, 0.12, size=2).astype(np.float32) if random_velocity else np.zeros(2, dtype=np.float32)
+            omega = float(self.rng.uniform(-0.15, 0.15)) if random_velocity else 0.0
+            self.state = np.zeros(6 + self.spec.action_dim, dtype=np.float32)
+            self.state[:6] = np.array([pos[0], pos[1], theta, vel[0], vel[1], omega], dtype=np.float32)
+
+        self.t = 0
+        self.time = 0.0
+        self.last_flow_velocity = self.flow.velocity(self.state[:2], self.time)
+        return self.observation(), self.info()
+
+    def observation(self) -> np.ndarray:
+        return obs_from_state(self.state[:6])
+
+    def full_state(self) -> np.ndarray:
+        return self.state.copy()
+
+    def flow_at(self, pos: np.ndarray) -> np.ndarray:
+        return self.flow.velocity(np.asarray(pos, dtype=np.float32), self.time)
+
+    def step(self, action: np.ndarray) -> tuple[np.ndarray, float, bool, dict[str, Any]]:
+        action = np.asarray(action, dtype=np.float32)[: self.action_dim]
+        flow_velocity = self.flow.velocity(self.state[:2], self.time)
+        self.last_flow_velocity = flow_velocity.astype(np.float32)
+        self.state, boundary_done = step_dynamics(
+            self.state,
+            action,
+            self.spec,
+            self.params,
+            flow_velocity,
+            self.config.dt,
+            self.config.workspace,
+            self.config.boundary,
+        )
+        self.flow.step(self.config.dt, self.rng)
+        self.t += 1
+        self.time += self.config.dt
+        timeout = self.t >= self.config.episode_steps
+        done = boundary_done or timeout
+        reward = 0.0
+        return self.observation(), reward, done, self.info()
+
+    def info(self) -> dict[str, Any]:
+        meta = {
+            "t": self.t,
+            "time": self.time,
+            "boat_type": self.spec.name,
+            "action_dim": self.action_dim,
+            "flow_velocity": self.last_flow_velocity.astype(float).tolist(),
+            "params": {k: float(v) for k, v in self.params.items()},
+        }
+        meta.update(self.flow.metadata())
+        return meta

driftwm/sim/flow.py

@@ -0,0 +1,428 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+
+
+PAPER_FLOW = {
+    "uniform_min": 0.03,
+    "uniform_max": 0.24,
+    "slow_max": 0.26,
+    "slow_noise": 0.0035,
+    "vortex_base_max": 0.12,
+    "vortex_gamma": 0.14,
+    "vortex_max": 0.34,
+    "gradient_base_min": 0.01,
+    "gradient_base_max": 0.16,
+    "gradient_matrix_std": 0.022,
+    "gradient_max": 0.34,
+    "turbulent_base_max": 0.12,
+    "turbulent_vector_std": 0.075,
+    "turbulent_max": 0.34,
+    "shear_max": 0.38,
+    "moving_vortex_max": 0.38,
+    "random_fourier_max": 0.38,
+}
+
+
+@dataclass
+class Flow:
+    name: str
+    flow_id: int
+
+    def velocity(self, pos: np.ndarray, t: float = 0.0) -> np.ndarray:
+        raise NotImplementedError
+
+    def step(self, dt: float, rng: np.random.Generator) -> None:
+        return None
+
+    def metadata(self) -> dict[str, Any]:
+        return {"flow_type": self.name, "flow_id": int(self.flow_id)}
+
+
+@dataclass
+class NoFlow(Flow):
+    def __init__(self, flow_id: int = 0):
+        super().__init__("noflow", flow_id)
+
+    def velocity(self, pos: np.ndarray, t: float = 0.0) -> np.ndarray:
+        return np.zeros_like(np.asarray(pos, dtype=np.float32))
+
+
+@dataclass
+class UniformFlow(Flow):
+    vector: np.ndarray
+
+    def __init__(self, vector: np.ndarray, flow_id: int):
+        super().__init__("uniform", flow_id)
+        self.vector = np.asarray(vector, dtype=np.float32)
+
+    def velocity(self, pos: np.ndarray, t: float = 0.0) -> np.ndarray:
+        pos = np.asarray(pos, dtype=np.float32)
+        return np.broadcast_to(self.vector, pos.shape).astype(np.float32)
+
+    def metadata(self) -> dict[str, Any]:
+        out = super().metadata()
+        out["vector"] = self.vector.astype(float).tolist()
+        return out
+
+
+@dataclass
+class SlowlyVaryingFlow(Flow):
+    vector: np.ndarray
+    rho: float
+    noise_std: float
+    max_speed: float
+
+    def __init__(
+        self,
+        vector: np.ndarray,
+        flow_id: int,
+        rho: float = 0.995,
+        noise_std: float = 0.005,
+        max_speed: float = 0.35,
+    ):
+        super().__init__("slowly_varying", flow_id)
+        self.vector = np.asarray(vector, dtype=np.float32)
+        self.rho = float(rho)
+        self.noise_std = float(noise_std)
+        self.max_speed = float(max_speed)
+
+    def velocity(self, pos: np.ndarray, t: float = 0.0) -> np.ndarray:
+        pos = np.asarray(pos, dtype=np.float32)
+        return np.broadcast_to(self.vector, pos.shape).astype(np.float32)
+
+    def step(self, dt: float, rng: np.random.Generator) -> None:
+        noise = rng.normal(0.0, self.noise_std, size=2).astype(np.float32)
+        self.vector = self.rho * self.vector + np.sqrt(max(0.0, 1.0 - self.rho**2)) * noise
+        speed = float(np.linalg.norm(self.vector))
+        if speed > self.max_speed:
+            self.vector = self.vector / speed * self.max_speed
+
+    def metadata(self) -> dict[str, Any]:
+        out = super().metadata()
+        out.update({"vector": self.vector.astype(float).tolist(), "rho": self.rho})
+        return out
+
+
+@dataclass
+class VortexFlow(Flow):
+    base: np.ndarray
+    center: np.ndarray
+    gamma: float
+    radius_eps: float
+    max_speed: float
+
+    def __init__(
+        self,
+        base: np.ndarray,
+        center: np.ndarray,
+        gamma: float,
+        flow_id: int,
+        radius_eps: float = 0.30,
+        max_speed: float = 0.50,
+        name: str = "vortex",
+    ):
+        super().__init__(name, flow_id)
+        self.base = np.asarray(base, dtype=np.float32)
+        self.center = np.asarray(center, dtype=np.float32)
+        self.gamma = float(gamma)
+        self.radius_eps = float(radius_eps)
+        self.max_speed = float(max_speed)
+
+    def velocity(self, pos: np.ndarray, t: float = 0.0) -> np.ndarray:
+        pos = np.asarray(pos, dtype=np.float32)
+        rel = pos - self.center
+        denom = np.sum(rel * rel, axis=-1, keepdims=True) + self.radius_eps**2
+        swirl = self.gamma * np.concatenate([-rel[..., 1:2], rel[..., 0:1]], axis=-1) / denom
+        vel = self.base + swirl
+        speed = np.linalg.norm(vel, axis=-1, keepdims=True)
+        scale = np.minimum(1.0, self.max_speed / np.maximum(speed, 1e-6))
+        return (vel * scale).astype(np.float32)
+
+    def metadata(self) -> dict[str, Any]:
+        out = super().metadata()
+        out.update(
+            {
+                "base": self.base.astype(float).tolist(),
+                "center": self.center.astype(float).tolist(),
+                "gamma": self.gamma,
+            }
+        )
+        return out
+
+
+@dataclass
+class GradientFlow(Flow):
+    base: np.ndarray
+    center: np.ndarray
+    matrix: np.ndarray
+    max_speed: float
+
+    def __init__(self, base: np.ndarray, center: np.ndarray, matrix: np.ndarray, flow_id: int, max_speed: float = 0.55):
+        super().__init__("gradient", flow_id)
+        self.base = np.asarray(base, dtype=np.float32)
+        self.center = np.asarray(center, dtype=np.float32)
+        self.matrix = np.asarray(matrix, dtype=np.float32)
+        self.max_speed = float(max_speed)
+
+    def velocity(self, pos: np.ndarray, t: float = 0.0) -> np.ndarray:
+        pos = np.asarray(pos, dtype=np.float32)
+        rel = pos - self.center
+        vel = self.base + rel @ self.matrix.T
+        speed = np.linalg.norm(vel, axis=-1, keepdims=True)
+        scale = np.minimum(1.0, self.max_speed / np.maximum(speed, 1e-6))
+        return (vel * scale).astype(np.float32)
+
+    def metadata(self) -> dict[str, Any]:
+        out = super().metadata()
+        out.update({"base": self.base.astype(float).tolist(), "matrix": self.matrix.astype(float).tolist()})
+        return out
+
+
+@dataclass
+class TurbulentPatchFlow(Flow):
+    base: np.ndarray
+    centers: np.ndarray
+    vectors: np.ndarray
+    sigma: float
+    max_speed: float
+
+    def __init__(
+        self,
+        base: np.ndarray,
+        centers: np.ndarray,
+        vectors: np.ndarray,
+        flow_id: int,
+        sigma: float = 1.15,
+        max_speed: float = 0.55,
+    ):
+        super().__init__("turbulent_patch", flow_id)
+        self.base = np.asarray(base, dtype=np.float32)
+        self.centers = np.asarray(centers, dtype=np.float32)
+        self.vectors = np.asarray(vectors, dtype=np.float32)
+        self.sigma = float(sigma)
+        self.max_speed = float(max_speed)
+
+    def velocity(self, pos: np.ndarray, t: float = 0.0) -> np.ndarray:
+        pos = np.asarray(pos, dtype=np.float32)
+        rel = pos[..., None, :] - self.centers
+        weights = np.exp(-np.sum(rel * rel, axis=-1, keepdims=True) / (2.0 * self.sigma**2))
+        perturb = np.sum(weights * self.vectors, axis=-2)
+        vel = self.base + perturb
+        speed = np.linalg.norm(vel, axis=-1, keepdims=True)
+        scale = np.minimum(1.0, self.max_speed / np.maximum(speed, 1e-6))
+        return (vel * scale).astype(np.float32)
+
+    def metadata(self) -> dict[str, Any]:
+        out = super().metadata()
+        out.update({"base": self.base.astype(float).tolist(), "centers": self.centers.astype(float).tolist()})
+        return out
+
+
+@dataclass
+class ShearFlow(Flow):
+    base: np.ndarray
+    center_y: float
+    shear: float
+    max_speed: float
+
+    def __init__(self, base: np.ndarray, center_y: float, shear: float, flow_id: int, max_speed: float = 0.60):
+        super().__init__("shear", flow_id)
+        self.base = np.asarray(base, dtype=np.float32)
+        self.center_y = float(center_y)
+        self.shear = float(shear)
+        self.max_speed = float(max_speed)
+
+    def velocity(self, pos: np.ndarray, t: float = 0.0) -> np.ndarray:
+        pos = np.asarray(pos, dtype=np.float32)
+        vel = np.broadcast_to(self.base, pos.shape).copy()
+        vel[..., 0] += self.shear * (pos[..., 1] - self.center_y)
+        speed = np.linalg.norm(vel, axis=-1, keepdims=True)
+        scale = np.minimum(1.0, self.max_speed / np.maximum(speed, 1e-6))
+        return (vel * scale).astype(np.float32)
+
+    def metadata(self) -> dict[str, Any]:
+        out = super().metadata()
+        out.update({"base": self.base.astype(float).tolist(), "center_y": self.center_y, "shear": self.shear})
+        return out
+
+
+@dataclass
+class MovingVortexFlow(VortexFlow):
+    center_velocity: np.ndarray
+    workspace: tuple[float, float, float, float]
+
+    def __init__(
+        self,
+        base: np.ndarray,
+        center: np.ndarray,
+        gamma: float,
+        center_velocity: np.ndarray,
+        flow_id: int,
+        workspace: tuple[float, float, float, float],
+        radius_eps: float = 0.35,
+        max_speed: float = 0.60,
+    ):
+        super().__init__(base=base, center=center, gamma=gamma, flow_id=flow_id, radius_eps=radius_eps, max_speed=max_speed, name="moving_vortex")
+        self.center_velocity = np.asarray(center_velocity, dtype=np.float32)
+        self.workspace = workspace
+
+    def step(self, dt: float, rng: np.random.Generator) -> None:
+        del rng
+        self.center = self.center + dt * self.center_velocity
+        xmin, xmax, ymin, ymax = self.workspace
+        if self.center[0] < xmin + 1.0 or self.center[0] > xmax - 1.0:
+            self.center_velocity[0] *= -1.0
+        if self.center[1] < ymin + 1.0 or self.center[1] > ymax - 1.0:
+            self.center_velocity[1] *= -1.0
+        self.center[0] = np.clip(self.center[0], xmin + 1.0, xmax - 1.0)
+        self.center[1] = np.clip(self.center[1], ymin + 1.0, ymax - 1.0)
+
+    def metadata(self) -> dict[str, Any]:
+        out = super().metadata()
+        out["center_velocity"] = self.center_velocity.astype(float).tolist()
+        return out
+
+
+@dataclass
+class RandomFourierFlow(Flow):
+    base: np.ndarray
+    k: np.ndarray
+    amp: np.ndarray
+    phase: np.ndarray
+    temporal: np.ndarray
+    max_speed: float
+
+    def __init__(
+        self,
+        base: np.ndarray,
+        k: np.ndarray,
+        amp: np.ndarray,
+        phase: np.ndarray,
+        temporal: np.ndarray,
+        flow_id: int,
+        max_speed: float = 0.60,
+    ):
+        super().__init__("random_fourier", flow_id)
+        self.base = np.asarray(base, dtype=np.float32)
+        self.k = np.asarray(k, dtype=np.float32)
+        self.amp = np.asarray(amp, dtype=np.float32)
+        self.phase = np.asarray(phase, dtype=np.float32)
+        self.temporal = np.asarray(temporal, dtype=np.float32)
+        self.max_speed = float(max_speed)
+
+    def velocity(self, pos: np.ndarray, t: float = 0.0) -> np.ndarray:
+        pos = np.asarray(pos, dtype=np.float32)
+        flat = pos.reshape(-1, 2)
+        arg = flat @ self.k.T + self.phase[None, :] + float(t) * self.temporal[None, :]
+        # Divergence-free field via stream function psi: v=(dpsi/dy, -dpsi/dx).
+        coeff = self.amp[None, :] * np.cos(arg)
+        vx = np.sum(coeff * self.k[None, :, 1], axis=1)
+        vy = -np.sum(coeff * self.k[None, :, 0], axis=1)
+        vel = np.stack([vx, vy], axis=-1).reshape(pos.shape) + self.base
+        speed = np.linalg.norm(vel, axis=-1, keepdims=True)
+        scale = np.minimum(1.0, self.max_speed / np.maximum(speed, 1e-6))
+        return (vel * scale).astype(np.float32)
+
+    def metadata(self) -> dict[str, Any]:
+        out = super().metadata()
+        out.update({"base": self.base.astype(float).tolist(), "num_modes": int(self.k.shape[0])})
+        return out
+
+
+def _sample_uniform_vector(rng: np.random.Generator, min_speed: float = 0.05, max_speed: float = 0.35) -> np.ndarray:
+    speed = rng.uniform(min_speed, max_speed)
+    direction = rng.uniform(0.0, 2.0 * np.pi)
+    return np.array([speed * np.cos(direction), speed * np.sin(direction)], dtype=np.float32)
+
+
+def sample_flow(
+    flow_type: str,
+    rng: np.random.Generator,
+    flow_id: int,
+    workspace: tuple[float, float, float, float] = (0.0, 10.0, 0.0, 10.0),
+) -> Flow:
+    flow_type = flow_type.lower()
+    profile = PAPER_FLOW
+    if flow_type == "noflow":
+        return NoFlow(flow_id=0)
+    if flow_type == "uniform":
+        return UniformFlow(_sample_uniform_vector(rng, profile["uniform_min"], profile["uniform_max"]), flow_id=flow_id)
+    if flow_type in {"slow", "slowly_varying", "ou"}:
+        return SlowlyVaryingFlow(
+            _sample_uniform_vector(rng, profile["uniform_min"], profile["uniform_max"]),
+            flow_id=flow_id,
+            noise_std=profile["slow_noise"],
+            max_speed=profile["slow_max"],
+        )
+    if flow_type in {"vortex", "vortex_center"}:
+        xmin, xmax, ymin, ymax = workspace
+        if flow_type == "vortex_center":
+            center = np.array([(xmin + xmax) / 2.0, (ymin + ymax) / 2.0], dtype=np.float32)
+        else:
+            center = np.array([rng.uniform(xmin + 2.0, xmax - 2.0), rng.uniform(ymin + 2.0, ymax - 2.0)], dtype=np.float32)
+        base = _sample_uniform_vector(rng, 0.0, profile["vortex_base_max"])
+        gamma = float(rng.uniform(-profile["vortex_gamma"], profile["vortex_gamma"]))
+        return VortexFlow(base=base, center=center, gamma=gamma, flow_id=flow_id, max_speed=profile["vortex_max"], name=flow_type)
+    if flow_type in {"gradient", "gradient_flow"}:
+        xmin, xmax, ymin, ymax = workspace
+        center = np.array([(xmin + xmax) / 2.0, (ymin + ymax) / 2.0], dtype=np.float32)
+        base = _sample_uniform_vector(rng, profile["gradient_base_min"], profile["gradient_base_max"])
+        mat = rng.normal(0.0, profile["gradient_matrix_std"], size=(2, 2)).astype(np.float32)
+        return GradientFlow(base=base, center=center, matrix=mat, flow_id=flow_id, max_speed=profile["gradient_max"])
+    if flow_type in {"turbulent", "turbulent_patch", "patch"}:
+        xmin, xmax, ymin, ymax = workspace
+        base = _sample_uniform_vector(rng, 0.0, profile["turbulent_base_max"])
+        centers = np.stack(
+            [
+                rng.uniform([xmin + 1.0, ymin + 1.0], [xmax - 1.0, ymax - 1.0])
+                for _ in range(5)
+            ],
+            axis=0,
+        ).astype(np.float32)
+        vectors = rng.normal(0.0, profile["turbulent_vector_std"], size=(5, 2)).astype(np.float32)
+        return TurbulentPatchFlow(base=base, centers=centers, vectors=vectors, flow_id=flow_id, max_speed=profile["turbulent_max"])
+    if flow_type in {"shear", "shear_flow"}:
+        xmin, xmax, ymin, ymax = workspace
+        base = _sample_uniform_vector(rng, 0.0, profile["turbulent_base_max"])
+        center_y = 0.5 * (ymin + ymax)
+        shear = float(rng.uniform(-0.08, 0.08))
+        return ShearFlow(base=base, center_y=center_y, shear=shear, flow_id=flow_id, max_speed=profile["shear_max"])
+    if flow_type in {"moving_vortex", "moving-vortex"}:
+        xmin, xmax, ymin, ymax = workspace
+        center = np.array([rng.uniform(xmin + 2.0, xmax - 2.0), rng.uniform(ymin + 2.0, ymax - 2.0)], dtype=np.float32)
+        base = _sample_uniform_vector(rng, 0.0, profile["vortex_base_max"])
+        gamma = float(rng.uniform(-profile["vortex_gamma"], profile["vortex_gamma"]))
+        center_velocity = _sample_uniform_vector(rng, 0.02, 0.08)
+        return MovingVortexFlow(
+            base=base,
+            center=center,
+            gamma=gamma,
+            center_velocity=center_velocity,
+            flow_id=flow_id,
+            workspace=workspace,
+            max_speed=profile["moving_vortex_max"],
+        )
+    if flow_type in {"random_fourier", "fourier", "divfree"}:
+        base = _sample_uniform_vector(rng, 0.0, profile["turbulent_base_max"])
+        modes = 8
+        k = rng.integers(1, 5, size=(modes, 2)).astype(np.float32)
+        signs = rng.choice([-1.0, 1.0], size=(modes, 2)).astype(np.float32)
+        k = signs * k * (2.0 * np.pi / 10.0)
+        amp_std = 0.028
+        amp = rng.normal(0.0, amp_std, size=(modes,)).astype(np.float32)
+        phase = rng.uniform(0.0, 2.0 * np.pi, size=(modes,)).astype(np.float32)
+        temporal = rng.normal(0.0, 0.12, size=(modes,)).astype(np.float32)
+        return RandomFourierFlow(
+            base=base,
+            k=k,
+            amp=amp,
+            phase=phase,
+            temporal=temporal,
+            flow_id=flow_id,
+            max_speed=profile["random_fourier_max"],
+        )
+    raise ValueError(f"unknown flow type: {flow_type}")

driftwm/sim/render.py

@@ -0,0 +1,145 @@
+from __future__ import annotations
+
+import math
+from pathlib import Path
+from typing import Iterable
+
+import numpy as np
+from PIL import Image, ImageDraw
+
+from driftwm.sim.boat import BoatSpec, get_boat_spec
+from driftwm.sim.flow import Flow
+from driftwm.sim.dynamics import rot_body_to_world
+from driftwm.utils import ensure_dir
+
+
+Color = tuple[int, int, int]
+
+
+def _world_to_px(point: np.ndarray, workspace: tuple[float, float, float, float], size: int, pad: int) -> tuple[int, int]:
+    xmin, xmax, ymin, ymax = workspace
+    x = (point[0] - xmin) / max(1e-6, xmax - xmin)
+    y = (point[1] - ymin) / max(1e-6, ymax - ymin)
+    px = int(pad + x * (size - 2 * pad))
+    py = int(size - pad - y * (size - 2 * pad))
+    return px, py
+
+
+def _draw_arrow(draw: ImageDraw.ImageDraw, p0: tuple[int, int], p1: tuple[int, int], color: Color, width: int = 2) -> None:
+    draw.line([p0, p1], fill=color, width=width)
+    dx = p1[0] - p0[0]
+    dy = p1[1] - p0[1]
+    angle = math.atan2(dy, dx)
+    head = 8
+    for sign in (-1, 1):
+        a = angle + sign * 2.55
+        p = (int(p1[0] + head * math.cos(a)), int(p1[1] + head * math.sin(a)))
+        draw.line([p1, p], fill=color, width=width)
+
+
+def draw_flow_field(
+    draw: ImageDraw.ImageDraw,
+    flow: Flow,
+    workspace: tuple[float, float, float, float],
+    size: int,
+    pad: int,
+    t: float = 0.0,
+    grid: int = 9,
+) -> None:
+    xmin, xmax, ymin, ymax = workspace
+    xs = np.linspace(xmin + 0.7, xmax - 0.7, grid)
+    ys = np.linspace(ymin + 0.7, ymax - 0.7, grid)
+    for x in xs:
+        for y in ys:
+            p = np.array([x, y], dtype=np.float32)
+            v = flow.velocity(p, t)
+            speed = float(np.linalg.norm(v))
+            if speed < 1e-4:
+                continue
+            q = p + 0.75 * v / max(0.15, speed)
+            p0 = _world_to_px(p, workspace, size, pad)
+            p1 = _world_to_px(q, workspace, size, pad)
+            _draw_arrow(draw, p0, p1, (160, 190, 218), width=1)
+
+
+def draw_boat(
+    draw: ImageDraw.ImageDraw,
+    state: np.ndarray,
+    spec: BoatSpec,
+    workspace: tuple[float, float, float, float],
+    size: int,
+    pad: int,
+    fill: Color = (38, 89, 133),
+    outline: Color = (15, 37, 61),
+) -> None:
+    pos = state[:2]
+    theta = float(state[2])
+    rot = rot_body_to_world(theta)
+    hull = (spec.hull_vertices @ rot.T) + pos
+    pts = [_world_to_px(p, workspace, size, pad) for p in hull]
+    draw.polygon(pts, fill=fill, outline=outline)
+    nose = pos + rot @ np.array([0.56, 0.0], dtype=np.float32)
+    _draw_arrow(draw, _world_to_px(pos, workspace, size, pad), _world_to_px(nose, workspace, size, pad), (230, 242, 255), width=2)
+    for r, d in zip(spec.thruster_positions, spec.thruster_dirs):
+        p = pos + rot @ r
+        q = p + rot @ (r - 0.16 * d)
+        _draw_arrow(draw, _world_to_px(q, workspace, size, pad), _world_to_px(p, workspace, size, pad), (215, 108, 71), width=2)
+
+
+def render_frame(
+    state: np.ndarray,
+    boat: str | BoatSpec,
+    flow: Flow,
+    workspace: tuple[float, float, float, float] = (0.0, 10.0, 0.0, 10.0),
+    trajectory: np.ndarray | None = None,
+    goal: np.ndarray | None = None,
+    planned: Iterable[np.ndarray] | None = None,
+    size: int = 512,
+    pad: int = 28,
+    t: float = 0.0,
+) -> Image.Image:
+    spec = get_boat_spec(boat) if isinstance(boat, str) else boat
+    img = Image.new("RGB", (size, size), (247, 250, 252))
+    draw = ImageDraw.Draw(img, "RGBA")
+    draw.rectangle([pad, pad, size - pad, size - pad], outline=(35, 54, 72, 255), width=2)
+    draw_flow_field(draw, flow, workspace, size, pad, t=t)
+    if planned is not None:
+        for rollout in planned:
+            pts = [_world_to_px(p[:2], workspace, size, pad) for p in rollout]
+            if len(pts) > 1:
+                draw.line(pts, fill=(110, 138, 183, 45), width=1)
+    if trajectory is not None and len(trajectory) > 1:
+        pts = [_world_to_px(p[:2], workspace, size, pad) for p in trajectory]
+        draw.line(pts, fill=(22, 131, 105, 230), width=3)
+    if goal is not None:
+        gx, gy = _world_to_px(np.asarray(goal, dtype=np.float32), workspace, size, pad)
+        r = 8
+        draw.ellipse([gx - r, gy - r, gx + r, gy + r], fill=(211, 67, 78, 230), outline=(120, 25, 33, 255), width=2)
+    draw_boat(draw, state, spec, workspace, size, pad)
+    return img
+
+
+def save_gif(frames: list[Image.Image], path: str | Path, duration_ms: int = 40) -> None:
+    path = Path(path)
+    ensure_dir(path.parent)
+    if not frames:
+        raise ValueError("no frames to save")
+    frames[0].save(path, save_all=True, append_images=frames[1:], duration=duration_ms, loop=0)
+
+
+def save_boat_geometry(boat: str, path: str | Path, size: int = 420) -> None:
+    from driftwm.sim.flow import NoFlow
+
+    state = np.array([5.0, 5.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float32)
+    img = render_frame(state, boat, NoFlow(), trajectory=None, size=size)
+    path = Path(path)
+    ensure_dir(path.parent)
+    img.save(path)
+
+
+def save_flow_quiver(flow: Flow, path: str | Path, workspace: tuple[float, float, float, float] = (0.0, 10.0, 0.0, 10.0)) -> None:
+    state = np.array([5.0, 5.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float32)
+    img = render_frame(state, "twin", flow, workspace=workspace, trajectory=None)
+    path = Path(path)
+    ensure_dir(path.parent)
+    img.save(path)

driftwm/sim/sanity.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+
+from driftwm.sim.env import SurfaceBoatEnv
+from driftwm.sim.render import render_frame, save_boat_geometry, save_flow_quiver, save_gif
+from driftwm.utils import ensure_dir
+
+
+def _action_for_scenario(boat: str, scenario: str, step: int) -> np.ndarray:
+    if scenario == "thruster":
+        if boat == "twin":
+            return np.array([0.9, 0.15], dtype=np.float32) if step < 80 else np.array([0.2, 0.8], dtype=np.float32)
+        return np.array([0.8, -0.2, 0.35], dtype=np.float32) if step < 80 else np.array([-0.35, 0.75, 0.1], dtype=np.float32)
+    if scenario == "random":
+        raise RuntimeError("random scenario is handled separately")
+    return np.zeros(3 if boat == "triangle" else 2, dtype=np.float32)
+
+
+def run_sanity(
+    boat: str,
+    flow_type: str,
+    out: str | Path,
+    scenario: str = "auto",
+    steps: int = 200,
+    seed: int = 0,
+    boundary: str = "bounce",
+) -> None:
+    if scenario == "auto":
+        scenario = "slide" if flow_type == "noflow" else "drift"
+    env = SurfaceBoatEnv(boat=boat, flow_type=flow_type, episode_steps=steps, boundary=boundary, randomize_params=False, seed=seed)
+    obs, _ = env.reset(random_velocity=False)
+    if scenario == "slide":
+        env.state[3:5] = np.array([0.95, 0.22], dtype=np.float32)
+        env.state[5] = 0.25
+    elif scenario == "drift":
+        env.state[3:6] = 0.0
+    elif scenario == "thruster":
+        env.state[3:6] = 0.0
+
+    frames = []
+    traj = [env.full_state()[:6].copy()]
+    rng = np.random.default_rng(seed + 17)
+    for k in range(steps):
+        if scenario == "random":
+            action = rng.uniform(-1.0, 1.0, size=env.action_dim).astype(np.float32)
+        else:
+            action = _action_for_scenario(boat, scenario, k)[: env.action_dim]
+        env.step(action)
+        traj.append(env.full_state()[:6].copy())
+        if k % 2 == 0:
+            frames.append(
+                render_frame(
+                    env.full_state()[:6],
+                    env.spec,
+                    env.flow,
+                    env.workspace,
+                    trajectory=np.asarray(traj),
+                    t=env.time,
+                )
+            )
+    save_gif(frames, out)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--boat", choices=["twin", "triangle"], default="twin")
+    parser.add_argument("--flow", choices=["noflow", "uniform", "slowly_varying", "vortex"], default="noflow")
+    parser.add_argument("--scenario", choices=["auto", "slide", "drift", "thruster", "random"], default="auto")
+    parser.add_argument("--steps", type=int, default=200)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--out", required=True)
+    parser.add_argument("--make-static", action="store_true")
+    args = parser.parse_args()
+    run_sanity(args.boat, args.flow, args.out, args.scenario, args.steps, args.seed)
+    if args.make_static:
+        out_dir = ensure_dir(Path(args.out).parent)
+        save_boat_geometry("twin", out_dir / "boat_geometry_twin.png")
+        save_boat_geometry("triangle", out_dir / "boat_geometry_triangle.png")
+        env = SurfaceBoatEnv(flow_type=args.flow, boundary="bounce", seed=args.seed)
+        env.reset()
+        save_flow_quiver(env.flow, out_dir / "flow_field_quiver.png")
+
+
+if __name__ == "__main__":
+    main()

driftwm/utils.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+import json
+import math
+import os
+import random
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+import yaml
+
+
+def ensure_dir(path: str | os.PathLike[str]) -> Path:
+    p = Path(path)
+    p.mkdir(parents=True, exist_ok=True)
+    return p
+
+
+def load_yaml(path: str | os.PathLike[str]) -> dict[str, Any]:
+    with open(path, "r", encoding="utf-8") as f:
+        return yaml.safe_load(f)
+
+
+def save_json(obj: Any, path: str | os.PathLike[str]) -> None:
+    path = Path(path)
+    ensure_dir(path.parent)
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(obj, f, indent=2, sort_keys=True)
+
+
+def set_seed(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def wrap_angle(theta: float | np.ndarray) -> float | np.ndarray:
+    return (theta + math.pi) % (2.0 * math.pi) - math.pi
+
+
+def angle_error(pred: np.ndarray, target: np.ndarray) -> np.ndarray:
+    return np.arctan2(np.sin(pred - target), np.cos(pred - target))
+
+
+def obs_from_state(state: np.ndarray) -> np.ndarray:
+    theta = state[..., 2]
+    return np.stack(
+        [state[..., 0], state[..., 1], np.cos(theta), np.sin(theta)],
+        axis=-1,
+    ).astype(np.float32)
+
+
+def state_dim_for_action_dim(action_dim: int) -> int:
+    return 6 + action_dim
+
+
+def pad_action(action: np.ndarray, max_dim: int = 3) -> np.ndarray:
+    out = np.zeros(max_dim, dtype=np.float32)
+    out[: action.shape[-1]] = action
+    return out
+
+
+def unpad_action(action: np.ndarray, action_dim: int) -> np.ndarray:
+    return np.asarray(action[..., :action_dim], dtype=np.float32)
+
+
+def device_from_arg(device: str | None = None) -> torch.device:
+    if device:
+        return torch.device(device)
+    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def configure_torch_runtime() -> None:
+    """Prefer stable kernels on local ROCm builds.
+
+    The local rocm712 environment can expose RDNA targets where MIOpen's GRU
+    kernel compilation fails. Disabling the cuDNN/MIOpen RNN backend keeps the
+    same ROCm device while using PyTorch's native kernels.
+    """
+    if torch.cuda.is_available() and getattr(torch.version, "hip", None):
+        torch.backends.cudnn.enabled = False

experiments/BASELINES.md

@@ -0,0 +1,44 @@
+# Baseline Scope
+
+The public benchmark has two formal comparison groups.
+
+## A. Learned World Models
+
+Purpose: compare image-input world-model architectures under the same data, optimizer budget, rollout target, and planning interface.
+
+| Directory | Report Name | Why It Is Included |
+|---|---|---|
+| `flowmo` | FlowMo | Proposed flow-momentum WM. Separates short object motion state from long ambient drift context. |
+| `leworldmodel` | LeWorldModel | Simple JEPA-style latent prediction baseline. Tests whether current-image latent dynamics are enough. |
+| `planet` | PlaNet RSSM | Recurrent state-space baseline. Tests whether generic recurrent memory can absorb momentum and drift. |
+| `tdmpc2` | TD-MPC2 Dynamics | Compact latent-dynamics baseline. Tests action-conditioned latent rollout with a task-oriented architecture. |
+
+Comparison outputs:
+
+```text
+rollout prediction error
+heading prediction error
+context ablation for FlowMo
+planning metrics when the learned WM is used inside the shared planner
+```
+
+## B. Traditional Non-WM Controllers
+
+Purpose: compare downstream behavior against hand-designed controllers that do not train a neural world model.
+
+| Directory | Report Name | Why It Is Included |
+|---|---|---|
+| `pid_los_controller` | PID/LOS controller | Simple classical waypoint tracking baseline. |
+| `physics_mpc_no_flow` | Physics MPC No-Flow | Nominal dynamics controller that ignores ambient current. |
+| `current_estimator_mpc` | Current-Estimator MPC | Strong classical baseline that estimates current from recent drift. |
+| `oracle_flow_mpc` | Oracle-Flow MPC | Reference bound using true local flow from the simulator. |
+
+Comparison outputs:
+
+```text
+success rate
+final distance
+trajectory length over successful episodes
+energy / thrust work over successful episodes
+time to goal over successful episodes
+```

experiments/EXPERIMENT_MATRIX.md

@@ -0,0 +1,151 @@
+# FlowMo Paper Experiment Matrix
+
+This document defines the paper-facing experiments. File and run names avoid version suffixes; regenerated artifacts replace the same public paths.
+
+## Shared Data And Observation Protocol
+
+All learned world models use the same simulator data and clean-image observation pipeline.
+
+```text
+Image input: clean top-down RGB boat image
+Image size: 160 x 160
+Visual scale: 2.5
+Forbidden image cues: flow arrows, velocity vectors, trajectory overlays, goal marker
+Train split: data/paper/train.npz
+Primary unseen-flow split: data/paper/test_unseen_flow.npz
+Primary unseen-boat-dynamics split: data/paper/test_unseen_boat_params.npz
+Diagnostic seen-flow-family split: data/paper/diagnostic_seen_flow.npz
+Config: experiments/shared/config/paper_image.json
+Checkpoint: paper.pt
+Intermediate checkpoints: paper_step_XXXXXX.pt
+```
+
+Formal training budget:
+
+```text
+train_episodes: 2400
+test_episodes: 480
+train_windows: 393216
+test_windows: 24576
+batch_size: 256
+steps: 20000
+checkpoint_interval: 2000
+num_workers: 4
+render_mode: device
+```
+
+Precision policy:
+
+```text
+training: bf16 model autocast, fp32 losses and metrics
+prediction_eval: bf16 model autocast, fp32 metrics
+planning_eval: fp32
+```
+
+## A. Learned World-Model Comparison
+
+Purpose: measure world-model quality directly. The key question is whether FlowMo's short object-motion state plus long ambient-drift context improves rollout prediction under hidden currents and momentum.
+
+| Method | Comparison Role | What It Tests |
+|---|---|---|
+| `flowmo` | Proposed WM | Explicit flow-momentum factorization: short state/momentum latent, long drift context, zero-context residual transition. |
+| `leworldmodel` | JEPA-style WM baseline | Whether simple image-latent prediction without explicit history/context can handle boat momentum and flow. |
+| `planet` | RSSM WM baseline | Whether generic recurrent latent memory can represent momentum and drift without a separate context factor. |
+| `tdmpc2` | Compact latent-dynamics WM baseline | Whether a compact action-conditioned latent transition matches FlowMo under equal supervision. |
+
+Prediction datasets:
+
+```text
+test_unseen_flow
+test_unseen_boat_params
+diagnostic_seen_flow
+```
+
+Prediction metrics:
+
+```text
+pos@1, pos@5, pos@10, pos@20, pos@40, pos@60
+heading@20, heading@60
+zero-action drift error
+no-flow momentum decay error
+same-action different-flow error
+```
+
+FlowMo context diagnostics:
+
+| Diagnostic | Operation | Evidence Sought |
+|---|---|---|
+| Inferred context | Normal rollout with inferred `c_t` | Best prediction under flow. |
+| Zero context | Set `c_t=0` | Degraded flow prediction, smaller change in no-flow. |
+| Shuffled context | Use context from another episode | Worse rollout when hidden flow differs. |
+| Same-flow transfer | Use context from another episode with the same hidden flow | Better than wrong-flow context transfer. |
+| No-flow context norm | Measure `||c_t||` on no-flow data | Smaller than flow context norm. |
+| Context PCA | Plot `c_t` by flow family / flow id | Flow-related organization. |
+
+FlowMo latent probes:
+
+| Probe Target | Feature Sets | Purpose |
+|---|---|---|
+| Object momentum `(vx, vy, omega)` | `z_t`, `c_t`, `[z_t,c_t]` | Tests whether short-history state contains object motion. |
+| Local flow vector | `z_t`, `c_t`, `[z_t,c_t]` | Tests whether state plus context exposes local ambient drift. |
+| Episode drift vector | `z_t`, `c_t`, `[z_t,c_t]` | Tests whether long context contains environment-level drift. |
+
+## B. Traditional Non-WM Control Comparison
+
+Purpose: provide downstream control references and report practical task behavior. The central WM claim still comes from A; B shows whether prediction differences matter for planning and control.
+
+Learned WM planners:
+
+```text
+flowmo
+leworldmodel
+planet
+tdmpc2
+```
+
+Traditional non-WM controllers:
+
+| Method | Comparison Role | What It Tests |
+|---|---|---|
+| `pid_los_controller` | Simple classical controller | Baseline waypoint tracking without learned dynamics. |
+| `physics_mpc_no_flow` | Nominal physics MPC | Effect of ignoring hidden current. |
+| `current_estimator_mpc` | Current-compensated classical MPC | Strength of a hand-designed drift estimator. |
+| `oracle_flow_mpc` | Oracle reference | Reference performance when true local flow is available. |
+
+Planning tasks:
+
+```text
+reach_uniform
+counterflow
+station_keeping
+passive_to_active
+waypoint_square
+waypoint_zigzag
+```
+
+Boats:
+
+```text
+twin
+triangle
+```
+
+Planning metrics:
+
+```text
+success rate
+final distance
+trajectory length over successful episodes
+energy / thrust work over successful episodes
+time to goal over successful episodes
+```
+
+Formal commands:
+
+```bash
+python -m experiments.run_paper_image_pipeline --stages train
+python -m experiments.run_paper_image_pipeline --stages prediction
+python -m experiments.run_paper_image_pipeline --stages probe
+python -m experiments.run_paper_image_pipeline --stages planning
+python -m experiments.run_paper_image_pipeline --stages report
+```

experiments/METHOD_AUDIT.md

@@ -0,0 +1,30 @@
+# FlowMo Method And Baseline Audit
+
+This document records what each paper-facing method implements and what comparison it supports.
+
+## A. Learned World Models
+
+| Method | Implemented Architecture | Rollout State | Transition Form | Comparison Purpose |
+|---|---|---|---|---|
+| FlowMo | Shared image encoder; short state GRU; long strided context GRU; base transition plus zero-context residual | `z_t` plus `c_t` | `z + F0(z,a) + R(z,c) - R(z,0)` | Tests whether separating short object motion from long ambient drift improves prediction and planning. |
+| LeWorldModel | Image encoder; current-frame latent state; action-conditioned residual transition | `z_t` | `z + F(z,a)` | Tests a simple JEPA-style image-latent predictor without temporal drift context. |
+| PlaNet RSSM | Image encoder; deterministic recurrent state; stochastic latent posterior/prior; decoder from recurrent plus stochastic state | deterministic recurrent state plus stochastic latent | recurrent transition plus prior/posterior latent | Tests generic recurrent state-space memory under the shared rollout protocol. |
+| TD-MPC2 Dynamics | Image encoder; short action-conditioned history encoder; residual latent dynamics | `z_t` | `z + F(z,a)` after history encoding | Tests compact action-conditioned latent dynamics without explicit drift context. |
+
+## B. Traditional Non-WM Controllers
+
+| Method | Input | Controller Type | Comparison Purpose |
+|---|---|---|---|
+| PID/LOS controller | Clean image pose estimate | Line-of-sight waypoint tracking | Simple hand-designed tracking baseline. |
+| Physics MPC No-Flow | Clean image pose estimate | MPC with nominal boat dynamics | Measures the cost of ignoring external current. |
+| Current-Estimator MPC | Clean image pose estimate and recent drift | MPC with estimated current | Strong classical current-compensation baseline. |
+| Oracle-Flow MPC | Clean image pose estimate and simulator local flow | MPC with true local current | Reference bound for downstream planning. |
+
+## Architecture Difference Checklist
+
+| Method | Uses History | Has Long Context | Has Explicit Drift Context | Distinct From FlowMo |
+|---|---|---|---|---|
+| FlowMo | Yes, short state history | Yes, strided long context | Yes, `c_t` | Proposed factorization. |
+| LeWorldModel | No, current image only | No | No | Tests no-history latent prediction. |
+| PlaNet RSSM | Yes | No | No | Tests generic recurrent memory. |
+| TD-MPC2 Dynamics | Yes, short history | No | No | Tests compact latent dynamics. |

experiments/README.md

@@ -0,0 +1,83 @@
+# Experiments
+
+This directory contains the paper-facing experiment code, checkpoints, results, figures, GIFs, tables, and reports.
+
+This directory contains two formal experiment categories:
+
+- **A. Learned world models**: trainable image-input WMs evaluated on rollout prediction and WM-based planning.
+- **B. Traditional non-WM controllers**: hand-designed control baselines evaluated on the same downstream tasks.
+
+Main method:
+
+- **FlowMo**: Flow-Momentum World Model, the proposed drift-aware world model for surface vehicles.
+
+Category A learned WM comparisons:
+
+- **LeWorldModel**: JEPA-style latent predictor under the shared clean-image protocol.
+- **PlaNet RSSM**: recurrent state-space world-model baseline under the shared clean-image protocol.
+- **TD-MPC2 Dynamics**: task-oriented latent dynamics baseline under the shared clean-image protocol.
+
+Purpose of Category A: compare world-model architectures under identical image data, optimizer budget, parameter budget, rollout target, and evaluation protocol.
+
+Category B traditional controllers:
+
+- **PID/LOS controller**
+- **Physics MPC No-Flow**
+- **Current-Estimator MPC**
+- **Oracle-Flow MPC**
+
+Purpose of Category B: compare downstream task behavior against non-neural controllers that do not train a world model.
+
+Baseline details are documented in `BASELINES.md`; the full experiment protocol is documented in `docs/EXPERIMENT_PROTOCOL.md`.
+
+Design principles:
+
+- Shared simulator, datasets, planning utilities, metrics, and visualization live in `shared/`.
+- Each method has its own directory with `src/`, `checkpoint/`, and `result/`.
+- Paper artifacts are collected in top-level `figures/`, `gifs/`, `tables/`, and `reports/`.
+- Method names should be explicit and readable. Avoid cryptic suffixes in paper-facing file names.
+
+Standard method interface:
+
+```text
+src/model.py      # build_model(), load_model()
+src/train.py      # train(config)
+src/predict.py    # rollout(model, batch)
+src/plan.py       # plan(model, env, task)
+src/config.py     # default_config()
+```
+
+Traditional controllers use:
+
+```text
+src/controller.py or src/mpc.py
+src/evaluate.py
+src/config.py
+```
+
+Formal clean-image configuration:
+
+```text
+image_size=160
+visual_scale=2.5
+train=data/paper/train.npz
+test=data/paper/test_unseen_flow.npz and data/paper/test_unseen_boat_params.npz
+```
+
+Full paper-facing image pipeline:
+
+```bash
+python -m experiments.run_paper_image_pipeline
+```
+
+The default command runs the paper configuration end to end: train all learned world models, evaluate long rollout prediction, run FlowMo latent probes, evaluate closed-loop planning against traditional controllers, generate GIFs, and write the final report. Images are rendered online from simulator states, so no separate image-cache preparation step is required.
+
+Manual image training:
+
+```bash
+python -m experiments.train_image_world_models
+python -m experiments.evaluate_image_world_models
+python -m experiments.evaluate_flowmo_latent_probes
+python -m experiments.evaluate_image_planning --task reach_uniform --boat twin
+python -m experiments.summarize_paper_image_results
+```

experiments/TASK_PLAN.md

@@ -0,0 +1,118 @@
+# Paper Task Plan
+
+This is the execution plan for the public FlowMo experiments. The plan has two parts: A evaluates world models directly, and B evaluates downstream control behavior with traditional non-WM references.
+
+## A. Learned World Models
+
+Purpose: test whether the FlowMo world-model architecture improves image-based prediction under hidden flow, boat momentum, actuator delay, and drag.
+
+Shared setup:
+
+```text
+Input: clean top-down boat images plus action history
+No image cues: no flow arrows, no velocity vector, no goal marker
+Training data: data/paper/train.npz
+Primary evaluation data: data/paper/test_unseen_flow.npz, data/paper/test_unseen_boat_params.npz
+Diagnostic data: data/paper/diagnostic_seen_flow.npz
+Training budget: shared optimizer, batch size, rollout horizon, step count, and checkpoint schedule
+Training precision: BF16 model autocast, FP32 losses and metrics
+Prediction precision: BF16 model autocast, FP32 metrics
+```
+
+Compared methods:
+
+| Method | Purpose |
+|---|---|
+| `flowmo` | Proposed flow-momentum WM. Tests explicit separation of short object-motion state and long ambient-drift context. |
+| `leworldmodel` | JEPA-style latent predictor. Tests whether a simple current-image latent transition is sufficient. |
+| `planet` | RSSM recurrent state-space WM. Tests whether generic recurrent latent memory can absorb momentum and flow effects without FlowMo's explicit context. |
+| `tdmpc2` | Compact latent-dynamics WM. Tests whether a task-oriented latent transition architecture matches FlowMo under the same rollout supervision. |
+
+Primary A metrics:
+
+```text
+pos@1, pos@5, pos@10, pos@20, pos@40, pos@60
+heading@20, heading@60
+zero-action drift prediction error
+no-flow momentum decay prediction error
+same-action different-flow prediction error
+FlowMo inferred-context vs c=0 vs shuffled-context error
+```
+
+Required A outputs:
+
+```text
+experiments/<method>/checkpoint/paper.pt
+experiments/<method>/checkpoint/paper_step_*.pt
+experiments/<method>/result/parameter_count.json
+experiments/<method>/result/paper_training.json
+experiments/reports/paper_prediction_seen_flow_diagnostic.json
+experiments/reports/paper_prediction_unseen_flow.json
+experiments/reports/paper_prediction_unseen_boat_params.json
+experiments/reports/paper_flowmo_latent_probes.json
+```
+
+Core A conclusions:
+
+```text
+1. Whether FlowMo has lower long-horizon rollout error.
+2. Whether the gain is strongest under unseen flow families and unseen boat dynamics.
+3. Whether explicit drift context helps beyond ordinary recurrent history.
+4. Whether the same architecture works for both twin and triangle boats.
+5. Whether frozen linear probes recover object momentum from `z_t` and ambient drift from `c_t`.
+```
+
+## B. Traditional Non-WM Controllers
+
+Purpose: evaluate downstream control behavior and provide non-neural-control reference points. These methods do not train a world model.
+
+Shared setup:
+
+```text
+Input: clean top-down images converted to pose for classical control
+Tasks: same simulator, same boats, same goals, same flow settings
+Metrics: success, final distance, successful-episode trajectory length, successful-episode thrust energy, successful-episode time-to-goal
+Planning precision: FP32
+```
+
+Compared methods:
+
+| Method | Purpose |
+|---|---|
+| `pid_los_controller` | Classical line-of-sight waypoint tracking. Tests a simple hand-designed controller. |
+| `physics_mpc_no_flow` | Physics MPC without external-current compensation. Tests how much hidden flow hurts a nominal dynamics controller. |
+| `current_estimator_mpc` | MPC with recent-drift current estimation. Tests a strong classical current-compensation baseline. |
+| `oracle_flow_mpc` | MPC with simulator true local flow. Provides a reference upper bound for classical planning. |
+
+Planning tasks:
+
+```text
+reach_uniform
+counterflow
+station_keeping
+passive_to_active
+waypoint_square
+waypoint_zigzag
+```
+
+Boats:
+
+```text
+twin
+triangle
+```
+
+Required B outputs:
+
+```text
+experiments/reports/paper_planning/*.json
+experiments/reports/paper_planning/gifs/*.gif
+```
+
+Core B conclusions:
+
+```text
+1. Whether WM-based planning is competitive with classical non-WM control.
+2. Whether FlowMo improves success, final distance, energy, and path length versus other learned WMs.
+3. How far FlowMo remains from the oracle-flow classical reference.
+```

experiments/__init__.py

@@ -0,0 +1 @@
+"""Paper-facing experiment package."""

experiments/current_estimator_mpc/README.md

@@ -0,0 +1,3 @@
+# Current-Estimator MPC Baseline
+
+Traditional MPC baseline with an online estimator for approximately uniform current drift.

experiments/current_estimator_mpc/__init__.py

@@ -0,0 +1 @@
+"""Current-estimator MPC experiment package."""

experiments/current_estimator_mpc/src/__init__.py

@@ -0,0 +1 @@
+"""Current-estimator MPC source package."""

experiments/current_estimator_mpc/src/config.py

@@ -0,0 +1,5 @@
+"""Current-estimator MPC config."""
+
+
+def default_config():
+    return {"goal": [8.0, 8.0], "gain": 0.65, "current_gain": 0.5}

experiments/current_estimator_mpc/src/estimator.py

@@ -0,0 +1,9 @@
+"""Online current estimator."""
+
+import numpy as np
+
+from experiments.shared.src.vision.pose_from_image import estimate_pose_from_clean_image
+
+def estimate_current(history, config):
+    poses = np.stack([estimate_pose_from_clean_image(img) for img in history], axis=0)
+    return poses[-1, :2] - poses[0, :2]

experiments/current_estimator_mpc/src/evaluate.py

@@ -0,0 +1,10 @@
+"""Current-estimator MPC evaluation."""
+
+from experiments.current_estimator_mpc.src.config import default_config
+from experiments.current_estimator_mpc.src.estimator import estimate_current
+from experiments.current_estimator_mpc.src.mpc import plan
+
+def evaluate(config):
+    cfg = default_config() | config
+    current = estimate_current(config["history"], cfg)
+    return plan(config["history"][-1], current, {"goal": cfg["goal"]}, cfg)

experiments/current_estimator_mpc/src/mpc.py

@@ -0,0 +1,19 @@
+"""MPC with online current estimate."""
+
+import numpy as np
+
+from experiments.shared.src.control.geometric import goal_action
+from experiments.shared.src.vision.pose_from_image import estimate_pose_from_clean_image
+
+def plan(state, current_estimate, task, config):
+    pose = estimate_pose_from_clean_image(state)
+    goal = np.asarray(task["goal"], dtype=np.float32)
+    return goal_action(
+        pose,
+        goal,
+        int(config["action_dim"]),
+        float(config["gain"]),
+        0.6,
+        drift=np.asarray(current_estimate, dtype=np.float32),
+        drift_gain=float(config["current_gain"]),
+    )

experiments/docs/EXPERIMENT_PROTOCOL.md

@@ -0,0 +1,211 @@
+# FlowMo Experiment Protocol
+
+This document is the single paper-facing record for the public FlowMo experiments. Regenerated artifacts should replace the same paths instead of introducing version suffixes.
+
+## Scope
+
+The project has two formal comparison groups.
+
+### A. Learned World Models
+
+Purpose: compare image-input world-model architectures under the same simulator data, optimizer budget, rollout target, parameter scale, and planning interface.
+
+| Directory | Report Name | Architecture | Comparison Purpose |
+|---|---|---|---|
+| `flowmo` | FlowMo | Shared image encoder; short object-motion state encoder; long strided ambient-drift context encoder; base transition plus zero-context residual | Proposed flow-momentum factorization. Tests whether separating endogenous motion from exogenous drift improves prediction and planning. |
+| `leworldmodel` | LeWorldModel | JEPA-style image-latent predictor with action-conditioned residual transition | Tests whether simple current-image latent prediction is sufficient. |
+| `planet` | RSSM | Recurrent state-space latent model with deterministic memory and stochastic latent state | Tests whether generic recurrent memory can absorb momentum and drift without a separate context factor. |
+| `tdmpc2` | TD-MPC2 Dynamics | Compact action-conditioned latent dynamics with shared image encoder and rollout heads | Tests task-oriented latent dynamics under equal supervision. |
+
+All learned methods receive clean top-down RGB boat images and action history. They do not receive flow labels, flow arrows, velocity vectors, trajectory overlays, or goal markers in the image.
+
+### B. Traditional Non-WM Controllers
+
+Purpose: compare downstream behavior against non-neural controllers that do not train a world model.
+
+| Directory | Report Name | Input | Comparison Purpose |
+|---|---|---|---|
+| `pid_los_controller` | PID/LOS | Clean image pose estimate | Simple hand-designed waypoint tracking baseline. |
+| `physics_mpc_no_flow` | Physics MPC No-Flow | Clean image pose estimate | Measures the cost of ignoring ambient current. |
+| `current_estimator_mpc` | Current-Estimator MPC | Clean image pose estimate and recent drift | Strong classical current-compensation baseline. |
+| `oracle_flow_mpc` | Oracle-Flow MPC | Clean image pose estimate and simulator local flow | Reference bound for control when true local flow is available. |
+
+## Data
+
+All methods use the same splits:
+
+```text
+train: data/paper/train.npz
+unseen_flow_test: data/paper/test_unseen_flow.npz
+unseen_boat_dynamics_test: data/paper/test_unseen_boat_params.npz
+seen_flow_diagnostic: data/paper/diagnostic_seen_flow.npz
+dataset_card: data/paper/dataset_card.md
+generation_config: data/paper/generation_config.json
+```
+
+Observation protocol:
+
+```text
+image_size: 160 x 160
+visual_scale: 2.5
+rendering: online clean top-down RGB images
+forbidden cues: flow arrows, velocity vectors, trajectory overlays, goal markers
+```
+
+Training budget:
+
+```text
+train_episodes: 2400
+test_episodes: 480
+train_windows: 393216
+test_windows: 24576
+batch_size: 256
+steps: 20000
+checkpoint_interval: 2000
+num_workers: 4
+render_mode: device
+```
+
+Precision policy:
+
+```text
+training: bf16 model autocast, fp32 losses and metrics
+prediction_eval: bf16 model autocast, fp32 metrics
+planning_eval: fp32
+```
+
+The precision split is intentional: BF16 speeds up image encoding and latent rollout on the RTX 5090 without measurable short-run loss drift, while CEM planning is dominated by small control tensors and did not improve under BF16.
+
+## Prediction Evaluation
+
+Datasets:
+
+```text
+test_unseen_flow
+test_unseen_boat_params
+diagnostic_seen_flow
+```
+
+Metrics:
+
+```text
+pos@1, pos@5, pos@10, pos@20, pos@40, pos@60
+heading@20, heading@60
+zero-action drift prediction error
+no-flow momentum decay prediction error
+same-action different-flow prediction error
+```
+
+FlowMo-only context diagnostics:
+
+| Diagnostic | Operation | Evidence Sought |
+|---|---|---|
+| Inferred context | Normal rollout with inferred `c_t` | Best prediction under flow. |
+| Zero context | Set `c_t=0` | Degraded flow prediction and limited change in no-flow. |
+| Shuffled context | Use context from another episode | Worse rollout when hidden flow differs. |
+| Same-flow transfer | Use context from another episode with the same hidden flow | Better transfer than wrong-flow context. |
+| Context norm | Compare no-flow and flow `||c_t||` | Flow context should be larger than no-flow context. |
+
+FlowMo latent probes:
+
+```text
+Train frozen linear probes from z_t, c_t, and [z_t,c_t].
+Targets: object momentum (vx, vy, omega), local flow vector, episode drift vector.
+Purpose: verify which latent carries object-motion information and which latent carries ambient-drift information.
+```
+
+## Planning Evaluation
+
+Learned WM planners:
+
+```text
+flowmo
+leworldmodel
+planet
+tdmpc2
+```
+
+Traditional non-WM controllers:
+
+```text
+pid_los_controller
+physics_mpc_no_flow
+current_estimator_mpc
+oracle_flow_mpc
+```
+
+Tasks:
+
+```text
+reach_uniform
+counterflow
+station_keeping
+passive_to_active
+waypoint_square
+waypoint_zigzag
+```
+
+Boats:
+
+```text
+twin
+triangle
+```
+
+Metrics:
+
+```text
+success rate
+final distance
+trajectory length over successful episodes
+energy / thrust work over successful episodes
+time to goal over successful episodes
+```
+
+## Required Outputs
+
+Training outputs:
+
+```text
+experiments/<method>/checkpoint/paper.pt
+experiments/<method>/checkpoint/paper_step_*.pt
+experiments/<method>/result/parameter_count.json
+experiments/<method>/result/paper_training.json
+experiments/<method>/result/paper_training_trace.jsonl
+```
+
+Evaluation outputs:
+
+```text
+experiments/reports/paper_prediction_unseen_flow.json
+experiments/reports/paper_prediction_unseen_boat_params.json
+experiments/reports/paper_prediction_seen_flow_diagnostic.json
+experiments/reports/paper_flowmo_latent_probes.json
+experiments/reports/paper_planning/*.json
+experiments/reports/paper_planning/gifs/*.gif
+experiments/reports/paper_report.md
+```
+
+## Commands
+
+Run the complete paper pipeline:
+
+```bash
+python -m experiments.run_paper_image_pipeline
+```
+
+Run stages separately:
+
+```bash
+python -m experiments.run_paper_image_pipeline --stages train
+python -m experiments.run_paper_image_pipeline --stages prediction
+python -m experiments.run_paper_image_pipeline --stages probe
+python -m experiments.run_paper_image_pipeline --stages planning
+python -m experiments.run_paper_image_pipeline --stages report
+```
+
+Run smoke tests:
+
+```bash
+python -m pytest -q
+```

experiments/evaluate_flowmo_latent_probes.py

@@ -0,0 +1,266 @@
+"""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=[
+        "unseen_flow:data/paper/test_unseen_flow.npz:480",
+        "unseen_boat_params:data/paper/test_unseen_boat_params.npz:480",
+        "seen_flow_diagnostic:data/paper/diagnostic_seen_flow.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()

experiments/evaluate_image_planning.py

@@ -0,0 +1,577 @@
+"""Closed-loop planning evaluation for clean-image world models and controllers."""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import json
+from collections import deque
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from driftwm.sim.env import SurfaceBoatEnv
+from driftwm.sim.flow import UniformFlow, sample_flow
+from driftwm.sim.render import render_frame, save_gif
+from experiments.shared.src.methods import PAPER_LEARNED_METHODS, TRADITIONAL_METHODS
+from experiments.shared.src.vision.clean_renderer import render_clean_boat_array
+from experiments.train_image_world_models import autocast_context
+
+
+LEARNED_METHODS = PAPER_LEARNED_METHODS
+
+POSITION_SCALE = 5.0
+
+
+def build_method(method: str):
+    config_module = importlib.import_module(f"experiments.{method}.src.config")
+    model_module = importlib.import_module(f"experiments.{method}.src.model")
+    cfg = config_module.default_config()
+    return cfg, model_module.build_model(cfg)
+
+
+def decode_absolute(prediction: torch.Tensor) -> torch.Tensor:
+    xy = prediction[..., :2] * POSITION_SCALE + POSITION_SCALE
+    return torch.cat([xy, prediction[..., 2:4]], dim=-1)
+
+
+def clean_observation(env: SurfaceBoatEnv, image_size: int, visual_scale: float) -> np.ndarray:
+    image = render_clean_boat_array(env.full_state()[:6], env.spec, image_size=image_size, visual_scale=visual_scale)
+    return np.transpose(image, (2, 0, 1))
+
+
+def pad_action(action: np.ndarray, action_dim: int) -> np.ndarray:
+    out = np.zeros((action_dim,), dtype=np.float32)
+    action = np.asarray(action, dtype=np.float32)
+    out[: min(len(action), action_dim)] = action[: min(len(action), action_dim)]
+    return out
+
+
+def task_goals(task: str, rng: np.random.Generator) -> np.ndarray:
+    if task == "waypoint_square":
+        return np.array([[2.5, 2.5], [7.5, 2.5], [7.5, 7.5], [2.5, 7.5]], dtype=np.float32)
+    if task == "waypoint_zigzag":
+        return np.array([[2.5, 7.0], [4.2, 3.0], [5.8, 7.0], [7.5, 3.0]], dtype=np.float32)
+    if task == "station_keeping":
+        return np.array([[5.0, 5.0]], dtype=np.float32)
+    if task == "counterflow":
+        return np.array([[8.4, 5.0]], dtype=np.float32)
+    return np.array([[8.0, 8.0]], dtype=np.float32)
+
+
+def reset_task(env: SurfaceBoatEnv, task: str, flow_type: str, rng: np.random.Generator) -> None:
+    if task == "counterflow":
+        env.reset(
+            flow_type="uniform",
+            flow=UniformFlow(np.array([-0.22, 0.0], dtype=np.float32), flow_id=7001),
+            random_velocity=False,
+        )
+        env.state[:6] = np.array([2.0, 5.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float32)
+        return
+    if task == "station_keeping":
+        env.reset(
+            flow_type="uniform",
+            flow=UniformFlow(np.array([0.16, 0.10], dtype=np.float32), flow_id=7002),
+            random_velocity=False,
+        )
+        env.state[:6] = np.array([5.0, 5.0, 0.3, 0.0, 0.0, 0.0], dtype=np.float32)
+        return
+    flow = sample_flow(flow_type, rng, flow_id=10_000 + int(rng.integers(1, 1_000_000)), workspace=env.workspace)
+    env.reset(flow_type=flow_type, flow=flow, random_velocity=False)
+    env.state[:6] = np.array([2.0, 2.0, float(rng.uniform(-np.pi, np.pi)), 0.0, 0.0, 0.0], dtype=np.float32)
+
+
+def rollout_latent(model, z: torch.Tensor, c: torch.Tensor, actions: torch.Tensor) -> torch.Tensor:
+    cur = z.repeat(actions.shape[0], 1)
+    ctx = c.repeat(actions.shape[0], 1) if c.numel() else c
+    preds = []
+    for t in range(actions.shape[1]):
+        cur = model.step(cur, actions[:, t], ctx)
+        preds.append(model.decoder(cur))
+    return decode_absolute(torch.stack(preds, dim=1)).float()
+
+
+def warm_start_mean(
+    previous_mean: np.ndarray | None,
+    horizon: int,
+    action_dim: int,
+    active_action_dim: int,
+    device: torch.device,
+) -> torch.Tensor:
+    mean = torch.zeros((horizon, action_dim), dtype=torch.float32, device=device)
+    if previous_mean is None:
+        return mean
+    previous = torch.as_tensor(previous_mean, dtype=torch.float32, device=device)
+    steps = min(horizon, max(0, previous.shape[0] - 1))
+    if steps > 0:
+        mean[:steps, :active_action_dim] = previous[1 : 1 + steps, :active_action_dim]
+    if previous.shape[0] > 0 and steps < horizon:
+        mean[steps:, :active_action_dim] = previous[-1, :active_action_dim]
+    return mean.clamp(-1.0, 1.0)
+
+
+def sample_action_sequences(mean: torch.Tensor, std: torch.Tensor, population: int, knots: int) -> torch.Tensor:
+    horizon, action_dim = mean.shape
+    if knots >= horizon:
+        noise = torch.randn(population, horizon, action_dim, device=mean.device)
+        return mean.unsqueeze(0) + std.unsqueeze(0) * noise
+    knots = max(2, knots)
+    knot_idx = torch.linspace(0, horizon - 1, knots, device=mean.device).round().long()
+    knot_mean = mean[knot_idx]
+    knot_std = std[knot_idx]
+    knot_samples = knot_mean.unsqueeze(0) + knot_std.unsqueeze(0) * torch.randn(
+        population,
+        knots,
+        action_dim,
+        device=mean.device,
+    )
+    samples = F.interpolate(
+        knot_samples.permute(0, 2, 1),
+        size=horizon,
+        mode="linear",
+        align_corners=True,
+    ).permute(0, 2, 1)
+    return samples
+
+
+def learned_plan(
+    model,
+    image_history: deque,
+    action_history: deque,
+    goal: np.ndarray,
+    active_action_dim: int,
+    args,
+    prev_action: np.ndarray,
+    previous_mean: np.ndarray | None,
+    context_mode: str,
+    donor_context: torch.Tensor | None,
+) -> tuple[np.ndarray, np.ndarray | None, np.ndarray]:
+    device = next(model.parameters()).device
+    images = torch.as_tensor(np.asarray(image_history, dtype=np.uint8), device=device).unsqueeze(0)
+    actions = torch.as_tensor(np.asarray(action_history, dtype=np.float32), device=device).unsqueeze(0)
+    with torch.no_grad(), autocast_context(device, args.precision):
+        z, c = model.encode(images, actions)
+    if c.numel() and context_mode == "zero":
+        c = torch.zeros_like(c)
+    if c.numel() and context_mode == "shuffled" and donor_context is not None:
+        c = donor_context.to(device=device, dtype=torch.float32)
+    z = z.detach()
+    c = c.detach()
+    goal_t = torch.as_tensor(goal, dtype=torch.float32, device=device).view(1, 2)
+    with torch.no_grad(), autocast_context(device, args.precision):
+        current_pos = decode_absolute(model.decoder(z)).float().detach()[..., :2]
+    mean = warm_start_mean(
+        previous_mean,
+        args.cem_horizon,
+        model.config.action_dim,
+        active_action_dim,
+        device,
+    )
+    std = torch.full_like(mean, args.cem_action_std)
+    prev = torch.zeros((model.config.action_dim,), dtype=torch.float32, device=device)
+    prev[:active_action_dim] = torch.as_tensor(prev_action, dtype=torch.float32, device=device)
+    best_candidates = None
+    if args.planner == "gradient":
+        action, best_candidates, mean = gradient_plan(
+            model,
+            z,
+            c,
+            mean,
+            goal_t,
+            current_pos,
+            prev,
+            active_action_dim,
+            args,
+        )
+        return action, best_candidates, mean
+    with torch.no_grad():
+        action, best_candidates, mean = cem_plan(
+            model,
+            z,
+            c,
+            mean,
+            std,
+            goal_t,
+            current_pos,
+            prev,
+            active_action_dim,
+            args,
+        )
+        return action, best_candidates, mean
+
+
+def planning_cost(
+    pred: torch.Tensor,
+    samples: torch.Tensor,
+    goal_t: torch.Tensor,
+    current_pos: torch.Tensor,
+    prev: torch.Tensor,
+    active_action_dim: int,
+    args,
+) -> torch.Tensor:
+    pos = pred[..., :2]
+    goal_delta = goal_t - current_pos
+    goal_dir = goal_delta / torch.linalg.norm(goal_delta, dim=-1, keepdim=True).clamp_min(1.0e-6)
+    progress = ((pos - current_pos[:, None]) * goal_dir[:, None]).sum(dim=-1).amax(dim=-1)
+    alpha = torch.linspace(1.0 / pos.shape[1], 1.0, pos.shape[1], device=pos.device, dtype=pos.dtype)
+    route = current_pos[:, None] + alpha.view(1, -1, 1) * goal_delta[:, None]
+    route_error = ((pos - route) ** 2).sum(dim=-1).mean(dim=-1)
+    goal_from_pos = goal_t[:, None] - pos
+    goal_from_pos = goal_from_pos / torch.linalg.norm(goal_from_pos, dim=-1, keepdim=True).clamp_min(1.0e-6)
+    heading = pred[..., 2:4]
+    heading = heading / torch.linalg.norm(heading, dim=-1, keepdim=True).clamp_min(1.0e-6)
+    heading_error = (1.0 - (heading * goal_from_pos).sum(dim=-1)).mean(dim=-1)
+    terminal = ((pos[:, -1] - goal_t) ** 2).sum(dim=-1)
+    path = ((pos - goal_t[:, None]) ** 2).sum(dim=-1).mean(dim=-1)
+    energy = (samples[..., :active_action_dim] ** 2).mean(dim=(1, 2))
+    smooth_prev = torch.cat([prev.view(1, 1, -1).repeat(samples.shape[0], 1, 1), samples[:, :-1]], dim=1)
+    smooth = ((samples - smooth_prev) ** 2).mean(dim=(1, 2))
+    boundary = (
+        torch.relu(-pos[..., 0])
+        + torch.relu(pos[..., 0] - 10.0)
+        + torch.relu(-pos[..., 1])
+        + torch.relu(pos[..., 1] - 10.0)
+    ).mean(dim=-1)
+    return (
+        args.cem_w_goal * terminal
+        + args.cem_w_path * path
+        + args.cem_w_route * route_error
+        + args.cem_w_heading_goal * heading_error
+        + args.cem_w_action * energy
+        + args.cem_w_smooth * smooth
+        + args.cem_w_boundary * boundary
+        - args.cem_w_progress * progress
+    )
+
+
+def cem_plan(
+    model,
+    z: torch.Tensor,
+    c: torch.Tensor,
+    mean: torch.Tensor,
+    std: torch.Tensor,
+    goal_t: torch.Tensor,
+    current_pos: torch.Tensor,
+    prev: torch.Tensor,
+    active_action_dim: int,
+    args,
+) -> tuple[np.ndarray, np.ndarray | None, np.ndarray]:
+    best_candidates = None
+    for _ in range(args.cem_iterations):
+        samples = sample_action_sequences(mean, std, args.cem_population, args.cem_knots)
+        samples[0] = mean
+        samples = samples.clamp(-1.0, 1.0)
+        if active_action_dim < model.config.action_dim:
+            samples[:, :, active_action_dim:] = 0.0
+        with autocast_context(mean.device, args.precision):
+            pred = rollout_latent(model, z, c, samples)
+        cost = planning_cost(pred, samples, goal_t, current_pos, prev, active_action_dim, args)
+        elite_idx = torch.topk(cost, k=args.cem_elites, largest=False).indices
+        elites = samples[elite_idx]
+        mean = elites.mean(dim=0)
+        std = elites.std(dim=0).clamp_min(0.05)
+        if args.make_gifs:
+            pos = pred[..., :2]
+            best_candidates = pos[elite_idx[:12]].detach().cpu().numpy()
+    action = mean[0, :active_action_dim].detach().cpu().numpy()
+    return (
+        np.clip(action, -1.0, 1.0).astype(np.float32),
+        best_candidates,
+        mean.detach().cpu().numpy(),
+    )
+
+
+def gradient_plan(
+    model,
+    z: torch.Tensor,
+    c: torch.Tensor,
+    mean: torch.Tensor,
+    goal_t: torch.Tensor,
+    current_pos: torch.Tensor,
+    prev: torch.Tensor,
+    active_action_dim: int,
+    args,
+) -> tuple[np.ndarray, np.ndarray | None, np.ndarray]:
+    init = mean.clamp(-0.95, 0.95)
+    logits = torch.atanh(init).detach().requires_grad_(True)
+    optimizer = torch.optim.Adam([logits], lr=args.planner_lr)
+    inactive = None
+    if active_action_dim < mean.shape[-1]:
+        inactive = torch.zeros_like(mean)
+        inactive[:, :active_action_dim] = 1.0
+    for _ in range(args.planner_iterations):
+        seq = torch.tanh(logits)
+        if inactive is not None:
+            seq = seq * inactive
+        with autocast_context(mean.device, args.precision):
+            pred = rollout_latent(model, z, c, seq.unsqueeze(0))
+        loss = planning_cost(pred, seq.unsqueeze(0), goal_t, current_pos, prev, active_action_dim, args).mean()
+        optimizer.zero_grad(set_to_none=True)
+        loss.backward()
+        optimizer.step()
+    with torch.no_grad():
+        seq = torch.tanh(logits)
+        if inactive is not None:
+            seq = seq * inactive
+        with autocast_context(mean.device, args.precision):
+            pred = rollout_latent(model, z, c, seq.unsqueeze(0))
+        candidates = pred[0, :, :2].detach().cpu().numpy()[None, ...] if args.make_gifs else None
+        action = seq[0, :active_action_dim].detach().cpu().numpy()
+        return (
+            np.clip(action, -1.0, 1.0).astype(np.float32),
+            candidates,
+            seq.detach().cpu().numpy(),
+        )
+
+
+@torch.no_grad()
+def donor_context_for_flowmo(model, env: SurfaceBoatEnv, args, seed: int) -> torch.Tensor | None:
+    if not hasattr(model, "to_c"):
+        return None
+    rng = np.random.default_rng(seed + 99_999)
+    donor = SurfaceBoatEnv(
+        boat=env.config.boat,
+        flow_type=env.config.flow_type,
+        boundary="terminate",
+        episode_steps=model.config.context_len + 8,
+        seed=seed + 99,
+    )
+    donor.reset(flow_type=env.config.flow_type, random_velocity=False)
+    image_history = deque(maxlen=args.history_len)
+    action_history = deque(maxlen=args.history_len)
+    action = np.zeros((model.config.action_dim,), dtype=np.float32)
+    for _ in range(args.history_len):
+        image_history.append(clean_observation(donor, args.image_size, args.visual_scale))
+        action_history.append(action.copy())
+        raw = rng.uniform(-0.5, 0.5, size=donor.action_dim).astype(np.float32)
+        donor.step(raw)
+        action = pad_action(raw, model.config.action_dim)
+    device = next(model.parameters()).device
+    images = torch.as_tensor(np.asarray(image_history, dtype=np.uint8), device=device).unsqueeze(0)
+    actions = torch.as_tensor(np.asarray(action_history, dtype=np.float32), device=device).unsqueeze(0)
+    with autocast_context(device, args.precision):
+        return model.encode(images, actions)[1].detach()
+
+
+def traditional_action(method: str, image_history: deque, env: SurfaceBoatEnv, goal: np.ndarray) -> np.ndarray:
+    evaluate_module = importlib.import_module(f"experiments.{method}.src.evaluate")
+    image = np.transpose(image_history[-1], (1, 2, 0))
+    history = [np.transpose(x, (1, 2, 0)) for x in image_history]
+    cfg = {
+        "image": image,
+        "history": history,
+        "true_flow": env.last_flow_velocity.copy(),
+        "goal": goal.astype(float).tolist(),
+        "action_dim": env.action_dim,
+        "boat": env.config.boat,
+    }
+    return evaluate_module.evaluate(cfg)[: env.action_dim].astype(np.float32)
+
+
+def evaluate_one_method(method: str, args) -> dict:
+    torch.manual_seed(args.seed)
+    learned = method in LEARNED_METHODS
+    model = None
+    if learned:
+        _cfg, model = build_method(method)
+        state = torch.load(Path("experiments") / method / "checkpoint" / args.checkpoint_name, map_location="cpu")
+        model.load_state_dict(state)
+        model.to(torch.device(args.device))
+        if torch.device(args.device).type == "cuda":
+            model.to(memory_format=torch.channels_last)
+        model.eval()
+        for param in model.parameters():
+            param.requires_grad_(False)
+    results = []
+    gif_dir = Path(args.out) / "gifs"
+    gif_dir.mkdir(parents=True, exist_ok=True)
+    context_modes = args.context_modes if method == "flowmo" else ["inferred"]
+    for context_mode in context_modes:
+        for ep in range(args.episodes):
+            episode_seed = int(args.seed + ep)
+            rng = np.random.default_rng(episode_seed)
+            env = SurfaceBoatEnv(
+                boat=args.boat,
+                flow_type=args.flow_type,
+                boundary="terminate",
+                episode_steps=args.max_steps,
+                seed=episode_seed,
+            )
+            reset_task(env, args.task, args.flow_type, rng)
+            goals = task_goals(args.task, rng)
+            goal_idx = 0
+            image_history = deque(maxlen=args.history_len)
+            action_history = deque(maxlen=args.history_len)
+            zero = np.zeros((model.config.action_dim if learned else 3,), dtype=np.float32)
+            first = clean_observation(env, args.image_size, args.visual_scale)
+            for _ in range(args.history_len):
+                image_history.append(first.copy())
+                action_history.append(zero.copy())
+            donor_context = donor_context_for_flowmo(model, env, args, episode_seed) if learned and context_mode == "shuffled" else None
+            trajectory = [env.full_state()[:6].copy()]
+            frames = []
+            prev_action = np.zeros((env.action_dim,), dtype=np.float32)
+            energy = 0.0
+            reached_times: list[int] = []
+            min_goal_dists = np.full((len(goals),), np.inf, dtype=np.float32)
+            passive_steps = args.passive_steps if args.task == "passive_to_active" else 0
+            planned = None
+            learned_plan_mean = None
+            for t in range(args.max_steps):
+                goal = goals[goal_idx]
+                if t < passive_steps:
+                    action = np.zeros((env.action_dim,), dtype=np.float32)
+                    planned = None
+                    learned_plan_mean = None
+                elif learned:
+                    action, planned, learned_plan_mean = learned_plan(
+                        model,
+                        image_history,
+                        action_history,
+                        goal,
+                        env.action_dim,
+                        args,
+                        prev_action,
+                        learned_plan_mean,
+                        context_mode,
+                        donor_context,
+                    )
+                else:
+                    action = traditional_action(method, image_history, env, goal)
+                    planned = None
+                prev_action = action.copy()
+                _obs, _reward, done, _info = env.step(action)
+                energy += float(np.sum(action * action))
+                trajectory.append(env.full_state()[:6].copy())
+                image_history.append(clean_observation(env, args.image_size, args.visual_scale))
+                action_history.append(pad_action(action, len(action_history[-1])))
+                dists = np.linalg.norm(goals - env.state[:2], axis=1)
+                min_goal_dists = np.minimum(min_goal_dists, dists)
+                if ep < args.make_gifs and t % args.gif_stride == 0:
+                    frames.append(
+                        render_frame(
+                            env.full_state()[:6],
+                            env.spec,
+                            env.flow,
+                            env.workspace,
+                            trajectory=np.asarray(trajectory),
+                            goal=goal,
+                            planned=planned,
+                            t=env.time,
+                        )
+                    )
+                if float(dists[goal_idx]) < args.success_radius:
+                    reached_times.append(t + 1)
+                    if args.task == "station_keeping":
+                        if t >= max(40, args.max_steps // 3):
+                            break
+                    else:
+                        goal_idx += 1
+                        learned_plan_mean = None
+                        if goal_idx >= len(goals):
+                            break
+                if done:
+                    break
+            path = np.asarray(trajectory)[:, :2]
+            final_goal = goals[min(goal_idx, len(goals) - 1)]
+            record = {
+                "method": method,
+                "context_mode": context_mode,
+                "episode": ep,
+                "success": bool(goal_idx >= len(goals) or (args.task == "station_keeping" and np.linalg.norm(env.state[:2] - goals[0]) < args.success_radius)),
+                "final_distance": float(np.linalg.norm(env.state[:2] - final_goal)),
+                "mean_min_goal_distance": float(min_goal_dists.mean()),
+                "path_length": float(np.linalg.norm(np.diff(path, axis=0), axis=-1).sum()) if len(path) > 1 else 0.0,
+                "energy": energy,
+                "steps": len(trajectory) - 1,
+                "reached_times": reached_times,
+            }
+            results.append(record)
+            if ep < args.make_gifs and frames:
+                name = f"image_planning_{method}_{context_mode}_{args.boat}_{args.task}_ep{ep:03d}.gif"
+                save_gif(frames, gif_dir / name, duration_ms=args.gif_duration_ms)
+    return summarize(method, args, results)
+
+
+def summarize(method: str, args, results: list[dict]) -> dict:
+    groups = sorted({r["context_mode"] for r in results})
+    by_context = {}
+    def success_mean(items: list[dict], key: str) -> float | None:
+        successful = [r[key] for r in items if r["success"]]
+        return float(np.mean(successful)) if successful else None
+
+    for context in groups:
+        items = [r for r in results if r["context_mode"] == context]
+        by_context[context] = {
+            "episodes": len(items),
+            "successes": len([r for r in items if r["success"]]),
+            "success_rate": float(np.mean([r["success"] for r in items])),
+            "final_distance_mean": float(np.mean([r["final_distance"] for r in items])),
+            "mean_min_goal_distance": float(np.mean([r["mean_min_goal_distance"] for r in items])),
+            "path_length_success_mean": success_mean(items, "path_length"),
+            "energy_success_mean": success_mean(items, "energy"),
+            "steps_success_mean": success_mean(items, "steps"),
+        }
+    return {
+        "method": method,
+        "task": args.task,
+        "boat": args.boat,
+        "flow_type": args.flow_type,
+        "by_context": by_context,
+        "results": results,
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--methods", nargs="+", default=LEARNED_METHODS + TRADITIONAL_METHODS)
+    parser.add_argument("--task", choices=["reach_uniform", "counterflow", "station_keeping", "passive_to_active", "waypoint_square", "waypoint_zigzag"], default="reach_uniform")
+    parser.add_argument("--boat", choices=["twin", "triangle"], default="twin")
+    parser.add_argument("--flow-type", choices=["uniform", "slowly_varying", "vortex_center", "gradient", "turbulent_patch"], default="uniform")
+    parser.add_argument("--episodes", type=int, default=50)
+    parser.add_argument("--max-steps", type=int, default=240)
+    parser.add_argument("--planner", choices=["gradient", "cem"], default="gradient")
+    parser.add_argument("--passive-steps", type=int, default=25)
+    parser.add_argument("--history-len", type=int, default=32)
+    parser.add_argument("--image-size", type=int, default=160)
+    parser.add_argument("--visual-scale", type=float, default=2.5)
+    parser.add_argument("--checkpoint-name", default="image_local.pt")
+    parser.add_argument("--context-modes", nargs="+", default=["inferred", "zero", "shuffled"])
+    parser.add_argument("--cem-horizon", type=int, default=15)
+    parser.add_argument("--cem-population", type=int, default=128)
+    parser.add_argument("--cem-elites", type=int, default=16)
+    parser.add_argument("--cem-iterations", type=int, default=3)
+    parser.add_argument("--cem-action-std", type=float, default=0.5)
+    parser.add_argument("--cem-knots", type=int, default=5)
+    parser.add_argument("--planner-iterations", type=int, default=30)
+    parser.add_argument("--planner-lr", type=float, default=0.08)
+    parser.add_argument("--cem-w-goal", type=float, default=8.0)
+    parser.add_argument("--cem-w-path", type=float, default=0.45)
+    parser.add_argument("--cem-w-route", type=float, default=1.0)
+    parser.add_argument("--cem-w-heading-goal", type=float, default=0.0)
+    parser.add_argument("--cem-w-action", type=float, default=0.04)
+    parser.add_argument("--cem-w-smooth", type=float, default=0.08)
+    parser.add_argument("--cem-w-boundary", type=float, default=10.0)
+    parser.add_argument("--cem-w-progress", type=float, default=2.0)
+    parser.add_argument("--success-radius", type=float, default=0.65)
+    parser.add_argument("--make-gifs", type=int, default=3)
+    parser.add_argument("--gif-stride", type=int, default=1)
+    parser.add_argument("--gif-duration-ms", type=int, default=55)
+    parser.add_argument("--seed", type=int, default=33)
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--precision", choices=["fp32", "bf16", "fp16"], default="fp32")
+    parser.add_argument("--out", default="experiments/reports/image_planning")
+    args = parser.parse_args()
+
+    out_dir = Path(args.out)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    payload = [evaluate_one_method(method, args) for method in args.methods]
+    out_path = out_dir / f"{args.task}_{args.boat}_{args.flow_type}.json"
+    out_path.write_text(json.dumps(payload, indent=2))
+    print(json.dumps(payload, indent=2))
+
+
+if __name__ == "__main__":
+    main()

experiments/evaluate_image_world_models.py

@@ -0,0 +1,269 @@
+"""Evaluate trained image-input world models on long open-loop rollouts."""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch.utils.data import DataLoader
+
+from experiments.shared.src.data.image_dataset import ImageTrajectoryDataset
+from experiments.shared.src.methods import PAPER_LEARNED_METHODS
+from experiments.shared.src.vision.clean_renderer import render_clean_boat_history_tensor
+from experiments.train_image_world_models import configure_training_runtime
+from experiments.train_image_world_models import autocast_context
+from experiments.train_image_world_models import decode_predictions
+from experiments.train_image_world_models import required_model_history
+from experiments.train_image_world_models import selected_history_indices
+
+
+METHODS = PAPER_LEARNED_METHODS
+
+
+def loader_kwargs(num_workers: int) -> dict:
+    if num_workers <= 0:
+        return {}
+    return {
+        "multiprocessing_context": "spawn",
+        "persistent_workers": True,
+        "prefetch_factor": 4,
+    }
+
+
+def prepare_batch(batch, args, device: torch.device):
+    observation_hist, actions, future_actions, targets, origin, prev_origin, flow_type_id, boat_id = batch
+    history_indices = getattr(args, "history_indices", None)
+    if history_indices is None:
+        model_history_len = int(getattr(args, "model_history_len", observation_hist.shape[1]))
+        observation_hist = observation_hist[:, -model_history_len:]
+        actions = actions[:, -model_history_len:]
+    else:
+        observation_hist = observation_hist[:, history_indices]
+        actions = actions[:, history_indices]
+    actions = actions.to(device, non_blocking=True)
+    future_actions = future_actions.to(device, non_blocking=True)
+    targets = targets.to(device, non_blocking=True)
+    origin = origin.to(device, non_blocking=True)
+    if args.render_mode == "device":
+        states = observation_hist.to(device, non_blocking=True)
+        boat_id_device = boat_id.to(device, non_blocking=True)
+        images = render_clean_boat_history_tensor(
+            states,
+            boat_id_device,
+            image_size=args.image_size,
+            visual_scale=args.visual_scale,
+        )
+    else:
+        images = observation_hist.to(device, non_blocking=True)
+    return images, actions, future_actions, targets, origin, prev_origin, flow_type_id, boat_id
+
+
+def build_method(method: str):
+    config_module = importlib.import_module(f"experiments.{method}.src.config")
+    model_module = importlib.import_module(f"experiments.{method}.src.model")
+    cfg = config_module.default_config()
+    return cfg, model_module.build_model(cfg)
+
+
+def load_flow_names(source_npz: str) -> dict[int, str]:
+    src = np.load(source_npz, allow_pickle=False)
+    metadata = json.loads(str(src["metadata"]))
+    return {int(v): str(k) for k, v in metadata["flows"].items()}
+
+
+def load_group_names(source_npz: str, key: str) -> dict[int, str]:
+    src = np.load(source_npz, allow_pickle=False)
+    metadata = json.loads(str(src["metadata"]))
+    return {int(v): str(k) for k, v in metadata[key].items()}
+
+
+@torch.no_grad()
+def rollout_with_context(model, images: torch.Tensor, actions: torch.Tensor, future_actions: torch.Tensor, mode: str) -> torch.Tensor:
+    z, c = model.encode(images, actions)
+    if mode == "zero":
+        c = torch.zeros_like(c)
+    elif mode == "shuffled":
+        c = c.roll(shifts=1, dims=0)
+    if hasattr(model, "rollout_with_context"):
+        return model.rollout_with_context(z, c, future_actions)
+    preds = []
+    cur = z
+    for t in range(future_actions.shape[1]):
+        cur = model.step(cur, future_actions[:, t], c)
+        preds.append(model.decoder(cur))
+    return torch.stack(preds, dim=1)
+
+
+@torch.no_grad()
+def evaluate_model(
+    model,
+    loader,
+    device: torch.device,
+    horizon: int,
+    target_mode: str,
+    flow_names: dict[int, str],
+    traj_names: dict[int, str],
+    boat_names: dict[int, str],
+    context_mode: str,
+    args,
+) -> dict:
+    model.eval()
+    steps = [s for s in [1, 3, 6, 8, 10, 20, 30, 40, 60] if s <= horizon]
+    pos_sum = np.zeros(horizon, dtype=np.float64)
+    heading_sum = np.zeros(horizon, dtype=np.float64)
+    flow_pos: dict[int, np.ndarray] = {}
+    flow_heading: dict[int, np.ndarray] = {}
+    flow_count: dict[int, int] = {}
+    traj_pos: dict[int, np.ndarray] = {}
+    traj_heading: dict[int, np.ndarray] = {}
+    traj_count: dict[int, int] = {}
+    boat_pos: dict[int, np.ndarray] = {}
+    boat_heading: dict[int, np.ndarray] = {}
+    boat_count: dict[int, int] = {}
+    count = 0
+    cursor = 0
+    for batch in loader:
+        images, actions, future_actions, targets, origin, _prev_origin, flow_type_id, _boat_id = prepare_batch(batch, args, device)
+        with autocast_context(device, args.precision):
+            if context_mode == "inferred":
+                encoded = model.rollout(images, actions, future_actions)
+            else:
+                encoded = rollout_with_context(model, images, actions, future_actions, context_mode)
+        pred = decode_predictions(encoded.float(), origin, target_mode)
+        pos = torch.linalg.norm(pred[..., :2] - targets[..., :2], dim=-1)
+        pred_angle = torch.atan2(pred[..., 3], pred[..., 2])
+        target_angle = torch.atan2(targets[..., 3], targets[..., 2])
+        heading = torch.atan2(torch.sin(pred_angle - target_angle), torch.cos(pred_angle - target_angle)).abs()
+        pos_np = pos.cpu().numpy()
+        heading_np = heading.cpu().numpy()
+        pos_sum += pos_np.sum(axis=0)
+        heading_sum += heading_np.sum(axis=0)
+        count += int(pos_np.shape[0])
+        flow_np = flow_type_id.numpy()
+        batch_indices = loader.dataset.indices[cursor : cursor + int(pos_np.shape[0])]
+        cursor += int(pos_np.shape[0])
+        traj_np = np.array([loader.dataset.traj_type_ids[ep] for ep, _t in batch_indices], dtype=np.int64)
+        boat_np = np.array([loader.dataset.boat_ids[ep] for ep, _t in batch_indices], dtype=np.int64)
+        for flow_id in np.unique(flow_np):
+            mask = flow_np == flow_id
+            fid = int(flow_id)
+            flow_pos.setdefault(fid, np.zeros(horizon, dtype=np.float64))
+            flow_heading.setdefault(fid, np.zeros(horizon, dtype=np.float64))
+            flow_count[fid] = flow_count.get(fid, 0) + int(mask.sum())
+            flow_pos[fid] += pos_np[mask].sum(axis=0)
+            flow_heading[fid] += heading_np[mask].sum(axis=0)
+        for traj_id in np.unique(traj_np):
+            mask = traj_np == traj_id
+            tid = int(traj_id)
+            traj_pos.setdefault(tid, np.zeros(horizon, dtype=np.float64))
+            traj_heading.setdefault(tid, np.zeros(horizon, dtype=np.float64))
+            traj_count[tid] = traj_count.get(tid, 0) + int(mask.sum())
+            traj_pos[tid] += pos_np[mask].sum(axis=0)
+            traj_heading[tid] += heading_np[mask].sum(axis=0)
+        for boat_id in np.unique(boat_np):
+            mask = boat_np == boat_id
+            bid = int(boat_id)
+            boat_pos.setdefault(bid, np.zeros(horizon, dtype=np.float64))
+            boat_heading.setdefault(bid, np.zeros(horizon, dtype=np.float64))
+            boat_count[bid] = boat_count.get(bid, 0) + int(mask.sum())
+            boat_pos[bid] += pos_np[mask].sum(axis=0)
+            boat_heading[bid] += heading_np[mask].sum(axis=0)
+    result = summarize(pos_sum / count, heading_sum / count, steps)
+    by_flow = {}
+    for fid, n in sorted(flow_count.items()):
+        by_flow[flow_names.get(fid, str(fid))] = summarize(flow_pos[fid] / n, flow_heading[fid] / n, steps)
+    result["by_flow"] = by_flow
+    result["by_trajectory"] = {
+        traj_names.get(tid, str(tid)): summarize(traj_pos[tid] / n, traj_heading[tid] / n, steps)
+        for tid, n in sorted(traj_count.items())
+    }
+    result["by_boat"] = {
+        boat_names.get(bid, str(bid)): summarize(boat_pos[bid] / n, boat_heading[bid] / n, steps)
+        for bid, n in sorted(boat_count.items())
+    }
+    return result
+
+
+def summarize(pos_mean: np.ndarray, heading_mean: np.ndarray, steps: list[int]) -> dict[str, float]:
+    result: dict[str, float] = {}
+    for step in steps:
+        result[f"pos{step}"] = float(pos_mean[step - 1])
+        result[f"heading{step}"] = float(heading_mean[step - 1])
+    return result
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--methods", nargs="+", default=METHODS)
+    parser.add_argument("--test-source", default="data/paper/test_unseen_flow.npz")
+    parser.add_argument("--test-episodes", type=int, default=256)
+    parser.add_argument("--history-len", type=int, default=32)
+    parser.add_argument("--horizon", type=int, default=60)
+    parser.add_argument("--test-windows", type=int, default=4096)
+    parser.add_argument("--batch-size", type=int, default=64)
+    parser.add_argument("--seed", type=int, default=20)
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--target-mode", choices=["absolute_normalized", "relative_motion"], default="absolute_normalized")
+    parser.add_argument("--checkpoint-name", default="image_local.pt")
+    parser.add_argument("--out", default="experiments/reports/image_long_rollout_eval.json")
+    parser.add_argument("--num-workers", type=int, default=4)
+    parser.add_argument("--image-size", type=int, default=160)
+    parser.add_argument("--visual-scale", type=float, default=2.5)
+    parser.add_argument("--render-mode", choices=["device", "dataset"], default="device")
+    parser.add_argument("--precision", choices=["fp32", "bf16", "fp16"], default="fp32")
+    args = parser.parse_args()
+    device = torch.device(args.device)
+    configure_training_runtime(device)
+    flow_names = load_flow_names(args.test_source)
+    traj_names = load_group_names(args.test_source, "trajectories")
+    boat_names = load_group_names(args.test_source, "boats")
+    ds = ImageTrajectoryDataset(
+        args.test_source,
+        history_len=args.history_len,
+        horizon=args.horizon,
+        episodes=args.test_episodes,
+        max_windows=args.test_windows,
+        seed=args.seed,
+        image_size=args.image_size,
+        visual_scale=args.visual_scale,
+        return_aux=True,
+        render_images=args.render_mode == "dataset",
+    )
+    loader = DataLoader(
+        ds,
+        batch_size=args.batch_size,
+        shuffle=False,
+        num_workers=args.num_workers,
+        pin_memory=device.type == "cuda",
+        **loader_kwargs(args.num_workers),
+    )
+    payload = []
+    for method in args.methods:
+        _cfg, model = build_method(method)
+        state = torch.load(Path("experiments") / method / "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)
+        args.model_history_len = required_model_history(model, args.history_len)
+        args.history_indices = selected_history_indices(model, args.history_len)
+        item = {
+            "method": method,
+            "inferred": evaluate_model(model, loader, device, args.horizon, args.target_mode, flow_names, traj_names, boat_names, "inferred", args),
+        }
+        if method == "flowmo":
+            item["context_zero"] = evaluate_model(model, loader, device, args.horizon, args.target_mode, flow_names, traj_names, boat_names, "zero", args)
+            item["context_shuffled"] = evaluate_model(model, loader, device, args.horizon, args.target_mode, flow_names, traj_names, boat_names, "shuffled", args)
+        payload.append(item)
+    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()

experiments/figures/README.md

@@ -0,0 +1,11 @@
+# Figures
+
+Paper-facing figures generated from method results.
+
+Subdirectories:
+
+- `prediction/`
+- `planning/`
+- `context/`
+- `energy/`
+- `trajectory/`

experiments/flowmo/README.md

@@ -0,0 +1,9 @@
+# FlowMo: Flow-Momentum World Model
+
+FlowMo is the proposed Flow-Momentum World Model for surface vehicles under hidden external drift.
+
+Core idea:
+
+- Short-history latent state for endogenous state, momentum, and actuator response.
+- Long-history drift context for exogenous flow effects.
+- Context-conditioned residual transition for drift-aware rollout and planning.

experiments/flowmo/__init__.py

@@ -0,0 +1 @@
+"""FlowMo experiment package."""

experiments/flowmo/src/__init__.py

@@ -0,0 +1 @@
+"""FlowMo source package."""

experiments/flowmo/src/config.py

@@ -0,0 +1,7 @@
+"""FlowMo default config."""
+
+from experiments.shared.src.models.image_world_models import ImageWorldModelConfig
+
+def default_config():
+    """Return the default FlowMo image-input config."""
+    return ImageWorldModelConfig(emb_dim=64, z_dim=112, c_dim=8, hidden_dim=128, history_len=8, context_len=32)

experiments/flowmo/src/model.py

@@ -0,0 +1,18 @@
+"""FlowMo model definition."""
+
+import torch
+
+from experiments.flowmo.src.config import default_config
+from experiments.shared.src.models.image_world_models import FlowMoImageWorldModel
+
+def build_model(config):
+    """Build the FlowMo world model."""
+    return FlowMoImageWorldModel(config)
+
+
+def load_model(checkpoint_path, config=None):
+    """Load a trained FlowMo checkpoint."""
+    cfg = default_config() if config is None else config
+    model = build_model(cfg)
+    model.load_state_dict(torch.load(checkpoint_path, map_location="cpu"))
+    return model

experiments/flowmo/src/plan.py

@@ -0,0 +1,6 @@
+"""FlowMo closed-loop planning interface."""
+
+
+def plan(model, env, task, config):
+    """Run closed-loop planning with FlowMo."""
+    return config["planner"](model, env, task, config)

experiments/flowmo/src/predict.py

@@ -0,0 +1,7 @@
+"""FlowMo rollout interface."""
+
+
+def rollout(model, batch, horizon):
+    """Run open-loop FlowMo prediction."""
+    images, actions, future_actions = batch
+    return model.rollout(images, actions, future_actions[:, :horizon])

experiments/flowmo/src/train.py

@@ -0,0 +1,16 @@
+"""FlowMo training entry point."""
+
+import torch.nn.functional as F
+
+
+def train(config):
+    """Train FlowMo on the configured dataset."""
+    model = config["model"]
+    optimizer = config["optimizer"]
+    images, actions, future_actions, targets = config["batch"]
+    pred = model.rollout(images, actions, future_actions)
+    loss = F.mse_loss(pred, targets)
+    optimizer.zero_grad(set_to_none=True)
+    loss.backward()
+    optimizer.step()
+    return float(loss.detach())