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

TL;DR

This paper introduces a flow-based policy parameterization for max-entropy reinforcement learning, enhancing expressiveness and robustness, with a theoretical analysis and a case study on LQR problems showing optimal policy learning.

Contribution

It proposes a novel flow-based policy with an online flow matching update method and provides theoretical insights into sampling distribution effects.

Findings

The flow-based policy achieves high expressiveness.

The ISFM algorithm effectively learns optimal policies.

Theoretical analysis links sampling choices to learning efficiency.

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.