Value Gradient weighted Model-Based Reinforcement Learning
Authors: Claas A Voelcker
Victor Liao
Animesh Garg
Amir-massoud Farahmand
Venue: ICLR 2022 Spotlight (2022)
Keywords: model-based reinforcement learning
reinforcment learning
objective mismatch
value function
sensitivity

Abstract: Model-based reinforcement learning (MBRL) is a sample efficient technique to obtain control policies, yet unavoidable modeling errors often lead performance deterioration. The model in MBRL is often solely fitted to reconstruct dynamics, state observations in particular, while the impact of model error on the policy is not captured by the training objective. This leads to a mismatch between the intended goal of MBRL, enabling good policy and value learning, and the target of the loss function employed in practice, future state prediction. Naive intuition would suggest that value-aware model learning would fix this problem and, indeed, several solutions to this objective mismatch problem have been proposed based on theoretical analysis. However, they tend to be inferior in practice to commonly used maximum likelihood (MLE) based approaches. In this paper we propose the Value-gradient weighted Model Learning (VaGraM), a novel method for value-aware model learning which improves the performance of MBRL in challenging settings, such as small model capacity and the presence of distracting state dimensions. We analyze both MLE and value-aware approaches and demonstrate how they fail to account for exploration and the behavior of function approximation when learning value-aware models and highlight the additional goals that must be met to stabilize optimization in the deep learning setting. We verify our analysis by showing that our loss function is able to achieve high returns on the Mujoco benchmark suite while being more robust than maximum likelihood based approaches.


Published as a conference paper at ICLR 2022
VALUE GRADIENT WEIGHTED MODEL-BASED REINFORCEMENT LEARNING
Claas A. Voelcker 2, Victor Liao Animesh Garg" 12,4, Amir-massoud Farahmand
1 Vector Institute, 2 University of Toronto, 3 University of Waterloo, 4 Nvidia Correspondence c. voelckerēcs toronto. edu
ABSTRACT
Model-basee reinforcement learning (MBRL) is a sample efficient technique to obtain control policies, yet unavoidable modeling errors often lead to performance deterioration. The model in MBRL is often solely fitted to reconstruct dynamics, state observations in particular, while the impact of model eπror on the policy is not captured by the training objective. This leads mismatch between the in- tended goal of MBRI, enabling good policy and value learning, and the target of the loss function employed in practice, future state prediction. Naive intuition suggests that value-aware model learning would fix this problem and, indeed, sev- eral solutions to this objective mismatch problem have been proposed based on theoretical analysis. However, they tend to be inferior in practice to commonly used maximum likelihood (MLE) based approaches. In this paper we propose the Gradient weighted Model loss (VaGraM), novel method for value-aware model learning which improves the performance of MBRL in challenging settings, such as small model capacity and the presence of distracting state dimensions. We analyze both MLE and value-aware approaches and demonstrate how they fail to account for sample coverage and the behavior of function approximation when learning value-awane aware models. Fom this, we highlight the additional goals that must be met to stabilize optimization in the deep learning setting. To achieve this, we leverage the gradient of the empirical value function as a measure the sensi- tivity of the RL algorithm to model errors. We verify analysis by showing that our loss function is able to achieve high returns on the Mujoco benchmark suite while being more robust than maximum likelihood based approaches.
INTRODUCTION
Model-based Reinforcement Learning (MBRL) is sample-efficient approach to obtain policy for a given control problem. It solves the control optimization into two interleaved stages: model learning and plamning. In the model learning stage, approximate model of the environment is learned which is then utilized 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 unseen regions of the state space. The approach has received 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 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 emoor 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-fres approaches. This is especially prevalent in settings with complex dynamics, such as robotic environments with discontinuities, which can be hard to model with common function approximation methods. These model approximation emror are nearly impossible avoid with current methods due to limits of

---

Published as a conference paper at ICLR 2022
function approximation in model and value learning algorithms, which cannot fully capture the full distribution over dynamics functions perfectly, and the use of finite datasets.
Hence, it is important that model is accurate where it counts for the planning procedure, by mod- elling dimensions and data points that have a higher impact on the planning. 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 plan- ning 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 rein- forcement 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; Zhang et al., 2021; Ayoub et al., 2020; Grimm al., 2020; 2021; Nikishin et al., 2022), but currently theoretical investigation and understanding of possible approaches do not perform well when applied complex deep learning based approaches (Lovatto 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-Gradics weighted Model loss (VaGraM) which rescales the mean squared error loss function with gradient information from the current value func- tion 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; Farahmand, 2018) 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 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.
BACKGROUND
We consider the discounted MDP setting A,p,,) (Puterman, 1994), where denotes the state space, AA the action space of an agent, p is transition probability kernel, S x AA - R is scalar reward function, and y denotes the reward discount factor. Following the standard setting of reinforcement learning, the goal is to obtain agent which maximizes the reward function while interacting with the environment by taking actions after an optimal (potentially stochastic) policy " 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 function V, S - IR which represent the (discounted) reward
obtained state s by following policy π from there: (s) = E (s...oo....) 1>> 7 (r(s,,, as) Su = s].
is also helpful to define an action-value function Q(,, a) = r(, a) +- p(s's(s a)) (s')ds' Many approaches (Watkins & Dayan, 1992; Mnih al., 2013; Wang al., 2016: Haarnoja et al., 2018) try to learn this function by minimizing the deviations of the value function approximation to bootstrap target: ming E [(Qu a) (r(s, a) + > p(s'ls, ( This equation forms the core motivation for our investigation of MBRL.
2.1 MODEL-BASED REINFORCEMENT LEARNING
In the MBRL framework, approximate model is trained from data to represent the unknown transition function p. We will use the word 'model' to refer to the learned approximation and "envi- romment' to refer to the unknown MDP transition function.'
We concentrate on the Dyna algorithm (Sutton, 1990) and specifically investigate the impact of model errors the planning procedure. Dyna uses a dataset D of past experiences from the envi-
'We limit the discussion in this paper only the model while assuming that the reward function is either known or leamed by mean squared error minimization. In all experiments, the reward function is learned using regression.
2

---

Published as a conference paper at ICLR 2022
romment D (si, as Fa. s A parametric model pa of the environment is learned by maximum likelihood estimate using D: to sample new next states S mmddet ~ ( (-1ss, 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 fo for clarity.
arg max; Σ log po ri|s,, This model po is then used
2.2 KEY INSIGHIT: 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 error A value function based method will use the model sample to compute biased bootstrap target r(s, a) + (V(+++e).
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. dimension has no consequence on the obtained policy, yet maximum likelihood loss for the model cannot properly capture this behavior without prior handcrafted features.
this case, large error in an irrelevant
We can motivate the use of MLE (such as the mean squared error for a Gaussian model with fixed variance) as loss function by an upper bound: supver (p - P. V) S
IP أll supvez || S / (KLpplb (T) SF ||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 mean squared error would penalize deviations equally by their L2 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 es- timation accuracy. The core idea behind VAML to penalize a model prediction by the resulting difference in a value function. Given distribution over the state-action space μ and value function V, is possible to define a value-aware loss function Ly (f, P. n))
environment value estimate
model value estimate
Lv(D.p.1) f ((ss,a)
- p(5) p
d(s,a)
(I)
and its empirical approximation Lv based on dataset D = (s, of samples from and p:
v(s)) f (s, (6(s)) V(s')ds
(o,
Σ (NL
It is worth noting that if the loss Lv is zero for 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 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. While this works well in the case of linear value function spaces, finding a supremum for function space parameterized by complex function approximators like neural networks is difficult. Furthermore, the supremum formulation is conservative and docs not account for the fact that knowledge about the value function is gained over the course of exploration and optimization a MBRL approach.
Instead of the supremum over value function class, Farahmand (2018) introduced modification of VAML called Iterative Value-Aware Model Learning VAML)) where the supremum replaced with the current estimate of the value function, 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
3

---

Published as a conference paper at ICLR 2022
the algorithm to ascertain whether the presented enoor bounds are sufficient to guarantee strong algorithm in practice. Notably IterAMML provides an intuitive fix to the model-mismatch problem, yet overlooks two key optimization issues which lead to empirical ineffectiveness.
3 VALUE-GRADIENT WEIGHTED MODEL LOSS (VAGRAM)
We present Value-Gradient weighted Model loss (VaGraM), loss which is value-aware and has stable optimization behavior even challenging domains with function approximation. To motivate the loss function, we highlight two causes for the lack of empirical improvements of IterVAML over MLE based approaches. These phenomena are investigated and verified in detail in section 4.
Value function evaluation outside of the empirical state-action distribution IterVAML suffers when randomly initialized models predict next states that are far away from the current data distribu- tion or if the optimization procedure leads the model's prediction outside of the covered state space. Since the value function has only been trained on the current data distribution, it will not have mean- ingful values points outside of its training set. Nonetheless, these points can still achieve small value prediction errors if, due to the optimization process, the value function outside the training dis- tribution happens to have the same value at the model prediction as at the environment sample, We therefore require that our value-aware Ioss function should not directly depend on the value function at the model prediction, since these might potentially meaningless.
Suboptimal local minima Since the model can converge to a solution that is far away from the environment sample if the values are equal, we find that the model-based value prediction often performs poorly after updating the value function. We expect that the updated model loss forces the model prediction to new solution, but due to the non-convex nature of the VAML loss, the model can get stuck or even diverge. This especially prevalent when the previous minimum is situated outside of the empirically covered state space. A stable value-aware loss function should therefore have only one minimum in the state space" that lies within the empirical state distribution.
3.1 APPROXIMATING VALUB-AWARY LOSS WITH THE VALUE FUNCTION GRADIENT
To derive a loss function that fulfils these requirements, we start from the assumption that the differ- ence between the model prediction and the environment next states 8' are small. This is implicitly required by many MBRL approaches, since an MLE model cannot be used to estimate the next state's value otherwise. We also assume that the model has small transition noise, akin to the model assumptions underlying MSE regression, otherwise the difference between model sample and the next state sample might be large. Under this assumption, the IterVAML loss can be approximated by Taylor expansion of the value function, where we denote the expansion of V around a reference point s' as V. and obtain Vyr(s
V(s') + (V (v,V(s)le (s 8'). Using this expansion at the next state sample s/ E 'D collected from the environment for each tuple independently instead of the original value function, the VAML error can be stated as:
cp= = Σ (v(e)) -
2
(3)
(fiololor.as ((V.V()))))T(( -[)) ds'
Σ
(4)
This objective function crucially does not depend on the value function at unknown state samples, all sf are in the dataset the value function is trained on, which solves the first of our major problems with the VAML paradigm.
We can simplify the objective above even further if we restrict ourselves to deterministic models of the form & = (o, a). Since VAML requires the expectation of the value function under the model
"The full loss function will likely still admit additional local minima due to the non- con-linear nature of the model itself, hut the global optimum should coincide with the true model and the loss function should be in the state space.
4

---

Published as a conference paper at ICLR 2022
Value finetion
VAML, loss
MSE loss
VaGra.
X
X
Figure 1: Visualization of discussed loss function with regards to reference point marked with the white cross and the corresponding value function on the Pendulum environment. For the value function, darker color indicates lower value. In the loss figures, darker color indicates how large the loss is if the model predicts (0, 0) instead of the reference sample marked in white. The VAML, loss has complex non-linear shape in the state space that follows isolines of the value function, while MSE and VaGraM are centered around the sample. For VaGraM, the rescaling the MSE in the direction of high gradient along the 0 axis is visible. Due to Equation 7, the scaling is aligned with the axis of the coordinate system and not rotated fit the value function closer.
and the environment be equal, we can exchange the probabilistic model with deterministic one as long as we assume that the mean value function under the true environment is close the empirical estimate of the value function from single sample. We explore the prerequisites and consequences of this assumption further in Appendix The model loss can then be expressed as:
Σ [(VV.V())))
x))²
(5)
We can see that the objective is similar to a mean squared error regression with vector that defines the local geometry of the objective function. This vector can be interpreted as a measure of sensi- tivity of the value function each data point and dimension. In regions where the value function changes significantly, the regression incentivizes the model to be very accurate.
3.2 PREVENTING SPURIOUS LOCAL MINIMA
The formulation above retains one problem, Equation 5 does not constrain the solution for each (s, a, s') tuple sufficiently. For each (s, a, tuple, the loss function only requires that the differ- ence between the model and environment sample be orthogonal the gradient of the value function, which describes a hyperplane of solutions. These predictions can lie arbitrarily far away from the en- vironment sample, which breaks the assumption underlying the Taylor approximation that the model prediction is within a small region of the expanded state point. For more details see Appendix A.
To prevent these suboptimal solutions and achieve our second design goal, we consider upper bound on the value-gradient loss by applying the Cauchy Schwartz inequality (Ei xx)" < n Σ x to change the square of the sum with a sum of squares. We denote the diagonal matrix with vector on the diagonal as diag(a) (a) and refer to the dimensionality of the state space as dim(S) and rephrase the sum vector-matric multiplication:
Σ
(VVV(()) x()) -
(6)
S dim((S)
Σ
((o(s), - (V, (V.V(s))e))
k())
This reformulation is equivalent scaling matrix. Because the scaling matrix is positive semi-definite by design, each summand the loss is quadratic function with single solution as long as the derivative of the value function does not become zero in any component. Therefore this upper bound assures our second requirement: the loss function does not admit spurious local minima."
squared enror loss function with per-sample diagonal
We note that state dimensions ignored for points in which components of the value function become zero, potentially leading additional solutions, but in practice this rarely happens for than few points.
5

---

Published as a conference paper at ICLR 2022
Linear model
NN hidden layer
VAML
6.3
Train epochss
Trai epocks
VaGraM
MSE
VAML
Figure 2: Evolution of the VAML loss changing value functions the Pendulum domain. Lines denote the mean and shaded areas show standard error over 8 model initialization and data set samples per model. In the linear setting, VAML achieves the lowest VAML error, while VaGraM able to significantly outperform MSE. In the NN setting, VAML diverges rapidly, while VaGraM and MSE converge to approximately the same solution.
To give intuitive insight into all the discussed loss functions, visualized each one for a ped- agogical environment, the Pendulum stabilization task. The resulting loss curves can be seen in Figure The VAML loss has complicated shape that depends on the exact values of the value function while both MSE and our proposal have paraboloid shape. Compared to MSE, our pro- posed loss function is rescaled account for the larger gradient of the value function in the axis.
4 EXPERIMENT: MODEL LEARNING LOW-DIMENSIONAL PROBLEM
We compare the performance of VaGraM, with both MSE and VAML on a pedagogical environ- ment with small state space and smooth dynamics to gain qualitative insight into the loss surfaces. We use the Pendulum environment, a canonical control problem in which an under-actuated pen- dulum must be swung and stabilized to upright position. We use the implementation provided by Brockman et al. (2016). To learn the policy and its value function, we use the SAC algorithm (Haarnoja al., 2018). The original IterVAMI paper assumed that the value function obtained using approximate value iteration (AVI) (Gordon, 1995; Ernst et al., 2005; Farahmand et al., 2010). We use SAC instead of full AVI for stability in large scale experiments and discuss proper ex- tension of the VAML loss to SAC in Appendix C. We find that the difference in loss is negligible and therefore use SAC together with VAML throughout our experiments. More information on the implementation and hyperparameters of all of our experiments can be found in Appendix E.
To simplify the setup of the evaluation, we decided investigate the model losses without model- based value function learning. This allows to focus solely on the loss functions, without taking into account the inter-dependency between model and value function updates. Instead of the model- based loop, used the SAC algorithm in a model-free setup to estimate the value function. We saved the intermediate value functions after each epoch training, corresponding to 200 envi- ronment steps, and optimized the models using stochastic gradient descent the respective loss function, updating the value function used for the loss every 1000 model training steps. As the MLE loss, we used the
squared error which assumes Gaussian model with fixed variance.
To compare the optimization, we used two architectures, a linear regression without feature transfor- mations and neural network with single hidden layer and 16 neurons. We sampled a dataset uni- formly over the whole state space and used the environment transition function to compute ground truth next state samples. Finally, we evaluated each models VAML error with regards the current value function on held out dataset and plotted the results in Figure 2.
Dependency untrained value function estimates. The first cause for lacking empirical per- formance with VAML that we discussed in section 3 was that the algorithm can move outside of the covered state space region where the value function prediction is often meaningless.

---

Published as a conference paper at ICLR 2022
Layers
Layers
Layers
2000
Hewaaris
m
Train epachs
Train epochs
Train epochs
VaGraM
MLE
SAC
Figure 3: Performance of VaGraM and MLE models with reduced model size. The dotted lines correspond to the final performance reported for model-free SAC (grey, approx. 3200). Shaded area represents standard error over 16 repeated runs. VaGraM continues to solve the task almost unim- peded, while MLE is unable to even stabilize the Hopper when using two layer neural network.
In our experiment we find that linear regression model remains stable under all three loss func- tions, VaGraM, MSE and VAMI. But when using flexible function approximation, the VAML loss converges in the first iteration with the given value function, but then rapidly diverges once the value function is updated. When investigating the mean squared error of the VAML solution, we find that the model finds stable minimum of the VAML loss outside of the reachable state space of the pen- dulum. This confirms our hypothesis that flexible VAML models can find solutions outside of the empirical state space distribution, which unstable once update the value function. VaGraM remains stable even with flexible function approximation and achieves a lower VAML error than the MSE baseline when using model with insufficient capacity represent the dynamics.
Single solution convergence. In the experiment, we see that the MSE and VaGraM models con- verge to a similar solution when using neural network. This leads us to the conclusion that our loss function really only admits single solution and that this solution coincides with the mean square error prediction when the function approximation has sufficient capacity to model the dynamics function with high precision. On the other hand, the VAML network converges to solutions that are far away from the environment sample measured in the La norm and it cannot recover from these spurious minima due to the complex optimization landscape.
EXPERIMENT: MODEL-BASED CONTINUOUS CONTROL
Due to the limited complexity of the pendulum environment, the quantitative differences between the mean squared error and VaGraM are at times insignificant in this setting. The dynamics function of the environment can be approximated sufficiently well with simple neural network and one hidden layer.
The underlying theory supporting VAML states that a value-aware loss preferable to mean squared error loss in setting where the capacity of the model is too small to fully approximate the problem or the state space contains dimensions that are irrelevant for the control problem. To test whether our loss function is superior maximum likelihood approach in these cases, used the Hopper environment from the OpenAl gym benchmark (Brockman et al., 2016). As deep learning chose Model-based Policy Optimization (MBPO) (Janner et al., 2019)
based Dyna algorithm, and ran all of our experiments using the implementation provided by Pineda et al. (2021). We kept the structure of the MBPO algorithm and models and replaced the model loss function with VaGraM.
5.1 HOPPER WITH REDUCED MODEL CAPACITY
In the first experiment, we decreased the network size of the used neural network ensemble. Janner al. (2019) use fully connected neural networks with four hidden layers and 200 neurons per layer.

---

Published as a conference paper at ICLR 2022
Hopper distracting
Hopper distracting
Hopper- distracting
2000
Hewaarss
Train eppchhs
Train epochs
Train epochs
VaGraM
MLE
SAC
Figure 4: Performance of VaGraM and MLE models with distracting state dimensions. The dotted lines correspond the final performance achieved by both algorithms on the Hopper task without distraction (grey, approx. 3200). Shaded area represents standard error over 16 repeated runs. VaGraM achieves significantly higher returns than the MLE baseline, especially in the most challenging setting with 15 distracting dimensions.
To test the performance of the algorithm under smaller models, we ran tests with two and three layer networks and 64 neurons per hidden layer. The results are shown in Figure As before, when using a sufficiently powerful function approximation, we see no difference between the maximum like- lihood approach and VaGraM, suggesting that the networks are flexible enough to capture the true environment's s dynamics sufficiently close for planning. But when reducing the model size, the max- imum likelihood models quickly lose performance, completely failing to even stabilize the Hopper for short period in the smallest setting, while VaGraM retains almost its original performance.
5.2 HOPPER WITH DISTRACTING DIMENSIONS
To show that VaGraM able to achieve good performance in a setting where there are additional dynamics in the environment that do not contribute to the control problem, appended distractor dimensions to the Hopper observations. These are independent of the original environment state space and reward function, and evolve under non-linear and discontinuous dynamics (details in Ap- pendix E). A setting with distracting dimensions is known to pose difficulty for model-based control algorithms (Stone al., 2021) and neural networks struggle model non-linear, discontinuous dynamics, so with an increasing number of dimensions the task becomes harder.
The results of this experiment are shown Figure When using five distractor dimensions, both models are able to capture the dynamics sufficiently well to achieve comparable reward to the orig- inal environment. When increasing the number of dimensions, the performance of the MLE model deteriorates, more and more of its capacity is used to model the added dynamics. VaGraM con- tinues to able to achieve reward even under the presence of distracting dimensions, since the gradient of the value function with regards to the state dimensions which are irrelevant to the control problem becomes small over training. Still, the performance VaGraM also suffers with increas- ing dimensions: when adding 20 distracting dimensions, neither algorithm is able stabilize the Hopper consistently. In this case, the value function approximation cannot differentiate sufficiently between the relevant and irrelevant dimensions with the amount of environment samples provided.
In summary, we find that VaGraM is able to deal with challenging distractions and reduced model capacity significantly better than MLE baseline. This validates that our algorithm is really value- aware and can use the value function information to improve the performance of a model-based controller in settings where the model is unable fully represent the environment.
RELATED WORK
Several authors have noted on the problem of learning models that align with the goal of obtaining a good policy. The proposed approaches fall into three broad categories: value-function or policy dependency, representation learning, and data resampling.
8

---

Published as a conference paper at ICLR 2022
Inspired by VAML, Abachi al. (2020) present method that seeks to align the policy gradients under a learned model with the gradients under the true environment. Similar to our proposal. D D'Oro et al. (2020) also proposed to reweigh samples in log likelihood loss, but used policy search the reinforcement learning approach and did not account for individual state dimensions. Nikishin et al. (2022) show an approach to directly optimizing the policy performance learning a model with implicit differentiation. Asadi et al. (2018) show that the original VAML loss coincides with a Wasserstein distance in the model space under the assumption that the value function is Lipschitz smooth. We find that all of these approaches suffer from similar scaling issues as VAML and have not been shown to lead strong empirical performance outside of toy settings.
the real environment by
Grimm et al. (2020) characterize the space of possible solutions to the model learning problem under different restrictions given by value functions and policies, and propose to learn models that are value-equivalent, similar to Farahmand et al. (2017). In follow-up work Grimm et al. (2021) expand on this idea and show that their principle can be used to improve the MuZero algorithm (Schritwwiserr et al., 2020). However, these works do not discuss the optimization challenges in actually finding such value-equivalent models, they mostly characterize the space under different value functions and policies, and present an orthogonal research direction to this paper.
An alternative to characterizing the problem in the state space of the MDP are representation learning approaches, which seek to map the original states to latent representation that is more amenable to control. Such approaches include Value Prediction Networks (Oh et al., 2017), 7),, Embed-to-Conte (Watter et al., 2015) and related proposals (Levine al., 2020; Cui et al., 2021). Zhang et al. (2021) build on the idea of bisimulation metrics (Ferns et al., 2004; 2011) which seeks to characterize the difference between MDPs by finding a state mapping that is reward-invariant under policy and value function. In this work, we did not investigate learning state embeddings, but combining our work with representation learning approaches is an exciting direction for future research.
Lambert et al. (2020) hypothesize that the objective mismatch could be solved by reweighing the training buffer for model learning to prioritize datapoints with high value. These are more likely to matter for obtaining an optimal policy. A similar proposal was evaluated empirically by Nair et al. (2020). Contrary to our work however, this technique cannot account for the differing impact of the state space dimensions and scaling, since the data points are weighted whole.
7 CONCLUSION
We presented the Value-Gradient weighted Model loss (VaGraM), novel loss function to train models that model a dynamics function where it matters for the control problem. We derived our loss function from the value-awan model learning framework, showing that previous work does not account for two important optimization phenomena that appear when learning models with empirical value function approximations. We highlighted how VaGraM counters these issues and showed the increased stability of the training procedure when using our loss in a pedagogical environment.
On the Mujoco benchmark, VaGraM performs on par with maximum likelihood estimation when using large neural network models. However, introducing additional complications to the problem results in drastic performance impacts for MLE based models, which highlights the necessity for value function aware losses in challenging environments and settings in which sufficient model capacity cannot be guaranteed. In these cases, value- awareness can greatly increase the performance of Dyna algorithms by focusing the model learning procedure on relevant aspects of the state space. In future work we seek to scale our loss function to image- based RI, where relevant state space dimensions can vary over task due to shifting camera angles. Furthermore, we seek to derive a related value-aware approach for partially observable domains that can take the state inference problem into account.

---

Published as a conference paper at ICLR 2022
ETHICAL CONCERNS AND LIMITATIONS
The proposed value-aware model learning approach is designed general purpose solution to the model mismatch problem in MBRL. While the reinforcement learning paradigm as whole has important ethical ramifications, especially in settings where automated decision making affects humans directly, we do not address these concerns specifically in our paper as our method focuses on algorithmic problems that are orthogonal to the question of proper reward design. We restrict our proposal to cases in which the reward design actually captures the intended task, which is common, yet rarely addressed, assumption in the RI. literature.
REPRODUCIBILITY
We provide an open-source version of code at httpp::/gggthhii cwwppairlab /vaab//aaaam....mm Furthermore we document all details on the implementation and evaluation setting in Appendix E and describe all necessary components of our loss in the main text. We also used the open source MBRL-Lib implementation for MBPO (Pineda et al., 2021) the basis for our code and for the evaluation of baselines and provide our model loss as an additional module in the framework for easy replication.
ACKNOWLEDGMENTS
AG and AMF acknowledge the funding from the Canada CIFAR AI Chairs program, as well as the support of the Natural Sciences and Engineering Research Council of Canada (NSERC) through the Discovery Grant program. AG is also supported by the University of Toronto XSeed, LG and Huawei. We thank members of the PAIR lab and AMF group for feedback on the paper and help with running the experiments. We are grateful to the anonymous reviewers for their construc- tive feedback and the valuable rebuttal discussion. Resources used preparing this research were provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, and com- panies sponsoring the Vector Institute.
10

---

Published as a conference paper at ICLR 2022
REFERENCES
Romina Abachi, Mohammad Ghavamzadeh, and Amir-massoud Farahmand. Policy-aware model learning for policy gradient methods. ArXiv, 0bs/20033000334 2020.
Kavosh Asadi, Evan Cater, Dipendra Misra, and Michael L Littman. Equivalence between wasser- stein and value-aware loss for model-based reinforcement learning. ArXiv, 0abs//806.ll 01..,,116,,, 2018.
Alex Ayoub, Zeyu Jia, Csaba Szepesvari, Mengdi Wang, and Lin F. Yang, Model-based reinforce- ment learning with value-targeted regression. In International Conference on Machine Learning, 2020.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. ArXiv, (oss/660600154,, 2016.
Brandon Cui, Yinlam Chow, and Mohammad Ghavamzadeh. Control-aware representations for model-based reinforcement learning. In International Conference on Learning Representations, 2021.
Mare Deisenroth and Carl E Rasmussen. Pilco: A model-based and data-efficient approach to policy search. In International Conference on Machine Learning. 2011.
Pierluca D'Oro, Alberto Maria Metelli, Andrea Tirinzoni, Matteo Papini, and Marcello Restelli. Gradient-aware model-based policy search. In AAAI Conference on Artificial Intelligence, 2020.
Damien Ernst, Pierre Geurts, and Louis Wehenkel. Tree-based batch mode reinforcement learning. Journal of Machine Learning Research, 6:503-556, 2005.
Amir-massout Farahmand. Iterative value-aware model learning. In Advances in Neural Informa- fion Processing Systems, volume 31, 2018.
Anir-massoud Farahmand, Csaba Szepesvári, and Rémi Munos. Error propagation for approximate policy and value iteration. Advances in Neural Information Processing Systems, 2010.
Anir-massoud Farahmand, André Barreto, and Daniel Nikovski. Value-Aware Loss Function for Model-based Reinforcement Learning. In International Conference on Artificial Intelligence and Statistics, 2017.
Norm Ferns, Prakash Panangaden, and Doina Precup. Metrics for finite markov decision processes. In Uncertainty AI, 2004.
Norm Ferns, Prakash Panangaden, and Doina Precup. Bisimulation metrics for continuous markov decision processes. SIAM Journal on Computing, 40(6), 2011.
Geoffrey J Gordon. Stable function approximation in dynamic programming. In International
Conference on Machine Learning, 1995.
Christopher Grimm, André Barreto, Satinder Singh, and David Silver. The value equivalence prin- ciple for model-based reinforcement learning. In Advances in Neural Information Processing Systems, 2020.
Christopher Grimm, André Barreto, Gregory Farquhar, David Silver, and Satinder Singh. Proper value equivalence. In Advances in Neural Information Processing Systems, 2021.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International Confer- ence on Machine Learning. 2018.
Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representations, 2020.
Michael Janner. Justin Fu. Marvin Zhang, and Sergey Levine. When to trust your model: Model- based policy optimization. In Advances in Neural Information Processing Systems, 2019.
11

---

Published as a conference paper at ICLR 2022
Joshua Joseph, Alborz Geramifard, John W Roberts, Jonathan P How, and Nicholas Roy. Reinforce- ment learning with misspecified model classes. In IEEE International Conference on Robotics and Automation, 2013.
Michael Kearns and Satinder Singh. Near-optimal reinforcement learning in polynomial time. Ma- chine learning, 49(2):209 232, 2002.
Nathan Lambert, Brandon Amos, Omry Yadan, and Roberto Calandra. Objective mismatch in learning. In Conference Learning for Dynamics and Control, 2020.
Nir Levine, Yinlam Chow, Rui Shu, Ang Li, Mohammad Ghavamzadeh, and Hung Bui. Predic- tion, consistency, curvature: Representation learning for locally-lineat control. In International Conference on Learning Representations, 2020.
Sergey Levine and Vladlen Koltun. Guided policy search. In International Conference Machine Learning. 2013.
Ángelo G. Lovatto, Thiago P. Bueno, Denis D. Mauá, and Leliane N. de Barros. Decision-awes aware model learning for actor-critic methods: When theory does not meet practice. In "I Can' Believe It's Not Better!" at NeurIPS Workshops, 2020.
Yuping Luo, Huazhe Xu, Yuanzhi Li, Yuandong Tian, Trevor Darrell, and Tengyu Ma. Algorithmic framework for model-based deep reinforcement learning with theoretical guarantees. In Interna- tional Conference on Learning Representations, 2019.
Michael Lutter, Leonard Hasenelever, Arunkumar Byravan, Gabriel Dulac-Amold, Piotr Trochim, Nicolas Heess, Josh Merel, and Yuval Tassa. Learning dynamics models for model predictive agents. ArXiv, 100///199.14311, 2021.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, loannis Antonoglou, Daan Wier- stra, and Martin Riedmiller. Playing atari with deep reinforcement learning. In NeurlPS Deep Learning Workshop. 2013.
Thomas M. Moerland, Joost Broekens, and Catholijn M. Jonker. Model-based reinforcement learn- ing: A survey, ArXiv, abs/2006. 16712, 2020.
Suraj Nair, Silvio Savarese, and Chelsea Finn. Goal-aware prediction: Learning model what matters. International Conference on Machine Learning, 2020.
Evgenii Nikishin, Romina Abachi, Rishabh Agarwal, and Pierre-Luc Bacon. Control-oriented model-based reinforcement learning with implicit differentiation. In AAAI Conference Ar- tificial Intelligence, 2022.
Junhyuk Oh, Satinder Singh, and Honglak Lee. Value prediction network. Advances in Neural Information Processing Systems, 2017.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vetoto Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems 32. 2019.
Jan Peters and Stefan Schaal. Natural actor-critic. Neurocomputing, 71(7-9): 100:11180-11110 1190, 2008
Luis Pineda, Brandon Amos, Amy Zhang, Nathan O. Lambert, and Roberto Calandra. Mbrl-lib: A modular library for model-based reinforcement learning. ArXiv, 10s/2104.10199) 2021.
Martin L. Puterman. Markov decision processes: Discrete stochastic dynamic programming. In
Wiley Series in Probability and Statistics, 1994.
Stéphane Ross and Drew Bagnell. Agnostic system identification for model-based reinforcement learning. In International Conference on Machine Learning, 2012.
Jeff Schneider. Exploiting model uncertainty estimates for safe dynamic control learning. In Advances in Neural Information Processing Systems, 1997.
12

---

Published as a conference paper at ICLR 2022
Julian Schrittweser, loannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with learned model. Nature, 588(7839):604-600 2020.
Austin Stone, Oscar Ramirez, Kurt Konolige, and Rico Jonschkowski. The distracting control suite a challenging benchmark for reinforcement learning from pixels. ArXiv, 0tls2101100222/ 2021.
Richard S Sutton. Integrated architectures for learning, planning, and reacting based on approximat- ing dynamic programming. In Machine learning Proceedings, pp. 216 224. 1990.
Erin Talvitie. colf-corg models for model-based reinforcement learning, AAAI Conference Artificial Intelligence, 2017.
Ziyu Wang, Tom Schaul, Matteo Hessel, Hado Hasselt, Marc Lanctot, and Nando Freitas. Dueling network architectures for deep reinforcement learning. In International Conference on Machine Learning. 2016.
Christopher JCH Watkins and Peter Dayan. Q-learning. Muchine learning, 8(3-4), 1992.
Manuel Watter, Jost Tobias Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. In Advances in Neural Information Processing Systems, 2015.
Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learn- ing invariant representations for reinforcement learning without reconstruction. In International Conference on Learning Representations, 2021.
13

---

Published as a conference paper at ICLR 2022
A BOUND BETWEEN VALUE-AWARE MODEL LEARNING AND MALUE-GRADIENT WEIGHTED MODEL LOSS
The error in Taylor approximation R(V, 8 [o(s, a)) is bounded by I1s' (o(s,a)) a) with M dependent on the Hessian of the value function. Plugging this into the VAML. loss and assuming worst case approximation errors, we obtain an upper bound on the VAML error:
E ,
((o(s,,))
(8)
= -E
-() -)) R(,,s,
(9)
+11+ (IVV.(())))((666..)) - - [R(V, $6 fo(s, a)))²
(10)
2
(IV.V( - M
2 | - sa(s,a)11
(11)
2 ||s / fo fo(s,a))**
< Eg' [(V.V(()))))(666s... -1))²] + Eg a
(12)
Our experiments show that we treat M like a tuneable hyperparameter, we obtain worse perfor- mance when optimizing this upper bound compared to VaGraM. The Hessian parameter is difficult to compute or estimate in practice and we find that most often, either the first or the second loss component will dominate when choosing heuristic values.
B ANALYZING THE ADDITIONAL LOCAL MINIMA OF THE TAYLOR APPROXIMATION LOSS
We noted in the main paper that the direct Taylor approximation the value function leads to a spurious local minimum. This is clear when looking at the loss for single datapoint:
min ((V (
(13)
dimi S)
min
Σ
(V.(()))<)) fo(s,a)n
(14)
Assuming that f is flexible and can predict any next state s' (i.e. by choosing f = 0), the optimal solution obtained from an undetermined lincar system of equations. This system admits far more solutions than either the corresponding IterVAMI loss will achieve arbitrary large value prediction errors. In fact, the equation describes a hyperplane of minimal solutions consisting of every weight vector that is orthogonal to the gradient of the value function at the reference sample, with dim(S) - free variables. Therefore we need to enforce closeness of the model prediction and the environment sample, since the Taylor approximation is only approximately valid in close ball around the reference sample.
mean squared error, and many of them
One way achieve this closeness is by adding the second order Taylor term, which results in an additional MSE loss term. As pointed out in Appendix A, we did not achieve good performance when testing out this version, since it is difficult to compute the Hessian in higher state spaces and heuristically choosing a value as a hyperparameter proved to difficult to tune in practice. Therefore, approached the solution to this problem outlined in the paper.
C FULL VAML FOR SAC
The formulation of VAML which derive for SAC is direct extension of the VAML loss to the SAC soft bellman backup. Specifically, given some distribution over the state-action space D, state-action value function Q and policy n, can define SAC-aware loss.
14

---

Published as a conference paper at ICLR 2022
Carppll
VaGraM
MIPO
VaCrM
BerVAMI
MEPO
Value weighted
VAML
SAC
SAC
Traio
Thaie
Figure 5: Comparison of VaGraM, MLE (MBPO baseline), IterVAML Farahmand (2018) and value-weighing ablation on two simple continuous control tasks. The shaded area represents stan- dard error estimated over 8 runs. While VaGraM, IterVAML and MBPO able to achieve satis- factory performance on the Pendulum Swingup task, VAML fails to stabilize the more difficult Cartpole balancing task.
/ ou -log logs(a'ls'))ds dD((,,)
(x /Vʃ -(())1,)))p(((1__00)))
(15)
((s'(ss,0)) V (s', d'sdss
/ (p(5'(s,a) p(s a)) E ((())) [loss(a'lss)]] ds' dD(s,a)
(16)
as well as its sample-based version:
(p. P.
Σ (T)) - as
(f (o - p(6 ai)aaa) log (og=ad''s')da'a dd''d' ds') )
(17)
We find that in most of our experiments, the entropy terms did not significantly contribute to the loss. Therefore we dropped it in our experiments and directly derived VaGraM for the value function prediction error only. This makes our loss applicable in all situations were the model used to estimate the value function of a policy.
Finally, we not account for the dependency of the policy function update on the model. Since SAC aims to minimize the KL divergence between the action distribution and the Gibbs distribution defined by the value function directly, only way model influences this update is via the induced state space distribution. We do not address this matching explicitly in this work, but hybrid policy-aware and value-aware model is an enticing direction for future research.
D ADDITIONAL EXPERIMENTS
D.1 ABLATIONS
To test the performance of VaGraM and MBPO against alternative models, used the classic Pendulum swingup and Cartpole benchmarks. As points of comparison, we used IterVAML (Farah- mand, 2018) and simple value-weigg value-weighted regression similar Nair al. (2020), where the MSE emror is multiplied by the inverse of the value function of the sample. The results are visualized in Figure 5. Hyperparameters follow the Cartpole task baseline in Pineda et al. (2021).
Since IterVAML suffers from strong destabilization on the Pendulum model learning problem, as discussed in section added an MSE loss the original formulation Lxiim LeerVAMI XLMSE
15

---

Published as a conference paper at ICLR 2022
VitGrabi
SAC
VsGmM MBPO
SAC
Tealn epechhs
Train spodts
Halliteetah
HalCChottaz UP composison
VaGraM I strg
13300
MBPO
VaGraM 5 stop
MBPO
VaCrra
МHРО
SAC
Train spocha
Train spedia
Figure 6: Comparison of VaGraM and MBPO on the Mujoco tasks presented in Janner et al. (2019). We see that VaGraM is able to perform on par with the MBPO implementation on all tasks, while outperforming both the model-based and model-free baseline on the Ant environment, For the HalfCHeetab-v2 task, different performances emerge when using different rollout lengths. While the MLE baseline is not able to benefit from longer rollouts, VaGraM increases in performance, although the statistical significance of this increase is low due the strong bimodality of solutions on the HalfCheetah task.
with a tradcoff' factor of = 0.01. We see that this is sufficient to stabilize the performance in the Pendulum swingup task and prevent catastrophic divergence, but the model still performs signifi- cantly below the MBPO and VaGraM results in the Cartpole stabilization task. This evidence that with additional stabilization mechanisms, the issues presented in section 3 prevents the easy adoption of VAML to complex domains.
The veluh-weighing baseline is unable achieve satisfactory performance dulum task, leading us to conclude that more complex reweighing schemes such as VaGraM are indeed necessary.
the simple pen-
D.2 PERFORMANCE IN MUJOCO BENCHMARK SUITE
The comparison runs on the Mujoco benchmark environments presented in Figure 6 are faulty due to a bug that was pointed out to us at the poster session. The bug only harms the performance of VaGraM, so we do not see as a reason to withdraw the paper, but due to limited time and computational resources, we were unable to repeat the full comparison presented here. All other experiments are corrected and to the best of our knowledge reflect the performance VaGraM. We will update the paper with updated comparisons once they are finished.
The results of VaGraM and MBPO are presented in Figure 6. We find that VaGraM is able to perform par with the MBPO baseline all the tasks. Due to stability issues, the Humanoid-v2 task excluded, no algorithm achieved satisfactory performance (compare Pineda et al. (2021) for discussion).
On the Ant-v2, we see small performance improvements above the results reported by Janner et al. (2019) and Pineda et al. (2021). We hypothesise that for this task, the canonical state space rep- resentation could be misaligned with the control problem in the sense that not all dimensions are informative for control. A further investigation of the phenomenon would be interesting direction
16

---

Published as a conference paper at ICLR 2022
for future research. All comparisons done the hyperparameter settings provided by Pineda et al. (2021).
Investigating the performance differences further, find that the HaliChectah-v2 domain, Va- GraM is able to profit from longer rollouts, while MBPO does not increase in performance when rolling out the model for more than step. The comparison is presented in Figure 6. The MLE based model completely destabilizes if trained beyond 180 200 epochs on longer rollouts, which is why we truncated the training length before convergence. The performance of VaGraM with longer rollouts is already beyond the final performance of the fully trained MBPO baseline after 180 steps.
The destabilization leads complete collapse of the policy performance, probably due increas- ing uncertainty of the model and in some cases to a gradient explosion that produces NaN values in the policy. We find that VaGraM does not suffer from this pattern and is able slightly outperform both MBPO and VaGraM trained on single step rollout. Similar patterns do not emerge in the Walker-v2 domain, here training on longer rollouts destabilizes all algorithms. We leave a more depth discussion of the impact of rollout length on policy performance under different models and losses for future work but highlight that significant performance might be gained by finding better tradeoffs than those discussed in Janner et al. (2019), especially when comparing deterministic and probabilistic models (compare Appendix F) as well as aware and value- agnussii models.
E EXPERIMENT DETAILS
As described in the main paper, we conducted our experiment in two domains from the OpenAI benchmark OpenAl Gym (Brockman et al., 2016), Pendulum-v2 and Hopper-vO. We used both environments as provided by the framework without further modification in our capacity tests. Our algorithm is shown in pseudocode algorithm
Expanded Hopper enviromment To test the performance of VaGraM in a setting with additional unnecessary state observations, we created random dynamical system that evolves according to randomly initialized dynamics function, where is a fixed matrix with entries randomly drawn from N (0., 10.), and a fixed initial state 8p with components randomly drawn from (O.,, 0.1). We do not resample the fixed components at reset.
The transitions then described by the following deterministic function:
Stt + sin(As) if [s(s, ar) 20 (90
f(s,
(18)
else
This dynamics function contains two attributes which are hard for neural networks to model: dis- continuity and non-linear dynamics. We find that this is sufficient to provides very challenging environment for the MLE trained neural network models. To account for varying difficulty over different random initialization of the environments, we made sure to the comparison runs with the same seeds over the different loss functions, but we did not find that the inter-seed variance was larger than the observed difference between the different loss functions.
E. ARCHITECTURE AND HYPERPARAMENT
For the Pendulum experiments, we use a simple fully connected neural network with single layer, and linear regression without feature transformations architectures. The used non-linearity is ReLU. All experiments were implemented the standard initialization of the library were kept. The exact versions of all used libraries documented in the provided source code.
PyTorch (Paszke al., 2019) and were not specified,
To assure fair comparison we used the hyperparameters provided by Janner et al. (2019) for all experiments with our approach and the NLL loss function used for the baseline. All models used were fully connected neural networks with SiLU aonlinnarssies and standard initialization. For all experiments with full model size we followed MBPO and used seven ensemble members with four layers and 200 neurons per layer.
17

---

Published as a conference paper at ICLR 2022
Algorithm 1: Value-Gradien weighted Model learning (VaGraM)
Initialize policy Td. value function Ugj, model fe. environment dataset D env, model dataset moddlli
for N epochs do
while pa not converged do Sample batch (8, a, r, 8') from Denvi
L.. = -((s' fol-so(s, T (+))))) (s' a))o: (o((oa))))
0 - ar 20 f
Train reward model
end
for E steps do
Take action in according Thi add Densi for M model rollouts do
Sample from Denvi sample (s);;;
8', r' fe(s,a), Add (s,a,r,s') Pmodal end for policy gradient updates do
Sample batch (s, a, r,,)'s 8') from Dew U dp (Up (8) (r yr'(ss))))); a I 2 (s)),
end
end
end
Even though our loss derivation does not support rolling out the model for more than a single step without accounting for this in the training setup, we find that VaGraM is still stable over the short rollout horizons used by MBPO. Therefore we used the rollout scheme from MBPO.
However, it was necessary to make a small alteration to the training setup: in the provided implemen- tation, the value function solely trained on model samples. Since our model is directly dependent on the value function, need to break the inter-dependency between model and value function in the early training iterations. Hence, we used both real environment data and model data to train the value function, linearly increasing the amount of model samples from 0 to 95% of the SAC replay buffer over the first 40 epochs of training (corresponding to 40.000 real environment steps). We did not transition fully using model data assure that the real environment samples to inform the value function learning. We found that this did not diminish the training returns of MBPO compared solely using model samples and so used this approach for both VaGraM and MLE.
still able
To estimate the gradient of the value function, we used the four empirical value functions, two direct estimates and two target value functions used in the SAC algorithm and summed the losses using each data tuple and value function gradient independently. Furthermore we calculated the Lę norm of all value function gradients and clipped these at the 95-th percentile. We found that this was necessary since the empirical value function gradients can get very sharp, which in rare cases leads to destabilization of the gradient descent algorithm used to update the model.
F DETERMINISTIC VS PROBABILISTIC MODELS
Our derivation of VaGraM lead us to the conclusion that a deterministic model was sufficient to achieve the goal of value-aware model learning in environments with small transition noise. This insight stands in contrast to the current literature, which often claims that probabilistic models are needed to achieve optimal performance in model-based reinforcement learning. However, there is no clear consensus among different authors whether probabilistic models are needed or a deter- ministic model can be sufficient for MBRL (compare Lutter al. (2021)).
18

---

Published as a conference paper at ICLR 2022
Limited capacity 2
Distracting dimensious:
VaGraM
VaGraM
MLE, probabilistic
MLE prohabilisttc
MLE deterministic
MLE deterministic.
Hewward
SAC
No distraction
1000
311
XX0
Troin epochs
Train epocla
Figure 7: Comparison of the empirical performance of deterministic and probabilistic MLE models VaGraM. Thick lines denote the mean and shaded area the standard error over 8 runs. In the limited capacity setting, the deterministic model is able to achieve significantly higher re- turns than the probabilistic baseline, however with stability issues during longer training. On the distracting benchmark, the deterministic model is not able to achieve any significant returns.
Our assumption that the model be replaced with deterministic relies on the assumption that the underlying model is not dominated by stochastic transitions, but is deterministic or near deterministic and unimodal. This follows from the requirement that the model prediction admits small error in the mean squared error sense, otherwise the Taylor approximation does not properly capture the behavior of the value function, as model prediction might have high likelihood under the environment and still be far away from the environment sample in mean squared sense other- wise. domains where capturing the stochasticity is erucial, we have to revisit this requirement in follow up work.
Furthermore, we are operating under the assumption that the mean value function of each distribu- tion over states can be represented as the value function of single state prediction. Since value function approximations represented by neural networks are continuous, this is true in our setting due to the mean value theorem for integrals. Implicitly due to the Taylor approximation, assume that this mean value lies close to or on the environment sample we obtained. If other approx- imation schemes used, this property needs to be checked and potentially probabilistic models needed to represent the expectation over all possible value functions and transition dynamics.
also
F.1 A ABLATION EXPERIMENT WITH DETERMINISTIC MODELS
In our experiments in the Hopper domain, we used probabilistic models following Janner et al. (2019). To verify that the improvement in performance did not result from using deterministic model instead of a probabilistic model, we repeated the experiment showing the impact of smaller model sizes, and replaced the probabilistic Gaussian ensemble with an ensemble of deterministic functions trained with the mean squared error loss between sample and state prediction. All other implementation details, architecture and hyperparameters were kept fixed. The results are shown in Figure 7. In this ablation, we do indeed see that a mean squared trained model capture the dynamics information better than model that is trained using negative log likelihood. To the best of knowledge, this phenomenon has not received attention in the literature, but we hypothesize that it is an artifact of training probabilistic model with gradient descent instead of natural gradient descent (compare Figure in Peters & Schaal (2008) for a visual intuition).
Nonetheless, VaGraM is still able achieve higher cumulative reward consistently than the MSE model. Especially on the distraction task we find that a deterministic MSE model performs on par with the probabilistic model and fails achieve any reward when faced with a challenging number of distractions. This validates our hypothesis that value awareness is important in settings with insufficient model capacity, but crucial in cases where the environment observations are not aligned with the control problem.
19


Reviews:
[{"review_id": "ic3HTSOrA-O", "reviewer": "Reviewer_dE85", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## General Comments\n- **writing quality**: This is a well written paper. Throughout, it was easy to follow where the authors are going and why.\n- **technical novelty**: The VaGraM method seems well motivated. The related works and introduction of the objective mismatch issue are well done. The mathematics are easy to follow throughout.\n\nA question I had / an intuition that I am curious how the authors would respond to: when training dynamics models, in my experience the models are pretty accurate on average and the error follows a normal distribution (with some points being very inaccurate), do you think your loss function and resulting algorithms would be able to handle these outlier error points? I think the Taylor approx. could get weird with big prediction errors?\n\n- I looked at the code and am familiar with the library the authors use. Everything they mention seems to be there, but I did not take the time to run all the experiments. \n- **experimental validation**: Due to a lot of other things the authors had to explain, there isn't a lot of room for experiments in this paper. Generally, the experiments support the authors claims, but some of them are somewhat questionable motivated (ie adding the irrelevant states in state-based MBRL). I've added some more questions on the experiments below, as they form the central way for the authors to improve the score of the paper.\n\nexper.1: the proposed \"challenges\" for MBRL do not seem to match up with the literature. Can the authors comment more on the motivation for small model capacity and distracting state dimensions? \n\nexper.2: small but crucial change -- it is near impossible to see the VaGraM line in figure 2. I saw it by zooming in a ton, but printing the paper would render it invisible. Please add markers to this and preferably all plots.\n\nexper.3: I think of a better test for models handling distracting states is to try it on high-dimensional tasks like humanoid. The paper the authors references (Stone et al 2021) focuses on pixel-based RL, which seems very different than what the authors are doing. Does this problem come up more in the literature? If the authors showed their loss function outperforming MBPO on a task like humanoid that would make me fight hard for its acceptance. \n\nMore detailed questions on the experiments are below, but these are the main points. \n\n## Comments By Section\n\n### 1. Intro\n- 1.1 do the authors have any citations for discontinuities being hard to model? It seems like something that is folk-wisdom but unproven?\n- 1.2 I really like the intro of the paper, good job authors. I would potentially add a hint that VaGraM uses value function gradients in the abstract. Now it just says value-aware learning. Maybe that's just preference!\n- 1.3 I think the last sentence of the contributions is oversold. At least in the paper and in a limited read of the appendix, the experiments are a little limited in breadth to say it is on par with existing algorithms. It's on par on a couple tasks with one algorithm. It is certainly promising, but that didn't quite match with the experiments. \n\n### 2. Background\n- 2.1 only comment is that I think the sum across i in equations 1,2, and on could be better described. Is that the sum across the state-dimension? \n\n### 3.Value-Gradient Weighted Model Loss\n- 3.1 The core issue for me in this section is figure 1. Either it seems like the VAML loss matches the shape / the color scales for higher is better / something is making the results hard to read. Are all the learning solutions showing loss that is too high around the area of interest? Why is value lower outside? Is there a simpler explanation to all of this - I don't see the VaGraM method being better as described in the text. \n- 3.2 is there a dropped negative sign in equation 4? Otherwise, how does the subtracted integral become positive? I see that one of the terms cancels out.\n- 3.3 in section 3.2 how is this bound implemented? Is this just a numerical bound on what the loss function can achieve or an implementation trick? That wasn't totally clear to me. Is there an experiment in the appendix showing the effects? \n- 3.4 what constitutes a \"small prediction error\" for the Taylor expansion to hold? Any examples to give on a simulated task?\n- 3.5 the authors speaking about dynamics models \"as there is no guarantee to which local optimum the algorithm will converge\" is there any reason to think that the \"optimum model\" wrt to loss will be helpful, in the light of objective mismatch? \n- 3.6 can the authors clarify footnote 2, this was confusing to me.\n\n### 4. Example: Cartpole\n- 4.1 AVI is used but not defined.\n- 4.2 Is the source of the value function across domains important? Ie the difference between IterVAML and SAC for generating it?\n- 4.3 I just wanted to confirm that the dataset generated and environment transition function used to generate next samples is to get the true value function for pendulum? End of second to last paragraph p6. \n- 4.4 the authors say \"a model with insufficient capacity to represent the dynamics.\" Isn't the standard pendulum task started vertical, then the dynamics are expressible by a linear function as long as the system is pretty stable?\n\n### 5 High-dimensional\n- 5.1 if the cited paper is for pixel based RL? What is the motivation here? The result is interesting, though it seems fairly limited in scope when it is easy to make these relatively small models bigger?\n- 5.2 the change by layer including a smaller layer size of 64 neurons is weird. What happens with only changing the amount of layers and having a width of 200 for all (4,3,2 layers)? As said above, what happens with the algorithm on Humanoid? \n\n## Nits / grammatical comments\n- n1 the abstract on open review has latex commands that didn't work, I am not sure if you can fix them :) \\algoNameFull (\\algoName)\n- n2 end of page 5, \"this , it\" is missing a word\n- n3 in non-generalizing local minima section, this reads weird \"Even worse, there examples where...\"\n- n4 in section 2.1 MBRL review,  the idea of an \"unknown true MDP's transition\" is a little weird. I get it, but non-experts may be confused \n\n## Other comments / ideas (not included in scoring of paper)\n- o1 I appreciate the author's comment on the ethical concerns regarding reward hacking and RL.\n- o2 I would be very interested in how the \n- o3 are the claims regarding \"Single solution convergence\" true across seeds? What would this look like? Is there interesting insight to be gained there?\n", "summary_of_the_paper": "The authors propose a new loss function for training dynamics models in MBRL that is aware of the environmental value. This loss function is a tool to mitigate the \"objective mismatch\" issue and seems well motivated. The authors show the usefulness of their VaGraM in simple experiments.\n", "summary_of_the_review": "A well written and motivated paper needs a bit of clarification on the motivation and execution of experiments. The motivation makes me lean towards accept, but I would appreciate if the authors fixed or clarified some of the limitations.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Y7I3QprUOB2", "reviewer": "Reviewer_kWH6", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Strengths:\n\n1. The paper is well-written. It has a clear motivation for improving upon the VAML model loss empirically and presents two important shortcomings of the VAML objective. The second issue on unstable local minimas is also highlighted empirically in the Inverted Pendulum environment.\n\n2. The proposed VaGraM loss is designed to specifically remedy the two presented issues in the original VAML loss. The derivation is sound and the final loss function is intuitively appealing.\n\n3. In the InvertedPendulum environment, two valuable empirical insights are presented: (a) Qualitatively (Figure 1), the proposed loss landscape is shown to be similar to MSE in it’s parabolic or “single-minima” nature, while also being sensitive to the value function landscape. This is in contrast to VAML whose loss landscape does not admit a single optimum. (b) In a controlled setting where value function and policy is selected at various stages of training of a model-free approach for the purpose of learning a model independently (with periodic value updates), the VAML loss is shown to converge to minimas that are unstable and cause the VAML loss to diverge significantly once the value network is updated (particularly in the NN - 1 hidden layer case). This supports the earlier presented hypothesis about the second shortcoming of VAML.\n\n4. In the complex continuous control MuJoCo domain, the Hopper environment is used to present two key benefits of the proposed objective over MLE: (a) it retains performance with reduced model capacity and (b) it is comparatively more robust to MLE with increasing distractor dimensions in the state space. The latter experiment supports the main motivation for using VAML -- the fact that it is theoretically supposed to ignore parts (or dimensions) in the state space that are irrelevant for the value function.\n\n\n## Weaknesses\n\n1. The empirical analysis in the paper leaves a lot to be desired for gaining a thorough understanding of the benefits of VaGraM.\nFor the complex continuous control environments, why was Hopper the only environment selected? It would be good to see the performance of the proposed model loss in other MuJoCo environments as well, especially HalfCheetah, Walker and Ant in which the open-source MBPO implementation by Pineda et al. (2021) is shown to work well. It is okay to skip Humanoid because they haven’t gotten it to work, but the others should be included. These environments need not be used for all the experiments on model-capacity and distracting dims -- just showing a benchmark of reward curves on each unmodified environment would be sufficient.\n\n2. Why is the VAML baseline not present in the Hopper experiments in Section 5? Even if it performs poorly, including it is important for the sake of completeness (as the earlier InvertedPendulum experiments do include the VAML baseline).\n\n3. The experiments on the Inverted Pendulum environment provide valuable insights. However, they do not include a comparison of the proposed method vs MLE and VAML trained end-to-end in an MBRL algorithm (i.e. the opposite of the controlled experimental setting where the policy, value and model all interact and influence each other). Such an experiment has been presented for the more difficult Hopper environment (left-most plot of Figure 3), but it is not presented for the easier InvertedPendulum setting.\n\n4. In the InvertedPendulum experiment represented by Figure 2, Tte paper shows that with the selected value update frequency, the VAML error diverges in the NN setting. I have two questions regarding this observation -- (a) Why not include a baseline that performs gradient clipping to prevent the exploding losses in VAML?, (b) Why not vary the value update frequency from once every 1000 model updates to a different value (possibly in conjunction with gradient clipping from my previous point), to test whether the VAML loss still diverges with higher value update frequency? Given these potential remedies, the presented second issue with VAML (“Non-generalizing local minima”) may actually be a non-issue, i.e. it may be an artificial problem that can be overcome in a simpler manner.\n\n5. Given that the taylor expansion in Eqn 4 only uses a linear approximation, this approximation is only accurate when consecutive states (s and s’) in the environment are close to each other. This assumption is in addition to the assumption presented in the paper that the model prediction should not be too different from the environment’s next state. It seems that this requirement for consecutive states to be close to each other is naturally met in most continuous state space environments such as the ones chosen in this paper. However, it may be the case that other environments have sharp state space transitions (or discrete state spaces) where the linear approximation may no longer hold. The paper should comment on the requirement of this assumption and when it may be violated, and whether or not it holds in the selected environments.\n\n## Factual Errors/Corrections\n\n1. The conclusion mentions “MuJoCo benchmark” -- this is not accurate as only the Hopper environment was selected.\nSection 5.2, last line: “improve the model learning performance” -- Here, ‘model learning performance’ is not defined. If it refers to the policy performance on Hopper as a result of the proposed model learning loss, then that is not really improved in the original unmodified architecture and state space setting, so this may be misleading. If it refers to actual performance of the model, then what is the metric? Is it the VAML error?\n\n1. Section 5 title: “High-Dimensional…” -- The selected Hopper environment isn’t high dimensional. This is misleading as it suggests that the environments use high dimensional state spaces such as RGB images as opposed to the joint and position information of the simulated MuJoCo robot.\n\n## Other Issues/Comments:\n\n* Figure 1: The “Ours” plot has a darker region that does not seem to be aligned with the value function axis in the left-most “Value function” plot. Is this an intended consequence or is there some error in this plot?\n\n", "summary_of_the_paper": "The paper proposes a new loss for model-learning in model-based RL, called “Value-Gradient weighted Model loss” (VaGraM). The loss belongs to the family of value-aware model learning losses and is inspired by the VAML and IterVAML losses by Farahmand et al. The paper designs this loss with two main issues of VAML in mind -- (1) value function being evaluated on irrelevant model predictions and (2) model converges to optimas that are sensitive to value function perturbation, being de-stabilized easily when value function is slightly updated. The proposed method is shown to fix both of these issues and is studied in the pedagogical Inverted Pendulum environment and in the MuJoCo Hopper environment. The key takeaways from the Hopper experiments is that the proposed method does better than the baseline MLE model loss in two scenarios -- (1) reduced model capacity and (2) distracting dimensions.", "summary_of_the_review": "The motivation and presentation of the proposed method are major strengths of the paper. The experiments presented on the simple (InvertedPendulum) and complex (Hopper) environments are insightful.\n\nThe paper has potential to be a strong contribution if the empirical analysis can be expanded based on the weaknesses I had mentioned in my review. \n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "McnzxXiEFAo", "reviewer": "Reviewer_VsMY", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n- The objective and the work is clearly motivated and grounded in previous work. \n- Paper is generally well written and easy to follow. \nWeaknesses:\n- The proposed objective is limited to deterministic dynamics; a more details discussion of this limitation is preferred in the main paper as opposed to the appendix. \n- The paper mainly lacks experimental evidence — \n    - how does VaGram fare compared to VAML[1] and other value-aware techniques (e.g. Modhe et al. [2])? \n    - do the trends in hopper hold in other mujoco environments? \n    - How does the proposed technique compare against state-of-the-art MBRL / MFRL methods? (Note that I do not expect outperforming them. But, a comparison will help better understand the contributions in better context). \n\n[1] Farahmand, Amir-massoud. \"Iterative Value-Aware Model Learning.\" In NeurIPS, pp. 9090-9101. 2018.\n[2] Modhe, Nirbhay, Harish Kamath, Dhruv Batra, and Ashwin Kalyan. \"Model-Advantage Optimization for Model-Based Reinforcement Learning.\" arXiv preprint arXiv:2106.14080 (2021).", "summary_of_the_paper": "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. ", "summary_of_the_review": "The proposed objective is straight-forward and can be of use to the MBRL community. However, I do think adding additional experiments to put the contribution in better context will make the work more useful for members of the community -- in both, understanding the significance of the work and potentially, improving value-aware objectives for MBRL. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "4DvoaSdmQ5F", "reviewer": "Reviewer_WgYP", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths\n\n- The paper is novel. It isolates the existing problems with VAML and proposes a loss that 1) has a unique optimum avoiding the shortcomings of VAML 2) addresses the objective mismatch of MLE-based algorithms, thus opening the possibility to get the best of both worlds.\n- MBRL algorithms that address the objective mismatch can learn the dynamics that are helpful for policy improvement but produce unrealistic samples. The VaGraM loss is an MLE weighted by sensitivity of value function to states. As a result, the learned model produces samples that are close to real yet focuses on components that are important for control.\n- The empirical part of the paper seems trustworthy: the practical implementation builds on top of standard baselines, the experiments perform controlled studies, and the results are averaged over a sufficient number of runs.\n- The paper is generally well written and easy to follow.\n\nWeaknesses\n\n- The experiments miss a few relevant baselines. Since the VaGraM loss is value-gradient-weighted MLE, natural baselines would be 1) return-weighted MLE 2) the method from [1] that uses policy-gradient-weighted MLE. While the reviewer appreciates that the authors mention [1] in related work and discuss that VaGraM is value-based and [1] is policy-based, the quality of the paper would be improved if the experiments had these or other related benchmarks that also address the objective mismatch.\n- The significance of the results could be increased by adding more environments. Right now, the continuous control experiments on MuJoCo are limited to only Hopper. Reporting the results on ~5 other tasks (like the baseline MBPO) would demonstrate whether the conclusions generalize beyond a single task.\n\nDetailed comments and questions\n\n- One of the motivations of the paper is that training models using value estimates outside of the empirical state distribution might result in inadequate VAML losses. While the proposed loss indeed removes the dependency on V(s), it still depends on $\\nabla_s V(s)$. Hence, VaGraM might perform inadequately when $\\nabla_s V(s)$ is evaluated outside of the state distribution. On the other hand, it is unclear whether it is tractable to ask for guarantees outside of the training distribution. Do you have any thoughts on that?\n- Figure 1 (right) shows the rescaling effect of VaGraM on the loss landspace. However, in addition to scaling, the value estimates in Figure 1 (left) are rotated. Would it be possible to extend the method with the curvature information e.g. using Hessians instead of $\\textrm{diag}(\\nabla_s V(s))$\n- Is it correct that on Figure 2 (right) MSE and VaGraM losses coincide? If so, that’s very interesting, but the reviewer is unsure whether it’s desirable or not. Could you elaborate on the observation?\n\n[1] D'Oro, Pierluca, Alberto Maria Metelli, Andrea Tirinzoni, Matteo Papini, and Marcello Restelli. \"Gradient-aware model-based policy search.\" In Proceedings of the AAAI Conference on Artificial Intelligence, vol. 34, no. 04, pp. 3801-3808. 2020.\n\n---------------------------------------------------------------\n\nPOST-REBUTTAL UPDATE:\n\nI increase the score to 8. The authors have addressed my major concerns. See the discussion for details.", "summary_of_the_paper": "The paper studies the model learning aspect in model-based reinforcement learning (MBRL). Standard Dyna-based MBRL algorithms rely on maximum likelihood estimation (MLE) or its variants for model learning. In contrast, the authors propose a Value-Gradient weighted Model (VaGraM) loss. Using the loss, the authors mitigate the objective mismatch in standard MBRL: models with high likelihood might not be good for maximizing returns since they optimize a different objective. The paper builds on top of value-aware model learning (VAML) but points out two problems with the original algorithm: 1) value function estimates can be arbitrary outside of the empirical state distribution 2) the entanglement of value and model learning creates the possibility for bad local optima. The experimental results provide evidence that the proposed VaGraM loss alleviates the problems and outperforms in the MLE baseline when the true model cannot be estimated accurately.", "summary_of_the_review": "The reviewer leans towards recommending the submission for acceptance. However, while the idea and the insights are sound, the experimental results could be strengthened by adding relevant baselines and testing on more environments. Addressing the outlined concerns might increase the overall score.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper studies model-based RL in the setting where the model can be misspecified. In this case, MLE of model parameters is a not necessarily a good idea because the error in the model estimate compounds when the model is used for planning. The authors solve this problem by optimizing a novel objective, which takes the quality of the next state prediction into account.

This paper studies an important problem and this was recognized by all reviewers. Its initial reviews were positive and improved to 8, 8, 6, and 6 after the rebuttal. The rebuttal was comprehensive and exemplary. For instance, one concern of this paper was limited empirical evaluation. The authors added 5 new benchmarks and also included stabilizing improvements in their original baselines. I strongly support acceptance of this paper.
Decision: Accept (Spotlight)