Neural Bayesian Networks as Continual World Models for Reinforcement Learning
Abstract
World models are central to sample-efficient Reinforcement Learning (RL), yet most are monolithic and, even when trained and updated online, they cannot answer the question that continual learning raises: ``which part of the world changed?'' Bayesian networks answer exactly that, splitting the dynamics into local mechanisms, one per variable, arranged on a Directed Acyclic Graph (DAG). However, their classical formulation is tabular, which has kept them out of modern world modeling: queries are answered one value at a time and continuous variables are poorly supported. We represent the environment dynamics with Neural Bayesian Networks (NBNs): Bayesian Networks whose mechanisms are learned conditional densities, which support both continuous and discrete variables and answer queries in batches. To make them continual world models, we build a pipeline that separates changes in a mechanism from shifts in the agent’s behavior, revises only the mechanisms that changed, and plans by inference in the network. In a preliminary experiment on a toy MDP with non-Gaussian dynamics, our pipeline flags only the changed mechanism and keeps its model of the environment dynamics accurate through behavior shifts and mechanism drifts. Released code: https://github.com/Giovannibriglia/nbn-continual-world-models.