Response to Reviewer hEJH

Thank you for the careful review and constructive comments. Our point-by-point responses are listed below.

Q1: Compatibility and comparison with Transformer models.

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.

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?

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.

Q3: Missing discussion on inference-time cost and rollout efficiency.

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.

Response to Reviewer kxo7

We 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.

Q1: Limited discussion on computational overhead of CsRSSM compared to baselines.

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.

Q2: Many components overlap with prior structured world models and the novelty primarily lies in combining them under a collider assumption.

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.

Q3: Assumption 4.2 is somewhat strong. How does the model handle cases where noise dynamics are action-dependent?

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.

Q4: Explanation for the outlier seed’s performance and unstable results.

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.

Q5: Beyond performance trends, is there a principled way to choose $\lambda$, or must it be tuned per environment

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.

[1] Kaiser et al. Model-Based Reinforcement Learning for Atari.

[2] Kostrikov et al. Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels.

[3] Hafner et al. Mastering Diverse Domains through World Models.

Response to Reviewer bcfS

Thank you for your valuable review and our detailed responses are provided below.

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?

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.

Q2: A clear, direct comparison between CsRSSM and RSSM.

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.

While 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.

Subsequently, we present the side-by-side algorithm comparison (some mathematical symbols are omitted for brevity).

Algorithm: Dreamer (RSSM)                   | Algorithm: CsDreamer (CsRSSM)
--------------------------------------------|------------------------------------------------
// Dynamics learning                        | // Dynamics learning
Sample {(o_t, a_t, r_t)}_{t=k}^{k+L} ~ D    | Sample {(o_t, a_t, r_t)}_{t=k}^{k+L} ~ D
// Single latent variable inference         | // Sequential dual latent inference
h_t = f_phi(h_{t-1},z_{t-1}, a_{t-1})       | h_t^c = f_\phi(h_{t-1}, c_{t-1})
                                            | h_t^s = f_\phi(h_{t-1}, s_{t-1}, a_{t-1})
Infer 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)
                                            | Infer s_t ~ p_\phi(s_t|h_t^c, h_t^s, c_t, o_t)
Update world model                          | Update world model and variational estimator
                                            |
// Behavior learning                        | // Behavior learning
Infer 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)
                                            | Infer s_t ~ p_\phi(s_t|h_t^c, h_t^s, c_t, o_t)
Imagine using p_\phi(hat{z}_t|h_t)          | Imagine using p_\phi(hat{s}_t|h_t^s)
Update policy and critic                    | Update policy and critic
                                            |
// Environment interaction                  | // Environment interaction
Infer 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)
                                            | Infer s_t ~ p_\phi(s_t|h_t^c, h_t^s, c_t, o_t)
Sample a_t ~ \pi(a_t|h_t, z_t)              | Sample a_t ~ \pi(a_t|h_t^s, s_t)
r_t, o_{t+1} <- env.step(a_t)               | r_t, o_{t+1} <- env.step(a_t)

Q3: How to interpret the decoupling regularization?

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.

Q4: Notation of $I$ and $q_\xi$.

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.

Q5: Could the authors provide ablation study of $\lambda$ in other datasets?

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.

Response to Reviewer LU9u

Thank you for the thorough and insightful review. We address the comments point-by-point as follows.

Q1: Lack of Computational Overhead Analysis.

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.

Q2: Limited Generality of Hyperparameter Analyses.

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.

In 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.

Q3: Concerns about Assumption 4.2.

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.

Q4: Concerns about Conclusion and Future Work.

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.

Limitations

The 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.

Societal Impacts

• 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.

• Potential Negative Societal Impacts

  • 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.
  • 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.

Future Work

We 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.