On Rollouts in Model-Based Reinforcement Learning
Authors: Bernd Frauenknecht
Devdutt Subhasish
Friedrich Solowjow
Sebastian Trimpe
Venue: ICLR 2025 Poster (2025)
Keywords: Model-Based Reinforcement Learning
Model Rollouts
Uncertainty Quantification
Primary area: reinforcement learning
TL;DR: We propose Infoprop a novel rollout mechanism for model-based reinforcement learning yielding substantially improved data consistency and long-term planning capabilities.

Abstract: Model-based reinforcement learning (MBRL) seeks to enhance data efficiency by learning a model of the environment and generating synthetic rollouts from it. However, accumulated model errors during these rollouts can distort the data distribution, negatively impacting policy learning and hindering long-term planning. Thus, the accumulation of model errors is a key bottleneck in current MBRL methods. We propose Infoprop, a model-based rollout mechanism that separates aleatoric from epistemic model uncertainty and reduces the influence of the latter on the data distribution. Further, Infoprop keeps track of accumulated model errors along a model rollout and provides termination criteria to limit data corruption. We demonstrate the capabilities of Infoprop in the Infoprop-Dyna algorithm, reporting state-of-the-art performance in Dyna-style MBRL on common MuJoCo benchmark tasks while substantially increasing rollout length and data quality.

Published as a conference paper at ICLR 2025
ON ROLLOUTS IN MODEL-BASED REINFORCEMENT LEARNING
Bernd Frauenknecht; Devdutt Subhasish; Friedrich Solowjow, and Sebastian Trimpe Institute for Data Science in Mechanical Engineering RWTH Aachen University
Aachen, 52062, Germany
firstname last name@dsme. rwth -aachen.de de
ABSTRACT
Model-based reinforcement learning (MBRL) seeks to enhance data efficiency by learning a model of the environment and generating synthetic rollouts from it. However, accumulated model errors during these rollouts can distort the data distribution, negatively impacting policy learning and hindering long-term plan- ning. Thus, the accumulation model errors a key bottleneck in current MBRL methods. We propose Infoprop, model-based rollout mechanism that separates aleatoric from epistemic model uncertainty and reduces the influence of the lat- ter on the data distribution. Further, Infoprop keeps track of accumulated model errors along a model rollout and provides termination criteria to limit data corrup- tion. We demonstrate the capabilities of Infoprop in the Infoprop-Dyya algorithm, reporting state-of-the-art performance in Dyna-style MBRL on common MuJoCo benchmark tasks while substantially increasing rollout length and data quality.
INTRODUCTION
Reinforcement learning (RL) has emerged as powerful framework for solving complex decision- making tasks like racing Vasco et al. (022); Kaufmann al. (2023) and gameplay OpenAl et al. 2211 Bi & D' Andrea (2024). However, when applying RL in real-world scenarios, significant challenge data ineffirreccy which hinders the practicality of standard RL methods. Model-based reinforcement learning (MBRL) addresses this issue by learning an internal model of the environ- ment Deisenroth & Rasmussen (20TT); (CCuaa et al.] (2018); Janner et al. (2019): [Hafner et al.) (2020). By generating simulated interactions through model rollouts, MBRL can make informed decisions while substantially reducing the need for real-world data collection.
The quality of data from baaed rollouts is critical for MBRL performance. Model errors can distort the data distribution and hurt policy learning. Long- nong-horizon planning is desirable, however, often infeasible model errors accumulate over time. This effect is demonstrated in Figure [ Even for a simple toy example (described in Appendix [B. we see the data distribution of based rollouts under the state-of-the-art Trajectory Sampling (TS) Chua al. (218)) scheme diverging quickly from the ground truth distribution of environment rollouts. Thus, data from TS rollouts can even be harmful to policy learning after a couple of time steps. This is largely because the TS mechanism does not explicitly address the effect of model errors on the propagated data distribution.
To tackle this challenge, we propose Infoprop a novel model-based rollout mechanism that miti- gates data distortion by addressing two key questions: How to propagate? and When to stop? We build our mechanism on explicitly leveraging the ability of common MBRL models to distinguish between aleatoric uncertainty due to process noise and epistemic uncertainty due to lack of data Lakshminarayanan et al. (2017): Becker & Neumann (2022). Making use of this property leads to substantially improved data consistency as depicted in Figure [ In particular, we
estimate and remove the stochasticity due to model error from the predictive distribution; formulate stopping criteria based on information loss to limit error accumulation; and
"Equal Contribution

---

Published as a conference paper at ICLR 2025
Trajectory Sampling
15 Infoprop
Ground Truth
5
Comparing Data Consistency of Model-based Rollouts. Trajectories under the pro- posed Infoprop mechanism follow the ground truth distribution of environment rollouts closely while rolling out the same model under the common TS scheme Chua al. (2018) results in distorted data.
Figure 1:
demonstrate the potential of Infoprop as a direct plugin to standard MBRL methods using the example of Dyna-style MBRL. The resulting Infoprop-Dyie algorithm yields state-of-the art performance in MBRL on common MuJoCo tasks, while substantially improving the data consistency of model-based rollouts and thus allowing for longer rollout horizons.
2
BACKGROUND
In the following, we introduce the fundamental concepts of information theory and MBRL. Ap- pendix A provides an overview of the notation introduced and used in the remainder of the paper.
2.1 INFORMATION THEORY
We will estimate the degree of data corruption in Infoprop rollouts using information -ifooration-thhoretiic ar- guments. Information theory serves to quantify the uncertainty of random variable (RV) [Shan-] + X' and Y : - Y, marginal entropy æ)) describes the average uncertainty about X in bits. Further, the conditional entropy [ y) -EEex [X|Y = y] log, (P(X v)) gives the uncertainty about X, given a realization of Y. Based on marginal and conditional entropy, the reduction in uncertainty about X given realization of Y is described by mutual information
non (1948). Given the discrete RVs X
- P[X 2| logg (PP
I(X; Y y) H(X) HX = y),
with I(X; = y) if the RVs are independent. In the following, we focus on Gaussian RVs and the notion of quantized entropy Cover & Thomas (006) with details provided in Appendix D.
2.2 REINFORCEMENT LEARNING
Reinforcement learning addresses sequential decision-makine problems where the environment is typically modeled discrete-time Markov decision process (MDP) represented by the tuple M (S, A, R, PR. Ps. Co, 7). Here, S C Rus denotes the state space with S, E S being the RV of the state time and s, its realization. Similarly, A S I n represents the action space with A, A the RV and a, the realization of the action as well as R C E the set of rewards with R₂ E R and rt the reward at time We make the common simplifying assumption Bellemare al. 2023) that the next state and reward are independent given the current state-action pair. Thus, a transition step in the environment can be expressed concerning reward kernel PR and a dynamics kernel Ps as
Ro+1 PR(-SS. Ar) and S Ps(.)S..
(2)
Further, initial states are distributed according to So Eu. and actions according to the policy A, ~ x(-(S)). We aim to learn an optimal policy πt" arg max, E, I> 7 R that maximizes the expected sum of rewards discounted by [0, 1), referred to return.
2

---

Published as a conference paper at ICLR 2025
2.3 MODEL-BASED REINFORCEMENT LEARNING
There are four main categories of MBRL that all build on model-based rollouts. (i) Dyna-style
methods Sutton [991]; : Janner et al. (019) model-based rollouts generate training data for
model-free agent. (i) Model-based planning approaches Chua et al. (2018); Williams et al. (2017);
Nagabandi al. (2018); Hafner al. 2019) do not learn an explicit policy but perform planning
via model rollouts during deployment. (it) Analytic gradient methods Deisenroth & Rasmussen
[011]: Hafner et al. (2020, 202T; 2023) optimize the policy by backpropagating the performance gradient through model rollouts. (iv) Value-expane expansion approaches Feinberg et al. (2018): Buckman ct al. 2018) stabilize the temporal difference target using model-based rollouts.
The model architecture of an MBRL algorithm determines the set of mechanisms for model rollouts. In this work, we focus on rolling out the particularly successful class of aleatoric epistemic separator (AES) models Lakshminarayanal et al. (2017): Becker & Neumann 2022) that can distinguish aleatoric uncertainty corresponding to the estimate of process noise from epistemic uncertainty.
2.4 ENVIRONMENT INTERACTION vs. MODEL-BASER ROLLOUTS
Model-basee rollouts aim to substitute environment interaction in MBRL. Thus, we compare the data generation process through environment interaction to the process of model-based rollouts.
We model environment dynamics as nonlinear function (1S, At) with additive heteroscedastic process noise that is normally distributed with variance (S Thus, environment rollouts, as depicted in Figure are generated by iterating the dynamics
(3) with L(S,, A,) L(S A,) T E(S, Ad) and the process noise W2 ~ I). Consequently, the transition kernel of the environment defined as Ps (-IS,, A.) = N (u(S,, E(S,, At)).
S = p(S,, At) + A.),,
In MBRL, however, we do not have access to Ps directly but typically rely a parametric model with the random parameters O E O. Besides estimates of nonlinear dynamics eee (St, A,) and process noise Ee, At), AES models provide an estimate of the parameter distribution O ~ Pe, e.g. via ensembling Lakshminaraaanann
al. (2017) or dropout Becker & Neumann (2022). These models are typically propagated using the TS Chua et al. (2018) rollout mechanism via iterating
S Me. (St, A.) (S,, A.)W,
(4)
with i.a, (St. A,) (S,, A, (S, A.). W, /N(0, 7), and 6, .P. This results in the - (Se, (S, Σe, (S,, A1)
TS rollouts in Figure and induces the kernel Ps.rs(-ISs,,A)
The majority of recent MBRL approaches use the TS rollout mechanism, e.g. Chua et al. (2018); Becker & Neumann (2022); Janner et al. (2019); Pan et al.] (2020); Yu et al. ) Luis et al. 2023). Pseudocode is provided Algorithm 2 of Appendix
PROBLEM STATEMENT
Revisiting Figure allows us to illustrate the effects of different sources of stochasticity by compar- ing environment interaction under Ps to TS rollouts under Ps.TS- While different realizations of process noise W/ ~ (0,,7) allow for keeping track of the environment distribution, the sampling process 0₂ ~ Pe introduces additional stochasticity that leads to overestimated total variance in the TS rollout distribution. This effect is amplified through the continued propagation of erroncous predictions making data later steps unfit for policy learning. We ask the following questions:
How can we construct predictive distribution closely resembling environment dynamics?
How can we quantify the degree of data corruption due to model error?
When should model-based rollouts be terminated due to data corruption?
We address these questions by proposing the Infoprop rollout mechanism. Infoprop isolates and removes epistemic uncertainty for an improved predictive distribution, keeps track of data corruption using
arguments, and terminates rollouts based on the degree of corruption.
As PR typically is known deterministic function in the context of MBRL, while Ps is the unknown object we model, the discussion henceforth focuses on approximating Ps without loss of generality.
3

---

Published as a conference paper at ICLR 2025
4 INFOPROP ROLLOUT MECHANISM
In the following, we introduce the Info- prop mechanism for model-based roll- outs. As depicted in Figure 2 we de- compose model predictions into a sig- nal fraction representing the environ- ment dynamies and noise fraction in- troduced by model enror. This perspec- tive allows to interpret model rollouts as communication through a noisy chan- nel. We estimate both the signal and the noise distribution and use these in- fer a belief over the environment state, given observation of the model state. This belief state represents the founda- tion of the Infoprop rollout mechanism.
Model
Maximum Likelihood Dynamics
S+11
Noisy Channel
St+1
f
Ground Truth Inference
$
H((S+1)
Figure 2: Infoprop block diagram
4.1 THEORETICAL SETUP
First, we introduce additional notation to specify RVs under different transition kernels.
Definition 1 (Environment state). We define the environment state the conditional expectation under the environment dynamics given realization of a state-action pair
S Eps [St+1|S, St, st, A, ae, Wi].
(5)
Thus, 5.+1 is an RV, where the randomness is induced by the process noise and has an aleatoric nature. we additionally condition on the realization W; wy, we obtain deterministic object.
Definition (Model state). We define the model state as the conditional expectation under Ps.TS
S == E [St+1|S
= st, AA W,, Oi)
As discussed in Section stochasticity in S441 is induced not only by W, but also by the randomness in the parameters θ, We project the uncertainty in the parameter space & to S via error process.
Definition (Model error process). We define model error process
Dr = S - SS+1
that, given a realization process noise W. = Wt. projects uncertainty in to S EE [A we] = E [S Se, ar We, θ.] Eps [St+1|31, at, wt].
(8)
We refer to the projected parameter uncertainty epistemic uncertainty.
Further, we restrict model usage to a sufficiently accurate subset & C S x A, as proposed in Frauenknecht et al. (2024). We define E amenable to the Infoprop setting in Section 4.4 and make the following assumptions when performing model-based rollouts in E:
Assumption (Consistent estimator aleatoric uncertainty). The model's predictive variance Le, consistent estimator of E following the definition of Julier Uhlmann (2007), i.e.
Ai) E Assumption (Unbiased estimator). The model bias A negligible. Thus Si1 according (4) is an unbiased estimator of Š141 according to {3 i.e.
b
IS, [S:+.[S,,
A(),,A)) EE
(10)
Figure D empirically shows that these assumptions are reasonable. The Infoprop distribution is
slightly more stochastic than the ground truth process, which indicates that Assumption hollds. As
(9) states, the model does not underestimate aleatoric uncertainty; the Infoprop rollouts should be at least stochastic the true process. Further, observe substantial bias of the Infoprop distribution underscoring the soundness Assumption Infoprop shows similar behavior high
dimensional problems reported in Section
4

---

Published as a conference paper at ICLR 2025
4.2 DECOMPOSING THE MODEL STATE SIGNAL AND NOISE
We aim to isolate the stochasticity due to parameter uncertainty in S.p1. We use the model error process (8) project the noise in & the same space as the signal, i.e. the dynamics, which is S. The parameter distribution θ ~ Pe can induce arbitrarily complex distributions A, ~ IPs. To simplify the analysis, we solely consider the first two moments of Pa, namely م and EA This allows to reformulate the propagation equation 4 of the model state
S4+1 S.+1 + A. 2 S.+1 (St, Ad) + (SA A.N,
(11)
concerning Si+1 and the model error A represented by μ (S,, At) the model bias, E A,) the epistemic variance with Cholesky decomposition (S: At), and N; the epistemic noise.
(S,, A,) By Assumption 2 we have (S,,
(S,, A,) &. Consequently, we can interpret the model rollout as communication through a Gaussian noise channel[Cover & Thomas (2006) via [1]
Based on the propagation equation 11 we aim to infer the maximum likelihood estimate of S from E realizations of [E[S++1| NN = n2 n to use it as the predictive distribution for our rollout scheme. As we cannot sample N, directly, we instead use an equivalent definition of $
Definition (Model state concerning epistemic uncertainty). Based on the model error process the model state is defined as
$41 p [S St, A. at, W,, A.] R Es [S St, A, ae, W,, M,] (12)
Reformulating (6) concerning A, does not change the information content or the induced sigma- algebra, as A, a measurable function of O. the simplified setting of solely considering the first two moments of FA. N₂ fully describes stochasticity due to model eиror. In reverse, we can obtain realizations { 11/4/11/ A fi]]_1 and interpret them samples [E[S,-1|N, nill
Lemma 1. Given E realizations TE 1 02] we estimate the environment state using maximum likelihood as
S = S441 { 0] Su41 plS,, A.) + A.),,
(13)
Proof. see Appendix D.2.1
Lemma 2. Following this line of thought, the maximum likelihood estimate of ΣA
given by
(Sc At) = E
Σ
- At)) ( -
P(S,, ( (14)
c=1
Proof. see Appendix D.2.2
Given the maximum likelihood estimates of the environment state S and the epistemic variance ΣA, we can decompose the model state $. in signal and noise fraction according to 11 in E.
4.3 CONSTRUCTING THE INFOPROP STATE
Having decomposed S441 into signal 5.+1 and noise EA, allows us to define the Infoprop state. Definition (Infoprop state). We define the Infoprop state
8 : (
SS41|S: =IE -
8t, A: at, S -841,
(15)
as the conditional expectation of the estimated environment state given a sample of the model state. We derive the corresponding Infoprop kernel Ps.rr-|SS,, A., S(n))
N A1, E(S,, A Se41) with the conditional noise U, ~ N(0, in Appendix D.3
Consequently, the Infoprop state aims to infer the signal 5 given a noisy observation 8 Prop agating model-based rollouts using S.41, yields favorable properties as stated in Theorem
5

---

Published as a conference paper at ICLR 2025
(S)
Sts - a.)
As
ac)
8511
S111
Sper
Se-1
B1-1
(a) Perfect model.
(b) Erroneous model.
(c) Rollout propagation.
Figure 3: Infoprop rollout mechanism. (a), (b): Generating the Infoprop state f
from the esti- mated predictive distribution S141 and the model sample $141- (c) Performing an Infoprop rollout.
Theorem (Infoprop state). By construction, S addresses questions (i) and (ii) of Section 3
(i) The distribution of Infoprop states is identical to the estimated environment distribution.
S dist S111
(16)
Proof. see Appendix D.4
(ii) The sum of marginal entropies of 5i+1 defines the information loss along an Infoprop rollout.
H S2. 80: Ao ag S1 -.,,... 88
H
Set1
(17)
1=
Proof. see Appendix D.5
Figure Billustrates the Infoprop rollout mechanism and provides intuition for Theorem | In the case of perfect model, i.c. Σ = O, depicted in Figure 3a the realization 6 provides the information about the process noise realization without ambiguity. Consequently, the belief about the environ- ment state given the sample from the model S.41 E[Se+1|$++1
$ is deterministic object and HH(Š++1) 0. In the general scenario of Σ depicted in Figure 3b the epistemic uncer- tainty results in ambiguity about the environment state given $141, such that ((S++)) > 0. Notably, conditioning S on 8141. results in Infoprop predictions S4,1 following estimated environment distribution S,+1 as stated in Theorem (i). This results in a data distribution that closely resem- bles the environment dynamics as desired in question of Section 3 Finally, Figure [3c depicts Infoprop rollout propagated via realization 1 We measure data corruption due to model eттor us- ing the conditional entropy of rollout under the estimated environment dynamics (51, . Sr)
given the realizations observed from the model (Su = sp, so Ao = ag. S $1.... Sr $T), ) ie. given the observed model trajectory, how
how the corresponding environment trajec- tory would look like?. As per Theorem T (ii), this trajectory-basee approach to uncertainty can be addressed with the accumulated marginal entropy S.+1, addressing question (ii) of Section
4.4 ROLLOUT TERMINATION CRITERIA
Having introduced how to propagate Infoprop rollouts, the question remains when to terminate them. In the following, propose two termination criteria to address question (iii) of Section 3
First, Infoprop rollouts build on the assumption that model usage is restricted to a sufficiently accu- rate subset & C S A, following the ideas of Frauenknecch et al. (2024).
Definition 6 (Sufficiently accurate subset). We define the sufficiently accurate subset
E = ((s,, at) S H{(())))) S A1, ~ (:),,)))
(18)
based threshold Xı for the single-step information loss H(()+1)) Second, restrict Infoprop rollouts sufficiently accurate paths limit uncertainty accumulation.

---

Published as a conference paper at ICLR 2025
Definition 7 (Sufficiently accurate path). Based on the estimated information loss along rollout 17 we define the set of sufficiently accurate paths of length {l,..., T) as
D'
(st E (S A)"
Σ HI (S++1) S X2 Noo
(19)
Heuristics for determining values of X1 and X2 depend on the class of AES model and MBRL algorithm hand with an example provided Section Combining the steps above yields the Infoprop rollout mechanism illustrated in Algorithm
Algorithm Infoprop
Require:
while < T do at T(1 (-(ss)
for e E {I...., F.) do
F.....,E.
Of Pe
S (st, ai)
from 13 and Σ (s,, at)
from 14
81+1 E
Se+1 W. Wh, O 0f with W2 NN(0, of ~ uu(tt,.... ...
S from (11) and 15(+1) from 24
[H[ (S(()))) > A1 then
break
else if Ei=
IHI (S(+)+1)) > A2 then
break
else
St E|$ ud] with u, 1)
5 AUGMENTING STATE-OF-THE-AA INFOPROP-DYNA
While the Infoprop rollout mechanism is applicable to different kinds of MBRL with AES models, illustrate its capabilities in Dyna-style architecture with probabilistic ensemble (PE) models Lakshminarayanan et al. (2017). We design Infoprop-Dyie -Dyn by integrating the Infoprop rollout mechanism in the fate-of-the-ard framework proposed in Janner et al. (2019) with minor adaptions.
As discussed in Section 4.4 heuristics for Xı and x2 depend on the algorithm at hand. In Infoprop- Dyna, we take the common approach [Chaa et al. (2018); Janner et al. (2019) of neglecting cross- correlations between state dimensions for computational reasons. Thus, we can consider data cor- ruption of each state dimension independently. As the predictive quality of different state dimensions can differ substantially, we choose both thresholds as ris dimensional vectors, such that a rollout is terminated as soon as the data corruption of any dimension overshoots the corresponding threshold.
In Dyna-style MBRL Janner et al. (219). the dynamics model is trained on the data distribution ob- served during environment interaction. The corresponding transitions are stored in an environment replay buffer Denv = {(x(0) a.") -(t) | where (b) indicates the index in the replay buffer. After a fixed number of interaction steps between a model-fres RL agent and the environ- ment, the dynamics model is retrained on the data in Denv. model- rollouts are performed, and the data stored stored replay buffer Dodd train the model-free RL agent. Consequently, assume the PE model to be accurate within the data distribution of Denv and build the heuristic for A and A on the predictive uncertainty within the environment buffer.
After each round of retraining the PE model, we compute set of dimension wise Infoprop state entropies for single-step predictions in Deny according to
['anel
² = (S =(6) A, a(6))
H (5 J
(20)
7

---

Published as a conference paper at ICLR 2025
where 1o E {1...., n indicates the corresponding state dimension. We define the dimension-wind thresholds Xi and X/ based on the cumulative distribution function of dimension -wie entropies
Σ
1 s].
(21)
771
MEM
The hthh element of A is defined as the Ç1 quantile of the single-step entropy set
X = inf {h H (h) 2 {1}
(22)
and limits model usage to the sufficiently accurate subset E. To restrict rollouts of length t' to pt we define the ht element of X2 as the (2 quantile of the entropy set scaled by
A { inf {h E 7² Fx (h) 2 {2)
(23)
Here, Ç2 denotes a quantile corresponding to precise predictions and & to the number of prediction
steps we are willing to accumulate the resulting data corruption. We choose (i 0.99, C2 0.01 and { = 100 for experiments Section without further hyperparameter tuning.
We use pink noise for environment exploration Eberhard et al. (2023) to quickly expand & Frauenknecht et al. (2024). Pseudocode is provided in Algorithm of Appendix
EXPERIMENTS AND DISCUSSION
To demonstrate the benefits of the Infoprop mechanism, compare Infoprop-Dyna to state-of-the- art Dyna-style MBRL algorithms on MuJoCo Todorov et al. (2012) benchmark tasks. We report
substantial improvements in the consistency of predicted data, especially over long horizons; effective rollout termination based on accumulated model error propagation; and state-of-the-art performance in Dyna-style MBRL on several MuloCo tasks.
Furthermore, we discuss the limitations of naively integrating Infoprop into the standard Dyna-style setupDanner Janner al. (2019) and point to further research questions.
6.1 EXPERIMENTAL SETUP
We compare Infoprop-Dyna to Model-Based Policy Optimization (MBPO) Janner et al. (2019) and Model-Based Actor-Critic with Uncertainty-Aware Rollout Adaption (MACURA) Frauenknecht et al. (2024) as well as to Soft Actor-Critic (SAC) et al. (0088) that represents the model- free learner of all the Dyna-style approaches above. We build our implementation on the code base provided by Frauenknecht et al. (2024). Further details are provided in Appendix E.1
6.2 PREDICTION QUALITY
To compare different rollout mechanisms, we train an Infoprop- agent on hopper for 120000 environment interactions and perform model rollouts from states in D
First, we evaluate the consistency of Infoprop and TS rollouts, propagating 20 steps without ter- mination. Figure 4a depicts the resulting distributions for the 11 dimension of the hopper state. Infoprop rollouts show substantially improved data consistency compared to TS rollouts, underscor- ing the ability of Infoprop to effectively mitigate model error propagation.
Next, we compare the rollout mechanisms of MBPO and MACURA based on TS sampling with Infoprop- -yna rollouts. Figure 4b shows the results for 11th dimension of the hopper and maxi- mum rollout length of 100 steps. MBPO rollouts are propagated for 11 steps following the schedule proposed in Janner al. (2019). resulting in widely spread distribution. In contrast, MACURA has an adaptive rollout length capped at 10 steps Frauenknecht et al.] (2024), leading to better data consistency. The improved predictive distribution and capability to estimate accumulated error of Infoprop allows for substantially longer rollouts up to 100 steps. The Infoprop termination criteria reliably stop distorted rollouts, resulting consistent rollouts over long horizons. Appendix E.2 provides additional results for setting the maximum rollout length of all three approaches to 100.
hittps://qithuie h Seneece-iieneceeecnn - Mechan cal htwpahhhaaaailiiiieeeeeeeoonnoon
https://githubs Data SasaaaScieecee Mechan cal Kngi www.kknnnninnnggee.eeeen www.gineeernnnrrrrr/r
8

---

Published as a conference paper at ICLR 2025
MBPO
Tacza
MACURA
Inloprep
(b) MBPO vs. MACURA vs. Infoprop- Figure 4: Predictive quality of rollouts the 11 state dimension of MuloCo hopper. (a) Rollouts according Trajectory Sampling (TS) and Infoprop (b) Rollout schemes of MBPO and MACURA based on TS compared to Infoprop-Dyna.
(a) Trajectory Sampling vs. Infoprop
1c3
tellcheetth
Welcer
Happer
Ic3
Antt
15
Return
30
20
10
20
10
2000
50
150
52
75
300
203
(a) Performance and average rollout length on MuJaCo tasks.
Indfronment 33P3)
Enmment B.Mer(AACLB))
Enstrommet Duffer nnoo Prop Cynss)
10
-20
-20
10
25
Mode MrTer16
75
Voda Sumggict Pog
100
10-9
10-4
-00
-40
Steps
1.0
2tpp
2tspe
162
(b) Adequacy of Dmod on the th state dimension of hopper.
Figure 5: Evaluation on MuJoCo tasks. (a) Infoprop- DDnopp--aaaaaa shows state-of-ther the-art performance for Dyna-style MBRL several MuJoCo tasks while considerably increasing average rollout length most tasks. (b) Infoprop-Dyme shows substantially improved consistency between 'eny and 'D mund-
6.3 PERFORMANCE EVALCATION
As depicted in the top row of Figure 5a Infoprop-Dyme performs on par with or better than MACURA, while substantially outperforming MBPO with respect data efficiency and asymptotic performance. Notably, Infoprop-Dyna consistently outperforms SAC with fraction of environment interaction. The bottom row of Figure depicts the average rollout lengths. Infoprop-Dyns shows substantially increased rollout lengths compared to prior methods in all environments but ant.
major concern of this work is the consistency of mosel-based rollouts with the environment dis- tribution. Figure 5b depicts the data distribution in Denv and Dmod of the respective Dyna-style approaches throughout training the th dimension of the hopper state. The distributions are illustrated via histograms of Infoprop-Dyren closely follows the distribution observed in the environment, while both the data from MBPO and MACURA show severe outliers. This is the case, Infoprop- is obtained from substantially longer rollouts as can be seen from Figure 5a which indicates the capabilities of the Infoprop rollout mechanism.
environment steps. It can be seen that the model data distribution
though the rollout data in

---

Published as a conference paper at ICLR 2025
6.4 LIMITATIONS AND OUTLOOK
Despite the excellent quality of model-generated data with the Infoprop rollout, the limitations of Infoprop-Dyna are most apparent on MuloCo humanoid with results provided in Appendix EE3 These show instabilities in learning and point to structural problems when integrating Infoprop roll- outs naively into standard Dyna-style architectures) Jannerr et al. (2019).
cause rapid distribution shifts in modd especially early in training. These nonstationary buffers are a well-known challenge to deep Q- learning methods MMnh et al. (2015). Another issue is primacy bias in model learning Qiao et al. 2023), where the model overfits to initial data and subsequently struggles to generalize, as seen in the decreasing rollout length for ant in Figure |5a| The main problem with Infoprop-Dyna is likely overfitting critics and plasticity loss Nikishin et al. (2022); D"Orro et al. (2023), as also reported by Frauenknecht al. (2024) for Dyna-style MBRL, trained on high-quality data. We provide an ablation on this observation and sketch methods to counteract this phenomenon in Appendix E.4
Figure (iinreebssoowws that the long rollouts of safoprop-Dyna
7 RELATED WORK
The negative effects of accumulated model error on the performance of MBRL methods is a long- studied problem Venkatraman et al. (015); Talvitie (2016); [Aadi al. (018b a)
Different model architectures have been proposed to mitigate this issue, such as trajectory models Asadi et al.; (2019); Lambert et al. (202I), bidirectional models Lai et al. (2020), temporal segment Mishras et al. (2017) or self-correcting models models(Talvitie) (016)). These architectures, however, imply substantial additional effort for model learning, such that state-of-the-art performance in the respective fields of MBRL is often reported for simpler single-step model architectures Chua et al. 20118
Janner et (2019) Buckman al. (0088).
These approaches address the problem of error accumulation by keeping model-based rollouts suf- ficiently short. Janner et al. (2019) introduce the concept of branched rollouts that allows to cover relevant parts of S with short model rollouts. Other methods weight rollouts of different lengths according their single-step uncertainty [Buckman et al. (2018) use single-step uncertainty to schedule rollout length |ann et al. (2020); Frauenknecht et al. (0024). Infoprop allows to infer model data consistent with the environment distribution over long rollout horizons using comparatively simple model architectures and computationally cheap conditioning operations.
Infoprop inspired by an cntemattion-thooretii view on RL et al. (202)). Thus far, information- theoretic arguments have been mostly used to improve the exploration Haarnoja et al.) (2018); Lu & Roy (2019; Ahmed et al. (2019); Mohamed & Rezende (2015) and generalization Tishby & Za- slavsky (2015): Lu et al. (2020): Igl et al. (2019); Islam al. (2023) of model-free RL methods. While aspects of dynamical systems such as causality, modeling. and control| Lozano-Duran & Ar- ranz (2021), predictability Kleeman (2011) dealing with noisy observations Gattami (2014) have been studied from an information theoretic perspective, these works do not directly apply to the
MBRL setup nor extend long model-based rollouts.
CONCLUDING REMARKS
Data consistency of model-based rollouts is key criterion for the performance of MBRL ap- proaches. This work proposes the novel Infoprop mechanism that substantially improves rollouts with common AES models. We reduce the influence of epistemic uncertainty on the predictive distribution of model-based rollouts, keep track of data corruption through propagated model error over long horizons, and terminate rollouts based data corruption. This allows for considerably increased rollout lengths while substantially improving data consistency simultancously.
While Infoprop is applicable to broad range of MBRL methods, we demonstrate its capabilities by naively integrating Infoprop into a standard Dyna-style MBRL architecture Janner et al. (2019) resulting in the Infoprop-Dyee algorithm. We report tfate-of-the-aa performance in several MuJoCo tasks while pointing to necessary adaptions to the existing algorithmic framework to fully unleash the potential of Infoprop rollouts.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
We thank Christian Fiedler, Pierre-François Massiani, and David Stenger for the fruitful discussions on the work presented in this paper. This work is funded in part under the Excellence Strategy of the Federal Government and the Länder (G:(DE-2)EEX) (ES--22EE---S--SSSS---- SF-OPSF854) and the German Federal Ministry of Education and Research (BMBF) under the Robotics Institute Germany (RIG), which the authors gratefully acknowledge. Friedrich Solowjow is supported by the KI-Starter grant by the state of NRW. Further, the authors gratefully acknowledge the computing time provided to them at the NHR Center NHR4CES at RWTH Aachen University (project number p0022301). This funded by the Federal Ministry of Education and Research, and the state governments participating on the basis of the resolutions of the GWK for national high performance computing at universities (www.nhrr
REFERENCES
Zafarali Ahmed, Nicolas Le Roux, Mohammad Norouzi, and Dale Schuurmans. Understanding the impact of entropy on policy optimization. Int. Conf. on Machine Learning, 2019.
Kavosh Asadi, Evan Cater, Dipendra Misra, and Michael L. Littman. Towards a Simple Approach to Multi-step Model-based Reinforcement Learning. arXiv, 2018a.
Kavosh Asadi, Dipendra Misra, and Michael L. Littman. Lipschitz Continuity in Model-based Reinforcement Learning. arXiv, 2018b.
Kavosh Asadi, Dipendra Misra, Seungchan Kim, and Michel L. Littman. Combating the Compounding- Problem with a Multi-step Model. arXiv, 2019.
Philipp Becker and Gerhard Neumann. On Uncertainty Deep State Space Models for Model- Based Reinforcement Learning. Transactions on Machine Learning Research, 2022.
Mare G. Bellemare, Will Dabney, and Mark Rowland. Distributional Reinforcement Learning. MIT Press, 2023.
Thomas Bi and Raffaello D' Andrea. Sample-efficient learning to solve a real-world labyrinth game using data-augmented model-based reinforcement learning. In IEEE Int. Conf. Robotics and Automation, 2024.
Jacob Buckman, Danijar Hafner, George Tucker, Eugene Brevdo, and Honglak Lee. Sample- efficient reinforcement learning with stochastic ensemble value expansion. In Int. Conf. on Neural Information Processing Systems. 201 8.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models. Adv. Neural Information Processing Systems, 2018.
Thomas M. Cover and Joy A. Thomas. Elements of Information Theory 2nd Edition (Wiley Series in Telecommunications and Signal Processing) by Thomas M. Cover Joy A. Thomas(2006-0771 Wiley-Interssience,, Hoboken, NJ, USA, 2006.
Mare Peter Deisenroth and Carl Edward Rasmussen. PILCO: model-based and data-efficient approach to policy search. In Int. Conf. on Machine Learning. 2011.
Pierluca D' D'Oro, Max Schwarzer. Evgenii Nikishin, Pierre-Luc Bacon, Marc G. Bellemare, and Aaron C. Courville. Sample-efficient reinforcement learning by breaking the replay ratio barrier. In Int. Conf. on Learning Representations, 2023.
Onno Eberhard, Jakob Hollenstein, Cristina Pinneri, and Georg Martius. Pink noise is all you need: Colored noise exploration in deep reinforcement learning. In Int. Conf. on Learning Representa- tions, 2023.
Vladimir Feinberg, Alvin Wan, Ion Stoica. Michael I. Jordan, Joseph E. Gonzalez, and Sergey Levine. Model- -ased Value Estimation for Efficient Model-Free Reinforcement Learning. Int. Conf. on Machine Learning, 2018
11

---

Published as a conference paper at ICLR 2025
Bernd Frauenknecht, Artur Eisele, Devdutt Subhasish, Friedrich Solowjow, and Sebastian Trimpe. Trust the Model Where Trusts Itself Model-Based Actor-Critic with Uncertainty-ANe Aware Roll- out Adaption. Int. Conf. on Machine Learning, 2024.
Ather Gattami. Kalman meets shannon. ArXiv, 2014.
Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, George Tucker, Sehoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, and Sergey Levine. Soft Actor-Crite Algo- rithms and Applications. arXiv, 2018.
Danijar Hafner, Timothy Lillicrap. Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning Latent Dynamics for Planning from Pixels. In Int. Conf. on Machine Learn- ing. PMLR, 2019.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to Control: Learn- ing Behaviors by Latent Imagination. In Inl. Conf. on Learning Representations, 2020.
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering Atari with Discrete World Models. Int. Conf. on Learning Representations, 2021.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering Diverse Domains through World Models. Int. Conf. on Learning Representations, 2023.
Maximilian Igl, Kamil Ciosek, Yingzhen Li, Sebastian Tschiatschek, Cheng Zhang, Sam Devlin, and Katja Hofmann. Generalization in reinforcement learning with selective noise injection and information bottleneck. In Int. Conf. on Neural Information Processing Systems, 2019.
Riashat Islam, Hongyu Zang, Manan Tomar, Aniket Didolkar, Md Mofijul Islam, Samin Yeasar Arnob, Tariq lqbal, Xin Li, Anirudh Goyal, Nicolas Heess, and Alex Lamb. Representation learning in deep rl via discrete information bottleneck. In Int. Conf. on Artificial Intelligence and Statistics, 2023.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: model- based policy optimization. In Int. Conf. on Neural Information Processing Systems. 2019.
Simon Julier and Jeffrey Uhlmann. General Decentralized Data Fusion with Covariance Intersec- fion (CI). 2001.
Elia Kaufmann, Leonard Bauersfeld, Antonio Loquercio, Matthias Müller, Vladlen Koltun, and Davide Scaramuzza. Champion-level drone racing using deep reinforcement learning. Nature, 2023.
Richard Kleeman. Information theory and dynamical system predictability. Entropy, 2011.
Hang Lai, Jian Shen, Weinan Zhang, and Yong Yu. Bidirectional Model-based Policy Optimization. In Int. Conf. on Machine Learning. PMLR, 2020.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In Int. Conf. on Neural Information Processing Systems. 2017.
Nathan Lambert, Albert Wilcox, Howard Zhang, Kristofer S. Pister, and Roberto Calandra. Learning Accurate Long-term Dynamics for Model-based Reinforcement Learning. IEEE Conf on Decision and Control, 2021.
Adrian Lozano-Duran and Gonzalo Arran. Information-theerrtiic formulation of dynamical sys- tems: causality, modeling, and control. ArXiv, 2021.
Xingyu Lu, Kimin Lee, P. Abbeel, and Stas Tiomkin. Dynamics generalization via information bottleneck in deep reinforcement learning. ArXiv. 2020.
Xiuyuan Lu and Benjamin Van Roy. Information thhooreti confidence bounds for reinforcement learning. Inf. Conf. Neural Information Processing Systems, 2019.
12

---

Published as a conference paper at ICLR 2025
Xiuyuan Lu, Benjamin Van Roy. Vikranth Dwaracherla, Morteza Ibrahimi, lan Osband, and Zheng Wen. Reinforcement learning, bit by bit. Foundations and Trends@ in Machine Learning, 2023.
Carlos E. Luis, Alessandro G. Bottero, Julia Vinogradska, Felix Berkenkamp, and Jan Peters. Model-Based Epistemic Variance of Values for Risk-Aware Policy Optimization. arXiv, 2023.
Nikhil Mishra, Pieter Abbeel, and Igor Mordatch. Prediction and Control with Temporal Segment Models. Int. Conf. on Machine Learning, 2017.
Volodymyr Mnih, Koray Kavukcuoglu. David Silver, Andrei A. Rusu, Joel Veness, Mare G. Belle- mare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wier- stra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature, 2015.
Shakir Mohamed and Danilo J. Rezende. Variational information maximisation for intrinsically motivated reinforcement learning. In Int. Conf. on Neural Information Processing Systems, 2015.
Anusha Nagabandi, Gregory Kahn, Ronald S. Fearing, and Sergey Levine. Neural Network Dynam- ies for Model-Based Deep Reinforcement Learning with Model-Free Fine-Tuning. In IEEE Int. Conf. on Robotics and Automation. 2018.
Michal Nauman, Michal Bortkiewicz, Piotr Mitoš, Tomasz Trzcinski, Mateusz Ostaszewski, and Marek Cygan. Overestimation, Overfitting, and Plasticity in Actor-Critic: the Bitter Lesson of Reinforcement Learning. Int. Conf. on Machine Learning. PMLR, 2024.
Evgenii Nikishin, Max Schwarzer, Pierluca D'Oro, Pierre-Lue Bacon, and Aaron C. Courville. The primacy bias in deep reinforcement learning. In Int. Conf. Machine Learning, 2022.
OpenAl, Christopher Berner, Greg Brockman, Brooke Chan, Vicki Cheung, Przemyslaw De- biak, Christy Dennison, David Farhi, Quirin Fischer, Shariq Hashme, Chris Hesse, Rafal Jozefowicz, Scott Gray, Catherine Olsson, Jakub Pachocki. Michael Petrov, Henrique Pondé de Oliveira Pinto, Jonathan Raiman, Tim Salimans, Jeremy Schlatter, Jonas Schneider, Szymon Sidor, Ilya Sutskever, Jie Tang. Filip Wolski, and Susan Zhang. Dota with large scale deep reinforcement learning. 2019.
Feiyang Pan, Jia He, Dandan Tu, and Qing He. Trust the model when it is confident: masked mosel-based acior. In Int. Conf. on Neural Information Processing Systems. 2020.
Zhongjian Qiao, Jiafei Lyu, and Xiu Li. Mind the Model, Not the Agent: The Primacy Bias in Model-based RL. arXiv, 2023.
C. Shannon. A mathematical theory of communication. Bell System Technical Journal, 1948.
Dan Simon. Optimal State Estimation. 2006. ISBN 978-0-47170855-2.
Laura Smith, Ilya Kostrikov, and Sergey Levine. DemonstratingA Walk in the Park: Learning to Walk in 20 Minutes With Model-Free Reinforcement Learning. Robotics, Science and Systems, 2023.
Richard S. Sutton. Dyna, an integrated architecture for learning, planning, and reacting. SIGART Bull., 1991.
Erik Talvitie. Ceff-Correcting Models for Model-Based Reinforcement Learning. arXiv, 2016.
Naftali Tishby and Noga Zaslavsky. Deep learning and the information bottleneck principle. 1EEE Information Theory Workshop (ITW), 2015.
Emanuel Todorov, Tom Erez, and Yuval Tassa. MuloCo: A physics engine for model-based control. In Int. Conf. on Intelligent Robots and Systems. IEFE, 2012.
Miguel Vasco, Takuma Seno, Kenta Kawamoto, Kaushik Subramanian, Peter R. Wurman, and Peter Stone. A Super-human based Reinforcement Learning Agent for Autonomous Racing in Gran Turismo. Reinforcement Learning Conference, 2024.
13

---

Published as a conference paper at ICLR 2025
Arun Venkatraman, Martial Hebert, and J.. Bagnell. Improving Multi-Step Prediction of Learned Time Series Models. AAAI, 2015.
Grady Williams, Nolan Wagener, Brian Goldfain, Paul Drews, James M. Rehg, Byron Boots, and Evangelos A. Theodorou. Information theoretic MPC for model-based reinforcement learning. In IEEE Int. Conf. on Robotics and Automation. IEEE, 2017.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. In Ady. in Neural Information Processing Systems, 2020.
14

---

Published as a conference paper at ICLR 2025
A NOTATION
A.1
RANDOM VARIABLES
S.
Random variable general state
S
Random variable the environment state
$
Random variable the model state
S.
Random variable of the estimated environment state
S.
Random variable of the Infoprop state
A
Random variable of the action
A,
Random variable of the model error
W.
Random variable of the aleatoric noise Random variable the epistemic noise Random variable of the conditional noise Random variable the model parameters
N,
U
A.2 REALIZATIONS
8t
Realization a general state Realization of the environment state
8t
Realization of the model state
St
Realization of the estimated environment state
5;
Realization of the Infoprop state
ag
Realization of the action
W:
Realization of the aleatoric noise
Rt
Realization the epistemic noise
th Realization of the conditional noise
Or.
Realization the model parameters
A.3
TRANSITION KERNELS
Ps (s = N ( Ps,rs(||Sr (-) = N (pe, At), Es, Ai)
Environment Transition Kernel
At),
Trajectory Sampling Kernel
Ps. Ar, St+1) = (ī(S,, A,, Ar, E(S, A,, Set1))
Infoprop Kernel
15

---

Published as a conference paper at ICLR 2025
B TOY EXAMPLE
In Figure ! we illustrate the data consistency of Trajectory Sampling Chua et al. al.1 and Infoprop one-dimensional random walk example with S and A R. The dynamics follow (3) with y (S;; A,) = S, + A, and L(St, A.) 0.01. Actions are distributed according to A, N (0, 0.1). All rollouts start from &u and are propagated for 100 steps. We perform 1000 rollouts under the environment dynamics and train Probabilistic Ensemble Lakshminarayana et al. (2017) model according the information provided in Table ! Subsequently, we perform 1000 model-based rollouts with this model and the respective rollout mechanism.
Hyperparameter
Value
number of ensemble members number of hidden neurons
5
number of layers learning rate weight decay
1
0.001
0.00001
number of epochs
Table Hyperparameters used for training the model on the random walk dataset.
C PSEUDOCODE ALGORITHMS
Algorithm Trajectory Sampling Chua et al. (2018)
Require: 80 while < T do
a, ~ T--(st)
Se+1 W wi, O₂ 0i] with wi 1), and Pe
81+1 E
81+1
Algorithm 3 Infoprop-Dyna (Pscudocode adapted from Janner al. (2019))
Require: Policy predictive AES model pe, environment buffer Denv, model buffer Dood rollout
parameters T, Ç1, C2. E
for N epochs do for J steps do
Interact with the environment according to æ; add Danv
Train model Pe on D Denv
Perform single-step predictions with pe in Denv
Compute À 22 and X2 (23)
for M model rollouts do
Sample sy uniformly from Deny
Perform Infoprop rollouts according Algorithm add Dmod
for G J gradient updates do
Update π on Donv U D mod
16

---

Published as a conference paper at ICLR 2025
D DERIVATIONS
D.1 QUANTIZED ENTROPY
For a RV R with Z M(uz, Ez) and discretization step size Ag(h) of the kth dimension, the quantized entropy Cover & Thomas 2006) is
H(Z) logg ((2Te))* (1(
Σ loga (A=(*))
(24)
D.2 MAXIMUM LIKELIHOOD PREDICTIVE DISTRIBUTION
D.2.1 PROOF LEMMA
Proof. We introduce the conditional expectation over the next state under the model, given a real-
ization fg
$4 E P S 0.]
(25)
Further, 3 == pe, Er = D0.=05 and i == Le. such that
$ = (S, Ar) + (S,,, Ar)Wa.
(26)
Given E RVs S8+1 we define their joint distribution
21
L
1411
~
E
E
(27)
D111
$ ~ N (1.
with Ge/ = Cov [5 1/411
We aim track S441 such that
HS41 ~ N (2.
(28)
where we use H = I.,....., We define the maximum likelihood loss
1 E Exns to project S.+1 to the dimension of the joint S.
(-2 (s HS++1) 8-1 1-1((- H(:+)))
exp
(29)
such that
log (C(S:+1)) = 1og (2mii))
(s HS 8-1 H (S-HSS41)
(30)
We aim to obtain the maximizer of the log- loglliiel such that
S arg max (L(S)-1)).
(31)
St
Consequently,
log (C(S++1)) = - - ($ (S-HHS++1)
OS,+1
HT5-15- - HT 1T511 HH,H
(32)
S = (HT H S.
As result, we obtain
Se+]] = =(HTS-HH) HTS-1
(33)
17

---

Published as a conference paper at ICLR 2025
and
Σ Var [Sc+1] Var ( (HTS-'H) H 1
SH H H H Var $ 1-11 H (HT
VVar
(34)
- -1 -T 11 (1TS~IN) -1 = (1T
-1
which corresponds to standard results in Kalman fusion. However, as the cross-correlations ΣN are unknown in practice, we approximate the Kalman fusion results (33) and (34) using covariance intersection fusion Julier & Uhlmann (200T) with uniform weights, making use of Assumption This results
Σ
2
(35)
and
1 Eイ
Σ
E Σ
110)
(36)
Hence, we can estimate the environment state as
544 (S,, At) + L(S., A.),,,
(37)
with LLT E and W₁ I).
D.2.2 PROOF LEMMA
Proof. We continue here using the quantities we estimated in the previous section. To estimate Σ we interpret { Jem samples from a distribution whose mean is known to be p. With this, the maximum likelihood estimate of د can be obtained trivially as
5
Σ
(i
(1) (it (1)
(38)
D.3 INFOPROP STATE
As introduced in [15.. the Infoprop state is defined
5.41 = S S 82-1
E Ps.I
St+1|S: = 8t, Ae a Se41 = 8441 Us
(39)
Combining 11 and Assumption have
S S.t1 + (S,, A.N,,
(40)
Plugging the respective maximum likelihood estimates into (40) yields
8.41 S.4 (S A.)
(41)
with
S A.) L(S,,
(42)
according to (3). As we can generally consider model uncertainty as independent from process noise, i.c. N₁ I W,, the Infoprop state 5 = 15 = $1+1] can be computed using standard Kalman update.
The general form of the Kalman update Simon (006 considers two Gaussian RVs X N (mx, Σx) and Y = x N with N N (0, EN) and X N. Then, given an observation
we can compute the conditional expectation of X
E X
N (μX|Y= Ex(Y==)
(43)
18

---

Published as a conference paper at ICLR 2025
with
Pxir=y = xx + K(y - -xx,,
(44)
Σx|Y= (1 K) Ex,
(45)
and
K= Ex (2xx + EN)
(46)
Following 115..
we can compute the Infoprop state via (43) choosing
px p(8t, (z),
(47)
Σx a();,
(48)
ΣN = (s),
(49)
and
y = [a(8,, 0e) + [ss, ar) + (s), This yields the propagation equation of the Infoprop state
(50)
1 (S: = St: A: = at, S = 8441) + i(s = st, Ar = at, Se+1 Stt1UU
(51)
with
ils at, 8(41) pa(sts az) + ((s,, a.) (L(se, anww, + (s): E(st, ac, 81+1) = (1 - K (st, az)) ((s),at), K(s;, az) E(s,, at) (E(s), a.) 8 (0), at)
(52)
(53)
(54)
İ( ((s, at) E(s a1);
(55)
and
[ at (L) E (st, at).
(56)
D.4 INDUCED STATE DISTRIBUTION BY THE INFOPROP ROLLOUT
Lemma 3. As introduced in 57 the next state distribution induced by the Infoprop rollout is the same as that given by the estimated ground truth:
$441 S
(57)
Proof. We show equality in distribution via comparison of the cumulative distribution functions
(CDF) of Sem1 and 5.41- If we can show that the CDFs are identical, i.c. P(S S 5,+1) =
P(S.+1 S 8t41) Vše41 E the equality in distribution follows.
We compute PPS S 8441) using S4+1 E[S141|$441 = $441] and marginalizing over S411
/
P(S.11 S 8t+1)
P( S4+1 S 8 1S $(41)fs, (6-1))d8+41
(58)
with fs... the probability density function of $4+1-
By construction, 1$++4 describes the behavior of 5.41 given S441. Consequently,
P(E ( I 1441 = $441) = P(55+1 S 5 = =8(+1)
(59)
and therefore
P(S.+1 S 3
P(Si+1 & 3141|S = $141)f (6e+ (d$111
(60)
The right hand side of 60 represents the law of total probability for P(S S Se+1)
/ P(S541 S 8e411Scn1 = 8e41)58, (8++ Id82+1
P(S0+1 S 3(41)
(61)
Therefore, we have
P(S B141) = PP(111 S 31+1) 49 E
(62)
and can conclude
Š4t dil S
(63)
19

---

Published as a conference paper at ICLR 2025
D.5 INFORMATION LOss ALONG A INFOPROP ROLLOUT
Lemma 4. As introduced in ( the total information loss incurred during a Infoprop equals the accumalated entropy the Infoprop state:
T-1
H S so, Au do. S₁ 81 S1...ST ST 5T)
Σ
H (Se+3)
(64)
Proof.
H S1, S2 ................................ So su, Ao 0o; S1 = Sr $1
H [141 |1,, S2, ..... 3R, So su, Au ao, S1 81 S1
Sr = 8r)
SS+1 |St, 12,82.... ......5. S So = 80, Ao = ag. S = SI,..., Sp = 81)
ST
(65)
H (sct1 - 0s, S:41 = =8tt1)
H (51+1)
where (a) follows from causality and (b) follows from the Markov property.
20

---

Published as a conference paper at ICLR 2025
E EXPERIMENTS
E. EXPERIMENTAL. SETUP
We used Weights&Biases for logging our experiments and run random seeds per experiment.
The respective hyperparameters for Infoprop-Dyna on MuJoCo are given below. Table 2 model learning, Table the Infoprop mechanism, and Table 4 training the model-free agent.
addresses
Table 2: Hyperparameters used to train the model of Infoprop-Dyni in the Mujoco Tasks.
Hyperparameter
Halfcheetah
Walker
Hopper
Ant
ensemble size E
number of hidden neurons
200
400
number hidden layers
learning rate
0.0003
0.0006
0.0004
0.001
weight decay
0.00005
0.0007
0.0008
0.00002
patience for early-stopping
10
9
retrain interval
250 environment steps
httpps:/wandd.. aitisiiiiieeeee
21

---

Published as a conference paper at ICLR 2025
Table 3: Hyperparameters of the Infoprop rollouts in the Mujoco Tasks.
Hyperparameter
Halfcheetah
Walker
Hopper
accurate quantile Çi
0.99
exceptionally accurate quantile $2
0.01
scaling factor g
100
rollout interval
250 environment steps
rollout batch size
100000
Table 4: Hyperparameters used to train the SAC agent of Infoprop- -yna in the Mujoco Tasks.
Hyperaarameter
Halfcheetah
Walker
Hopper
Ant
number of hidden neurons
1024
512
1024
number of hidden layers
2
learning rate
0.0003
0.0002
0.0004
0.0005
SAC target entropy
-6
target update interval
6
update steps G'
10
20
The results for SAC, MBPO and MACURA are obtained from Frauenknecht et al. (2024).
E.2 PREDICTION QUALITY
We provide additional results for the rollout consistency experiments introduced in Section 6.2 Fig- ure 6 depicts model-based rollouts for the 10th dimension of hopper under MBPO, MACURA and Infoprop-Dyy -oona when setting the maximum rollout length of all approaches 100. In the original ex- periment depicted in Figure maximum rollout length was 11 for MBPO and 10 for MACURA, following the hyperparameter settings reported in the respective publications Janner al. 2019):
Frauenknecht et al. (2024).
We observe vastly spread distribution MBPO rollouts, as every rollout is propagated for 100 steps, irrespective of model uncertainty, long it does not reach terminal state of the hop- per task. MACURA rollouts have an improved consistency compared to MBPO, especially in the beginning of the rollouts. Over long horizons, however, the TS propagation mechanism and the single-step termination criterion cannot produce consistent data. In contrast, Infoprop-Dyme is able propagate consistent rollouts over long horizons.
22

---

Published as a conference paper at ICLR 2025
10
5 -5
-10
MBPO
MACURA
Infoprop
-20
Ground Truth
Figure 6: Rollout consistency MBPO vs. MACURA vs. Infoprop-Dye for 100 steps. Comparison of the respective rollout mechanisms similar Figure 4b but with maximum rollout length of 100 for all approaches.
Humanoid
Info-Prop Dyna
MACURA
MIPO
SAC
Asymgratte SAC
Return
50
100
150 Steps
200
250
300
1e3
Figure 7: Performance Humanoid
E.3 PERFORMANCE ON HUMANOID
Figure depicts the return on MuJoCo humanoid. We observe instabilities in the performance of Infoprop-Dy) Dyna towards the end of training. We assume this occurs due to overfitting and plasticity loss in the critic of the model-free learner Nikishin atal, (2022): D'Oro (2023). This reflected in the peaking critic loss depicted in Figure concurrently with the performance drops. We set the update ratio G (see Algorithm 3 to a relatively low value of 10 which explains the slower learning
behavior than MACURA. For higher values of G, instabilities occur even earlier in the training process, underscoring our assumption of overfitting critics.
Model rollout inconsistency does not appear to be the destabilizing factor, rollout data is consis- tent with the environment distribution as depicted in Figure 10 and the rollout adaption mechanism seems to react to policy shifts induced by high critic losses through reducing the average rollout length as depicted in Figure 10
E.4 INVESTIGATING INSTABILITIES LEARNING
Although Infoprop gives better quality data over longer rollout horizons than TS rollouts, observe instabilities learning when naively integrating Infoprop into the conventional Dyna setting. We hypothesize that the main cause of these instabilities is due the agent overfitting the higher
23

---

Published as a conference paper at ICLR 2025
Humanoid
Prop Dyti
200
150
Critic Loss
100
50
0.0
0.5
1.0
1.5
2.0
2.5
3.0
Agent Updates
x10
Figure 8: Critic Loss on Humanoid
Prop Dyna)
20
Left Elbow Angular Velediy
10
-10
10-3
50
100
150
200
250
300
Density
163
Model Buffer|info Prop Dyna)
20
10-5
Left Eboow Angular Velocity
10-
-20
50
100
150
200
250
300
1e3
Figure 9: Comparison between Deny and Dmod for the 45th dimension of Humanoid
quality data produced by Infoprop rollouts, followed by loss of plasticity Nikishin et al. (2022); D'oro et al. (2023). To investigate this, carried out an ablation by varying the values of which we introduced in Equation 22 This hyperparameter controls the size the subset where the model is considered sufficiently accurate. The smaller the value of (1. the more aggressive the filtering of single-step information losses, leading to a smaller E.
Figure II shows the returns obtained the Hopper task for three values of S1. For Ç1 0.97, we see that the returns are unstable throughout training, even though this setting gives the best quality data. On the other hand, Çi 0.9999 produces more stable learning curve compared to
24

---

Published as a conference paper at ICLR 2025
Humanoid
40
Prop
MACURAA
MBPO
Average Rollout Length
30
20
50
100
150 Steps
200
250
300
x10'
Figure 10: Average Rollout Length on Humanoid
1e3
Hopper
Return
1
Ç1 0.97
< 0.99
Ç 0.9999
25
50
75
100
125
Steps
1e3
Figure 11: Ablation Study on Hopper.
Ç1 = 0.99, which was used for all the experiments in Section Ø This shows that better data quality does necessarily lead better training performance since that was the case, Ci 0.97 would have produced the best performance. A similar observation is reported in Frauenknecht et al. (024)) where low values the scaling factor ç, corresponding accurate model rollouts, led instabilities in learning.
Our observations show that producing high-quality synthetic data the conventional Dyna setting leads to issues seen in MFRL. when using high update-to-datt (UTD) ratio. There have been recent works on regularization methods counteract agent overfitting and loss of plasticity. One such approach is applying layer normalization Smmth et al. (223) Nauman et al. (2024). Figure T2 shows the same settings as in Figure [] but with layer normalization applied to the critic and actor networks. It can be seen that even for 0.97, the learning is stable.
The primary aim of this paper is introduce the conceptual framework of Infoprop rollout, well as show its application to MBRL. Hence, do not spend additional effort tuning the hyperparameters or adding regularizations since this takes us away from the main objective. We defer such enhancements for future work.
25

---

Published as a conference paper at ICLR 2025
1e3
Hopper with Layer Normalization
0.97
Ç 0.99
ç 0.9999
Return
2
1
25
50
75
100
125
Steps
1e3
Figure 12: Ablation Study on Hopper with layer normalization.
26


Reviews:
[{"review_id": "yFUMaOBxA6", "reviewer": "Reviewer_j8sE", "summary": "This paper proposes Infoprop, a rollout mechanism designed to improve sample efficiency and asymptotic performance in model-based reinforcement learning (MBRL). Infoprop addresses the issue of epistemic model uncertainty typically introduced by ensemble methods or dropouts. It projects sampled next states into the maximum likelihood distribution of the predicted next state distribution and uses an information-theoretic criterion to determine the termination point of rollouts. When applied to Dyna-style MBRL, Infoprop demonstrates state-of-the-art performance, as shown in experiment results.", "questions": "1. Is the model error come solely from epistemic noise?\n2. What is the advantage of using AES models, and why do you believe Infoprop does not impact their effectiveness?", "limitations": "", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "1. **Significance**: Model error remains a significant challenge in MBRL, and Infoprop addresses it through the rollout process.\n2. **Extensive Experimentation**: The authors provide comprehensive experiments and analysis, including detailed evaluations of algorithm trajectories and environment buffers, which validate the effectiveness of the Infoprop algorithm.", "weaknesses": "1. **Assumptions**: The approach relies on two key assumptions. It is important to empirically assess the extent to which these assumptions hold in experimental environments.\n2. **Organization**: The theoretical section could benefit from clearer structure. Presenting the main results in the form of definitions or theorems could improve readability. Simplifying technical details and focusing on the implementation—particularly regarding Equation (14)—would make the logical flow smoother.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "XUDU1dZaaK", "reviewer": "Reviewer_NHA5", "summary": "This paper presents a rollout mechanism that uses information theory to reduce the impact of epistemic uncertainty in MBRL. The paper presents some experimental analysis and a derivation of a result related to information loss.  The paper uses the idea of marginal entropy and mutual information to determine data corruption and stop a rollout.", "questions": "The main limitations of the work, as identified above, are in the quality and significance as discussed above. These could be alleviated with a more in-depth experimentation which would lead to a more mature solution, and a more in-depth discussion of related work. All of them require space, and a such my questions relate to the guidance for the discussion section.\n\nSpecifically:\n1. where is overfitting most prominent with Infoprop?\n2. where there other information theoretic metrics (e.g redundancy) considered in the Infoprop design? If no, why not? if yes, why were they discarded?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "Soundness\n======\nThe paper uses the idea of marginal entropy and mutual information to determine data corruption and stop a rollout in MBRL. The idea is presented carefully with a derivation for the main result.\n  \nSignificance & Related work\n=========\nThe paper has a brief related work section, as part of the main paper, which is appreciated. However, there is a need for a deeper analysis of existing work, including  more in-depth comparisons that show the benefit of Infoprop to others, specifically including work in using simulated rollouts to guarantee performance, robust MBRL (Sung 2024, Kuo 2021), risk sensitivity (e.g. Webster 2021) and other uses of information theory in MBRL.\n\nPresentation\n=========\nThe paper is well written with few typos, usually around empty spaces in the text.\n\nQuality\n======\nMy main concerns are around the quality of the experimentation (see below) and on the need for authors to explore the limitations of their approach a bit deeper before publication. For example, the authors mention that the rollouts exhibit excellent data quality and yet there are instabilities: is this a function of the integration mechanism or of overfitting?", "weaknesses": "Experimentation\n=========\nThe experimental analysis show the benefits of Infoprop when it runs in simple environments, however deeper experimentation is needed. For example, I would have expected to see some ablation studies showing the impact of the choice of lambda, instability, overfitting and other issues. \n\nOverall the paper presents a very promising initial idea with several avenues of experimentation that need to be explored to ensure its maximum potential is reached.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "ulyMfkU99A", "reviewer": "Reviewer_uFDs", "summary": "This paper proposes a model rollout method for model-based reinforcement learning called infoprop. The main idea is to control for epistemic uncertainty using model rollout to reduce model error accumulation and in turn improve value or policy optimization. The concrete method is a re-weighting of ensemble predictions and an entropy-based criteria to stop rollout. The Mujoco experiments demonstrated substantial improvement in reducing error accumulation as well as convergence speed.", "questions": "* How does eq 14 actually follow from eq 4 and 9? \n* When the authors say eliminate epistemic uncertainty, does that mean setting the last term in eq 9 to zero? If so is the rest of the derivations simply accounting for $\\mu^{\\Delta}$?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "The proposed method is original and sound. The experiment hypotheses are very clear the the results clearly demonstrated the advantage of the proposed method.", "weaknesses": "The paper could benefit a lot from intuitive explanations of the proposed approach. For example, what's the intuition of the covariance intersection fusion result for readers that are not familiar with this literature.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "C86KnsGR3u", "reviewer": "Reviewer_Anzx", "summary": "The author presents a new method called *Infoprop*, which can be used as a plugin for Dyna-style model-based RL. The plugin increases the quality of data rollouts from the estimated model, thus enhancing the Dyna-style model-based RL. The experiment in Mujoco showed the SOTA performance except for the humanoid environment. The author also provides a solid theoretical analysis of why the plugin can increase the quality of rollouts based on the estimated model.", "questions": "In Figure 6, INFOPROP Dyna does not outperform MACURA. However, if I understand correctly, the proposed method can be used as a plugin within MACURA, which makes it puzzling that INFOPROP Dyna does not surpass MACURA.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The author developed a new method, Infoprop, that enhances the performance of Dyna-style model-based RL:\n\n*Versatility*: The Infoprop method can be integrated as a plugin for Dyna-style model-based RL, making it a flexible and adaptable enhancement for existing frameworks.\n\n*Improved Data Quality*: By increasing the quality of data rollouts from the estimated model, Infoprop enhances the overall performance of Dyna-style RL approaches.\n\n*Experimental Results*: The experiments conducted in Mujoco environments demonstrate state-of-the-art (SOTA) performance, showcasing the effectiveness of Infoprop in most environments, with the exception of the humanoid scenario.\n\n*Theoretical Foundation*: The method is backed by robust theoretical analysis, explaining how the plugin improves rollout quality from the estimated model, which adds credibility to the approach.", "weaknesses": "1. The theoretical results in Section 4 are difficult to follow. A summary of the main result should be given at the start of the section. Also, I suggest the author clearly present the main result as a theorem. \n2. Too many notations have been introduced, which could be simplified. For example, Section 2.4 could be condensed into a single sentence, accompanied by equation (5) and relevant citations since it represents a well-known formulation in model-based RL.\n3. The enhancement applies specifically to Dyna-style RL, not all model-based RL approaches. Making the title misleading.\n4. The Assumptions at the end of Section 4.2 are stated without any explainations.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (Poster)