Leveraging Conditional Dependence for Efficient World Model Denoising
Authors: Shaowei Zhang
Jiahan Cao
Dian Cheng
Xunlan Zhou
Shenghua Wan
Le Gan
De-Chuan Zhan
Venue: NeurIPS 2025 poster ()
Keywords: Model-Based Reinforcement Learning
Visual Inputs with Distractors
Conditional Dependence
Primary area: reinforcement_learning
TL;DR: A model-based reinforcement learning framework that effectively extracts task-relevant information from noisy observations by leveraging conditional dependence.

Abstract: Effective denoising is critical for managing complex visual inputs contaminated with noisy distractors in model-based reinforcement learning (RL). Current methods often oversimplify the decomposition of observations by neglecting the conditional dependence between task-relevant and task-irrelevant components given an observation. To address this limitation, we introduce CsDreamer, a model-based RL approach built upon the world model of Collider-structure Recurrent State-Space Model (CsRSSM). CsRSSM incorporates colliders to comprehensively model the denoising inference process and explicitly capture the conditional dependence. Furthermore, it employs a decoupling regularization to balance the influence of this conditional dependence. By accurately inferring a task-relevant state space, CsDreamer improves learning efficiency during rollouts. Experimental results demonstrate the effectiveness of CsRSSM in extracting task-relevant information, leading to CsDreamer outperforming existing approaches in environments characterized by complex noise interference.

Leveraging Conditional Dependence for Efficient World Model Denoising
Shaowei Zhang! Jiahan Cao Dian Cheng! Xunlan Zhou', Shenghua Wan Le Gan 1.2 and De-Chuan Zhan
National Key Laboratory for Novel Software Technology, Nanjing University, China d School of Artificial Intelligence, Nanjing University, China
School of Intelligence Science and Technology, Nanjing University, China
(zhangaw, caojh, chengd, wamsh)-0lamd... njj.. nju. edu. en, wyattzhoux1@email. iopeiinnjjuu... .nju. eju.edu.cn, {ganl, zhandc]Onju (aaddd)@jjueeuu.... ed cn.
Abstract
Effective denoising is critical for managing complex visual inputs contaminated with noisy distractors in model- reinforcement learning (RL). Current meth- ods often oversimplify the decomposition of observations by neglecting the con- ditional dependence between task-relevant and task-irrelevant components given an observation. To address this limitation, we introduce CsDreamer, a model- based RL approach built upon the world model of Collider-structure Recurrent State-Space Model (CsRSSM). CsRSSM incorporates colliders to comprehensively model the denoising inference process and explicitly capture the conditional de- pendence. Furthermore, it employs decoupling regularization to balance the influence this conditional dependence. By accurately inferring task-relevant state space, CsDreamer improves learning efficiency during rollouts. Experimental results demonstrate the effectiveness of CsRSSM in extracting task-relevant infor- mation, leading to CsDreamer outperforming existing approaches in environments characterized by complex noise interference.
Introduction
Reinforcement Learning has achieved remarkable success in complex applications such as au- tonomous driving []] and conversational interactions 3 A key factor in these advancements is the agent's ability to accurately perceive and interpret observations from its environment. However, real-world observations are often corrupted by noise and extraneous information, which can severely impair an agent's 's ability to make sound decisions. This challenge intensifies significantly when dealing with high- diteenninnii observations like images [4] 5 6). Therefore, effectively managing such noisy observations is essential for developing robust and reliable RL agents.
From generative model perspective 77, at timestep an observation or is generated by two distinct sets of latent variables: task-relevant variables 8t, which directly influence the agent's rewards or actions, and task-irrelevant variables ct, which introduce noise and distractions without contributing to the task (Figure [ (a)). Existing model-based RL approaches have sought to extract task-relevant information from observations g 9]. Some methods utilize two separate dynamics models to extract task-relevant and task-irrelevant information from observations [00 while others decompose observations into task-relevant and task-irrelevant branches [2] or pursue more fine-grained decomposition [3] [4]. However, these methods oversimplify the decomposition of observations by neglecting the conditional dependence between s; and given o₁. The importance of this conditional dependence can illustrated with a simple example: consider random variables
Corresponding author.
'hhe code is available at //tttts://ggthhiii http:hhhww..cn/haaaaaaaaaaae///e
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

A, B. and C' where A B = C. C is unknown, A and B are independent. However, once C' is observed (e.g., C c), their joint probability becomes P(a, d =c = P(a. P(a.B)((++b=o) P(Cme)
This implies that if determine then B must be
a, Thus, A and B become conditionally dependent given C. This conditional dependence can be leveraged for efficient denoising. For instance, as illustrated in Figure T (b), previous approaches perform separate inferences for 8, (the walker's 's state) and e (the background) by assuming condi- tional independence. However, the back- ground information can significantly in- fluence the inference of the walker's ' state. By capturing the conditional de- pendence between 8; and Ct. we can sig-
Conditinnal Depenseane
Previous
Ours
Gencration
(b) Inferemee
Figure I: (a) Generation process of noisy observations. (b) Inference process where we consider the conditional dependence.
nificantly enhance the inference process, as the walker's 's state can be rapidly deduced by comparing differences between the background and the noisy observation. Therefore, capturing this conditional dependence is beneficial and can improve the inference process.
Inspired by this, we propose Collider-struclute Recurrent State-Space Model (CsRSSM), world model designed for denoising from a generative modeling perspective. CsRSSM extends Recurrent State-Space Model (RSSM) [15] by modeling the noisy observation generation process sequen- tial collider-structurg model. By explicitly utilizing the collider structures, CsRSSM exploits the conditional dependence between 8, and o given o, to effectively extract task-relevant information. Specifically, CsRSSM first infers the task-irrelevant variable c from Og and then utilizes both oa and O₂ to infer the task-relevant latent variable St. This sequential inference approach utilizes the contextual information provided by task-irreevaat variables to enhance the efficiency of task-relevant information extraction. Additionally, regularize the model using conditional mutual information to balance the conditional dependence between and ea during learning. Furthermore, existing methods typically employ masking mechanisms (TO, 12] tailored optimization objectives [16), which depend on prior knowledge or assumptions, to distinguish telk-relevant information from task-irrelevant information. Compared with these strategies which may constrain model flexibilityy and introduce biases, CsRSSM adopts a generative modeling perspective, employing dual latent variables to model the generation of noisy observations and the environment dynamics. relies solely on the decoupling regularization and the inherent optimization objectives of the generative model, without incorporating any additional prior knowledge or supplementary optimization objectives. Experimental results demonstrate the feasibility and effectiveness of this approach.
Finally, we introduce CsDreamer, extension of the Dreamer framework 17), that employs CsRSSM to effectively extract task-relevant information from noisy observations. By conducting rollouts exclusively within the space st. CsDreamer enhances the agent's 's ability efficiently learn policies in high- -igh-dimensional visual environments characterized by complex noise distractors. Our contributions are threefold:
We propose CsRSSM, novel generative model for environment dynamics that explicitly utilizes conditional dependence to model the noisy observation generation.
We introduce CsDreamer, which builds upon CsRSSM to effectively denoise observations and train policies within the task-relevant space.
We conduct extensive experiments demonstrating that our approach outperforms existing state-of-the-an methods in handling noisy observations.
2 Preliminaries
Recurrent State-Space Model. Recurrent State-Sp Saate-Space Model (RSSM) 15] can capture the dynamic relationships in the latent space. Dreamer [8
17 adopts a world model based on RSSM to construct the dynamies model of the environment perceived by the agent in the latent space. the real environment, the agent receives observation input o and performs an action d,, after which it receives the next observation 0t+1 following the environment's 's state transition governed by the proba- bility distribution P(o++1 lot, oz). The world model captures the dynamics in observation space based

---

on a dynamics model in the latent space 2, which primarily comprises the following components: the trajectory history model h, = fo 2:1 at-1), the encoder model go |h,, o.), the transition model in latent space P& p& (εg(/4) and the reward model (F l zt), where & represents the parameter vector of the world model. The dynamies model in 2 uses the same action space A as the real environment. In the dynamics model, h records the historical trajectory information, (o Ihe 0e) extracts the useful information from the observation, and po (fel/t) represents the transition in the latent space.
3 Related Work
Reinforcement Learning with Noisy Observations. Many recent studies have focused on the reinforcement learning problem with noisy observations. Some studies adopt metrie-based repre-
sentation learning approaches to deal with noisy observations [20] 21 22
23
MIIR [24] adopts
information -hoomation-thhoreticc principles to learn invariant representations. InfoPower 8 prioritizes action- correlated factors. RePo 66 encourages reward-predictivy yet compact encodings. Another line of Additionally, structured latent variables have been explored for denoising. SEAR (28 partitions agent- and environment-relate
work relies on data augmentation to reduce distractors [25 26 27
components via segmentation, while DEAR (29 applies a similar bask- approach without re- constructing the entire observation. Other methods decompose observations by training two separate dynamics models to separate task-relevant and task-irrelevant information [10 11. splitting them
into distinct branches [12] adopting more fine-grained decomposition I3 14 However, these
decomposition- approaches typically assume conditional independence between factors given the observation.
Generative Models in Reinforcement Learning. Generative models have played a pivotal role in tackling nigh-dimensional observations for RI. Some works employ the Variational Autoen- coder (VAE) 30 31 to learn lower-dimensional latent variables from images, facilitating latent-space policy optimization 332 [5]. Flow-based generative models [33 34 also provide way to learn flexible distributions with explicit likelihoods, though balancing model complexity with real-time efficiency can be challenging in RL settings. Diffusion models B5 performance in image generation, and their potential for RL has begun to attract attention [37]. However, most of these generative approaches focus primarily on high-fidelin reconstruction or predictive modeling, as opposed to explicitly separating out different observational factors such as background noise irrelevant distractors.
36) have demonstrated impressive
Disentanglement in Reinforcement Learning. Several studies concentrate on disentanglement within reinforcement learning to derive efficient representations for behavioral learning. TED (38) and CMID 39 introduce disentanglement techniques into feature learning process for RL. Our method draws inspiration from them while focusing on decoupling between two sets of latent variables rather than disentangling the original feature into individual feature dimensions. Some causal reinforcement learning approaches target the disentanglement of different state components [40] [6], while they
pay less attention to the observation generation process or employ counterfactual and intervention mechanisms 42 43 Instead, we treat noisy observations from a generative model perspective and focus on conditional dependence between task-relevant and task- aask-irrelevant variables during inference. Although Cao et al. [44] also notice the conditional dependence, they combine two separate generative models rather than utilizing a unified model directly.
4 Method
In this section, we first introduce the underlying assumption regarding the structure of noisy ob- servations and then present the stlleider-strucuure Recurrent State-Space Model (CsRSSM) from a generative modeling perspective in Section 4.1, Subsequently, we introduce a decoupling reg- ularization according to the characteristics of the network in Section 4.2 to facilitate subsequent reinforcement learning. We present the overall loss objective for the CsRSSM world model in Section 4.3 and introduce the policy learning in Section 4.4
The overall framework of the CsRSSM
world model is shown in Figure 2
4.1 Collider-Structure Recurrent State-Space Model
In real-world scenarios, agent observations are often contaminated by noise and irrelevant information, which degrades decision-making performance. From generative modeling perspective, we propose the following assumption to elucidate the relationship between latent variables and the generation of observations in such environments.
3

---

11
Q2
62
6x
Observation decoder
F
F2
$3
Task info predictor
11
Q3
C
C2
C3
02
02 03
Observation encoder
(a)
(b)
Figure 2: Framework of the CsRSSM world model. For brevity, content related to historical information in Eq. (3) is omitted. (a) Probabilistic graphical model of CsRSSM. Shaded nodes represent observed variables such actions and observations. Among the unshaded nodes, 02 and og are future variables that have not yet been observed at timestep / = 2, while the remaining unshaded nodes are latent variables. The black solid lines represent the generative model, and the blue dotted lines illustrate the inference process, which accounts for conditional dependence. (b) The training framework of the world model.
Assumption 4.1. (Collider structure assumption) There exist two distinct sets of latent variables that jointly determine the generation of the observation ot E O: the tusk-relevant latent variable S. which is associated with the agent's ss actions and rewards, and the task-irrelevant latent variable C C, which pertains solely to interference.
This assumption posits that the observations an agent perceives are composite two underlying factors and acknowledges the conditional dependence between the latent variables s, and e given the observation o₂. This nuanced modeling ensures that the interplay between task-relevant and musk-irrelevant factors is accurately captured, thereby enhancing the agent's ability to make informed and robust decisions despite the presence of noise and irrelevant information. Under Assumption we can obtain the following generative process
p(s,, ct, 0z) = p(c))((s)pp(olle st).
(1)
which is similar the former methods [T0
12 13
Due the presence of collider structures, condi-
tional independence typically does not hold during inference; that is, plct, st)o:) F p(ce o, )(ss lo)).
Instead, recognize the dependence between 8, and c, and employ the chain rule of probability rather than the arbitrary factorization based on conditional independence to factorize the inference probability. There are two ways implement the probability chain rule: either by inferring c first or inferring 8t first. In scenarios where agents receive noisy observations, the provision of additional task iireevrrllevat information can aid in inferring task-relevant information given the observations of an agent, as indicated by reduction in conditional entropy: H(ss|c,, oz) I H()) Therefore, we utilize
p(c, selo,) p(cclo:)p(ssle...))
(2)
where the main difference lies in the inclusion of c for the inference of st. We first use the observation o, to infer qt, the task iss-virreeevntt information, and then utilize both e and og infer st. the task-relevant information.
To better capture the temporal information in noise, model the task-irrelevant variables as time-varying. Constant noise can be regarded as a special case time-varying noise where e remains constant. In practice, the noise at the next moment is primarily influenced by the noise at the current moment. Hence, we posit the following assumption.
Assumption 4.2. The transitions of the bask-irrelevanl latent variables et satisfy the Markov property. i.e., Ct p(eelce)) = p(cx(ot-1).

---

Remark 4.3. Assumption 4.2 is analogous to the Markov property in Reinforcement learning. with the distinction that transitions of e, are independent of the agent's 'ss actions.
Consider the sequences {ot, at, Ft, 2e ) where is the timestep, a, is the action taken by the agent, Ts is the reward signal, 2z is the episode continuation flag and is the length of the sequence. Under Assumption 4.1 and Assumption 4.2 we propose the CsRSSM world model, which leverages the conditional dependence of the collider structures to model the noisy observation generation process from generative modeling perspective:
iask-irrelevvat history:
hf = folhi-11 c1-1)
Task-relevant history:
he = 11(1-111 11111111111 12-11
Task-relevant transition:
Po (8g11)))
Task-irrelevant transition:
Po((e(lif)
irreeeitel encoder:
20(c he, 0t)
(3)
Task-relevant encoder:
(0 14,, 14, 0t) Po 8r,
Reconstruction model:
Reward model: Continuation predictor:
(f st),
p (
where é is the parameter vector of the world model. Analogous to the RSSM [55, our approach incorporates historical information. We denote the task-irrelevant history as h2 and the task-relevant history as hę- The transition models provide the transition priors: P predicts the next task-relevant state based on the current task-relevant history, while Po(2 predicts the next task-irrelevant state from the current task-irrelevant history. The two observation encoders operate similarly to the transition models but incorporate observations inputs.
The inclusion of historical information introduces conditional dependence, and the posterior distri- butions of s, and C capture all conditional dependence between task-relevant and task-irrelevar information within the collider structures given the observations. Specifically, for each collider structure, the model first uses the history and the current observation to infer the task-irrelevan information et. Subsequently, utilizes the history along with both o, and e to infer the tesk-nelevant information 8y. We use the features of both c and s; to reconstruct the observation. In contrast, we only use the features of & to predict the task information, such as rewards and episode continuation flags. These task information reconstructions facilitate the learning of task-relevant latent variables 8,. Finally, we derive the evidence lower bound (ELBO) loss for the CsRSSM world model from generative modeling perspective:
CELBO
Eiin po (o hig st (g,c)) Inpo(re(((sss)) + (n ln:o(z, st)l
(4)
E (C (c.. E[KL (go (sllp(
where we maximize the ligelike of observations, rewards and continuation flags. The two KL divergence losses serve separately for task-relevant and task-irrelevant information, simultaneously training the priors toward the representations and regularizing the representations toward the priors. The derivation of Eq. 4 can refer to Appendix I Notably, when faced with noisy observations, the presence of conditional dependence mitigates trivial solutions where o. and s; Oz, ensuring that both et and 8, are effectively utilized.
4.2 Decoupling Regularization in CsRSSM
Directly training the CsRSSM world model may present some issues. Specifically, since the model training primarily relies on the reconstruction objectives, there is a risk of conflating the task- relevant and task-irrelevant variables during the training process. This conflation can hinder the model's s ability to accurately disentangle the underlying factors of the observations, ultimately de- grading performance and slowing convergence. To mitigate these issues, we propose explicitly decoupling the task-relevant latent variables s, from the task inreiitelevaat latent variables ct. We
5

---

measure the information shared between s and ot given using the conditional mutual informa- tion, denoted as [sy; co)or). Noting the relationship that the KL divergence between conditional dependence and conditional independence is equivalent to the conditional mutual information, i.e., KT.[p(ee) (or )p((s & [ct, ot) ()p((p((())))))))))) (c)))))))(( = T(x;; ee (o:), we achieve decoupling by introducing a regularization loss that minimizes the conditional mutual information between &t and er given the condition variables {o<e, ).e). denoted as [srr eq hhg hi, or). Minimizing this mutual information encourages the task-relevant and task-irrelevant variables to capture distinct and non-overlapping aspects of the observations, thereby facilitating more effective learning for CsRSSM. Directly min- imizing the conditional mutual information is intractable. By incorporating qg gg(st [hG hū, ot), variational distribution with parameters & that approximates the true posterior, we can obtain an upper bound loss for the conditional mutual information:
L = E ET NG, hf, ot) (ag h4, ot)]];
(5)
where qa (8x ce MG, h2, 0.) is the task-relevant encoder of CsRSSM. The Derivation can refer to Appendix A.2 We can minimize the conditional mutual information for the task-relevant latent variables st and the task-irrelevant latent variables e by minimizing the upper bound in Eq. (5).
4.3 Overall Objective for CsRSSM
Integrating the discussions in Sections world model
4.1 and 4.2 we obtain the final loss function for CsRSSM
CCARSSM LELBO + ALMI.
(6)
where the second item serves as the regularization of the ELBO loss to balance the conditional dependence between and Cz. The hyperparameter À controls the weight of the regularization. larger enforces greater separation between 8 and ct. promoting distinct representations of task-relevant and task-irrelevant information within the latent space, while smaller À leverages the conditional dependence to facilitate more efficient inference.
Previous studies operate under the conditional independence assumption for 8, and or 10 given an observation. This assumption is special case of our formulation where the regularization coefficient approaches infinity. In such scenarios, LII dominates the loss function, effectively reducing the mutual information between 8t and c to zero. This enforces strict independence between the task- relevant and iiseeirtelevant variables given an observation, aligning with the conditional independence assumption. Our approach generalizes existing methods by introducing finite A, enabling flexible trade-off for the degree of decoupling enforced between and c,. This extension incorporates the conditional dependence between &t and et for world model denoising. More experimental discussions
are provided in Section 5.3
Moreover, existing methods typically employ masking mechanisms [T0 12]
tailored optimization objectives [6). which depend on prior knowledge or assumptions, to distinguish task-relevant infor- mation from task-irrelevant information. However, these strategies may constrain model flexibility and introduce biases. Instead, we adopt a generative modeling perspective, employing dual latent variables to model the generation of noisy observations and the environment dynamics. The model is trained solely based on the inherent optimization objectives of the generative model (the reconstruc- tion errors and the prior regularizations in Eq. (4), and the decoupling regularization according to the model network (Eq. (5). It does not incorporate any additional prior knowledge or supplementary optimization objectives.
4.4 Policy Learning in CsDreamer
Finally, we propose CsDreamer, which is an extension of Dreamer I7 built upon CsRSSM. The agent's 's policy is learned during the imagination phase of the CsRSSM. By focusing on task-relevant information, we only utilize s; to choose the next action. i.e., a ~ Te(a, hς. st). where the policy of the agent. During the imagination phase, the agent interacts solely with the relk-atl transition model, (r [h.], to generate rollout trajectories, with rewards predicted by the reward model pa (frlhç,8)) We predict the value utilizing V4 (h+, 8t) based on 8;. The agent policy is trained analogously to Dreamer [17. Focusing on task-relevant information and leveraging the strengths of model-basee reinforcement learning ensure efficient and targeted policy development.

---

Walker Run GB
Cheerah Ruc. GB
Finger Spon GB
Hopper Hop + GB
130
Newn 32kord Bature
300
20-
30-
200
105.0
3003
3050
4008
$30.0
206.0
2002
100.
500.
106.0
3015
3050
<000
100.0
200
200.
100.0
5000
Sept:
Stepe 00
Sepp. [X)
Steys 00
Dovase
DAC
Figure 3: Performance DMC using gray natural videos background.
Wilkee Run
Cheeuah Run
CB
Finger Spla nn CB
Hoppes Hop CB
1030
200.
500.
750.
1000.
250.2
500.
750.
1000.
0.0
252.0
520.0
7520
10000
0.0
250.0
530.0
752.0
1000
Stea 00
Stzze 00
Stepe 00
Stpi 022
Draazer
Figure 4: Performance DMC using colorful natural videos as background.
5
Experiments
We begin by describing the experimental setup in Section 5.1 Subsequently, we address the following research questions: (1) Can CsDreamer enhance the training performance in complex environments with distractors? (Section 5.2]: (2) What role does the decoupling regularization loss play in the overall objective, and how do the proposed modules affect performance? (Section [53); (3) Do the latent variables in the proposed framework exhibit interpretable semantics upon reconstruction visualization? (Section 5.4).
5.1 Experiment Setup
Benchmarks. We evaluate our model and baselines on visual control tasks. First, we assess performance on four DeepMind Control Suite (DMC) 45 tasks: Walker Run, Cheetah Run, Finger Spin and Hopper Hop. To introduce noise distractors, we replace the original backgrounds with two types of task-irrelevant information. The first is gray background composed of natural videos from the Kinetics 400 dataset [46], following the DBC [2]] configuration (denoted as GB). The second is colorful background derived from DAVIS 2017 videos (47), adhering the background distractor settings in Distracting Control Suite 48] (denoted as CB). These benchmarks require the agent to extract task-relevant information, identify the target entity within the DMC environment, and effectively filter out background distractions. Subsequently. we evaluate these approaches in the more realistic simulated driving environment, CARLA 49). Here, the agent must extract task-relevant information from visual perception while mitigating distractions such as trees and dynamic sunlight. We also conduct experiments on the complex Atari 100K benchmark (50) in Appendix F.2 F.2 Further details on the experimental setup can be found in Appendix C
Baselines. To evaluate the effectiveness of our proposed CsDreamer framework, compare against several state-of-the-an model-based RI. algorithms known for their strong performance in high-dimenniooall observation environments. Our primary baselines include Dreamer (17). TIA (10), Iso-Dream [12] and Denoised MDP [13]. Dreamer is classic work in model- based RL but does not specifically address environments with noisy distractors. TIA targets settings with significant distractions by employing separate dynamics models for task-relevant and task-irrelevan features and utilizing masking techniques isolate essential information. Iso-Dream incorporates separate task-relevant and task-irrelevant branches within single world model to focus on aspects of the environment that the agent can influence. Denoised MDP decomposes observations into more fine-grained components, enhancing performance in complex environments with noisy distractors. Additionally, we compare our method with DBC (2], a model-free RI. method, in the DMC environment with gray natural videos background.
7

---

Mui Epiata Katas
Figure 5: Performance the Figure 6: Ablation study across different À Cheetah Run using autonomous driving simulator gray driving car as background. Left is the relationship between CARLA.
mutual information upper bound and mean episode return 500K timesteps. Right the mean episode return over varying steps.
Walteer Run GB
Cheerah Bun GB
Finger Spia GB
Hopper Hop GB
Maun Totece Beere
40-
30
Galonia V0
Cdloorent apocch
Figure Ablation Performance on DMC using gray natural videos as background.
5.2 Performance in Visual Inputs with Complex Noisy Distractors
Figure presents the performance of various methods on DMC with gray natural video backgrounds. Solid lines represent mean episode returns, while shaded areas represent the 95% confidence intervals. All methods show some performance approaches performed well on the Finger Spin task, the model-free method, DBC, lags behind. Only CsDreamer, TIA, and DenoisedMDP demonstrate substantial performance on the more complex Hopper Hop task. Notably, CsDreamer consistently outperforms all baseline methods across all four benchmark tasks. Figure presents the performance on DMC with colorful backgrounds. The results also show that CsDreamer consistently outperforms all baseline methods across the four tasks. Iso-Dream exhibits the lowest performance in every task, likely due to its difficulty in managing complex and colorful background distraction. In the Cheetah Run and Finger Spin tasks, TIA shows competitive performance, approaching that of CsDreamer. In addition, Dreamer also delivers promising results within this benchmark, which may be attributed to the characteristics of the dataset used in Distracting Control Suite (48). Specifically, it introduces background variations with multiple non-continuous images, making it easier for Dreamer to capture and disregard background-relle relltted features compared to scenarios using original videos as backgrounds, as shown in Figure 3
the Walker Run and Cheetah Run tasks. While most
We also evaluate these algorithms on the CARLA simulator, more realistic benchmark for au- tonomous driving. The results are shown Figure 5 All the methods exhibit similar performance at the first 200K timesteps. Although the interference-handlie MDP perform well in the DMC environment with noisy background distractors, they fail to achieve comparable performance in the CARLA environment. This discrepancy arises from CARLA'S sub- stantial complexity, where these methods struggle to effectively distinguish task-relevant information from irrelevant distractions in first-person views and fail extract task-relevant features from the highly intricate and volatile input observation. In contrast, CsDreamer demonstrates superior adapt- ability, outperforming all baselines and highlighting its strength in addressing more realistic and complex reinforcement learning tasks. Notably, Dreamer achieves impressive performance in the CARLA environment, slightly trailing behind CsDreamer. This strong performance is attributable to Dreamer's 's approach of encoding all environmental information into its latent space, which ensures the preservation of critical task-relevant features essential for autonomous driving. However, the inclusion of irrelevant information may slightly hinder its efficiency and robustness compared to CsDreamer.
baselines such as TIA and Denoised-
5.3 Ablation Study
The Role of Regularization. Figure illustrates the effect of varying the coefficient À in Eq. (6). As À increases, the upper bound of mutual information progressively decreases, signaling a reduced conditional dependence between task-relevant and task-irrelevant variables. Mcanwhile, the mean

---

episode return exhibits a non-monotonic trend, initially increasing and then decreasing. A small weakens the constraint imposed by the mutual information upper bound, heightening the risk of confounding task-relevant with task-irrelevant information, as previously discussed. Conversely, a large À minimizes the association between these variables. The near independence of e, and 8t, given the observation Oz, results in poorer performance, thereby validating the analysis presented in Section 4.3
Effectiveness of Different Modules. We conduct ablation studies to evaluate the contributions of the conditional dependence and decoupling regularization modules in CsDreamer. Specifically, compare the performance of CsDreamer and Dreamer against:
CsDreamer w/o CD: CsDreamer using the CsRSSM framework without conditional depen- dence for st and o given or in collider structures.
CsDreamer w/o MT: CsDreamer with the CsRSSM framework excluding decoupling regularization.
The implementation details can refer to Appendix D.3 As shown in Figure 7 CsDreamer w/o CD achieves mean episode returns comparable or higher than those of Dreamer, demonstrating the effectiveness of structural decomposition. However, it underperforms compared to CsDreamer w/o MI, highlighting the importance of conditional dependence. Introducing decoupling regularization, CsDreamer performs better than CsDreamer w/o MI. It demonstrates that the decoupling regular- ization can mitigate confusion between task-relevant and task-irrelevant variables, as mentioned in Section 4.2 CsDreamer consistently outperforms all other models across all tasks, achieving the best performance. It validates the effectiveness of CsRSSM and the efficient utilization of task-relevant information within the behavior policy of CsDreamer, ultimately enhancing overall performance. Additional ablation experiments are provided in Appendix
E.5
5.4 Visualization of The Latent Variables In this section, we reconstruct the visual input observation to interpret semantic information extracted by and G, Figure 8 presents the reconstruction visualization results for Chee- tah Run + GB. We sample trajectory and record both ground truth and reconstructed images at selected timesteps. The first row displays the ground truth of the trajectory. The second row shows reconstructions using the model p& (o, [if, st, hf, ce) from CsRSSM, which concatenates features from both pos- terior 8t and G. To isolate the semantic in- formation captured by ct, we set all feature dimensions of 8, to zero and concatenate it with posterior e; as the features for reconstruction, denoted as (o 00 0, hi, ct) in the third row. Conversely, to focus on information from St, we set features of cr to zero and reconstruct using (o( st, 0, 0), shown in the fourth row. In the fifth row, we obtain the prior & by imagining within the task-relevant transition model p (ht) using the first five posterior 5; features and then reconstruct the subsequent observations accordingly.
20
30
40
50
60
Ground Truth
5u Resorntmacion 1y Sy 16.r))
Roometration Pa 6.
Recoramettr
O, (66, a, O, 0)
Imagined es Recon s,
Figure 8: Reconstruction visualization of Cheetah Run + GB.
The second row demonstrates that combining features from 8 and e₁ facilitates effective reconstruc- tion. In the third row, the cheetah entity is nearly imperceptible, leaving only background automobile details. Conversely, the fourth row exhibits highly blurred background while the cheetah remains clear. This indicates that a effectively captures task-irrelevant information, whereas 8, captures task-relevant information. In the fifth row, the cheetah in the first 30 frames closely matches the ground truth, and the background is significantly blurred. These observations suggest that performing imagination in S space allows the model to focus on task-relevant dynamics. These results illus- trate why CsDreamer achieves outstanding performance. Additional visualization experiments are provided in Appendix E.8
Conclusion and Future Work
In this work, we adopt a generative modeling perspective and exploit the conditional dependence between task-relevant and task-irrelevant latent variables for both observation generation and environ- ment dynamics in scenarios with complex noisy distractors. By effectively separating task-relevant features from irrelevant background interference, CsDreamer enhances the agent's 's decision-makin

---

capability, Our experiments confirm its superior performance in noisy environments. We also ac- knowledge the limitations, including the assumption of Markovian, action -ndeeedddnn noise and simplified binary partition of latent variables, which may not capture the complexity of real-world scenarios. The model's capability carries some societal implications. While it can improve safety in applications like autonomous navigation and robotics, it also presents some risks. These include overreliance, where a critical safety signal could be erroneously ignored, and dual-use potential, where the technology could be repurposed for invasive surveillance. Our future work will address these challenges on two fronts. Technically, we will focus on modeling more complex non-Markovian dynamics and exploring richer, more structured latent representations. Concurrently, oritize research into enhancing model transparency to mitigate overreliance and will investigate technical safeguards, such as built-in privacy privacy-preserving preserving mechanisms, to deter misuse. We believe this dual focus on technical advancement and ethical considerations is essential for the responsible development of robust autonomous agents.
will pri-
Acknowledgments and Disclosure of Funding
We would like to thank Yucen Wang and the anonymous reviewers for their helpful discussions and support. This work was supported by the National Science Foundation of China (62476123) and the Young Scientists Fund of the National Natural Science Foundation of China (PhD Candidate) (624B200197).
References
[1] Noriaki Hirose, Dhruv Shah, Kyle Stachowicz, Ajay Sridhar, and Sergey Levine. SELFI: autonomous self-improvement with reinforcement learning for social navigation. CoRR, 2024.
[2] Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll L.. Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul F. Christiano, Jan Leike, and Ryan Lowe. Training language models to follow instructions with human feedback. In Advances in Neural Information Processing Systems, 2022.
[3] DeepSeck-Al. Deepscck-rl: Incentivizing reasoning capability in llms via reinforcement learning. CoRR, 2025.
[4] Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Mare G. Bellemare. Deepmdp: Learning continuous latent space models for representation learning. In International Conference on Machine Learning, 2019.
[5] Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In International Conference on Learning Representutions, 2021.
[6] Chuning Zhu, Max Simchowitz, Siri Gadipudi, and Abhishck Gupta. Repo: Resilient model-basee rein- forcement learning by regularizing posterior predictability. In Advances in Neural Information Processing Systems, 2023.
[7] Mohammad Ghavamzadeh, Shic Mannor, Joelle Pineau, and Aviy Tamar. Bayesian reinforcement learning: survey, CORR, 2016.
[8] Homanga Bharadhwaj. Mohammad Babaeizadeh, Dumitru Erhan, and Sergey Levine Information prioritization through empowerment in visual model-base RL. In International Conference on Learning Representations, 2022.
[9] Fei Deng, Ingook Jang, and Sungjin Ahn. Dreamerpro: Reconstruction-fret model-based reinforcement learning with prototypical representations. In International Conference on Machine Learning, 2022.
[10] Xiang Fu, Ge Yang, Pulkit Agrawal, and Tommi S. Jaakkola. Learning task informed abstractions. In International Conference on Machine Learning, 2021.
[11] Yucen Wang. Shenghua Wan, Le Gan, Shuai Feng, and De- De-Chuan Zhan. AD3: implicit action is the key for world models to distinguish the diverse visual distractors. In International Conference on Machine Learning, 2024.
10

---

[12] Minting Pan, Xiangming Zhu, Yunbo Wang, and Xiaokang Yang. Iso-dream: Isolating and leveraging noncontrollable visual dynamics in world models. In Advances in Neural Information Processing Systems, 2022.
[13] Tongzhou Wang. Simon S. Du, Antonio Torralba, Phillip Isola, Amy Zhang. and Yuandong Tian. Denoised mdps: Learning world models better than the world itself. In International Conference on Machine Learning, 2022.
[14] Yuren Liu, Biwei Huang, Zhengmao Zhu, Hong-Long Tian, Mingming Gong, Yang Yu, and Kun Zhang. Learning world models with identifiable factorization. In Advances Neural Information Processing Systems, 2023.
[15] Danijar Hafner, Timothy P. Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning, 2019.
[16] Biwei Huang, Chaochao Lu, Liu Leqi, José Miguel Hernández-Lobato, Clark Glymour, Bernhard Schölkopf, and Kun Zhang. Action-sufficient state representation learning for control with structural constraints. International Conference on Machine Learning, 2022.
[17] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy P. Lillicrap. Mastering diverse domains through world models. CoRR, 2023.
[18] Danijar Hafner. Timothy P. Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representations, 2020.
[19] Danijar Hafner, Timothy P. Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. International Conference on Learning Representations, 2021.
[20] Rishabh Agarwal, Marlos C. Machado, Pablo Samuel Castro, and Mare G. Bellemare. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. In International Conference on Leurning Reprexentations, 2021.
[21] Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. In International Conference on Learning Representations, 2021.
[22] Mete Kemertas and Tristan Aumentado-Armsst Advances in Neural Information Processing Systems, 2021.
Towards robust bisimulation metric learning. In
[23] Jianda Chen and Sinno Jialin Pan. Learning representations via a robust behavioral metric for deep reinforcement learning. In Advances in Neural Information Processing Systems, 2022.
[24] Shuo Wang, Zhihao Wu, Jinwen Wang. Xiaobo Hu, Youfang Lin, and Kai Lv. How learn domain- invariant representations for visual reinforcement learning: An information -thoreetcal perspective. In International Joint Conference on Artificial Intelligence, 2024,
[25] Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforce- ment learning with augmented data. In Advances in Neural Information Processing Systems, 2020.
[26] Nicklas Hansen, Hao Su, and Xiaolong Wang. Stabilizing deep q-learning with convnets and vision transformers under data augmentation. In Advances in Neural Information Processing Systems, 2021.
[27] Denis Yarats, Rob Fergus, Alessandro Lazaric, and Letrel Pinto. Mastering visual continuous control: Im- proved data-augmented reinforcement learning. International Conference on Learning Representations, 2022.
[28] Kevin Gmelin, Shikhar Bahl, Russell Mendonca, and Deepak Pathak. Efficient RI. via disentangled environment and agent representations. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, International Conference Machine Learning, 2023.
[29] Ameya Pore, Riccardo Muradore, and Diego Dall' Alba. DEAR: disentangled environment and agent representations for reinforcement learning without reconstruction. CoRR, 2024.
[30] Diederik P. Kingma and Max Welling, Auto-encoding variational bayes. In International Conference on Learning Representations, 2014.
11

---

[31] Irina Higgins, Loic Matthey, Arka Pal. Christopher P. Burgess, Xavier Glorot, Matthew M. Botvinick.
Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with constrained variational framework. International Conference on Learning Representationss 2017.
[32] David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In Advances in Neural Information Processing Systems, 2018.
[33] Danilo Jimenez Rezende and Shakir Mohamed. Variational inference with normalizing flows. In Interna- tional Conference on Machine Learning, 2015.
[34] Diederik P. Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1xl convolutions. In Advances in Neural Information Processing Systems, 2018.
[35] Jascha Sohl- Dickstein, Eric A. Weiss, Niru Maheswaranathan. and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. International Conference on Machine Learning, 2015.
[36] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabílistic models. In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6- 12, 2020, virtual, 2020.
[37] Michael Janner, Yilun Du, Joshua B. Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning, 2022.
[38] Mhairi Dunion, Trevor McInroe, Kevin Sebastian Luck, Josiah P. Hanna, and Stefano V. Albrecht. Temporal disentanglement of representations for improved generalisation in reinforcement learning. In International Conference on Learning Representations, 2023.
[39] Mhairi Dunion, Trevor McInroe, Kevin Schastian Luck, Josiah Hanna, and Stefano V. Albrecht. Conditional mutual information for disentangled representations in reinforcement learning. In Advances in Neurul Information Procensing Systems. 2023.
[40] Amy Zhang, Clare Lyle, Shagun Sodhani, Angelos Filos, Marta Kwiatkowska, Joelle Pineau, Yarin Gal, and Doina Precup. Invariant causal prediction for block mdps. In International Conference on Machine Learning. 2020.
[41] Andrew Bennett, Nathan Kallus, Lihong Li, and Ali Mousavi. Off-policy evaluation in infinite-horizon reinforcement learning with latent confounders. In International Conference on Artificial Intelligence and Statistics. PMLR, 2021.
[42] Gokul Swamy, Sanjiban Choudhury, Drew Bagnell, and Steven Wu. Causal imitation learning under temporally correlated noise. In International Conference on Machine Learning, 2022.
[43] Lars Buesing, Theophane Weber, Yori Zwols, Nicolas Heess, Sébastien Racanière, Arthur Guez, and Jean- Baptiste Lespiau. Woulda, coulda, shoulda: Counterfactually-guee guidddd policy search, In International Conference on Learning Representations, 2019.
[44] Haiyao Cao, Zhen Zhang, Panpan Cai, Yuhang Liu, Jinan Zou, Ehsan Abhasnejad, Biwei Huang, Mingming Gong, Anton van den Hengel, and Javen Qinfeng Shi. Rethinking state disentanglement in causal reinforcement learninn CoRR, 2024.
[45] Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li. Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, Timothy P. Lillicrap, and Martin A. Riedmiller. Deepmind control suite. CoRR, 2018
[46] Will Kay, João Carreira, Karen Simonyan, Brian Zhang, Chloe Hillier, Sudheendra Vijayanaraimmhan, Fabio Viola, Tim Green Trevor Back, Paul Natsev, Mustafa Suleyman, and Andrew Zisserman. The kinetics human action video dataset. CoRR, 2017.
[47] Jordi Pont-Tuset, Federico Perazzi. Sergi Caelles, Pablo Arbeláez, Alexander Sorkine Morkung, and Luc Van Gool. The 2017 DAVIS challenge on video object segmentation. CoRR, 2017.
[48] Austin Stone, Oscar Ramirez, Kurt Konolige, and Rico Jonschkowski. The distracting control suite- challenging benchmark for reinforeement learning from pixels. CoRR, 2021.
[49] Alexey Dosovitskiy, Germán Ros, Felipe Codevilla, Antonio M. López, and Vladlen Koltun. CARLA: an open urban driving simulator. In Annual Conference on Robot Learning, 2017.
12

---

[50] Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H. Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohiuddin, Ryan Sepassi, George Tucker, and Henryk Michalewski. Model based reinforcement learning atari. In International Conference on Learning Representations, 2020.
[51] Chang Chen, Yi-Fu Wu, Jaesik Yoon, and Sungjin Ahn. Transdreamer: Reinforcement learning with transformer world models. CoRR, 2022.
13

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and introduction describe that we propose to leverage conditional dependence to implement efficient world model denoising.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We mainly describe the limitations in Section 6
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
14

---

Justification:
We provide full set of assumptions for each theoretical result in the paper.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide all the information to reproduce the results in the paper, and we also provide the code in the supplementary materials.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
15

---

Answer: [Yes]
Justification: We provide the code the supplementary materials, and give the instructions the code in the README.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips... cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to to reproduce the results. See the NeurIPS code and data submission guidelines (https
wttp://iiggii/////ee/eee wtt.gguddde/eeeSSuunrrnn ssionPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We give all the training and test details the experiment part and the related appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[Yes]
Justification: We use the confidence intervals and the error bars in our experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
16

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide the information in Appendix D.4
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification:
The research conforms with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss the potential positive and negative societal impacts in Section 6
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no socictal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
17

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper poses such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We cite them explicitly.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
18

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: This paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
19

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality the research, declaration is not required.
Answer: [NA]
Justification: The core method development in this research does not involve LLMs any important, original, or non-standard components.
Guidelines:
The answer NA
that the method development in this research does involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy for what should or should be described.
20

---

A Detailed Derivations
A.1 Derivation of ELBO Loss in CsRSSM
We use 01,7, 41:7, F1t and Fi.T for observation sequence, action sequence, reward sequence and episode continuation flag sequence. Similarly, also $1.7 and C1:r for task-relevant and task-irrelevant latent variables. After introducing the variational posterior (88:2, cu.r (o:r, (1.r) = obtain the variational bound for the dynamics models p((O1.T T1:T), 81:7, C1:7(117) = 11, p(s: 81-1, az- (ppppp/pp-)pppppppppp e (s)(rs)))) CsRSSM using Jensen's inequality
H. g(cloog, a<x) (<)) [c, Ogr: aet),
In Inp((1i Po:r, (err, (a.r)
pl p(OoT, T1:T 91 a1T)
=ln B arto
((81.T, Curlorr (1.)
pl (i
2B careora
((115,
[1, -1)),, 4 (oo on, e (p(rr))))))())))
oe. (c) (g osa, aci)
[mp((1/(-_1, (ntaaai))nnnlllleeeeee Implc hacl:/e-i-)nnnplllllllaap Inp(o.)s)))) + Imp(ra(s)) + Imp(aals))
aeg -I -ng(sil
(7)
II derlooe (ce)s, ogr, aze) [inplo]ss Imp(r)(s)) Inp((((s))
+ (Inp(s)l:-1n,
lngei))nnn((((ggnnnnooo (Inp(clle-) Ing(cllos,, (Inpalll dayrderr
T
Em [ln a)) + E., (1r Es, [Imp((r(s)) -E
E c
KKL (sele, oga, a a (1(())11-11
By applying the history models in Eq. ( we can ultimately obtain the CELBO in Eq.
A.2 Derivation of the upper bound for conditional mutual information
By incorporating the variational approximate distribution ((se//f h2, oz). the original conditional mutual information is given by:
In p(s, celhi, o.)
(185 M2 0z) =Eag,s
p(s hf, or (.) hi oz)
In p(sele, hi, hi, oz)g of) o.)
p( hi; (1 (x)q())) 16 hū, o.)
p(salce. hi, oz)
In p(s by. hi, oz)
((selhE, hz.
g he, hi, o.)
(8)
In p(selce. hf hf,
In p(s hhf oz)
E Eo
a( 1b, o.) Eo,,, JKL (KLppssele. by, hū, ( Mi, or)]] - EE, (KLpppspssssp
( hi; oz)
14,
<Eo, [L[ppss hg. hi (s: 16. or)]]-
Then obtain the upper bound regularization loss in Eq. (
21

---

B Pseudo Code
The whole algorithm is shown in Algorithm For brevity, we omit the episode continuation flag st-
Algorithm 1 CsDreamer
Initialize: Dataset collected by random policy, the policy parameters A, the critic parameters yA the variational estimator parameters and the parameters & in the CsRSSM world model for training step 1 = I...f do
for update step /2 = 1...72 do
// Dynami cs learning
Sample minibatch {(o,, at rr) rT) k+L D
Compute loss according to Eq. (6)
Update CsRSSM world model parameters & and the variational estimator parameters ę // Behavior learning
Infer the task-irrelevant information сt ~ (cilhh, hç, Infer the task-relevant information ( h4, 01,
Imagine trajectories {(s, at, ++)+ using (5( (1f),, Po(r hhg s.) and the policy x
Update the policy parameters and the critic parameters v using the imagined trajectories end for for rollout step t3 = 1...T3 do
Infer the task-irrelevant information ct ~ qo (o(celhf hi, oz) Infer the task-relevant information ~ qo (s(sellf, f lif Ct, 0t) Sample action from exploration policy at Texp (ee hi, Tt. O++1 env (a;)
end for
Add experience to dataset D + DU (Or,, end for
C Details about the Benchmark
(c)
(d)
Figure 9: Example observations of the benchmarks. (a) The observations in DMC with gray videos in the background and the four tasks from left to right are Walker Run, Cheetah Run, Finger Spin and Hopper Hop. (b) The observations in DMC with colourful videos in the background and the four tasks from left to right are Walker Run, Cheetah Run, Finger Spin and Hopper Hop. (c) The observation in CARLA simulator. (d) The observation in the game Alien, which is one of the 26 games in the Atari 100K benchmark.
DMC with Complex Backgrounds. We introduce the natural video background as the noisy distractor for the widely used robotic locomotion benchmark, DeepMind Control Suite [45). We select four tasks for evaluation. The Walker Run task assesses stability and coordination of bipedal locomotion, Cheetah Run examines high-speed movement and agility, Finger Spin evaluates
22

---

precise motor control and object manipulation capabilities, and Hopper Hop tests dynamic balance and energy management in single-leg hopping. We choose Walker Run and Hopper Hop instead of Walker Walk and Hopper Stand for more challenging evaluation. For each task, we adopt two types of noisy distractors. We first introduce the natural video from the Kinetics dataset 466 We only use the videos in 'riiiing car' class and set it to gray as that in 21 [0] (Figure (a)). Then we utilize the colourful videos of the DAVIS 2017 dataset [47] as that in
[48). We utilize all the 90 train- val videos and adopt the dynamic setting, where the video plays forwards backwards until the last or first frame is reached at which point the playing direction is reversed, thereby the background motion is always smooth and without 'cuts' (Figure (b)). The height and the width of the input observation image are 64 x 64.
CARLA. To evaluate on a more real-world control system, we use the CARLA simulator [49] to conduct photo-realistic visual observations. The agent's goal is to drive as far as possible along CARLA's Town04's highway in 1000 timesteps in a first-person way without colliding with other moving vehicles or barriers as the setting in DBC 21 To increase the task difficulty, we make two modifications as that in Iso-Dream setting [2]. We use one camera which obtains images of 64 64 pixels instead of five to limit the field of view and we include 30 other moving vehicles or obstacles instead of 20 to increase the likelihood of collisions. The example observation is shown in Figure (c).
Atari 100K Benchmark. The Atari 100K benchmark [50] comprises 26 distinct Atari games. Within this benchmark, an agent permitted 100K interaction steps for each game environment. Due to a frameskip setting of 4, this translates to 400K frames. This interaction volume is roughly equivalent to about two hours of real-time gameplay. The game environments within this benchmark notably complex, making it a practical testbed for evaluating the algorithms' robustness and data efficiency. We keep all implementation details the same as Dreamer [17,, and the example observation is shown in Figure (d).
D Implementation Details
D.1 Base Method
CsDreamer is implemented based on the classic model-based reinforcement learning method,
Dreamer I8
19
17 Given that Dreamer has multiple versions, we first evaluate DreamerV2 19
and DreamerV3 17 on DMC using gray natural videos as background.
The performance results are in Figure 10 both DreamerV2 and DreamerV3 demonstrate overall similar performance across the four tasks. Specifically, DreamerV3 shows slight advantage in the Finger Spin GB tasks, while DreamerV2 performs marginally better in the Walker Run GB and Hopper Hop + GB tasks. They have comparable performance in the Cheetah Run + GB task. Overall, the performance differences between the two versions are minimal. The minimal performance differences between DreamerV2 and DreamerV3 can likely be attributed to several factors. Firstly, both models share a significant degree of architectural and methodological overlap. Secondly, some of the improvements in DreamerV3 are primarily designed to enhance adaptability across various domains and tasks. However, these enhancements may compromise the model's fundamental capabilities. This trade-off means that while DreamerV3 can perform effectively in broader range of scenarios, it might not significantly outperform DreamerV2, resulting in similar performance across the evaluated tasks. Considering that DreamerV3 has much smaller world model loss (on the order of tens) than DreamerV2 (over ten thousand) and in our experiments LMI is typically within single digits, have a similar order of magnitude to avoid problems. Unless otherwise specified, in the experiments, Dreamer refers to DreamerV3, and CsDreamer is based on DreamerV3.
DreamerV3 as the base method so that LCaRSSM and LMI
In this paper, utilize the 'S' size model for DreamerV3 in https:/wwxiiiongg//rr/000.1111/1
23

---

Walker Run 6B
Cheerah Ruc. GB
Finger 6pon + GB
Hopper Hop + GB
Tinane
Main Tohoc Betre
00
7:00
1000 2013 3000
$35.
2060 2010
5000
100.0
2010
305.1
2000
50003
Figure 10: Performance of DreamerV2 and DreamerV3 on DMC using the gray natural videos as background.
D.2 Key Component Implementation
The RSSM world model in Dreamer implements the KL divergence as two separate components in different stop-gradient operator places and loss scales (17). Similarly, the two KL divergences in Eq. (4) are each composed of two distinct loss terms. Specifically, we have
[[KL(as(eelh hi (n))lllllc Pay (annfyl))) + Prep (O) EKLL (a(sslhe hi, et ox) (x)llo(sr)1(]]]]]] Bayn Layu (0) + Ceep (1))
(9)
where
Ldyn (6) = max (1, KT. [sg (go (e hū, o.)) Il [k))
Ciep (Ø) = max (1, KT. (co 1y Ni, ot) || sg (p& (ce (b: [)])))
(10)
Layn (Ø) max (1, KT. [sg (1o ( (4g. hy Ct ot)) IP (st (n))
D (Ø) max (1, KT. (Se 16, Ct o:) sg (s [+])])
and sg(+) denotes the stop-gradient operator. In the CsRSSM world model, the task-relevant com- ponent utilizes the same network architecture as Dreamer. The primary distinction arises during the computation of the observation posterior, where both the observation o, and the inferred latent variable e concatenated and provided input, rather than using only the observation o, in Dreamer. Additionally, when reconstructing the observation, the model concatenates the inferred a. with the variable 8; reconstruct o, jointly. The task-irrelevant component employs network structure similar to that of Dreamer but with some modifications. It does not take actions as input, ensuring that it remains unaffected by action-related information. Moreover, it does not predict task-relevant variables such as rewards, thereby focusing solely on modeling the noise distractors independent of the task-relevant objectives.
For the regularization loss LMI. we use a variational estimator q for the variational approximate distribution g( [4. hū oz). The variational estimator takes the history (yÂ. h2), and the embedding of o as the input. The gradients from these variables will be blocked from backpropagation during the training of the variational estimator, i.e., 95 (sxl sg(h;, hit o,)). This ensures that the learning of the variational estimator does not influence the parameters of the preceding models. Then, we utilize the part network of o, get the embedding of (hς, ot), and we the embedding to get the variational approximate distribution after a linear layer.
D.3 Module Ablation Implementation
Here, we detail the implementation of the two ablation methods described in Section 5.3 CsDreamer w/o MI implemented by simply removing the decoupling regularization from CsDrmamer. Cs- Dreamer w/o CD adopts the CsRSSM framework with some modifications. Specifically, since CsDreamer w/o CD does not account for the conditional dependence between 8; and 6 given o, in collider structures, we adjust its encoders: the task-irrelevant encoder is defined go(ca [if oz) and the relk-ntle encoder defined as go ot). The world model is then trained using LELBO- Because CsDreamer w/o CD ignores all conditional dependence, the LMI term is omitted.
24

---

D.4 Hyperparameters and Time Cost
Table I preeeets the primary hyperparameters of CsDreamer. Since the behavior policy relies solely on the feature of st, and the hyperparameters for 9; in CsRSSM closely resemble those of the latent variables in RSSM of DreamerV3, adopt the same hyperparameters for the behavior policy as in DreamerV3. The experiments are mainly conducted on NVIDIA RTX 4090 GPUs. With each GPU, we are able to train each environment rate of approximately 24K timesteps per hour.
Table 1: Hyperparameters for CsDreamer
Hyperparameter
Value
Action Repeat
4 for CARLA and Atari, and for others 10.0 for Hopper Hop+GB, 0.2 for CARLA and Atari, and 1.0 for others
A
0.5
Siep
0.1
P
0.5
0.1
Discrete latent dimensions of 8:
Discrete latent classes of
32
Discrete latent dimensions of ot
16 for CARLA, for Atari, and 32 for others
Discrete latent classes of ct GRU recurrent units of st GRU recurrent units of c Dense hidden units of 8; Dense hidden units of MLP layers
16 for CARLA, 8 for Atari, and 32 for others
512
256 for CARLA, 128 for Atari, and 512 for others
512
512
E Additional Experiment Results
E.1
Performance Scores
Table we summarize the final mean episode return and their corresponding standard deviations for various model-based RL methods across multiple environments in the main text, evaluated using four distinct random seeds, each associated with 10 evaluation episodes. The results consistently show that CsDreamer outperforms the other methods in the majority of environments, thereby demonstrating the superior effectiveness our proposed approach.
Table Final performance across model-based RL methods in different environments.
Environment
CsDreamer (Ours)
Dreamer
TIA
Denoised MDP
Iso-Dream
Walker Run GB
533 + 98
162 + 32
293 + 129
117 +
131 + 23
Cheetah Run GB
547 + 159
171 1 84
432 + 172
215 + 191
109 + 12
Finger Spin + GB
408 45
287 54
123 + 201
331 146
415 118
Hopper Hop + GB
75 58
0=0
50 48
13
0
Walker Run CB
678
41
588 135
88
211 95
Cheetah Run CB
821 58
549 126
758 - 80
392 195
6 4
Finger Spin + CB
576 241
466 = 130
490 - 136
460 123
4
Hopper Hop + CB
70 + 52
23 + 28
42 4:
21 19
0
CARLA
235 + 118
202 + 144
44
21
24 + 25
E.2 Performance Atari 100K Benchmark
In order to assess the denoising capabilities of CsDreamer under highly complex visual inputs, we conduct experiments on the Atari 100K benchmark, and the results are presented in Figure T Table For each game, we use at least five seeds. During our experiments, we encounter a random seed that produces performance approximately 20 times higher than the current baseline, which
25

---

EJM 2M0 6050
BAXK CD
XUL
6.12
20p1 30
2015
-0L8
saga
202
1002 2003 2002
1000
20pe 190
2018
1003
Sapa
Japs
Sirpe
Stye
Bick 1D
zagt SC
Zapa
Toull
BIX 60
20g1 80
saga
pagt
Up
Dpp0 with Guay Sond
Dowe Sead
Ganar
S4pp 60
JKC L0
Japt 90
sage
Sagu190 (0)
Figure 11: CsDreamer performance on Atari 100K Benchmark. In the Up N Down environment, we encounter a random seed that yields performance approximately 20 times higher than the current baseline, which we call the "crazy seed'. We use seeds this environment to ensure a more accurate evaluation. The Up N Down results reflect the performance of the 9 seeds, excluding the crazy seed, while Up Down with Crazy Seed presents the outcomes incorporating 10 seeds.
call the 'crazy seed'. To facilitate more robust and accurate evaluation, we employ total of 10 seeds this environment.
In Figure I1 the results labeled Up N Down reflect the performance obtained from nine seeds (excluding the crazy seed), whereas Up Down with Cruzy Seed encompasses outcomes derived from all 10 seeds. In Table values within parentheses indicate results achieved using the crazy seed, while results outside the parentheses correspond to evaluations without its inclusion. The Dreamer results in Table are sourced directly from the official DreamerV3 paper [17] The experimental outcomes clearly indicate that CsDreamer outperforms the bascline approach. This improvement demonstrates that leveraging conditional dependencies enables CsDreamer to effectively denoise complex visual inputs, thereby enhancing overall task performance.
E.3 Performance on Standard DMC
We conduct experiments using the standard DeepMind Control (DMC) suite evaluate performance on clean benchmarks. Figure 12 illustrates performance across four tasks. As depicted, CsDreamer generally achieves performance on par with the original Dreamer algorithm across all four evaluated tasks. These results indicate that the CsDreamer algorithm effectively handles noisy scenarios and maintains strong performance in noise-free environments.
In this paper, utilize the data in https://ani..oo///drr20....4
26

---

Table 3: CsDreamer's performance on the Atari 100K Benchmark. Values in parentheses indicate the
use of the "crazy seed'.
Game
Random
Human
Dreamer (official)
CsDreamer
Alien
228
7128
959
888 + 234
Amidar
6
1720
139
184 + 43
Assault
222
742
706
748 + 262
Asterix
210
8503
932
1114 + 266
Bank Heist
14
753
649
946 + 465
Battle Zone
2360
37188
12250
3960 + 2537
Boxing
0
12
78
80 = 13
Breakout
2
30
31
31 51
Chopper Com.
811
7388
420
864 + 480
Crazy Climber
10780
35829
97190
90196 + 22281
Demon Attack
152
1971
303
207 + 124
Freeway
0
30
0
0
Frostbite
65
4335
909
2177 + 995
Gopher
258
2412
3730
7771 + 18563
Hero
1027
30826
11161
8124 + 3345
James Bond
29
303
445
292 2 147
Kangaroo
52
3035
4098
6590 + 4797
Krull
1598
2666
7782
9636 - 3531
Kung Fu Master
258
22736
21420
24847 + 8928
Ms Pacman
307
6952
1327
2170 + 1306
Pong
-21
15
18
20 + 1 977 1813
Private Eye
25
69571
882
Qbert
164
13455
3405
1050 + 692
Road Runner
12
7845
15565
11980 + 3946
Seaquest
68
42055
618
578 + 252
Up N Down
533
11693
7667
9800 13122 32205 13122 69863)
Human Mean
0%
100%
112%
129% (136%)
Human Median
0%
100%
49%
66% (73%)
Walker
Chertsh Run
Plager Spir
Hopper Hop
2:30
Vien Dlace beemm
k0-
Voen Ditekk
200
2000
200.
2003
500.0
250.0
1000 3000 5560
Figure 12: Performance on Standard DMC.
27

---

E.4 Performance on DMC + GB Using Train-Eval Split
In the experiments above on DMC GB, we follow previous work by training and testing using the same dataset as background. Specifically, we utilize 16 car videos beginning with the letter 'A' from the Kinetics-400 dataset, mentioned in TIA [TO] We download the complete' driving car' class using the GitHub repository ascertain whether CsDreamer's 's robustness is a result of merely memorizing background information or if it genuinely possesses an intrinsic denoising capability. Due to the reasons mentioned in the repository, we succeed in obtaining 641 videos. Then we alphabetically split the dataset into a training dataset (512 videos) and an evaluation dataset (129 videos) using an 80 20 ratio. We choose the best-performing baseline TIA in DMC + GB the baseline, and the experimental results are presented in Figure 13 Compared to the results in DMC + GB above, the outcomes for Walker Run, Cheetah Run, and Finger Spin are slightly inferior, while the outcome for Hopper Hop is somewhat superior. The results depicted in the figure demonstrate that CsDreamer outperforms TIA across all tasks. This confirms that our approach effectively denoises rather than merely memorizing backgrounds.
Trsin/kkll
Chanrah
Train/'bl
Finge Spmn
Hupper Hop
130
2:30
240
300.
2000
1000
Stnge (K)
Sirge O9)
Figure 13: Performance the train-eval dataset, with 641 driving car videos in total, 512 videos for the training dataset, and 129 videos for the evaluation dataset.
E,5 Additional Ablation Results
We conduct more module ablation studies the DMC benchmark using colorful natural videos as the background. Similar to that in Section [53]] we compare the performance of the baseline Dreamer, CsDreamer w/o MI (CsDreamer with CsRSSM framework without the conditional mutual information-basee regularization), CsDreamer w/o CD (CsDreamer with CsRSSM framework without the conditional dependence) and CsDreamer. Shown in Figure 14] the results are consistent with those in Section 5.3 The comparison between Dreamer and CsDreamer w/o MI reveals that the CsRSSM can significantly boost performance visual input with complex distractors by utilizing the conditional dependence. By introducing the decoupling regularization, CsDreamer consistently outperforms baselines on all tasks, achieving the highest performance. This ablation study also demonstrates that CsRSSM with the decoupling information regularization significantly enhances learning efficiency and task performance.
Waller Run CB
Cheetah Ran CB
Finger Spis CB
Hopper Hup CBl
k0-
Mox
2000
1260
Gallrouncc w/t Mi
Drtaner
Figure 14: Ablation Performance on DMC using colorful background.
We also examine the role of regularization in the Walker Run + GB in Section 5.3 in Figure [5] the results echo those from Section 5.3 as the regularization parameter À increases,
As illustrated
We use the Github repository in http://gghhhccoooSStoooo//eteeee
28

---

the upper bound of mutual information progressively decreases, while the mean return follows mon-monot trend initially rising and then falling. All these experiments validate the analysis presented in Section
Mean Egisade
300
2=0.11
1=1.0
Monn Rppovdee kenur
4001
2-50.0
1-100.0
303
300
2=1000.0
203
200
130
0.01
0.1
1.0
10.0
50.0
100.0
1000.0
100
300
500
Scale
Stepe(K)
Figure 15: Ablation study across different À on Walker Run using gray driving car as background.
E.6
Correlation analysis
Whlkke: Run :GB
Cheersh Bun GB
Whlker Run -CB
Cheersh Run CB
0.75
6.75
6.75
0.75
0/0
600
6:0
1uan Epbock Saun
Mtan Ipback Sauur
Siea Eptrado Sartum
Misa Eptidd Raturn
0.0 300.0
2000 2000
4060 5000
0.0
305.0
2000 3000
406n 5000
0.0
105.0
305.0
300.0
406n 5000
100.0
302.0
300.0 400n 5000
Stegi (0)
Stege 60
Stepe 60
Stege 60
Figure 16: Coefficient of determination (1t) on DMC tasks with background distractors.
In this section, we assess the coefficient of determination (R²) within our framework. We find that the distribution of the test data significantly influences the R² values. Specifically, evaluating R² using the data sampled from the current policy may result in artificially high values. To obtain more reliable assessment, we use the sampled data from the final phase prior training session as dataset because, this stage, the policy has sufficiently learned perform the task effectively, leading to the observation continuously undergoing noticeable changes. In detail, we utilize the last 20 sampled episodes from a previous training process, with each episode comprising 1,000 timesteps (derived from 500 data points with an action repeat factor of 2). Each data point has one observation and one corresponding agent state. We utilize the observation these episodes to yield the features for variables 8, and er. We use the feature of 8, and the feature of e as the independent variables separately predict the agent's real state. To mitigate potential interference from newly introduced training parameters, we utilize a parameter-free k-nearest neighbors (KNN) regressor for prediction. The results are shown in Figure 16. The R² value for s; rapidly increases in the early stages and stabilizes, demonstrating a strong and consistent correlation. In contrast, the R² value for decreases and remains low throughout the experiments. The results highlight that s; exhibits much stronger correlation to the real agent state than c, suggesting that s, can effectively capture the task-relevant information.
E.7 Additional Results Using Transformer-Based Framework
To clarify whether the proposed collider-structurg approach restricted to RSSM-based designs or if it is compatible with transformer-syyle sequence encoders, we extended our methodology to TransDreamer [5T]. TransDreamer utilizes Transformer State-Space Model (TSSM) to capture long- term dependencies, replacing the RSSM of Dreamer. We integrate our approach into this baseline, resulting the Collider-structun Transformer State-Space Model (CsTSSM) and the corresponding Cs STTansDrramer A primary adaptation challenge stems from TransDreames architecture, which
29

---

is optimized for parallel training. To achieve this, TSSM employs Myopic Representation Model, where the posterior inference g(2, (o,) is independent of the deterministic history state h, CsTSSM adapts the collider framework within this constraint.
The key modifications in CsTSSM
follows:
Dual Transformer Dynamies, where we replace the single dynamic model with two independent Transformer networks. One models the task-relevant dynamics, generating hy from past relevant states and actions; the other models the task-irrelevant dynamics, generating h; solely from past task-irrelevant states.
Parallel Collider Inference, where
implement the collider-structure inference Pa (Ct, 8t ox) Ps (c, o²) Pa (s, Ct 0.). Crucially, because this inference path does not depend on the history states, the entire sequence can encoded in parallel across the time dimension before the Transformer dynamics are computed, preserving the efficiency of the architecture.
Decoupling Regularization, where adapt the decoupling regularization loss to this archi- tecture. We introduce the variational estimator q5 (&t 0s) to approximate the relevvel distribution without conditioning Ct, thereby balancing the conditional dependence within the parallel framework.
Since TransDreamer open-sources its code for the Atari benchmark, conduct experiments on the first three environments of the Atari 100K benchmark (shown in Figure The results demonstrate that the benefits of leveraging conditional dependence via the collider structure are architecture- agnostic. The proposed methodology is not limited to RSSM-based designs and can effectively enhance transformer-basee world models in environments with complex noise interference.
Alien
Amidar
Assaul
500
DditansUreamer
DelransUreamer
DolransDreamer
9.0
TransDreumerr
TransDreamter
TransDreatter
600
Mezn Eguoods Bertm
Moan Bptoode Detum
Mean Equode Betum
0.30
900
600
200
300
100-
100.0
Stegs
Stegs
Stegs
Figure 17: Transformer-bute based results on three Atari environments.
E.8 Additional Reconstruction Visualization
We extend the reconstruction visualization additional environments following the configurations presented in Section 5.4, yielding consistent results (Figure 18(25). Furthermore, in the complex autonomous driving scenario within CARLA (Figure 18). find that the reconstruction of task- irrelevant variables excluded nearby blue and red vehicles (as shown in the third row). In contrast, the reconstruction of task-relevant variables retains these vehicles, as their proximity influences the decision-making of the autonomous vehicle. These experiments the interpret semantic information extracted by the latent variables and qualitatively demonstrate that our method can effectively ex- tract task-relevant information from noisy observations in complex environments. We also conduct the reconstruction visualization for the evin-eval-split experiments in Appendix [E.4] In each fig- ure (Figure 26 29) the fourth and fifth rows exhibit heavily blurred background while maintaining clear task-relevant details. These visualizations demonstrate that CsDreamer effectively learns to extract task-relevant information from noisy observations rather than memorizing the background information, thereby accounting for its superior performance.
30

---

18
20
4 Recorstmation Pg (...000
Imagined % Px
Figure 18: Reconstruction visualization CARLA
1
10
20
30
40
50
60
Ground Truth
Se Reconstruction Po (Gellff c)
Reconstruction Po(ā,00, 0,
9; Reconstmation s,, 0,0)
Imagined s, Recon Pg (6 0,
Figure 19: Reconstruction visualization of Walker Run using gray videos as background.
1
5
10
20
30
40
50
60
Ground Truth
5,5, Reconstruction (6.
Reconstmuction 0,, c.)
5t Reconstruction Po (b, 0,
Imagined 5, Recon Pe (p 0,
Figure 20: Reconstruction visualization of Finger Spin using gray videos as background.
31

---

10
20
30
40
50
Ground Truth
Sr Reconstruction Pa (o,lhf cg)
Cr Recorrstructin P4[0, 0, ct)
Reconstruction Pa sq. 0,0)
Imagined S Recon Pe Sp 0,0)
Figure 21: Reconstruction visualization of Hopper Hop using gray videos as background.
5
10
20
30
40
50
60
Ground Truth
5, Reconstruction (6 (o(6,1h h(,
c Reconstrution Ps (, 4)
Reconstraction Po 56. 0)
Imagined Recon Po (o
Figure 22: Reconstruction visualization of Cheetah Run using colorful videos as background.
5
10
20
30
40
50
60
Ground Truth
q. Reconstruction Pe (E
c Reconstruction 0,,10,0,45
Sy Reconstruction
Pa s,
Imagined s, Recon Po (6, 0)
Figure 23: Reconstruction visualization of Walker Run using colorful videos as background.
32

---

5
10
20
30
40
50
60
Ground Truth
Su Reconstruction P4 (6, Ky. c,)
c Reconstruction (6 10,10,000 0, /f.
8 Reconstruction Ps St, 0,
Imagined Po s4,
Figure 24: Reconstruction visualization of Finger Spin using colorful videos background.
20
30
50
Ground Truth
Ce Reconstruction
Pa (6, AG, G)
Reconstruction P& c))
S; Reconstruction Pe (6.1 0, 0)
Imagined St Recon (,,11,0,,,00000
Figure 25: Reconstruction visualization of Hopper Hop using colorful videos as background.
10
20
30
50
60
Ground Truth
5 Reconstruction Pa (o,1h, 48.
c Reconstruction Pe (6, N, 4
Reconstruction Po (Q.lhE 0)
Imagined Recon Pe (o, 0,0)
Figure 26: Reconstruction visualization of Cheetah Run using train-eval gray videos as background.
33

---

5
10
20
30
40
50
Ground Truth
Su Reconstruction Pp st. hg Ce)
o Reconstruction (6, 0,0, ce)
9; Reconstruction ( [h, St 0,0)
Imagined Recon Po(6elhE.s,.00g
Figure 27: Reconstruction visualization of Walker Run using train-eval gray videos as background.
10
20
30
40
50
60
Ground Truth
G. St Reconstruction
Po (6,1 c))
Reconstruction Ps (6: 0,,
9; Reconstruction Pe (6, 0,0)
Imagined & Recon Pq (h, 0,
Figure 28: Reconstruction visualization of Finger Spin using train-eval gray videos as background.
5
10
20
30
40
50
60
Ground Truth
G, St Reconstruction
Reconstruction (6, c.)
5, Reconstruction Pg (6, (hg 0,
Imagined 5, Recon Po 0)
Figure 29: Reconstruction visualization of Hopper Hop using train-eval gray videos as background.
34


Reviews:
[{"review_id": "IoXrktlJTz", "reviewer": "Reviewer_LU9u", "summary": "This work introduces CsDreamer, a model-based reinforcement learning agent designed for environments with visual distractors. Its world model, CsRSSM, reframes the denoising problem by explicitly modeling the conditional dependence between task-relevant state and irrelevant noise, which it posits forms a \"collider\" structure. By leveraging this dependence in a sequential inference process, the work aims to achieve more robust and efficient state estimation compared to methods that assume conditional independence.", "questions": "Q1. The proposed CsRSSM architecture is more complex than its predecessors. Could the submission provide a quantitative comparison of its computational cost (e.g., wall-clock training time or steps per second) against the main baselines like Dreamer and TIA?\n\nQ2. The analysis of the critical hyperparameter λ is presented for a single environment in Figure 6. How sensitive is the final performance to this parameter across the other evaluation domains, and are there any general heuristics for setting it?\n\nQ3. Assumption 4.2 posits that the noise process is Markovian. How would the proposed method be expected to perform if this assumption were violated—for instance, in environments with structured, non-Markovian distractors that have long-term temporal correlations?", "limitations": "The \"Conclusion and Future Work\" section is brief and does not sufficiently address the potential limitations or negative societal impacts of this research. A work that significantly improves an agent's robustness yo visual distractors has clear dual-use potential (e.g., in surveillance). A more thorough discussion of these implications is expected and **necessary** for a submission of this quality. If the paper does not include a significant Conclusions, Future Work, and Limitations that isn't simply taken as a formality, it is my duty to reverse my rating into a clear Rejection.", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "This is a strong submission presenting a novel and well-reasoned contribution to representation learning in model-based RL, supported by convincing empirical results. The work is technically sound and clearly presented.\n\n**Strengths:**\n\n**Principled Methodological Advance:** The core contribution is the conceptual reframing of the denoising problem. Moving beyond the standard assumption of conditional independence to explicitly model the collider structure is a non-trivial, prinicpled insight that better reflects the underlying generative process. This represents a significant and elegant theoretical advance over prior art.\n\n**Strong Empirical Results: **The proposed agent, CsDreamer, demonstrates consistent and substantial outperformance against a strong set of relevant baselines (Dreamer, TIA, Iso-Dream) across multiple challenging visual control domains, including DMC with dynamic backgrounds and the CARLA simulator .\n\n**Effective Ablation and Visualization.** The main body provides clear evidence of the method's efficacy. The ablation study in Figure 7 convincingly demonstrates the importance of both the conditional dependence mechanism and the mutual information regularization . Furthermore, the disentanglement visualizations in Figure 8 provide strong qualitative support for the claimed mechanism.\n\n**Weaknesses:**\n\n**Lack of Computational Overhead Analysis.** The paper introduces a more complex architecture with dual history streams and an additional variational estimator but provides no analysis of the computational overhead. A comparison of training and inference times against the baselines is a critical missing detail needed to evaluate the practical trade- offs of this approach.\n\n**Limited Generality of Hyperparameter Analyses.** The analysis of the key hyperparameter λ is limited to a single task (Cheetah Run) in Figure 6. Without a broader study, it is difficult to assess how sensitive the method is to this parameter and how much per-task tuning is required, which may limit practical applicability., and limits confidence in this from this work alone.", "was_revised": true, "final_justification": "Authors have address all crucial points satisfactorily to maintain my accept score, especially regarding limitations, societal impact, and future work. I wish the Authors the best, and I am happy to help them improve their work.", "rebuttal": "# Response to Reviewer LU9u\n\nThank you for the thorough and insightful review. We address the comments point-by-point as follows.\n\n\n**Q1**: Lack of Computational Overhead Analysis.\n\n**A1**: Thank you for the valuable comment. A detailed performance analysis comparing CsDreamer to DreamerV3 is conducted on the Cheetah-Run+GB task using an AMD 7950X CPU and an NVIDIA RTX 4090 GPU, with quantitative results presented in the table below. The data confirms that our model has a larger parameter count (+74.1%) and exhibits a higher single-step inference latency (+39.86%), leading to a longer total training time (+48.91%). The critical imagination rollout speed, which dictates policy optimization efficiency, remains nearly identical (+0.21%). This is achieved by performing imagination exclusively in the compact, task-relevant state space $s_t$. And to ensure a fair comparison, we match $s_t$ to the dimensionality and complexity of the latent space in DreamerV3. Furthermore, we employ an actor-critic network identical to that of DreamerV3, ensuring that the superior final task performance is attributable to the quality of the learned representations rather than an increase in policy capacity. Thus, we trade a modest increase in world model learning cost for better performance while preserving policy learning efficiency.\n\n|  | DreamerV3 | **CsDreamer (Ours)** | Relative Change |\n| - | :-: | :-: | :-: |\n| **Model Complexity** |  | |   |\n| Model Parameters (M)  | 19.11   | 33.27   | +74.10%    |\n| **GPU Memory Usage** |  |  |  |\n| Peak GPU Memory Allocated (GB) | 2.89 | 3.33| +15.22%  |\n| Peak GPU Memory Reserved (GB) | 4.78  | 4.96 | +3.77%  |\n| **Training & Rollout Performance**   |  |  |    |\n| Training Speed (steps per second, SPS)  | 18.50  | 12.33 | -33.35%  |\n| Imagination Rollout Speed (steps per second, SPS) | 987,895   | 989,988 | +0.21%  |\n| Total Time Usage (500K steps, hours)  | 7.81  | 11.63   | +48.91%  |\n| **Inference  Performance**   |   |    |   |\n| Inference FLOPs (M)  | 33.10  | 39.67   | +19.85% |\n| Inference Latency (ms) | 1.38 | 1.93   | +39.86%   |\n\n\n**Q2**: Limited Generality of Hyperparameter Analyses.\n\n**A2**: Thank you for this important comment. The experiment in Figure 6 is designed not only to illustrate the sensitivity but more critically, to provide empirical support for our central hypothesis: that explicitly modeling conditional dependence enhances inference.  The hyperparameter $\\lambda$ mediates a trade-off between leveraging this dependence and preventing the confounding of task-relevant and task-irrelevant representations. Our analysis shows that performance degrades at both extremes of this trade-off. As $\\lambda\\to 0$, the model risks conflating these representations, whereas as $\\lambda \\to \\infty$, it degenerates to assuming conditional independence, as in prior work, thereby forfeiting the gains from our conditional dependence modeling. The non-monotonic performance curve in Figure 6 empirically validates this trade-off. We conduct a supplementary hyperparameter analysis in two Atari 100k environments, Alien and Amidar, and in the CARLA autonomous driving simulator. The results are presented in the tables below. Due to submission constraints that prohibit figures, the tables also include intermediate results to provide a comprehensive view. The supplementary experimental results exhibit a trend consistent with that presented in Figure 6 of the original paper.\n\n| Alien   |  100K   |  200K   |  400K   |\n| :------ | :----: | :----: | :----: |\n| $\\lambda=0.01$   | $516\\pm247$ | $598\\pm219$ | $709\\pm113$ |\n| $\\lambda=0.1$    | $460\\pm227$ | $527\\pm246$ | $790\\pm414$ |\n| $\\lambda=0.2$    | $520\\pm203$ | $716\\pm323$ | $888\\pm234$ |\n| $\\lambda=1.0$    | $564\\pm396$ | $592\\pm248$ | $838\\pm345$ |\n| $\\lambda=2.0$    | $431\\pm180$ | $570\\pm142$ | $757\\pm216$ |\n| $\\lambda=10.0$   | $338\\pm150$ | $282\\pm73$  | $290\\pm174$ |\n| $\\lambda=100.0$  | $370\\pm199$ | $223\\pm77$  | $214\\pm109$ |\n| $\\lambda=1000.0$ | $202\\pm88$  | $200\\pm76$  | $210\\pm84$  |\n\n\n| Amidar    | 100K  |  200K  | 400K  |\n| :------ | :----: | :----: | :--------: |\n| $\\lambda=0.01$   | $63\\pm32$ | $81\\pm20$  | $125\\pm29$ |\n| $\\lambda=0.1$    | $52\\pm17$ | $90\\pm49$  | $146\\pm24$ |\n| $\\lambda=0.2$    | $73\\pm52$ | $110\\pm52$ | $184\\pm43$ |\n| $\\lambda=1.0$    | $50\\pm33$ | $79\\pm29$  | $108\\pm35$ |\n| $\\lambda=2.0$    | $59\\pm44$ | $77\\pm42$  | $82\\pm19$  |\n| $\\lambda=10.0$   | $36\\pm24$ | $61\\pm23$  | $50\\pm24$  |\n| $\\lambda=100.0$  | $34\\pm29$ | $36\\pm21$  | $43\\pm24$  |\n| $\\lambda=1000.0$ | $24\\pm14$ | $41\\pm19$  | $43\\pm30$  |\n\n| CARLA    | 100K  |  300K  |  500K  |\n| :------- | :----: | :----: | :----: |\n| $\\lambda=0.1$   | $28\\pm37$ | $80\\pm83$  | $94\\pm105$  |\n| $\\lambda=0.2$   | $32\\pm46$ | $86\\pm100$ | $121\\pm113$ |\n| $\\lambda=1.0$   | $33\\pm39$ | $91\\pm81$  | $128\\pm107$ |\n| $\\lambda=2.0$   | $20\\pm34$ | $80\\pm84$  | $116\\pm112$ |\n| $\\lambda=10.0$  | $14\\pm24$ | $22\\pm34$  |  $40\\pm40$  |\n| $\\lambda=100.0$ |  $0\\pm0$  | $16\\pm33$  |  $10\\pm15$  |\n\nIn terms of general guidance, our experiments reveal that $\\lambda=1$ consistently yields strong performance, and we propose it as an effective default when switching to new domains.  For instance, although $\\lambda=10$  achieves a slightly higher final performance in Figure 6, we find that $\\lambda=1$ already performs sufficiently well. We therefore adopt $\\lambda=1$ for this environment in our final reported hyperparameters to reduce the hyperparameter tuning. The optimal $\\lambda$ is correlated with the complexity of environmental distractors and their visual similarity to the agent. For domains with complex or highly similar distractors, a larger $\\lambda$ is advisable to enforce stronger decoupling. Conversely, in environments with simpler distractors, a smaller $\\lambda$ can preserve more of the benefits from the conditional structure. Importantly, while the optimal $\\lambda$ may vary across domains, we find it to be robust across different tasks within the same domain, often eliminating the need for task-specific fine-tuning.\n\n\n**Q3**: Concerns about Assumption 4.2.\n\n**A3**: Thank you for the insightful comment. While Assumption 4.2 posits a Markovian noise process, our CsRSSM model can partially mitigate the effects of non-Markovian noise. This is achieved through the recurrent structure in CsRSSM, which can implicitly capture and propagate long-term temporal correlations. Consequently, when confronted with structured distractors exhibiting long-term correlations, our proposed method maintains a degree of robustness and performance, although a performance degradation may be observed, particularly when the noise patterns feature complex periodicity or long-range dependence. We consider modeling distractors that have long-term temporal correlations an important and promising extension of this work, which we plan to investigate in future work.\n\n\n**Q4**: Concerns about Conclusion and Future Work.\n\n**A4**: We sincerely thank you for this important and valuable comment. We add a comprehensive discussion on the model's limitations, societal impact, and future work. Due to character limit, we have distilled our key points and will include detailed elaborations in the final version.\n\n**Limitations**\n\nThe primary limitations are the assumptions of Markovian and action-independent noise, which challenge the model with non-Markovian distractors or noise correlated with the agent's actions.  While conditioning noise on actions (i.e., $c_t∼p(c_t∣c_{t−1}, a_{t−1})$) is a potential solution, issues with non-Markovian dynamics would likely persist.  Furthermore, our binary partition of latent variables is an oversimplification of complex real-world scenarios. Developing methods to model the real world in a more structured manner, while striking a balance between leveraging conditional dependence and maintaining computational tractability, remains a valuable avenue for future research.\n\n**Societal Impacts**\n\n- **Positive:** Our work can enhance the robustness and safety of autonomous agents. This could lead to practical benefits such as improving the autonomous vehicle safety in adverse weather conditions, the operational capabilities of industrial robots in dynamic or cluttered settings and the efficiency of search-and-rescue robots in disaster scenarios.\n\n- **Potential Negative Societal Impacts**\n  - **Overreliance:** The model could still misclassify a rare but critical safety signal as task-irrelevant noise and ignore it, leading to catastrophic failure. **Mitigation:** Deployment must be preceded by extremely thorough and diverse testing. Maintaining a human-in-the-loop supervisory mechanism and conducting exhaustive testing on edge cases are crucial.\n  - **Misuse for Surveillance:** The technology could be adapted to track individuals by filtering environmental noise. **Mitigation:** We urge policymakers and ethics committees to establish clear guidelines restricting the application of such highly robust autonomous technologies in invasive surveillance. We recommend developing usage protocols that limit deployment in sensitive domains.\n\n**Future Work**\n\nWe will focus on addressing the above limitations by developing models for non-Markovian and action-dependent noise and exploring more structured latent representations. Furthermore, future work should continue to focus on enhancing model robustness and the transparency of its decision-making process. In terms of social impact, we will improve model interpretability to mitigate the risk of overreliance. We will also explore the design of built-in privacy-preserving mechanisms that limit the processing of personally identifiable information and establish benchmarks to assess the potential for misuse. We will remain attentive to these issues and actively investigate how to maximize the societal benefits of this technology while minimizing its potential negative impacts in our ongoing research."}, {"review_id": "irpb94aZHs", "reviewer": "Reviewer_bcfS", "summary": "This paper addresses a critical issue for Dreamer-based world modeling: the conditional dependence between task-relevant variables $s_t$ and task-irrelevant variables $c_t$. Previous works assume they are independent given observations $o_t$, which can lead to suboptimal denoising. As a solution, the authors propose CsRSSM, which generalizes the Dreamer framework to capture the conditional dependence $p(s_t|c_t,o_t)$ and further introduces decoupling regularization to avoid conflation of task-relevant and task-irrelevant variables. Experiment results show that CsDreamer, built on proposed CsRSSM, outperforms Dreamer and other baselines in terms of noisy observation denoising.", "questions": "**1. How should $s_t$ and $c_t$ be conditional dependent intuitively in general perspective?** In the introduction 1 the author introduce the ``conditional dependence'' concept in a mathematical manner, and further assert such conditional dependence is crucial for $s_t$ and $c_t$ in world modeling. However, the intuition behind is not clearly explained. Could the authors provide a more intuitive explanation of how $s_t$ and $c_t$ should be conditionally dependent in general world modeling scenarios, or when this conditional dependence is crucial?\n\n**2. Could the authors provide a clear, direct comparison between CsRSSM and RSSM?** In Eq. (3) the authors provide the framework of CsRSSM. However, despite discussed in an overall view in Eq. (2), the specific differences between CsRSSM and RSSM are not clearly highlighted. Could the authors provide a more direct comparison, such as a side-by-side comparison of the two frameworks, to clarify how CsRSSM differs from RSSM in terms of capturing conditional dependence? Furthermore, it would be better if the authors further provide a similar comparison in algorithms (in Appendix B).\n\n**3. How to interpret the decoupling regularization?** In Section 4.2, the authors aim to introduce decoupling regularization for avoiding conflation of task-relevant and task-irrelevant variables. However, the definition and motivation of this regularization are, in some words, not clearly explained:\n- *\"This conflation can hinder the model’s ability to accurately disentangle the underlying factors of the observations, ultimately degrading performance and slowing convergence.\"* How does it make sense? Why cannot such conflation be avoided by the insymmetricity itself of $s_t$ and $c_t$ in the model? Are there any references that can support this claim, or how can the authors provide a more intuitive explanation?\n- *\"The conditional mutual information, ..., $I(s_t,c_t|o_t)$* is somehow unclear. A possible suggestion is to clarity *\"the conditional mutual information between $s_t$ and $c_t$\"* and then explain the notation of $I$.\n- *\"By incorporating the variational approximate distribution $q_\\xi$...\"* The notation of $q_\\xi$ is not defined in the main text, which may lead to confusion. It would be better if the authors could rewrite in a form without using $q_\\xi$.\n\n\n**4. Could the authors provide ablation study of $\\lambda$ in other datasets?** In Secrtion 5.3 the authors conduct an analysis of $\\lambda$ on the manually constructed dataset. While the construction is a reasonable approach to introduce noise distractors, the validity in real world is not clearly justified and corresponding experiments are better to be said as \"toy experiments\". Therefore, it would be better if the authors could always provide ablation studies in real-world datasets to further validate the effectiveness of $\\lambda$.", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "- The paper focuses on the probabilistic graphical model, a significant aspect of world modeling.\n- The proposed CsRSSM is clearly formulated and theoretically sound.\n- The experiment results demonstrate the effectiveness of CsDreamer.\n\n### Weaknesses\n\n- The motivation for the conditional dependence should be better highlighted.\n- There lacks a clear, direct comparison between CsRSSM and RSSM.\n- The definition and motivation of decoupling regularization are not clearly explained.\n- The ablation study of $\\lambda$ are only conducted on toy datasets.", "was_revised": true, "final_justification": "My previous concerns are addressed.", "rebuttal": "# Response to Reviewer bcfS\n\nThank you for your valuable review and our detailed responses are provided below.\n\n\n**Q1**: The motivation for the conditional dependence should be better highlighted. How should $s_t$ and $c_t$ be conditional dependent intuitively in general perspective?\n\n**A1**: Thank you for this insightful question. The conditional dependence between the task-relevant state $s_t$ and task-irrelevant state $c_t$ is motivated by the fact that they are distinct underlying factors that jointly generate the observation $o_t$. This creates a natural collider structure ($c_t \\rightarrow o_t \\leftarrow s_t$) in the graphical model, where conditioning on the observation $o_t$ renders its parents, $s_t$ and $c_t$, conditionally dependent. For instance, in a visual scene ($o_t$), once the pixels occupied by the background (e.g., a vehicle, $c_t$) are known, the pixels available for the foreground agent ($s_t$) are constrained. Ignoring this dependence leads to an ill-posed inference problem with high ambiguity. This forces the model to search for $s_t$ and $c_t$ without mutual constraints. By explicitly modeling it, we can leverage inferred context $c_t$ to constrain the estimation of $s_t$, thereby reducing the conditional entropy $H(s_t∣c_t,o_t)$. This results in a more accurate and stable state representation, which is crucial for improving the sample efficiency and final performance of the downstream policy.\n\n\n\n**Q2**: A clear, direct comparison between CsRSSM and RSSM.\n\n**A2**: Thank you for this valuable comment. First, we present a table that provides a side-by-side comparison of the model architectures, directly highlighting how CsRSSM's design facilitates the capture of conditional dependence.\n\n| Component                  | RSSM                                              | CsRSSM                                                       |\n| -------------------------- | ------------------------------------------------- | ------------------------------------------------------------ |\n| Latent Variables           | Single variable $z_t$ which mixes all information | Dual variables, $s_t$ (task-relevant) and $c_t$ (task-irrelevant) |\n| History Model              | $h_t = f(h_{t-1}, z_{t-1}, a_{t-1})$              | $h_t^s = f(h_{t-1}, s_{t-1}, a_{t-1})$, $h_t^c = f(h_{t-1}, c_{t-1})$ |\n| Encoder                    | $p(z_t \\mid h_t, o_t)$                            | $p(c_t \\mid h_t^c, h_t^s, o_t)$, $p(s_t \\mid h_t^c, h_t^s, c_t, o_t)$ |\n| Dynamics Transition Model  | $p(\\hat{z}_t \\mid h_t)$                           | $p(\\hat{s}_t \\mid h_t^s)$, $p(\\hat{c}_t \\mid h_t^c)$         |\n| Observation Reconstruction | $p(\\hat{o}_t \\mid h_t, z_t)$                      | $p(\\hat{o}_t \\mid h_t^s, s_t, h_t^c, c_t)$                   |\n| Inference Process          | $o_t \\rightarrow z_t$                             | $o_t \\rightarrow c_t \\rightarrow s_t$ (sequential inference, leveraging conditional dependence) |\n\nWhile RSSM is trained with a standard Evidence Lower Bound (ELBO) loss containing a single KL divergence term, the loss for CsRSSM includes two separate KL divergence terms (for $s_t$ and $c_t$, respectively) and an additional conditional mutual information regularizer, $\\mathcal{L}_{MI}$, to balance the conditional dependence. CsRSSM performs a two-step inference at each timestep (first inferring $c_t$, then $s_t$), whereas RSSM performs only one. This design enables CsRSSM to infer task-relevant state $s_t$ more accurately by using the context provided by $c_t$, achieving more effective denoising.\n\nSubsequently, we present the side-by-side algorithm comparison (some mathematical symbols are omitted for brevity).\n\n```\nAlgorithm: Dreamer (RSSM)                   | Algorithm: CsDreamer (CsRSSM)\n--------------------------------------------|------------------------------------------------\n// Dynamics learning                        | // Dynamics learning\nSample {(o_t, a_t, r_t)}_{t=k}^{k+L} ~ D    | Sample {(o_t, a_t, r_t)}_{t=k}^{k+L} ~ D\n// Single latent variable inference         | // Sequential dual latent inference\nh_t = f_phi(h_{t-1},z_{t-1}, a_{t-1})       | h_t^c = f_\\phi(h_{t-1}, c_{t-1})\n                                            | h_t^s = f_\\phi(h_{t-1}, s_{t-1}, a_{t-1})\nInfer z_t ~ p_\\phi(z_t|h_t, o_t)            | Infer c_t ~ p_\\phi(c_t|h_t^c, h_t^s, o_t)\n                                            | Infer s_t ~ p_\\phi(s_t|h_t^c, h_t^s, c_t, o_t)\nUpdate world model                          | Update world model and variational estimator\n                                            |\n// Behavior learning                        | // Behavior learning\nInfer z_t ~ p_\\phi(z_t|h_t, o_t)            | Infer c_t ~ p_\\phi(c_t|h_t^c, h_t^s, o_t)\n                                            | Infer s_t ~ p_\\phi(s_t|h_t^c, h_t^s, c_t, o_t)\nImagine using p_\\phi(hat{z}_t|h_t)          | Imagine using p_\\phi(hat{s}_t|h_t^s)\nUpdate policy and critic                    | Update policy and critic\n                                            |\n// Environment interaction                  | // Environment interaction\nInfer z_t ~ p_\\phi(z_t|h_t, o_t)            | Infer c_t ~ p_\\phi(c_t|h_t^c, h_t^s, o_t)\n                                            | Infer s_t ~ p_\\phi(s_t|h_t^c, h_t^s, c_t, o_t)\nSample a_t ~ \\pi(a_t|h_t, z_t)              | Sample a_t ~ \\pi(a_t|h_t^s, s_t)\nr_t, o_{t+1} <- env.step(a_t)               | r_t, o_{t+1} <- env.step(a_t)\n```\n\n\n\n\n\n**Q3**: How to interpret the decoupling regularization?\n\n**A3**: Thank you for this perceptive comment. The introduction of the decoupling regularization is intended to address a critical issue: despite the architectural asymmetry between $s_t$ and $c_t$ (where $s_t$ is used for predicting task-relevant information, while $c_t$ contributes only to observation reconstruction), a potential problem arises during training. Since the reconstruction loss is the primary optimization objective and both variables are involved in the reconstruction process, the model may encode overlapping information into both latent variables to minimize reconstruction error. This can lead to representations that are highly correlated and mutually redundant. Consequently, task-irrelevant information may become entangled within $s_t$, which diminishes the specificity of $s_t$ to the task-relevant signals and complicates the optimization process. By introducing the decoupling regularization term, we explicitly encourage the two variables to capture distinct and complementary aspects of the observation: $c_t$ is guided to focus on task-irrelevant information, while $s_t$ is guided to concentrate on task-relevant information pertinent to rewards and actions. This explicit information partitioning enables more efficient policy learning in the latent space of $s_t$. We have also empirically validated this point in the ablation study of our paper (Figure 7) with the CsDreamer w/o MI variant. The performance degradation observed in this variant demonstrates that the absence of appropriate regularization may impair the model's decoupling capability and, ultimately, its final performance.\n\n\n\n**Q4**: Notation of $I$ and $q_\\xi$.\n\n**A4**: Thank you for the careful reading and valuable suggestions. We fully agree that the two statements in question could indeed lead to confusion. In the revised version, we will (1) add an explicit explanation of the mutual-information notation $I$; (2) rewrite the parameterized variational distribution $q_\\xi$ and give a clear accompanying description. These changes should make the presentation clearer and help readers understand our approach more easily.\n\n\n\n**Q5**: Could the authors provide ablation study of $\\lambda$ in other datasets?\n\n**A5**: Thank you for the valuable suggestion. We run additional experiments on the first two Atari100k environments, Alien and Amidar, and on the CARLA autonomous-driving simulator. The results are presented in the table below. As can be seen, the supplementary ablation study exhibits the same phenomenon as observed in the original paper: as $\\lambda$ increases, the performance first improves and then degrades.\n\n\n\n\n\n| Alien | 100K | 200K | 400K |\n|:-|:----:|:----:|:----:|\n| $\\lambda=0.01$ | $516\\pm247$ | $598\\pm219$ | $709\\pm113$ |\n| $\\lambda=0.1$ | $460\\pm227$ | $527\\pm246$ | $790\\pm414$ |\n| $\\lambda=0.2$ | $520\\pm203$ | $716\\pm323$ | $888\\pm234$ |\n| $\\lambda=1.0$ | $564\\pm396$ | $592\\pm248$ | $838\\pm345$ |\n| $\\lambda=2.0$ | $431\\pm180$ | $570\\pm142$ | $757\\pm216$ |\n| $\\lambda=10.0$ | $338\\pm150$ | $282\\pm73$ | $290\\pm174$ |\n| $\\lambda=100.0$ | $370\\pm199$ | $223\\pm77$ | $214\\pm109$ |\n| $\\lambda=1000.0$ | $202\\pm88$ | $200\\pm76$ | $210\\pm84$ |\n\n| Amidar | 100K | 200K | 400K |\n|:-|:----:|:----:|:----:|\n| $\\lambda=0.01$ | $63\\pm32$ | $81\\pm20$ | $125\\pm29$ |\n| $\\lambda=0.1$ | $52\\pm17$ | $90\\pm49$ | $146\\pm24$ |\n| $\\lambda=0.2$ | $73\\pm52$ | $110\\pm52$ | $184\\pm43$ |\n| $\\lambda=1.0$ | $50\\pm33$ | $79\\pm29$ | $108\\pm35$ |\n| $\\lambda=2.0$ | $59\\pm44$ | $77\\pm42$ | $82\\pm19$ |\n| $\\lambda=10.0$ | $36\\pm24$ | $61\\pm23$ | $50\\pm24$ |\n| $\\lambda=100.0$ | $34\\pm29$ | $36\\pm21$ | $43\\pm24$ |\n| $\\lambda=1000.0$ | $24\\pm14$ | $41\\pm19$ | $43\\pm30$ |\n\n| CARLA | 100K | 300K | 500K |\n|:-|:----:|:----:|:----:|\n| $\\lambda=0.1$ | $28\\pm37$ | $80\\pm83$ | $94\\pm105$ |\n| $\\lambda=0.2$ | $32\\pm46$ | $86\\pm100$ | $121\\pm113$ |\n| $\\lambda=1.0$ | $33\\pm39$ | $91\\pm81$ | $128\\pm107$ |\n| $\\lambda=2.0$ | $20\\pm34$ | $80\\pm84$ | $116\\pm112$ |\n| $\\lambda=10.0$ | $14\\pm24$ | $22\\pm34$ | $40\\pm40$ |\n| $\\lambda=100.0$ | $0\\pm0$ | $16\\pm33$ | $10\\pm15$ |"}, {"review_id": "0Zg1VUIyZR", "reviewer": "Reviewer_kxo7", "summary": "This paper introduces CsDreamer, a model-based reinforcement learning (RL) approach designed to effectively handle noisy observations in visual environments. The key innovation lies in the Collider-structure Recurrent State-Space Model (CsRSSM), which explicitly models the conditional dependence between task-relevant and task-irrelevant latent variables, denoted as s_t and c_t, respectively. Unlike prior works that assume conditional independence between these variables given an observation, CsRSSM leverages the collider structure to first infer c_t from the observation o_t, and then uses both c_t and o_t to infer s_t.\n\nTo further aid in representation learning, the authors introduce a decoupling regularization via conditional mutual information to minimize overlap between the information captured by s_t and c_t. The final CsDreamer architecture integrates CsRSSM into the Dreamer RL framework and uses rollouts in the task-relevant latent space to improve policy learning.", "questions": "1. How does the runtime/memory footprint of CsRSSM compare to DreamerV3, especially given the dual latent variables (s_t, c_t) and additional regularization?\n\n2. Assumption 4.2 states c_t is Markovian and action-independent. How does the model handle cases where noise dynamics are action-dependent (e.g., adversarial perturbations)?\n\n3. For Atari 100K, what explains the outlier seed’s performance in Up N Down? Is this indicative of sensitivity to initialization?\n\n4. As shown in Table 3, some results are highly unstable, e.g. 7771 +/- 18563 for Gopher and 9800 +/- 13122 for Up N Down. Are there any explanations to them?\n\n5. Beyond performance trends, is there a principled way to choose \\lamda, or must it be tuned per environment?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "** Strengths\n\n1. The derivation of the ELBO and the conditional mutual information-based regularization are mathematically sound and well-justified.\n\n2. The paper clearly articulates why conditional dependence matters in noisy settings, using intuitive examples and diagrams.\n\n3. The paper addresses a critical challenge in RL: robustness to noisy observations, which is vital for real-world applications like autonomous driving. It introduces a principled approach to exploit conditional dependence, advancing the state of the art in model-based RL.\n\n4. The collider-structure-based generative modeling for denoising is original and rarely explored in model-based RL. Applying mutual information bounds to control latent variable interaction is a meaningful contribution.\n\n** Weakness\n\n1. Limited discussion on computational overhead of CsRSSM compared to baselines (e.g., DreamerV3).\n\n2. While CsRSSM offers a unified framework, many components--e.g., disentangled latent spaces, decoders, KL regularization--overlap with prior structured world models. The novelty primarily lies in combining them under a collider assumption.\n\n3. Assumption 4.2 is somewhat strong. It is reasonable only in settings where distractors evolve independently and may be violated in real-world or agent-interactive scenarios.", "was_revised": true, "final_justification": "The authors' responses address most of my concerns, and I understand that some of my concerns are open challenges in the field. Although there are no perfect solutions, the authors' answers prove they have deep insights into the topic. Therefore, I'd like to raise my rating to \"Borderline accept\".", "rebuttal": "# Response to Reviewer kxo7\n\nWe would like to extend our sincere thanks for the time and effort dedicated to reviewing our submission and for the highly valuable comments. Please find our point-by-point responses below.\n\n\n\n\n\n**Q1**: Limited discussion on computational overhead of CsRSSM compared to baselines.\n\n**A1**: Thank you for raising this important point. A detailed analysis for the Cheetah-Run+GB task, conducted under an identical experimental setup with an AMD Ryzen 9 7950X CPU and an NVIDIA GeForce RTX 4090 GPU is summarized in the provided table. It confirms that the CsDreamer, with its dual latent variables $(s_t, c_t) $ and additional regularization, naturally incurs a higher computational cost than DreamerV3. Specifically, this is reflected by a 74.1% increase in model parameters, a 15.2% increase in peak allocated GPU memory, and a 39.9% increase in single-step inference latency. These factors culminate in a 48.9% increase in total time usage.  However, the imagination rollout speed, a key determinant of training efficiency in world models, remains virtually identical to that of DreamerV3 (+0.21%). This is because the imagination process in CsDreamer is performed exclusively within the compact, task-relevant state space $s_t$. To ensure a fair comparison of performance, we intentionally design its dynamics model to match the dimensionality and complexity of the latent space in DreamerV3. We also keep the actor-critic networks identical to DreamerV3's, ensuring that any improvements in final task performance are attributable to the quality of the learned representations, not an increase in policy capacity. In essence, CsRSSM makes a calculated trade-off, accepting a moderate increase in world model update costs to achieve a more robust state representation, while preserving the rollout efficiency essential for effective policy learning.\n\n\n|                                                   | DreamerV3 | **CsDreamer (Ours)** | Relative Change |\n| :-------------------------------------------------| :---------: | :--------------------: | :---------------: |\n| **Model Complexity**                              |           |                      |                 |\n| Model Parameters (M)                              | 19.11     | 33.27                | +74.10%         |\n| **GPU Memory Usage**                              |           |                      |                 |\n| Peak GPU Memory Allocated (GB)                    | 2.89      | 3.33                 | +15.22%         |\n| Peak GPU Memory Reserved (GB)                     | 4.78      | 4.96                 | +3.77%          |\n| **Training & Rollout Performance**                |           |                      |                 |\n| Training Speed (steps per second, SPS)            | 18.50     | 12.33                | -33.35%         |\n| Imagination Rollout Speed (steps per second, SPS) | 987,895   | 989,988              | +0.21%          |\n| Total Time Usage (500K steps, hours)              | 7.81      | 11.63                | +48.91%         |\n| **Inference  Performance**                        |           |                      |                 |\n| Inference FLOPs (M)                               | 33.10     | 39.67                | +19.85%         |\n| Inference Latency (ms)                            | 1.38      | 1.93                 | +39.86%         |\n\n\n\n\n\n**Q2**: Many components overlap with prior structured world models and the novelty primarily lies in combining them under a collider assumption.\n\n**A2**: Thank you for your valuable comment. While we acknowledge the use of established components, we would like to respectfully clarify that our main contribution lies in a more fundamental conceptual advance. we challenge the prevailing assumption of conditional independence in prior work and instead propose a collider-structure model that more accurately reflects the real-world generative process. Previous methods often discard valuable information by assuming that task-relevant and task-irrelevant factors are independent given an observation. Our model, by contrast, leverages the conditional dependence inherent in the collider structure. This enables it to utilize task-irrelevant variables as essential context to infer task-relevant states more accurately and efficiently. We are encouraged that Reviewer LU9u also highlights the significance of this shift, stating, \"Moving beyond the standard assumption of conditional independence to explicitly model the collider structure is a non-trivial, principled insight that better reflects the underlying generative process.\" Therefore, we believe our contribution is not merely a novel combination of modules, but the introduction of a new, principled, and more effective modeling paradigm for reinforcement learning in noisy environments.\n\n\n\n**Q3**: Assumption 4.2 is somewhat strong. How does the model handle cases where noise dynamics are action-dependent?\n\n**A3**: Thank you for this perceptive question. We acknowledge that Assumption 4.2 is a relatively strong condition. This assumption is motivated by many real-world scenarios where distractors evolve independently of the agent's actions, such as changing weather, lighting conditions, or background traffic in autonomous driving. Regarding the more complex case where noise dynamics are action-dependent, we would like to clarify that our model architecture offers a degree of adaptability. Critically, while the prior transition model for the task-irrelevant state, $p(c_t|c_{t-1})$, is action-independent, the posterior inference, $p(c_t|h_t^c, h_t^s, o_t)$, is conditioned on the task-relevant history $h_t^s$, which encapsulates past actions. This mechanism allows the model to indirectly leverage action information to refine its estimate of  $c_t$, thereby providing a degree of robustness to disturbances that are weakly correlated with actions. Nevertheless, we agree that the current model would be challenged by noise that is strongly dependent on the agent's actions. We consider modeling such scenarios an important and promising direction for future research. A natural and promising extension would be to model the transition of $c_t$ as explicitly action-conditioned, e.g., $c_t \\sim p(c_t|c_{t-1}, a_{t-1})$. We plan to explore this avenue in future work.\n\n\n\n\n\n**Q4**: Explanation for the outlier seed’s performance and unstable results.\n\n**A4**: Thank you for this valuable question. The high variance you have pointed out, particularly in exploration-hard games like Up N Down and Gopher, is indeed a core challenge inherent to the extreme data-efficiency constraints of the Atari 100k benchmark. The root cause is the stringent budget of 100,000 interaction steps, which prevents most RL agents from exploring the state space adequately, often leading them to converge to suboptimal policies. Our method, by modeling the environment's conditional dependencies, learns the core game dynamics more efficiently. This enhanced efficiency provides the agent, under certain random seeds, the unique capability to overcome exploration bottlenecks and exploit high-reward strategies that are otherwise inaccessible within such a limited budget. Therefore, the outlier seed's performance in Up N Down and the resulting high standard deviation are not merely signs of instability but are direct manifestations of our method's effectiveness.  While some classical works [1, 2, 3] omit standard deviations on this benchmark in their results table, we choose to report them transparently. We believe this provides the community with an honest and realistic view of CsDreamer's performance on this challenging benchmark. To ensure our evaluation is as rigorous as possible, we utilize an increased number of random seeds to mitigate this inherent stochasticity.\n\n\n\n\n\n**Q5**: Beyond performance trends, is there a principled way to choose $\\lambda$, or must it be tuned per environment\n\n**A5**: Thank you for this important question. The choice of $\\lambda$ is guided by the principle of balancing a fundamental trade-off: leveraging the conditional dependence between task-relevant $s_t$ and task-irrelevant $c_t$ variables for more efficient inference versus enforcing their decoupling to prevent representational confounding. The experiment in Figure 6 provides empirical validation for this principle. A very small $\\lambda$ risks conflating the two representations, while a very large $\\lambda$ effectively degenerates to the conditional independence assumption of prior methods, losing the benefits of our model's structure. The non-monotonic curve in Figure 6 demonstrates that optimal performance is achieved by balancing these two forces. From a practical standpoint, this principle provides clear guidance. While the ideal $\\lambda$ can depend on the complexity of the environmental distractors, our experiments show that $\\lambda=1$ serves as a robust and effective default that delivers strong performance across diverse environments. More importantly, we find that $\\lambda$ does not need to be tuned for every individual task. Once set for a specific domain, it often generalizes well to other tasks within that same domain, drastically reducing the tuning overhead and enhancing the practical applicability of our method.\n\n\n\n\n\n\n\n[1] Kaiser et al. Model-Based Reinforcement Learning for Atari.\n\n[2] Kostrikov et al. Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels.\n\n[3] Hafner et al. Mastering Diverse Domains through World Models."}, {"review_id": "5tdv28Le9X", "reviewer": "Reviewer_hEJH", "summary": "This work introduces CsRSSM, a generative world model that explicitly leverages the conditional dependence between task-relevant and task-irrelevant latent variables. By sequentially inferring the task-irrelevant component before the task-relevant one and employing a conditional mutual information regularizer, the model effectively disentangles noisy observations. Built upon this, the proposed CsDreamer framework significantly enhances policy learning performance in visually distracting environments compared to state-of-the-art approaches.", "questions": "1. Could the authors clarify whether CsRSSM is compatible with transformer-style sequence encoders or if its benefits are limited to RSSM-based designs? A controlled comparison or ablation with transformer-based baselines would strengthen the contribution.\n\n2. Could the authors provide further insights or heuristics on how to select $\\lambda$ in practice, especially when switching to new domains?\n\n3. How does the runtime cost of CsDreamer compare to DreamerV3 in terms of rollout speed and training time per iteration?", "limitations": "Yes.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths: 1. Novel structural modeling of conditional dependence: The paper introduces CsRSSM, which explicitly models the conditional dependence between task-relevant and task-irrelevant latent variables ($s_t$, $c_t$) via a collider structure. This goes beyond prior works that assume conditional independence, providing a theoretically grounded improvement.\n2. Principled regularization via conditional mutual information: The use of a variational upper bound on conditional mutual information I($s_t$; $c_t$|$o_t$) as a regularizer is novel and supports the structural disentanglement of latent variables in a meaningful and interpretable way.\n3. Strong empirical results in noisy environments: The proposed CsDreamer achieves consistent improvements over competitive baselines (e.g., DreamerV3, TIA, Iso-Dream) in DMC, CARLA, and Distracting Control Suite benchmarks, especially under challenging visual distractors.\n4. Well writing.\n\nWeakness: 1. Limited evaluation on transformer-based baselines. Given the recent trend of replacing RSSM with Transformer-based architectures in world models (e.g., TransDreamer, Trajectory Transformer), the current evaluation lacks a clear comparison with such models.\n2. Sensitivity to the decoupling regularization weight $\\lambda$. As shown in Fig. 6, the model performance is non-monotonically sensitive to the value of $\\lambda$, indicating the importance of hyperparameter tuning.\n3. Missing discussion on inference-time cost and rollout efficiency. While the paper discusses representation quality, it omits inference-time analysis. Since CsRSSM sequentially infers $c_t$ then $s_t$, it may increase latency.", "was_revised": true, "final_justification": "The author's response addressed most of my questions (including regularization weight and inference time, etc.), but the author mentioned that due to time constraints, no comparative experiments with the transformer as I hoped were provided in the rebuttal stage. Therefore, I believe this paper still has areas that need improvement. I will maintain my score of borderline accept. I hope the author can conduct further experiments to supplement and enrich this work.", "rebuttal": "# Response to Reviewer hEJH\n\nThank you for the careful review and constructive comments. Our point-by-point responses are listed below.\n\n\n\n**Q1**: Compatibility and comparison with Transformer models.\n\n**A1**: Thank you for this valuable suggestion. The core idea of our work, using conditional dependence for denoising, is orthogonal to the choice of the sequence modeling backbone. Therefore, our method is indeed compatible with Transformer-based architectures, which can serve as a direct replacement for the RSSM component. Due to time constraints during the rebuttal period, we are unable to provide these experiments immediately. However, we commit to incorporating a controlled comparison against Transformer-based baselines in the final camera-ready version. This will better highlight the generality of our contribution. Thank you again for the suggestion, which helps us improve the quality of our paper.\n\n\n\n**Q2**: Sensitivity to the decoupling regularization weight $\\lambda$. Could the authors provide further insights or heuristics on how to select $\\lambda$ in practice, especially when switching to new domains?\n\n**A2**: Thank you for this insightful question. The experiment in Figure 6 is designed not only to illustrate this sensitivity but, more critically, to provide empirical support for our central hypothesis: that explicitly modeling conditional dependence enhances inference.  The weight $\\lambda$ mediates a trade-off between leveraging this dependence and preventing the confounding of task-relevant and task-irrelevant representations. As $\\lambda\\to 0$, the model risks conflating these representations, whereas as $\\lambda \\to \\infty$, it degenerates to assuming conditional independence, as in prior work, thereby forfeiting the gains from our conditional dependence modeling. The non-monotonic performance curve in Figure 6 empirically validates this trade-off.  In practice, the optimal $\\lambda$ is correlated with the complexity of environmental distractors and their visual similarity to the agent. For domains with complex or highly similar distractors, a larger $\\lambda$ is advisable to enforce stronger decoupling. Conversely, in environments with simpler distractors, a smaller $\\lambda$ can preserve more of the benefits from the conditional structure. Importantly, while the optimal $\\lambda$ may vary across domains, we find it to be robust across different tasks within the same domain, often eliminating the need for task-specific fine-tuning. Our experiments reveal that $\\lambda=1$ consistently yields strong performance. We therefore propose it as an effective default when switching to new domains, to be fine-tuned only if necessary.\n\n\n\n**Q3**: Missing discussion on inference-time cost and rollout efficiency.\n\n**A3**: Thank you for the important question regarding the computational overhead. We conduct a detailed performance analysis comparing CsDreamer to DreamerV3 under an identical experimental setup (Cheetah-Run+GB task on an AMD Ryzen 9 7950X CPU and an NVIDIA GeForce RTX 4090 GPU), with the results presented in the table. Our analysis confirms that due to its larger model size (+74.1% parameters) and the sequential inference of the task-irrelevant state $c_t$ and the task-relevant state $s_t$, CsDreamer exhibits a higher single-step inference latency (+39.86%), leading to a longer total training time (+48.91%). Crucially, however, the imagination rollout speed, a key determinant of training efficiency in world models, remains virtually identical to that of DreamerV3 (+0.21%). This is because the imagination process in CsDreamer is performed exclusively within the compact, task-relevant state space $s_t$. To ensure a fair comparison of performance, we intentionally design its dynamics model to match the dimensionality and complexity of the latent space in DreamerV3. We also kept the actor-critic networks identical to DreamerV3's, ensuring that any improvements in final task performance are attributable to the quality of the learned representations, not an increase in policy capacity. Therefore, while there is a modest increase in the per-step inference cost, the critical rollout efficiency for policy learning is preserved. \n\n\n|                                                   | DreamerV3 | **CsDreamer (Ours)** | Relative Change |\n| ------------------------------------------------- | :---------: | :--------------------: | :---------------: |\n| **Model Complexity**                              |           |                      |                 |\n| Model Parameters (M)                              | 19.11     | 33.27                | +74.10%         |\n| **GPU Memory Usage**                              |           |                      |                 |\n| Peak GPU Memory Allocated (GB)                    | 2.89      | 3.33                 | +15.22%         |\n| Peak GPU Memory Reserved (GB)                     | 4.78      | 4.96                 | +3.77%          |\n| **Training & Rollout Performance**                |           |                      |                 |\n| Training Speed (steps per second, SPS)            | 18.50     | 12.33                | -33.35%         |\n| Imagination Rollout Speed (steps per second, SPS) | 987,895   | 989,988              | +0.21%          |\n| Total Time Usage (500K steps, hours)              | 7.81      | 11.63                | +48.91%         |\n| **Inference  Performance**                        |           |                      |                 |\n| Inference FLOPs (M)                               | 33.10     | 39.67                | +19.85%         |\n| Inference Latency (ms)                            | 1.38      | 1.93                 | +39.86%         |"}]

Author rebuttal:

Meta-review: The paper presents CsDreamer, an agent based on model-based reinforcement learning that is able to learn in distracting environments. The novelties proposed are centered around the world model, called CsRSSM, where the task of denoising observations is reformulated. Instead of treating the conditional independence between the task-relevant state, the approach uses the information that models their conditional dependence, also using the collider structure in the model. The results of the proposed approach are consistently strong and compelling, with a clear superiority over a number of pertinent state-of-the-art baselines in challenging and diverse benchmarks, such as the DeepMind Control Suite with video distractors or CARLA.

The authors have managed to address concerns and doubts made by the reviewers, related to the initial computational costs, hyperparameter sensitivity, and the presentation of the novelty of the method clearly. Particularly, a well detailed response regarding computational costs was provided, by showing that although the update of the world model is indeed more expensive, the crucial imagination rollout speed has been maintained. On the regularization hyperparameter, additional analysis was likewise provided. Also, clear comparisons were made between the results obtained using proposed approach and those obtained using standard RSSM framework, presented in a side-by-side manner. The first two ablation studies were satisfactory with respect to issues raised by reviewers bcfs and kxo7. The paper had been praised for its strong ablation studies, and for visualization clarity, allowing reviewers to understand clearly what is the mechanism at play. Eventually, all four reviewers voted for acceptance.

There are still some trivial issues, like comparison to state-of-the-art transformer-based world models (highlighted by reviewer hEJH), but they are justifiable for a single contribution and represent solid future work. Overall, one can say this paper makes a clear, novel and significant contribution to representation learning in model-based RL. A sound theoretical insight, solid empirical evidence with extensive ablations and an earnest rebuttal make the case for acceptance very clear.
Decision: Accept (poster)