Compete and Compose: Learning Independent Mechanisms for Modular World Models

TL;DR

COMET introduces a modular world model that learns independent, reusable mechanisms through competition and composition, enabling efficient transfer and adaptation across diverse environments with minimal supervision.

Contribution

The paper proposes COMET, a novel two-phase training method that encourages the emergence of independent mechanisms for better transferability and interpretability in world models.

Findings

COMET captures recognizable mechanisms without supervision.

COMET adapts efficiently to new environments with varying objects.

Outperforms baseline models in transfer and adaptation tasks.

Abstract

We present COmpetitive Mechanisms for Efficient Transfer (COMET), a modular world model which leverages reusable, independent mechanisms across different environments. COMET is trained on multiple environments with varying dynamics via a two-step process: competition and composition. This enables the model to recognise and learn transferable mechanisms. Specifically, in the competition phase, COMET is trained with a winner-takes-all gradient allocation, encouraging the emergence of independent mechanisms. These are then re-used in the composition phase, where COMET learns to re-compose learnt mechanisms in ways that capture the dynamics of intervened environments. In so doing, COMET explicitly reuses prior knowledge, enabling efficient and interpretable adaptation. We evaluate COMET on environments with image-based observations. In contrast to competitive baselines, we demonstrate that COMET captures recognisable mechanisms without supervision. Moreover, we show that COMET is able to adapt to new environments with varying numbers of objects with improved sample efficiency compared to more conventional finetuning approaches.