Neurosymbolic World Models for Sequential Decision Making

We thank Reviewer pyge for their thoughtful comments and are committed to incorporating your feedback into the manuscript.

Note: Since our submission, we have demonstrated stronger RL performance (see response to B4Kd).

Please excuse our brevity due to the character limit.

Mode Identifiability Assumption

Please refer to the response to Reviewer bZQ6 ("Mode Identifiability Assumption").

Mode Sharing and Diverse Environments

In our original submission, we stated a stronger assumption than needed. Assumption 4.5 should instead be:

“The modes in the active POMDP must be a subset of the modes in the offline POMDPs.”

E.g, consider two offline POMDPs: one with modes A and B and one with modes C and D; the active POMDP could consist of modes A and C which we argue is not very restrictive. We will update our manuscript to reflect this correction.

We also note that the Salamander environment already exhibits high diversity, as it is a 3D simulation of a real robotic platform [5, 6] with computational fluid dynamics and rigid-body contacts.

Real-World Applicability

Please see response to bZQ6 ("Real-World Applicability").

Framework Implementation Challenging

To help the community build on our results:

• We open-source our source code
• The pseudocode typeset in the manuscript closely follows our implementation
• As mentioned, we leverage off-the-shelf, well-established implementations of key algorithms

Computational Complexity

Let $n$ be the number of experiences $(o_t, a_t, o_{t+1})$ of dimension $d$, $m$ the number of modes. The computational complexity of training is driven by these steps:

1. Solving Eq. 1: To solve Eq. 1 we use SGD, which is $O(G_1 n d K_1)$, where $G_1$ is the number of SGD steps, and $K_1$ accounts for the architecture of the model (see [9] for discussion). The variables corresponding to the amount of data, the number of iterations and the model complexity needed for training, are not independent. E.g., a complex system (e.g., high number of modes) might need more data. However, a monolithic model would also require more data because it too must implicitly learn the same complex dynamics. Thus, we do not expect this step to present significant more overhead than a monolithic approach.

2. Training the $m$ neural primitives: We assume datasets will have in expectation $O(\frac{n}{m})$ elements. If models are trained sequentially, this is $O(m G_2 \frac{n}{m} d K_2)$, where $G_2$ and $K_2$ account for the number of SGD steps and the architecture of the primitives. The caveats in 1. apply, but each step is at most as expensive as training a monolithic model. Primitives can be trained in parallel if wall-clock performance is critical.

3. Transition Predicate Synthesis: This step trains $2m^2$ decision trees with CART [12]. For the average case, we assume each dataset will be of size $O(\frac{n}{m})$. Therefore, the overall cost in expectation is $O(m^3 d \frac{n}{m} \log^2(\frac{n}{m}))$.

Applying dimensionality reduction techniques may help with high-dimensional spaces. In some high-dimensional domains (e.g., visual domains, which are not the focus of our work), we would expect a different class of models (e.g., CNNs) to be more effective.

Simplifying, the average complexity of training a SWMPO model is

$O(GnmdK + nmd + nmd\log^2(\frac{n}{m})),$

where $G = \max(G_1, G_2)$ and $K = \max(K_1, K_2)$. Thus, the expected overhead of training in SWMPO compared to a monolithic model is only a linear factor on $m$.

Then, during each model-based RL step, SWMPO evaluates the predicates of the current mode, which are small functions that could be evaluated in parallel, and the active mode's neural network. Thus, the expected run time is only an added small constant factor slower than a traditional monolithic world model.

Pruning Guidelines

We manually tuned $\epsilon$ by inspecting error plots (e.g., Figure 4), but did not find the state machines overly sensitive to it.

As a guideline, we suggest (1) inspecting model error plots to see if there are transitions that should be pruned in the first place, and (2) using error between models as a starting value for $\epsilon$.

Additionally, hyperparameter tuning techniques could automate the process.

Domain Knowledge

While out of scope for this manuscript, we consider incorporating domain knowledge explicitly a promising future direction. For instance:

1. Replace models in Eq. 1 with physics-informed neural networks or other models that enforce physical constraints
2. Integrate syntactic/semantic constraints into the predicate synthesis using systems like Sketch [12]

We hypothesize that each approach would improve the local models and the FSM transition dynamics, respectively.

We note that the assumption that systems can be decomposed into modes is a form of domain knowledge in itself.

References

Please refer to the response to Reviewer bZQ6.

We thank reviewer B4Kd for their thoughtful comments and appreciate your recognition of SWMPO’s extension of FSM-based modeling to non-linear continuous domains and of its practical competitiveness. We are committed to incorporating your feedback into the final version of the manuscript.

Note: Since the results in our initial submission, we have further demonstrated a stronger performance of our approach over baseline methods in the RL experiment (see below).

Performance over Monolithic Models

We understand your interest in a stronger empirical performance gap between SWMPO and the monolithic model-based RL baseline. We decided to submit our manuscript with early results, but since the submission we have been able to make some small improvements to the algorithm that in turn have resulted in significant improvement in performance. We also spent comparable time improving the performance of the baseline.

Specifically, this is what we have done for both SWMPO and the monolithic model-based RL baseline:

1. We ran the RL training 64 times, instead of the 16 times in our original experiment.
2. We further tuned the hyperparameters of all algorithms.
3. We applied a simple modification to the model-based rollout logic for both SWMPO and the baseline: instead of relying on the models for long trajectories (150 timesteps) ---which can lead to compounding errors--- we sampled a random observation from an environment trajectory and performed a shorter (30-step) model rollout. This reduces reliance on the model for long-horizon prediction, improving sample efficiency and stability for both methods. Note that the performance of the low-level neural models for long horizons is fully orthogonal to our claims, and the model-based RL community has mechanisms to scale neural forecasting performance. We will add this information to the final manuscript, along with our updated results.

These adjustments increase the performance of both SWMPO and the model-based baseline. However, SWMPO sees a greater gain, resulting in a larger performance gap between SWMPO and the monolithic model of approximately ~40% (~21 mean reward for SWMPO vs ~15 mean reward for vanilla model-based RL by timestep 75,000), a significant improvement over the much smaller gap in our original submission . Furthermore, SWMPO now reaches this performance level within 75,000 timesteps, compared to 200,000 steps in our original submission. We stopped at 75,000 timesteps due to rebuttal time and compute constraints; however, we will include the experiment up to 200,000 steps in our revised manuscript.

We encourage reviewers to check the updated results here (corresponding to an updated Figure 7): https://anonymous.5open.science/r/anonymousfig-30AA/improvedrewards.png

We summarize these updated results in the table below:

Extreme Case Scenario

We hope that our new results above demonstrating improved performance over monolithic models renders the use of an “extreme case” where SWMPO would outperform unnecessary. Our terrain-mass environment is already a scenario that closely matches the assumptions of the framework. We did not construct a more extreme synthetic case, as the updated results already show a significant performance gap in our original setting --- without needing to simplify the task.

Additional References

We appreciate the reviewer's suggestion of additional references, and will include a discussion of all of them in the final manuscript.

We sincerely thank Reviewer bZQ6 for their thoughtful comments, and we are committed to incorporating your feedback into the final manuscript.

Note: Since our initial submission, we have demonstrated stronger RL performance of SWMPO over baselines (see response to B4Kd).

Please excuse our brevity due to the character limit.

Scalability

We understand your concern about the scalability of SWMPO in environments with a high number of modes or high dimensionality.

The computational complexity of SWMPO is driven by three steps: (1) solving Eq. 1, (2) training of neural primitives, (3) synthesizing predicates.

It can be shown that, in expectation, the overhead of training a SWMPO model is a constant factor dependent on the number of modes (training the end-to-end model while solving Eq. 1 is asymptotically equivalent to training a standard monolithic model). Additionally, this can be amortized by training the neural primitives in parallel if wall-clock performance is critical. We argue that this overhead is not prohibitive.

For more details, please refer to our response to pyge (”Computational Complexity").

Additionally, the “forward pass” of the SWMPO model consists only of evaluating the transition predicates of the current mode (small functions that can be evaluated in parallel if needed) and the active mode's neural network. Thus, runtime overhead is only a small constant factor slower than a standard monolithic world model.

Mode Identifiability Assumption

As part of our follow-up work, we are working on generalizing the framework to cases where mode inference may require longer temporal context or probabilistic reasoning. However, we believe that for many useful applications this assumption holds. Indeed, many state-of-the-art systems continue to assume fully Markovian dynamics (e.g., Bhatt et al., ICLR 2024 [7]; Kuznetsov et al., ICML 2020 [8]). We argue that this assumption does not preclude the application of our method to useful problems.

Real-World Applicability

Our evaluation focuses on challenging but controlled simulated environments, which allow for systematic study of the components of SWMPO. These environments are aligned with recent work in the field, which are often benchmarked on systems of similar or lower complexity---e.g., simulated Cartpole in [0] (ICML 2020), simulated low-dimensional three-mode systems in [3] (NeurIPS 2021), and simulated grid-world environments in [4] (NeurIPS 2024). We thus argue that simulation is a valuable setting for benchmarking novel structure learning frameworks.

Moreover, we note that the Salamander environment already features high complexity, including 3D locomotion of a simulated real robotic platform [5, 6], rigid-body dynamics, and computational fluid dynamics to simulate water. We believe this makes it a strong intermediate testbed bridging controlled and real-world scenarios.

Nonetheless, we acknowledge the limitations of not including physical robotic platforms. While real-world deployment introduces important challenges, we believe it is important to first validate the core elements of the framework in simulation. Furthermore, there is precedent for transferring small learned automata to hardware: Liu et al. (preprint, 2025) [1] demonstrate a three-mode FSM controlling a quadruped robot. We do not claim that our current results are directly applicable to real-world systems, but we strongly believe that SWMPO is a promising foundation for future real-world deployment.

Writing and Figure

We will incorporate your feedback into the final manuscript by modifying the teaser figure to highlight environment modes. Primitives are presented early in the introduction (e.g., paragraphs 2–5), but we will add more detail to their description.

References

[0] Zhang et al., Invariant Causal Prediction for Block MDPs, ICML 2020
[1] Liu et al., Discrete-Time Hybrid Automata Learning: Legged Locomotion Meets Skateboarding, arXiv preprint, 2025
[3] Poli et al., Neural Hybrid Automata: Learning Dynamics with Multiple Modes and Stochastic Transitions, NeurIPS 2021
[4] WorldCoder: a Model-Based LLM Agent, NeurIPS 2024
[5] https://www.epfl.ch/labs/biorob/research/amphibious/salamandra/
[6] https://www.cyberbotics.com/doc/guide/salamander?version=R2021a#salamander-wbt
[7] Bhatt et al., CrossQ: Batch Normalization in Deep Reinforcement Learning for Greater Sample Efficiency and Simplicity, ICLR 2024
[8] Kuznetsov et al., Controlling Overestimation Bias with Truncated Mixture of Continuous Distributional Quantile Critics, ICML 2020
[9] Goodfellow, I., Bengio, Y., & Courville, A. (2016). Deep Learning. MIT Press
[10] David Arthur, Sergei Vassilvitskii, How slow is the k-means method?, SCG 2006
[11] Klusowski & Tian, Large Scale Prediction with Decision Trees, Journal of the American Statistical Association, 2023
[12] Solar-Lezama, The Sketching Approach to Program Synthesis, APLAS 2009