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

TL;DR

This paper introduces WM3C, a novel reinforcement learning framework that uses language-guided compositional causal components to improve generalization to unseen environments by disentangling dynamics and leveraging causal reasoning.

Contribution

WM3C is the first approach to incorporate language-guided compositional causal components for enhanced RL generalization and provides theoretical guarantees for component identification.

Findings

WM3C outperforms existing methods in identifying latent processes.

It improves policy learning and generalization to unseen tasks.

Demonstrated effectiveness on simulations and robotic manipulation.

Abstract

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). This novel framework enhances RL generalization by learning and leveraging compositional causal components. Unlike previous approaches focusing on invariant representation learning or meta-learning, WM3C identifies and utilizes causal dynamics among composable elements, facilitating robust adaptation to new tasks. Our approach 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. Experiments on numerical simulations and real-world robotic manipulation tasks demonstrate that WM3C significantly outperforms existing methods in identifying latent processes, improving policy learning, and generalizing to unseen tasks.