Provably Efficient Reinforcement Learning in
Partially Observable Dynamical Systems

Large state space and partial observability are two key challenges of Reinforcement Learning (RL). While recent advances in RL for fully observable systems have focused on the challenge of scaling RL to large state space in both theory and in practice using rich function approximation, the understanding of large scale RL under partial observability is still limited. In POMDPs, for example, a core issue is that the optimal policy is not necessarily Markovian since the observations are not Markovian.

A common heuristic to tackle large scale RL with partial observability in practice is to simply maintain a time window of the history of observations, which is treated as a state to feed into the policy and the value function. Such a window of history can be often maintained explicitly via truncating away older history (e.g., DQN uses a window with length 4 for playing video games [49]; Open AI Five uses a window with length 16 for LSTMs [4]). Since even for planning under partial observations and known dynamics, finding the globally optimal policy conditional on the entire history is generally NP-hard (due to the curse of the history) [44, 55, 24], searching for a short memory-based policy can be understood as a reasonable middle ground that balances computation and optimality. The impressive empirical results of these prior works also demonstrate that in practice, there often exists a high-quality policy (not necessarily the globally optimal) that is only a function of a short window of recent observations. However, these prior works that search for the best memory-based policy unfortunately cannot ensure sample efficient PAC guarantees due to the difficulty of strategic exploration in POMDPs. The key question that we aim to answer in this work is:

Can we design provably efficient RL algorithms that agnostically learn the best fixed-length memory based policy with function approximation?

We provide affirmative answers to the above question. More formally, we study RL for partially observable dynamical systems that include not only the classic Partially Observable MDPs (POMDPs) [52, 56, 64], but also a more general model called Predictive State Representations (PSRs) [46]. We design a model-free actor-critic framework, named PO-Bilinear Actor-Critic Class, where we have a policy class (i.e., actors) that consists of policies that take a fixed-length window of observations as input (memory-based policy), and a newly introduced value link function class (i.e., critics) that consists of functions that take the fixed-length window of history and (possibly multi-step if the system is overcomplete) future observations as inputs. A value link function class is an analog of the value function class tailored to partially observable systems that only relies on observable quantities (i.e., past and future observations and actions). In our algorithm, we agnostically search for the best memory-based policy from the given policy class.

Our framework is based on the idea of a newly introduced notion of value link function equipped with future observations. While the idea of using future observations has been used in the literature on POMDPs, our work is the first to use this idea to learn a high-quality policy in a model-free manner. Existing works discuss how to use future observations only in a model-based manner [8, 25]. By leveraging these model-based viewpoints, while recent works discuss strategic exploration to learn near-optimal policies, their results are either limited to the tabular setting (and are not scalable for large state spaces) [36, 22, 3, 74, 43] or are tailored to specific non-tabular models and unclear how to incorporate general function approximation [65, 42, 12]. We break these barriers by devising a new actor-critic-based model-free view on POMDPs. We demonstrate the scalability and generality of our PO-bilinear actor-critic framework by showing PAC-guarantee on many models as follows (see Table 1 for a summary).

Observable Tabular POMDPs. In tabular observable POMDPs, i.e., POMDPs where multi-step future observations retain information about latent states, the PO-bilinear rank decomposition holds. We can ensure the sample complexity is $\mathrm{\mathrm{Poly}}(S,A^{M},O^{M},A^{K},O^{K},H,1/\sigma_{1})$ where $\sigma_{1}=\min_{x}\|\mathbb{O}x\|_{1}/\|x\|_{1}$ ($\mathbb{O}$ is an emission matrix),and $S,A,O$ are the cardinality of state, action, observation space, respectively, $H$ is the horizon, and $K$ is the number of future observations.¹¹1In Section 8, we discuss how to get rid of $O^{M},O^{K}$ using a model-based learning perspective. The intuition is that a tabular POMDP’s model complexity has nothing to do with $M$ or $K$, i.e., number of parameters in transition and omission distribution is $S^{2}A+OA$ (even if we consider the time-inhomogeneous setting, it scales with $H(S^{2}A+OA)$, but no $O^{M}$ and $O^{K}$) and the PO-bilinear rank is still $S$. In the special undercomplete ($O\geq S$) case, our framework is also flexible enough to set the memory length according to the property of the problems in order to search for the globally optimal policy. More specifically, using the latest result from [24] about belief contraction, we can set $M=\tilde{O}((1/\sigma^{4}_{1})\ln(SH/\epsilon))$ with $\epsilon$ being the optimality threshold. This allows us to compete against the globally optimal policy without paying an exponential dependence on $H$.

Observable Linear Quadratic Gaussian (LQG). In observable LQG – a classic partial observable linear dynamical system, our algorithm can compete against the globally optimal policy with a sample complexity scaling polynomially with respect to the horizon, dimensions of the state, observation, and action spaces (and other system parameters). This is achieved by simply setting the memory length $M$ to $H$. The special linear structures of the problem allow us to avoid exponential dependence on $H$ even when using the full history as a memory. While the global optimality results in tabular POMDPs and LQG exist by using different algorithms, to the best of our knowledge, this is the first unified algorithm that can solve both tabular POMDPs and LQG simultaneously without paying an exponential dependence on horizon $H$.

Observable Hilbert Space Embedding POMDPs (HSE-POMDPs). Our framework ensures the agnostic PAC guarantee on HSE-POMDPs where policy induced transitions and omission distributions have condition mean embeddings [8, 6]. This model naturally generalizes tabular POMDPs and LQG. We show that the sample complexity scales polynomially with respect to the dimensions of the embeddings. This is the first PAC guarantee in HSE-POMDPs.

Predictive State Representations (PSRs). We give the first PAC-guarantee on PSRs. PSRs model partially observable dynamical systems without even using the concept of latent states and strictly generalize the POMDP model. Our work significantly generalizes a prior PAC learning result for reactive PSRs (i.e., reactive PSRs require a strong condition that the optimal policy only depends on the latest observation) which is a much more restricted setting [33].

$M$-step decodable POMDPs [20]. Our framework can capture $M$-step decodable POMDPs where there is a (unknown) decoder that can perfectly decode the latent state by looking at the latest $M$-memory. Our algorithm can compete against the globally optimal policy with the sample complexity scaling polynomially with respect to horizon $H$, $S,A^{M}$, and the statistical complexities of function classes, without any explicit dependence on $O$. This PAC result is similar with the one from [20].

Observable POMDPs with low-rank latent transition. Our framework captures observable POMDPs where the latent transition is low-rank. This is the first PAC guarantee in this model. Under this model, we first show that with $M=\tilde{O}\left((1/\sigma_{1}^{4})\ln(dH/\epsilon)\right)$ where $d$ is the rank of the latent transition matrix, there exists an $M$-memory policy that is $\epsilon$-near optimal with respect to the globally optimal policy. Then, starting with a general model class that contains the ground truth transition and omission distribution (i.e., realizability in model class), we first convert the model class to a policy class and a value link function class, and we then show that our algorithm competes against the globally optimal policy with a sample complexity scaling polynomially with respect to $H,d,|\mathcal{A}|^{(1/\sigma_{1}^{4})\ln(dH/\epsilon)},1/\sigma_{1}$, and the statistical complexity of the model class. Particularly, the sample complexity has no explicit dependence on the size of the state and observation space, instead it just depends on the statistical complexity of the given model class.

We introduce background for POMDPs here and defer the introduction of PSRs to Section 6. We consider an episodic POMDP specified by $\mathcal{M}=\langle{\mathcal{S}},\mathcal{O},\mathcal{A},H,\mathbb{T},\mathbb{O}\rangle$, where ${\mathcal{S}}$ is the unobserved state space, $\mathcal{O}$ is the observation space, $\mathcal{A}$ is the action space, $H$ is the horizon, $\mathbb{T}:{\mathcal{S}}\times\mathcal{A}\to\Delta({\mathcal{S}})$ is the transition probability, $\mathbb{O}:{\mathcal{S}}\to\Delta(\mathcal{O})$ is the emission probability, and $r:\mathcal{O}\times\mathcal{A}\to\mathbb{R}$ is the reward. Here, $\mathbb{T},\mathbb{O}$ are unknown distributions. For notation simplicity, we consider the time-homogeneous case in this paper; Extension to the time-inhomogeneous setting is straightforward.

In our work, we consider $M$-memory policies. Let $\mathcal{Z}_{h}=(\mathcal{O}\times\mathcal{A})^{\min\{h,M\}}$ and $\bar{\mathcal{Z}}_{h}=\mathcal{Z}_{h-1}\times\mathcal{O}$. An element $z_{h}\in\mathcal{Z}_{h}$ is represented as $z_{h}=[o_{\max(h-M+1,1):h},a_{\max(h-M+1,1):h}]$, and an element $\bar{z}_{h}\in\bar{\mathcal{Z}}_{h}$ is represented as $\bar{z}_{h}=[o_{\max(h-M,1):h},a_{\max(h-M,1):h-1}]$ (thus, $\bar{z}_{h}=[z_{h-1},o_{h}]$). Figure 1 illustrates this situation. An $M$-memory policy is defined as $\pi=\{\pi_{h}\}_{h=1}^{H}$ where each $\pi_{h}$ is a mapping from $\bar{\mathcal{Z}}_{h}$ to a distribution over actions $\Delta(\mathcal{A})$.

In a POMDP, an $M$-memory policy generates the data as follows. Each episode starts with the initial state $s_{1}$ sampled from some unknown distribution. At each step $h\in[H]$, from $s_{h}\in{\mathcal{S}}$, the agent observes $o_{h}\sim\mathbb{O}(\cdot|s_{h})$, executes action $a_{h}\sim\pi_{h}(\cdot|\bar{z}_{h})$, receives reward $r(s_{h},a_{h})$, and transits to the next latent state $s_{h+1}\sim\mathbb{T}(\cdot|s_{h},a_{h})$. Note that the agent does not observe the underlying states but only the observations $\{o_{h}\}_{h\leq H}$. We denote $J(\pi)$ as the value of the policy $\pi$, i.e., $\mathbb{E}[\sum_{h=1}^{H}r_{h};a_{1:H}\sim\pi]$ where the expectation is taken w.r.t. the stochasticity of the policy $\pi$, emissions distribution $\mathbb{O}$ and transition dynamics $\mathbb{T}$.

We define a value function for a policy $\pi$ at step $h$ to be the expected cumulative reward to go under the policy $\pi$ starting from a $z\in\mathcal{Z}_{h-1}$ and $s\in{\mathcal{S}}$, i.e. $V^{\pi}_{h}:\mathcal{Z}_{h-1}\times{\mathcal{S}}\to\mathbb{R}$ where $V^{\pi}_{h}(z,s)=\mathbb{E}[\sum_{h^{\prime}=h}^{H}r_{h^{\prime}}\mid z_{h-1}=z,s_{h}=s;a_{h:H}\sim\pi]$. The notation $\mathbb{E}[\cdot\,;a_{h:H}\sim\pi]$ means the expectation is taken under a policy $\pi$ from $h$ to $H$. Compared to the standard MDP setting, the expectation is conditional on not only $s_{h}$ but also $z_{h-1}$ since we consider $M$-memory policies. The corresponding Bellman equation for $V^{\pi}_{h}$ is $V^{\pi}_{h}(z_{h-1},s_{h})=\mathbb{E}\left[r_{h}+V^{\pi}_{h+1}(z_{h},s_{h+1})\mid z_{h-1},s_{h};a_{h}\sim\pi\right]$.

Unlike MDPs, we cannot directly work with value functions $V^{\pi}_{h}(s)$ (or Q functions) in POMDPs, since they depend on the unobserved state $s$. To handle this issue, below we first introduce new value link functions by using future observations, and then discuss the PO-bilinear framework.

We consider three examples (observable tabular POMDPs, LQG, HSE-POMDPs) that admit PO-bilinear rank decomposition. Our framework can also capture PSRs and $M$-step decodable POMDPs, of which the discussions are deferred to Section 6 and Section 8, respectively. We mainly focus on one-step future, i.e., $K=1$, and briefly discuss extension to $K>1$ in the the tabular case. In this section, except for LQG, we assume $r_{h}\in[0,1]$ for any $h\in[H]$. All the missing proofs are deferred to Section B.

In this section, we first give our algorithm followed by a general sample complexity analysis. We then instantiate our analysis to specific models considered in Section 4.