← PROJECTS
Feb 2025 · World Models · Causal Learning · ICLR 2025

WM3C: World Modeling with Compositional Causal Components

A novel reinforcement learning framework that enhances generalization to unseen environments through language-guided compositional causal components. Accepted at ICLR 2025.

Xinyue Wang · Biwei Huang
Compositional generalization in reinforcement learning.
Compositional generalization in reinforcement learning. Language guides the decomposition of the latent space into causal components that can be recombined for environments the agent has never seen.

The problem

Generalization in reinforcement learning remains a significant challenge, especially when agents encounter novel environments with unseen dynamics. Drawing inspiration from human compositional reasoning — where known components are reconfigured to handle new situations — we introduce World Modeling with Compositional Causal Components (WM3C), a framework that enhances RL generalization by learning and leveraging compositional causal components.

The method

WM3C integrates language as a compositional modality to decompose the latent space into meaningful components, and provides theoretical guarantees for their unique identification under mild assumptions. Our practical implementation uses a masked autoencoder with mutual information constraints and adaptive sparsity regularization to capture high-level semantic information and effectively disentangle transition dynamics.

Key highlights

Identifiability. Theoretical guarantees provide a solid foundation for compositional causal learning.

State of the art. WM3C achieves state-of-the-art performance in generalizing to unseen environments.

Simulated and real. Significant improvements in both simulated and robot manipulation tasks.