Source-linked AI summary
Consistency Models as a Rich and Efficient Policy Class for Reinforcement Learning
Zihan Ding, Chi Jin
TL;DR
Diffusion policies can represent multi-modal actions, but iterative denoising makes inference slow for reinforcement learning. This paper applies consistency models as expressive, efficient policies within actor-critic methods across offline, offline-to-online, and online RL. The consistency policy achieves comparable or sometimes higher performance with significant speedups across these settings.
Problem
Diffusion policies represent multi-modal action distributions, but their many denoising steps slow inference and make online RL policy-gradient optimization costly in time and memory.
Source: §1 · p. 1Method
The paper uses consistency models as RL policies in behavior-cloning and actor-critic algorithms, comparing them with diffusion policies across offline, offline-to-online, and online settings.
Source: §1 · pp. 1–2Results
Takeaways & Limitations
Consistency policies balance computational efficiency with multi-modal distribution modeling and offer a faster policy alternative to diffusion policies in the evaluated RL settings.
Source: §6 · p. 9
Abstract
from arXiv · showhide
Score-based generative models like the diffusion model have been testified to be effective in modeling multi-modal data from image generation to reinforcement learning (RL). However, the inference process of diffusion model can be slow, which hinders its usage in RL with iterative sampling. We propose to apply the consistency model as an efficient yet expressive policy representation, namely consistency policy, with an actor-critic style algorithm for three typical RL settings: offline, offline-to-online and online. For offline RL, we demonstrate the expressiveness of generative models as policies from multi-modal data. For offline-to-online RL, the consistency policy is shown to be more computational efficient than diffusion policy, with a comparable performance. For online RL, the consistency policy demonstrates significant speedup and even higher average performances than the diffusion policy.
1 INTRODUCTION
The paper introduces consistency models as expressive, efficient policy representations for deep reinforcement learning. It develops consistency-policy algorithms and evaluates them alongside diffusion policies across offline, offline-to-online, and online settings.
1 INTRODUCTION
- The paper presents multi-modal consistency and diffusion policies as expressive alternatives to unimodal Gaussian policies for demonstrations drawn from mixtures of behavior policies.It also proposes loss scaling and policy regularization for offline consistency-policy training.Source: §1 · pp. 1–2
- The paper adapts consistency models as expressive, efficient policy representations and embeds them in Consistency-BC and Consistency-AC.The evaluation compares consistency and diffusion policies in offline, offline-to-online, and online RL on D4RL.Source: §1 · pp. 1–2
- Consistency policies achieve comparable or sometimes higher performance than diffusion policies in offline-to-online and online RL, with shorter training and inference time.For offline behavior cloning, consistency sampling reduces training time by 43%.Source: §1 · p. 2
2 RELATED WORKS
Related work covers approaches to offline-RL value overestimation, offline-to-online fine-tuning, and the use of score-based generative models as RL policies. These strands include policy constraints, pessimistic value estimation, and diffusion-based policy or planning methods.
Offline and Offline-to-Online RL
- Offline-RL methods address value overestimation through policy constraints, such as BCQ (Fujimoto et al., 2019), or pessimistic value estimation, such as CQL (Kumar et al., 2020); MoRel constructs pessimistic MDPs (Kidambi et al., 2020).The passage also lists batch-constraining and divergence-constraint methods, alongside other pessimistic estimators.Source: §2 · p. 2
- Prior offline-to-online work studies fine-tuning with offline data or pretrained policies, including Hybrid Q learning (Song et al., 2022), RLPD (Ball et al., 2023), and Cal-QL (Nakamoto et al., 2023).The passage identifies distribution shift as a cause of performance degradation early in online training.Source: §2 · p. 2
Score-based Generative Model for RL
- Diffusion-based RL work uses generative models for policy representation, imitation learning, model-based RL, and guided sampling, including Diffusion QL (Wang et al., 2022) and Diffuser (Janner et al., 2022).Other cited examples include IDQL, diffusion policies for robotics, Decision Diffuser, and QGPO.Source: §2 · p. 2
3 PRELIMINARIES
The preliminaries define offline and online RL in an MDP framework, then describe consistency models as mappings from noisy samples to near-clean samples. Unlike stepwise reverse diffusion, this mapping supports sampling with fewer steps.
3.1 OFFLINE AND ONLINE RL
- Offline RL optimizes a policy from a fixed dataset, whereas online RL collects environment interaction data to optimize its policy.The setup defines policies over actions given states and frames optimization around discounted cumulative reward.Source: §3.1 · p. 3
3.2 CONSISTENCY MODEL
- A consistency model follows a probability-flow ODE and generates data by reversing its trajectory from an initial random sample.The diffusion model instead solves distribution matching with an SDE.Source: §3.2 · p. 3
- The consistency function maps a noisy sample at time τ to the original near-clean sample, replacing step-by-step reverse diffusion with fewer sampling steps.For conditional distributions, the function also takes the condition variable as input.Source: §3.2 · p. 3
4 CONSISTENCY MODEL AS RL POLICY
The consistency model is adapted as a state-conditioned action policy, trained through behavior cloning or an actor-critic method. The section describes iterative or single-step action inference and a training-time loss weighting for offline RL.
4 CONSISTENCY MODEL AS RL POLICY
- Consistency inference conditions on state and iteratively denoises noisy actions, with a single-step version obtained by using only the endpoint time points.The consistency function maps state to action according to the offline dataset's conditional action distribution.Source: §4 · p. 3
- Consistency behavior cloning trains the conditional consistency policy on the offline dataset using a step-weighted consistency loss and an exponential-moving-average target.The target stabilizes estimation during training.Source: §4 · p. 3
- Consistency-AC learns a double-Q critic and updates the policy by maximizing expected Q-values with consistency-loss behavior-cloning regularization.Action inference uses N denoising steps, and policy gradients backpropagate through the consistency model N times, adding computation.Source: §4 · p. 4
- A k-dependent weight scales the consistency loss by the gap between adjacent time points, balancing it against the Q-policy objective during offline training.The authors report that this weighting improves consistency-based policy regularization on offline RL tasks.Source: §4 · p. 4
5 EXPERIMENTAL EVALUATION
Across four D4RL task suites and offline, offline-to-online, and online settings, the evaluation compares consistency and diffusion policies in behavior-cloning and actor-critic methods. Consistency policies reduce computation while retaining competitive performance, with outcomes varying by setting.
5 EXPERIMENTAL EVALUATION
- The evaluation spans D4RL Gym, AntMaze, Adroit, and Kitchen suites, comparing behavior-cloning and actor-critic generative policies across offline, offline-to-online, and online RL.Consistency policies generally use N = 2 denoising steps, while diffusion policies use N = 5; the study compares performance and computational time, including ablations of Consistency-AC.Source: §5 · pp. 4–5
5.1 OFFLINE RL: BEHAVIOR CLONING WITH EXPRESSIVE POLICY REPRESENTATION
- 42.97% less computational time on average across 20 tasks for Consistency-BC, which performs comparably to many offline RL methods using behavior cloning alone.Consistency-BC outperforms Gaussian-policy BC on 14/20 tasks; Diffusion-BC is better than or equivalent to Gaussian-policy BC on 20/20 tasks. Table 1 covers D4RL Gym, AntMaze, Adroit, and Kitchen tasks; Figure 1 compares average training time per epoch.Source: §5.1 · pp. 5–6, Table 1 · p. 5, Figure 1 · p. 6
5.2 OFFLINE RL: CONSISTENCY ACTOR-CRITIC
- 72.1 for Consistency-AC versus 80.3 for Diffusion-QL in average normalized score; Consistency-AC outperforms other baselines on most Gym and Adroit tasks but struggles on AntMaze.Consistency-AC uses BC policy regularization to avoid generating out-of-distribution actions; the authors suggest sparse rewards may make Q-learning, rather than multimodal modeling, the AntMaze bottleneck.Source: Table 2 · p. 6, §5.2 · p. 7
- Consistency-AC saturates at N = 2, while Diffusion-QL saturates at N = 5; at those defaults, Diffusion-QL uses about 1.786× more training time while yielding 1.057× better performance.The comparison is on hopper-medium-expert-v2; Figure 2 shows normalized scores and training time versus N, and Table 3 compares computational time across N values.Source: §5.2 · pp. 7–8, Figure 2 · p. 7, Table 3 · p. 7
- Without loss scaling, Consistency-AC performance degrades by 37.8% on average for most tasks, though it can improve some tasks such as AntMaze.LN-Resnet does not consistently improve over MLP across tasks, but benefits Adroit; for most tasks except AntMaze, Consistency-BC does not match the best performances, while a properly scaled Q-learning loss helps improve scores.Source: §5.2 · p. 8, Figure 3 · p. 8
5.3 OFFLINE-TO-ONLINE AND ONLINE RL
- Offline-to-online scores are slightly lower for Consistency-AC than Diffusion-QL, while online-from-scratch scores are higher; consistency policies finish training sooner and show more efficient learning on most tasks.Offline-to-online methods are usually more sample-efficient than online methods, except on three hopper tasks; online consistency policies have sharper score-increasing slopes on 8/9 tasks. Online results use fixed training steps and a learning rate chosen for fine-tuning, and the authors report that these models have not yet converged to optimal performance.Source: Table 4 · p. 8, §5.3 · p. 9, Figure 4 · p. 9
6 CONCLUSION
The consistency policy balances computational efficiency with modeling accuracy for multimodal offline data and achieves comparable performance with diffusion policies while significantly speeding up across three RL settings.
6 CONCLUSION
- The consistency policy balances computational efficiency with modeling accuracy for multimodal distributions in offline RL and achieves comparable performance with significant speedup across three RL settings.The conclusion characterizes Consistency-AC as using policy regularization and notes that it is orthogonal to other offline RL techniques.Source: §6 · p. 9
- Future work could combine consistency policy with conservative Q-value estimation, improve offline-to-online alignment and initialization, or develop advanced online exploration methods.The conclusion also proposes scaling task complexity to test whether consistency policies retain expressiveness while reducing computational cost.Source: §6 · p. 9
- Consistency-AC uses a novel policy representation with policy regularization for offline RL, which the authors describe as orthogonal to other offline RL techniques.Source: §6 · p. 9
A D4RL DATASET VISUALIZATION
Figure 5 visualizes randomly selected D4RL state-action samples with t-SNE, coloring them by normalized reward.
A D4RL DATASET VISUALIZATION
- Figure 5 uses t-SNE to visualize 10,000 randomly selected state-action samples, or 3,000 for pen-human-v1 and kitchen-complete-v0, colored by normalized reward from −1 to 1.Source: Figure 5 · p. 13
B CONSISTENCY MODEL TRAINING AND INFERENCE DETAILS
The appendix describes the consistency model’s training loss and inference options, then gives training settings and hyperparameters for the D4RL experiments.
B CONSISTENCY MODEL TRAINING AND INFERENCE DETAILS
- The consistency model is trained to model the data distribution with a loss that uses L2 distance between outputs.The passage specifies d(x, y) = ∥x − y∥₂.Source: B CONSISTENCY MODEL TRAINING AND… · p. 14
- The training noise-level subsequence differs from inference and follows a Karras boundary schedule.Source: B CONSISTENCY MODEL TRAINING AND… · p. 14
- After training, the model generates samples from noisy input using either single-step sampling or iterative multistep sampling.Multistep sampling follows a time sequence; the inference sequence is a linspace from epsilon to T with N−1 sub-intervals.Source: B CONSISTENCY MODEL TRAINING AND… · p. 14
C.1 HYPERPARAMETERS
- Offline Consistency-BC and Consistency-AC training uses batch size 256, cosine-annealed learning rates, and task-dependent epoch counts on D4RL.The experiments set ξ = 100.0 for Consistency-AC loss scaling; listed noise-schedule constants are ϵ = 0.002, T = 80, ρ = 7, s0 = 2, and s1 = 150.Source: §C.1 · p. 14
- Table 5 lists Consistency-AC hyperparameters for offline training, including behavior cloning, across D4RL Gym, AntMaze, Adroit, and Kitchen tasks.Source: Table 5 · p. 14
C.2 COMPUTATIONAL TIME
The section compares training and inference time for consistency and diffusion methods, including scaling coefficients and a per-epoch comparison for behavior-cloning methods.
C.2 COMPUTATIONAL TIME
- The per-epoch comparison accounts for the methods’ different total epoch counts by comparing time consumption per epoch.Source: §C.2 · p. 15
- The Diffusion-BC and Consistency-BC training-time comparison reports seconds per epoch across D4RL Gym, AntMaze, Adroit, and Kitchen tasks.The comparison averages five random seeds and uses the same batch size, iterations per epoch, and network architecture for both methods.Source: Table 6 · p. 15, §C.2 · p. 15
- 2.47 vs. 3.54 for training and 0.515 vs. 0.598 for inference: Consistency-AC has smaller time-scaling coefficients than Diffusion-QL as N increases.The scaling-law comparison is for offline RL on hopper-medium-expert-v2.Source: §C.2 · p. 15
C.3 ABLATION STUDIES
The appendix examines Consistency-AC variants, gives offline-to-online and online training procedures and hyperparameters, and compares computational time with Diffusion-QL.
C.3 ABLATION STUDIES
- Four Consistency-AC variants test behavior cloning, loss scaling, and MLP versus LN-Resnet parameterizations, with scores reported across four D4RL task groups.Figure 6 plots training and inference time against denoising steps on hopper-medium-expert-v2.Source: §C.3 · pp. 15–16, Table 7 · p. 16, Figure 6 · p. 16
D.1 ALGORITHMS
- Offline-to-online Consistency-AC starts from an offline policy and critics, collects transitions into an online dataset, then updates critics and policy.The policy update uses the stated Q-loss; the critic update uses Eq. 3.Source: §D.1 · p. 17
D.2 HYPERPARAMETERS
- Online Consistency-AC follows the same interaction-and-update loop while initializing its policy online; Table 8 lists its offline-to-online and online Gym-task hyperparameters.The loop stores environment transitions and updates Q networks and the policy.Source: §D.2 · p. 17
D.3 COMPUTATIONAL TIME
- Online training-time savings are less pronounced than offline because environment simulation remains a large time component; Table 9 and Figure 7 summarize one-million-step costs.The passage attributes this to simulation following action inference, which is unaffected by improvements to model inference and updates.Source: §D.3 · p. 17, Table 9 · p. 18, Figure 7 · p. 18
D.4 MORE RESULTS
Additional results compare methods in offline-to-online and online RL. They include score comparisons and learning curves under offline and online model selection, shown against both time and environment steps.
D.4 MORE RESULTS
- Table 10 compares maximum-over-epochs, unnormalized scores for methods in offline-to-online and online RL.Scores are reported for both settings.Source: Table 10 · p. 18
- Figures 8 and 9 show offline-model-selection learning curves against time and environment steps, respectively, averaged over five seeds with 95% confidence intervals.The runs use one million environment steps.Source: Figure 8 · p. 19, Figure 9 · p. 19
- Figures 10 and 11 show online-model-selection learning curves against time and environment steps, respectively, averaged over five seeds with 95% confidence intervals.The runs use one million environment steps.Source: Figure 10 · p. 20, Figure 11 · p. 20