Modeling Unseen Environments with Language-guided Composable Causal Components in Reinforcement Learning

Generalization in reinforcement learning (RL) 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 novel framework that enhances RL generalization by learning and leveraging compositional causal components.

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.

Coming Soon!

Thank you for your interest in the WM3C project. We’re currently working on expanding this website!

Our paper has been accepted at ICLR 2025, and we’re excited to share more details about our work in the coming weeks. In the meantime, you can check the links above for our paper.

Key Highlights

• Theoretical guarantees provide a solid foundation for compositional causal learning
• WM3C achieves state-of-the-art performance in generalizing to unseen environments
• Experimental results demonstrate significant improvements in both simulated and robot manipulation tasks

Stay tuned for updates as we continue to develop this resource!