MoMA: Model-based Mirror Ascent for Offline Reinforcement Learning

Thank you for your constructive feedback. We would like to address you concerns point by point in this response.

Regarding novelty and significance

We acknowledge that theoretical techniques involved in conservative policy evaluation, policy improvement, and their combination have been explored in related work. However, we believe that the integration of these techniques in our study offers valuable insights, particularly in the context of model-based offline RL under general function approximation settings without the need for explicit policy parameterization.

The theoretical results of our approach, as detailed in Theorem 1, enables the separation of model learning and policy learning. This separation allows for the use of an arbitrarily large function class for the (augmented) value function and an unlimited policy class. This aspect of our work is distinct from existing studies such as those by Xie et al., Cheng et al., Bhardwaj et al., and Rashidinejad et al. In particular, our method does not require explicit policy parameterization, setting it apart from these prior works.

The empirical performance...

We understand your concern about its comparative performance relative to RAMBO, IQL, and CQL. Nevertheless, it is important to note that the primary focus of our work is on the theoretical aspects of the algorithm. The empirical experiments were designed to support the theoretical findings rather than to showcase state-of-the-art performance. Additionally, the hyperparameters used in our experiments were tuned based on heuristics rather than through rigorous optimization.

Regarding Bellman completeness

Thank you for pointing out the reference and its implications regarding the realizability of the transition model in relation to Bellman completeness. Upon revisiting this reference and considering your observation, we recognize that the condition of model realizability is not a milder condition compared to Bellman completeness. In light of this, we have amended our paper by removing the aforementioned statement from the list of contributions in the revised version.

The paper claims that the size of the function can be arbitrarily large. However,...

We are sorry for the confusion. We clarify that we consider two notions of sample sizes in this paper: the sample size $n$ for offline data and the sample size for the Monte-Carlo evaluation $N$. While the offline sample size $n$ is fixed, the sample size of Monte-Carlo simulation $N$ can be picked by the user and arbitrarily large. In Theorem 1, the size of function class $|\mathcal{F}_{t,i}|$ only depends on $N$, and is not affected by $n$. The ability of $N\to\infty$ allows the size of |\mathcal{F} _\\{t,i\\}| to be arbitrarily large.

The paper does not need a parametric policy class. However, in general the computation complexity of the optimization problem Equation 7 will depend on the size of the state space, which can be infinite.

In algorithm 2, we only need to solve equation (7) for finitely-many states $s$ by using the strategy of Monte-Carlo estimation. In fact, we have illustrated in Appendix A.3 that it takes polynomial time to run Algorithm 2.

The authors claim that Theorem 1 characterizes the performance of Algorithm 2, but it seems not.

Thank you. We clarify that we provide an approximate solution to (1) by implementing the primal-dual algorithm (3)(4). The regularizer in the objective function of (3) forces the transition model to be nearly inside the confidence set. We note that the idea of incorporating a Lagrangian multiplier into the objective function for a constrained optimization problem is common in the optimization literature.

We hope we have addressed your concerns, and we sincerely appreciate the time and effort you have dedicated to reviewing our paper.

We are grateful for your insightful comment and would like to address your concerns in this response.

Regarding the PD algorithm (3) and (4)

Thank you. The PD algorithm (3) and (4) is an implementable version of equation (1) - the constrained minimization problem. Regarding your concern on the non-convexity of (1), we would like to discuss on the objective function $V_P^{\pi}$ and the empirical loss function $\mathcal{E}_n(P)$ with respect to $P$ in the following.

Regarding $V_P^{\pi}$. While $V_P^{\pi}$ is in general not a convex function in $P$, we would like to present here a promising way of handling the non-convex optimization problem to address your concern. We view it as a future direction, which is out of the scope of our current work.

Motivated by that DL and RL both involve optimizing non-convex objectives, non-convex optimization is a rising field and many recent works demonstrate the global convergence of gradient descent (GD): [Zou et al., 2020] and [Ding et al., 2022] show that GD methods converge to globally optimal solutions in over-parameterized neural networks. [Agarwal et al., 2021] and [Bhandari and Russo, 2019] demonstrate the global optimality of policy gradient methods. In our case, the proposed GD method for problem (1) can be considered a counterpart of the policy gradient method, i.e., "model gradient method." Specifically, we can view the fixed policy $\pi_t$ as the transition model, and minimizing $V_P^{\pi_t}$ w.r.t. $P$ in Eq. (1) may be viewed as finding the best "policy" that minimizes the overall cost. Techniques from [Agarwal et al., 2021] and [Bhandari and Russo, 2019] can then be naturally employed to prove that the model gradient method for Eq. (1) can find the global minimizer. In particular, their key results arise from the form of policy gradient and performance difference lemma, providing a gradient dominance condition typically needed for the global optimality of first-order methods for non-convex objectives. Similarly, by considering the form of model gradient (w.r.t. $P$) instead of policy gradient (w.r.t. $\pi$), and considering the simulation lemma (difference in values between two $P^{\prime}$ s) instead of the performance difference lemma (difference in values between two $\pi^{\prime}$ ), an analogous gradient dominance condition for $P$ could be naturally derived. As a result, the global minimizer for Eq. (1) can be found efficiently using primal-dual gradient descent Eq.(3)(4).

Regarding $\mathcal E(P)$. Though the empirical loss function $\mathcal E(P)$ is non-convex in general, it is indeed convex for a large class of distributions. For example, when $P$ is in an exponential family, then the negative likelihood is convex (Corollary 1.6.2. of [5]). Given that the empirical loss function is convex, the level set of the empirical loss function (i.e. the constraint set) is also convex.

References:

[1] Zou, Difan, et al. "Gradient descent optimizes over-parameterized deep ReLU networks." Machine learning 109 (2020): 467-492.

[2] Ding, Zhiyan, et al. "Overparameterization of deep ResNet: zero loss and mean-field analysis." Journal of machine learning research (2022).

[3] Agarwal, Alekh, et al. "On the theory of policy gradient methods: Optimality, approximation, and distribution shift." The Journal of Machine Learning Research 22.1 (2021): 4431-4506.

[4] Bhandari, Jalaj, and Daniel Russo. "Global optimality guarantees for policy gradient methods." arXiv preprint arXiv:1906.01786 (2019).

[5] Bickel, Peter J., and Kjell A. Doksum. Mathematical statistics: basic ideas and selected topics, volumes I-II package. CRC Press, 2015.

Thank you for your constructive feedback. We appreciate the opportunity to address each of your concerns.

The contribution of this work...

Thank you. We would like to clarify this aspect of our work.

In our study, we indeed present a theoretical framework for pessimistic model learning as outlined in (1), and alongside, we offer a practical implementation through the primal-dual algorithm as detailed in (3) and (4). We understand your concern about the potential differences between these two approaches. However, we would like to emphasize that the transition from the theoretical model to its practical implementation is grounded in a well-established technique in the optimization literature. This technique involves the incorporation of a Lagrangian multiplier into the objective function to address constrained optimization problems.

This approach ensures that while the practical implementation may appear as an approximate solution to the theoretical model, it fundamentally retains the core principles and objectives of the initial theoretical design. Therefore, we assert that the theoretically analyzed algorithm and its practical counterpart are aligned in their essence and are not significantly different as might be perceived.

In terms of the theoretically sound version...

Thank you for your comments regarding the theoretical aspects of our work. We acknowledge that theoretical techniques involved in conservative policy evaluation, policy improvement, and their combination have been explored in related work. However, we believe that the integration of these techniques in our study offers valuable insights, particularly in the context of model-based offline RL under general function approximation settings without the need for explicit policy parameterization.

The theoretical results of our approach, as detailed in Theorem 1, enable the separation of model learning and policy learning. This separation allows for the use of an arbitrarily large function class for the (augmented) value function and an unlimited policy class. This aspect of our work is distinct from existing studies. In particular, our method does not require explicit policy parameterization, setting it apart from these prior works.

Addressing your reference to "Provable Benefits of Actor-Critic Methods for Offline Reinforcement Learning," it's noteworthy that this work primarily focuses on linear models. In contrast, our research extends to more complex and general function approximation settings.

The practically implemented version...

Thank you for the references. We note that the policy classes in the works you mentioned are parameterized. While effective in many scenarios, this approach can be potentially restrictive, particularly when the optimal policy lies outside the pre-specified parametric policy class. Our work, on the other hand, does not need an explicit policy parameterization. By employing mirror ascent in conjunction with our model-based design, we enable our policy class to be sufficiently large, potentially containing the true optimal policy. This flexibility is a fundamental advantage of our approach, addressing limitations inherent in parameterized policy classes.

The implemented version of algorithm is only...

Thank you for the suggestions. Since this work mainly focuses on the theoretical part, we believe the numerical experiments in the D4RL tasks are enough to demonstrate the efficacy of MoMA.

How to implement the policy update when there is a normalization factor?

Thank you. When evaluating a policy, we need to calculate the normalization factor by summing the policy probabilities over all actions. Then we use the evaluated policy to generate Monte-Carlo trajectories.

How do you update the multiplier $\lambda$?

The update rule for $\lambda$ is described in equation (4). See section 5.3 for an example. In particular, we evaluate the gradient of the empirical loss function, and use it as an ascent direction for $\lambda$.

Is it possible to unify these two different versions of algorithms?

Thank you for the idea. It is an interesting and challenging problem, and we leave it as a future study.

We hope we have addressed your concerns. Thank you again for your time and valuable comments.

We would like to express our gratitude for your time and effort in reviewing our paper. Your constructive feedback is greatly appreciated. In this response, we address your concerns point by point.

One of the main contributions is stated as...

Thank you for pointing out the reference [1] and its implications regarding the realizability of the transition model in relation to Bellman completeness. Upon revisiting this reference and considering your observation, we recognize that the condition of model realizability is not a milder condition compared to Bellman completeness. In light of this, we have amended our paper by removing the aforementioned statement from the list of contributions in the revised version. We appreciate your insight, which has helped clarify this aspect of our research and improve the accuracy of our claims.

It is stated that “To our knowledge...

Thank you for pointing out these two references [2] and [4]. We acknowledge that our study is not the first one to offer a practically implementable model-based offline RL algorithm with theoretical guarantees under general function approximations. Accordingly, we have revised our manuscript to remove this claim. Additionally, we have included comparisons with the algorithms presented in [2] and [4] to provide a more comprehensive overview.

However, we note that there is a specific distinction between our approach and the ones in [2] and [4]. Our methodology does not necessitate the assumption of a parametric policy class in advance. This advantage stems from our utilization of mirror ascent, which we believe provides an alternative perspective to the field of model-based offline RL.

Although the work is focused on general function approximation, it requires the concentrability definition...

Thank you for the comment. We recognize that our current assumption may indeed be more stringent compared to the Bellman-consistent variant outlined in [3]. In light of this, we plan to revisit and revise our assumptions and proofs. This revision will explore the possibility of modifying our concentrability coefficient to align more closely with the Bellman-consistent variant mentioned in [3]. We aim to include these changes and any resultant findings in our next revision.

Theory does not seem to be particularly challenging and/or offer new insights or techniques.

Thank you for your comments regarding the theoretical aspects of our work. We acknowledge that theoretical techniques involved in conservative policy evaluation, policy improvement, and their combination have been explored in related work. However, we believe that the integration of these techniques in our study offers valuable insights, particularly in the context of model-based offline RL under general function approximation settings without the need for explicit policy parameterization.

The theoretical results of our approach, as detailed in Theorem 1, enables the separation of model learning and policy learning. This separation allows for the use of an arbitrarily large function class for the (augmented) value function and an unlimited policy class. This aspect of our work is distinct from existing studies such as those by Xie et al., Cheng et al., Bhardwaj et al., and Rashidinejad et al. In particular, our method does not require explicit policy parameterization, setting it apart from these prior works.

Regarding numerical results

Thank you. The synthetic data experiments mainly aim to demonstrate the theoretical results and the efficacy of the proposed algorithm. Therefore, we compared the proposed algorithm to the ones that do not incorporate pessimism, and showed that the proposed algorithm can indeed address the distribution shift issue. In addition, it can be seen from table 2 that we have already compared MoMA to both model-based and model-free offline RL methods for the D4RL benchmark.

Thank you again for your time and valuable suggestions.

References:

[1] Chen, Jinglin, and Nan Jiang. "Information-theoretic considerations in batch reinforcement learning." In International Conference on Machine Learning, pp. 1042-1051. PMLR, 2019.

[2] Rashidinejad, Paria, Hanlin Zhu, Kunhe Yang, Stuart Russell, and Jiantao Jiao. "Optimal Conservative Offline RL with General Function Approximation via Augmented Lagrangian." In The Eleventh International Conference on Learning Representations. 2022.

[3] Cheng, Ching-An, Tengyang Xie, Nan Jiang, and Alekh Agarwal. "Adversarially trained actor critic for offline reinforcement learning." In International Conference on Machine Learning, pp. 3852-3878. PMLR, 2022.

[4] Bhardwaj, Mohak, Tengyang Xie, Byron Boots, Nan Jiang, and Ching-An Cheng. "Adversarial model for offline reinforcement learning." arXiv preprint arXiv:2302.11048 (2023).