Diffusion Policies creating a Trust Region for Offline Reinforcement Learning

Summary

The paper proposes a new diffusion-based training loss for offline reinforcement learning, in which the empirical actions from the offline datasets are replaced with the actions sampled from a parameterized policy. The paper calls this new loss “trust region loss” and uses it as one of the objectives to minimize in the diffusion policy learning. The paper points out the mode seeking behavior induced by the minimization of such “trust region loss” and compares it with forward KL-regularization behavior cloning. The paper also evaluates the methods on D4RL datasets with some offline RL methods.

Strengths

### One-step policy Introducing a one-step policy seems a promising way to address the challenges in the diffusion sampling. This idea is nicely motivated and well-justified. The discussion on the mode-seeking behaviors is interesting, especially with the empirical results as an extra evidence to support.

Weaknesses

I found the paper shares great similarities to the SRPO paper, in terms of the idea, the formulation, and the presentations. The empirical results of the proposed method are also close to SRPO method. Beyond that, I have several other concerns: ### **Unknown trust region** 1. A bit surprisingly, **the paper gives no explicit definition of the crucial trust region used in Eq. (4)**, the diffusion-based trust region loss, even though trust region is the key for the whole paper. Particularly, the loss function defined by Eq. (4) comes directly from a simple modification of Eq. (3) by replacing the empirical action $a_0$ with the action generated by the policy $\pi_\theta$. It is unclear what the trust region this simple modification induces. 2. Regarding “the loss effectively creates a trust region defined by the diffusion-based behavior cloning policy, within which the one-step policy $\pi_\theta$ can move freely”, the paper does not explain how “to generate actions (data) that reside in the high-density region of the data manifold specified by $\mu_\theta$ through optimizing $\theta$”. Note that $a_\theta$ generated by $\pi_\theta$ will be fed into $\mu_\phi$. **So a trivial solution of this optimization will be $\pi_\theta(\cdot|s) = a_0$, i.e., $\pi_\theta$ just returns $a_0$ for every $(a_0, s)\sim D$ as in Eq. (2). Then loss in Eq. (4) is exactly the same as Eq. (3)**. So how does the optimization of Eq. (4) avoid such trivial solutions? Especially when the paper says Eq. (4) is to encourage mode seeking behaviors. 3. **the “one-step policy”, another important factor in the trust region, is introduced a bit inconsistently section 2.3 and 2.4**. I understand that Eq. (4) is to introduce a policy that seeks to model only one mode, i.e., one action only, from the training samples. But then line 133-137, “the generated action $a_\theta$ appear similar to in-sample actions and penalizes those that differ”, which means the action should be similar to any of the training samples to avoid penalties. Also, this “mode-seeking” contradicts with the statement “explore freely” in line 149: how does this policy move freely? And how the penalization is triggered? In addition, the paper mentioned an implicit parameterization of $\pi_\theta$: $a_\theta=\pi_\theta(s, \epsilon$. For this implicit parameterization, what’s the different between optimization Eq. (4) and the simple behavior cloning? One even more inconsistent description is given in line 183 “one-step Implicit policy”: why $\pi_\theta$ here is instantiated implicitly? ### A bit trivial discussion on mode seeking behaviors **I found the discussion on mode seeking behaviors in Section 3 a bit trivial and pretty incremental as these have been pointed by SRPO already in its section 2.2**. 1. **the notations are not quite clear**: $p_{fake}$ and $p_{real}$ are not defined explicitly with respect to $\pi_\theta$ and $\mu_\phi$. 2. the forward KL divergence is mode-covering in that sense it learns a distribution that tries to cover all modes, while reverse KL divergence is mode-seeking, behaving similarly to the optimization of $L_{TR}$. This discussion has been provided by SRPO paper its equations (2) and (3). So what’s new? More importantly, **since the paper focuses more on the mode seeking behavior, why does it not compare against the reverse KL divergence minimization?** 3. **Theorem 1 tells very little about the mode seeking behavior of $\pi_\theta$ except the simple substitution of mode actions**. I don’t its importance in the training, especially the mode seeking behaviors. ### Hard to understand **There are quite many sentences and descriptions with language issues or typos, which make the whole paper rather hard to understand**. Examples of language issues and typos: * Line 192-193: “KL divergence is employed in this context, it is designed to recover all modalities of the data distribution” * Line 199: “the loss ensures that the generated action .$a_\theta$ lies within the in-sample datasets” * Line 288: “a Jacobian term” is different from what in Eq. (11). * Line 126, “minimize $L_{PB}$ in Equation 4”. There is no such $L_{PB}$ term. * Line 183, “is instantiates as an” **Importantly, the hyperparameters used in the experiments are a bit hard to understand**. Particularly, in the experiments, the paper mentioned that “the variation between datasets significantly impacts the algorithm performance”. Then why did the paper still “employed the official SRPO code on Antemze-v0 and maintained identical hyperparameters used for Antmaze-v2”? Why wouldn’t the paper just report the performance on both versions, separating DQL, IDQL from SRPO?

Questions

1. Can the authors explain which empirical results support that the assertion “diversity, while valued in image and 3D generation, is not essential in offline RL”? as speculated in Line 234 - 235. 2. Can the authors explain “discourage out-of-distribution sampling” in line 107? 3. What does it mean by “policy $\pi_\theta$ can move freely” in line 121?

Rating

4

Confidence

4

Soundness

1

Presentation

1

Contribution

2

Limitations

n/a

Summary

This paper introduces Diffusion Trusted Q-Learning (DTQL) for offline reinforcement learning. DTQL employs a dual-policy representation, a diffusion policy trained by behaviour cloning, and a one-step Gaussian policy trained by RL and distilling the diffusion model. Specifically, DTQL introduces a diffusion trust region loss for the distillation. Such a scheme allows both one-step generation, and achieves strong performance on the D4RL benchmark.

Strengths

1/ The proposed method is interesting. 2/ The proposed DTQL is well-motivated. It addresses one of the big issues of diffusion policy, i.e., the multi-step inference, through policy distillation. 3/ The paper has conducted solid experiments, for both intuitive understanding and empirical results on benchmarks.

Weaknesses

Overall I don’t see many weak points of this work. 1/ The distilled policy captures some certain modes, but is actually not multi-modal. It also missed certain modes as in Fig. 2, which is inconsistent with what the authors have claimed in the caption. I would suggest the authors to be careful and avoid over-claiming. 2/ The writing has certain room for improvement. The introduction could be reorganised for better clarity. There are some typos, e.g., in Line 183, “instantiates”, and inconsistent wording, e.g., “cooperative policy” and “dual policy”.

Questions

1/ It seems to me from Fig. 2 that this approach encourages interpolation of the distribution. Have you tried to design tasks to validate such behaviours? 2/ It is a bit unclear how to guarantee the distilled Gaussian policy captures all relevant information of the diffusion process, as its policy class naturally limits the expressiveness. Could you provide more explanations for this?

Rating

7

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

The experiments are done only on simple D4RL benchmarks. The policy distribution is quite simple and tasks are relatively short horizon. As a result, I’m not fully convinced when applied to more challenging scenarios, distilling the diffusion policy into a Gaussian with trust-region loss can still work given more complex action distributions.

Summary

The paper introduces Diffusion Trusted Q-Learning (DTQL), a novel approach for offline reinforcement learning that enhances performance and efficiency. DTQL utilizes a dual policy framework to bridge the gap between diffusion policies and one-step policies, improving expressiveness without the need for iterative denoising sampling. By incorporating a diffusion trust region loss, DTQL ensures stability and effectiveness in training. The method outperforms traditional distillation methods in various scenarios, offering a promising solution for offline RL challenges.

Strengths

* To my knowledge this paper presents a new approach to make Diffusion policy computationally efficient compared with previous method. * Based on the experimental results, it is evident that DTQL outperforms the previous method in terms of performance.

Weaknesses

* This is an incremental work that relies on Diffusion Q-learning (DQL) and adds a weighted loss to the diffusion model objectives(Equation 4) to accelerate model inference, without clearly explaining the necessity of this combination. * The explanation of the proposed method in Section 2 is too brief, with details only provided in subsection 2.3. Additionally, subsection 2.3 does not clearly explain the meaning of Trust Region in Equation 4. The subsequent Theorem and Remark do not effectively justify the necessity of using this update formula. The overall logic lacks coherence. * The introduction of other methods for accelerating training and inference in diffusion-based policy learning in Section 3 is overly lengthy and lacks relevance to the methods used in this paper. It is recommended to move most of this content to the appendix.

Questions

* How should the term "trust region" in the article's title and method name be understood? * How is the weight function $w(t)$ in Equation 4 set in this paper, and why is it configured in this manner? A detailed discussion is needed. * How does DTQL eliminate the need for iterative denoising sampling during both training and inference, making it computationally efficient?

Rating

4

Confidence

3

Soundness

2

Presentation

2

Contribution

2

Limitations

* The writing of this paper has issues—there is a lack of coherence between the motivation, method, and experiments. Section 2 provides insufficient detail on the proposed method and does not effectively explain why it improves computational efficiency. Additionally, Section 3 dedicates too much space to the comparison experiments replacing $\mathcal{L}_{\mathcal{TR}}$, lacking a comparative analysis with the baseline method DQL. * The method presented in this paper is merely a simple combination of previous works [1] and [2], which is lack of novelty. [1] Wang, Zhendong, Jonathan J. Hunt, and Mingyuan Zhou. "Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning." The Eleventh International Conference on Learning Representations. [2 ] Kingma, Diederik, and Ruiqi Gao. "Understanding diffusion objectives as the elbo with simple data augmentation." Advances in Neural Information Processing Systems 36 (2024).

Summary

This paper introduces DTQL, a novel offline reinforcement learning (offline RL) method. With a newly introduced diffusion trust region loss, DTQL constrained policy update within a predefined trust region near a diffusion policy trained by behavior cloning (BC). Through the empirical experiments, DTQL demonstrated improved performance against BC and offline RL baselines, including recent baselines that utilize diffusion models in offline RL settings.

Strengths

1. The diffusion-based trust region loss looks novel and interesting. 2. The visualization of trust region loss and the toy-tasks provide an intuitive explanation of the proposed algorithm. 3. The empirical experiments show improved performance compared to standard offline RL methods and faster inference time than most offline RL with diffusions. The reviewer appreciates that the author also shows the results where DTQL performs worse than IDQL (among the Antmaze tasks), which provides a more comprehensive view of the algorithm's performance. 4. The paper is easy to follow, and the connections and differences to related works are clearly outlined.

Weaknesses

1. The performance improvements compared with DQL, IDQL, and SRPO are not very significant; in some experiments (Antmaze), IDQL outperformed the proposed algorithms. 2. The ablation studies only show one seed result, which reduces the results' plausibility, especially for Figure 5 (a), where the learning curves are quite noisy.

Questions

1. The reverse KL objective should also be mode-seeking instead of mode-averaging. I do not quite understand why the KL regularization does not show this property in the example provided in Figure 3. As in the SRPO[1], a similar illustration was shown in Figure 5. Could the author provide more insights on what caused this difference in behavior? 2. When DTQL only utilizes four layers of MLP and SRPO requires several resnet blocks, why is SRPO inference faster than DTQL? [1] Chen H, Lu C, Wang Z, Su H, Zhu J. Score regularized policy optimization through diffusion behavior. arXiv preprint arXiv:2310.07297. 2023 Oct 11.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The limitations were discussed in the conclusion sections, along with possible future works.

Response to rebuttal

Thanks for providing the rebuttal response. I still have concerns and some of my questions are also not answered directly. I do think the paper in its current version needs some substantial improvement. I would thus keep my original assessment.

Response

I thank the authors for the detailed rebuttal. Most of by concerns has been addressed. I would like to raise the score. However, the organization of this paper can be improved. For the current form, the reviewer found it not easy to follow.