Contextual Decision Processes with Low Bellman Rank are PAC-Learnable

Introduction

How can we tractably solve sequential decision making problems where the agent receives rich observations?

In this paper, we study this question by considering reinforcement learning (RL) problems where the agent receives rich sensory observations from the environment, forms complex contexts from the sensorimotor streams, uses function approximation to generalize to unseen contexts, and must engage in systematic exploration to efficiently learn to complete tasks. Such problems are at the core of empirical reinforcement learning research (e.g., (Mnih et al., 2015; Bellemare et al., 2016)), yet no existing theory provides rigorous and satisfactory guarantees in a general setting.

To answer the question, we propose a new formulation, which we call Contextual Decision Processes (CDPs), to capture a large class of sequential decision-making problems: CDPs generalize MDPs where the state forms the context (Example 1), and POMDPs where the history forms the context (Example 2). We describe CDPs formally in Section 2, and the learning goal is to find a near-optimal policy for a CDP in a sample-efficient manner.Throughout the paper, by sample-efficient we mean a number of trajectories that is polynomial in the problem horizon, number of actions, Bellman rank (to be introduced), and polylogarithmic in the number of candidate value-functions.

When the context space is very large or infinite, as is common in practice, lower bounds that are exponential in the problem horizon preclude efficient learning of CDPs, even when simple function approximators are used. However, RL problems arising in applications are often far more benign than the pathological lower bound instances, and we identify a structural assumption capturing this intuition. As our first major contribution, we define a notion of Bellman factorization (Definition 5) in Section 3, and focus on problems with low Bellman rank.

At a high level, Bellman rank is a form of algebraic dimension on the interplay between the CDP and the value-function approximator that we show is small for many natural settings. For example, every MDP with a tabular value-function has Bellman rank bounded by the rank of its transition matrix, which is at most the number of states but can be considerably smaller. For a POMDP with reactive value-functions, the Bellman rank is at most the number of hidden states and has no dependence on the observation space. We provide other instances of low Bellman rank including Linear Quadratic Regulators and Predictive State Representations. Overall, CDPs with a small Bellman rank yield a unified framework for a large class of sequential decision making problems.

Our second contribution is a new algorithm for episodic reinforcement learning called Olive (Optimism Led Iterative Value-function Elimination), detailed in Section 4.1. Olive combines optimism-driven exploration and Bellman error-based search in a new way crucial for theoretical guarantees.

The algorithm is an iterative procedure that successively refines a space of candidate Q-value functions $\mathcal{F}$. At iteration $t$, it first finds the surviving value function $f\in\mathcal{F}_{t}$ that predicts the highest value on the initial context distribution. By collecting a few trajectories according to $f$’s greedy policy, $\pi_{f}$, we can verify this prediction. If the attained value is close to the prediction, our algorithm terminates and outputs $f$. If not, we eliminate all surviving $f^{\prime}\in\mathcal{F}_{t}$ which violate certain Bellman equations on trajectories sampled using $\pi_{f}$. $\mathcal{F}_{t+1}$ is set to all the surviving functions and this process repeats.

Importantly, the sample complexity bound has a logarithmic dependence on $\mathcal{F}$, thus enabling powerful function approximation, and no direct dependence on the size of the context space, which can be very large or even infinite. As many existing models, including the ones mentioned above, have low Bellman rank, the result immediately implies sample-efficient learning in all of these settings,Our algorithm requires discrete action spaces and does not immediately apply to LQRs; see more discussion in Section 3. as highlighted in Table 1.

We also present several extensions of the main result, showing robustness to the failure of the assumption that the optimal value-function is captured by the function approximator, adaptivity to unknown Bellman rank, and extension to infinite function classes of bounded statistical complexity. Altogether, these results show that the notion of Bellman rank robustly captures the difficulty of exploration in sequential decision-making problems.

To summarize, this work advances our understanding in reinforcement learning with complex observations where long-term planning and exploration are critical. There are, of course, several additional questions that must be resolved before we have satisfactory tools for these problems. The biggest drawback of Olive is its computational complexity, which is polynomial in the number of value functions and hence intractable for the powerful classes of interest. This issue must be addressed before we can empirically evaluate the effectiveness of the proposed algorithm. We leave this and other open questions for future work.

There is a rich body of theoretical literature on learning Markov Decision Processes (MDPs) with small state spaces (Kearns and Singh, 2002; Brafman and Tennenholtz, 2003; Strehl et al., 2006), with an emphasis on sophisticated exploration techniques that find near-optimal policies in a sample-efficient manner. While there have been attempts to extend these techniques to large state spaces (Kakade et al., 2003; Jong and Stone, 2007; Pazis and Parr, 2016), these approaches fail to be a good fit for practical scenarios where the environment is typically perceived through complex sensory observations such as image, text, or audio signals. Alternatively, Monte Carlo Tree Search (MCTS) methods can handle arbitrarily large state spaces, but only at the cost of exponential dependence on the planning horizon (Kearns et al., 2002; Kocsis and Szepesvári, 2006). Our work departs from these existing efforts by aiming for a sample complexity that is independent of the size of the context space and at most polynomial in the horizon. Similar goals have been attempted by Wen and Van Roy (2013) and Krishnamurthy et al. (2016) where attention is restricted to decision processes with deterministic dynamics and special structures. In contrast, we study a much broader class of problems with relatively mild conditions.

On the empirical side, the prominent recent success on both the Atari platform (Mnih et al., 2015; Wang et al., 2015) and Go (Silver et al., 2016) have sparked a flurry of research interest. These approaches leverage advances in deep learning for powerful function approximation, while, in most cases, using simple heuristic strategies, such as $\epsilon$-greedy, for exploration. More advanced exploration strategies include extending the methods for small state spaces (e.g., the use of pseudo-counts in Bellemare et al. (2016)), and combining MCTS with function approximation (e.g., Silver et al. (2016)). Unfortunately, both types of approaches often require strong domain knowledge and large amounts of data to be successful.

Hallak et al. (2015) have proposed a setting called Contextual MDPs, where a context refers to some static information that can be used to generalize across many similar MDPs. In this paper, a context is most similar to state features in the RL literature and is a natural generalization of the notion of context as in the contextual bandit literature (Langford and Zhang, 2008).

Contextual Decision Processes (CDPs)

We introduce a new model, called a Contextual Decision Process, as a unified framework for reinforcement learning with rich observations. We first present the model, before the relevant notation and definitions.

Contextual Decision Processes make minimal assumptions to capture a very general class of RL problems and are defined as follows.

In a CDP, the agent’s interaction with the environment proceeds in episodes. In each episode, the agent observes a context $x_{1}$, takes action $a_{1}$, receives reward $r_{1}$ and observes $x_{2}$, repeating $H$ times. A policy $\pi:\mathcal{X}\to\mathcal{A}$ specifies the decision-making strategy of an agent, that is $a_{h}=\pi(x_{h}),~{}\forall h\in[H]$, and induces a distribution over the trajectory $(x_{1},a_{1},r_{1},\ldots,x_{H},a_{H},r_{H},x_{H+1})$ according to the system descriptor $P$. More generally, a sequence of stochastic policies $\pi_{1},\ldots,\pi_{H}:\mathcal{X}\to\Delta(\mathcal{A})$ induces a distribution over trajectories in a similar way, where $a_{h}\sim\pi_{h}(x_{h})~{}\forall h\in[H]$. The value of a policy, $V^{\pi}$, is defined as

Below we show that CDPs capture classical RL models, including MDPs and POMDPs, and the optimal policies can be expressed as a function of appropriately chosen contexts.

Consider a finite-horizon MDP $({\mathcal{S}},\mathcal{A},H,\Gamma_{1},\Gamma,R)$, where ${\mathcal{S}}$ is the state space, $\mathcal{A}$ is the action space, $H$ is the horizon, $\Gamma_{1}\in\Delta({\mathcal{S}})$ is the initial state distribution, $\Gamma:{\mathcal{S}}\times\mathcal{A}\to\Delta({\mathcal{S}})$ is the state transition function, $R:{\mathcal{S}}\times\mathcal{A}\to\Delta()$ is the reward function, and an episode takes the form of $(s_{1},a_{1},r_{1},\ldots,s_{H},a_{H},r_{H})$. We can convert the MDP to a CDP $(\mathcal{X},\mathcal{A},H,P)$ by letting $\mathcal{X}={\mathcal{S}}\times[H]$ and $x_{h}=(s_{h},h)$, which allows the set of policies $\{\mathcal{X}\to\mathcal{A}\}$ to contain the optimal policy (Puterman, 1994). The system descriptor is $P=(P_{\varnothing},P_{+})$, where $P_{\varnothing}(x_{1})=\Gamma_{1}(s_{1})$, and $P_{+}(r_{h},x_{h+1}\,|\,x_{1},a_{1},r_{1},\ldots,x_{h},a_{h})=R(r_{h}|s_{h},a_{h})\,\Gamma(s_{h+1}|s_{h},a_{h})$.

The system descriptor for a particular model is usually obvious from the definitions, and here we give its explicit form as an illustration. We omit the specification of system descriptor in the remaining examples.

Next we turn to POMDPs. It might seem that a CDP describes a similar process as a POMDP but limits the agent’s decision-making strategies to memoryless (or reactive) policies, as we only consider policies in $\{\mathcal{X}\to\mathcal{A}\}$. This is not true. We clarify this issue by showing that we can use the history as context, and the induced CDP suffers no loss in the ability to represent optimal policies.

It is also clear from this example that we can assume contexts are Markovian in CDPs without loss of generality, as we can always use history as context. While we do not commit to this assumption to allow for a flexible framework with simple notation (see Example 3), we later connect to well-known results in MDP literature based on this observation so that readers can transfer insights from MDPs to CDPs.

In some application scenarios, partial observability can be resolved by using a small sliding window: for example, in Atari games, it is common to keep track of the last $4$ frames of images (Mnih et al., 2015). In this case, we can represent the problem as a CDP by letting $x_{h}=(o_{h-3},o_{h-2},o_{h-1},o_{h})$.

We hope these examples convince the reader of the flexibility of keeping contexts separate from intrinsic quantities such as states or observations. Finally, we introduce a regularity assumption on the rewards.

We assume that regardless of how actions are chosen, for any $h=1,\ldots,H$, $r_{h}\geq 0$ and $\sum_{h=1}^{H}r_{h}\leq 1$ almost surely.

2 Value-based RL and Function Approximation

Now that we have a model in place, we turn to some important solution concepts.

A CDP makes no assumption on the cardinality of the context space, which makes it critical to generalize across contexts, since the agent might not encounter the same context twice. Therefore, we consider value-based RL with function approximation. That is, the agent is given a set of functions $\mathcal{F}\subseteq\mathcal{X}\times\mathcal{A}\to$ and uses it to approximate an action-value function (or Q-value function). Without loss of generality we assume that $f(x_{H+1},a)\equiv 0$.This frees us from having to treat the last level ($h=H$) differently in the Bellman equations. For the purpose of presentation, we assume that $\mathcal{F}$ is a finite space with $|\mathcal{F}|=N<\infty$ for most of the paper. In Section 5.3 we relax this assumption and allow infinite function classes with bounded complexity.

Given any policy $\pi:\mathcal{X}\to\mathcal{A}$ and a function $f:\mathcal{X}\times\mathcal{A}\to$, the average Bellman error of $f$ under roll-in policy $\pi$ at level $h$ is defined as

In words, the average Bellman error measures the self-consistency of a function $f$ between its predictions at levels $h$ and $h+1$ when all the previous actions are taken according to some policy $\pi$. In many existing approaches (e.g., LSPI (Lagoudakis and Parr, 2003) and FQI (Ernst et al., 2005)), the Bellman errors are defined as taking the expectation of a squared error unlike this definition. Given this definition, we now define a set of Bellman equations.

Given an $(f,\pi,h)$ triple, a Bellman equation posits $\mathcal{E}(f,\pi,h)=0$. We say $f\in\mathcal{F}$ is valid if the Bellman equation on $(f,\pi_{f^{\prime}},h)$ holds for every $f^{\prime}\in\mathcal{F},h\in[H]$.

Note that the validity assumption only considers roll-ins according to the greedy policies $\pi_{f}$, which is the natural policy class in a function approximation setting. In Section 5.2, we show how to incorporate a separate policy class in these definitions. In the MDP setting, each Bellman equation can be viewed as the linear combination of the standard Bellman optimality equations for $Q^{\star}$, Readers who are not familiar with the definition of $Q^{\star}$ are advised to consult a textbook, such as (Sutton and Barto, 1998). where the coefficients are the probabilities with which the roll-in policy $\pi$ visits each state. This leads to the following consequence.

Given an MDP and a space of functions $\mathcal{F}:{\mathcal{S}}\times[H]\times\mathcal{A}\to$, if the optimal Q-value function of the MDP $Q^{\star}$ lies in $\mathcal{F}$, then in the corresponding CDP with $\mathcal{X}={\mathcal{S}}\times[H]$, $Q^{\star}$ is valid.

While $Q^{\star}$ satisfies the Bellman equations and yields the optimal policy $\pi^{\star}=\pi_{Q^{\star}}$, there can be other functions which also satisfy the equations while yielding suboptimal policies. This happens because Eq. (2) only considers $a_{h}$ drawn according to $\pi_{f}$ and does not use the values on other actions. For instance, consider a CDP where at every context, action $a$ always gets a reward of 0 and action $a^{\prime}$ always gets a reward of 1. A function that predicts $f(x,a)=f(x,a^{\prime})=0~{}\forall x,a$ is trivially valid as long as tie-breaks always favor $a$.

Since validity alone does not imply that we get a good policy, it is natural to search for a valid value function which also induces a high-value policy. We formalize this goal in the next definition.

For the same setting as in Fact 1, when $Q^{\star}\in\mathcal{F}$, we have $f^{\star}=Q^{\star}$, and $V_{\mathcal{F}}^{\star}=V^{\star}$, which is the optimal long-term value.

Definition 4 implicitly assumes that there is at least one valid $f\in\mathcal{F}$. This is weaker than the realizability assumption made in the value-based RL literature, that $\mathcal{F}$ contains the optimal Q-value function of an MDP $Q^{\star}$ (Antos et al., 2008; Krishnamurthy et al., 2016). Indeed, the setup subsumes realizability, as evidenced by Fact 1 and 2. When $Q^{\star}\in\mathcal{F}$, the algorithm aims to identify a policy achieving value close to $V^{\star}$, the optimal value achievable by any agent. When no functions in $\mathcal{F}$ approximate $Q^{\star}$ well, finding the best valid value function is still a meaningful and non-trivial objective. In this sense, our work makes substantially weaker realizability-type assumptions than prior theoretical results for value-based RL (Antos et al., 2008; Krishnamurthy et al., 2016), which assume $Q^{\star}\in\mathcal{F}$ often in addition to several stronger requirements.

In general, $\mathcal{F}$ may not contain $Q^{*}$, or any valid functions at all, which makes our learning goal trivial. It is desirable to have an algorithm robust to such a scenario, and we show how our algorithm requires only an approximate notion of validity in Section 5.4, implying a graceful degradation in the results.

Bellman Factorization and Bellman Rank

CDPs are general models for sequential decision making, but are there efficient RL algorithms for them?

Unfortunately, without further assumptions, learning in CDPs is generally hard, since they subsume MDPs and POMDPs with arbitrarily large state/observation spaces. Moreover, a function class $\mathcal{F}$ with low statistical complexity, which would generalize effectively in a standard supervised learning setting with a fixed data distribution, does not overcome this difficulty in CDPs where the data distribution crucially depends on the agent’s policy. In particular, even when $\log N$, the statistical complexity for finite classes, is small, there exists an $\Omega(K^{H})$ lower bound on the sample complexity of learning CDPs. The result is due to Krishnamurthy et al. (2016), and is included in Appendix A.1 for completeness.

While exponential lower bounds for learning CDPs exist, they are fairly pathological, and most real problems have substantially more structure. To capture these realistic instances and circumnavigate the lower bounds, we propose a new complexity measure and restrict our attention to settings where this measure is low. As we will see, this measure is naturally small for many existing models, and, when it is small, efficient reinforcement learning is possible.

The complexity measure we propose is a structural characterization of the set of Bellman equations induced by the CDP and the value-function class (recall Definition 2) that we need to check to find valid functions. While checking validity by enumeration is intractable for large $\mathcal{F}$, observe that the Bellman equations are structured in tabular MDPs: the average Bellman error under any roll-in policy is a stochastic combination of the single-state errors, and checking the single-state errors (which is tractable) is sufficient to guarantee validity. This observation hints toward a more general phenomenon: whenever the collection of Bellman errors across all roll-in policies can be concisely represented, we may be able to check the validity of all functions in a tractable way.

This intuition motivates a new complexity measure that we call the Bellman rank. Define the Bellman error matrices, one for each $h$, to be $|\mathcal{F}|\times|\mathcal{F}|$ matrices where the $(f,f^{\prime})^{\textrm{th}}$ entry is the Bellman error $\mathcal{E}(f,\pi_{f^{\prime}},h)$. Informally, the Bellman rank for a CDP and a given value-function class $\mathcal{F}$ is a uniform upper bound on the rank of these $H$ Bellman error matrices.

and $\|\nu_{h}(f^{\prime})\|_{2}\cdot\|\xi_{h}(f)\|_{2}\leq\zeta<\infty$.

The exact factorization in Eq. (3) can be relaxed to an approximate version as is discussed in Section 5.4. In the remaining sections of this paper we introduce the main algorithm, and analyze its sample-efficiency in problems with low Bellman rank. In the remainder of this section we showcase the generality of Definition 5 by describing a number of common RL settings that have a small Bellman rank. Throughout, we see how the Bellman rank captures the process-specific structures that allow for efficient exploration. Proofs of all claims in this section are deferred to Appendix B.

We start with the tabular MDP setting, and show that the Bellman rank is at most the number of states.

Consider the MDP setting of Example 1 with the corresponding CDP. With any $\mathcal{F}\subset\mathcal{X}\times\mathcal{A}\to$, this model admits a Bellman factorization with $M=|{\mathcal{S}}|$ and $\zeta=2\sqrt{M}$.

A related model introduced by Li (2009) for extending tabular PAC-MDP methods to large MDPs using a form of state abstractions also has low Bellman rank (See Appendix B.2).

The MDP example is particularly simple as each coordinate of the $M$-dimensional space corresponds to a state, which is observable. Our next few examples show that this is not necessary, and that Bellman factorization can be based on latent properties of the process. We next consider large MDPs whose transition dynamics have a low-rank structure. A closely related setting has been considered by Barreto et al. (2011, 2014) where the low-rank structure is exploited to speed up MDP planning, but prior to this work, no sample-efficient RL algorithms are known for this setting.

Consider the MDP setting of Example 1 with a transition matrix $\Gamma$ having rank at most $M$. The induced CDP along with any $\mathcal{F}\subset\mathcal{X}\times\mathcal{A}\to$ admits a Bellman factorization with Bellman rank $M$.

The next example considers POMDPs with large observations spaces and reactive value functions, where the Bellman rank is at most the number of hidden states.

Consider the POMDP setting of Example 3 with $|{\mathcal{S}}|<\infty$ and a sliding window of size 1 along with the induced CDP. Given any $\mathcal{F}\subset\mathcal{X}\times\mathcal{A}\to$, this model admits a Bellman factorization with $M=|{\mathcal{S}}|$ and $\zeta=2\sqrt{M}$.

Proposition 2 and 3 can be proved under a unified model that generalizes POMDPs by allowing the transition function and the reward function to depend on the observation (See Figure 1 (a) – (c) for graphical representations of these models). This unified model captures the experimental settings considered in state-of-the-art empirical RL work (Figure 1 (d)), where agents act in a grid-world ($|{\mathcal{S}}|$ is small) and receives complex and rich observations such as raw pixel images ($|\mathcal{O}|$ is large). The model also subsumes and generalizes the setting of Krishnamurthy et al. (2016) which requires deterministic transitions in the underlying MDP. Our new algorithm eliminates the need for determinism, and still guarantee sample-efficient learning.

Next, we consider Predictive State Representations (PSRs), which are models of partially observable systems with parameters grounded in observable quantities (Littman et al., 2001). Similar to the case of POMDPs, we can bound the Bellman rank in terms of the rank of the PSREvery POMDP has an equivalent PSR whose rank is bounded by the number of hidden states (Singh et al., 2004). when the candidate value functions are reactive.

Consider a partially observable system with observation space $\mathcal{O}$, and the induced CDP $(\mathcal{X},\mathcal{A},H,P)$ with $x_{h}=(o_{h},h)$. If the linear dimension of the system (i.e., rank of its PSR model) is at most $L$, then given any $\mathcal{F}:\mathcal{X}\times\mathcal{A}\to$, the Bellman rank is bounded by $LK$.

The last example considers a class of linear control problems well studied in control theory, called Linear Quadratic Regulators (LQRs). We show that the Bellman rank in LQRs is bounded by the dimension of the state space. Exploration in this class of problems has been previously considered by Osband and Van Roy (2014). Note that the algorithm to be introduced in the next section does not directly apply to LQRs due to the continuous action space, and adaptations that exploit the structure of the action space may be needed, which we leave for future work.

Algorithm and Main Results

In this section we present the algorithm for learning CDPs that have a Bellman factorization with a small Bellman rank and the main sample complexity guarantee. To aid presentation and help convey the main ideas, we make three simplifying assumptions:

We assume the Bellman rank parameter $M$ is known to the agent.We also assume knowledge of the corresponding norm parameter, but this is relatively minor.

We assume the function class $\mathcal{F}$ is finite with $|\mathcal{F}|=N$.

We assume exact validity (Definition 3) and exact Bellman factorization (Definition 5).

All three assumptions can be relaxed, and we sketch these relaxations in Section 5.

We are interested in designing an algorithm for PAC Learning CDPs. We say that an algorithm PAC learns if given $\mathcal{F}$, two parameters $\epsilon,\delta\in(0,1)$, and access to a CDP, the algorithm outputs a policy $\hat{\pi}$ with $V^{\hat{\pi}}\geq V_{\mathcal{F}}^{\star}-\epsilon$ with probability at least $1-\delta$. The sample complexity is the number of episodes needed to achieve such a guarantee, and is typically expressed in terms of $\epsilon$, $\delta$ and other relevant parameters. The goal is to design an algorithm with sample complexity that is $\textrm{Poly}(M,K,H,1/\epsilon,\log(N),\log(1/\delta))$ where $M$ is the Bellman rank, $K$ is the number of actions, and $H$ is the time horizon. Importantly, the bound has no dependence on the number of unique contexts $|\mathcal{X}|$.

Pseudocode for the algorithm, which we call Olive (Optimism Led Iterative Value-function Elimination), is displayed in Algorithm 1. Theorem 1 describes how to set the parameters $n_{\textrm{est}},n_{\textrm{eval}},n$, and $\phi$.

At a high level, the algorithm aims to eliminate functions $f\in\mathcal{F}$ that fail to satisfy the validity condition in Definition 3. This is done by Lines 13 and 14 inside the loop of the algorithm. Observe that, since the actions $a_{h_{t}}$ are chosen uniformly at random, Eq. (5) produces an unbiased estimate of $\mathcal{E}(f,\pi_{t},h_{t})$, the average Bellman error for function $f$ on roll-in policy $\pi_{t}$ at time $h_{t}$. Thus, Eq. (6) eliminates functions that have high average Bellman error on this distribution, which means they fail to satisfy the validity criteria.

The other major component of the algorithm involves choosing the roll-in policy and level on which to do the learning step. At iteration $t$, we choose the roll-in policy $\pi_{t}$ optimistically, by choosing $f_{t}$ that predicts the highest value at the starting context distribution, and letting $\pi_{t}=\pi_{f_{t}}$. To pick the level, we compute $f_{t}$’s average Bellman error on its own roll-in distribution (Eq. (4)), and set $h_{t}$ to be any level for which this average Bellman error is high (See Line 11). As we will show, these choices ensure that substantial learning happens on each iteration, guaranteeing that the algorithm uses polynomially many episodes.

The last component is the termination criterion. The algorithm terminates if $f_{t}$ has small average Bellman error on its own roll-in distribution at all levels. This criteria guarantees that $\pi_{t}$ is near optimal.

Computationally, the algorithm requires enumeration of the value-function class, which we expect to be extremely large or infinite in practice. A computationally efficient implementation is essential for a practical algorithm, which is left to future work. We focus on the sample efficiency of the algorithm in this paper.

Intuition for OLIVE. To convey intuition, it is helpful to ignore any sampling effects by replacing all empirical estimates with population values, and set $\epsilon$ to . The first important fact is that the algorithm never eliminates a valid function, since the learning step in Eq. (6) only eliminates a function $f$ if we can find a distribution on which it has a large average Bellman error. If $f$ is valid, then $\mathcal{E}(f,\pi,h)=0$ for all $\pi,h$, so $f$ is never eliminated.

The more challenging component is ensuring that the algorithm terminates in polynomially many iterations, which is critical for obtaining a polynomial sample complexity bound. This argument crucially relies on the Bellman factorization (recall Definition 5), which enables us to embed the distributions into $M$ dimensions and measure progress in this low-dimensional space.

For now, fix some $h$ and focus on the iterations when $h_{t}=h$. If we ignore sampling effects we can set $\phi=0$, and, by using the Bellman factorization to write $\mathcal{E}(f,\pi_{f_{t}},h)$ as an inner product, we can think of the learning step in Line 14 as introducing a homogeneous linear constraint on the set of $\xi_{h}(f)$ vectors, that is, $\langle\nu_{h}(f_{t}),\xi_{h}(f)\rangle=0$. Now, if we execute the learning step at level $h$ again in a later iteration $t^{\prime}$, we have $\langle\nu_{h}(f_{t^{\prime}}),\xi_{h}(f_{t^{\prime}})\rangle\neq 0$ from Line 11. Importantly, this means that $\nu_{h}(f_{t^{\prime}})$ must be linearly independent from previous $\nu_{h}(f_{t})$ since $\langle\nu_{h}(f_{t}),\xi_{h}(f_{t^{\prime}})\rangle=0$. In general, every time $h_{t}=h$, the number of linearly independent constraints increases by 1, and therefore the number of iterations where $h_{t}=h$ is at most the dimension of the space, which is $M$. Thus the Bellman rank leads to a bound on the number of iterations.

The above heuristic reasoning, despite relying on the brittle notion of linear independence, can be made robust. With sampling effects, rather than homogeneous linear equalities, the learning step for level $h$ introduces linear inequality constraints to the $\xi_{h}(f)$ vectors. But if $f^{\prime}$ is a surviving function that forces us to train at level $h$, it means that $\langle\nu_{h}(f^{\prime}),\xi_{h}(f^{\prime})\rangle$ is very large, while $\langle\nu_{h}(\cdot),\xi_{h}(f^{\prime})\rangle$ is very small for all previous $\nu_{h}(\cdot)$ vectors used in the learning step. Intuitively this means that the new $\nu(f^{\prime})$ vector is quite different from all of the previous ones. In our proof, we use a volumetric argument to show that this suffices to guarantee substantial learning takes place.

The optimistic choice for $f_{t}$ is critical for driving the agent’s exploration. With this choice, if $f_{t}$ is valid, then the algorithm terminates correctly, and if $f_{t}$ is not valid, then substantial progress is made. Thus the agent does not get stuck exploring with many valid but suboptimal functions, which could result in exponential sample complexity.

2 Sample Complexity

We now turn to the main result, which guarantees that Olive PAC-learns Contextual Decision Processes with polynomial sample complexity.

For any $\epsilon,\delta\in(0,1)$, any Contextual Decision Process and function class $\mathcal{F}$ that admits a Bellman factorization with parameters $M,\zeta$, run Olive with the following parameters:

According to this theorem, if a Contextual Decision Process and function class $\mathcal{F}$ admit a Bellman factorization with small Bellman rank and $\mathcal{F}$ contains valid functions, Olive is guaranteed to find a near optimal valid function using only polynomially many episodes. To our knowledge, our result is the most general polynomial sample complexity bound for reinforcement learning with rich observation spaces and function approximation, as many popular models are shown to admit small Bellman rank (see Section 3). The result also certifies that the notion of Bellman factorization, which is quite general, is sufficient for efficient exploration and learning in sequential decision making problems.

It is worth briefly comparing this result with prior work.

The most closely related result is the recent work of Krishnamurthy et al. (2016), who also consider episodic reinforcement learning with infinite observation spaces and function approximation. The model studied there is a form of Contextual Decision Process with Bellman rank $M$, so the result applies as is to that setting. Importantly, we eliminate the need for deterministic transitions, resolving one of their open problems. Moreover, the sample complexity bound improves the dependence on $H$ and $\epsilon$, at the cost of a worse dependence on $M$. We emphasize that this result applies to a much more general class of models.

Another related body of work provides sample complexity bounds for fitted value/policy iteration methods (e.g., (Munos, 2003; Antos et al., 2008; Munos and Szepesvári, 2008)). These works consider the infinite-horizon discounted MDP setting, and impose much stronger assumptions than we do including not only that the function class captures $Q^{\star}$, but also that it is approximately closed under Bellman update operators. More importantly, the analyses rely on the so-called concentrability coefficients to correct the mismatch between training and test distributions (Farahmand et al., 2010; Lazaric et al., 2012), implicitly assuming that an exploration distribution is given, hence these results do not address the exploration issue which is the main focus here.

This approach applies to learning reactive policies in POMDPs (see Proposition 3). Azizzadenesheli et al. (2016) provides a sample-efficient algorithm in a closely related setting, where both the observation space and the hidden-state space are small in cardinality. While their approach does not require realizable value-functions, the sample complexity depends polynomially on the number of unique observations, and the method relies on additional mixing assumptions, which we do not require.

Finally, Contextual Decision Processes also encompass contextual bandits, where the optimal sample complexity is $O(K\log(N)/\epsilon^{2})$ (Agarwal et al., 2012). As contextual bandits have $M=1$ and $H=1$, Olive achieves optimal sample complexity in this special case.

Turning briefly to lower bounds, since the CDP setting with Bellman factorization is new, general lower bounds for the broad class do not exist. However, we can use MDP lower bounds for guidance on the question of optimality, since the small-state MDPs in Example 1 are a special case. While no existing MDP lower bounds apply as is (because formulations vary), in Appendix A.2 we adapt ideas from Auer et al. (2002) to obtain a $\Omega(MKH/\epsilon^{2})$ sample complexity lower bound for learning the MDPs in Example 1.

Comparing with this lower bound, the sample complexity in Theorem 1 is worse in $M,H$, and $\log(N)$ factors, but of course the small-state MDP is a significantly simpler special case. We leave as future work the question of optimal sample complexity for learning CDPs with low Bellman rank.

Extensions

We introduce four important extensions to the algorithm and analysis.

The first extension eliminates the need to know $M$ in advance (note that Algorithm 1 requires $M$ as an input parameter). A simple procedure described in Algorithm 2, can guess the value of $M$ on a doubling schedule and handle this situation with no consequences to asymptotic sample complexity.In Algorithm 2 we assume that $\zeta$ is known. In the examples provided in Proposition 1, 2, and 3, however, $\zeta$ grows with $M$ in the form of $\zeta=2\sqrt{M}$. In this case, we can compute $\zeta^{\prime}=2\sqrt{M^{\prime}}$ and call Olive with $\zeta^{\prime}$ instead of $\zeta$. As long as $\zeta$ is a polynomial term and non-decreasing in $M$ the same analysis applies and Theorem 2 holds.

For any $\epsilon,\delta\in(0,1)$, any Contextual Decision Process and function class $\mathcal{F}$ that admits a Bellman factorization with parameters $M,\zeta$, if we run GuessM$(\mathcal{F},\epsilon,\delta)$, then with probability at least $1-\delta$, Olive halts and returns a policy which satisfies $V^{\hat{\pi}}\geq V_{\mathcal{F}}^{\star}-\epsilon$, and the number of episodes required is at most

We give some intuition about the proof here, with details in Appendix D. In Algorithm 2, $M^{\prime}$ is a guess for $M$ which grows exponentially. When $M^{\prime}\geq M$, analysis of the main algorithm shows that $\textsc{Olive}(\mathcal{F},M^{\prime},\epsilon,\frac{\delta}{i(i+1)})$ terminates and returns a near-optimal policy with high probability. The doubling schedule implies that the largest guess is at most $2M$, which has negligible effect on the sample complexity. On the other hand, Olive may not explore effectively when $M^{\prime}<M$, because not enough samples (chosen according to $M^{\prime}$) are used to estimate the average Bellman errors in Eq. (5). This worse accuracy does not guarantee sufficient progress in learning.

However, the high-probability guarantee that $f^{\star}$ is not eliminated is unaffected, because the threshold $\phi$ on Line 6 of Olive is set in accordance with the sample size $n$ specified in Theorem 1, regardless of $M$. Consequently, if the algorithm ever terminates when $M^{\prime}<M$, we still get a near-optimal policy. When $M^{\prime}<M$ the Olive subroutine may not terminate, which the explicit termination on line 4 in Algorithm 2 addresses. Finally, by splitting the failure probability $\delta$ appropriately among all guesses of $M^{\prime}$, we obtain the same order of sample complexity as in Theorem 1.

2 Separation of Policy Class and V-value Class

So far, we have assumed that the agent has access to a class of Q-value functions $\mathcal{F}\subset\mathcal{X}\times\mathcal{A}\to$. In this section, we show the algorithm allows separate representations of policies and V-value functions.

For every $f\in\mathcal{F}$, and any $x\in\mathcal{X},a\neq\pi_{f}(x)$, we note that the value of $f(x,a)$ is not used by Algorithm 1, and changing it to arbitrary values does not affect the execution of the algorithm as long as $f(x,a)\leq f(x,\pi_{f}(x))$ (so that $\pi_{f}$ does not change). In other words, the algorithm only interacts with $f$ in two forms:

A mapping $g_{f}:x\mapsto f(x,\pi_{f}(x))$. We call such mappings V-value functions to contrast the previous use of Q-value functions.In the MDP setting, such functions are also known as state-value functions.

Hence, supplying $\mathcal{F}$ is equivalent to supplying the following space of (policy, V-value function) pairs:

This observation provides further evidence that Definition 3 is significantly less restrictive than standard realizability assumptions. Validity of $f$ means that $(\pi_{f},g_{f})$ obeys the Bellman Equations for Policy Evaluation (i.e., $g_{f}$ predicts the long-term value of following $\pi_{f}$), as opposed to the more common Bellman Optimality Equations. In MDPs, there are many ways to satisfy the policy evaluation equations at every state simultaneously, while $Q^{\star}$ is the only function that satisfies all optimality equations.

More generally, instead of using a Q-value function class, we can run Olive with a policy space $\Pi\subset\mathcal{X}\to\mathcal{A}$ and a V-value function class $\mathcal{G}\subset\mathcal{X}\to$ where we assemble (policy,V-value function) pairs by taking the Cartesian product of $\Pi$ and $\mathcal{G}$. Olive can be run here with the understanding that each Q-value function $f$ in Olive is associated with a $(\pi,g)$ pair, and the algorithm uses $\pi$ instead of $\pi_{f}$ and $g(x)$ instead of $f(x,\pi_{f}(x))$. All the analysis applies directly with this transformation, and the $\log|\mathcal{F}|$ dependence in sample complexity is replaced by $\log|\Pi|+\log|\mathcal{G}|$. Note also that the definition of Bellman factorization also extends naturally to this case, where the first argument is the $(\pi,g)$ pair and the second argument is a roll-in policy, $\pi^{\prime}$.

3 Infinite Hypothesis Classes

The arguments in Section 4 assume that $|\mathcal{F}|=N<\infty$. However, almost all commonly used function approximators are infinite classes, which restricts the applicability of the algorithm. On the other hand, the size of the function class appears in the analysis only through deviation bounds, so techniques from empirical process theory can be used to generalize the results to infinite classes. This section establishes parallel versions of those deviation bounds for function classes with finite combinatorial dimensions, and together with the rest of the original analysis we can show the algorithm enjoys similar guarantees when working with infinite hypothesis classes.

Specifically, we consider the setting where $\Pi$ and $\mathcal{G}$ are given (see Section 5.2), and they are infinite classes with finite combinatorial dimensions. We assume that $\Pi$ has finite Natarajan dimension (Definition 6), and $\mathcal{G}$ has finite pseudo dimension (Definition 7). These two dimensions are standard extensions of VC-dimension to multi-class classification and regression respectively.

Suppose $\mathcal{X}$ is a feature space and $\mathcal{Y}$ is a finite label space. Given hypothesis class $\mathcal{H}\subset\mathcal{X}\to\mathcal{Y}$, its Natarajan dimension $\textrm{Ndim}(\mathcal{H})$ is defined as the maximum cardinality of a set $A\subseteq\mathcal{X}$ that satisfies the following: there exists $h_{1},h_{2}:A\to\mathcal{Y}$ such that (1) $\forall x\in A$, $h_{1}(x)\neq h_{2}(x)$, and (2) $\forall B\subseteq A$, $\exists h\in\mathcal{H}$ such that $\forall x\in B$, $h(x)=h_{1}(x)$ and $\forall x\in A\setminus B$, $h(x)=h_{2}(x)$.

The definition of pseudo dimension relies on that of VC-dimension, whose definition and basic properties are recalled in the appendix. We state the final sample complexity result here. Since the algorithm parameters are somewhat complex expressions, we omit them in the theorem statement and provide specification in the proof, which is deferred to Appendix D.

Let $\Pi\subset\mathcal{X}\to\mathcal{A}$ with $\textrm{Ndim}(\Pi)\leq d_{\Pi}<\infty$ and $\mathcal{G}\subset\mathcal{X}\to$ with $\textrm{Pdim}(\mathcal{G})\leq d_{\mathcal{G}}<\infty$. For any $\epsilon,\delta\in(0,1)$, any Contextual Decision Process with policy space $\Pi$ and function space $\mathcal{G}$ that admits a Bellman factorization with parameters $M,\zeta$, if we run Olive with appropriate parameters, then with probability at least $1-\delta$, Olive halts and returns a policy $\hat{\pi}$ that satisfies $V^{\hat{\pi}}\geq V_{\mathcal{F}}^{\star}-\epsilon$, and the number of episodes required is at most

Compared to Theorem 1, the sample complexity we get for infinite hypothesis classes has two differences: (1) $\log N$ is replaced by $d_{\Pi}+d_{\mathcal{G}}$, which is expected, based on the discussion in Section 5.2, and (2) the dependence on $K$ is quadratic as opposed to linear. In fact, in the proof of Theorem 1, we exploited the low-variance property of importance weights in Eq. (5), and applied Bernstein’s inequality to avoid a factor of $K$. With infinite hypothesis classes, the same approach does not apply directly. However, this may only be a technical issue, and a more refined analysis may recover a linear dependence (e.g., using tools from Panchenko (2002)).

4 Approximate Validity and Approximate Bellman Rank

Recall that the sample-efficiency guarantee of Olive relies on two major assumptions:

We assumed that $\mathcal{F}$ contains valid functions (Definition 3). In practice, however, it is hard to specify a function class that contains strictly valid functions, as the notion of validity depends on the environment dynamics, which are unknown. A much more realistic situation is that some functions in $\mathcal{F}$ satisfy validity only approximately.

We assumed that the average Bellman errors have an exact low-rank factorization (Definition 5). While this is true for a number of RL models (Section 3), it is worth keeping in mind that these are only models of the environments, which are different from and only approximations to the real environments themselves. Therefore, it is more realistic to assume that an approximate factorization exists when defining Bellman factorization.

In this section, we show that the algorithmic ideas of Olive are indeed robust against both types of approximation errors, and degrades gracefully as the two assumptions are violated. Below we introduce the approximate versions of Definition 3 and 5, give a slightly extended version of the algorithm, Oliver (for Optimism-Led Iterative Value-function Elimination with Robustness, see Algorithm 3), and state its sample complexity guarantee in Theorem 4.

Given any CDP and function class $\mathcal{F}$, we say $f\in\mathcal{F}$ is $\theta$-valid if for any $f^{\prime}\in\mathcal{F}$ and any $h\in[H]$, $\left|\mathcal{E}(f,\pi_{f^{\prime}},h)\right|\leq\theta.$

The approximation error $\theta$ introduced in Definition 8 allows the algorithm to compete against a broader range of functions; hence the notions of optimal function and value need to be re-defined accordingly.

By definition, $V_{\mathcal{F},\theta}^{\star}$ is non-decreasing in $\theta$ with Definition 3 being a special case where $\theta=0$. When $\theta>0$, we compete against some functions that do not obey Bellman equations, breaking an essential element of value-based RL. As a consequence, returning a policy with value close to $V_{\mathcal{F},\theta}^{\star}$ in a sample-efficient manner is very challenging, so the value that Oliver can guarantee is suboptimal to $V_{\mathcal{F},\theta}^{\star}$ by a term that is proportional to $\theta$ and does not diminish with more data.

and $\|\nu_{h}(f^{\prime})\|_{2}\cdot\|\xi_{h}(f)\|_{2}\leq\zeta<\infty$.

A modified version of Olive that deals with these approximation errors, Oliver, is specified in Algorithm 3. Here, we use $\epsilon$ to denote the component of the suboptimality that diminish as more data is collected, and the total suboptimality that we can guarantee is $\epsilon$ plus a term proportional to $\theta$ and $\eta$ (see Eq. (13) in Theorem 4). The algorithm is almost identical to Olive except in two places: (1) it uses $\epsilon^{\prime}$ (defined on Line 1) in the termination condition (Line 9) as opposed to $\epsilon$, and (2) it uses a higher threshold that depends on $\theta$ in Eq. (12) to avoid eliminating $\theta$-valid functions.

For any $\epsilon,\delta\in(0,1)$, any Contextual Decision Process and function class $\mathcal{F}$ that admits a Bellman factorization with parameters $M$, $\zeta$, and $\eta$, suppose we run Oliver with any $\theta\in$, and $n_{\textrm{est}},n_{\textrm{eval}},n,\phi$ as specified in Theorem 1. Then with probability at least $1-\delta$, Oliver halts and returns a policy $\hat{\pi}$ which is at most

suboptimal compared to $V_{\mathcal{F},\theta}^{\star}$ defined in Definition 9, and the number of episodes required is at most

Proofs of Main Results

In this section, we provide the main ideas as well as the key lemmas involved in proving Theorem 1. We also show how the lemmas are assembled to prove the theorem. Detailed proofs of the lemmas are in Appendix C.

We begin by decomposing a policy-loss-like term into the sum of Bellman errors.

The structure of this lemma is similar to many existing results in RL that upper-bound the loss of following an approximate value function greedily using the function’s Bellman errors (e.g., Singh and Yee (1994)). However, most existing results are inequalities that use max-norm relaxations to deal with mismatch in distributions; hence, they are likely to be loose. This lemma, on the other hand, is an equality, thanks to the fact that we are comparing $V^{\pi_{f}}$ to $V_{f}$, not $V^{\star}$. As the remaining analysis shows, this simple equation allows us to relate policy loss (from the LHS) with the average Bellman error (the RHS) that we use to drive exploration. In particular, this lemma implies an explore-or-terminate behavior for the algorithm.

throughout the execution of the algorithm. Assume further that $f^{\star}$ is never eliminated. Then in any iteration $t$, one of the following two statements holds:

the algorithm terminates and the output policy $\pi_{t}$ satisfies $V^{\pi_{t}}\geq V_{\mathcal{F}}^{\star}-\epsilon$.

The lemma guarantees that the policy $\pi_{t}$ used at iteration $t$ in Olive has sufficiently large Bellman error on at least one of the levels, provided that the two conditions in Equation (16) are met. These conditions require that (1) we have reasonably accurate value function estimates from Line 1, and (2) we collect enough samples in Line 6 to form reliable Bellman error estimates under $f_{t}$ at each level $h$. The result of Theorem 1 can then be obtained using two further ingredients. First, we need to make sure that the first case in Lemma 2 does not happen too many times. Second, we need to collect enough samples in Lines 1 and 6 to ensure the preconditions in Equation (16). We first establish a bound on the number of iterations using the Bellman rank of the problem, before moving on to sample complexity questions.

If $\hat{\mathcal{E}}(f,\pi_{t},h_{t})$ in Eq. (5) always satisfies

throughout the execution of the algorithm ($\phi$ is the threshold in the elimination criterion), then $f^{\star}$ is never eliminated. Furthermore, for any particular level $h$, if whenever $h_{t}=h$ we have

then the number of iterations that $h_{t}=h$ is at most $M\log\left(\frac{\zeta}{2\phi}\right)/\log\frac{5}{3}.$

Precondition (18) simply posits that we collect enough samples for reliable Bellman error estimation in Line 12. Intuitively, since $f^{\star}$ has no Bellman error, this is sufficient to ensure that it is never eliminated. Precondition (19) is naturally satisfied by the exploration policies $\pi_{t}$ given Lemma 2. Given this, the above lemma bounds the number of iterations at which we can find a large Bellman error at any particular level.

Finally, we need to ensure that the number of samples collected in each of Lines 1, 6, and 12 of Olive can be upper bounded, which yields the overall PAC learning result in Theorem 1. The next three lemmas present precisely the deviation bounds required for this argument. The first two follow from simple applications of Hoeffding’s inequality.

holds for all $f\in\mathcal{F}$ simultaneously. Hence, we can set $n_{\textrm{est}}\geq\frac{32}{\epsilon^{2}}\log\frac{2N}{\delta}$ to guarantee that $|\hat{V}_{f}-V_{f}|\leq\epsilon/8$.

This controls the number of samples required in Line 1.

For any fixed $f_{t}$, with probability at least $1-\delta$,

This lemma can be seen as the sample complexity at each iteration in Line 6. Note that no union bound over $\mathcal{F}$ is needed here, since Line 6 only estimates the average Bellman error for a single function, which is fixed before data is collected. Finally, we bound the sample complexity of the learning step.

For any fixed $\pi_{t}$ and $h_{t}$, with probability at least $1-\delta$,

holds for all $f\in\mathcal{F}$ simultaneously. Hence, we can set $n\geq\frac{32K}{\phi^{2}}\log\frac{2N}{\delta}$ to guarantee that $|\hat{\mathcal{E}}(f,\pi_{t},h_{t})-\mathcal{E}(f,\pi_{t},h_{t})|\leq\phi$ as long as $\phi\leq 4$.

This lemma uses Bernstein’s inequality to exploit the small variance of the importance weighted estimates.

2 Proof of Theorem 1

Suppose the preconditions of Lemma 2 (Eq. (16)) and Lemma 3 (Eq. (18)) hold; we show them via concentration inequalities later. Applying Lemma 2, in every iteration $t$ before the algorithm terminates,

due to the choice of $\phi$. For level $h=h_{t}$, Eq. (19) is satisfied. According to Lemma 3, the event $h_{t}=h$ can happen at most $M\log\left(\frac{\zeta}{2\phi}\right)/\log\frac{5}{3}$ times for every $h\in[H]$. Hence, the total number of iterations in the algorithm is at most

Now we are ready to apply the concentration inequalities to show that Eq. (16) and (18) hold with high probability. We split the total failure probability $\delta$ among the following estimation events:

Estimation of $\hat{V}_{f}$ (Lemma 4; only once): $\delta/3$.

Estimation of $\hat{\mathcal{E}}(f,\pi_{t},h_{t})$ (Lemma 6; every iteration): same as above.

Since these events happen in a particular sequence, the proof actually bounds the probability of these failure events conditioned on all previous events succeeding. This imposes no technical challenge as fresh data is collected for every event, so it effectively reduces to a standard union bound.

Applying Lemmas 4, 5, and 6 with the above failure probabilities, we can verify that the choices of $n_{\textrm{est}},n_{\textrm{eval}},$ and $n$ in the algorithm statement satisfy the preconditions of Lemmas 2 and 3. Finally, we upper bound the total number of episodes as

Conclusions and Discussions

In this paper, we presented a new model for RL with rich observations, called Contextual Decision Processes, and a structural property, the Bellman factorization, of these models that enables sample-efficient learning. The unified approach allows us to address several settings of practical interest that have largely eluded RL theory to date. Via extensions of the main result, we also demonstrated that the techniques are quite robust and degrade gracefully with violation of assumptions. These results also elicit several further questions:

Can we obtain a computationally efficient algorithm for some form of this setting? Prior related work (for instance in contextual bandits (Dudik et al., 2011; Agarwal et al., 2014)) used supervised learning oracles for computationally efficient approaches. Is there a suitable oracle for this setting?

The sample complexity depends polynomially on the cardinality of the action space. Can we extend the results to handle large or continuous action spaces?

Can we address sample-efficient RL given only a policy class rather than a value function class? Empirical approaches in policy search often rely on policy gradients, which are subject to local optima. Are there parallel results to this work, without access to value functions?

Understanding these questions is a key bridge between RL theory and practice, and is critical to success in challenging reinforcement learning problems.

Appendix

Appendix A Lower Bounds

We include a result from Krishnamurthy et al. (2016) to formally show that, without making additional assumptions, the sample complexity of value-based RL for CDPs as introduced in Section 2 has a lower bound of order $K^{H}$.

The proof relies on the fact that CDPs include MDPs where the state space is arbitrarily large. Each instance of the MDP family is a complete tree with branching factor $K$ and depth $H$. Transition dynamics are deterministic, and only leaf nodes have non-zero rewards. All leaves give Ber$(1/2)$ rewards, except for one that gives Ber$(1/2+\epsilon)$. Changing the position of the most rewarding leaf node yields a family of $K^{H}$ MDP instances so collecting optimal Q-value functions forms the desired function class $\mathcal{F}$. Since $\mathcal{F}$ provides no information other than the fact that the true MDP lies in this family, the problem is equivalent to identifying the best arm in a multi-arm bandit with $K^{H}$ arms, and the remaining analysis follows exactly the same as in Krishnamurthy et al. (2016). ∎

A.2 A Polynomial Lower Bound that Depends on Bellman Rank

In this section, we prove a new lower bound for layered episodic MDPs that meet the assumptions we make in this paper.

We first recall some definitions. A layered episodic MDP is defined by a time horizon $H$, a state space ${\mathcal{S}}$, partitioned into sets ${\mathcal{S}}_{1},\ldots,{\mathcal{S}}_{H}$, each of size at most $M$, and an action space $\mathcal{A}$ of size $K$. The system descriptor is replaced with a transition function $\Gamma$ that associates a distribution over states with each state action pair. More formally, for any $s_{h}\in{\mathcal{S}}_{h}$, and $a\in\mathcal{A}$, $\Gamma(s_{h},a)\in\Delta({\mathcal{S}}_{h+1})$. The starting state is drawn from $\Gamma_{1}\in\Delta({\mathcal{S}}_{1})$, and all transitions from ${\mathcal{S}}_{H}$ are terminal.

There is also a reward distribution $R$ that associates a random reward with each state-action pair. We use $r\sim R(s,a)$ to denote the random instantaneous reward for taking action $a$ at state $s$. We assume that the cumulative reward $\sum_{h=1}^{H}r_{h}\in$, where $r_{h}$ is the reward received at level $h^{\textrm{th}}$ as in Assumption 1.

Observe that this process is a special case of the finite-horizon Contextual Decision Process and moreover, with the set of all value functions $\mathcal{F}=({\mathcal{S}}\times\mathcal{A}\rightarrow)$, admits a Bellman factorization with Bellman rank at most $M$ (by Proposition 1). Thus the upper bounds for PAC learning apply directly to this setting.

Fix $M\geq 4,H,K\geq 2$ and $\epsilon\in(0,\frac{1}{48\sqrt{8}})$. For any algorithm and any $n\leq cMKH/\epsilon^{2}$, there exists a layered episodic MDP with $H$ layers, $M$ states per layer, and $K$ actions, such that the probability that the algorithm outputs a policy $\hat{\pi}$ with $V(\hat{\pi})\geq V^{\star}-\epsilon$ after collecting $n$ trajectories is at most $11/12$. Here $c>0$ is a universal constant.

The result precludes a $o(MKH/\epsilon^{2})$ PAC-learning sample complexity bound since in this case the algorithm must fail with constant probability. The result is similar in spirit to other lower bounds for PAC-learning MDPs (Dann and Brunskill, 2015; Krishnamurthy et al., 2016), but we are not aware of any lower bound that applies directly to the setting. There are two main differences between this bound and the lower bound due to Dann and Brunskill (2015) for episodic MDPs. First, that bound assumes that the total reward is in $[0,H]$, so the $H^{2}$ dependence in the sample complexity is a consequence of scaling the rewards. Second, that MDP is not layered, but instead has $M$ total states shared across all layers. In contrast, the process is layered with $M$ distinct states per layer and total reward bounded in $$. Intuitively, the additional$H$dependence arises simply from having$MH$ total states.

At a high level, the proof is based on embedding $\Theta(MH)$ independent multi-arm bandit instances into a MDP and requiring that the algorithm identify the best action in $\Omega(MH)$ of them to produce a near-optimal policy. By appealing to a sample complexity lower bound for best arm identification, this implies that the algorithm requires $\Omega(MHK/\epsilon^{2})$ samples to identify a near-optimal policy.

We rely on a fairly standard lower bound for best arm identification. We reproduce the formal statement from Krishnamurthy et al. (2016), although the proof is based on earlier lower bounds due to Auer et al. (2002).

In particular, the problem instance used in this lower bound is one where the best arm $a^{\star}$ has reward $\textrm{Ber}(1/2+\epsilon)$, while all other arms have reward $\textrm{Ber}(1/2)$. The construction embeds precisely these instances into the MDP.

We construct an MDP with $M$ states per level, $H$ levels, and $K$ actions per state. At each level, we allocate three special states, $w_{h},g_{h}$, and $b_{h}$ for “waiting”, “good”, and “bad.” The remaining $M-3$ “bandit” states are denoted $s_{h,i},i\in[M-3]$. Each bandit state has an unknown optimal action $a_{h,i}^{\star}$.

For waiting states $w_{h}$, all actions are equivalent and with probability $1-1/H$ they transition to the next waiting state $w_{h+1}$. With the remaining $1/H$ probability, they transition randomly to one of the bandit state $s_{h+1,i}$ so each subsequent bandit state is visited with probability $\frac{1}{H(M-3)}$.

For bandit states $s_{h,i}$, the optimal action $a^{\star}_{h,i}$ transitions to the good state $g_{h+1}$ with probability $1/2+\tau$ and otherwise to the bad state $b_{h+1}$. All other actions transition to $g_{h+1}$ and $b_{h+1}$ with probability $1/2$. Here $\tau$ is a parameter we set toward the end of the proof.

Good states always transition to the next good state and bad states always transition to bad states.

The starting state is $w_{1}$ with probability $1-1/H$ and $s_{1,i}$ with probability $\frac{1}{H(M-3)}$ for each $i\in[M-3]$.

The reward at all states except $g_{H}$ is zero, and the reward at $g_{H}$ is one. Clearly the optimal policy takes actions $a_{h,i}^{\star}$ for each bandit state, and takes arbitrary actions at the waiting, good, and bad states.

This construction embeds $H(M-3)$ best arm identification problems that are identical to the one used in Proposition 7 into the MDP. Moreover, these problems are independent in the sense that samples collected from one provides no information about any others. Appealing to Proposition 7, for each bandit state $(h,i)$, unless $\frac{K}{72\tau^{2}}$ samples are collected from that state, the learning algorithm fails to identify the optimal action $a^{\star}_{h,i}$ with probability at least $1/3$.

After the execution of the algorithm, let $B$ be the set of $(h,s)$ pairs for which the algorithm identifies the correct action. Let $C$ be the set of $(h,s)$ pairs for which the algorithm collects fewer than $\frac{K}{72\tau^{2}}$ samples. For a set $S$, we use $S^{C}$ to denote the complement.

The second inequality is based on Proposition 7. Now, by the pigeonhole principle, if $n\leq\frac{(M-3)H}{2}\times\frac{K}{72\tau^{2}}$, then the algorithm can collect $\frac{K}{72\tau^{2}}$ samples from at most half of the bandit problems. Thus $|C|\geq(M-3)H/2$, which implies,

Thus with probability at least $1/11$ we know that $|B|\leq\frac{11}{12}(M-3)H$, so the algorithm failed to identify the optimal action on $1/12$ fraction of the bandit problems. Under this event, the suboptimality of the policy produced by the algorithm is,

Here we use the fact that the probability of visiting a bandit state is independent of the policy and that the policy can only visit one bandit state per episode, so the events are disjoint. Moreover, if we visit a bandit state for which the algorithm failed to identify the optimal action, the difference in value is $\tau$, since the optimal action visits the good state with $\tau$ more probability than a suboptimal one. The remainder of the calculation uses the transition model, the fact that $H\geq 2$, and finally the fact that $|B|\leq\frac{11}{12}(M-3)H$. Setting $\tau=48\epsilon$ and using the requirement on $\tau$ gives a stricter requirement on $\epsilon$ and proves the result. ∎

Appendix B Models with Low Bellman Rank

B.2 Generalization of Li (2009)’s Setting

Li (2009, Section 8.2.3) considers the setting where the learner is given an abstraction $\phi$ that maps the large state space ${\mathcal{S}}$ in an MDP to some finite abstract state space $\bar{\mathcal{S}}$ in an MDP. $|\bar{\mathcal{S}}|$ is potentially much smaller than $|{\mathcal{S}}|$, and it is guaranteed that $Q^{\star}$ can be expressed as a function of $(\phi(s),a)$. Li shows that when delayed Q-learning is applied to this setting, the sample complexity has polynomial dependence on $|\bar{\mathcal{S}}|$ with no direct dependence on $|{\mathcal{S}}|$.

In the next proposition, we show that a similar setting for finite-horizon problems admits Bellman factorization with low Bellman rank. In particular, we subsume Li’s setting by viewing it as a POMDP, where $\phi$ is a deterministic emission process that maps hidden state $s\in{\mathcal{S}}$ to discrete observations $\phi(s)\in\bar{\mathcal{S}}=\mathcal{O}$, and the candidate value functions are reactive so they depend on $\phi(s)$ but not directly on $s$ or any previous state. More generally, Proposition 8 claims that for POMDPs with large hidden-state spaces and finite observation spaces, the Bellman rank is polynomial in the number of observations if the function class is reactive.

Consider a POMDP introduced in Example 2 with $|\mathcal{O}|<\infty$, and assume that rewards can only take $C_{R}$ different discrete values.The discrete reward assumption is made to simplify presentation and can be relaxed. For arbitrary rewards, we can always discretize the reward distribution onto a grid of resolution $C_{R}$, which incurs $\eta=O(1/C_{R})$ approximation error in Definition 10. The CDP $(\mathcal{X},\mathcal{A},H,P)$ induced by letting $\mathcal{X}=\mathcal{O}\times[H]$ and $x_{h}=(o_{h},h)$, with any $\mathcal{F}:\mathcal{X}\times\mathcal{A}\to$, admits a Bellman factorization with $M=|\mathcal{O}|^{2}C_{R}K$ and $\zeta=2|\mathcal{O}|K\sqrt{C_{R}}$.

For any $f,f^{\prime}\in\mathcal{F},h\in[H]$, let $\nu_{h}(f^{\prime})$ and $\xi_{h}(f)$ be vectors of length $|\mathcal{O}|^{2}C_{R}K$. Let the entry of $\nu_{h}(f^{\prime})$ indexed by $(o_{h},a_{h},r_{h},o_{h+1})$ be

interpreted as the following: conditioned on the fact that the first $h-1$ actions are chosen according to $\pi_{f^{\prime}}$, what is the probability of seeing a particular tuple of $(o_{h},r_{h},o_{h+1})$ when taking a particular action for $a_{h}$? For $\xi_{h}(f)$, let the corresponding entry be (with $x_{h}=(o_{h},h)$ and $x_{h+1}=(o_{h+1},h+1)$ as the corresponding contexts in the CDP)

It is not hard to verify that $\mathcal{E}(f,\pi_{f^{\prime}},h)=\langle\nu_{h}(f^{\prime}),\xi_{h}(f)\rangle$. Since fixing $a_{h}$ to any non-adaptive choice of action induces a valid distribution over $(o_{h},r_{h},o_{h+1})$, we have $\|\nu_{h}(f^{\prime})\|_{1}=K$ and $\|\nu_{h}(f^{\prime})\|_{2}\leq K$. On the other hand, $\|\xi_{h}(f)\|_{\infty}\leq 2$ but the vector only has $|\mathcal{O}|^{2}C_{R}$ non-zero entries, so $\|\xi_{h}(f)\|_{2}\leq 2|\mathcal{O}|\sqrt{C_{R}}$. Together the norm bound follows. ∎

B.3 POMDP-like Models

Here we first state the formal version of Proposition 2, and prove Propositions 2 and 3 together by studying a slightly more general model (See Figure 1(c)).

Consider an MDP introduced in Example 1. With a slight abuse of notation let $\Gamma$ denote its transition matrix of size $|{\mathcal{S}}\times\mathcal{A}|\times|{\mathcal{S}}|$, whose element indexed by $((s,a),s^{\prime})$ is $\Gamma(s^{\prime}|s,a)$. Assume that there are two row-stochastic matrices $\Gamma^{(1)}$ and $\Gamma^{(2)}$ with sizes $|{\mathcal{S}}\times\mathcal{A}|\times M$ and $M\times|{\mathcal{S}}|$ respectively, such that $\Gamma=\Gamma^{(1)}\Gamma^{(2)}.$ Recall that we convert an MDP into a CDP by letting $\mathcal{X}={\mathcal{S}}\times[H]$, $x_{h}=(s_{h},h)$. For any $\mathcal{F}\subset\mathcal{X}\times\mathcal{A}\to$, this model admits a Bellman factorization with Bellman rank $M$ and $\zeta=2\sqrt{M}$.

The model that we use to study Proposition 2 and 3 simultaneously behaves like a POMDP except that both the transition function and the reward depends also on the observation, that is $\Gamma:{\mathcal{S}}\times\mathcal{O}\times\mathcal{A}\rightarrow\Delta({\mathcal{S}})$ and $R:{\mathcal{S}}\times\mathcal{O}\times\mathcal{A}\rightarrow\Delta()$. Clearly this model generalizes standard POMDPs, where the transition and reward are both assumed to be independent of the current observation.

This model also generalizes the MDP with low-rank dynamics described in Proposition 2: if the future hidden-state is independent of the current hidden-state conditioned on the observation (i.e., $\Gamma(s^{\prime}|s,o,a)$ does not depend on $s$), the observations themselves become Markovian, and we can treat $o$ as the observed state $s$ in Proposition 9, and the hidden-state $s$ as the low-rank factor in Proposition 9 (see Figure 1). Hence, Proposition 2 follows as a special case of the analysis for this more general model.

As in Proposition 3, we consider a class $\mathcal{F}$ reactive value functions. Observe that for the MDP with low rank dynamics, this provides essentially no loss of generality, since the optimal value function is reactive.

Let $(\mathcal{X},\mathcal{A},H,P)$ be the CDP induced by the above model which generalizes POMDPs, with $\mathcal{X}=\mathcal{O}\times[H]$ and $x_{h}=(o_{h},h)$. Given any $\mathcal{F}:\mathcal{X}\times\mathcal{A}\to$, the Bellman rank $M\leq|{\mathcal{S}}|$ with $\zeta=2\sqrt{|{\mathcal{S}}|}$.

For any $f,f^{\prime}\in\mathcal{F},h\in[H]$, consider

which is how actions are chosen in the definition of $\mathcal{E}(f,\pi_{f^{\prime}},h)$ (see Definition 2). Such a decision-making strategy induces a distribution over the following set of variables

That is, we first sample $s_{h}$ from the marginal $\mu_{f^{\prime}}$, and then sample the remaining variables conditioned on $s_{h}$. Notice that once we condition on $s_{h}$, the sampling of the remaining variable has no dependence on $f^{\prime}$, so we denote the joint distribution over the remaining variables (conditioned on the value of $s_{h}$) $\mu_{f|s_{h}}$.

Finally, we express the factorization of $\mathcal{E}(f,\pi_{f^{\prime}},h)$ as follows:

B.4 Predictive State Representations

In this subsection we state and prove the formal version of Proposition 4. We first recall the definitions and some basic properties of PSRs, which can be found in Singh et al. (2004); Boots et al. (2011). Consider dynamical systems with discrete and finite observation space $\mathcal{O}$ and action space $\mathcal{A}$. Such systems can be fully specified by moment matrices $P_{{\mathcal{T}}|\mathcal{H}}$, where $\mathcal{H}$ is a set of histories (past events) and ${\mathcal{T}}$ is a set of tests (future events). Elements of ${\mathcal{T}}$ and $\mathcal{H}$ are sequences of alternating actions and observations, and the entry of $P_{{\mathcal{T}}|\mathcal{H}}$ indexed by $t\in{\mathcal{T}}$ on the row and $\tau\in\mathcal{H}$ on the column is $P_{t|\tau}$, the probability that the test $t$ succeeds conditioned on a particular past $\tau$. For example, if $t=aoa^{\prime}o^{\prime}$, success of $t$ means seeing $o$ and $o^{\prime}$ in the next two steps after $\tau$ is observed, if interventions $a$ and $a^{\prime}$ were to be taken.

Among all such systems, we are concerned about those that have finite linear dimension, defined as $\sup_{{\mathcal{T}},\mathcal{H}}\textrm{rank}(P_{{\mathcal{T}}|\mathcal{H}})$. As an example, the linear dimension of a POMDP is bounded by the number of hidden-states. Systems with finite linear dimension have many nice properties, which allow them to be expressed by compact models, namely PSRs. In particular, fixing any ${\mathcal{T}}$ and $\mathcal{H}$ such that $\textrm{rank}(P_{{\mathcal{T}}|\mathcal{H}})$ is equal to the linear dimension (such $(\mathcal{H},{\mathcal{T}})$ are called core histories and core tests), we have:

For any history $\tau\in(\mathcal{A}\times\mathcal{O})^{*}$, the conditional predictions of core tests $P_{{\mathcal{T}}|\{\tau\}}$ (we also write $P_{{\mathcal{T}}|\tau}$) is always a state, that is, a sufficient statistics of history. This gives rise to the name “predictive state representation”.

Based on $P_{{\mathcal{T}}|\tau}$, the conditional prediction of any test $t$ can be computed from a PSR model, parameterized by square matrices $\{B_{ao}\}$ and a vector $b_{\infty}$ with dimension $|{\mathcal{T}}|$. Letting $t^{(i)}$ be the $i$-th (action, observation) pair in $t$, and $|t|$ be the number of such pairs, the prediction rule is

$P_{{\mathcal{T}},\mathcal{H}}$ is a matrix whose element indexed by $(t\in{\mathcal{T}},\tau\in\mathcal{H})$ is $P_{\tau t|\varnothing}$, where $\tau t$ is the concatenation of $\tau$ and $t$ and $\varnothing$ is the null history.

$P_{\mathcal{H}}=P_{\{\varnothing\},\mathcal{H}}$.

$P_{{\mathcal{T}},ao,\mathcal{H}}=P_{{\mathcal{T}},\mathcal{H}_{ao}}$, where $\mathcal{H}_{ao}=\{\tau ao:\tau\in\mathcal{H}\}$.

Now we are ready to state and prove the formal version of Proposition 4.

Consider a partially observable system with observation space $\mathcal{O}$, and the induced CDP $(\mathcal{X},\mathcal{A},H,P)$ with $x_{h}=(o_{h},h)$. To handle some subtleties, we assume that

$|\mathcal{O}|<\infty$ (classical PSR results assume discrete observations).

$o_{1}$ is deterministic (PSR trajectories always start with an action), and $r_{h}$ is a deterministic function of $o_{h+1}$ (reward is usually omitted or assumed to be part of the observation).

If the linear dimension of the original system is at most $L$, then with any $\mathcal{F}:\mathcal{X}\times\mathcal{A}\to$, this model admits a Bellman factorization with $M=LK$. Assuming further that the PSR’s parameters are non-negative under some choice of core histories and tests $(\mathcal{H},{\mathcal{T}})$ of size $|\mathcal{H}|=|{\mathcal{T}}|=L$, then we have $\zeta\leq 2K^{2}L^{3}\sqrt{L}/\sigma_{\min}^{3}$, where $\sigma_{\min}$ is the minimal non-zero singular value of $P_{{\mathcal{T}},\mathcal{H}}$.

For any $f,f^{\prime}\in\mathcal{F},h\in[H]$, define

$\mu_{f^{\prime},h}$ as the distribution vector over $(a_{1},o_{2},\ldots,o_{h-1},a_{h-1})\in(\mathcal{A}\times\mathcal{O})^{h-2}\times\mathcal{A}$ induced by $a_{1:h-1}\sim\pi_{f^{\prime}}$. (Recall that $o_{1}$ is deterministic.)

$P_{2|h-1}$ as a moment matrix whose element with column index $(o_{h},a_{h},o_{h+1})\in\mathcal{O}\times\mathcal{A}\times\mathcal{O}$ and row index $(a_{1},o_{2},\ldots,o_{h-1},a_{h-1})\in(\mathcal{A}\times\mathcal{O})^{h-2}\times\mathcal{A}$ is

$F_{f,h}$ as a vector whose element indexed by $(o_{h},a_{h},o_{h+1})\in\mathcal{O}\times\mathcal{A}\times\mathcal{O}$ is (recall that $x_{h}=(o_{h},h)$ and $r_{h}$ is function of $o_{h+1}$)

To show this, first observe that $\mu_{f^{\prime},h}^{\top}P_{2|h-1}$ is a row vector whose element indexed by $(o_{h},a_{h},o_{h+1})$ is

Next, we explicit construct $\xi_{h}(f)$ and $\nu_{h}(f^{\prime})$ by factorizing $P_{2|h-1}=P_{1}\times P_{2}$, where both $P_{1}$ and $P_{2}$ have no dependence on either $f$ or $f^{\prime}$. Recall that for PSRs, any history $(a_{1},o_{2},\ldots,o_{h-1})$ has a sufficient statistics $P_{{\mathcal{T}}|a_{1},o_{2},\ldots,o_{h-1}}$, that is a vector of predictions over the selected core tests ${\mathcal{T}}$ conditioned on the observed history. $P_{1}$ consists of row vectors of length $LK$, and for the row indexed by $(a_{1},o_{2},\ldots,o_{h-1},a_{h-1})$ the vector is