Learning State Representations via Retracing in Reinforcement Learning

1 INTRODUCTION . Recent developments in deep reinforcement learning ( RL ) have made great progress in solving complex control tasks ( Mnih et al. , 2013 ; Levine et al. , 2016 ; Silver et al. , 2017 ; Vinyals et al. , 2017 ; Schrittwieser et al. , 2020 ) . With the increasing capacity of deep RL algorithms , the problems of interests become increasingly complex . An immediate challenge is that the observation space becomes unprecedentedly high-dimensional , and often the perceived observations have significant redundancy , negatively impacting the policy learning . The field of representation learning offers a wide range of approaches for extracting useful information from high-dimensional data ( with potentially sequential dependency structure ) ( Bengio et al. , 2013 ) . Many recent works have explored the application of representation learning in RL ( Ha and Schmidhuber , 2018 ; Hafner et al. , 2019 ; 2020a ; Schrittwieser et al. , 2020 ; Schwarzer et al. , 2021 ; Zhang et al. , 2020 ) , which lead to superior performance comparing to naive embedding . Many of such algorithms rely on predictive ( reconstruction ) supervision for representation learning , i.e. , the state representations are learned such that transitions in the observation space are maximally preserved in the embedding space . Here we argue that existing methods do not fully exploit the supervisory signals inherent in the data . Additionally valid supervision can often be obtained for representation learning by including temporally “ backward ” transitions in situations in which the same set of rules govern both temporally forward and backward transitions . In these cases we are able to use such temporally backward transitions as “ free ” supervision for training the representation/model without having to perform any additional interaction with the environment . Hence , by “ learning via retracing ” , we obtain more training samples for representation learning ( twice as much as existing approaches in tasks that admit perfect reversibility across all transitions ) , hence improve the sample efficiency of representation learning . Moreover , the quality of the representation is strongly correlated with the quality of policy learning ( Gelada et al. , 2019 ) . Hence with faster representation learning , we can subsequently achieve more efficient policy learning . Thus we hypothesise that by augmenting representation learning with “ learning via retracing ” in RL , we can significantly improve the sample efficiency of learning , which is a long-standing issue that plagues the practical applicability of deep RL algorithms . Beyond improved sample efficiency , joint predictive supervision in temporally forward and backward directions conveys information from both the future and past to the target state , similar to the smoothing operation for latent state inference in state-space models ( Kalman , 1960 ; Murphy , 2012 ) , leading to more accurate latent state inference , hence achieving stronger representation learning . “ Learning via retracing ” can be integrated into any representation learning that utilises a transition model , under both the model-free and model-based RL frameworks . One problem that plagues the successful application of “ learning via retracing ” is that the reversibility might not be preserved across all valid transitions , i.e. , there exists transitions such that s→ s′ for some action a , but s′ 6→ s for all possible actions a ( Figure 1b ) . Under these situations , naively forcing the representations for s and š ( the retraced state given s′ and a , see Section 3 ) leads to suboptimal representation space , hence potentially hinders or even completely interrupts the overall RL training . Hence in order to maximally preserve the advantages brought by “ learning via retracing ” , it is essential to identify the “ irreversible ” transitions and rule them out from representation learning via retracing . To this end , we propose a novel dynamically regulated approach for identifying such “ irreversible ” states , which we term as adaptive truncation ( Section 3.3 ) . As a motivating example , consider a rat navigating towards a cheese in a cluttered environment ( Figure 1a ) . Upon first visit to the goal state , multiple imaginative retracing trajectories can be randomly simulated . By constraining the temporal cycle-consistency of the retracing transitions , the rat quickly builds a state representation that accurately preserves the transitions in the area around the actual forward trajectory taken by the rat . Moreover , all retracing simulations pass through the two “ doors ” , allowing the rat to quickly identify the key bottleneck states that are essential for successful navigation towards the goal , hence facilitating stronger transferability across environments ( Section 5.3 ) . We conjecture that such imaginative retracing could be physiologically implemented as the reversed hippocampal “ replay ” that has been observed in both rodents and humans ( Foster and Wilson , 2006 ; Mattar and Daw , 2018 ) ( see Section 6 for further discussion ) . We propose Cycle-Consistency World Model ( CCWM ) , a self-supervised instantiation of “ learning via retracing ” , for joint representation learning and generative model learning under the model-based RL settings . We empirically evaluate the performance of CCWM on challenging visual-based continuous control benchmarks . Experimental results show that CCWM achieves state-of-the-art performance in terms of sample efficiency and asymptotic performance , whilst providing additional advantages such as stronger generalisability and extended planning horizon . 2 PRELIMINARIES . 2.1 PROBLEM FORMULATION . We consider reinforcement learning problems in Markov Decision Processes ( MDPs ) . An MDP can be charaterised by the tuple , M = 〈S , A , R , P , γ〉 , where S , A are the state and action spaces , respectively ; R : S → R is the reward function ( we assume determinisitc reward functions unless stated otherwise ) , P : S ×A×S → [ 0 , 1 ] is the transition distribution of the task dynamics ; γ ∈ R is the discounting factor . The control policy , π : S ×A → [ 0 , 1 ] , represents a distribution over actions at each state . The goal is to learn the optimal policy , π∗ , such that the expected future reward is maximised across all states , i.e. , π∗ = arg max π∈Π Eπ [ ∑ t γtR ( st , at ) |s0 = s ] , ∀s ∈ S , where st ∼ p ( ·|st−1 , at−1 ) for t = 1 , 2 , . . . ( 1 ) Here we consider tasks in which the perceivable observation space , O , is high-dimensional , due to either redundant information or simply because only visual inputs are available . Hence it is necessary to learn an embedding function , φ : O → Z , such that the embedded observation space , φO , could act as S in the MDP , to support efficient learning of the optimal policy using existing RL algorithms . 2.2 GENERATIVE MODELLING OF DYNAMICS . Modelling the transition dynamics using a sequential VAE-like structure enables joint learning of the latent representation and the associated latent transitions . Specifically , the dynamics model is defined in terms of the following components ( see also top part in Figure 2b ) . Observation ( context ) embedding : et = qφ ( Ot ) , Latent posterior distribution : p ( zt+1|zt , at , Ot+1 ) , Latent variational transition ( prior ) distribution : qψ1 ( zt+1|zt , at ) , Latent variational posterior distribution : qψ2 ( zt+1|zt , at , et+1 ) , Generative distribution : pθ ( Ot+1|zt+1 ) , ( 2 ) where ψ = { ψ1 , ψ2 , φ } and θ represent the parameters associated with the recognition and generative models respectively . Latent variables z ∈ Z are introduced for more flexible modelling of the distribution of the observed variables . A variational approximation is employed since the true posterior p ( zt+1|zt , at , Ot+1 ) is usually intractable in practice . At each time step t , the agent receives an observationOt , which is then embedded into a context vector et = qφ ( Ot ) . After initialisation , the latent space vector is rolled out in a forward fashion given the action at , yielding one-step prediction into the future following the variational transition distribution qψ1 ( zt+1|zt , at ) ( we assume standard first-order Markovian structure on the latent variables ) . The dynamics model is trained following the maximum likelihood paradigm , and due to intractability , the parameters of the recognition and generative models in Eq . 2 are learned by maximising the variational free energy ( also known as the ELBO ; Wainwright and Jordan ( 2008 ) . LELBO ( Ot+1 ) = Ezt+1∼qψ2 [ log pθ ( Ot+1|zt+1 ) − βDKL [ qψ2 ( zt+1|zt , at , et+1 ) ||qψ1 ( zt+1|zt , at ) ] ] , ( 3 ) The variational free energy loss function consists of the reconstruction error , pθ ( Ot+1|zt+1 ) , and the KL-divergence between the variational posterior distributions and the predictive prior as regularisation . The intuitive autoencoder-like structure nicely separates the generative process from the inference process , with the variational posterior ( qψ2 ( zt+1|zt , at , et+1 ) ) serving as the main inference engine of the optimal latent representations . Note that the β parameter controls the degree of factored structure ( disentanglement ) of the latent code , which by default is set to 1 ( Higgins et al. , 2016 ) . 3 METHOD . We firstly introduce learning via retracing in its most general format , then provide a concrete instantiation based on generative dynamics modelling under the model-based framework . We finally propose a novel adaptive truncation scheme for dealing with the “ irreversibility ” issue in “ learning via retracing ” . 3.1 LEARNING VIA RETRACING . We always assume the usage of an approximate dynamics model , regardless of the overall RL agent being model-free or model-based ( the dynamics model would only be used for representation learning under the model-free setting , hence it is possible to have separate dynamics model for forward and “ reversed ” transitions ) . Given a set observations , O = { O1 , . . . , OT } , a dynamics model , T : S ×A× S → [ 0 , 1 ] , the classical methods of self-supervised learning of representations involve learning an encoder ( Eφ ) and a decoder ( Dθ ) , trained via minimising the predictive reconstruction . L ( φ , θ ) = d ( O1 : T , D ( f ( E ( O1 : T ) ; T ) ) ( 4 ) where d is some metric of the observable space ( e.g. , the L2 distance ) . The function f ( e ; T ) is some function specifying the schedule for predictions , e.g. , f ( et , T ) = T ( et ) = êt+1 corresponds to learning the representations based on one-step predictive reconstruction . Existing methods explore predictive supervision in a temporally forward direction , however , we argue that the “ reversed ” transitions can also contain useful signals for learning . Consider a transition tuple , ( s , a , s′ ) , we define the “ reversed ” transitions being the tuple ( s′ , a′ , s ) , where a′ is the “ reversed ” action . In situations where the same set of rules apply to forward and backward transitions , the reversed transition given a′ could contribute to representation learning via Eq . 4 as an additional training sample ( utilising a potentially different loss function from the forward supervision , see Figure 2a ) . Hence by utilising the additional reversed transition for representation learning , we improve the sample efficiency of learning without any additional interactions with the environment . As mentioned earlier , in situations where perfect reversibility is not preserved across all transitions , such “ irreversible ” transitions could negatively impact the overall learning . Correct identification of such states is hence essential for the successful implementation of “ learning via retracing ” . To this end , we propose a novel adaptive truncation scheme in Section 3.3 . Despite the intuitive simplicity of learning via retracing , it offers a number of advantages comparing to existing representation learning methods . In addition to the improved sample efficiency , we can interpret learning with “ reversed ” transitions as explicit inference of the current latent state given future information . In combination with the forward predictive supervision , the joint learning dynamics is similar to the smoothing operation in dynamical systems , which is often superior than filtering ( corresponds to using solely the forward predictive supervision ) in terms of inference accuracy ( Kalman , 1960 ; Murphy , 2012 ) . Hence learning via retracing could support stronger representation learning . We have introduced learning via retracing in its most general form , where a large degree of freedom exists such that the method can be tailored and integrated with many of the existing representation learning approaches . There are many free model choices , such as being model-free or model-based ; whether or not using a separate “ reversed ” dynamics model ; loss function for constraining the “ retraced ” transitions ; deterministic or probabilistic dynamics model , just to name a few . Below we provide one possible instantiation of learning via retracing , the Cycle-Consistency World Model ( CCWM ) , under the model-based framework based on generative dynamics modelling .

Summary: The paper investigates a self-supervised approach for learning the state representation in RL tasks. In addition to the predictive (reconstruction) supervision in the forward direction, the authors include a “retraced” transitions for representation/model learning, by enforcing the cycle-consistency constraint between the original and retraced states. The authors claim that this facilitates stronger representation learning and improve upon the sample efﬁciency of learning. As it is not always possible to find such a cycle consistency between two states and for counteracting the negative impacts brought by the “irreversible” transitions, the authors add a novel adaptive “truncation” mechanism.