Minimax Optimal and Computationally Efficient Algorithms for Distributionally Robust Offline Reinforcement Learning

We thank the reviewer for your positive feedback on our work. We hope our response fully addresses all your questions.

---

1. Requirement of the number of trajectories $K$ to be poly(d,H).

This is an interesting question. We acknowledge that our current analysis requires the sample size $K$ has a large order dependence on $d$ and $H$. Similar requirements are also necessary for standard offline RL with linear function approximation (Yin et al., 2022, Xiong et al. 2022). It is still unsure if we could construct upper and lower bounds for full-range $K$ in the robust linear function approximation setting. We will leave this problem for future study.

---

2. Comparison with Ma et al. (2022).

We acknowledge that Ma et al. (2022)’s work is the most closely related to ours. However, we note that there are several technical flaws in their paper. We summarize the main issues here:

• The proofs of main lemmas (Lemma D.1 and Lemma D.2 in the arXiv version) related to suboptimality decomposition and the proof of theorems are incorrect.

• The sample complexity results for the two algorithms are both incorrect on the order of $d$.

• The Assumption 4.4 on the dual variable of the dual formulation of the KL-divergence is too strong to be realistic.

• The concept of mis-specification in Section 5.2 is not well-defined, which makes all results in Section 5 vacuous.

Given these technique flaws, we believe the fundamental challenges of $d$-rectangular linear DRMDPs are not properly addressed in their work. It is also unsure whether $d$-rectangular linear DRMDPs with KL-divergence uncertainty set are solvable by their current methodologies. Moreover, in our work, we start with the setting of $d$-rectangular linear DRMDPs with TV-divergence uncertainty set. We address the essential challenges and explore the fundamental limit of $d$-rectangular linear DRMDPs.

---

3. Reasons for considering TV-divergence and challenges in the analysis.

There are two main reasons we consider the TV-divergence: 1) TV-divergence in a commonly studied distance in the field of DRMDPs, it serves as a proper “start” setting to address the fundamental challenges of DRMDPs with linear function approximation; 2) We also tried the KL-divergence uncertainty set, however the analysis heavily relies on an unrealistic assumption on the dual variable of the duality for KL. Thus, it requires further effort to figure out if the linear DRMDP with KL-divergence uncertainty set is both theoretically and practically solvable.

Thanks to the simple format of the duality for TV, we can simply focus on the essential challenges of the linear DRMDP setting. To see the main challenge of the analysis, we first note that existing analysis of standard linear MDPs highly relies on the linear dependency of the Bellman equation on the nominal kernel. However, the consideration of model uncertainty of DRMDPs disrupts the linear dependency (see (2.1a)) of the robust Bellman equation on the nominal kernel, where the offline dataset is collected. This nonlinearity leads to unique challenges in both upper bound and lower bound analysis, and we are the first to solve those challenges.

---

4. Extending the results to general linear MDPs.

Generalizing the results to standard linear MDPs presents challenges primarily due to the construction of uncertainty sets. In particular, the design of the $d$-rectangular uncertainty set relies on the fact that the factor uncertainty sets are defined around probability distributions, since the total variation distance $D_{\text{TV}}(\cdot||\cdot)$ is a distance measure for probability distributions. Further, it is non-trivial to define uncertainty sets for the standard linear MDP to ensure (1) it contains valid probability distributions and (2) the optimization among the uncertainty set can be efficiently solved (by duality for example). We leave this extension for future research.

---

5. Wording suggestion.

Thanks, we will take your advice and revise the “minimax optimal” to “near-optimal”.

---

We hope we have addressed all of your questions/concerns. If you have any further questions, we would be more than happy to answer them and if you don’t, would you kindly consider increasing your score?

---

References

[1] Yin, Ming, Yaqi Duan, Mengdi Wang, and Yu-Xiang Wang. "Near-optimal Offline Reinforcement Learning with Linear Representation: Leveraging Variance Information with Pessimism." In International Conference on Learning Representation. 2022.

[2] Xiong, Wei, Han Zhong, Chengshuai Shi, Cong Shen, Liwei Wang, and Tong Zhang. Nearly Minimax Optimal Offline Reinforcement Learning with Linear Function Approximation: Single-Agent MDP and Markov Game. In International Conference on Learning Representations (ICLR). 2023.

We thank the reviewer for your valuable time and effort in providing detailed feedback on our work. We hope our response will fully address all your questions.

---

1. Computational tractability.

We would like to note that we discussed the computational tractability of our algorithm on Line 178 - Line 203. Specifically, if the ‘fail-state’ assumption holds, then the minimal value is trivially equal to zero hence no computation is needed. If the ‘fail-state’ does not hold, we can use some heuristic methods, such as the Nelder-Mead method, to search for the minimal value. Thus, our algorithms are in general computationally tractable. We have also conducted numerical experiments to show the computational tractability of our algorithm.

In terms of computation efficiency, 1) our algorithms leverage the linear function approximation, and do not need to estimate the robust value at each state-action pair as those methods developed for the tabular DRMDP setting; 2) Moreover, thanks to the structure of linear DRMDPs, we do not need to solve the dual problem at each state-action pair. Thus, our algorithm is computationally efficient compared to those methods developed for tabular DRMDPs and general $(s, a)$-rectangular DRMDPs. This is consistent with the claim of computationally efficient algorithms in theoretical RL (Jin et al., 2019, Xiong et al. 2022, Yin et al. 2022), which stating the runtime and the sample complexity should not depend on the number of states, but should depend instead on an intrinsic complexity measure of the function class.

---

2. In Line 202, I don't find Assumption 4.1 and Remark 4.2 in the paper. Additionally, the "fail-state" assumption doesn't necessarily hold in multiple MuJoCo environments, i.e., some environments have negative minimum values.

We believe there is a misunderstanding. Here we are discussing previous work’s assumption 4.1 and remark 4.2 (Liu and Xu, 2024: https://proceedings.mlr.press/v238/liu24d/liu24d.pdf). We highlight that our work does not need the fail-state assumption. For environments that the fail-state does not hold, we can 1) either modify the reward of the environment to rescale the minimal value (of an absorbing state) to zero, which creates a fail-state, or 2) use heuristic methods to search for/approximate the minimal value.

---

3. Numerical experiments

Thanks for your suggestion. We have done some experiments to verify our proposed algorithm. Please see the overall response for more details.

---

4. Clarify the linear function approximation in the open question.

We will revise the open question as follows: Is it possible to design a computationally efficient and minimax optimal algorithm for robust offline RL with linear function approximation?

---

5. Typos and references.

Thanks for pointing out these issues. We have revised the typo and added the references.

---

6. Justification of why linear function approximation rather than the general function approximation is considered in this work.

In this work, we focus on offline DRMDPs with linear function approximation under the setting of $d$-rectangular linear DRMDPs with TV-divergence uncertainty sets, which is an understudied area with unique fundamental challenges compared to standard offline RL. It would be an interesting future research direction to incorporate general function approximation techniques, such as those in Chen and Jiang (2019) into offline DRMDPs. We would add the following paragraph in the introduction section to motivate our setting:

Although standard offline MDPs based on linear function approximation have exhibited theoretical success (Jin et al., 2021; Yin et al., 2022; Xiong et al., 2022), DRMDPs with linear function approximation remain understudied. In particular, DRMDP encounters unique difficulties when applying linear function approximations, even when the source domain transition kernel is linear. The dual formulation in worst-case analyses induces extra nonlinearity for the function approximation. Consequently, the theoretical understanding of offline DRMDPs with function approximation remains elusive, even when the approximation is linear. In this work, we aim to address the essential challenges and explore the fundamental limit of this setting.

In the conclusion section, we will add the following sentences to explore future research directions:

Moreover, we would like to extend the pessimism principle utilized in this work beyond linear function approximation to explore general function approximation in DRMDPs. Leveraging the techniques for general function approximation in standard offline RL (Chen and Jiang, 2021), we would explore the unique challenges and fundamental limits of practically solving DRMDPs with general function approximation.

---

We hope we have addressed all of your questions/concerns. If you have any further questions, we would be more than happy to answer them and if you don’t, would you kindly consider increasing your score?

---

References

[1] Zhishuai Liu, and Pan Xu. "Distributionally robust off-dynamics reinforcement learning: Provable efficiency with linear function approximation." In International Conference on Artificial Intelligence and Statistics, pp. 2719-2727. PMLR, 2024.

[2] Yin, Ming, Yaqi Duan, Mengdi Wang, and Yu-Xiang Wang. (2022). "Near-optimal Offline Reinforcement Learning with Linear Representation: Leveraging Variance Information with Pessimism."

[4] Ying Jin, Zhuoran Yang, and Zhaoran Wang. (2021). Is pessimism provably efficient for offline RL?

[4] Xiong, Wei, Han Zhong, Chengshuai Shi, Cong Shen, Liwei Wang, and Tong Zhang. (2022). Nearly Minimax Optimal Offline Reinforcement Learning with Linear Function Approximation: Single-Agent MDP and Markov Game.

[5] Jin, Chi, Zhuoran Yang, Zhaoran Wang, and Michael I. Jordan. (2019). Provably efficient reinforcement learning with linear function approximation.

We thank the reviewer for your positive feedback on our work. We hope our response will fully address all your questions.

---

1. Comparison with Ma et al. (2022).

We acknowledge that Ma et al. (2022)’s work is most closely related to ours. However, we note that there are several technical flaws in their paper. We summarize the main issues here:

(i) The proofs of their main lemmas (Lemma D.1 and Lemma D.2 in the arXiv version) related to suboptimality decomposition and the proof of theorems are incorrect.

(ii) The sample complexity results for the two algorithms are both incorrect on the order of $d$.

(iii) Assumption 4.4 on the dual variable of the dual formulation of the KL-divergence is too strong to be realistic.

(iv) The concept of mis-specification in Section 5.2 is not well-defined, which makes results in Section 5 vacuous.

Given these technique flaws, we believe the fundamental challenges of $d$-rectangular linear DRMDPs are not properly addressed in their work. It is also unsure whether $d$-rectangular linear DRMDPs with KL-divergence uncertainty set are solvable by their current methodologies. Moreover, in our work, we start with the setting of $d$-rectangular linear DRMDPs with TV-divergence uncertainty sets. We address the essential challenges and explore the fundamental limit of $d$-rectangular linear DRMDPs. In the next version of our manuscript, we plan to add a paragraph discussing the comparison with Ma et al. (2022) in detail and further elaborate on our contributions.

---

2. Stronger coverage assumption.

First of all, Assumption 5.1 in Ma et al. (2022) is problematic, since the right hand side has an additional dependence on $d$, which makes the right hand side unreasonably large. The assumption 5.1 should be in this form: there exists some absolute constant $c>0$, such that for any $(i, h, s, P) \in [d]\times[H] \times \mathcal{S} \times \mathcal{U}^{\rho}(P^0)$, we have

$$\Lambda_h \succeq \lambda I + K\cdot c \cdot \mathbb{E}^{\pi^{\star},P}\big[(\phi_i(s,a)\mathbf{1}_i)(\phi_i(s,a)\mathbf{1}_i)^\top |s_1=s \big],$$

which is exactly the robust partial coverage assumption in Blanchet et al. (2023, Assumption 6.3).
On the one hand, a closer examination of Assumption 3.3 in our paper reveals that it guarantees a weaker version of the robust partial coverage assumption: there exists some constant $c^{\dagger}>0$, such that for any $(i, h, s, P)\in[d]\times[H]\times \mathcal{S} \times \mathcal{U}^{\rho}(P^0)$, we have

$

 \Lambda_h \succeq \lambda I + K\cdot c^{\dagger}/d \cdot \mathbb{E}^{\pi^{\star},P}\big[(\phi_i(s,a)\mathbf{1}_i)(\phi_i(s,a)\mathbf{1}_i)^\top |s_1=s \big].

$

Nevertheless, Assumption 3.3 does not directly imply the robust partial coverage assumption. On the other hand, assumption 3.3 is essential in achieving the minimax optimality of our algorithm. A similar phenomenon also appears for the non-robust offline RL with linear function approximation. Specifically, leveraging the pessimism principle, Jin et al. (2021) show that the partial coverage assumption is enough for the provable efficiency of offline RL with linear function approximation. However, to achieve minimax optimality, Yin et al. (2022) and Xiong et al. (2022) show that a stronger full-type coverage is needed. Lastly, in a concurrent work (Wang et al, 2024) of ours, the same assumption is also required to achieve tighter dependences on $d$ and $H$ (see Assumption 4 and Theorem 2 of their paper).

---

3. Question marks on line 618.

Thanks for pointing out this typo. We have fixed it in our manuscript.

---

We hope we have addressed all of your questions/concerns. If you have any further questions, we would be more than happy to answer them and if you don’t, would you kindly consider increasing your score?

---

References

[1] Ying Jin, Zhuoran Yang, and Zhaoran Wang. Is pessimism provably efficient for offline rl? In International Conference on Machine Learning, pages 5084–5096. PMLR, 2021.

[2] Xiong, Wei, Han Zhong, Chengshuai Shi, Cong Shen, Liwei Wang, and Tong Zhang. Nearly Minimax Optimal Offline Reinforcement Learning with Linear Function Approximation: Single-Agent MDP and Markov Game. In International Conference on Learning Representations (ICLR). 2023.

[3] Blanchet, Jose, Miao Lu, Tong Zhang, and Han Zhong. Double Pessimism is Provably Efficient for Distributionally Robust Offline Reinforcement Learning: Generic Algorithm and Robust Partial Coverage. In Thirty-seventh Conference on Neural Information Processing Systems. 2023.

[4] Wang, He, Laixi Shi, and Yuejie Chi. Sample complexity of offline distributionally robust linear Markov decision processes. arXiv preprint arXiv:2403.12946 (2024).

We thank the reviewer for your positive feedback on our work. We hope our response fully addresses all of your questions.

---

1. More motivation on practical scenarios where DRMDPs are applicable.

In many real-world applications, the agent only has access to a single source domain, which can be different from target domains where the policy is deployed, assuming the same task is being performed. Importantly the target domains are unknown and can vary during deployment. Given the limited information and large uncertainty, one may hold a pessimistic perspective hoping the learned policy is robust enough that would perform well even under the worst-case transition. This basic idea leads to the DRMDP framework, which models the unknown target domain in an uncertainty set defined around the source domain.

We take the control of infectious diseases for an example. Suppose we have access to a simulator (source domain), which emulates one specific real-world disease infection process (target domain). However, the simulator does not have perfect knowledge about the complex real-world environment, but it is reasonable to assume that the real-world environment is close to the simulator, thus lying in an uncertainty set around it. Notably, the real-world environment can be any environment in the uncertainty set, and the DRMDP framework provides a worst-case guarantee for policies learned merely through the simulator.

---

2. Sample simulation examples.

Thanks for your suggestion. We have done some experiments to verify our proposed algorithm. Please see the overall response for more details.

---

3. Literature on theoretical multi-task RL.

Thank you for bringing this work to our attention. We will add the following paragraphs to the related work section.

Besides the distributionally robust perspective to solve the planning problem in a nearly unknown target environment, another line of work focuses on transfer learning in low-rank MDPs (Cheng et al., 2022; Lu et al., 2022; Agarwal et al., 2023; Bose et al., 2024). Specifically, the problem setup assumes that the agent has access to information of several source tasks. The agent learns a common representation from the source domains and then leverages the learned representation to learn a policy performing well in the target tasks with limited information. This setting is in stark contrast to the setting of DRMDP, where the agent only has access to the information of a single source domain, without any available information of the target domain, assuming the same task is being performed. This motivates the pessimistic attitude of the distributionally robust perspective.

Among the aforementioned works, Bose et al. (2024) studied the offline multi-task RL, which is the most closely related to our setting. In particular, they investigate the representation transfer error in their Theorem 1, stating that the learned representation can lead to a transition kernel that is close to the target kernel in terms of the TV divergence. Note that the uncertainty is induced by the representation estimation error, which is different from our setting assuming that the uncertainty comes from perturbations on underlying factor distributions. Nevertheless, this work provides evidence that TV divergence is a reasonable measure to quantify the uncertainty in transition kernels and motivates a future research direction in learning robust policies that are robust to the uncertainty induced by the representation estimation error.

---

We hope we have addressed all of your questions/concerns. If you have any further questions, we would be more than happy to answer them and if you don’t, would you kindly consider increasing your score?

---

References

[1] Rui Lu, Andrew Zhao, Simon S Du, and Gao Huang. Provable general function class representation learning in multitask bandits and mdp. Advances in Neural Information Processing Systems, 35:11507–11519, 2022.

[2] Yuan Cheng, Songtao Feng, Jing Yang, Hong Zhang, and Yingbin Liang. Provable benefit of multitask representation learning in reinforcement learning. Advances in Neural Information Processing Systems, 35: 31741–31754, 2022.

[3] Alekh Agarwal, Yuda Song, Wen Sun, Kaiwen Wang, Mengdi Wang, and Xuezhou Zhang. Provable benefits of representational transfer in reinforcement learning. In The Thirty Sixth Annual Conference on Learning Theory, pages 2114–2187. PMLR, 2023.

[4] Bose, Avinandan, Simon Shaolei Du, and Maryam Fazel. "Offline multi-task transfer RL with representational penalization." arXiv preprint arXiv:2402.12570 (2024).

[5] Zhishuai Liu, and Pan Xu. "Distributionally robust off-dynamics reinforcement learning: Provable efficiency with linear function approximation." In International Conference on Artificial Intelligence and Statistics, pp. 2719-2727. PMLR, 2024.

[6] Ying Jin, Zhuoran Yang, and Zhaoran Wang. Is pessimism provably efficient for offline rl? In International Conference on Machine Learning, pages 5084–5096. PMLR, 2021.