Neurosymbolic World Models for Sequential Decision Making
Authors: Leonardo Hernandez Cano
Maxine Perroni-Scharf
Neil Dhir
Arun Ramamurthy
Armando Solar-Lezama
Venue: ICML 2025 poster (2025)
Keywords: world models
finite state machines
structure learning
model-based reinforcement learning
neurosymbolic
Primary area: general_machine_learning->unsupervised_and_semisupervised_learning
TL;DR: We synthesize state machines from low-level continuous observations and apply them to policy optimization

Abstract: We present Structured World Modeling for Policy Optimization (SWMPO), a framework for unsupervised learning of neurosymbolic Finite State Machines (FSM) that capture environmental structure for policy optimization. Traditional unsupervised world modeling methods rely on unstructured representations, such as neural networks, that do not explicitly represent high-level patterns within the system (e.g., patterns in the dynamics of regions such as \emph{water} and \emph{land}).
Instead, SWMPO models the environment as a finite state machine (FSM), where each state corresponds to a specific region with distinct dynamics. This structured representation can then be leveraged for tasks like policy optimization. Previous works that synthesize FSMs for this purpose have been limited to discrete spaces, not continuous spaces. Instead, our proposed FSM synthesis algorithm operates in an unsupervised manner, leveraging low-level features from unprocessed, non-visual data, making it adaptable across various domains. 
The synthesized FSM models are expressive enough to be used in a model-based Reinforcement Learning scheme that leverages offline data to efficiently synthesize environment-specific world models.
We demonstrate the advantages of SWMPO by benchmarking its environment modeling capabilities in simulated environments.

Neurosymbolic World Models for Sequential Decision Making
Leonardo Hernández Cano
Maxine Perroni-Scharf'
Neil Dhir? Arun Ramamurthy
Armando Solar-Lezama
Abstract
Reward
Initial state
Trajector
Emv modes water,
We present Structured World Modeling for Policy Optimization (SWMPO), a framework for unsu- pervised learning of neurosymbolic Finite State Machines (FSM) that capture environmental struc- ture for policy optimization. SWMPO models the environment as a FSM, where each state cor- responds to a specific region of the state space with distinct dynamics (e.g., water and land). This structured representation can be leveraged for tasks like policy optimization. Our proposed FSM synthesis algorithm operates in unsuper- vised manner, leveraging low-level features from unprocessed, non-visual data to learn non-linear models, making it adaptable across various do- mains. The synthesized FSM models are expres- sive enough to be used in a model-based Rein- forcement Learning scheme that leverages offline data to efficiently synthesize environment--pecifee world models. We demonstrate the advantages of SWMPO by benchmarking its environment mod- cling capabilities in number of simulation tasks.
POMOP
Offline data
Collect data
NeuratPrinitives()
FSwSynth()
(1) Neural world- model primitives
naurosymbalic
speciipccific
Model-based
Control
world-model
RL
policy
Figure I. Proposed codicy-optimization framework: (1) data is col- lected from the input POMDP by interacting with the environment in standard RL loop; (2) world-model specific to the POMDP is synthesized leveraging neural primitives: (3) the world model and the POMDP's reward function used perform model-based optimization. These steps are repeated until a termination con- dition met. The neural primitives are learned from data from previous POMDPs. The output is a policy for the new POMPD.
environments, these modes may appear in different config- urations (e.g., one environment may be a lake surrounded by land, another may be mostly land separated by a small stream). Instead of modeling each environment from scratch, better approach might be to first construct 'primitives' for water and land, and then, given a specific environment, as- semble those two primitives into complete model. This only requires the identification of the boundaries between the two modes relevant to a task, which in principle requires less effort than the construction of the model from scratch.
1. Introduction
This work focuses on learned approximations of environ- ment dynamics- known world models (Ha Schmid- huber, 2018)-in in the special case where these models must explicitly encode the high-level structure of the environment dynamics.
Following this reasoning. our goal is to develop a method that (1) learns supervised manner from offline data and (2) reuses those primitives in the synthesis of a model specific to new en- vironment. This model should be accurate enough to be used for model-based policy-optimizationn enviromments with continuous, low-level and non- non-visual observations (e.g.. LiDAR measurements or joint positions).
of world-model primitives in an
We are motivated by the observation that cruetured-world models may facilitate the reuse of components across dif- ferent environments. For example, consider modeling a LiDAR-based amphibious robot that must navigate both water and hand-whidd
We focus on
call modes. Across different
'Computer Science and Artificial Intelligence Laboratory, Massachusetts Institute Of Technology, Massachusetts, United States of America 'SIEMENS, New Jersey, United States of America. Correspondence to: Leonardo Hernández Cano Elcohe@mit.cdux..
To this end, we propose Structured World Modeling for Pol- icy Optimization (SWMPO), neurosymbolic framework for model-based Reinforcement Learning (RL). SWMPO first uses offline data to learn a set of neural world-model primitives (e.g., water and land). SWMPO then assembles these primitives into a complete world model specific to new environment. Finally, SWMPO uses this model for
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Neurosymbolic World Models for Sequential Decision Making
model-based policy optimization (see Figure 1).
o++1 the source and next observations, respectively.
Definition 2.3 (Finite State Machine World Model), We define a finite state machine world model (FSMWM) to be a tuple F & (F,, O, 6, fo), where O and A are re- spectively the observation and action spaces of a POMDP: F (fs O x A - O) is a set of models ""mode""o---"" of the environment; and ė {6a. O x A + {0, 1)),, is set of mode-transition predicates. The initial mode is fo-
Internally, SWMPO models environment with a neu- rosymbolic FSM. The synthesized FSM consists of world- model primitives and transitions between them. Each primi- tive is a neural network that approximates the environment dynamics within a subset of the state space, The transitions are logical rules that determine when to switch between primitives based on observations of the environment. The modular structure of an FSM allows our system to reuse set of neural primitives to efficiently synthesize environment- specific FSMs, as then learning the boundaries between the modes is sufficient to assemble a complete world-model.
For a given active mode indexed by i and an observation Og E O, the predicted next observation is o++1 = Se(or, oz). The next active mode index is
argmax, 01,5(11, u/) if 6i, u4) > 0
We evaluate components of SWMPO across various bench- marks and environments with continuous dynamics, includ- ing two and three -dimensional simulations.
5 ue i)
otherwise.
We define argmax to choose the first matching index in case of tie. This definition mirrors previous use of FSMs as policies in the MDP setting (Inala et al., 2020), but here we are using them as world models.
Contributions Our contributions are follows:
1. An unsupervised learning algorithm that trains set of modular world-model primitives (e.g., water and land models) from a dataset of low-level continuous observations.
Given an observation and policy, the FSMWM can simu- late a rollout by using the local models to predict the next observation and transitioning between the models accord- ingly. Unlike a POMDP, the FSMWM operates purely on observations. For conciseness, we write FSM shorthand for FSMWM.
2. A state-machine synthesis algorithm that leverages set of primitives learned from a collection of environ- ments and new data from a specific environment to construct an environment-specifee FSM world model. 3. An evaluation that demonstrates these primitives and the FSMs can be used to perform model-based RL in a meta-learning setup.
3. Problem statement
We want to synthesize a world model whose structure cor- responds to the high-level structure of POMDP, and then leverage this model within RL training loop. With this in mind, we restate our goals as follows:
2. Background
We operate in the standard discrete-time RL framework, where an agent interacts with an environment according to some control policy (Sutton & Barto, 2018). A summary of the notation used in this manuscript can be found in Table in the supplementary material.
Goal 1: Use a dataset of trajectories collected by an agent, possibly across several 'offline' POMDPs, to learn a set of world model primitives.
Goal 2: Use dataset of trajectories from single 'ac- tive' POMDP and set of world model primitives to synthesize an FSMWM of the active POMDP.
Definition 2.1 (Partially Observable Markov Decision Pro- cess). A discrete-time Partially Observable Markov Deci- sion Process (POMDP) is a tuple M é (S, A, So, O, O), where S C is the set of states, .A is a set of actions, T: S x A + S is a transition function, Sy is a distribution of initial states, O is the set of observations the agent can make, each observation O E O made under some state s E S and action a E A with probability ( s, a). We associate POMDP with a reward function R: O A O R.
These tasks must be performed in unsupervised manner (i.e., without mode labels in the input data). Our proposed method is described in Section 4.
Our core assumption is that each of the POMDPs of interest has an (unknown) FSMWM that models it. This implies the existence of a POMDP-specifi latent categorical variable Mr, whose value corresponds to the index of the mode of the system at time t f- it further implies that M, can be characterized by a function (o--1, 4t114411,00 (0:-1,4t-15 0t) - m,, We also assume that the number of modes is known in advance. Further assumptions are motivated and stated in Section 4.1 and Section 4.4. We may use examples of what modes may represent for clarity (e.g., m, = water instead of m, = 1).
For given policy To O - A, we use St, O2 and A, to denote the random variables of the state and action at time respectively. We use lowercase letters to denote values of a random variable.
Definition 2.2 (Trajectory). A trajectory is a iime-indexed sequence of transition tuples (ot, az, 0)+1). We call og and
2

---

Neurosymbolic World Models for Sequential Decision Making
Prediction error
0.4
Transition predicate synthesis: Learn when to transition between modes (Section 4.3).
Mode
Mode
0.2
4.1. Neural world-model primitives
0.0
We first address the problem of decomposing enviromment dynamics by assigning each transition in dataset D of trajectories to one of n disjoint subsets, with each subset corresponding to a mode of the POMDP. We assume that n is known in advance. This is standard task in time-series analysis, although here we focus specifically on discovering the (unobserved) categorical mode variable M, of a system.
Ground truth
Min-loss
20
80
100
Timestep
Figure 2. Prediction error of the learned modes (top) and corre- sponding mode classification (bottom) for test trajectory from Point Mass - an idealized amphibious robot that travels on land (pink) and water (blue) (see Section 6). Each local model is special- ized in a specific mode, resulting in low prediction error throughout the entire episode.
We focus on the case where observations conditioned on state are deterministic, so o, O(s)). Our method revolves around learning Mr as an intermediate computation of a learned first-order model of the form
foz (me, (0:(((0)))))) (0:-1, ::,1; oz), Or, a4) R 0++1 Ot,
4. Structured world models
where fa, and ma, are functions parametrized by B,, 0, E R In principle, fe, could be a standard forward model, but in preliminary experiments we observed that first order model yields better performance. We hypothesize this hap- pens because first-order model weakens the dependence on Or-
This section describes our proposed FSM synthesis algo- rithm. Our method receives a dataset of trajectories and synthesizes an FSM to model the structure of the environ- ment (see Figure 3), where each state in the FSM represents a distinct mode of the data (e.g. water or land).
Ultimately, we characterize modes as a categorical vari- able (i.e., the robot is either on water or land), but here m&, (o:-1, oz) E Ra is a continuous approximation of the true latent m4,
Twater
Land
Data collection
FSM synthesis
(unsupervided)
We now describe a series of assumptions about the POMDP and the mode variable that allows us to design an algorithm to predict Mc.
Figure 3. An illustration of the proposed neurosymbolic world- model synthesis method in an environment comprising land and water. First, a controller is used to gather data (left). Then, this data labeled according to modes of the system in an unsupervised fashion (middle). These labels are then used to supervise the synthesis of a neural state machine world-mode (right), that can be leveraged for policy optimization.
Assumption 4.1 (Mode Identifiability). We assume that Mt can be modeled as a function m, R m((0-1, (100-,,0;1111 ot). Intuitively, this means that the current mode can be identified by observing how the world changes in response to the latest action. Thus, we interpret M, as an abstraction of the observed system change given particular action and state.
A key challenge is discovering these modes in an unsuper- vised manner, where only the number of modes is known. is important to achieve clean separation between modes at different stages of the algorithm to avoid cascading errors from misclassified data, which can progressively degrade the model's is performance. To address these challenges, we synthesize our state machine follows:
Assumption 4.2 (Predictability of Change Conditioned on Mode), We assume that the POMDP reduces to a determin- istic MDP when conditioned on the mode. More precisely, we assume the existence of a function T" M x such that T" (m,, 0t, o.) O (st, az)).
Thus far, our constraints allow the trivial solution M, S., which is not useful. There may be many other variables which satisfy our assumptions. Consequently, we add an- other assumption that allows us to uniquely identify Mc;
Neural world-model primitives: Divide the transitions in a given dataset into different mode subsets and train neural primitives (Section 4.)..
Assumption 4.3 (Minimality of Modes). Let M be the set of random variables that satisfy the previous assumptions, Then, the mode variable M, is the unique solution
Pruning:
dynamics of the par- tition by removing spurious transitions between modes (Section 4.2).
Simplify the
M = arg min T(Me, Ot+1);
M.E
3

---

Neurosymbolic World Models for Sequential Decision Making
where I(:))) stands for the mutual information between two random variables.
Algorithm 1 NeuralPrimitives
Require:
List of trajectories D, latent space M, number of modes m, and, observation and action spaces O and A respectively.
Under the assumptions so far, we can conclude that if we find a variable Mr that allows us to predict the change in the environment given an action and has minimal mutual information with Oz+1, then M, must be the mode variable. However, our approximation of M, takes values in a vector space, but we ulltimatly want to model it as a categorical variable. We therefore add our last assumption:
I: Use gradient-descent to find neural networks So, : M x + M that
A x o + O and me. O x A
O
approximately solve Equation (1).
Embed the transitions into mode vectors, m(D).
3: Cluster the embeddings into m disjoint subsets using k-means.
4: Let D; + {T D cluster(r) = i].
Assumption 4.4 (Mode Vectors Form Clusters). The mode variable M, corresponds to partition of the state space, where, in expectation, the within-subsee sum of squares to the centroid of each subset is minimized. That is, we assume a strict partitioning and a centroid model clustering scheme. Consequently, if k-means is applied to vectors of M,, the resulting clusters will, in expectation, correspond the different modes.
5: Assemble partition D &- --i,...., Dm
6: Reorder D so that D1 is the set with the most initial states.
Let f. < arg orgminn B Blo. where 0,41 + -01, 0t).
o+1lll
8: Assemble the set of local models F + fm]
9: return F
Together, these assumptions imply that if a variable M, is predictive of the change in observed state for any action and has minimal mutual information with Ot+1, then that variable corresponds to the mode variable. In other words, consider So,- M x A x O + O and mæ, : O x A x O + .M under the joint optimization problem deseribed in Equation (1) therein || || is the Euclidean norm, or-1 ((s-1) and o2 = (st). From the assumptions stated above, it follows that the solution to Equation (1) is such that that (.) corresponds to the mode variable, which can be clustered with k-means to obtain mode labels for a set of transitions. In practice, we parametrize m(.) and f()) with neural networks and approximate the solution of Equation (1) with gradient-based search. To compute the mutual-ine nntuulliiffrrmatio I(: I(::-). we assume independence of features and fit Gaussian distributions to compute Monte- Carlo approximation.
4.2. Pruning
The aforementioned partitioning process can create overly complex transitions. While the FSM globally approximates the environment dynamics, some state regions may have multiple models with similar accuracy, resulting in spurious transitions between states. In such cases, simplifying transi- tion dynamics with minimal impact on forecasting accuracy might be possible.
To address this, we propose pruning mechanism to elim- inate these unwanted transitions. This helps balance the complexity-
trade-off in the state machine search space: while more complex transition patterns can improve accuracy, they also increase the risk of overfitting.
Pruning Approach The system begins by labeling each transition in the dataset with the index of the neural network from the ensemble that best predicts the system's 's evolution in that state. A mode transition occurs when this label changes between consecutive states. For example, in the sequence 1 113322, we transition from mode 1 to 3, then from 3 to 2. Pruning the transition to mode 3 yields two possible sequences: 111122 (forward-prums or 112222 (backward-prune))
Fitting local models to the data Once the approximation of the latent mode variable, mo,, is trained, the system uses the k-means algorithm to cluster the induced set of latent vectors into disjoint subsets. Then, for each subset, the system fits a local forward model to model the correspond- ing transitions, as illustrated in Figure 2, where each local model has higher performance for a particular mode of the environment.
To prune a mode transition, the system relabels the transi- tion so that a different neural network handles it (c.g.. we relabel transition originally marked land to water). If the increase in prediction error in the transition is within the user-defined tolerance factor e, the move is considered e-valid relative to the original partition. A mode transition is c-prunable if all the associated moves are c-valid. Note that there may be multiple e-prunable mode transitions for a given trajectory and partition. The system greedily prunes
The entire process is shown in Algorithm 1. The algorithm takes dataset of trajectories D, and partitions it by solving Equation (1) and clustering the resulting mode vectors. The resulting partition induces a set of neural primitives, each specialized in modeling the dynamics of a particular mode of the environment.
4

---

Neurosymbolic World Models for Sequential Decision Making
arg min B fo, (00, (ot+1 oz)1l + 0,,02
(1)
the first prunable mode transitions with the strategy that re- sults in the smallest prediction error increase, until there are no more e-prunable transitions (this process is summarized in Algorithm 4 in the supplementary material).
of D; containing transitions where the next state is a source state in Dj. referred to the 'positive' set. The 'negative' set is the complement of the positive set with respect to D;. The task is reduced to standard classification problem, where the goal is to find predicate that outputs True for the positive set and False for the negative set. We use wetteesiiiiieeeeeeeeeeee .) to denote the routine that solves this supervised learning problem. our experiments, we use an off-the-shelf tree-learning algorithm.
0.20
Prediction error
Mode 0 Mode Mode
0.15
0.10
Full FSM
0.05
Once transition predicates are synthesized, the FSM has been defined. See Algorithm 2 for the entire FSM synthe- sis algorithm. The computational complexity of training a SWMPO model is linearly slower compared to a stan- dard monolithic world model and depends on the number of modes; see Appendix for further discussion.
0.00
Unprunned
Prunned
20
40 60 100
120
140
Figure 4, Illustrative comparison of a prunned and an unprunned partition from preliminary experiments. Pruning simplifies the tran- sition dynamics by removing spurious transitions whose removal does not substantially degrade the performance of the resulting FSM. For instance, at timestep 95, model 0 exhibits a lower pre- diction error than model 2. In the unpruned FSM, there is a brief transition to mode 2 and back, which is removed through pruning.
Algorithm FSMSynth
Require: Primitives F (Algorithm 1), dataset of trajecto- ries D, observation and action spaces O and A respec- tively
1: for fi E F do
2: D, + ft D f. = arg min min jer C(f+t)
4.3. Transition predicate synthesis
3: end for
4: for i, E [l......].. x (1.....m) mm do
We now describe the mechanism by which the FSM learns when to transition from one mode to another given a parti- tioned dataset trajectories and corresponding primitives.
positive + Ti E D. 3rz E D, s.t. follows( T2)
6: negative <- D. \ {positive)
To construct a TOMDP-specifi FSM, the system leverages the set of primitives and new POMDP-speciff data to iden- tify the subsets of the POMDP where each primitive should be used. Specifically, for each transition (5t, qg $1+1) ob- served in the new POMDP the FSM must use the primitive which results in the lowest prediction error. That is, given a set of primitives F = [/i...., Jn), the resulting state ma- chine must use fo for transition Ti (ot: de, 0++1) where
7: ōijj + peruttesizeerediiatepppssiii negative))
8: end for
Let <- {ōe. | 1 isijs [m)
10: F - (F,O,A, A,
11: return F
4.4. Neurosymbolic model-based policy optimization
We now describe our approach to integrating the synthesis of FSMWMs into a model-based RL loop. Given a POMDP AS, T, So, S2 O), reward function R, and horizon length E N, our goal is now to find closed-loop policy π O + A that achieves a maximal cumulative reward. That is, the goal is to solve
fi = argg min [|o++1 f(o,, a:)l fEF
(2)
See Figure 2 for an example of how the primitives lead to accurate labeling.
Labeling induces a partitioning of a dataset of transitions. Each subset of partition corresponds to state of the FSM being synthesized. Each pair of FSM states (fr, fz) will be associated with a transition predicate ō.g that corresponds to the following question: given that the state machine was in state fü and the agent observed 8t took action u, and then observed 8t+1. should the FSM transition to state f To synthesize this predicate, the system identifies the subset
π = arg max Egg~ [[(
(3)
HEll
where
R(x) = Σ R(s,, (O(ss))
05t5R
Our algorithm should reuse offline data from similar 'offline' POMDPs, while only using relatively little data specific of
5

---

Neurosymbolic World Models for Sequential Decision Making
"active' POMDP to optimize a policy. This approach aims to reduce the number of agent-environment interactions that must occur within the POMDP to find optimal policy by alternating optimization in the model and in the POMDP. Specifically, the offiine data is leveraged to synthesize a neural set of primitives which will be used in the POMDP- specifie FSM.
5. Related work
Graphical Models Hidden Markov models (HMMs) are standard approach to capture temporal dependencies and mode-switching behavior in sequential data (Li & Biswas, 2002; Bouguila et al., 2022). In robotics, HMMs have been leveraged to segment trajectories into discrete modes (Goh et al., 2012) and used during policy learning for multi- modal or hierarchical tasks (Marturi et al., 2019). Recent advances have extended HMMs using deep neural network architectures (neural HMMs) to handle continuous and high- dimensional observation spaces and to learn more complex transition dynamics (Tran et al., 2016). For example, neural HMMs have been used in unsupervised settings to model complex sensory streams for trajectory clustering (Vakanski et al., 2012) and to predict latent modes during task execu- tion (Wu et al., 2019). However, HMMs suffer from the fundamental limitation that the transition between modes is determined by a probability distribution that is only con- ditioned on the latent state, this means that the observed evolution of the system itself is only indirectly used to up- date the active mode. State of the art graphical models in- clude Hidden Parameter recurrent State Space Models (Shaj et al., 2023), and recurrent Switching Dynamical Systems Models (Glaser et al., 2020). Generalizing these classes of models to include non- non-linear dynamics makes inference computationally expensive (Frigola al., 2013).
At the core of our approach is the following assumption:
Assumption 4.5 (Global Modes), The modes in the ac- tive POMDP must be subset of the modes in the offline POMDPs.
For example, 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. In the example of the am- phibious robot, Assumption 4.5 is satisfied because each POMDP corresponds to a different terrain configuration. Thus, transitions between land and water dynamics change between POMDPs, but the dynamics of each mode are con- stant across all POMDPs.
Thus, given new POMDP and offline data from similar POMDPs, our proposed algorithm proceeds as follows: first, a set of neural primitives is trained from the offline data in an unsupervised manner. Then, a short standard RL loop (observe, act, receive reward) is run over the POMDP to collect data. The new data and the set of primitives are used synthesize POMDP-specific FSM, which is used to perform model-basee RL. These steps are iterated until termination criterion is met. The algorithm is described in Algorithm 3.
Leveraging structure in RL Many prior works focus on leveraging structure to solve control problems with RL (Mohan et al., 2024). Hierarchical RL (Xu & Fekri, 2021; Botviniek, 2012; Li al., 2006) and modular RL (Simpkins & Isbell, 2019; Andreas et al., 2017; Devin et al., 2017) encode structure directly into the policy architecture, here we instead consider the synthesis of a structured model. Model-based RL approaches leverage neural world models to optimize policies more efficiently (Moerland et al., 2023; Ha & Schmidhuber, 2018). However, traditional neural models lack distinct boundaries between the representation of different modes in the environment; this lack of modular- ity makes reuse challenging. Reward machines (Toro Icarte et al., 2019; Icarte et al., 2018) leverage structured models of the reward function to guide the policy optimization pro- cess. Large Language Models have been used to synthesize programmatic world models (Tang et al., 2024). but not in continuous systems with low-level observations. Hasanbeig et al. (2021) demonstrate that FSMs can be synthesized to model discrete environments, improving performance in RL tasks. However, this method is reliant on fully symbolic representations obtained from pre-trained vision models in gord settings. In contrast, our approach extracts struc- ture from continuous, low-level, non- visual observations.
Algorithm 3 SWMPO
Require: POMDP and reward function (M, R). trajectory datasets from previous POMDPs Di...., Dn, obser- vation, and action spaces O and A respectively, la- tent space M, number of modes m. T Tisssionn T Tssioo the length of real- and model-based RL phases, respectively, off-the-shelf RL algorithm RL, step budget Twall 1: Initialize parametric policy Tg.
2: D = Dj U U.. D,,
3: F- = NeuralPrimitives D, M, m. A)
4: DM + +0; 40
5: while S Tiotal do
6: TT. D'M + RL( M, R, T )
7:
t+ +-t+ Tression
8:
DM + DM S D'M
9:
F FSMSYnth(P DM,, A,
10: TB,- + RL((TO, F, R, 7 )
11: end while
12: return n0
6

---

Neurosymbolic World Models for Sequential Decision Making
Structure induction and hybrid systems In the broader field of hybrid systems, modeling environments collec- tion of modes with distinct dynamics is standard practice (Alur et al., 1995; Paoletti et al., 2007; Ferrari- Trecate et al., 2003; Devin et al., 2017; Camacho et al., 2010). Soto et al. (2021) have shown that automata with affine dynamics can be synthesized from time-series data. Our approach instead more general, using non-lineat neural models to both ex- tract structure and represent dynamics within the automata.
actions in the land terrain are inverted. We use MPC controller to gather offline data.
LiDAR Racing Adapted from the work by Ivanov et al. (2021). Tasks in this environment consist of driving a two-dimensional vehicle with bicycle dynamics and LiDAR sensors through track randomly assembled from pieces of five different types. which correspond to the modes. We use a pre-trained controller provided by the authors to gather offline data.
Other approaches to discovering structure which are not di- rectly applicable to time-series data include the use of graph neural networks (Cranmer et al., 2020) and sparse networks (Gupta et al., 2024). Similarly, methods that leverage recur- rent neural networks (RNNs) and FSMs to model linguistic structures (Kolen, 1993; Koul et al., 2018; Jacobsson, 2005) share conceptual foundation with our work, but focus formal languages and are not directly applicable to the class of problems we study in this work.
Salamander A locomotion task in which amphibious
salamander must navigate through water and land, which correspond to the modes. This environment is implemented in the Webots 3D simulator (Michel, 2004), in which the Salamandra Robotica If (Crespi et al., 2013) robot available. This environment is a scaled-up version of Point Mass. For observations, we use the motor positions, the LiDAR readings and the GPS position. We use the controller provided in Webots to generate offline data.
6. Experiments
Our experiments aim to answer two research questions: (1) how effectively does SWMPO leverage offline data in the synthesis of environment-speen sppecific FSM?: (2) is the result- ing FSM accurate enough for model-based RL?
Bipedal Walker (BipedalWH) A locomotion task in which a bipedal robot has to loco mote over uneven terrain with four different types of obstacles, which cor- respond to the modes (see Figure 6). This is standard benchmark in the Gymnasium library (Towers et al., 2024) and employs the Box2D rigid body simulator (Catto, 2024). We use a pre-trained controller from the RL Baselines3 Zoo library (Raffin et al., 2021) to generate offline data.
The algorithm' to approximate the solution to Equation (1) is written with Pytorch (Ansel et al., 2024) using the Adam optimizer (Kingma & Ba, 2017). We use multi-layer per- ceptrons with ReLU activations for all the neural networks involved in the algorithm. Predicate synthesis is performed with the tree-learning algorithm implemented in Scikit-lear (Pedregosa et al., 2011). The underlying RL algorithm is Soft-Actor Critic, as implemented in StableBasclines3 (Raf- fin et al., 2021). We manually tune the hyperparameters of all algorithms, including baselines. We test SWMPO on four simulated environments of varying complexity. Hyper- parameters for all experiments can be found in Appendix C.
See Figure 5 for an illustration of each environment.
6.2. FSM synthesis
We evaluate the performance of our FSM synthesis algo- rithm across all four environments. As described before, our goal is to train a set of neural primitives from a dataset of 'offline" POMDPs, and finally assemble the primitives into FSM specific to 'active' POMDP. Crucially, we want the resulting FSMWM accurately track the mode variable (e.g.. the FSMWM should have modes that correspond to water and land and accurately switch between them as the robot locomotes through the 'active' POMDP).
6.1. Simulation environments
All of our experiments are performed in simulation on the following environments:
Point Mass These tasks are idealized model of the am- phibious robot running example, and consist of apply- ing a sequence of thrusts to a two-dimensional point mass to take to a target position from an initial posi- tion (see Figure 6). Crucially, the environments consist of randomly generated terrains with different character- isties: land and water, which correspond to the modes. To simulate the different dynamics in different terrains,
To test the accuracy of the structure of the FSMWM, we compare the states visited by the synthesized FSMWM in unseen data from the "active" POMDP against ground truth mode labels, as well as the visited states predicted by base- line models. The baseline models are (1) a Hidden Markov model (HMM) with Gaussian emissions (Baum & Petrie, 1966) implemented in the HMM Learn library (Learn, 2024), (2) a Switching Linear Dynamical System (SLDS) (Acker- son & Fu, 1970) and (3) a recurrent SLDS (rSLDS) (Glaser
IImplementation details can be found at: https: ///::////attt com/ da www.e/oooorreeeeeeeee /swmpo
7

---

Neurosymbolic World Models for Sequential Decision Making
(a) Point Mass
(b) LiDAR Racing
(c) Salamander
(d) Bipedal Walker
Figure The four simulated evaluation environments.
Shaanc)) (stance)
100
300
2000
80
800
1500
60
200
600
1000
40
20
100
500
(a) Point Mass
(b) LiDAR Racing
(c)
Salamander
(d) Bipedal Walker
Global HMM
Local HMM
Global SLDS
Local SLDS
Global rSLDS
Local rSLDS
FSM (ours)
Figure 6. Bottom row: Box plots comparing the Levenshtein distance between predicted and ground truth labels for unseen trajectories (lower is hetter) for each evaluation environment. Systems that accurately track the latent mode variable of a dynamical system more accurately match the ground-truth labels. The hox-plots show that SWMPO (right) outperforms haselines (left) in all environments in tracking the latent mode variable.
al., 2020), the latter two implemented in the SSM library (Linderman, 2020).
the induced partition with the ground-truth partition Figure 7.
Since the FSM undergoes two training stages-one one using data from the 'offline' POMDPs multiple terrains to syn- thesize primitives, and another on the 'active' POMDP to synthesize transitions-we train two versions of each base- trajectories
6.3. Reinforcement Learning
We evaluate the performance of the FSM models within the model-based RL loop described in Algorithm 3, shown in Figure 8. We use the data from the 'offline' POMDPs (described Section 6.2) as the input. We use Soft-Actor Critic (Haarnoja et al., 2018) as the RL routine. We compare with the following baselines: (2) standard model-free RL, and (2) mode-based RI. using neural forward-model of the environment synthesized from the 'active' POMDP data the environment is explored. The model-based RL baseline is the same as Algorithm except that it learns online a monolithic neural network instead of refining the FSMWM transitions. Note that, for the model-based RL baseline, we do not pre-train the neural network on the 'offline' data because the data comes from other POMDPs with different dynamics. We augment the observations with time variable and normalized them with running mean wrapper. Due to resource constraints, we benchmark only on Point Mass.
line: 'local' version, trained exclusively from the 'active' POMDP (excluding the test trajectories), and "global' version, trained data from both the 'offline" and "active' POMDPs.
To evaluate each model, we compute the Levenshtein dis- tance between the predicted labels and the ground-truth labels of unseen trajectories. The final errors ag- gregated across four test trajectories for each of eight dif- ferent terrains (see Figure 6). To synthesize the FSMWM, we use Algorithm to train the neural primitives from the 'offline' data. Together with this set of neural primitives, data from the "active' POMDP is used to synthesize the POMDP FSMWM with Algorithm 2. We use an expert policy to generate all data.
In Point Mass and Salamander, SWMPO outperforms all baselines. In LiDAR Racing, SWMPO significantly outper- forms SLDS and rSLDS and only marginally outperforms HMM. In Bi,, SWMPO marginally beats the base- lines, but all models struggle to capture the structure and dynamics of the environment.
The model-based rollout logic for both SWMPO and the model-based baseline consist of sampling a random obser- vation from an environment trajectory and performing shorter (30 steps) model-based rollout. This instead of re- lying the models for long trajectories (150 steps), which can lead to compounding errors. This reduces reliance on
We plot the latent mode variable and qualitatively compare
8

---

Neurosymbolic World Models for Sequential Decision Making
UMAP Dimension 2
UMAP Dimennsion.
UMAP Dimension
UMAP Dimension
(a) Ground Truth
(b) Predicted
Figure UMAP-procsssed plot of the ground truth (top) and approximation (bottom) of the latent mode variable mæz (Ot, Oe, Oc+1) across a collection of trajectorics in the Point Mass environment.
Our experiments indicate that the synthesized FSMWMs effectively capture the underlying structure of the simulated evaluation environments, as they match or surpass the per- formance of the baselines on modeling the latent variable of the tested dynamical systems.
20
15
Reward
In the RL experiment, our model-based RL approach outper- formed both the model-free and the model-basee monolithic neural baselines, the latter of which learned the environment model online. We highlight that SWMPO did not train any neural model online, instead using the online data to synthe- size transitions between the neural primitives trained from unlabeled offline data.
5
Method SAC (baseline) SWMPO (ours)
0
Model-based (baseline)
25 50
75
100 125
150
175 k 200
Avenues for future work broadly include (1) weakening the limitations of the framework, (2) leveraging neural primi- tives in a different class of models (i.e., instead of FSMs), (3) explicitly encoding priors (e.g.. through Physics Informed Neural Networks), and (4) leveraging the structure of the FSM explicitly during policy optimization (e.g.. by using the 'product MDP done by Hasanbeig et al. (2021)).
Timestep
Figure 8. Reward curves in the Point Mass environment for SWMPO (ours), baseline model-free RI, and baseline model- based RL with neural feed-forward model learned online. Shaded area are 95% confidence intervals computed over 64 different ran- dom sceds.
Impact Statement
the model for long-horizon prediction, improving sample efficiency and stability for both methods. Note that the per- formance of the low-level neural models for long horizons orthogonal to our claims.
Given the relatively narrow scope of our experiments, we do not believe there are any direct societal consequences of this work that must be specifically highlighted.
References
7. Conclusion
Ackerson, G. and Fu, K. On state estimation switching environments. IEEE Transactions on Automatic Control, 15(1), 1970.
We presented novel framework for synthesizing Finite State Machine World Models (FSMWMs) in unsuper- vised manner using low-level, non-visual continuous obser- vations. We outlined the key assumptions underpinning our approach and demonstrated its applicability. The main limi- tations of the framework are stated explicitly as assumptions throughout the text.
Alur, R., Courcoubetis, C., Halbwachs, N., Henzinger, T. A., Ho, H., Nicollin, X., Olivero, A., Sifakis, J., and Yovine, S. The algorithmic analysis of hybrid systems. Theoretical Computer Science, 138(1), 1995.

---

Neurosymbolic World Models for Sequential Decision Making
Andreas, J., Klein, D., and Levine, S. Modular Multitask Re- inforcement Learning with Policy Sketches. In Proceed- ings of the 34th International Conference on Machine Learning. PMLR, 2017.
Devin, C., Gupta, A., Darrell, T., Abbeel, P., and Levine, S. Learning modular neural network policies for multi- task and multi-robot transfer. In 2017 IEEE International Conference on Robotics and Automation (ICRA), 2017.
Ansel, J., Yang, E., He, H., Gimelshein, N., Jain, A., Voz- nesensky, M., Bao, B., Bell, P., Berard, D., Burovski, E.,, Chauhan, G., Chourdia, A., Constable, W., Desmaison, A., Vitoo, Z., Ellison, E., Feng, W., Gong. J., Gschwind, M., Hirsh, B., Huang, S., Kalambarkar, K., Kirsch, L., Lazos, M., Lezcano, M., Liang, Y. Liang, J., Lu, Y., Luk, C., Maher, B., Pan, Y., Puhrsch, C.,, Reso, M., Saroufim, M., Siraichi, M. Y., Suk, H., Suo, M., Tillet, P., Wang, E., Wang. X., Wen, W., Zhang, S., Zhao, X., Zhou, K., Zou, R., Mathews, A., Chanan, G., Wu, P., and Chin- tala, S. PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation. In 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume (ASPLOS '24). ACM, 2024.
Perrari-Trecaate G., Muselli, M., Liberati, D., and Morari, M. A clustering technique for the identification of piecewise affine systems. Automatica, 39(2), 2003.
Frigola, R., Lindsten, F., Schön, T. B., and Rasmussen, C. E. Bayesian inference and learning in gaussian process state-space models with particle memc. In Advances in Neural Information Processing Systems, volume 26. Curran Associates, Inc., 2013.
Glaser, J., Whiteway, M., Cunningham, P., Paninski, L., and Linderman, Recurrent Switching Dynamical Sys- tems Models for Multiple Interacting Neural Populations. In Advances in Neural Information Processing Systems, volume 33, Curran Associates, Inc., 2020.
Goh, C. Y., Dauwels, J., Mitrovic, N., Asif, M. T., Oran, A., and Jaillet, P. Online map-matching based on hidden markov model for real-time traffic sensing applications. 2012 15th International IEEE Conference on Intelli- gent Transportation Systems. IEEE, 2012.
Baum, L. E. and Petrie, T. Statistical Inference for Prob- abilistic Functions of Finite State Markov Chains. The Annals of Mathematical Statistics, 37(6), 1966.
Botvinick, M. M. Hierarchical reinforcement learning and decision making. Current Opinion in Neurobiology, 22 (6). 2012.
Goodfellow, I., Bengio, Y., Courville, A., and Bengio, Y. Deep learning, volume 1. MIT Press, 2016.
Bouguila, N., Fan, W., and Amayri, M. Hidden Markov models and applications. Springer, 2022.
Gupta, K., Yang, C.,, McCue, K., Bastani, O., Sharp, P. A., Burge, C. B.,, and Solar-Lezama, A. Improved modeling of RNA-binding protein motifs in interpretable neural model of RNA splicing. Genome Biology, 25(1), 2024.
Breiman, L.., Friedman, J., Olshen, R. A., and Stone, C. J. Classification and Regression Trees. Chapman and Hall/CRC, 1984.
Ha, D. and Schmidhuber, J. Recurrent World Models Facili- tate Policy Evolution. In Advances in Neural Information Processing Systems, volume 31, Curran Associates, Inc., 2018.
Camacho, H. F., Ramirez, D. R., Limon, D., Muñoz de la Peña, D., and Alamo, T. Model predictive control tech- niques for hybrid systems. Annual Reviews in Control, 34(1), 2010.
Haarnoja, T., Zhou, A., Abbeel, P., and Levine, S. Soft actor-critic: Off-policy maximum entropy deep reinforce- ment learning with a stochastic actor. In International conference on machine learning. PMLR, 2018.
Catto, E. erincatto/box2d: a 2D physics engine for games, 2024. URL https: //g/thubbbb//////// com/erincatto/ box2d. Accessed: 2025-01-31.
Hasanbeig, M., Jeppu, N. Y., Abate, A., Melham, T., and Kroening, D. DeepSynth: Automata Synthesis for Auto- matic Task Segmentation in Deep Reinforcement Leam- ing. Proceedings of the AAAI Conference on Artificial Intelligence, 35(9), 2021.
Cranmer, M., Sanchez Gonzalez, A., Battaglia, P., Xu, R., Cranmer, K., Spergel, D., and Ho, S. Discovering Sym- bolic Models from Deep Learning with Inductive Biases. In Advances in Neural Information Processing Systems, volume 33, Curran Associates, Inc., 2020.
Icarte, R. T., Klassen, T., Valenzano, R., and Mcllraith, S. Using Reward Machines for High-Level Task Specifi- cation and Decomposition in Reinforcement Learning. Proceedings of the 35th International Conference on Machine Learning. PMLR, 2018
Crespi, A., Karakasiliotis, K., Guignard, A., and Ijspeert, A. Salamandra Robotica II: An Amphibious Robot to Study Salamander-Lite Swimming and Walking Gaits. IEEE Transactions on Robotics, 29(2), 2013.
10

---

Neurosymbolic World Models for Sequential Decision Making
Inala, J. P., Bastani, O., Tavares, Z., and Solar-Lezama, A. Synthesizing Programmatic Policies that Inductively Generalize. In 8th International Conference on Learn- ing Representations. ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
Moerland, T. M., Broekens, J., Plaat, A., and Jonker, C. M. Model-based Reinforcement Learning: A Survey. Foun- dations and Trends@ in Machine Learning, 16(1), 2023.
Mohan, A., Zhang, A., and Lindauer, M. Structure in Deep Reinforcement Learning: A Survey and Open Problems. J. Artif. Int. Res., 79, 2024.
Ivanov, R., Jothimurugan, K., Hsu, S., Vaidya. S., Alur, R., and Bastani, O. Compositional Learning and Verification of Neural Network Controllers. ACM Transactions Embedded Computing Systems, 20(5s), 2021.
Paoletti, S., Juloski, A. L.., Perrari-Trecatee G., and Vidal, R. Identification of Hybrid Systems A Tutorial. European Journal of Control, 13(2), 2007.
Jacobsson, H. Rule Extraction from Recurrent Neural Net- works: ATaxonomy and Review, Neural Computation, 17 (6), 2005.
Pedregosa, F, Varoquaux, G., Gramfort, A., Michel, V. Thirion, B., Grisel, O., Blondel, M., Prettenhofer, P., Weiss, R., Dubourg, V., Vanderplas, J., Passos, A., Cour- napeau, D.,, Brucher, M., Perrot, M., and Duchesnay, Scikit-learn: Machine Learning in Python. Journal of Machine Learning Research, 12, 2011.
Kingma, D. P. and Ba, J. Adam: A Method for Stochastic Optimization, 2017.
Klusowski, J. M.,, and Tian, P. M. Large Scale Prediction with Decision Trees. Journal of the American Statistical Association, 19(545), 2024. Publisher: ASA Website.
Raffin, A., Hill, A., Gleave, A., Kanervisto, A., Ernstts,, M., and Dormann, N. Muuble-aseelines3. Reliable Reinforce- ment Learning Implementations. Journal of Machine Learning Research, 22(268), 2021.
Kolen, J. Fool's Gold: Extracting Finite State Machines from Recurrent Network Dynamics. In Advances in Neu- ral Information Processing Systems, volume 6. Morgan- Kaufmann, 1993.
Shaj, V., Buchler, D., Sonker, R., Becker, P., and Neumann, G. Hidden Parameter Recurrent State Space Models For Changing Dynamics Scenarios, 2023.
Koul, A., Greydanus, S., and Femm A. Learning Finite State Representations of Recurrent Policy Networks, 2018.
Simpkins, C. and Isbell, C. Composable Modular Reinforce- ment Learning. Proceedings of the AAAI Conference on Artificial Intelligence, 33(01), 2019.
Learn, H. hmmllaa/hhmmlear: Unsupervised learning and inference of Hidden Markov Models, 2024. URL htttp:: //github.com htmm eerellarn//ii http/hmmlelarrr Accessed: 2025-01-31.
Soto, M. G., Henzinger, T. A., and Schilling, C. Synthesis of hybrid automata with affine dynamics from time-series data. In Proceedings of the 24th International Conference Hybrid Systems: Computation and Control, HSCC 221. Association for Computing Machinery, 2021.
Li, and Biswas, G. Applying the hidden Markov model methodology for unsupervised learning of temporal data. International Journal of Knowledge Based Intelligent Engineering Systems, 6(3), 2002.
Sutton, R. S. and Barto, A. G. Reinforcement learning: An introduction, 2nd ed. Reinforcement learning: An introduction, 2nd ed. The MIT Press, 2018.
Li, L., Walsh, T. J., and Littman, M. L.. Towards Unified Theory of State Abstraction for MDPs. In International Symposium on Artificial Intelligence and Mathematics, Al&Math 2006, Fort Lauderdale, Florida, USA, January 4-6, 2006, 2006.
Tang, H., Key, D., and Ellis, K. WorldCoder, Model- Based LLM Agent: Building World Models by Writ- ing Code and Interacting with the Environment. CoRR, Abs/2402.12273. 2024.
Linderman, L. lindermanlab/ssm: Bayesian learning and inference for state space models, 2020. URL https: //github. com/ h waaeeraaanaaaaaaeaam ssm. Accessed: 2025-01-31.
Toro Icarte, R., Waldie, E., Klassen, T., Valenzano, R., Cas- tro, M., and Mellraith, Learning Reward Machines for Partially Observable Reinforcement Learning. In Advances in Neural Information Processing Systems, vol- ume 32. Curran Associates, Inc., 2019.
Marturi, N., Kopicki, M., Rastegarpanah, A., Rajasekaran, V., Adjigble, M., Stolkin, R., Leonardis, A., and Bekiroglu, Y. Dynamic grasp and trajectory planning for moving objects. Autonomous Robots, 43, 2019.
Towers, M., Terry, J. K., Kwiatkowski, A., Balis, J. U.,
Cola, G., Deleu, T., Goulão, M., Kallinteris, A., KG, A., Krimmel, M., Vicent,e R., Pierré, A., Schul- hoff, S., Tai, J. J., Tan, A. J. S., and Younis, O. G. Gymnasium, 2024. URL https: // landdooorg/ records/1123255
Michel, O. Webots: Professional Mobile Robot Simulation. Journal of Advanced Robotics Systems, 1(1), 2004.

---

Neurosymbolic World Models for Sequential Decision Making
Tran, K., Bisk, Y., Vaswani, A., Marcu, D., and Knight, K. Unsupervised neural hidden Markov models. arXiv preprint arXiv:I609.000
2016.
Vakanski, A., Mantegh, I., Irish, A., and Janabi-Sharifi, F. Trajectory learning for robot programming by demon- stration using hidden Markov model and dynamic time warping. IEEE Transactions Systems, Man, and Cy- bernetics, Part B (Cybernetics), 42(4), 2012.
Wu, H., Guan, Y., and Rojas, J. latent state-based mul- timodal execution monitor with anomaly detection and classification for robot introspection. Applied Sciences, 9 (6), 2019.
Xu, D. and Fekri, F. Interpretable Model-based Hicrarchical Reinforcement Learning using Inductive Logic Program- ming, 2021.
12

---

Neurosymbolic World Models for Sequential Decision Making
A. Notation
Table Notation Table
Symbol
Description
M
Markov Decision Process (MDP)
Sell I2*
Set of observations
A
Set of actions
S: x A S
Transition function
So
Distribution of initial states
R X R
Reward function
:
Policy parametrized by Value function for state s under policy x0 Source state at time t Action at time t
(s)
81+1
Next state after taking action Finite State Machine World Model (FSMWM)
F
f
Set of environment models
6
Set mode- transition predicates
81+1 fi(St, us)
Predicted next observation
8(st, ue.
Mode transition function
B. Partition prunning
Algorithm 4 greedyPrune
Require: Partition 5 = (D....., Dm) of trajectory dataset D, error tolerance factor c,
1: Du (D).
2: for / D do
3: while exists e-prunnable (relative to Do) mode transition in t do
4: Prune the first e-prunable (relative to Do) mode transition D, updating D.
5: end while
6: end for
7: return D
13

---

Neurosymbolic World Models for Sequential Decision Making
C. Hyperparameters
Table 2, Parameters for Point Mass.
Parameter
Value
hidden_sizee sizes
32 32 32
learning-rate
0.0001
partition_sie _iie batch_size
1000
min island size autoencoder latent prunning_eree trrlecrrec model_mode_iter_.n
8
0.04
1024
cluster dimensionality reduce information contennt.eggularzzttooo scale
0
le-05
mutual infformattio..eegllazzztttioo
0.00001
nutual_informall _mini miii_batchh_siee
predicate.hyprppaae hyperparameters.n max .eeppett _eepth
8
local model hyperparameters..idee hidden layer sizes
[32, 32]
locel_ mooood__yypppprrmmtess.aaa. mix.ir
500
Table Parameters for Autonomous Driving.
Parameter
Value
hidden_sizes
32 32 32
learning rate
0.01
partition_sia batch size min_island _siee autoencoder Jatent-siiiziiii
2
1024
8
aaaniing_erro_toleeeee
0.04
model mode iter
1000
cluster diasterrdimensionalityy reduce
0
ronteent.rguulaizzationno scale
1e-05
mutual information regularization scale
0.00001
mutual_information.nne www.ibatcccsiieeee
predicate hyperparameters..me depth
local. hww.ehaaaraaaateessshheeeeeeeeee
[32, 32]
local.mode model hyperparameters.max. iter
500
14

---

Neurosymbolic World Models for Sequential Decision Making
Table 4. Parameters for Salamander.
Parameter
Value
hidden_size
32
learning-rate
0.001
partition size
2
batch_size
1024
min size
4
autoencoder lateentee_ltttttiiiiiiii
8
prunning error tolerance
0.04
model_mode.ite .neen Custer_.immnsioonllin reduce
1000
information_conteettrgggge content.regularrizttoo _sclle
1e-05
mutual inwormatioooooonnn....oo reffomaatoonreggulaizaaaiionnnoo scale
0.00001
mutual information mini batch
32
predicate_hy www.aaaaaatterrss.....mmmm max.smmmaacoommmmmmmmm _depth
8
local model hyperparameters.hiddee laaeeee.ereiideeeeeeeeer..ereere lazerssizes sizes
[32, 32]
local_mode leaal_moeel_iig hyprlllpppppprrmeeeesmmaa
500
Table 5, Parameters for Bipedal Walker.
Parameter
Value
machine hidden.ddeee sizes machine learning rate machine cichine.partitionnsiii size machine batchn size machine_min.it min__sland_sizze
32 32 32
0.001
1024
4
machine autoencoder latent size
8
machine_pmine machine.prumming www.lpuuuunnnnnneeee wwweerrtteeraaeee
0.04
machine machine_model_mmee mode iiter.n
1000
machine cluster dimensionality reduce
0
machine information content regularization scale
lc-05
machine, mutual information wewarrtiooeggllliiattioooo.oooo scale machine .mttealluaaaaaaaa informatio. mini batch_bstize
machine_mutual.
0.00001
machine predicate hyperparameters. depth
8
Lacaee_oooal__mleee moode_.yyerpprrameeers.hht hidden. www.eennayeeeriiieeeee........
[32, 32]
machine local model hyperparameters.mog max iter
500
D. Complexity
Let be the number of experiences (ot, 0t, 0t+1), 1,, each of dimension d, and let m the number of modes.
The computational complexity of training is driven by these steps:
Solving Eq. To solve Eq. we use SGD, which (G))).)) where is the number of SGD steps, and Ki accounts for the architecture of the model (see work by Goodfellow et (2016) for further discussion). The variables corresponding the amount data, the number iterations and the model complexity needed for training, are not independent. R.g., system with high number of modes might need more data. However, 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 monolithic approach.
Training the neural primitives: We assume datasets will have in expectation ( elements. If models are trained sequentially, this (mG2 a dK2), where G2 and K, account for the number of SGD steps and the architecture of the primitives. The caveats in I. regarding the non- non-independence of features 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.
15

---

Neurosymbolic World Models for Sequential Decision Making
Transition Predicate Synthesis: This step trains 2m decision with CART (Breiman et al., 1984; Klusowski et al., 2024). For the average case, assume each dataset will be of size O m Therefore, the overall cost in expectation is O dm log" 2
m
Applying dimensionality reduction techniques may help with high-dimensional spaces. In high-dimensional domains (e.g., visual domains, which are not the focus of our work), we would expect different class of models (e.g.. CNNs) to be more effective.
Simplifying, the average complexity of training a SWMPO model is
(Gnmdk + nmd + mod
(m)) m
where G = m G2) and K = max K1, K2). Thus, the expected overhead of training in SWMPO compared to 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 added small constant factor slower than traditional monolithic world model.
16


Reviews:
[{"review_id": "ikuLBBHgaG", "reviewer": "Reviewer_bZQ6", "summary": "The paper presents Structured World Modeling for Policy Optimization (SWMPO), a framework for unsupervised learning of neurosymbolic Finite State Machines (FSM) that capture environmental structure for policy optimization. The method operates in two main stages: (1) learning local “world-model primitives” that specialize in modeling different “modes” of a partially observed system, and (2) assembling those primitives into an FSM that captures transitions among those modes. These local world-model primitives are trained in an unsupervised fashion from offline data. Then, with limited new data from the current task, the system stitches the primitives into a new FSM that is specialized to the particular environment. Finally, the authors leverage this FSM representation for model-based policy optimization.\n\nIn terms of experiments, the paper reports results on four main environments: (1) PointMass (2) LiDAR-Racing (3) Salamander and (4) BipedalWalkerHardcore. The empirical results suggest that (a) the approach can learn to recover latent modes with reasonable accuracy, sometimes outperforming classical switching system baselines like HMMs and switching linear dynamical systems, and (b) the resulting FSM-based world models can be used effectively for model-based RL, matching or marginally outperforming a comparable monolithic neural model in some test settings.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank Reviewer **bZQ6** for their thoughtful comments, and we are committed to incorporating your feedback into the final manuscript.\n\n> **Note:** Since our initial submission, we have demonstrated stronger RL performance of SWMPO over baselines (see response to **B4Kd**).\n\nPlease excuse our brevity due to the character limit.\n\n## Scalability\n\nWe understand your concern about the scalability of SWMPO in environments with a high number of modes or high dimensionality.\n\nThe computational complexity of SWMPO is driven by three steps: (1) solving Eq. 1, (2) training of neural primitives, (3) synthesizing predicates.\n\nIt 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.\n\nFor more details, please refer to our response to **pyge** (*”Computational Complexity\"*).\n\nAdditionally, 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.\n\n## Mode Identifiability Assumption\n\nAs 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.\n\n## Real-World Applicability\n\nOur 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.\n\nMoreover, 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.\n\nNonetheless, 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.\n\n## Writing and Figure\n\nWe 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.\n\n## References\n\n[0] Zhang et al., *Invariant Causal Prediction for Block MDPs*, ICML 2020  \n[1] Liu et al., *Discrete-Time Hybrid Automata Learning: Legged Locomotion Meets Skateboarding*, arXiv preprint, 2025  \n[3] Poli et al., *Neural Hybrid Automata: Learning Dynamics with Multiple Modes and Stochastic Transitions*, NeurIPS 2021  \n[4] *WorldCoder: a Model-Based LLM Agent*, NeurIPS 2024  \n[5] https://www.epfl.ch/labs/biorob/research/amphibious/salamandra/  \n[6] https://www.cyberbotics.com/doc/guide/salamander?version=R2021a#salamander-wbt  \n[7] Bhatt et al., *CrossQ: Batch Normalization in Deep Reinforcement Learning for Greater Sample Efficiency and Simplicity*, ICLR 2024  \n[8] Kuznetsov et al., *Controlling Overestimation Bias with Truncated Mixture of Continuous Distributional Quantile Critics*, ICML 2020  \n[9] Goodfellow, I., Bengio, Y., & Courville, A. (2016). *Deep Learning*. MIT Press  \n[10] David Arthur, Sergei Vassilvitskii, *How slow is the k-means method?*, SCG 2006  \n[11] Klusowski & Tian, *Large Scale Prediction with Decision Trees*, Journal of the American Statistical Association, 2023  \n[12] Solar-Lezama, *The Sketching Approach to Program Synthesis*, APLAS 2009", "claims_and_evidence": "(1) The paper claims that modeling the environment with a finite set of local dynamics (modes) can lead to better structural world models.\n    \nThrough experiments (e.g., Figure 5 and related discussion), the authors show that the learned finite-state structure captures mode switching more accurately than certain baselines like HMM or switching linear models, especially in some of the environments (PointMass, LiDAR-Racing, Salamander).\n\n\n(2) The paper argues that offline training of local dynamics primitives, followed by environment-specific stitching, can improve efficiency in policy optimization.\n\nThe authors’ experimental results (especially in PointMass) demonstrate that their approach, SWMPO, can use a short amount of new data in the active environment to synthesize an FSM that is then used for policy optimization. They compare a purely online-learned forward model vs. offline-learned primitives combined with minimal environment interactions, showing they achieve roughly similar or slightly better performance with fewer interactions.", "methods_and_evaluation_criteria": "Methods\n\n(1) Neural Primitives: A neural network is trained to embed each observed transition into a continuous latent vector that captures the local mode. Then, an additional network predicts the next observation from the latent mode.\n        \n(2) Clustering and Pruning: Those embeddings are then clustered (k-means), and the resulting assignment to clusters is used to label each transition with a symbolic “mode.” A pruning step removes spurious transitions among modes.\n        \n(3) FSM Predicate Synthesis: Decision trees are used to learn conditions under which the FSM transitions from one mode to another, given the environment’s (observation, action) pairs.\n        \n(4) Model-Based RL: The final FSM is used in a model-based policy optimization loop (Soft-Actor Critic).\n\nEvaluation Criteria\n    \n(1) Mode Accuracy: Measured via Levenshtein distance between the predicted mode sequence and ground-truth mode labels for new episodes.\n        \n(2) Policy Performance: Cumulative rewards achieved by the learned policy.\n\nOverall, the method and evaluation are comprehensive.", "theoretical_claims": "No detailed proofs (e.g., in an appendix) of correctness or identifiability are provided beyond references to known standard assumptions (like the existence of a unique minimal partition).", "experimental_designs_or_analyses": "The experimental design is centered around four environments of increasing complexity, each with a known “ground-truth” notion of mode (e.g., land vs. water, track type, or different obstacle types). This design is appropriate.", "relation_to_broader_scientific_literature": "The references provided do give an overall sense that the authors are aware of classical baselines in this area.", "essential_references_not_discussed": "The references included do not appear incomplete for an initial demonstration.", "other_strengths_and_weaknesses": "Strength:\n\n1. Structured Representation: The approach explicitly partitions the environment’s dynamics into interpretable “modes,” making it easier to reason about or visualize transitions.\n    \n2. Modular Reuse: Potentially very useful if the same local dynamics (e.g., “land locomotion,” “water locomotion,” “curved track,” etc.) reappear across multiple tasks.\n\nWeaknesses:\n    \n1. Scalability: If there are many different modes or if the environment is very complex high-dimensional, the overhead of separate networks plus a large transition graph may become expensive and potentially complicated.\n        \n2. Partial Observability: The method’s success hinges on the assumption that a single time step (previous observation, action, new observation) is enough to identify the latent mode. This assumption may fail in more complex partially observed tasks.\n\n3. More analysis on extending to complex environment and real-world data. The mode presented in the paper is too simple, and might not be practicle for more complex setting.\n        \n4. Writing and Figure. (1) The authors should give examples of primitives at the very early stage for better readibility. (2) Mode should be presented in teaser figure,  as it is a key in this paper.", "other_comments_or_suggestions": "Analyzing or providing examples of failure cases Cocould help future readers understand the method’s limits.", "questions_for_authors": "Thank you for the paper submission. Overall, it presents a promising neurosymbolic approach that merges mode discovery and WM.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "X1BHT32mGJ", "reviewer": "Reviewer_B4Kd", "summary": "In the setting of POMDP, the paper introduces SWMPO, a framework based on a Markov Decision space model in which each transition can be characterized by mode (FSM).  That is, at each $t$, the transition occurs by \n$(o_t, a_t) \\mapsto o_{t+1}  = f(o_t, a_t | M_t) $ where $M_t is the mode at time $t$. \nThe mode $M_t$ is to change based on $o_t, a_t$ through what they call \"transition predicate\". \n\nThey train this model in two steps: \n(1) the first step in which the proxi $m_{\\theta1}$ is included as \"soft\" modes so that the fuction of the form  \n$$f_{\\theta1}(m_{\\theta2}(o_{t-1}, o_t, a_t) , o_t, a_t) $$ \ncan predict the the state well and that $m_{\\theta2}$ is not correlated to the future. \n\nBased on this $m_{\\theta2}$, the clustering is conducted, so that any $(o_{t-1}, o_t, a_t)$ can be assigned to a particular mode. \nThen, for each mode $m$,  $f(o_t, a_t | m)$ is trained to predict the observation whenever $(o_{t-1}, o_t, a_t)$ is of mode $m$. \n\nThe transition predicate is trained so that the model transition to the best predicting state at any given $(o, a)$. \n\nAt each round, SWMPO then uses small number of new environmental rollouts, and then uses the FSM that is newly synthized from the dataset.  \n\nTheir methods are compared against Monolithic Neural Model based RL as well as model-free RL and is shown to outperform them on several benchmark examples.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n> **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).\n\n## Performance over Monolithic Models\n\nWe 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.\n\nSpecifically, this is what we have done for both SWMPO and the monolithic model-based RL baseline:\n\n1. We ran the RL training 64 times, instead of the 16 times in our original experiment.  \n2. We further tuned the hyperparameters of all algorithms.  \n3. 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.\n\nThese 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.\n\nWe encourage reviewers to check the updated results here (corresponding to an updated Figure 7): [https://anonymous.5open.science/r/anonymousfig-30AA/improvedrewards.png](https://anonymous.4open.science/r/anonymousfig-30AA/improvedrewards.png)\n\nWe summarize these updated results in the table below:\n\n| Iteration      | Baseline RL (Mean Reward) | Model-based RL (Mean Reward) | SWMPO (Mean Reward) |\n|----------------|---------------------------|-------------------------------|----------------------|\n| 20,000         | ~-1                       | ~4.8                          | **~5.2**             |\n| 40,000         | ~0.5                      | ~9                            | **~14**              |\n| 60,000         | ~0.5                      | ~9                            | **~17.5**            |\n| 75,000         | ~4                        | ~15                           | **~21**              |\n\n## Extreme Case Scenario\n\nWe 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.\n\n## Additional References\n\nWe appreciate the reviewer's suggestion of additional references, and will include a discussion of all of them in the final manuscript.", "claims_and_evidence": "While it is true that their approach extended the FSM based approach to continuous domain, it is unfortunately a hard call to approve the claim that they demonstrate the advantage of the approach based on the model-based RL using a single monolithic neural network with no structure.", "methods_and_evaluation_criteria": "The evaluation criteria of the work is based on mode-label accuracy and the sheer policy performance. I believe that these criteria are fair, with the first one validating the claim that they can nail down the mode when there is one.", "theoretical_claims": "There seems to be limited theoretical discussion.", "experimental_designs_or_analyses": "The experimental designs seems conventional, and I believe they are sound.", "relation_to_broader_scientific_literature": "N.A", "essential_references_not_discussed": "Nothing in particular that comes to my mind, but \nPlaNet(Hafner et al 2019), Dreamer (Hafner et al 2021), SimPLe (Kaiser et al. 2020) might be worthy of mention as the \nthe members of the kind that uses future-prediction model in a rather \"meta\" way. \n\nAlso, in Neural Fourier Transform (Koyama et al) and Unsupervised Learning of Equivariant Structure from Sequences\n(Miyato et al) they also use the two-stage framework that consists of the first step of making the \"prediction model\" with the similar loss and the second \"block-diagonalizing step\" of decomposing the prediction into disentangled features. They might be somewhat related as well?", "other_strengths_and_weaknesses": "## Strengths \n- They extended FSM  and its application to RL to the tasks in  continuous domain, and brought it up to competitive level \n- They have shown an inspiring framework of including \"mode/gear change\"in the prediction, and realized it at the level that \nRL can be competitively performed.\n\n## Weaknesses\n- It was sincerely my hope as the reviewer that this approach strongly outperforms the blackbox monolithic  version, but it was not the case. \n- In the similar note to above, the paper is not too convincing in that the introduction of \"mode\" and of building sepeparate model is a beneficial approach.", "other_comments_or_suggestions": "It is hard to believe that, when the inductive bias of \"neurosymbolic world model\" is truely valid at the level of data-generation, the approach like this cannot strongly outperform the blackbox counter part neither in terms of  training speed nor on sheer reward-based performance. Has there been any effort building an \"extreme case\" scenario? \n\nClearly, the efficacy of the approach in terms of labeling accuracy / Labenshtein distance is proven to be solid. I wonder if there is no Reward that more directly depends on these labels. \n\nI value that you highlight the fact that SWMPO did not train any neural model online. I wonder if there is more quantitative way to highlight the SWMPO's gray-box approach.", "questions_for_authors": "I think I have asked questions in the comments/suggestions section.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dd6SvdQcWd", "reviewer": "Reviewer_pyge", "summary": "The paper presents SWMPO, a framework for learning neurosymbolic Finite State Machines (FSMs) to model environmental structures for policy optimization. The key contributions are an unsupervised learning algorithm for training modular world-model primitives from low-level continuous observations, a state-machine synthesis algorithm for constructing environment-specific FSMs, and evaluations showing these FSMs effectively support model-based Reinforcement Learning (RL). The framework was tested in environments like PointMass, LiDAR-Racing, Salamander, and BipedalWalkerHardcore, demonstrating accurate FSM synthesis and competitive RL performance.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer **pyge** for their thoughtful comments and are committed to incorporating your feedback into the manuscript.\n\n> **Note:** Since our submission, we have demonstrated stronger RL performance (see response to **B4Kd**).\n\nPlease excuse our brevity due to the character limit.\n\n## Mode Identifiability Assumption\n\nPlease refer to the response to Reviewer **bZQ6** (*\"Mode Identifiability Assumption\"*).\n\n## Mode Sharing and Diverse Environments\n\nIn our original submission, we stated a stronger assumption than needed. Assumption 4.5 should instead be:  \n\n> “The modes in the active POMDP must be a subset of the modes in the offline POMDPs.”  \n\nE.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.\n\nWe 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.\n\n## Real-World Applicability\n\nPlease see response to **bZQ6** (*\"Real-World Applicability\"*).\n\n## Framework Implementation Challenging\n\nTo help the community build on our results:\n\n- We open-source our source code  \n- The pseudocode typeset in the manuscript closely follows our implementation  \n- As mentioned, we leverage off-the-shelf, well-established implementations of key algorithms\n\n## Computational Complexity\n\nLet $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: \n\n1. **Solving Eq. 1:** \nTo 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.\n\n2. **Training the $m$ neural primitives:** \nWe 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.\n\n3. **Transition Predicate Synthesis:** \nThis 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}))$.\n\nApplying 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.\n\nSimplifying, the average complexity of training a SWMPO model is\n\n$$O(GnmdK + nmd + nmd\\log^2(\\frac{n}{m})),$$\n\nwhere $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$.\n\nThen, 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.\n\n## Pruning Guidelines\n\nWe manually tuned $\\epsilon$ by inspecting error plots (e.g., Figure 4), but did not find the state machines overly sensitive to it.\n\nAs 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$.\n\nAdditionally, hyperparameter tuning techniques could automate the process.\n\n## Domain Knowledge\n\nWhile out of scope for this manuscript, we consider incorporating domain knowledge explicitly a promising future direction. For instance:\n\n1. Replace models in Eq. 1 with physics-informed neural networks or other models that enforce physical constraints \n2. Integrate syntactic/semantic constraints into the predicate synthesis using systems like *Sketch* [12]\n\nWe hypothesize that each approach would improve the local models and the FSM transition dynamics, respectively.\n\nWe note that the assumption that systems can be decomposed into modes is a form of domain knowledge in itself.\n\n## References\n\nPlease refer to the response to Reviewer **bZQ6**.", "claims_and_evidence": "The claims in the submission are generally supported by clear and convincing evidence. The authors provide a detailed description of their framework, SWMPO, and its components, including the unsupervised learning algorithm for training modular world-model primitives and the state-machine synthesis algorithm for constructing environment-specific FSMs. They also present evaluations in various simulated environments, such as PointMass, LiDAR-Racing, Salamander, and BipedalWalkerHardcore, demonstrating the effectiveness of their approach in synthesizing accurate FSMs and its competitive performance in model-based RL tasks. However, some claims could be considered problematic due to the limitations of the framework. For instance, the assumption that the latent categorical variable Mt​ can be characterized by a function $m:(o_{t−1​},a_{t−1​},o_t​)\\to m_t$​ might be too restrictive for more complex environments where the mode variable is not easily identifiable from a single transition. Additionally, the pruning mechanism, while helpful in simplifying the state machine, might lead to removing transitions that are important for accurate modeling in specific scenarios.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria are well-suited for the problem of synthesizing neurosymbolic world models for sequential decision making. The framework’s ability to learn structured world models from low-level observations and use them for efficient policy optimization is demonstrated through appropriate methods and comprehensive evaluations. The choice of benchmark environments and metrics effectively supports the claims made in the paper.", "theoretical_claims": "There is no theoretical analysis in this paper.", "experimental_designs_or_analyses": "Yes.", "relation_to_broader_scientific_literature": "The concept of neural world models has been explored in previous works, such as the recurrent world models by Ha and Schmidhuber (2018). These models use neural networks to predict future states and facilitate policy evolution. The proposed SWMPO framework extends this by incorporating a structured representation through FSMs, allowing for more interpretable and modular world models.\n\nThe use of structured models in RL has been explored in various contexts, such as hierarchical RL (Xu & Fekri, 2021; Botvinick, 2012) and modular RL (Simpkins & Isbell, 2019). These approaches aim to improve policy learning by encoding structure into the policy architecture. The SWMPO framework differs by focusing on the synthesis of a structured world model rather than directly encoding structure into the policy.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "**Strengths**\n\n1. The paper proposes a creative combination of neurosymbolic approaches, integrating neural networks with finite-state machines to model complex environments. This hybrid approach is innovative and addresses the limitations of purely neural or symbolic models.\n2. The paper is well-structured and clearly presents the proposed methods, experimental designs, and results.\n\n**Weaknesses**\n\n1. The paper relies on several assumptions, such as the identifiability of the latent categorical variable $M_t$ and the minimality of modes. These assumptions might be too restrictive for more complex environments where the mode variable is not easily identifiable from a single transition.\n\n2. The assumption that all POMDPs share the same set of modes but have different mode-transition dynamics might limit the framework's applicability to environments with highly diverse dynamics.\n\n3. While the benchmark environments are relevant, they might not be sufficiently challenging to fully test the framework's capabilities. More complex, real-world scenarios could better assess the framework's performance.\n\n4. The framework involves multiple components, including neural primitives, state-machine synthesis, and transition predicate synthesis. Its complexity might make it challenging to implement and optimize, especially for less experienced researchers or practitioners.", "other_comments_or_suggestions": "See questions.", "questions_for_authors": "1. The framework relies on assumptions about the identifiability and minimality of modes. How would the method perform when these assumptions are violated, e.g., in environments with overlapping mode dynamics or non-minimal modes? Can you provide strategies to mitigate such issues?\n\n2. The evaluations are conducted in simulated environments. What challenges do you foresee when applying SWMPO to real-world tasks (e.g., robotics, autonomous systems)? Have you considered case studies or initial experiments with real data?\n\n3. How does the framework’s computational complexity scale with the number of modes or environment complexity? Are there strategies to optimize efficiency for high-dimensional state spaces?\n\n4. The pruning approach removes spurious transitions. How sensitive is the pruned FSM’s performance to the choice of the error tolerance factor ε? Can you provide an analysis or guidelines for selecting ε?\n\n5. How does SWMPO compare to hybrid neuro-symbolic models that explicitly encode domain knowledge (e.g., physics-based constraints)? Could combining such knowledge with SWMPO further improve performance?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This work proposes an interesting idea by learning a neurosymbolic Finite State Machines (FSMs) that model environmental dynamics in continuous spaces for policy optimization. The authors provide compelling evidence for SWMPO’s effectiveness through experiments in four simulated environments (PointMass, LiDAR-Racing, Salamander, BipedalWalkerHardcore), demonstrating competitive model-based RL performance. The rebuttal addresses reviewer concerns, particularly those regarding mode identifiability, scalability, and real-world applicability, by clarifying assumptions, providing detailed computational complexity analysis, and open-sourcing code to enhance accessibility. Significant improvements in RL performance (~40% over baselines) further strengthen the paper’s claims. Therefore, the work makes meaningful contribution to ICML.
Decision: Accept (poster)