Behavior-Regularized Diffusion Policy Optimization for Offline Reinforcement Learning

We sincerely appreciate your insightful feedback, which has greatly improved our work. We hope the following clarifications can further address your concerns and enhance your evaluation of our paper.

Q1: About the proof in Theorem C.1 First, we would like to restate the question to address any misunderstandings. Let $Q(s, \cdot)$ denote the ending reward in the diffusion MDP with a discount factor of 1. The reviewer considers the following relationships between the value functions:

V^{\star, s}\_0&=Q(s, a^0)\\\\ Q^{\star, s}\_0&=Q(s, a^0)\\\\ V^{\star, s}\_n(a^n) &= \eta\log\mathbb{E}\_{a^{n-1}}\left[\exp(Q^{\star, s}\_n(a^n)/\eta)\right]\\\\ Q^{\star, s}\_n(a^n)&=0 + 1.0\mathbb{E}\_{a^{n-1}}\left[V^{*, s}\_{n-1}(a^{n-1})\right] \end{aligned}$$ which seems to contradict our results in line 820. However, we believe the main confusion here is that in the diffusion MDP, the “state” is actually the tuple $(s, a^{n})$ and the state transition is implicit. The policy follows the Gaussian distribution $p^{\pi,s}(a^{n-1}|a^n)$, and upon selecting action $a^{n-1}$, the state instantly transitions to $(s, a^{n-1})$ deterministically. Therefore, the $Q$ function of the diffusion MDP should be defined w.r.t the “state” $(s, a^n)$ and action $a^{n-1}$, leading to the corrected relationships: $$\begin{aligned} V^{\star, s}\_0(a^0)&=Q(s, a^0)\\\\ V^{\star, s}\_n(a^n) &= \eta\log\mathbb{E}\_{a^{n-1}}\left[\exp(Q^{\star, s}\_n(a^n, a^{n-1})/\eta)\right]\\\\ Q^{\star, s}\_{n}(a^n, a^{n-1})&=0 + 1.0V^{\star, s}\_{n-1}(a^{n-1}) \end{aligned}$$ Substituting the last equation into the second leads to the formulation in line 820 $$V^{\star, s}\_n(a^n) = \eta\log\mathbb{E}\_{a^{n-1}}\left[\exp(V^{\star, s}\_{n-1}(a^{n-1})/\eta)\right]。$$ As for question (2), expanding the recursion gives $$\begin{aligned} V^{\star, s}\_n(a^n) &= \eta\log\mathbb{E}\_{a^{n-1}\sim p^{\nu,s}\_{n-1|n}}\left[\exp(V^{\star, s}_{n-1}(a^{n-1})/\eta)\right]\\\\ &= \eta\log\mathbb{E}\_{a^{n-1}\sim p^{\nu,s}\_{n-1|n},a^{n-2}\sim p^{\nu,s}\_{n-2|n-1}}\left[\exp(V^{\star, s}\_{n-2}(a^{n-2})/\eta)\right]\\\\ &=\ldots\\\\ &=\eta\log\mathbb{E}\_{a^{n-1}\sim p^{\nu,s}\_{n-1|n}, \ldots, a^0\sim p^{\nu,s}\_{0|1}}\left[\exp(V^{\star, s}\_{0}(a^{0})/\eta)\right]\\\\ &=\eta\log\int p^{\nu,s}\_{0,1,\ldots,n-1|n}(a^0, a^1, \ldots, a^{n-1}|a^n)\exp(V^{\star, s}\_0(a^0))\mathrm{d} a^{n-1}a^{n-2}\ldots a^{0}\\\\ &=\eta\log\int p^{\nu,s}\_{0|n}(a^0|a^n)\exp(V^{\star, s}\_0(a^0))\mathrm{d}a^{0}\\\\ &=\eta\log\mathbb{E}\_{a^{0}\sim p^{\nu,s}\_{0|n}}\left[\exp(V^{\star, s}\_{0}(a^0)/\eta)\right] \end{aligned}$$ where the last but one equation is due to the marginalization over intermediate actions $a^1, a^2, \ldots, a^{n-1}$. We acknowledge that the derivation in the appendix is somewhat vague and will revise it to explicitly incorporate a proof akin to MaxEnt RL theory for clarity. **Q2: About the update of value networks.** The detailed update procedure of the value networks consists of two steps: 1 ) The first step is updating $Q^{\pi}$. To calculate the target, we use the actor diffusion policy to generate paths $a’^{0:N}$ at the next state $s’$, calculate the target value $Q(s’, a’)$ and accumulated penalties along the path $\sum_{n=1}^N\ell^{\pi,s'}_{n}(a^n)$, and perform temporal difference update as per Eq. 12; 2 ) The second step is updating $V^{\pi,s}$. This is achieved by sampling $n$ and $a^n$ using forward process, performing one-step diffusion to obtain $a^{n-1}$, and update according to Eq. 14. For $n=1$, we directly use $Q(s, a^0)$ as the target, rather than additionally regressing the output of $V^{s}_0(\cdot)$ to $Q(s, \cdot)$. Crucially, the update of $Q^{\pi}$ does not depend on $V^{\pi,s}_N$. The diagram of the update is presented in Figure 3. Compared to other methods, the additional cost w.r.t. value function training comes from the second step, which is a constant that does not scale with diffusion steps (see Figure 10). **Q3: Eq. 7 should be $q\_{n-1|n}(x^{n-1}|x^n,x^0)$** We appreciate the suggestion and would like to note that the actual distribution we wish to approximate is the posterior without conditioning on $a^0$, as $a^0$ is unknown during generation. When $a^0$ is given, the posterior distribution $q\_{n-1|n, 0}$ is tractable, with its analytical form being: $$q\_{n|n-1, 0}(a^n|a^{n-1}, a^0)=\mathcal{N}(a^{t-1}; \frac{\sqrt{\bar{\alpha}\_{n-1}}\beta_n}{1-\bar{\alpha}\_n}a^0+\frac{\sqrt{\alpha\_n}(1-\bar{\alpha}\_{n-1})}{(1-\bar{\alpha}\_n)}a^n, \sigma\_n I).$$ Since exact $q\_{n-1|n}$ is intractable, we turn to a parameterized distribution $p^\theta\_{n-1|n}$ and optimize it towards $q\_{n-1|n, 0}$ via in Eq. 10. As shown in DDPM, this training objective yields $p^\theta\_{n-1|n} \approx q\_{n|n-1}$. We recognize that the current presentation lacks clarity and will revise the text to explicitly articulate the connection between $q\_{n-1|n, 0}$ and $q\_{n-1|n, 0}$ in the next version.

Thank you for providing valuable feedback. Below, we provide further clarification and results to address your concern and we hope these materials can enhance your evaluation of our paper. Due to the space constraint, we will post our discussion of the connection between BDPO and broader literatures (Q6 and Q8) in follow-up responses. You can also access the responses in the anonymous link: https://f005.backblazeb2.com/file/bdpo-review/Q6-andQ8.txt .

Q1: Can an optimal policy be delivered by optimizing eq.11?

We would like to clarify that Eq. 1 is also a well-established RL objective widely adopted in the literature, and therefore, optimizing Eq. 11 delivers optimal policies with respect to Eq. 1. Please refer to our response to Q5 for details.

Q2: The pathwise KL regularization is stronger, making the policy suffer from sub-optimal data?

Yes, BDPO additionally regularizes intermediate actions along the diffusion path. However, we emphasize that this is not a stronger regularization, as we have established the equivalence between pathwise and action-wise constraint in Theorem 4.2. Given this guarantee, the pathwise constraint is in fact preferable as it enables a finer-grained control over the generation process. Furthermore, when the behavior policy is suboptimal, the regularization strength $\eta$ can be decreased to allow greater exploitation from the value function.

Q3: Intermediate actions are not present in the dataset.

When training value functions $V^{\pi, s}\_n$ for $n>0$, we optimize them using actions $a^n$ generated by first sampling clean actions $a^0$ from the dataset and then perturbing them using the forward diffusion process $a^n\sim q\_{n|0}(a^0)$. Consequently, the action support for $V^{\pi,s}\_{n}$ with $n > 0$ is infinite and spans the entire action space due to the unbounded support of the Gaussian noise distribution. The only risk of OOD evaluation comes from querying $V^{\pi,s}_0(a^0)$ for some $a^0$ generated by the actor. However, this challenge is inherent to all policy iteration methods and can be effectively mitigated by tuning the regularization strength.

Q4: The assumption about concentrability is stronger?

We clarify that Assumption 4.3 is made only to ensure the boundedness of the KL divergence in our theory, and it is not the concentrability used in papers like the reviewer mentioned. Similar assumptions are also made in papers that incorporate behavior regularization -- for example, SAC assumes $|\mathcal{A}| < \infty$ to ensure the policy entropy is bounded.

Regarding the concentrability, as discussed in Q3, the marginal distribution $p^{\nu,s}_n$ at $n>0$ has infinite support over the action space due to Gaussian perturbations. Consequently, the concentrability requirement reduces to the base case of $n=0$, meaning BDPO does not impose stronger concentrability assumptions than other methods. Our primary contribution is providing a practical implementation of the behavior-regularization RL for diffusion-based policies, rather than introducing new theoretical bounds.

Q5: The optimal policy differs from other works.

The optimal policy differs because we consider the behavior-regularized RL framework (Eq. 1), which augments the standard RL objective with KL divergence. This framework helps us to shape the policy and is widely adopted in RL. For example, when $\nu$ is specified as the uniform distribution or the dataset collection policy, the framework becomes MaxEnt RL (used by SAC) or regularized RL (used by offline RL methods like ReBRAC and XQL), respectively. Other offline RL algorithms (e.g., TD3-BC, AWAC) can also be categorized into this framework if we omit the KL divergence term during policy evaluation. This framework is also adopted by applications like RLHF, where KL regularization towards the reference model prevents model collapse. Therefore, we believe it is important to study this framework.

Q6: Lack of theoretical comparison with other diffusion-based works.

See follow-up responses or the link.

Q7: Experiment results on random and expert datasets

We compare BDPO against several baselines, including a previous model-free SOTA, ReBRAC and the best performing diffusion-based method, DAC. All results are averaged over 4 seeds.

Overall, BDPO matches SOTA performance (ReBRAC) and outperforms DAC, though the margin is narrow. Note that model-free offline RL methods tend to struggle on random datasets due to the inferior data quality.

Q8: Lack of references about imitation learning

See follow-up responses or the link.

We would like to thank the reviewer for his/her constructive feedback. Below is further clarification for the reviewer's concern.

Q1: Using EDP for policy optimization?

Our policy improvement objective is to maximize the expected Q-values while also minimizing the KL divergence: $\max\_{p^{s,\pi}}\ \mathbb{E}\_{a^0\sim p^{s,\pi}\_0}[Q(s, a^0)]-\eta \mathrm{KL}\left[p^{s,\pi}\_{0:N}\|p^{s,\nu}\_{0:N}\right]$ Let us ignore the second KL term for now, since EDP implements this constraint using diffusion loss. For the first term, in order to circumvent backpropagating the gradient of Q through the diffusion path, EDP introduces the following action approximation: $\hat{a} = \frac{1}{\sqrt{\bar{\alpha}\_n}}a^n - \frac{\sqrt{1-\bar{\alpha}^n}}{\sqrt{\bar{\alpha}^n}}\epsilon^{\pi,s}(a^n, n)$ where $\epsilon^{\pi,s}$ is the output of the score network. Afterwards, they use $\hat{a}$ as the approximation for $\mathbb{E}\_{a^0\sim p^{s,\pi}\_{0|n}}[Q(s, a^0)]$. However, we emphasize that this approximation is inexact and biased. The action approximation is essentially the mean of the posterior distribution: $\hat{a}=\mathbb{E}\_{a^0\sim p^{\pi,s}\_{0|n}}\left[a^0\right]$ meaning that they are using $Q(s, \mathbb{E}\_{a^0\sim p^{\pi,s}\_{0|n}}\left[a^0\right])$ to approximate $\mathbb{E}\_{a^0\sim p^{s,\pi}\_{0|n}}[Q(s, a^0)]$, which is biased and inconsistent with our theory.

Q2: How to select the action with the highest Q value? What about the inference time?

Yes, we use the average of the ensemble Q networks to select the actions. During inference, BDPO first generates $N_a=10$ candidate actions in parallel, calculates their $Q$-values and selects the best one. The following table presents the inference latency per state, averaged over 100K trials:

We found that the inference cost of BDPO is comparable to DAC, which also generates 10 actions and selects the best one with the highest Q-value. The JAX implementation of Diffusion-QL is faster than BDPO, since it only generates one action and does not query the Q-values. The PyTorch version of Diffusion-QL is much slower. For DTQL, since its policy network is simply a one-step policy, its inference time is comparable to BDPO. Finally, QGPO requires taking the gradient of the Q-value network to calculate the guidance, which results in heavy computation overhead during inference.

Q3: About hyperparameters

• The range of parameter sweeping for eta and rho?

  For locomotion tasks, we swept eta and rho among {0.5, 1.0, 1.5, 2.0}. For antmaze tasks, we swept rho among {0.5, 0.8, 1.0} and eta among {1.0, 5.0, 10.0}.

• The different trend w.r.t. the hyper-parameters in halfcheetah and walker2d tasks, and what about other datasets like medium-replay?

  The sensitivity analysis of eta and rho for medium-replay tasks is provided in the following anonymous links: https://f005.backblazeb2.com/file/bdpo-review/rebuttal_abla_eta.pdf and https://f005.backblazeb2.com/file/bdpo-review/rebuttal_abla_rho.pdf . Overall we found that the trend is similar across most of the datasets, that is, excessively large or small $\rho$ and $\eta$ may result in fluctuation or degradation in the performance. The halfcheetah-medium-v2 dataset seems to be an exception in that it is more tolerant to extreme parameter values.

• Setting $\eta=0$ or $\rho=0$?
  For results with $\rho=0$, please refer to the above link to ablation w.r.t. rho. For results with $\eta=0$, we additionally provide curves on medium-expert datasets, and the results are provided in the anonymous link: https://f005.backblazeb2.com/file/bdpo-review/rebuttal_eta0.pdf . The lower confidence bound technique is conceptually similar to the commonly adopted double Q-network trick that penalizes the Q-values of OOD actions, and therefore when there is no LCB penalty, the performance drops sharply due to severe over-estimation. Setting $\eta=0$ also results in performance fluctuation and decrease for walker datasets due to insufficient constraint. However, in halfcheetah datasets, $\eta=0$ improves performance, likely because LCB already penalizes out-of-distribution (OOD) actions effectively.

• About the sensitivity to rho and advice on choosing hyper-parameters?
  After inspecting the training details, we observed that an excessively small $\rho$ leads to severe over-estimation in Q-values, while an excessively large $\rho$ causes severe underestimation instead, both of which result in performance degradation. We found that a $\rho$ value yielding stable value estimations generally correlates with strong performance. Thus, we recommend to first adjust $\rho$ until stable value estimations are achieved, and then gradually decrease $\eta$ to strike a balance between robustness and performance.