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

TL;DR

This work tackles reinforcement learning generalization to unseen environments by introducing WM3C, a framework that learns composable causal components guided by language. It establishes block-wise identifiability for language-controlled latent factors and presents a learning objective that combines reconstruction, dynamics, mutual information regularization, and sparsity to disentangle causal components. Empirical results on synthetic data and Meta-World robotic tasks show WM3C identifies latent components more accurately, generalizes to novel task compositions, and enables faster policy learning with interpretable interventions on language-controlled latents. Overall, WM3C demonstrates that language-guided, modular causal representations can significantly improve adaptation and transfer in complex RL settings.

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.

Figures (12)

• Figure 1: Illustrations of environment models
• Figure 2: Illustrations of environment models.
• Figure 3: Language-controlled composable components identification results on known tasks and novel tasks. Left: Coefficient of determination ($R^2$) using kernel ridge regression to regress estimated latents on true latents. Middle: Average $R^2$ over the three language-controlled composable components during training, regressing estimated latents on true latents (the shaded areas represent the standard deviation across three runs). Right: Average $R^2$ of model imagination over time during test-time on unseen tasks, where latents are new combinations of known composable components. All results are reported across three runs with different seeds.
• Figure 4: Learning curves on Meta-world: Average success rate and $9$ specific task success rates. Results are reported with three seed runs and $10$ episodes for each task evaluation.
• Figure 5: Adaptation curves on 9 unseen tasks, including 7 tasks that are recombination of known components and 2 tasks have unknown components (handle-press-side and coffee-button). DreamerV3 and MT-SAC are full-parameter finetuned while WM3C is only tuned with the dynamics module. Results are reported with three seed runs and $10$ episodes for each task evaluation.