DiffCPS: Diffusion Model based Constrained Policy Search for Offline Reinforcement Learning

We thank reviewer 4X5N for pointing out the drawbacks of our paper. Below, we address the concerns raised in your review point by point. Please let us know if you have any further concerns or whether this adequately addresses all the issues that you raised with the paper.

Q1: Corollary 3.1.1 and Theorem 3.2 are flawed.

A1: Thanks for pointing this out. We have rigorously re-proven the strong duality of Equation 12 in Appendix A.2. New proof relies on a well-known result from perturbation theory connecting strong duality to the convexity of the perturbation function. Here, we formally state the lemma used in the proof, which is referred to [1] Cor. 30.2.2 and [2] Ch 3.

If (i) $r$ is bounded; (ii) Slater's condition holds for the diffusion-based constrained policy search and (iii) its perturbation function~$P(\xi)$ is concave, then strong duality holds for equation 12.

The proof consists of two main steps (See details in Appendix A.2. in our revision): (i) demonstrating that the perturbation problem of the simplified diffusion-based constrained search problem, as per Theorem 3.1, is a concave function, and (ii) proving the Slater's condition holds for the diffusion-based constrained search problem. We have also rephrased the main theorems of the paper to ensure a more coherent logical flow.

Q2: The paper lacks clarity. Multiple mathematical objects are introduced without clear definitions. If the reader is not familiar with the literature on RL and diffusion models, she/he will definitely get confused by the undefined notations. Also, I think the authors should provide more explanations about the exact methods used to update $\mu$ and $\lambda$ . It would be appreciated if the authors could give a complete and precise description of DiffCPSS. (If there are problems with page limits the authors can put it in the Appendix.)

A2: Thanks for suggesting this. We have added the Notation table in Appendix B and rephrased our DiffCPS algorithm in Chapter 3. Here, we reiterate the logic of our algorithm. We first establish the strong duality for Equation 12, allowing us to use Equation 14 to find the optimal policy. The approach to solving Equation 14 follows SAC [3], using alternating optimization to optimize the policy and Lagrange multipliers separately. Equation 19 and Equation 20 represent the loss functions for the Lagrange multipliers and the policy, respectively.

Q3: The empirical performance of the proposed algorithm exhibits only marginal improvement when compared to the baselines.

A3:

Table 1: D4RL performance of different offline RL methods. The table above illustrates the performance of recent ICLR offline RL algorithms. SfBC (ICLR 2023) shows a similar performance to IQL (ICLR 2022), with an average decrease of 1.7% in locomotion tasks and a 17.8% improvement in antmaze tasks. The improvements of our DiffCPS compared to these algorithms are shown in the table below.

Table 2: The improvements of our DiffCPS compared to recent ICLR offline RL algorithms.

Compared with other offline reinforcement learning algorithms, it's evident that our algorithm exhibits a significant improvement over the baseline.

Q4: In the toy example, the authors claim that "SfBC incorrectly models the dataset as a circle instead of a noisy circle". I checked the image and did not find such a difference. Can the authors give an explanation?

A4: Thanks for pointing this out. We've realized that our statement might lead to misunderstandings. We've revised the statement in our revision. What we intended to convey is that due to the architecture of SfBC introducing unnecessary sampling errors at T=15, it generates points that differ significantly from the dataset.

Thanks for the authors' response. I find the revised version more readable. The revised proof appears to be sound. Nevertheless, I still have several questions after checking the proof in detail.

1, About equation (34). In equation (34), the authors use the relationship $\rho^{\pi_b}(s)\approx\rho^{\mu_t}(s)\approx\rho^{\mu_1}(s)\approx\rho^{\mu_2}(s)$. What does the $\approx$ mean here? Why does such a relationship hold?

2, About the Slater's condition. When verifying the Slater's condition, the authors only verify that for a specific $\kappa_0$ the Slater's condition holds. My interpretation is that this means the Slater's condition might only be satisfied for certain choices of $\epsilon$. Am I correct? However, it seems the authors conclude that the Slater's condition holds for all $\epsilon$.

If the authors could address my concerns I would consider raising my score.

(Some random thoughts: I understand this paper may not be a theory-oriented work. And I agree that intuitively it is totally fine to apply the dual methods even if the strong duality is only "approximately true". But I think if the authors choose to state their intuitions in the form of mathematical theorems, then they should ensure full technical rigor. )

We thank reviewer hHF8 for providing positive feedback and helpful suggestions. Below please find a response.

Q1: Qualitative behavior of DiffCPS compared to other baselines

A1: Thanks for suggesting this. We have added an experiment to show the behavior of DiffCPS in Appendix F.1

Q2: missing some relevant work

A2: Thanks for pointing that out. We have supplemented information about the related work on model-based offline reinforcement learning in Appendix E. And we will emphasize here the distinctions between our approach and [1]. [1] introduces a method that utilizes the score function of the Diffusion Model to approximate the gradient of the proposed uncertainty estimation. Our method, DiffCPS, diverges from this approach in several key aspects: (i) DiffCPS tackles the constrained policy search problem in offline RL through Lagrange dual and recursive optimization, while the method in the paper employs first-order optimizers directly. (ii) DiffCPS is a model-free method, in contrast to the model-based nature of the method in [1]. (iii) The objective function in DiffCPS differs as it incorporates the loss from DDPM instead of score matching, and our coefficient is auto-adaptive. (iv) Our experimental results in D4RL demonstrate the superior performance of DiffCPS over the method in the paper, attributed to the design of our algorithm.

Q3: Maybe Diffusion based might be a better title?

A3: Thanks for suggesting this. We notice that "diffusion model-based" could be ambiguous to some extent. Following your advice, we have changed the title and algorithm names to Diffusion-based constrained policy search to eliminate ambiguity.

Q4: Have the authors considered using the Augmented Lagrangian Method (ALM)?

A4: Initially, for simplicity, we directly applied the alternating optimization (Dual Ascent) of SAC [2] to solve the dual problem. Theoretically, using ALM could offer better convergence performance. Our experiments on HalfCheetah demonstrate that using ALM increases training stability. However, the final score is slightly lower than the Dual Ascent method. This could be attributed to (i) ALM introducing penalty terms, making the optimal target KL divergence for Dual Ascent not necessarily optimal for ALM, and (ii) the penalty function making the policy too conservative, thereby preventing the adoption of high-reward actions. We plan to further explore the use of ALM to optimize our problem in future research. You can find our discussion on the application of ALM in our algorithm in Appendix G.

Q5: question about $d_{\pi_b(s)}$

A5: Thanks for pointing this out. $d_{\pi_b(s)}$ denotes the unnormalized discounted state visitation frequencies of the behavior policy. We acknowledge the issues with the representation of some notations in the paper. In the revised version, we use $\rho^{\pi_b}(s)$ to represent the unnormalized discounted state visitation frequencies of the behavior policy and provide a notation table in Appendix B.

Q6: How specific are the authors' claims to diffusion models? For instance, if we had trained a denoising autoencoder model (assuming they can be trained well), would the theorems still hold?

A6: In fact, when using a denoising autoencoder instead of the diffusion model, the theoretical properties of DiffCPS are no longer guaranteed. This is because (i) Constraint $\int_a \mu(a\vert s)da =1$ is challenging to establish for other models; it holds for the diffusion model because $\mu(a\vert s) = \int \mu(a^{0:T}\vert s)da^{1:T}$ holds. (ii) Denoising autoencoders struggle to accurately approximate the constraint on KL divergence.

[1] Suh et al., "Fighting Uncertainty with Gradients: Offline Reinforcement Learning with Diffusion Score Matching", CoRL 2023

[2] Haarnoja, Tuomas, Aurick Zhou, Pieter Abbeel, and Sergey Levine. “Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor.” arXiv, August 8, 2018. https://doi.org/10.48550/arXiv.1801.01290.

I would like to thank the authors for the detailed response, I mostly agree with the points and will keep the score as is.

We thank reviewer UbgT for the comments and suggestions. Below, we address the concerns raised in your review point by point. Please let us know if you have any further concerns or whether this adequately addresses all the issues that you raised with the paper.

Q1: The importance or motivation of propositions, theorems, and corollaries is not well explained. Most results can be more directly obtained using simple calculations that are known in the diffusion model.

A1: In fact, the logic of our theorems, proposition, and corollary is as follows: Theorem 3.1 and Corollary 3.1.1 establish conditions for Theorem 3.2 to hold. Theorem 3.2 proves the strong duality for the constrained policy search problem with policies based on the diffusion model. This allows us to use the dual problem to find the optimal solution (optimal policy) for the original problem (Theorem 3.2 in the revised paper). Proposition 3.1 is introduced for solving the dual problem, stating that the loss of DDPM can serve as an approximation to $H(\pi_b,\mu)$. In the revised version, we have reorganized the theorems, proposition, and collary for a clearer presentation.

Q2: A large paragraph of this paper introduces known results. Can the authors highlight more new developments compared with existing methods?

A2: Compared to other methods using the diffusion model, we directly employ diffusion-based policy to solve constrained policy search. We demonstrate the previously unexplored strong duality for the constrained policy search problem with policies based on the diffusion model. Our research indicates that the favorable properties of the diffusion model can be leveraged for constrained policy search in offline RL, addressing the limited expressivity of Gaussian-based models.

Q3: It is not clear to me the strong duality of Equation (12). Can the authors justify it?

A3: We have rigorously re-proven the strong duality of Equation 12 in Appendix A.2. New proof relies on a well-known result from perturbation theory connecting strong duality to the convexity of the perturbation function. Here, we formally state the lemma used in the proof, which is referred to [1] Cor. 30.2.2 and [2] Ch 3.

If (i) $r$ is bounded; (ii) Slater's condition holds for the diffusion-based constrained policy search and (iii) its perturbation function~$P(\xi)$ is concave, then strong duality holds for equation 12.

The proof consists of two main steps (See details in Appendix A.2. in our revision): (i) demonstrating that the perturbation problem of the simplified diffusion-based constrained search problem, as per Theorem 3.1, is a concave function, and (ii) proving the Slater's condition holds for the diffusion-based constrained search problem. And we have also rephrased the main theorems of the paper to ensure a more coherent logical flow.

Q4: How does the problem (18)-(19) can solve the original problem (12)?

A4: The dual function for Problem 12 is given by Equation 14. Following SAC[3], we can use alternating optimization to solve the dual problem. Theorem 3.2 also proves that the policy obtained through the dual problem is optimal for the original problem. In fact, Equation 18 and Equation 19 represent the objective functions for the policy and Lagrange multipliers. By applying stochastic gradient descent to Equation 18 and Equation 19, we can approximately solve the dual problem Equation 14, and consequently address Problem 12. The updated paper includes a clearer pseudocode (Ch3) for DiffCPS, highlighting this process.

Q5: Are there other diffusion models useful for constrained policy search? It is useful if the authors could discuss the generalization of this approach.

A5: Other methods based on the diffusion model can also leverage our framework. This is because other diffusion model-based methods, such as Score-Based Generative Models (SGMs) and Stochastic Differential Equations (Score SDEs), define a common distribution for all actions. Thus, $\mu(a\vert s) = \int \mu(a^{0:T}\vert s)da^{1:T}$ holds for other diffusion model-based methods, allowing the elimination of constraints $\int_a \mu(a\vert s)da =1$ and the application of our theoretical framework.