Robustify the Latent Space: Offline Distributionally Robust Reinforcement Learning with Linear Function Approximation

Thank you for your detailed reading and valuable comments. Please find our response to the comments below:

Q. The approach to tackle the uncertainty from the dataset.

A. We adopt a similar pessimism principle by subtracting $\gamma_h \sum_{i=1}^d \lVert \phi_i(s,a)1_i\rVert_{\Lambda_h^{-1}}$ from the estimated $Q$-value, which is stated in the last paragraph on Page 7. We apologize for any confusion in our presentation and will provide a clearer outline in our next version.

Q. The novelty of the results.

A. We acknowledge the reviewer's observation that, after understanding the proper construction of the ambiguity set, the subsequent results can generally hold under many favorable properties. The novelty of our work lies in uncovering the proper construction, a previously unknown aspect until recent progress by Goyal, Vineet, and Julien Grand-Clement in 2023. As our paper is the first to establish finite-sample results in this setting, we rely on some strong conditions, like Definition 2.1, to outline the main message. Additionally, we introduce the novel value-shift trick and the corresponding technical skills in proving the theoretical results, which themselves contribute valuable insights.

[1] Goyal, Vineet, and Julien Grand-Clement. "Robust Markov decision processes: Beyond rectangularity." Mathematics of Operations Research 48.1 (2023): 203-226.

Thank you for your detailed reading and valuable comments. Please find our response to the comments below:

Q. The necessity of using value shift to the improved theoretical results.

A. The vanilla estimator itself can take any real number. The simple clip can ensure the estimator to be nonnegative, which is still not sufficient to prevent the approach to infinity, as $\log(Z)$ will still go to infinity when $Z$ approaches zero. Thus, the value shift is necessary to ensure the numerical safety of the estimation.

In terms of the theoretical results, without the value shift technique, one can obtain $O(\underline{\beta}e^{H/\underline{\beta}} H^2 / N^{1/2})$ result, and it would approach infinity when $\underline{\beta}$ approaches zero (when the algorithm is learning a nearly non-robust view). In fact, Zhou et al. 2021 establishes the finite-sample guarantee under the KL divergence without using the value shift and ends up with $\underline{\beta}e^{H/\underline{\beta}}$ as a coefficient of their dominating term (See Lemma 5 in Zhou et al. 2021). By introducing the value shift technique, we can obtain $O(\underline{\beta}(e^{H/\underline{\beta}}-1) H^2 / N^{1/2})$ and can succeed in recovering the same dependence on H as the non-robust PEVI algorithm.

Q. Uniformly well-explored dataset assumption.

A. Corollary 4.6 in Jin et al. 2021, Theorem 2 in Duan et al. 2020, and Theorem 3.2 in Xie, Tengyang, et al. 2021 all depend on $1/\lambda_{\min}(\Sigma_h)$, which, in fact, requires the dataset to be uniformly well-explored.

In terms of the correctness of Assumption 4.3, note that in the proof of Theorem 5.2 on Page 13 in the Appendix, $\lVert \Lambda_h^{-1}\rVert_{\operatorname{tr}(\phi(s,a))}\le \lVert \Lambda_h^{-1}\rVert^{1/2} \sum_{i=1}^d \lVert \phi_i(s,a) 1_i\rVert.$ In fact, $\sum_{i=1}^d \lVert \phi_i(s,a) 1_i\rVert = 1$ due to our Definition 2.1, and thus the key to this proof is to upper bound the $\lVert \Lambda_h^{-1}\rVert^{1/2}$ with sufficient high probability. Note that $\Lambda_h$ is N samples with mean as $E[\phi(s_h, a_h)\phi(s_h, a_h)^\top]$ and thus we can apply concentration argument to upper bound it, which is the same as the proof of Corollary 4.6 in Appendix B.4 in Jin et al. 2021. We will add the full proof to show the correctness of this assumption in the next version.

Q. The justification of Assumption 4.4.

A. Assumption 4.4 is common in literature when considering KL divergence. In fact, Zhou et al. (2021b) heavily relies on Assumption 4.4, as their results involve e^{R_\max/\underline{\beta}} in their constant $C_{\pi}$ and $C$ (See their proof of Lemma 5, Theorem 1, and Theorem 2 in the supplement material). Panaganti & Kalathil (2022) only consider the total-variation distance instead of KL divergence. We agree with the reviewer's understanding of the relationship between $\rho$ and $\beta^*$. By the completion of this paper, all existing literature relies on the same or similar assumption for finite-sample results under KL divergence. Yang, Wenhao et al. 2021 does not rely on this assumption. Instead, their results depend on $1/\underline{p}$ where $\underline{p} = \min_{p(s'\lvert s,a)>0} p(s'\lvert s,a)$ and is also restricted in some cases.

Q. Correctness of the dependence on the dimension.

A. The dependence on the dimension should be $d^{2}$ for the Sufficient Coverage case, which matches Blanchet et al. (2023). This error comes from the fifth inequality on Page 19, where $\sum_{i=1}^d \sqrt{\frac{1}{1+c^{\dagger}N}}$ should be upper bounded by $d$, rather than $d^{1/2}$. However, for the well-explored setting, the dependence on $d$ is indeed $d^{1/2}$, as the condition is stronger than sufficient-coverage, and by Definition 2.1, the dependence reduces to bounding $\sum_{i=1}^d \lVert \phi_i(s,a)1_i\rVert =1$.

Q. Correctness of the dependence on the model misspecification.

A. We agree with the reviewer's understanding. The transformation from the error in the transition model $\xi$ to the final suboptimality is not a one-step proof. In fact, we need to follow some parts of the proof when non-model misspecification happens, and only finite-sample error occupies. We will update the results accordingly.

Q. Typos.

A. We will correct the definition of the value function and the $Q$ function as $V_h^{\pi}(s,a) = E_{P}^{\pi}[\sum_{h'=h}^H r_{h'}(s_{h'}, a_{h'})\lvert s_{h} = s]$ and $Q_h^{\pi}(s,a) = E_{P}^{\pi}[\sum_{h'=h}^H r_{h'}(s_{h'}, a_{h'})\lvert s_{h} = s, a_{h} = a]$, as well as the misuse of RLS as RTS.

Q. Stronger data condition vs. weak data condition.

A. A strong data condition refers to our Assumption 4.3 (Uniformly Well-explored Dataset). Weak data condition refers to Assumption 5.1 (Sufficient Coverage).

Thank you for your detailed reading and valuable comments. Please find our response to the comments below:

Q. Confusion about the notations

A. We sincerely apologize for the confusion caused by the excessive use of notations. We will clarify them one by one in the following:

1. $R^{rob}(\pi, P)$ is abbreviated and used as $R(\pi, P)$ since we no longer discuss the non-robust return and only focus on the distributionally robust one.

2. We acknowledge the reviewer's understanding of $P_{1,h}^*$. To be precise, $P_1^*$ is a transition kernel in $\mathcal{P}_{1,h}(s,a)$ satisfying the last equality of Equation (14). The two arg min notions below (13) and the $P_2^*$ notion are supposed to be removed. In fact, Lemma E.2 should be corrected, as $P_h^*$ and $\widehat{P}_h^*$ are transition kernels in $\mathcal{P}_h(s,a)$ and $\widehat{\mathcal{P}}_h(s,a).$ Sorry for the typo from the older version, but the proof is established for the correct statement.

3. We agree with the reviewer's understanding of the expectation over $\mathcal{P}_{1,\pi}^{*}$ and will polish the notation in the next version.

4. $P^*$ denotes a transition kernel in $\mathcal{P}$.

Thank you for your detailed reading and valuable comments. Please find our response to the comments below:

Q. Explain $\psi$ in Definition 2.1.

A. $\phi_{h,i}(s)$ represents the $i$-th coordinate of $\psi_h(s)$. To understand the second equation in Definition 2.1, it can also be understood as the i-th hidden state among the total d hidden states. $\psi_{h,i}(s)$ can be understood as the probability to transition from the i-th hidden state to the original state $s$.

Q. What is the difference between pessimistic and robustness in the algorithm PDRVI-L.

A. Pessimism is the general learning principle for designing sample-efficient algorithms that can overcome poor data coverage, i.e., the finite samples in the offline dataset are not diverse enough to cover sufficient state-action pairs. This is achieved by subtracting $\gamma_h \sum_{i=1}^d \lVert \phi_i(s,a) 1_i\rVert_{\Lambda_h^{-1}}$ from the estimated $Q$-value. Robustness is aimed at tackling the possible distribution mismatch from the training dataset to the test environment.

Q. Extension to the infinite horizon discounted MDP.

A. Most of the techniques, including the algorithmic design as well as the finite-sample analysis of the KL divergence, can be extended to the infinite horizon discounted MDP. The gap comes from transferring the finite-sample error from the estimation of the KL ambiguity set to the policy value. While in the finite horizon settings, each stage is independent and can be analyzed backward, in the infinite discounted MDP, we need to analyze the sensitivity of the fix-point solution of the Bellman optimality equation to the finite-sample error .

Thank you for your review and positive comments on our paper. Please find our response to the comments below:

Q. What if it is not KL? What motivates it to be KL in practice?

A. Our algorithm, as well as the theoretical results, can easily be extended to other divergences, such as $\chi^2$ divergence and total variation distance, by simply replacing the dual function of KL distance in Equation (3) with those of other distances. The analysis of other $f$-divergences is easier than that of the KL divergence, as KL involves high nonlinearity concerning expectation and exhibits no Lipschitz when the dual variable gets close to zero. A similar example to generalize the results under KL divergence to other distances can be found in Yang, Wenhao et al. 2022. We present the results under KL divergence to facilitate the presentation and keep the reader's focus on our main contributions, primarily in our algorithmic design and various data coverage extensions.

Q. How do we know $\rho$ in practice? What if we $\rho$ misspecified?

A. Finding the optimal ambiguity set radius in distributionally robust optimization (DRO) involves a balance between robustness and conservatism. There isn't a one-size-fits-all approach, as the choice of the ambiguity set radius depends on the specific problem and the available information about the uncertainty in the data. Statistical methods such as cross-validation or bootstrapping can be employed to estimate the uncertainty in your data, as seen in Wang, Jie et al. 2021. A more advanced approach is to adopt Bayesian inference, and in general, tuning $\rho$ itself is an interesting research direction that we leave for further work. If $\rho$ is misspecified, the algorithm may fail to achieve sufficient robustness or may tend to be overly conservative.

Q. In what problems does this kind of model mismatch actually happen? Is there any motivating real example?

A. Any domain where the dynamics of the data could change gradually or even rapidly would be an ideal testbed for our algorithm. Even when the dynamics remain consistent between the training environment and the test one, there may be insufficient data for accurate enough estimation of the dynamics, leading to non-negligible finite-sample error. For example, in power control schemes, the algorithm is optimized for estimated parameters, such as wireless channel quality (Zhang, Haijun, et al. 2020); or in financial portfolio management, financial data always exhibits high non-stationarity (Huang, Biwei, et al. 2020).

[1] Yang, Wenhao, Liangyu Zhang, and Zhihua Zhang. "Toward theoretical understandings of robust Markov decision processes: Sample complexity and asymptotics." The Annals of Statistics 50.6 (2022): 3223-3248.

[2] Wang, Jie, Rui Gao, and Yao Xie. "Sinkhorn distributionally robust optimization." arXiv preprint arXiv:2109.11926 (2021).

[3] Gupta, Vishal. "Near-optimal Bayesian ambiguity sets for distributionally robust optimization." Management Science 65.9 (2019): 4242-4260.

[4] Zhang, Haijun, et al. "Power control based on deep reinforcement learning for spectrum sharing." IEEE Transactions on Wireless Communications 19.6 (2020): 4209-4219.

[5] Huang, Biwei, et al. "Causal discovery from heterogeneous/nonstationary data." The Journal of Machine Learning Research 21.1 (2020): 3482-3534.

Thank you for the rebuttal. The assumption of the knowledge of $\rho$ is the main issue I’ve generally had with DRO. This parameter is sensitive and misspecifying this parameter basically makes the framework not useful anymore. I encourage the authors to discuss and emphasize this issue clearly in the manuscript.

My third question was about the model mismatch of $\psi$ in the linear mdp. Why does distribution shift only happen with $\psi$ but not $\phi$? This assumption looks quite unnatural to me, more as a way to sidestep the results of Tamar et al 14 and to fit it the linear mdp, rather than motivated by real problems about model mismatch.