Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Max-Entropy Reinforcement Learning with Flow Matching and A Case Study on LQR

Yuyang Zhang, Yang Hu, Bo Dai, Na Li

TL;DR

This work addresses the expressiveness gap in max-entropy RL by parameterizing SAC policies with flow-based models trained via flow matching. It introduces an online variant, importance sampling flow matching (ISFM), for policy improvement that uses samples from a flexible tilde p and provides a bound on the resulting policy approximation that depends on the fourth-order Renyi divergence $D_4(p^{\theta^*}_{x,1}\|\tilde p)$. Policy evaluation leverages the instantaneous change-of-variable formula to compute entropy-regularized returns efficiently. The approach is validated on max-entropy LQR, where the optimal policy is Gaussian with closed-form parameters; the theory and experiments show that SAC-ISFM with flow-based policies recovers the optimal Gaussian policy, confirming both convergence and practical viability.

Abstract

Soft actor-critic (SAC) is a popular algorithm for max-entropy reinforcement learning. In practice, the energy-based policies in SAC are often approximated using simple policy classes for efficiency, sacrificing the expressiveness and robustness. In this paper, we propose a variant of the SAC algorithm that parameterizes the policy with flow-based models, leveraging their rich expressiveness. In the algorithm, we evaluate the flow-based policy utilizing the instantaneous change-of-variable technique and update the policy with an online variant of flow matching developed in this paper. This online variant, termed importance sampling flow matching (ISFM), enables policy update with only samples from a user-specified sampling distribution rather than the unknown target distribution. We develop a theoretical analysis of ISFM, characterizing how different choices of sampling distributions affect the learning efficiency. Finally, we conduct a case study of our algorithm on the max-entropy linear quadratic regulator problems, demonstrating that the proposed algorithm learns the optimal action distribution.

Max-Entropy Reinforcement Learning with Flow Matching and A Case Study on LQR

TL;DR

This work addresses the expressiveness gap in max-entropy RL by parameterizing SAC policies with flow-based models trained via flow matching. It introduces an online variant, importance sampling flow matching (ISFM), for policy improvement that uses samples from a flexible tilde p and provides a bound on the resulting policy approximation that depends on the fourth-order Renyi divergence . Policy evaluation leverages the instantaneous change-of-variable formula to compute entropy-regularized returns efficiently. The approach is validated on max-entropy LQR, where the optimal policy is Gaussian with closed-form parameters; the theory and experiments show that SAC-ISFM with flow-based policies recovers the optimal Gaussian policy, confirming both convergence and practical viability.

Abstract

Soft actor-critic (SAC) is a popular algorithm for max-entropy reinforcement learning. In practice, the energy-based policies in SAC are often approximated using simple policy classes for efficiency, sacrificing the expressiveness and robustness. In this paper, we propose a variant of the SAC algorithm that parameterizes the policy with flow-based models, leveraging their rich expressiveness. In the algorithm, we evaluate the flow-based policy utilizing the instantaneous change-of-variable technique and update the policy with an online variant of flow matching developed in this paper. This online variant, termed importance sampling flow matching (ISFM), enables policy update with only samples from a user-specified sampling distribution rather than the unknown target distribution. We develop a theoretical analysis of ISFM, characterizing how different choices of sampling distributions affect the learning efficiency. Finally, we conduct a case study of our algorithm on the max-entropy linear quadratic regulator problems, demonstrating that the proposed algorithm learns the optimal action distribution.
Paper Structure (17 sections, 4 theorems, 43 equations, 2 figures, 2 algorithms)

This paper contains 17 sections, 4 theorems, 43 equations, 2 figures, 2 algorithms.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Max-Entropy Reinforcement Learning
  4. Flow Matching
  5. Training Flow-based Policies Online: From Soft Policy Iteration to Soft Actor-Critic
  6. SPI with Flow-based Policies
  7. Calculating Action Probability in Policy Evaluation
  8. Importance Sampling Flow Matching in Policy Improvement
  9. SAC with Flow-based Policies
  10. A Case Study --- Max-Entropy LQR
  11. The Max-Entropy LQR Problem
  12. Applying SPI to Max-Entropy LQR Problems
  13. Applying SAC-ISFM to Max-Entropy LQR Problems
  14. Conclusion
  15. Proofs
  16. ...and 2 more sections

Key Result

theorem 1

Consider vector field $v_{\tau}(u_\tau)$ and the generated probability $p_\tau(u_\tau)$. Suppose both of them are continuously differentiable in $(\tau,x)$. The following holds for all $\tau\in[0,1]$:

Figures (2)

  • Figure 1: Discounted LQR return (without entropy regularization) during training for different $\alpha$.
  • Figure 2: Mean differences and covariances difference between $\pi$ in episode $K=80000$ and the optimal policy (\ref{['thm:lqr_ent']}), averaged across states from $50$ trajectories with length 100.

Theorems & Definitions (7)

  • theorem 1: Theorem 4.7 in prelim_ode_intro
  • theorem 2
  • remark 1
  • theorem 3: maxent_lqr
  • lemma 1
  • proof
  • proof : Proof of \ref{['lem:lqr']}