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 problem of learning representations for RL. On the theory side, the paper considers the setting where the state-transition is a nonlinear function of the past state-action plus additive noise, and develops a no-regret algorithm. On the empirical side, the paper shows that an adaptation of this algorithm on real-world RL tasks could perform better than existing model-based algorithms.