Value Gradient weighted Model-Based Reinforcement Learning

1 INTRODUCTION . Model-based Reinforcement Learning ( MBRL ) is a sample-efficient approach to obtain a policy for a given control problem in two stages : model learning and planning . In the model learning stage , an approximate model of the environment is learned which is then utilized in the planning stage to generate new experience without having to query the original environment . Often this process is repeated by continuously updating the model with new experience and replanning based on the updated model . MBRL is an enticing paradigm for scenarios in which samples of the true environment are difficult or expensive to obtain , such as computationally intensive simulators , real-world robots , or environments involving humans , since the model can be used to generalize the policy to unseen regions of the state space . The approach has received a lot of attention and significant progress has been made in the field ( Sutton , 1990 ; Deisenroth & Rasmussen , 2011 ; Levine & Koltun , 2013 ; Hafner et al. , 2020 ; Moerland et al. , 2020 ; Schrittwieser et al. , 2020 ) . One of the core problems of model-based policy learning methods however is that the accuracy of the model directly influences the quality of the learned policy or plan ( Schneider , 1997 ; Kearns & Singh , 2002 ; Ross & Bagnell , 2012 ; Talvitie , 2017 ; Luo et al. , 2019 ; Janner et al. , 2019 ) . Model errors tend to accumulate over time , and therefore long-term planning under approximate models can lead to suboptimal policies compared to the performance achievable with model-free approaches . This is especially prevalent in settings with complex dynamics , such as robotic environments with discontinuities , which are hard to capture with common function approximation methods . These model approximation errors are nearly impossible to avoid with current methods due to limits of function approximation in model and value learning algorithms and the use of finite datasets which can not fully capture the full distribution over dynamics functions perfectly . Hence , it is important that a model is accurate where it counts for the planning procedure , by modelling dimensions and data points more precisely which have a higher impact on the planning outcome . But this objective is not captured in most current MBRL methods , which generally use maximum likelihood estimation ( MLE ) to learn a parametric model of the environment without involving information from the planning process . The misalignment between the model learning and planning stages of MBRL has recently received renewed interest and is now commonly termed the objective mismatch of reinforcement learning ( Lambert et al. , 2020 ) , but the problems has been investigated in earlier works ( Joseph et al. , 2013 ) . Several recent papers have investigated the objective mismatch ( Abachi et al. , 2020 ; Nikishin et al. , 2021 ; Zhang et al. , 2021 ; Grimm et al. , 2020 ; 2021 ) , but currently theoretical investigation and understanding of possible approaches do not perform well when applied to complex deep learning based approaches ( Lovatto et al. , 2020 ) or the proposed approaches rely on heuristics which might not be applicable generally ( Nair et al. , 2020 ) . Summary of Contributions . We present the Value-Gradient weighted Model loss ( VaGraM ) which rescales the mean squared error loss function with gradient information from the current value function estimate . We demonstrate the advantage of the VaGraM loss over previous approaches via the analysis of the optimization behavior of the value-aware model learning framework ( Farahmand et al. , 2017 ) and form two hypotheses for the lack of empirical performance gain despite theoretical intuition : ( a ) the theory does not account for the optimization trajectory induced by the loss function and ( b ) it also does not address how to counter problems that arise when the state-space is yet insufficiently explored in early stages of the model training . Our experiments show , qualitatively and quantitatively , that the VaGraM loss impacts the resulting state and value prediction accuracy , and that it solves the optimization problems of previously published approaches . Beyond pedagogical domains , we show that VaGraM performs on par with a current state-of-the art MBRL algorithms in more complex continuous control domains , while improving robustness to irrelevant dimensions in the state-space and smaller model sizes . 2 BACKGROUND . We consider the discounted MDP setting ( S , A , p , r , γ ) ( Puterman , 1994 ) , where S denotes the state space , A the action space of an agent , p is a transition probability kernel , r : S × A → R is a scalar reward function , and γ denotes the reward discount factor . Following the standard setting of reinforcement learning , the goal is to obtain an agent which maximizes the reward function while interacting with the environment by taking actions after an optimal ( potentially stochastic ) policy π∗ ( s ) without knowledge of the true transition kernel . We will concentrate on value function based methods to solve the reinforcement learning problem . With these , the aim is to learn a function Vπ : S → R which represent the ( discounted ) reward obtained in state s by following policy π from there : Vπ ( s ) = E ( s0 , a0 , ... ) [ ∑∞ t=0 γ tr ( st , at ) |s0 = s ] . It is also helpful to define an action value functionQ ( s , a ) = r ( s , a ) +γ ∫ p ( s′|s , a ) V ( s′ ) ds′ . Many approaches ( Watkins & Dayan , 1992 ; Mnih et al. , 2013 ; Wang et al. , 2016 ; Haarnoja et al. , 2018 ) try to learn this function by minimizing the deviations of the value function approximation to a bootstrap target : minφ E [ ( Qφ ( s , a ) − ( r ( s , a ) + γ ∫ p ( s′|s , a ) V ( s′ ) ds′ ) ) 2 ] . This equation forms the core motivation for our investigation of MBRL . 2.1 MODEL-BASED REINFORCEMENT LEARNING . In the MBRL framework , an approximate model p̂ is trained from data to represent the unknown transition function p. We will use the word ‘ model ’ to refer to the learned approximation and ‘ environment ’ to refer to the unknown MDP transition function.1 We concentrate on the Dyna algorithm ( Sutton , 1990 ) and specifically investigate the impact of model errors on the planning procedure . Dyna uses a dataset D of past experiences from the environmentD = ( si , ai , ri , s′i ) Ni=0 . A parametric model p̂θ of the environment is learned by a maximum likelihood estimate using D : θ∗ = arg maxθ ∑N i=0 log p̂θ ( s ′ i , ri|si , ai ) . This model p̂θ is then used 1We limit the discussion in this paper to only the model while assuming that the reward function is either known or learned by mean squared error minimization . In all experiments , the reward function is learned using regression . to sample new next states s′model ∼ p̂θ ( ·|s , a ) to obtain better coverage of the state-action space . The samples are used to train the value function and policy as if they were samples from the environment . It is also possible to learn deterministic models , which we will denote as fθ for clarity . 2.2 KEY INSIGHT : MODEL MISMATCH PROBLEM . One of the main drawbacks of model-based reinforcement learning is the fact that model errors propagate and compound when the model is used for planning ( Schneider , 1997 ; Kearns & Singh , 2002 ; Talvitie , 2017 ) . As a simple example , assume that a sample is collected from a deterministic model and has an error . A value function based method will use the model sample to compute a biased bootstrap target r ( s , a ) + γV ( s′+ ) . The impact of the modelling error on the value function therefore depends on the size of the error and the local behavior of the value function . As an extreme example take a value function that only depends on a subset of all state observation dimensions . In this case , a large error in an irrelevant dimension has no consequence on the obtained policy , yet a maximum likelihood loss for the model can not properly capture this behavior without prior handcrafted features . We can motivate the use of MLE ( such as the mean squared error for a Gaussian model with fixed variance ) as a loss function by an upper bound : supV ∈F |〈p− p̂ , V 〉| ≤ ||p− p̂||1 supV ∈F ||V ||∞ ≤√ KL ( p||p̂ ) supV ∈F ||V ||∞ ( Farahmand et al. , 2017 ) , but this bound is loose and does not account for the geometry of the problem ’ s value function . In our example above a mean squared error would penalize deviations equally by theirL2 norm without accounting for the relevance of the dimensions . 2.3 VALUE-AWARE MODEL LEARNING . To address the model mismatch , Farahmand et al . ( 2017 ) proposed Value-aware Model Learning ( VAML ) , a loss function that captures the impact the model errors have on the one-step value estimation accuracy . The core idea behind VAML is to penalize a model prediction by the resulting difference in a value function . Given a distribution over the state-actions space µ ∈ M̄ ( S ×A ) and a value function V , it is possible to define a value-aware loss function LV ( p̂ , p , µ ) : LV ( p̂ , p , µ ) = ∫ µ ( s , a ) ∣∣∣∣ environment value estimate︷ ︸︸ ︷∫ p ( s′|s , a ) V ( s′ ) ds′− model value estimate︷ ︸︸ ︷∫ p̂ ( s′|s , a ) V ( s′ ) ds′ ∣∣∣∣2ds , a ( 1 ) and its empirical approximation L̂V based on a dataset D = ( s , a , s′ ) Ni=0 of samples from µ and p : L̂V ( p̂ , D ) = ∑ ( si , ai , s′i ) ∈D ∣∣∣∣V ( s′i ) − ∫ ( p̂ ( s′|si , ai ) ) V ( s′ ) ds′∣∣∣∣2 ( 2 ) It is worth noting that if the loss LV is zero for a given model , environment and corresponding value function , then estimating the bootstrap target based on the model will result in the exact same update as if the environment were used . However , this is rarely the case in practice ! The main problem of this approach is that it relies on the value function , which is not known a priori while learning the model . In the original formulation by Farahmand et al . ( 2017 ) , the value function is replaced with the supremum over a function space to enable analysis . While this works well in the case of linear value function spaces , finding a supremum for a function space parameterized by complex function approximators like neural networks is difficult . Furthermore , the supremum formulation is conservative and does not account for the fact that knowledge about the value function is gained over the course of exploration and optimization in a MBRL approach . Instead of the supremum over a value function class , Farahmand ( 2018 ) introduced an extension of VAML where the supremum is replaced with the current estimate of the value function , Iterative Value-Aware Model Learning ( IterVAML ) . In each iteration , the value function is updated based on the model , and the model is trained using the loss function based on the last iteration ’ s value function . The author presents error bounds for both steps of the iteration , but did not test the algorithm to ascertain whether the presented error bounds are sufficient to guarantee a strong algorithm in practice . Notably both VAML and IterVAML provide an intuitive fix to the model-mismatch problem , yet overlook two key optimization issues which lead to empirical ineffectiveness .

Summary: The paper introduces a new value-aware objective for model learning that directly builds on VAML — specifically, using a Taylor series approximation for the next state value prediction w.r.t the state under consideration. This yields a somewhat intuitive form upon simplification — essentially, dimension-wise weighted L2 distance between predicted and actual states of the model where the weights correspond to the gradient. The experimental results in restricted domains indicate that the proposed VaGram method outperforms MLE.