DMWM: Dual-Mind World Model with Long-Term Imagination

We are grateful to Reviewer Y1Jx for spending time on reviewing our manuscript and raising many valuable improvement suggestions. In the response and our updated manuscript, we have made every effort to answer and address your insightful concerns. Our detailed responses are given as follows.

W1: The reliance on predefined logical rules may limit adaptability to dynamic environments where such rules are not clear-cut. Although the authors have mentioned that in limitations, this remains a big concern about the generalizability and scalability of the proposed dual-process methods.

A: We thank you for this comment that deserves further discussion. We first note that our method is general since the logical thinking ability of the System 2 component can be obtained by a learning-based method. In particular, the System 2 component be designed to learn logical rules from the actual environment changes with one-step logic $s_t \land a_t \rightarrow s_{t+1}$ and deep logical chain $s_{t-\alpha} \land a_{t-\alpha} \cdots \land s_t \land a_t \rightarrow s_{t+1}$. This learning-based method is applicable to any environment and any task without reliance on known and predefined logical rules. If there is an environment with clear rules, such as traffic rules for vehicles in autonomous driving, these rules can also scale to our proposed method as the logical regularization. Hence, our learning-based System 2 component is general and scalable to more complex applications. With regards to our claim regarding the reliance on domain-specific logic, this relates to the fact that action $a_t$ in $s_t \land a_t \rightarrow s_{t+1}$ follows the specific task $g$, i.e., $a \in \mathcal{A}_g \subsetneq \mathcal{A}$, and it leads to $s \in \mathcal{S}_g \subsetneq \mathcal{S}$, where $S$ is the state space of the open world, $\mathcal{A}$ is the action space of the agent, and $\mathcal{S}_g$ and $\mathcal{A}_g$ are respectively the subset of the state space and the action space with the task $g$. In this context, we can use the goal-conditioned method, similar to TD-MPC2 [1], to overcome this particular limitation and scale DMWM to different tasks with generalization. In particular, DMWM will learn a task vector $g$ as an additional input to the System 1 component, such as $\hat{z}_t \sim p\left(\hat{z}_t \mid h_t, g\right)$. To demonstrate the generalization and scalability of the proposed method, we have added two experiments that extend DMWM to multi-task environments (scalability) and few-shot learning for new tasks (generalization). Due to the image and link limitations, we present data tables below for the experiment results. We first investigate the task performance in multi-task complex environments.

In Table 1, the value for DMC tasks is episode return, and the value for ManiSkill2 and MyoSuite is success rate. Moreover, to support the multi-task environments, varying action dimensions across tasks by zero-padding are handled to a unified maximum dimension, and the action masking is applied to ignore invalid dimensions during training and inference [1]. The simulation results demonstrate that the proposed DMWM can be extended to multi-task environments based on the goal-conditioned method [1]. Compared to TD-MPC2, DMWM can achieve a performance improvement of 8% on DMC tasks, 30% on ManiSkill2, and 6% on MyoSuite at 1M environment steps by using the same hyperparameters across all tasks.

Next, we investigate the generalization of the DMWM to new tasks, where the DMWM is trained on 16 DMC tasks and then finetuned to 4 new DMC tasks. The task vector $g$ is initialized by a semantically similar seen task.

The performance improvement across unseen tasks with few-shot learning demonstrates that DMWM can generalize effectively to novel environments. Compared to training from scratch, the finetuned model achieves higher scores under limited environment steps.

References: [1] N. Hansen, H. Su, and X. Wang. Td-mpc2: Scalable, robust world models for continuous control. International Conference on Learning Representations, 2024.

W2: Missing references. There are more relevant studies about dual-process models from human cognition and physical inference [1-4].

A: We have carefully read these references and have introduced them in our updated manuscript. These references are also valuable to our future work.

W3: While the model performs well in controlled environments, its effectiveness in dynamic and unpredictable real-world scenarios remains uncertain.

A: While our experiments focus on controlled benchmarks such as the DMControl suite and robotic platforms, such as MyoSuite and ManiSkill2, we acknowledge that these environments, despite their complexity, may not fully capture the variability and stochasticity present in open-ended real-world scenarios. We have provided a more detailed discussion of this limitation and the potential solutions in the appendix (Please see the response to W2 of Reviewer Jdzu). In particular, for future work, to address the effectiveness of WMWM, we plan to develop the probabilistic logic to handle the ambiguous logical rules in real-world scenarios and develop the mechanisms for online logical rule adaptation to handle evolving and uncertain environments.

Q1: How do the authors plan to extend the DMWM framework to environments with ambiguous or evolving logical rules?

A: For ambiguous logical rules, we can further develop the probabilistic logic for the System 2 component. In particular, the existing logic similarity metric $\text{Sim}(\varphi, \text{True}) \in [0,1]$ for DMWM can be considered as a probabilistic estimate $\mathbb{P}(\varphi)$ of the truth value of a logical proposition $\varphi$. In this way, the probabilistic logic can be incorporated by introducing soft constraints of the form $\mathbb{P}(\varphi) \geq \tau$, where $\tau$ is a required confidence threshold. These constraints can be integrated into the logic loss through regularization terms or logic-enhanced variational objectives that reflect the probabilistic consistency of imagined trajectories. To address the evolving logical rules, DMWM can incorporate meta-learning or continual-learning mechanisms to dynamically update logic rules.

Q2: What strategies will be adopted to mitigate the complexity introduced by the dual-system architecture in real-world applications?

A: One of the most promising strategies can be selective logical Inference through uncertainty estimation. In particular, to reduce the computational overhead of always-on logical reasoning, we can use uncertainty-aware control mechanisms, such as the mentioned probabilistic logic in the response to your Q1 comment, to selectively activate the System 2 component. In particular, when the System 1 component exhibits high epistemic or aleatoric uncertainty in prediction (e.g., via ensemble variance or Monte Carlo dropout), logical reasoning from the System 2 component will be activated to guide or correct the imagined trajectory. This adaptive triggering mimics the cognitive allocation of deliberative reasoning in humans and significantly reduces redundant logical inference in seen situations. Moreover, we can adopt sparse and pruned logic representations. In particular, as the logic heatmap in Figure 3, we have found that many logical rules can be sparsely activated since the logical operations within the System 2 component are executed over a structured set of vectorized embeddings. This enables effective pruning or attention-based selection of relevant rules during inference, which can provide a reduced computation graph without losing consistency.

W1: The logical rules have to be known and predefined, which is fine for some applications but a limiting factor for others.

A: Thank you for this comment. Our proposed method is general and scalable since the logical thinking ability of the System 2 component can be obtained by a learning-based method instead of hand-crafted rules. In particular, the System 2 component is designed to learn logical rules from the actual environment trajectories $s_t \land a_t \rightarrow s_{t+1}$. This method is applicable to any environment and any task without reliance on prior knowledge of logic. To demonstrate the generalization and scalability of the proposed method, we have added two experiments that extend DMWM to multi-task environments and few-shot learning for new tasks (Please see our response to W1 of Reviewer Y1Jx).

W2: Limitations of the approach could be better discussed. When does the framework perform particularly good/ not good?

A: We first detail the cases in which DMWM performs well:

• Environments with Clear Structural Rules: DMWM uses logic rules most effectively in domains where domain knowledge (e.g., physical constraints or control dependencies) can be formalized through first-order logic. For instance, in the Cartpole Balance task from the DMControl suite, domain rules like “if the pole angle exceeds a threshold, the agent must apply a force in the opposite direction to maintain balance” can be encoded as logical implications, i.e., angle > θ → force = left. In such environments with clear structural rules, logical regularization and learned logical rules can improve both interpretability and generalization.
• Complex Tasks Requiring Long-Horizon Planning: RSSM-S1 performs well in short-term prediction but cannot handle the error accumulation over longer horizons. LINN-S2 enhances recursive implication reasoning to construct long-term logical chains and aligns the predicted trajectories with logical inference, as shown in (10) and (11). This reliable imagination ability is important for complex tasks with the requirements of long-horizon planning. For instance, in tasks such as Quadruped Run in Figure 13 and Key Turn Hard in Figure 12, DMWM achieves more stable long-horizon planning compared to Dreamer.
• Limited-Data or Limited-Trial Settings: Unlike conventional RSSM-based models that rely solely on statistical patterns in observed transitions, LINN-S2 imposes logical constraints that function as supervisory signals. These logic-driven regularizations make the model generalize effectively from sparse or partial observations, which enables more efficient learning of long-horizon dynamics. As shown in Figures 9–12, DMWM outperforms baselines with substantially fewer environment interactions.

Next, we discuss some of the scenarios in which DMWM may have limited performance gains:

• Simple Tasks with Short-Term Dependencies: In simple tasks that are characterized by short-term dependencies and low-dimensional dynamics, such as Cartpole Balance or Cup Catch, DMWM takes limited advantage over baseline models, as shown in Figure 10. The additional logical reasoning from LINN-S2 does not significantly enhance prediction quality, as the short-term dynamics can be effectively captured by RSSM-S1 alone. In such cases, the complexity and computational overhead introduced by hierarchical logical inference are over the task demands.
• Highly Stochastic Environments: In highly stochastic environments, LINN-S2 faces difficulty in establishing effective logical dependencies between states and actions. Although LINN-S2 learn logical rules from general environment transitions $s_t \land a_t \rightarrow s_{t+1}$, this still relies on the presence of stable, interpretable relationships, which are often absent in high-stochastic settings with sparse logic rules. In this context, the logic consistency loss $C(\phi)$ can be noisy and unreliable. To address this, future work can incorporate uncertainty-aware logic modeling or probabilistic logic reasoning to selectively apply logical constraints only where they are supported by sufficient structural confidence.
• Non-Stationary Environments: In non-stationary environments where underlying logical dependencies evolve over time, the static formulation of LINN-S2 fails to maintain accurate inference since it cannot adapt to abrupt changes in environment dynamics. This failure can lead to outdated or incorrect logical constraints being imposed on RSSM-S1 and lead to degraded imagination quality, along with unstable planning. Future extensions can incorporate continual learning to dynamically update logic rules.

We will add this discussion to the appendix.

Q1 What are the main insights? Under which conditions (e.g. size/dimensions of the state space or action space) the framework performs good / not good? Where does the good performance come from? Which components are responsible?

A: We agree that this issue deserves more discussions, as explained next. Our work proposes a dual cognitive system with fast, data-driven System 1 and logic-driven System 2, which enables robust long-horizon imagination by mitigating the accumulated prediction errors. This is achieved through a recursive logic reasoning framework that imposes local and global logical consistency on temporal predictions to improve long-term planning. The inter-system feedback mechanism ensures that symbolic logical rules can refine intuitive transitions and adapt logic based on environmental feedback. The conditions of performance can be seen in our response to W2. Performance gains primarily arise from the LINN-S2 component, which learns logical rules from the environment through hierarchical deep reasoning, thereby ensuring logical consistency in imagined rollouts. The inter-system feedback mechanism enables the intuitive predictions of RSSM-S1 to align with logical constraints, thus creating a closed-loop dual-mind learning process. Overall, the dual-mind design and the inter-system feedback are responsible for performance gains.

Q2: Does the proposed technique scale to larger problems (dimensions of the state space or action space)? How many training will be required? Please discuss.

A: Yes, DMWM can scale to larger problems due to its modular dual-system architecture. In particular, RSSM-S1 captures environment dynamics through compact latent representations to generalize across high-dimensional observations. Meanwhile, the LINN-S2 enhances scalability by reasoning over logic-structured embeddings rather than raw state-action inputs, which reduces the dimensional complexity of inference. The use of logic operations in LINN-S2 abstracts high-dimensional dependencies into compact, symbolic representations. It maintains reasoning tractability even as the state/action dimension increases. Moreover, the proposed hierarchical reasoning mechanism (10) enables the model to recursively infer over long horizons without unrolling full trajectories, which also alleviates the computational overhead in large-scale environments.

Q3: What are the required runtimes for the performed experiments?

A: All experiments are conducted on a single NVIDIA RTX 3090 GPU. Training a single DMC task with DMWM requires 1.2 GPU days, compared to 0.35 GPU days for a standard RSSM. It reflects the additional cost of dual-system optimization and inter-system feedback. However, DMWM introduces no additional computational overhead for inference after training since S1 and S2 function as a unified model, with logical reasoning distilled into RSSM-S1. Hence, DMWM maintains runtime complexity comparable to RSSM.

Q4: How problematic is the tuning of the hyperparameters used?

A: The tuning of hyperparameters is moderately complex. In particular, hyperparameters stem from RSSM-S1, LINN-S2, and inter-system feedback. While RSSM-S1 builds on existing hyperparameters from architectures like DreamerV3, LINN-S2 introduces additional tuning needs, such as reasoning depth and logic regularization weights. However, these are guided by logical rules, which reduce arbitrariness and improve stability. The inter-system feedback mechanism requires balancing, but is supported by clearly defined loss components. Overall, although tuning DMWM involves more components than standard world models, the process remains feasible and justified.

Q5: Could the operations used to describe logical rules be extended?

A: Yes, they could. As we have implemented the propositional logic system with $\neg$, $\land$, $\lor$, and $\to$ in a modular way, it is easy to form more complex, compositional logical operations within propositional logic. Moreover, the operations can also be extended to first-order logic ($\forall$ and $\exists$) and probabilistic logic ($P(\varphi)$). In particular, the universal quantification $\forall$ can be extended by $\forall x: \varphi(x) \to \min_{x \in D} \text{Sim}(\varphi(x), \text{True})$, and the existential quantification $\forall$ can be extended by $\exists x: \varphi(x) \Rightarrow \max_{x \in D} \text{Sim}(\varphi(x), \text{True})$, where $D$ represents the domain of variable $x$ and $\varphi(x)$ represents a logical formula evaluated over $x$. The min-pool and max-pool can be added to the logic loss to train the $\forall$ and $\exists$, respectively. To extend operations to probabilistic logic, the existing logic similarity metric $\text{Sim}(\varphi, \text{True}) \in [0,1]$ for DMWM can be considered as a probabilistic estimate $P(\varphi)$ of the truth value of a logical proposition $\varphi$. In this way, the probabilistic logic can be incorporated by introducing soft constraints of the form $P(\varphi) \geq \tau$, where $\tau$ is a required confidence threshold. These constraints can be integrated into the logic loss through variational objectives that reflect the probabilistic consistency of imagination.

We are grateful to Reviewer 2YL1 for spending time on reviewing our manuscript and raising many valuable improvement suggestions. Our detailed responses are given as follows.

W1: Heavy reliance on manually defined logic rules.

A: Thank you for this feedback. First of all, our method is general since the logical thinking ability of the System 2 component can be obtained by a learning-based method instead of hand-crafted rules. In particular, the System 2 component learns logics from the actual environment changes with one-step logic $s_t \land a_t \rightarrow s_{t+1}$ and multi-step logical chain $s_{t-\alpha} \land a_{t-\alpha} \cdots \land s_t \land a_t \rightarrow s_{t+1}$. This learning-based method is applicable to any environment and any task without reliance on known and predefined logical rules. To demonstrate the generalization and scalability of the proposed method, we have added two experiments that extend DMWM to multi-task environments (scalability) and few-shot learning for new tasks (generalization) (Please see our response to W1 of Reviewer Y1Jx). Moreover, we provide more analysis of when DMWM performs well/not well (Please see our response to W2 of Reviewer Jdzu).

W2: The dual-module design, including Kronecker-product feature interactions and recursive logical rollouts, adds significant architectural and computational complexity. The paper does not explore how well this framework scales to high-dimensional or more open-ended tasks where symbolic abstraction may be harder to define. It can be hard to scale this work up to a modern scale.

A: While the integration of Kronecker product-based feature interactions and recursive logical reasoning adds architectural complexity, these operations are confined to compact logical embeddings rather than raw high-dimensional spaces. Moreover, System 2 operates at an abstract symbolic level that scales with task logic rather than input dimensionality. Hence, the architectural and computational complexity of DMWM is controllable and moderate. Moreover, as mentioned in response to W1, the logical rules learned by the System 2 component are general and can be extended to an open-ended environment while keeping the scalability for specific rules.

W3: Authors should provide more quantitative and qualitative comparisons to the line of work mentioned in Appendix F2 (complex task planning [71, 72, 25, 73, 26, 74, 75]), including Logic Tensor Networks. Also, LS‑Imagine (Li et al., ICLR 2025) and other jumpy state-space methods offer similar long-range planning benefits.

A: We are grateful for this feedback. In the revised manuscript, we have first added more detailed introductions of S‑Imagine (Li et al., ICLR 2025) and other jumpy state-space methods as the important complex planning methods in Appendix F2. Regarding to your concerns about quantitative and qualitative comparisons with [71, 72, 25, 73, 26, 74, 75], one of our main contributions is to propose a general dual-mind world model, in which we can replace LINN with any logical method for the System 2 component. Hence, the choice of the best logical inference method is not our research focus. The jumpy state-space methods are important to realize the long-term planning, especially the long short-term imagination proposed in (Li et al., ICLR 2025). For control tasks that require accurate single-step actions, the quantitative comparison of DMWM and LS-imagine via environment trails in 4 representative DMC tasks is shown as follows, where environment trials represent the number of environment exploration opportunities.

The simulation results demonstrate that the proposed DMWM has better exploration efficiency compared to LS-imagine since LS-image is still a pattern-driven Intuitive approach, i.e., System 1 defined in our paper. Next, the qualitative comparison is given as follows. DMWM and LS‑Imagine both aim to improve long-horizon imagination, with different paradigms. LS‑Imagine relies on jumpy transitions and affordance maps to extend the imagination horizon in visually grounded environments. It enables efficient exploration through value-guided behaviors for an open world. In contrast, DMWM introduces a dual-mind architecture inspired by human cognition, which combines intuitive RSSM-based modeling with a logic-integrated neural network that enforces symbolic reasoning and consistency over an extended horizon. While LS‑Imagine performs spatially-aware exploration using visual semantics, DMWM provides interpretable and semantically coherent predictions through structured logical inference.

Q1: Regarding the use of the Kronecker product in LINN-S2, this work applies Kronecker product to combine the latent state and action vectors before logic processing. What motivated this choice over more conventional fusion strategies like concatenation or bilinear pooling? Did you experiment with alternatives, and how does the Kronecker product impact downstream logical reasoning or model stability?

A: We agree that the motivation needs a more detailed discussion. We must first point out that we had a typo in the manuscript. In particular, $\otimes$ represents outer product instead of Kronecker product since the state embedding $v \in \mathbb{R}^d$ and the action embedding $m \in \mathbb{R}^d$ are vectors. In other words, for vectors, $\text{OP}(v,m) = \text{KP}(v,m).\text{reshape}(d,d) \in \mathbb{R}^{d \times d}$, where $\text{OP}$ is the outer product operator and $\text{KP}$ is the outer product operator. In this context, our method of the LINN-S2 component that uses outer product and the convolutional neural network (CNN) to capture the cross-space logical relationships is similar to bilinear pooling [1], which also captures multiplicative interactions between feature pairs. Moreover, in the appendix, to verify the stability and logical consistency of the outer product, we have added an ablation study to demonstrate the superior performance of the outer product compared to concatenation (Please see our response to W1 of Reviewer YoYe). The simulation results demonstrate that the outer product of DMWM can improve logical consistency compared to the concatenation method across all tasks. This is because, compared to the concatenation method that overlooks structural dependencies, the outer product captures second-order interactions between latent state and action embeddings for cross-space logic modeling. Hence, the cross-space representation by outer product and CNN can enable the System 2 component to express more complex logical dependencies essential for long-horizon imagination. As demonstrated in the experiment results, this design is effective and stable for logic processing in our framework.

References: [1] Kong S, et al. Low-rank bilinear pooling for fine-grained classification. Proc IEEE CVPR, 2017: 365–374.

Q2: Did you explore how DMWM performs when some logic rules are systematically incorrect, redundant, or even contradictory?

A: Yes. As our response to W1, the System 2 component is a learning-based method to capture logic rules from the environment changes. Hence, if there are systematically incorrect, redundant, or contradictory logic rules, it means that the observations from the environment are noisy or extremely stochastic. In this context, we can incorporate meta-learning or continual-learning mechanisms to dynamically update logic rules or introduce probabilistic logic to reflect the probabilistic consistency of trajectories. We have provided more details and potential solutions for this discussion in the appendix (Please see the response to W2 of Reviewer Jdzu).

Q3: Why apply logical regularization on joint embeddings rather than over predicted transitions? Could you elaborate on why you chose to apply logic at this representational level, and whether other points of intervention (e.g., such as over pairs $\boldsymbol{(s_t, s_{t+1})}$ or full imagined trajectories) were considered or tested?

A: To clarify, the logic learning can be divided into two parts: one is logical rules from actual environment trajectories, and the other is logical regularization from basic logical laws. In particular, the purpose of logical regularization is to ensure that the neural modules implement logical operations (i.e., $\text{AND}$, $\text{OR}$, and $\text{NOT}$) following basic logical laws such as identity, annihilator and idempotence, as shown in Appendix G. The logical regularizers are applied over individual vectors $w$ rather than over joint embeddings. Theoretically, $w$ can be any vector, since logical regularization only aims to ensure that the logical networks behave consistently under logical transformations, regardless of the specific origin of input vectors. Hence, logical operations require semantically structured inputs, while the raw state lacks learned logic structures. Moreover, embedding both actions and states into a shared logic space enables compositional reasoning across spaces, which can be explicitly optimized for symbolic logical rules.

We are grateful to Reviewer YoYe for spending time on reviewing our manuscript and raising many valuable improvement suggestions. In the response and our updated manuscript, we have made every effort to answer and address your insightful concerns. Our detailed responses are given as follows.

W1: As discussed in the limitation section, the method relies on hand-crafted rules, which are difficult to scale or fully cover in more complex environments.

A: Thank you for this feedback. Our proposed method is general and scalable since the logical thinking ability of the System 2 component can be obtained by a learning-based method instead of hand-crafted rules. In particular, the System 2 component is designed to learn logical rules from the actual environment trajectories $s_t \land a_t \rightarrow s_{t+1}$. This method is applicable to any environment and any task without reliance on prior knowledge of logic. With regards to our claim regarding the reliance on domain-specific logic, this relates to the fact that action $a_t$ in $s_t \land a_t \rightarrow s_{t+1}$ follows the specific task $g$, i.e., $a \in \mathcal{A}_g \subsetneq \mathcal{A}$, and it leads to $s \in \mathcal{S}_g \subsetneq \mathcal{S}$, where $S$ is the state space of the open world, $\mathcal{A}$ is the action space of the agent, and $\mathcal{S}_g$ and $\mathcal{A}_g$ are respectively the subset of the state space and the action space with the task $g$. In this context, we can use a goal-conditioned method, similar to TD-MPC2 [1], to overcome this particular limitation and scale to different tasks. To demonstrate the generalization and scalability of the proposed method, we have added two experiments that extend DMWM to multi-task environments (scalability) and few-shot learning for new tasks (generalization) (Please see our response to W1 of Reviewer Y1Jx). The simulation results demonstrate that DMWM can achieve a performance improvement of 8% on DMC tasks, 30% on ManiSkill2, and 6% on MyoSuite compared to TD-MPC2 at 1M environment steps by using the same hyperparameters across all tasks. Moreover, we have added one paragraph that discusses, in detail, the limitations of our method in the appendix for clarifying these points. In a nutshell, our approach can be general and scalable for complex environments.

References:

[1] N. Hansen, H. Su, and X. Wang. Td-mpc2: Scalable, robust world models for continuous control. International Conference on Learning Representations, 2024.

W2: The logical rules considered are restricted to propositional logic. It's nontrivial to extend this approach to more expressive logic with existential quantifiers and uncertainty.

A: We have adopted the propositional logic system since it is enough to capture the general environment logical rules by $s_t \land a_t \rightarrow s_{t+1}$. However, DMWM can also be extended to existential quantifiers ($\forall$ and $\exists$) and probabilistic logic ($\mathbb{P}(\varphi)$). In particular, the universal quantification $\forall$ can be extended by $\forall x: \varphi(x) \Rightarrow \min_{x \in D} \text{Sim}(\varphi(x), \text{True})$, and the existential quantification $\forall$ can be extended by $\exists x: \varphi(x) \Rightarrow \max_{x \in D} \text{Sim}(\varphi(x), \text{True})$, where $D$ represents the domain of variable $x$ and $\varphi(x)$ represents a logical formula evaluated over $x$. The min-pool and max-pool can be added to the logic loss to train $\forall$ and $\exists$, respectively. To enable DMWM with logic uncertainty, the existing logic similarity metric $\text{Sim}(\varphi, \text{True}) \in [0,1]$ in (12) can be considered as a probabilistic estimate $\mathbb{P}(\varphi)$ of the truth value of a logical proposition $\varphi$. In this way, probabilistic logic can be incorporated by introducing soft constraints of the form $\mathbb{P}(\varphi) \geq \tau$, where $\tau$ is a required confidence threshold. These constraints can be integrated into the logic loss through regularization terms or logic-enhanced variational objectives that reflect the probabilistic consistency of imagined trajectories.

W3: The authors claim that their application of Kronecker product can facilitate action-state alignment, but there lacks an ablation study to support this claim.

A: We agree with the need for an ablation study. We must first point out that we had a typo in the manuscript. In particular, $\otimes$ represents the outer product rather than the Kronecker product since the state embedding $v \in \mathbb{R}^d$ and the action embedding $m \in \mathbb{R}^d$ are vectors. In other words, for vectors, $\text{OP}(v,m) = \text{KP}(v,m).\text{reshape}(d,d) \in \mathbb{R}^{d \times d}$, where $\text{OP}$ is the outer product operator, and $\text{KP}$ is the Kronecker product operator. Following your suggestion, in the updated appendix, we have added an ablation study to support the advantages of the outer product for action-state alignment compared to the concatenation method as follows.

The simulation results demonstrate that the outer product of DMWM can improve logical consistency compared to the concatenation method across all tasks. In particular, the outer product-based method achieves 3.9% improvements on DMC Tasks, 4.2% on ManiSkill2, and 5.1% on MyoSuite in terms of logical consistency. This is because, compared to the concatenation method that overlooks structural dependencies, the outer product can capture second-order interactions between latent state and action embeddings, which enables cross-space logic modeling.

Q: How did you measure the logical consistency for evaluation? Did you use the trained system II?

A: Yes, we use the trained System 2 components to evaluate and verify the logical consistency across different approaches during the test stage. In particular, for a fair comparison, all the LINN-S2 components are trained with the same network architecture, parameter scale, and training steps. For the baselines, including Dreamer, Hieros, and HRSSM, the LINN-S2 components only capture and verify the logic of state transitions and will not be used to impose the logical constraints on the System 1 components during training.