README.md

---
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/<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.