Paper: Score Regularized Policy Optimization through Diffusion Behavior (ICLR-2024)

Reviewer comment (Reproducibility):
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?)

Author response:
**A1:** The actual time cost for training was actually provided in Section 6.2 of the paper (Figure 7), where we compared training and evaluation time between our work and several baselines in detail. We present Figure 7 in tabular form below:

| Algorithm    | Training Time (h)  | Evaluation Time (h) | Total Time (h)  | Evaluation/Total (\%)  |
|--------------|:-----------------:|:-------------------:|:-------------------:|:-----------------:|
| IQL          | 4.36            | 0.40               | 4.76                 | 8.4            |
| DiffusionQL  | 14.3            | 8.66               | 22.96                 | 37.7            |
| IDQL         | 6.57            | 16.9               | 23.47                |  72.0            |
| SRPO(Ours)   | 9.32            | 0.38               | 9.7                |  3.9            |

(Suppose each experiment requires 2,000K gradient steps for training, 50(times)*20(seeds) episodes of evaluation, test on Nvidia A40 GPUs)

Reviewer comment (Theory):
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.

Author response:
**A2:** We appreciate the opportunity to clarify the confusion, primarily regarding the definitions of $\pi_\theta(s)$ and $\pi_{\theta, t}(\cdot|s)$. **We have updated the manuscript, adding Figure 5 and Figure 10 to help better illustrate our idea.** Below is a summarization of the main distributions from the Diffusion Models:


|   Diffusion time  |  &emsp;     $t=0$                       |       | &emsp;&emsp;&emsp;&emsp; $t \in (0,1)$  |
|:--------:|:---------:|:-----:|:-----:|
|                       |           |         diffusion (add noise)       |               |
| behavior data:  |  $ \mu(a\|s)   $ | ----------------------------------> | $\mu_t(a_t\|s)$ | 
|                       |           |                |          ↓     |               | 
|  *diffusion''modeling*    |           |               |         supervise(Eq.13)      |   
|                       |           |                |         ↓      |               | 
|behavior model:  |  $ \epsilon_\psi(a\|s)   $ |  | $\epsilon_\psi(a_t\|s,t)$ |
|                       |↓  |         <--------------  +$Q_\phi(s,a)$  **-------------->**   |        **↓**      |               | 
|  *score''regularization*    |        supervise(Eq.10)   |               |         **supervise(Eq.15)**      |   
|                       |         ↓  |                |         ↓      |               | 
| policy model:  |  $ \pi_\theta(s)   $ | **------------(define by diffusion)---------->** | **$\pi_{\theta, t}(\cdot\|s) $:=$\mathcal{N}(\cdot\|\alpha\_t \pi_\theta(s),\sigma_t^2 I ) $** |
|                       |           |      **<---------------(supervise)---------------**           |               |               | 


(The method proposed in Section 4.2 of the paper is highlighted in **boldface**.)

- $\mu(a|s)$: The underlying behavior policy for behavior dataset $D^{\mu}$.
- $\mu_t(a|s)$: The noisy version of $\mu(a|s)$, defined by the forward diffusion process (Eq.4 in paper).
- $\epsilon\_\psi(a\_t|s,t)$: The diffusion behavior model learned to represent $\nabla_a \log \mu\_t(a|s)$ for a series of $t\in(0,1)$. We need $\epsilon\_\psi$ to regularize gradient when optimizing policy $\pi\_\theta$. 
- $\pi_\theta(s)$: A deterministic policy which is actually **used for evaluation**.
- $\pi\_{\theta, t}(\cdot|s)$: The noisy (diffused) version of $\pi\_\theta(s)$. **Note that** $\pi\_{\theta, t}(\cdot|s)$ **is defined by** $\pi\_\theta(s)$. Its mean is $\alpha\_t \pi\_\theta(s)$. It has fixed variance. As a result, **when we optimize $\pi\_{\theta, t}(\cdot|s)$, we are essentially optimizing the underlying evaluation policy $\pi\_\theta(s)$.** $\pi\_{\theta, t}(\cdot|s)$ is **not** used for evaluation.


**Why optimize $\pi\_{\theta, t}(\cdot|s)$ as in Eq.15? Why not just directly optimize $\pi\_\theta(s)$ as in Eq.10?** (since we do not actually use $\pi\_{\theta, t}(\cdot|s)$)

1. Our goal is to leverage the pretrained diffusion model $\epsilon_\psi(a_t|s,t)$ at every $t \in (0,1)$, not just at $t=0$ (Eq. 15). Proposition 1 in our paper (proof in Appendix B) demonstrates the equivalence between minimizing $KL(\pi_\theta(s) || \mu(a|s))$ and minimizing $KL(\pi_{\theta, t}(\cdot|s) || \mu_t(a|s))$.
2. When $t$ is always 0, Eq.15 and Eq.10 become identical, making Eq.10 a special case of Eq.15.
3. Utilizing pretrained diffusion models across all $t \in (0,1)$ empirically improves performance. See Section 6.3 for ablations, and see Figure 5,10 for illustrations.

Reviewer comment (Theory):
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?

Author response:
**A5:** As clarified in **A4**, our evaluation policy is either deterministic or Gaussian, so the reviewer's understanding of its usage is correct. Therefore, the entropy term in Eq. 9 remains valid.

Reviewer comment (Reproducibility):
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？

Author response:
**Q4: In Table 1, where do the results of baselines come from?**

**A4** They all come from the original paper except for BEAR, which we refer to D4RL benchmark's white paper due to unavailability in the original paper.

Reviewer comment (Experiments):
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.

Author response:
**Q5: It is considered that the pretrained diffusion model for behavior policy still requires iterative sampling during the pretraining phase which is visualized in Fig 5 when compared to IDQL.**

**A5:** Neither IDQL nor our method requires iterative sampling during the pretraining phase. 

Reference code for IDQL:  https://github.com/philippe-eecs/IDQL/blob/6b071e7a1eea57bd93956f2b0a9e07a68081a096/jaxrl5/agents/ddpm_iql/ddpm_iql_learner.py#L426

The pipeline (simplified) is 

1. Sample $s, a_0$ from dataset, $t$ from $(0,1)$, Gaussian noise $\epsilon$

2. Add noise to $a\_0$ based on $t$ and $\epsilon$ to form $a\_t$

3. Optimize MSE loss between $\epsilon$ and $\epsilon\_\theta(a\_t)$

Reviewer comment (Presentation):
In Fig 5, are the results averaged across all locomotion tasks?

Author response:
**Q6: In Fig 5, are the results averaged across all locomotion tasks?**

**A6** Yes, they are. 

Note that for hopper and walker2d environments, the episode may terminate in less than 1000 steps during evaluation (e.g., the robot falls to the ground). This causes too much randomness for evaluation, so we manually normalize the evaluation time hypothesizing that all episodes have a length of 1000.

Reviewer comment (Theory):
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.

Author response:
We partially agree with the reviewer. In our opinion, diffusion models are beneficial for offline RL for mainly two reasons:

1. The behavior dataset $\mu(a|s)$ is complex.

2. The optimal policy $\pi^* \propto  \mu(a|s) \ e^{\beta Q(s, a)}$ is potentially complex (multi-modal).

If referring to problem 1, our proposed method captures complex behavior distributions effectively and **efficiently** by utilizing a diffusion behavior model. No Gaussian model is used.

If referring to problem 2, we use a Dirac (gaussian) policy to capture the mode of $\pi^*$:

$\pi^{\text{gauss}}\_{\theta} (s) = \arg \max\_a \mu(a|s) \ e^{\beta Q(s, a)} $ instead of $\pi_\theta^{\text{diff}}(a|s) \propto \mu(a|s) \ e^{\beta Q(s, a)}$.

**We agree with the reviewer that the Gaussian policy cannot capture complex $\pi^{\*}$. This is simply not our intention.** 
If there are two optimal policies for one single task, our method will simply choose either one of them deterministically. Our observation is that in many scenarios, inference speed is of much higher priority than policy diversity. For instance, in robotics, we often do not care about how many ways our robot can solve a task. Contrastively, we do care about whether our policy is optimal and efficient (hardware-friendly). Also, current evaluation metrics of the D4RL benchmark only emphasize policy optimality and ignore policy diversity.

Still, we believe a better trade-off between inference speed and policy diversity exists and could be a future direction of our work.

Reviewer comment (Theory):
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.

Author response:
We really appreciate this insightful comment. 

Theoretically, since SRPO can be thought of as performing gradient descent in the action space: $\arg \max\_a [\log \mu(a|s) + \beta Q(s, a)]$, it is possible that the learned policy gets stuck in a local optimal and thus fails to capture the mode of $\pi^*$. 

To address this problem, our work proposes to ensemble diffusion score estimators at diffusion times $t\in (0,1)$ to make the gradient landscape smoother. This part is explained in section 4.2 and analyzed in Appendix B of the paper. **We have newly added an illustration in Figure 5 and Figure 10 in the revisited paper.** Ablation studies in Section 6.3 also support the effectiveness of our proposed method. 

**Empirically, we feel confident about SRPO's ability to capture complex distribution information modeled by the pre-trained diffusion in offline RL.** Besides reported evaluation numbers, our confidence also comes from DreamFusion (a text-to-3D research winning the Outstanding Paper Award of ICLR2023), which inspires our work. DreamFusion seeks to learn a Dirac 3D model from a pre-trained large-scale image diffusion model. Their method similarly adopts a score distillation method to supervise the training of another Dirac (Gaussian) 3D model. Experimental results show that Gaussian-based policy can already capture complex and high dimensional distribution successfully **even in complex image space** (https://dreamfusion3d.github.io/).

Reviewer comment (Writing):
DIQL should be IDQL in the first and second paragraph of section 6.1.

Author response:
We sincerely thank the reviewer for the careful reading. The typo is fixed now.

Reviewer comment (Evaluation):
Why do different Gaussian-based policies have significantly varying inference times, as shown in Figure 1?

Author response:
Actually inference times for Gaussian-based policies are very similar. They only look that way because the x-axis is in the log scale in Figure 1.

|                | Ours | BEAR | IQL  | CQL  | AWAC | BCQ | DT   | IDQL  | Diffuser |
|:---------------|:-----|:-----|:-----|:-----|:-----|:----|:-----|:------|:---------|
| Inference time | 1.28 | 1.28 | 1.43 | 2.32 | 2.86 | 9.0 | 31.6 | 60.83 | 2136.4   |

ps. AWAC takes a longer time for action sampling simply because it is a 4-layer MLP, while ours is a 2-layer MLP.