Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Controllable Flow Matching for Online Reinforcement Learning

Bin Wang, Boxiang Tao, Haifeng Jing, Hongbo Dou, Zijian Wang

TL;DR

CtrlFlow introduces trajectory-level data synthesis for online reinforcement learning by leveraging Conditional Flow Matching (CFM) to generate whole trajectories without explicitly modeling environment dynamics. A non-linear Controllability Gramian Matrix $N(u)$ is used to minimize control energy during sampling, ensuring global controllability and robustness to noise. The framework further augments data with a value-guided energy vector field, biasing trajectories toward high-return regions and improving policy learning. Empirical results on MuJoCo benchmarks show superior sample efficiency and competitive or superior asymptotic performance compared to model-based and model-free baselines, with ablations validating the roles of trajectory length, control, and value guidance.

Abstract

Model-based reinforcement learning (MBRL) typically relies on modeling environment dynamics for data efficiency. However, due to the accumulation of model errors over long-horizon rollouts, such methods often face challenges in maintaining modeling stability. To address this, we propose CtrlFlow, a trajectory-level synthetic method using conditional flow matching (CFM), which directly modeling the distribution of trajectories from initial states to high-return terminal states without explicitly modeling the environment transition function. Our method ensures optimal trajectory sampling by minimizing the control energy governed by the non-linear Controllability Gramian Matrix, while the generated diverse trajectory data significantly enhances the robustness and cross-task generalization of policy learning. In online settings, CtrlFlow demonstrates the better performance on common MuJoCo benchmark tasks than dynamics models and achieves superior sample efficiency compared to standard MBRL methods.

Controllable Flow Matching for Online Reinforcement Learning

TL;DR

CtrlFlow introduces trajectory-level data synthesis for online reinforcement learning by leveraging Conditional Flow Matching (CFM) to generate whole trajectories without explicitly modeling environment dynamics. A non-linear Controllability Gramian Matrix is used to minimize control energy during sampling, ensuring global controllability and robustness to noise. The framework further augments data with a value-guided energy vector field, biasing trajectories toward high-return regions and improving policy learning. Empirical results on MuJoCo benchmarks show superior sample efficiency and competitive or superior asymptotic performance compared to model-based and model-free baselines, with ablations validating the roles of trajectory length, control, and value guidance.

Abstract

Model-based reinforcement learning (MBRL) typically relies on modeling environment dynamics for data efficiency. However, due to the accumulation of model errors over long-horizon rollouts, such methods often face challenges in maintaining modeling stability. To address this, we propose CtrlFlow, a trajectory-level synthetic method using conditional flow matching (CFM), which directly modeling the distribution of trajectories from initial states to high-return terminal states without explicitly modeling the environment transition function. Our method ensures optimal trajectory sampling by minimizing the control energy governed by the non-linear Controllability Gramian Matrix, while the generated diverse trajectory data significantly enhances the robustness and cross-task generalization of policy learning. In online settings, CtrlFlow demonstrates the better performance on common MuJoCo benchmark tasks than dynamics models and achieves superior sample efficiency compared to standard MBRL methods.

Paper Structure

This paper contains 14 sections, 4 theorems, 20 equations, 6 figures, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Normalizing Flows and Flow Matching Models
  4. Flow Matching in Decision Making and RL
  5. Preliminaries
  6. Problem Setup
  7. Conditional Flow Matching Models
  8. Controllability Gramian Matrix
  9. Method
  10. Trajectory-Level Flow Matching Model
  11. Controllable Flow for Distribution Adaptation
  12. Value Guidance with Energy Vector Field
  13. Experiments
  14. Conclusion

Key Result

Theorem 5

For any $\widehat{\tau}^0\in\mathbb{R}^d$ and every $b\in B\in L^\infty$, the solution $\widehat{\tau}$ of $\dot{\widehat{\tau}}^t=v^t(\widehat{\tau}^t)+b^t$, $\widehat{\tau}^{t_0}=\widehat{\tau}^0$ can be expressed in following form: where $\mathbf{D}$ denotes the Jacobian matrix of the flow $\Phi$ with respect to time at the initial state $\widehat{\tau}^0$, and is defined for every fixed $\wid

Figures (6)

  • Figure 1: A validation episode in the Walker2d benchmark task of MuJoCo. The Soft Actor-Critic (SAC) algorithm is employed to optimize the policy in the same seed. In the first row (FailedReward: 23), the trajectory data directly generated by the conditional flow matching model is used for policy training. In the second and third rows (SuccessReward: 392), on the basis of the flow matching model, data with control correction and value guidance is utilized.
  • Figure 2: The overall of CtrlFlow. The CFM model $\mathcal{M}$ is trained using data from environment replay buffer $\mathcal{B}_\text{env}$ (Left Top). During sampling (Left Bottom), $\mathcal{M}$ generates trajectory data while being guided by both the value guidance vector field $\mathcal{G}$ and control vector field $\mathcal{V}_c$. These components work in concert to refine the original vector field through an ODE-based process from $t=0$ to $t=1$. The resulting trajectories are then stored in the model replay buffer $\mathcal{B}_\text{mod}$, which is subsequently combined with $\mathcal{B}_\text{env}$ for policy learning (Right SAC).
  • Figure 3: Performance on the MuJoCo Benchmark. CtrlFlow (red) and other five baselines on control tasks. The blue dashed lines indicate the asymptotic performance of SAC for reference. The solid lines indicate the mean while the shaded areas indicate the standard error over five different seeds.
  • Figure 4: Study on generated length and value guidance on MuJoCo Hopper-v3 task with five random seeds. (Left) Compare the performance of six different lengths with $h=2,5,8,10,30,50$ by model generation. (Right) The model with and without the value guidance vector field $\mathcal{G}$.
  • Figure 5: (Top Left) Study on control theory on MoJoCo Hopper-v3 task. (Others) The comparison of controllable and cosine similarity between the generated final state and the ground truth.
  • ...and 1 more figures

Theorems & Definitions (11)

  • Definition 1: Trajectory-Level Flow Matching Model
  • Definition 2: The Input of Model $\mathcal{M}$
  • Definition 4
  • Theorem 5
  • Definition 6: Non-linear Controllability Gramian Matrix
  • Remark 7
  • Theorem 8: Control Input $u$
  • Proposition 9
  • Definition 10: Controllability Vector Field
  • Theorem 11: Theorem 3.1. in feng2025guidance
  • ...and 1 more