Source-linked AI summary
Towards Zero-Shot Task Transfer with Neurosymbolic World Models
Isidoro Tamassia, Lennert De Smet, Giuseppe Marra
TL;DR
Task-dependent neural world models are difficult to reuse when only the reward changes. NeSy-WMs expose reward-relevant symbolic state properties, enabling zero-shot adaptation to new rewards and outperforming DreamerV3 in the reported experiments.
Problem
Neural world models are poorly suited to transferring across tasks with shared dynamics but changing reward functions because objectives become entangled with latent dynamics.
Method
NeSy-WMs combine learned latent dynamics with explicit symbolic reward and continuation predictors over reward-relevant state properties.
Results
NeSy-WMs train more sample-efficiently than DreamerV3 and support effective zero-shot planning and imagination finetuning on novel tasks without test-environment interaction.
Takeaways & Limitations
An explicit symbolic interface supports task transfer when new reward functions are defined over the same symbolic state properties.
Takeaways & Limitations
Reliable task transfer depends on aligning predicted symbolic distributions with the intended ground-truth properties; misalignment can cause adaptation failure.
Abstract
from arXiv · showhide
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
The paper introduces Neurosymbolic World Models, which separate learned latent dynamics from reward prediction through explicit symbolic, reward-relevant state properties, enabling zero-shot transfer across tasks with shared dynamics. It formalizes this setting and evaluates supervision regimes, training efficiency, and zero-shot adaptation against DreamerV3.
- Motivation: Neurosymbolic World Models address task transfer under shared dynamics, where the environment remains unchanged but the reward function changes.Existing world models learn reward predictors jointly with latent states and dynamics, tying representations and rewards to the training objective.
- Motivation: Reward functions often express user-specified goals through high-level semantic state abstractions rather than intrinsic environment properties.Examples include reaching target locations, achieving object configurations, and avoiding unsafe states.
- Method: Because rewards depend on symbolic properties rather than the full latent state, the learned model can be reused for new reward functions defined over the same properties.The symbolic state need expose only reward-relevant properties, not a complete symbolic description of the environment.
- Method: NeSy-WMs combine RSSM-based latent dynamics with an explicit symbolic interface that predicts reward-relevant state properties.Rewards are predicted only from these symbolic properties, while other information remains in the learned latent state.
- Contributions: The paper formalizes shared-dynamics task transfer and evaluates symbolic supervision regimes, training performance, sample efficiency, and zero-shot adaptation against DreamerV3.The evaluations cover two different settings and test adaptation to new tasks without additional environment interactions.
2 Related Work
The work connects model-based reinforcement-learning world models with probabilistic neurosymbolic learning. It also relates to neurosymbolic world modeling and goal-conditioned reinforcement learning while differing in how task adaptation is trained and performed.
- World Models: World models predict environment evolution, while task-dependent MBRL models learn latent dynamics through reward, value, and policy-improvement objectives.Reward-free models are typically trained from observations using predictive or generative objectives.
- Probabilistic Neurosymbolic Learning: Probabilistic neurosymbolic models integrate symbolic knowledge and reasoning with generative models, enabling constrained generation and test-time interventions.Their probabilistic formulation also supports end-to-end training via probabilistic inference.
- Neurosymbolic World Models: Prior neurosymbolic world models learn reusable neural primitives for finite-state control or augment object representations with symbolic attributes.Cano et al. use offline data, whereas Sehgal et al. introduce an object-centric formulation.
- Goal-Conditioned Reinforcement Learning: Goal-conditioned RL conditions policies or value functions on goal representations and typically amortizes adaptation by training jointly across many goals.The cited work instead learns a NeSy-WM on a single training task and adapts in imagination without multi-task pre-training.
3 Methodology
NeSy-WMs augment DreamerV3 with explicit symbolic properties that bottleneck reward and continuation prediction, enabling zero-shot replacement of reward functions under shared dynamics and vocabulary assumptions. Their transfer reliability depends on aligning learned symbols with intended semantics and ensuring new rewards are expressible over the available symbolic properties.
- Assumptions: Task transfer assumes training and test rewards use a user-provided vocabulary of high-level symbolic properties, while all tasks share the same underlying dynamics.These are the two stated assumptions, A1 and A2.
- Model: NeSy-WMs augment DreamerV3 with neurosymbolic reward and continuation predictors, making symbolic properties an explicit bottleneck while preserving representation learning.The symbolic bottleneck allows reward and continuation functions to be substituted over the same properties.
- Neurosymbolic predictors: The reward predictor neurally infers a distribution over symbolic properties, then combines it with a symbolic reward function by marginalizing over possible symbolic states.The symbolic reward function receives only the symbolic properties under Assumption A1.
- Task transfer: At test time, replacing the symbolic reward and continuation functions enables zero-shot planning or imagination training for a new actor-critic without learning new reward predictors.The same symbolic interface supports new reward functions defined over the existing properties.
- Learning and supervision: Exact neurosymbolic inference is end-to-end differentiable, making the predictors drop-in replacements for DreamerV3 MLPs; symbolic supervision losses can additionally align learned representations.The method considers full supervision and other supervision schemes because reward-only learning may not identify the intended symbols.
- Limitations: Misaligned symbolic predictions can assign test-time rewards to the wrong interpretations, while rewards outside the available symbolic vocabulary cannot be wired to an interpretable transfer mechanism.Symbolic supervision facilitates reliable transfer, but the formulation remains limited to rewards expressible over shared symbolic properties.
4 Experimental Evaluation
The evaluation tests sample-efficient training and two forms of zero-shot task adaptation: imagination planning and imagination finetuning. NeSy-WMs generally improve training stability and efficiency under symbolic supervision, support zero-shot planning with limitations, and outperform DreamerV3 in imagination finetuning.
- Experimental setup: The experiments test sample efficiency, zero-shot imagination planning, and zero-shot imagination finetuning across MiniGrid, MiniWorld, and Sokoban tasks.Training compares fully supervised, partially supervised, and unsupervised NeSy-WMs with DreamerV3; planning uses MCTS without further training, while finetuning adapts models in imagination.
- Training sample efficiency: NeSy-WMs with full or partial symbolic supervision mostly outperform DreamerV3 in sample efficiency and stability across five training tasks.Supervised NeSy-WMs converge significantly faster in MiniGrid-small and outperform DreamerV3 in MiniGrid-large and Sokoban, where DreamerV3 is unstable and does not converge; partial supervision converges later only in MiniWorld-large.
- Zero-shot adaptation by planning: 75.6% and 69.6% success rates for fully and partially supervised NeSy-WMs, respectively, trail 79.3% for perfect-simulator MCTS planning.With 256 tree expansions per step, fully supervised models solve most Easy and Medium challenges optimally or near-optimally, but pure planning is not consistently successful across all challenges.
- Zero-shot adaptation by imagination finetuning: NeSy-WMs outperform DreamerV3 in zero-shot imagination finetuning, while supervised models require only a newly trained actor-critic in imagination.DreamerV3 baselines and unsupervised NeSy-WMs additionally require replay-buffer relabeling, whereas NeSy-WMs freeze their latent dynamics model during adaptation.
5 Conclusion
Neurosymbolic World Models combine generative world models with symbolic reward and continuation predictors, enabling sample-efficient training and zero-shot planning on new reward functions. Key extensions include more robust long-horizon dynamics learning, alternative symbolic alignment constraints, and exploration strategies for difficult task transfer.
- Contributions: NeSy-WMs incorporate symbolic reward and continuation predictors over learned, task-relevant symbolic state properties.They build on generative world models while structuring reward and continuation prediction around symbolic components.
- Contributions: NeSy-WMs achieve consistently more sample-efficient training than DreamerV3 on the sparse-reward tasks studied.The comparison applies to the sparse-reward tasks considered in this work.
- Contributions: NeSy-WMs support zero-shot test-time planning on tasks different from their training task, including new reward functions.This transfer occurs without further environment interactions, as established in the paper context.
- Future directions: Future work should improve dynamics learning for reliable test-time planning over long horizons.The conclusion identifies robust dynamics learning as a promising extension.
- Future directions: Other priorities are temporal consistency constraints for symbolic alignment and exploration strategies that expose symbolic states needed for reliable transfer in hard-exploration environments.These directions target broader symbolic alignment and more reliable task transfer.