Provably Efficient UCB-type Algorithms For Learning Predictive State Representations
Authors: Ruiquan Huang
Yingbin Liang
Jing Yang
Venue: ICLR 2024 poster (2024)
Keywords: Reinforcement learning
Sequential decision-making problem
Predictive state representation
POMDP
UCB
online and offline
Primary area: reinforcement learning
TL;DR: We developed a provably efficient upper confidence bound algorithm for online and offline low rank decision-making problem, which is modeled by predictive state representation.

Abstract: The general sequential decision-making problem, which includes Markov decision processes (MDPs) and partially observable MDPs (POMDPs) as special cases, aims at maximizing a cumulative reward by making a sequence of decisions based on a history of observations and actions over time. Recent studies have shown that the sequential decision-making problem is statistically learnable if it admits a low-rank structure modeled by predictive state representations (PSRs). Despite these advancements, existing approaches typically involve oracles or steps that are computationally intractable. On the other hand, the upper confidence bound (UCB) based approaches, which have served successfully as computationally efficient methods in bandits and MDPs, have not been investigated for more general PSRs, due to the difficulty of optimistic bonus design in these more challenging settings. This paper proposes the first known UCB-type approach for PSRs, featuring a novel bonus term that upper bounds the total variation distance between the estimated and true models. We further characterize the sample complexity bounds for our designed UCB-type algorithms for both online and offline PSRs. In contrast to existing approaches for PSRs, our UCB-type algorithms enjoy computational tractability, last-iterate guaranteed near-optimal policy, and guaranteed model accuracy.

Published as a conference paper at ICLR 2024
PROVABLY EFFICIENT UCB-TYPE ALGORITTMS FOR LEARNING PREDICTIVE STATE REPRESENTATIONS
Ruiquan Huang"
Yingbin Liang!
Jing Yang"
ABSTRACT
The general sequential decision-making problem, which includes Markov de- cision processes (MDPs) and partially observable MDPs (POMDPs) as special cases, aims at maximizing a cumulative reward by making a sequence of deci- sions based on history of observations and actions time. Recent studies have shown that the sequential decision-making problem is statistically learn- able if it admits low-rank structure modeled by predictive state representations (PSRs). Despite these advancements, existing approaches typically involve ora- cles steps that are computationally intractable. On the other hand, the upper confidence bound (UCB) based approaches, which have served successfully as computationally efficient methods in bandits and MDPs, have not been investi- gated for more general PSRs, due to the difficulty of optimistic bonus design in these more challenging settings. This paper proposes the first known UCB-type approach for PSRs, featuring a novel bonus term that upper bounds the total vari- ation distance between the estimated and true models. We further characterize the sample complexity bounds for our designed UCB-type algorithms for both on- line and offline PSRs. contrast to existing approaches for PSRs, our UCB-type algorithms enjoy computational tractability, last-iterate guaranteed near-optimal policy, and guaranteed model accuracy.
INTRODUCTION
As general framework of reinforcement learning (RL), the sequential decision-maie making problem aims at maximizing a cumulative reward by making a sequence of decisions based on a history of observations and actions over time. This framework is powerful to include and generalize Markov Decision Processes (MDPs) and Partially Observable Markov Decision Processes (POMDPs), and
captures wide range of real-world applications such recommender systems (Li al. 2010 Wu
et al.
2021> 1, business management (De Brito & Van Der Laan
2009). economic simulation (Zheng
al.
2020 robotics (Akkaya et al.) (219)), strategic games Brrwnn& & Sandholm 2018 Vinyals
et al.
2019). and medical diagnostic systems (Hauskrecht & Fraser 2000).
However, tackling POMDPs alone presents significant challenges, not to mention general sequential decision-making problems. Many hardness results have been developed (Mossel & Roch, 2005 Mundhenk et al. 2000, Paaaddiiiriiou 2016). showing that learning POMDPs is already computationally and statistically intractable in the worst case. The reason is that the non-Markovian property of these problems implies that the sufficient statistics or belief about the current environmental state encompasses all observations and actions from past interactions with the environment. This dramatically increases the computational burden and statistical complexity, since even for a finite Abservation-action space, the possibilities of the beliefs are exponentially large in terms of the number of observation-actin pairs.
Tsitsiklis 1987 Vlassis et al. [012 |012] Krishnamurthy al,
In order to tackle these challenges, recent research has introduced various structural conditions for POMDPs and general sequential decision-making problems, such as reactiveness (Jiang et al. (017)). decodability (Du et al. 2019), revealing conditions (Liu al. 2022a), hindsight observability (Lee
let al.
2023 and low-rank representations with regularization conditions (Zhan et al.] 2022 Liu
al.
2022b Chen al. 2022). These conditions have opened up new possibilities for achieving
Ohio State University, Columbus, OH 43210, USA. 1 aang. 88900au. Apsu.eddpppuueuu
Peen State University, State College, PA 16801, USA. (#zh514, (=zh514,yangjiing aar

---

Published as a conference paper at ICLR 2024
polynomial sample complexities in general sequential Socision-making problems. Among them, predictive state representations (PSRs) (Littman & Sutton 2001) has been proved to capture and generalize rich subclass of sequential peciision-maiing problems such as MDPs and observable POMDPs. Yet, most existing solutions for PSRs involve oracles that might not be computationally tractable. For instance, Optimistic MILE (OMLE) (Liu al.. 2022b) involves step that maximizes the optimal value function over a confidence set of models. Typically, such a confidence set does not exhibit advantageous structures, resulting in a potentially combinatorial search within the set. Another popular posterior sampling based approach (Agarwal & Zhang 2022] Zhong et al. 2022) requires to maintain a distribution over the entire set of models, which is highly memory inefficient. In addition, most existing results lack a last-iterate guarantee and produce only mixture policies, which often exhibit a very large variance in practical applications. On the other hand, the upper confidence bound (UCB) based approach has been proved to be computationally efficient and pro- vide last-iterate guarantce in many cetiisionmmaking problems such as bandits (Auer et al.) 2002) and MDPs (Ménard et al., 2021). However, due to the non-Markovian property of POMDPs and general sequential To the best of our knowledge, such design has been seldomly explored in POMDPs and beyond. Thus, we are motivated to address the following important open question:
problems, designing an explicit UCB is extremely challenging.
Q1: Can we design UCB-type algorithm for learning PSRs that (a) is both computationally tractable and statistically efficient, and (b) enjoys the last-iterate guarantee?
Another important research direction in RL is offline learning (Lange al. 2012), where the learn- ing agent has access to a pre-collected dataset and aims to design a favorable policy without any interaction with the environment. While offline MDPs have been extensively studied (Jin et al. 20216, Xiong et al. 2022a; Xie et al., 2. there exist very limited studies of offline POMDPs from the theoretical perspective (Guo et al.] [2022] Lu et al. 2022). To our best knowledge, offline learning for more general model of PSRs has never been explored. Thus, we will further address the following research question:
Q2: Can we design a UCB-rype algorithm for offline learning of PSRs with guaranteed policy performance and sample efficiency?
Main Contributions: We provide affirmative answers to both aforementioned questions by making the following contributions.
We introduce the first known UCB-type approach to learning PSRs with only regularization assumption, characterized by a novel bonus term that upper bounds the total variation distance between the estimated and true models. The bonus term is designed based new confidence bound induced by a new model estimation guarantee for PSRs, and is computationally tractable.
We theoretically characterize the performance of our UCB-type algorithm for learning PSRs online, called PSR-UCB. In contrast to existing approaches, PSR-UCH is computationally tractable with only supervised learning oracles, guarantees near-optimal policy in the last iter- ation, and ensures model accuracy. When the rank of the PSR is small, our sample complexity matches the best known upper bound in terms of the rank and the accuracy level.
We further extend our UCB-type approach to the offline setting, and propose the PSR-LCB algorithm. We then develop an upper bound on the performance difference between the output policy of PSR-LCB and any policy covered by the behavior policy. The performance difference scales in OCC.oV/K where C.o is the coverage coefficient and is the size of the offline dataset. This is the first known sample complexity result on offline PSRs.
Technically, we develop two key properties for PSRs to establish the sample complexity guar- antees: (a) a new estimation guarantee on the distribution of future observations conditioned on empirical samples, enabled by the stable model estimation step, and (b) a new relationship between the empirical UCB and the ground-trui -ruth UCB. We believe these insights advance the current understanding of PSRs, and will benefit future studies on this topic.
RELATED WORK
Learning MDPs and POMDPs. The MDP is basic model in RL that assumes Markovian prop- erty in the model dynamics, i.c. the distribution of the future states only depends on the current system state. Researchers show that learning tabular MDPs (with finite state and action spaces) is
2

---

Published as a conference paper at ICLR 2024
both computationally and statistically efficient in both online setting (Auer et al.) 2008: Azar et al.
2017 Dann et al. 2017, Agrawal & Jial 2017; Jin et al. 2018 Li et al. 2021 and offliee setting
(Jin et al. [2021b] Rashidinejad et al.] [021] Yin et al.) [221]] Chang et al. 202I).
Learning MDPs
with function approximations is also well-studied by establishing favorable statistical complexity
and computation efficiency (Jin et al., 2020; Wagenmaker & Jamieson 2022 Zanette et al. 2020
Zhou et al. 2021 Agarwal et al. 2020; [eehara ct al. 2021; Du et al. [021 [ Foster et al. 202
al. 202 la Jiang et al. 2017: Wang al. 2020 Jin al. b02 Xiong et al. [022a
Xie al.
2021 Notably, several algorithms designed for learning MDPs with general function approxima-
tions can be extended
solve a subclass of POMDPs. In particular, OLIVE (Tiagg al.) (017)
and GOLF (Jin et al.) (021a), which are originally designed for MDPs with low Bellman rank and low Bellman-Eluder dimension, respectively, can efficiently learn reactive POMDPs, where the op-
timal policy only depends on the current observation. Besides, Du et al. (2019); Efroni et al. (2022)
study decodable RL where the observations determine the underlying states, and Kwon et al. (2021) investigate latent MDPs where there are multiple MDPs determined by some latent variables.
To directly address the partial observability in POMDPs, some works assume exploratory data or reachability property and provide polynomial sample complexity for learning these POMDPs (Guo) et al.] 2016 Azizzadenesseli et al. 2016] Xiong et al. 2022b). Others tackle the challenge of explo-
ration and exploitation tradeoff in POMDPs by considering various sub-classes of POMDPs such as low-rank POMDPs (Wang et al. 2022). observable POMDPs (Golowich et al. 2022a b), hindsight observability (Lee et al.] (023). and weakly-revealing POMDPs (Liu et al.J [2022c a). Furthermore, Liu et al. (2022c): Uchaara et al. (2022a) propose computationally efficient algorithms for POMDPs with deterministic latent transitions. Notably, Golowich et al. (2022a) propose a provably efficient algorithm for learning observable tabular POMDPs without computationally intractable oracles. Fi- nally, [u al.] (2022) > study the offline POMDPs in the presence of confounders, and [Guo et al.) (2022) provide provably efficient algorithm for offline linear POMDPs. After the initial submission of this work, we notice that a UCB-type algorithm has been studied by Guo et al. (2023) under low-rank L-step decodable POMDPs, which subclass of the PSRs considered here.
Learning PSRs and general sequential recision-making problems. The PSR is first introduced by Littman & Sutton 2001);: Singh et al. (2012) and considered as general representation to model dynamic systems. A line of research (Boots et al., 2011 Hefny et al. 2015; Jiang et al., 2018 Zhang et al. 2022) obtain polynomial sample complexity with observability ity assumption and spectral techniques. Later, Zhhan et al. (2022) demonstrate that learning regular PSRs is sample efficient and can avoid poly(IO) (")) in the sample complexity. [eehara al. (2022b) propose PO-bilinear class that captures a rich class of tractable RL problems with partial observations, including weakly revealing POMDPs and PSRs and design an actor-critic style algorithm. For the works most closely related to ours, Liu et al. (2022b) propose a universal algorithm known as OMLE, which is capable of learning PSRs and its generalizations under certain conditions: Chen et al. (2022) enhance the sample complexity upper bounds for three distinct algorithms, including OMLE, E, the model-based posterior sampling, and the estimation-to-deciis (022) address the general sequential decision-making problem by posterior sampling under newly proposed low generalized Eluder coefficient. However, as elaborated in Section T those approaches not efficient in terms of computational complexity or memory.
type algorithm (Foster et al. (02T]; Zhong et al.)
3 PRELIMINARIES
Problem Setting. We consider a finite horizon episodic sequential decision-making problem, de- fined by tuple?? P = (O,A, A, H, P, Rt,, where O represents the observation space, A is a finite action space, H is the number of time steps within an episode, IP = (Px} determines the model dynamics, i.c., (o, (o,,oo..... lo...........................l 0h- 1. A1,..... .n-1), where o, E O is the observation at time step t, and a, E A the action taken by the agent at time step & for all t E (1....,I) h), and R (O x A) -> [0, 1] is the reward function defined on trajectories of one episode. We denote historical trajectory at time step as TA := (01, ai,..........................a Ob, an), , and denote a future trajectory as wh := (ON+1, ........, ,0,, ah).
The set of all Th is denoted by Hr, (O x A)a and the set of all future trajectories is denoted by Ωr = (O x A) "r-hh In addition, let wh- = (o.t1...... oH) and wi = (a)+1..... ax) be the observation sequence and the action sequence contained in Whe respectively. Similarly, for his- tory The we denote T and rf as the observation and action sequences in Tr respectively. Notably, the general framework of sequential decision-making problem subsume not only fully observable
3

---

Published as a conference paper at ICLR 2024
MDPs but also POMDPs as special cases, because in MDPs, Pn Th-1) = Pn 0h-1, ag a/-1) and in POMDPs, PA, (o) 7--1) can be factorized as (on 7b-1) ) = Σ. (onlonls) Th-1) where represents unobserved states.
The interaction between an agent and proceeds as follows. At the beginning of cach episode, the environment initializes fixed observation on at time step 1. After observing 01, the agent takes action 0j. and the environment transits to 02, which is sampled according to the distribution Pi (02 l0,, ai) 1). Then, any time step h 2 2. due the non-Markovian nature of the problem, the agent takes action ag based on all past information (Ta-1, Ob.,, and the environment transits to 0 sampled from PA, (oh++1/TT) The interaction terminates after time step H.
The policy ग = { } of the agent is collection of H distributions where Ty (ax/7e-11 0%) is the probability of choosing action an at time step given the history Th-1 and the current observation on. For simplicity, we use (T%) = m (oh, Th-1) 11..... ... T (a1 (1)) to denote the probability of the sequence of actions 7h given the observations T.. We denote P" as the distribution of trajectories
induced by policy π under dynamics IP. The value of a policy πT under IP and the reward R is denoted by VF,R = E-g- (P((T))
The goal of the agent is to find an e-optimal policy テ that satisfies max; VE. finding a near-optimal policy for general decision- makinn problem incurs exponentially large sam- ple complexity in the worst case, in this paper we follow the line of research in Zhan et al. (2022); Chen et al. (2022); Zhong et al. (2022) and focus on the low-rank class of problems. To define a low-rank problem, we introduce the dynamic matrix Dy R Hx Ix Pal for each h, where the entry at the TV th row and -th column of D2 is P(wh, T (
V S e. Since
Definition (Rank-r sequential decision-making problem) A sequential decision-making prob- lem is rank if for any h, the model dynamic matrix Dj has rank r.
Predictive State Representation (PSR). To exploit the low-rank structure, we assume that for each h, there exists a set of future trajectories, namely, core tests (known to the agent) On = {qh: .... 9% } C S4, such that the submatrix restricted to these tests [x(2n) has rank r, where dn 2 is positive integer. This special set Q allows the system dynamics to be factorized as P(Wg, Tg wx) ) = m(wx)" ((()))). where m(wn)..(77) E IRdh, and the E-th coordinate of (TT) is the joint probability of To and the f-th core test g.. Mathematically, if we use of, and af, to denote the observation sequence and the action sequence of Lh respectively, then P(o T/ ( af,) = [Tx)ee By Theorem C.1 in Liu et al.) (2022b), any low-rank decision-male making problem admits (self- consistent) predictive state representation 0- = {Φr, Ma given core tests { H-1 such that for any Th E Hr, Wh. E Si.
(TT) = Mι(on, an) (110,111)00 m(w)) OH aH) M1+110n41, Σ Mh an-1) = oh, P(o 01 oila,........................... an) = 25 (TT), 0
where M, O x A xd, xdz Qr E Rd, and vi do For ease of presentation, we assume do is known to the agen | Notably, the normalized version of v'(T.) with respect to Φ (TT)), denoted as (Tx)) = / (TT)), is known as the prediction vector (Littman Sutton 2001 or prediction feature of The since [v((T)]e = Po |T,, a.). As illustrated in Appendices and the prediction feature plays an important role in our algorithm design.
In the following context, we use Po to indicate the model determined by the representation (dr, M For simplicity, we denote V = Vg Moreover, let O = {af, l be the set of action sequences that are part of core tests, constructed by eliminating any repeated action sequence. O known as the set of core action sequences, plays crucial role during the exploration process. Selecting from these sequences is sufficient to sample core tests, leading to accurate estimate of 9.
We further assume that the PSRs studied in this paper are well-conditioned, as specified in Assump- tion I Such an assumption and its variants are commonly adopted in the study of PSRs (Liu et al.
2022b
Chen et al. 2022; Zhong et al. 2022).
'The sample complexity of learning tho if it is unknown is relatively small compared with learning other parameters.
4

---

Published as a conference paper at ICLR 2024
Assumption (y-well-conditioned PSR) A PSR 0 is said to be --welll-oonditooee if
Vh,
max JER S
max max
Σ [a(walTa)mll Txl ?
(I)
7h
Assumption requires that the error of estimating 0 does not significantly blow up when the estima- tion error of estimating the probability of core tests is small.
Notations. We denote the complete set of model parameters as and the true model parameter as 0" For a vector 2, || stands for Vx Az, and the i-th coordinate of is represented as [æ] For functions IP and Q (not necessarily probability measures) over a set X, the total variation distance between them is Drv (P(z), Q(x)Q((x))))))))) Σ. P(z) - Q(z)), while the hellinger-squared distance is de- Q(=))2 Note that definition of the total variation distance is slightly different from convention by a constant factor. We define d = max, d and QA max, O We use en (Tt, œ²) denote the policy that takes at the initial h - 1 steps and switches to n from the h-th step. Lastly, uxy represents the uniform distribution over the set X'
fined as D (a()), Q(r)) = ) (vP(a)
ONLINE LEARNING FOR PREDIVTIVE STATE REPRESENTATIONS
In this section, we propose a model-based algorithm PSR-UCB which features three main novel designs: (a) stable model estimation step controlling the quality of the estimated model such that its prediction features of the empirical data are useful in the design of UCB, (b) an upper confidence bound that captures the uncertainty of the estimated model, and (c) termination condition that guar- antees the last-iterate model is near-accurate and the corresponding greedy policy is near-optimal.
4.1 ALGORITHM
The pseudo-code for the PSR-UCB algorithm is presented in Algorithm We highlight the key idea of PSR-UCB: in contrast to the design of UCB-type algorithms in MDPs that leverage low-rank structures, we exploit the physical meaning of the prediction feature (TT)) to design bonus terms that enable efficient exploration. In particular, "physical meaning" refers to that each coordinate of a prediction feature of Th represents the probability of visiting a core test conditioned on T and taking the corresponding core action sequence. Therefore, it suffices to explore a set of whose prediction features can span the entire feature space, and use core action sequences to learn those "base" features. We next elaborate the main steps of the algorithm in greater detail as follows.
Exploration. At each iteration k, PSR-UCB constructs greedy policy = based on a previous dataset D = (D), 11 together with an estimated model Øc-11 Intuitively, to enable efficient exploration, x= --1 is expected to sample Th that "differs" the most from previous collected samples Th D How to quantify such differences forms the foundation of our algorithm design and will be elaborated later.
J==0
Then, for each E [H] PSR-UCB first uses 7 k-1 to get a sample 24Ah 7--1 then follows the policy uou get wh-11 where 2 (A 2 O Here, superscripts k, h represent the
k.b
index of episodes. In other words, PSR-UCB adopts the policy (n uugez) ) to collect sample trajectory (T (/-1), which, together with D/= forms the dataset D
The importance of uniformly selecting actions from 2ex,, can be explained follows. First, the actions in Q assist to learn the prediction feature 20 (T/--1 as its f-th coordiate equals to Pe. (of 17h af af.1)) Second, the action sequence (an, af)) E A OA helps us to estimate ME (on, (T, because its l-th coordinate represents Po- (of, 0z 1--11 ab; a.)) Therefore, by uniform exploration on QueP, given that T differs from previous dataset D/-1) PSR-UCB collects the most informative samples for estimating the true model 6*.
AAA
Stable model estimation. With the updated dataset D {D]], we estimate the model by maxi- mizing the log-likelihood functions with constraints. Specifically, PSR-UCB extracts any model @k
5

---

Published as a conference paper at ICLR 2024
from B defined as:
e = {0: wh, (Th, R) E Di, Pā 2 Puiin
B
e min
Σ
log 0
6 (Tn)
max
log Por (TH) -
(2)
(TA,
(TH.
where the estimation margin is pre-defined constant. Here the threshold probability Pmii is sufficiently small to guarantee that, with high probability, 0 E @ Note that compared to existing vanilla maximum likelihood estimators (Liu al., 2022b Chen et al. 2022). PSR-UCB has an k,h+l sampled infer the conditional probability Pa- (T .wh. or the prediction feature v/ (Th (Tk,t+1) and learning this feature is
additional constraint O min- This is crucial condition, as T
useless even harmful if F d+ (T) 1. too small, is proved that 0 E O (Proposition 7
Design of UCB with prediction features. From the discussion of the previous two steps, that the prediction features are vital since (a) actions in 2% can efficiently explore the coordi- nates of these features, and (b) constraint O ensures the significance of the learned features of the collected samples. Therefore, in the next round of exploration, our objective is to sample ты whose prediction feature exhibits the greatest dissimilarity compared with those of the previously collected samples 7 D Towards that, PSR-UCB constructs an upper confidence bound V j
for Drv (PPg. (TH), (T((H)) where the bonus 8 is defined as:
H Σ
Ü = +
(TAN (THT)) min T
Σ
(T)
(3)
with pre-specified regularizer and UCB coefficient a. Note that large (TH) indicates that ти "perceived" to be under explored since the estimated prediction feature *h (Tx) is significantly different from (Th))mDD
Algorithm Learning Predictive State Representation with Upper Confidence Bound (PSR-UCH UCB)
1: Input: threshold probability Pmin. estimation margin regularizer X, UCB coefficient a.
2: for k = 1,..., K do
3: for 1, 1,... H do
4:
Use u( uopu, to collect data k. (w k.h T-1).
5:
D +- D ((TH ( (uppr ))}-
6:
end for
D = H- H-1 Extract any 0² E B according to Equation 2
7:
8:
Define bonus function 6* (TH) according to Equation
9:
Solve x arg max, V
$*,6**
10:
if V B+ S c/2 then
0= 0==00 0k break.
12:
end if
end for
14:
Output: म = max, V V
Design of the greedy policy and last-iterate guaranteed termination. The construction of UCB implies that 0 highly uncertain on the trajectories ти with large bonuses. Thus, PSR-UCB finds greedy policy n = arg max, Va ft, and terminates if Va is sufficiently small, indicating that @ .e
the estimated model 6 sufficiently accurate on any trajectory. Otherwise, n² serves for the next iteration A: + 1, as it tries to sample the most dissimilar (Tx) compared with the previous samples ( h,++1 for efficient exploration. We remark that the termination condition favors the last-iterate guarantee, where single model and greedy policy are identified. Compared with algorithms that output mixture policy (e.g. uniform selection in a large policy set), this guarantee may present lower variance in practical applications.

---

Published as a conference paper at ICLR 2024
Remark 1 (Computation) Our algorithu only calls the MLE oracle H times in each episode to construct the confidence model class in Equation (2). Similar to Cuo et al. (2023), our UCB based planning avoids searching for optimistic policy that maximizes the optimal value over large model class, thus is more amendable to practical implementation.
Remark 2 (Reward-free PSRs) Our algorithm naturally handles reward-free PSRs since during
the exploration, no reward information needed. In the reward-free -ree setting, Algorithm 7 exhibits greater advantage compared with existing reward-free algorithms (Liu et al. [022b] Chen et al. 2022] for PSRs, which require a potentially combinatorial optimization oracle over a pair of models in the (non-convex) confidence set and examination of the total variance distance and are compu- tationally intractable in general (Golowich et al. 2022a
4.2 THEORETICAL RESULTS
In this section, we present the theoretical results for PSR-UCB. To make a general statement, we first introduce the notion of the optimistic net of the parameter space.
Definition (Optimistic net) Consider two bounded functions IP and IP over a set X. Then, P is
e-optimistic over IP if (a) P(x) 2 (x) Vx E X, and (b) Dry (P, IP) E. The s-optimistic net of θ is a smallest finite space O so that for all 6 E O there exists f O when Pg is e-optimistic over Po-
Note that if θ is the parameter space for tabular PSRs (including finite observation and action spaces)
with rank r. we have |+,1 S 1O log HIOIAA| (see Proposition 3 or Theorem 4.7 in Liu et al. )022b)). Now, we are ready to present the main theorem for PSR-UCB.
Theorem Suppose Assumption holds. Let Pmin = O O K LA
B = O (el), where
Then, with probubility at least 5, PSR-UCB outputs a model 6 and policy ग that satisfy
0
Pmin X 2 2 max (T.QAVB//) and O QAv KH
VX
LAVE
vdH
72t
V Vo-.re S. c, and Vx, Drv (PT (TH), PT. (TH)) Se
(4)
In addition, PSR-UCB terminates with a sample complexity of
ò ((1+ (2H 72
22H
rdH T
7/62
The following remarks highlight few insights conveyed by Theorem I. First, Theorem D explic- itly states that PSR-UCB features last-iterate guarantee, i.e., the guaranteed performance is on the last output of the algorithm. This is in contrast to the previous studies (Liu et al. 2022a b Chen et al. 2022) POMDP and/or PSRs, where the performance guarantee is on mixture of policies obtained over the entire execution of algorithms. Such policies often have large variance. Second,
in the regime with a low PSR rank such that r < Li our result matches the best known sample complexity (Chen al. 2022) in terms of rank r and e. Third, thanks the explicitly constructed bonus function b, whose computation complexity is polynomial in terms of the iteration number and the dimension d, PSR-UCB is computationally tractable, provided that there exist oracles for plan-
ning in POMDPs (e.g. Line 9 in PSR-UCB) (Papadimitriou & Tsitsiklis 1987 and maximum likelihood estimation.
Guo et al. 2023)
Proof Sketch of Theorem 7 The proof relies on three main steps corresponding to three new tech- nical developments, respectively. (a) A new MIEE guarantee. Due the novel stable model esti- mation design, we establish a new MLE guarantee (which has not been developed in the previous studies) that the total variation distance between the conditional distributions of future trajectories conditioned on (Tr. TT) E D is small. Mathematically, (T..))ED) Drv (IP ğa. (w), (m),, P... (T.)
*Computation complexity of calculating total variation distance is still open question (Bhattacharyya al. 2022 and only some specific forms can be approximated in polynomial time. [ can he the cardinality of O if it finite, the measure of O if it is measurable set with positive and bounded measure.
7

---

Published as a conference paper at ICLR 2024
upper bounded by constant. (b) A new confidence bound. Based on (a) and the obser- vation that the estimation error of the prediction feature (n,) is upper bounded by the total variation distance DTv jz (wi, (V)]TT),P (gn,((/TTT) (wn/TTh) the estimation enror of other prediction features captured by the function || (Tw)l(opp)) up to some constant, leading to the valid UCB de- sign. (c) A new relationship between the empirical bonus and the ground-truth bonus. To characterize the sample complexity, need to show that V 0r,b* can be small for some k. One approach to prove that Σx V is sub-linear. We validate this sub-linearity by establishing j0,
8 $ O (En (TW)ll where U A E (T))) This inequality bridges the empirical and the ground-truth bonuses, and yields the final result when combined with the elliptical potential lemma (Carpentier et al., 2020). The complete proof can be found in Ap- pendix
When PSR is specialized to m-step decodable POMDPs (Liu et al. 2022b), as elaborated in Ap- pendix C.5 the sample complexity of PSR-UCB does not depend d with slight modification of the algorithm.
Corollary When PSR specialized to m-step decodable POMDPs, for each k. there exists set
of matrices (GR hi==1 C Rdxr such that if we replace & (Tn) by ( & (T,) in Equation 3
for any The then PSR-UCB terminates with a sample complexity of poly(r, (afbyrr,1/// 1/7, QA, B)/E
OFFLINE LEARNING FOR PREDICTIVE STATE REPRESENTATIONS
In this section, we develop a computationally tractable algorithm PSR-LCB to learn PSRs in the offline setting. In offline PSRs, dataset DD contains K pre-collected trajectories that are indepen- dently sampled from a behavior policy nhh We slightly generalize the learning goal to finding a policy that can compete with any target policy whose coverage coefficient is finite. We highlight that PSR-LCB the first offline algorithm for learning PSRs. When specialized to POMDPs, our pro- posed algorithm enjoys better computational complexity than confidence region based algorithms (Guo et al. 2022).
5.1 ALGORITHM
We design PSR-LCB algorithm for offline PSRs, whose pscudo-code is presented in Algorithm 2 In the offline setting, we also leverage the prediction feature (b(T)) to design lower confidence bound (LCB) of the true value function v.... We explain the main steps of the algorithm in detail as follows.
stable model estimation. Inspired by PSR-UCB, where the prediction features (T(TT) ) at each step are learned through separate datasets, PSR-LCB first randomly and evenly divides D into H datasets Do Do...., DH D The goal of this division is to separately learn " (Tr) ) for each h E {0, I,...,ft 1}. Then, PSR-LCB extracts a model 0 from Bp defined as:
{0: Vh, Th Dr. P7 (Th) 2 Pmii )
Omin
BD A Omin
Σ
log Pa (TH)
max 8'e6main
log pad (TH) $3.
(5)
THED
ED
where the estimation margin @ is pre-specified constant, and Puuin guarantees that, with high prob- ability, E emin-
Following similar reasons of the design of PSR-UCB, the constraint O controls the quality of the estimated model such that the prediction features of the behavior samples are non-negligible and useful for learning.
8

---

Published as a conference paper at ICLR 2024
Design of LCB with prediction features. Given the estimated model 9, PSR-LCB constructs upper confidence bound of Drv
(Pg ( (THT) in the form of V where b(TH) is defined as:
= Σ (T.)) b(TH) = min Σ 1 (TTT))
((00)-
(6)
TRE
with pre-defined regularizer À and LCB coefficient é. Note that we refer the L.B. coefficient, as we adopt the pessimism principle in offline learning. Differently from PSR-UCB, we aim to select policy where the estimated model exhibits the least uncertainty. Therefore, the output policy should allocate a high probability to тH if b(TH) is small.
Output policy design. Building on the discussion above, PSR-LCB outputs policy क arg max= V V which maximizes a lower confidence bound of V
6.R
0,5'
Algorithm Offline Predictive State Representations with Lower Confidence Bound (PSR-LCB)
1: Input: Offline dataset D, threshold probability Pmin. estimation margin 3, regularizer X, LCB coefficient å
2: Estimate model 0 E B according to Equation
3: Construct b according to Equation (
4: Output = arg max, 0.R 0.8'
5.2 THEORETICAL RESULTS
In this section, we develop the theoretical guarantee for PSR-LCB. To capture the distribution shift in offline PSRs, we require that the foo norm of the ratio between the probabilities the entire trajectory under the target policy and under the behavior policy is finite. Mathematically, if π is
Pg. (na)
the target policy, then, $0000 max max
Dc. We note that, for general sequential g- (Tx)
decision-main making problems, no Markovian property hidden state is assumed. Therefore, defining the coverage assumption on a single observation-acee actioon pair as in Xie et al. (202T) for offline MDP,
or on hidden states as in Guo et al. 2022) for offline POMDP, is not suitable for PSRs.
Theorem Suppose Assumption holds. Let 1. = mina,
ag E 2 n (an)) Punin
5 O KH OIAA)
E O Priin ) β Ollog ( l). X 70 3 mux/ /F.QA WH/ and ā = O
QAVAH 7
VX /2
(QAVAH Then, with probability at least 1 5, the oulput of Algorithm 2 salisfies thut
KH
((
((F++
VT+ QAVH
C
QAH2 Va
r
VT, Vē - V SO
Theorem states that for any target policy with finite coverage coefficient C
70,00" the performance C -0 which is
degradation of the output policy म with respect to the target policy π is at most Ö( negligible if the size of the offline dataset K is sufficiently large. Thee full proof is in Appendix D
vR
CONCLUSION
We studied learning predictive state representations (PSRs) for low-rank sequential decision-making problems. We developed a novel upper confidence bound for the total variation distance of the estimated model and the true model that enables both computationally tractable with MLE oracles and statistically efficient learning for PSRs with only supervised learning oracles. Specifically, we proposed PSR-UCB for online learning for PSRs with last-iterate guarantee, i.c. producing not only near-optimal policy, but also near-accurate model. The statistical efficiency was validated by a polynomial sample complexity in terms of the model parameters. addition, we extended this UCB-type approach to the offline setting and proposed PSR-LCB. Our theoretical result offers an initial perspective on offline PSRs by demonstrating that PSR-LCB outputs a policy that can compete with any policy with finite coverage coefficient.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
The work of Huang and Yang was supported in part by the U.S. National Science Founda- tion under the grants
CNS-200313/ and CNS-2030026. The work of Y. Liang was supported in part by the U.S. National Science Foundation under the grants RINGS-2148253, CCF- 1900145, and CNS-2112471.
IMPACT STATEMENT
This paper presents work whose goal is to advance the field of Machine Learning- There are many potential societal consequences of our work, none which we feel must specifically highlighted here.
REFERENCES
Yasin Abbasi- Yadkori, Dávid Pál, and Csaba Szepesvári. Improved algorithms for linear stochastic bandits. Advances in neural information processing systems, 24, 2011.
Alekh Agarwal and Tong Zhang. Model-based with optimistic posterior sampling: Structural conditions and sample complexity. arXiv preprint arXiv:2206.07665) 2022.
Alckh Agarwal, Sham Kakade, Akshay Krishnamurthy, and Wen Sun. Flambe: Structural com-
plexity and representation learning of low rank mdps. Advances in neural information processing systems, 33:20095-20111 2020.
Shipra Agrawal and Randy Jia. Optimistic posterior sampling for reinforcement learning: worst- case regret bounds. Advances in Neural Information Processing Systems, 30, 2017.
Ilge Akkaya, Marcin Andrychowicz. Maciek Chociej, Mateusz Litwin, Bob McGrew, Arthur Petron, Alex Paino, Matthias Plappert, Glenn Powell, Raphael Ribas, et al. Solving rubik's cube with a robot hand. arXiv preprint 1900:0910.07113 2019.
Peter Auer, Nicolo Cesa-Bianchi, and Paul Fischer. Finite-time analysis of the multiarmed bandit problem. Machine learning, 47:235-256, 2002.
Peter Auer, Thomas Jaksch, and Ronald Ortner. Near-optimal regret bounds for reinforcement learning, Advances neural information processing systems, 21, 2008.
Mohammad Gheshlaghi Azar, Ian Osband, and Rémi Munos. Minimax regret bounds for reinforce- ment learning. In International Conference on Machine Learning, pp. 263 272. PMLR, 2017.
Kamyar Azizzadenesheli, Alessandro Lazaric, and Animashree Anandkumar. Reinforcement learn- ing of pomdps using spectral methods. In Conference on Learning Theory, pp. 193-256 PMLR, 2016.
Arnab Bhattacharyya, Sutanu Gayen, Kuldeep Meel, Dimitrios Myrisiotis, Pavan, and NV Vin- odchandran. On approximating total variation distance. arXiv preprint arXiv:2206.07209 2022.
Byron Boots, Sajid M Siddiqi, and Geoffrey J Gordon. Closing the Saaming-planniingg loop with predictive state representations. The International Journal of Robotics Reseurch, 30(7):954-966. 2011.
Noam Brown and Tuomas Sandholm. Superhuman ai for heads-up no-limit poker: top professionals. Science, 359(6374):418 424, 2018.
Libratus beats
Alexandra Carpentier, Claire Vernade, and Yasin Abbasi-Yan Yadkori. The elliptical potential lemma revisited. arXiv preprint arXiv:2010.10182 2020.
Jonathan Chang, Masatoshi Uchara, Dhruv Sreenivas. Rahul Kidambi, and Wen Sun. Mitigating covariate shift imitation learning via offline data with partial coverage. Advances in Neural Information Processing Systems, 34:965-979, 2021.
10

---

Published as a conference paper at ICLR 2024
Fan Chen, Yu Bai, and Song Mei. Partially observable rl with b-stability: Unified structural condition and sharp sample-effiiien algorithms. arXiv preprint arXiv:2209. 2022.
Varsha Dani, Thomas P Hayes, and Sham M Kakade. Stochastic linear optimization under bandit feedback. 2008.
Christoph Dann, Tor Lattimore, and Emma Brunskill. Unifying pac and regret: Uniform pac bounds for episodic reinforcement learning. Advances in Neural Information Processing Systems, 30, 2017.
Marisa P De Brito and Erwin A Van Der Laan. Inventory control with product returns: The impact of imperfect information. European journal of operational research, 194(1):85-101, 2009.
Simon Du, Akshay Krishnamurthy, Nan Jiang, Alekh Agarwal, Miroslav Dudik, and John Langford. Provably efficient with rich observations via latent state decoding. In International Conference on Machine Learning, pp. 1665-1674. PMLR, 2019.
Simon Du Sham Kakade, Jason Lee, Shachar Lovett, Gaurav Mahajan, Wen Sun, and Ruosong Wang. Bilinear classes: A structural framework for provable generalization in rl. In International Conference on Machine Learning, pp. 2836. PMLR, 2021.
Yonathan Efroni, Chi Jin, Akshay Krishnamurthy, and Sobhan Miryoosefi. Provable reinforcement learning with a short-term memory. In International Conference on Machine Learning, pp. 5832- 5850. PMLR, 2022.
Dylan J Foster, Sham M Kakade, Jian Qian, and Alexander Rakhlin. The statistical complexity of interactive decision making. arXiv preprint arXiv:2112.13487, 2021.
Sara A Geer. Empirical Processes in M-estimation, volume 6. Cambridge university press, 2000.
Noah Golowich, Ankur Moitra, and Dhruv Rohatgi. Learning observable pomdps, without com- putationally intractable oracles. arXiv preprint arXiv:2206.03444 2022a.
Noah Golowich, Ankur Moitra, and Dhruv Rohatgi. Planning in observable pomdps in quasipoly- nomial time. arXiv preprint wrXiv:2201.04733 2022b.
Hongyi Guo, Qi Cai, Yufeng Zhang, Zhuoran Yang. and Zhaoran Wang. Provably efficient of- fline reinforcement learning for partially observable markov decision processes. In International Conference on Machine Learning, pp. 8016-8038. PMLR, 2022.
Jiacheng Guo, Zihao Li, Huazheng Wang, Mengdi Wang. Zhuoran Yang, and Xuezhou Zhang. Prov- ably efficient representation learning with tractable planning in low-rank pomdp. arXiv preprint arXiv:2306. 12356, 2023.
Zhaohan Daniel Guo, Shayan Doroudi, and Emma Brunskill. A pac rl algorithm for episodic pomdps. In Artificial Intelligence and Statistics, pp. 510-518. PMLR, 2016.
Milos Hauskrecht and Hamish Fraser. Planning treatment of ischemic heart disease with partially observable markov decision processes. Artificial intelligence in medicine, 18(3):221-244, 2000.
Ahmed Hefny, Carlton Downey, and Geoffrey J Gordon. Supervised learning for dynamical system learning. Advances neural information processing systems, 28, 2015.
Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, John Langford, and Robert E Schapire. Con- textual decision processes with low bellman rank are pac-learnable. International Conference on Machine Learning, pp. 1704-1713. PMLR, 2017.
Nan Jiang, Alex Kulesza, and Satinder Singh. Completing state representations using spectral learn- ing. Advances in Neural Information Processing Systems, 31, 2018.
Chi Jin, Zeyuan Allen Zhu, Sebastien Bubeck, and Michael Jordan. Is q-learning provably effi- cient? Advances in neural information processing systems, 31, 2018.

---

Published as a conference paper at ICLR 2024
Chi Jin, Zhuoran Yang, Zhaoran Wang, and Michael Jordan. Provably efficient reinforcement learning with linear function approximation. In Conference on Learning Theory, pp. 2137-2143. PMLR. 2020.
Chi Jin, Qinghua Liu, and Sobhan Miryoosefi. Bellman eluder dimension: New rich classes of problems, and sample-efficien algorithms. Advances in neural information processing systems, 34:13406-13418 2021a.
Ying Jin, Zhuoran Yang, and Zhaoran Wang. Is pessimism provably efficient for offline rl? In International Conference on Machine Learning, pp. 5084-5096, PMLR, 2021b.
Akshay Krishnamurthy, Alekh Agarwal, and John Langford. Pac reinforcement learning with rich observations. Advances in Neural Information Processing Systems, 29, 2016.
Jeongyeol Kwon, Yonathan Efroni, Constantine Caramanis, and Shie Mannor. RI for latent mdps: Regret guarantees and a lower bound. Advances in Neural Information Processing Systems, 34; 2021.
Sascha Lange, Thomas Gabel, and Martin Riedmiller. Batch reinforcement learning. Reinforcement learning: State-of-the-art, pp. 45-73, 2012.
Jonathan N Lee, Alekh Agarwal, Christoph Dann, and Tong Zhang. Learning in pomdps is sample- efficient with hindsight observability. arXiv preprint arXiv:2301.1333 2023.
Gen Li, Laixi Shi, Yuxin Chen, Yuantao Gu, and Yuejie Chi. Breaking the sample complexity barrier to regret-optimal model-free reinforcement learning. Advances Neural Information Processing Systems, 34:17762-17776. 2021.
Lihong Li, Wei Chu, John Langford, and Robert Schapire. nantextual-bandeee approach to personalized news article recommendation. In Proceedings of the 19th international conference on World wide web, pp. 661-670, 2010.
Michael Littman and Richard S Sutton. Predictive representations of state. Advances in neural information processing systems, 14, 2001.
Qinghua Liu, Alan Chung, Csaba Szepesvári, and Chi Jin. When is partially observable reinforce- ment learning not scary? In Conference on Learning Theory, pp. 5175-5220. PMLR, 2022a.
Qinghua Liu, Praneeth Netrapalli, Csaba Szepesvari, and Chi Jin. Optimistic mle-a generic model-based algorithm for partially observable sequential decision making. arXiv preprint arXiv:2209. 14997, 2022b.
Qinghua Liu, Csaba Szepesvári, and Chi Jin. Sample-efficient reinforcement learning of partially observable markov games. arXiv preprint 2.XXi:22006003155. 2022c.
Miao Lu, Yifei Min, Zhaoran Wang, and Zhuoran Yang. Pessimism in the face of confounders: Prov- ably efficient offline reinforcement learning in partially observable markov decision processes. arXiv preprint 2205. 2205. 13589, 2022.
Pierre Ménard, Omar Darwiche Domingues, Anders Jonsson, Emilie Kaufmann, Edouard Leurent, and Michal Valko. Fast active learning for pure exploration in reinforcement learning. In Inter- national Conference on Machine Learning, pp. 7599-7608. PMLR, 2021.
Elchanan Mossel and Sébastien Roch. Leamning nonsingular phylogenies and hidden markov mod- els. In Proceedings of the sevrrt- annual ACM symposium on Theory of computing, pp- 366-37 166-375, 2005.
Martin Mundhenk, Judy Goldsmith, Christopher Lusena, and Eric Allender. Complexity of finite- horizon markov decision process problems. Journal of the ACM (JACM), 47(4))681-720 2000.
Christos Papadimitriou and John N Tsitsiklis. The complexity of markov decision processes. Mathematics of operations research, 12(3):411440, 1987.
12

---

Published as a conference paper at ICLR 2024
Paria Rashidinejad, Banghua Zhu, Cong Ma, Jiantao Jiao, and Stuart Russell. Bridging offline rein- forcement learning and imitation learning: A tale of pessimism. Advances in Neural Information Processing Systems, 14411170--1716, 2021.
Satinder Singh, Michael James, and Matthew Rudary. Predictive state representations: new theory for modeling dynamical systems. arXiv preprint arXiv: arXiv: 1207,4167, 2012.
Masatoshi Uehara, Xuezhou Zhang, and Wen Sun. Representation learning for online and offline rl in low-rank mdps. arXiv preprint arXiv:21 2021.
Masatoshi Uehara, Ayush Sekhari, Jason D Lee, Nathan Kallus, and Wen Sun. Computationally efficient pac in pomdps with latent determinism and conditional embeddings. arXiv preprint [ttT:2206.1208// 2022a.
Masatoshi Uehara, Ayush Sckhari, Jason D Lee, Nathan Kallus, and Wen Sun. Provably efficient reinforcement learning in partially observable dynamical systems. arXiv preprint arXiv:2206. 12020, 2022b.
Oriol Vinyals, Igor Babuschkin, Wojciech M Czamecki, Michaël Mathieu, Andrew Dudzik, Juny- oung Chung, David H Choi, Richard Powell, Timo Ewalds, Petko Georgiev, et al. Grandmaster level in starcraft ii using multi-agent reinforcement learning. Nature, 575(7782)3350-3555 2019.
Nikos Vlassis, Michael L Littman, and David Barber. On the computational complexity of stochastic controller optimization in pomdps. ACM Transactions on Computation Theory (TOCT), 4(4):1-8, 2012.
Andrew Wagenmaker and Kevin G Jamieson. Instance-dependent near-optimal policy identification in linear mdps via online experiment design. Advances in Neural Information Processing Systems, 35:5968 5981, 2022.
Lingxiao Wang, Qi Cai, Zhuoran Yang, and Zhaoran Wang. Embed to control partially observed sys- tems: Representation learning with provable sample efficiency. arXiv preprint 23847 2022.
Ruosong Wang, Russ R Salakhutdinov, and Lin Yang. Reinforcement learning with general value function approximation: Provably efficient approach via bounded eluder dimension. Advances in Neural Information Processing Systems, 33:6123-6135, 2020.
Yaxiong Wu, Craig Macdonald, and Iadh Ounis. Partially observable reinforcement learning for dialog-based interactive recommendation. In Proceedings of the 15th ACM Conference on Rec- ommender Systems, pp. 241-251, 2021.
Tengyang Xie, Ching-An An Cheng, Nan Jiang, Paul Minciro, and Alekh Agarwal. pessimism for offline reinforcement learning. Advances in neural information processing systems, 34:6683-6694, 2021.
Wei Xiong, Han Zhong, Chengshuai Shi, Cong Shen, Liwei Wang, and Tong Zhang. Nearly mini- max optimal offline reinforcement learning with linear function approximation: Single-agent mdp and markov game. arXiv preprint 1xXiv:2205./55122 2022a.
Yi Xiong, Ningyuan Chen, Xuefeng Gao, and Xiang Zhou. Sublinear regret for learning pomdps. Production and Operations Management, 319))3341-3304, 2022b.
Ming Yin, Yu Bai, and Yu-Xiang Wang. Near-optimal offline reinforcement learning via double variance reduction. Advances in neural information processing systems, 34:7677-7688, 2021.
Andrea Zanette, Alessandro Lazaric, Mykel Kochenderfer, and Emma Brunskill. Learning near op- timal policies with low inherent bellman error. In International Conference on Machine Learning, pp. 10978-10989, PMLR, 2020.
Andrea Zanette, Ching-An Cheng, and Alekh Agarwal. Cautiously optimistic policy optimization and exploration with linear function approximation. In Conference on Learning Theory, pp. 4473- 4525. PMLR, 2021.
13

---

Published as a conference paper at ICLR 2024
Wenhao Zhan, Masatoshi Uchara, Wen Sun, and Jason D Lee. Pac reinforcement learning for predictive state representations. arXiv preprint arXiv:2207.05738, 2022.
Zhi Zhang, Zhuoran Yang, Han Liu, Pratap Tokekar, and Furong Huang. Reinforcement learning under a multi-agent predictive state representation model: Method and theory. In The Tenth International Conference on Learning Representations (ICLR 2022), 2022.
Stephan Zheng, Alexander Trott, Sunil Srinivasa, Nikhil Naik, Melvin Gruesbeck, David C Parkes, and Richard Socher. The ai economist: Improving equality and productivity with ai-driven policies, arXiv preprint arXiv:2004. 13332, 2020.
Han Zhong, Wei Xiong, Sirui Zheng, Liwei Wang, Zhaoran Wang. Zhuoran Yang, and Tong Zhang. posterior sampling framework for interactive decision making. arXiv preprint arX1v:2211.01962, 2022.
Dongruo Zhou, Quanquan Gu, and Csaba Szepesvari. Nearly minimax optimal reinforcement learn- ing for linear mixture markov decision processes. In Conference on Learning Theory, pp. 4532- 4576. PMLR, 2021.
14

---

Published as a conference paper at ICLR 2024
Provably Efficient UCB-type Algorithms For Learning PSRs: Supplementary Materials
A PROPERTIES OF PSRs
In this section, we present few important properties of PSRs, which will be intensively used the algorithm analysis.
First, for any model {ōn. (o,, ax)). we have the following identity
(o,, ) (n))(((T11) = Po(o, (TT)).
(8)
The following proposition is directly adapted from Lemma C.3 in Liu et al. (202bb). Note that tho known to the agent.
Proposition 1 (TV-distance Estimation error) Consider two g-well-conditoned PSRs 0,0 O We have
H
DTv
P
S
( (M(on, an) - M(on, an)) (h(Tn- ((TH),
1=1
TM
Drv
P
po
m(wa) MMon, an) M(on, an) (1) ((H)).
The next proposition characterizes well-conditioned PSRs Zhong et al. 2022 which is obtained by noting that m(q)) is the /-th row of M (0141100011000 1h+1).
Proposition? 2 For well-conditioned (self- ceefssoonsstentt PSR 8 we have
QA
max
max
T(WN)) [M (410000111111)))1))) a++1)=11, S
SER f -el
The following proposition characterizes the log-cardinality of the minimal optimistic net of the rank- r PSRs. We note that this notion closely related to the bracketing number, typical complexity measure in the maximum likelihood estimation (MLE) anslysis (Geer) 2000). The proof follows directly from Theorem C.9 in Liu et al. (002b))
Proposition 3 (Optimistic net for tabular PSRs) Given any there exists a finite parameter space θ. satisfying the following property: for any 0 O can find 0 E O associated with measure Py such that
Vn, h,
(T(TT) (Tn):
Vn, h,
Σ (( (T)) Sc.
Th
Moreover, log |0,1 < 2 log 0A
B GENERAL MLE ANALYSIS
In this section, we present four general propositions that characterize the performance of MLE.
We start with a proposition that states that the log-likelihood of the true model is relatively high compared any model.
Proposition 4 e < Rे.. With probability at least for any 0 E θ2 and any k E [K], the following two inequalities hold:
Kl
Yo E Se,
Σ h (TWD))DD
Σ
log (Pgh)TT)))) (Thh)=+- 3log
Σ
Ig (Th),
5
(T.,))
K
YO E Oe,
Σ
log (TH) 3 log
Σ
log (TH).
8
(TH.))ED
15

---

Published as a conference paper at ICLR 2024
Proof: We start with the first inequality. Suppose the data in Dj is indexed by t. Then,
0 (Th)
E
exp
log
Pp- (Th)
h
(TA.
B (TT) (T))
EE
Ps
E
P (++)
ISk-1
= EE
PP (Tr)
1Sk-1
h Th
(a)
Pg (TH) ( (T))
(1 + E) "E
(1 2)+ K" =)"""
0/ e,
where (a) follows because Σ., IP (Th) - PT S e, and (b) follows because S
KH
By the Chernoff bound and the union bound over θ, and ls E [K]. with probability at least 1 6, have,
log (TT) p- (Tx)
o. 8
Ak E [K],0E θ., Σ Σ
< 3 log
h (TR.))EDY
which yields the first result of this proposition.
To show the second inequality, we follow an argument similar to that for the first inequality. We have
Pr. (TH) P (TH)
P (TH) log
exp
Σ
log
S exp
#- (TH)
(THAT)DD
(a)
(1 8) KH Se.
where (a) follows from the tower rule of the expectation and because Σ., ( S.c
Thus, with probability at least - 6, for any k: E [K] and any 0 E O, the following inequality holds
P (TH) p 8- (TH)
< 3 log K]0.
log
which completes the proof.
The following proposition upper bounds the total variation distance between the conditional distri- butions over the future trajectory conditioned on the empirical history trajectories. This proposition crucial to ensure that the model estimated by PSR-UCB is accurate on those sample trajectories.
Puin and S Consider the following event
Proposition 5
Pain KH Let O O
{0 wh, (To, x) E Dj. TP (Tx) 2 Pmin).
E.-
Vk: [K], E * thin
Σ
Σ
Dfy (wn((w]T)) Pg. (walTT))
h (T.D))ED)
16

---

Published as a conference paper at ICLR 2024
26 Σ Σ
log
pr (TH)
31 log
K
0 5
P (TH)
(TR,=JED)
Then, P (E 2 1 6.
Proof:
We start with a general upper bound on the total variation distance between two conditional distributions. Note that for any 0, 0' U θ. and fixed (Th, T), we have
DTv (IP (Whl (T),, P7 (w), (T))
po (wh, Th (Th) (Whn Th PP (Th)
Σ
Po (T)) P. (Th)
wh
P (wh: Th) (wn T()) (T.) + PT (Wh, Th) (P)(TT) P (Tk))
(T Pg (Th)
PT (T))
PT (Tw)]
Σ
(Pgg(w,, Tn)) -P77 mn))
(Th)
Py (Ta)
2
DTv (T (TH), (gg((
Pg (Th)
By symmetry, we also have
DTv (IP5 (walTT)) 17 ( (wn(Th)) max (17((T). PT (Tn)}
Drv (L (T
We replace 0' by a 0 E θ₂ that is s-optimistic over (recall Definition 2 i.c. DTv (177, P7) S E, and Pg (T.) 2 Pa (Th) holds for any æ and The Then, due the construction of @ min' we have
2c
2
V(Th, T) E D Drv (PT mx), 1Pg (walTn)) S
Prnin
K H
which implies
Σ
Σ
Div (17 (wn PT. (w)lTT))
(Ta))eD
(a)
2D
+ 2D (195 (wnlTT)) (wa(Th)
4
2
Σ
Σ
PP)
KH
(DUD))DD
Here (a) follows because the total variation distance satisfies the triangle inequality and (a + 6)2 s 2a² + 26
Moreover, note that
D P (wh (h), Pg. (w)lTH))
(0)
<
4(2 2/(K)) D (F77 (wa TT), 17.. (WN(TH))
po (WW(TT)
6
1
KH
wh Pu
p... (W)(TH)
(b)
P. (wonlTm (Th)
6
6 log
I
+
Wa
(177))
pst
(w/ (Th)
KH
where (a) is due to Lemma
12 and (b) follows because
-log for any x. >> 0.
17

---

Published as a conference paper at ICLR 2024
Thus, the summation of the total variation distance between conditional distributions conditioned on (Thy т) E Dj can be upper bounded by
Σ
Σ (TA,T)EDY
Div (Pă (wh (T), Pg- (walTh))
18 KH
P (wh |m ) g- (wh) (ih
12
Σ
log
P. (m
(TUT))D
In addition, we have
Σ Σ h (H)
Pp Ww Fh) p (wh (T))
PPz (walTT)) (Th) P (w/ Ta)
E
exp
log
log
E we warpg. (-17a)
mE Wh -o (11Ta)
2
h (T
(T Whn PanPPo (-17a)
Ih II (w..w.nNJD
Pī (w) mm) Fi. mm))
Ih (T Ewn -Pp (1m))
p3 (wn Ta) Fg. (wn (Th)
where the last equality is due to the conditional independence of wh Dý given (Th, T) D Therefore, by the Chernoff bound, with probability 1 ō, we have
log
E (17a))
P
wwh (h) gg (wh, (Th)
Σ
log
P (wi (T)) p (wh (Th)
+ log 1 8
h (TK.T)EDE
Taking the union bound over θ., [K]. and rescaling ó, we have, with probability least 6, Ak E [K], the following inequality holds:
Σ
Σ
DVv (IT (o)) (wnlTh))
(Th, (EA,
18
Σ
log
Po (walTh) (Th)
log
K [
KH
po (W/ Th)
(W)
P. (wh, Tn) I. (wh, Th)
PT (Th)
K
Σ
log
+6 +6
Σ log p (Tr)
13 13log
8
h (E)
Note that, following from Proposition
4 with probability at least
5, we have for any l: E [K].
PT (Th)
K 62
Σ
Σ
log
S 3llog
P (TT)
h (T.D))D
Hence, combining with the optimistic property of 0 and rescaling 8, have that the following inequality holds with probability at least 1 ō:
PT (TH) + 31 log
Σ Σ
Diy (P) Py (T)) <6
Σ
log
PT(TH)
8
(TH.))E))
which yields the final result.
The following proposition is standard in the MLE analysis, and we provide the full analysis here for completeness.
18

---

Published as a conference paper at ICLR 2024
Proposition 6 Fixe K2H2-
Define the following event:
4B O,4 [K], E
(THH))
- 2 (TI,))EDI
Σ
log
0- (TH) P8 (TH)
2log
K|O
T TEDE
We have P(E+) 21 8.
Proof:
First, by the construction of O., for any 0, let 0 be optimistie 0, i.c., E. (Tm) - (T((T)] S e. We translate the distance between 0 and 0" to the distance between
1 and 0 as follows.
(1P (TH), P. (TH))
=1- Σ
Pg (TH (TH)
1
Σ
P (T F... (TH) + (TH) 175 (TH)) PT. (TH)
(0)
Σ
PT (T PT (TH)
P (TH) (TH)
Pg- (TH)
TH
TM
(b)
P (TH P7- (TH)
5
log
1o (TH) - PP
Pa-P3. (-)
P (TH)
< - log
E 6.
Ve,
P. (TH)
TH TH-
where (a) follows because va +
Va Vbl if > 0 and a + b 0, and (6) follows from the
Cauchy's inequality and the fact that x S log
Hence, in order to upper bound
ExED.
(TH)), it suffices to upper bound
(g(T)) Fi. (TH)
ΣTED, log ETH Pp. (-)
To this end, we observe that,
P (TH) PP (TH)
P (TH) g- (TH)
E exp
log
log
E TH
THP (-)
(TH.M)
RED*
F7 (Tn)
IE (a) [THH,)))
P.. (T2)
(TR)
E ()
Pg- (Tw)
where (a) follows because (TH, ω) D form filtration. Then, by the Chernoff bound, we have
P& (TH)
P (TH)
1
P
log
Px (TH)
Σ
log
TH wil (-) V
E
P (TH)
> log
8
(TD
DEDi
P (TH
P (TH)
exp
Σ
log
log
E PH
IV
2
To 6- (TH)
Pot (TH)
(DD
TED
TH
(a)
1
1 (TH)
F7 (TH)
SE
exp
log
log
2
po 8" (TH)
Pg. (TH)
(TH,T)EDK
TEDD
< 6,
19

---

Published as a conference paper at ICLR 2024
where (a) is due to the Markov inequality.
Finally, rescaling to (/( and taking the union bound over θ. and & E [K], we conclude that, with probability least 1 5, v0 E 0, k e [K],
Σ (PP Pg- MED
S KHVE++
Σ
log
P.
O" (TH)
log
K|
K|
2
IP (TH)
6
(E
F (TH)
K|O,1
log
+ 2 log
P (TH)
6
(TR,
where (a) follows because 8 < K"
The following proposition states that the constraint O% does not rule out the true model.
Proposition 7 Fix Pmii I KH{oŌ]AA) Consider the following event:
8
Emin = =(vR E [K], Vh, (T), E DA, PT-(TN) (Tr) 2 Pmin
n kE(K]
{0r E @ @min)-
have P(Emii)
2 21
Proof: For any & E [K]. h E {0, 1......., H -1-1 1} and (Th, T) E Dj, we have ,,...,,-
IP (1P7- (Th) < Prnin)
E
TP (1... (Tn) < Penin [r]
E
po (Th) p (Th) < Punin }
EE
Pmin
8
K
Thus, taking the union bound over is E [K], E that
H - 1} and (Th, T) E Dr, we conclude
PPEEin)) 2 1
C PROOF OF THBOREM 1 (FOR ONLINE PSR-UCB)
In this section, we present the full analysis for the online algorithm PSR-UCB to show Theorem T In particular, the proof of Theorem [ consists of three main steps. Step 1. We prove that the estimated model jk not only accurate over the exploration policies, but also accurate conditioned on emipircal samples Th E DÍ. Step 2. Building up on the first step, we are able to show that V j
valid upper bound the total variation distance between 6 and 6*. Step 3. Based on a newly developed inequality that translates g to the ground-truth prediction feature vi (7%), we show that the summation of V g over the iteration l grows sublinear in time. This finally characterizes the optimality of the output policy, the accuracy of the output model and the sample complexity of PSR- UCB, and finishes the proof.
Before we proceed, we use to denote the event & E, nmmin where these three events are defined in Appendix B Due Propositions to 7 and union bound, immediately have TP(E) 38.
20

---

Published as a conference paper at ICLR 2024
C.I STEP 1: ESTIMATION GUARANTEE
We show that the estimated model is accurate with the past exploration policies and dataset.
Lemma Under event E, the following two inequalities hold:
Σ
Σ
DTv
[ (wwlTa),PP- (wo) (o0 S TB,
h (ThAT)DD
Σ
D
ge (TH): - 73,
where B = =31 log Kl and E S K-H=(O--A)**
Proof: To show the first inequality, note that by the selection of 0%, we have 0 B (defined in Equation ()).
Following from Proposition
we have
Σ
Σ
Div
p
PT- (WnlTT)
(Th.
K|
6
log PT. (TH) - 6 Σ
Σ
log IP (TH) + 31 log
h (TH
K 6,]
6 max
Σ
Σ
log P (TH) - Σ
Σ
log P (TH) + 31 log
8'E min
8
(TH,))EDt
(TH,T)EDt
< 7B,
where (a) follows from 0* Oh... (Proposition 7 To show the second inequality, Proposition 6 implies that
Σ D (TH). (P (TH)
RED
K K
log po- (TH)
Σ TED4
log F (TH) + +lo
8
DE
(4)
K K
max
Σ
log Pu (TH)
Σ
log Pa (TH) 2og
5
min xEDk
$ 7B,
where (a) follows from 0* Oh khin (Proposition
C.2 STEP 2: UCB FOR TOTAL VARIATION DISTANCE
Following from Proposition π the total variation distance between PSRs is controlled by the estimation error. Hence, we first characterize the estimation error of M (on, ax) in the following lemma.
Lemma Under event for any ls E [K] and any policy n, we have
Σ m" (Wa)T (Mi an) (o, an)) (Tn x(TH) TW
# ph ( 1100))
BTN-I
(9)
21

---

Published as a conference paper at ICLR 2024
where
- +
Σ
(T
TA-IEDA
42Q2d hh-1 7p
IARQa
Σ
Dîv (P "on-1 (w) (wn-217h-1))-
-
Th-1ED
Proof: We index future trajectory wh-1 (0h, ap...., OH an) by and history trajectory Th-1 by
For simplicity, denote m² (wx) [M (on; an) M2 (on: (on:)))))) an)
w.", and ((k) (Th-1)
( 1) o (Th-1) as .j. We also denote π (wh-11Th-1) as Tal-
Then, the LHS of Equation ( can be written as
LHS =
Σ
m (wn)² [on, an) M7 as)) i
( (Ta- 1) (TH))
7M
(
we Ty 1T (j)
(Tala sgn(w, 2z)) wi) Jazz P. (i)
Tay sgn(w, xz)) w))
D3 P.. (f)
(8)
Σ Tas sga(w] x;) w/
2
E
|| 110
j-Pp
where (a) follows because (T.) - of (Tr). and (b) follows from the Cauchy's inequality.
Fix Th-1 = jo- We aim to upper bound the term / == > Tijjo sgn(w, xz)) wa lōo.. Then,
we have
1 A
Σ Tijjo
su, x wi
Σ Tajjo
sw, xjo) wi Ij
JEDT_
We first upper bound the first term 12 as follows:
VT2 - V
Tilo sgn( æzo )os
max
2ER :||a]//b===
V
m (wa)T (NI - (aaaaaaMMM(ooo an) x ((w)-
max
ZER 11
m" (wx) "M (onn an)or x Twwn (wwnnnnl 15o)
max
IER -1=====]=||=>
V
mm (w.) (n)) an)er/ ( 1500
max
SEB 4- -illlllaaa wh
(6)
V
T(anlon.Jo) VT
max
Σ
M an)o
ZER - ||0||1
max
7 ZER dn |||
on oh
22

---

Published as a conference paper at ICLR 2024
( 2QAVdx
where (a) follows because la + bj I |a| + (b], (b) follows from Assumption and (c) follows from Proposition 2 and because max # 1=I|=|la=1 ||0|| S vd.
We next upper bound the second term I3 as follows.
2
13
Σ
Σ
m" (wn)
(o an) Mo an)) 3 (Ta- b(wn-(/)))) 11j0)
Th-1ED
A-1
W/
I
Σ m (w))
(on, an) (Tn-1) (Tn-) -M ag) (o, (Ta-))) x xw2n-
Th-1ED h-1
Wh-1
T T(wn-115)) Σ m (wn) (wn) an) (n) (T-1)) -5 (TA-1))
(on
Σ
Σ
Pon (on 1h (o - Po(on1tn-1))3
//1 I(an (o)
EDE
h-1
on
Σ
hh (T - 35 (Ta 1
IED
10
(c)
Σ
Σ
Σ
Po (0 on on ah) - Po- (0%, on 0 an, an) (annoon.
72
TA-IED
Oh 8m1
1-
P 1 Po- (of 1 -. ah-il))
1=1
2
Σ
Σ
a an - iP-- (wh-11th an an-1)] 10
22
T--EED)
1 w
Σ
DTv
IP
Poi (wh-117h-1) GF
No Pp (wn-1lTn-1))
TA-
where (a) follows because la + bl S |a| + (6]. (b) follows from Equation (8) and Assumption and (c) follows from the physical meaning of the prediction feature, i.e. [e((T)]ee Po [Th, af..
By combining the upper bounds for Ig and Iz, we conclude that
I 4XQ3d +
4/21222
Σ
D
(P
"a gt
(wh-1 Th
1). (h-1),)11)))))))(
(wh)11/7-1)) of
72
T
which completes the proof.
The following lemma validates that V is an upper bound on the total variation distance between 0t ga
the estimated model 0 and the true model 0*.
Lemma 3 Under event E, for any have
H-1
DTv [pT (TH), pg- (TH) S aIETTRP de
Σ
Th
(Tm)
(10)
23

---

Published as a conference paper at ICLR 2024
where
4AHQAd
28/
a
722
Proof: We proceed the proof as follows:
Σ h 7
12
4
A
Σ
Σ
Div (Pg (wh-11mh-1), p. (60-1/7m-11)
(Th-
(a) PAHQAd
28,
A
722
where (a) follows from Lemma and the Cauchy's inequality.
1
The proof then follows directly from Proposition Lemma
Note that the reward function R is within [0, 1]. We hence obtain the following corollary.
Corollary (UCB) Under event for any [K] and any reward R, we have
V
S Ve V
gr,
where (TH) min a
Σn
2et (Th)
(0t)
C.3 STEP 3: SUBLINEAR SUMMATION
To prove that E Vat _a is sublinear, i.e., scales ((vK), we first prove the following lemma that relates the estimated feature and the ground-truth feature via the total variation distance between the estimated model and the true model.
Lemma 4 Under the event for any k E [K] and any policy n, we have
E TH
Juk (Th)
(0:))
s i 2| A v 7
Σ
| [TW)
211 DTv (PH), (TH)
V
Pa
h=0
Proof: Recall that
Ü X Σ (T))
TEDİ
We define the ground-truth counterpart of Üj as follows:
U
AI Σ (T.)
TED
Then, following from Lemma we have
|
(Tr)
(0:)-
(Th) ((w))
24

---

Published as a conference paper at ICLR 2024
(TA) - कه 6will
H
H-1
(
- (Ta))/
Σ
1+
11 (w)]((o))
VA
h=0
Furthermore, note that
1 (Tn) - (Tx) - (T)) - (TT)]
(a) S ATr (F "
(w). 2o (wnlTh))
where (a) follows from the physical meaning of the prediction feature.
Following from Lemma
we conclude that
(Tm) (U+)-*
(Th) - 30 (TN) + (1+ 21A 7r
H-1
Σ
(T)))(())))
Vi
h=0
For the first term, taking expectation, we have
H-1
(Tr) 35 (Th)
E
Bra
h=0
H-1
(| (PT) - (Th)) + & (P.) (Tn) - 35 (T.) (Tr) (Tri)l))
h=D Th
(1187 (Tx)) ||, (P5(T)) -P(())))P% (((T))PP7. (T)) + 3 (Tx) - o (m))) π(Ta))
h=D
H-1
Σ QADTV (T))
h=
< ArV (P (TH), (TH))
where (a) follows from (Th)l1 < 2 QA and the physical meaning of (Tx.).
Thus,
H
Σ
J (Th) (0:2)-
E TH--E
1
2 2|A|Q A
Tr
Σ h=
(T)))))
: +
2HQ DTv (Pg- P (TH) v
V
Th Pp
The following lemma can be proved via the l Eluder argument (Chen et al. 2022 Zhong et al., 2022 Since we have a slightly different estimation oracle and guarantee, we provide the full proof here for completeness.
Lemma Under event E, for any h E (O..... H 1}. we have
LAIQAVB VrHK log + dQAK/14).
Σ
Drv Laa' (Th). (T))
K
Here, a. S. b indicates that there is absolute positive constant e such that S c h.
25

---

Published as a conference paper at ICLR 2024
Proof: First, by the first inequality in Proposition I we have
Drv (Pg (Tn), (TN)
m (wn)
(M (on, an) Mi an)) v0 (Tn-1)) x² (TH).
-
TH
suffices to upper bound Σ i (Wa) (o,, an) - M2; (on: an) 30" (Tn-1) т(TH) for any policy For simplicity, Wh-1 = (on, ah.-... OH, aw) by i, and the history trajectory Th-1 by j. We represent of (7h-1) as Zy and r² (wa) [M (on, an) - (on, an)) as w,. We also denote 7(wn (own--17e (h-1) 1) by msly-
similar notations in Lemma 2 We index the future trajectory
Let λo be constant determined later and define the matrix
A do + ΣE [x,=]]]
tck
Then, for any n, we have
Σ m (wh) T [o an) - on) 20 (Ta-1) )(TH)
Σ ue 2zy
E Bwaay
-B
Σ Til r Bingggg(( (w) 2z)mss
Σ T. swg xy)Ws
E j-pat
An
where (a) follows from the Cauchy's inequality. Now fix j ju and consider the following term.
Σ Tilso Tilso sgn(w, 13 ho
Aa
27
Ao
Σ Tiljoo sgn i Ijo Wwi
+
> E
Tiliio sgn(w, x )w? x;j
S2
For the first term I, we have
V4+ = Ao
Tijjo sgn im yw!
max
"h 1:]|m|/2=1
An
max
n 1 Im : (wn-1) T xo|
he -
W
+
V
max
Σ ((wn-11jo) I (w)) ) M on)
-ER=M-1:11a
(a)
VdAa QAvani
7
72
AQQAddM
72
where (a) follows from Assumption
Proposition
2
and the fact that max maRk 1:1|0]|2=1 ||0/|/ s
dd Vd.
26

---

Published as a conference paper at ICLR 2024
For the second term I2, we have
12
Σ
7(wn-1)jo) (wx) (Mi an) (n(on, an)) 30 (n~-1))
t The-1
4//
Σ n(wn-116)) i (w) (T--)) (Tn-1) (T+-)))))
t<c Th 1 Pr
Σ
x niwn-1l66) (wa) (on (n an) (Ta-)))
15 (Th-1) 3 (Ta-1))l
27
Σ
(an (o) 10) (on; (n) M (on, an) (ma-1)])) (Fnn)))))
On t
Σ
E
Σ IP 1
Po 10h a
27
12
n (ap jon Pg.. (of on Th- a1,) Po- onl 0f at)
.A 8=1
1
DTv
[ (wh-1 1h- Po (wn-I1TA-I))
72
B<c Th Th-Irep 6m
4AA02 Q
un( Moppp (TH))
D
IP gh
(TH) Po-
72
[<c
where (a) follows from Assumption 2|A AA-
and (b) follows because exp
O + 1
Thus, we have
Σ m (wan)
(on, an) M (i and & (Ta-1) ((TH)
E
a
1| (Th-
where
Q2d +
412122 72
en( mogen (TH)
(ā )
D
(TH); Fo-
g+
t<k
3 By choosing Ao 1 and recalling Lemma
!
we further have
Eatti133 28| 8
ā.
72
Hence, with the Cauchy's inequality, we have
H
Drv (Pnt) (TH), (TH) < min
ā
Σ
E
(T~- 1>112 ||اا1 2
1=
27

---

Published as a conference paper at ICLR 2024
Taking the summation and applying the elliptical potential lemma (i.e., Lemma 14
we have
DTv
[ (TH).
P7 (TH)
K
H1
S
K
min
á
E Th-
1 (Th- 1) 1 1 .4
h=1
6*
rHK3logh dQAK/T),
which yields the final result.
The next lemma shows that the summation E
V** bh ji grows sublinearly in K.
Lemma 6 Under the event E, with probability at least 1 - 8, we huve
Vo 2 (v++ QAVH
VarHaKK
7
72
where Bo max(1og(1 + K/A), log(1 + dQaK/h)}, and = >/4l² QAI A max/
ma../yF....... vF.q 9. A//11
VdH
Proof: First, since 8t (TH) [0, 1], we have
v S
Va+
(pa* pgg.)
0-,8"
2 Drv
Then, following from Lemma
we have
K
Vx
8+ .c
2|A|Q Tr
H-1
aHQAA
min
a
1+
[TT
DTV (1Nz" (TH), po gu (TH)
.1
TH
K
2|4|Q
Tr
H
(T) ((a))
K
aHQQ
DTv (pg* (TH)), Pat (TH)
min
1+
E
+
VA
h=0 Ta TANPE
k=1
We next analyze the term J. Recall that UØ AT + Emery 25 (TA) (TH) Further note that
{ E Th pgt (TW)ller)))) 11 (++t1,tt))) || k=1
forms a martingale. Applying the Azuma Houff--g inequality (see Lemma probability at least 1
we have, with
5,
4 2K 12/0) +
min a 1+ (1
2/..1 TrB
Σ
(Th
k+1,+11
VA
k=1
2 (og((/2/ + 1 2|A|QAV Tr
(4)
H w e/rKlog( + K/X)
VA
Sa 1+ AIQAVTB
HyrKlug log(1 + K/A),
where (a) follows from Lemma 14
28

---

Published as a conference paper at ICLR 2024
Let 3 : max K/A), + dQAK//)) for simplicity. By Lemma have,
Σ
V bb
AlQAVTB
aH |AIQAVB VrHK30
H VrKB3
VA
VA
?
1 AIQAVTT VA
LAQQAVBH
H HvrK30
VVA
QAVTH AQQAB +
LAIQAV max Vr. QAVH/A)
1+
H HVrK30
??
V
VdH
vA) (+
LAIQAV m Vr. QAVH/A))
IIQAHYT333
AAVBY
VA
QAVCH maxt QAVH/T)
LIQAH
+
VF+ QAVH
?
72
where (a) follows by choosing
>lA AB max /r, QA QAVH/T}
VdH
This completes the proof.
C.4 PROOF OF THEOREM
Theorem (Restatement of Theorem I Suppose Assumption O & ). B Ollog (e)). where O KH
holds. Let Pmiin
Af max (axx(v (VrQAV//1))
X
VaH
and a O QAyTTH + LAIQaVT) Then, with probability at least 8, PSR-UCB outputs model 6 and a policy that satisfy
V V S and Vz, Drv (PT- (TH), Pi.. (TH)) S..
In addition, PSR-UCB terminates with a sample complexity of
δ ((1+ H2H
rdH LAIQ
72
7122
Proof: The proof is under event E, which occurs with probability least 1 36. Following from Corollary 2 if PSR-UCB terminates, we have
YT, Drv (Pg (Tn), ) pg. (TH)) max V Vē_ 2V 0r S e,
R
where the last inequality follows from the termination condition of PSR-UCB. In addition,
V
T V V + V
V2
V
V
(a)
<
2 max V Se
where (a) is due the design of मे and Corollary
29

---

We remark that the key difference between our definition and the definition introduced in Liu et al. (022a) is that we allow h > H m. This indicates that after step > H m, the process (H ht)-stp decodable, which is reasonable in most real-world examples, since the closer Published as a conference paper at ICLR 2024
Finally, recall that Lemma states that
6
Σ Vosi (v++ QAVH
K
72
By the pigeon- hole principle and the termination condition of PSR-UCB,
K K ((
2 rdH2|AA|RQ2
72
7/222
must terminate within K episodes, implying that the sample complexity of PSR-UCB is most
ò ((++
224 72
rdH rdH
7²22
C.5 PROOF OF COROLLARY 1 (FOR POMDPS)
When PSR specialized to an m-step decodable POMDP, we follow the setup described in Section D in Liu et al. (2022b). In POMDPs, there exists a state space S such that the system dynamics proceeds as follows. When = 1, the system is at the state 81 and the agent observes 01 that is determined by $1. At any step h > 1, if the system is state 8/, S and the agent takes action an E A. the system transits to a state 84+1 sampled from a transition distribution Th, (-1s)), and then the agent observes 0++1 sampled from emission distribution On(-Isa+1). Notably, all states are unobservable to the agent.
To define m-step decodability, introduce the m-step emission-action matrix {Ga) where [Gil(,,), = ,(. ( = O S = s, ah (m (h++m-2.FF) = a).
Therefore, for h H h, have G₁ RR--I 10 sss To make unified argument, define m-step decodability as follows.
Definition 3 (m-step decodable POMDPs) A POMDP with m-step emmision-action matrix {Gn) is said to be m-step decodable if there exists an a >> O such that minAC (is)(Gn) 2 a. Here an (A) is the n-th singular value of matrix A.
approach to the end, the clearer our understanding the current state becomes. Moreover, the new definition admits more convenient notations with almost not impact on the logic of the proof.
In this section, a PSR 0 = {p,, (o )x)} is reparametrized by (Tha, Onhne, More importantly, have the following equations and properties.
Yy arg
min Y.
G Yallı
MA on) G Than di (onl)))) YY)
MH aH) - (on.a)) OxA
2n (0 A) mmi
(11)
d <<< | LA
S
Q = |A|m- m--1
m
r QA
30

---

Published as a conference paper at ICLR 2024
Corollary 3 (Restatement of Corollary 1 When PSR is specialized to m-step decodable POMDP, for each k, there exists set of matrices (h C [R such that if we replace u/* (Tr)
by (G)) ((T)) in Equation then PSR-UCB terminates with sample complexity of poly(r, 1/y, QA, H. B).
Proof: We use the notations defined in this section and particularly those in Equation
11
The proof differentiates from the proof of Theorem in two lemmas.
First, in Lemma 2 the meaning of z changes from e (T)-1) to (Q+)) (TL-)), and w changes from m
an) - M (on, ax)) to m (w), (w,)(M) (on. an) - M, (on, aA))
Then, the proof of the upper bound of the term 12 will become QA TTV (without dependency d), simply because the dimension of bonus term 8 is now S instead of d.
Second, in Lemma 4. we need to show that the new bonus term ( can still be bounded in a similar way, namely, (I(G*+))) ( á (1I (I( (5 (Tw)l(c)))))
We note that the matrices 0% and Uz now have new definitions:
Ü =1 (1)
Σ
(T))
(G
U = XI + (G
Σ
(m)) (Tx)
TRED
To show such an upper bound, we next provide the key steps. First, apply Lemma
13 and obtain
(G h+ 1) 1t (Th 1(0e)"
G ( (G) (T))
vrv E EDD EDR
+ t
+
1+
For any The we can derive
(G (GiH))))))))))))))
((Gh+)) + 8htt))
(G)))
((Ght))* + +1it1)) (5 (Th) 0 (m))) || + (G + 4ht (Gi+))) (mW)]1
(Gi+)) + +Y || || | (T)) = "((na)]| + -
+ (1 G ( (T))1
< 1 (m))
- - (1 + (l t - (Gi (
where the last inequality is due to the m-step decodability.
The first term above has been bounded in Lemma 4 Cauchy's 's inequality as follows:
Then, we bound the second term by the
|| ((GH+))* + rhtt) (Gi-1- G%-,) (Gi To (Th)
31

---

Published as a conference paper at ICLR 2024
rQ2//++ +
Brannt
( ((
+ (h+1)
(Gi+1 Gh+t) (G (+0
ick
< ( V (TWl(()))))
vrQnd/++ IAQ2
where Ug denotes the expected matrix of Uh, and has the following form U = = = XI + (G)+1) Σ
Enow
(+ (Ta)e² (Ta) (G++1)*
In particular, the last inequality above follows from two facts. First, use Lemma 39 in Zanette ct al. (2021) to transform the expected matrix to the empirical matrix U/A Second, we use Lemma
Lemma 12 to upper-bound E || ((C + th+1) (Gi -/+1)) (G (T)) by
the Hellinger squared distance between the estimated model 6h and the true model 0*, Then, due to the estimation guarantee (Lemma
[]] we have
ΣE E bock
((()++)* + h (Gh+1 - -Gh+1) (
(Gi++ Ghit1) (G/+1) (T)))
E
Tap
t<
2
E
Σ
E
(o 81+1 s) -o. (ola, sh41 = =)) Po. (s(TH)
72 t<k
Th Por
OERA-IaEQ SE8
E
(ReNPor PTV
po (.182+1 = =), (), Pg-
poft11 (.188+1 =s) 2
Q7
"oft1 1 (.18h41 = =s) "odt1 (-1 (-18h41 = =))
Ppg (- Div
t<k
|AQQ2
A
D
P Gk:
py. ) = ) 0(3).
?
EEDI
Combining above key steps, we obtain that
E E [G A0 (Tw) lo3)-
I polly(r, |A], QA 1/7)
ET M (G (TWlloz)) DTv [D (TH), Po (TH)))
The rest of the proof follow the same argument as that for Theorem
D PROOF OF THEOREM 2 (FOR OFFLINE PSR-LCB)
In this section, we present the full analysis for the offline algorithm PSR-LCB to show Theorem 2 In particular, the proof of Theorem 2 consists of three main steps. Step 1: We provide the offline estimation guarantee for the estimated model ô. Step 2: Building up on the first step, we are able to show that V is valid upper bound of the total variation distance between Ô and 0*. Hence, V 9.6 valid LCB for the true value VØ-,u- Step 3: We translate V/8 the ground- 6.6
truth prediction feature J (Tr). Finally, we show that the Va scales in the order of 1/v K, which characterizes the performance of PSR-LCB and completes the proof.
32

---

Published as a conference paper at ICLR 2024
We first introduce the following definitions for good events.
Recall that Omi = {o vh, Th E Dr, F7 (Ta) 2 Pmiin where Pmin &
Good Events.
Let e S KH Analogous to the proof for online learning, we introduce three events
A))
defined as follows.
EW
40 E Owin:
ΣΣ TD.
Div (WalTT)) Pp. (wal (Th)TT))
p (TH)
3K |0
Σ
log
+ 31 log
P (TH)
8
h THEDA
Pg (TH) P (TH)
3K|
40 E O,K, Da (TH), (TH)) v
Σ
log
2 log
8
TH ED
8 min
vh, Th E Dns Py (Th) 2 Pmin
Following the proof steps similar those in Appendix B occurs with probability at least
we conclude that E" ED SEn SenEgnEE NE
D.I STEP 1: ESTIMATION GUARANTEE
Lemma (MLE guarantee) Under event E", the estimated model ê by PSR-LCB satisfies
Σ
DIv
P. (w) (WnlTh),PPt (a (wnlTa)) 7B,
TAED
D
(p=" (TH), PT- (TH)) $ 273/K,
where 3 = -311 log e
Proof:
The proof follows the steps similar to those in Lemma ration policies Ma (x* uogu)) by the behavior policy.
[ except that we replace the explo-
D.2 STEP 2: UCB FOR TOTAL VARIATION DISTANCE AND LCB FOR VALUE FUNCTION
The following lemma provides an explicit upper bound on the total variation distance between the estimated model and the true model.
Lemma Under event & for any policy n. we have
3
Drv (Pg. (TH).PT(TAT)
E
[| (Th 11
(12)
H2222
Th Impy.
hm
where Ap-1 AoI
K E [15 (Th-T)
(Tn-1) and Au
74 4Q2 d'
- H
Proof: Similarly the analysis in that for Lemma
9 we index wh-1 (on, ah. ... on an) by
i, and Th-1 by j. In addition, we denote m(w))
Mr an) - MM, an)) by w... denote
J (Th-1) by Zj. and denote π (wm--1/7-- 1) by Tal-- Then, we have
m( T (M,(on, an) Mr an)) $3 (Tn- =((=)))(TT)))))))))))))))
Io xy 1a ())
33

---

Published as a conference paper at ICLR 2024
Σ
E mayy sgn(wc) xz)) "ws"" "xy ipz. (i)
T
Σ
Σ Tals sgn(w} 2z)) -us) Iy . pg. (i)
S E jrPog
| -لllاam
Σ may sgn(wc xz)) wi
where An-1 Xn +
K H
[45 (Th 1)86 (Th-1) ]]
and Ao will be determined later. We fix
Th-1 = jo and aim to analyze the coefficient of | As We have
Σ Ti jo sgn 2z0)=
Wi
An-1
2
T 2
do Σ Adjio
sgn(w, xzo) Wi
sgn
K
E
Σ Tojio
sgn
xza)
wi
It
H
juuggh
For the first term J,, we have
Vh
2o
Σ Tijo sgn (ur, 3zo ) x
max
#ER th- -
Ao
max
m" (wh)
T (Mn( and Mī (on, an)) x (wnn-1ljo)
#ER -=||=||==1 wh
max
m(ww) MMron, an )x π (wn-11)0)
mEB dh- A|larllaa w
Σ (mn(w))) M2 an)ar ((wn-1)jo)
max
zE =||+|||== Wa
(a) v d
Σ
Mi an)o T(anlon.
+
max
hERth-1 1u|====:||=||>>> |||=]||==
(0) 2QAVdAa dAg
where (a) follows from Assumption | and (b) follows from Proposition! 2 For the second term Iz. we have
1 <
E
Σ
[ T (o)) - Mo an)) 35 (Tn-1) क (wwn- 1150)
H
The
Wh
Σ T an - Mr an) (Tc-1) ((nn-115)) whh-1
E
H
PP- gm
27
Σ i((w)) (M(on.an) (5(Ta-1) 35 (Ta- a(wa- -150)
(a)
K
(
- - -Po(onla-- a(anlon.Jo)
34

---

Published as a conference paper at ICLR 2024
1 [ 20 (Pn-1))1)
27
K
E
Σ
P. (o on 04, ai) Po. (of, o an, an) x (ann Jo)
Hy
uppd Th-Inupg gm
aa
27
12
1 Po- 1o o
K
E
Σ
an .Po (wh-11 Th a--1) 1)1
H7 Th-inpp
Bh-1 Q
QoAP W
(6)
K
D IPg" (wh- (wn-117h-1 1In pg 1/
E
H2272
1h-1
2A IAS
64K
Dj
(Pg' (TH), 1g (g-()n)
H2222
3
H22_2
where (a) follows from Assumption and Equation (8, (b) follows because (an--) 2 1. for all
a E 2 (c) follows from Lemma 12 and (d) follows from Lemma 7
Thus, by choosing Àu 3a 1
we have
Σ m(wn) (M,(on) an) - My(on. an) w0 (Ta-1) x(TH)
(T-
E 1--1-
H222
Following from Proposition
! we conclude that
B H222
DTv P
(TH),
P7
(TH)
S min
(T
(Th-2) |لااا 2
h=l
Following from Proposition error by the bonus function.
[
we provide the following lemma that upper bounds the estimation
Lemma's 9 Under event E", for any policy we have
m" (wn) (Mn(on, - an)) d(mn-1) m(TH) I Bmn-inpg @n- [T(Tn-n (T) (((m))))))))
where
Úh-1= År+
Σ
((T-1))((Tm-17
Th
iQ3 d
Σ
Dry (PY' (wn-al/T- 1), po" (wn-alTh-1)
ôp-1= 7
122
Th- Tb-1EDn-1
Proof: We index Wh-1 (on, ap.......................... ,H, by i, and Th-1 by In addition, we de- note m² (Th) (M,(on:))) MM an) by wt (nhi(Th_1) )n- )ー
by 2j, and
35

---

Published as a conference paper at ICLR 2024
((Wn-1(Th-1) by mal-- Then, we have
Σ TM
m (wn) (Ma(on, an) M2 (Oh: an) v (Tn=)) ((TH)
Σ Σ l Ty (
Σ E Tayy sgn(w xz)) ws) Jzy (P)())
Σ
Σ Fsla sgn(w, 2z)) mi Iy (T(
Σ F sgn(w, xz)) wi
< E jr
Cr
We fix an index j = Jo, and aim to analyze E #d/30 sgn(w, xzn) ws lloz= which can be written
as
Σ Tijio
(w x55) wi
Σ do sgn(w, 2z0) wi
Σ d sgn(w) [5)) wz Dij
+ JEDa-1
For the first term I, we have
===Vā
Σ Tijoo sgn w 2z 20
max * L:]la| 2=1
m" (wn) T (Ma an) Mī Mi(ooooo)) x (wn-
max
zh 1--h--1 / 2
wa
Σ m (wn)
Mn an) ((wn-1ljo)
max
2 Wh-1
+ VÁ max
Σ mm (wn)
T (os. and TT (wn-1ljo)
a: |||==
Wh-1
[ (1)))a T(an maanlon.Jo) + V ||rl1
Σ
max
s: ==
an
2Q 2Q aa
72
For the second term Iz, we have
I < Σ
Σ
m" (wh)T
(Ma(on, and - M2 aw) (( π awnn- 120)
Th-1EDp-1
Σ
Σ Wh-1
m" (wn)
(Malon.an)o(tn-) - (o, an) (n))"((n-))) I(wn-1l50)
TE-IEDp-1
36

---

Published as a conference paper at ICLR 2024
Σ m (wr)
MM (on, an) -55 m(wn-1l5))
15
(a)
Σ
(Po(onITo-))((mm) - Io( Tonnoon:
Th-1EDn-1
- 13 - (m^-))]
Qn1
Σ
Σ
Polo on 0, a/,) Fo. (o onl 1 ahs an. af) (aanlonn
72
Th T--1ED,
oh,an
Qn-1l
[ 11. Po- (0% ah- 1
Emll
Σ
1. Po- (wh-11TA-11 ap 8h 1E wh-1
Σ
Σ
2
Th IED, 1
Dîv [pg' (wn-almn-)) (wn-alTn-1)).
322
LED,
where (a) follows from Assumption and (b) follows from the condition (aa))) t for any ag E 2h-1
Thus, we conclude that
Σ m (wn)"
(M (Ma(on, an) (o an)) ((Tn-1) x(TH)
â=-1 3 (((()_)))) I 1) 2
< Er, 1N2
where
4iQ2d (bn-1)²=
Σ
Div 1 (wh-1l7n-1)) Pp" (wa-alTh-1))
7t + 222
Th Th-1EDn
Corollary 4 Under event E", for any reward R, we have,
Va OR
S Vil
where (TH) min { a ag Σn
(
( } and â
PAHQU 78
2,,3
Proof: By the definition of the total variation distance, we have
v 0.
V8=
< Drv
[o 1P7.)
m (WiiT (Ma(on,an) M (on, (n)) b( ((TH)
h=
(6)
H
< min
E E
way
ân-1 [(
37

---

Published as a conference paper at ICLR 2024
GV
H
3n (TT)
min
(a))
A=1
A==
where (a) follows from Proposition (c) follows from the Cauchy's inequality.
(b) follows from Lemma 9 and because R(TH) E [0, 1], and
By Lemmal 7 we further have
Σ
Div (Pg (wh- 1 Th 1), Pp (wh- 1T6-
h Th -1 EDa-1
4XHQ3d
73
71
(272
which concludes the proof.
D.3
STEP 3: RELATIONSHIP BETWEEN EMPIRICAL BONUS AND GROUND-TRUTI BONUS
Lemma 10 (Offline empirical bonus and true bonus term) Under event .or any n, we have
H-1 M
E TH-
h=0
1
2
3
Σ h=0
H-1
E
2HQA V
Drv (F (F
vV
Proof: Recall that
Ü = dr+ Σ H(T))(TH)
TAED,
We define the ground-truth counterpart of Ü, as follows:
Un = 7 + Σ T(T))ō(T))
TAD
Then, by Lemma 13
we have
(Tm)) 1'a.>
( (
ViVEn.eD
(((m)) ( 3TH
(5(T)) - ( (T)))1 Σ 1+
1 (w)
Ví
Furthermore, note that
(TTh) - 4 (Tx) 2
S ((a(TT))) (T)
Drv (Dg' (walTn)) pa (wal (wnlTh))
where the second inequality follows because ('(a))) 2 e for all E 2 By Lemma 7 we conclude that
(TT B
38

---

Published as a conference paper at ICLR 2024
(1t - +
Σ (m()) #
Σ 1 (T
For the first term, following argument similar to those in Lemma 4 and taking the expectation, we have
Σ ETA 2p
-
[(TN)
(11((T)) -P)()) + (T))) - (p (TN) (TT)),)
Σ
(1(((T)) |, (T.)).) (T() + - {(w)]| π T(T))
h
Th
< 2 Σ 1 Drv [g (T.),
h
S ATT
PT (TTH)
which completes the proof.
D.4 PROOF OF THEOREM
Now, we are ready prove Theorem
Theorem (Restatement of Theorem Suppose
Assumption
7 holds.
Lel Ollog 1+=l),
min 2M (an). Pmin
0 KHOooAAaE ). E
0 Puiin
B
y0 3 m
and á O Qavan
VX
yB
Then, with probability at
QQav
least 1 - 8, the output R of Algorithm 2 satisfies that
((v++ QAVH
CT CF..QCAAF
QAH2
rd
VT, V8_
V SO
K
Proof: Under event 8 the performance difference can be upper bounded as follows.
V .r
Vøt.r
(V.
0,r - 0.6 +
+V=
0.r
-
Vx
(V
0.r
V
+
1
V
0.b
(a)
DTv
P.g PT- + 0,b
Drv
+
V
8. .
where (a) follows from the design that results in J, S and from Corollary which implies T2 < 0, and (b) follows from Corollary
By Lemma we have
Drv
[ 8:
S
Σ
E
||
(Th-)))a)
Hi?a2
h=1
(a)
B
C -
,00| V
|| (T--1)l
Hi?22
39

---

Published as a conference paper at ICLR 2024
B Hi223
rH C
r
SC
S
H
V Ki292 70,
H
x
K
where (a) follows from the definition of the coverage coefficient.
By Lemma
10
have
V
2 7rB
H-1
+ 2áHQAD DrV
<min
á
1
Σ
E
[| (wwll(cc))
,1
Th Aar
Ví
h=0
2 2/77D TrB
#-1
AaH Drv (TH)
S Smin
Σ
E
[1 (Tw]loww))
Th Ta
Ví
h=0
1+ 2 7r
H-1 Σ
24HQA
<min aC
E
(T))l
+
Drv (P%- (TH), Py (TH))
vVA
h=0 ThOUPN
1s
Recall that Un 'I + EnED, Therefore, by the Hoefddnggss inequality (Lemma 11]. with probability at least 1 - 5, have
J (T. (Tr) and the distribution of Th E Dn follows Pg..
2v
8r
H-1
KI3 S 2K log(2/0) + min ac
1
Σ h=O TAEDs
Σ
vVA
1+ Vrß
K H
K Klog(2/0) + a0
LV
H
Vr OVA
1+
V /rHK 12g()/
where (a) follows from the Cauchy's inequality. Combining the above results, we have
V
Voh.r
Vr rB
TH 1og(2/0)
aHQA
O H
r
S &CT $0000
K
Vi
K
VB
QAVdH
VX+
(1+
max VF. VHQA/T
TH 3 log(2/8)
S O C
HN
17
IV
K
QAVAH
Vā+1)
(1+ max Vr, VHQA/T3
H V
3
rH log(2/0) (2/0)
C ,0
TBB
iVā
K
Vr, VHQA 7
QaH'Va
r (ogg0(2//8)
max
u72
K
where (a) follows because
TV B Vr, QAVH/T}
2QAVdH
40

---

Published as a conference paper at ICLR 2024
This completes the proof.
E AUXILIARY LEMMAS
Lemma 11 (Azuma-Hoeffding inequality) Consider a domain X. and filtration Fi C. C.... C FR C... on the domain X. Suppose that (X ......} C [[B, B) is adapted the filtration (Fi) Le., Xk is Fe-meuuuralle. Then
k-1
Σ X [[X t=1
2 22 log(2/5) S 8.
(13)
The following lemma characterize the relationship between the total variation distance and the Hellinger-squaree distance. Note that the result for probability measures has been proved Lemma H.1 in Zhong et al. (2022). Since consider more general bounded measures, provide the full proof for completeness.
Lemma 12 Given two bounded measures P and Q defined the set X. Let [P| Enex P(x)
and Q ExEx ((x). We have
(P, (
In addition, Pr|x Qrix are two conditional distributions over a random variable Y, and Px.y P, Qx QrixQ Q are the joint distributions when X follows the distributions P and Q respectively, we have
E
[D Qylx)] s (((PXxXxxxx)))
Proof: We first prove the first inequality. By the definition of total variation distance, we have
(P,Q) =
Σ P((x) -
Σ (VP(z) - v((x)) (VP(a)+ (Q(O))
Σ (Q(=))
( (P(x) + Q((1)
(((()))))()))
< 4(1P)+ 1Q))
where (a) follows from the Cauchy's inequality and because (a + b)2 < 2a + 262.
For the second inequality, we have,
X~P [P Qrix)]
= pP(x) Σ P - Verix(e)))
Σ (v (x.r(x,)) (x. (rrxx())(() - VerIN(W)Pa))
2 Σ (xxr(2, /xxrlawu) +2 +2 Qri((v) ( [VQ(=)-VPi)
(D{(PxY, +2( [Q) 2 Σ
vp(z)((z))
(a)
(r((Pxy.xxxx)+22
Σ Σ /P8 (BXxW)(()0xxx000)))
41

---

Published as a conference paper at ICLR 2024
= 8D Qx,y),
where (a) follows from the Cauchy's inequality that applies on E, (ylx(())xyn
Lemma 13 Consider two vector sequences fzihez and (ystiez and index subset J I. Suppose A = AI+ 2 and = =EEo UzUli. and rank (Ha)lez) = rank (fys ((y))ez)) = Then
V |,, 3ill + 1+ 2V EJes |a
15112
ViEI, I, | S
lIve ll
VT
Proof: We first write
|| - ||
|| - 11vbella=1 11villa- - -111:11 B-1
|| 11uil1a-1
1luil
|1 11
||aa||e- + 1Balla-1 llyella= + 11 18:
For the first term 4, we repeatedly apply the Cauchy's 's inequality, and have
|| > 112 - 11 112--
2.7 A (2: - yi) + y, Axx yi)
s ||اا 11|23 - -villa 1 + Ilvilla 2110 villA
(1|=:|la-1 + (lualla-)) || -vll2-
For the second term Iz, repeatedly apply the Cauchy's inequality, and have
|| 1/2 A- 11:11
= v.- A ' A) y/
Σ
vs AA''Iz(15 - xz5 B y + vt A- - 2z)u, B B-
| ( [1yj -,,18=>11 |1 (1velll-11 || 1--
JEJ
|luj xzy all |1 bowssll=i/lie
JEC
Σ ||,r - 13 113 FEJT
Σ
| 112-1
VX
xEJ
Vx A 1lvill
Σ JEJ
1,
80113 -1y11,
Σ
1
jE.J
2 2 || lA-ll/ullll-- VX
||ر|
Myhh 112.
JEJ
Therefore, the lemma follows from the fact that I|æallA- S l1ualla- + |1 B-1.
The following lemma and its variants has been developed in Dani et al. (2008); Abbasi- Yadkori
ct al. (201T); Carpentier et al.) (2020). We slightly generalize the padding term from to an arbitrary positive number B and provide the full proof for completeness.
42

---

Published as a conference paper at ICLR 2024
Lemma 14 (Elliptical potential lemma) For any sequence vectors X {a1,..., Iny.... C R. let Ug = XI Eick T where X is positive constant, and B > 0 is real number. If the rank of X is at most r, then, we have
min [1e 1, B} S (1 B))
K
min [Ioxllee 1, VB S v(1+ B)rk 1ogr1
(1 K/A).
Note that the second inequality an immediate result from the first inequality by the Cauchy's 's inequality. Hence, it suffices to prove the first inequality. To this end, we have
Proof:
min 2
BJ
(1+ B)
log (1
| (/
= =1++ B)
log (1 +1+ trace ((U++1- - 11141-Ux)(0 UN) U)
k=1
K
=(1 B)
Σ
log (1
trace (UE (Uk+1 Uk) Uk 1/2)
k=1
S B)
Σ
log det [Ia U- (Uk+1 U2) 1/2)
K
det (Uk+1)
= (1 + B)
log
det(Ux)
(1 + B) log det(Uk+1)
det(Ui)
K
(1 + log det 1
k=1
( S (1 + log(1 + K/A),
where (a) follows because
S (1 log(1 x) if 0
B, and (b) follows because
rank(X)) S r.
43


Reviews:
[{"review_id": "c5vpHhkq1s", "reviewer": "Reviewer_YPuy", "summary": "This paper presents an UCB-based algorithm to learn parameters in POMDP using predictive state representations (PSRs) to simplify the number of potential historical and future trajectories. PSRs assume a smaller subset of possible trajectories with a low-rank basis. Therefore, PSRs can efficiently reduce the dimensionality required to represent all the historical or future trajectories.\n\nUsing PSRs, the authors show how to use maximum likelihood method to find the most likely PSR parameters, and then use an additional bonus function to quantify the uncertainty of the learned PSR parameters. The algorithm leads to no-regret guarantees with a sample complexity bound induced from the regret analysis. The UCB construction is quite different from other UCB algorithms that I have seen in the literature of multi-armed bandits.", "questions": "**Questions**: please answer my questions mentioned in the weaknesses section. I am happy to raise the score if some of my questions are clarified.\n\n\n**Typos or confusions**:\n- Equation 1: what does $\\pi(\\omega_h)$ mean? Your definition of $\\pi$ measures the probability of choosing a sequence of actions conditioned on a sequence of observations. $\\pi(\\omega_h)$ only specifies the future trajectory after time $h$ but does not specify the history prior to $h$. Please clarify.\n- After finishing reading the paper, I still don’t see what does the “physical meaning” mean in Section 4.1.\n- Line 4 in Algorithm 1: notations of $\\omega_{h-1}^{k,h}$ and $\\tau_{h-1}^{k,h}$ are confusing. You have a different use of superscript on $\\omega$ and $\\tau$ defined previously. The following discussion in Page 6 is helpful though.\n- Equation 3: $\\tau_h \\in D_h^k$, abused notation which collides with the definition of the bonus function with a different $\\tau_h$ involved. Similarly, there is a typo in the following sentence: $\\{ \\hat{\\phi}^k(\\tau_h) \\}_{\\tau \\in D^k_h}$ where $tau_h$ should be $\\tau$ unless I misunderstood.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "**Solid theoretical analysis**: I checked the main paper but didn’t check the appendix. The main paper and proof sketches are theoretically sound to me.\n\n**UCB algorithm for learning parameters in POMDP (reduction in learning space)**: careful and detailed analysis of a complex problem of learning parameters of POMDP from historical trajectories. Previous online learning in sequential problems usually learns the tabular-from parameters directly with easily constructed UCBs. This paper analyzes UCBs in predictive state representations (PSRs), which can reduce the number of potential historical trajectories in POMDP to maintain UCBs.\n\n**New analysis**: the analysis is indeed new to me. I appreciate the novelty in theoretical contributions.", "weaknesses": "**Computation**: the computation cost is still expensive with a dependency on how good the PSR is (the size of $Q_h$). Could you elaborate more on the overall computation complexity and point out what is the bottleneck?\n\n**Reward-free claim**: the theoretical results and the analysis are reward-free by design. The major learning problem is on learning the transition probability using PSR. The reward function is bounded within [0,1] so I am not surprised that the algorithm and results are reward-free.\n\n**Overly complicated notations**: I understand this is a difficult problem with many notations needed for the analysis. But the presentation is very difficult to follow and its clarity can be significantly improved. I believe there is a better way to present and help readers understand the analysis better.\n\n**Sample complexity**: I thought a major advantage of PSRs is to discard the dependency on the number of all possible history with size $O(|A|^H |O|^H)$. However, this term still appears in your sample complexity with a dependency on $O(\\epsilon^{-2})$. The dependency on the action and observation space is hidden in $\\epsilon = p_{min}$ where $p_{min} = O(\\delta / {|O|^H |A|^H})$. Could you explain how your algorithm improves and compares with the online learning algorithms without using PSRs?\n\n**Offline learning**: the offline learning part directly follows the same proof as the online learning part. Could you elaborate if there is any new insights brought by the offline learning analysis?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yV2bkvmFMn", "reviewer": "Reviewer_4GGq", "summary": "This paper studies learning predictive state representations (PSRs) using UCB-type approaches. The proposed PSR-UCB algorithm is based on an upper bound for the total variance distance between estimated and true models. The algorithm requires supervised learning oracles based on which it is computationally efficient and has a last iterate guarantee. An offline variant PSR-LCB is also proposed.", "questions": "I've detailed my primary concerns and suggestions above. However, I have a couple of specific questions regarding certain aspects:\n\n1. On Page 6, just below equation (2), could the authors elaborate on why learning \\$\\bar{\\psi}^*(\\tau_h^{k, h+1}) $ might be harmful when $\\mathbb{P}_{\\theta^*}^{\\pi^{k-1}}(\\tau_h^{k, h+1}) $ is very small?\n\n2. In Algorithm 1, the UCB-approach seems to serve as a pure exploration mechanism aimed at enhancing the estimation of $\\theta $. Yet, in Algorithm 2, the upper confidence bound for $ D_{TV} $ is directly employed to construct a lower confidence bound for $V_{\\theta^*, R}^\\pi $, which the output policy then optimizes. Could you shed light on what drives this distinction between the designs of the two algorithms? Or, is there an underlying equivalence between them?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This paper is in general well-written and smooth to follow, with the exception of some proofs in the supplementary material. \n- The originality of the paper is high as it proposes a novel UCB-based algorithm for PSRs. Related works are covered in detail.\n- The theoretical proofs seem to be rigorous.\n- The proposed UCB-based algorithm enjoys strong theoretical guarantees and might be of interest to researchers.", "weaknesses": "My main concerns with this paper are as follows.\n- First, the theoretical results crucially depend on the MLE oracle, which is called $H$ times during the implementation. However, the paper falls short in discussing the difficulty of solving the MLE estimation. At least, it would be better if, for some special class of $\\theta$'s, the authors could show how the estimation could be implemented and what's computationally costly. \n- This paper is purely theoretical and lacks experiments even on synthetic data. I would appreciate simulation results validating the sample complexity presented in the theory.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7t4VOtjIJW", "reviewer": "Reviewer_dtbr", "summary": "This paper adds the UCB (or LCB) term to online or offline POMDP. The modification compared with (Liu et al., 2022) is that 1) a hard threshold $p_\\min$ adding to the selection of $\\Theta$ and 2) the UCB term used in exploration.", "questions": "Please see weakness", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- This paper is well-written and easy to follow\n- The online algorithm can provide a `last-iteration' bound, and the bonus term can be efficiently calculated.", "weaknesses": "- Compare with (Liu et al., 2023), why during the exploration, the action trajectory needs to be sampled from $(\\mathcal A \\times \\mathcal Q_h^{\\mathcal A}) \\cup Q_{h-1}^{\\mathcal A}$, why we need the last term $Q_{h-1}^{\\mathcal A}$? In addition, why we do not need to guarantee $\\mathcal B_{k+1} \\subseteq \\mathcal B_{k}$ here?\n- The sample complexity of the online algorithm is confusing, it looks more like a 'reward-free' paradigm, in which we are concerned about the *sample complexity* instead of the *regret*. When we talk about the 'last iteration regret', we should assume this regret bound works for all $k$ and the algorithm will not change regarding different $k$, which is not the case in this paper: we need to predefine a $K$ and get the sample complexity bound. Thus I think this result is kind of strange.\n- I wonder what will the algorithm look like if we do not only use $V_{theta, b}$ but use the optimistic value function $V_{theta, R} + V_{theta, b}$ in the online algorithm, should the result be the same with (Liu et al., 2023)?\n- Why the constrain related to $p_\\min$ is easy to hold? Is there any theoretical justification? In other words, what will happen if that constrain does not hold for any $\\theta$?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "B2ZDstmsCW", "reviewer": "Reviewer_eKGX", "summary": "This work establishes the first UCB-type algorithms for PSR problems. The proposed algorithm is computationally tractable compared with previous methods and obtains sample complexity which matches the previous best upper bound w.r.t. the rank of PSR and accuracy level. The counterpart of the algorithm and analysis in the offline case has also been established.", "questions": "Please see the weakness part above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Motivation: Given previous methods studying PSR problems that are not computationally tractable or sample from posteriors maintained over a large model set, studying a computationally tractable method for PSR is well-motivated.\n2. Novelty: This work proposes the first UCB bonus to enable efficient exploration of PSR problems, the core design of which is a new MLE guarantee. The newly proposed MLE guarantee and UCB bonus might be of independent interest.", "weaknesses": "1. The proof of Theorem 1 relies on some sort of MLE guarantee, a confidence bound constructed by the empirical features and relating the empirical bonus to the ground-truth bonus. Such an approach shares similar spirits as previous works studying MDP with low-rank structures (say [1]). I would like to suggest the authors give more comparisons on this.\n2. The result in Theorem 1 seems to critically depend on the newly proposed MLE guarantee in this work. It might be better if discussions about how this MLE guarantee is established and to what extent it differs from previous ones are given.\n3. By “computationally efficient”, it in most cases means that the algorithm enjoys polynomial computation efficiency. In some parts of this work, the authors claim that the proposed algorithms are computationally efficient. In my opinion, the proposed algorithm is indeed computationally tractable or oracle-efficient but might not be “computation efficient”, since the polynomial efficiency of the proposed algorithm is not guaranteed (correct me if I am wrong).\n\n[1] Uehara et al. Representation Learning for Online and Offline RL in Low-rank MDPs. ICLR, 2022.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)