Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Boosting Continuous Control with Consistency Policy

Yuhui Chen, Haoran Li, Dongbin Zhao

TL;DR

The paper tackles offline RL's value overestimation and the slow, multi-step nature of diffusion-model policies. It introduces CPQL, a one-step, consistency-based policy learning approach that maps probability-flow ODE trajectories to actions, avoiding inaccurate intermediate Q-value guidance. The authors provide theoretical guarantees of policy improvement with accurate guidance and extend the method to online RL. Empirically, CPQL achieves state-of-the-art performance across 11 offline and 21 online tasks and delivers substantial speedups (approximately 15x training and 45x inference) over diffusion-based baselines, highlighting its potential for real-time control applications.

Abstract

Due to its training stability and strong expression, the diffusion model has attracted considerable attention in offline reinforcement learning. However, several challenges have also come with it: 1) The demand for a large number of diffusion steps makes the diffusion-model-based methods time inefficient and limits their applications in real-time control; 2) How to achieve policy improvement with accurate guidance for diffusion model-based policy is still an open problem. Inspired by the consistency model, we propose a novel time-efficiency method named Consistency Policy with Q-Learning (CPQL), which derives action from noise by a single step. By establishing a mapping from the reverse diffusion trajectories to the desired policy, we simultaneously address the issues of time efficiency and inaccurate guidance when updating diffusion model-based policy with the learned Q-function. We demonstrate that CPQL can achieve policy improvement with accurate guidance for offline reinforcement learning, and can be seamlessly extended for online RL tasks. Experimental results indicate that CPQL achieves new state-of-the-art performance on 11 offline and 21 online tasks, significantly improving inference speed by nearly 45 times compared to Diffusion-QL. We will release our code later.

Boosting Continuous Control with Consistency Policy

TL;DR

The paper tackles offline RL's value overestimation and the slow, multi-step nature of diffusion-model policies. It introduces CPQL, a one-step, consistency-based policy learning approach that maps probability-flow ODE trajectories to actions, avoiding inaccurate intermediate Q-value guidance. The authors provide theoretical guarantees of policy improvement with accurate guidance and extend the method to online RL. Empirically, CPQL achieves state-of-the-art performance across 11 offline and 21 online tasks and delivers substantial speedups (approximately 15x training and 45x inference) over diffusion-based baselines, highlighting its potential for real-time control applications.

Abstract

Due to its training stability and strong expression, the diffusion model has attracted considerable attention in offline reinforcement learning. However, several challenges have also come with it: 1) The demand for a large number of diffusion steps makes the diffusion-model-based methods time inefficient and limits their applications in real-time control; 2) How to achieve policy improvement with accurate guidance for diffusion model-based policy is still an open problem. Inspired by the consistency model, we propose a novel time-efficiency method named Consistency Policy with Q-Learning (CPQL), which derives action from noise by a single step. By establishing a mapping from the reverse diffusion trajectories to the desired policy, we simultaneously address the issues of time efficiency and inaccurate guidance when updating diffusion model-based policy with the learned Q-function. We demonstrate that CPQL can achieve policy improvement with accurate guidance for offline reinforcement learning, and can be seamlessly extended for online RL tasks. Experimental results indicate that CPQL achieves new state-of-the-art performance on 11 offline and 21 online tasks, significantly improving inference speed by nearly 45 times compared to Diffusion-QL. We will release our code later.
Paper Structure (36 sections, 2 theorems, 22 equations, 9 figures, 6 tables, 2 algorithms)

This paper contains 36 sections, 2 theorems, 22 equations, 9 figures, 6 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Offline RL
  4. Diffusion Models for Imitation Learning
  5. Diffusion Models for RL
  6. Diffusion policy for offline RL
  7. Offline RL
  8. Policy Constraint for Offline Learning
  9. Policy Improvement with Constraint Policy Search
  10. Diffusion Policy
  11. Consistency Policy with Q-Learning
  12. Consistency Policy
  13. Training Loss for Consistency Policy
  14. Policy Evaluation
  15. Extension for Online RL
  16. ...and 21 more sections

Key Result

Theorem 1

Let $\Delta k=\max_{m\in [1, M-1]}\{|k_{m+1}-k_m|\}$. We have the assumptions: 1) Distance function $d$, value function $Q$ and $f_{\theta^{-}}$ are all twice continuously differentiable with bounded second derivatives; 2) There is a pre-trained score function representing the desired policy: $\fora where $\hat{a}_{\phi^*}^{k_m}$ is calculated with Euler solver and the optimal score function $s_{\

Figures (9)

  • Figure 1: Given an ODE that smoothly converts from actions of the reference policy (e.g., $x^0, y^0\in A$) to Gaussian noises, the consistency policy $f_{\theta}$ maps any point (e.g., $x^{k}, x^{k'}, y^{k"}, y^K$) on the PF ODE trajectory to the desired actions (e.g., $\hat{x}^{\epsilon}, \hat{y}^{\epsilon}$). Since consistency policy generates the actions from the noise by one step, it reduces an enormous amount of time for policy training and inference.
  • Figure 2: Training curves for 8 online tasks, including 6 tasks from dm_control (top row, and leftmost two of the bottom row) and 2 tasks from Gym Mujoco (rightmost two of the bottom row). CPQL, PPO, SAC and TD3 are compared on each task with 5 random seeds.
  • Figure 3: Training and inference speedup on D4RL locomotion tasks. We choose the Diffusion-QL (Pytorch) as the baseline with all original data provided in Appendix \ref{['apx:more d4rl']}. Implementation for Diffusion-QL(Pytorch), DIPO(Pytorch), EDP(Jax) are from the official repository, with websites listed in Appendix \ref{['apx: exp setup']}. And we implement Diffusion-QL(Jax) by ourselves for comparison.
  • Figure 4: The effect of the different hyperparameters $\eta$ and $\alpha$ on the online and offline policy learning.
  • Figure 5: The effect of the different training losses on policy learning.
  • ...and 4 more figures

Theorems & Definitions (2)

  • Theorem 1
  • Theorem 2