# Value Gradient weighted Model-Based Reinforcement Learning

arXiv: https://arxiv.org/abs/2204.01464 · HTML: https://mlsanity.com/paper/2204.01464/

## Key facts

- Cited by 32 papers in the map (data through 14 Oct 2026).
- Builds on 27 of its 51 references that are in the map.
- First posted on arXiv 4 Apr 2022.

## Builds on (most cited first, 27 of 27)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Dueling Network Architectures for Deep Reinforcement Learning](https://mlsanity.com/paper/1511.06581/) (2015)
- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [Learning Invariant Representations for Reinforcement Learning without Reconstruction](https://mlsanity.com/paper/2006.10742/) (2020)
- [Embed to Control: A Locally Linear Latent Dynamics Model for Control from Raw Images](https://mlsanity.com/paper/1506.07365/) (2015)
- [Model-Based Reinforcement Learning with Value-Targeted Regression](https://mlsanity.com/paper/2006.01107/) (2020)
- [Algorithmic Framework for Model-based Deep Reinforcement Learning with Theoretical Guarantees](https://mlsanity.com/paper/1807.03858/) (2018)
- [Value Prediction Network](https://mlsanity.com/paper/1707.03497/) (2017)
- [The Distracting Control Suite -- A Challenging Benchmark for Reinforcement Learning from Pixels](https://mlsanity.com/paper/2101.02722/) (2021)
- [Objective Mismatch in Model-based Reinforcement Learning](https://mlsanity.com/paper/2002.04523/) (2020)
- [The Value Equivalence Principle for Model-Based Reinforcement Learning](https://mlsanity.com/paper/2011.03506/) (2020)
- [Goal-Aware Prediction: Learning to Model What Matters](https://mlsanity.com/paper/2007.07170/) (2020)
- [Self-Correcting Models for Model-Based Reinforcement Learning](https://mlsanity.com/paper/1612.06018/) (2016)
- [MBRL-Lib: A Modular Library for Model-based Reinforcement Learning](https://mlsanity.com/paper/2104.10159/) (2021)
- [Control-Oriented Model-Based Reinforcement Learning with Implicit Differentiation](https://mlsanity.com/paper/2106.03273/) (2021)
- [Proper Value Equivalence](https://mlsanity.com/paper/2106.10316/) (2021)
- [Is High Variance Unavoidable in RL? A Case Study in Continuous Control](https://mlsanity.com/paper/2110.11222/) (2021)
- [Policy-Aware Model Learning for Policy Gradient Methods](https://mlsanity.com/paper/2003.00030/) (2020)
- [Learning Dynamics Models for Model Predictive Agents](https://mlsanity.com/paper/2109.14311/) (2021)
- [Gradient-Aware Model-based Policy Search](https://mlsanity.com/paper/1909.04115/) (2019)
- [Towards Deeper Deep Reinforcement Learning with Spectral Normalization](https://mlsanity.com/paper/2106.01151/) (2021)
- [Prediction, Consistency, Curvature: Representation Learning for Locally-Linear Control](https://mlsanity.com/paper/1909.01506/) (2019)
- [Control-Aware Representations for Model-based Reinforcement Learning](https://mlsanity.com/paper/2006.13408/) (2020)

## Built on by (most cited first, 30 of 32)

- [A Survey on Model-based Reinforcement Learning](https://mlsanity.com/paper/2206.09328/) (2022)
- [Simplifying Model-based RL: Learning Representations, Latent-space Models, and Policies with One Objective](https://mlsanity.com/paper/2209.08466/) (2022)
- [Refining Diffusion Planner for Reliable Behavior Synthesis by Automatic Detection of Infeasible Plans](https://mlsanity.com/paper/2310.19427/) (2023)
- [Deciding What to Model: Value-Equivalent Sampling for Reinforcement Learning](https://mlsanity.com/paper/2206.02072/) (2022)
- [Live in the Moment: Learning Dynamics Model Adapted to Evolving Policy](https://mlsanity.com/paper/2207.12141/) (2022)
- [MAD-TD: Model-Augmented Data stabilizes High Update Ratio RL](https://mlsanity.com/paper/2410.08896/) (2024)
- [A Unified Framework for Alternating Offline Model Training and Policy Learning](https://mlsanity.com/paper/2210.05922/) (2022)
- [When does Self-Prediction help? Understanding Auxiliary Tasks in Reinforcement Learning](https://mlsanity.com/paper/2406.17718/) (2024)
- [State-Covering Trajectory Stitching for Diffusion Planners](https://mlsanity.com/paper/2506.00895/) (2025)
- [Off-Policy RL Algorithms Can be Sample-Efficient for Continuous Control via Sample Multiple Reuse](https://mlsanity.com/paper/2305.18443/) (2023)
- [TaskMet: Task-Driven Metric Learning for Model Learning](https://mlsanity.com/paper/2312.05250/) (2023)
- [TOM: Learning Policy-Aware Models for Model-Based Reinforcement Learning via Transition Occupancy Matching](https://mlsanity.com/paper/2305.12663/) (2023)
- [Distributional Model Equivalence for Risk-Sensitive Reinforcement Learning](https://mlsanity.com/paper/2307.01708/) (2023)
- [A Unified View on Solving Objective Mismatch in Model-Based Reinforcement Learning](https://mlsanity.com/paper/2310.06253/) (2023)
- [Local Manifold Approximation and Projection for Manifold-Aware Diffusion Planning](https://mlsanity.com/paper/2506.00867/) (2025)
- [Operator Splitting Value Iteration](https://mlsanity.com/paper/2211.13937/) (2022)
- [On Rate-Distortion Theory in Capacity-Limited Cognition & Reinforcement Learning](https://mlsanity.com/paper/2210.16877/) (2022)
- [Policy-shaped prediction: avoiding distractions in model-based reinforcement learning](https://mlsanity.com/paper/2412.05766/) (2024)
- [Bayesian Reinforcement Learning with Limited Cognitive Load](https://mlsanity.com/paper/2305.03263/) (2023)
- [Model-based Offline RL via Robust Value-Aware Model Learning with Implicitly Differentiable Adaptive Weighting](https://mlsanity.com/paper/2603.08118/) (2026)
- [Maximum Entropy Model Correction in Reinforcement Learning](https://mlsanity.com/paper/2311.17855/) (2023)
- [Calibrated Value-Aware Model Learning with Probabilistic Environment Models](https://mlsanity.com/paper/2505.22772/) (2025)
- [λ-models: Effective Decision-Aware Reinforcement Learning with Latent Models](https://mlsanity.com/paper/2306.17366/) (2023)
- [On Task-Relevant Loss Functions in Meta-Reinforcement Learning and Online LQR](https://mlsanity.com/paper/2312.05465/) (2023)
- [Bi-Level Reinforcement Learning Pathway for Sim-to-Real Optimality](https://mlsanity.com/paper/2510.17709/) (2025)
- [Taylor TD-learning](https://mlsanity.com/paper/2302.14182/) (2023)
- [MrCoM: A Meta-Regularized World-Model Generalizing Across Multi-Scenarios](https://mlsanity.com/paper/2511.06252/) (2025)
- [CORL: Reinforcement Learning of MILP Policies Solved via Branch and Bound](https://mlsanity.com/paper/2512.11169/) (2025)
- [Representation Learning Enables Scalable Multitask Deep Reinforcement Learning](https://mlsanity.com/paper/2606.05555/) (2026)
- [DT2: Decision-Targeted Digital Twins](https://mlsanity.com/paper/2606.25923/) (2026)

## Abstract (arXiv)

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.