FlowMo-WM

FlowMo-WM is a visual world model for surface vehicles under hidden ambient drift. The model uses clean top-down image histories and action histories to infer a short-history object-motion state and a long-history drift context, then predicts future latent states with a zero-context residual transition.

This repository contains the code, canonical datasets, trained checkpoints, evaluation outputs, GIFs, tables, and report files needed to reproduce the paper experiments.

Repository Layout

data/paper/                         canonical train/test splits
driftwm/                            simulator and flow-field code
experiments/shared/                 shared data, renderer, planner, metrics, model utilities
experiments/flowmo/                 proposed model
experiments/leworldmodel/           JEPA-style latent world-model baseline
experiments/planet/                 PlaNet RSSM baseline
experiments/tdmpc2/                 TD-MPC2-style latent dynamics baseline
experiments/*_los_controller/       traditional non-WM controllers
experiments/reports/                prediction, probes, planning JSON, GIFs, paper tables/figures
tests/                              interface and pipeline tests

Data

The canonical paper data are:

data/paper/train.npz
data/paper/test.npz
data/paper/generation_config.json
data/paper/dataset_card.md

All learned models use the same splits, image renderer, optimizer budget, rollout targets, and evaluation protocol. Images are rendered online from simulator states as clean RGB frames. Model inputs do not include flow arrows, goal markers, velocity vectors, trajectory overlays, flow labels, or low-dimensional ground-truth state.

The train split, test split, and final planning tasks use the same static flow-family set:

noflow
uniform
vortex_center
double_gyre
source_sink
source_sink_pair
gradient
shear
turbulent_patch
random_fourier

Methods

Learned world models:

Directory	Report name	Role
experiments/flowmo	FlowMo	Proposed short-state / long-context residual world model.
experiments/leworldmodel	LeWorldModel	JEPA-style image latent prediction baseline.
experiments/planet	PlaNet RSSM	Recurrent state-space world-model baseline.
experiments/tdmpc2	TD-MPC2 Dynamics	Compact task-oriented latent dynamics baseline.

Traditional non-WM controllers:

Directory	Report name	Role
experiments/pid_los_controller	PID/LOS	Hand-designed waypoint tracking baseline.
experiments/no_flow_los_controller	No-Flow LOS	LOS controller that ignores ambient current.
experiments/current_estimator_los_controller	Current-Estimator LOS	LOS controller with recent-drift current compensation.
experiments/oracle_flow_los_controller	Oracle-Flow LOS	LOS controller with privileged true local-flow feed-forward.

Baseline definitions are in experiments/BASELINES.md. The complete experiment matrix is in experiments/EXPERIMENT_MATRIX.md.

Install

python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install -e .

Run tests:

python -m pytest -q tests

Reproduce

Run the complete paper pipeline:

python -m experiments.run_paper_image_pipeline

The command trains all learned models, evaluates prediction, runs FlowMo latent probes, evaluates planning for learned and traditional methods, generates GIFs, exports paper artifacts, and writes:

experiments/reports/paper_prediction.json
experiments/reports/paper_flowmo_latent_probes.json
experiments/reports/paper_planning/*.json
experiments/reports/paper_planning/gifs/*.gif
experiments/reports/paper_artifacts/
experiments/reports/paper_report.md

Stages can be rerun separately:

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

To rebuild tables and figure-ready artifacts from existing results:

python -m experiments.summarize_paper_image_results
python -m experiments.export_paper_artifacts

Formal Outputs

The repository includes the current formal outputs:

experiments/<learned_method>/checkpoint/paper.pt
experiments/<learned_method>/checkpoint/paper_step_*.pt
experiments/<learned_method>/result/parameter_count.json
experiments/<learned_method>/result/paper_training.json
experiments/reports/paper_prediction.json
experiments/reports/paper_flowmo_latent_probes.json
experiments/reports/paper_planning/
experiments/reports/paper_artifacts/
experiments/reports/paper_report.md

Planning JSON files record every episode, including success, final distance, trajectory length, control effort sum_t ||a_t||_2^2, time to goal, path, actions, and per-frame metadata. Planning GIFs render the task background with flow arrows, the oriented boat, targets, and executed trajectory.