Modeling Unseen Environments with Language-guided Composable Causal Components in Reinforcement Learning
Authors: Xinyue Wang
Biwei Huang
Venue: ICLR 2025 Poster (2025)
Keywords: Model-based Reinforcement Learning
World Model
Causal Representation Learning
Generalization in Reinforcement Learning
Primary area: reinforcement learning

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.

Published as a conference paper at ICLR 2025
MODELING UNSEEN ENVIRONMENTS LANGUAGE-GUIDED COMPOSABLE CAUSAL PONENTS IN REINFORCEMENT LEARNING
WITH
COM-
Xinyue Wang', Biwei Huang!
University of California San Diego {xiw159, bih007Jeucsd. edu
ABSTRACT
Generalization in reinforcement learning (RL) remains significant challenge, es- pecially when agents encounter novel environments with unseen dynamics. Draw- ing inspiration from human compositional reasoning-wiir where known components are reconfigured to handle new situations-a we introduce World Modeling with Compositional Causal Components (WM3C). This novel framework enhances RL generalization by learning and leveraging compositional causal components. Un- like previous approaches focusing invariant representation learning meta- learning, WM3C identifies and utilizes causal dynamics among composable ele- ments, facilitating robust adaptation
tasks. Our approach integrates lan- guage as compositional modality to decompose the latent space into meaningful components and provides theoretical guarantees for their unique identification un- der mild assumptions. Our practical implementation uses masked autoencoder with mutual information constraints and adaptive sparsity regularization 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
INTRODUCTION
Reinforcement learning (RL) has rapidly progressed, driving innovations in domains such as game playing, robotics, and autonomous driving (Silver al., 2018 Vinyals et 2019 Shi al. 2022 Kiran et al. 2020). Deep reinforcement learning (DRL) methods, including Deep Q-Networks (DQN), Soft Actor-Critic (SAC), and Proximal Policy Optimization (PPO), have addressed vari- ous challenges in RI., such stability in training, exploration in large state spaces, and efficient policy optimization (Haarnoja et al.) 20188; Schulman et al.) 2017 Mnih et al. 2015 2016 Fuji- moto et al. 20188. These breakthroughs underscore the pivotal role of DRL in advancing artificial intelligence.
Despite these substantial advancements, of the most pressing issues of DRL is the generaliza- tion of learned policies to novel, unseen environments (Gamrian & Goldberg 2018 Song et al.) 2019 Cobbe al. 2018). For example, the policy excels in push ball to place A might perform notoriously poorly in the task push ball to place B. This limitation is primarily due to overfitting to specific training environments. Especially when the agent only receive visual input in partially observable environment, capturing the changes of the observation function and reward function in the new environment becomes even harder. Thus, locating the change in unseen environment and adapting learned knowledge to accommodate it are crucial for generalizable agent.
Previous methods address the challenge of accommodating changes from different perspectives. Methods like data augmentation and visual encoders improve model robustness to observation func- tion changes by attempting to incorporate potential visual changes in the training domain (Lee al.)
2019
Hansen & Wang 2020: Yuan et al. 2022 Nair et al. 2022 However, these methods of-
Project wehsite: https: : www. http://wwwc.arwwaagggg.gggg wttp://..oooooooooooooooo http://wwwwcrrraaaagoooooooorrrr
wm33c

---

Published as a conference paper at ICLR 2025
ten rely on extensive domain knowledge to design effective augmentations and can struggle with changes that were not anticipated during training, e.g. changes in the state space and its dynamic structure. Additionally, approaches such as invariant repressnttioi learning and meta-reinforceemet learning learn task-agnostic information that is stable across multiple training tasks to generalize to novel observations (Zhang et al.] 2020b at Duan al.) 2016: Finn et al. 2017). While the former extracts invariant features that improve the agent's 's performance in new environments by leverag- ing consistent and transferable information, it may be limited in its ability to handle environments with fundamentally different underlying dynamics. The latter optimizes meta- parameters that enable rapid adaptation and efficient learning in novel tasks, but it can be computationally expensive and may not always guarantee quick convergence in highly variable environments, Meanwhile, advances in model-based reinforcement learning, including DreamerV3 (Hafner et al. 2023) and TD-MPC2 (Hansen et al. 2023), exhibit hyper-parameter robustness abilities across different kinds of environ- ments. Despite their data efficiency, the learned world model may not accurately capture all aspects of the underlying data generation process, leading to suboptimal planning and poor generalization to new situations.
To this end, ask the question:
What does it take to learn generalizable world model? Humans understand new concepts and generalize from known structures to unknown situations in a compo- sitional way (Tenenbaum al.) 2011] Marcus 2003 Gentner & Markman, 1997). For example, we learn push ball to place A as learning the composable components push, ball, place A and their relationships rather than as single, indivisible task. This compositional understanding allows us to apply the learned components to new contexts, such as push puck to place A, reusing familiar ele- ments in a new configuration. By learning these components and their dynamics, can efficiently adapt to new tasks and environments.
A critical aspect of achieving this compositional generalization is understanding the environment's causal structure. A causal system is characterized by modularity and sparsity (Pearl, 2009; PPteer eet al. 2017 Glymour et al., 2019), which naturally compositionally generalizable. This means that each component can be learned independently, and then recombined with minimal changes to handle new tasks. For instance, understanding the causal relationship between the action push and the object ball in one scenario be reused to understand the relationship between 'push' and 'puck' in another. This modular approach aligns with human learning, where we reuse learned components across different contexts, enhancing our ability to generalize.
By drawing the connection between the causal system and compositional generalization, we propose World Modeling with Compositional Causal Components (WM3C), framework that enables the agent to identify composable components, learn the causal dynamics among them, and utilize them for efficient training and adaptation. While previous work has explored causal representation learn- ing to enhance reinforcement learning (Huang et al., [022] Liu et al.) 2023 Feng & Magliacane) 2024), none have focused on improving generalization via learning composable causal components and understanding their dynamics. To identify these composable causal components, we leverage another compositional modality, language (Fodor & Pylyshyn| 1988). We theoretically show that under mild and reasonable assumptions, composable causal components and their dynamics can be uniquely identified. We further provide an algorithm for learning world model that incorporates these components and demonstrate how this approach generalizes to unseen tasks. The effectiveness of our approach is validated by achieving staeeof-thhe-a performance on a numerical simulation dataset and collection of robot manipulation in the Meta-Woo World environment.
2 WORLD MODEL WITH COMPOSITIONAL CAUSAL COMPONENTS
In the following, we start by giving the motivation and intuition behind our formulation. Then, we present the identification theory that guarantees their correct identification. We further demonstrate the world model learning framework based on the environment model and our identification results.
2.1 MOTIVATION AND INTUITIONS
Consider transfer learning scenario in visual-based RL, where agent interacts with various fa- miliar environments and then adapts to new, unseen ones. The challenge here is not just mastering individual tasks but leveraging past experiences to quickly adapt to novel tasks. Humans excel at this
2

---

Published as a conference paper at ICLR 2025
Lears chieeppon
(b) Causal graphical view of the environ- model. While the number of language- controlled components can be arbitrary, we present the framework two components without loss of generality.
Becowfigere Iearoed
combrnrannts
(a) Composable view of enviromment models.
Figure 1: Illustrations environment models
by abstracting composable concepts from known tasks and reusing them in new contexts through compositional generalization. In this process, language often plays a key role, naturally encapsu- lating the causal relationships and compositional structures of the environment (Fodor & Pylyshyn 1988). For example, from tasks like pick-place ball and pull puck, we abstract the concepts of pick- place and puck. When faced with the task pick-place puck, only need few samples to combine these learned components-such as adjusting our focus and gesture to pick up the puck and place in the desired location. This can be and object (puck), whose values can be flexibly recombined in new domains for efficient adaptation.
system with two latent components: verb (pick-place)
This process mirrors the core principles of causal system: sparsity and modularity (Pearl 2009) Peters et al. 2017 Glymour et al.) 2019) Schölkopf et al. 2021). Modularity enables us to inde- pendently identify causal components and recombine them in novel ways, facilitating adaptation to different systems. Sparsity ensures that when a distribution shift occurs, it minimally alters the system, leaving most components unchanged. Thus, only few new samples are needed adapt to new compositions of familiar components. For instance, to generalize efficiently to pick-place puck, we only need to adjust the causal dynamics between pick-place and puck, while the rest of the system remains largely intact.
Motivated by modeling the composable environment model from causal view, characterize the environment model using an augmented graph of the partially observable Markov decision process (See Figure lbp. We denote sequence of observations as 0t, Gt rt >>== referring to image, action, and reward. We denote the underlying latents by components as {l..... m. further assume that 8, = (cI,t...., Cm,t) can be uniquely parti- tioned into disjoint language-controlledd components. The dimension of the language-control components c, is denoted as ne,, The following proposition establishes how we identify these com- ponents:
(51.t, $2,1 ... Sd,t) and N language
Proposition Language-Controlle Components). Under the ussumption that the graphical rep- resentation of the environment model is Markov and faithful (Spirees et al.) [007]] Glymour el al.) 2019 Pearl 2009) to the data, c,
[2] is a minimal subset of state dimensions that are directly con- trolled by the language component and Sj C.a and only if sj L l 0t 81-1, and 8j, IL. {1e)k 4ú, a1114 8t-1.
"ii, and e₁ used interchangeably in this paper.
3

---

Published as a conference paper at ICLR 2025
Intuitively speaking, language-controlled components set of state dimensions that are directly controlled by the individual language components, and this allows us to identify a modular and interpretable environment model. We achieve compositional generalization in the latent space with the guidance of language components, in the unseen test domains. We further relax the con- ditions allow the individual component be influenced by all previous states instead of only the previous itself (see Equation 2). The data generation process based on composable components is mathematically formulated follows:
[oe, rz] g($t; Et), 8: = (Ci,s Cm,)), L (1,,,,,m
C₁, ~ p(cs,elle, St-1, 0--1) for 1....,
2.2 IDENTIHIABILITY THEORY
Accurately identifying composable components and causal dynamics is crucial to developing robust models capable of generalizing across diverse environments. However, this task is challenging due to the inherent uncertainty and complexity of the data generation process. Previous work in non-
linear ICA [Khemakhem et al. 2019 Hoyer et al.] 2008 Hyyääinen & Pajunen 1999. Klindt et al.) 2020 Hyvarinen & Morioka 2016) develops methods for identifying latent variables dimension-
wise using strong assumptions like independence of noise terms and specific functional form priors. However, they often struggle with complex systems having interdependent latent structures and are limited by their reliance on a single auxiliary variable, parametric assumptions, and simple graphical models. Additionally, their sophisticated optimization procedures are difficult to scale up in real- world applications.
We focus on block-wise identifiability instead of dimension-wise identifiability achieve bet- ter trade-off between scalability and estimation accuracy. Our approach allows separate language components to independently control their corresponding latent variables, unlike previous temporal methods (Yao et al. 2 2022; Song al.) (024) that rely single auxiliary variable con- necting with all latent variables. By focusing on identifying language-controlled composable com- ponents without parametric assumptions, can maintain theoretical guarantees while requiring significantly fewer language component values for identification. Note that this framework extends beyond language compositional modality it applies to any latent system with multiple inter- mittent control signals, such as decomposing robotics tasks into components controlled by various signals. We believe this is the first work demonstrate the identifiability of disentangled compo- nents separately controlled by different intermittent control signals in a general non-linear case for reinforcement learning tasks.
We define block-wise identifiability concerning the identification of the language-controlled com- ponent as follows. Some proof techniques and notations are related Liu al. (2023): Sun et al. 2024).
Definition 1 (Block-wiss Identifiability) The true components of changing variables c; are block- wise identifiable if. for the estimated component of changing variables è; and each component of changing variables Ci there exists an invertible function h4 [8 + g.e,, such that c; h.(C.).
Theorem Suppose that the data generation process follows Equation 7 2 and the following assumptions are fulfilled, then the language-controle component c; is block-wise identifiable:
1. The mixing function (oe, rı] = g(st, 6t) is invertible and smooth.
2. The set {st E p(s)) = 0} has measure zero.
3. The conditional probability density should be sufficiently smooth, i.e., p(si,)44, 8:-1, 1)-1)
is at least first-order differentiable.
4. Given language components l,, previous state 8:-1 and previous action a; every element of latent variable 8322 should be independent of each other, i.e., 83, II. 8k,t 44, 81-1. a
for j. k E {I,...,,u n. and j + k.
5. For any C, E Ct,, we assume that there exist nee + values of such that j = 1,.... ne. and k 1, 1,.... Deg the following matrix is invertible:
4

---

Published as a conference paper at ICLR 2025
(1.
1)(1,0)
p (1,0)
4 (),
())
p (k,0)
gi 0)
(mat 0)
where &2) (k, = log pla kl.
1.a1 1) a llg p(a li is the difference of the first-
88g4 order derivative of the log density of 88, between the kth value and Oth value of the language component 4T
Most of these assumptions are commonly made in the field of causal representation learning (Knng al., 2023a; Von Kügelgen et al. 2021; Huang et al. 2022 Liu al.) 2023] Sunn al. 2024). They help prevent degenerate cases, ensuring that the composable components in the model are in generic conditions, while enabling causal structure recovery once latent variables are identified, their causal relationships can be discovered with standard causal discovery methods as used in Yao jet al. [2021: 2022). The above theorem introduces a relaxed form of identifiability, showing that for each language-controlled composable component C;, with fee + distinct values in a given lan- guage component (c.g., nee + objects object component), each true changing variable can be represented function of all estimated changing variables. It suggests that the estimated chang- ing variables encapsulate all the necessary information for the true changing variables, effectively separating the component corresponding the language component Li and the other components not controlled by it. The minimum number of tasks for identifying all m language-controlled com- ponents can be as low E; ne. More discussions about the requirements of identifying all language-controlled components {cI,..., Cm. and detailed proof are in Appendix A.3
2.3
LEARNING COMPOSABLE WORLD MODELS
Seaprig Occporst
Stuppl Miseculas
Componat
P
O
Viebs
Object
Object
(a) Composable world model learning. The world model learns language-controlle components and the dynamics among them from language compo- nents and prior experiences. The learned latent are expected be mapped to the observation parts that have corresponding semantic meaning.
(b) World model imagination and policy learning. Prompting by the language component values and the intial observations. the world model composes imaginations with language-controlle, components and only the reward-related states are used for pol- icy learning.
Figure 2: Illustrations of environment models.
Our world model identifies composable components and learns composition rules through agent in- teractions. The learning procedure is grounded on the identifiability results and incorporates natural properties of causal systems to improve learning and generalization. Notably, our learning frame- work is agnostic. To create more efficient prototype, we utilize the state-of-the-art model-based reinforcement learning algorithm, DreamerV3 (Hafner et al., (023), and apply insights from the identifiability results to guide the learning process.
5

---

Published as a conference paper at ICLR 2025
Approximating invertible mixing function. The traditional encoder-decoder architecture is adopted to fulfill the invertibility of the mixing function given in Assumption [ We also learn to predict the continuation signal for policy learning. To enhance practicality in the multi-task set- ting, we use a learnable task embedding z to encode each task u, helping these decoding functions capture variations across domains (see Figure 2a). We present
uses convolutional neural network (CNN) as the default observation encoder and decoder, and the other employs a Masked Auto-encoder (MAE) following (Soo et al. (2023). MAE is particularly effec- tive at extracting high-level semantic information, which may help bridge the gap between language components and their corresponding latent components (Kong al.) 2023b). The model components
variants of WM3C,
are:
Task Encoder:
(u(u)
Observation Encoder:
he pa (hg oz)
Observation Decoder:
ōe pe pol (ū
mo ? 8c: z)
(3)
Reward Decoder:
pe (f: my 2 8t: 2)
Continuation Decoder:
& ~ pe (ce
m. ① 82,
and the corresponding loss function is:
T
Tep
log ma O (0, log | my ④ (0, (cc me 2 8t,
Facilitating modular dynamics. We decompose both the transition model and the representation model into disentangled modules, following the causal structure of the environment model. The KL regularization is factorized into component-wise KL divergence well, acting as soft indepen- dence constraint that encourages the disentanglement of composable components. Moreover, the language components I are encoded token embeddings e; to increase model flexibility.
Language Component Encoder:
e; = (4:4)
Component C,2 Representation Model:
C ~ 4- (Cit hes Est 82-1, 4--1)
(4)
Component C₁, Transition Model:
ci, ~ pb 8t, a1-1)
Irans
Σ KL (c.(c,,) he, e, 81-1, a1-1) Po(Ci, S1,1 (1--1,
Enhancing causal structure. The identification results imply that we have enough values for each component, we can identify the component without additional constraints. However, we observe that it can be challenging in practice when the effect of the language signal is small. To address this, add mutual information (MI) constraint to strengthen the conditional independence described the environment model (Belghazi et al. 2018). Specifically, for each aanuagge-cnntoolee component, mutual information neural estimator is used to maximize the joint MI between the component and its corresponding language component, while MI between the component and other language components; see below. Detailed derivation and discussion can be found in Appendix A.4.2
composable
estimators minimize the joint
Mutual Information Maximization:
1(4: Ci, 1t
Mutual Information Minimization:
Σ
101i 81 12 (t-lit-))
(5)
T
m Σ
mm
Σ ifi
1=1
Enforcing sparse interactions. Real-world causal systems are usually sparse, with most variables not directly influencing each other (Scholkopf et al. 2021 Zhang Hyvärinen 2009). To reflect this characteristic and facilitate learning and identification, we incorporate learnable masks into the world model. We assume that the observation decoding. reward decoding, and continuous signal decoding only subset of latents, characterized by binary masks m., m,, M.. respectively. To alleviate the shrinkage effect of L1 loss and prevent excessive sparsity, which could result in infor- mation loss during the early stages, we apply modified gating mask from Rajamanoharan et al.

---

Published as a conference paper at ICLR 2025
2024) and adaptive L1 loss that dynamically controls the sparsity during the optimization process, accounting for the sparsity ratio of the latent variables, with the corresponding loss function given below (see implementation details in Appendix A.5.2)
T
Spar Σ
Σ La (ms). (
Therefore, the total objective for learning the composable world model is expressed as weighted summation:
Local aleep + Aluans + Blau Usper
Compact states for policy learning. We use standard soft actor-critic structure our policy module, following DreamerV3 (Hafner et al. 2023). The task embedding 2, learned from the world model, is used as condition in both the actor and critic networks to accommodate the multi-task setting. Instead of using latent states 8r, only the compact reward-relevant states, selected via reward mask M,y, are fed into the actor and critic networks. The policy module and world model are optimized alternately, with the world model utilizing the updated policy module to gather new interactions (see Figure 2b
Quick adaptation to tasks of new compositions. Synthesis effect might happen in the test time. For instance, the action open within the verb component can have different meanings when com- bined with objects that have distinct affordances, such as door and window. Instead of standard fine-tuning or retraining on all parameters of the world model, we only require minor adjustments to dynamics-related parts. This is feasible because the individual components have already been learned in previous tasks-an is only the interactions among them that remain unclear. Hence, we fine-tune the task encoder, representation model, transition model, decoding masks, and policy mod- ule account for the synthesis effect of these recombinations while keeping others fixed.
EXPERIMENTS
In this section, we address the following questions:
How effectively can our learning framework identify language-conttrolled components in known environments?
How accurately can the learned world model predict language-controlled components in new environments with novel combinations?
Can our framework enhance RL training and generalization in real-world applications?
Are the learned language-controlleed components interpretable?
We aim to answer these questions through series of experiments, including tests on synthetic data where we have access to ground truth latent states, as well as real-world robotic environments involving a collection of manipulation tasks.
3.1 SYNTHETIC DATA
To validate the accuracy of our canguage--onttrlle composable component identification, we con- duct simulation study based on the assumed data generation process described Section 2.1 We initialize three language-controlles components {c1, C2, C2,C3), c3). each independently controlled by three groups of discrete language tokens (where language component takes fe. + values). We split all possible combinations into training set for the i.i.d component identification test and test set for the o.o.d component prediction test. Following the experimental setting block-wise identification compute the coefficient of determination R²
in Von Kügelgen et al. (2021): b: Liu al. (2023), between the estimated components and the ground truth.
We compare the identification accuracy of our method with other causal representation learning approaches, such as iVAE (Khemakhem et al. (2019). which does not incorporate temporal infor- mation, TCL (Hyvarinen & Moriokal 2016)), and TDRL and NCTRL (Yao al. 2022 Song al.

---

Published as a conference paper at ICLR 2025
2024). which the independence of the Jacobian matrix to identify dimension-wiss temporal de- pendencies. Additionally, we evaluate our method against state-of-the-ae world models, including DreamerV3 (Hafner et al. 2023) and TD-MPC2 (Hansen et al. 2023), which do not account for causal representation learning or generalization optimization.
3.1.1 COMPONENT IDENTIFIATIOON KNOWN TASKS
WM3C effectively identifies controle components in known tasks. The plot the left of Figure 3 shows that our method (WM3C) accurately captures the underlying latent structure of the data, with diagonal R² > 0.9 and off-diagonal values around 0.1. The middle figure plots the average R² values of several baseline models against training steps, where WM3C significantly outperforms the other models, demonstrating both higher accuracy and minimal standard deviation, indicating consistent performance.
We attribute the inferior performance of other methods to differences in model priors and optimiza- tion objectives. For example, iVAE uses standard domain index rather than language components, which limits its ability identify and utilize composable elements effectively. TDRL, which relies on the independent noise assumption to capture conditional independence among latents, fails to recognize the composable components tied to language. Although NCTRL models domain shifts through an underlying hidden Markov model, the prior is too simplistic to capture complex latent structure changes across domains. Both DreamerV3 and TD-MPC2 focus on fitting models recon- struct observations, resulting in entangled representation space that does not prioritize identifying distinct latent components.
0.9286
Average n.
Acarage g Over
0.91
0.12
0.12
Estimated latents
0.05
0.92
0.04
Value
0.09
0.12
0.95
Gur
True latents
Troining Stops
Tian Pouns
Figure 3: Language-conttrolee composable components identification results on known tasks and novel tasks. Left: Coefficient of determination (R2²) using kernel ridge regression regress esti- mated latents on true latents. Middle: Average 12 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² 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.
3.1.2 COMPONENT PREDICTION IN NOVEL TASKS
We evaluate the generalization of the learned model by comparing the latent rollouts generated by the world model (through imagination without consecutive observation) with the true latent in unseen tasks. Here, the unseen tasks are new combinations of language-controlled composable components from the known tasks. In the plot on the right of Figure 3 see that WM3C consistently maintains the highest R² values over different time points, reinforcing the effectiveness observed in the middle figure. It is important to note that this evaluation is based on imagination (latent rollouts) in an unseen environment without observations, which highlights WM3C's s ability to generalize through learning composable components in the world model.
8

---

Published as a conference paper at ICLR 2025
3.2 ROBOT MANIPULATION
To answer whether our framework facilitates RL training and generalization in real-world applica- tions, we conduct experiments in the robotic simulation environment, Meta-world (Yu al. 2019). is a benchmark suite of robotic manipulation environments, where each task is paired with a corre- sponding language description. Assuming language components verb and object in the data gener- ation process, we take 18 training tasks that follow this structure and 9 test tasks that either are new combinations of known language components or mixtures of known and unknown language com- ponents. Models are trained in multitask setting, including all 18 training tasks, for total of 5M steps, while the adaptation performed on each test task for 250 steps. We compare the training and adaptation efficiency of WM3C (MAE) and WM3C (CNN) with DreamerV3 and visual-based Multi-tast SAC (MT-SAC). Additional training and test details are provided in Appendix A.3.2
3.2.1 TRAINING AND ADAPTATION
Training efficiency. In Figure we report the average success rate across all tasks during the training process, along with the learning curves of 9 tasks. By learning a composable world model guided by language, our WM3C demonstrates significantly better data efficiency and performance compared to DreamerV3 and MT-SAC, with further improvements achieved through masked image modeling. Additionally, in tasks such as box-close and door-unlock, where DreamerV3 struggles in learning meaningful representation, ing. This highlights the WM3C's ability to handle diverse task dynamics and confirms its advantage in learning composable world models causal system compared conventional world models. We provide an ablation study in Appendix A.7.2 to investigate the effects of the integrated modules.
WM3C framework consistently exhibits rapid policy learn-
WMBC MAE
WABCC CNN
Dreamner V2
MESSCC
0.6
Success
spen
0.3
Figure 4: Learning curves on Meta-world: Average success rate and 9 specific task success rates. Results reported with three seed runs and 10 episodes for each task evaluation.
Adaptation efficiency. To evaluate the generalization advantage of the composable world model, we test its adaptation ability on new, unseen tasks, including both the compositions of known com- ponents and compositions involving both known and unknown components (See adaptation details in Appendix
Figure we see that the first 7 tasks that recombine learned components, the WM3C variants, particularly WM3C MAE, consistently achieve higher success rates compared to full-pam pall-aereeter tuned DreamerV3. This aligns with our assumption that after lcarning composable components, adaptation should focus learning compositional dynamics. For tasks coffe-button (coffee button is unknown, press known) and mandle-press-siie (handle known, press-side is unknown). WM3C performs better in namics module adaptation alone may not always suffice for novel components, it can be surprisingly effective certain similar scenarios.
and worse in both cases, suggesting that while dy-
3.2.2 INTERVENTION IN THE MANGUAGGF-CONTROIIA COMPONENTS
In this section. interpret the language-controllee components identified in Meta-world by inter- vening corresponding latent components and see how this affects the reconstruction. Although not all identifiability conditions be easily met in practical applications, estimated latent states

---

Published as a conference paper at ICLR 2025
1.5
Aneruge Suecena Halt
Doee opan
Wanlow soen
Husdloo
0.5
peom
MY
MADAY
0.5
Step. 1.0
Figure 5: Adaptation curves on 9 unseen tasks, including 7 tasks that are recombination known components and 2 tasks have unknown components Mandle--presssstee erV3 and MT-SAC are full-parameter finetuned while WM3C is only tuned with the dynamics mod- ule. Results reported with three seed runs and 10 episodes for each task evaluation.
and coffee-button). Dream-
are still reasonable and meaningful (see Figure 6). We find that intervention on the verb controlled component does not affect the object in the image but causes the end effector of the robot arm and its adjacent joint to disappear, as expected for the executor of different verbs. On the other hand, intervention on the object does affect the robot arm, but changes the appearance of objects signif- icantly. This demonstrates that our model successfully learns modular and compositional represen- tations aligned with the semantic meaning of language components, enhancing both interpretability and efficient task manipulation. Additional examples can be found in Appendix A.7.3
Task: Coffee-pull
Took: Window-ckose
Task:
Task:
Drawe-opon
Doot-mblokk
Observation
Object
Reconstruction after Intervening on Object Component
Verb
Reconstruction
after Intervening Componcet
Figure 6: Intervention on the language-controlled components. The first row is the original image observation WM3C receives, the second row is the reconstruction of WM3C's observation decoder after intervening on its object latent component and the third row is the reconstruction of WM3C's observation decoder after intervening on its verb latent component.
CONCLUSION
In this work, we introduced World Modeling with Compositional Causal Components (WM3C), framework aimed at improving generalization in reinforcement learning by utilizing composi- tional causal dynamics. By integrating modalities like language and offering theoretical guarantees for identifying distinct causal components, WM3C enhances adaptability to unseen environments. Extensive experiments on synthetic data and real-world robotic tasks show WM3C outperforms ex- isting methods in uncovering underlying processes and improving policy learning. Future research can extend WM3C to offline learning and explore scalable modularity and sparsity constraints, con- sidering the large number of language components involved.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
The authors extend their sincere gratitude to the anonymous reviewers for their insightful feedback and constructive suggestions. Biwei Huang would like acknowledge the support of the National Science Foundation (Grant No. DMS-2428058).
REFERENCES
Andrew G. Barto and Sridhar Mahadevan. Recent advances in hierarchical reinforcement
learning. Discrete Event Dynamic Systems, 13:41-77, 2003. URL https://api.
semant www.icccchl...... wwtpohooooaoog//oruuuuttasaa 38:822
Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeswar, Sherjil Ozair, Yoshua Bengio, R. De- von Hjelm, and Aaron C. Courville. Mutual information neural estimation. In International Conference on Machine Learning, 2018. URL ott. org/
CorpusID:44220111
Karl Cobbe, Oleg Klimov, Christopher Hesse, Taehoon Kim, and John Schulman. Quantifying gen- eralization in reinforcement learning. In International Conference Machine Learning, 2018.
URL httthtttp///ppl www.nticccannnnne.n...... :ttg:coopps11::54400000
Yan Duan, John Schulman, Xi Chen, Peter L. Bartlett, lly Sutskever, and P. Abbeel. RI$ http:///eeerrenooacaaaaeeeeeeeee
wbw///lsl...///... 2016. R1
Fan Feng and Sara Magliacane. Learning dynamic attribute-factored world models for efficient multi-object reinforcement learning. Advances in Neural Information Processing Systems, 36, 2024.
Chelsea Finn, P. Abbeel, and Sergey Levine Model-agnostic meta-learning for fast adaptation of
deep networks. In International Conference on Machine Learning. 2017. URL httpps//aapi.
semanticschol aat.aaaa....c/oooaaaDD:7........
Jerry A Fodor and Zenon W Pylyshyn. Connectionism and cognitive architecture: critical analy- sis. Cognition, 28(1-)):3-71, 1988.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International Conference on Machine Learning, 2018. URL https:://pp. semanticscie bemanticscholar.com hww.choo.....oogooooaaatt..::111
Shani Gamrian and Yoav Goldberg. Transfer learning for related reinforcement learning tasks via image-to-image translation. In International Conference on Machine Learning, 2018. URL
https://api. www.nt htt.ooaa....rr/ccrorr111::99911
Dedre Gentner and Arthur B Markman. Structure mapping in analogy and similarity. American psychologist, 52(1):45, 1997.
Clark Glymour, Kun Zhang, and Peter Spirtes. Review of causal discovery methods based on graphical models. Frontiers genetics, 10:524, 2019.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv: 10122, 2018.
Tuomas Haarnoja, Aurick Zhou, P. Abbeel, and Sergey Levine. Soft actor-critic: Off-policy max- imum entropy deep reinforcement learning with stochastic actor. ArXiv, htts/1801.01290 2018. UtLhhttp:://eet sewant aar ootpttsoii.oooooooooooooooooooo. 28202811
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv: 1912.0/603 2019.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXIv:2301.04110 2023.
Nicklas Hansen and Xiaolong Wang. Generalization in reinforcement learning by soft data augmen- tation. 2021 IEEE International Conference on Robotics and Automation (ICRA), pp. 13611-13617, 2020. URL httph:ttt/tt..... //ppi. semant csch http://:/...aaaaa....ooooooooooa

---

Published as a conference paper at ICLR 2025
Nicklas Hansen, Hao Su, and Xiaolong Wang. Td-mpe2: Scalable, robust world models for contin- uous control. arXiv preprint arXiv:2310.16828, 2023.
Patrik Hoyer, Dominik Janzing, Joris M Mooij, Jonas Peters, and Bernhard Schölkopf. Nonlinear causal discovery with additive noise models. Advances in neural information processing systems, 21, 2008.
Biwei Huang. Fan Feng, Chaochao Lu, Sara Magliacane, and Kun Zhang. Adarl: What, where, and how to adapt in transfer reinforcement learning. arXiv preprint arXiv:2107.02729 2021.
Biwei Huang, Chaochao Lu, Liu Leqi. José Miguel Hermández-Lobato, Clark Glymour, Bernhard Schölkopf, and Kun Zhang. Action-sufficient state representation learning for control with structural constraints. In International Conference on Machine Learning, pp.
PMLR, 2022.
Aapo Hyvarinen and Hiroshi Morioka. Unsupervised feature extraction by time-contrastive learning and nonlinear ica. Advances in neural information processing systems, 29, 2016.
Aapo Hyvärinen and Petteri Pajunen. Nonlinear independent component analysis: Existence and uniqueness results. Neural networks, 12(3):429-439, 1999.
llyes Khemakhem, Diederik Kingma, and Aapo Hyvärinen. Variational autoencoders and nonlin- ica: A unifying framework. In Internatioonal Conference on Artificial Intelligence and Statistics, 2019. URL URL /tttLs::///// //api..///pp...e/ /tt.t:::/ap..ssaaaaaa.a.m
Ilyes Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen. Variational autoen- coders and nonlinear ica: unifying framework. In International conference on artificial intelli- gence and statistics, pp. 2207-2217. PMLR, 2020.
Bangalore Ravi Kiran, Ibrahim Sobh, Victor Talpaert, Patrick Mannion, Ahmad A. Al Sallab, Senthil Kumar Yogamani, and Patrick P erez. Deep reinforcement learning for autonomous driving: 2020. URL
A survey. IEEE Transactions Intelligent Transportation Systems, 033
https: /app. semant. cscho ar. org Cttpg//CrpuusSDDD
David Klindt, Lukas Schott, Yash Sharma, Ivan Ustyuzhaninov, Wieland Brendel, Matthias Bethge, and Dylan Paiton. Towards nonlinear disentanglement in natural data with temporal sparse coding. arXiv preprint arXiv:2007. 10930, 2020.
Lingjing Kong, Biwei Huang, Feng Xie, Eric Xing, Yuejie Chi, and Kun Zhang, Identification of nonlinear latent hierarchical models. Advances in Neural Information Processing Systems, 36: 2010-2032, 2023a.
Lingjing Kong, Martin Q Ma, Guangyi Chen, Eric Xing. Yuejie Chi, Louis-Philippe Morency, and Kun Zhang. Understanding masked autoencoders via hierarchical latent variable models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 7918- 7928, 2023b.
Kimin Lee, Kibok Lee, Jinwoo Shin, and Honglak Lee. Network randomization: A simple technique for generalization in deep reinforcement learning. In International Conference on Learning Repre- sentations. 2019. URL htttts:///appp www.:essssactcccc............... www.:laar..rrooopppoeeeee ottp/CorrpusssDD2233300000
Yuren Liu, Biwei Huang, Zhengmao Zhu, Honglong Tian, Mingming Gong, Yang Yu, and Kun Zhang. Learning world models with identifiable factorization. Advances in Neural Information Processing Systems, 36:31831-31864 2023.
Gary F Marcus. The algebraic mind: Integrating connectionism and cognitive science. MIT press, 2003.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Belle- mare, Alex Graves, Martin A. Riedmiller, Andreas Kirkeby Fidjeland, Georg Ostrovski. Stig Pe- tersen, Charlie Beattie, Amir Sadik, Ioannis Antonoglou, Helen King. Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature, 518:529 533,, 2015.
URL https:/ /api. semanticschon .rr/
Corpus ID::01242440
12

---

Published as a conference paper at ICLR 2025
Volodymyr Mnih, Adrià Puigdomènech Badia, Mehdi Mirza, Alex Graves, Timothy P. Lillierap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep rein- forcement learning. In International Conference Machine Learning, 2016. hthps::
//api. sema nt i cacchlllaaaoooocccccoo wwwpllaa..rrgooooaaeeeeee ID: 68 6875312
Hiroshi Morioka and Aapo Hyvärinen. Causal representation learning made identifiable by grouping of observational variables. arXiv preprint arXiv:2310.1570% 2023.
Ofir Nachum, Shixiang Shane Gu, Honglak Lee, and Sergey Levine. Data-efficient hierarchical reinforcement learning. Advances in neural information processing systems, 31, 2018.
Suraj Nair, Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, and Abhinav Gupta. R3m: A uni- versal visual representation for robot manipulation. arXiv preprint arXiv:2203. 12601, 2022.
Ronald E. Parr and Stuart J. Russell. Reinforcement learning with hierarchies of machines. In Neu- ral Information Processing Systems, 1997. URL https://api semant www.cicccchcoooooooooooooo.o
Cerpus ID: TD: 6760236
Judea Pearl. Causality. Cambridge university press, 2009.
Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Elements of causal inference: foundations and learming algorithms. The MIT Press, 2017.
foundaact.oundaaiioon
Senthooran Rajamanoharan, Arthur Conmy, Lewis Smith, Tom Lieberum, Vikrant Varma, János Kramár, Rohin Shah, and Neel Nanda. Improving dictionary learning with gated sparse autoen- coders. arXiv preprint arXiv:2404. 16014, 2024.
Bernhard Scholkopf, Francesco Locatello, Stefan Bauer, Nan Rosemary Ke, Nal Kalchbren- ner, Anirudh Goyal, and Yoshua Bengio. Towards causal representation learning. ArXiv, abs/2102.11107. 2021.
URL https::/app..eenntiicssss www.:/aru.ooooorroosoo.i..
231986372
Bernhard Schölkopf. Francesco Locatello, Stefan Bauer, Nan Rosemary Ke, Nal Kalchbrenner, Anirudh Goyal, and Yoshua Bengio. Toward causal representation learning. Proceedings of the IEEE, 10965):622-634 2021.
Proxi- mal policy optimization algorithms. ArXiv, 0bs/1707.0647, 2017. URL https://appl
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov.
Younggyo Seo, Danijar Hafner, Hao Liu, Fangchen Liu, Stephen James, Kimin Lee, and Pieter Abbeel. Masked world models for visual control. Conference on Robot Learning, pp. 1332- 1344. PMLR, 2023.
Lucy Xiaoyang Shi, Joseph J Lim, and Youngwoon Lee. Skill-based model-based reinforcement learning. arXiv preprint arXiv:2207.0756) 2022.
David Silver, Thomas Hubert, Julian Schrittwieser, loannis Antonoglou, Matthew Lai, Arthur Guez, Mare Lanctot, L. Sifre. Dharshan Kumaran, Thore Graepel, Timothy P. Lillicrap. Karen Simonyan, and Demis Hassabis. A general reinforcement learning algorithm that masters chess, shogi. and go through seay- Science, 362: 140 1 144, 2018. URL stpp:...smmaticcccooll....
: org/CorpusID: I:: 54: 54457125
Shagun Sodhani, Amy Zhang, and Joelle Pineau. Multi-task reinforcement learning with context- based representations. In International Conference on Machine Learning, pp. 9767-9779. PMLR, 2021.
Xiangchen Song, Weiran Yao, Yewen Fan, Xinshuai Dong, Guangyi Chen, Juan Carlos Niebles, Eric Xing, and Kun Zhang. Temporally disentangled representation learning under unknown non- stationarity. Advances in Neural Information Processing Systems, 36, 2024.
Xingyou Song, Yiding Jiang, Stephen Tu, Yilun Du, and Behnam Neyshabur. Observational overfitting in reinforcement learning. ArXiv, abs/1912.02975, 2019. URL lhtpss: //app. semant: cscho arr. wacphhllr.rrgg/aaCo arrpaarrcccoorrssit008558533355
13

---

Published as a conference paper at ICLR 2025
Peter Spirtes, Clark Glymour, and Richard Scheines. Causation, prediction, and search. MIT press, 2001.
Boyang Sun, Ignavier Ng. Guangyi Chen, Yifan Shen, Qirong Ho, and Kun Zhang, Continual non- linear ICA-based representation learning, 2024. URL https:://pennreveew..e wwwp:e..eeefrrrre
IddxTXaJmMXXB
Joshua B Tenenbaum, Charles Kemp, Thomas L Griffiths, and Noah D Goodman. How to grow mind: Statistics, structure, and abstraction, science, 331 (6022) 2011.
Oriol Vinyals, Igor Babuschkin, Wojeiech M. Czarnecki, Michaël Mathicu, Andrew Dudzik, Juny- oung Chung, David Choi, Richard Powell, Timo Ewalds, Petko Georgiev, Junhyuk Oh, Dan Horgan, Manuel Kroiss, Ivo Danihelka, Aja Huang, L.. Sifre, Trevor Cai, John P. Agapiou, Max Jaderberg, Alexander Sasha Vezhnevets, Rémi Leblond, Tobias Pohlen, Valentin Dalibard, David Budden, Yury Sulsky, James Molloy, Tom Le Paine, Caglar Gulcehre, Ziyun Wang, Tobias Pfaff, Yuhuai Wu, Ro- man Ring, Dani Yogatama, Dario Wünsch, Katrina McKinney, Oliver Smith, Tom Schaul, Timo- thy P. Lillicrap, Koray Kavukcuoglu, Demis Hassabis, Chris Apps, and David Silver. Grandmaster level in starcraft using multi-agent reinforcement learning. Nature, 575:350- 354, 2019. URL https: / ap semant. ascho ar org rpus ID: 204 972004
Julius Von Kügelgen, Yash Sharma, Luigi Gresele, Wieland Brendel, Bernhard Schölkopf, Michel Besserve, and Francesco Locatello. Self-supervised learning with data augmentations provably isolates content from style. Advances in neural information processing systems, 34:1 34:16451-16467, 2021.
Weiran Yao, Yuewen Sun, Alex Ho, Changyin Sun, and Kun Zhang. Learning temporally causal latent processes from general temporal data. arXiv preprint arXiv:2110.05428, 2021.
Weiran Yao, Guangyi Chen, and Kun Zhang. Temporally disentangled representation learning. Advances in Neural Information Processing Systems, 35:26492-26503, 2022.
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on Robot Learning (CoRL), 2019.
Zhecheng Yuan, Zhengrong Xue, Bo Yuan, Xueqian Wang, Yi Wu, Yang Gao, and Huazhe Xu. Pre-trained image encoder for generalizable visual reinforcement learning, Advances in Neural Information Processing Systems, 15:12222-1303, 2022.
Amy Zhang, Clare Lyle, Shagun Sodhani, Angelos Filos, Marta Z. Kwiatkowska, Joelle Pineau, Yarin Gal, and Doina Precup. Invariant causal prediction for block mdps. ArXiv, abs/2003.065 00001, 2020a. URL httpt: semanticschi lar. orrg :ttp://ooruuII::2211:333
Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. ArXiv, 200./0000..0722 2020b. Utt.httss:://// http::/ppp..emttt......eoooarrr/
Kun Zhang and Aapo Hyvärinen. On the identifiability of the post-nonlincar causal model. In Conference on Uncertainty in Artificial Intelligence, 2009. URL httpps://api.
semant. cwcholar.c...oooco waccho/la...oor///oopoatttttaaa.
Kun Zhang and Aapo Hyvarinen. On the identifiability of the post-nonlinear causal model. arXiv preprint 1rXiv:1205.200 2012.
14

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1 RELATED WORK
Generalization in Reinforcement Learning. Generalization in reinforcement learning (RL) mains a core challenge, particularly when agents need to adapt to unseen task variations or environ- ments with minimal additional training. Early approaches like Meta-RL and invariant representa- tion learning (Lee et al.; 2019 Hansen & Wang [020 Yuan et al., 2022] Nair et al. (022) improved adaptability by learning task-agnostic policies that can be efficiently reused across multiple domains.
In addition, Model-based approaches, particularly world models (Ha & Schmidhuber 2018 Hatner
al. 2023; Hansen et al., 2023), have also demonstrated success in generalization by building latent representations of the environment that allow for planning and imagination baseod learning. However, these approaches often ignore the compositional nature of tasks, limiting their ability to generalize across tasks that share modular components. In contrast, our work leverages language as a natural guide identify and control composable task components for a better environment
modeling.
Causal Representation Learning. Causal representation learning has been a pivotal develop- ment in machine learning, focusing uncovering the underlying causal mechanisms within data (Schölkopf et al.] 202T). Prior methods such as nonlinear ICA (Zhang Hyvarinen 2012
Hyvärinen & Pajunen| [1999) and temporal causal representation learning (Yao al. |021 2022
Song et al., 2024) have sought to identify causal relationships in latent variables using priors such as independent noise conditions and auxiliary variables. However, these approaches often focus on dimension-we identifiability. which can be difficult to scale in real-world applications with complex causal dynamics. Our framework advances this line of work by introducing block-wise identifiability across multiple coexisting language components, which establishes a connection to a compositional modality. This further enables the identification of language composable components as distinct blocks, improving scalability and generalization.
A.2 PROOF OF PROPOSITION 1 (LANGUAGE- BAANGUAGE-CONTROLLEED COMPONENTS)
Here we present the proof for identifiability of the language-controlled composable components. The content is arranged as, firstly deriving the relationship between the estimated latent and ground truth latent variables, then proving that under some kunguage-controllld composable components can be uniquely and correctly identified. The identifiability theory is suitable for the system in- cluding arbitrary number of language components and language-controlle latent components. For simplicity, we present the proof with a system having two language components and it can be easily extended to arbitrary number of language components.
We first give the definitions of d-separation, global Markov condition, faithfulness assumption, which are used in the proof (Spirtes et al. 2001 Glymour et al., 2019 Pearl 2009).
Definition (d-separation). In a directed acyelic graph (DAG), a path between two nodes X and Y is said to be blocked (or d-separated) by set of nodes Z if and only if:
For any node on the path:
If the path includes chain X + M + +Y or fork X t M +,, the path is blocked if M Z.
If the path includes collider X M + Y, the path blocked unless M Z or any of
its descendants are in Z.
Definition (Global Markov Condition). The distribution p over set of variables V satisfies the global Markov property on directed acyclic graph (DAG) G for any partition (X,Z,Y, Y) of the variables such that Z d-separates X from Y in G. we have:
p(X,Y 2) p(X (Y Z)
Definition 3 (Faithfulness Assumption). A distribution
set of variables V is faithful to a directed acyclic graph G if all and only the conditional independencies in are implied by the global Markov condition on G.
15

---

Published as a conference paper at ICLR 2025
This implies that if two variables are conditionally independent in p. this must be reflected by d- separation in G, and there are no additional independencies beyond those imposed by the graph structure.
Proposition 1 (Language-Contten Controlled Components). Under the assumption that the graphical rep- resentation of the environment model is Markov and faithful to the data, Citt is a minimal subset of state dimensions that are directly controlled by the language component 4 and $J Ci,t if and only if S,a AL 4 Qt-1:t, St-1, and j.t 11 (4x Kegt l2, 04-1:t: St-1-
Proof. The proof is done with twice contradiction follows:
Step 1: We first show that if 8j,t E Ci.t, then 8, 4 at-1:t; 8
We prove this by contradiction. Assume $j,t _l 1 dt-1:1, 8t-1, meaning that $j.t is independent of 4 given the previous actions at-1st and state dimensions 8:-1. According to the faithfulness assumption, this implies that there is no directed path from 4 to 8.. in the graphical model, contra- dicting the assumption that 8j, is part Ci which consists state dimensions controlled by 4 Thus, by contradiction, it follows that $,, Al de-1:t; St-1, meaning that $3,0 has directed path from 4, and is thereby controlled by 4,
Step 2: Next, we show that 8j.t IL feat Li Gc- Ge, 81-1.
Again, suppose by contradiction that 83,8 M. (la)ayi 44 1:t1 81-1, implying that 8j.7 depen- dent another language component l (for some k of i) even after conditioning on 4. the action sequence, and the previous state. According to the faithfulness assumption, this would mean there a direct path from l to sj However, this contradicts the proposition's condition that $j E Ci which requires that $j.t is controlled specifically by I and not influenced by any other lg for k + i. Therefore, by contradiction, it must be true that S3,a 11 (Tx hept 14, 1:t; 8t-1.
From the two steps, we have shown that $3. E C. if and only if $j. P 14 at-1:t S1-1 and Sj, IL {1x Jeys 4,, 01:: 81-1. Thus, Ci is minimal subset of state dimensions controlled by the language component 4.
A.3 PROOF OF IDENTIFYING LANGUAGE- COOTRROLLED LATENT COMPONENTS
Consider the data generation process described by Equations (1) and (2). To simplify the notation, we demonstrate the proof for system with two language components, (2,,2), and their corre- sponding language-controlle latent components, (CI,,,C,.1. The proof generalizes to any number of language components and associated latent variables.
[ot, ri] c() (cI.t...,,Cm.. ci p(p((ll 9t-1,0-1) for 1, I,...,m
(8)
We begin by matching the observation distributions:
p rt 11,2, 9:-1, = p(ōr, Pr 11, 81-1, p(g(s (8, ()) 1, 12, 8:-1, = ((0( ét) 7, 12, 8, ar
(9)
(10)
The proof the elimation of noise can be referred to Khemakhem et al. (002). Assuming that the observation function g is invertible and differentiable, and applying the change of variables formula, we have:
p(sr 41, 12, St-1, a,-1) = p(ãr 44, 42, 8, az-1) det J
(1I)
where h gl is the invertible transformation from the estimated latent variables 8, the true latent variables 8., and Jn is the Jacobian of ht
Taking the logarithm of both sides:
16

---

Published as a conference paper at ICLR 2025
logp(st 11,2, 21-1,0(11) = logp(š; 11,2, a(-1) + log det J/=
(12)
Under Assumption 3, given the language token I and 12, the previous state 8t-1 and the previous action a;-1, each latent variable 85, is independent of the other latent variables. Thus, the log- density of given (1, l2 can be decomposed as:
logp(s: 2,,22, 81-1, 01-1)
Σ log lsg 4,12 81-1,
(13)
and similarly for the estimated latent variables:
logp(š; 41,2, 81 ar-1)
log ( 4, 8, at-1).
(14)
Substituting these into the previous equation:
p(sa 11,2, 8c-1, a2-1)
log p(gp | 11, 1,, 8 a,-1) + log det J=
(15)
Here, we first identify C₁, the l controlled latent component. We use dimension index { .... Ray to indicate the latent variables belonging to Ci, and dimension idex {ne, + 1,... n to indi- cate the latent variables belonging to C2, where represents the total dimensions of the latent space. We further simplify the notation that qi (1,5,,,11112) = p(si 1,12, 8t-1, 4x-1) and ds (8,,,, 12) == log p(i,e 14,2, 8,-1, 41)
We begin by taking the derivative with the estimated latent variable 83 where j E {ne, + 1....n..
0q. (1, 12) Usa. Osis
(8,(8 log det J
(16)
083.0
083,1
Here build the connection between the dimensions in the true I controlled component st, and the estimated latent in the L2 controlled component, $j,t- Although is helpful in the sense of identification of C1, is challenging in estimation because do not have knowledge about the invertible function, h, making the estimation intractable.
Fortunately, if we have multiple values of 4,, leverage them make the estimation tractable by constructing the difference terms. We assume l 's value can be taken from (1,....., J1
dgi(8i, (1, 12) Osi. 083,
dq: 00,0 Usi. G8
(8,e, Ü8,,
(ūj(S 11,0, 12) (17)
O8it
ÜSit
08,,2
Recall that the independent assumpation that the latent component controlled by la not trolled by 4,- qj ($),,, 4., 12) does not change when only 1, changes. However, this also adds some constraints and requirements discuss later. Here we can remove the l2 from the conditions since according to the proposition of $1,2 E Ci it is independent from given l1, St-1 and a;-1-
the number of tasks that combined ( and 12, which are going
Igs (gr, (0((8i,,, Osi. 0 Osit
(18)
Osi,
Suppose we have Mey values in 4, from 11,0 o to l, we construct linear system as follows:
17

---

Published as a conference paper at ICLR 2025
Oyn (st.c11)))
aq (11.
Baney (Smey ad1a) Demey it
Bines (Sneg d11.0) Oe it
S01,1
881.
0s1.e
0. (19)
Oq (3)
Diwey 8an.
(Ane adt..))
Boney (Sme) dd1.a) 88neg of
Osnni D63.t
0=1.1
88l.1
To achieve Oss.t
0 for any E (1..... ney and j E {nes -,...,, n), the left matrix has to be
full rank. This further implies that if the cardinality of I range (k) is larger than ne, + 1, dei as well as
0,
de2
de 8c, O.
Next, we turn to the Jacobian matrix of h that describes the relationship between the true latent variables 8, and the estimated latent variables $y:
Des.t
acL,1 Des.t
Jn
Sei.t
(20)
8eg.t
deg.1
Bei.
0eg.
By assuming the transformation h is invertible, the Jacobian matrix Jn. is full rank. We can safely meet conclude that, Tiez
daej 0, and the non-zero entries can only appear in dse,, De., The here can be further extended to the concatenation of components that are not controlled by 4, namely ei in the cases of having language components more than 2.
After proving the number I for identifying Cı, it straightforward know the number I, needs and the number of tasks (combinations of 1,) for identifying all the language-controoledd components.
We can use the same procedure to prove the number of l2 needed for identifying C₁ the same as proving the identification of C1, which requires Reg + 1 values of 12.
Below, we discuss two estimation procedures for identifying all the language-conttrolldd compo- nents, related to the tradeoff between practicality and minimal number of tasks.
For the case having two language components, suppose we have enough data, e.g. sufficient vari- ability in I and lz. We have two ways to identify e₁ and C2.
One by one: We can first identify one component with model and then identify another component with another model in one-by-one order. The least number of tasks we need
Ex=1 Re: + for language components and E Re. + for m language components. However, it is limited in practical use since we can not identify all language-controoledd components simultaneously in the estimation procedure. in Equation| 17, need to control the ( (2, does not change when only 4, It requires us to have tasks that allow us to control the l2 to be a fixed value when 4 is changing, and vice versa. Then, we can use these tasks to estimate e₁ and in turn estimate the C2 with another set of tasks where I is fixed and 22 is changing. While requires the minimal number of tasks and language tokens, the estimation procedure is sophisticated.
order to meet the requirement
All in one: We can learn a model to simultaneously identify and estimate all the com- ponents with more needs on variability. To make sure the the Ω ($), I, (2) does not change when only (1,s, we need enough values of (2 presented with each value of 4, That helps because when I is changing, there also enough l2 values to let us identify C2. is the same in the case where more than 2 language components evolved, where to have enough combinations quire I ne, tasks for 2 language components and m ne. + tasks for m language components.
want the rest language components' values. This would re-
Once the latent variables are identified using Theorem 1 the causal structure among these latent variables can be directly recovered using standard causal discovery methods, such as constraint- based score-based approaches used in prior work (Yao et al., 2021 2022 [ong et al. 2024). These methods leverage the identified latent space to ensure the moderlying causal relationships are correctly estimated, providing clear path to causal structure identifiabiliyy.
18

---

Published as a conference paper at ICLR 2025
A.4 OBJECTIVE DERIVATIONS
A.4.1 REPRESENTATION AND DYNAMICS OBJECTIVES
We define the extended information bottleneck objective based Dreamer (Hatner et al. 2019).
now conditioning on L = }. as follows:
max (S11:3 (owr, Frrr))or: L) (I(r, harlor, L),
where 3 is scalar and i, are dataset indices that determine the observations p(orli)) 8(or ōx) as in
Maximizing this objective encourages the model to reconstruct each image by relying on information extracted at preceding time steps to the extent possible, and only accessing additional information from the current image when necessary.
Lower bounding first term gives us the objective optimizing the latent representation.
(ourr L.
Σ
Implour, inplour
E
In mallan a1:,,
E
In mallsn
KL plous. arrr || II ((00)s)))()))))) (s))
E
(ng(o//))) + Ing(rals)))
Upper bounding the second gives us the objective of optimizing the latent transition dynamics.
E P
Σ inp(ss/s--n,a-nnne L) (1,,
E
I1 -anenl)nppppepeee Inp(
< E
Σ Inp(8)/8--11.01100.. 11
= E
KL (p(s)/51-110011000,11 1 00-1,,
This lower bounds the objective. It can be further decomposed to the summation of independent KL terms according to the environment model we assume Section. 2.1 and the Proposition A.2 of the language-controlledd component.
19

---

Published as a conference paper at ICLR 2025
(sur;hhrl0arr L.)
a1-1, KL (p( (( [
m
KL (p(ci,(ss-) 06,11, 14 1(
A.4.2 MUTUAL INFORMATION CONSTRAINT DERIVATIONS
In order enhance the conditional independence in the latent space, also adopt the mutual information constraints, which are as follows:
I(4; St-10-r; 41-18-) - Σ
I(4: Cja 12..14,,, 11-,--1,,, 8#
It is used characterize the language-controlled component C. by enhancing the dependence be- tween e; and I,, given dt-1:t7, 8t-1:t-: and enhancing the independence between 4 and each сj conditioning on corresponding (j, 8t-1:tr: lttl-: The 7 here refers the length of history that considering conditioning on, which is normally when assumed markov assumption. By applying the chain rule of mutual information for the first term and second term, can have the following:
T(4; Ci.. 8t-11-7, (0 T(,;; 1,,8811 at (I- T(1;; 8:-1:-111111 11 1t 0t-1:: I(4: Cja 8t-1:t77; 4:-1:; [11;; Cjrts lj, 81 1at1, (t-1:-7) 11--111---111-7-- I(4:; 1j, lit-r ag 11t-r)
In order to simultaneously optimize the representation and dynamics objectives, we apply the stop gradient to 8;-1, which converts the optimization of the conditional mutual information to the joint mutual information since the second terms in both directions are constant.
I(4; Cia 91-14-1;0---- $ C1 a(-1:--)) T(1: Cja Gja 14, 14 $ el,
Cja, 44 1:- 01-1:-7)))
Then we obtain the mutual information constraints 5 For each languauge component 4, we maximize the conditional mutual information between it and its corresponding language-controlle component 4 and minimize the summation of the conditional mutual information between it and the other language-controlled component Cj.
T
Il4: Ca,, (1-1-1 ΣΣ
Σ S i7i
81-10-1,
Concretely, we learn groups of mutual information estimators estimate these mutual information values by maximizing the Donsker- Varadhan representation Lower Bound with mutual information neural estimator Belghazi et al. (2018), then these estimators to help constrain the latent space of world model.
A.5 EXPERIMENT SETTINGS
A.5.1 NUMERAICAL. SIMULATION
In the simulation process, we follow the data generation process in Ioquatioo[aand and identifiabil- ity conditions in Thereom I While our framework is suitable for any number of language-controlleed
20

---

Published as a conference paper at ICLR 2025
components, we use three language components {41, 1s] for simplicity. The latent variables s have 6 dimensions, where nee = Rc, = ne. = 2. We set the number of values in each 4 as 3 and generate all 27 possible combinations of tasks, indicated by (4, 22, (2), e.g. (0, 2.). We take three tasks that are the recombinations of values that appeared in the other tasks as test tasks, and use the rest tasks as training tasks. Even though this makes cannot have an optimal number of tasks to identify all components simultaneously in perfect way, the R² shows that it does not affect lot. To make sure that different language components have different effects on the transition dynamics, we further embed them using embedding layers of different sizes before incorporating them in the transition functions. At each time step, a one-hot action of dimension 3 is taken. The functions in the data generation process are initialized with MLPs. This setting allows us to take the all-in-one strategy to identify C₁, C2 and C3.
A.5.2 META-WORLD
Environemnt Details Meta-World is benchmark suite of 50 robotic manipulation environments designed for multitask and meta-reinforccmment learning, where each task is accompanied by a corre- sponding language description. To best achieve the identifiability condition, we choose the common language component system that include most tasks, verb and object. We use these 18 tasks as train- ing tasks and the other 9 tasks that haven't i shown in the training tasks but can be either presented as the recombination of language components in the training tasks or tasks that have one known component in the training tasks, as test tasks. The tasks are follows:
Set
Task
Language Components (Verb, Object)
Train
Box-Close
Pick-Place, Cover
Bin- Bin-Picking
Pick-Place, Bin
Basketball
Pick-Place, Basketball
Soccer
Push, Soccer
Button-Press
Press, Button
Coffee-Pull
Pull, Mug
Dial-
Open, Dial
Door-Close
Close, Door
Door-Lock
Lock, Door
Door-Unlock
Unlock, Door
Faucet-Open
Open, Faucet
Handle-Pull
Pull, Handle
Push-Back
Pull, Puck
Push
Push, Puck
Plate-Slide-Batt
Retrieve, Plate
Sweep
Sweep, Puck
Window-Close
Close, Window
Drawer-Open
Open, Drawer
Test
Faucet-Close
Close, Faucet
Coffee-Pusii
Push, Mug
Handle-Press
Press, Handle
Drawer-Close!
Close, Drawer
Window-Open'
Open, Window
Door-Open'
Open, Door
Plate-Slide
Push, Plate
Handle-Press-Sicee
Press-Side, Handle
Coffee- Button"
Press, Coffee Button
Table 1: Task descriptions in Meta-World training and test sets. Test tasks that are recombinations of verb-object components from the training set. "Test tasks that contain either a novel verb variant (Press-Side) or a novel object (Coffee Button) not present in the training set.
Model Details For WM3C CNN, we build upon the JAX implementation of DreamerV3- small and use the medium-size visual encoder and policy module, to balance the extra parameters introduced by the modifications, For WM3C MAF, we substitute the CNN encoder and decoder in the WM3C
21

---

Published as a conference paper at ICLR 2025
CNN with MAE and keep the others to be the same WM3C CNN, referring Seo al. (2023). The details in implementing mutual information constraints and learnable masks are described as below:
MI The mutual information estimators implemented in JAX, referring MINE (Belg- hazi et al. 2018). An cosine annealing scheduler as below used to smoothly adjust the coefficient of mutual information constraints in the objective from the value to prevent the training instability brought by inaccurate estimation and variances in min-max optimization in the early stages.
Ustart
(Vend Usnn)(1 cost cos(Ttt) Mend > Ustart
scale(t)
(21)
Uend + (Ustan Vend) cos(TT)) otherwise
Mask We apply modified gated masks inspired by (Rajamanoharan et al. 2024) have more flexibility sparsity and alleviate the shrinkge effect of L1 loss, which potentially affects the exploration efficiency of policy. To accurately control the sparsity rate in the latents and mitigate the ineffective learning because of over-sparse, we also apply cosine annealing scheduler smoothly increase the sparsity rate threshold from zero, jointly with the adaptive LI loss. We apply the masks only to the deterministic part and keep the stochastic part of latent to be unchanged. The masking formula and L1 loss are as follows:
:= In
[ + bgate
(cxp(Fmmag x + bmag)
(22)
(a)((a)
froe(
[i(m) = [ReLL(
(23)
For DreamerV3, take the official JAX implementation of DreamerV3 (Hafner et al. 2023) from http://llbbbooom/mmmjjarrrreeee and use the medium version for all experiments.
For visual-based multi-task SAC (MT-SAC), we take the visual-based SAC implementation from Sostul--ootinouu-oooee
and modified it to the multi-task SAC by including the contextual embedding of the task language description, as well as the hyperparam- eters for Meta-worid, referring CARE (Sodhani et al. 2021).
https:/github.com/Kaaat
Training Details In the training stage, we train all the models (WM3C, DreamerV3, MT-SAC) for 5M environment steps in total in a multi-task online learning way, which on average for each task around 280K steps.
Adaptation Details In the adaptation stage, we fine-tune all the pretrained models (WM3C, Dream- erV3, MT-SAC) for 250K steps separately for each test task. The 250K steps can be reduced since for most tasks, the adaptation does not require much change and converges much earlier than that. For WM3C CNN and MAE, we only fine-tune the factorized dynamics module (representation mod- els and transition models), task encoder, decoding masks, and policy module, while the rest of the model frozen. We improve the threshold the sparsity rate to encourage more strict selection of the latent for decoding since for specific task, the ground truth latent should be less than the latent related to multiple tasks. For Dreamer V3 and MT-SAC, we finettune all the parameters.
Hyperparameters We summarized the important hyperparameters of WM3C CNN as follows: The WM3C MAE the same as WM3C CNN except the visual encoder and decoder are substituted with vanilla MAE. The important hyperparameters follows:
A.5.3 COMPUTATIONAL RESOURCES
All experiments are conducted 4 x Nvidia 3090 GPU. Training from scratch on the simulation experiment takes hours for one run, training on the 18 Meta-world tasks takes days for 5M steps and each fine-tuning task takes 8 hours for 250 steps.
A.6 EXTENDED DISCUSSIONS
A.6.1 COMPARISON WITH HIERARCHICAL RI
WM3C and hierarchical reinforcement learning (HRL) (Barto & Mahadevan 2003 Nachum et al.]
2018 Parr & Russell 1997) share conceptual similarity in their focus decomposing tasks into
22

---

Published as a conference paper at ICLR 2025
Category
Hyperparameter
Value
General
Batch size
50
Batch length
50
Train ratio
500
Training fill
5000
Image size
(64, 64, 3)
World Model
Component deterministic size
512
Component stochastic size
16
Component classes
32
CNN depth
48
Reward&Con layers
2
Reward& Cont units
512
Actor&Critic layers
3
Actor&Critic units
640
Task embedding dim
128
Token embedding dim
128
Mask types
[decoder, reward, cont]
Training
Training steps
x 106
Sparsity rate
0.25
Optimize MINE after steps
0
Maximize MINE after steps Minimize MINE after steps
le4
le4
Adaptation
Adaptation steps
2.5 x 10a
Sparsity rate
0.35
Maximize MINE after steps
0
Minimize MINE after steps
0
Table 2: Hyperparameters for WM3C CNN implementation in Meta-world.
Component
Hyperparameter
Value
MAE Encoder
Patch size Embedding dimension
8
256
Encoder depth
4
Number of heads
4
Mask ratio
0.75
MAE Decoder
Decoder embedding dim
256
Decoder depth
3
Decoder heads
4
Early convolution
True
VIT Parameters
Image size
8
Patch size Embedding dimension
1
128
Depth
2
Number of heads
4
Input channels
256
Table 3: MAE and VIT hyperparameters for WM3C MAE in Meta-world. The architecture decou-
ples the visual learning and dynamics learning, and small ViTs added in the dynamics module to align visual and dynamics information.
23

---

Published as a conference paper at ICLR 2025
smaller, more manageable components. However, their methodologies and applications diverge significantly. HRL typically structures policies into hierarchies, such as high-level managers setting subgoals for low-level controllers. These subgoals are often space-specific and lack explicit modeling of compositionality or causal dynamics among task components.
In contrast, WM3C explicitly learns compositional causal components grounded in the environ- ment's causal structure, enabling robust adaptation to unseen tasks by recombining previously learned components. This distinction is particularly evident in WM3C's use of language as com- positional modality, which aligns causal components with interpretable semantic elements. While HRL might excel in navigating hierarchical tasks, it may struggle with environments requiring flex- ible recombination of causal dynamics a strength of WM3C.
A.6.2 COMPARISON WITH PREVIOUS NON-LINEAR ICA AND CAUSAL REPRESENTATION LEARNING
WM3C builds upon the principles of non- non-linear independent component analysis (ICA) but departs from prior approaches in two critical ways. First, non-linear ICA methods, such as iVAE (Khe- makhem et al. 2019), Slow VAE (Klindt et al. (0200) and TCL. (Hyvarinen & Morioka 2016). focus on identifying latent variables under strong assumptions like independence of noise terms specific priors on functional forms. These methods are often dimension-wist and fail to scale effectively to complex systems with interdependent latent structures. Furthermore, WM3C considers the case of multiple auxiliary variables (language components) co-existing, while prior results only allow one auxiliary variable and a much simpler graph rather than the graphical model in the POMDP process. WM3C does not assume a parametric expression to acquire identification, enabling more flexible and practical estimation.
WM3C introduces block-wis identifiability, enabling the identification of aiguugee-contooller com- positional components rather than individual latent variables. We allow the separate language com- ponents connecting to the latents they control rather than having only one auxiliary variable con- necting to all the latents, different from previous non- non-linear ICA for time-series like LEAP(Yao) et al. 202I), TDRL(Yao et al.) (2022), NCTRL (Song et al. 2024). The formalization of block- identihability on language-control guarantees, which requires much smaller number of changes in auxiliary variables identify the latent variables of interest than the previous work.
components improves scalability and maintains theoretical
A.6.3 APPLICABILITY
The applicability of WM3C is currently demonstrated in multi-task training and single-task adap- tation settings in Meta- World. By pre-training the world model on a set of tasks, WM3C learns composable causal components while simultaneously training multi-task policy. This composable representation not only makes policy learning more efficient but also enables quick adaptation to new tasks with shared components through minimal adjustments of components' dynamics. This makes WM3C particularly suitable for domains where tasks can be naturally decomposed into modular components, such robotic manipulation Furthermore, while the current framework primarily uses language as the compositional modality, WM3C is expected generalize to other modalities, such as visual or auditory signals. For exam- plc, auditory signals could be decomposed into frequency domains or audio patterns, allowing the framework leverage the stability and generality of the composition system to handle multi-modal tasks effectively.
other scenarios with well-defined causal structures.
Current formulation of WM3C has limitations when applied to environments with highly overlap- ping causal components or complex language instructions (c.g., long or ambiguous sentences), where disentangling and identifying independent components becomes more challenging. How- ever, these challenges are not insurmountable. One potential solution is to leverage large language models (LLMs) or human annotations to first convert long and intricate language descriptions into clearer, more structured formats. This preprocessing step simplifies the identification of composi- tional components by aligning the input with the framework's 'ss assumptions. Furthermore, causal components overlapping and complex structure learning (c.g. hierarchy) have been explored in the field of causal representation learning (Liu et al., 2203 Kong et al. 2023a] Morioka & Hyvärinen) 2023). By drawing these existing studies, the WM3C framework can be extended to effectively
24

---

Published as a conference paper at ICLR 2025
handle such scenarios. Jointly with the extension to offline learning, these enhancements would enable WM3C to be scale-up and learn a more complex and generalized causal component system.
A.7 EXTENDED RESULTS
A.7.1 SINGLE TASK PERFORMANCE META-WORLD
We present the success rate curves of all 18 training tasks, comparing WM3C, DreamerV3, and visual-based multi-task SAC. Results show that WM3C outperforms DreamerV3 in most tasks, with CNN.
WABC MAB
WMICC CNN
Breamer V21
ME-SAC
Plat-dide-bec
MM
Dear-book
Binglcking
Windme-looo
Pasckcchll
Prsch
1.0
Surtiea Ratis
0.0
Pusb-bach bubb-back
0.0
Box ciose
Drawer open
Door wnlock
1.0
0.2
Stop (x199)
Figure 7: Training success rates Meta-world: success rates curve of all 18 tasks
A.7.2 ABLATION STUDY
To further investigate the effects of integrated modules, conduct ablation study on scales of training in Meta-world, a 5 tasks training in 1M steps and a 18 tasks training in 2M steps. The tasks are push. faucet-open. handle-pull, door-lock and drawer-open, subset complete 18 tasks. The 18 tasks are the same those we use for training. We compare WM3C with MAE and CNN, and remove the mutual information constraints and decoding masks. The difference between the WM3C without masks and mutual information constraints and DreamerV3 is still present. WM3C features a language-conditioned, factorized dynamics module and learnable task embedding, which sets it apart from DreamerV3.
We see that all modules contribute to certain parts of the model's robust performance (see Figure 10 Using MAE the visual module consistently improves the representation quality in both small and large-scale training regimes. The incorporated mutual information and sparsity constraints significantly improve sample efficiency consistently in both scales. Interestingly, find that when
25

---

Published as a conference paper at ICLR 2025
the task number increases, the benefits brought by the mutual information and sparsity constraints are larger.
20000 Ave
Figure 8: Ablation study on two scales of training data: tasks in 1M steps training and 18 tasks in 2M steps training.
Dase-lock
Handle pull
Dawwer-per
Flucett-peen
Push
Sunonee Rldeo
0.6
0.0
1.0
(X10P)
Figure 9: Individual tasks plot tasks and 1M steps training.
A.7.3 INTERVENTION EFFECT ON THE -ONIIAGG--COTOOLLEEE COMPONENTS
Here we present more examples of the effect of intervention on the language components different tasks.
26

---

Published as a conference paper at ICLR 2025
WMBC MAE
WSEC CNN
WXEB CNN WID MASK & MI
Bocesunleck
Browwrsppen
Twurebsrgen
Boce-loos
Windeotsclooe
088
04
0.0
Placesllde-bacc
Biak-turn
Binpicking
Hondle pull
Soccer
1.0
0.6
Eucoles Rate
0.
Pushsbarke
Netton-grwo
Doardock
Baokkichaal
0.8
0.4
0.0
Push
Bonecloww
Colneepul
1.0
0.6
0.2
Step (x19f)
Figure 10: Individual tasks plot in 18 tasks and 2M steps training.
27

---

Published as a conference paper at ICLR 2025
Anwwica deredde
Figure 11: Intervention effect on the object component. The first row is the image observation WM3C receives, and the second row is the observation reconstruction from the latent whose object component has been intervened.
28

---

Published as a conference paper at ICLR 2025
wondinaiee dbe akromla
Figure 12: Intervention effect on the verh component. The first row is the image observation WM3C receives, and the second row the observation reconstruction from the latent whose verb component has been intervened.
29


Reviews:
[{"review_id": "uooZsmrUWf", "reviewer": "Reviewer_svqL", "summary": "After reading the rebuttal, I agree that language information could be somewhat generally available and I appreciate the effort in fixing writing issues with the paper, so I decided to slightly increase my score.  \n\nThis paper proposes a verb/object decomposed world model where the verb and object for a task are given beforehand.  There is improved sample complexity on metaworld, but the method seems like it’s much more complicated and using quite a bit of domain knowledge.  The analytical experiments are good.  \n\nnotes from reading the paper: \n  -RL generalization to new environments is difficult, especially if there are unseen dynamics.  \n  -World modeling with composition causal components.  \n  -Causal dynamics over composable elements.  \n  -Language as a compositional modality.  \n  -Masked autoencoder with MI constraint and sparsity to capture high-level semantic information and disentangle transition dynamics.  \n  -Experiments on real robots and numerical simulations, method is called WM3C.  \n  -DRL algorithms have a difficult time to adapt to changes in the environment.  Many ways of getting around this problem require domain-specific knowledge.  \n  -Composable causal components and dynamics can be uniquely identified (this seems surprising to me).  \n  -Language-controlled components are a subset of state dimensions that are directly controlled by individual language components.  \n  -[o_t, r_t] = g(s_t, \\eps_t).  s_t is defined as set of composable components.  The c_{it} is defined as a distribution conditioned on the language l_t and the previous state and action.  \n  -Blockwise identifiability requires the mixing function to be invertible and smooth.  Also every state must have density.  \n  -Utilize model-based RL algorithm DreamerV3.  \n  -Additional sparsity constraint. \n  -Only reward-relevant states are used for the policy model.", "questions": "-Typo on first page, poorly is spelled “pooly”.  \n  -How much effort needs to go into figuring out the “controllable composable components”?  This might be less desirable than a system which works for any general MDP, by processing the observations.  \n   -Experiments show effective component identification over the course of training in a synthetic task.  \n  -The gains in sample complexity on meta-world still seem to be significant.  \n  -Figure 6 looks fairly unprofessional and is hard to make sense of.  Intervention the verb component makes the arm blurry, which is good to verify.  However, it doesn’t really clearly show a change what the action or the object are.  This may be related to the nature of pixel-space reconstructions.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "Good analytical experiments, some improvements in sample complexity on metaworld.", "weaknesses": "Method adds substantial complexity and also seems to require much more domain knowledge than the baseline.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Gz3XIRmutP", "reviewer": "Reviewer_ByBw", "summary": "The paper presents World Modeling with Compositional Causal Components (WM3C), a novel approach in reinforcement learning (RL) aimed at improving generalisation in novel environments. Inspired by human reasoning, WM3C identifies and utilizes causal, compositional elements, leveraging language as a guiding modality. The approach builds on the causal structure of RL tasks by combining theoretical identifiability guarantees with a masked autoencoder architecture and adaptive sparsity. WM3C outperforms existing methods in identifying causal components and supports efficient policy learning and generalization, as shown through synthetic data simulations and robotic manipulation tasks in the Meta-World environment. Although promising, the results’ scalability in more diverse real-world scenarios remains unclear.", "questions": "1. How is the presented setup related to hierarchical RL?\n2. How is it possible that WM3C can acquire the latent language structure with only 20 tasks? It feels very nice but somewhat implausible.   \n3. Adaptation results in Fig 5 feel somewhat poor. Could you comment on that?\n\t1. I'm wondering if pure Dreamer is a proper baseline. E.g., recent [1] shows that transfer is better if forgetting mitigation is applied\n4. Could you describe the applicability of the proposed approach? Is it a multi-task setting like Meta-World or perhaps pertaining? \n\n[1] Fine-tuning Reinforcement Learning Models is Secretly a Forgetting Mitigation Problem, Wolczyk et al.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. An exciting and important problem is being tackled. \n2. Comprehensive conceptual and theoretical analysis.", "weaknesses": "My main concern is a relatively lean experimental section; see the questions and points below.\n\n1. As the proposed method is quite complicated, it is unclear if the effects can be due to the described mechanisms.  \n2. Only one real-world environment is tested (I am aware that it might be partially due to a lack of proper benchmarks).\n3. Only one algorithm is tested.\n4. The transfer and generalization results are relatively weak (at the same time, I am not quite sure if the current scale of experiments is large enough to reveal strong effects.)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1MDcpLH6WI", "reviewer": "Reviewer_Ytw5", "summary": "This paper proposes a method to identify the underlying composable causal components in reinforcement learning (RL) environments to facilitate policy generalization. The main idea is to leverage language descriptions of the task to identify the underlying latents that govern the data generation process. The paper implements this theoretical result based on the model-based RL framework Dreamer V3, with additional constraints to foster sparsity and modularity of the learned components. The proposed method is empirically verified on synthetic data and robot manipulation tasks in Meta-World.", "questions": "- What is the relation between the identifiability result in Theorem 1 and prior results in non-linear ICA?\n\n- What is the effect of each constraint in the algorithm? It seems that no ablation on them is presented in the experiments section.\n\nMinor:\n\n- Typo: Line 252 -> bfrom.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The identification result in Theorem 1 is interesting and implies an advantage of language-driven RL agents compared to non-language counterparts.\n\n- Empirical results are generally good. I also like the intervention results in Figure 6.", "weaknesses": "I think the biggest weakness of the paper in its current form is the discrepancy between the motivation, theoretical results, and the algorithm.\n- The authors motivate their algorithm by the need to identify composition causal components in the underlying data generation model, which possesses some ideal traits like modularity and sparsity, as introduced in the causal learning literature. However, if our aim is to identify \"causal\" components, then we need to recover not only those components but also their causal structure (as also noted by the authors in Lines 74-75), or the word \"causal\" itself would become meaningless. Yet, the identification result presented in Theorem 1 gives no guarantee of recovering such a causal structure.\n\n- On the other hand, the algorithm indeed uses modularity and sparsity constraints. However, such constraints are not reflected in the identification results, making the theory and the algorithm not really consistent. It feels to me that those constraints are more heuristics inspired by the causal learning literature, rather than motivated by Theorem 1.\n\nI think to rectify the potential understanding, the authors should explicitly make it clear that the theoretical results do not guarantee \"causal\" identification.\n\nAlso, I think the authors should make the relationship between their theoretical results and prior results in representation identifiability more clear.\n\n- The authors have claimed that prior methods for identifying latent variables use \"different priors like functional form prior, independent noise conditions and auxiliary variables\" (Line 176), which limits their scalability. Yet, the \"language components\" used by the paper can also be modeled as \"auxiliary variables\", which indeed play a crucial role in the identifiability result.\n\n- Technically, the high-level idea of proving identifiability is similar to the results in [1, 2], which rely heavily on the conditional independence between latent variables and observable variables. The authors could include a discussion of those results.\n\n[1] Nonlinear ICA Using Auxiliary Variables and Generalized Contrastive Learning.\n\n[2] Variational Autoencoders and Nonlinear ICA: A Unifying Framework.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Bn5rE0QeDp", "reviewer": "Reviewer_mPjx", "summary": "This paper introduces the World Modeling with Compositional Causal Components (WM3C) framework for reinforcement learning (RL). The framework aims to address the problem of generalization to unseen environments by leveraging compositional causal components by incorporating language-guided compositional causal reasoning inspired by human reasoning. Using language as a compositional modality, the framework decomposes latent space into semantically meaningful modular components with causal dynamics. Unlike prior methods focusing on invariant representation learning or meta-learning, WM3C aims to learn compositional components and utilize them to generalize effectively to new environments. The paper provides theoretical guarantees for the identifiability of these components under certain assumptions and implements the concept using a masked autoencoder (MAE) with mutual information constraints and adaptive sparsity regularization. WM3C's effectiveness is demonstrated through experiments on (synthetic) numerical simulations and (real-world) robotic manipulation tasks on the MetaWorld environment, where it outperforms existing approaches on generalization/adaptation to unseen tasks.", "questions": "**Questions**:\n\n1. Can the authors elaborate on the train/test methodology of the baseline model (Dreamer V3) on MetaWorld? Were textual task descriptions included in the encodings? Did the authors use any existing implementation of Dreamer V3 on MetaWorld, or is it novel for this paper?\n\n2. Why didn’t the authors include performance comparisons with generic RL models like Multitask-[SAC|PPO], which are provided out-of-the-box with the MetaWorld environment?\n\n3. Could the authors provide more intuition or real-world examples where the faithfulness assumption might fail? How sensitive is the performance of WM3C to violations of this assumption?\n\n4. What are the key limitations of WM3C in environments with highly overlapping causal components (more likely in richer/complex language instructions)? How might the model be adjusted to handle such scenarios?\n\n5. How sensitive is the framework to the choice of language components? Would other compositional modalities (e.g., visual <-> auditory signals) work as effectively?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- **Originality**: WM3C's approach of integrating causal reasoning with compositional learning in RL is interesting. The use of language-guided decomposition adds a unique dimension to generalization in RL.\n\n- **Quality**: Theoretical grounding is generally robust, with guarantees on the identifiability of compositional components. While the presented experimental results support the overarching claim, the strength and suitability of the experiments have room for improvements.\n\n- **Clarity**: The clarity of the technical writing could be improved. Discussed further in the ‘Presentation’ section of the review.\n\n- **Significance**: Generalization in RL is a well-known challenge. The paper tackles this critical problem in RL—generalization to unseen environments—and presents a solution with potential for significant impact in real-world applications.", "weaknesses": "**Discussion** \nThe crux of the work's novelty lies in the mapping between orthographic components (as we understand it: textual, visual, etc.) and latent space blocks, thereby allowing generalization to unseen tasks through the composition of these learned invariant primitive blocks. Such compositional generalization is a well-researched (and active) area, especially in the NLP, VQA, and EAI domains. For example, numerous papers using the CLEVR images and captions dataset, along with mixing the attributes in the VQA domain, explore such generalization using text and image modalities. One key enabler of these models was applying masked object encoders, thereby infusing the models with knowledge of what the objects are.\nIn this work, the effect of using MAE (vs. CNN) is not entirely clear or conclusive (from Fig. 5). However, the point remains: for the visual modality, the object and latent space mapping is fairly well-established and by no means novel, especially if object masking is in play.\nThe mapping between textual semantics — especially explicit mapping of complex verbs or actions with latent space blocks — is, however, quite promising and, to my knowledge, novel for this specific domain. However, the experiments conducted to show the efficacy use very simple instantiations of the language modality. There are a fairly large number of environments in the embodied AI space (like AI2Thor's ALFRED, Habitat, etc.) that have more realistic textual instructions for task completion scenarios—which, in my opinion (IMO), would be much more convincing to demonstrate the efficacy of this approach, the effectiveness of textual component identifiability, and the overall claim of the effectiveness of using causal RL.\n\n\n\n* **Limited Evaluation** The presented experiments are weak. Experiment 1 using synthetic data (3 tokens, 3 control groups) should not be a major experiment with supporting figures in the main paper, but rather should have been a preliminary ‘smoke-test’ \\- maybe getting a subsection of an appendix section at best. I have major concerns about the baselines of Experiment 2 conducted on MetaWorld’s robotic manipulation tasks.  The only baseline comparison is against DreamerV3 (2 variants).  I am wondering why any of the available, known, easily replicable RL models (e.g. MT-\\[SAC|PPO\\]) were not included as baselines? \n\n* **Evaluation environments**: Given the claim of mapping ‘textual components’, the choice of environments and models could be more commensurate. Specifically, the proposed approach should be evaluated on environments that allow richer text+vision modality and agent interaction. There are a slew of such multi-modal environments across subdomains, like in the embodied AI domain: RoboThor (the ALFRED task), Habitat or, simpler HomeGrid. Given the richer, realistic text instructions on these domains, the claim of unique identification of causal textual components would be far better convincing. \n\n* **Scalability**: The scalability of WM3C to environments with more complex dynamics and a higher number of compositional components has not been fully explored. Future work should test the framework in larger-scale environments.\n\n* **Limited Analysis of Failure Cases**: It would be helpful to include a discussion on the limitations of WM3C, particularly regarding situations where the causal components may overlap or cannot be distinctly identified. Understanding where and why the model fails is crucial for improving the framework.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)