Normalizing Flows are Capable Models for RL

TL;DR

This paper demonstrates that normalizing flows are powerful and expressive models suitable for reinforcement learning, offering simpler algorithms and improved performance across various RL tasks.

Contribution

The authors introduce a novel NF architecture that integrates into RL as a policy, Q-function, and occupancy measure, challenging the belief that NFs lack expressivity.

Findings

Achieves higher performance in imitation learning

Performs well in offline and goal-conditioned RL

Enables simpler RL algorithms

Abstract

Modern reinforcement learning (RL) algorithms have found success by using powerful probabilistic models, such as transformers, energy-based models, and diffusion/flow-based models. To this end, RL researchers often choose to pay the price of accommodating these models into their algorithms -- diffusion models are expressive, but are computationally intensive due to their reliance on solving differential equations, while autoregressive transformer models are scalable but typically require learning discrete representations. Normalizing flows (NFs), by contrast, seem to provide an appealing alternative, as they enable likelihoods and sampling without solving differential equations or autoregressive architectures. However, their potential in RL has received limited attention, partly due to the prevailing belief that normalizing flows lack sufficient expressivity. We show that this is not the case. Building on recent work in NFs, we propose a single NF architecture which integrates seamlessly into RL algorithms, serving as a policy, Q-function, and occupancy measure. Our approach leads to much simpler algorithms, and achieves higher performance in imitation learning, offline, goal conditioned RL and unsupervised RL.