Paper deep dive
Towards Zero-Shot Task Transfer with Neurosymbolic World Models
Isidoro Tamassia, Lennert De Smet, Giuseppe Marra
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 91%
Last extracted: 8/25/2026, 6:20:49 AM
Summary
The paper introduces Neurosymbolic World Models (NeSy-WMs), a reinforcement learning framework that decouples latent dynamics from reward prediction by using explicit symbolic state components. This architecture allows for zero-shot task transfer under shared dynamics, enabling agents to adapt to new reward functions without further environment interaction or model finetuning, unlike purely neural world models like DreamerV3.
Entities (6)
Relation Signals (4)
Neurosymbolic World Models → enables → Zero-Shot Task Transfer
confidence 95% · Decoupling observation reconstruction and reward prediction allows us to learn world models that can adapt zero-shot... to new reward functions
Neurosymbolic World Models → buildson → Recurrent State-Space Models
confidence 90% · NeSy-WMs build on Recurrent State-Space Models (RSSMs) (17), which learn compact latent dynamics from high-dimensional observations.
Neurosymbolic World Models → outperforms → DreamerV3
confidence 88% · demonstrate the strong generalisation properties of our approach over purely neural methods... evaluate their training performance and sample-efficiency against the state-of-the-art DreamerV3
Neurosymbolic World Models → implementedin → PyTorch
confidence 85% · We provide an open-source and highly configurable PyTorch (35) implementation of NeSy-WMs.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:State-of-the-art model-based reinforcement learning methods learn neural world models that allow policy improvement by planning in a latent space, without assumptions on the structure of the underlying environment. While expressive, these models are generally task-dependent: they learn uninterpretable latent representations that are tied to the training task and thus hard to generalize to new tasks. In this work, we present a novel world model formulation where the reward prediction only depends on a subset of structured, symbolic components of the whole latent state. Decoupling observation reconstruction and reward prediction allows us to learn world models that can adapt zero-shot, i.e. without further environment interactions, to new reward functions defined over the same symbolic state space. We discuss the main advantages and challenges of learning these neurosymbolic world models and demonstrate the strong generalisation properties of our approach over purely neural methods.
Tags
Links
- Source: https://arxiv.org/abs/2608.17959v2
- Canonical: https://arxiv.org/abs/2608.17959v2
Trouble viewing inline? Open PDF directly →
Full Text
48,972 characters extracted from source content.
Expand or collapse full text
Towards Zero-Shot Task Transfer with Neurosymbolic World Models Isidoro Tamassia Lennert De Smet Giuseppe Marra Abstract State-of-the-art model-based reinforcement learning methods learn neural world models that allow policy improvement by planning in a latent space, without assumptions on the structure of the underlying environment. While expressive, these models are generally task-dependent: they learn uninterpretable latent representations that are tied to the training task and thus hard to generalize to new tasks. In this work, we present a novel world model formulation where the reward prediction only depends on a subset of structured, symbolic components of the whole latent state. Decoupling observation reconstruction and reward prediction allows us to learn world models that can adapt zero-shot, i.e. without further environment interactions, to new reward functions defined over the same symbolic state space. We discuss the main advantages and challenges of learning these neurosymbolic world models and demonstrate the strong generalisation properties of our approach over purely neural methods. 1 Introduction World models are motivated by the idea that agents can learn compact internal representations of their environments from experience and use them to guide decision-making (15). In reinforcement learning (RL), popular methods typically learn latent representations of states, dynamics, and reward functions directly from high-dimensional observations such as sequences of images (40; 17; 16; 18; 19; 20; 21). These learned representations can be used to optimise behaviour in several ways, including model-predictive control (20; 21), Monte-Carlo Tree Search (40), or actor-critic learning on imagined trajectories (16; 18; 19). Despite their ability to learn compact representations of environment dynamics, RL world models are not straightforwardly reusable across tasks. A particularly important case is task transfer under shared dynamics: the environment and its transition structure remain the same, but the reward function changes. This setting arises naturally in domains where the same environment supports multiple objectives, such as reaching different goals, placing objects in different configurations, or avoiding different hazards. Current RL world models are poorly suited to this setting because they typically learn a reward predictor jointly with the latent state and dynamics model. As a result, the learned representations may depend fully (40; 20; 21) or partially (16; 18; 19) on the rewards observed during training. When the task changes, the learned reward predictor remains tied to the training objective and is therefore not aligned with the new reward function, even though the underlying dynamics are unchanged. At the same time, the latent states learned by these models do not provide an interpretable interface on which a new reward function can be specified directly. Consequently, reusing a task-dependent world model for a new task typically requires learning or finetuning a new reward predictor using data from the test environment, or by relabeling training trajectories with the downstream task rewards (42). Unlike transition dynamics, reward functions often encode the desiderata of a user or task designer rather than an intrinsic property of the environment. Hence, they are frequently specified in terms of high-level semantic abstractions of the state, such as an agent reaching a target location, objects satisfying a desired configuration, or unsafe states being avoided (23). When a reward function is given, replacing it with a learned predictor can unnecessarily entangle the task objective with the latent dynamics model. However, directly using an explicit reward function requires the variables on which it depends to be available to the model. If these symbolic properties were available as part of the model state, the same learned model could be reused for new tasks whose reward functions are defined over the same properties, even when the reward specification itself changes. Crucially, exposing such symbolic properties should not require replacing learned latent dynamics with a fully symbolic model. The symbolic component only needs to capture the reward-relevant aspects of the state, while the remaining information required for control can remain encoded in latent representations learned from raw, subsymbolic observations. We instantiate this principle with Neurosymbolic World Models (NeSy-WMs), a new class of world models that combine latent dynamics learning with an explicit symbolic interface for reward evaluation. NeSy-WMs build on Recurrent State-Space Models (RSSMs) (17), which learn compact latent dynamics from high-dimensional observations. In addition, a NeSy-WM predicts a set of symbolic state properties that are chosen to capture the variables on which task rewards depend. Rewards are then predicted only through these symbolic properties, rather than directly from the latent state. This design disentangles latent modeling of environment dynamics from symbolic, interpretable reward prediction. Importantly, the symbolic state does not need to provide a complete symbolic description of the environment. It only needs to expose the reward-relevant properties, leaving the remaining information needed for prediction and control in the learned latent representation. Our contributions are as follows. (1) We formalise NeSy-WMs and their use for task transfer under shared dynamics, where new reward functions are defined over the same symbolic state properties as the training task. (2) We propose different symbolic supervision regimes for NeSy-WMs and evaluate their training performance and sample-efficiency against the state-of-the-art DreamerV3 (19). (3) We evaluate the zero-shot adaptation of NeSy-WMs to new test-time tasks in two different settings: one where pure imagination planning is performed to solve the task (no further learning), and one where the model is finetuned in imagination. (4) We provide an open-source and highly configurable PyTorch (35) implementation of NeSy-WMs. 2 Related Work Our work lies at the intersection of world models for model-based RL (MBRL) and probabilistic neurosymbolic learning. World models aim to predict, explicitly or implicitly, how an environment (the world) evolves under actions. While reward-free world models are typically trained from observations alone through predictive or generative objectives (4; 46; 33), several task-dependent world models used in MBRL instead aim for optimal control by learning latent dynamics directly through reward prediction, value prediction, and policy improvement objectives (40; 11; 20; 21). Generative world models (15; 1; 16; 18; 19) additionally shape their representations by reconstructing observations or sensory inputs. In particular, Recurrent State-Space Models (RSSMs) (17) can successfully model high-dimensional, partially observable environments and form the backbone of the Dreamer family of algorithms (16; 18; 19). Moreover, their probabilistic formulation is naturally compatible with probabilistic neurosymbolic methods. Neurosymbolic AI (14; 32) integrates symbolic knowledge and reasoning into machine learning models. Various probabilistic neurosymbolic models (29; 44) now integrate generative models to constrain generations and perform test-time interventions. They mostly use variational autoencoders (24) with a static (34; 12) or dynamic (13) encoding space combining latent and symbolic features. Their probabilistic nature is not only used for compatibility with generative models, but also to enable end-to-end training via probabilistic inference. A number of neurosymbolic approaches to world modeling have recently been proposed. 6 learn reusable neural world-model primitives from offline data and assembles them into a finite-state-machine world model to facilitate control across environments. 41 introduce an object-centric neurosymbolic world model that augments object representations with symbolic attributes to improve compositional generalisation in dynamics prediction. In both cases, the primary goal is to leverage symbolic structure to improve modeling and generalisation of environment dynamics, rather than our goal of enabling adaptation to new tasks under shared dynamics. Finally, our work is related to goal-conditioned RL (GCRL), which learns policies or value functions conditioned on a goal representation to solve a family of tasks (38; 2; 27). GCRL typically amortizes adaptation by jointly training on many goals and executing a goal-conditioned policy at test time. In contrast, we learn a NeSy-WM on a single training task and perform the desired task adaptation in the imagination of the world model, without the need to pre-train in a multi-task setting. 3 Methodology 3.1 Overview Our goal is task transfer under two main assumptions. (A1) Training and test reward functions are explicitly defined in terms of a limited number of high-level abstractions of the state. Such abstractions will be called the symbolic properties of the state and the vocabulary of such abstractions is assumed to be given by the user. (A2) The underlying dynamics of all tasks is the same. For instance, consider a navigation environment where the training task is to reach a fixed goal. Whether this task has been achieved only depends on the location of the agent. No other properties of the environment, e.g. agent orientation, matter for assessing task-achievement. A possible test task would be to shift the goal to any other position in the map; the reward function still depends only on the agent location, despite inducing a task that is completely different from the training task. Current world models learn purely latent representations of the state which cannot directly exploit user-specified reward functions, even when the test tasks are naturally defined on symbolic properties. As a result, a new reward predictor has to be learned, e.g. by manually relabeling replayed observations (42). NeSy-WMs avoid learning new reward functions without harming representation learning. They instead make the symbolic properties an explicit part of the world model. The symbolic properties form a bottleneck for the reward prediction to facilitate substitution of any other reward function defined on the same properties. More specifically, the model augments DreamerV3 (19) with neurosymbolic reward and continue predictors (Figure 1). Figure 1: Overview of NeSy-WMs. The symbolic-bottlenecked reward and continue predictors are shown in blue; the rest follows the RSSM structure of (17). At test time, new reward and continue functions can be defined over the same symbols to enable zero-shot planning or imagination training of a new actor-critic. Sequence model: ht=fϕ(zt−1,ht−1,at−1) h_t=f_φ(z_t-1,h_t-1,a_t-1) Encoder: zt∼qϕ(zt∣ht,xt) z_t q_φ(z_t h_t,x_t) Dynamics predictor: z^t∼pϕ(z^t∣ht) z_t p_φ( z_t h_t) NeSy reward predictor: r^t∼pϕsym(r^t∣zt,ht) r_t p^sym_φ( r_t z_t,h_t) NeSy continue predictor: c^t∼pϕsym(c^t∣zt,ht) c_t p^sym_φ( c_t z_t,h_t) Decoder: x^t∼pϕ(x^t∣zt,ht) x_t p_φ( x_t z_t,h_t) Here, xtx_t denotes the observation at time t, ata_t the action, hth_t the deterministic recurrent hidden state, ztz_t the stochastic latent state. The "continue" predictor models the probability that the episode continues at the next step, i.e. that the current state is non-terminal. The function fϕf_φ is the deterministic recurrent transition of the RSSM. Next, we specify the neurosymbolic predictors and how to compute their probabilities. Neurosymbolic predictors combine neural parametrisations with functions on symbols. The neurosymbolic reward predictor first neurally predicts a distribution pϕ(st∣zt,ht)p_φ(s_t z_t,h_t) over the symbolic properties using a simple MLP. The symbolic properties are assumed to be the only input to a given reward function psym(rt∣st)p^sym(r_t s_t) (Assumption (A1)). The uncertainty on sts_t parametrised by pϕ(st∣zt,ht)p_φ(s_t z_t,h_t) induces a probability distribution over the rewards rtr_t. Indeed, the distribution pϕsym(rt∣zt,ht)p_φ^sym(r_t z_t,h_t) is given by marginalising over the symbolic states sts_t as ∑stpsym(rt∣st)pϕ(st∣zt,ht) _s_tp^sym(r_t s_t)p_φ(s_t z_t,h_t) (1) or∫psym(rt∣st)pϕ(st∣zt,ht)dst, p^sym(r_t s_t)p_φ(s_t z_t,h_t)\ ds_t, (2) for discrete or continuous symbolic properties, respectively. Equations 1 and 2 also explain the generalisation ability of NeSy-WMs; psym(rt∣st)p^sym(r_t s_t) can simply be replaced with a new reward function. The same holds for the neurosymbolic continue predictor, which can be defined analogously. Example 3.1 (Discrete Navigation). Consider a discrete grid world where the symbolic property st=(xt,yt)s_t=(x_t,y_t) models the discrete coordinates of the player at time step t. The reward function psym(rt∣st)p^sym(r_t s_t) takes the coordinates sts_t as input and returns 11 if sts_t is the goal (xg,yg)(x_g,y_g) and 00 otherwise. If we model the distribution over the symbolic properties pϕ(st∣zt,ht)p_φ(s_t z_t,h_t) as two independent categorical distributions, then pϕsym(rt=1∣zt,ht)p_φ^sym(r_t=1 z_t,h_t) (Equation 1) reduces to Prϕ(xt=xg∣zt,ht)⋅Prϕ(yt=yg∣zt,ht). _φ(x_t=x_g z_t,h_t)· _φ(y_t=y_g z_t,h_t). Example 3.2 (Continuous Navigation). Consider a continuous grid world that has a continuous symbolic state st=(xt,yt)∈ℝ×ℝs_t=(x_t,y_t) ×R. Here, a goal is reached when the symbolic state sts_t lies in a square target area (x,y)∣ax−≤x≤ax+,ay−≤y≤ay+, \(x,y) a_x^-≤ x≤ a_x^+,a_y^-≤ y≤ a_y^+, \ . If we model the distribution over the symbolic properties pϕ(st∣zt,ht)p_φ(s_t z_t,h_t) as two independent Gaussian distributions xt x_t ∣zt,ht∼(μx(zt,ht),σx2(zt,ht)), z_t,h_t ( _x(z_t,h_t), _x^2(z_t,h_t)), yt y_t ∣zt,ht∼(μy(zt,ht),σy2(zt,ht)), z_t,h_t ( _y(z_t,h_t), _y^2(z_t,h_t)), then pϕsym(rt=1∣zt,ht)p_φ^sym(r_t=1 z_t,h_t) (Equation 2) is given by Prϕ(ax−<xt<ax+)⋅Prϕ(ay−<yt<ay+), _φ(a_x^-<x_t<a_x^+)· _φ(a_y^-<y_t<a_y^+), where Prϕ(ax−<xt<ax+) _φ(a_x^-<x_t<a_x^+) is computable using the Gaussian CDF (supplementary material). In the relevant case where the reward function checks whether sts_t satisfies logical properties, computing Equations 1 and 2 reduces to the standard inference tasks in probabilistic neurosymbolic AI: Weighted Model Counting (WMC) (7) for discrete sts_t or Weighted Model Integration (WMI) (5) for continuous sts_t. WMC and WMI have highly-optimised inference strategies (7; 45; 28) to better handle the intractability of computing Equations 1 and 2 (36). 3.2 Model Learning Our model-learning setup differs from DreamerV3 in two ways. (1) The neurosymbolic predictors pϕsym(rt∣zt,ht)p^sym_φ(r_t z_t,h_t) and pϕsym(ct∣zt,ht)p^sym_φ(c_t z_t,h_t) replace the standard MLPs and (2) symbolic supervision losses that can be added to ensure aligned symbolic representations. Standard Loss Given a sequence of observations x1:Tx_1:T, actions a1:Ta_1:T, rewards r1:Tr_1:T and continuation flags c1:Tc_1:T, we minimize the overall world model loss ℒW(ϕ) _W(φ) =qϕ[∑t=1T(βpredℒpredt(ϕ)+βdynℒdynt(ϕ) =E_q_φ [ _t=1^T ( _predL_pred^t(φ)+ _dynL_dyn^t(φ) +βrepℒrept(ϕ)+βrecℒrect(ϕ))], 48.00009pt+ _repL_rep^t(φ)+ _recL_rec^t(φ) ) ], where βpred,βdyn,βrep,βrec _pred, _dyn, _rep, _rec are fixed hyperparameters and the individual losses are defined as ℒrect(ϕ) _rec^t(φ) =−logpϕ(xt∣zt,ht), =- p_φ(x_t z_t,h_t), ℒpredt(ϕ) ^t_pred(φ) =−logpϕsym(rt∣zt,ht)−logpϕsym(ct∣zt,ht), =- p^sym_φ(r_t z_t,h_t)- p^sym_φ(c_t z_t,h_t), ℒdynt(ϕ) ^t_dyn(φ) =max(1,KL[sg(qϕ(zt∣ht,xt))||pϕ(zt∣ht)]), = (1,KL[sg(q_φ(z_t h_t,x_t))\,||\,p_φ(z_t h_t)]), ℒrept(ϕ) ^t_rep(φ) =max(1,KL[qϕ(zt∣ht,xt)||sg(pϕ(zt∣ht))]), = (1,KL[q_φ(z_t h_t,x_t)\,||\,sg(p_φ(z_t h_t))]), where sg(⋅)sg(·) is the stop-gradient operator. The dynamics loss ℒdynt(ϕ)L_dyn^t(φ) and representation loss ℒrept(ϕ)L_rep^t(φ) optimise the transitions via a KL divergence between prior and posterior. Reconstruction loss ℒrect(ϕ)L_rec^t(φ) and prediction loss ℒpredt(ϕ)L_pred^t(φ) ensure accurate observations and reward/continuation signals. Neurosymbolic predictors are differentiable. The neurosymbolic reward and continuation predictors are computed exactly via Equation 1 or 2. Exact neurosymbolic inference remains end-to-end differentiable (29) to propagate gradients from sts_t to ztz_t and hth_t. Consequently, the neurosymbolic pϕsym(rt∣zt,ht)p^sym_φ(r_t z_t,h_t) and pϕsym(ct∣zt,ht)p^sym_φ(c_t z_t,h_t) are drop-in replacements for the MLPs used by DreamerV3. Symbolic supervision losses make symbolic states identifiable. Solely relying on task reward prediction as a source of distant supervision for learning symbolic states can be insufficient to learn aligned representations (31). Identifiability of the symbols impacts task generalisation; if the symbols are learned wrongly, symbolic states may be assigned wrong rewards at test-time. For this reason, we consider and test three different supervision schemes (S). (S1) Full supervision. Direct symbolic supervision on all the states visited throughout world model training. This requires the agent to directly observe the symbolic properties s1:Ts_1:T. The per-step symbolic supervision loss is the negative log-likelihood of the ground-truth symbol ℒsymt(ϕ)=−logpϕ(st∣zt,ht).L^t_sym(φ)=- p_φ(s_t z_t,h_t). This kind of supervision can be seen as a form of privileged-information training (26; 22). However, such approaches either provide the full Markovian state during training, or a surrogate that is sufficient for reconstructing image observations (26). This is a stronger assumption than observing the symbolic properties considered in our work, which neither constitute a Markovian state representation nor suffice to reconstruct image observations. (S2) Partial supervision. A weaker assumption is to supervise only a subset ′S of all symbolic states. This limits the generalisation of NeSy-WMs to novel tasks involving relevant symbolic states, e.g. potential goal states, in ′S . Concretely, partial supervision adds the log-likelihood of the ground-truth symbols only when a state in ′S is visited: ℒsymt(ϕ)=−st∈′⋅logpϕ(st∣zt,ht).L^t_sym(φ)=- 1_ \s_t \· p_φ(s_t z_t,h_t). Intuitively, this corresponds to having sensors in limited parts of the environment. A potential issue is that the predictor may collapse to always predicting supervised symbols; after a reward intervention, a state incorrectly mapped to a rewarded supervised symbol may produce a spurious reward and cause transfer failure. A simple solution inspired by 30 is to add an entropy regularisation term to discourage overly confident predictions. For a symbolic state composed of M categorical variables, we minimise ℒentt(ϕ)=−βent1M∑j=1Mℋ[pϕ(st(j)∣zt,ht)]log|j|,L^t_ent(φ)=- _ent 1M _j=1^M H\! [p_φ(s_t^(j) z_t,h_t) ] |S_j|, where ℋH denotes Shannon entropy, jS_j is the domain of the j-th symbolic variable, and βent _ent controls the regularisation. (S3) No supervision. When no supervision is provided, the symbolic predictors can learn any mapping as long as it suffices to solve the training task. Nonetheless, we still consider this setting because the imposed symbolic structure may still help stabilise training and facilitate downstream finetuning. 3.3 Test-Time Interventions and Generalisation At training time, control follows the actor-critic scheme of DreamerV3, summarised in the supplementary material. Once a NeSy-WM is trained, it can be adapted to different tasks without access to the test environments. To do so, it is sufficient to replace the training predictors psym(rt∣st)p^sym(r_t s_t) and psym(ct∣st)p^sym(c_t s_t) in Equation 1 or 2 with the given test predictors. This requires new tasks to be defined over the same symbols (Assumption (A1)) and assumes fixed dynamics (Assumption (A2)). Given the replaced reward function, there are two options (O) to adapt a NeSy-WM to a novel test task. In both cases, there is no interaction with the test environment. (O1) Adaptation by pure planning. Monte-Carlo Tree Search (MCTS) (10; 25) can be employed to plan at test time using the learned model and the new reward/continue predictors without any further learning. We stress that this is not possible for purely neural models such as Dreamer, since they cannot replace the previously learned reward/continue predictors without a finetuning phase. While we experiment on environments with discrete action spaces, NeSy-WMs can also be used for planning in continuous action spaces using compatible algorithms, such as the Cross-Entropy Method (CEM) (37) or appropriate extensions of MCTS (9). (O2) Adaptation by imagination finetuning. NeSy-WMs can be finetuned on a new task entirely in the world model’s imagination without any further interaction with the environment. By manually replacing the predictors psym(rt∣st)p^sym(r_t s_t) and psym(ct∣st)p^sym(c_t s_t) as described above, a new actor-critic can be trained without collecting new task-specific data or performing the time-consuming replay-buffer relabeling required by previous Dreamer adaptation approaches (42). 3.4 Limitations The main limitation of our approach surrounds the alignment of the symbolic predictions pϕ(st∣zt,ht)p_φ(s_t z_t,h_t) with the ground-truth properties they are meant to represent. If the predicted symbolic distribution is misaligned, then intervening on the reward function can assign rewards according to the wrong symbolic interpretation, leading to adaptation failure. Symbolic supervision aligns the symbols with the intended semantics and facilitates reliable task transfer, while purely reward-based training may lead to symbolic shortcuts sufficient only for the training task. That is, NeSy-WMs make explicit the relation between the symbolic information provided during learning and the range of reward interventions that can be trusted at test time. Purely neural world models do not provide this possibility because they have no interpretable interface for directly specifying new reward functions. A separate limitation concerns the availability and sufficiency of the symbolic language. Our formulation assumes that training and test rewards are defined over the same symbolic properties. If a new task depends on properties outside this vocabulary, the latent RSSM state may still contain useful predictive information, but provides no interpretable mechanism for wiring the new reward function on top of it. Transfer is therefore limited to tasks whose rewards can be expressed sufficiently in the chosen symbolic vocabulary. 4 Experimental Evaluation In this section, we propose experiments to test the following three main claims (C) of our work. (C1) Sample-efficient training. We expect NeSy-WMs to improve sample efficiency with respect to DreamerV3 due to the provided structure of the reward/continue predictors, especially when some symbolic supervision is provided. (C2) Zero-shot adaptation by imagination planning. We expect NeSy-WMs to enable planning on new test tasks under the same dynamics without further learning, unlike existing finetuning approaches. (C3) Zero-shot adaptation by imagination finetuning. We expect NeSy-WMs to enable zero-shot finetuning on novel test tasks completely in the model imagination. This ability contrasts with existing adaptation settings of DreamerV3, which either require online test-task experience or replay buffer relabeling (42). To test our claims, we propose a set of corresponding experiments (E): (E1) Training sample efficiency. We train NeSy-WMs in fully supervised (S1), partially supervised (S2), and unsupervised (S3) settings to assess their training sample efficiency compared to standard DreamerV3. (E2) Zero-shot adaptation by planning. The trained NeSy-WMs from E1 are adapted to a variety of test tasks of increasing difficulty without any further training through MCTS planning in model imagination (O1). Since Dreamer adaptation requires further training, it cannot be compared in this setting. Instead, we compare against the same MCTS planning algorithm using a perfect environment simulator. (E3) Zero-shot adaptation by imagination finetuning. Using the models trained in E1, NeSy-WMs are adapted to a selected set of hard test tasks through imagination finetuning (O2). Specifically, NeSy-WMs always freeze their latent dynamics model and train a new actor-critic if symbols were fully supervised (S1) or partially supervised (S2) during training (new AC). If symbols were unsupervised during training (S3), NeSy-WMs instead learn a new actor-critic and reward/continue predictors through replay buffer relabeling following 42. For DreamerV3, we employ the same adaptation scheme by training a new actor-critic and reward/continue predictors from the relabeled buffer (new AC + R + C). We also evaluate a baseline where the entire DreamerV3 world model is unfrozen (WM unfrozen). Figure 2: Training tasks. The initial position of the agent is random. From left to right: MiniGrid-small, MiniGrid-large, a top view of MiniWorld-small and MiniWorld-large differing in movement granularity and goal area size, and Sokoban. Figure 3: Training results. Curves show the average return across three seeds and their standard error. Both NeSy-WMs and Dreamer start by an exploration phase in the MiniWorld environments and in MiniGrid-large (see supplementary material). MG-small MG-large MW-small MW-large Sokoban Challenge Full Partial Full Partial Full Partial Full Partial Full Partial Easy ±0.001.00\!±\!0.00 Δ=0.0 =0.0 ±0.001.00\!±\!0.00 Δ=2.7 =2.7 ±0.001.00\!±\!0.00 Δ=0.0 = 00.0 ±0.001.00\!±\!0.00 Δ=0.0 = 00.0 ±0.150.78\!±\!0.15 Δ=8.7 = 08.7 ±0.110.89\!±\!0.11 Δ=7.4 = 07.4 ±0.110.89\!±\!0.11 Δ=5.1 = 05.1 ±0.180.44\!±\!0.18 Δ=24.5 =24.5 ±0.001.00\!±\!0.00 Δ=0.0 = 00.0 ±0.001.00\!±\!0.00 Δ=1.1 = 01.1 Medium ±0.001.00\!±\!0.00 Δ=0.0 =0.0 ±0.170.67\!±\!0.17 Δ=0.0 =0.0 ±0.001.00\!±\!0.00 Δ=0.0 = 00.0 ±0.001.00\!±\!0.00 Δ=8.2 = 08.2 ±0.001.00\!±\!0.00 Δ=8.9 = 08.9 ±0.001.00\!±\!0.00 Δ=8.2 = 08.2 ±0.110.11\!±\!0.11 Δ=33.0 =33.0 ±0.000.00\!±\!0.00 Δ=0.–0 = 0.-- 0 ±0.001.00\!±\!0.00 Δ=1.8 = 01.8 ±0.001.00\!±\!0.00 Δ=0.2 = 00.2 Hard ±0.001.00\!±\!0.00 Δ=9.8 =9.8 ±0.170.67\!±\!0.17 Δ=9.0 =9.0 ±0.150.22\!±\!0.15 Δ=33.0 =33.0 ±0.180.44\!±\!0.18 Δ=20.0 =20.0 ±0.170.67\!±\!0.17 Δ=35.3 =35.3 ±0.170.67\!±\!0.17 Δ=16.5 =16.5 ±0.000.00\!±\!0.00 Δ=0.–0 = 0.-- 0 ±0.000.00\!±\!0.00 Δ=0.–0 = 0.-- 0 ±0.170.67\!±\!0.17 Δ=72.5 =72.5 ±0.170.67\!±\!0.17 Δ=21.8 =21.8 Table 1: MCTS evaluation of fully and partially supervised NeSy-WMs. MG and MW denote MiniGrid and MiniWorld, respectively. We report success rate ± standard error over nine episodes (three episodes for each of three trained models) and the optimality gap Δ=L−L⋆ =L-L , where L is the mean length of successful plans and L⋆L is the optimal plan length for that challenge. Figure 4: Imagination finetuning results. Curves show the average return across three seeds and their standard error. Adaptation is entirely offline; during finetuning, we periodically evaluate the policy in the test environment solely to measure performance. 4.1 Training Environments We describe the environments used for the training experiments in E1 (Figure 2), and their symbolic abstractions. MiniGrid 2D-FourRooms. We employ two versions of a 2D Four-Rooms environment implemented using MiniGrid (8). In particular, MiniGrid-small has four 3×33× 3 rooms, whereas MiniGrid-large has 5×55× 5 rooms. The observations are images of the whole grid, and the agent has 3 available actions: forward, turn-left, and turn-right. The agent only receives a positive reward of +1 if the goal is reached and 00 otherwise. At training time, the initial agent position is randomized, and the goal position is fixed as the center of the top-right room. The (x,y)(x,y) coordinate of the agent is used as the task-relevant symbolic abstraction. This abstraction is sufficient for reward prediction, but incomplete for control since the orientation of the agent also matters for choosing the correct action. In experiments with partial symbolic supervision, we only supervise the center of each room. MiniWorld 3D-FourRooms. We employ two versions of a 3D partially observable four-rooms environment with continuous state space, implemented in MiniWorld (8). Figure 2 shows a top-down view (the egocentric view is provided in the supplementary material). The action space matches MiniGrid, except that forward moves the agent by 0.30.3 and 0.150.15 world units in MiniWorld-small and MiniWorld-large, respectively, while turning rotates the agent by 30∘30 and 15∘15 . A reward of +1+1 is received only when the agent enters a goal region. During training, the initial position is randomized and the goal is fixed at the center of the top-right room. We use the continuous (x,y)(x,y) position as the symbolic abstraction. In experiments with partial symbolic supervision, we only supervise the four potential goal areas centered in each room. Sokoban. Here, the agent needs to push all the boxes to the indicated goal positions in a customized version of the game Sokoban (39). In this version, it only receives a reward (+1+1) when the task is completed. The agent has 9 available actions: move-dir and push-dir for each direction dir ∈ left, right, up, down, and a null action which does nothing. We employ a task-relevant symbolic abstraction comprising only the coordinates of the boxes, since the reward depends only on their locations. In experiments with partial symbolic supervision, we only provide supervision when all boxes are on training or test targets. 4.2 Test Tasks Planning Tasks For the experiments in E2, we select a set of test tasks for each training environment, visualized in the supplementary material. In particular, we take the center of the bottom-left room as test goal, opposite the training goal, for MiniGrid and MiniWorld, and move the boxes to opposite (top) target positions in Sokoban. The tasks are crafted to identify at what point the planning-adapted model fails as the distance between the initial position of the agent and the goal increases. Specifically, the Easy, Medium and Hard challenges for MiniGrid and MiniWorld feature an increasingly distant starting position of the agent from the goal position. Similarly, the Easy task in Sokoban requires placing only one box in a correct test-task configuration, while Hard requires moving all boxes to new target positions. Imagination Finetuning Tasks For the experiments in E3, we consider again the opposite map position with respect to training as the goal for MiniGrid and MiniWorld, and moving all three boxes to opposite target positions in Sokoban. While finetuning the model, we take snapshots of the policy to evaluate it on the test environment with randomized agent initial position at each evaluation episode. 4.3 Results (C1) NeSy-WMs improve the stability and sample efficiency of world-model training in sparse-reward environments, particularly under symbolic supervision. Figure 3 shows the training performance of NeSy-WMs under different supervision regimes S1-S3, compared to DreamerV3. Across the five training tasks, NeSy-WMs with either full or partial supervision mostly outperform the baseline in terms of sample-efficiency and stability. The only exception is the partially supervised agent in MiniWorld-large that converged some steps later. Supervised NeSy-WMs also converge significantly faster in MiniGrid-small, and outperform the baseline in MiniGrid-large and Sokoban, where Dreamer is unstable and does not converge. In the case where no symbolic supervision is provided, NeSy-WMs still perform comparably to Dreamer, outperforming it on MiniGrid-small but matching Dreamer’s failures on Sokoban and MiniGrid-large. The benefits provided by the supervision in these experiments suggest that symbol prediction may by itself be a strong learning signal for RL world models. We further investigated this hypothesis in an ablation reported in the supplementary material, where we show that our performance under symbolic supervision remains equally strong without reconstruction gradients, unlike Dreamer which dramatically fails. (C2) NeSy-WMs enable zero-shot MCTS test-time adaptation, but are constrained by the limitations of uninformed planning. Table 1 shows the results of the zero-shot MCTS planning with fully supervised and partially supervised NeSy-WMs on the selected test tasks. With the same planning budget of 256 tree expansions per step, the fully supervised model solves most of the Easy and Medium challenges in an optimal or close-to-optimal number of steps. However, pure planning is not enough to consistently solve the Hard tasks in MiniGrid-large, MiniWorld-small/large, and Sokoban. The partially supervised model generally achieves comparable performance, despite lower success rates in MiniGrid-small. These results were compared to planning with perfect environment simulators using the same planning budget and MCTS parameters. Aggregated across all challenges, fully and partially supervised NeSy-WMs achieve success rates of 75.6%75.6\% and 69.6%69.6\%, respectively, compared with 79.3%79.3\% for perfect-simulator planning. The complete comparison is reported in the supplementary material. While increasing the planning budget may improve the results, doing so with a learned model without value bootstrapping is slow (40) and prone to dynamics degradation (43; 3). (C3) NeSy-WMs outperform DreamerV3 in the zero-shot imagination finetuning setting. Figure 4 shows the adaptation performance of NeSy-WMs against the two Dreamer adaptation regimes. We remark that while supervised NeSy-WMs only need to train a fresh actor-critic in imagination without any replay buffer relabeling, both Dreamer baselines and unsupervised NeSy-WMs need to train new reward/continue predictors on the relabeled buffer. The results clearly show how Dreamer consistently fails to adapt to the test-time tasks, with the only exception being MiniGrid-small for the WM unfrozen regime. Conversely, fully supervised NeSy-WMs adapt almost immediately to all the test-time tasks, followed by partially supervised NeSy-WMs which only struggle in MiniWorld-large. Unsupervised NeSy-WMs also result in a substantially more reliable adaptation approach than Dreamer, successfully adapting in four out of the five challenges. This shows that even when symbolic supervision is not available, the structure of the neurosymbolic predictors can still make NeSy-WMs a convenient alternative for efficient task-adaptation by imagination finetuning. 5 Conclusion In this work, we presented Neurosymbolic World Models (NeSy-WMs) that build on top of generative world models by incorporating symbolic reward and continuation predictors over learned, task-relevant symbolic properties of the state. NeSy-WMs demonstrate a consistently more sample-efficient training than DreamerV3 on the sparse-reward tasks considered in this work. They can also be employed for zero-shot test-time planning on different tasks than the one they were trained on, i.e. new reward functions defined over the same set of symbolic properties. Finally, NeSy-WMs demonstrated effective zero-shot finetuning on novel tasks in imagination without any interaction with the test environment. Importantly, they consistently outperformed the common approach of finetuning DreamerV3 with a relabeled replay buffer. Overall, we showed how prior symbolic knowledge of the reward function can be incorporated into generative world models in a way that is simple, benefits the training process, and provides distinctive generalisation properties. Promising directions for extending NeSy-WMs include robust dynamics learning for reliable test-time planning over long horizons, testing other forms of symbolic alignment such as temporal consistency constraints, and exploration strategies that expose the model to the symbolic states needed for reliable task transfer in hard-exploration environments. 6 Acknowledgments This research has received funding from the KU Leuven Research Funds (C14/24/092) and from the Flemish Government under the "Onderzoeksprogramma Artificiële Intelligentie (AI) Vlaanderen" programme. References Alonso et al. (2024) E. Alonso, A. Jelley, V. Micheli, A. Kanervisto, A. Storkey, T. Pearce, and F. Fleuret Diffusion for world modeling: visual details matter in atari. Advances in Neural Information Processing Systems 37, p. 58757–58791. Cited by: §2. Andrychowicz et al. (2017) M. Andrychowicz, F. Wolski, A. Ray, J. Schneider, R. Fong, P. Welinder, B. McGrew, J. Tobin, P. Abbeel, and W. Zaremba Hindsight experience replay. Advances in neural information processing systems 30. Cited by: §2. Asadi et al. (2019) K. Asadi, D. Misra, S. Kim, and M. L. Littman Combating the compounding-error problem with a multi-step model. arXiv preprint arXiv:1905.13320. Cited by: §4.3. Assran et al. (2023) M. Assran, Q. Duval, I. Misra, P. Bojanowski, P. Vincent, M. Rabbat, Y. LeCun, and N. Ballas Self-supervised learning from images with a joint-embedding predictive architecture. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, p. 15619–15629. Cited by: §2. Belle et al. (2015) V. Belle, A. Passerini, and G. Van den Broeck Probabilistic inference in hybrid domains by weighted model integration. In Proceedings of the Twenty-Fourth International Joint Conference on Artificial Intelligence, IJCAI 2015, Buenos Aires, Argentina, July 25-31, 2015, p. 2770–2776. Cited by: §3.1. Cano et al. (2025) L. H. Cano, M. Perroni-Scharf, N. Dhir, A. Ramamurthy, and A. Solar-Lezama Neurosymbolic world models for sequential decision making. In Forty-second International Conference on Machine Learning, Cited by: §2. Chavira and Darwiche (2008) M. Chavira and A. Darwiche On probabilistic inference by weighted model counting. Artificial Intelligence 172 (6-7), p. 772–799. Cited by: §3.1. Chevalier-Boisvert et al. (2023) M. Chevalier-Boisvert, B. Dai, M. Towers, R. de Lazcano, L. Willems, S. Lahlou, S. Pal, P. S. Castro, and J. Terry Minigrid & miniworld: modular & customizable reinforcement learning environments for goal-oriented tasks. CoRR abs/2306.13831. Cited by: §4.1, §4.1. Couetoux (2013) A. Couetoux Monte carlo tree search for continuous and stochastic sequential decision making problems. Ph.D. Thesis, Université Paris Sud-Paris XI. Cited by: item O1. Coulom (2006) R. Coulom Efficient selectivity and backup operators in monte-carlo tree search. In International conference on computers and games, p. 72–83. Cited by: item O1. Danihelka et al. (2022) I. Danihelka, A. Guez, J. Schrittwieser, and D. Silver Policy improvement by planning with gumbel. In International Conference on Learning Representations, Cited by: §2. De Smet et al. (2023) L. De Smet, P. Z. Dos Martires, R. Manhaeve, G. Marra, A. Kimmig, and L. De Raedt Neural probabilistic logic programming in discrete-continuous domains. In Uncertainty in Artificial Intelligence, p. 529–538. Cited by: §2. De Smet et al. (2025) L. De Smet, G. Venturato, L. De Raedt, and G. Marra Relational neurosymbolic markov models. In Proceedings of the AAAI Conference on Artificial Intelligence, Vol. 39, p. 16181–16189. Cited by: §2. Garcez and Lamb (2023) A. d. Garcez and L. C. Lamb Neurosymbolic ai: the 3rd wave. Artificial Intelligence Review 56 (11), p. 12387–12406. Cited by: §2. Ha and Schmidhuber (2018) D. Ha and J. Schmidhuber Recurrent world models facilitate policy evolution. Advances in neural information processing systems 31. Cited by: §1, §2. Hafner et al. (2020) D. Hafner, T. Lillicrap, J. Ba, and M. Norouzi Dream to control: learning behaviors by latent imagination. In International Conference on Learning Representations, Cited by: §1, §1, §2. Hafner et al. (2019) D. Hafner, T. Lillicrap, I. Fischer, R. Villegas, D. Ha, H. Lee, and J. Davidson Learning latent dynamics for planning from pixels. In International conference on machine learning, p. 2555–2565. Cited by: §1, §1, §2, Figure 1. Hafner et al. (2021) D. Hafner, T. P. Lillicrap, M. Norouzi, and J. Ba Mastering atari with discrete world models. In International Conference on Learning Representations, Cited by: §1, §1, §2. Hafner et al. (2025) D. Hafner, J. Pasukonis, J. Ba, and T. Lillicrap Mastering diverse control tasks through world models. Nature 640 (8059), p. 647–653. Cited by: item (2), §1, §1, §2, §3.1. Hansen et al. (2022) N. Hansen, X. Wang, and H. Su Temporal difference learning for model predictive control. In International Conference on Machine Learning, PMLR, Cited by: §1, §1, §2. Hansen et al. (2024) N. Hansen, H. Su, and X. Wang TD-mpc2: scalable, robust world models for continuous control. In The Twelfth International Conference on Learning Representations, Cited by: §1, §1, §2. Huang et al. (2025) D. Huang, J. Wang, Y. Li, C. Xia, T. Zhang, and K. Zhang PIGDreamer: privileged information guided world models for safe partially observable reinforcement learning. In Forty-second International Conference on Machine Learning, Cited by: item S1. Icarte et al. (2022) R. T. Icarte, T. Q. Klassen, R. Valenzano, and S. A. McIlraith Reward machines: exploiting reward function structure in reinforcement learning. Journal of Artificial Intelligence Research 73, p. 173–208. Cited by: §1. Kingma and Welling (2014) D. P. Kingma and M. Welling Auto-encoding variational bayes. In 2nd International Conference on Learning Representations ICLR, Cited by: §2. Kocsis and Szepesvári (2006) L. Kocsis and C. Szepesvári Bandit based monte-carlo planning. In European conference on machine learning, p. 282–293. Cited by: item O1. Lambrechts et al. (2024) G. Lambrechts, A. Bolland, and D. Ernst Informed pomdp: leveraging additional information in model-based rl. Reinforcement Learning Journal. Cited by: item S1. Liu et al. (2022) M. Liu, M. Zhu, and W. Zhang Goal-conditioned reinforcement learning: problems and solutions. In Proceedings of the Thirty-First International Joint Conference on Artificial Intelligence, p. 5502–5511. Cited by: §2. Maene et al. (2025) J. Maene, V. Derkinderen, and P. Zuidberg Dos Martires KLay: accelerating arithmetic circuits for neurosymbolic ai. In The Thirteenth International Conference on Learning Representations, Cited by: §3.1. Manhaeve et al. (2018) R. Manhaeve, S. Dumancic, A. Kimmig, T. Demeester, and L. De Raedt Deepproblog: neural probabilistic logic programming. Advances in neural information processing systems 31. Cited by: §2, §3.2. Marconato et al. (2024) E. Marconato, S. Bortolotti, E. van Krieken, A. Vergari, A. Passerini, and S. Teso BEARS make neuro-symbolic models aware of their reasoning shortcuts. In Uncertainty in Artificial Intelligence, p. 2399–2433. Cited by: item S2. Marconato et al. (2023) E. Marconato, S. Teso, A. Vergari, and A. Passerini Not all neuro-symbolic concepts are created equal: analysis and mitigation of reasoning shortcuts. Advances in Neural Information Processing Systems 36, p. 72507–72539. Cited by: §3.2. Marra et al. (2024) G. Marra, S. Dumančić, R. Manhaeve, and L. De Raedt From statistical relational to neurosymbolic artificial intelligence: a survey. Artificial Intelligence 328, p. 104062. Cited by: §2. Micheli et al. (2023) V. Micheli, E. Alonso, and F. Fleuret Transformers are sample-efficient world models. In The Eleventh International Conference on Learning Representations, Cited by: §2. Misino et al. (2022) E. Misino, G. Marra, and E. Sansone Vael: bridging variational autoencoders and probabilistic logic programming. Advances in Neural Information Processing Systems 35, p. 4667–4679. Cited by: §2. Paszke et al. (2019) A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, et al. Pytorch: an imperative style, high-performance deep learning library. Advances in neural information processing systems 32. Cited by: item (4). Roth (1996) D. Roth On the hardness of approximate reasoning. Artificial intelligence 82 (1-2), p. 273–302. Cited by: §3.1. Rubinstein (1997) R. Y. Rubinstein Optimization of computer simulation models with rare events. European Journal of Operational Research 99 (1), p. 89–112. Cited by: item O1. Schaul et al. (2015) T. Schaul, D. Horgan, K. Gregor, and D. Silver Universal value function approximators. In International conference on machine learning, p. 1312–1320. Cited by: §2. Schrader (2018) M. B. Schrader Gym-sokoban. GitHub. Note: https://github.com/mpSchrader/gym-sokoban Cited by: §4.1. Schrittwieser et al. (2020) J. Schrittwieser, I. Antonoglou, T. Hubert, K. Simonyan, L. Sifre, S. Schmitt, A. Guez, E. Lockhart, D. Hassabis, T. Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature 588 (7839), p. 604–609. Cited by: §1, §1, §2, §4.3. Sehgal et al. (2024) A. Sehgal, A. Grayeli, J. J. Sun, and S. Chaudhuri Neurosymbolic grounding for compositional world models. In The Twelfth International Conference on Learning Representations, Cited by: §2. Sekar et al. (2020) R. Sekar, O. Rybkin, K. Daniilidis, P. Abbeel, D. Hafner, and D. Pathak Planning to explore via self-supervised world models. In International conference on machine learning, p. 8583–8592. Cited by: §1, item O2, §3.1, item C3, item E3. Talvitie (2017) E. Talvitie Self-correcting models for model-based reinforcement learning. In Proceedings of the AAAI conference on artificial intelligence, Vol. 31. Cited by: §4.3. Yang et al. (2020) Z. Yang, A. Ishay, and J. Lee NeurASP: embracing neural networks into answer set programming. In 29th International Joint Conference on Artificial Intelligence, IJCAI 2020, p. 1755–1762. Cited by: §2. Zeng et al. (2020) Z. Zeng, P. Morettin, F. Yan, A. Vergari, and G. Van den Broeck Scaling up hybrid probabilistic inference with logical and arithmetic constraints via message passing. In International Conference on Machine Learning, p. 10990–11000. Cited by: §3.1. Zhou et al. (2025) G. Zhou, H. Pan, Y. LeCun, and L. Pinto DINO-wm: world models on pre-trained visual features enable zero-shot planning. In International Conference on Machine Learning, p. 79115–79135. Cited by: §2.