Variational Causal Dynamics: Discovering Modular World Models from Interventions
Latent world models allow agents to reason about complex environments with
high-dimensional observations. However, adapting to new environments and
effectively leveraging previous knowledge remain significant challenges. We
present variational causal dynamics (VCD), a structured world model that
exploits the invariance of causal mechanisms across environments to achieve
fast and modular adaptation. By causally factorising a transition model, VCD is
able to identify reusable components across different environments. This is
achieved by combining causal discovery and variational inference to learn a
latent representation and transition model jointly in an unsupervised manner.
Specifically, we optimise the evidence lower bound jointly over a
representation model and a transition model structured as a causal graphical
model. In evaluations on simulated environments with state and image
observations, we show that VCD is able to successfully identify causal
variables, and to discover consistent causal structures across different
environments. Moreover, given a small number of observations in a previously
unseen, intervened environment, VCD is able to identify the sparse changes in
the dynamics and to adapt efficiently. In doing so, VCD significantly extends
the capabilities of the current state-of-the-art in latent world models while
also comparing favourably in terms of prediction accuracy.