Paper: Robustify the Latent Space: Offline Distributionally Robust Reinforcement Learning with Linear Function Approximation (ICLR-2024)

Reviewer comment (Theory):
I don't get the method in the work, how do you tackle the uncertainty from the dataset? Like in the tabular setting (Shi & Chi, 2022) and previous offline RL works, a penalty term is subtracted for those less visited states. How do you handle it in your method?

Author response:
**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.

Reviewer comment (Theory):
The result seems not too surprising to me. The work is under the linear MDP setting, and the uncertainty set is defined w.r.t. $\psi$. This setting is a little bit easy to me, as many nice properties hold under it, like the completeness of the linear function class.

Author response:
**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.

Reviewer comment (Theory):
For the uniformly well-explored dataset assumption, the authors state that it is adopted in Jin et al. (2021); Duan et al. (2019) and Xie et al. (2021). Can you specify respectively in which assumptions/theorems they use the uniformly well-explored dataset assumption? Moreover, in the proof of Th 5.2, in the end there is a term like this, $\Vert\sum_{i=1}^d\phi_i(s,a)\bf{1}_i\Vert$, which is unique in your setting. But the assumption is in the form of $E[\phi(s_h,a_h)\phi(s_h,a_h)^{\top}]$. The proof is omitted by claiming using similar steps in the proof of Corollar 4.6 in Jin et al. (2021). More details should be provided to show the correctness of this assumption.

Author response:
**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.

Reviewer comment (Theory):
On page 8, it seems to me that section 5.2 of model misspecification comes out of nowhere. There is no model and formulation here, such as what is that uncertainty set for the non-linear transition kernel. In the proof, the authors simply bound the infimum of two uncertainty sets like as follows:  
$$\inf_{P_{h+1} \in \mathcal{P_{h+1}}}E_{P_{h+1}}[V_{h+1}]-\inf_{P_{h+1} \in \tilde{\mathcal{P}}}E_{P_{h+1}}[V_{h+1}] \leq (H-h)\xi.$$  
I don't think this is correct. Please provide more argument.

Author response:
**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.

Reviewer comment (Theory):
At the beginning of page 3, the definition of value function and Q-function are wrong. There is no definition of $N$, $d$, $\delta$.

Author response:
**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.

Reviewer comment (Writing):
Misuse of RLS as RTS on page 4.

Author response:
**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.

Reviewer comment (Theory):
Should $\phi(s,a)\geq 0$, which is necessary in the definition of $d$-rectangularity? The current definition in linear MDP cannot ensure that.

Author response:
**Q.** The necessity of the non-negativity of the feature map $\phi(s,a)$.

**A.** Yes, we should impose $\phi(s,a)\ge 0$ in Definition 2.1, and we will add it in the next version.

Reviewer comment (Theory):
I suppose $\psi_{h,i}(s)$ represents the $i$-th coordinate of $\psi_{h}(s)$? Can the author explain why $\psi$ satisfies the second equation in Definition 2.1?

Author response:
**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$.

Reviewer comment (Theory):
Robustness by definition already means being pessimistic. I wonder what is the different of pessimistic and robustness in the algorithm PDRVI-L?

Author response:
**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.

Reviewer comment (Theory):
Does the technique and result extend to infinite horizon discounted MDP and why?

Author response:
**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 .

Reviewer comment (Writing):
There exist many confusing notations or not proper abuse of notations. For example, $R^{rob} (\pi, P) $ is abbreviated as $R(\pi, P)$, which actually already represents the non-robust expected return, which causes some confusion about the notation $R$ later. More are listed later in details.

Author response:
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.

Reviewer comment (Theory):
The technical correctness is not clear to the reviewer since the Lemma E.2 in Appendix seems not correct. The reason is that: firstly, in Lemma E.1, there exists some transition kernel $P\_{exist} \in \mathcal{P}\_1$ (denoted as $P_{1,h}^*$ in the paper) that can leads to the equality in (14). However, this applying $P_{1,h}^*$ also defined in a $\arg\min$ notion below (13). I guess actually $P_{1,h}^*$ is just some transition kernel in $\mathcal{P}_{1}$ but not in a $\arg\min$ notion. Then applying Lemma E.1 to Lemma E.2 won't leads to the results that this work claims. Since Lemma E.1 and E.2 are essential part in the proof, it makes the technical correctness doubtful for the reviewer.

Author response:
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.

Reviewer comment (Writing):
As mentioned, what does the notation $P_{1,h}^*$ and $P_{2,h}^*$ mean in the proof of Lemma E.1, it has been re-defined at least twice.

Author response:
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.

Reviewer comment (Writing):
The expectation over $P_{1,\pi}^*$ in Lemma E.1 is not clear. At each time step $h$, the expectation is determined by $P_{1,h,\pi}^*$ according to the reviewer's understanding. Carefully polishing the notation and presentation will be very helpful for the reader to check and follow the proof.

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

Reviewer comment (Writing):
What does $P^*$ mean in the proof of Lemma E.2, which has not been mentioned?

Author response:
$P^*$ denotes a transition kernel in $\mathcal{P}$.

Reviewer comment (Theory):
The problem setting is limited — What if $D$ is not KL? What motivates $D$ to be KL in practice?

Author response:
**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.

Reviewer comment (Experiments):
How do we know $\rho$ in practice? What if we misspecified $\rho$?

Author response:
**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.