Self-Consistent Models and Values

Learned models of the environment provide reinforcement learning ( RL ) agents with flexible ways of making predictions about the environment . In particular , models enable planning , i.e . using more computation to improve value functions or policies , without requiring additional environment interactions . In this work , we investigate a way of augmenting model-based RL , by additionally encouraging a learned model and value function to be jointly self-consistent . Our approach differs from classic planning methods such as Dyna , which only update values to be consistent with the model . We propose multiple self-consistency updates , evaluate these in both tabular and function approximation settings , and find that , with appropriate choices , self-consistency helps both policy evaluation and control . 1 Introduction . Models of the environment provide reinforcement learning ( RL ) agents with flexible ways of making predictions about the environment . They have been used to great effect in planning for action selection [ 45 , 29 , 49 ] , and for learning policies or value functions more efficiently [ 51 ] . Learning models can also assist representation learning , serving as an auxiliary task , even if not used for planning [ 25 , 48 ] . Traditionally , models are trained to be consistent with experience gathered in the environment . For instance , an agent may learn maximum likelihood estimates of the reward function and state-transition probabilities , based on the observed rewards and state transitions . Alternatively , an agent may learn a model that only predicts behaviourally-relevant quantities like rewards , values , and policies [ 47 ] . In this work , we study a possible way to augment model-learning , by additionally encouraging a learned model m̂ and value function v̂ to be jointly self-consistent , in the sense of jointly satisfying the Bellman equation with respect to m̂ and v̂ for the agent ’ s policy π . Typical methods for using models in learning , like Dyna [ 51 ] , treat the model as a fixed best estimate of the environment , and only update the value to be consistent with the model . Self-consistency , by contrast , jointly updates the model and value to be consistent with each other . This may allow information to flow more flexibly between the learned reward function , transition model , and approximate value . Since the true model and value are self-consistent , this type of update may also serve as a useful regulariser . We investigate self-consistency both in a tabular setting and at scale , in the context of deep RL . There are many ways to formulate a model-value update based on the principle of self-consistency , but we find that naive updates may be useless , or even detrimental . However , one variant based on a semi-gradient temporal difference objective can accelerate value learning and policy optimisation . We evaluate different search-control strategies ( i.e . the choice of states and actions used in the selfconsistency update ) , and show that self-consistency can improve sample efficiency in environments such as Atari , Sokoban and Go . We conclude with experiments designed to shed light on the mechanisms by which our proposed self-consistency update aids learning . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . 2 Background . We adopt the Markov Decision Process ( MDP ) formalism for an agent interacting with an environment [ 7 , 42 ] . The agent selects an action At ∈ A in state St ∈ S on each time step t , then observes a reward Rt+1 and transitions to the successor state St+1 . The environment dynamics is given by a “ true ” model m∗ = ( r , P ) , consisting of reward function r ( s , a ) and transition dynamics P ( s′|s , a ) . The behaviour of the agent is characterized by a policy π : S → ∆A , a mapping from states to the space of probability distributions over actions . The agent ’ s objective is to update its policy to maximise the value vπ ( s ) = Em∗ , π [ ∑∞ j=0 γ jr ( St+j , At+j ) | St = s ] of each state . We are interested in model-based approaches to this problem , where an agent uses a learned model of the environment m̂ = ( r̂ , P̂ ) , possibly together with a learned value function v̂ ( s ) ≈ vπ , to optimise π . The simplest approach to learning such a model is to make a maximum likelihood estimate of the true reward and transition dynamics , based on the agent ’ s experience [ 32 ] . The model can then be used to perform updates for arbitrary states and action pairs , entirely in imagination . Dyna [ 51 ] is one algorithm with this property , that has proven effective in improving the data efficiency of RL [ 22 , 31 ] . An alternative model-learning objective is value equivalence1 i.e . equivalence of the true and learned models in terms of the induced values for the agent ’ s policy . This can be formalised as requiring the induced Bellman operator of the model to match that of the environment : T πm∗vπ = T πm̂vπ , ( 1 ) where the Bellman operator for a model m̂ is defined by T πm̂v = Em̂ , π [ r ( s , a ) + γv ( s′ ) ] . Such models do not need to capture every detail of the environment , but must be consistent with it in terms of the induced Bellman operators applied to certain functions ( in this case , the agent ’ s value vπ ) . We take a particular interest in value equivalent models for two reasons . First , they are used in stateof-the-art agents , as described in the next section and used as baselines in our deep RL experiments . Second , these algorithms update the model using a value-based loss , like the self-consistency updates we will study in this work . We describe this connection further in Section 3.3 . 2.1 Value equivalent models in deep RL . A specific flavour of value equivalent model is used by the MuZero [ 47 ] and Muesli [ 25 ] agents . Both use a representation function ( encoder ) hθ to construct a state zt = hθ ( o1 : t ) from a history of observations o1 : t2 . We define a deterministic recurrent model m̂ ≡ mθ in this latent space . On each step k the model , conditioned on an action , predicts a new latent state zkt and reward r̂ k t ; the agent also predicts from each zkt a value v̂ k t and a policy π̂ k t . We use superscripts to denote time steps in the model : i.e. , zkt is the latent state after taking k steps with the model , starting in the root state z 0 t ≡ zt . Muesli and MuZero unroll such a model for K steps , conditioning on a sequence of actions taken in the environment . The learned components are then trained end-to-end , using data gathered from interaction with the environment , D∗π , by minimizing the loss function Lbaset ( m̂ , v̂ , π̂|D∗π ) = ED∗π K∑ k=0 [ ` r ( rtargett+k , r̂ k t ) + ` v ( vtargett+k , v̂ k t ) + ` π ( πtargett+k , π̂ k t ) ] . ( 2 ) This loss is the sum of a reward loss ` r , a value loss ` v and a policy loss ` π . Note that there is no loss on the latent states zkt : e.g . these are not required to match zt+k = hθ ( o1 : t+k ) . The targets r target t+k for the reward loss are the true rewards Rt+k observed in the environment . The value targets are constructed from sequences of rewards and n-bootstrap value estimates vtargett+k = ∑n j=1 γ j−1rt+k+j + γ nṽt+k+n . The bootstrap value estimates ṽ , as well as the policy targets , are constructed using the model m̂ . This is achieved in MuZero by applying Monte-Carlo tree search , and , in Muesli , by using one-step look-ahead to create MPO-like [ 1 ] targets . Muesli also employs a policy-gradient objective which does not update the model . We provide some further details about Muesli in Appendix B.2 , as we use 1We borrow the terminology from Grimm et al . [ 19 ] , who provided a theoretical framework to reason about such models . Related ideas were also introduced under the name ‘ value aware ’ model learning [ 15 , 14 ] 2We denote sequences of outcomes of random variables O1 = o1 , . . . , Ot = ot as o1 : t for simplicity . Algorithm 1 : Model-based RL with joint grounded and self-consistency updates . input : initial m̂ , v̂ , π output : estimated value v̂ ≈ vπ and/or optimal policy π∗ repeat Collect D∗π from m∗ following π Compute grounded loss : Lbase ( m̂ , v̂ , π|D∗π ) // e.g . Eq . ( 2 ) Generate D̂µ from m̂ following µ Compute self-consistency loss : Lsc ( m̂ , v̂|D̂µ ) // see Eqs . ( 4,5,6 ) Update m̂ , v̂ , π by minimising ( e.g. , with SGD ) : L = Lbase + Lsc until convergence ; it as a baseline for our deep RL experiments . For comprehensive descriptions of MuZero and Muesli we refer to the respective publications [ 47 , 25 ] . In both cases , the model is used to construct targets based on multiple imagined actions , but the latent states whose values are updated by optimising the objective ( 2 ) always correspond to real states that were actually encountered by the agent when executing the corresponding action sequence in its environment . Value equivalent models could also , in principle , update state and action pairs entirely in imagination ( similarly to Dyna ) but , to the best of our knowledge , this has not been investigated in the literature . Our proposed self-consistency updates provide one possible mechanism to do so . 3 Self-consistent models and values . A particular model m̂ and policy π induce a corresponding value function vπm̂ , which satisfies a Bellman equation T πm̂vπm̂ = vπm̂ . We describe model-value pairs which satisfy this condition as self-consistent . Note that the true model m∗ and value function vπ are self-consistent by definition . If an approximate model m̂ and value v̂ are learned independently , they may only be asymptotically self-consistent , in that the model is trained to converge to the true model , and the estimated value to the true value . Model-based RL algorithms , such as Dyna , introduce an explicit drive towards consistency throughout learning : the value is updated to be consistent with the model , in addition to the environment . However , the model is only updated to be consistent with transitions in the real environment D∗π , and not to be consistent with the approximate values . Instead , we propose to update both the model and values so that they are self-consistent with respect to trajectories D̂µ sampled by rolling out a model m̂ , under action sequences sampled from some policy µ ( with µ 6= π in general ) . We conjecture that allowing information to flow more freely between the learned rewards , transitions , and values , may make learning more efficient . Self-consistency may also implement a useful form of regularisation – in the absence of sufficient data , we might prefer a model and value to be selfconsistent . Finally , in the context of function approximation , it may help representation learning by constraining it to support self-consistent model-value pairs . In this way , self-consistency may also be valuable as an auxiliary task even if the model is not used for planning .

Summary: This paper presents the idea of self-consistent reinforcement learning.  By enforcing "consistency" between value functions and approximate  (learned) models, a model-based learning algorithm can use synthetic experience generated from the approximate model.  This paper presents a few minor variations on this theme, and explores the results empirically.