Towards Generalizable Reinforcement Learning via Causality-Guided Self-Adaptive Representations

TL;DR

CSR addresses generalizable reinforcement learning across tasks with evolving dynamics, extending beyond simple distribution shifts to include environment-space changes. It builds a causal, self-adaptive representation framework that introduces a domain-specific change factor $\bm{\theta}_i$ and potential state additions $\bm{s}^{\text{add}}$, guided by a structured causal graph with masks $D$. The approach employs a three-step adaptation (detection, expansion, pruning) within a Dreamer-based objective $\mathcal{J} = \mathcal{J}_{\text{rec}} - \mathcal{J}_{\text{KL}} + \mathcal{J}_{\text{reg}}$, achieving low-cost policy transfer and strong generalization across CartPole, CoinRun, and Atari. Empirical results demonstrate that incorporating explicit causal structure and self-adaptive expansion yields superior adaptation speed and performance compared to multiple baselines, highlighting the practical impact for robust, scalable generalization in RL.

Abstract

General intelligence requires quick adaption across tasks. While existing reinforcement learning (RL) methods have made progress in generalization, they typically assume only distribution changes between source and target domains. In this paper, we explore a wider range of scenarios where not only the distribution but also the environment spaces may change. For example, in the CoinRun environment, we train agents from easy levels and generalize them to difficulty levels where there could be new enemies that have never occurred before. To address this challenging setting, we introduce a causality-guided self-adaptive representation-based approach, called CSR, that equips the agent to generalize effectively across tasks with evolving dynamics. Specifically, we employ causal representation learning to characterize the latent causal variables within the RL system. Such compact causal representations uncover the structural relationships among variables, enabling the agent to autonomously determine whether changes in the environment stem from distribution shifts or variations in space, and to precisely locate these changes. We then devise a three-step strategy to fine-tune the causal model under different scenarios accordingly. Empirical experiments show that CSR efficiently adapts to the target domains with only a few samples and outperforms state-of-the-art baselines on a wide range of scenarios, including our simulated environments, CartPole, CoinRun and Atari games.

Figures (21)

• Figure 1: Environmental changes may or may not necessitate retraining RL agents, as illustrated on different variations of CoinRun. Changes in the amount and shape of obstacles from (a) to (b) do not prevent the agent from completing the task, while deadly holes and enemies introduced in (c) necessitate retraining.
• Figure 2: Efficient policy adaptation through the CSR framework. For each target task, we first use the prediction error, $\mathcal{L}_{\text{pred}}$, to determine whether it involves distribution shifts or space shifts. We then adjust the model accordingly by updating the task-specific change factor ${\bm{\theta}}_i$, or by adding new variables. Finally, we conduct causal graph pruning that removes variables unnecessary for the current task. Based on such compact causal representations, we can efficiently implement policy adaptation in a self-adaptive manner.
• Figure 3: Experimental results that answer the key questions in Section \ref{['sec:exp']}: (Q1) CSR demonstrates the best generalization capability compared to baseline methods in (a) Simulation and (b) CoinRun; (Q2) CSR with structural embeddings $D$ significantly outperforms CSR without $D$ in Atari games; (Q3) The SA expansion strategy yields the highest normalized average training episodic return in our experiments.
• Figure 4: The estimated $\hat{\theta}_i^{{\bm{s}}}$ is a monotonic function of the ground-truth values in our simulations.
• Figure 5: A graphical illustration of the generative environment model and the two types of changes addressed by CSR. (a) Source task; (b) Distribution shift scenario where the causal diagram remains unchanged but the value of ${\bm{\theta}}^r$ differs; (c) Space expansion scenario involving the emergence of new variable ${\bm{s}}_{2,t}$. Grey nodes denote observed variables, white nodes represent unobserved variables, and red nodes highlight the changing components in the target task compared to the source task.

Theorems & Definitions (5)

• proof
• proof
• proof
• proof
• proof