Trust Region Constrained Measure Transport in Path Space for Stochastic Optimal Control and Inference
This work addresses a specific bottleneck in stochastic optimal control and inference for researchers and practitioners, offering a principled approach to improve optimization when target measures differ substantially from priors.
The paper tackles the challenge of approximating target path space measures in stochastic optimal control by proposing a trust region constrained method that gradually approaches the target measure through geometric annealing, demonstrating significant performance improvements in applications like diffusion-based sampling and fine-tuning of diffusion models.
Solving stochastic optimal control problems with quadratic control costs can be viewed as approximating a target path space measure, e.g. via gradient-based optimization. In practice, however, this optimization is challenging in particular if the target measure differs substantially from the prior. In this work, we therefore approach the problem by iteratively solving constrained problems incorporating trust regions that aim for approaching the target measure gradually in a systematic way. It turns out that this trust region based strategy can be understood as a geometric annealing from the prior to the target measure, where, however, the incorporated trust regions lead to a principled and educated way of choosing the time steps in the annealing path. We demonstrate in multiple optimal control applications that our novel method can improve performance significantly, including tasks in diffusion-based sampling, transition path sampling, and fine-tuning of diffusion models.