Diffusion-DICE: In-Sample Diffusion Guidance for Offline Reinforcement Learning

1 Introduction

• Offline RL learns effective policies solely from offline data, avoiding additional online interactions when exploration with untrained policies is costly or dangerous.The introduction notes robotics and industrial control as examples where sufficient prior data may be available.Source: §1 · p. 1
• DICE can be viewed as transforming the behavior distribution into the optimal policy distribution, motivating diffusion models to perform that transformation during sampling.The proposed score combines the behavior-policy score with the gradient of a guidance term dependent on the optimal distribution ratio.Source: §1 · p. 2
• The guidance term is learned through an in-sample objective, while a diffusion model trained on offline data supplies the behavior-policy score.The authors describe a mathematical transformation that makes learning the guidance term equivalent to solving a convex optimization problem.Source: §1 · p. 2
• Because the optimal policy may be multi-modal, Diffusion-DICE generates candidate actions and uses the value function to select among them.The authors describe this as a way to move beyond local-optimal modes toward the global optimum.Source: §1 · p. 2
• Diffusion-DICE reports outperforming prior state-of-the-art methods across varied datasets, especially those with multi-modal complex behavior distributions.The method uses guide-then-select to choose in-support actions while limiting value-function error exploitation.Source: §1 · p. 2
• Unlike prior diffusion-based offline RL methods that are guide-only or select-only, Diffusion-DICE combines both stages and reports state-of-the-art D4RL performance.The authors say its in-sample guide step generates in-distribution, high-value actions, so selection needs only a few candidates.Source: §1 · p. 2

2 Preliminaries

• Offline RL seeks a policy maximizing expected return using a static dataset, without interaction with the environment.The RL problem is represented as an MDP with state and action spaces, transition dynamics, an initial-state distribution, rewards, and a discount factor.Source: §2 · p. 2

2.1 Distribution Correction Estimation

• DICE optimizes expected return with an f-divergence regularizer between the policy’s stationary distribution and the data distribution.Its dual objective is tractable, and practical updates use dataset samples and estimated Bellman quantities; the learning objectives can be computed in-sample.Source: §2.1 · pp. 2–3
• DICE methods connect Q* and V* to the optimal stationary distribution ratio, but common policy-extraction approaches use unimodal Gaussian policies that limit expressivity.The cited extraction approaches include weighted behavior cloning, information projection, and policy gradient.Source: §2.1 · p. 3

2.2 Diffusion Models in Offline Reinforcement Learning

• Diffusion models learn a score function for reversing a gradual Gaussian-noising process and can model multi-modal distributions.Denoising score matching trains a network to estimate the score; applying it to RL is challenging because the target distribution is the optimal policy.Source: §2.2 · p. 3
• Guide-based methods add a learned guidance term to the behavior-policy score, while select-based methods sample multiple behavior-policy actions and resample using critic values.These approaches rely on critic-computed values for sampled actions.Source: §2.2 · pp. 3–4
• Diffusion approximation errors can produce out-of-distribution actions, and critics may overestimate those actions, contributing to error exploitation.The passage says this can affect the generation of high-quality actions in existing methods.Source: §2.2 · p. 4

3.1 An Optimal Policy Transformation View of DICE

• DICE's optimal stationary distribution ratio also defines a transformation from the behavior policy toward the optimal policy, motivating diffusion models to fit and transform multimodal behavior distributions.The policy relationship follows from a bijection between a policy and its stationary distribution.Source: §3.1 · p. 4

3.2 In-sample Guidance Learning for Accurate Policy Transformation

• In-sample Guidance Learning estimates the transformation's guidance term from offline data using one behavior-policy sample for an unbiased objective that is convex in the guidance network output.The method relies on in-sample actions and value predictions, which the paper says exploit minimal critic error and guide samples toward high-ratio in-support actions.Source: §3.2 · p. 5
• A piecewise f-divergence combines Reverse KL-like behavior for small inputs with Pearson χ²-like behavior for large inputs to stabilize guidance gradients.The proposed divergence keeps the inverse derivative positive and its convex conjugate at most quadratic, addressing instability and non-positive weights.Source: §3.2 · pp. 5–6

3.3 Boost Performance with Optimal Action Selection

• Because the optimal policy can be multimodal, Diffusion-DICE generates a few in-support candidates and selects the highest-value action, aiming to limit out-of-distribution error exploitation.The guide stage steers candidates toward high-value in-support actions before selection.Source: §3.3 · pp. 6–7
• In the 2-D bandit toycase, Diffusion-DICE guides actions toward true high-reward outer-circle regions, while QGPO guidance and IDQL selection exploit overestimated out-of-distribution actions.The behavior data occupy an annular region, and the learned reward is erroneous in out-of-distribution regions.Source: §3.3 · p. 7, Figure 2 · p. 7
• Diffusion-DICE demonstrates superior performance to all baselines on 13 of 15 D4RL tasks, especially the more challenging tasks.The table caption reports average normalized scores at training end with standard deviation across five random seeds.Source: Table 1 · p. 8

4 Experiments

• The evaluation covers D4RL MuJoCo locomotion and AntMaze navigation tasks, alongside a separate comparison of error exploitation.The baselines include Gaussian-policy, DICE-based, and diffusion-based algorithms.Source: §4.1 · p. 8, §4 · p. 8

4.1 D4RL Benchmark Datasets:

• Diffusion-DICE outperforms all baselines, especially on MuJoCo medium, medium-replay, and AntMaze datasets.The authors attribute its substantial gap over O-DICE to the multi-modality of DICE’s optimal policy.Source: §4.1 · pp. 8–9

4.2 Further Experiments on Error Exploitation

• Across 10 episodes with the same critic and behavior policy, guide-then-select achieves better performance with less overestimated Q(s, a) than behavior-policy selection.The authors treat this comparison as evidence of minimal error exploitation; the guide stage uses only in-sample actions for critic training and guidance learning.Source: §4.2 · p. 9, Figure 3 · p. 9

• Offline RL methods address distribution shift through behavior regularization, including divergence penalties, weighted behavior cloning, policy parameterization, value-learning regularization, uncertainty, or distance.Source: §5 · p. 9
• Generative policies, including diffusion models, can represent multi-modal policies, but few prior methods consider whether generation exploits model errors.Earlier actor-critic approaches commonly use unimodal Gaussian policies, whose limited ability to model multi-modal distributions can lead to suboptimal performance.Source: §5 · p. 9
• This work extends DICE’s stationary-distribution ratio into a behavior-to-optimal-policy transformation and uses diffusion models instead of Gaussian-based policy extraction.Source: §5 · p. 9

• Diffusion-DICE combines guide-then-select with diffusion-based policy extraction, reporting stronger performance than prior state-of-the-art methods across datasets, especially with multi-modal complex behavior distributions.The authors identify costly, slow diffusion sampling and possible training difficulties in low-data regimes as limitations.Source: §6 · p. 10

• DICE formulates regularized reinforcement learning as a convex program with Bellman-flow constraints and applies Fenchel-Rockafeller or Lagrangian duality.Source: A A More Detailed Discussion of DICE… · p. 14
• A bijection between policies and stationary distributions permits changing the optimization variable from policy π to distribution dπ.The resulting optimization is expressed over d with Bellman-flow constraints.Source: A A More Detailed Discussion of DICE… · p. 14
• The feasible distribution must be nonnegative to correspond to a valid policy; duality then yields an inner optimum using the ratio ω(s, a) = d(s, a)/dD(s, a).Source: A A More Detailed Discussion of DICE… · p. 14
• Because the next-state expectation is generally intractable with limited offline data, DICE methods use Q(s, a) to fit the reward-plus-discounted-value target with an MSE objective.The text says this approximation is used to handle the intractability and increase training stability.Source: A A More Detailed Discussion of DICE… · p. 14
• Diffusion-policy methods may model the behavior or optimal policy, while trajectory-level methods generate high-return trajectories and take the action for the current state.The trajectory approach requires additional effort to generate high-quality trajectories.Source: A A More Detailed Discussion of DICE… · pp. 14–15

B Additional Proofs

• The proof derives the optimal-policy score relation from the behavior-policy distribution and the stationary distribution ratio.It uses the forward diffusion transition’s independence from the state to derive the relationship between score functions.Source: B Additional Proofs · pp. 15–16
• For a non-negative function f, the scalar optimization in Lemma 1 is convex and has optimizer log E[f(X)].The proof establishes convexity through a non-negative second derivative and obtains the optimizer from the first-order condition.Source: B Additional Proofs · p. 16
• The guidance-learning objective is convex when the distribution ratio is non-negative, with a convergence guarantee under unlimited model capacity.The proof checks the second derivative with respect to gθ(s, at, t).Source: B Additional Proofs · p. 17
• The appendix verifies that the proposed piecewise function is a well-defined f-divergence and derives its conjugate by taking the larger supremum across its two intervals.The proof checks convexity and boundary properties, then treats [0, 1) and [1, +∞) separately.Source: B Additional Proofs · pp. 17–18

C In-depth Discussion of In-sample Guidance Learning and Other Guidance Methods

• IGL can induce the exact optimal policy when that policy is representable as a weighted behavior policy, unlike the cited inexact guidance methods [17] [6].Those methods use MSE guidance or training-free guidance based on a pretrained diffusion model.Source: C In-depth Discussion of In-sample… · p. 18
• CEP can induce exactly the desired policy if representable as a weighted behavior policy, yet remains prone to critic overestimation, which the passage attributes to CEP’s inherent contrastive-loss property.CEP fits gθ(s, at, t) to log Ep(a0|at,s)[w(s, a0)] and then takes its gradient; K must satisfy K > 1 to ensure the stated optimal condition.Source: C In-depth Discussion of In-sample… · p. 18
• If CEP’s behavior model generates OOD actions, guidance can favor them when the critic overestimates their distribution-ratio values.The resulting samples may be steered toward OOD actions rather than genuinely high-quality actions.Source: C In-depth Discussion of In-sample… · p. 19
• IGL’s objective can be estimated using one sample from the behavior policy because it contains no expectation over a joint distribution.CEP’s optimality condition requires independent multiple fake actions, which one repeated sample cannot provide.Source: C In-depth Discussion of In-sample… · p. 19

D Experimental Details

• Implementation uses a U-net score network, double Q-learning, and DPM-solver sampling, with α and K as the two Diffusion-DICE hyperparameters.The reported Q and V networks are 3-layer MLPs with 256 hidden units, and the guidance network is a 4-layer MLP with 256 hidden units.Source: D Experimental Details · pp. 19–20
• The reported D4RL results use 10^6 training steps, five seeds, and averages over 10 MuJoCo or 50 AntMaze evaluations.Evaluation occurs every 4 · 10^4 training steps; returns are standardized for MuJoCo and adjusted by subtracting 1 for AntMaze.Source: D Experimental Details · p. 20
• With piecewise f-divergence retained, the Gaussian policy in traditional DICE performs worse than the diffusion policy.This is the policy-replacement ablation shown in Figure 4.Source: Figure 4 · p. 20, D Experimental Details · p. 20
• Guide-then-select performs better even though its average Q(s, a) curves remain below those of select-from-behavior.The comparison uses normalized returns and average Q(s, a) curves.Source: D Experimental Details · p. 20, Figure 5 · p. 22
• The appendix presents D4RL learning curves for MuJoCo and AntMaze and studies candidate-count choices, including comparisons with IDQL.Table 2 lists chosen α and K, Table 3 compares K with IDQL, and Table 4 studies K.Source: D Experimental Details · p. 20, Table 2 · p. 21, Table 3 · p. 21

E Experiments Compute Resources

• A complete experiment takes about 8 hours on an NVIDIA RTX 3080Ti or 4 hours on an NVIDIA RTX 4090.Training and evaluation are parallelized; training takes about 6 or 4 hours, respectively, and evaluation about 3 or 2 hours.Source: E Experiments Compute Resources · p. 23

F Broader Impacts

• Diffusion-DICE may help solve practical offline RL tasks in robotics, health care, and industrial control.Source: F Broader Impacts · p. 23
• In-sample Guidance Learning can also be applied to generative tasks such as image generation.Source: F Broader Impacts · p. 23
• Generated fake images could mislead the public if spread online.Source: F Broader Impacts · p. 23