--- library_name: pytorch tags: - robotics - world-model - visual-world-model - model-based-control - surface-vehicle - hidden-drift --- # 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 ```text 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: ```text 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: ```text 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 ```bash python -m venv .venv source .venv/bin/activate python -m pip install --upgrade pip python -m pip install -e . ``` Run tests: ```bash python -m pytest -q tests ``` ## Reproduce Run the complete paper pipeline: ```bash 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: ```text 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: ```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 ``` To rebuild tables and figure-ready artifacts from existing results: ```bash python -m experiments.summarize_paper_image_results python -m experiments.export_paper_artifacts ``` ## Formal Outputs The repository includes the current formal outputs: ```text experiments//checkpoint/paper.pt experiments//checkpoint/paper_step_*.pt experiments//result/parameter_count.json experiments//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.