Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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. 
main_review: Strengths:
- The objective and the work is clearly motivated and grounded in previous work. 
- Paper is generally well written and easy to follow. 
Weaknesses:
- 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. 
- The paper mainly lacks experimental evidence — 
    - how does VaGram fare compared to VAML[1] and other value-aware techniques (e.g. Modhe et al. [2])? 
    - do the trends in hopper hold in other mujoco environments? 
    - 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). 

[1] Farahmand, Amir-massoud. "Iterative Value-Aware Model Learning." In NeurIPS, pp. 9090-9101. 2018.
[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_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.
review_title: null

final_score: rating: 8: accept, good paper
review_content: 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.
main_review: Strengths

- 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.
- 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.
- 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.
- The paper is generally well written and easy to follow.

Weaknesses

- 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.
- 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.

Detailed comments and questions

- 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?
- 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))$
- 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?

[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.

---------------------------------------------------------------

POST-REBUTTAL UPDATE:

I increase the score to 8. The authors have addressed my major concerns. See the discussion for details.
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.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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.
main_review: ## Strengths:

1. 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.

2. 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.

3. 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.

4. 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.


## Weaknesses

1. The empirical analysis in the paper leaves a lot to be desired for gaining a thorough understanding of the benefits of VaGraM.
For 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.

2. 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).

3. 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.

4. 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.

5. 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.

## Factual Errors/Corrections

1. The conclusion mentions “MuJoCo benchmark” -- this is not accurate as only the Hopper environment was selected.
Section 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?

1. 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.

## Other Issues/Comments:

* 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?


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.

The 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.
review_title: null

final_score: rating: 8: accept, good paper
review_content: 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.

main_review: ## General Comments
- **writing quality**: This is a well written paper. Throughout, it was easy to follow where the authors are going and why.
- **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.

A 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?

- 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. 
- **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.

exper.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? 

exper.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.

exper.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. 

More detailed questions on the experiments are below, but these are the main points. 

## Comments By Section

### 1. Intro
- 1.1 do the authors have any citations for discontinuities being hard to model? It seems like something that is folk-wisdom but unproven?
- 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!
- 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. 

### 2. Background
- 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? 

### 3.Value-Gradient Weighted Model Loss
- 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. 
- 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.
- 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? 
- 3.4 what constitutes a "small prediction error" for the Taylor expansion to hold? Any examples to give on a simulated task?
- 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? 
- 3.6 can the authors clarify footnote 2, this was confusing to me.

### 4. Example: Cartpole
- 4.1 AVI is used but not defined.
- 4.2 Is the source of the value function across domains important? Ie the difference between IterVAML and SAC for generating it?
- 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. 
- 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?

### 5 High-dimensional
- 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?
- 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? 

## Nits / grammatical comments
- n1 the abstract on open review has latex commands that didn't work, I am not sure if you can fix them :) \algoNameFull (\algoName)
- n2 end of page 5, "this , it" is missing a word
- n3 in non-generalizing local minima section, this reads weird "Even worse, there examples where..."
- 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 

## Other comments / ideas (not included in scoring of paper)
- o1 I appreciate the author's comment on the ethical concerns regarding reward hacking and RL.
- o2 I would be very interested in how the 
- o3 are the claims regarding "Single solution convergence" true across seeds? What would this look like? Is there interesting insight to be gained there?

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.
review_title: null

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)