Value-Driven Hindsight Modelling

Value estimation is a critical component of the reinforcement learning ( RL ) paradigm . The question of how to effectively learn predictors for value from data is one of the major problems studied by the RL community , and different approaches exploit structure in the problem domain in different ways . Model learning can make use of the rich transition structure present in sequences of observations , but this approach is usually not sensitive to the reward function . In contrast , model-free methods directly leverage the quantity of interest from the future but have to compose with a potentially weak scalar signal ( an estimate of the return ) . In this paper we develop an approach for representation learning in RL that sits in between these two extremes : we propose to learn what to model in a way that can directly help value prediction . To this end we determine which features of the future trajectory provide useful information to predict the associated return . This provides us with tractable prediction targets that are directly relevant for a task , and can thus accelerate learning of the value function . The idea can be understood as reasoning , in hindsight , about which aspects of the future observations could help past value prediction . We show how this can help dramatically even in simple policy evaluation settings . We then test our approach at scale in challenging domains , including on 57 Atari 2600 games . 1 INTRODUCTION . Consider a baseball player trying to perfect their pitch . The player performs an arm motion and releases the ball towards the batter , but suppose that instead of observing where the ball lands and the reaction of the batter , the player only gets told the result of the play in terms of points or , worse , only gets told the final result of the game . Improving their pitch from this experience appears hard and inefficient , yet this is essentially the paradigm we employ when optimizing policies in model-free reinforcement learning . The scalar feedback that estimates the return from a state ( and action ) , encoding how well things went , drives the learning while the accompanying observations that may explain that result ( e.g . flight path of the ball or the way the batter anticipated and struck the incoming baseball ) are ignored . To intuitively understand how such information could help value prediction , consider a simple discrete Markov chain X → Y → Z , where Z is the scalar return and X is the observation from which we are trying to predict Z . If the space of possible values of Y is smaller than X , then it may be more efficient to estimate both P ( Y |X ) and P ( Z|Y ) rather than directly estimating P ( Z|X ) .1 In other words observing and then predicting Y can be advantageous to directly estimating the signal of interestZ . Model-based RL approaches would duly exploit the observed Y ( by modeling the transition Y |X ) , but Y would , in general scenarios , contain information that is irrelevant to Z and hard to predict . Building a full high-dimensional predictive model to indiscriminately estimate all possible future observations , including potentially chaotic details of the ball trajectory and the spectators ’ response , is a challenge that may not pay off if the task-relevant predictions ( e.g. , was the throw accepted , was the batter surprised ) are error-ridden . Model-free RL methods directly consider the relation X to Z , and focus solely upon predicting and optimising this goal , rather than attempting to learn the full dynamics . These methods have recently dominated the literature , and have attained the best performance in a wide array of complex problems with high-dimensional observations ( Mnih et al. , 2015 ; Schulman et al. , 2017 ; Haarnoja et al. , 2018 ; Guez et al. , 2019 ) . 1In the discrete case , this follows from a counting argument from the size of the probability tables involved . In this paper , we propose to augment model-free methods with a lightweight model of future quantities of interest . The motivation is to model only those parts of the future observations ( Y ) that are needed to obtain better value predictions . The major research challenge is to learn , from observational data , which aspects of the future are important to model ( i.e . what Y should be ) . To this end , we propose to learn a special value function in hindsight that receives future observations as an additional input . This learning process reveals features of the future observations that would be most useful for value prediction ( e.g . flight path of the ball or the reaction of the batter ) , if provided by an oracle . These important features are then predicted , in advance , using only information available at test time ( at the time of releasing the baseball , we knew the identity of the batter , the type of throw and spin given to the ball ) . Learning these value-relevant features can help representation learning for an agent and provide an additional useful input to its value and policy . Experimentally , hindsight value functions surpassed the performance of model-free RL methods in a challenging association task ( Portal Choice ) . When hindsight value functions were added to the prior state-of-the-art RL method for Atari games , they significantly increased median performance from 833 % to 965 % . 2 BACKGROUND AND NOTATION . We consider a reinforcement learning setting whereby an agent learns from interaction in a sequential decision-making environment ( Sutton & Barto , 2011 ) . An agent ’ s policy π , mapping states to an action distribution , is executed to obtain a sequence of rewards and observations as follows . At each step t , after observing state st , the policy outputs an action at , sampled from π ( A|st ) , and obtains a scalar reward rt and the next-state st+1 from the environment . The sum of discounted rewards from state s is the return denoted by G = ∑∞ t=0 γ tRt , with γ < 1 denoting the discount factor . Its expectation , as a function of the starting state , is called the value function , vπ ( s ) = Eπ [ G|S0 = s ] . An important related quantity is the action-value , or Q-value , which corresponds to the same expectation with a particular action executed first : qπ ( s , a ) = Eπ [ G|S0 = s , A0 = a ] . The learning problem consists in adapting the policy π in order to achieve a higher value vπ . This usually entails learning an estimate of vπ for the current policy π , this is the problem we focus on in this paper . Note that in practice we are interested in partially-observed environments where the state of the world is not directly accessible . For this case , we can think of replacing the observed state s in the case of the fully-observed case by a learned function that depends on past observations . 3 VALUE LEARNING . 3.1 DIRECT LEARNING . A common approach to estimate v ( or q ) is to represent it as a parametric function vθ ( or qθ ) and directly update its parameters based on sample returns of the policy of interest . Value-based RL algorithms vary in how they construct a value target Y from a single trajectory . They may regress vθ towards the Monte-Carlo return ( Yt = Gt ) , or exploit sequentiality in the reward process by relying on a form of temporal-difference learning to reduce variance ( e.g . the TD ( 0 ) target Yt = Rt + γvθ ( St+1 ) ) . For a given target definition Y , the value loss Lv to derive an update for θ is : Lv ( θ ) = 12Es [ ( vθ ( s ) − Y ) 2 ] . In constructing a target Yt based on a trajectory of observations and rewards from time t , the observations are either unused ( for a Monte Carlo return ) or only indirectly exploited ( when bootstrapping to obtain a value estimate ) . In all cases , the trajectory is distilled into a scalar signal that estimates the return of a policy , and other relevant aspects of future observations are discarded . In particular in partially observed domains or domains with highdimensional observation spaces it can be difficult to discover correlations with this noisy signal . 3.2 MODEL-BASED APPROACH . An indirect way of estimating the value is to first learn a model of the dynamics . For example a 1-step observation model mθ learns to predict the conditional distribution st+1 , rt|st , at . Then a value estimate v ( s ) for state s can be obtained by autoregressively rolling out the model ( until the end of the episode or to a fixed depth with a parametric value bootstrap ) . The model is trained on potentially much richer data than the return signal since it exploits all information in the trajectory . Indeed , the observed transitions between states can reveal the structure behind a sparse reward signal . A drawback of classic model-based approaches is that they predict a high-dimensional signal , a task which may be costly and harder than directly predicting the scalar value . As a result , the approximation of the dynamics mθ may contain errors where it matters most for predicting the value ( Talvitie , 2014 ) . Although the observations carry all the data from the environment , most of it is not essential to the task ( Gelada et al. , 2019 ) . The concern that modeling all observations is expensive also applies when the model is not used for actual rollouts but merely for representation learning . So while classic model-based methods fully use this highdimensional signal at some cost , model-free methods take the other extreme to focus only on the most relevant low-dimensional signal ( the scalar return ) . Below we propose a method that strikes a balance between these paradigms . 3.3 HINDSIGHT VALUE AND MODEL . We introduce a new value function estimate that can only be computed at training time , the hindsight value function v+ . This value still estimates the expected return from a state st but it is further conditioned on k additional observations τ+t = st+1 , st+2 , . . . st+k occurring after time t : 2 v+ ( st , τ + t ) ≈ E [ G|S0 = st , . . . , Sk = st+k ] . ( 1 ) Furthermore , we require v+ to follow this particular parametric structure : v+ ( st , τ + t ; θ ) = ψθ1 ( f ( st ) , φθ2 ( τ + t ) ) , ( 2 ) where θ = ( θ1 , θ2 ) , which forces information about the future trajectory through some vector-valued function φ ∈ Rd . Intuitively , v+ is estimating the expected return from a past time point using privileged access to future observations . Note that if k is large enough , then v+ simply estimates the empirical return from time t given access to the state trajectory . However , if k is small and φ is lowdimensional , then φ becomes a bottleneck representation of the future trajectory τ+t . By learning in hindsight , we identify features that are maximally useful to predict the return on the trajectory from time t. The hindsight value function is not a useful quantity by itself , since – because of its use of privileged future observations – we can not readily use it at test time . Furthermore , it can not be used as a baseline either , as when computing the policy gradient it will yield a biased gradient estimator. . Instead , the idea is to learn a model φ̂ of φ , that can be used at test time . We conjecture that if privileged features φ are useful for estimating the value , then the model of those features will also be useful for estimating the value function . We propose to learn the approximate expectation model φ̂η2 ( s ) conditioned on the current state s and parametrized by η2 , minimizing the following squared loss : Lmodel ( η2 ) = Es , τ+ [ ‖φθ2 ( s , τ+ ) − φ̂η2 ( s ) ‖22 ] ( 3 ) where the expectation is taken over the distribution of states and partial trajectories τ+ resulting from that state . The approximate model φ̂ can then be leveraged to obtain a better model-based value estimate vm ( s ; η ) = ψη1 ( f ( s ) , φ̂η2 ( s ) ) . Although φ̂ ( s ) can not contain more information than included already in the state s , it can still benefit from having being trained using a richer signal before the value converges . Figure 3 summarizes the relation between the different quantities .

Summary: This paper presents a new model-based reinforcement learning method, termed hindsight modelling. The method works by training a value function which, in addition to depending on information available at the present time is conditioned on some learned embedding of a partial future trajectory. A model is then trained to predict the learned embedding based on information available at the current time-step. This predicted value is fed in place of the actual embedding to the same value model, to generate a value prediction for the current time-step. So instead of just learning a value function based on future returns, the method uses a two-step process of learning an embedding of value relevant information from the future and then learns to predict that embedding.