PaperlayerView exact revision on arXiv

Source-linked AI summary

Efficient Diffusion Policies for Offline Reinforcement Learning

Bingyi Kang, Xiao Ma, Chao Du, Tianyu Pang, Shuicheng Yan

arXiv:2305.20081v2cs.LGcs.AI

TL;DR

Offline RL needs expressive policy parameterizations, but Diffusion-QL is expensive to train and incompatible with likelihood-based algorithms. EDP constructs actions from corrupted dataset actions to avoid the sampling chain, supports multiple offline RL algorithms, and achieves strong D4RL results while reducing gym-locomotion training from five days to five hours.

  • Problem

    Diffusion policies require costly forward and backward passes through long Markov chains and cannot directly support likelihood-based offline RL algorithms because their likelihood is intractable.

  • Method

    EDP uses action approximation to construct actions from corrupted dataset actions and extends diffusion-policy training to direct and likelihood-based offline RL algorithms.

  • Results

    EDP reduces gym-locomotion diffusion-policy training time from five days to five hours and sets new state-of-the-art results across all four D4RL domains.

  • Takeaways & Limitations

    EDP makes diffusion policies faster and compatible with TD3, CRR, and IQL across D4RL domains.

  • Takeaways & Limitations

    The paper notes that EDP's generalization ability could enable harmful strategic agents in weapon or war robots, with possible domestic privacy and safety risks.

Abstract

from arXiv · show

Offline reinforcement learning (RL) aims to learn optimal policies from offline datasets, where the parameterization of policies is crucial but often overlooked. Recently, Diffsuion-QL significantly boosts the performance of offline RL by representing a policy with a diffusion model, whose success relies on a parametrized Markov Chain with hundreds of steps for sampling. However, Diffusion-QL suffers from two critical limitations. 1) It is computationally inefficient to forward and backward through the whole Markov chain during training. 2) It is incompatible with maximum likelihood-based RL algorithms (e.g., policy gradient methods) as the likelihood of diffusion models is intractable. Therefore, we propose efficient diffusion policy (EDP) to overcome these two challenges. EDP approximately constructs actions from corrupted ones at training to avoid running the sampling chain. We conduct extensive experiments on the D4RL benchmark. The results show that EDP can reduce the diffusion policy training time from 5 days to 5 hours on gym-locomotion tasks. Moreover, we show that EDP is compatible with various offline RL algorithms (TD3, CRR, and IQL) and achieves new state-of-the-art on D4RL by large margins over previous methods. Our code is available at https://github.com/sail-sg/edp.

1 Introduction

Offline RL policies must model multimodal action data, but Diffusion-QL incurs costly long-chain training and cannot support likelihood-based algorithms. EDP uses action approximation to improve efficiency and generality, reducing training time from five days to five hours and reaching new D4RL state-of-the-art.

  • Offline RL policy parameterization is important because collected data can exhibit strong multimodality.
  • Diffusion-QL improves TD3+BC performance by replacing diagonal Gaussian policies with diffusion models.
  • Diffusion-QL is computationally inefficient because training forwards and backpropagates through a long Markov chain, and its likelihood is intractable for likelihood-based RL.
  • EDP approximates clean actions from corrupted actions using DDPM noise prediction and the reparametrization trick.
  • Five days to five hours: replacing Diffusion-QL with EDP substantially reduces gym-locomotion training time.
  • EDP works with TD3+BC, CRR, and IQL and establishes new state-of-the-art results across all four D4RL domains.

2 Related Work

Offline RL policy classes must support either differentiable sampling or tractable likelihoods, while standard Gaussian policies may inadequately represent complex data. EDP targets both requirements by making diffusion policies efficient and compatible with maximum-likelihood-based algorithms.

  • Offline RL primarily addresses distributional shift through policy constraints, conservative Q-learning, or combinations of these strategies.
  • Policy parameterizations differ in whether sampling is differentiable or sample log-likelihood is tractable.
  • Diagonal Gaussian policies are common because they satisfy these requirements, while CVAEs and diffusion models provide alternative policy representations.
  • EDP aims to reduce diffusion-policy training time to Gaussian-policy levels and support maximum-likelihood-based offline RL methods.

3 Preliminaries

Offline RL learns policies from fixed transition datasets using actor-critic policy evaluation and improvement, while diffusion models provide a Markov-chain framework for generating actions by reversing progressive Gaussian corruption.

  • An MDP specifies state and action spaces, transition probabilities, rewards, and a discount factor for sequential decision-making.
  • Value and action-value functions represent expected discounted returns from states and state-action pairs.
  • Offline RL learns from a previously collected transition dataset rather than environment interaction.
  • Continuous-control offline RL alternates Q-network policy evaluation with policy improvement while addressing distributional shift.
  • Diffusion forward processes add Gaussian noise over K Markov-chain steps, and learned reverse transitions generate samples from noise.
  • The policy may be deterministic or stochastic; this paper uses the stochastic form without loss of generality.

4 Efficient Diffusion Policy

EDP trains diffusion policies efficiently by constructing approximate actions from corrupted dataset actions and accelerating evaluation sampling, then extends diffusion policies to direct and likelihood-based offline RL algorithms.

  • 4 Efficient Diffusion Policy: EDP uses a conditional DDPM reverse process as a parametric policy and develops Reinforcement-Guided Diffusion Policy Learning.
  • 4.1 Diffusion Policy: DDPM action generation normally samples sequentially from K Gaussian distributions, requiring a long reverse diffusion process.
  • 4.2 Reinforcement-Guided Diffusion Policy Learning: Action approximation corrupts a dataset action to a_k and applies one-step denoising to construct an approximate clean action without reverse-chain sampling.
  • 4.3 Generalization to Various RL algorithms: EDP supports direct policy optimization because both expensive DDPM sampling and efficient approximation can be reparameterized for gradient-based updates.
  • 4.3 Generalization to Various RL algorithms: EDP supports likelihood-based methods by replacing intractable diffusion likelihoods with a DDPM lower bound or an approximate Gaussian policy.
  • 4.3 Generalization to Various RL algorithms: The approximate Gaussian realization is easier to implement and is used mainly for CRR and IQL experiments.
  • 4.4 Comparison to Diffusion-QL: EDP improves diffusion-policy computational efficiency and generality relative to Diffusion-QL.
  • 4.4 Comparison to Diffusion-QL: DPM-Solver replaces DDPM sampling for policy evaluation, reducing sampling steps from 1000 to 15.

5 Experiments

Experiments evaluate EDP's efficiency, performance, generality, evaluation metrics, and ablations across D4RL domains. EDP accelerates diffusion-policy training while remaining competitive or superior across offline RL algorithms and tasks.

  • Experimental setup: EDP experiments cover D4RL domains including Gym-locomotion, AntMaze, Adroit, and Kitchen, using shared network and optimization settings.The evaluation includes extensive baselines and reports both OMS and RAT metrics.
  • Efficiency and reproducibility: 2.3x training and 3.3x sampling speedups come from action approximation, while DPM-Solver adds another 2.3x training speedup over DQL (JAX).The comparison uses 10,000 policy updates for IPS and 10,000 environment transitions for SPS on walker2d-medium-expert-v2.
  • Efficiency and reproducibility: EDP and DQL (JAX) are comparable with Diffusion-QL on Gym-locomotion and substantially better on the other three D4RL domains.EDP enables training with K = 1000, whereas Diffusion-QL uses 5∼100 and cannot efficiently forward and backward through larger K.
  • Generality and overall results: EDP beats feed-forward Gaussian policies across all tested domains and algorithms, with EDP + TD3 reaching an 85.5 average score on Gym-locomotion.EDP + IQL exceeds competing methods by more than 10 average scores on Kitchen, Adroit, and AntMaze.
  • Ablation study: OMS and RAT remain close under stable training but diverge when AntMaze training succeeds and then crashes, making OMS unreliable for evaluation.Figure 3 compares RAT as the average and OMS as the best score across representative environments.
  • Ablation study: Normalized performance increases with the number of EAS actions on 8 of 9 locomotion tasks, while the main experiments use 10 actions as a performance-efficiency trade-off.The EAS sweep varies the number of actions from 1 to 200.

6 Conclusion

EDP makes diffusion policies faster and more general for offline RL while retaining their expressive policy representation. It achieves a 25× training-time speedup, supports TD3, CRR, and IQL, and reaches new state-of-the-art performance across all four D4RL domains.

  • 25× training-time speedup over Diffusion-QL is achieved on D4RL gym-locomotion tasks.EDP avoids running the Markov chain for action sampling during training by constructing actions from corrupted ones through approximation.
  • EDP supports offline RL algorithms including TD3, CRR, and IQL.The paper evaluates EDP with multiple offline RL algorithms to establish its generality as a policy class.
  • EDP sets new state-of-the-art results on all four D4RL domains.

Broader Impact

The section describes how EDP is integrated with TD3, CRR, and IQL, including action approximation and diffusion-based policy improvement. It also notes that EDP may be applied in strategic-agent settings with potential privacy and safety risks.

  • Broader Impact: The paper notes that EDP could be used by strategic agents in war or weapon robots, with application-dependent risks to domestic privacy and safety.The passage frames these as potential harms rather than established outcomes.
  • TD3 and TD3+BC: TD3 uses double Q-learning with two value networks, while TD3 policy improvement uses only Q1.The target uses the minimum of Q1 and Q2, with the next action supplied by the deterministic policy.
  • TD3 and TD3+BC: EDP uses diffusion loss for behavior cloning and action approximation instead of sampled actions during TD3+BC policy improvement.It samples from Q1 and Q2 with equal probability for each policy improvement step.
  • CRR: CRR follows AWR-style constrained policy improvement and replaces the sampling distribution with a fixed dataset.This replacement naturally imposes an implicit constraint on policy improvement.
  • CRR: Because diffusion-model log-likelihood is intractable, EDP uses Eqn. 13 in practice and estimates advantages with sampled actions around the approximated action.The implementation uses σ = I, β = 1, and sample size N = 10 as generally effective settings.
  • IQL: IQL stays in-distribution by using dataset actions for advantage computation while retaining multi-step dynamic programming during policy evaluation.Its expectile regression approximates the Bellman optimality equation without querying out-of-distribution actions.

C.2 More Results

This section reports detailed EDP results across evaluation metrics and finds substantial differences between OMS and RAT, with IQL and CRR more stable than TD3. On antmaze, IQL has a much smaller gap between its best and average scores than TD3.

  • Evaluation metrics: OMS and RAT show huge score gaps across all domains and algorithms.The comparison uses Tab. 5 scores for both evaluation metrics.
  • Algorithm comparison: IQL and CRR are relatively more stable than TD3 across the reported evaluations.The passage illustrates this stability difference using antmaze scores.
  • Antmaze results: 80.5 best versus 29.8 average is reported for TD3 on antmaze, compared with 89.2 best versus 73.4 average for IQL.These figures show the reported best and average scores for the two algorithms on antmaze.

C.3 More Results on EAS

The section introduces a comparison among normal TD3+BC, TD3+BC with EAS during evaluation, and TD3 + EDP.

  • Compared methods: The comparison includes normal TD3+BC, TD3+BC with EAS for evaluation, and TD3 + EDP.The passage identifies the compared methods but does not report their outcomes.
  • Compared methods: EAS is evaluated as an addition to TD3+BC rather than as a standalone comparison condition.The reported setup specifically contrasts TD3+BC with and without EAS during evaluation.
  • Compared methods: The section is designed to compare evaluation-time EAS against the TD3 + EDP approach.No quantitative ranking or performance difference is stated in the supplied passage.

C.4 More Experiments on Controlled Sampling

Controlled sampling experiments find that policy scaling does not improve performance, while reducing initial sampling noise generally helps but remains inferior to EAS.

  • Policy scaling: Policy scaling achieves its best performance at τ = 1.0, indicating that sampling from a scaled policy does not work.Experiments vary τ from 0.5 to 2.0 on D4RL gym-locomotion tasks.
  • Deterministic Sampling: Reducing initial noise improves diffusion-policy performance in most cases, with zero noise performing best and normal noise performing worst.The initial noise scale is varied from 0.0 to 1.0.
  • Deterministic Sampling: The best performance from reduced initial noise still falls behind EAS.This comparison limits the benefit of deterministic sampling relative to EAS.

C.5 Computational Cost Comparison between EDP and Feed-Forward Policy networks

EDP is slower than a feed-forward policy during direct TD3+BC training but avoids repeated diffusion-network passes during action approximation, enabling larger diffusion chains that can recover performance.

  • Training Speed: 689 IPS for TD3+BC versus 412 IPS for EDP makes training take around 3 hours versus 5 hours, respectively.The comparison uses walker2d-medium-expert-v2 with 10,000 policy-update iterations.
  • Action Approximation: Action approximation forwards and backwards through the policy network once during training, whereas DDPM does so 100 times.Using the same number of diffusion steps can slightly harm performance, but approximation supports larger K.
Loading 2305.20081v2…