---
library_name: pytorch
tags:
  - robotics
  - world-model
  - visual-world-model
  - model-based-control
  - surface-vehicle
  - hidden-drift
---

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

This repository currently contains the public code, tests, configuration, and canonical paper datasets. Official checkpoints, generated GIFs, tables, and full experiment reports will be uploaded after the paper-scale training and evaluation runs finish.

## 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
```