A Free Lunch from the Noise: Provable and Practical Exploration for Representation Learning

1 INTRODUCTION . Reinforcement learning ( RL ) dedicates to solve the sequential decision making problem , where an agent is interacting with an unknown environment to find the best policy that maximizes the expected cumulative rewards ( Sutton & Barto , 2018 ) . It is known that the tabular algorithms direct controlling over the original state and action achieve the minimax-optimal regret depending on the cardinality of the state and action space ( Jaksch et al. , 2010 ; Osband & Van Roy , 2016 ; Azar et al. , 2017 ; Jin et al. , 2018 ) . However , these algorithms become intractable for the real-world problems with an enormous number of states , due to the curse of dimensionality . Learning with function approximation upon good representation is a natural idea to tackle the curse and serving as the key for the success of deep learning ( Bengio et al. , 2013 ) . In fact , representation learning lies at the heart of the empirical successes of deep RL in video games ( Mnih et al. , 2013 ) , robotics ( Levine et al. , 2016 ) , Go ( Silver et al. , 2017 ) , dialogue systems ( Jiang et al. , 2021 ) to name a few . Meanwhile , the importance and benefits of the representation in RL is rigorously justified ( Jin et al. , 2020 ; Yang & Wang , 2020 ) , which quantifies the regret in terms of the dimension of the known representation based on a subclass in Markov decision processes ( MDPs ) ( Puterman , 2014 ) . A natural question raises : How to design provably efficient and practical algorithm for representation learning in RL ? Here , by “ provably efficient ” we mean the sample complexity of the algorithm can be rigorously characterized only in terms of the complexity of representation class , without explicit dependency on the number of states and actions , while by “ practical ” we mean the algorithm can be implemented and deployed for the real-world applications . Therefore , we not only require the representation learned is expressive enough for handling complex practical environments , but also require the operations in the algorithm tractable and computation/memory efficient . The major difficulty of this question lies in two-fold : i ) The trade-off between the expressiveness and the tractability in the design of the representations ; ii ) The learning of representation is intimately coupled with exploration . Specifically , a desired representation should be sufficiently expressive1 to capture the dynamic system , while still tractable in practice . However , in general , expressive representation leads to complicated optimization in learning . For example , the representation in the linear MDP is exponential 1For a formal definition of expressiveness , see ( Agarwal et al. , 2020a ) . stronger than the latent variable MDPs in terms of expressiveness ( Agarwal et al. , 2020a ) . However , its representation learning is either depending on a MLE oracle that is computational intractable due to the constraint on the regularity of conditional density ( Agarwal et al. , 2020a ) , or a complicated constrained min-max-min-max optimization ( Modi et al. , 2021 ) . On the other hand , Misra et al . ( 2020 ) considers the representation introduced by an encoder in block MDP ( Du et al. , 2019 ) , in which the learning problem can be completed by a regression , but with the payoff that the representations in block MDP is even weaker than latent variable MDP ( Agarwal et al. , 2020a ) . Meanwhile , the coupling of the learning of representation and exploration induces the difficulty in practical algorithm design and analysis . One can not learn a precise representation without enough experiences from a comprehensive exploration , while the exploration depends on an reliable estimation of the representation . Most of the known results depends on a policy-cover-based exploration ( Du et al. , 2019 ; Misra et al. , 2020 ; Agarwal et al. , 2020a ; Modi et al. , 2021 ) , which maintains and samples a set of policies during training for systematic exploration , that significant increases the computation and memory cost in implementation . In this work , we propose Spectral Dynamics Embedding ( SPEDE ) , bypassing the aforementioned difficulties and answering the question affirmatively . SPEDE is established on an observation that connects the stochastic control dynamics ( Osband & Van Roy , 2014 ; Kakade et al. , 2020 ) with linear MDPs in Section 3 . Specifically , by exploiting the property of the noise in the stochastic control dynamics , we can recover the factorization of its corresponding Markov transition operator in closed-form without extra computation . This equivalency immediately overcome the computational intractability in the model estimation in Agarwal et al . ( 2020a ) , which breaks the trade-off between expressiveness and tractability . More importantly , the connection unifies two highly-related but commonly-known different models , i.e. , stochastic nonlinear control models ( Kakade et al. , 2020 ) and linear MDPs ( Jin et al. , 2020 ) , therefore , provides the opportunity to share benefits from both sides : i ) , it sheds the light on exploiting optimistic control for exploratory representation learning , instead of expensive policycover-based exploration ; and ii ) , it introduces the linear sufficient feature from the spectral space of Markov operator , in which the planning can be completed efficiently . We rigorously characterize the statistical property of SPEDE in terms of regret w.r.t . the complexity of representation class in Section 4 , without explicit dependence on the raw feature of state and action . With the established unified view , our results generalize online control ( Kakade et al. , 2020 ) and linear MDP ( Jin et al. , 2020 ) beyond known features . We finally demonstrate the superior of SPEDE on the MuJoCo benchmarks in Section 5 . It significantly outperforms the empirical stateof-the-art RL algorithms . To our knowledge , SPEDE is the first representation learning algorithm achieving statistical , computational , and memory efficiency with sufficient expressiveness . 1.1 RELATED WORK . There have been many great attempts to learn a variety of algorithmic representation learning in RL designed for different purposes , e.g. , bisimulation ( Ferns et al. , 2004 ; Gelada et al. , 2019 ; Castro , 2020 ) , reconstruction ( Watter et al. , 2015 ; Hafner et al. , 2019 ) . Recently , there are also several works considering the spectral features based on decomposing different variants of the transition operator , including successor features ( Dayan , 1993 ; Kulkarni et al. , 2016 ) , proto-value functions ( Mahadevan & Maggioni , 2007 ; Wu et al. , 2018 ) , spectral state-aggregation ( Duan et al. , 2018 ; Zhang & Wang , 2019 ) , and contrastive fourier features ( Nachum & Yang , 2021 ) . These works are highly-related to the proposed SPEDE . Besides these features focus on state-only representation , the major difference between SPEDE and these spectral features lies in i ) , the target operators in existing spectral features are state-state transition , which cancel the effect of action ; ii ) , the target operators are estimated based on empirical data from a fixed behavior policy under the implicit assumption that the estimated operator is uniformly accurate , ignoring the major difficulty in exploration , while SPEDE carefully designed the systematic exploration with theoretical guarantee ; iii ) , most of the existing spectral features rely on explicitly decomposion of the operators , while SPEDE obtains the spectral for free . Turning to the rigorously-justified representation learning , a large body of effort focuses on policy-cover-based methods in an “ explore-then-commit ” strategy ( Du et al. , 2019 ; Misra et al. , 2020 ; Agarwal et al. , 2020a ; Modi et al. , 2021 ) . These algorithms learn a uniformly accurate model/representation through a reward-free exploration , upon which decouple the learning from the exploration . Then , the representation from the learned model can be used for optimal policy seeking for the particular reward . The major difficulty impedes their practical application is the computa- tion and memory cost : the policy-cover-based exploration requires a set of exploratory polices to be maintained and sampled from during training , which can be extremely expensive . Another two related lines of research are model-based RL and online control , which are commonly known overlapped but separate communities considering different formulations of the dynamics . Our finding establishes the equivalency among the model , so that bridges these communities . Osband & Van Roy ( 2014 ) and Kakade et al . ( 2020 ) are the most related to our work in each community . These models generalize their corresponded linear models , i.e. , Jin et al . ( 2020 ) and Cohen et al . ( 2019 ) , with general nonlinear model and kernel function within a known RKHS , respectively . The regret of the optimstic ( pessimistic ) algorithm has been carefully characterized for these models . However , both of the proposed algorithms in Osband & Van Roy ( 2014 ) and Kakade et al . ( 2020 ) require an planning oracle to seek the optimal policy , which might be computational intractable . In SPEDE , this is easily handled in the equivalent linear MDP . 2 PRELIMINARIES . Markov Decision Process ( MDP ) is one of the most standard models studied in the reinforcement learning that can be denoted by the tuple M = ( S , A , r , T , µ , H ) , where S is the state space , A is the action space , r : S × A → R+ is the reward function2 ( where R+ denotes the set of nonnegative real numbers ) , T : S × A → ∆ ( S ) is the transition and µ is an initial state distribution and H is the horizon3 ( i.e . the length of each episode ) . A ( potentially non-stationary ) policy π can be defined as { πh } h∈ [ H ] where πh : S → A , ∀h ∈ [ H ] . Following the standard notation , we define the value function V πh ( sh ) : = ET , π [ ∑H−1 t=h r ( st , at ) |sh = s ] and the action-value function ( i.e . the Q function ) Qπh ( sh , ah ) = ET , π [ ∑H−1 t=h r ( st , at ) |sh = s , ah = a ] , which are the expected cumulative rewards under transition T when executing policy π starting from sh and ( sh , ah ) . With these two definitions at hand , it is straightforward to show the following Bellman equation : Qπh ( sh , ah ) = r ( sh , ah ) + Esh+1∼T ( ·|sh , ah ) V πh+1 ( sh+1 ) . Reinforcement learning aims at finding the optimal policy π∗ = arg maxπ Es∼µV π0 ( s ) . It is well known that in the tabular setting when the state space and action space are finite , we can provably identify the optimal policy with both sample-efficient and computational-efficient optimism-based methods ( e.g . Azar et al. , 2017 ; Jin et al. , 2018 ; Zhang et al. , 2021 ) with the complexity proportion to |S||A| . However , in practice , the cardinality of state and action space can be large or even infinite . Hence , we need to incorporate function approximation into the learning algorithm when we deal with such cases . The linear MDP ( Jin et al. , 2020 ) or low-rank MDP ( Agarwal et al. , 2020a ; Modi et al. , 2021 ) is the most well-known reinforcement learning model that can incorporate linear function approximation with theoretical guarantee , thanks to the following assumption on the transition : T ( s′|s , a ) = 〈φ ( s , a ) , µ ( s′ ) 〉H , ( 1 ) where φ : S × A → H , µ : S → H are two feature maps and H is a Hilbert space . The most essential observation for them is that , Qπh ( s , a ) for any policy π is linear w.r.t φ ( sh , ah ) , due to the following observation ( Jin et al. , 2020 ) : ∫ V πh+1 ( sh+1 ) T ( sh+1|sh , ah ) dsh+1 = 〈 φ ( sh , ah ) , ∫ V πh+1 ( sh+1 ) µ ( sh+1 ) dsh+1 〉 H . ( 2 ) Therefore , φ serves as a sufficient representation for the estimation of Qπh , that can provide uncertainty estimation with standard linear model analysis and eventually lead to sample-efficient learning when φ is fixed and known to the agent ( see Theorem 3.1 in Jin et al. , 2020 ) . However , we in general do not have such representation in advance4 and we need to learn the representation from the data , which constraints the applicability of the algorithms derived with fixed and known representation .

Summary: This paper studies the single reward episodic MDP problem with the model-based TS algorithm. Assuming the given model class satisfies realizability, regularization property, bounded Eluder dimension, low covering number, and the true dynamics is the stochastic control model under gaussian noise transition (eq 5), the authors show the polynomial Bayesian regret guarantee. In addition, experiments on OPENAI MuJoCo is conducted.