Summary
The paper mainly addresses the problem of obtaining the optimal policies in the distribution correction estimation (DICE) setting, which is one popular offline RL approach. Note that the offline RL assumes that we can’t interact with an environment and only have access to a dataset—sets of (state, action, and next state) tuples collected by some behavior policy. For the given history, DICE estimates the optimal value functions $V^*$ and $Q^*$ and obtains the optimal stationary distribution ratio $w^*$, which is the rate between the optimal policy’s state occupancy distribution $d^*$ and the one of the behavior policy $d^D$. Even if it is possible to get the optimal value functions and the ratio via DICE, obtaining the corresponding optimal policy $\pi^*$ from them is still challenging.
In this context, the paper proposes a novel training method for diffusion-based models to learn the optimal policy from the DICE-driven optimal value functions and the ratio. Unlike common density modeling, learning the diffusion-based models from the optimal value functions is not straightforward. Such a problem is closer to variational inference than density modeling; specifically, we don’t have samples from the target distribution to diffuse. Moreover, unlike a typical variational inference problem, sampling from the models is not available due to the nature of the offline RLs.
To do that the paper, the paper first introduces that the optimal policy can be represented by a product of the optimal stationary distribution ratio $w^*$ and the behavior policy $\pi^D$. Since the product of two distributions is unnormalized, the complete expression includes the normalization constant, which is the expected value of $w^*$ under the behavior policy (Line 155).
Next, the paper assumes that the behavior policy can be represented by a diffusion-based model with the forward process, which will also be used to perturb the policy of interest. Then, the authors show that the optimal policy at the perturbation level $t$ can be represented by using the behavior’s perturbed distribution at the same $t$ and the DICE ratio $w^*$ (Equation 6). In particular, this representation includes the logarithm of the expectation of the DICE ratio $w^*$ under the posterior distribution of the clean action $a_0$ for a given perturbed action $a_t$. While this novel representation doesn’t require sampling from the model policy anymore, computing the logarithm of the Monte Carlo estimate of any expectation is biased in general.
To circumvent this, the authors propose using the tangent transform, one of the variational approximation methods; thus, the quantity inside the logarithm is represented by an optimization problem, as in Equation 7. Interestingly, the new training objective only requires sampling from the behavior policy, which is suitable for the offline RLs.
In summary, the paper introduces a new representation of the diffusion-based optimal policy model by using the diffusion-based behavior policy and time-dependent learnable terms that will be learned by the convex problem described in Equation 7. The authors refer to this approach as In-sample Guidance Learning (IGL). While there have been a few previous approaches to learning diffusion-based optimal policy models, the authors point out that the IGL shows some benefits. For example, IGL only requires a single sample from the behavior policy, which can be obtained from the history, while some previous approaches require more than one sample from the behavior, which is not favorable in the offline RLs.
In addition, the authors suggest a few techniques to improve the stability of IGLs, such as using piecewise $f$-divergence for the DICE regularization term.
Finally, the paper discusses some failure cases of previous approaches and how the proposed method bypasses them. It also demonstrates the efficacy of the proposed method via experiments on several benchmark datasets.
-----------------------------------------------------
Updated the rating from 7 to 8 after the authors' rebuttal.
Strengths
In my understanding, the paper's contributions are clear, and I also consider that the results are essential for several reasons.
The paper introduces a novel representation of the diffusion-based optimal policy model and its training method, IGL. In addition, the paper also motivates the solution well. For example, this approach circumvents the drawbacks of previous approaches, which require more than one sample from the behavior policy, and such actions may be out-of-distribution of the environments.
Moreover, the authors provide extensive discussions to help potential readers comprehend the characteristics of previous approaches and the proposed method.
Finally, the paper demonstrates the effectiveness of the proposed method via various experiments, which further supports the authors' claim.
Weaknesses
Overall, the paper presents a novel method for learning optimal policy in DICE, which is a valuable contribution to the field. However, improvements in the presentation would greatly enhance the clarity and comprehensibility of the manuscript.
In particular, several equations within the paper omit the definitions of variables, which can lead to confusion—for example, $a_0$ in Line 119. In addition, some variables overlap while they are independent. For instance, in Line 155, the variable of the integration overlaps with the $a$ in the numerator.
I recommend revising the paper to address these issues.