EXPO: Stable Reinforcement Learning with Expressive Policies

Perry Dong, Qiyang Li, Dorsa Sadigh, Chelsea Finn
Stanford University and UC Berkeley
Paper ArXiv Code

Abstract

We study the problem of training and fine-tuning expressive policies with online reinforcement learning (RL) given an offline dataset. Training expressive policy classes with online RL present a unique challenge of stable value maximization. Unlike simpler Gaussian policies commonly used in online RL, expressive policies like diffusion and flow-matching policies are parameterized by a long denoising chain, which hinders stable gradient propagation from actions to policy parameters when optimizing against some value function. Our key insight is that we can address stable value maximization by avoiding direct optimization over value with the expressive policy and instead construct an on-the-fly RL policy to maximize Q-value. We propose Expressive Policy Optimization (EXPO), a sample-efficient online RL algorithm that utilizes an on-the-fly policy to maximize value with two parameterized policies -- a larger expressive base policy trained with a stable imitation learning objective and a light-weight Gaussian edit policy that edits the actions sampled from the base policy toward a higher value distribution. The on-the-fly policy optimizes the actions from the base policy with the learned edit policy and chooses the value maximizing action from the base and edited actions for both sampling and temporal-difference (TD) backup. Our approach yields up to 2-3x improvement in sample efficiency on average over prior methods both in the setting of fine-tuning a pretrained policy given offline data and in leveraging offline data to train online.

Teaser image

Expressive Policy Optimization (EXPO) is a stable, sample efficient method for training expressive policies with reinforcement learning by avoiding direct optimization over the value function with the expressive policy. Instead, EXPO trains the base expressive policy using a stable supervised learning objective and an edit policy with the reinforcement learning objective and construct an on-the-fly policy to maximize value through (1) editing the action samples from the base expressive policy to higher value distributions, and (2) non-parametric post-processing that takes multiple action candidates from the base and edit policy and selects the highest-value action among base and edited actions.

Intuition of Action Edits

Action image


The base expressive policy by itself cannot gaurantee the sampled actions maximize Q-value. The edit policy transforms actions of the base policy into actions that further maximize Q-value while encouraging exploration. The blue contour represents the Q-values of actions of a single state and the orange contours represent the Gaussian distributions of actions the edit policy changes the base actions into.

Experiment Results

Online RL (No Pre-training). Online RL results on 12 challenging sparse-reward tasks. Across almost every task, EXPO consistently exceeds or matches the performance of the best baseline—even without any pretraining.

Performance of Different Algorithms in Online RL with an Offline Dataset


Offline-to-Online RL (Offline Pre-training Followed by Online Fine-tuning). Offline-to-online RL results on 12 challenging sparse-reward tasks. EXPO consistently exceeds or matches the performance of the best baseline. The relative benefit of EXPO over baselines is especially large on the manipulation tasks, where prior methods often struggle to improve in performance. Importantly, EXPO does not drop in performance going from pre-training to fine-tuning.

Performance of Different Algorithms in Offline-to-Online RL

BibTeX


@article{dong2025expo,
    title   = {EXPO: Stable Reinforcement Learning with Expressive Policies},
    author  = {Perry Dong and Qiyang Li and Dorsa Sadigh and Chelsea Finn},
    journal = {arXiv},
    year    = {2025},
}