Research map / Reinforcement learning / Policy gradient and actor-critic
Calibrated Value-Aware Model Learning with Probabilistic Environment Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Analyzes MuZero-style and IterVAML value-aware model losses, shows sample-based versions are uncalibrated for stochastic models, proposes a variance-corrected CVAML loss, and tests calibration on Garnet MDPs and DMC humanoid/dog latent models.
Problem
Value-aware model learning losses (MuZero, IterVAML) are used as surrogate losses for model and value learning, but their calibration, i.e. whether they recover the correct model and value function, is not theoretically understood, especially with stochastic models.
Approach
Unifies MuZero and IterVAML into (m,b)-VAML losses, proves sample-based versions are uncalibrated with stochastic models because they bias toward low-variance models, and corrects this with a sample-variance subtraction (CVAML); recommends using CVAML for the model and a TD loss for value learning.
Key results
- Calibrated (CVAML) losses give the lowest value prediction error in Garnet MDPs for all stochastic (non-deterministic) settings.
- On DMC humanoid, calibrating (1,1)-VAML gives a clear performance improvement over uncalibrated MuZero-style loss.
- Deterministic models are theoretically sufficient for value-equivalence in compact settings, yet stochastic models outperform deterministic ones on the dog domain.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Apply CVAML calibration to the MuZero-style (1,1)-VAML loss | (1,1)-VAML · Humanoid | performance | Helped |
| Apply CVAML calibration to the (1,0)-VAML loss | (1,0)-VAML · Dog and humanoid | performance | No clear effect |
| Use a probabilistic latent model instead of a deterministic latent model | Latent model · Dog | performance | Helped |
| Use a probabilistic latent model instead of a deterministic latent model | Latent model · Humanoid | performance | Hurt |
| Use (1,0)-VAML updates instead of (1,1)-VAML updates | Latent model · DMC dog and humanoid tasks | performance | Helped |
| Add the BYOL auxiliary loss (ablation) | IterVAML latent model · Humanoid | performance | Helped |
| Use CVAML instead of uncalibrated VAML | Stochastic latent model · Garnet stochastic environments | value prediction error | Helped |
| Use calibrated (1,1)-VAML instead of uncalibrated (1,1)-VAML | (1,1)-VAML · Garnet environments | value prediction | Helped |
| Use calibrated VAML instead of the baseline KL-based model loss | Garnet model · Near-deterministic Garnet environments | value estimation | No clear effect |
| Use calibrated (1,1)-VAML instead of uncalibrated (1,1)-VAML | (1,1)-VAML · Garnet environments with high-entropy model initialization | value prediction | Hurt |
| Add the variance correction to the sample-based IterVAML loss | Stochastic model learning · Stochastic environments | model calibration | Helped |
| Use a stochastic model class with uncorrected (m,b)-VAML instead of a deterministic model class | (m,b)-VAML · Stochastic environments | model calibration | Hurt |
Surprising findings
- IterVAML (1,0) updates are stable and outperform MuZero-style (1,1) updates on continuous control when paired with the right architecture and auxiliary losses.
- Standard MuZero-style (m,b) losses can be uncalibrated even when the ground-truth environment is used, due to value-function variance.
Limitations
- Calibration proofs rely on specific assumptions such as the non-constant value function and non-degenerate distributions.
- Stochastic model advantage is environment-dependent and not uniform across benchmarks.
- Calibrated (1,1)-VAML still struggles in low-stochasticity environments due to local plateaus.
Experimental setup
Datasets: Garnet MDPs DeepMind Control humanoid and dog tasks
Benchmarks: Garnet MDPs DMC humanoid DMC dog
Metrics: mean squared error of value prediction normalized return aggregate performance with 95% bootstrap CI
Baselines: MuZero loss (m,b)-VAML IterVAML (1,0)-VAML KL-based model loss model-free TD
Builds on: IterVAML MuZero TD-MPC BYOL TD3
Closest papers by meaning
- Value Gradient weighted Model-Based Reinforcement Learning (ICLR 2022 · similarity 0.90)
- Deciding What to Model: Value-Equivalent Sampling for Reinforcement Learning (NeurIPS 2022 · similarity 0.89)
- The Value Equivalence Principle for Model-Based Reinforcement Learning (NeurIPS 2020 · similarity 0.89)
- Inverting the Bellman Equation: From $Q$-Values to World Models (NeurIPS 2026 · similarity 0.88)
- Long-Horizon Q-Learning: Accurate Value Learning via n-Step Inequalities (NeurIPS 2026 · similarity 0.87)
- The Edge-of-Reach Problem in Offline Model-Based Reinforcement Learning (NeurIPS 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.