Agnostic Reinforcement Learning with Low-Rank MDPs and Rich Observations

Abstract: Ayush Sekhari, Christoph Dann, Mehryar Mohri, Yishay Mansour, Karthik Sridharan

Reviews and author responses:

The paper studies reinforcement learning in low-rank MDPs under an agnostic perspective, where the agent does not assume realizability of the optimal value function in some given functional space but rather has to find a near-optimal policy in a given policy class. For this problem, the authors propose a policy search algorithm based on uniform exploration and careful estimation of the expected returns of each policy. They derive an upper bound on its sample complexity for computing an \epsilon-optimal (within the class) policy that scales logarithmically with the number of policies but exponentially in the MDP rank. They further show a lower bound which confirms that this exponential dependence is unavoidable in the worst-case. Finally, they propose a variant of this algorithm that enjoys a more adaptive sample complexity scaling with the eigenvalues of the true transitions matrices.  Significance, novelty, and relevance
---

The paper considers a relevant problem, agnostic RL in MDPs with low-rank structure. While recently there have been many works on both low-rank MDPs and on dealing with misspecified models (i.e., non-realizable function classes), to my knowledge the perspective studied in this paper is novel. I found this contribution interesting and significant. In particular, I appreciated that the authors derived lower bounds which provide good insights into the complexity of this setting. While the proposed algorithm is not very exciting (it is simply an instance of uniform exploration), I found quite interesting the way expected returns are estimated by leveraging the low-rank structure. Overall I think that this work is a good starting point for studying this setting and tend to vote for acceptance.
 

Relation with prior work
---

To my knowledge, relevant prior works are discussed.

It might be good to also discuss [1], where the authors consider a fixed policy class and the goal is to compete with the best-in-class policy as in this agnostic setting. In that case, the MDP is assumed to be adversarially generated and the authors still show guarantees that scale with \log|\Pi|. Another paper with a similar aim is [2].

[1] Abbasi-Yadkori, Yasin, Peter L. Bartlett, and Csaba Szepesvári. "Online learning in Markov decision processes with adversarially chosen transition probability distributions." arXiv preprint arXiv:1303.3055 (2013).

[2] Azar, Mohammad Gheshlaghi, Alessandro Lazaric, and Emma Brunskill. "Regret bounds for reinforcement learning with policy advice." Joint European Conference on Machine Learning and Knowledge Discovery in Databases. Springer, Berlin, Heidelberg, 2013.


Clarity of writing
---

I found the paper very well-written and easy to read. I think the authors did a good job in providing intuition behind all results, in motivating their design choices, and in sketching the proofs. The paper seems mostly self-contained, except maybe for Section 6 which requires going back and forth from the appendix.


Detailed comments/questions
---

1. I was wondering how the agnostic setting studied here compares to the one of linear (low-rank) MDPs with misspecification (the one studied by the papers mentioned in Section 3). In the latter setting, not only the value functions become non-realizable under the given linear function class, but the MDP itself is no longer describable by a linear model, and hence it is no longer low-rank, since it is assumed linear only up to a (possibly small) non-linear perturbation. Here, instead, while the agent is agnostic to the representation of the optimal value function, the MDP is still describable by a low-rank model and the agent can leverage such structure. How do the two settings compare? Can we say that one is more general than the other? Or easier than the other?

2. Related to the previous point, the second-last paragraph of Section 3 mentions that the setting studied in this paper is more ambitious. I am not sure that I fully understand why. It seems to me that competing with the optimal policy of the given MDP under a non-realizability assumption (as in the works mentioned there) is more challenging than competing with the best policy in a given class. 

3. A possible limitation of the proposed algorithm is that it needs to enumerate all policies in the given class in order to estimate their returns, which is very inefficient when the number of policies is large. Related question: any idea on how to deal with infinite/continuous policy spaces?

4. One limitation of the lower bound (Theorem 2) is that it does not show what is the optimal dependence on K in the worst-case. The upper bound suffers a dependence of K^d, which I guess is due to importance sampling with a uniform exploration policy. Do you believe that it would be possible to extend the lower bound to more than 2 actions and show that such dependence is unavoidable in the worst-case?

5. Regarding uniform exploration, is there anything better you can do without changing the analysis much? For instance, if the final purpose is to estimate the return of the policies in \Pi only, can you take, e.g., a mixture of these for exploration? For instance, if no policy in \Pi plays a certain action in a given state, do we still need to explore that action?

Minor:
- line 106: h should be h' in r()
- line 133: an optimal policy
- equation below line 180: is that \Omega() of that factor?
- I would suggest moving Alg. 1 and 2 to page 4, closer to where they are described to ease readability
- are the policies in \Pi stochastic or deterministic? Sec. 2 mentions stochastic, but the importance sampling estimator in Alg. 2 is for deterministic
- (5) in Alg. 2: why not using the estimator \hat{R} for all stages up to 3d instead of d? Lemma 2 gives concentration for all of them
- line 240: figure 2 should be figure 1
- line 279: predict The authors discuss the limitations of their algorithm and leave their solution to future work. I do not see any societal impact.

 Dear authors,

Thank you for the detailed response and for clarifying all my doubts. I have updated my score accordingly.

 Thanks for the response. The authors have addressed my questions. This is a solid work and my recommendation remains acceptance. 

 Thank you for your review and positive support for our paper. 

 Thanks for a detailed review. We will add more details to make Section 6 clearer in the final version of the paper. 


&nbsp;


***It might be good to also discuss [1]  (and [2]) ....***

Thank you for pointing out these works, we will cite and discuss them. Note that the guarantees in [2] scale linearly with $|\Pi|$ as opposed to logarithmically, and while [1] achieves a strong $\sqrt{\ln |\Pi| T}$ regret bound they assume that the agent fully observed the current transition kernel and reward, instead of just a sample from it.

&nbsp;

## Response to questions ## 

1. ***How the agnostic setting studied here compares to the one of linear (low-rank) MDPs with misspecification***

   In the linear MDP setting with misspecification, as for example studied in Section 3.1 of Jin et al. 2020, the learner is assumed to be given access to $d$-dimentional state-action features $\phi$ which can represent the transitions and rewards in each state-action pair *approximately* linearly. Existing results show that learning with regret $O(\text{poly}(d, H) \sqrt{T} + \zeta d H^2 T)$ is possible where $\zeta$ is the maximum approximation error in any state-action pair. 
    Algorithms with this guarantee (implicitly) learn in the class $\Pi$ of policies that are linear in $\phi$ and by the approximation assumption, the best policy in this class $\Pi$ is at most $\zeta H$ suboptimal. Importantly, the linear in $T$ regret term has an additional $dH$ factor compared to $\zeta H T$ and, thus, these existing methods are not guaranteed to find the best policy in the class. The error within the class can be $d H$ times as large as the suboptimality of the best policy in $\Pi$ compared to the optimal policy. If, for example, the optimal policy can achieve a total value of $1$ but the best policy in $\Pi$ only achieves $0.75$, then guarantees of the form above are meaningless when $d H \geq 4$, even though a policy that achieves $0.75$ reward may still be highly desirable.
    Instead, in our setting, we aim for learning guarantees that are independent of how good the best policy in $\Pi$ is and which are still meaningful even when there is a significant gap between $\pi^*$ and the best policy in $\Pi$. 

   &nbsp;

   On a technical note, the approximate linear assumption in the works above only implies an *approximate* rank of $d$ of the transition kernel while we assume an exact rank of $d$ in our setting. However, our eigenspectrum-based guarantees in combination with our extension that adapts to an unknown rank suggest that our approach is robust to an approximate notion of rank. 


&nbsp;


2. ***Related to the previous point, the second-last paragraph of Section 3 mentions that the setting studied in this paper is more ambitious. I am not sure that I fully understand why....***

   Note that in our setting, the algorithms are guaranteed to return a policy $\pi'$ that is at most $\epsilon$ worse than $\tilde \pi = \arg\max_{\pi \in \Pi } V^\pi$, the best policy in the class. This means that $\pi'$ is $(\epsilon + \xi)$-optimal compared to $\pi^*$ where $\xi = V^\star - V^{\tilde \pi}$ is the error of best policy in the class.

   &nbsp;
 
    Conversely, the results provided in the cited works **do not** imply that their methods converge to a policy that is $\epsilon$ close to the best in class. Taking the example of LSVI-UCB mentioned in the response to the previous question, the algorithm is only guaranteed to converge to a policy that is $\epsilon + \zeta H (Hd - 1)$ close to $\tilde \pi$. When the error $\xi \leq \zeta H$ of $\tilde \pi$ is significant, e.g. $\xi > 1 / d$, then this guarantee is meaningless.
    In essence, the works referred to in that paragraph only provide non-trivial guarantees when the approximation error $\xi$ of the policy class is very small, while our guarantees hold independent of that error. 


&nbsp;


3. ***How to deal with infinite/continuous policy spaces?***

   We can run our algorithm with infinite / continuous policy spaces by discretizing the policy class (e.g., when given a Lipschitz policy parameterization). However, agnostic learning in general with continuous policy classes is an interesting research direction, and we leave it open for future work. 



&nbsp;

4. ***Lower bound dependence on $K^d$***

   We think that the lower bound should be extendable to give the factor of $K^d$ when there can be $K$ actions at every observation; however the technical details need to be worked out. The key idea is to change the current lower bound construction where we follow the blue arrow when the chosen action is same as that of $\pi^*$ on the given observation, and the red arrow for the other $K-1$ actions. We will add a discussion on this in the final version of the paper. 

&nbsp;

5. ***Regarding uniform exploration: Is there anything better you can do without changing the analysis much?***

     Our algorithm decouples estimation of the expected reward for the first $3d$ time steps, and value prediction via extrapolation. For policy class with additional structure, one may be able to improve over uniform exploration by sampling actions as suggested by the reviewer. However, as our lower bound suggests, such methods will not give any substantial improvement over uniform exploration in the worst case.


&nbsp;


***Response to minor comments***

Thanks for pointing the typos and writing suggestions; we will correct them. The importance sampling estimator used in Algorithm 2 can be easily adapted to work with stochastic policies as in Kearns et. al. 2000. We will add a discussion on this in the final version of the paper. 

 Thanks for your detailed review. Please find below response to all your questions below.


&nbsp;



* ***Can $3d$ (i.e., the number of steps performing uniform exploration) be shortened to $2d$ or perhaps $d$ ?***

   Our algorithm requires expected reward estimates for the first $3d$ time steps and we do not see any slack to improve this further in the current analysis. However, a different error propagation bound or improvements in the lower bound may bridge the gap further. 


&nbsp;


* ***Does the lower bound still hold if we have access to a generative model?***

    The lower bound should also hold under access to a generative model; however the technical details need to be worked out. At an intuitive level, in our lower bound construction, the only way to be able to find an $\epsilon$-suboptimal policy is by reaching the $+$-state along a path along the top row latent states (Figure 1). However, finding such a path already requires $\Omega(H^d)$ number of samples even under a generative model. Note that there is no other signal available to the learner (rich observation setting) except for the rewards in the lower bound construction. 


&nbsp;



* ***Response to minor comments:*** Thanks for pointing out all these typos. We will fix them in the final version of the paper. 

 Thanks for your detailed review. Please find below response to all your concerns. 


1. ***Do you think the $K^d$ factor  in the upper bound is necessary?*** 

     We think that the lower bound should be extendable to give the factor of $K^d$ when there can be $K$ actions at every observation; however the technical details need to be worked out. The key idea is to change the current lower bound construction where we follow the blue arrow when the chosen action is same as that of $\pi^*$ on the given observation, and the red arrow for the other $K-1$ actions. We will add a discussion on this in the final version of the paper.


&nbsp;


2. ***Is the optimization problem in (4) computationally hard in general?***

   Clearly, the optimization problem in $(4)$ can be solved in time exponential in $d$ by brute search over a $d-$dimension grid over $[-1, 1]^d$ to find $\lambda$. Note that the sample complexity is already $\Theta(H^d)$ and this dominates the running time.  It is a very interesting research direction whether this optimization problem can be implemented efficiently, e.g., using regression based techniques, and we hope to pursue this in future work.     

&nbsp;


3. ***Why the sample complexity doesn't depend on the magnitude of the eigenvalues?***

    The eigenvalues of a stochastic transition matrix (as in our case) are bounded by $1$ in magnitude, and thus our worst case upper bounds given in Theorem 1 **do not depend** on the norm of the eigenvalues. An improvement over this worst case upper bound is possible when the eigenvalues have small magnitude. We present this improvement in our adaptive upper bound in Theorem 3 on page 8.   


This paper studies RL with rich observation and low rank structure. It designs an algorithm that can work without the realizability assumption. The algorithm is guaranteed to find a good policy within a number of samples exponential in rank $d$. An almost matching lower bound is also provided.  This paper is well written and easy to follow. The key idea is cute. The most interesting point to me is they put no assumption on the policy class and they don’t assume access to any feature representation of the underlying low rank structure. The exponential dependence is also well justified. 


Questions:

Do you think the factor $K^d$ in the upper bound is necessary? The reason I am asking this question is that in linear MDPs, the optimal sample complexity doesn’t depend on $K$. Although the assumption in this paper is strictly weaker than linear MDPs, I am still curious if it’s possible to avoid this dependence or prove it’s indeed unavoidable. 

Is the optimization problem in (4) computationally hard in general?

Why the sample complexity doesn't depend on the magnitude of the eigenvalues? Is it hided in the $\mathcal{O}$-notation or you simply assume it's upper bounded by some constant? If so, please write it as an assumption or state the dependence explicitly in the theorem.  This work is purely theoretical. I am not aware of any potential negative societal impact.

This paper studies RL with linear function approximation when transition is low rank, but the features are unknown. The goal is to compete against a policy class that may not contain any near-optimal policy. The paper proposes an algorithm based on autoregression and shows a sample complexity bound that is exponential in the rank; a corresponding lower bound suggests that this dependence is unavoidable.  This paper studies the setting where: (1) only realizability under unknown features (factorizations) is assumed; (2) only an agnostic policy class is given. Compared to existing work on linear function approximation RL (e.g. [Agarwal et al. 2020]), the assumption made in this paper is among the weakest and least restrictive. It is then perhaps not surprising that the final bounds have an exponential dependence on $d$, which is further showed to be necessary. Overall I think the results and approach in this paper are novel and would be very interesting to RL theory community.

The basic algorithmic idea is to collect samples via uniform exploration for the first $3d$ time steps, and then evaluate each policy via autoregression. This framework of data collection-evaluating all policies is a bit similar to some prior work (e.g. [Wang et al. 2020]), but the key innovation is using autoregression relations for policy evaluation, which is a nice way to exploit linearity when the feature/factorization is unknown.

I have quickly checked the analysis in this paper and have found no issues with correctness. Also most of my questions have been well addressed by the discussions. I only have two small questions for the authors:

- Can $3d$ (i.e. the number of steps performing uniform exploration) be shortened to $2d$ or perhaps $d$? It would be nice if the exponents can be closer to that of the lower bound.

- Does the lower bound still hold if we have access to a generative model?

----

Minor comments/typos:

- In Lemma 13: Should it be assumed that $|\lambda_i|<1$ for all $i\in [d]$ in the statement? Otherwise $\hat\lambda=\lambda$ may not be feasible.
- In the proof of Theorem 1, should it be mentioned that all eigenvalues of $T^\pi$ have magnitude $\le 1$ when evoking Lemma 13? It just seems a bit odd to me if the proof works without using the fact.
- In Lemma 11: is $\lambda_k$ supposed to be $\lambda_d$?
- On Line 803 and 819: should references to Definition 2 be Definition 3?
- Line 847, Eq. (21): should $\lambda$ be $\hat\lambda$?
- Line 293: Should $K$ it be $H$?
- Line 293: Probably should be $1+\lambda+...+\lambda^{H-1}$.

-----

References:

[Agarwal et al. 2020] Flambe: Structural complexity and representation learning of low rank mdps.

[Wang et al. 2020] Is Long Horizon Reinforcement Learning More Difficult Than Short Horizon Reinforcement Learning? This work is purely theoretical and do not present foreseeable negative social impact.

This paper firstly discusses the possibility of RL without assuming optimal value function realizability in a rich observation setting. Relying on a uniform low-rank assumption, a novel algorithm is proposed to return a near-sub-optimal policy and the sample complexity scales polynomially in log(|\Pi|) and H while exponentially in d, which is further proved to be necessary.  This paper discusses an interesting but less-studied question and provides systematic answers (sufficient, necessary, and improved sufficient results). I appreciate its technical novelty, theoretical completeness, and well-structured writing.  Yes.

Meta-review: The paper considers the problem of learning near-optimal policies in low-rank MDPs and arbitrary policy classes, and proposes an algorithm with nontrivial guarantees in this setting, constituting one of the first theoretical guarantees in any agnostic learning scenario within the context of reinforcement learning. This is a problem of significant interest, and the reviewers all agreed that the paper provides some interesting technical innovations and clearly advances the state of the art in the area.

Overall, this is an excellent paper that clearly deserves to be published at the conference.

Decision: accept