Back to AI Research

AI Research

Towards Zero-Shot Task Transfer with Neurosymbolic... | AI Research

Key Takeaways

  • Towards Zero-Shot Task Transfer with Neurosymbolic World Models introduces a method to enable reinforcement learning agents to adapt to new tasks without fur...
  • 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.
Paper AbstractExpand

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.

Towards Zero-Shot Task Transfer with Neurosymbolic World Models introduces a method to enable reinforcement learning agents to adapt to new tasks without further environment interaction. By decoupling the learning of environment dynamics from reward prediction, the authors allow agents to reuse a single world model for multiple objectives, provided those objectives are defined over the same symbolic state properties.

The Problem with Task-Dependent Models

State-of-the-art reinforcement learning models typically learn latent representations that are tightly coupled with the specific reward function used during training. Because these latent states are uninterpretable, they cannot be easily repurposed when a task changes, even if the underlying environment dynamics remain identical. To adapt to a new goal, current methods usually require either retraining or expensive data relabeling. This paper addresses this by proposing a way to separate the "what" of the environment (dynamics) from the "why" of the task (rewards).

How Neurosymbolic World Models Work

The authors propose Neurosymbolic World Models (NeSy-WMs), which build upon the Recurrent State-Space Model (RSSM) architecture. The key innovation is the introduction of a symbolic bottleneck for reward and continuation prediction.
Instead of predicting rewards directly from the entire latent state, the model is trained to predict a set of user-defined symbolic properties—such as an agent's coordinates or object configurations. These symbolic properties then serve as the sole input for the reward function. Because the reward predictor is now a separate, explicit component, a user can swap it out for a new reward function at test time without needing to retrain the underlying dynamics model. This allows for zero-shot adaptation through either pure planning or imagination-based finetuning.

Supervision and Training

The researchers identify that relying solely on reward signals to learn these symbolic states can be insufficient, as the model might learn "shortcuts" that only work for the training task. To ensure the symbolic states are identifiable and aligned with the intended semantics, they evaluate three supervision regimes:

  • Full supervision: Providing ground-truth symbolic labels for all states during training.

  • Partial supervision: Providing labels only for a subset of states, which the authors suggest mimics having sensors in specific parts of an environment.

  • No supervision: Relying on the model to learn the mapping, which may still provide stability for downstream tasks.

Limitations and Considerations

The authors note that the primary limitation of this approach is the alignment between the predicted symbolic distribution and the ground-truth properties. If the model incorrectly maps a state to a symbolic property, the agent will receive incorrect rewards during test-time interventions, leading to failure. Furthermore, the approach assumes that both training and test tasks are defined over the same vocabulary of symbolic properties. If a new task requires understanding properties outside of this predefined set, the latent model may not be able to accommodate the change.

Comments (0)

No comments yet

Be the first to share your thoughts!