On Rollouts in Model-Based Reinforcement Learning

**W1:** Thank you for pointing this out. We agree and have revised Section 4 accordingly, as discussed in point 1 of the general rebuttal. **W2:** Thank you for pointing this out. We strongly believe that all introduced notation is necessary to precisely distinguish the different rollout configurations which is crucial for a clear presentation of our results. We have put in a lot of effort to make our notation as compact as possible. However, we agree that there is a lot of notation one needs to keep track of. We plan to have a notation list that includes descriptions, definitions, and update rules of different objects in Appendix A. This Appendix can be opened on a second screen and can help navigate our notation. In order to upload the revised version in a timely manner, we have not finished this list yet but will do so for a potential camera-ready version. **W3:** The main contribution of the paper is the Infoprop rollout mechanism that is applicable to different types of MBRL algorithms using AES models. We illustrate the capabilities in a Dyna-style architecture, however, Infoprop is not limited to these architectures. Infoprop could benefit all MBRL algorithms using TS rollouts, e.g. the ones mentioned at the end of Section 2.4. these comprise model-based planning, analytic-gradient and Dyna-style architectures. Infoprop should likewise be applicable to model-based value expansion methods, e.g. a STEVE [1] -style approach using probabilistic ensemble models. [1] Sample-Efficient Reinforcement Learning with Stochastic Ensemble Value Expansion, Buckman et al. 2018 **W4:** Thank you for pointing this out, we have addressed this concern in point 2 of the general rebuttal. **Q1:** Thank you for your question. The reason we present Infoprop-Dyna is threefold: (1) to give an example on how to define $\lambda_1$ and $\lambda_2$ in practice, (2) to show that naively integrating Infoprop into a given algorithmic architecture performs reasonably well, and (3) to showcase the substantially improved quality of predictive data in a realistic scenario. We do not aim to propose an extremely high-performing Dyna-style algorithm. Thus, Infoprop-Dyna builds on the simpler MBPO architecture (e.g. we do not have an adaptive number of critic updates as MACURA). Further, we did not put additional effort into fine-tuning the algorithm or counteract the reported critic overfitting issues, as this would distract from the intended main contribution of the Infoprop rollout scheme. We have added an ablation discussing tuning and stabilizing the algorithm in Appendix E4. Consequently, we do not consider it surprising or critical that MACURA outperforms Infoprop-Dyna on humanoid.

**W1:** Thank you for pointing this out. As discussed in point 1 of the general rebuttal, we did a major revision on how Section 4 is presented. We now provide an intuitive explanation of the Infoprop approach itself in the beginning of Section 4. Concerning covariance intersection fusion, we rather consider this a technical detail that is not crucial for understanding the general approach. Thus, we have moved this to Appendix D2. If requested, we can provide a geometrical interpretation of CI fusion. For now, we have decided not to do so to keep the appendix as compact as possible. **Q1:** Thank you for the valuable remark. We now provide the full derivation of the Infoprop state in Appendix D3. **Q2:** The term 'eliminating epistemic uncertainty' was a misleading formulation. Pointing us to this helped improve the presentation of Section 4. What we meant is that we want to remove the influence of epistemic uncertainty from the predictive distribution. However, we cannot set N to 0 in practice as we don't have access to N directly. Instead, we have to sample parameter realizations and interpret them as epistemic noise realizations (see discussion around Definition 4). From these samples, we estimate the environment distribution and the noise distribution. Finally, we want to make sure that Infoprop follows the estimated environment distribution (which it does, see Theorem 1 (i)) while neglecting the noise distribution. We hope the presentation is clearer now.

**Related Work:** Thank you for the remark, we are happy to discuss how to incorporate these papers. However, while the proposed works address interesting problems, we are not exactly sure how to incorporate them into the related work section. While these methods could potentially benefit from using Infoprop rollouts, e.g. being less conservative wrt. safety certification as the uncertainty can be specified more precisely, their key contributions are different from our approach. So from our understanding, these works are potential application fields of Infoprop but not primarily works we need to differentiate our approach from. Do you agree on this or did we get you wrong? Further, we are not aware of further literature regarding information theory in MBRL to discuss in the related work section. In case you have works in mind we missed, we would be happy to incorporate them. **W1:** Thank you for well thought through feedback. We have added an ablation study in Appendix E.4, investigating the instabilities in learning resulting from training the agent using high-quality data. We performed experiments on the Hopper task with three different values for the hyperparameter $\zeta_1$. The smaller the value of $\zeta_1$, the more aggressive the filtering of single-step information losses, resulting in fewer inaccurate model transitions being added to the replay buffer. We observe that the training destabilizes for the lower value of $\zeta_1=0.97$. In fact, the higher value $\zeta_1=0.9999$ gives a more stable learning performance. We hypothesize that using higher-quality data in MBRL causes the same destabilizing issues faced by model-free RL (MFRL) when using a high update-to-data (UTD) ratio[1][2]. To counteract this effect, different regularization approaches have been developed for MFRL, the most prominent of which is layer normalization[3][4]. We performed a separate set of experiments using layer normalization for the actor and critic networks and observed that this does stabilize training even for $\zeta_1=0.97$. Hence, in our opinion, the instabilities we report are not due to our proposed rollout scheme rather due to pre-existing failure cases with the model-free agent. Since the objective of this paper is to lay out the mathematical framework for Infoprop, and show its applicability to MBRL, we do not do extensive hyperparameter tuning or incorporate regularization in the main results. [1] The primacy bias in deep reinforcement learning, Nikishin et al., 2022. [2] Sample-efficient reinforcement learning by breaking the replay ratio barrier, D'Oro et al., 2023. [3] Demonstrating a Walk in the Park: Learning to Walk in 20 Minutes With Model-Free Reinforcement Learning, Smith et al., 2023. [4] Overestimation, Overfitting, and Plasticity in Actor-Critic: the Bitter Lesson of Reinforcement Learning, Nauman et al., 2024. **Q1:** As mentioned in W1, the overfitting is more prominent when using more aggressive filtering of model transitions, which gives better quality data. We have shown that regularization approaches developed for using high UTD in MFRL can be used to mitigate the instabilities in learning. **Q2:** Q3: There is a set of requirements for the metrics used for such a rollout: - The resulting RV needs to follow the environment distribution. - The resulting sequence of RVs needs to provide a notion of information loss. - This notion of information loss needs to allow for a feasible cutoff criterion (e.g. as the conditional entropy sums up over the trajectory we can simply use thresholding). - The mechanism needs to be numerically tractable. Generally, there might be different metrics fulfilling these requirements. However, we consider it nontrivial to come up with a method building on a different metric that fulfills all of the requirements above.

Thank you for your reply and for considering our revision. From what we understood, the main concern mentioned in your review was insufficient ablation studies concerning critic overfitting in the naive integration scenario. We were hoping to address this by adding Appendix E4 to the revised version. Here, we provide several ablations that clearly show when overfitting occurs and that this is a problem of the downstream Q-learning architecture and not the presented rollout scheme. Do you agree that this part is now clearer and does not leave ambiguity?

**W1:** Thank you for pointing this out. I would kindly refer you to point 2 of our general rebuttal. **W2:** We agree and have addressed this in the revised version (see points 1 and 3 of the general rebuttal). **Q1:** Yes, in our set of assumptions that hold well in practice as discussed above, epistemic noise is the only source of model error. **Q2:** The major benefit of AES models over other model architectures is that they can distinguish aleatoric uncertainty due to process noise from epistemic uncertainty due to model error. This allows to model stochastic systems while having a notion of model uncertainty. However, AES models combined with the standard TS rollout [1] are only capable of making a reliable distinction of this kind for a single transition. As depicted in Figure 1, the TS mechanism mixes both types of uncertainty when using the model for multi-step predictions. Infoprop strengthens the core benefit of AES models (namely the separation of aleatoric and epistemic uncertainty) as it enables this separation for multi-step rollouts instead of just single-step predictions. [1] Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models, Chua et al. 2018