WTF?! Simulation-Free Reinforcement Learning with Wasserstein-Tilted Flow Maps

Abstract

Reward fine-tuning aims to update a pre-trained flow-based generative model to improve the downstream reward of its generated samples. Existing methods typically formulate this problem as sampling from a reward-tilted distribution, the solution to a KL-regularized reward-maximization problem. Here, we introduce an optimal transport regularizer built directly from the pre-trained drift. Unlike KL reward tilting, the resulting objective transports individual samples toward higher reward rather than reweighting the base distribution. We show that the resulting problem is equivalent to a deterministic optimal control problem on the flow. Given a pre-trained flow map, this equivalence yields a simulation-free reinforcement learning algorithm for fine-tuning generative flows. We call the resulting framework Wasserstein-Tilted Flow Maps (WTF), the first end-to-end fine-tuning recipe native to flow maps. The output is a fine-tuned flow map that retains strong reward-aligned performance at few-step inference budgets without post-hoc distillation. Experiments on ImageNet-256 and text-to-image show that WTF achieves higher reward with comparable or higher diversity than baselines, while requiring up to 280× less training compute. More broadly, we argue that accelerated samplers such as flow maps are essential infrastructure for efficient post-training, and that the dominant KL-regularized formulation is only one of many choices worth revisiting.

Simulation-Free Learning of Population Dynamics with Wasserstein Lagrangian Residuals
Fedor Sergeev, Markus Heinonen, Daniel Waxman +4
arXiv2 Oct 2026PreprintCC BY 4.0Full textStatisticsarXiv
Exact Fast Batch Simulation for Tabular Reinforcement Learning
Haochen Zhang, Lingzhou Xue, Zhong Zheng
arXiv3 Oct 2026PreprintCC BY 4.0Full textStatisticsarXiv
Integrating YOLOv8-based disaster detection with deep reinforcement learning for closed-loop UAV emergency response scheduling: A simulation study
Chunsheng Zhang, Ruichun He
PLOS One29 Oct 2026Peer-reviewedCC BYFull textNeuroscience StatisticsPLOSEurope PMCPMC
Model-free Reinforcement Learning for Continuous Time and State: A Stochastic Maximum Principle Approach
Lijun Bo, Yijie Huang, Jingfei Wang
arXiv9 Oct 2026PreprintarXiv licenceFull textOptimisation & controlarXiv
Hamiltonian Primal-Dual Policy Reinforcement Learning for Derivative-Free Model-Based Control
K. Rabiei, Z. Rui, X. Lu +1
arXiv13 Oct 2026PreprintarXiv licenceFull textOptimisation & controlarXiv

Cite

BibTeX RIS

BibTeX

@misc{mammadov2026simulationfree,
  title = {{WTF?! Simulation-Free Reinforcement Learning with Wasserstein-Tilted Flow Maps}},
  author = {Abbas Mammadov and Jerry Y. Huang and Justin Lin and Partha Kaushik and Sheel Shah and Kartik Nair and Yee Whye Teh and Nicholas M. Boffi},
  howpublished = {arXiv},
  year = {2026},
  month = oct,
  doi = {10.48550/arxiv.2609.27033},
  eprint = {2609.27033},
  archivePrefix = {arXiv},
  url = {https://arxiv.org/abs/2609.27033},
  note = {CC BY 4.0},
}

RIS

TY  - UNPB
TI  - WTF?! Simulation-Free Reinforcement Learning with Wasserstein-Tilted Flow Maps
AU  - Abbas Mammadov
AU  - Jerry Y. Huang
AU  - Justin Lin
AU  - Partha Kaushik
AU  - Sheel Shah
AU  - Kartik Nair
AU  - Yee Whye Teh
AU  - Nicholas M. Boffi
PY  - 2026
DA  - 2026/10/12
PB  - arXiv
DO  - 10.48550/arxiv.2609.27033
UR  - https://arxiv.org/abs/2609.27033
AB  - Reward fine-tuning aims to update a pre-trained flow-based generative model to improve the downstream reward of its generated samples. Existing methods typically formulate this problem as sampling from a reward-tilted distribution, the solution to a KL-regularized reward-maximization problem. Here, we introduce an optimal transport regularizer built directly from the pre-trained drift. Unlike KL reward tilting, the resulting objective transports individual samples toward higher reward rather than reweighting the base distribution. We show that the resulting problem is equivalent to a deterministic optimal control problem on the flow. Given a pre-trained flow map, this equivalence yields a simulation-free reinforcement learning algorithm for fine-tuning generative flows. We call the resulting framework Wasserstein-Tilted Flow Maps (WTF), the first end-to-end fine-tuning recipe native to flow maps. The output is a fine-tuned flow map that retains strong reward-aligned performance at few-step inference budgets without post-hoc distillation. Experiments on ImageNet-256 and text-to-image show that WTF achieves higher reward with comparable or higher diversity than baselines, while requiring up to 280× less training compute. More broadly, we argue that accelerated samplers such as flow maps are essential infrastructure for efficient post-training, and that the dominant KL-regularized formulation is only one of many choices worth revisiting.
ER  -

Abstract

Related

Cite