Dichotomous Diffusion Policy Optimization

TL;DR

DIPOLE tackles the instability of training diffusion-based policies with RL by introducing a greedified KL-regularized objective that decomposes the optimal policy into two dichotomous components: a positive policy aimed at reward maximization and a negative policy aimed at minimization. The optimal policy can be sampled by linearly combining the scores of these two dichotomous policies, controlled by a greediness parameter $\omega$, and implemented via diffusion-model CFG-like guidance. This decomposition yields stable training with bounded, sigmoid-based weights and enables leveraging both high- and low-return samples for learning. Empirical results on ExORL and OGBench show strong offline and offline-to-online performance, and a large VLA autonomous driving model trained with DIPOLE demonstrates scalable improvements on NAVSIM, highlighting practical potential for complex real-world decision-making.

Abstract

Diffusion-based policies have gained growing popularity in solving a wide range of decision-making tasks due to their superior expressiveness and controllable generation during inference. However, effectively training large diffusion policies using reinforcement learning (RL) remains challenging. Existing methods either suffer from unstable training due to directly maximizing value objectives, or face computational issues due to relying on crude Gaussian likelihood approximation, which requires a large amount of sufficiently small denoising steps. In this work, we propose DIPOLE (Dichotomous diffusion Policy improvement), a novel RL algorithm designed for stable and controllable diffusion policy optimization. We begin by revisiting the KL-regularized objective in RL, which offers a desirable weighted regression objective for diffusion policy extraction, but often struggles to balance greediness and stability. We then formulate a greedified policy regularization scheme, which naturally enables decomposing the optimal policy into a pair of stably learned dichotomous policies: one aims at reward maximization, and the other focuses on reward minimization. Under such a design, optimized actions can be generated by linearly combining the scores of dichotomous policies during inference, thereby enabling flexible control over the level of greediness.Evaluations in offline and offline-to-online RL settings on ExORL and OGBench demonstrate the effectiveness of our approach. We also use DIPOLE to train a large vision-language-action (VLA) model for end-to-end autonomous driving (AD) and evaluate it on the large-scale real-world AD benchmark NAVSIM, highlighting its potential for complex real-world applications.

Figures (8)

• Figure 1: Illustration of the policy weighting scheme in DIPOLE. Based on our greedified policy optimization objective, the regression weight of the optimal policy can be decomposed into a pair of dichotomous terms, and the greediness for reward/value maximization can be flexibly controlled by $\omega$.
• Figure 2: NAVSIM Results: DP-VLA w/ DIPOLE fine-tuned model trajectory; ground truth ego trajectory; DP-VLA imitation pretrained model trajectory.
• Figure 3: ExORL environments. We experiment on 4 high-dimensional complex domains: Walker, Cheetah, Quadruped, and Jaco Arm.
• Figure 4: OGBench environments. We experiment on 5 complex domains: antsoccer-arena, humanoidmaze, scene, cube-single, and cube-double.
• Figure 5: Top-left: ablation on $\beta$ and $k$; Top-right: ablation on expectile $\tau$; Bottom: ablation on $w$ and $\beta$

Theorems & Definitions (4)

• Lemma 1
• Theorem 1
• Theorem 1
• proof