# Calibrated Value-Aware Model Learning with Probabilistic Environment Models

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in RL theory and sample complexity, one of 15 directions in Deep reinforcement learning methods.
- Builds on 41 of its 62 references that are in the map.
- First posted on arXiv 28 May 2025.

## Placed in

- Deep reinforcement learning methods / [RL theory and sample complexity](https://mlsanity.com/deeprl/rl-theory-and-sample-complexity/)

## Builds on (most cited first, 30 of 41)

- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Addressing Function Approximation Error in Actor-Critic Methods](https://mlsanity.com/paper/1802.09477/) (2018)
- [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)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [Deep Reinforcement Learning at the Edge of the Statistical Precipice](https://mlsanity.com/paper/2108.13264/) (2021)
- [Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models](https://mlsanity.com/paper/1805.12114/) (2018)
- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [TD-MPC2: Scalable, Robust World Models for Continuous Control](https://mlsanity.com/paper/2310.16828/) (2023)
- [Temporal Difference Learning for Model Predictive Control](https://mlsanity.com/paper/2203.04955/) (2022)
- [dm_control: Software and Tasks for Continuous Control](https://mlsanity.com/paper/2006.12983/) (2020)
- [Mastering Atari Games with Limited Data](https://mlsanity.com/paper/2111.00210/) (2021)
- [Model-Based Reinforcement Learning with Value-Targeted Regression](https://mlsanity.com/paper/2006.01107/) (2020)
- [Sample-Efficient Reinforcement Learning with Stochastic Ensemble Value Expansion](https://mlsanity.com/paper/1807.01675/) (2018)
- [Understanding and Preventing Capacity Loss in Reinforcement Learning](https://mlsanity.com/paper/2204.09560/) (2022)
- [Value Prediction Network](https://mlsanity.com/paper/1707.03497/) (2017)
- [Bigger, Regularized, Optimistic: scaling for compute and sample-efficient continuous control](https://mlsanity.com/paper/2405.16158/) (2024)
- [Sample Complexity of Reinforcement Learning using Linearly Combined Model Ensembles](https://mlsanity.com/paper/1910.10597/) (2019)
- [The Value Equivalence Principle for Model-Based Reinforcement Learning](https://mlsanity.com/paper/2011.03506/) (2020)
- [The Predictron: End-To-End Learning and Planning](https://mlsanity.com/paper/1612.08810/) (2016)
- [BYOL-Explore: Exploration by Bootstrapped Prediction](https://mlsanity.com/paper/2206.08332/) (2022)
- [Empirical Design in Reinforcement Learning](https://mlsanity.com/paper/2304.01315/) (2023)
- [Bridging State and History Representations: Understanding Self-Predictive RL](https://mlsanity.com/paper/2401.08898/) (2024)
- [Understanding Self-Predictive Learning for Reinforcement Learning](https://mlsanity.com/paper/2212.03319/) (2022)
- [Self-Correcting Models for Model-Based Reinforcement Learning](https://mlsanity.com/paper/1612.06018/) (2016)
- [Towards General-Purpose Model-Free Reinforcement Learning](https://mlsanity.com/paper/2501.16142/) (2025)
- [Mismatched No More: Joint Model-Policy Optimization for Model-Based RL](https://mlsanity.com/paper/2110.02758/) (2021)
- [Efficient Deep Reinforcement Learning Requires Regulating Overfitting](https://mlsanity.com/paper/2304.10466/) (2023)
- [Control-Oriented Model-Based Reinforcement Learning with Implicit Differentiation](https://mlsanity.com/paper/2106.03273/) (2021)

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

- [Diagnosing JEPA World Models with Action-Conditioned Predictive Consistency](https://mlsanity.com/paper/2608.12939/) (2026)
- [Not All Errors Matter: Decision-Relevant Prediction Error Predicts Planning Quality](https://mlsanity.com/paper/2609.32322/) (2026)

## Abstract (arXiv)

The idea of value-aware model learning, that models should produce accurate value estimates, has gained prominence in model-based reinforcement learning. The MuZero loss, which penalizes a model's value function prediction compared to the ground-truth value function, has been utilized in several prominent empirical works in the literature. However, theoretical investigation into its strengths and weaknesses is limited. In this paper, we analyze the family of value-aware model learning losses, which includes the popular MuZero loss. We show that these losses, as normally used, are uncalibrated surrogate losses, which means that they do not always recover the correct model and value function. Building on this insight, we propose corrections to solve this issue. Furthermore, we investigate the interplay between the loss calibration, latent model architectures, and auxiliary losses that are commonly employed when training MuZero-style agents. We show that while deterministic models can be sufficient to predict accurate values, learning calibrated stochastic models is still advantageous.