# Gradient-Aware Model-based Policy Search

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

## Key facts

- Cited by 19 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 9 Sep 2019.

## Built on by (most cited first, 19 of 19)

- [Motif: Intrinsic Motivation from Artificial Intelligence Feedback](https://mlsanity.com/paper/2310.00166/) (2023)
- [Mismatched No More: Joint Model-Policy Optimization for Model-Based RL](https://mlsanity.com/paper/2110.02758/) (2021)
- [Control-Oriented Model-Based Reinforcement Learning with Implicit Differentiation](https://mlsanity.com/paper/2106.03273/) (2021)
- [Value Gradient weighted Model-Based Reinforcement Learning](https://mlsanity.com/paper/2204.01464/) (2022)
- [Simplifying Model-based RL: Learning Representations, Latent-space Models, and Policies with One Objective](https://mlsanity.com/paper/2209.08466/) (2022)
- [Deciding What to Model: Value-Equivalent Sampling for Reinforcement Learning](https://mlsanity.com/paper/2206.02072/) (2022)
- [Plan To Predict: Learning an Uncertainty-Foreseeing Model for Model-Based Reinforcement Learning](https://mlsanity.com/paper/2301.08502/) (2023)
- [A Unified Framework for Alternating Offline Model Training and Policy Learning](https://mlsanity.com/paper/2210.05922/) (2022)
- [Do Transformer World Models Give Better Policy Gradients?](https://mlsanity.com/paper/2402.05290/) (2024)
- [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)
- [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)
- [Bayesian Reinforcement Learning with Limited Cognitive Load](https://mlsanity.com/paper/2305.03263/) (2023)
- [Decision-Aware Actor-Critic with Function Approximation and Theoretical Guarantees](https://mlsanity.com/paper/2305.15249/) (2023)
- [BECAUSE: Bilinear Causal Representation for Generalizable Offline Model-based Reinforcement Learning](https://mlsanity.com/paper/2407.10967/) (2024)
- [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)
- [Bi-Level Reinforcement Learning Pathway for Sim-to-Real Optimality](https://mlsanity.com/paper/2510.17709/) (2025)

## Abstract (arXiv)

Traditional model-based reinforcement learning approaches learn a model of the environment dynamics without explicitly considering how it will be used by the agent. In the presence of misspecified model classes, this can lead to poor estimates, as some relevant available information is ignored. In this paper, we introduce a novel model-based policy search approach that exploits the knowledge of the current agent policy to learn an approximate transition model, focusing on the portions of the environment that are most relevant for policy improvement. We leverage a weighting scheme, derived from the minimization of the error on the model-based policy gradient estimator, in order to define a suitable objective function that is optimized for learning the approximate transition model. Then, we integrate this procedure into a batch policy improvement algorithm, named Gradient-Aware Model-based Policy Search (GAMPS), which iteratively learns a transition model and uses it, together with the collected trajectories, to compute the new policy parameters. Finally, we empirically validate GAMPS on benchmark domains analyzing and discussing its properties.