Score Regularized Policy Optimization through Diffusion Behavior
Authors: Huayu Chen
Cheng Lu
Zhengyi Wang
Hang Su
Jun Zhu
Venue: ICLR 2024 poster (2024)
Keywords: offline reinforcement learning
generative models
diffusion models
behavior modeling
computational efficiency
Primary area: reinforcement learning
TL;DR: We propose a new diffusion-based offline RL algorithm which circumvents the time-consuming diffusion sampling scheme.

Abstract: Recent developments in offline reinforcement learning have uncovered the immense potential of diffusion modeling, which excels at representing heterogeneous behavior policies. However, sampling from diffusion policies is considerably slow because it necessitates tens to hundreds of iterative inference steps for one action. To address this issue, we propose to extract an efficient deterministic inference policy from critic models and pretrained diffusion behavior models, leveraging the latter to directly regularize the policy gradient with the behavior
distribution’s score function during optimization. Our method enjoys powerful generative capabilities of diffusion modeling while completely circumventing the computationally intensive and time-consuming diffusion sampling scheme, both during training and evaluation. Extensive results on D4RL tasks show that our method boosts action sampling speed by more than 25 times compared with various leading diffusion-based methods in locomotion tasks, while still maintaining state-of-the-art performance.

Published as a conference paper at ICLR 2024
SCORE REGULARIZED POLICY OPTIMIZATION THROUGH DIFFUSION BEHAVIOR
Huayu Chen', Cheng Lu' Zhengyi Wang', Hang Su Su Jun Zhu 1.2+
'Department of Computer Science & Technology, Institute for AI, BNRist Center, Tsinghua-Bose Bosch Joint ML Center, THBI Lab, Tsinghua University PPazhou Laboratory (Huangpu), Guangzhou, Guangdong (chenhuay21, wang zy21 @@mails..... wwl.ssssinghua. edu.
lucheng. 1cl 50 gma il...... com {suhangss, des2 j settsinghua edu.
ABSTRACT
Recent developments offline reinforcement learning have uncovered the im- mense potential of diffusion modeling, which excels at representing heterogeneous behavior policies. However, sampling from diffusion policies is considerably slow because it necessitates tens to hundreds of iterative inference steps for one action. To address this issue, we propose to extract an efficient deterministic inference policy from critic models and pretrained diffusion behavior models, leveraging the latter to directly regularize the policy gradient with the behavior distribution's score function during optimization. Our method enjoys powerful generative capabilities of diffusion modeling while completely circumventing the computationally intensive and time-consuming diffusion sampling scheme, both during training and evaluation. Extensive results on D4RL tasks show that our method boosts action sampling speed by more than 25 times compared with vari- ous leading iiffusiin-based methods in locomotion tasks, while still maintaining state-of- the-art performance. Code: fattps://gite
com/ thu-mlhuumsaa SRPO
Ours
Diflusion TnQL.
QGPO
80
Normalized score
IỘL CºL
MoRel
70
SAB
Difliser
Inference computation
Policy
60
100k FLOPs
/aussian/Dme
Diffusion
10M
VAE
BCQ
50
AWAC
Transformer
BEAR
10'
10
Inference time (seconds' IK samples)
Figure 1: Performance and computational efficiency of different algorithms in D4RL Locomotion tasks. Computation time is assessed using a consistent hardware setup and PyTorch backend.
INTRODUCTION
Offline reinforcement learning (RL) aims to tackle decision decision-making problems by solely utilizing pre-collected behavior dataset. This offers practical solution for tasks where data collection can be associated with substantial risks exorbitant costs. A central challenge for offline RL is the realization of behavior regularization, which entails ensuring the learned policy stays in support of the behavior distribution. Weighted regression (Peng et al. 2019; Kostrikov et al. (222)) provides a promising approach that directly utilizes behavioral actions as sources supervision for policy
training. Another prevalent approach is behavior-regularizz policy optimization (Kumar et al. 2019
"Corresponding author.

---

Published as a conference paper at ICLR 2024
Wu et al. 2019), which builds generative behavior model, followed by constraining the divergence between the learned and the behavior model during policy optimization.
An expressive generative model holds pivotal role in the aforementioned regularization methods. In weighted regression, a unimodal actor is prone to suffer from the "mode covering" issue, a phenomenon where policies end up selecting out-of-support actions in the middle region between two behavioral modes (Wang et al. [023a]] Hansen- et al., 2023). Expressive policy classes, with diffusion models (Ho al.) 2020) as prime choices, help resolve this issue. In behavior- regularized policy optimization, diffusion modeling can also be significantly advantageous for an accurate estimate of policy divergence, due to its strong ability to represent heterogeneous behavior datasets, outperforming conventional methods like Gaussians or variational auto-encoders (VAEs)
(Goo & Niekum 2022 Chen et al., 2023).
However, a major drawback of utilizing diffusion models in offline RL is the considerably slow sampling speed diffusion policies usually require 5-100 iterative inference steps to create an action sample. Moreover, diffusion policies tend to be excessively stochastic, forcing the generation of dozens of action candidates in parallel to pinpoint the final optimal one (Wang et al. 2023a). As existing methods necessitate sampling from or backpropagating through diffusion policies during training and evaluation, it has significantly slowed down experimentation and limited the application in fields that are computationally sensitive require high control frequency, such as robotics. Therefore, it is critical to systematically investigate the question: is it feasible to fully exploit the generative capabilities of diffusion models without directly sampling actions from them?
In this paper, we propose Score Regularized Policy Optimization (SRPO) with a positive answer to the above question. The basic idea is to extract simple deterministic inference policy from critic and diffusion behavior models to avoid the iterative diffusion sampling process during evaluation. To achieve this, we show that the gradient of the divergence term in regularized policy optimization is essentially related to the score function of behavior distribution. The latter can be effectively approximated by any pretrained score-based model including diffusion models (Song et al., (02T). This allows us to directly regularize the policy gradient instead of the policy loss, removing the need to generate fake behavioral actions for volicy-divergence estimation (Section B
We develop a practical algorithm to solve continuous control tasks (Section 4 by combining SRPO with implicit Q-learning (Kostrikov et al. 2022) and continuous-time diffusion behavior modeling (Lu et al. 2.23). For policy extraction, we incorporate similar techniques that have facilitated recent advances in text-to-3D research such as DreamFusion (Poole et al. 2023}. These include leveraging an ensemble of score approximations under different diffusion times to exploit the pretrained behavior model and a baseline term to reduce variance for gradient estimation. We empirically show that these techniques successfully help improve performance and stabilize training for policy extraction.
We evaluate our method in D4RL tasks (Fu et al., 2020). Results demonstrate that our method enjoys more than 25x boost in action sampling speed and less than 1% of computational cost for evaluation compared with several leading diffusion-based methods while maintaining similar overall performance in locomotion tasks (Figure [ We also conduct 2D experiments better illustrate that SRPO successfully constrains the learned policy close to various complex behavior distributions.
2
BACKGROUND
2.1 OFFLINE REINFORCEMEN LEARNING
Consider a typical Markov Decision Process (MDP) described by the tuple (S, P,r,)) where S is the state space, A the action space, P(s' (s a) the transition function, r(s, a) the reward function and y the discount factor. The goal of reinforcement learning (RL) is to train a parameterized policy To(als 8) which maximizes the expected episode return. Offline RL relies solely on static dataset DA containing interacting history {s, as.,,,, a, s'} between a behavior policy p(als) and the environment to train the parameterized policy.
Suppose we can evaluate the quality of a given action by estimating its expected return-to-go using Q-network Qs(s, a) (=(ssa) := E.,=..r==::7 D r an)], we can formulate the training objective of offline RL as max, Ea- (a) a) DKL [+((|s)]|(([ss]]] (Wu et al. 2019). Note that KI. regularization term is added mainly ensure the learned policy stays in
2

---

Published as a conference paper at ICLR 2024
20
mode covering
1,0
1.
mode socking
drocsiian
93
0.5
0,5
0.0
30
0.0-
0.0
B-05
1110
8.0
9-13
-05
-6.5
-0.5
8:40
0:43
8.4.
-10
-1.0
-1.0
0
0.5
0o
0.5 1.
1.0
Action dinession
Adtion
dimeetsion
(a) Behavior dataset
Forward KI. &
(c) Forward KT.
(d) Reverse KT. &
Unimodal policy
Expressive policy
Unimodal policy
Figure 2: Comparison of different policy extraction methods under bandit settings. Forward KL policy extraction prone generate out-of-support actions if the policy sufficiently expressive (e.g., Gaussians). This can be mitigated either by employing more expressive policy class or by switching to a reverse objective (our method), which demonstrates mode-seeking nature.
support of explored actions. 3 is some temperature coefficient. Previous work (Peters et al. 2010 Peng et al. 2019) has shown that the optimal policy for such an optimization problem
T (als)
(ss) p(als) exp (3Qx(s.a))
(1)
where (s) is the partition function. The core problem for offlin RL now becomes how to efficiently model and sample from the optimal policy distribution x² (-(ss.
2.2 OPTIMAL POLICY EXTRACTION
Existing methods to explicitly model x² with a parameterized policy we be roughly divided into two main categories-weightee weighted regression and behavior-regularizee policy optimization:
min DR DKL [xt (-1s) (To{)))) + max (a 0
log To Ta(as)) Qo((t))
(2)
Z(s)
Forward KI.
Weighted Regression
min DR DKL [To (1)(())) max R aa Qo 1 DKL. (To (mo((1))]
(3)
Reverse KL
Sehavior-Regularin Policy Optimization
Weighted regression directly utilizes behavioral actions as sources of supervision for policy training. This circumvents the necessity to explicitly model the intricate behavior policy but leads another mode-covering issue due to the objective's forward-Ko KL nature (Eq. (2)) The direct consequence is that weighted regression algorithms display sensitivity to the proportion of suboptimal data in the dataset Yue et al. 2022), especially when the policy model lacks distributional expressivity (Chen
jet al. (023), as is depicted in Figure 2 Recent work (Wang et al., 2023a Lu et al. 2023) attempts to alleviate this by employing more expressive policy classes, such diffusion models (Ho et al. 2020).
However, these methods usually compromise on computational efficiency (Kang et al. 2023
In comparison, behavior-regularized policy optimization (Wu et al. 2019) emerges approach for training simpler policy models such as Gaussian models. This fundamentally stems
suitable
from its basis on reverse-KL. objective (Eq. (3)) which inherently encourages mode-seeking behavior. However, approximating the second KL term in Eq. (3) is usually difficult. Regarding this, practical implementation previous studies (Kumar al., 2019 Wu al. 2019 Xu et al., 2021 usually first construct generative behavior models to approximate the policy-behavior divergence.
2.3 DIFFUSION MODELS FOR SCORE FUNCTION ESTIMATION
Diffusion models (Sohl-Dickstein et al. 2015 Ho et al. 2020 Song et al. 2021 are powerful gener- ative models. They operate by defining forward diffusion process to perturb the data distribution into a noise distribution for training the diffusion model. Subsequently. this model is employed to reverse the diffusion process, thereby generating data samples from pure noise.

---

Published as a conference paper at ICLR 2024
In particular, the forward process is conducted by gradually adding Gaussian noise to samples zo from an unknown data distribution (o(Zu) == ((x) ) at time 0, forming a series of diffused distributions qs (x,) at time t. The transition distribution 9:0 (g)))) is:
gro (xo) (a, of
which implies
3t - 0, + ore,
Here, Ct2, de 0 are manually defined, and is random Gaussian noise.
For the reverse process, Ho et al. (020)) train a diffusion model eø (az(() to predict the noise added to the difffused sample æ, in order to iteratively reconstruct æo. The optimization problem is
min E 2o..
[1) [t) =
(5)
More formally, SSng et al. (2021) show that diffusion models are in essence estimating the score function Vz, log qe (z;) of the diffused data distribution qt. such that:
m, log qr (xx) - (are1)/a, Ea (2//o/)/
(6)
and the reverse diffusion process can alternatively be interpreted discretizing an ODE:
dæ. f(t) / (t) log (x(:),
(7)
df
where f(() dlog (() dof 2 2 leaving V æz log (xx) the only unknown term.
dt
In offline RL, diffusion models have been discovered as an effective tool for modeling heterogeneous behavior policies. Usually states considered as conditions while actions considered as data points æ, such that a conditional diffusion model e(a,/s [s, t) can be constructed to represent ju (s)ss).
3 SCORE REGULARIZED POLICY OPTIMIZATION
In this paper, we seek to learn a deterministic policy To to capture the mode of a potentially complex policy distribution x² introduced in Eq. ( To achieve this, employ reverse-KL policy extraction scheme (Eq. (3) given its mode-seeking nature:
L. Bg amy Qa(s, a) DEL (To( (-s))))ll(((1)))
Solving the above optimization problem requires estimating DKL. (- | [[[[ss)] Regarding this, previous research (Kumar al. 2019 Wu et al., 2019; Xu al. 2021 Wu et al. (022)) use sample-
based methods: first constructing behavioral model ps 2= μr, followed by sampling fake actions from (-(s) and (-(s)) to approximate the policy divergence. However, this approach necessitates sampling from the behavior model during training, imposing substantial computational burden. This drawback is exacerbated when employing expressive yet sampling-expensive behavior models such diffusion models.
We propose an alternative way to solve Eq. S
By decomposing the KL term, we can get
(0)
BENDD ao a0(s.
EgnDr ammy log p(als)
E. EANDE H ((
Policy optiiization
Behavior rogularization
Entropy (aften constan
Then we calculate the gradient of Eq. under the condition that ny is deterministie. Applying the chain rule and the reparameterization trick, we have:
(0) = EaNDI VaQu(s, a)l==o())) + Va ((
V (s).
(10)
(by Eq
noted that the only unknown term above is the score function V log u(a]s) (s)s of the behavior distribution. Our key insight is that a pretrained diffusion behavior model can be leveraged to
In this paper only consider the cases where is isotropic Gaussian with fixed variance. For brevity, informally view Dirac as Gaussian whose variance is infinitesimally small.
4

---

Published as a conference paper at ICLR 2024
Data samples (from behavior dataset)
(b) Behavior density ji(a) (c) Quadratic Q-functions
(d) Behavior regularization
(e) Result policy shift (by varying 3
(by diffusion models)
(stacked)
1)
w/o behavior regularizution
w/ behavior regularization
Illustration of SRPO in 2D bandit settings. (a) A predefined complex data distribution, which represents the potentially heterogeneous behavior policy p(a). (b) A diffusion model j a is trained to fit the behavior distribution. The data density can be analytically calculated based on Song al. (2021). (c) The Q-function manually defined a quadratic function: Q(a) == aur where represents the 2D point with the highest estimated Q-value and is selected from set of grid intersections. These individual Q-functions with different aur are depicted together stacked way in Figure (c). (e)&(e) By optimizing deterministic policies TT-) areg according to Eq. (10) and tuning the temperature coefficient 3, resulting policies shift from greedy ones which tend to maximize corresponding Q-functions conservative which are successfully constrained close to the behavior distribution. See more experimental results in Appendix A
Figure 3:
Behavior Density (VAFs)
BCQ (Fujimoto et al. 2019
BEAR (Kumar et al. 2019)
TD3+BC Fujimoto & Gu 202
Figure 4: Performance of other behavior regularization methods. See more results in Appendix A
effectively estimate this term. This is because diffusion models e(xltt essentially approximating the score function V log (x) of the diffused data distribution (xx)) (Eq. (6))
Specifically, first pretrain diffusion behavior model, denoted <(a, [s,(), approximate
Va log (((as)) By doing so, we can regularize the optimization process of another deterministic actor иT.. We term our method as Score Regularized Policy Optimization (SRPO), given its distinctive feature of performing regularization at the gradient level, as opposed to the loss function level. Figure provides 2D bandit example of SRPO.
Compared with previous work (Wang et al. 2023a Lu et al., 2023 Hansen- Estruch et al. 2023 that directly trains diffusion policy for inference in evaluation, the main advantage SRPO its computational efficiency. SRPO entirely circumvents the computationally demanding action sampling scheme associated with the diffusion process. Yet, still taps the robust generative strengths of diffusion models, especially their ability to represent potentially diverse behavior datasets.
PRACTICAL ALGORITHM
In this section, we derive a practical algorithm for applying SRPO offline RL (Algorithm [T. The algorithm includes three parts: implicit Q-learning (Section 4.1l; diffusion-based behavior modeling (Section 1); and score-regularized policy extraction (Section 4.2
4.1 PRETRAINING THE DIFFUSION BEHAVIOR MODEL AND Q-NETWORKS
For Q-networks, we choose to use implicit Q-learning (Kostrikov et al. 2022) to decouple critic training from actor training. The core ingredient of the training pipeline is expectile regression, which
5

---

Published as a conference paper at ICLR 2024
diffusion
diffusion 0.1
diffusion 0.3
diffusion
ensemble (0,
density Kensity-map -aa
regularization
Figure 5: Empirical benefits of ensembling multiple diffusion times. See Remark for detailed explanation.
in Appendix B
requires only sampling actions from existing datasets for bootstrapping:
min (C) = (D (L..)) a) - V((s))], where (z = =|r- 1(u < 0)/2²,
(11)
min (o( = =B(,,,,"') [l a) + +V(('') (s, a
(12)
When the expectile parameter E (0, 1) is larger than 0.5, the asymmetric L2-objective Ly (C) would downweight suboptimal actions which have lower Q-values, removing the need of an explicit policy.
For behavior models, in order to represent the behavior distribution with high fidelity and estimate its score function, we follow previous work (Chen al. 2023 Hansen- et al. 2023) and train conditional behavior cloning model:
min (1b) = (s,,, [alew = =aiatoce */
(13)
where (0, and 1). The model chitecture of cg is consistent with the one proposed by Mansen-Estruce et al. (2023), with the sole difference be- ing that our model uses continuous-time inputs similar to Lu et al. (2023) instead of discrete ones as used by Wang al. (2023a) and Hansenn-nsttte et al. (2023).
Algorithm 1 SRPO
Initialize parameters vh, S. Φ, 6.
// Critic truining (IQL) for each gradient step do < xv Lv (Eq.
- Φ Xo V Lo (Eq // Behavior training for each gradient step do x 3 A, L.) (Eq.
Once we have finished training the behavior model 6 we can use it estimate the score function of pe, the diffused distribution of p(als) at time We have Va, log (s,))sstt la,,asss (arlal tt/ as shown by Song et al. (222))
// Policy extraction for each gradient step do
A V (um
(Bq.
4.2 POLICY EXTRACTION FROM PRETRAINED MODELS
The policy extraction scheme proposed in Section 3 only leverages the pretrained diffusion behavior model 6ø (ar)s, I f) time 0, where the behavior distribution j has not been diffused. However, ca(a,)s (s, t) is trained to represent series of diffused behavior distributions #4 at various times t E (0, 1). In order to exploit the generative capacity Es, we replace the original training objective C, (0) with new surrogate objective:
max T (0) = Egganto Qo(s, a) ( w(t Tf DeL (m))(()))) az
(14)
where 4(0002, 0.98). DH. Both Pe and To. follow the same forward diffu- sion process in Eq. 4 where (4(aa(s) Banp(le:Vaa (aaaaaa of I), and (a.(s))))

---

Published as a conference paper at ICLR 2024
Ba--so(iaa of (t) is weighting function that adjusts the importance of each time t.
We can nearly recover (0) by setting (() to 5(t - 0.02) a2.02 (ablation studies in Section 6.3)
00.02
Empirically, the surrogate objective - (0) ensembles various diffused behavior policies Ar to
regularize training of the same parameterized policy Tg. This is supported by an observation (Proposition I in Appendix B) arg min, DKL [T(((1s)]lx<((sss]] arg min, DKL [=(-|s)]|a((|ss]
Similarly to Section we can optimize Eq. (14) by calculating its gradient:
Proposition 2. (Proof in Appendix B) Given that is deterministic (a TT(S)) such that ne, is Gaussian (a, = 0,a + oie, 6 ~~ N0, T)), the gradient for optimizing max ( satisfies
VeL*" E. (oalols,a)aa-nae Brr (ttleelll 1a VOTO(S). (15)
1
mudracctdd bazeline
Note that we additionally subtract the action noise from (a, 8s f) in the above equation. This term does not influence the expected value of V VoLuurg (0) but could reduce the estimation variance since it is correlated with c. (arls, t) (See Appendix B). We refer to it as the baseline term because it
similar the subtracted baseline in Policy Gradient (Sutton Barto 1998) algorithms.
The idea of ensembling diffused behavior policies for regularization and subtracting the baseline term to reduce estimation variance both draw inspiration from the latest developments in text-to-3I) research such as DreamFusion (Poole et al. 2023). We elaborate more on the connection between the two work in Section and ablate these techniques in the realm of continuous control in Section 6.3
RELATED WORK
Behavior Regularization in Offline Reinforcement Learning. Behavior regularization can be achieved either implicitly or explicitly. Explicit methods usually necessitate the construction of behavior model to regularize the learned policy. For example, TD3+BC (Fujimoto & Gu 2021) implicitly views the behavior as Gaussians and introduces an auxiliary L2-loss term to realize the regularization. SBAC (Xu et al. 202I and Fisher- Fisher-BRR (Kostrikov et al. 2021) leverage explicit Gaussian (mixture) behavior models. BCQ (Fujimoto et al. 2019) and BEAR (Kumar et al. 2019) leverage VAE behavior models (Kingma & Welling, 2014). Diffusion-QL. (Wang et al. 2023a) is similar TD3+BC but swaps the auxiliary loss with diffusion-centre objective. QGPO (Lu et al. 2023
views a pretrained diffusion behavior as the Bayesian prior in energy-guided sampling.
Diffusion Models in Offfine Reinforcement Learning. Recent advancements in offline RL have identified diffusion models as an impactful tool. A primary strength of diffusion modeling lies in its robust generative capability, combined with a straightforward training pipeline (Dhariwal &
Nichol 2021 Xu et al., 2022 Poole et al. (023). This renders it particularly suitable for modeling
heterogeneous behavior datasets (Janner et al.J 2022 Wang et al., 2023a Ajay et al. 2023 Hansen-
Estruch et al. 2023), generating in-support actions for O-lernning (Goo & Niekum
2022 Lu et al.
2023). and representing multimodal policies (Chen et al. 2023 Pearce et al. 2022).
However, significant concern for integrating diffusion models into RL is the considerable time taken for sampling. Various strategies have been proposed to address this challenge. including employing parallel sampling scheme during both training and evaluation (Chen al.) (023). utilizing a specialized diffusion ODE solver such as DPM-solver (Lu et al.) (202]]202)) adopting an approximate diffusion sampling scheme to minimize sampling steps required (Kang al. 2023), and crafting high-throughput network architectures (Hansen-Esiin Estruch et al. (023). While these techniques offer improvements, they don't 't entirely eliminate the need for iterative sampling.
Score Distillation Methods. Recent developments in text-to-31) generation have enabled the transformation of textual information into 3D content without requiring any 3D training data (Poole) al. 2023 Wang et al., 2023b). This is realized by distilling knowledge from text-to-image diffusion models. A representative method in this domain is DreamFusion (Poole et al. (023). It optimizes a 3D NeRF model (Mildenhall et al. 2222 by ensuring its projected 2D gradient follows the score direction of large-scale, pretrained 2D diffusion model (Sahania et al. 2022). Similar to DreamFusion, SRPO also employs diffusion model guide the training of a subsequent network. However, our method
7

---

Published as a conference paper at ICLR 2024
Dataset
Environment
BEAR
TD3+BC
IQL
SIBC
Diffuser
DiffusionQLL
QGPO
IDQL
SRPO
(Ours)
Mediam- Expert
HallCheetah
53.1
90.7
86.7
92.6 6
79.8
96.8
95.5
95 9
92 2 3.0
Expert
Hopper
96.0
107.2 2
108.6
Medium-Fxpen
Walker2d
40.1
110.1
109.6
109.8
108.4 4
110.1
110.7
112.7
114.0 2.1
Medium
HallCheetahh
41.7
48.3
47.4
45.9
44.2
51.1
54.1
51.0
60.4
0.8
Medium
Hopper
62.1
59.3
66.3
67.1
68.5
90.3
98.0
65.4
95 5 2.0
Medium
Walker2d
59.1
63.7
78.3
77.9
79.7
67.0
86.0
82.5
84.4 1: 4.4
Replay
HaJlCheetah
41.6
37.1
42.2
45.9
51 3.4
Medium-Reet Replay
Hopper
33.7
60.9
94.7
86.2
101.3
100.7
96.9
92.1
101.2 1.0
Medium Replay
Walker2d
19.2
81.8
73.9
65.1
61.2
95.5
64.4
85.1
S4.6 = 7.1
Average (Locomotion)
75.3
75.6
75.3
Default
KaiMaze-umaze
73.0
78.0
67.5
92.0
93.4
96.4
94.00
97.1 1 2.7
Diverse
AntMaze
61.0
T1.4
62.2
85.3
66.2
74.4
80.2
82.1 10.8
Play
AntMaze medium
0.0
10.6
71.2
81.3
76.6
83.6
84.5
80.7 + 7.1
Diverse
AnlMaze-metium
8.0
3.0
70.0
82
78.6
83.8
84.8
75.0 +: 12.3
Play
AntMazzelarge
0.0
0:22
39.6
50.3
46.4
66.6
63.
53.6 12.5
Diverse
AniMaze-large
0.0
0.0
17.5
45.5
56.6
64.8
9
63.6 = 0.3
Average (AntMaze)
27.3
74.2
79.1
Table I: Evaluation numbers of D4RL benchmarks (normalized suggested by Fu et al. (2020))
We report mean standard deviation of algorithm performance across random seeds at the end of training. Numbers within 5 % of the maximum in every individual task are highlighted.
emphasizes score regularization as opposed to score distillation. The behavior score is additionally incorporated to regularize the Q-gradient instead of being the only supervising signal.
6 EVALUATION
6.1 D4RL PERFORMANCE
In Table we evaluate the D4RL performance (Fu et al. 2020) of SRPO against other offline RL algorithms. Our chosen benchmarks include conventional methods like BEAR (Kumar et al.
2019). TD3+BC (Fujimoto & Gu (021). and IQL (Kostrikov et al. (022), which feature extracting Gaussian/Dirac policy for evaluation. We also look at newer diffusion- basoo- offline RI. techniques, such as Diffuser (Janner et al. 2022), DIffusion-QUI (Wang et al. 2023a), SfBC (Chen et al., 2023)..
QGPO (Lu et al., 2023), and IDQL Sltansnn-Estucce et al., 2023). These methods tend to be more computationally intensive but generally offer better results.
Of all the baselines, our comparison with IDQL is particularly informative. This is because SRPO shares a virtually identical training pipeline and model architecture for critic and behavior models with IDQL, as is deliberately crafted. The most significant distinction lies in their approaches for extracting policy: IDQL. skips the policy extraction step, choosing to evaluate directly with the behavior policy, using
technique (Chen et al. 2023). In contrast, SRPO extracts a Dirac policy from behavior and critic models.
Overall, SRPO consistently surpasses referenced baselines that also utilize a Gaussian (or Dirac) inference policy, leading by large margins in the majority of tasks. It also comes close to match- ing the benchmarks set by other state-of-the-ar diffusion-based methods, such as Diffusion-Q!! and IDQI,, though it features a much simpler inference policy. Moreover, in the case of the convergence rate and training stability, we show- case training plots of SRPO alongside several baselines in Figure 6 Results also suggest that compared with both weighted regression meth- ods like IQL and reverse-KI -ba-ee optimization strategies like Diffusion-QI., SRPO exhibits more favorable attributes. We attribute the performance gain of SRPO the utilization of diffusion behav- ior modeling, which facilitates a more effective realization of the training formulation.
weluch
Felyy Clad and
Eellay Cord
Foke Gralee
Figure 6: Training curves of SRPO (ours) and sev- eral baselines. See complete experimental results in Appendix
8

---

Published as a conference paper at ICLR 2024
6.2 COMPUTATIONAL EFFICIENCY
As indicated by Figure and 7 SRPO maintains
high computational efficiency, especially fast in- ference speed, while enabling the use of power- ful diffusion model. Notably, the action sampling speed of SRPO is 25 to 1000 times faster than that of other diffusion- based methods. Addition- ally, its computational FLOPS amount to only 0.25% to 0.01% of other methods. This makes SRPO ideal for computation--senittve contexts such as robotics. also speeds up experimenta- tion, since policy evaluation typically consumes over half of the experiment duration for previ- diffusion-based approaches. This efficiency stems from SRPO's design, which completely avoids diffusion sampling throughout both train-
Thiiitg Time (aconderi gowter apo
Pokey Traniig
Baaaii Transea
(scoedi ego
Piio
Cakic Tainiing
122. DifferionQL IPQU
OhEc
12"
101. D.t.tloogOLLUI
Figure 7: Training and inference (evaluation) time required for different algorithms in D4RL loco- motion tasks. All experiments conducted with the same PyTorc backend and the same comput- ing hardware setup.
ing and evaluation procedures.
6.3
A ABLATION STUDIES
Weighting function w(t). As is explained in Section 4.2 if (t) x d(t- 0.02), the surrogate objective nearly recovers the original one, which could guarantee convergence to the wanted opti- mal policy x². Adjusting w()) to ensemble dif- fused behavior score various times / (0, 1) might yield a smoother and more robust gradi- landscape. However, it biases the original training objective. In an extreme case where wt) x 6(t 0.98), the estimated behavior score almost becomes the score function of pure Gaussian distribution, making it entirely help- less. From Figure we empirically observe that Antmaze tasks are sensitive variation of ((). while Locomotion tasks not. Overall, find w(t) = of a suitable choice for all tested tasks and choose it the default hyperparameter throughout our experiments. This hyperparam- eter choice is consistent with previous literature (Poole et al. 2023).
MulaCo Locometion
NomalizedSoo Scors
12 10 -6
subtructigg & MAU suboading
0d/%
Const1
610
Nomalized Scom Soom
subouctiss suberattin
WIB: AI 0.02)
ds/d
Const1l.0
6/C 1 0.98)
wit
Figure 8: Ablation for implementation details of Eq. (15). All results are averaged under 6 random seeds.
Subtracting baselines 6. Figure we show that subtracting baselines from the estimated consistently offers slightly better performance.
behavior gradient, as is described in Eq. 15
Temperature coefficient | and others. Varying 3 directly controls the conservativness of policy, and thus influences the performance. Experimental results with respect to 3 and other implementation choices are presented in Appendix
7 CONCLUSION
In this paper, we introduce Score Regularized Policy Optimization (SRPO), an innovative offline RL algorithm harnessing the capabilities of diffusion models while circumventing the time-consuming diffusion sampling scheme. SRPO tackles the behavior-reegularized policy optimization problem and provides a computationally efficient way realize behavior regularization through diffusion behavior modeling. The fusion of SRPO with techniques like implicit Q-leanning can further solidify its application computationally sensitive domains, such as roboties.

---

Published as a conference paper at ICLR 2024
REPRODUCIBILITY
To ensure that our work is reproducible, we submit the source code as supplementary material. Re- ported experimental numbers are averaged under multiple random seeds. We provide implementation details of our work in Appendix
ACKNOWLEDGEMENT
This work was supported by the National Key Research and Development Program of China (No. (220AAA0103002), NSFC Projects (Nos. 62350080, 92370124, 92248303, 62106123, 61972224), Tsinghua Institute for Guo Qiang, and the High Performance Comput- ing Center, Tsinghua University. J.Z. was also supported by the New Cornerttone Science Foundation through the XPLORER PRIZE.
BNRist
REFERENCES
Anurag Ajay, Yilun Du. Abhi Gupta, Joshua B. Tenenbaum, Tommi S. Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision making? In The Eleventh International Conference on Learning Representations, 2023.
Huayu Chen, Cheng Lu, Chengyang Ying, Hang Su, and Jun Zhu. Offline reinforcement learning via high-fidelity generative behavior modeling. In The Eleventh International Conference Learning Representations, 2023.
Prafulla Dhariwal and Alexander Quinn Nichol, Diffusion models beat GANs on image synthesis. In Advances in Neural Information Processing Systems, 2021.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020.
Scott Fujimoto and Shixiang Shane Gu. minimalist approach to offline reinforcement learning. Advances in neural information processing systems, 34:20132-201453 2021.
Scott Fujimoto. David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In Proceedings of the 36th International Conference on Machine Learning, 2019.
Wonjoon Goo and Scott Nickum. Know your boundaries: The necessity of explicit behavioral cloning in offline arXiv preprint arXiv:2206.0066 00695, 2022.
Philippe Hansen- Estrucc,, llya Kostrikov, Michael Janner, Jakub Grudzien Kuba, and Sergey Levine. Idql: Implicit q- gelaarning as an actor-eritic method with diffusion policies. arXiv preprint arXiv:2304. 10573, 2023.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, 2020.
Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning, 2022.
Bingyi Kang, Xiao Ma, Chao Du, Tianyu Pang, and Shuicheng Yan. Efficient diffusion policies for
offline reinforcement learning, arXiv preprint arXiv:2305.2000 200000 2023.
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. 2nd International
Conference Learning Representations, 2014.
Ilya Kostrikov, Rob Fergus, Jonathan Tompson, and Ofir Nachum. Offline reinforcement learning with fisher divergence critic regularization. In International Conference on Machine Learning, pp. 5774-5783. PMLR, 2021.
Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit Q-learning. In International Conference on Learning Representations, 2022.
10

---

Published as a conference paper at ICLR 2024
Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off-policy q-learning via bootstrapping error reduction. Advances in Neural Information Processing Systems, 2019.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. arXiv preprint arXiv:2206. 00022 2022.
Cheng Lu, Huayu Chen, Jianfei Chen, Hang Su, Chongxuan Li, and Jun Zhu. Contrastive energy prediction for exact energy-guided diffusion sampling in offline reinforcement learning. In Proceedings of the 40th International Conference on Machine Learning. 2023.
Ben Mildenhall, Pratul P Srinivasan, Matthew Tancik, Jonathan T Barron, Ravi Ramamoorthi, and Ren Ng. Nerf: Representing scenes as neural radiance fields for view synthesis. Communications of the ACM, 65(1):99-111 106, 2021.
Tim Pearce, Tabish Rashid, Anssi Kanervisto, David Bignell, Mingfei Sun, Raluca Georgescu, Sergio Valcarcel Macua, Shan Zheng Tan, Ida Momennejad, Katja Hofmann, et al. Imitating human behaviour with diffusion models. In Deep Reinforcement Learning Workshop NeurIPS, 2022.
Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. Advantago-wcighted regression: Simple and scalable off-policy reinforcement learning. arXiv preprint arXiv: 191.00//77 2019.
Jan Peters, Katharina Mulling, and Yasemin Altun. Relative entropy policy search. In Twenty-Fourth AAAI Conference on Artificial Intelligence, 2010.
Ben Poole, Ajay Jain, Jonathan T. Barron, and Ben Mildenhall. Dreamfusion:
Text-to-3d using 2d diffusion. In The Eleventh International Conference on Learning Representations, 2023.
Geoffrey Roeder, Yuhuai Wu, and David K Duvenaud. Sticking the landing: Simple, lower-variance gradient estimators for variational inference. Advances Neural Information Processing Systems, 30, 2017.
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L. Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Bureu Karagol Ayan, Tim Salimans, et Photorealistic text-to-image diffusion models with deep language understanding. Advances in Neural Information Processing Systems, 35:36479-36494 2022.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamies. In International Conference Machine Learning, pp. 2256-2265, 2015.
Yang Song, Jascha
Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021.
Richard S. Sutton and Andrew G. Barto. Introduction to Reinforcement Learning. MIT Press, Cambridge, MA, USA, 1st edition, 1998. ISBN 0262193981.
Zhendong Wang, Jonathan Hunt, and Mingyuan Zhou. Diffusion policies as an expressive policy class for offline reinforcement learning. In The Eleventh International Conference on Learning Representations, 2023a.
Zhengyi Wang, Cheng Lu, Yikai Wang, Fan Bao, Chongxuan Li, Hang Su, and Jun Zhu. Prolific- dreamer: High-fidelity and diverse text-to-3d generation with variational score distillation. arXiv preprint ar05 16213, 2023b.
Jialong Wu, Haixu Wu, Zihan Qiu, Jianmin Wang, and Mingsheng Long. Supported policy opti- mization for offline reinforeement learning, Advances in Neural Information Processing Systems, 2022.
Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv: /9 1361, 2019.
11

---

Published as a conference paper at ICLR 2024
Haoran Xu, Xianyuan Zhan, Jianxiong Li, and Honglei Yin. Offliee reinforcement learning with soft behavior regularization. arXiv preprint arXiv:2110.07395, 2021.
Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. In International Conference on Learning Representations, 2022.
Yang Yue, Bingyi Kang, Xiao Ma, Zhongwen Xu, Gao Huang, and Shuicheng Yan. Boosting offline reinforcement learning via data rebalancing. arXiv preprint 2Xiv:2210.09241 2022.
12

---

Published as a conference paper at ICLR 2024
A
COMPLETE 2D EXPERIMENTS
8gaussians
swissroll
checkboard
2spirals
rings
moons
(a) data
(b) diffusion
(c) VAE
(e) BCQ (d) SRPO (ours)
( BEAR
(g) TD3+BC
wio behavior regularization
w: behavior regularization
Figure 9: Experimental results of SRPO and other baselines in 2D bandit settings.
13

---

Published as a conference paper at ICLR 2024
diffusion
diffusion 0.1
diffusion 0.3
diffusion
ensemble (0,
density map
regularization
density-map
regularization
density-min -ap
regularization
density-man
regularization
wio behavior regularization
w:' behavior regularization
Figure 10: Empirical benefits of ensembling multiple diffusion times
14

---

Published as a conference paper at ICLR 2024
diffusion t
diffusion 0.1
diffusion 0.3
diffusion
ensemble (0,
density- Sensity-map
regularization
wio behavior regularization
w/ behavior regularization
Figure 11: Empirical benefits of ensembling multiple diffusion times
Figure 12: Effect of varying temperature /3.
15

---

Published as a conference paper at ICLR 2024
B THEORETICAL ANALYSIS
In this section, we provide some theoretical analysis related to the SPRO algorithm. First, aim provide some insights into why is reasonable to replace
1
- (0)= = EN am Q0(8, a) DKL. [m0([))]|aa({s{])
with the surrogate objective
a (0) Es Qo(s, Be at
DKL (
(16)
a4
Proposition
Given that is sufficiently expressive, for any time 1, any state 8, we have arg min DKL (Te(-s))]pee((ess arg min DKL [=((ss)]a((|s]]],
where both jt and ne. follow the same predefined diffusion process in Eq. ④
Proof. Our proof is inspired by Wang al. (2023b). Regarding the property of the KL divergence, we have arg min, DKL. (-(()) [ [8] = (4((8)) and (-18) := arg min, DRL. [Te(-18)[lu<(((s)) (-1s)))
We conclude that nद (-|s) = ( (-18) is equivalent to n= r=(s[) = pa-(s) (s) by transforming all distributions
into their characteristic functions.
According to the forward diffusion process defined by Eq. for any prespecified state 8, we have a -aaa toe a,a" + oge,
such that
Te Te(a, [s) -
/ NV(alaaaa of (a)ssaa
Therefore, the characteristic function of (a,ls)) satisfies
(+1e)())))) = =er(ialaawoxxooooou = (arl))(ass) oi
Similarly, we also have
(u) = =x(olaat))xxxoon(()) = (asl))(asse oi2
Finally, we can see that
t/ (-18) pa(+1s) + (M) = Qui( s)(u) t> )n (p(()1)(()) & (118) = m(-1s)
is imperative note that although arg min, DRL. (e(els))l [-()) and arg min, DKL [n[-]88 [m((s)]lu(-|ss] converge to the same global optimal solution. max Luim (0) does not converge the wanted optimal policy x² (als) p(a)s s)) exp (BQs (s, a)) while max L (0) does. This discrepancy arises primarily from the inclusion of the additional Q-function term. Fur- thermore, in practical scenarios, the parameterized policy ny lacks expressivity, which hinders its ability to truly attain the global optimum. Nevertheless, ensembling a series of I might be beneficial for the optimization to get to a better minimum in practice. Since the distribution p(-(s) is usually complex and high-dimensional, it is easy to become trapped in a local minimum while maximizing its likelihood. However, the distribution of (+(+s)) gets smoother as gets larger.
Remark
Next, we derive the gradient for (= under the condition that лe deterministic actor.
Proposition 2. Given that ma is deterministic actor (a TT(8)) and e" is optimal diffusion behavior model (E" (ar (, -0, Va, log pe (a, 18), the gradient for optimizing maxę L (0) in
Eq. 16 satisfies
1 1
surT (0) = R. ( Puattc)))
-) (S)
16

---

Published as a conference paper at ICLR 2024
Proof.
According the predefined forward diffusion process (Bq. for any state s, we have
To.((ai)))= = N(ailora, (ooolaasad N Dia =N
Therefore (-18) is a Gaussian with expected value (8T and variance of We rewrite the training objective below:
(0) E ss Qu(s,
2 3 [lee
1
DKL DKL (Teo
ar
Q0) +
log (orls())))) - log ne,0 (a: s)] at 3 B ((.am)))()))())
ae
- =o.Qo a) (m) + ww(()
(oo un (aaaaaaa (a,)s) log ((r)s)
az
Then we derive the gradient of (*) by applying the chain rule and the parameterization trick:
V (0)
OR, (o,,Q0(800 a)
(m)(s)
at + 3 (w(t)
E. aa,
aloog pr (a; 8))
da.
da a=
00
axe
a0
behuviue
(l
1
02
E.
llog Te, (ac(s) da.
00 la==aro()))trcc + BaINnr
i log Te, (arls)
(17)
B
Da.
00
policy score
parameter
The behavior score term in the above equation can be represented by the optimal diffusion behavior model:
alog pe(ar [s] E" ('(ar(8,1)
Da.
Ot
The policy score term is the score function of x-0((18) (N(LTTO))) of have
i log te, [)]
ô ||a: (1)
Da.
da,
20²
0.=0)
The parameter score term equals 0 for any distribution T, regardless of whether it is Gaussian:
Ø Olog Te,
i log ne,o Te,
E Ea
Tt,0 (a:ls)
da.
00
00
da;
00
a
/ [ s)das)
00
=0
Continue on VoL (0) and substitute the conclusions from above:
ZE, (s, a)
@To(s)
- E. E'(alsstt) 0a:
Vol ASUI (0)
da
a=
00
+ 3
ar
00 (s)
da,
(t) Ott
E.
+0
a: 00
DE. a)
(T)(s)
1 E. (E" (arls,
da,
Be,
-c]
da
a=TT(())
00
at
20
OE,Qo(8 a) da
@m((s)
1
OTo
Be, ar
R. (arls,)) -c c) 00
a=
00
a,=0rm0())+or
E, VaQQos alss,a)l=ooo0
(e ( VOTO(s)
subtracted baseline
17

---

Published as a conference paper at ICLR 2024
The subtracted baseline above corresponds the policy score term in Eq. (17.. It does not inftuence the expected value of the empirical surrogate gradient 19Lx"(0) ) To see this, consider isolating the baseline term L.E.,.cC. Given that is random Gaussian noise independent of both state and time I, we can prove ,,, O. Still, previous work (Roeder et al. [017]] Poole et al. 2023) shows that keeping the policy score term can reduce the variance of the gradient estimate and thus speed up training.
Remark 2.
C EXPERIMENTAL DETAILS FOR D4RL BENCHMARKS
Critic training. We train our critic models following Kostrikov al. 2022 For the convenience of readers, we recap some key hyperparameters: All networks are 2-layer MLPs with 256 hidden units and ReLU activations. We train them for 1.5M gradient steps using Adam optimizer with learning rate of 3e-4. Batchsize is 256. Temperature: T 0.7 (MuloCo locomotion) and T = 0.9 (Antmaze).
Behavior training. We adopt the model architecture proposed by Haasen-EEstuuh et al. 2023)) with modification to accommodate continuous- input. A single scalar time input is mapped to a high-dimensional feature using Gaussian Fourier Projection before concatenated with other inputs. The network is basically a 6-layer MI.P with residual connections, layer normalizations, and dropout regularizations. We train the behavior model for 2.0M gradient steps using AdamW optimizer with a learning rate of 3e-4 and a batchsize of 2048. Empirical observations suggest that much fewer pretraining iterations (e-g.. 0.5M steps) do not cause a drastic performance drop, but we want to ensure training convergence in this work. The diffusion data perturbation method follows the default VPSDE setting in Song al. (202T) and consistent with prior work (Lu al. 2023).
Policy extraction (Locomotion). The policy model is 2-layer MLP with 256 hidden units and ReLU activations. is trained for 1.0M gradient steps using Adam optimizer with a learning rate of 3e-4 and batchsize of 256. For all tasks w(t) o². For the temperature coefficient, sweep over 3 E (0.01, 0.02, 0.05, 0.2, 0.5) and observe large variances in appropriate values across different tasks (Figure 15 We speculate this might be due to / being closely intertwined with the behavior distribution and the variance of the Q-value. These factors might exhibit entirely different characteristics across diverse tasks. Our choices for are detailed Table 2
Policy extraction (Antmaze). We empirically find that deeper policy network improves over- all performance in Antmaze tasks (Figure 14). As a result, we employ a 4-layer MLP as the policy model. Additionally, adopted implicit Q-learning method sometimes has training instability issue in Antmaze tasks, resulting in highly divergent estimated Q-values (Figure 13). To stabilize training for policy ex- traction, we replace the temperature coefficient 3 with Snomm (8, a) = [V][(s.)]]] We sweep 3 E (0.01, 0.02, ..., 0.05) for umaze vironments, and B E (0.03, 0.04, ...,, 0.08} for other environments. Our choices for are de- tailed in Table 2 Other hyperparameters remain consistent with those the locomotion tasks.
Average Q-value battmaz--mazz- umaze diverse)
200
sced: 6
100
seed: 3
observe that the
1.0
Figure 13: The training instability issue.
Normalized scores
80
70
60-
2
5 4 (wio norm)
Policy network layers
Figure 14: Ablation for Antmaze tasks.
Evaluation. We all experiments
independent trials. For each trial, we additionally collect the evaluation score averaged across 20 test episodes at regular intervals for plots in Figure [66] The average performance at the end of training is reported in Table We use NVIDIA A40 GPUs for reporting computing results in Figure]
18

---

Published as a conference paper at ICLR 2024
Mediam-Expest
Modi
Medlimn-appay
1.0
Noma waw.alirdl...oo Scare
1,00
0.75
Nomaliat Scoes
1.0
Harper
Halkbootal
Wilkerd
0.5
0.56
Hoppor
Wolkiell
Walke20
Hopper
0.6
0.0
[talfchorach
0.08
0.11
0.02
0.16
0.1
0.2
0.
0.01
3.02
0.05
0.1
02
031
0.00
10.55
18.1
0.2
0.5
Animaze-Timaze
Antimas- Mediues
Animaze-1 aarge
Natvviiod Scotes
Nazalleed Scon Sooe
9.6
08
9,4
Defaal: Divere
Flaz
0.2
Plar
Divene
Divene
6:05
Figure 15: Ablation of the temperature coefficient / in D4RL benchmarks.
Iitomotion-Medium--ae Walker2d
Halfcheetah
Hopper
0.1
0.01
0.01
Walker2d 0.05
Halfcheetah 0.2
Hopper 0.05
Locomotion-Menn
Walker2d 0.5
Halfcheetah 0.2
Hopper 0.2
Locomotion-Mcdium-Reee
Umaze
Medium 0.08
Large 0.06
AntMaze-Fixed
0.02
Umaze
Medium
Large
AntMaze-Diverse
0.04
0.05
0.05
Table 2: Temperature coefficient / for every individual task.
19

---

Published as a conference paper at ICLR 2024
D TRAINING CURVES FOR OFFLINE REINFORCEMENT LEARNING
balfeheetab medium expert
hopper medium expert
walkerdd medium expert
120
80-
Nomallld Saroe
Nattlledd Seoole
1000
23
80
40-
60-
Doet
Ouet
TQL
[CC
Dettamion-Qe
Drtiumsn-QU
Pelicy Stape (k)
Pelicy Crandenn Stape
Gradien Stape
ballcheztah-m-diinn
welle:24-mediim
Nittllaat Suees
Nuzalaal Ruems
Necolood
60
60-
IQL
Difluoee Qf
Uilleaize-QU
Dillaaion-
3300
1200
2006
20
353
1000
1200
2006
1500
2000
Polacy Gradent Steps (1)
Police Crankenn Stape S41
Police Grudina Stope (41
balfchermhh medium roplay
hopper medium ceplay
Walker2d medium replay
100
60
1300
1300
2006
500
1600
1200
2066
1000
1500
2000
Policy Steps
Poñiey
Suge
100
63-
Notrated Sooms Scens
Nottafadd Scanes
60-
%
20
Ouas
Oura
Oura
IQU.
1998
1500
Poliay Gleudens
Poblcy Graalicn
Policy Cradicnt Skezn Mi)
Biverae
30-
66-
Nomalaed Saves
Nonsted Swers
QF.
Nomasted
40.
Cuara IOS: Dilftaion-Q
IQL
900
753
Erio
1299
1998
250
760
1800
1250
14000
260
110
1000
1250
1550
Pollay Ceacions
Polley Cradiee Shage
Pollay Cindlira Whipe-00
Figure 16: Training curves of SRPO (ours) and several baselines. Scores normalized according to Fu et al. (2020).
20


Reviews:
[{"review_id": "UNtvxW5VVc", "reviewer": "Reviewer_bNEQ", "summary": "This paper introduces a score regularized policy optimization algorithm (SRPO) based on behavior diffusion, aiming to address the problem of heterogeneous behavior distributions in offline reinforcement learning. The algorithm utilizes a pre-trained behavior diffusion model to score-normalize policy gradients and presents a practical method based on behavior diffusion and implicit Q-learning. The paper provides a detailed explanation of SRPO's principles, implementation, and optimization techniques, and it validates the effectiveness of the algorithm through experiments.", "questions": "1. DIQL should be IDQL in the first and second paragraph of section 6.1.\n2. Why do different Gaussian-based policies have significantly varying inference times, as shown in Figure 1?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "In the offline reinforcement learning setting, the paper introduces a novel approach that leverages the diffusion model. This method uses the powerful modeling capabilities of the diffusion model while avoiding the extensive time-consuming iterative inference stage.", "weaknesses": "The final policy used by the algorithm is still based on a Gaussian distribution. This Gaussian policy might not capture complex distributions as effectively as the diffusion model when dealing with complex offline datasets. The key concern here is whether the complex distribution information modeled by the pretrained diffusion behavior can be adequately captured by a policy based on a Gaussian distribution.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "U3GQvgP6Ow", "reviewer": "Reviewer_fYZE", "summary": "This paper proposes the use of pre-trained diffusion model for the behavior regularized offline policy optimization objective. The work shows a careful derivation of how pre-trained diffusion models, using the score function, can be replaced in the existing behavior regularized objective. The trick is to use the behavior policy distribution’s score function, and the paper claims to show faster compute time compared to existing approaches. Experimental results are demonstrated using the D4RL benchmark.", "questions": "See questions in the weakness section. \n\nThe key algorithmic pipeline is to integrate the SRPO technique with implicit Q learning. I wonder what happens if the SRPO objective is used in other behavior regularized offline rl objectives? Since the key algorithmic novelty comes through the derivation, it would be helpful if the authors can do more thorough experimental evaluation. Otherwise, the novel contributions of the paper are unclear.", "limitations": "", "rating": 3, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "Algorithmically, the paper provides an interesting insight showing that in behavior regularized policy optimization objective, the gradient of the diverse term is indeed related to the score function of the behavior policy distribution. This therefore allows the use of pre-trained diffusion models to be used in these objectives.\n\nThe challenge of measuring the divergence term in offline regularized objective is generally difficult, where typically a separate model is needed to approximate the behavior model. Equation (9) provides a simple trick following equation (8) to show exactly where pre-trained diffusion models can be used for the behavior distribution, where the diffusion model e(a|s, t) can approximate the grad log term of the behavior policy. This is an interesting and novel insight to derive the algorithm, making use of the widely available pre-trained diffusion models these days.", "weaknesses": "The paper is a bit hard to follow; while the claims are justified, the paper is not so well written and seems convoluted. I believe this is also because the key idea/trick of the paper is to use pre-trained diffusion models in existing offline rl objectives, so the paper tries to lay out the context for that. However, it makes the paper rather difficult to follow, to completely understand the full contribution of the work. \n\n\nSince the key idea is to use existing pre-trained diffusion models, I expected that other than the algorithmic contribution, the paper can do a much thorough job at more experimental evaluations? It would be useful to see existing setups where behavior regularized policy optimization is used, including toy examples, and perhaps provide a comprehensive qualitative study of the use of different pre-trained diffusion models in this context?\n\nExperimental results probably need to be more thorough; there are only marginal benefits from the SRPO objective and it is not clear whether the proposed approach leads to empirical benefits. It would be helpful to do more qualitative studies on the objective and the use of different pre-trained diffusion models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "384fAwl08y", "reviewer": "Reviewer_cVKT", "summary": "The paper addresses the slow inference issue of diffusion policies and proposes an offline RL algorithm to avoid iterative diffusion sampling process during policy evaluation by leveraging a critic model and a pretrained diffusion model. \nThe proposed method, SRPO, consists of three components: a critic model realized by the IQL, a pretrained diffusion model to explicitly model the behavior policy, and a policy extraction module utilizing the pretrained behavior model to regularize the policy gradient.\nExperimental results on D4RL benchmark suggest the proposed method can achieve higher computational efficiency especially during inference compared to other diffusion-based policies, while maintaining comparable performance by exploiting the modeling expressiveness of diffusion models.\nThe main contribution of this paper is to employ a pretrained diffusion behavior model, which can approximate the score function of the behavior distribution in offline datasets, to regularize the policy gradient during the optimization of the actor.", "questions": "1. In Table 1, where do the results of baselines come from? Are they reported from original papers or all have been re-implemented for comparison？\n2. In section 6.2, the paper claims that SRPO \"completely avoids diffusion sampling throughout both training and evaluation procedures\". But it is considered that the pretrained diffusion model for behavior policy still requires iterative sampling during the pretraining phase, which is also visualized in Fig 5 when compared to IDQL. This claim needs further explanation.\n3. In Fig 5, are the results averaged across all locomotion tasks?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 2, "presentation": 4, "contribution": 3, "strengths": "1. The paper clearly states its motivation and presents a clear illustration to demonstrate the derivation of the proposed method, SRPO.\n2. The paper provides reproducible details for its experiments, and make a relatively comprehensive comparison with both conventional behavior regularization methods and recent diffusion-based policies in offline RL, in terms of task performance and computational efficiency.", "weaknesses": "1. The paper mainly aims to improve the computational efficiency of diffusion-based polices, which is highlighted in computation-sensitive contexts such as robotics as stated in the paper, yet there is no experiment concerning the robot scenarios especially with real data. If such experimental results are provided, the central claim made by the paper can be more convincing.\n2. The novelty of the proposed method is limited as it seems to be a combination of previous work, and especially an incremental work based on IDQL (Hansen-Estruch et al., 2023), though it has been compared in the experiments. The main difference between the two work is the policy extraction process, which is the main contribution of this paper but does not provide adequate contribution to the community.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7832Yekkn7", "reviewer": "Reviewer_f3qQ", "summary": "This paper aims to improve the inference efficiency of diffusion models in the context of offline reinforcement learning, by proposing a deterministic inference policy from critic models and pretrained diffusion behavior models. The pretrained diffusion behavior models are then utilized to regularize the policy gradient. Experiments are conducted on D4RL in locomotion tasks, the speed is improved without an evident performance drop.", "questions": "Please see the weaknesses for details. Overall, I think I get confused by the distributions from the DMs and the actual deployment in the context of RL in this work, as well as the rationale/motivation behind it.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The research problem of improving inference efficiency for diffusion models in offline reinforcement learning is interesting and worth investigating.\n\n- The paper is well organized and written in general, with motivations and methods well explained (but somewhat unclear to me).\n\n- The experiments show evident improvement in time cost and demonstrate the effectiveness.", "weaknesses": "- Some implementation details are unclear, while the Fig.1 shows the comparison of the inference efficiency, the actual time cost for training is not mentioned. (In Appendix C, it seems the actual time used for training is not specified?)\n\n- I get a bit confused by some high-level assumptions in this work. Essentially, I believe there are two types of distributions that matter in the context of Diffusion models for offline RL, namely the actual data distribution and the learnable Gaussian kernel (noise) distribution (which is pre-scheduled on the mean values with fixed variance). If the actual learnable $\\pi_\\theta$ is used as the appropriation/predictor of the policy distribution (which I believe is the case according to Eq. 9?), then what is the difference between this proposed paradigm and the Gaussian case, while the latter is believed to be lacking in expressivity, as mentioned by the authors in the intro and Fig.2, however the $\\pi_\\theta$ that you are sampling from is also a Gaussian with known variance and learnable mean value.\n\n- Following my previous point, which I think is also related to the choice of $t$ in Sec. 4.2 and also in Eq. 9. If I misunderstood the actual usage of distribution here, and the proposed SRPO method actually uses the actual distribution, which corresponds to $x_0$ (or somewhere close to $x_0$), as the optimal policy that you are trying to sample from, then the expressivity concern does not exist, but then in this case, the entropy term in Eq. 9 does not hold, because we don’t have any prior on the actual distribution of $x_0$ also on its variance?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)