# Policy-Aware Model Learning for Policy Gradient Methods

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

## Key facts

- Cited by 23 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 28 Feb 2020.

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

- [A Survey on Model-based Reinforcement Learning](https://mlsanity.com/paper/2206.09328/) (2022)
- [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)
- [Control-Aware Prediction Objectives for Autonomous Driving](https://mlsanity.com/paper/2204.13319/) (2022)
- [Minimax Model Learning](https://mlsanity.com/paper/2103.02084/) (2021)
- [Complex Momentum for Optimization in Games](https://mlsanity.com/paper/2102.08431/) (2021)
- [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)
- [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)
- [Model-Advantage and Value-Aware Models for Model-Based Reinforcement Learning: Bridging the Gap in Theory and Practice](https://mlsanity.com/paper/2106.14080/) (2021)
- [λ-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)
- [Dual-Frontier: When Can an Agent Trust Its World Model?](https://mlsanity.com/paper/2609.26293/) (2026)
- [What a World Model Represents Is Three Questions](https://mlsanity.com/paper/2607.06640/) (2026)

## Abstract (arXiv)

This paper considers the problem of learning a model in model-based reinforcement learning (MBRL). We examine how the planning module of an MBRL algorithm uses the model, and propose that the model learning module should incorporate the way the planner is going to use the model. This is in contrast to conventional model learning approaches, such as those based on maximum likelihood estimate, that learn a predictive model of the environment without explicitly considering the interaction of the model and the planner. We focus on policy gradient type of planning algorithms and derive new loss functions for model learning that incorporate how the planner uses the model. We call this approach Policy-Aware Model Learning (PAML). We theoretically analyze a generic model-based policy gradient algorithm and provide a convergence guarantee for the optimized policy. We also empirically evaluate PAML on some benchmark problems, showing promising results.