Structured World Models From Low-Level Observations

Overall, the paper reads more like an experimental report or an essay rather than a fully developed research contribution.

1. The proposed algorithm depends heavily on the initial expert policy $\pi_0$, which limits its applicability in online learning settings.

2. While the motivation centres on improving RL control via FSM synthesis, the paper does not demonstrate how this method performs when integrated with RL algorithms. The experiments are limited to unsupervised labelling and FSM generation, which undermines the credibility of the claimed benefits for RL control.

3. Although the authors compare their approach with an HMM-based method, the details of the latter are insufficiently explained, with no descriptions or references provided. Generally speaking, this paper lacks of comparison with other methods (experimentally or analytically).

4. The caption for Figure 4 is incomplete, and there is a typographical error on line 053.

The paper devotes a large portion of space to formal definitions and algorithms but lacks comprehensive experimentation. Key elements, such as ablation studies on hyperparameters (e.g., epsilon, terms in the loss function, and pruning), are missing, which would help understand the impact of various components and parameters in SWMPO.

Testing SWMPO on only four environments limits its generalizability. Notably, three of these environments are custom-designed, raising questions about reproducibility and relevance. Including standardized benchmarks (e.g., MiniGrid or Four Rooms) would provide a more robust evaluation. Comparisons are limited, with only one baseline, Hidden Markov Models (HMMs). Introducing additional baselines like DreamerV3 or simpler models (e.g., behavioral cloning) would strengthen the experimental results and clarify the unique strengths of FSMs over other architectures. The results lack interpretation, and some visual elements are unclear. For instance, Figure 6d seems to illustrate a failure mode, yet this is not discussed. Similarly, Figure 4's caption is cluttered, and details like the number of seeds for Figure 8 are missing. This lack of detail reduces the clarity and interpretability of the findings.

Most importantly, the FSM's performance advantage over HMMs is sometimes minimal, particularly in Figure 8, where the FSM's results lie within HMM’s performance range in three of the four cases. This raises questions about the FSM's robustness as a general-purpose world model.

I suggest revising the paper and resubmit with a stronger experiment section, including: ablation studies on key hyperparameters and components, pruning etc., as well as broader environmental benchmarks and standardized baselines to demonstrate SWMPO's applicability across various scenarios. An expanded discussion section interpreting results, especially on observed failure modes and performance variations between SWMPO and HMM.

These are the weaknesses of the paper in my opinion.

1. The paper ignores a plethora of literature in this domain, like switching state-space models[1][2], hidden-paramter SSMs[3] etc which tackles the same problem but with more principled approaches and in more complex settings. These should be made as baselines.

2. The main purpose of using this as a world model, is for making action conditional multi-step ahead predictions. No experiments on this regard are seen.

3. The paper could be made even more stronger, if control experiments (model-based RL) with this world model are performed.

4. There are a few typos, which need to be corrected, including in the objective function in equation (1). It should be $f\left(m\left(o_{t-1}, a_{t-1}, o_t\right), a_t, o_t \right)$ in Equaton 1? what is e and d ?

• The paper makes references (L158, L166, L180, L344) to how SWMPO has been used in the paper in an RL training loop, but there are no experiments to back up that claim. Either these experiments should be added or the claims removed from the paper.
• All the experiments assume access to an expert policy.
• Assumption 2 is restrictive. It is not trivial to assume that a POMDP can become a deterministic MDP conditioned on the mode.
• There are only two environments tested where the number of modes is >2 and LiDAR racing appears to have simple dynamics. Most environments where this system can be used have more than two modes.
• The paper assumes knowledge of the number of modes in the environment.
• Presentation is poor, making the paper difficult to read (see changes requested in the section below)

I believe the paper shows promise but in its current state is not ready for acceptance. I would be willing to reconsider if the authors are able to provide the requested changes.