Paper: A Policy Gradient Method for Confounded POMDPs (ICLR-2024)

Reviewer comment (Novelty):
From my viewpoint, by looking into the previous line of works, the main theory (Sections 6.1 & 6.2) of this work is mostly based on **(i)** the theoretical understanding of statistical analysis for using minimax estimator to solve bridge functions in confounded offline POMDP settings, e.g., [1, 2, 3]; **(ii)** the analysis of global convergence of policy gradient ascent methods in standard MDP settings, e.g., [4, 5]. So the technical contributions of the main theory part are somehow weakened given these prior works.

Author response:
We are grateful for your valuable feedback and the recognition you've given to our work. We would like to address your concerns in this response.

> Weakness 1 and question 1 regarding technical contributions in Sections 6.1 & 6.2

Thank you for acknowledging **one of our main theoretical contributions in the derivation of the identification results** for the policy gradient.  In addition to this, we would like to highlight some crucial technical contributions in estimation and optimization presented in Sections 6.1 and 6.2. Below we provide a detailed summary. 

In the theoretical results related to the estimation of the policy gradient (section 6.1), our main
technical contributions are as follows. 

- **Novel Error Decomposition.** We derived a novel decomposition of error for policy gradient estimation, outlined in Appendix F.1 or equations (63)(64)(65). Existing works such as [1, 2, 3] typically consider the decomposition of the estimation error of *policy value* into multiple one-step errors caused by min-max estimation for the conditional operator of the *value* bridge function. In contrast, we study the decomposition of the estimation error of *policy gradient*, which turns out to involve one-step errors caused by min-max estimation for *both* the conditional operator of the *value* bridge function and the *gradient* bridge functions. This unique decomposition arises from our method of policy gradient identification, incorporating policy value identification results as assistance. Importantly, the derivation of this decomposition is not a trivial extension from existing works. Such error decomposition enhances the understanding of the source of all involved statistical errors, and can be potentially used to design more efficient or robust estimators for policy gradient in confounded POMDPs. For example, a direct idea is to simply consider more efficient/robust off-policy evaluation methods in confounded POMDPs and see if their incorporation can improve the efficiency/robustness of the policy gradient estimators by looking at the proposed error decomposition of the policy gradient.

- **Focus on Policy Gradient Estimation with Dimension $>1$.** In contrast to existing works [1, 2, 3], which primarily focused on policy value estimation with dimension equal to 1, our work uniquely concentrates on policy gradient estimation with dimensions greater than 1. This emphasis, driven by the significance of history-dependent policies, introduces an additional layer of complexity to the analysis. Specifically, the dimensionality of the policy parameter space, denoted as $d_\Theta$, becomes a crucial factor in our theoretical results. Therefore, our derivation of upper bounds with explicit dependence on $d_\Theta$ facilitates understanding of the complexity in higher-dimensional min-max estimation.

- **Upper Bounds with Explicit Dependence on Key Parameters (in the Nonparametric Settings).** In contrast to [1,2,3], our work provides upper bounds in terms of all key parameters when function classes are assumed to be  RKHSs (see details in Appendix B). This allows for non-parametric estimation. Existing works often provide examples only under linear settings or with finite VC dimensions, making our results more versatile and applicable to a broader range of scenarios.

In the theoretical results related to the global convergence of the proposed algorithm (section 6.2), our main technical contributions are as follows. 

- **Extension from MDPs to POMDPs.**  Building upon the theoretical results of [4, 5], we extend their results from MDPs to POMDPs. This extension necessitates the incorporation of specific lemmas tailored for POMDPs, including the policy gradient lemma and the performance difference lemma (Lemmas 3, 8, 10). Our work thus contributes to advancing theoretical understanding in the context of POMDPs.

- **Finite-horizon v.s. Infinite-horizon**. While [4,5] focus on the infinite-horizon discounted MDPs, we study the finite-horizon settings. Therefore, our work requires us to study the upper bound with explicit dependence on the length of horizon $T$, especially when history is included in our POMDP settings.

References:

[1] Bennett, Andrew, and Nathan Kallus. "Proximal Reinforcement Learning: Efficient Off-policy Evaluation in Partially Observed Markov Decision Processes." 

[2] Shi, Chengchun, et al. "A Minimax Learning Approach to Off-policy Evaluation in Confounded Partially Observable Markov Decision Processes." 

[3] Lu, Miao, et al. "Pessimism in the Face of Confounders: Provably Efficient Offline Reinforcement Learning in Partially Observable Markov Decision Processes." 

[4] Agarwal, Alekh, et al. "On the Theory of Policy Gradient Methods: Optimality, Approximation, and Distribution Shift." 

[5] Liu, Yanli, et al. "An Improved Analysis of (Variance-reduced) Policy Gradient and Natural Policy Gradient Methods."

Reviewer comment (Theory):
As stated in 1., a consequence is that the theoretical assumptions regarding the policy gradient analysis (Section 6.2) are mostly adapted directly from those for MDP setups. How to understand these assumptions in POMDP settings with a history dependent policy class is less discussed.

Author response:
> Weakness 2 and question 2 regarding discussions on assumptions (made in Section 6.2) in POMDP settings with a history dependent policy class and how the dependence on history changes (or does not change) the difficulty of doing policy gradient

Thank you for your suggestions. We have included such a discussion in Appendix B.7 of the revised manuscript. To address your concern, below we discuss how the inclusion of history impacts the complexity of policy gradient ascent for confounded POMDPs under the offline settings across statistical, optimization, and computational aspects. Specifically, we illustrate this with the example of a log-linear policy: 
$$\pi_{\theta_t}(a_t\mid o_{t},h_{t-1}) \propto \exp(\theta_t^\top \phi_t(a_t,o_t,h_{t-1})).$$

- **Statistical Aspect.** In terms of statistical estimation, we examine the upper bound presented in Theorem 2 for the policy gradient estimation error. First, the dimension of the policy space $d_\Theta$ implicitly depends on $T$. At each step $t$, $\phi_t$ is a feature map from a space of dimension $t|\mathcal{A}|dim(\mathcal{O})$ to a space of dimension $d_{\Theta_t}$. To preserve information adequately, it's reasonable to assume that $d_{\Theta_t}$ grows with $t$. In contrast, in MDP settings, a fixed $d_{\Theta_t}$ assumption may suffice for all $t$. Second, the function classes for the bridge functions and test functions should also be rich enough because they are functions of histories that scale with $t$ at each stage $t$. Therefore the complexity of the function classes $\gamma(\mathcal{F},\mathcal{B})$ also grows when the number of stages $T$ increases. These factors collectively contribute to the complexity of estimation when dealing with history-dependent policies.

- **Optimization Aspect.** We discuss the assumptions in section 6.2 in response to your concern. Firstly, assumption 5(a) (optimal policy coverage) can be dropped with the introduction of a notion we term ``transferred compatible approximation error." We respectfully direct you to our response regarding "Regarding Optimal Policy Coverage (Assumption 5(a))" to Reviewer KDFK for more details. We then focus on the revised manuscript in the rest. Regarding Assumption 5(a) $\beta$-smoothness, when $\phi _t(a _t,o _t,h _{t-1})$ is in a compact region with $\\|\phi _t(a _t,o _t,h _{t-1})\\| _{\ell _2}\le B$, it is straightforward to show that $\left\\|\nabla _{\theta _t} \log \pi _{\theta _t}\left(a _t \mid o _t, h _{t-1}\right)-\nabla _{\theta _t} \log \pi _{\theta _t^{\prime}}\left(a _t \mid o _t, h _{t-1}\right)\right\\| _{\ell^2} \leq B^2\left\\|\theta _t-\theta _t^{\prime}\right\\| _{\ell^2}$. This assures that the Lipschitz constant remains unaffected by the historical dependence. Assumption 5(b) (Positive definiteness of Fisher information matrix)) requires the positive definiteness of Fisher information matrix, where the constant $\mu$ implicitly depends on the number of stage $T$. Intuitively, obtaining a large $\mu$ becomes more challenging in the context of history-dependent policies due to the high dimensionality of the history space. A potential approach to mitigate this challenge involves mapping the history to a lower-dimensional space that retains sufficient information. For Assumption 5(c) (L-smoothness), the scale of the constant $L$ increases when considering history, compared to the standard MDP settings. See Appendix B.3 for more details. It's evident that the historical dependence amplifies the complexity through Assumptions 5(b) and 5(c). Furthermore, the dimension of the parameter space $d _\Theta$, implicitly depending on $T$, heightens the challenge of gradient ascent for the same reasons elucidated in the statistical aspect.

-  **Computational Aspect.** We have added the analysis of the computational complexity of Algorithm 1 in Appendix I.2 of the revised manuscript, yielding a time complexity of $O(Kd _{\Theta}N^2 T\max\\{T,N\\})$. Compared to the standard MDP settings, the introduction of history-dependence primarily increases the computational complexities in two steps: the evaluation of the empirical kernel matrix to derive the closed-form solution for the min-max optimization problem (Equation 9) and the update of the policy parameter. For the empirical kernel matrix evaluation, kernel functions must be computed in the history space, a task that scales with $T$. Furthermore, we need $d _{\Theta}$ operations to update each coordinate, where $d _\Theta$ implicitly depends on $T$.

We hope that we have addressed your concerns, and are open to any further questions or feedback. If our revised manuscript and rebuttal more closely meet your expectations for the paper, we respectfully ask you to reconsider your initial rating.

Reviewer comment (Novelty):
Continued from Weakness 1., I would appreciate it if the authors could highlight more on the technical contributions behind the main theory, especially when compared with the previous line of works listed.

Author response:
We are grateful for your valuable feedback and the recognition you've given to our work. We would like to address your concerns in this response.

> Weakness 1 and question 1 regarding technical contributions in Sections 6.1 & 6.2

Thank you for acknowledging **one of our main theoretical contributions in the derivation of the identification results** for the policy gradient.  In addition to this, we would like to highlight some crucial technical contributions in estimation and optimization presented in Sections 6.1 and 6.2. Below we provide a detailed summary. 

In the theoretical results related to the estimation of the policy gradient (section 6.1), our main
technical contributions are as follows. 

- **Novel Error Decomposition.** We derived a novel decomposition of error for policy gradient estimation, outlined in Appendix F.1 or equations (63)(64)(65). Existing works such as [1, 2, 3] typically consider the decomposition of the estimation error of *policy value* into multiple one-step errors caused by min-max estimation for the conditional operator of the *value* bridge function. In contrast, we study the decomposition of the estimation error of *policy gradient*, which turns out to involve one-step errors caused by min-max estimation for *both* the conditional operator of the *value* bridge function and the *gradient* bridge functions. This unique decomposition arises from our method of policy gradient identification, incorporating policy value identification results as assistance. Importantly, the derivation of this decomposition is not a trivial extension from existing works. Such error decomposition enhances the understanding of the source of all involved statistical errors, and can be potentially used to design more efficient or robust estimators for policy gradient in confounded POMDPs. For example, a direct idea is to simply consider more efficient/robust off-policy evaluation methods in confounded POMDPs and see if their incorporation can improve the efficiency/robustness of the policy gradient estimators by looking at the proposed error decomposition of the policy gradient.

- **Focus on Policy Gradient Estimation with Dimension $>1$.** In contrast to existing works [1, 2, 3], which primarily focused on policy value estimation with dimension equal to 1, our work uniquely concentrates on policy gradient estimation with dimensions greater than 1. This emphasis, driven by the significance of history-dependent policies, introduces an additional layer of complexity to the analysis. Specifically, the dimensionality of the policy parameter space, denoted as $d_\Theta$, becomes a crucial factor in our theoretical results. Therefore, our derivation of upper bounds with explicit dependence on $d_\Theta$ facilitates understanding of the complexity in higher-dimensional min-max estimation.

- **Upper Bounds with Explicit Dependence on Key Parameters (in the Nonparametric Settings).** In contrast to [1,2,3], our work provides upper bounds in terms of all key parameters when function classes are assumed to be  RKHSs (see details in Appendix B). This allows for non-parametric estimation. Existing works often provide examples only under linear settings or with finite VC dimensions, making our results more versatile and applicable to a broader range of scenarios.

In the theoretical results related to the global convergence of the proposed algorithm (section 6.2), our main technical contributions are as follows. 

- **Extension from MDPs to POMDPs.**  Building upon the theoretical results of [4, 5], we extend their results from MDPs to POMDPs. This extension necessitates the incorporation of specific lemmas tailored for POMDPs, including the policy gradient lemma and the performance difference lemma (Lemmas 3, 8, 10). Our work thus contributes to advancing theoretical understanding in the context of POMDPs.

- **Finite-horizon v.s. Infinite-horizon**. While [4,5] focus on the infinite-horizon discounted MDPs, we study the finite-horizon settings. Therefore, our work requires us to study the upper bound with explicit dependence on the length of horizon $T$, especially when history is included in our POMDP settings.

References:

[1] Bennett, Andrew, and Nathan Kallus. "Proximal Reinforcement Learning: Efficient Off-policy Evaluation in Partially Observed Markov Decision Processes." 

[2] Shi, Chengchun, et al. "A Minimax Learning Approach to Off-policy Evaluation in Confounded Partially Observable Markov Decision Processes." 

[3] Lu, Miao, et al. "Pessimism in the Face of Confounders: Provably Efficient Offline Reinforcement Learning in Partially Observable Markov Decision Processes." 

[4] Agarwal, Alekh, et al. "On the Theory of Policy Gradient Methods: Optimality, Approximation, and Distribution Shift." 

[5] Liu, Yanli, et al. "An Improved Analysis of (Variance-reduced) Policy Gradient and Natural Policy Gradient Methods."

Reviewer comment (Writing):
The paper introduces a significant number of symbols and notations, which might overwhelm readers, especially those who are less familiar with the topic. To improve clarity, the authors could consider providing a table or appendix that lists all the symbols used along with their definitions. Additionally, they might simplify the notation where possible and ensure that each symbol is clearly defined upon its first use. For example, the notation seems to be quite heavy at the end of Page 5 where $\mathcal{Z_t}$ and $\mathcal{W_t}$ are introduced.

Author response:
We sincerely appreciate your constructive feedback and your acknowledgment of our work. 

>Regarding notations

 Thank you for the valuable suggestion. We acknowledge the importance of clarity in notation. In response to your feedback, we have included a table in Appendix A that lists all the main notations used along with their definitions. Additionally, we have simplified the notations where feasible and ensured that each symbol is explicitly defined upon its initial use in the manuscript. See Section 5 in the revised version.

Reviewer comment (Theory):
The full coverage assumptions stated in Assumption 4(a) and Assumption 5(a) are indeed common in offline RL literature, but they bring about challenges and potential limitations to the proposed method in the paper.

Author response:
We sincerely appreciate your constructive feedback and your acknowledgment of our work. 

> Regarding Full Coverage Assumption (Assumption 4(a))

Thank you for your suggestions. In the following, we discuss on practical scenarios in which this assumption is either met or not, and propose a potential way to address this concern. We have also added this discussion in Appendix B.4 of the revised manuscript.

In certain real-world applications, such as the sequential multiple assignment randomized trials (SMART) designed to build optimal adaptive interventions, the assumption of full coverage is usually satisfied. This is because data collection in these trials is typically randomized, ensuring a comprehensive representation by the behavior policy. However, we acknowledge that in domains such as electronic medical records, meeting the full coverage assumption may pose challenges due to ethical or logistical constraints. 

To address scenarios where the full coverage assumption might not hold, we could incorporate the principle of pessimism into our approach. This involves penalizing state-action pairs that are rarely visited under the offline distribution. The idea of incorporating pessimism has been widely used in the offline RL literature for MDPs. For example, a practical implementation of this idea can be adapted from [Zhan et al., 2022], where a regularity term is added to the objective function to measure the discrepancy between the policy of interest and the behavior policy. By identifying and estimating the gradient of this modified objective function, we could potentially provide an upper bound on suboptimality and maintain a similar theoretical result by only assuming partial coverage, i.e. 
$$C_{\pi^b}^{\pi_{\theta^*}}:=\max_{t=1,..,T}(\mathbb E^{\pi^b}[(\frac{p_t^{\pi_{\theta^*}}(S_t,H_{t-1})}{p_t^{\pi^b}(S_t,H_{t-1})\pi^b_t(A_t\mid S_t)})^2])^{\frac{1}{2}}<\infty.$$
This partial coverage assumption only needs that the offline distribution $\mathbb{P}^{\pi^b}$ can calibrate the distribution $\mathbb{P}^{\pi_{\theta^*}}$ induced by the in-class optimal policy, which is a milder condition compared to the full coverage assumption.

Reviewer comment (Theory):
**Optimal Policy Coverage (Assumption 5(a))**: This assumption requires that the optimal policy $\pi^*$ is covered by all the policies in the class. While this condition ensures that the policy class is rich enough to contain the optimal policy, it might be too restrictive in practice.

Author response:
We sincerely appreciate your constructive feedback and your acknowledgment of our work. 

>Regarding Optimal Policy Coverage (Assumption 5(a))

Thank you for your comment on this assumption. Inspired by [Liu et al.][Masiha et al.], we can relax this assumption by introducing a transferred compatible function approximation error in the following way:

$$\varepsilon _{approx}:=\sup  _{\theta \in \Theta}  \sum _{t=1}^T\left(\mathbb{E}^{\pi _{\theta^*}}\left[\left(\mathbb{A} _t^{\pi _\theta}\left(A _t, O _t, S _t, H _{t-1}\right)-w _t^*(\theta)^{\top} \nabla _{\theta _t} \log \pi _{\theta _t}\left(A _t \mid O _t, H _{t-1}\right)\right)^2\right]\right)^{\frac{1}{2}},$$
where $w _t^*(\theta)$ denotes the optimal solution of the following optimization problem for each $\theta$:
$$\arg\min _{w _t}  \sum _{t=1}^T\left(\mathbb{E}^{\pi _{\theta}}\left[\left(\mathbb{A} _t^{\pi _\theta}\left(A _t, O _t, S _t, H _{t-1}\right)-w _t^{\top} \nabla _{\theta _t} \log \pi _{\theta _t}\left(A _t \mid O _t, H _{t-1}\right)\right)^2\right]\right)^{\frac{1}{2}}.$$

Note that this notion is similar to the definition (Def. C.2 in Appendix C) in our original paper, which also quantifies the richness of the policy class. It tends to be small when the policy class is sufficiently rich. In such cases, the requirement for optimal policy coverage (Assumption 5(a)) becomes unnecessary. In the revised manuscript, we have introduced this transferred compatible approximation error in Definition D.2, deleted the original optimal policy coverage assumption, and updated the results \& proof of Theorem 3.

Reviewer comment (Experiments):
While I typically do not complain about the empirical results of a theory paper, I do expect that the authors could show how to implement such estimation and algorithm on real practical RL problems.

Author response:
We appreciate your feedback and understand the importance of practical implementation in reinforcing the applicability of our proposed estimation and algorithm. In response to your suggestion, we have added discussion about the practical implementation of the proposed algorithm on real-world RL problems in Appendix I.3 of the revised manuscript. 

First of all, an exact implementation method has been  provided in Appendix I.1 of the manuscript. Specifically, assuming all function classes are reproducing kernel Hilbert spaces (RKHSs), the min-max optimization problem (9) comprises a quadratic concave inner maximization problem and a quadratic convex outer minimization problem. Consequently, a closed-form solution can be computed, as demonstrated in Appendix I.1.

In practical implementations on real RL problems, the exact execution of the proposed algorithm may sometimes computationally costly. For instance, when dealing with large sample sizes or a high number of stages, solving the min-max optimization problem (9) precisely in each iteration may become computationally expensive. Particularly, exact solutions with RKHSs involve calculating the inverse of an $N\times N$ (which takes $O(N^3)$) and evaluating the empirical kernel matrix on the space of history (taking approximately $O(N^2 T)$).  Additionally, when assuming other function classes for bridge functions and test functions, such as neural networks, achieving an exact solution for the min-max optimization problem (9) is not feasible.

To enhance computational efficiency and accommodate general function classes, an alternative approach is to compute the (stochastic) gradient of the empirical loss function in the min-max optimization problem (9) with respect to $b$ and $f$ respectively. Subsequently, updating $f$ and $b$ 
alternately using gradient ascent and gradient descent only once in each iteration can be performed. In this case, gradient estimation is expedited, contributing to faster computations. We acknowledge the need for further exploration to assess its efficacy in real-world applications, and we consider this an avenue for future research.

Reviewer comment (Writing):
In the third paragraph of Section 3, does the definition of history $\mathcal{H}_{t}$ lack the subscript on $\mathcal{O}$ and $\mathcal{A}$?

Author response:
Thank you for your question. We assume the observation space $\mathcal{O}$ and the action space $\mathcal{A}$ are the same for each stage $t=1,2,...,T$, eliminating the need for subscripts on 
$\mathcal{O}$ and $\mathcal{A}$. The space of history 
$\mathcal{H}_t$ can be conceptualized as a product space. One can generalize the setting by allowing $\mathcal{O}$ and $\mathcal{A}$ to be time-dependent.

Reviewer comment (Writing):
The notation $Z_{t}$ has already been defined in Section 3; however, another $\mathcal{z}$ is used in Section 5 as additional notation. Are they the same things?

Author response:
Thank you for pointing this out. They are not the same thing, and we have corrected the notations in the revised manuscript.  We now use $\mathcal E_t$ rather than $Z_t$ to denote the emission kernel in Section 3.