Calibrated Value-Aware Model Learning with Probabilistic Environment Models
Authors: Claas A Voelcker
Anastasiia Pedan
Arash Ahmadian
Romina Abachi
Igor Gilitschenski
Amir-massoud Farahmand
Venue: ICML 2025 poster (2025)
Keywords: model-based rl
muzero
itervaml
theory
calibration
latent models
Primary area: reinforcement_learning
TL;DR: We prove a flaw in the popular VAML family of algorithms, propose a fix, and validate it empirically.

Abstract: 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.

Calibrated Value-Aware Model Learning with Probabilistic Environment Models
Claas Voelcker Anastasiia Pedan 312 Arash Ahmadian Romina Abachi Igor Gilitschenski 12
Amir-massoud Farahmand 671
Abstract
errors which propagate to value function learning (Schnei- der, 1997: Kearns & Singh, 2002: Talvitie, 2017; Lambert et al., 2020). In such cases, deciding what aspects of the environment to model is crucial.
The idea of value-aware model learning, that mod- els 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-trute truth value function, has been utilized in several prominent empirical works in the liter- ature. However, theoretical investigation into its strengths and weaknesses is limited. In this paper, we analyze the family of value-aware model learn- ing losses, which includes the popular MuZero loss. We show that these losses, normally used, are uncalibrated surrogate losses, which means that they do not always recover the cor- rect model and value function. Building on this insight, we propose corrections to solve this is- sue. Furthermore, we investigate the interplay between the loss calibration, latent model architec- tures, and auxiliary losses that are commonly em- ployed when training MuZero-style agents. We show that while deterministic models can be suf- ficient to predict accurate values, learning cali- brated stochastic models is still advantagcous.
The paradigm of value-aware model learning (VAML) (Farahmand et al., 2017) and value- equivalence (Grimm et al., 2020; 2021) addresses this by training models that lead to accurate value estima- tion. Prominent value-aware model learning approaches are MuZero (Schrittwieser al., 2020) and IterVAML (Farah- mand, 2018). The MuZero loss has been shown to perform well in discrete (Schrittwieser et al., 2020; Ye et al., 2021) and continuous control tasks (Hansen et al., 2022; 2024), but has received little theoretical investigation. On the other hand, IterVAML is a theoretically motivated algorithm but not commonly used in empirical work. We show that MuZero and IterVAML. can be unified in a family of losses, which we term (m, baaaluee-AAarr Model Losses ((m, b)-VALL.. The name stresses the two core hyperparameters: the model rollout steps, m, and steps used to estimate the bootstrapped value function target, b.
The (m, b)- -ALLL losses are used as surrogate losses in place of other value- or model learning losses. Therefore, is important to ask whether they are calibrated (Steinwart & Christmann, 2008). A calibrated loss does not lead to suboptimal minima when the function class includes optimal functions for the original target loss.
1. Introduction
In model-based reinforcement learning, an agent collects in- formation in environment and uses it to learn a model of the world. This model is used to improve value estimation and policy learning (Sutton, 1990; Deisenroth & Rasmussen, 2011; 1; Hafner et al., 2020; Schrittwieser et al., 2020). How- ever, as environment complexity increases, learning a model becomes more and more challenging. This leads to model
This paper has two parts, each with a theoretical and empirical section. We answer two questions about the (m, b)-L family: (a) What variants of the (m, VALL losses well-calibrated to recover correct models and value functions? (b) Do we observe problems with uncalibrated losses when using standard architectures, especially deterministic latent-space models?
Research question:
'Department of Computer Science, University of Toronto, Canada VVector Institute, Toronto, Canada 'Igor Sikorsky Kyiv Polytechnic Institute, Kyiv, Ukraine "Cohere, Toronto, Canada Ubisoff, Montreal, Canada 'Polytechnique Montreal, Canada "MILA, Montreal, Canada. Correspondence Claas Voeleker <evoelcker@ es. toronto.cdu>>.
Contributions: As our main theoretical contribution, we mathematically analyze the family of (m, VAML algo- rithms. We prove that all sampled-based loss variants from this family uncalibrated when used with stochastic environment model. Minimizing the losses with regard to data samples will result in value functions and models with lower variance than the correct ground-truth functions, even
Proceedings of the 4 International Conference on Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Calibrated Value-Aware Model Learning with Probabilistic Environment Models
if the dataset adequately covers the state-action space, To counter this issue, we derive novel loss variant.
where E Z is representation of state x E X' given by ( X' + 2. Observation-space models predict next states in the representation of the environment, while latent-space models reconstruct learned features of these states.
In the second part, we address issues arising from the way currett algorithms in the (m, ---V family are com- monly implemented. We prove that a stochastic model class is not necessary to learn a single-step decision-equivalen model in stochastic environments. This validates the prac- tice of primarily using deterministic models in empirical work (Oh et al., 2017; Schrittwieser al., 2020; Hansen et al., 2022). Empirically we find that using stochastic mod- els can still lead to improved performance, although this is environment-dependent.
The notation xx()) refers to the n-th step in a rollout starting in state x in the environment. We will use g(n)) to refer to samples from the n-th step model prediction and write Egm H as a shorthand for (ptm) (-Ja) [-]
3. The Value-Aware Model Learning framework
The losses of the decision-aware learning framework share the goal of finding models that provide good value func- tion estimates. Instead of simply learning model using maximum likelihood estimation, the losses are based on differences in value prediction.
2. Background
Reinforcement Learning: We consider a standard Markov Decision Process (MDP) (Puterman, 1994) A, P,r,)) with state space X' action space A. transition kernel P(z'Iz,a), a). reward function T X' x A + I., and discount factor 7 E [0, 1). A policy ((als) maps a state to a distribution over actions. The goal of a re- inforcement learning agent is to find policy maximiz- ing the expected discounted infinite horizon reward, i.e., max, E E y'r(x,, at) Øo]-
3.1. Iterative Value Aware Model Learning
Iterative Value Aware Model Learning (IteVAMM) (Farah- mand, 2018) computes the difference between the expected value under the model and samples in the environment
P (Ø, V(x)
The value function is defined as the expected return of the policy T: (*) = Ez.p Eizor'r(aa (z) = T The policy-conditioned transition kernel is P (x'(a) P [ar, (a(aaaa)d The value is the unique fixed point of the Bellman operator
Emm V (I(mn)
-
[[Epr Epr
V (2( (m(mm)
(1)
V (a(m))
} [v (2((600) (z((m)
(Tp= VI)))= = E, [r(zo, + Bpp [V(x1) [2 = x]-
We will use LiterVAMLLm to refer to the expectation-based version, and 2 L to refer to the sampling- based version. [-] denotes the stop- gradient operation. To re- duce notational complexity, we drop the action dependence in all following propositions; all results hold without loss of generality for the action-conditioned case as well. The expectation-basee version of the IterVAML loss has an im- portant relationship to the error of computing the model's s Bellman operator compared to the true environments Bell- man operator.
This operator can be extended to a multi-step version as
(TB. V)() = ET.P Σ r" r" an) + VVV( x
We also define [Tp- VI(z) V(x).
Model-based RL: An environment model is function that approximates the transition kernel." Learned models are used to augment RL algorithms in several ways. For a comprehensive survey, refer to Moerland et al. (2023). In this paper, we focus on value learning with model data, which is commonly referred to as Dyna (Sutton, 1990).
Proposition 1. Farahmand (2018) Let pa be the policy- conditioned transition kernel of an MDP and let V X' + IR be a function. Let p(-(x) be a model so that The VAML. (1.V)z) = 0. Then [Tp= V] (x) = [TpV](x)) (z).
We use P to refer a stochastic model, and f for determinis- tic models. When model is used to predict the observation r' from I, a (such as the model used in MBPO (Janner et al., 2019) we call it an observation-space models. Alter- natively, latent-space models of the form p(='|=, |z, a) are used,
We give short proof in Subsection A.1
Intuitively, model achieving 0 loss can be used instead of the ground truth environment when computing the Bellman operator. The set of models that achieve 0 Iter VAML error is equal to the set of equiveequi models for the value estimate V- (Grimm et al., 2021).
In this paper, we will generally use the term model to refer to an environment model, to neural network, to keep consistent with the reinforcement learning nomenclature.
2

---

Calibrated Value-Aware Model Learning with Probabilistic Environment Models
D
④ (m)
Model trujectory
Emp (D. V(x))
x: (0)
Environment trajectory
V (8(m)
TTp v (m(()) Bg
D m
x (0)
x (m)
(m)
(m++))
(m-+))
model
environment
Figure 1. Sketch of (m. bnn,b)-Vann VA-M.. The loss is computed from an m-step model and (m + b)-step environment trajectory. It is the difference between the estimated value of the m-th model state, and the b-step Bellman operator starting from the m-tth environment state.
3.2. The (m,b)-VAML Family
4. Analysis of decision-aware losses in stochastic environments
The MuZero loss (Schrittwieser et al., 2020) was introduced to unify the value function and model learning components of an MBRL algorithm. It can be interpreted as a variation of the Iter loss that uses a single sample estimate of the Bellman Operator and bootstrapped target estimate. To unify the losses, we present them in single equation with two important hyperparameters: mi the number of steps in the trajectory that the loss is computed over, and 6, the number of steps for the multi-step Bellman operator. We refer to this unified family of losses as (m, b)- Value Aware Model Losses (VAML)
The goal of learning MBRL is to recover an (approxi- mately) optimal model and to learn a correct value function. As we have shown, minimizing the IterVAML loss perfectly results in model that leads to correct Bellman Operator. However, in practice, the inner expectation of the IterVAMI. loss has to be replaced by a sampling- based approximation. In addition, (m, b)- VAML is used in the MuZero algorithm to update the value function directly in addition to the model. We show when this leads to learning correct models and value functions asymptotically. An overview of conclu- sions be found in Table 1.
4.1. Calibration of surrogate loss functions
Cm. (p. VIViars
Formally, we ask whether the surrogate (m, 6A-VAM is calibrated. Intuitively, calibrated surrogate loss does not select a suboptimal function for the target loss. Formally, we require that the surrogate loss is not perfectly minimized by a function that does not perfectly minimize the target loss. Therefore, we call losses that do not have this prop- erty comiimum-uncalbraati losses, or simply uncalibrated outside of formal statements. We borrow the concept from Steinwart & Christmann (2008), however, we use slightly more restricted definition of uncalibrated here. This means that all losses which we consider uncalibrated are also un- calibrated in the definition given in Steinwart & Christmann (2008), but not vice versa.
|
(z(m))
[T Vear (a(m)) 5x
= E
x
(2)
where Vear is a target network (Mnih et al., 2013) and (m)) and (mm) are sampled independently from 7d' and ğ respec- tively. Note that samples from the real environment are used to approximate [Tp- V.nl]. The loss function and the rela- tion of the model and environments rollout are visualized in Figure 1.
Several works use variations of this loss: The original MuZero algorithm (Schrittwieser et al., 2020) and follow-up work (Ye et al., 2021; Antonoglou et al., 2022) use m 2 and b 2 1. In continuous control, m 2 1,b = 1 has been used in the TD-MPC line of work (Hansen al., 2022: 2024). When using only a single sample based VAML loss is equal to - Note that the (m, b)- loss can easily be extended to a k sample variant analogous to the IteL loss, which we will use later. We dropped the summation over k samples here to simplify the (already dense) notation. Finally, regular model-free TD learning corresponds to m = 0, b 2211 1.
surrogate losses). Let Cia (f x, (m)) be a loss function defined over samples from an MDP. Let Laur be surrogate function for the loss. Let F be set of (perfect) minima of fras so that lfarlf" (m)) 0 for all f E F² A surrogate func- tion is minimum-uncalibrutee for the target loss if there exists an MDP, a function class F with F C F, and a state x so that
Definition 1
1, the sample-
arg min EE (p Coue 2, am))) ZF. JEF
3

---

Calibrated Value-Aware Model Learning with Probabilistic Environment Models
An example of siniimun-uncaliirattd--sss is the double sampling issue in Bellman Residual Minimization (BRM). While the target loss is minimized by the ground truth value function, BRM actually chooses a function that minimizes the Bellman error and and additional variance term.
We refer to the loss as Calibrated VAML (CVAML).
Proposition 3. The variance-corrected loss
E LE (P, Vlr, x (m) is a calibrated surrogate loss for f
LoeVAAML. (P, Vlrr,
We will show that some issues introduced by naively using (m, b)- 1)-V can be fixed with a tractable modification to the loss. We call such cases resolvable.
Analogous to the IterVAML case, the sample-based (m, b Z 1)-VAML loss is an uncalibrated loss for learning model.
4.3. Value learning bias in stochastic models
4.2. Model learning bias with stochastic models
We now show that this issue also affects the value function learning with the MuZero loss. The problem here lies in the use of the bootstrapped Bellman target together with a multi- step value rollout. In an MDP, the values of two states x and y are not guaranteed to be equal (or even particularly close) just because they share an ancestor state, unless we make assumptions about the variance of the value function over successor states. However, the MuZero loss still minimizes the difference in value functions between these two states.
Most prior work uses (m, b)- VAML with deterministic mod- els. Exceptions are Voelcker al. (2022) and Antonoglou et al. (2022), but neither changes the loss functions to account for the model parametrization. For stochastic model class, (m, V)-VLM is uncalibrated for the target loss VAMLL
We begin with analyzing the loss for m 2 and = For simplicity we assume that Viar =..
Proposition 2. Let 1m (ê, Vla, x(m)) be the surrogate
We show that the MuZero loss therefore is not guaranteed to recover the correct value function, even when we have a perfect (stochastic) model. To formalize the issue, we compare the solution found with (m, 6)- 6AML with the regular TD loss function. Note that for this loss, we assume that I(m) is fixed sample from the environment, not random variable like ,.(m), with 4(m++) being its successor state in the ground truth environment. This distinction is important, as it is exactly what leads to the bias in the (m, b 2 1) loss.
loss for ChervAML (D. V(z). Let P be the set of all dis- tributions p for which Ep [f(mm))]
Ep- (x(mm))) There exist an MDP and a class of distributions P with pa E P C so that
Argmin E& (e & p
PEP
Therefore. Emodel is mininumm-tmcllibraae
Proofs for this section can be found in subsubsection A.2.2. When using samples from a stochastic model to compute the loss function in Equation 2, we are left with a variance error term that is closely related to the double-sampling problem
LLD (m(), (((+1), ((()))
= (V(alm))) - [p(mm) + (o
(4)
Ep Ep (p, V(2)
To drive home that this problem is independent of the model error, we show that the (m, b 2 1)-VAML loss is uncali- brated even if we substitute the ground truth environment for the learned model, and focus solely on the value learning component.
- (Ø, V|x,,') + Vara (V())).
(3)
In the classic somplegsa problem, we cannot correct the variance term as we do not have oracle access to the en- vironment. Here the issue is our model, and we can generate multiple samples from it. Therefore, we can estimate this variance term from samples and correct the loss. This cor- rection reminiscent of Antos et al. (2008) but is simpler to obtain as we only need model samples.
Proposition 4. Let LVD(V xlm)) be the target loss, (PT,V and let E (P=, V/Vaar X, (((m),r(m)) be the surrogate
loss. There exists a set of functions V for any Viar that is not a constant function, functiio,, for which two conditions hold:
mm.k To obtain the correction, we define P ((( the empirical estimator of the ex- pected return in Equation 1. The variance of this estimator -
1. The set is complete, meaning that (Tpn Viarl E V for some target function Viar.
2. and,
16 m With this we can define a new loss which can be computed with at least two samples from the model
can be estimated as Var, -..m
(V
[ Vi x, a (m))] Z [Tp= Viarl-
Arg min Epp Vev Therefore, CP is minimum-uncalibrun
L
Varg 1.k
4

---

Calibrated Value-Aware Model Learning with Probabilistic Environment Models
IterVAML
MuZero
Model-free TD
Model
Env. det. env
(m 2 b = 0)
(m 2 1,6 2 1)
(m = 0,b 2 1)
Det. model,
calibrated
calibrated
calibrated
Stoch. model,
det. env
uncalibrated (resolvable)
uncalibrated (for VF updates)
calibrated
Det. model, Stoch. model,
stoch. env
calibrated
calibrated
calibrated
stoch. env uncalibrated (resolvable)
uncalibrated (for VF updates)
calibrated
Can update model
X
Can update VF
X
v (for det. model)
Table 1. Comparison of the major design choices in the (m, b)- VA-I. framework. VAMI. and MuZero can be used to update the model, while MuZero and Model-free TD learning can used to update the value function. All model-hased losses are uncalibrated when applied to stochastic model classes. In addition, MuZero suffers a bias when used for value function prediction which cannot be surmounted with an easy modification to the loss function.
The problem in the loss function again depends on the vari- ance of the value function with regard to the model. Proofs for this section can be found in subsubsection A.2.3.
regular model-based TD loss to update the value function. For the baseline, we similarly use a model-based TD loss together with a KL-based loss.
To overcome this issue, we can introduce another variance correction term, similar LCVAMI.. This retains the advan- tage that the MuZero loss is used for both model and value updates. However, as we average over the value prediction from the model, we only guarantee that
We use n to denote the size of the state space, for the number of successor states in the garnet, and j for the rank of the model. Every problem is generated by sampling k successor states for each state Ii, and parameterizing the transitions with weights wi, sampled iid from a standard normal distribution for all successor states T... The transi- tion distribution is given by p(x)|zz) = e/t,/7 /E; emt,//T Varying the temperature T we can interpolate between a deterministic transition and uniform one.
Eg (((amm)) ~ EEp [T (oto)()
(5)
not that the values for each state are correct. However, these are necessary for accurate planning or policy improvement. Therefore, it is still important to use a loss variant with m 0 for learning accurate values for each state.
The models are parameterized with two learnable matri- ces 4 E Rixn and vo E Ri'== so that Ps softmax (@v,,) = softmax(yp) vt). By varying we create a low-rank constraint. As Farahmand et al. (2017) shows, (m, b)-VAML should be used when the model has insuffi- cient capacity to represent the environment. As a baseline. we use the model ÅKL, E arg min, (Llp))), which we approximate with gradient descent. The matrices o and &o are initialized by drawing weights randomly from a normal distribution with a very small standard deviation.
4.4. Discussion of theoretical results
While the problem resulting from the use of stochastic mod- els is resolvable, the MuZero loss is still insufficient for learning per state value function with stochastic models. We observe issues with learning accurate value functions with the (m 2 1,b b 2 1)-VAML loss even after correctly calibrating it. If the calibrated loss is used to update the model, and standard model-free or model-based value estimate is used as the value function learning target, the correct value function can be learned.
We focus solely on value estimation in the Garnet MDPs with a fixed policy to simplify the experimental setup. We also use the ground truth reward function.
The numerical results are graphed in Figure 2. For the near-deterministic ground truth environment, we see no benefit from using a calibrated (m, -AML, We find that the algorithms are able to exploit the non-linear softmax function to achieve very accurate models in deterministic environments. However, even small amounts of stochasticity prevent this solution.
Results:
Insight (Calibrated losses). To obtain a calibrated loss, we propose using LCVAMI. to update the model and a regular model-based or model-free TD loss to update the value function.
5. Calibration impact on finite state MDPs
As the stochasticity increases, we see advantage for cal- ibrated losses. (1. 1A-V is not able to achieve good results even in a deterministic environment. As we initial- ize the models with high entropy, (1, VA-MLL is unable to learn a correct value function. This corroborates the
To test our findings, we run the (m, b)-MMM- losses on small, finite state Garnet problems (Bhatnagar et al., 2007). We use the MuZero-style (1, 1)- VAML loss to update both model and value function, while the Iter VAr---L-ttylle (1, 0)-
VAML loss is only used to train the model, and we use a
5

---

Calibrated Value-Aware Model Learning with Probabilistic Environment Models
Temperature: 0.01
Temperature: 0.1
Temperature: 10.0
V loes
10
10-2
10
10
10
10' Model rank
10
10'
Model rank:
Model rank:
Model rank: 50
V loss
10"
10"
10 10
10
10'
10-a
10
10
10'
10
10-1
10 10'
Temperature
(1.0)-CVAL
(1,0)VVAM
(L.I)-VVAL
(1,1)-VAML
MLE
Figure 2. Results for the Garnet experiments. In the top row, we show the mean squared error of the value function prediction using bottom row, we vary the temperature and show results for three values
different latent sizes over three different temperatures. k. Shaded regions are bootstrapped confidence intervals of the mean at 95% over 1000 independent problems. With the exception of deterministic problems (left), the CVAML loss reliably results in the lowest value prediction error.
theoretical finding that the error depends on the model's stochasticity and not on the environment.
are necessary to prove the existence and measurability of standard functions such as the value function (Bertsekas Shreve, 1978).
Calibrated (1, VAML still struggles in environments with low stochasticity. We find that the model is prone get stuck in local plateaus it reduces the value function dif- ference per state faster than the model predictive variance. This suggests that a (m 2 1, 1,0)-LL loss is preferable with stochastic environments.
Proposition 5. Let X be compact, connected, metrizable space. Let p be a continuous kernel from to probability measures over X. Let Z be a metrizable space, Consider a bijective latent mapping : Z and any V + R. Assume that they are both continuous. Denote Vx Vo
Then there exists measurable functtionn f Z + Z such that we have (o(x))) = Ep [Vx(æ(()) for all X.
6. Stochasticity and auxiliary models
All previous results hold independent of the model archi- tecture. However, in practice, most implementations of decision-aware models use deterministic latent model struc- tures (Schrittwieser et al., 2020; Ye et al., 2021; Hansen et al., 2022; Antonoglou et al., 2022).
Furthermore, the same f is a minimizer of the expected IterVAML loss over any distribution x ~
f E arg min Ep, [FuerVAML. Vx
Proofs for this section can be found in Subsection A.3. In general, deterministic function approximations cannot capture the full transition distribution in stochastic envi- ronments. However, it is open question whether de- terministic latent model can be sufficient for learning equivalettn model, as conjectured by Oh et al. (2017). We settle this question now.
We can conclude that given a sufficiently flexible function class F, (1, BVAML can recover an optimal deterministic model for value function prediction. Note that our condi- tions solely ensure the existence of measurable function; the learning problem might still be very challenging.
6.1. Deterministic models for stochastic environments
Note that the conditions for our proof here are slightly dif- ferent than those for our definition of an uncalibrated loss. Here we show that given a flexible enough function class, a deterministic model can sufficiently minimize the Iter- VAMI. loss. In the previous section, we showed that if the
Showing the existence of deterministic value-aware model relies on the continuity of the transition kernel and involved functions ip and V. These are standard assumptions that
6

---

Calibrated Yalue-Aware Model Learning with Probabilistic Environment Models
function class admits perfect stochastic model, but not perfect deterministic one, the (m, b--VAML loss will incur calibration error. Here we show the conditions under which perfect deterministic model exists.
their model and auxiliary loss rely on biased straight- through gradient estimation. This introduces complications for finding the optima of the loss function and model class.
Therefore, it is still important to have a calibrated surro- gate loss for stochastic models, as many use cases make stochastic models attractive. For example, a pre-trained LLM backbone might be used, which is naturally stochastic due to the sampling strategies used query it.
7. Experiments with latent-space models
We examine the impact of the calibrated losses on sub- set of DMC environments (Tunyasuvunakool et al., 2020) encompassing 7 total tasks across the humanoid and dog domains. These two domains are the most challenging in the DMC suite, and the standard comparison for current methods (Voelcker et al., 2025; Nauman et al., 2024).
6.2. Auxiliary losses for latent-space models
While the original MuZero algorithm only uses the MuZero loss update the latent embedding, dynamics model, and value function estimation, several more recent works add auxiliary stabilizing losses (Ye et al., 2021; Hafner et al., 2021; Hansen et al., 2024; Voelcker et al., 2025). These allow the model to learn meaningful transitions even before the value function is properly approximated, which helps e.g. in sparse-reward environments.
Architecture and Setup: We present a comparison between (m, b)-CAAMI. and (m, b)- VAMI. families of losses. To conduct our experiments we use two neural network archi- tectures for the model, stochastic and a deterministic latent model, while keeping the rest of the training setup consistent for a clean comparison. All experiments are conducted with latent model architectures composed of an encoder, a latent dynamics model, and a value and policy function head.
Most prior works use some form of Bootstrap-your-ove latent (B YOL) (Grill et al., 2020) loss, which minimizes the difference between the next states encoding (()(m) and the model prediction (*) In the following we write p) = # when we want to highlight the different components of the learned model.
For the stochastic case, the latent models are multivariate, diagonal Gaussian distributions where mean and variance are parameterized by the latent network (Chua et al., 2018; 8;; Janner et al., 2019; Paster et al., 2021). In the deterministic case we simply use a feed-forward network.
Model rollouts are produced by sequentially sampling latent states from the model conditional on the initial state and an action sequence using the reparametrization trick.
In the case of linear features, models, and value functions, recent works have shown that such a loss can greatly aid in learning meaningful features for value function prediction (Lyle et al., 2021; Tang al., 2023; Ni et al., 2024; Voelcker et al., 2024). When the difference is computed with an T.2 the auxiliary loss is
((2'|2,, a) ()z, + +(22,) e, e e~ N~0
(8)
The auxiliary loss for stochastic models computed as a negative log-likelihood of the next states' latent representa- tions under the current dynamics model.
(Uf+)))))) = (m(())))) | - [ (6)
2PT aux,m
Laux,m (p. sp x. a (00mm))) mm))
The introduction of a stabilizing loss poses a challenge for the result in Proposition 5. For a flexible enough model and embedding function class, the perfect model fäux would predict Ep- [x((')]=)) However, fäux only coincides with the optimal model under the VAMI. loss if
= log (m)( (o(())) a (0m--))
(9)
where (0:m-1) is a sequence of actions of length m start- ing from state z0()) In the case of deterministic models, the auxiliary loss is simply the MSE version introduced in Equation 7.
Epp [M(o(a())]= (Ep [EE))((((())))
The full model loss used is
ppm Emodel, (IF.), V1z))
This is the case if and only if the value function affine in the embedding features sp which is also referred to as linear expectation model (Wan et al., 2019). However, learning an embedding in which the value function is linear can be difficult and may not lead to stable model predictions in complex, high-dimensional environments.
2pr m.b (1),s), V
+
auxx (1))))
(10)
(m,,)) VAML
auxiliary
replacing f with e for the stochastic model version.
An exception to the use of deterministic models is the ar- chitecture proposed by Antonoglou et al. (2022). However,
For training the policy and value function, we follow the data-mixing protocol suggested by Voelcker et al. (2025).
7

---

Calibrated Value-Aware Model Learning with Probabilistic Environment Models
dog
humanold
Mean
IQM
Mean
IQM
(1,0)-OVALL (Stoch.)
(1,0)VAMML (Stoch.)
(1,0)VAMM (Det.)
(1,1)-CVAML (Stoch.)
(1,1)-AML (Stoch.)
(1,1)VAML (Det.)
0.75 0.80 0.85 0.90
0.78 0.84
0,90
0.96
0.75
0.60
0.73 0.90
Normalized Score
Figure 3. Results for the latent-space model experiments. We show the aggregate metrics of the final performance on dog (left) and humanoid (right) environments, following (Agarwal ct al., 2021). Calibrating (1, 1)-MML lcads to clear improvement in performance the humanoid environment, while (1, 0)-(C)AA the difference less noticeable. This is consistent with our theoretical findings as learning smaller variance model can be sufficient, but wrong value functions will be learned with stochastic model when b 2 1.
We use both model-generatle on-policy and real environ- ment data from the replay buffer to train value and policy head using the (m, b)- VAML loss with twinned critic parametrization similar to TD3 (Fujimoto et al., 8b). At every timestep, half of the next states in the minibatch are re- placed with model-simulated states resulting from on-policy actions, the other half ground-truie environment data. This means the model influences the value function both through the encoder and through on-policy model-generade data.
style loss. The dog domain is less affected by the difference in the calibration.
In addition to the calibration effect, we observe that prob- abilistic models outperform deterministic ones the dog domain, even though the simulator is deterministic. This is in line with previous work (Chua et al., 2018; Janner et al., 2019), stochastic models can reduce the tendency of the critic exploit model errors. However, this advan- tage seems to be Nomain-specific and is less pronounced in the humanoid suite, where deterministic models with a (1, VA-ML loss perform best overall.
The actions for environment interactions are obtained from the actor and the environment model using MPC (Hansen et al., 2022): initialized at the action produced by the ac- tor for the current state, this algorithm iteratively refines the action to maximize the expected return. In total our setup uses the model for training the shared encoder, gener- ating data for value and policy improvement, and for online model-based search.
Finally, we observe an advantage of (1, Ouppattes over the (I, 1)-updates. While previous work claimed that IterVAML is unstable (Lovatto al., 2020; Voelcker et al., 2022), com- parisons were not made with SOTA model architectures and auxiliary tasks. As additional experiments (see Figure 8) show, the Iter loss is stable even without additional auxiliary losses, although adding the B BOL loss is neces- sary to achieve non-trivia performance in humanoid tasks.
For (m, b)-CA, the calibration term is computed by sampling multiple trajectories from the model and calculat- ing the means and variances of next-state values produced by the critic across the different samples. Additional imple- mentation details are provided in Appendix C.
Insight (Remarks for practitioners). While deter- ministic models are theoretically sufficient for learn- ing calue-equivalent models, we observe benefits in some benchmarks from the use of stochastic models. Using calibrated losses is empirically especially im- portant for Mulzero-stye model and value updates.
Results: We plot aggregated performance over 20 random seeds with 95% CI, estimated with stratified percentile boot- strap (Patterson et al., 2024). Aggregations of final perfor- mance over several environments are visualized using the RLiable library (Agarwal et al., 2021).
The difference in performance is most noticeable between the (1, 1)-VAML and -CVAML losses on the humanoid benchmark, as the uncalibrated loss impacts both model and value function learning in this case. This effect is less prominent but a trend is still visible with (1, OVCAMM, where only the model learning affected.
8. Related work
Farahmand (2018) established Iter- VAML based on earlier work (Farahmand et al., 2017). Several extensions have been proposed, such as a VAML- regularized MSE loss (Voelcker et al., 2022) and a policy- gradient aware loss (Abachi et al., 2020). Combining Iter- VAML with latent spaces was first explored by Abachi et al. (2022), but no experimental results were provided.
VAML and MuZero:
In the humanoid domain, observe performance im- provement when using the calibration with the MuZero-
8

---

Calibrated Value-Aware Model Learning with Probabilistic Environment Models
dog
humanood
800
Reward
600
400
200
0.0
1.0
2.0
0.0
1.0 Step
2.0
Step
x10
x10
(1.1)-CVAML (Stoch.)
(1.0)-VAML (Stoch.)
(1,1)-VAML. (Det.)
(V.O).VVMLL (Stoch.) (1,0) VAML (Det.)
(1,1)- VALL (Stoch.)
Figure 4. Sample efficiency curve for both dog (left) and humanoid (right). Per-task normalized return is aggregated over 30 seeds per environment and tasks for the humanoid domain, and 4 - for the dog domain, with 95% bootstrapped confidence intervals shaded. In addition to higher final return, observe significantly carlier learning for the (1, -CCCVALL on the dog task.
MuZero (Schrittwieser et al., 2020; Ye al., 2021) is built based on earlier works that introduce the ideas of learning latent model jointly with the value function (Silver al., 2017; Oh et al., 2017). However, none of these works in- vestigate the calibration of the loss function. Antonoglou et al. (2022) propose an extension to MuZero in stochastic environments but focus on the model architecture, not the value function loss. Hansen et al. (2022) and Hansen al. (2024) adapted the MuZero loss continuous control envi- ronments but did not extend their formulation to stochastic models. Grimm et al. (2020) and Grimm et al. (2021) con- sider how the set of value equivalent models relates to value functions. They are the first to show the close connection between the notions of value-awareness and MuZero.
limit interaction with wrong models (Buckman et al., 2018; Janner al., 2019; Abbas et al., 2020). Several of these techniques can be applied together with (m, b), bVCAMLL to improve the model and value function learning further. Finally, we do not focus on exploration, but Guo et al. (2022) show how auxiliary losses can be used not only to stabilize learning but also to improve exploration.
9. Conclusions
We theoretically analyze commonly used value-aware losses such as the MuZero and IterVAML loss and show that they are uncalibrated surrogate losses. When using (m, b)- VAMI. losses, such as the popular VAML and MuZero algorithms, with stochastic environment models, the loss learns low variance models, even those do not recover the correct value function. Building on our proofs, we propose novel variant of the loss to stabilize learning with stochastic environments and evaluate its efficacy in practice.
Other decision-aware algorithms: Several other works propose decision-aware model learning algorithms that do not directly minimize value function difference. D'Oro et al. (2020) weigh the samples used for model learning by their impact on the policy gradient. Nikishin et al. (2021) uses implicit differentiation to obtain a loss for the model function with regard to the policy performance measure. To achieve the same goal, Eysenbach et al. (2022) and Ghugare al. (2023) choose variational formulation. Modhe al. (2021) proposes compute the advantage function resulting from different models instead of using the value function. Ayoub et al. (2020) presents an algorithm for selecting mod- els based on their ability to predict value function estimates and provide regret bounds with this algorithm.
Our experiments further show that the calibration of the (m, b)-VAM losses is important for obtaining strong learn- ing with stochastic environment models. In addition, while previous papers showed that IterVAML losses can be unsta- ble in practice (Lovatto et al., 2020; Voelcker et al., 2022), we find that this can be overcome by adopting the latent model architecture used by Schritwieesr et al. (2020) and the auxiliary losses established by Li et al. (2023); Hansen et al. (2022). When combined with a suitable value learn- ing procedure, IterVAML performs on par with or better than MuZero in continuous control tasks. Finally, while (m, b)-VAML losses have mostly been used with determin- istic environments, our work enables the community to stochastic models with calibrated loss and shows the po- tential merits of this approach in a number of environments.
Several works have focused on the broader goal of using models with errors without addressing the loss functions of the model. Among these, some attempt to correct models using information obtained during exploration (Joseph et al., 2013; Talvitie, 2017; Modi et al., 2020; Rakhsha et al., 2022; 2024), or to
Learning with suboptimal models:

---

Calibrated Yalue-Aware Model Learning with Probabilistic Environment Models
Acknowledgments
Ayoub, A., Jia, Z., Szepesvári, Cs., Wang, M., and Yang, L. Model-based reinforcement learning with value-targeted regression. In International Conference on Machine Learning, 2020.
We would like to thank the members of the Adaptive Agents Lab and the Toronto Intelligent Systems Lab who provided feedback on the ideas, the draft of this paper, and pro- vided additional resources for running all experiments. We acknowledge Amin Rakhsha for helping with the formal statement of Lemma 4. We would also like to thank the anonymous reviewers for providing valuable feedback to improve this work. AMF acknowledges the funding from the Natural Sciences and Engineering Research Council of Canada (NSERC) through the Discovery Grant program (2021-03701). Resources used in preparing this research provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute.
Bellemare, M. G., Naddaf, Y., Veness, J., and Bowling, M. The arcade learning environment: An evaluation plat- form for general agents. Journal of Artificial Intelligence Research, 47, 2013.
Bertsekas, D. P. and Shreve, S. E. Stochastic Optimal Con- trol: The Discrete-Time Case, Academic Press, 1978.
Bhatnagar, S., Ghavamzadeh, M., Lee, M., and Sutton, R. S. Incremental natural actor-critic algorithms. Advances in Neural Information Processing Systems, 2007.
Buckman, J.. Hafner, D., Tucker, G., Brevdo, E., and Lee, H. Sample-efficient reinforcement learning with stochastic ensemble value expansion. Advances in neural informa- tion processing systems, 2018.
Impact Statement
Our paper is first and foremost theoretical analysis of loss function used in reinforcement learning. As such, it presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
Chua, K., Calandra, R., McAllister, R., and Levine, S. Deep reinforcement learning in a handful of trials using proba- bilistic dynamics models. In Advances in Neural Infor- mation Processing Systems, 2018.
References
Deisenroth, M. and Rasmussen, C. E. Pilco: model- based and data-efficient approach to policy search. In International Conference on Machine Learning. 2011.
Abachi, R., Ghavamzadeh, M., and Farahmand, A.mm Policy-aware model learning for policy gradient meth- ods. ArXiv, abs/2003.00053 2020.
D'Oro, P., Metelli, A. M., Tirinzoni, A., Papini, M., and Restelli, M. Gradient-aware model-based policy search. In AAAI Conference Artificial Intelligence, 2020.
Abachi, R., Voelcker, C. A., Garg, A., and Farahmand, A.-m VIPer: Iterative value-aware model learning on the value improvement path. In Decision Awareness in Reinforcement Learning Workshop at ICML 2022, 2022.
Eysenbach, B., Khazatsky, A., Levine, S., and Salakhut- dinov, R. R. Mismatched more: Joint model-policy optimization for model-based rl. In Advances in Neural Information Processing Systems, 2022.
Abbas, Z., Sokota, S., Talvitie, E.,, and White, M. Selective dyna-style planning under limited model capacity. In International Conference on Machine Learning, 2020.
Farahmand, A.m. Iterative value-aware model learning. In Advances in Neural Information Processing Systems, 2018.
Agarwal, R., Schwarzer, M., Castro, S., Courville, A., and Bellemare, M. G. Deep reinforcement learning at the edge of the statistical precipice. In Advances in Neural Information Processing Systems, 2021.
Farahmand, A.m., Barreto, A., and Nikovski, D. Value- Aware Loss Function for Model-based Reinforcement Learning. In International Conference on Artificial Intel- ligence and Statistics, 2017.
Antonoglou, I., Schrittwieser, J., Ozair, S., Hubert, T. K., and Silver, Planning in stochastic environments with learned model. In International Conference on Learning Representations, 2022.
Fujimoto, S., Hoof, H., and Meger, D. Addressing function approximation error in actor-critic methods. In Interna- tional Conference on Machine Learning, 2018a.
Antos, A., Szepesvári, Cs., and Munos, R. Learning near-optimal policies with bellman-residual minimiza- tion based fitted policy iteration and a single sample path. Machine Learning. 71:89-129, 2008.
Fujimoto, S., van Hoof, H., and Meger, D. Addressing function approximation error in actor-critic methods. In International Conference on Machine Learning, 2018h
10

---

Calibrated Value-Aware Model Learning with Probabilistic Environment Models
Fujimoto, S., D'Oro, P., Zhang, A., Tian, Y., and Rabbat, M. Towards general-purpose model-free reinforcement learning. In The Thirteenth International Conference Learning Representations, 2025.
Joseph, J., Geramifard, A., Roberts, J. W., How, J. P., and Roy, N. Reinforcement learning with misspecified model classes. In IEEE International Conference Robotics and Automation, 2013.
Ghugare, R., Bharadhwaj, H., Eysenbach, B., Levine, S., and Salakhutdinov, R. Simplifying model-based RI.: Learning representations, latent-space models, and poli- cies with one objective. In International Conference on Learning Reprexentations, 2023.
Kearns, M. and Singh, S. Near-optimal reinforcement learn- ing in polynomial time. Machine learning, 49(2). 2002.
Lambert, N., Amos, B., Yadan, O., and Calandra, R. Ob- jective mismatch in model-based reinforcement learning. Conference on Learning for Dynamics and Control, 2020.
Grill, J.-B., Strub, F., Altché F., Tallec, C.,, Richemond, P., Buchatskaya, E., Doersch, C., Avila Pires, B., Guo, Z., Gheshlaghi Azar, M., et al. Bootstrap your own latent-a new approach to self-supervised learning. In Advances in neural information procexsing systems, 2020.
Li, Q. Kumar, A., Kostrikov, I., and Levine, S. Efficient deep reinforcement learning requires regulating overfit- ting. In International Conference on Learning Represen- tations, 2023.
Grimm, C., Barreto, A., Singh, S., and Silver, D. The value equivalence principle for model-based reinforce- ment learning. In Advances in Neural Information Pro- cessing Systems, 2020.
Lovatto, A. G., Bueno, T. P., Mauá, D. D., and de Barros, L. N. Decision-aware model learning for actor-critic methods: When theory does not meet practice. In "I Can't 't Believe It's Not Better!" at NeurIPS Workshops, 2020.
Grimm, C., Barreto, A., Farquhar, G., Silver, D., and Singh, S. Proper value equivalence. In Advances in Neural Information Processing Systems, 2021.
Lyle, C., Rowland, M., and Dabney, W. Understanding and preventing capacity loss in reinforcement learning. In International Conference on Learning Representations, 2021.
Guo, Z. D., Thakoor, S., Pislar, M., Pires, B. A.. Altché, E,, Tallec, C., Saade, A., Calandriello, D., Grill, J.-B., Tang, Y., Valko, M., Munos, R., Azar, M. G., and Piot, B. B BYOL-explore: Exploration by bootstrapped prediction. In Advances in Neural Information Processing Systems, 2022.
Mnih, V. Kavukcuoglu, K., Silver, D., Graves, A.. Antonoglou, I., Wierstra, D., and Riedmiller, M. Playing atari with deep reinforcement learning, In NeurIPS Deep Learning Workshop. 2013.
Györfi, L., Kohler, M., Krzyžak, A., and Walk, H. A distribution-free theory of nonparametric regression. In Springer Series in Statistics, 2002.
Modhe, N., Kamath, H., Batra, D., and Kalyan, A. Model- advantage optimization for model-based reinforcement learning. ArXiv, abs/2106.1408 2021.
Hafner, D., Lillicrap, T., Ba, J., and Norouzi, M. Dream to control: Learning behaviors by latent imagination. In International Conference Learning Representations, 2020.
Modi, A., Jiang, N., Tewari, A., and Singh, S. Sample complexity reinforcement learning using linearly com- bined model ensembles. In International Conference on Artificial Intelligence and Statistics, 2020.
Hafner, D.,. Lillicrap, P., Norouzi, M., and Ba, J. Mas- tering atari with discrete world models. In International Conference Learning Representations, 2021.
Moerland. T. M., Brockens, J., Plaat, A., and Jonker, C. M. Model-based reinforcement learning: A survey. Founda- tions and Trends Machine Learning, 16(1), 2023.
Hansen, N., Su, H., and Wang, X. Temporal difference learning for model predictive control. In International Conference on Machine Learning, 2022.
Munkres, J. R. Topology- Pearson Modern Classic, 2nd edition, 2018.
Hansen, N., Su, H., and Wang, X. TD-MPC2: Scalable, ro- bust world models for continuous control. In The Twelfth International Conference on Learning Representations, 2024.
Nauman, M., Ostaszewski, M., Jankowski, K., Mitos, P.,
and Cygan, M. Bigger, regularized, optimistic: scaling for compute and sample-efficient continuous control. Ad- vances in Neural Information Processing Systems, 2024.
Janner, M., Fu, J., Zhang, M., and Levine, S. When to trust your model: Model-based policy optimization. Ad- vances Neural Information Processing Systems, 2019.
Ni, T., Eysenbach, B., Seyedsalehi, E., Ma, M., Gehring, C., Mahajan, A., and Bacon, P.-L. Bridging state and history representations: Understanding preditivee

---

Calibrated Yalue-Aware Model Learning with Probabilistic Environment Models
rl. To appear International Conference Learning Representations, 2024.
Talvitie, E. Self-correcting models for model-based rein- forcement learning. In AAAI Conference on Artificial Intelligence, 2017.
Nikishin, E., Abachi, R., Agarwal, R., and Bacon, P.- L. Control-oriented model-based reinforcement learn- ing with implicit differentiation. ArXiv, abs/2106.03273, 2021.
Tang. Y., Guo, D., Richemond, P. H., Pires, B. Á., Chan- dak, Y., Munos, R., Rowland, M., Azar, M. G., Lan, C. L., Lyle, C., et al. Understanding self-predictive learning for reinforcement learning. International Conference Ma- chine Learning, 2023.
Oh, J., Singh, S., and Lee, H. Value prediction network. Advances in neural information processing systems, 30, 2017.
Tunyasuvunakool, Muldal, A., Doron, Y., Liu, S., Bohez, S., Merel, J., Erez, T., Lillicrap, T., Heess, N., and Tassa, Y. dm_control: Software and tasks for continuous control. Software Impacts, 6, 2020.
Paster, K., McKinney, L. E., Mcllraith, S. A., and Ba, J. BLAST: Latent dynamies models from bootstrapping. In Deep RL Workshop NeurIPS 2021, 2021.
Voelcker, C. Liao, V., Garg, A., and Farahmand, A.-m. Value Gradient weighted Model-Based Reinforcement Learning. International Conference on Learning Repre- sentations, 2022.
Patterson, A., Neumann, S., White, M., and White, A. Em- pirical design in reinforcement learning. Journal of Ma- chine Learning Research, 25, 2024.
Voelcker, C. A., Kastner, T., Gilitschenski, I., and Farah- mand, A.m. When does self-prediction help? Under- standing Auxiliary Tasks in Reinforcement Learning. In Reinforcement Learning Conference, 2024.
Puterman, M. L. Markov decision processes: Discrete stochastic dynamic programming. In Wiley Series in Probability and Statistics, 1994.
Rakhsha, A., Wang, A., Ghavamzadeh, M., and Farahmand, A.-m. Operator splitting value iteration. In Advances in Neural Information Processing Systems, 2022.
Voelcker, C. A., Hussing, M., Eaton, E., Farahmand, A.-m., and Gilitschenski, MAD-TD: Model -ugdee-d data stabilizes high update ratio RL. International Conference Learning Representations, 2025.
Rakhsha, A., Kemertas, M., Ghavamzadeh, M., and Farah- mand, A.-m Maximum entropy model correction in reinforcement learning. In International Conference on Learning Representations, 2024.
Wan, Y., Abbas, Z., White, A., White, M., and Sutton, R. Planning with expectation models. International Joint Conference Artificial Intelligence, 2019.
Schneider, G. Exploiting model uncertainty estimates for safe dynamic control learning. In Advances in Neural Information Processing Systems, 1997.
Ye, W., Liu, S., Kurutach, T., Abbeel, P., and Gao, Y. Mas- tering atari games with limited data. Advances Neural Information Processing Systems, 2021.
Schrittwieser, J., Antonoglou, I., Hubert, T., Simonyan, K., Sifre, L., Schmitt, S., Guez, A., Lockhart, E., Hassabis, D., Graepel, T., et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839), 2020.
Silver, D.. Hasselt, H., Hessel, M.. Schaul, T., Guez, A., Harley, T., Dulac-Arnold, G., Reichert, D., Rabinowitz, N., Barreto, A., et al. The predictron: end-to-end learning and planning. In International Conference on Machine Learning, 2017.
Steinwart, I. and Christmann, A. Support vector machines: Information Science and Statistics. Springer New York, 2008
Sutton, R. S. Integrated architectures for learning, planning, and reacting based on approximating dynamic program- ming. Machine learning Proceedings. 1990.
12

---

Calibrated Value-Aware Model Learning with Probabilistic Environment Models
A. Proofs and mathematical clarifications
We provide helper lemmata for Section and Section 6 this section.
All of our proofs rely heavily on standard expansion technique which is used prove that sample-based losses correctly approximate population losses. This is found in standard textbooks on learning theory such as Györfi al. (2002).
[I((=) - =1²] with an expected target (x). In our case, this is mostly the
It proceeds by expanding a loss Ba-p..r-g (m) minimizer of the target loss when evaluating surrogate losses. Then we obtain
E (t [[((e) - u1² Em-psra(taa) [If(x) (1(()) (z)) -
and continue expanding from there. Issues generally arise when y depends on f(x) in some way, or when f(x) itself involves sampling procedure.
A.1. Main propositions: Section
Proposition 1. Farahmand (2018) Let pa be the policy-conditioned transition kernel of an MDP and ler V :x R bea
function. Let p(-(a:) be model that fi (f, V (r) Then (Tp- V1 (a) = [TpV]
Proof. Assume ý fulfills fieervAML, (i, k1) operator, have
0. Then Epp [V (+(4) = -Epr [V (z(t1)] By definition of the Bellman
[P'))(>= = + (1 =
r(x) ( (0 (Vx(0)) = [T*)]<x).
(11)
A.2. Main propositions: Section
1. HELPER LEMMATA
For the following results, we will use the following notation. Let X' be a discrete sample space and p distribution over it. Let
f: + R be random variable, Let p be the set of distributions for which both p Arg min,. (f and for all
p" E P E,, [f(z) = E, [=)]. Let & be an integer. Finally, let g()) Exap (Ey [f(v)] f())}² + _ (f)g
We assume the following condition on distribution p and a function f for all the lemmata in this section. Assumption 1. Let p be a probability distribution over X with E, [=)], for which no exists that f E, (f()) This is assumption both and p.
This is an important assumption as it guarantees that there is no distribution with O variance such that E,y [f())]
Ep [f()]] This excludes a corner case of our proof, as fully deterministic environments and models do not lead to an uncalibrated Iter VAML loss.
We now obtain three simple lemmata about the minima of the function g- Lemma 1. There does not exist distribution p' such that p' & P E. [f(z)] E, [f(x)), and g(p') S g(p") for any p" P.
Proof. To prove the lemma, first evaluate g for any distribution 5" with Ee- [f(x)] = E, [f(x))
9(E*)= Hgpp (By-g: [f(v)] - f(x))) Varg. [f())]
(12)
Erup (Ey-p f(x)) + Varg- [f())
(13)
Varp(f(x)) + k Varg. [f()))
(14)
As we constructed the set P* so that all members have equal (minimum) variance, we obtain the same (())) for all p² Pt
13

---

Calibrated Value-Aware Model Learning with Probabilistic Environment Models
We now up contradiction by assuming p' exists. By this assumption
s(v') < g(p")
(15)
Var,(f(x)) [f(]] < Var, (fr,( + F -aar--1f(0) [f()]
(16)
Vary [x() < Var, [f())
(17)
For the function p' to cxist so that g(p') < g(p"), we would therefore require Var, [f(]]] < Varp- which is impossible by the definition of p This is contradiction with the requirements for p' in the lemma, which concludes our proof.
In the following lemma, we show that under some conditions distribution q exists which has g(g) g(0)) and for which Ep [x)) * Bg [f(x)] We can rewrite g
s(a) = Bp (B, [f(v]] - f(=))² + [ar,If(x)]
(18)
p [[(x}]] 2T, [xx) Egg [fx))) + B, [x)) + k Var
(19)
S SEp [[(x)]] 2T, [(x)) F,g [ + E, [xx) + Var,[[(x)]
(20)
-ar,If((=) h Var, [f(]] (E [f())) - B. [f()))
(21)
Equation 20 follows from Jensen's inequality. Intuitively, the function g depends on the squared deviation of the expectation
and the variance of both p and q. If the variance of q can be reduced more than the squared deviation of the means, then p* will not contain the minimum of g-
Note that the conditions y are sufficient but not necessary, as we use Jensen's 's inequality to obtain our bound. In addition, to simplify the proof, we assume that is a distribution with zero variance. In practice, any distribution with k+1 [f)] + (E, [=)] - E, (=))) < Varp-If [f)) and Ey [f(x)] * Ep [x(] will suffice, but this requirement is somewhat less intuitive to grasp. As our definition of calibration does not require us to exhaustively characterize all cases for p. p", and q, we have chosen this set of conditions which simplifies the proof. As our Garnet experiments show, many randomly generated transition distributions admit distributions where the minimizer of g does not match p in expectation.
&
Lemma 2. If there exists distribution q with (E,, [f)) Ep [f(r)]}² < Var,, [f()), and [f)] 0, then g(g) < g(p") ) for all p" and, by the assumptions on p and f. Exng [f(z)] + Ezup [z)].
Proof. Choose any p" E P' As [fr][((x)] = 00 = Ezup [E B [f (x))). We can now decompose g(4) as
((q)= Ep [E, [f(y)] - Ep [f(x)}]] - 2E, [f B. [f(x) B, [f()) Varp[[(x)] + (LE (f B, (f))))
(22)
(23)
Var,[f(x)] + + Var,- [f(x)] =((²).
(24)
Equation 24 follows from the assumption on q.
By the assumptions on p and f there does not exist a E X so that f(x) R. [fx)]. However, as Var,[f(x)] = 0. E, [f()]] = f(x)) for some E X. Therefore, E, [f()] + E, [f)], which concludes the proof.
As consequence, we obtain the following, final lemma. Lemma 3. Let P be the set of all distributions & so that E, [f(x)] = Egg [f(z)) Assume satisfying Lemma 2 exists. Then arg ming g(E) g PB.
Proof. The lemma is a direct consequence of Lemma and Lemma 2. By Lemma 1, for all p' E p3 g(p') 2 g(p"). As a consequence of Lemma 2, there exists at least one g with g(g) < g(p") and that q has E. [f(x)] + Ep [f(] Then the statement follows directly as 9 & PE.
14

---

Calibrated Value-Aware Model Learning with Probabilistic Environment Models
A.2.2. MAIN RESULTS: ITERVAML
Proposition 2. Let L V Var, VIr, x(mi be the surrogate loss for C LiterAMM. (b. V|2). Let Pe be the set of all distribu- tions for which Ep [V(elmi)] = Epr [V(xlmi)] There exist an MDP and a class of distributions P with P E P CP so that
Arg Ep [mm21ol0, V (/lc,llmm)) & p
PEP
Therefore. Lmodel is minimum--umcllibraate
Proof. Expanding the empirical IterVAML loss with & samples, obtain
Fm21. (p. Vlr,,, ((m)
Epn
(25)
(I(m)) (v(())))
V (afmi)) 2
= p+
(26)
v (m(m)
Bp v (s((m) ]] + 1Ep ( (s)m))) V (z(m²) )
px
(27)
v (s(m))
BBp v (m)))) )) 27
px pr
(28)
1p V (s/mi)] v (z!mml) )
(v (2(ml) Ep v [a(mii]
2 Bx
(29)
=0
(Ep v (a(mre) v (a))
)
27
Ep,pr
(30)
= Epppo (v(25
EB V (a)) ))) 27
(Epp [v (stmi)) (2)(me) 27
+
Ep.po
(31)
=Var
(2)
(ifm)) Ex IV [i(m))]] = O, and samples from f and pr are independent. Note that can separate the first factor and the second factor since the first one only contains (mm) random variable, and the second one only contains (mm) The second factor only includes ((m² inside of an expectation, and the expected value is not random varaible anymore.
Equation 29 is O since Eg,P= [V
The final term can again be decomposed as
Eg.p*
(Ep
V
(m(m)) } V (xfm))) 27
(32)
Ep V ( }] - Epo v (x(me)
+ Epr [r (almi)) V (ac(me)) )) (Epr [v (artmm) ]] - V (x((m) )) ))²}
Ep.pr
(33)
v (m)))) ] Epp vv (z((m))
27
Ep.p
Ep
Ep.pr
(34)
=IterVAN VAML
independent of f
The middle term of the binomial expansion again vanishes as Ep- and pd are independent.
EEp
[V (zo(ma))
V (z(m))] = 0, and samples from ò
When we drop the term that is independent of the model, the loss has the form of g in Lemma
Epr
[ (p. V Vlzc, (m)))
(35)
(v (2)6m) - EB V (s(m))] 27
EEp V ) Epn [v (arlm)) 27
= Ep.p.
+ Ep.p-
(36)
VVar
=IterVAML
15

---

Calibrated Yalue-Aware Model Learning with Probabilistic Environment Models
Let pa((12) be distribution that Lemma 3 holds. Then there exists with m (e. V lz, xl)) f (P*, Vla, V|r,a!!m) al)) and Ep [V (z(m))] + Ep- [V(x())]
If we use the & sample version of the empirical IterVAML loss instead, obtain the following decomposition
( (p, V|rr (l(r,
Epn
(37)
(V(((m)
Ep v (s(mi) ))² Ep [(Ep) V (a(ml) v (z((m) ))²
Ep
(38)
$ Var
IterVAML
The only difference when using more samples that the variance term is scaled by the factor E.- as the variance of the mean estimator Σ V [2 (mm) decreases. Consequently, for larger values of k; the condition in Lemma 2 for g becomes stricter.
As > oc, the condition becomes unfulfillable. This is also intuitive, as x E V (z.") almost surely (assuming standard conditions hold).
R& [V (2/(m)) as k - oo
Proposition The variance-corrected loss LCVAML, (P, vl),)(m)) is a calibruted surrogate loss for L LiserVAML. (P, (1,,,(()).
Proof. Reusing the previous derivation, we obtain
Epr [ (P, V/22, x(1))
Ep- [E2T,0 (p. V|2, Vizr,o(10 mm(
1
Ea, (v(()l
(a(m))
Ep v (H(m)) )) 27
(39)
k
Vaar
(EEp v (a(m))
] Eps
V (z((m)))
Bp
(40)
Iter VAML
Therefore, the surrogate loss and target loss are equivalent in expectation.
A.2.3. MAIN RESULTS: MUZERO
The following lemma shows that function that minimizes quadratic and variance term cannot be the minimum function of the quadratic. This is used to show that the minimum of the MuZero value function learning term is not the same as applying the model-based Bellman operator.
Lemma 4. Let X + R be function that is not constant and let μ. be non-degenerate probability distribution over X. Let (f) = Exey (f(x) - g(x))²
+ Ezza [f))g(=)] Ezzn [f()) E, [g(z)]. There exists a function space F with
g E so that g & arg min C(f).
Proof. The proof follows by showing that there a descent direction from g that improves upon L. For this, we construct the auxiliary function g(z) = g(x) =g(x). Evaluating L(ğ) yields
E E [9(x)² + E, [(((x) - =g(x)) 9(x)]
Ep [(((=)) =g(x))] E,, [g(x)]
= RR [=] (1 - s) E, [o(x)²] - (1 - -s) (2
Taking the derivative of this function wrt to yields
- E [o(x)*] + (1 - E,, [9(x) - (1 E, [o(x)))
de
2s E, [g(x)*] En [g(x)²] E, [=(x))
16

---

Calibrated Yalue-Aware Model Learning with Probabilistic Environment Models
Setting e to O, we obtain
B [x(x)]² Ep [x(x)?] = Var [o(2;)]
By the Cauchy-Schwarz inequality, the variance is only for a g(x) constant almost everywhere. However, this violates the assumption. Therefore there exists an E > 0 so that & (g) S C (g). We now can construct the function space F so that includes at least and + Eg.
Proposition 4. Let CTD "VD(V[Vuar ml) be the target loss, and let f (P", VIVaar I, (m(m))r(m) be the surrogate loss. There exists set of functions V for any Viar that is not constant function, for which two conditions hold:
1. The set is complete, meaning that [Tpa Vearl E V for some target function Viar
2. and,
Arg min Epx [mm (PE, V tars Fr tr slmi) g varl-
Vev
Therefore, (n, is minimum-umcalibrated
Proof. By assumption, let ş the MuZero loss be the true transition kernel p. Expand the MuZero loss by and take its expectation:
Tp- Viar
(z(())
Epn
f Vivar x., x: sfmis]]
(41)
v (x(m))
rr (x(m)) + Voar (slmati)]]
= pp
(42)
v (m(m))
To Vear) (alm)
Tpp Vear (s(m)) - [[r ( (a(mm) + Vear (afmm1)) )]"
Ex
+
(43)
(atm)) (fatm))
(stm²))
27 +
= Ep.po
Tor Viar
(44)
(((((m)
- Tip Vear (s(m)) )) TT Vear (I(m))
[r (2((m)) + Vear (z (z((+ 1))] + (45)
2 Ep.pr
[(Tion ([TonViar] (2fmm)) - (x)(m)) + year (2lomat))] ))
Ep.pn
(46)
We aim study the minimizer of this term. The first term (Equation 44) is the bootstrapped Bellman residual with target Vaar- The third term (Equation 46) is independent of V. can drop it when analyzing the minimization problem.
The second term (Equation 45) simplifies to
Epppa v (a(m)) ([TRm Viar (z(m)) - [r (I(m²) + Viar (2(m+1))
(47)
as the remainder is independent of V again.
This remaining term however is not independent of 인 and not equal to either. Instead, it decomposes into a variance-like term, using the conditional independence of &(1) and (1)) given (.());
Ep.p
v
(z(m))
Tpp Vear
(f(m)
(m(
+ Vear (a(mmtt)])
(48)
Esp
v
(m()
Tpp Vear
(m)))
Ep.rr
vv (2(1)
(r (2)(m) + VVar (1(m+t)))
(49)
Ep
aa (m)
Tp- Viar
(p()))
Ep v (a(m)
Epr r (zlmm) + Vear (m(++1)
(50)
17

---

Calibrated Yalue-Aware Model Learning with Probabilistic Environment Models
Combining this with Equation 44, we obtaín
Bx
(im VIVears I, m))))
(51)
Epppr (iv (almi) (Tp= Viar) () BB [V (s(m) Bpp [[r (alm)) + 7 Year (alm+1))
+
(52)
Bp V (stm) (Tre Vear (atm)) -Ep
(53)
The first summand the Bellman residual, for which the only minimizer is Tps Viar- However, by Lemma 4, we can construct function class so that Arg min Ep- [ (P", VV 2, xm)) E {Tpr Viar)
While our proof only discusses the case b = 1, the same issue also appears with larger b. In many cases, the problem will be exacerbated by longer rollout horizons m and b, as the variance of all functions involved grows with the time horizon.
A.3. Main propositions: Section 6
A.3.1. PROPOSITIONS FROM BERTSEKAS & SHREVE (1978)
For convenience, we quote some results from Bertsekas & Shreve (1978). These are used in the proof of Lemma 5. While some of the definitions are rather technical, it is mostly sufficient to see a stochastic kernel as the continuous generalization of the transition matrix in finite MDPs. The projection used Proposition 2 is simply a restriction of set of tuples (æ, y) on z values. Other topological statements are standard and can be found in textbooks such as Munkres (2018). In following refers to sets of continuous functions.
Proposition (Proposition 7.30). Let X and y be separable metrizable spaces and let q(dy|z) be a continuous stochastic kernel on given X. f C(X Y), the function X X -> IR. defined by
(xx / f(1,)))(d)(z)
continuous.
Proposition 2 (Proposition 7.33). Let X be metrizable space, y a compact metrizable space, D a closed subset of X x y, D. = [ul(x, y) E D), and let f D + TR. be lower semicontinuous. Let f (D) + R." be given by
(*
min f(x v). VED,
Then x (D) is closed X, f is lower semicontinuous, and there exists a Bevel--meauurable 4p proj + y such that rangely) C and
[xx] Vr (
In our proof, construct l the minimum of IterVAML. style loss and equate &p with the function call f in proof. The change in notation is chosen reflect the modern notation in MBRL in the textbook, older notation is used.
A.3.2. HELPER LEMMA
The first proposition relies on the existence of a deterministic mapping, which we prove here as a lemma.
Lemma 5 (Deterministic Representation Lemma). Let X' be compact, connected, metrizable space. Let p be a continuous kernel from X to probability measures over X. Let Z be metrizable space. Consider bijective mapping 4p X + Z and any V > TR. Assume that they are both continuous. Denote Vx = V
Then there exists measurable function f + such that we have (((((<(z)) = E,
all X.
Proof. Since @ is a bijective continuous function over a compact space and maps to a Hausdorff space (Z is metrizable, which implics Hausdorff), it is a homeomorphism. The image of X' under s. Zx is then connected and compact. Since X' metrizable and compact and y is homeomorphism, Zx is metrizable and compact. Let ( = (ppp))()))) (x('))
18

---

Calibrated Value-Aware Model Learning with Probabilistic Environment Models
Then, x is continuous (Proposition 1). function 2 can be represented composition of continuous functions Pv,z θv,x o yp and is therefore continuous.
Define 0y,x 0v.2 g. Since yp is homeomorphism, sp is continuous. The
As 2x is compact, the continuous function V takes maximum and minimum over the set Zx. This follows from the compactness of Zx and the extreme value theorem. Furthermore Viin I (z) I Vmaxx for every 2 E Zx. By the intermediate value theorem over compact, connected spaces, and the continuity of V, for every value Vmin I v S Vnax. there exists a 2 E Zx so that (z)) = v.
Let h: 2x x xx + be the function 222 [0v.z(z) - V== As h is a composition of continuous functions, it is itself continuous. Let h (z) = ming'ezy M2),z'). For any 2 E Zx, by the intermediate value argument, there exist 2' such that V(=')= = v. Therefore h" (z) can be minimized perfectly for all 2 Zx.
Since Xx is compact, h is defined over compact subset of Z. By Proposition 2, there exists measurable function f*()) so that ming h(z, =²) fzz, l (z)) = 0. Therefore, the function f has the property that (=)( Ey (=)=)
this minimizes the function h.
Now consider any x E X and its corresponding = = (xx). As h(2, (0v. 0 for any 2 E 2x. (f (f(x))) = 2 = Ep [Vx as desired.
A.3.3. MAIN RESULT: DETERMINISTIC MODEL
Proposition 5. Let X be compact, connected, metrizable space. Let p be continuous kernel X to probability measures over X. Let Z be metrizable space. Consider bijective latent mapping & X Z and any V Z - R. Assume that they are both continuous. Denote Vx = =V ob.
Then there exists a measurable function 8 Z + Z such that we have V(F((()))) EEp [Vx for all x E X.
Furthermore, the same f* is minimizer of the expected IteVAMM loss over any distribution x ~ p
S E arg min Ep,p- (f Vx
Proof. As we only deal with single steps here, we use ((1) and x' interchangeably as the latter is simpler and more common nomenclature. The statement of the proposition itself is written with the x(1) notation to remain consistent with the main body of the paper.
The existence of '" follows under the stated assumptions (compact, connected, and metrizable state space, metrizable latent space, continuity of all involved functions) from Lemma 5.
First, expand the equation to obtain:
Epp (J.
(54)
Epp [IV(F((()))) - -xx(a)))))V((((( Vx(
(55)
Ep- [(V Ep-
Ep- [Me(co(ig]] + Eps Vxx VV(el-Va(())))) Ve((((o))) )²)
(56)
After expanding the square, we obtain three terms:
Eg.p
[GuevvaLL (f,, Vf,V
-Ep,p [IV (f((())) Ep- [V((1')11
(57)
2mp=
[[V (f (+(x))) Ep- [Vx(2')) Ep- [Vx(=')] Vx()
(58)
Ep-
Epp -Vx(a') (=')
-
(59)
19

---

Calibrated Value-Aware Model Learning with Probabilistic Environment Models
Apply the tower property to the inner term obtain:
2Ep [IV (f (((x))) Ep- [Vx(<')]] Ep- [V<(z')] Vx((')]]]
(60)
2pp
[f E [Vx(2')] Epn (pp, Vx(a0) Vx(x |2] 0.
(61)
Since the statement we are proving only applies to the minimum of the IteL loss, we will work with the arg min of the loss function above. The resulting equation contains a term dependent on f and one independent of f:
VV (f(())) Ep- [Vx(<')]]
Ep- [IEp~ [ Vx(2'))
arg min Ep-
(62)
= arg min Ep-
VV (f (f(x))) - Ep- [Vx(<')11
(63)
Finally, it is easy notice that V (f" (p(x))) Ep- [Vx(<')) by the definition of f" Therefore f minimizes the final loss term and, due to that, the IterVAML loss.
Comparing these results to Lemma 3, we can clarify one potential confusion about our results. The assumptions which are required for Proposition 5 ensure that the condition in Lemma 2 cannot he fulfilled in a continuous space, as there always exists a single point y for which f(y) = Ep [=)] In this case, the optimal deterministic IterVAML model f* and the minimum variance distribution p' coincide. However, our definition of calibration requires the loss to recover optimal minima for any function class that includes them.
20

---

Calibrated Value-Aware Model Learning with Probabilistic Environment Models
Hyperparameter
Value
Hyperparameter (cont.)
Value (cont.)
Batch size
128
Model rollout depth m
Discount ?
0.99
Model bootstrap depth b
Varied (0 and 1)
Actor learning rate orr
0.0003
Model samples A
Varied (1 and 4)
Critic learning rate aq
0.0003
Proportion real p
0.9
Model learning rate Op
0.0003
Latent dimension
512
Encoder learning rate ap
0.0001
Gradient clipping
10
Table 2. Hyperparameter of the deep learning experiments. Aside from the (m, b)-CCVAM relevant hyperparameters, we keep all others
consistent across environments and algorithms.
B. Implementation Details Section 5: Garnet
Each Garnet is defined by a randomly generated transition matrix and random reward function. Across all experiments we state spaces of size [ 50. For each state Vi, sample k 10 successor states random without replacement. Then assign each successor state Iy weight Wij (O, 1) and set the weight of all non-successor states to the floating point minimum. The transition matrix is computed via the softmax function. The reward is sampled as r r(z;) ~ NO0 1).
This setup allows us to vary the stochasticity of the environment naturally. The softmax temperature parameter T controls the stochasticity, As r + oo, the problem becomes deterministic with pıx)|æ;) = 1 for the state æ, with Wiy = arg max, Wii+ For + oo, the probabilities for all successor states become equal p(zy)az (xx) - In our experiments, we vary T such that for almost all Garnets we obtain full determinism and equal transition probabilities, which empirically happens in the range T E 10.001,0.00)
C. Implementation Details - Section 7: DM Control
C.1. Architecture
The original MuZero paper (Schrittwieser al., 2020) provides large-scale experiments the ALE benchmark (Bellemare et al., 2013). However, the original algorithm and the ALE suite are computationally extremely expensive to run. In addition, to the best of our knowledge, no reliable open-source replication of the closed-source MuZero and Probabilistic MuZero models and algorithms exists. Instead, we use the TD-MPC architecture and algorithm (Hansen et al., 2022) as our main reference. This implementation uses the MuZero loss but is significantly cheaper and faster to run in a small-scale experimental setting, while still solving non- non-trivial control tasks.
We adopt the architecture and model-based MPC search procedure of Hansen et al. (2022) as baseline model. In addition to the Le based deterministic latent setup, we also use Gaussian latent model as our stochastic environment model. This implementation follows similar ideas used in observation-space models, such as Chua et al. (2018); Janner et al. (2019). While the stochastic environment models used in Dreamer (Hafner et al., 2021) would have been an alternative, it has not shown strong performance in continuous control tasks and is significantly more expensive to run. In addition, Voelcker et al. (2024) shows that observation-space reconstruction can be a suboptimal auxiliary task for aligning a latent space with value function prediction. Finally, we adopt the joint model-free and model-based training proposed in Voelcker et al. (2025). For the MuZero setting 1) we update the value function both for the state and the predicted latent state. In the IterVAML setting, there is no meaningful way to update the model's 's predicted value function, and therefore only update the current state's 's value function.
Code is provided in https: / / thub. com. haat//ddptiieee agents seeentsslabb//a CCAML.
Hyperparameters can be found in Table 2.
C.2. Environment
As several recent papers have shown, the vast majority of environments in the DM Control suite are effectively saturated. We, therefore, restrict our experiments to the most difficult domains, humanoid and dog, on which large performance differences still appear (Nauman al., 2024; Voelcker al., 2025; Fujimoto et al., 2025).
21

---

Calibrated Value-Aware Model Learning with Probabilistic Environment Models
C.3. Additional results
In addition the raw performance graphs in the main paper, we present the average model entropy and VAML loss curves in Figure 5. As expected. we find that the CVAML loss leads to higher entropy model compared to the VAML loss. We compute the empirical differential entropy of the model prediction over the replay buffer
Σ p(<(lp(z), a) log Z,EED
with z' sampled from the latent model. This however does not lead a noticeable difference in the VAMI error itself. We also present the individual reward curves for each task in both domains in Figure 6.
Metric = Model Entropy
Metric
VAML
Valne
-200
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Step
x 10
Step
x10
(1,1) CVAML
(1,1) VAML
(1,0) VAML
(1.0) CVAMML
Figure 5. Model entropy (left) and loss value for the VAML-losses (right) aggregated across all environments. The model entropy of (1, 0)-CAAML differs significantly. However, this does not translate to a pronounced difference in the VAMI. error itself. Therefore, while the calibration term does lead a higher variance, this does not necessarily translate into an improvement performance. In (1, 1)-VAML compared with the calibrated version, the model entropy does not differ significantly. However, without calibration, the VAML error seems to be growing slightly stronger. Given the large deviation in performance we observe in the main results for (1, 1A-M, we conjecture that correcting the value function learning term is significantly more important for performance.
dogustand
dog_walk
dog_trot
dog.rrm
1000
750
Reward
500
0
Step
x10
humumoodd
Inmanaid
lumanoid
1000
(1,1)CVAML (Stoch.)
750
(1,0)CVAML (Stoch.)
Reward
(1.0)VAMM (Stoch.)
500
(AMO) (Des.)
(1,1)-VMML, (Des.)
250
(1.1)VAML, (Stoch.)
2
2
Step
x10
Step
x10"
Step
100
Figure 6. Episode returns per environment averaged over 30 seeds.
22

---

Calibrated Value-Aware Model Learning with Probabilistic Environment Models
C.4. Ablations
We provide additional ablations on the model loss here.
In the first set of experiments (Figure 7), we remove the CVAML loss from the model. We observe deteriorating agent performance, especially with the stochastic models, and in the humanoid tasks. In addition, we compare against a model-free TD3 baseline that uses the same architecture as our model based versions.
In the second set (Figure 8), we ablate the auxiliary loss. Again, we find that the auxiliary loss helps especially in the humanoid experiments. In addition, the (1, 1.-AMML baseline suffers significantly more from the removal of the auxiliary loss than the (1, (1.0)) loss. Upon further investigation we observe that the bootstrapped TD target of the (1, ().AAAC loss leads to less stable learning compared to the value function target.
dog Stand
dog.walk
dog.
dog
1000
750
Rewaad
300
230
0
Step
x10
homsenoid_stand
humoncld walk
humonold_ym
1000
750
(1.0)-CVMML (Stoch.) VAML (Stoch.
Rewaed
300
VAML (Det.)
TD3 (Model free)
250
0.
2 Step x10
0
Step x10f
Step
x10"
Figure 7. Performance of model-free baseline (Fujimoto et al., 201 8a), and the model-based algorithm using only the auxiliary loss for model training (No VAML), compared to (1,0)-CVAML. The model-free haseline fails to achieve strong returns on any problem. The No VAML baseline achieves slightly better returns on the dog task, but is unable to achieve any performance in the humanoid tasks.
dog.stand
dog.wolk
dog.troe
dog.rm
1000
750
Roward
250
Stop
X30
Iumassidastand
Immauodd.oalk
Inmamodd.zus
1000
750
Rewid
(100)-VAM,, No (Stech.)
(1.I)-CVABLL No Aux Loss (Stoch.)
(1.0) CVAML (Stoch.)
250
Stop
X10
Step
XII
Step
X10
Figure& 8. Performance comparison hetween agents trained using the (1,0)-CVMML and (1, 1)-CAAM. without the auxiliary loss to (1, (VVAM- trained with the auxiliary BYOL-style loss. The inclusion of the auxiliary loss is crucial for performance, both (1, 0)- CVAML and (1, 1-11)-CC --CAMML without the auxiliary struggle on all the dog tasks, compared to the agent trained with both losses, the effect is even more pronounced for the (1, 1)- CV-CALL variant. On the humanoid task both variants without the auxiliary loss exhibit a complete failure, which further exemplifies the importance of this loss.
23

---

Calibrated Value-Aware Model Learning with Probabilistic Environment Models
C.5. Garnet Policy Iteration
Here we extend the Garnet experiments in Section 5 to the control case. In each iteration, the policy's is value function is estimated with the model-based losses, following the Garnet experiments in the main paper. The policy is improved in greedy fashion. The problem is a modified 5x5 cliffwall environment https:/ / /gyynaaslsss i um. farama org/ anvironmente / ttt_text//ttttt/tt/// / cllffi cliff_ weiff__waaknn walking/ where the likelihood of moving in the intended direction is given by
temp. Results are presented in Figure 9
Results are consistent with the garnet experiments in the main paper, except for the surprising outlier of the uncalibrated (1, 1)-AML loss in the almost-deterministic case. This might be due to the different structures of the cliffwalk environment compared the more diverse Garnet environments. However, this only holds for low rank deterministic case (top right corner).
Triree 6.55, MR:
Teapge 666 MR:
Terrye (100), MIT-
0.2
CIAAE
1.16
(1-)) VAME (saner)
CVAML
(L21 CHAAE
(L. VAME. (ssoee)
EVAMI
(1.13 CIANNI
MELE
VASIL
0.34
2.0
2.5
003
1.0
6.58, MTT:
Teap::
3ET
[3AG
12.3) CHAML
[100 VAME (
0.30
(1.1) VAMI (oncer)))
222
HLSS
(LAI CIVAN
CHADEL
0.20
(1.1
CHAME
0L1 VIXI (KEOe:))
(UA VAMI, (sanee))
(11) VAMEL (SEDNC))
(1:3 VANIL (Sacer))
0.5
MILK
0.10
2.5 Trape: 0.83, MR. 10
4.0
Temp:: 0.60, MER: 10
Teapp: 0.05, ME: 10
8.40
12.1 CSAML
11A0 CVAMG
VAME-
2.3
VAHI
8 0.30
CLU CHAME
1316 CHAML
0.20
(1.1 CINANE
CHAMM.
VAMI,
VAMI,
0.15
0.10
0.0
3.35
2.1
311
4.3
012
2.5
30
030
1.5
3.5
301
4.0
Ternge 0:33, MII: 15
Terag: 9.66, MT: R5
Tempee 0.95, MR;: 35
0.40
0.20
3.3
8.30
3.5
2.50
ORAME
0.20
[1:3 COAME
1.3
CHAME
IC VAMI (ow0n)
(1-3) VIAM (enno))
(1:0 VAMI. (sncer))
0.15
[LJ] VADEL (1202)))
(L1 VAKEL (ssovr))
(1.3 VAMIL (sammm)) ML.
MUS
0.5
Figure 9. Return curves for cliffwalk policy iteration. Each curve shows confidence intervals over 1000 seeds. Each row is a different model rank (sec description in the main paper), each column a different temperature.
24


Reviews:
[{"review_id": "9Z6GKuv1MA", "reviewer": "Reviewer_v9zD", "summary": "The paper studies the family of value-aware model learning losses in model-based reinforcement learning, including MuZero loss. By theoretical analysis, it shows these losses are essentially uncalibrated surrogate losses. Then it proposes corrections for the losses. Experiments are conducted to show the correction to the losses is effective to obtain strong models.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. We are using this reply as a general reply since all reviewers raised similar concerns and this review is the first one that appears on open review. As ICML does not allow an updated manuscript, but all reviewers are mostly concerned about empirical questions, we provide additional experiments under this URL https://drive.google.com/file/d/178cVcy05grmQ-dZCFu1p8ixIItgghoxG/view?usp=sharing\n\n### Number of environments\n\nWe acknowledge that our experiments cover fewer environments than other, more empirically focused papers. However, we see our main contribution in the theoretical aspects of our work. Therefore, our focus was not on providing a broad benchmarking comparison.\nWe want to point out that we did not compare on 2, but on the 7 most challenging DMC tasks. We acknowledge that this was not stressed clearly in the paper and will update the writing accordingly. The graphs in Figure 4 show aggregated results over several tasks on each locomotion robot. We will provide individual curves in an updated document in the rebuttal URL. We would like to stress that this is already a broader comparison than e.g papers which compare on the standard Open AI Gym Mujoco environments [1,2,3], and the DMC dog and humanoid tasks are significantly more challenging.\n\n[1] CrossQ https://openreview.net/forum?id=PczQtTsTIX\n\n[2] MBPO https://openreview.net/forum?id=BJg8cHBxUS\n\n[3] ALM https://openreview.net/forum?id=MQcmfgRxf7a\n\nTo provide additional evidence for our claims, we ran the following additional experiments:\n\n- We ran all the environments to 2,000,000 environment steps. This creates a clearer picture: the humanoid confidence intervals for (1,1)-VAML and CVAML do not overlap anymore.\n- We ran additional experiments as requested by the reviewers:\n- We ablated the TD-MPC1 model used by removing the auxiliary loss and the decision-aware loss respectively. We show that both variants decrease in performance.\n- We added a model-free TD3 baseline using the TD-MPC1 architecture (without model) for a full comparison. We show that this does not result in strong performance, especially on hard humanoid tasks.\n\nAs requested by reviewer 9s7Q, we added a grid world experiment in which we conduct policy iteration.\n\nA general note regarding the scope of the empirical comparison: If the reviewers feel that specific environments would greatly aid in understanding the implications of our findings, we are happy to add them. We did so in the case of the gridworlds mentioned by reviewer 9s7Q. Otherwise, without concrete recommendations, it is hard to scope a proper reply to a general call for “more experiments”. Thus, we were wondering whether you had any particular environments in mind.", "claims_and_evidence": "The claims are well supported by evidence.", "methods_and_evaluation_criteria": "The proposed methods and the evaluation criteria make sense.", "theoretical_claims": "Yes. The proofs for the theoretical claims are correct.", "experimental_designs_or_analyses": "The experiment designs and analyses are generally sound. However, in Figure 4 the CVAML is not significantly better than VAML, which makes the claim less convincing.", "relation_to_broader_scientific_literature": "The paper is related to the losses used in value-aware model training.", "essential_references_not_discussed": "None.", "other_strengths_and_weaknesses": "Strength\n- The paper provides sound theoretical proof of the nature of losses used in the value-aware model learning, and proposes effective correction to them. \n- The paper is well-written and easy to follow.\n\nWeakness\n- More experiments are needed to support the claims, e.g. about stochastic and deterministic models.", "other_comments_or_suggestions": "None.", "questions_for_authors": "None.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "W5T85IZ7b8", "reviewer": "Reviewer_ABNg", "summary": "Value aware model learning losses penalizes if the model's value function prediction goes wrong. This works theoretically investigates the losses and shows that generally these losses are uncalibrated surrogate losses. They devise corrective measure for the losses. They provide experimental results on DMC control suite for their loss variant. They also show that learning a deterministic model can be sufficient but learning a calibrated stochastic model is more beneficial.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. For your concerns on the amount of environments, please refer to our reply to reviewer v9zD. Additional graphs can be found here https://drive.google.com/file/d/178cVcy05grmQ-dZCFu1p8ixIItgghoxG/view?usp=sharing\n\nRegarding your concern about novelty: while we agree that our focus is on model-based RL and therefore “narrow” in the sense that it targets a well-defined subarea of research, our results are, to the best of our knowledge, novel and unpublished. Model-based RL is also a mainstay topic at all major machine learning conferences, and we are therefore confident that this venue is an excellent place to disseminate our research.", "claims_and_evidence": "The claims are well supported by the theoretical evidence. They clearly poses the research questions and conduct relevant analysis.", "methods_and_evaluation_criteria": "As core of the Value-Aware Model Learning (VAML) framework, they consider the (m, b)-VAML family of losses and exhaustively consider the value categories for m and b. The analysis of variance and bias for stochastic model of the environment is crucial to understand loss performance. Further, the discussion on auxiliary losses makes the discussion comprehensive.", "theoretical_claims": "They claim that (m, b)-VAML family of losses are uncalibrated under a stochastic model. They successfully demonstrate that the error is dependent on the model’s stochasticity rather than the environment. They nicely carried out the analysis to show the advantage of the calibrated losses.", "experimental_designs_or_analyses": "The experiments are relevant with the research questions. However, the experiments are limited to two DMC environment. To make the analysis robust, more experiments on diverse environments are crucial.", "relation_to_broader_scientific_literature": "The scope of the paper is a bit narrow. It considers a certain aspect of the model-based RL.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "The paper identifies key insights that would be very helpful to understand the relation between the model representation and loss functions. The proposed corrective loss components are of marginal novelty.", "other_comments_or_suggestions": "The paper should clearly mentions the definition of calibration early in the paper.", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3KDZYAddRg", "reviewer": "Reviewer_Eb87", "summary": "This paper examines a systematic issue in value-aware model learning (VAML) for reinforcement learning (RL).\n\nCore Idea of VAML:\nUnlike standard model learning, which maximizes log-likelihood, VAML optimizes for a model that results in zero value-function error. In other words, even if the model does not accurately capture the true transition dynamics of the MDP, it is considered \"perfect\" as long as it produces no value error.\n\nKey Contribution of the Paper:\nThe paper highlights a critical issue with VAML-style objectives: when naively estimated empirically, they introduce additional bias. Specifically, the VAML objective requires an expectation over next states inside the squared error term, but empirical estimation places the expectation outside the square. This shift leads to overestimation due to Jensen’s inequality (since the squared function is convex). While the paper describes this as an \"uncalibrated\" estimate, a more standard term would be \"biased.\"\n\nConnection to Value Function Learning:\nThe paper makes a useful connection between this issue and value function learning in RL. A similar overestimation problem arises in residual learning, as described in Baird’s seminal work. Baird’s solution—using two next-state samples—mitigates this bias, though it is impractical in a model-free setting. However, in model-based RL, we can sample from the learned model and explicitly correct for the bias, effectively estimating and subtracting the overestimation error.\n\nPractical Implications:\nWhile the proposed bias mitigation approach is promising, its practical benefits remain uncertain. Reducing bias often comes at the cost of increased variance, and it is not always clear whether a zero-bias estimator is preferable to one with lower variance. The paper presents empirical evidence suggesting that this tradeoff is worthwhile in certain domains, but further theoretical and empirical validation is needed to confirm its general effectiveness.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your kind and thorough review. We are grateful for your comments. For your concerns on the number of environments, please refer to our reply to reviewer v9zD. Graphs can be found here https://drive.google.com/file/d/178cVcy05grmQ-dZCFu1p8ixIItgghoxG/view?usp=sharing\n\nBias-variance tradeoff: This is an excellent question and indeed the reason we chose to use the more technical term “calibration” instead of the term bias. We didn’t want to use “bias” as we felt it would be conflated too much with the phenomenon of bias and variance due to model architecture. However, in our case, we are talking about the bias of an estimator, not one stemming from a model class. As ML and statistical terminology are not always used in a consistent way across the literature, we felt that using the more specific idea of a calibrated surrogate loss would improve understanding.\n\nOur method does not induce a classic bias-variance tradeoff, as we can easily drive both terms to 0. Instead, we have a variance-compute tradeoff, as we can always draw more samples from our model to reduce the variance of the estimator, but that comes at the expense of additional computational steps. However, even for statistical estimators like ours, the unbiased estimator is indeed not guaranteed to be a minimum variance one for a fixed sample size. So there might be model classes and learning problems where the biased estimate does better with a fixed number of samples. However, as our method allows us to decrease the variance arbitrarily (with additional computational resources) we believe that it is still an important extension to the uncorrected estimator, which will result in non-zero error even in the limit of infinite model samples.\n\nDeterministic models: You are correct that the majority of published algorithms use a deterministic model. However, this is more of a historical issue than a reflection of the superiority of deterministic models. Most benchmark environments have a negligible amount of transition noise, and so using a stochastic model is often unnecessary. However, this is an artifact of the benchmarks the community has settled on, not really a condemnation of stochastic models. \n\nGeneral concerns about VAML: The questions you raise are important and we believe unsettled in the literature so far. To the best of our knowledge, all current SOTA model-based RL approaches use some form of DAML loss [1][2][3][4]. Note that Dreamer [5] is an outlier, but seems to be outperformed in many cases by one of these approaches.\nHowever, constantly changing the target is indeed a concern. There is one interesting thing about the difference of (corrected) MuZero and IterVAML: For the true value function, the ground truth model is also a perfect model under the (corrected) VAML (MuZero and IterVAML) loss. However, for any other value function, this is still true for the IterVAML loss, but not for the MuZero loss. This might be an interesting insight as the VAML loss has at least one stable solution, while the MuZero loss does not enjoy this property.\n\nWe added ablations to investigate your question further and were unable to achieve strong performance without the VAML or auxiliary components, so both seem important for these architectures.\n\n[1] Efficient MuZero https://openreview.net/forum?id=OKrNPg3xR3T\n\n[2] Efficient MuZero v2 https://arxiv.org/abs/2403.00564\n\n[2] TD-MPC1 https://www.nicklashansen.com/td-mpc/\n\n[3] TD-MPC2 https://openreview.net/forum?id=Oxh5CstDJU\n\n[4] MAD-TD https://openreview.net/forum?id=6RtRsg8ZV1", "claims_and_evidence": "The main claim—that VAML-style losses suffer from bias—is well-supported by both theoretical analysis and empirical results. The authors make a sharp observation about how the placement of the expectation affects estimation and introduce a clear argument grounded in Jensen’s inequality. This insight is particularly valuable because it connects VAML’s estimation issues to a well-known challenge in RL: overestimation bias in value function learning. The link to Baird’s residual learning framework strengthens the argument and provides a historical precedent for this type of error. While the paper does not explicitly frame the issue in terms of Jensen’s inequality, doing so makes the reasoning even more intuitive.\n\nHowever, the secondary claim—that removing this bias in the prescribed manner is necessarily beneficial—is less well-supported. While the paper demonstrates that the proposed bias correction technique improves estimates in certain domains, the broader implications remain unclear. A key concern is the classic tradeoff between bias and variance: reducing bias can increase variance, sometimes to the detriment of overall performance. The authors acknowledge this tradeoff but do not thoroughly analyze how their correction method affects variance in different settings. While empirical results show an advantage in some cases, they do not provide a comprehensive theoretical justification or broader empirical validation across a wide range of environments.\n\nFurther, it is not always the case that a lower-bias estimator is preferable to one with slightly higher bias but significantly lower variance. Some RL methods explicitly tolerate bias in favor of stability and better long-term learning dynamics. A deeper analysis of this tradeoff—both theoretically and through more diverse empirical settings—would strengthen the paper’s argument for the practical benefits of its proposed correction.", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "I did check the proofs to the best of my ability. I think sometimes the notation is overly complicated. The main text could have presented things in the simplest case possible, and also not hold off to the main result on overestimation so much.\n\nIn equation (1), the sg notation is not immediately defined (deferred to the next page for some reason). It is also somewhat confusing, because on the term we are applying sg, there is no dependency to the model itself, so what does that even mean to say we are sg ing that term?", "experimental_designs_or_analyses": "My intuition about these VAML models is that while the point about likelihood models being an overkill is a fair one, in practice the correct VAML model might vary radically from one state to the other. This is because the correct next state to choose is the one with appropriate value function, and this can be a confounding factor. For example two very different next states have the same value function. Thus, when doing VAML in the function approximation case, generalization might become a big issue (two input states having radically different output).\n\nThus, I find it interesting that this paper is applying VAML to the latent space case where the issue explained above may be less severe.\n\nOn the negative side, a small portion of previous work in this space actually uses stochastic VAMLs. If the model is deterministic the bias issue is vanished, as identified by the paper. So while the bias issue is indeed present, it is not quite applicable in practice in light of the majority of papers using deterministic models anyways.\n\n\nI think that the experimental results are quite limited. I would have loved to see the impact of this bias reduction idea in a much larger set of benchmarks. Other than the toy setting, the paper is evaluating the effectiveness on just two domains, which I feel is not sufficient to make a strong argument about the usefulness of the idea in practice. I still lean on the acceptance side in light of the nice observation and despite the weakness on the experimental side of things.", "relation_to_broader_scientific_literature": "The connection to residual gradient method makes the result better situated.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "NA", "other_comments_or_suggestions": "NA", "questions_for_authors": "I have an ongoing internal debate about the fundamental difference between VAML-based models and classical likelihood-based models, and I would love to hear the authors' perspective on it.\n\nTo me, VAML seems to introduce a kind of chicken-and-egg problem. The goal in VAML is to learn a model that minimizes value function approximation error when used for planning or learning. However, in the control setting, once the policy changes, the optimal model should also change. This is because the value function under the new policy might have different structure, meaning the previously learned model may no longer be optimal for minimizing value estimation error in the new setting. As a result, every policy update implies a corresponding update to the model, and every model update affects the policy, creating a tightly coupled iterative process.\n\nThis stands in contrast to classical likelihood-based models, where the objective is to approximate the true environment dynamics, which remain fixed regardless of policy updates. In this case, learning the model is a separate, more stable problem, and once the model is learned, it does not need to be modified with every policy update.\n\nGiven this difference, I’m curious how the authors view the stability and convergence properties of VAML-style models in iterative control settings. Does the need to repeatedly adapt the model with every policy change introduce instability or inefficiencies? Additionally, do the authors see scenarios where this iterative adaptation of the model and policy could be an advantage rather than a liability?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "F0N2h5cr4G", "reviewer": "Reviewer_9s7Q", "summary": "The paper analytically investigates *value aware model learning* (VAML) and puts it into relation with the MuZero loss, by defining a generalizing (m,b)-VAML loss. The authors show that both approaches, that is, (1,0) and (1,1)-VAML, are *uncalibrated* because averaging over samples introduces a variance term, akin to Bellman Residual Minimization (BRM). However, in difference to BRM this variance term can be estimated with samples from the learned model, and the authors call the corresponding correction the *corrected VAML loss* (CVAML). CVAML outperforms VAML in a variety of studied cases of (m,b)-VAML on randomly generated toy MDP. Furthermore, the authors prove that (1,0)-VAML generally allows to learn deterministic transition models in latent spaces, even if the true model is stochastic. Classical auxiliary losses for latent space models introduce a bias for all but linear value functions, though. Experiments on two DMC environment demonstrate that using stochastic models can in some cases still be empirically beneficial, although the difference between VAML and CVAML auxiliary losses does not seem to be significant.\n\nIn summary, the paper discusses an important aspect of stochastic models in RL. It is generally well written and decently evaluated, but both points suffer in some parts which makes some contributions very hard to follow. I would recommend acceptance if some ambiguities in the formal notation and in the experimental description were rectified. I have not read the proofs in detail, though, so if other reviewers find faults, I would also be fine with rejecting the paper.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your kind and thorough review. We are grateful you took the time to engage very thoroughly with our paper! For your concerns on the number of environments and experiments, please refer to our reply to reviewer v9zD. These extended experiments also addresses concerns about significance. Graphs can be found here https://drive.google.com/file/d/178cVcy05grmQ-dZCFu1p8ixIItgghoxG/view?usp=sharing\n\nNotation: We share your concerns about the complexity of the notation. The loss has a lot of moving parts, and several past papers struggled to write it out properly in a non-confusing manner. We truly appreciate your suggestions for improvement and will adapt them where possible. As ICMl sadly doesn’t allow us to upload an updated PDF, we can only sketch out what we plan to change. All unmentioned typos will, of course, be corrected. Thank you for providing a thorough list.\n\nEquation 2: We are happy to simplify the equation as suggested!\n\nEquation 4: You are correct about $\\hat{V}$. We will change that. The right-hand side is wrt to the target network and does not contain the model, so it does not have a double sampling issue. We will clarify this further by adding $[]_\\mathrm{sg}$\n\nEquation 6: We dropped the dependency on $\\hat{x}^{(m)}$ (as you suggested for Eq. 2) here already. We will unify this across the paper. We also split the model explicitly into an encoder and latent transition model, we’ll clarify this.\n\nSome concrete replies to confusing bits\n\nProposition 4: We are analysing the value learning part of the MuZero loss here, so we simplified the setup by assuming that the learned model matches the ground-truth model. Our result specifically states that even if you assume that the model is perfect (equal to the ground-truth model), the MuZero value learning component will not learn the correct value function.\n\nEquation 5: We will clarify the writing here. Our concern is that the MuZero loss only enforces that the expectation of the next model state value function will match the bootstrapped target. Your insight is correct if each state x appears in the dataset as x^{(0)}. However, this might not be true due to the dataset being used. Then we cannot guarantee that model-generated states will have correct value estimates. This is not an empirical finding, simply a consequence of the loss formulation.\n\nProposition 5: As we are looking at a deterministic model here, $k$ doesn’t change the loss (every sample would be the same), so we dropped it to declutter the presentation.\n\nInsight 1: Indeed, we will clarify the writing here. We meant a regular model-based target estimate with a j-step model rollout, similar to the one used in [1][2][3]\n\nL. 248: The ground truth for calculating the MSE is computed via the analytical solution over the transition matrix. \n\nl.260R: Why does (1,1)-VAML perform so poorly in deterministic environments?\n\nNote that while the environment is deterministic, the model is not. For the corrected (1,1) loss, we find that it is much more prone to get stuck in local minima due to the averaging effects of the MuZero style value update.\n\nOn the large-scale experiments, we see a similar pattern even with deterministic models. We debated this question a bit among ourselves. Our conclusion is that the (1,1) loss updates the value function on a model-generated sample $\\hat{x}^{(1)}$ while the (1,0) loss only updates the value functions on samples which actually appear in the replay buffer. As the model will have errors, this can cause additional difficulties with the loss function even on deterministic environment, as the sample $\\hat{x}^{(1)}$ and target $T[V](x^{(1)}$ don’t match even if the model is not stochastic (simply due to model errors). Updating values on generated samples (not just with model-generated targets) \n\nRunning TD3:\n\nThis seems to be a misunderstanding. We will improve the writing, thanks for noticing! We did not run TD3, but we understand why this might appear to be the case. We ran TD-MPC1 with two modifications: we replaced SAC in TD-MPC1 with TD3 (which is likely where the confusion arises), and we used the model for value learning (similar to the way [1],[2],[3] use the model). This means the model here is used both for value improvement and also at inference time, as the MPC procedure uses the model and the value estimate for planning. We chose TD-MPC1 over a MuZero variant for a very simple reason: computational efficiency. MuZero is notoriously slow due to the overhead of MCTS and reanalyze, and few efficient open-source implementations exist. As our contribution is first and foremost a theoretical one, we decided to use the more lightweight TD-MPC algorithm, which still contains all the relevant parts of the MuZero algorithm (loss, model-based search, + added model-based value improvement).\n\n[1] MBPO https://openreview.net/forum?id=BJg8cHBxUS\n\n[2] Dreamer https://openreview.net/forum?id=S1lOTC4tDS\n\n[3] MAD-TD https://openreview.net/forum?id=6RtRsg8ZV1", "claims_and_evidence": "1. \"Iterative VAML and MuZero value-aware model losses are not calibrated, but can be\". This claim is shown analytically for simple cases of both losses and evaluated on many randomly generated toy MDP. The corrected Iterative VAML loss outperforms the corrected MuZero loss significantly, but only on unrealistic toy tasks.\n2. \"Deterministic (1,0)-VAML losses in latent spaces are sufficient to be value aware, even if the true dynamics are stochastic\". This claim is analytically proven. The authors also find empirically that stochastic models can improve performance in one out of two DMC environments, but attribute this to induced robustness against deviations. Practical differences between CVAML and VAML seem to be insignificant, though.\n3. \"Auxiliary losses for latent dynamics models help learning, but introduce biases except for value functions that are linear in the latent space\". This claim is mostly argued at an example. The empirical evidence only compares model-learning as an auxiliary task without clarifying wether the value function was linear in the last layer. The practical relevance of this claim is therefore questionable.", "methods_and_evaluation_criteria": "The claims are theoretically derived and empirically validated. Both parts can be improved by being more precise in formalism and details.", "theoretical_claims": "The first theoretical claim is based on a simple, but elegant and to my knowledge unpublished, insight that learned stochastic models suffer from an additional variance term. The presented formalism, in particular the unified (m,b)-VAML loss (Equation 2), is very complex and in some parts unclear. I understand that the authors wanted to present a unified loss, but it is a bit unclear over which training sets the loss is defined, and how certain inputs and indices affect the loss (like in Proposition 4). The claim that it unifies the IterVAML loss is also quite questionable, as this is only true for $k=1$. However, the main claims on calibration (Proposition 2, 3 and 4) are interesting and correct (with some createve reformulations, please add links to the proofs in the appendix, though). Again, the definition of the TD loss in Equation 4 is hard to parse without any expectation over $x^{(m)}$.\n\nThe second theoretical claim (Proposition 5) is interesting, and I have not seen it stated like this before. It might be worth mentioning that bijective latent mapping are hard (or impossible) to come by. I did not follow the proof in detail here, though.", "experimental_designs_or_analyses": "The empirical evaluation of the first claim is a bit underwhelming, as randomly generated MDP rarely have properties similar to real applications. Even an evaluation on gridworlds could improve the analysis. Nonetheless, the results are fairly clear and support the claim.\n\nThe evaluation of the second claim is a bit unclear: the authors train a model-free algorithm (TD3) with a VAML model-based auxiliary loss. This contradicts a bit the insight from their claim 3, unless the auxiliary loss is applied to the last layer of the value function (I could not find a mention thereof). Details are generally scarce here. MuZero would have made a much clearer test example, as it actually uses a value-equivalent model for inference. The results are also not very convincing, as the CVAML is not significantly better than VAML (except maybe for humanoid, but there the deterministic version is very strong). Please also add the results of vanilla TD3, as it is currently unclear whether VAML improves the performance at all.", "relation_to_broader_scientific_literature": "The paper is in parts very well written and discusses a wide range of relevant literature. The authors are commended for discussing many variants and differences to other methods in the main text.", "essential_references_not_discussed": "None that I know of.", "other_strengths_and_weaknesses": "The paper contains a small number of typos and inaccuracies that need to be fixed (see detailed comments). I also believe the formal notation can be massively improved, as sometimes the authors use losses which have not been explicitly defined, or where the general definition is ambiguous of how they should be read.", "other_comments_or_suggestions": "- l.61L: the MDP is missing an initial state distribution\n- l.105L: it must be $\\mathbb E_{\\hat x^{(m)} \\cdots}$\n- Equation 2 is very (and IMHO unnecessarily) confusing. There are 8 parameters, but it seems you don't actually need $x^{(m)}$ (is always drawn from $\\mathcal P^\\pi$). Furthermore, this loss only makes sense if you define an expectation over it. How about this: \n$\\hat{\\mathcal L}(m, b, p, p', V, V'|x) := \\mathbb E[ (V(x^{(m)}) - [[T^b_{P^\\pi} V'](x'^{(m)}))^2 | x \\sim p, x' \\sim p']$?\n- Table 1: the (m,b) tuples in the second row have to shifted one column to the right\n- The TD-loss in Equation 4 was very confusing to me. First, it should use $\\hat V$, not $V$, as the value is learned. Second, an expectation over it (in the real model) still contains a variance term that cannot be reduced without double-sampling. I think you wanted to define $(\\mathbb E_{\\hat p}[V(x^{(m)})] - \\mathbb E_{\\mathcal P^\\pi}[r^{(m)} + \\gamma V_{tar}(x^{(m+1)})])^2$ to use in Proposition 4.\n- Proposition 4 uses a surrogate loss with $\\hat p = \\mathcal P^\\pi$, which does not contain a learned model $\\hat p$! I assume you wanted to still use $\\hat p$ of MuZero here. With this the (m,1) loss is indeed uncalibrated. \n- Please refer to the appendix where the proofs can be found!\n- I don't fully understand the issue exemplified in Equation 5. I would assume that the value loss is minimized for all $m \\in \\{0, \\ldots\\}$, so also for $m=0$. Your conclusion seems to be motivated by an empirical finding, not by this theoretical insight. Maybe there is another cause for it?\n- Figure 2: I assume these are (1,0)-VAML losses, not (0,1) which would make no sense?\n- Insight 1 mentions $\\hat{\\mathcal L}^{\\hat p}_{0,j}$, but it is unclear what changing the superscript would do in Equation 2. Do you mean you learn the value by unrolling the model $\\hat p$ for $j$ steps?\n- l.265L: mention which distirbution the weights $\\omega_{i,j}$ are drawn from\n- l.248R: how is the value estimated? Analytically from the transition model?\n- l.260R: why does (1,1)-VAML perform so poorly in deterministic environments?\n- Proposition 5: It is not clear which loss $\\hat{\\mathcal L}_{IterVAML,1}$ refers to. The expected loss? The approximated? For which $k$?\n- Equation 6: the (m,b)-VAML loss has a different signature than before. What does it mean specifically?\n- l.363R: you mean (1,0)-updates!\n- Please clarify that the performance improvement of the corrected loss (not calibrated loss) in Section 7 is **not** significant. A trend, not a reliable result!\n- l.836, You are proving Proposition 4, not 2! Following propositions have the wrong numbers (except 5).", "questions_for_authors": "- Why did you test latent CVAML on a model-free algorithm? Why not MuZero?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "s7jOHszAs8", "reviewer": "Reviewer_6JPt", "summary": "The paper investigates value-aware model learning (VAML), particularly examining the MuZero loss and Iterative VAML (IterVAML) within a unified framework termed (m, b)-VAML. The authors present theoretical insights, showing that standard (m, b)-VAML losses are generally uncalibrated surrogates when applied to stochastic models. The authors introduce a correction term leading to a calibrated version called CVAML to handle this. Additionally, the authors explore deterministic versus stochastic models, empirically demonstrating stochastic models' advantages in certain settings.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. We are thankful you took the time to engage with our paper! For your concerns on the number of environments and experiments, please refer to our reply to reviewer v9zD. Graphs can be found here https://drive.google.com/file/d/178cVcy05grmQ-dZCFu1p8ixIItgghoxG/view?usp=sharing", "claims_and_evidence": "The claims are supported by evidence on a few environments: Garnet MDPs and two environments from the DMC suite. Although these results are insightful, a broader set of environments, possibly including more complex ones, would strengthen the claims.", "methods_and_evaluation_criteria": "The proposed method is simple yet makes sense for the given problem. As noted above, the provided evaluation criteria are not very thorough, but I think they are enough for proof-of-concept.", "theoretical_claims": "I have checked the theoretical claims' correctness and found no issues.", "experimental_designs_or_analyses": "I have checked the soundness of the experimental designs and found no issues.", "relation_to_broader_scientific_literature": "Addressing calibration in surrogate losses for stochastic models is highly relevant to the reinforcement learning community, especially for those focusing on model-based approaches.", "essential_references_not_discussed": "I did not find any essential references that were not discussed.", "other_strengths_and_weaknesses": "Strengths\n- The paper provides a nice unifying view of the value-aware model learning methods.\n- The paper identifies a calibration issue that many value-aware model learning methods can suffer from.\n- The proposed method is simple yet seems to address the issue effectively, given the empirical results.\n\nWeaknesses\n- Although the empirical experiments are insightful, they are somewhat limited to specific domains and scenarios. A broader set of environments or more extensive empirical validation might strengthen the claims about stochastic vs deterministic models.", "other_comments_or_suggestions": "- The term sg is not defined until equation (2), even though it is first used in equation (1).\n- In Table 1, the $m$ and $b$ conditions in the header columns seem to be shifted by one cell.", "questions_for_authors": "- L271: Is it correct that the indices $i$ and $j$ are inside the softmax, not outside?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper investigates value-aware model learning (VAML) for reinforcement learning (RL) and highlights a problem with the current method. Specifically, it argues that the family of (m,b)-VAML algorithms are uncalibrated when applied to a stochastic environment model. The authors propose a novel loss variant to address this issue and provide both theoretical and empirical evidence to support their solution.

All reviewers acknowledge that the main claim—that VAML-style losses are uncalibrated and that this issue can be corrected—is well-supported by both theoretical analysis and empirical results. However, some reviewers raised concerns that other claims in the paper were not sufficiently substantiated.

During the rebuttal period, the authors clarified ambiguities in the experimental section and provided additional experimental results. In summary, the authors’ responses adequately addressed the concerns raised by the reviewers.

In light of the reviewers’ assessments following the rebuttal, this paper is deemed suitable for acceptance.
Decision: Accept (poster)