mlsanityData to 14 Oct 2026

Calibrated Value-Aware Model Learning with Probabilistic Environment Models

arXiv 2505.22772 ↗First posted 28 May 2025cs.LGcited by 2

On the map: Deep reinforcement learning methods → RL theory and sample complexity

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.

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.

Builds on

41 of 62 references.

References · 41

  1. 2020Bootstrap your own latent: A new approach to self-supervised Learningcited by 4,750
  2. 2013Playing Atari with Deep Reinforcement Learningcited by 3,274
  3. 2018Addressing Function Approximation Error in Actor-Critic Methodscited by 2,417
  4. 2019Dream to Control: Learning Behaviors by Latent Imaginationcited by 1,653
  5. 2019Mastering Atari, Go, Chess and Shogi by Planning with a Learned Modelcited by 1,353
  6. 2020Mastering Atari with Discrete World Modelscited by 1,007
  7. 2021Deep Reinforcement Learning at the Edge of the Statistical Precipicecited by 807
  8. 2018Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Modelscited by 774
  9. 2019When to Trust Your Model: Model-Based Policy Optimizationcited by 773
  10. 2023TD-MPC2: Scalable, Robust World Models for Continuous Controlcited by 512
  11. 2022Temporal Difference Learning for Model Predictive Controlcited by 397
  12. 2020dm_control: Software and Tasks for Continuous Controlcited by 375
29 more
  1. 2021Mastering Atari Games with Limited Datacited by 237
  2. 2020Model-Based Reinforcement Learning with Value-Targeted Regressioncited by 234
  3. 2018Sample-Efficient Reinforcement Learning with Stochastic Ensemble Value Expansioncited by 168
  4. 2022Understanding and Preventing Capacity Loss in Reinforcement Learningcited by 125
  5. 2017Value Prediction Networkcited by 117
  6. 2024Bigger, Regularized, Optimistic: scaling for compute and sample-efficient continuous controlcited by 108
  7. 2019Sample Complexity of Reinforcement Learning using Linearly Combined Model Ensemblescited by 100
  8. 2020The Value Equivalence Principle for Model-Based Reinforcement Learningcited by 94
  9. 2016The Predictron: End-To-End Learning and Planningcited by 93
  10. 2022BYOL-Explore: Exploration by Bootstrapped Predictioncited by 73
  11. 2023Empirical Design in Reinforcement Learningcited by 55
  12. 2024Bridging State and History Representations: Understanding Self-Predictive RLcited by 54
  13. 2022Understanding Self-Predictive Learning for Reinforcement Learningcited by 44
  14. 2016Self-Correcting Models for Model-Based Reinforcement Learningcited by 43
  15. 2025Towards General-Purpose Model-Free Reinforcement Learningcited by 41
  16. 2021Mismatched No More: Joint Model-Policy Optimization for Model-Based RLcited by 36
  17. 2023Efficient Deep Reinforcement Learning Requires Regulating Overfittingcited by 36
  18. 2021Control-Oriented Model-Based Reinforcement Learning with Implicit Differentiationcited by 35
  19. 2021Proper Value Equivalencecited by 33
  20. 2022Value Gradient weighted Model-Based Reinforcement Learningcited by 32
  21. 2020Selective Dyna-style Planning Under Limited Model Capacitycited by 30
  22. 2022Simplifying Model-based RL: Learning Representations, Latent-space Models, and Policies with One Objectivecited by 28
  23. 2020Policy-Aware Model Learning for Policy Gradient Methodscited by 23
  24. 2019Gradient-Aware Model-based Policy Searchcited by 19
  25. 2024MAD-TD: Model-Augmented Data stabilizes High Update Ratio RLcited by 16
  26. 2024When does Self-Prediction help? Understanding Auxiliary Tasks in Reinforcement Learningcited by 14
  27. 2019Planning with Expectation Modelscited by 12
  28. 2022Operator Splitting Value Iterationcited by 8
  29. 2023Maximum Entropy Model Correction in Reinforcement Learningcited by 2

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2026: 2.

2 papers.

  1. 2026Diagnosing JEPA World Models with Action-Conditioned Predictive Consistencycited by 3
  2. 2026Not All Errors Matter: Decision-Relevant Prediction Error Predicts Planning Qualitycited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2024A Note on Loss Functions and Error Compounding in Model-based Reinforcement Learning
  2. 2026Inverting the Bellman Equation: From Q-Values to World Models
  3. 2026Robo-ValueRL: Reliable Value Estimation for Offline-to-Online Reinforcement Learning
  4. 2024On the Limited Representational Power of Value Functions and its Links to Statistical (In)Efficiency
  5. 2025Rectifying Regression in Reinforcement Learning
  6. 2024A Multi-step Loss Function for Robust Learning of the Dynamics in Model-based Reinforcement Learning
  7. 2026Operator-on-F complements value-equivalence: a planning-time diagnostic for latent world models
  8. 2026A Pontryagin Method of Model-based Reinforcement Learning via Hamiltonian Actor-Critic