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: The paper proposes a way to learn better representation for RL by employing a hindsight model-based approach. The reasoning is that during training, we can observe the future trajectory and use features from it to better predict past values/returns. However to make this practical, the proposed approach fits an approximator to predict these features of the future trajectory from the current state and then subsequently, use them to predict the value. The authors claim that this extra information can be used to learn a better representation in some problems and lead to faster learning of good policies (or optimal value functions)