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 work proposes to improve model-based RL methods by enforcing the model and value learned to be “self-consistent”, i.e. to satisfy the Bellman eqn where both the Bellman operator and value function are defined in terms of the model. The motivation behind this is that in classic MBRL, the model and value function are learned separately, and the value function is typically learned to be consistent with the model but not vice versa. The authors propose to enforce self-consistency by adding a "self-consistency loss” (L_SC) to the base loss and jointly learning the model and value functions. L_SC is the squared TD error with a stop-gradient on the target and is calculated using rollouts from the model. Although in MuZero/Muesli the model and value functions are learned jointly as well, the targets of the loss are all defined with respect to data from the true environment and there is not term enforcing the “self-consistency” between model and value functions. Experimental evaluations on a selection of Atari games show that jointly updating the model and value functions with the additional L_SC term leads to improved sample efficiency as compared to Dyna and a model-free baseline. The authors also empirically study the search control problem.