Towards Empowerment Gain through Causal Structure Learning in Model-Based Reinforcement Learning
Authors: Hongye Cao
Fan Feng
Meng Fang
Shaokang Dong
Tianpei Yang
Jing Huo
Yang Gao
Venue: ICLR 2025 Poster (2025)
Keywords: Causal RL
MBRL
Empowerment
Intrinsic Motivation
Primary area: reinforcement learning
TL;DR: We propose a framework, Empowerment through Causal Learning , where an agent with the awareness of causal models achieves empowerment-driven exploration and utilize its structured causal perception and control for task learning.

Abstract: In Model-Based Reinforcement Learning (MBRL), incorporating causal structures into dynamics models provides agents with a structured understanding of the environments, enabling efficient decision. 
Empowerment as an intrinsic motivation enhances the ability of agents to actively control their environments by maximizing the mutual information between future states and actions. 
We posit that empowerment coupled with causal understanding can improve controllability, while enhanced empowerment gain can further facilitate causal reasoning in MBRL. 
To improve learning efficiency and controllability, we propose a novel framework, Empowerment through Causal Learning (ECL), where an agent with the awareness of causal dynamics models achieves empowerment-driven exploration and optimizes its causal structure for task learning. 
Specifically, ECL operates by first training a causal dynamics model of the environment based on collected data. We then maximize empowerment under the causal structure for exploration, simultaneously using data gathered through exploration to update causal dynamics model to be more controllable than dense dynamics model without causal structure. In downstream task learning, an intrinsic curiosity reward is included to balance the causality, mitigating overfitting. 
Importantly, ECL is method-agnostic and is capable of integrating various causal discovery methods. 
We evaluate ECL combined with $3$ causal discovery methods across $6$ environments including pixel-based tasks, demonstrating its superior performance compared to other causal MBRL methods, in terms of causal discovery, sample efficiency, and asymptotic performance.

Published as a conference paper at ICLR 2025
TOWARDS EMPOWERMENT GAIN THROUGH CAUSAL STRUCTURE LEARNING IN MODEL-BASED REINFORCE-
MENT LEARNING
Hongye Cao Fan 2
Meng Fang*
Shaokang Dong*
Tianpei Yang
Jing Huo! Yang Gao
National Key Laboratory for Novel Software Technology, Nanjing University
2University of California, San Diego *MBZUAI *University of Liverpool
School of Intelligence Science and Technology, Nanjing University
ABSTRACT
In Model Reinforcement Learning (MBRL), incorporating causal structures into dynamics models provides agents with the structured understanding of envi- ronments, enabling more efficient and effective decisions. Empowerment, intrinsic motivation, enhances the ability of agents to actively control environments by maximizing mutual information between future states and actions. We posit that empowerment coupled with the causal understanding of the environment can im- prove the agent's 's controllability over environments, while enhanced empowerment gain can further facilitate causal reasoning. To this end, we propose the framework that pioneers the integration of empowerment with causal reasoning. Empowerment through Causal Learning (ECL), where an agent with the awareness of the causal dynamics model achieves empowerment-drive
an
exploration and optimizes its causal structure for task learning. Specifically, we first train causal dynamics model of the environment based on collected data. Next, we maximize empowerment under the causal structure for exploration, simultaneously using data gathered through exploration to update the causal dynamics model, which could be more controllable than dynamics models without the causal structure. We also design an intrinsic cu- riosity reward to mitigate overfitting during downstream task learning. Importantly, ECL is method-agnoste and can integrate diverse causal discovery methods. We evaluate ECL combined with causal discovery methods across environments including both state- based and pixel-based tasks, demonstrating its performance gain compared to other causal MBRL methods, in terms of causal structure discov- ery, sample efficiency, and asymptotic performance in policy learning. The project page is https:/ /sites..s...//////////e/////tte/ google. com/ viii/vieweeeeeeeeeeeeeee hec/eellla 1112-194./..
INTRODUCTION
Model-Based Reinforcement Learning (MBRL) uses predictive dynamics models to enhance decision- making and planning (Moerland et al., 2023). Recent advances in integrating causal structures into MBRL have provided more accurate description of systems, achieve better adaptation (Huang et al., 2021; 2022; Feng & Magliacane, 2023), generalization (Pitis et al., 2022; Zhang et al., 2020; Wang et al., 2022c; Richens & Everitt, 2024; Lu et al., 2021), and avoiding spurious correlations (Ding al., 2022; 2024; Liu et al., 2024; Mutti et al., 2023a).
However, these methods often passively rely on pre-existing or learned causal structures for policy learning generalization. In this work, aim to enable the agent actively leverage causal structures, guiding more efficient exploration of the environment. The agent can then refine its causal structure through newly acquired data, resulting in improvements in both the causal model and policy, This could further enhance the agent's 's controllability over the environment and its learning efficiency,
We hypothesize that agents equipped with learned causal structures will have better controllability than those using traditional dynamics models without causal modeling. This is because causal
Equul contribution. Corresponding to Jing Huo (huoj i www.nnnnnnnuu.....eee.mmm

---

Published as a conference paper at ICLR 2025
(a)
(b)
Target noce
Notse node (movable)
S
Naise node (unmovable)
Controllable efficient X
target
3
F42
C2
[+]
a
1+2
Controllable efficient v
True target X
a
[+]
1+2
Fi+1
Controllsble efficient
True target
Figure 1: (a). An example of robot manipulation task with three trajectories and three nodes: one target node (movable) and two noisy nodes (one movable, one unmovable). (b). Underlying causal structures with a factored MDP, Different nodes represent different dimensional states and actions.
structures inform agents to explore the environment more efficiently by nulling out the irrelevant system variables. This assumption serves as intrinsic motivation to guide the policy in exploring higher-quality data, which in turn improves both causal and policy learning. Specifically, we employ empowerment gain, an information-theoreti framework where agents maximize mutual information between their actions and future states to improve control (Leibfried et al., 2019; Klyubin et al., 2005; 2008; Bharadhwaj et al., 2022; Eysenbach et al., 2018 88; Mohamed & Jimenez Rezende, 2015), as the intrinsic motivation to measure the agent's ' controllability. Concurrently, through empowerment, agents develop more nuanced comprehension of their actions' consequences, implicitly discovering the causal relationships within their environment. Hence, by iteratively improving empowerment gain with causal structure for exploration, refining causal structure with data gathered through the exploration, the agent should be able to develop robust causal model for effective policy learning.
We give a motivating example (Fig. (a)) in a manipulation task, where the robot aims move a target node while avoiding noisy nodes. Three possible trajectories (rows 1-3) are shown with different levels of control, efficiency, and success. Row (irrelevant states) represents the least effective trajectory that can not control nodes and find the target, while rows 2 and 3 (controllable states) demonstrate learned control and efficiency, with high empowerment focusing on movable objects. In the corresponding causal graphs, represented as Dynamics Bayesian Network in Fig. 1(b), s², s?. Id denote the states of three objects. For simplicity and clarity, we assume each object is represented by single variable. The graph illustrates the causal relationships between these states, actions, and rewards. Assuming the agent follows the causal structure (Fig. (b)), it will likely execute actions similar to rows 2 and 3 since there causal relationships between actions and states of movable objects, effectively improving controllability. Through exploration with better control, agents can facilitate improved causal discovery of the task, leading to high-reward outcomes and resulting in more efficient task completion like row 3.
To this end, we propose an Empowerment through Causal Learning (ECL) framework that actively leverages causal structure to maximize empowerment gain, improving controllability and learning efficiency. ECL consists of three main steps: model learning, model optimization, and policy learning. In model learning (step 1), we learn the causal dynamics model with a causal mask and a reward model. We then integrate an empowerment -rrienr--eeddeiee exploration policy with the learned causal structure to better control the environment (step 2). We alternately update the causal structure with the collected data through exploration and policy of empowerment maximization. Finally, the optimized causal dynamics and reward models used to learn policies for downstream tasks with curiosity reward to maintain robustness and prevent overfitting (step 3). Importantly, ECL is method- -aahhooie,, being able to integrate diverse causal discovery (i.e., score-based and constraint-basedl methods. The main contributions of this work can be summarized as follows:
To improve controllability and learning efficiency, propose ECL, novel method-agnostic framework that actively leverages causal structures to boost empowerment gain, facilitating efficient exploration and causal discovery.
2

---

Published as a conference paper at ICLR 2025
ECL leverages causal dynamics model to conduct empowerment-based exploration. It also utilizes controllable data gathered through exploration to optimize causal structure and reward models, thereby delving deeper into the causal relationships among states, actions, and rewards.
We evaluate ECL combined with 3 causal discovery methods across 6 environments, encompassing both In-Distribution (ID) and nut-Of-Distributi (OOD) settings, well pixel-based tasks. Our results demonstrate that ECL outperforms other causal MBRL methods, exhibiting superior performance in terms of causal discovery accuracy, sample efficiency, and asymptotic performance.
2
PRELIMINARIES
2.1 MDP WITH CAUSAL STRUCTURES
Markov Decision Process In MBRL, the interaction between the agent and the environment is formalized
Markov Decision Process (MDP). The standard MDP is defined by the tuple M = (S,A, T, Po: r, 7). where S denotes the state space, A represents the action space, T(s'|s. a) is the transition dynamies model, r(s, a) is the reward function, and μ4 is the distribution of the initial state 80. The discount factor y E [0, 1) is also included. The objective of RL is to learn a [0, 1] that maximizes the expected discounted cumulative reward (M((T)
policy : S x A
Bgg Beg~ aar
Di= t= ('r(st,a)) a.).
Structural Causal Model A Structural Causal Model (SCM) (Pearl, 2009) defined by a distri- bution over random variables V {s]. ... St ...... ai Graph (DAG) g = (V,&) with a conditional distribution (P(()) for node v, E V. Then the distribution can be specified as:
st-1.-. , } and Directed Acyelic
M
p(u
MI) = II p(u ((
(1)
where PA(U;) is the set of parents of the node v in the graph G.
Causal Structures in MDP We model a factored MDP (Guestrin et al., 2003; 2001) with the underlying SCM between states, actions, and rewards (Fig.Ib). lb). In this factored MDP, nodes represent system variables (different dimensions of the state, action, and reward), while edges denote their relationships within the MDP. We employ causal discovery methods to learn the structures of G. We identify the graph structures in G, which can be represented as the causal mask M. Hence, the dynamics transitions and reward functions in MDP with causal structures are defined as follows:
st {
= (M St, Ma- o at Cr.i,t) rr = R(o.(s M).a,) a.)
where s/م represents the next state in dimension i, M {0, 1] and M/~+ {0, 1]/0|x|s| are the causal masks indicating the influence of current states and actions on the next state, respectively, denotes the element-wing product, and cs,i, represents i.i.d. Gaussian noise. Each entry in the causal mask M (represented as the adjacency matrix of the causal graph 9) indicates the presence (1) or absence (0) of a causal relationship between elements. The reward r, is a function of the state abstraction de | M) under the learned causal mask M, which filters out the state dimensions without direct edges to the target state dimension, and the action at. We list the assumptions and propositions in Appendix C.
2.2 EMPOWERMENT
Empowerment is to quantify the influence an agent has over its environment and the extent which this influence can be perceived by the agent (Klyubin et al., 2005; Salge et al., 2014; Jung et al., 201 Within our framework, the empowerment the mutual information between the agent action a, and its subsequent state $141 under the causal mask M as follows:
E == max I(84415 M). x((3))
where E is used to represent the channel capacity from the action to state observation. <--1st) the conditional distribution of actions given states.
3

---

Published as a conference paper at ICLR 2025
Step Model Learinng
Step 2 &: Model Optimization and Pollcy Learring
[
Peete Peotnal
A(a
Caussi Dynemice Model
Reward Model
Causal Dynamice Model
max (1 (N.13 4/ M) (511)
1 = Fisk + Arour
Reward Model
anxt Empowerment- Exploration
Policy Leamning
Reward
Optimization
(s, Ap Suae Beat findl
(8p ap S A Fins)
(5, Fp Sug Beaty food
dynamics encoder
Figure 2: The framework overview ECL. Gold lines: model learning. Blue lines: model optimiza-
tion alternating with empowerment-driven exploration (yellow lines). Green lines: policy learning.
EMPOWERMENT THROUGH CAUSAL LEARNING
An illustration of the ECL framework is shown in Fig. 2, comprising three main steps: model learning, model optimization, and policy learning. In model learning (step 1), we learn causal dynamics model with the causal mask and reward model. This causal dynamics model is trained using collected data to identify causal structures (i.e., causal masks M). by maximizing the likelihood of observed trajectories. The reward model is trained based on state abstraction that masks irrelevant state dimensions with the causal structure. With the learned causal structure, we integrate empowerment- driven exploration for model optimization (step 2). This process involves learning the empowerment policy ne that enhances the agent's controllability by actively leveraging the causal mask. We alternately update the policy He for empowerment maximization and generate data with ne to optimize the causal mask and reward model Po. Finally, step 3, the learned causal dynamics and reward models are used learn policies for the downstream tasks. In addition the task reward, to maintain robustness and prevent overfitting, an intrinsic curiosity reward is incorporated to balance the causality.
3.1 STEP 1: MODEL LEARNING WITH CAUSAL DISCOVERY
We first learn causal dynamics model with the causal mask and reward model for the empowerment and downstream task learning. Specifically, a dynamics encoder is trained by maximizing the likelihood of observed trajectories D. Then, the causal mask is learned based the dynamics model and reward model is trained with the state abstraction under the causal mask and action.
Causal Dynamies Model The causal dynamics model is composed with dynamics model P and a causal mask M. The dynamics model maximizes the likelihood of observed trajectories as follows:
ds Σ
Layn = (e
log Po, (5/
18, ari 6c)
where dy is the dimension of the state space, and &e denotes the parameters of the dynamics model. We train the dynamics model as a dense dynamics model that incorporates all state dimensions to capture the state transitions within the environment, facilitating subsequent causal discovery and empowerment. Additionally, we assess the performance of the dense model, specifically the baseline MLP, within the experimental evaluations detailed in Section 5. Next, we this learned dynamics model for causal discovery.
Causal Discovery For causal discovery, with the learned dynamics model Pa_. further embed the causal masks = and Ma-s into the learning objective. To learn the causal mask, we employ both conditional independence testing (constraint-based) (Wang et al., 2022c) and mask learning by sparse regularization (score-basee -bsedd (Huang et al., 2022). We further maximize the likelihood of states by updating the dynamics model and learned masks. Thus, the learning objective for the causal dynamics model is as follows:
ds
Le-dyn E(
- St, Ma-ss O as Cajee) + Leausal log Po.
|
(5)
4

---

Published as a conference paper at ICLR 2025
where Leausal represents the objective term associated with learning the causal structure. C
log (ogg(5(411le 8, \ { and L -A |M]] represent constraint-basel and score-
based objectives respectively. AM is regularization coefficient.
Reward Model After obtaining the causal dynamics model, we process states using the causal mask M to derive state abstractions de(- MM for the reward model learning, effectively filtering out irrelevant state dimensions. Simultaneously, the reward model Pg, maximizes the likelihood of observed rewards sampled from trajectories D):
Lrew R (.) [lgg.p (relde(s, M))
In this way. ECL leverages causal understanding enhance both state representation and reward prediction accuracy. Finally, the overall objective of the model learning with the causal structure is to maximize C Ldyn + Ly-dya frew-
3.2 STEP MODEL OPTIMIZATION WITH EMPOWERMENT-DRIVEE EXPLORATION
In Step 2, we optimize the learning of the causal structure and empowerment. As depicted in Fig. this procedure alternates between optimizing the impowerment-drivee exploration policy Te and update the causal mask M using data gathered through exploration. Furthermore, ensure the stability, we update the reward model to adapt to changes in state abstraction induced by updates to the causal mask M. Note that the dynamics model Po. learned in Step remains fixed, allowing for focused optimization of both the causal structure and the empowerment in alternating manner. The causal structure is optimized by the causal mask through maximizing Leausai (Eq. 5), while keeping the parameters of de fixed during this learning step.
Empowerment-drriee -drivee Exploration To enhance the agent's control and efficiency given the causal structure. instead of maximizing I (st+1, Qe (r]s;) at each step, we consider baseline that uses the dense dynamics model Ps. without the causal mask M. We then prioritize causal information by maximizing the difference empowerment gain between the causal and dense dynamies models.
We first denote the empowerment gain of the causal dynamics model and dense dynamics model as Eg., (sIM)= = max, I ($t+1:00 0t 8t; &e. M) and &g, (s) max, (st+1; 0e sy; Φc), respectively. Here, Eg. (s) corresponds to the dynamics model without considering causal structures.
Then, we have the following learning objective:
max R' (s) Ess (+)] a~mr, (a|s)
In practice, we employ the estimated Eg_ (8 M) and (s) with the policy me for computing, specifically;
(+) = max ET (arls)) (s [log Po Sts az; M, Φe) - log P(s [s)]+ (a/s)
(8)
and:
(x))
max (alls)))) (*41) [log Po. (8441 81s 455 6e) - log P(s a>wx(als)
(9)
where P(s (8t) is the conditional distribution of the current state. Hence, the objective function Eq. is derived as:
max (als)
H(8841 | sy; -((441 + + [KL (Po (8e41 8t, az; Po. (8t41 |1, a)),
(10)
where M(s141 st; M) and H(s141 st) denote the entropy at time ++ 1 under the causal dynamics model and dense dynamics model, respectively. For simplicity, we update ne by optimizing the KL term.
Model Optimization In Step 2, we fix the dynamics model Po. and further fine-tune the causal mask M and the reward model Pp.. We adopt alternating optimization with the policy me to optimize the causal mask. Specifically, given M, we first optimize т.. The policy ne is designed to collect controllable trajectories by maximizing the distance of empowerment between causal and dense models. These collected trajectories are then used to optimize both the causal structure M and reward model Pg.-
5

---

Published as a conference paper at ICLR 2025
3.3 STEP 3: POLICY LEARNING WITH CURIOSITY REWARD
We learn the downstream task policy based on the optimized causal structure. To mitigate potential overfitting of the causality learned in Steps 1 &2, we incorporate curiosity-based reward an intrinsic motivation objective or exploration bonus, in conjunction with a task-specific reward, to prevent overfitting during task learning:
KIL (Pen( Po (8c.1 Ocr M))
Feur (8, a) ( (cee))))))))
(1I)
-KIL -KIL (Pour (sttl))))))))
Po. Po. (o-
where Peny the ground truth dynamics collected from the environment. By taking account of Yeurs we encourage the agent explore states that the causal dynamics cannot capture but the dense dynamics can from the true environment dynamics, thus preventing the policy from being overly conservative due to model learning with trajectories. Hence, the shaped reward function is:
r(s, a) Tinsk (s. a) + Ar Arcar (s,
(12)
where (s a) is the task reward, is a balancing hyperparameter, In section D.8, we conduct ablation experiments to thoroughly analyze the impact of different shaped rewards, including curiosity, causality and original task rewards.
4 PRACTICAL IMPLEMENTATION
We introduce the practical implementation of ECL for casual dynamics learning with empowerment- driven exploration and task learning. The proposed framework for the entire learning process is illustrated in Figure 2, comprising three steps and the full pipeline is listed in Algorithm 1.
transition collection policy Teollet by formulating reward function that incentivizes selecting transitions that cover more state-action pairs to expose causal relationships thoroughly. We train the dynamics model Ps. by maximizing the log- liglliiel Layn. following Eq. 4. Then, we employ the causal discovery approach for learning causal mask M by maximizing the log-likelihood Le-dyn followed Eq. 5. Subsequently, we train the reward model P.. with the state abstraction 6.(s M) by maximizing the likelihood.
Step 1: Model Learning Initially, following (Wang et al., 2022c),
exploration by max (als) Es,,,, S141 (Eg, (s|M) - &s, (s)] followed Eq. 7 with causal dynamies model and dense dynamics model for policy ne learning. Furthermore, the learned policy Te is used to sample transitions for updating casual mask M and reward model. We alternately perform empowerment- driven exploration for policy learning and causal model optimization.
Step 2: Model Optimization We execute the empowerment-drivee
Step 3: Policy Learning During downstream task learning, we incorporate the causal effects of different actions as curiosity rewards combined with the task reward, following Eq. 12. We maximize the discounted cumulative reward to learn the policy by the cross entropy method (CEM) (Rubinstein, 1997). Specifically, The causal model used execute dynamic state transitions defined in Eq. 2. The reward model evaluates these transitions and provides feedback in the form of rewards. The CEM handles the planning process by leveraging the predictions from the causal and reward models to optimize the task's 's objectives effectively.
EXPERIMENTS
We aim answer the following questions in experimental evaluation: (i) How does the performance of ECL compare to other causal and dense models across different environments for tasks and dynamics learning, including pixel-based tasks? (ii) Does ECL improve causal discovery by eliminating more imreevant state dimensions interference, thereby enhancing learning efficiency and generalization towards the empowerment gain? (ii) Whether different causal discovery methods in step I and 2, impact policy performance? What are the effects when combine the step and 2? (iv) What are the effects of the components and hyperparameters in ECL?
5.1 SETUP
Environments. We select different environments for basic experimental evaluation. Chemical (Ke al., 2021): The task is to discover the causal relationship (Chain, Collider & Full) of chemical

---

Published as a conference paper at ICLR 2025
Chemical (Chaii)
Chemical (Pull)
Manipulziio (Stack)
Plyvical
Exporelte Reward Moin
Epoaedie Roward Moa
Eplodi Rowad Moen
Epiandie Maen
500
20
40
600
20-
20
200
A80
MLP
COL
10.5
100
Figure 3: The task learning of episodic reward in three environments of ECL-Con (ECL-C) and ECL-Sco (ECL-S).
Chemical (Full)
Physical
Manipulation (Pick)
E Rewwad Mean
40
Episode Reward Mean
Reisode/eeeess Reward Mear
75
500
70-
20-
ECL-Con
400
ECI-Con
Epiisnde
CDL
CDL
CDL
0
200
400
600
800
1000
250
500 Episodes
750
1000
200
400
600
800
1000
Episodes
Episodes
Figure 4: The learning curves of episodic reward in three different environments and the shadow is the standard error.
items which proves the learned dynamics and explains the behavior without spurious correlations. Manipulation (Wang et al., 2022c): The task to prove dynamics and policy for difficult settings with spurious correlations and multi-dimennion action causal influence. Physical (Ke et al., 2021) a dense mode Physical environment. Furthermore, we also include pixel-based environments of Modified Cartpole (Liu et al., 2024), Robedesk (Wang et al., 2022a) and Deep Mind Control (DMC) (Wang et al., 2022a) for evaluation in latent state environments. For the details of the environment setup, please refer to Appendix D.2.
Baselines. We compare ECL with causal and standard MBRL methods. CDL (Wang al., 2022c): infers causal relationships between the variables for dynamics learning with Conditional Inde- pendence Test (CIT) of constraint- based causal discovery. REG (Wang al., 2021): Action-sufficien state representation based on regularization of boore-based causal discovery. GRADER (Ding et al., 2022): generalizing goal-conditioned RI. with CIT by variational causal reasoning. IFactor (Liu et al., 2024): causal framework to model four distinct categories of latent state variables within the RL system for pixel-based environments. GNN (Ke et al., 2021): a graph neural network with dense dependence for each state variable. Monolithic (Wang et al., 2022c): a Multi-Layer Perceptron (MLP) network that takes all state variables and actions for prediction. For ECL, we employ both conditional independence testing (constraint-basee (ECI-Con)) used in (Wang et al., 2022c) and mask learning by sparse regularization (score-based (ECL-Sco)) used in (Huang et al., 2022). We also combine [Factor (Liu et al., 2024) for pixel-based tasks learning detailed in Appendix D.2.2.
Evaluation Metrics. tasks learning, we utilize episodic reward and task success as evaluation criteria for downstream tasks. For causal dynamics learning, we employ five metrics to evaluate the learned causal graph and assess the mean accuracy for dynamics predictions of future states both ID and OOD. For pixel-based tasks, we use average return and visualization results for evaluation!
5.2 RESULTS
5.2.1 TASK LEARNING
We evaluate each method with the following 7 downstream tasks in the chemical (C), physical (P) and the manipulation (M) environments. Match (C): match the object colors with goal colors individually. Push (P): use the heavier object to push the lighter object the goal position. Reach (M): move the end-effector to the goal position. Pick (M): pick the movable object to the goal position. Stack (M): stack the movable object on the top of the unmovable object.
We conduct each experiment using random seeds.

---

Published as a conference paper at ICLR 2025
As shown in Fig. 3, compared to dense dynamics models GNN and MLP, as well as the causal approaches CDL and REG, ECL-Con attains the highest reward across 3 environments. Notably, ECL-Con outperforms other methods in the intricate manipulation tasks. Furthermore, ECL-Sco surpasses REG, elevating model performance and achieving a reward comparable to CDL. The proposed curiosity reward encourages exploration and avoids local optimality during the policy learning process. For full results, please refer to Appendix D.5.
Additionally, Figure depicts the learning curves across three environments. Across these diverse settings, ECL exhibits elevated sample efficiency compared to CDL and higher reward attainment. The introduction of curiosity reward bonus enables efficient exploration of strategies, thus averting causal empowerment
the risk of falling into local optima. Overall, our proposed purrisicc-mtivatee learning framework demonstrates improved stability and learning efficiency. We also evaluate the effect of combining steps and 2, as shown in Appendix D.8. For full experimental results in property analysis and ablation studies, please refer to Appendix D.7 and D.8.
Sample Efficiency Analysis. After validat- ing the effectiveness of ECL in reward learning, we further substantiate the improvements in sample efficiency of ECL for task execution. As depicted in Figure 5, we illustrate task suc- cess in both collider and manipulation reach tasks. The compared experimental results un- derscore the efficiency of our approach, demon- strating enhanced sample efficiency across dif- ferent environments.
Chemical (Callider)
Manipulation (Reach)
Success Rate
1.0
Success Rate
1.0
0.5
Eta--Da CDL
0.5
500 # Episodes
500 Episodes
1000
Figure 5: Success rate in collider and manipulation environments and the shadow is the standard error.
5.2.2 CAUSAL DYNAMICS LEARNING
Causal Graph Learning. To evaluate the efficacy of our proposed method for learning causal relationships, we first conduct experimental analyses across three chemical environments, employing five evaluation metrics. We conduct causal learning based on the causal discovery with Con and Sco respectively. The comparative results using the same causal discovery methods are presented in Table I, with each cell containing the comparative results for that method across different sce- narios. These results demonstrate the superior performance of our approach in causal reasoning, exhibiting both effectiveness and robustness as evinced by the evaluation metrics of F1 score and ROC AUC (Wang al., 2022c). All results exceed 0.90. Notably, approach exhibits exceptional learning capabilities in chemical chain and collider environments. Moreover, it significantly enhances models performance when handling more complex full causal relationships, underscoring its remark- able capability in grasping intricate causal structures. This proposed causal empowerment framework facilitates more precise uncovering of causal relationships by actively using the causal structure.
Visualization. Moreover, we visually compare the inferred causal graph with the ground truth graph in terms of edge accuracy. The results depicted in Figure 6 illustrate the causal graphs of ECL-Sco compared to REG and GRADER in the collider environment. For nodes exhibiting strong causality, ECL-Sco achieves fully accurate learning and substantial accuracy enhancements compared to REG. Concurrently, ECL-Sco elucidates the causality between action and state more effectively. Furthermore, ECL-Sco mitigates interference from irrelevant causal nodes more proficiently than GRADER. The causal graph learned in the complex manipulation environment shown in Figure 15, demonstrates that ECL effectively excludes irrelevant state dimensions to avoid the influence of spuri- ous correlations. These findings substantiate that the proposed method attains superior performance compared to other causal discovery methods in causal learning.
Predicting Future States. Given the current state and sequence of actions, evaluate the accuracy of each method's prediction, for states both ID and OOD. We evaluate each method for one step prediction 5K transitions, for both ID and OOD states. To create OOD states, change object positions in the chemical environment and marker positions in the manipulation environment to unseen values, followed (Wang et al., 2022c).
Figure illustrates the prediction results across four environments. In the ID settings, our proposed methods, based both Sco and Con, achieve performance on par with GNNs and MLPs, while
8

---

Published as a conference paper at ICLR 2025
Table Experimental results on causal graph learning three chemical environments.
Metries
Methods
Chain
Collider
Full
ECL/CDL
1. 1.00=0.000 10100.00/...00 1011.00e 000=0.00
100+1,00/1, .0000000000 .00
1.00=0,00/0. 00000000000999:0.000 0900
Accuracy
ECL/REG
0.99+0. 0000099=0.00
0.99+0.000 99a0.00
0.99#0.01/0099:00.00
ECL/CDL
000+0.00/0.99=00000
1.0000.00/ 0wt.:0.0011.00000 1000=00000
0.97:0.01/0.924..02
Recall
ECL/REG
000+0.0000.94400..0
0.9940.01/00 1100.89=0.09
0.90=0.02/0..9440.01
ECL/CDI.
1. 1.00=0000/.0000000 0000=00000
1. 1.0+0.000/ 100400.0/1.0000:00 100=0000000
0.96=#000// 0.97=0.02
Precision
ECL/REG
099*00.1/0.99=00.00
0.99+0.01 (0140001/0.9990 01/0.99=0.01
0.97=0.0330.92200
ECL/CDL
L.0000.0000.9922000 10000
1. 1.0+0.00011 00/1.00=0.00
0.97=0.01/0.9440..0
F1 Score
ECL/REG
0.9940.00% 009=0.00/009000 0000009400000666
0.93=*00020.85550.
ECL/CDL
1.0=0.00/00 000.0000099900000
1. 1.0+00.00/ 0110000/11100 00/100000.00
0.98=0.01/0.9664
ROC AUC
ECL/REG
0. 99+0.0110.99=0000
399=0.01/0.93:1000
0.95=0.01/0.55000
The causal graph
ECI Seo
REG
GRADER
Figure 6: The causal graph comparison in the chemical collider environment.
significantly elevating performance in the intricate manipulation environment. These findings validate the efficacy of our proposed approach for causal learning. For the OOD settings, method attains comparable performance to the ID setting. These results demonstrate strong generalization and robustness capabilities compared to GNNs and MLPs. Moreover, it outperforms CDL and REG. The comprehensive experimental results substantiate the proficiency of our proposed method in accurately uncovering causal relationships and enhancing generalization abilities. For full results of causal dynamics learning, please refer to Appendix D.3 and D.4.
5.2.3 PIXEL- HISEE-BASED TASK LEARNING
In complex pixel-based robodesk task, where video backgrounds serve as distractors, ECL effectively learns controllable policies for changing background colors to green, as shown in Figure Addition- ally, ECL surpasses IFactor in terms of average return. These results further validate ECL'S 's efficacy in pixel-based tasks and its ability to overcome spurious correlations (video backgrounds). For more results in pixel- based tasks, please refer to Appendix D.6.
RELATED WORK
Causal MBRL MBRL involves training a dynamics model by maximizing the likelihood of collected transitions, known the world model (Moerland et al., 2023; Janner al., 2019; Nguyen et al., 2021; Zhao et al., 2021). Due to the exclusion of irrelevant factors from the environment through state abstraction, the application of causal inference in MBRL effectively improve sample efficiency and generalization (Ke et al.. 2021; Mutti et al., 2023b; Hwang et al., 2023). Wang et al. (2021) propose a constraint-basee causal dynamics learning that explicitly learns causal dependencies by action-sufficient state representations. GRADER (Ding et al., 2022) executes variational inference by regarding the causal graph as latent variable. CDL (Wang al., 2022c) causal dynamics learning method based on CIT. CDL employs conditional mutual information to compute the causal relationships between different dimensions states and actions. For additional related work, please refer to Appendix B.
Empowerment in RL Empowerment is intrinsic motivation to improve the controllability over the environment (Klyubin et al., 2005; Salge et al., 2014). This concept is from the information- theoretic framework, wherein actions and future states viewed channels for information

---

Published as a conference paper at ICLR 2025
In Distribution
Out Of Distribution
Distribution
Of Distribution
Moan MeanAccure Avcurscy
80
Mean Moan Accurac
100
97
94
60
75
80
28
50
(a) Chemical (Chain)
(b) Chemical (Collider)
In Distribution
Out Of Distribution
Distribuzion
Out Of Distribution
78
77
77
76
100
55
Mean Accuracy
logs 13k calloed
16
51
40
20
6004 COL ECLS REG GNN RDP
ECL COL ECLS REG GAN MD
-50
ECC COL ECLS GAN MIP
E COL ECDS REG CNN MLP
(c) Chemical (Full)
(d) Manipulation
Figure 7: Prediction performance (%) on and OOD states ECL-Con (ECL-C) and ECL-Sco (ECL-S). The mean score is marked on the top of each bar.
Rebk
FFactor
800
700
AXA
Avcrage
600
Rellout
Factoo
500
Rellout
400
250
500
750
1000
Enviromment steps (Ic3)
Average reward
Visualization
Figure 8: The compared results with IFactor and visualized trajectories in Robodesk environment.
transmission. In RI,, empowerment is applied to uncover more controllable associations between states and actions or skills (Mohamed & Jimenez Rezende, 2015; Bharadhwaj et al., 2022; Choi et al., 2021; Eysenbach et al., 2018). By quantifying the influence of different behaviors state transitions, empowerment encourages the agent to explore further to enhance controllability over the system (Leibfried et al., 2019; Seitzer al., 2021). Maximizing empowerment max, I can be used as the learning objective, empowering agents to demonstrate intelligent behavior without requiring predefined external goals.
7 CONCLUSION
This paper proposes method-agnostie framework of empowerment through causal structure learning in MBRL to improve controllability and learning efficiency by iterative policy learning and causal structure optimization. We maximize empowerment under causal structure prioritize controllable information and optimize causal dynamics and reward models to guide downstream task learning. Extensive experiments across environments included pixel-based tasks substantiate the remarkable performance of the proposed framework.
Limitation and Future Work ECL implicitly enhances the controllability but does explicitly tease apart different behavioral dimensions. In our future work, we plan to extend this framework in several directions. First, we aim disentangle behaviors and explore entropy relaxation methods to enhance empowerment, particularly for real-world robotics tasks (Collaboration et al., 2023). Second, while the current framework does not account for changing dynamics, we intend to incorporate insights from recent advancements in local causal discovery (Hwang al., 2023) and modeling non non-stationary change factors (Huang et al., 2020) to enhance the causal discovery component. Third, we plan to leverage pre-trained 3D object-centric visual dynamics models (Shi et al., 2024; Wang et al., 2023b; Luo et al., 2024; Team et al., 2024) to scale our approach to roal-world robotics applications. These directions will be pursued in future work.
10

---

Published as a conference paper at ICLR 2025
REPRODUCIBILITY STATEMENT
We provide the source code of ECL in the supplementary material. The implementation details of experimental settings and platforms are shown in Appendix D.
ACKNOWLEDGMENT
This work was supported in part by the National Natural Science Foundation of China under Grant 62276128, Grant 62192783; part by the Jiangsu Science and Technology Major Project BG2024031; in part by the Natural Science Foundation of Jiangsu Province under Grant BK20243051; the Fundamental Research Funds for the Central Universitiex(14380122 in part by the Collaborative Innovation Center of Novel Software Technology and Industrialization.
REFERENCES
Brandon Amos, Samuel Stanton, Denis Yarats, and Andrew Gordon Wilson. On the model-based stochastic value gradient for continuous reinforcement learning. In Learning for Dynamics and Control, pp. 6-20. PMLR, 2021.
Homanga Bharadhwaj, Mohammad Babaeizadeh, Dumitru Erhan, and Sergey Levine. Information prioritization through empowerment in visual model-based rl. arXiv preprint arXiv:2204.08585, 2022.
Yushi Cao, Zhiming Li, Tianpei Yang, Hao Zhang, Yan Zheng, Yi Li, Jianye Hao, and Yang Liu. Galois: boosting deep reinforcement learning via generalizable logic synthesis. Advances in Neural Information Processing Systems, 35: 15930-99994,, 2022.
Daniel Castro, Ian Walker, and Ben Glocker. Causality matters in medical imaging. Nature Communications, 111)):6673, 2020.
Jongwook Choi, Archit Sharma, Honglak Lee, Sergey Levine, and Shixiang Shane Gu. Variational empowerment representation learning for goal-based reinforcement learning. arXiv preprint arXiv:2/06601104, 2021.
Open-X Embodiment Collaboration, A Padalkar, A Pooley, A Jain, A Bewley, A Herzog, A Irpan, A Khazatsky, A Rai, Singh, et al. Open x-embodiment: Robotic learning datasets and rt-x models. arXiv preprint arXhv:2310.0886. 1(2), 2023.
Wenhao Ding, Haohong Lin, Bo Li, and Ding Zhao. Generalizing goal-conditioned reinforcement learning with variational causal reasoning. Advances in Neural Information Processing Systems, 2--26532-26548, 2022.
Wenhao Ding, Laixi Shi, Yuejie Chi, and Ding Zhao. Seeing is not believing: Robust reinforcement learning against spurious correlation. Advances Neural Information Processing Systems, 36, 2024.
Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine Diversity is you need: Learning skills without a reward function. arXiv preprint arXiv: / 1802006070 2018.
Fan Feng and Sara Magliacane. Learning dynamic attribute-factored world models for efficient multi-object reinforcement learning. Advances in Neural Information Processing Systems, 36, 2023.
Carlos Guestrin. Daphne Koller, and Ronald Parr. Multiagent planning with factored mdps. Advances in neural information processing systems, 14, 2001.
Carlos Guestrin, Daphne Koller, Ronald Parr, and Shobha Venkataraman. Efficient solution algorithms for factored mdps. Journal of Artificial Intelligence Research, 19:399-468. 2003.
Danijar Hafner. Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.041 2023.

---

Published as a conference paper at ICLR 2025
Jianye Hao, Tianpei Yang, Hongyao Tang, Chenjia Bai, Jinyi Liu, Zhaopeng Meng, Peng Liu, and Zhen Wang. Exploration in deep reinforcement learning: From single-agent to multiagent domain. IEEE Transactions Neural Networks and Learning Systems, 2023.
Biwei Huang, Kun Zhang, Jiji Zhang, Joseph Ramsey, Ruben Sanchez-Romens Clark Glymour, and Bernhard Schölkopf. Causal discovery from belerogeneous//onstttiionny data. Journal of Machine Learning Research, 21 1-5,, 2020.
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 Hernández-Lobute 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. 2270-9279. PMLR, 2022.
Inwoo Hwang, Yunhyeok Kwak, Suhyung Choi, Byoung-Tak Zhang, and Sanghack Lee. Quantized local independence discovery for fine-grained causal dynamics learning in reinforcement learning. 2023.
Julian Ibarz, Jie Tan, Chelsea Finn, Mrinal Kalakrishnan, Peter Pastor, and Sergey Levine. How to train your robot with deep reinforcement learning: lessons have learned. The International Journal of Robotics Research, 40(4-5))698-711 2021.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. Advances in neural information processing systems, 32, 2019.
Tobias Jung, Daniel Polani, and Peter Stone. Empowerment for continuous agent- envirnenn systems. Adaptive Behavior, 19(1): 2011.
Harini Kannan, Danijar Hafner, Chelsea Finn, and Dumitru Erhan. Robodesk: A multi-task reinforce-
ment learning benchmark, 2021.
Nan Rosemary Ke, Aniket Didolkar, Sarthak Mittal, Anirudh Goyal, Guillaume Lajoie, Stefan Bauer, Danilo Rezende, Yoshua Bengio, Michael Mozer, and Christopher Pal. Systematic evaluation of causal discovery in visual model based reinforcement learning. arXiv preprint 1eXiv:2107.00k6 2021.
Alexander S Klyubin, Daniel Polani, and Chrystopher L. Nehaniv. Empowerment: A universal agent-centric measure of control. In 2005 ieee congress on evolutionary computation, volume 1, pp. 128-135. IEEE, 2005.
Alexander S Klyubin, Daniel Polani, and Chrystopher L Nehaniv. Keep your options open: An information-based driving principle for sensorimotor systems. PloS one, 3(12)=e4018, 2008
Tabitha E Lee, Jialiang Alan Zhao, Amrita S Sawhney, Siddharth Girdhar, and Oliver Kroemer. Causal reasoning in simulation for structure and transfer learning of robot manipulation policies. In 2021 IEEE International Conference on Robotics and Automation (ICRA), pp. 4776-4782. IEEE, 2021.
Felix Leibfried, Sergio Pascual-Dias and Jordi Grau-Moya. A unified bellman optimality principle combining reward maximization and empowerment. Advances in Neural Information Processing Systems, 32, 2019.
Qi Liu, Yuanqi Du, Fan Feng, Qiwei Ye, and Jie Fu. Structural causal model for molecular dynamics simulation. In NeurlPS 2022 AI for Science: Progress and Promises, 2022.
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, 2024.
Chaochao Lu, Yuhuai Wu, José Miguel Hernández-Lobato, and Bernhard Schölkopf. Invariant causal representation learning for out-of-distribution generalization. In International Conference on Learning Representations, 2021.
12

---

Published as a conference paper at ICLR 2025
Hao Luo, Bohan Zhou, and Zongqing Lu. Pre-trained visual dynamics representations for efficient policy learning. In European Conference on Computer Vision, pp. 249-267. Springer, 2024.
Thomas M Moerland, Joost Broekens, Aske Plaat, Catholijn M Jonker, et al. Model-based rein- forcement learning: survey. Foundations and Trends& in Machine Leurning, 16(1):1-118, 2023.
Shakir Mohamed and Danilo Jimenez Rezende. Variational information maximisation for intrinsically motivated reinforcement learning. Advances in neural information processing systems, 28, 2015.
Mirco Mutti, Riccardo De Santi, Marcello Restelli, Alexander Marx, and Giorgia Ramponi. Ex- ploiting causal graph priors with posterior sampling for reinforcement learning. arXiv preprint arXiv:2310.07518. 2023a.
Mirco Mutti, Riccardo De Santi, Emanuele Rossi, Juan Felipe Calderon, Michael Bronstein, and Marcello Restelli. Provably efficient causal model-basee reinforcement learning for systematic generalization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pp-
9251-9259, 2023b. 2023b..
Tung D Nguyen, Rui Shu, Tuan Pham, Hung Bui, and Stefano Ermon. Temporal predictive coding for model-based planning in latent space. In International Conference on Machine Learning, pp. 8130-8139. PMLR, 2021.
Peer Nowack, Jakob Runge, Veronika Eyring, and Joanna D Haigh. Causal networks for climate model evaluation and constrained projections. Nature communications, 11(1):1415, 2020.
Masashi Okada and Tadahiro Taniguchi. Dreaming: Model-based reinforcement learning by la- tent imagination without reconstruction. In 2021 ieee international conference robotics and automation (icra), pp. 4209-4215. IEEE, 2021.
Judea Pearl. Causality. Cambridge university press, 2009.
Silviu Pitis, Elliot Creager, Ajay Mandlekar, and Animesh Garg. Mocoda: Model-based counterfac- tual data augmentation. Advances in Neural Information Processing Systems, 15:18143-18156, 2022.
Jonathan Richens and Tom Everitt. Robust agents learn causal world models. arXiv preprint arXiv:2402.10877 2024.
Reuven Rubinstein. Optimization of computer simulation models with rare events. European Journal of Operational Research, 99(1)89-1122 1997.
Jakob Runge. Sebastian Bathiany, Erik Bollt, Gustau Camps-Vall Valls, Dim Coumou, Ethan Deyle, Clark Glymour, Marlene Kretschmer, Miguel D Mahecha, Jordi Muñoz-Mari, et al. Inferring causation from time series in earth system sciences. Nature communications, 10(1):2553, 2019.
Christoph Salge, Cornelius Glackin, and Daniel Polani. Empowerment-an introduction. Guided
Self-Organization: Inception, pp- 67-114, 2014.
Maximilian Seitzer, Bernhard Schölkopf, and Georg Martius. Causal influence detection for improv- ing efficiency in reinforcement learning. Advances in Neural Information Processing Systems, 34: 12900522991 2021.
Junyao Shi, Jianing Qian, Yecheng Jason Ma, and Dinesh Jayaraman. Composing pre-trained object- centric representations for roboties from" what" and" where" foundation models. arXiv preprint arXiv:2404. 13474, 2024.
Peter Spirtes, Clark Glymour, and Richard Scheines. Causation, prediction, and search. MIT press, 2001.
Shengpu Tang, Maggie Makar, Michael Sjoding, Finale Doshi-Velez, and Jenna Wiens. Leveraging factored action spaces for efficient offline reinforcement learning in healthcare. Advances in neural information processing systems, 35:34272-34288 2022.
13

---

Published as a conference paper at ICLR 2025
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmalcki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv e-prints, pp. arXiv-1801, 2018.
Octo Model Team, Dibya Ghosh, Homer Walke, Karl Pertsch, Kevin Black, Oier Mees, Sudeep Dasari, Joey Hejna, Tobias Kreiman, Charles Xu, et al. Octo: An open-source generalist robot policy. arXiv preprint 2xXi::2405.12213 2024.
Núria Armengol Urpi, Marco Bagatella, Marin Vlastelica, and Georg Martius. Causal action influence aware counterfactual data augmentation. In Forty-first International Conference Machine Learning, 2024.
Hanchen Wang, Tianfan Fu, Yuanqi Du, Wenhao Gao, Kexin Huang, Ziming Liu, Payal Chandak, Shengchao Liu, Peter Van Katwyk, Andreea Deac, al. Scientific discovery in the age of artificial intelligence. Nature, 620(7972):47-60, 2023a.
Jianren Wang, Sudeep Dasari, Mohan Kumar Srirama, Shubham Tulsiani, and Abhinav Gupta. Manipulate by seeing: Creating manipulation controllers from pre-trained representations. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 3859-3800 2023b.
Tongzhou Wang. Robodesk with a diverse set distractors, 2022.
Tongzhou Wang, Simon Du, Antonio Torralba, Phillip Isola, Amy Zhang, and Yuandong Tian. Denoised mdps: Learning world models better than the world itself. In International Conference on Machine Learning, pp. 22591-22612. PMLR, 2022a.
Zhihai Wang, Jie Wang, Qi Zhou, Bin Li, and Houqiang Li. Sample-efficient reinforcement learning via conservative model-based actor-critic. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pp. 8612 8620, 2022b.
Zizhao Wang, Xuesu Xiao, Yuke Zhu, and Peter Stone. Task fusk-independent causal state abstraction. In Proceedings of the 35th International Conference on Neural Information Processing Systems, Robot Learning workshop. 2021.
Zizhao Wang, Xuesu Xiao, Zifan Xu, Yuke Zhu, and Peter Stone. Causal dynamics learning for task-independent state abstraction. arXiv preprint 2rXiv:2206.13452 2022c.
Zizhao Wang, Caroline Wang, Xuesu Xiao, Yuke Zhu, and Peter Stone. Building minimal and reusable causal state abstractions for reinforcement learning. arXiv preprint arXhv:2401.12499 2024.
Shangdong Yang and Yang Gao. An optimal algorithm for the stochastic bandits while knowing the near-optimal mean reward. IEEE Transactions on Neural Networks and Learning Systems, 32(5): 2285-2291, 2020.
Chao Yu, Jiming Liu, Shamim Nemati, and Guosheng Yin. Reinforcement learning in healtheare: survey. ACM Computing Surveys (CSUR), 55(1):1-36, 2021.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. Advances in Neural Information Processing Systems, 33:14129-14144 2020.
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. arXiv preprint 22Xi:20006.1072.. 2020.
Mingde Zhao, Zhen Liu, Sitao Luan, Shuyuan Zhang, Doina Precup, and Yoshua Bengio. A consciousnes--inspree planning agent for model-based reinforcement learning. Advances in neural information processing systems, 34:1569-155 1581, 2021.
14

---

Published as a conference paper at ICLR 2025
SUPPLEMENTARY MATERIAL TOWARDS EMPOWERMENT GAIN THROUGH
CAUSAL STRUCTURE LEARNING IN MODEL-BASED REINFORCEMENT LEARNING
TABLE OF CONTENTS
Discussions and Broader Impact
16
A.1 Discussions
16
A.2 Broader Impact
16
Additional Related Works
16
B.1 Model-Based Reinforcement Learning
16
B.2 Causality in MBRL
16
C Notations, Assumptions and Propositions
17
C.I Notations
17
C.2 Detailed objective functions
17
C Assumptions and Propositions
17
D Details on Experimental Design and Results
18
D.I Experimental Environments
18
D.2 Experimental setup
20
D.3 Results of Causal Dynamics Learning
D.4 Visualization on the Learned Causal Graphs
22
D.5 Downstream Tasks Learning
29
D.6 Pixel-Based Tasks Learning
30
D.7 Property Analysis
31
D.8 Ablation Studies
33
Details on the Proposed Framework
34
F Experimental Platforms and Licenses
34
F.1 Platforms
34
F22 Licenses
34
15

---

Published as a conference paper at ICLR 2025
A DISCUSSIONS AND BROADER IMPACT
A.1 DISCUSSIONS
From the technical perspective, our work explores leveraging causal structure to enhance empower- ment for efficient policy learning, enabling better control of the environment in MBRL. We propose framework that can effectively combine diverse causal discovery methods. This holistic approach not only refines policy learning but also ensures that the causal model remains adaptable and accurate, even when faced with novel or shifting environmental conditions. ECL demonstrates improved learning efficiency and generalization compared to other causal MBRL methods across six differ- ent RI. environments, including pixel-based tasks. Simultaneously. ECL achieves more accurate causal relationship discovery, overcoming spurious correlation present in the environment. While ECL demonstrated strengths in the accurate causal discovery and overcoming spurious correlation, disentangling controllable behavioral dimensions remains a limitation. Our implicit empowerment approach enhances the policy's control over the environment but does not explicitly tease apart different behavioral axes. Explicitly disentangling controllable behavioral dimensions could be an important future work to further improve behavioral control and empowerment. Additionally, our current approach involves substantial data collection and model optimization efforts, which can hinder training efficiency. Moving forward, aim further streamline framework to enable more efficient policy training and causal structure learning. Enhancing computational performance while maintaining accuracy will be a key focus area for future iterations of this work. In the empowerment maximization described by Eq. 10, we currently omit two entropy terms. In our future work, we plan to explore additional entropy relaxation methods to further optimize this causal empowerment learning objective.
A.2 BROADER IMPACT
We believe that our work contributes to developing a more interpretable model via causality for model-based reinforcement learning (MBRL), which could benefit various domains, including robotics (Ibarz et al., 2021; Lee et al., 2021), healthcare (Yu et al., 2021; Tang et al., 2022), and scientific discovery (Runge al., 2019; Nowack al., 2020; Castro al., 2020; Liu al., 2022; Wang et al., 2023a), particularly in contexts involving causal modeling and decision- While RL-based approaches may present potential risks in these domains, we do not identify any risks uniquely introduced by this work that require specific discussion.
agents.
B ADDITIONAL RELATED WORKS
B. MODEL-BASED REINFORCEMENT LEARNING
MBRL involves training a dynamics model by maximizing the likelihood of collected transitions, known as the world model, well learning reward model (Moerland al., 2023; Janner et al., 2019; Hao et al., 2023; Cao et al., 2022; Yang & Gao, 2020). Based on learned models, MBRL can execute downstream task planning (Nguyen et al., 2021: Zhao al., 202.). data augmentation (Pitis et al., 2022; Okada & Taniguchi, 2021; Yu et al., 2020), and Q-value estimation (Wang et al., 2022b; Amos et al., 2021). MBRL can easily leverage prior knowledge of dynamics, making it more effective at enhancing policy stability and generalization. However, when faced with high-dimensions spaces and confounders in complex environments, the dense models learned by MBRL suffer from spurious correlations and poor generalization (Wang al., 2022c; Bharadhwaj et al., 2022). To tackle these issues, causal inference approaches are applied to MBRL for state abstraction, removing unrelated components (Hwang et al., 2023; Ding et al., 2022; Wang al., 2024).
state
B.2 CAUSALITY IN MBRL
Due to the exclusion of irrelevant factors from the environment through causality, the application of causal inference in MBRL can effectively improve sample efficiency and generalization (Ke et al., 2021; Mutti et al., 2023b; Liu et al., 2024; Urpí et al., 2024). Wang (Wang et al., 2021) proposes regularization-batt causal dynamics learning method that explicitly learns causal dependencies by regularizing the number of variables used when predicting each state variable. GRADER (Ding et al.,
16

---

Published as a conference paper at ICLR 2025
2022) execute variational inference by regarding the causal graph latent variable. IFactor (Liu et al., 2024) is general framework to model four distinct categories of latent state variables, capturing various aspects of information. CDL (Wang al., 2022c) is causal dynamics learning method based on conditional independence testing. CDL employs conditional mutual information to compute the causal relationships between different dimensions of states and actions, thereby explicitly removing unrelated components. However, it is challenging to strike a balance between explicit causal discovery and prediction performance, and the learned policy has lower controllability over the system. In this work, we aim to actively leverage learned causal structures to achieve effective exploration of the environment through empowerment, thereby learning controllable policies that generate data to further optimize causal structures.
C NOTATIONS, ASSUMPTIONS AND PROPOSITIONS
C.1 NOTATIONS
Symbol
Description
Details
of
i-th state at time action at time
dr
Tt
reward time
M
Causal masks between actions and states
Trainable parameters in Eq. 5 Trainable parameters in Eq. Dynamics function of the MDPs Reward function the MDPs
M
Causal masks between actions and states Dynamics function
f
Reward function
I
Mutual information
3
Empowerment gain
de
Parameters of dynamics model
Trainable parameters in Eq.
Pr
Parameters of reward model
Trainable parameters in Eq. 6
He
Parameters of empowerment-drrvee policy
Trainable parameters in Eq. 7
Parameters of task policy
Trainable parameters for task policy learning
Table 2: Notations used throughout the paper.
C.2 DETAILED OBJECTIVE FUNCTIONS
To better illustrate the trainable parameters in each objective function, we mark the trainable ones in red as follows.
de log (o. 31c ari dri
Layn = B(
(13)
dy
Le-dyn Bis
Σ log P (s/4 St, Ma-+o 2 (c;0e) + Leausal
(14)
Lrew E E M.D llogP,. (rel0e(st M),
(15)
max (a)s(a)s)
E B( (s) (s|M) - (x))
(16)
C.3 ASSUMPTIONS AND PROPOSITIONS
Assumption (d-separation (Pearl, 2009)) d-separation is graphical criterion used to determine, from given causal graph, a set of variables X is conditionally independent another set Y, given
17

---

Published as a conference paper at ICLR 2025
third set of variables Z. In a directed acyelic graph (DAG) 9, path between nodes nı and n'm is said to be blocked by a set S if there exists a node nk. for k = 2, m 1, that satisfies one of the following two conditions:
(i) ny E S, and the path between nk-1 and ng+1 forms (Tk-1 + n2 + nx+1). (Ax-1 t n2 f-
nk+1), or (nk-1 <- n& (k++)).
(ii) Neither ny nor any of its descendants is in S, and the path between ng-1 and n&+1 forms (ng-1 ni f- n2-1))
In a DAG, we say that two nodes n, and nu are d-separated by a third node n. if every path between nodes na and n& is blocked by nee denoted na Il. n6 Ma-
Assumption (Global Markov Condition (Spirtes et al., 2001; Pearl, 2009)) The state is fully observable and the dynamics is Markovian. The distribution p over a set of variables V (s), ........ St 0,," ...... aq. rr)" satisfies the global Markov condition on the graph if for any parti- tion (S, A, R. in V such that if A d-separates S from R. then p(S, R] = p(SIA) ((R(A)
Assumption 3 (Faithfulness Assumption (Spirtes et al., 2001; Pearl, 2009)) For of variables V (8), ,,, implied by the Markovian Condition.
,.4 0t, ... ,a. rt)"" there are no independencies between variables that are not
Assumption Under the assumptions that the causal graph is Markov and faithful to the observa- tions, the edge 8) + 8841 exists for all state variables s'.
Assumption No simultaneous backward edges in time.
Theorem Based on above assumptions, we define the conditioning set far, 8t \ {ae; st.... .6 f-1
186 3+1 }. If ot MK. 2411(ar, 8t \ 88), then of + 82-1- Similarly, if az t a 1441 \ at set, then at - 5/-1-
Under the assumptions that the causal graph is Markov and faithful to the observa- tions, there exists an edge from of + s4 if and only if of H slaIla, \ ots 8t then of + s
Proposition 1
Proof. We first prove that if there exists edge from at to s44 then at 14 $441 ) at, se). We prove it by contradiction. Suppose that of is independent of s/41 given {a, \ at, 8t). According to the faithfulness assumption, we can infer this independence from the graph structure. If of is independent of 1 given {a, \ at. ),), then there cannot be directed path from of to 8441 in the graph. Hence, there is no edge between af and s'+1- This contradicts our initial statement about the existence of this edge.
Now, we prove the converse: if at 3L s/41/(a; \\ a4 }t}, then there exists an edge from of to s/41-
Again, we use proof by contradiction. Suppose there no edge between of and s441 in the graph. Due to the Markov assumption, the lack of an edge between these variables implies their conditional independence given {a, a/ }. This contradicts our initial statement that oj UL 14 \ of ). Therefore, there must exist an edge from of to 5'41.
Proposition 2
Under the assumptions that the causal graph Markov and faithful to the observa- tions, there exists an edge from si + s and only if s) V4.
St \ (s).
The proof of Proposition follows a similar line of reasoning to that of Proposition 1. Consequently, the two propositions collectively serve as the foundation for deriving Theorem
D DETAILS ON EXPERIMENTAL DESIGN AND RESULTS
D. EXPERIMENTAL ENVIRONMENTS
We select three different types environments for basic experimental evaluation, as shown in Figure 9.
18

---

Published as a conference paper at ICLR 2025
Calbbe
Uhukelyjng connel graph
(a) Chemical
(b) Physical
(c) Manipulation
Figure Three basic experimental environments.
Chemical In chemical environment, we aim to discover the causal relationship (Chain, Collider Full) of chemical items which will prove the learned dynamics and explain the behavior without spurious correlations. Meanwhile, in the downstream tasks, evaluate the proposed methods by episodic reward and success rate. The reward function is defined as follows:
Match: match the object colors with goal colors individually:
match
Σ
[mg = 5')
(17)
where 1 is the indicator function, m. is the current color of the i-object, and g" is the goal color of the i-object.
Manipulation the manipulation environment, we aim to prove the learned dynamics and policy for difficult settings with spurious correlations and multi-dimension action causal influence. The state space consists of the robot end-effector (EEF) location (IR3), gripper (grp) joint angles (LR²), and locations of objects and markers (6 x I The action space includes EEF location displacement (IR3) and the degree which the gripper is opened ([O, I]). each episode, the objects and markers are reset to randomly sampled poses on the table. The task reward functions of Reach, Pick and Stack are followed (Wang et al., 2022c).
Physical In addition to the chemical and manipulation environment, also evaluate method the physical environment. In a 5 5 grid-world, there are objects and each of them has a unique weight. The state space is 10-dimensional, consisting of x, positions (a categorical variable over 5 possible values) of all objects. At each step, the action selects one object, moves it in one of directions lets stay at the same position categorical variable over 25 possible actions). During the movement, only the heavier object can push the lighter object (the object won't 't move if it tries to push an object heavier than itself). Meanwhile, the object cannot move of the grid-world nor can it push other lighter objects out of the grid-world. Moreover, the object cannot push two objects together, even when both of them are lighter than itself (Dense model mode). The task reward function is defined follows:
Push: calculate the average distance between the current node and the target location:
match
dis(o:. tx)
(18)
where dis(-) is the distance between two objects position. 0; is the position of current node and f, is the position of target node.
D.1.1 PIXEL-BASHE ENVIRONMENTS
Importantly, to evaluate the performance our proposed ECL framework in latent state environments, select three distinct categories of pixel-based environments with distractors for assessment, as shown in Figure 10. We employ IFactor (Liu et al., 2024) as our baseline method and used its encoders to process visual inputs. Subsequently, we apply the proposed ECL framework for policy
19

---

Published as a conference paper at ICLR 2025
learning. The parameter settings for these three environments are kept consistent with the default configurations of IFactor.
(a) Cartpole
(b) Robodesk
(c) Cheetah Run
(d) Reacher Easy
(e) Walker Walk
Figure 10: pixel- experimental environments with tasks.
Modified Cartpole We select variant of the original Cartpole environment by incorporating two distractors (Liu et al., 2024), as shown Figure 10(a). The first distractor is an uncontrollable Cartpole located in the upper portion of the image, which is irrelevant to the rewards. The second distractor is a controllable but reward-irrelevant green light positioned below the reward-relevant Cartpole in the lower part of the image.
Robodesk We select variant of Robodesk (Kannan al., 2021; Wang, 2022), which includes realistic noise element with dynamic video background, as shown in Figure 10(b). In this task, the objective for the agent is change the hue of TV screen green using button press, while ignoring the distractions from the video background.
Deep Mind Control We also consider variants of DMC (Wang et al., 2022a; Tassa et al., 2018), where dynamic video background is introduced to the original DMC environment distractor. We select cheetah Run, reacher Easy and walker Walk three specific tasks for evaluation, as shown in Figure 10(c, d, e).
D.2 EXPERIMENTAL SETUP
D.2.1 DYNAMICS LEARNING IMPLEMENTATION DETAILS
We present the architectures of the proposed method across all environments in Table 3. For all activation functions, the Rectified Linear Unit (ReLU) is employed. Additionally, we summarize the hyperparameters for causal mask learning used in all environments for ECL-Con and ECL-Sco in Table Regarding the other parameter settings, we adhered to the parameter configurations established in CDL (Wang et al., 2022c) and ASR (Huang et al., 2022). Moreover, The policy collect p(s) (((,0))))))))
trained with reward function tanh (E log p(= ) This reward function measures
the prediction difference between the dense predictor and the current causal predictor, following the approach described CDL (Wang et al., 2022c).
chemical, manipulation, and physical environments, utilize well-defined feature spaces for states and actions, which are explicitly designed for causal structure learning. For pixel-based environments such DMC, Cartpole, and RoboDesk, ECL operates latent states extracted by visual encoders. These encoders are supported by the identifiability theory proposed in IFactor (Liu et al., 2024), which ensures that these latent states effectively map the states. While establishing identifiability is not the primary focus of our work, we leverage IFactor's 's encoders and include comparisons with IFactor in our experiments. Importantly, even in these settings, we can learn meaningful causal graphs.
D.2.2 TASK LEARNING IMPLEMENTATION DETAILS
We list the downstream task learning architectures of the proposed method across all environments in Table 5. We outline the parameter configurations for the reward predictor, as well as the settings employed for the cross-entropy method that applied. For pixel-hased task learning, leverage the
20

---

Published as a conference paper at ICLR 2025
Table Architecture settings in all environments.
Environments
Architecture
Chemical
Physical
Manipulation
feature dimension
64
128
128
predictive networks
[64,32]
[128,128)
[128.64]
training steps max step of environment
500K
500K
32M
50
100
250
batch size learing rate max sample time
64
le-4
128
prediction step during training
2
Table 4: Hyperparameters for causal mask learning in all environments.
Environments
Method
hyperparameters
Chemical
Physical
Manipulation
CMI threshold
0.02
0.01
0.002
optimization frequency
10
evaluation frequency
10
ECL-Con
evaluation batch size
32
evaluation step prediction reward weight coefficient
1
1.0
0.002
0.02
0.001
ECL-Sco
regularization starts after N steps
100K
100K
750K
four distinct categories of latent state variables by IFactor conduct empowerment maximization for policy learning. Moreover, we follow the same parameter settings in IFactor, and used the same video background all tasks.
Table 5: Hyperparameters for downstream task learning in all environments.
Environments
Method
hyperparameters
Chemical
Physical
Manipulation
training steps optimizer learing rate batch size number of candidates
300K
1.5M
2M
Adam
Reward Predictor
3e-4
CEM
64
128
number of iterations
5
10
number of top candidates action_noise
0.03
D.3 RESULTS OF CAUSAL DYNAMICS LEARNING
We compare the performance of causal dynamics learning with score-based method GRADER (Ding et al., 2022), CDL (Wang et al., 2022c) and constraint-bu method REG (Wang et al.. 2021) across different environments. The experimental results, presented in Table 6, reveal that although
21

---

Published as a conference paper at ICLR 2025
Table 6: Compared results of causal graph learning on three chemical and physical environments.
Metrics
Methods
Chain
Collider
Full
Physical
ECL-Con
1.00=0.00
1.00=0.00
1.00=0.00
1.00=0.00
Accuracy
ECL-Sco
0.9920.00
0.99=0.00
0.99#0.01
1.00=0.00
GRADER
0.9920.00
0.99=0.00
0.99a0.00
ECL-Con
1.00=0.00
1.00=0.00
0.97=0.00
1.0000.00
Recall
ECL-Sco
1.00=0.00
0.99=0.01
0.90=1.02
1.00=0.00
GRADER
0.96#0.03
0.99=0.02
0.9620.02
ECL-Con
1.00=0.00
1.00=0.00
0.96+0.02
1.00=0.00
Precision
ECL-Sco
0.99#0.01
0.99=0.01
0.97+0.03
1.00=0.00
GRADER
0.94+0.04
0.90=0.05
1.0020.00
ECL-Con
1.00=0.00
1.00=0.00
0.97=0.01
1.00=0.00
F1 Score
ECL-Sco
0.99#0.00
0.99=0.00
0.93+0.02
1.0000.00
GRADER
0,95+0.03
0.94+0.03
0.98
ECL-Con
1.00=0.00
1.00=0.00
0.98+0.01
1.00=0.00
ROC AUC
ECL-Sco
0.99=0.01
0.99=0.01
0.95+0.01
1.00=0.00
GRADER
0.94+0.02
0.99#0.01
0.96+0.01
GRADER exhibits superior performance in the chemical full environment, ECL-based methods overall achieve better results than GRADER across three chemical environments. In the accuracy assessment metrics, ECL-Con attains 100% precision, and across the chain and collider environments, all evaluation metrics achieve perfect 100% scores. Furthermore, in the physical environment, our proposed methods attain 100% performance. The result of rigorous evaluation metrics substantiates that incorporating ECL has boosted the dynamics model performance. These experimental results further validate the effectiveness of the proposed ECL approach in both sparse and dense modal environments.
Furthermore, analyze the prediction accuracy performance of the causal dynamics constructed by our proposed method. The multi-step (1-5 steps) prediction experimental results across four environments are illustrated in Figure 11. ECL-Con and CDL exhibit smaller declines in accuracy as the prediction steps increase, benefiting from the causal discovery realized based on conditional mutual information. Compared to REG, ECL-Sco achieves a significant improvement in accuracy under different settings. Concurrently, we find that the outstanding aut-of-distribution experimental results further corroborate the strong generalization capability our proposed method. By actively leveraging the learned causal structure for empowerment-driven exploration, ECL facilitates more accurate causal discovery. Overall, we can demonstrate that the proposed ECL framework realizes efficient and robust causal dynamics learning.
D.4 VISUALIZATION ON THE LEARNED CAUSAL GRAPHS
We conduct a detailed comparative analysis by visualizing the learned causal graphs. each causal graph, these are ds rows and ds + columns, and the element the j-th row and i-th column represents whether the variable و4, depends on the variable s141 if < ds + a if = ds measured by CMI for score-based methods and Bernoulli success probability for Reg. First, the causal graph learning scenario the chemical chain environment shown Figure 12. Compared to CDL and REG, ECL-Con accurately uncovers the causal relationships among crucial elements, such as all different dimensions between states and actions, outperforming the other two methods. Moreover, achieve extensive climination of causality between irrelevant factors. These results demonstrate the accuracy of the proposed method in causal inference within the chemical chain environment.
Furthermore, for the chemical collider environment, the compared causal graphs are depicted in Figure 13. We observe that both CDL and ECL-Con achieved optimal discovery of causal relationships. Moreover, in contrast to the REG method, ECL-Con is not impeded by interference from irrelevant causal factors. For the chemical full environment, the causal graph is illustrated in
22

---

Published as a conference paper at ICLR 2025
Distribution
of Distribution
Distribution
Distriuttin
Mein Aeguracy
70
Moin Accuray
85
ECL Con ICL-Sco
ECL con
R2-
ECL-StD
55
CDL
63
COL
Prediction
Prediction
Prodiction
Prediction stop
(a) Chemical (Chain)
(b) Chemical (Collider)
In Distribution
Out of Distribution
In Distribution
Out of Distribtion
8D-
Moan Accuncy
lag- Wekllhee
n0-
70
60-
82-
60-
20- RB6
RB5
Prodiction step
Prediction slup
Prodiction step
Prodiction stup
(c) Chemical (Full)
(d) Manipulation
Figure 11: Multi-step prediction performance for four basic environments. (Left) prediction in distribution states. (Right) prediction on OOD states.
Figure 14. Compared to CDL, ECL-Con better excludes interference from irrelevant causal factors. In comparison with the REG method, ECL-Con attains superior overall performance in discovering causal relationships. Additionally, ECL-Con reaches optimal learning performance when provided the true causal graph.
Moreover, for the manipulation environment, the experimental results are presented in Figures 15 and 16. From the results in Figure 6, we can discern that ECL-Con achieves around 90% overall fitting degree with the true causal graph and accurately learns the causal association between state and action. Compared to CDL shown Figure 16, ECL-Con learns more causal associations from relevant causal components related to the gripper, movable states, and actions. Conversely, in contrast REG, ECL-Con better excludes interference from irrelevant causal factors, such as unmovable and marker states. In summary, the proposed method achieves more accurate and efficient learning performance in causal dynamics learning. analyzing the enhanced performance of ECL in optimizing causal dynamics and reward models, and how these optimizations manifest in the learning policies for downstream tasks, including complex pixel- bose-based tasks.
the subsequent section, we will delve further into
23

---

Published as a conference paper at ICLR 2025
True causal graph
ECL-Con causal graph
1.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
1.0
1.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
1.0
ong?
0.0
1.0
0.0
0.9
0.0
0.0
0.0
0.0
0.0
0.0
1.0
0.0
1.3
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
1.0
083
0.0
0.0
0.0
0.0
0.0
0.0
0,0
0.0
1.0
d6f
1.0
0.0
0.0
0.0
0.5
dbf
0.0
0.0
0.0
0.0
0.0
1.0
0.0
0.0
0.0
0.0
1.0
ukys
0.0
0.0
0.0
0.0
0.0
1.0
0.0
0.0
0.0
0.0
1.0
0ky
0.0
0.0
0.0
0.0
0.0
0.0
1.0
0.0
0.0
0.0
1.0
0.0
0.0
0.0
0.0
0.5
0.0
1.0
0.0
0.0
0.0
1.0
08
0.0
0.0
0.0
0.0
0.0
0.0
1.0
0.0
0.0
0.0
0.0
0.0
0.0
1.0
0.0
0.0
1.0
d6p
1.0
0811
0.0
1.0
1.0
oby1
saj2
dbj4
Jaj5
0b68
dbj7
ogg8
daj10
action
oby
deg2
ogg
ob,4
0b/6
oag8
Jaj10
action
CDL causal graph
REG causal graph
0.0
0.0
0.0
0.0
1.0
1.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
1.0
abp
0.2
1.0
0.2
0.0
0.0
0.0
0.0
0.2
0.0
0.0
1.0
de3 dap
1.0
1.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
1.3
dbg
0.0
0.2
1.0
0.2
0.0
0.0
0.0
0.0
0.0
0.0
1.0
0.0
0.0
1.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
1.0
04/4
0.0
0.0
1.0
0.3
0.0
0.0
0.0
0.0
0.0
1.9
0.0
0.0
0.0
1.0
0.3
0.0
0.0
0.0
0.0
0.0
1.0
dbfs
0.0
0.0
0.4
1.0
0.0
0.0
0.0
0.0
0.0
1.9
dej5
0.0
0.0
0.0
0.0
1.0
0.0
0.0
0.0
0.0
0.0
1.0
0b
0.2
0.0
0.0
0.0
1.0
0.0
0.0
0.0
0.0
0.0
1.0
0.0
0.0
0.0
0.0
1.0
ol7
1.0
1.0
dsf
0949 abp
obg
1.0
0.0
0.0
0.0
1.0
oN1
onj3
0bj4
onj5
objs
onj7
obja
oejo
02/10 action
otj
onj2
0nj3
0n44
0nj5
obys
onj7
onjs
onjo
ooj10 action
0.0
02
0.4
0.6
0.8
1.0
Figure 12: Causal graph for the chemical chain environment learned by the ECL, CDL and REG.
24

---

Published as a conference paper at ICLR 2025
True causal graph
ECL-Con causal graph
1.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
1.0
1.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
1.0
ong?
0.0
1.0
0.0
0.9
0.0
0.0
0.0
0.0
0.0
0.0
1.0
0.0
1.3
0.0
0.9
0.0
0.0
0.0
0.0
0.0
0.0
1.0
083
0.0
0.0
0.0
0.0
0.0
0.0
0,0
0.0
1.0
ddf
1.0
0.0
0.0
0.0
dbf
0.0
0.0
0.0
0.3
0.0
1.0
0.0
0.0
0.0
0.0
1.0
dky
0.0
0.0
0.0
0.0
0.0
1.0
0.0
0.0
0.0
0.0
1.0
0.0
0.0
0.0
0.0
0.0
0.0
1.0
0.0
0.0
0.0
1.0
0.0
0.0
0.0
0.0
0.0
0.0
1.0
0.0
0.0
0.0
1.0
08
0.0
0.0
0.0
0.0
0.0
0.0
1.0
0.0
0.0
0.0
0.0
0.0
1.0
0.0
0.0
1.0
d6
1.0
0911
1.3
1.0
1.0
1.0
1.3
1.0
1.0
1.0
oby1
Jaj2
dbi4
Jbj5
obj7
o8g8
Jaj10
antion
oby1
deg2
ogg
ob/4
ob.6
Jaj10
action
CDL causal graph
REG causal graph
0.0
0.0
0.0
0.0
0.0
0.0
1.0
1.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
1.0
abp
0.0
1.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
1.0
db3 dbp
0.0
1.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.1
1.3
dbg
0.0
0.0
1.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
1.0
0.0
0.0
1.0
0.0
0.0
0.0
0.4
0.0
0.0
0.0
1.0
04/4
0.0
0.0
1.0
0.0
0.0
0.0
0.0
0.0
0.0
1.9
0.0
0.0
0.0
1.0
0.0
0.0
0.0
0.0
0.2
0.4
1.0
d8fs
0.0
0.0
0.0
1.0
0.0
0.0
0.0
0.0
0.0
1.9
0.0
0.0
0.0
0.0
1.0
0.0
0.0
0.0
0.9
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
1.0
0.0
0.0
0.0
0.2
1.0
al7
1.0
0.0
0.0
dsf
1.0
091 akj
o0y13 obg
0.0
1.0
1.0
1.0
1.0
1.0
1.0
1.9
obj3
obj4
onj5
objs
onj?
obja
objo
o2j10 action
obj1
onj2
0nj3
onj4
onj5
onys
onj7
onj9
onja
onj10 action
0.0
02
0.4
0.6
0.8
1.0
Figure 13: Causal graph for the chemical collider environment learned by the ECL, CDL and REG.
25

---

Published as a conference paper at ICLR 2025
True causal graph
ECL-Con causal graph
1.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
1.0
1.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
1.0
ong?
0.0
1.0
0.0
0.9
0.0
0.0
0.0
0.0
0.0
0.0
1.0
0.0
1.3
0.0
0.2
0.5
0.5
0.0
0.0
0.0
0.2
1.0
083
0.0
0.0
0.0
0.0
0.0
1.9
0.9
1.0
0.5
0.0
1.0
1.0
0.0
days Jbjs
1.0
1.0
dbf
1.0
1.0
1.0
1.3
0.0
1.0
0.0
0.0
0.0
0.0
1.0
1.0
1.0
1.0
0.5
0.0
1.0
0.0
0.0
0.0
0.0
1.0
1.0
1.0
1.0
1.3
1.0
0.0
1.0
0.0
0.0
0.0
1.0
1.0
1.0
0.9
1.0
0.5
0.0
1.0
0.0
0.0
0.2
1.0
08
1.0
1.0
1.0
1.0
1.0
0.0
1.0
0.0
1.0
1.0
1.0
1.0
0.0
1.0
0.0
0.0
1.0
d6p
1.0
091O
1.3
1.0
1.0
1.0
1.0
1.0
1.0
1.0
oby1
dej2
dby4
dbj5
dbj7
o8g8
daj10
antion
oby1
deg2
olg8
ob/4
0b/6
oag8
Jaj10
action
CDL causal graph
REG causal graph
0.0
0.0
0.0
0.0
1.0
1.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
1.0
abp
0.0
1.0
0.0
0.0
0.5
0.0
0.0
0.0
0.0
0.0
1.0
0.5
1.0
0.1
0.8
0.6
0.1
0.4
0.1
0.2
0.0
1.3
0.5
0.0
1.0
0.5
0.5
0.0
0.0
0.0
0.0
0.0
1.0
1.0
0.5
1.0
1.0
0.1
0.4
0.0
0.1
0.2
0.0
1.0
04/4
1.0
0.0
1.0
1.0
1.0
0.0
1.0
0.0
0.0
1.9
1.0
1.0
1.0
1.0
0.3
0.2
1.0
0.1
0.1
0.1
1.0
defs
1.0
1.0
0.5
1.0
0.0
0.0
0.0
0.0
0.0
1.9
1.0
1.0
0.5
0.2
1.0
0.0
0.1
0.0
0.0
0.0
1.0
1.0
0.0
0.5
0.0
0.0
1.0
1.0
1.0
0.5
0.4
1.0
0.5
0.2
0.2
0.4
1.0
0/7
1.0
1.0
1.0
1.0
0.6
1.0
0.0
dsj
abj
1.0
091O
1.0
1.0
0.3
1.0
0.0
1.0
011
obj3
obj4
obj5
objs
obj7
obja
oajo
02/10 action
obj1
onj2
0nj3
ony4
0nj5
obys
on/7
onjs
onjo
ooj10 action
0.0
02
0.4
0.6
0.8
1.0
Figure 14: Causal graph for the chemical full environment learned by the ECL, CDL and REG.
26

---

Published as a conference paper at ICLR 2025
Trme causal graph
ecty
1.0
1.0
1.3
1.0
1.0
1.0
10000
1.0
1.3 1.0 1.0
0.0
eefz
1.0
1.0
1.0
1.0
1.0
1.0
1.0
1.0
1.0
1.0
1.0
1.0
1.0
0.0
erppert
1.0
1.0
1.0
erpper 2
1.0
1.0
1.3
1.0
1.0
100
1.0
1.0
1.0 1.0
0.0
1.0 1.0
1.0 1.0
1.0
1.0 1.0
1.0
1.0
1.0
1.0
0.0
movabley
y
1.0
1.0
1.0
1.0
movable
1.0
1.0
1.3
1.0
1.0
1.0
1.0
1.0
1.0
1.0 1.0
0.0
u
0.0 0.0 0.0
0.0
0.0
0.0
0.0 O.C
0.0
0.0
0.0 0.0 0.0
y
unmovable 0.0 0.0 0.0
0.0
1.0
0.0
0.0
0.0
0.0
random x
0.0 0.0 0.0
1.0
0.0
0.0
randomy y
1.0
random 0.0 0.00
markort 0.0 0.0 0.0
0.0
markorty
markert 0.0 0.0 00000000000000000000000000000000
marker2 x 0.0
marker2 y 00000000000000000000000000000000
0.0
0.0 0.0
marker3 x 0.0 0.0 0.0
1.0
0.0
0.0
marker3y y 0.0 0.0
1.0
0.0
marker32 0.0
0.0
-0.0
odtx
groper
novabl
novabley
novolle
unmoretble
unnovabie
unmorabio
maddom
randon
maraaft
merertty
mat motarr1
morkat22
materry
marterze
marier3X
mararry
adton
ECL-Con causal graph
oofx
1.0
1.0
1.0
1.0
0.1
1.0
0.1 0.0
1.0
0.5
0.0
0.0
0.0
0.5
0.0
0.0
1.0
0.0
0.0
0.0
0.0
0.0
0.0
eefy
1.0
1.0
1.0
1.0
1.0
0.5 0.5
1.0
0.0
1.0
0.5
1.0
0.5
0.0
0.0
1.0
0.0
0.0
1.0
1.0
0.0
1.3
1.0
0.0
prpper1
1.0
1.0
1.9
1.0
1.0
0.
1.0
1.0
0.0
0.0
0.0
0.0
0.0 0.
0.0
0.0
0.0
0.0
0.0
0.0
0.0
gr pper
1.0
1.0
1.9
1.0
1.0
0.0
1.0
0.0
0.0
1.3
0.0 0.0
0.0 0.0
0.5
0.0
0.0
0.0
1.0
0.0
0.0
movabley
1.0
1:0
1.3
1.0
1.0
1.0
D.0 9.0
1.0 1.0
0.0 0.0
0.0
0.0
0.0
movable
1.0
1.0
1.0
1.0
1.0
1.0
1.0 1.0
0.0 0.0
1.0
0.0 0.0
1.0
0.0
0.0
0.0
1.0
1.0
1.0
1.0
unmovablex
1.0
0.0
1.0
0.0
unmovabley
0.5
1.0
0.0
1.0
0.0
1.0 1.0
0.0
0.0 0.0
0.0
0.0
0.5
0.0
0.0
-0.8
unmovable
z
1.0
0.0
0.0 0.0
0.0 0.0
0.0
0.0 0.0
1.0
1.0
0.7
0.6
0.0
1.0
randomx
rindomy
0.0 0.0 0.0 0.0
0.0
0.0
0.0
1.0
0.0
0.0
randomz
0.0 0.0 0.0 0.0
0.0
0.0
0.0
044
0.0
marker1y
0.0 0.0 0.0
0.2
0.0
0.0
0.0
0.0
0.0
markert
0.0 0.0 0.0 0.0 0.0
0.0
0.0
0.0
0.0
0.0
0.0
1.0
0.0
0.0
0.1 0.0
marker2x
0.0
0.0
0.0
marker2y
0.5 0.4 0.0 0.0
0.0
1.0
0.2
0.0 0.2
0.0
-0.2
marker2 z 0.0 0.0 0.0 0.0
0.0
1.0
0.0 0.0
0.0 0.0
marker3x
0.0 0.0
1,0
0.0
0.0
markeraly
0.0 0.3 0.0 0.0
0.0
1.0
0.0
marker32 0.0 0.0 0.0 0.0 0.3 0.0
0.0
0.0
0.0
0.0 0.0 0.0 0.2
0.2
0.0
0.0
0.0
0.0
1.0
0.0
0.0
edxx
1000
novabey
Figure 15: Causal graph for the manipulation environment learned by the true graph and ECL.
27

---

Published as a conference paper at ICLR 2025
CDL causal graph
1.0
1.0
0.0 0.0
0.0
1.0 0.0
ecty
1.0
1.0
1.3
1.0
1.0
0.0
0.0
0.0
1.3
1.0
0.0 0.0
0.0 0.0 0.0
0.0
eefz
1.0
1.0
1.0
1.0 1.0
0.0 0.0
0.5
0.0
0.0
0.0
0.0
0.0
0.0
1.0 0.0
0.0
0.0
1.0
erpper1
1.0
1.0
0.0
erpper
1.0 0.0
1.3
1.0 1.0
0.0
0.0
1.0
0.0
0.0
0.0
0.0
1.0
0.0
0.0
0.0
movable x
1.0 1.0
1.0
1.0 1.0 1.0
1.0
1.0
0.0
0.0
0.0
0.5
0.0
1.0 0.0
0.0
0.0
0.0
0.0
0.0
0.0
movabley y
1.0
1.0
1.0
0.0
movable 2
1.0
1.0 0.0
1.0
1.0
1.0
1.0
0:5
0.0
0.0
1.0
1.0
1.0
1.0
3.1
0.0
0.0
1.0
1.0
u x
0.0 0.0 0.0
0.5
0.0
0.0
0.0
1.0
0.0
0.0
0.0
0.0
0.5
0.0
1.0 0.5
0.0 0.0
0.5
0.0
0.0
0.0
0.0
ynmovabley y
unmovable 0.0 0.0 0.0
0.0
0.0
0.0
1.0
0.5
0.0
0.5
0.0
1.0 0.5
0.0
0.0
0.5
0.0
random x
0.0 0.0 0.0 0.0
0.0
0.0
0.0
0.0
1.0
0.O
0.0
0.0
1.0
0.0
0.0
0.0
0.0
0.0
0.D
0.0
randomy y 0.0
1.0
0.0
0.0
0.0
0.0
random 0.0 0.0 0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
markort 0.0 0.0 0.0
1.0 0.5
0.0 0.0
0.0
0.0
0.0
0.0
markorty
markert 0.0 0.0 0.0
1.0
0.0 0.0
0.0 0.0
1.0
marker2 x
0.0 0.0 0.0
0.8
0.0
0.0
0.0
0.0
0.0
0.0
0.0
marker2 y
0.0
0.0
markor2z 0.0 0.0
1:0
0.5
marker3 x 0.0 0.0 0.0
0.0
0.0
0.0
0.0
0.0
1.0
0.0
0.0
0.0
marker3y y 0.0 0.0
0.0
0.0
1.0
0.0
marker3z 0.0 0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
-0.0
odtx
impps
groper
novabl
novabley
novelll
unmovelber
unmovable
unmorabios
mandom
randomn
markar1
mar morarr12 y
morkot22
materry
marterzz
marier3X
manerry
adton
REG causal graph
oofx
1.0
1.0
1.0
1.0
1.0
0.6 0.5
0.7 0.7
0.3
0.0
0.5 0.5
0.0 0.0 0.0 0.0
0.1
0.0
0.0
0.0
0.0
eefy
1.0
1.0
1.0
1.0
1.0
0.2 0.0
0.3
0.3
0.3
0.0
0.8 0.0
0.0
0.0
0.0
1.0
1.5
1.0
0.8
0.0
prpper1
0.6
0.1
1.3
1.0 1.0
0.1 0.0
1.0
0.0 0.2
0.0
0.0
0.0 0.0
0.0 0.0 0.0
gr pper
0.7
0.1
1.9
1.0
1.0
1
0.0
1.0
0.0 0.0
0.0
0.0
0.2
0.0 0.0
0.0
1.0
movabley y
1.0
1.0
0.7
1.0
1.0
1.0
1.0
0.4
0.6
0.0 0.2 0.3
0.0 1
0.0
0.1
0.0
0.0
movable
1.0
0.9 0.7
1.0
1.0
1.0
0.9
1.0
0.6
0.3 0.0
0.8 0.5
0.0
0.6
0.0
0.0
0.0
0.9
unmovablex
unmovabley
0.0 0.0 0.0
0.0
0.0
0.0
0.0 0.0
0.0
0.0
0.0
0.0
0.0
0.0
-0.8
unmovablez
0.0 0.0 0.1
0.9
0.9
0.0
0.9 0.7
0.9
1.0
0.9
0.0 0.0
0:9
0.8
0.0
randomx
0.0
rindomy
0.0 0.0 0.0
0.0
0.0
0.0
0.0
0.0
0.0
1.0
0.0
0.0
0.0
random z
0.3 0.0 0.5 0.9 0.9
0.0
0.7
0.0
0
0.5
0.0
0.0
0.0
0.0
044
marker1y
0.0 0.0 0.0
0.0 0.0 0.0
0.0
0.0
0.
0.0
0.0
0.0
0.0
0.0
0.0
0.0 0.0
0.0 0.0
markert
0.9 0.8 0.1 1.0 1.00000000 0.8
8
0.1
0
0.9
0.8
0
0.15
1.0
0.5
0.0
0.0
0.7 0.6
0.2 0.0
marker2x
0.0
0.0
marker2y
0.0 0.0 0.0
0.0 0.0
0.0
0.0
0.0
0.0
0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
-0.2
marker2
1.0 0.7 0.1 1.0 1.0
8.8 0.8 0.6
1.0
0.9
0.8
0
C 8
0.9
0.8
0.7
1.0 1.0
0.9 0.8
0.3
marker3x
0.0 0.0 0.0
0.0 0.0
1,0
0.0
0.0
markeraly
0.0 0.0 0.0
0.0 0.0 0.0 0.0 0.0
0.0
0.0
0.0
0.0
0.0
1.0
0.0 0.0
marker32
0.9 1.0 0.1 0.8 1.0 0.9 0.9 0.8 0.9 0.8
0.9
0.9
0.9
0.9
0.2
0.2
0.0
0.7
0.8
0.1
0.7
0.0
1.0
0.0
0.0
extx
wiy
enf
grpper
movala
novaBey moovltt
unmovale
Handomy
mathert
meher1
Figure 16: Causal graph for the manipulation environment learned by CDL and REG.
28

---

Published as a conference paper at ICLR 2025
Chemical (Chain)
Chemical (Pull)
Chemical Chemical(Collider
Episcdite Reward Mean
Episodic Reward Moan
Episodic Reward Mean
48
47
46
BCL-C COL ICL-S 6NN
BCL-CCCL REG MLP
XXL BCL-S GNN MLP
Figure 17: The task learning of episodic reward three environments with ECL-Con (ECL-C), ECL-Sco (ECL-S) and baselines.
(Reach)
Manipulation
Manipubation
Physical
Episodie Reward Mein
200.
Episodie Mean
200
Episadie Roward MaooaddMMos Mean
40
Moa
600
.25
359-
600
12
Episodic
200
030
06
M1
80.0
K70
M0P
60000
101.
100
Figure 18: The task learning of episodic reward in three manipulation and physical environments.
D.5 DOWNSTREAM TASKS LEARNING
As illustrated in Figures 17 and 18, ECL-Con attains the highest reward across three environments when compared to dense models like GNN and MLP, as well as causal approaches such CDL and REG. Notably, ECL-Con outperforms other methods in intricate manipulation tasks. Furthermore, ECL-Sco surpasses REG, enhancing model performance and achieving a reward comparable to CDL. The proposed curiosity reward encourages exploration and avoids local optimality during the policy learning process. Moreover, ECL excels not only in accurately uncovering causal relationships but also enabling efficient learning for downstream tasks.
Sample efficiency analysis. We perform comparative analysis of downstream tasks learning across all environments. As depicted in Figure 19 for experiments in three chemical environments, we can find that ECL-Con and ECL-Sco achieve outstanding performance in all three environments. Furthermore, the policy learning exhibits relative stability, reaching steady state after approximately 400 episodes. Additionally, Figure 20 illustrates the reward learning scenarios in the other four environments. Within the intricate manipulation environment, ECL-Con facilitates more expeditious policy learning. Moreover, in the dense physical environment, ECL-Con and ECL-Sco also exhibit the most expeditious learning efficiency. The experimental results demonstrate that the proposed methods outperform CDL. Moreover, compared to CDL, ECL enhances sample efficiency, further corroborating the effectiveness of the proposed intrinsic-motivated empowerment method.
Chemical (Chain)
Chemical (Collider)
Chemical (Full)
Episode Roward Mean
Episodo Roward Mean
Episode Reward Meen
40
30
25
30
20
ECL-Con
ECL-Con
20
ECL-Con
ECL-Sco
ECL-Sco
ECL-Scc
CDL
200
400
600
800
1000
200
800
1000
200
600
800
1000
Episodes
Episndes
Episodes
Figure 19: The task learning curves of episodic reward in three chemical environments and the shadow is the standard error.
Causal Discovery with FCIT We further conduct causal discovery using the explicit condi- tional independence test, specifically the Fast Conditional Independence Test (FCIT) employed in GRADER (Ding al., 2022), for task learning evaluation. The comparative task learning results
29

---

Published as a conference paper at ICLR 2025
Manipulativn (Reach)
Monppitation (Pick)
Maripuileion (Sacc))
Phyrical
Eppcod (eeerr X000
Eppcode BeNcciL 3O
020
30-
1.0
10-
Epeods
EC
4.0
830
ECL-Can
800
600
1000
Figure 20: The task learning curves of episodic reward in four environments and the shadow the standard error.
Episodic Reward Mean
Chemical (Chain)
Episodic Reward Mean
Chemical (Collider)
Episodic Reward RewardMMaan Mean
Chemical (Fall)
40-
35
50-
48
36
80
30
22
30-
30
20
20
10
BCL4
60C
BCLS
X60
FOAADEK
366
BE-FFORDDE
B0L CDC
ECLS
X50
ELFORADER
Figure 21: The task learning of episodic reward in three chemical environments. ECL-S represents ECL with score-based causal discovery. ECL-C represents ECL with L1-norm regularization of constrant-based causal discovery. ECL-F represents ECL with FCIT (used in GRADER for causal
discovery).
are presented in Figure 21. These findings demonstrate that ECI-FCIT, achieves improved policy learning performance than GRADER, further validating the effectiveness of our proposed learning framework ECL.
D.6 PIXEL-BASED TASKS LEARNING
We evaluate ECL on 5 pixel-input tasks across 3 latent state environments. Figure 22 presents comparative experimental results and visualized trajectories in the modified cartpole task. Our findings reveal that ECL achieves superior sample efficiency compared to IFactor. Furthermore, the visualized results demonstrate ECL's 's effectiveness in controlling the target cartpole, successfully overcoming distractions from both the upper cartpole and the lower green light, which are not controlled in the Faccor policy.
Moreover, we conduct evaluations on three DMC tasks. The visualized results in Figure 23 confirm effective control for all three agents. Moreover, as shown in Figure 24, ECL achieves average return results, corroborating the enhanced controllability provided by our proposed causal empowerment approach. Finally, we evaluate our method against DreamerV3 (Hafner al., 2023), a current state-of state-of-the-ar approach, across three DMC tasks under noiseless settings. As shown in Figure 25, ECL consistently outperforms DreamerV3 V3 in all 3 tasks.
stable
CartPole
250
IFactor
ECL
200
Average
150
Feator Rollout
100
50
BCL Rollot
250
500
750
1000
Environment (1c2)
(a) Average reward
(b) Visualization
Figure 22: The results of average return compared with IFactor and visualized trajectories in Modified Cartpole environment.
30

---

Published as a conference paper at ICLR 2025
Rollut
BCu.
Rellot
(a) Cheetah Run
(b) Reacher Easy
(c) Walker Walk
Figure 23: The results of visualization in three pixel- pixel-based tasks of DMC environment.
Choctah Run
Reacher Easy
Walker Walk
1000
IFactor
FFactor
IFactor
Anarge seturs
800
ECL
Avczage setura
1000
ECL
ECL
600
Average:
500
500
200
1000 Bavironment stops
Ravironment stops
Eavironment stops
Figure 24: The results average return compared with IFactor in three pixel- tasks of DMC environment under video background setting.
Cheetah
Reacher Easy
1000
Dreamerv3
1000
Dreamerv3
1000
DreamerV3
Avcrage retum
800
ECL
Averago retum
800
800
PPP
600
Avcrage
600
400
400
400
200
200
200
0
1000
2000
1000
2000
1000
2000
Environment
Environment steps
Environment steps
Figure 25: The results of average return compared with Dreamer in three pixel-based tasks of DMC environment under noiseless setting.
D.7 PROPERTY ANALYSIS
Training steps analysis. For property analysis, we set different training steps for causal dynamics learning of ECL-Con. As depicted in Figure 26, in the chemical chain environment, observe that the mean prediction accuracy reaches its peak at 300k training steps. A similar trend observed collider environment, where the maximum accuracy is achieved at 150k training steps. Although in the full environment, ECL attains its maximum accuracy at 600k steps, which higher than the 500k steps used for training CDL, notice that at 500k steps, ECL has already achieved performance comparable to CDL. These results substantiate that our proposed causal action empowerment method effectively enhances sample efficiency and dynamics performance.
31

---

Published as a conference paper at ICLR 2025
Chemical (Chain)
Chemical (Collider)
Chemical (Full)
75
Mean securscy
Mcan accuracy
Mean accuracy
90
70
65-
85
60
60
ECL-Con
80
ECL Con
ECL Con
100K
200K
300K
400K
100KK
150K
200K
300K
400K
500K
600K
700K
training steps
#| training steps
# training steps
Figure 26: The mean accuracy of prediction with different training steps in chemical environments.
Hyperparameter analysis. We further analyze the impact of the hyperparameter À introduced in the downstream task reward function with CUR. We compare four different threshold settings, and the experimental results are depicted in Figure 27. From the results, observe that when the parameter set to 1, the policy learning performance is optimal. When the parameter is set to O, the introduced curiosity cannot encourage exploratory behavior in the policy. Nonetheless, it still achieves reward performance comparable to CDL. This finding further corroborates the effectiveness of method for dynamies learning. Conversely, when this parameter causes the policy to explore too broadly, subjecting it to increased risks, and thus more easily leading policy divergence. Through comparative analysis, ultimately set this parameter 1. In our future work, we will further optimize the improvement scheme for the reward function.
excessively high,
Computation cost. To consider the computation cost, we calculate the computation time for two chemical tasks of Chain, and Collider. The experimental results shown in Figure 28 demonstrate that ECL achieves its performance improvements with minimal additional computational burden specifically less than 10% increase compared to CDL and REG. These results demonstrate that ECL's enhanced performance comes without significant computational cost. All experiments were conducted the same computing platform with the same computational resources detailed in Appendix F.
Epiisdii Reward Mean
Chemical (Chain)
Episodic Roward Mean
Chemical (Collider)
Episodic Reward Mean
Chemical (Full)
40
50
35
30
30
25
ECL-Con
30
ECL-Con
ECL-Con
CDL
20
20
# 2
#
Figure 27: The episodic reward with different hyperparameter À in three chemical environments.
Chemical (Chain)
Chemical (Collider)
Computation Time (1)
30
20.
30
30.
29.5
Computation Time
20
20
BCL-C
CDL
BCL-8
REG
ECL-C
CDL
BCL.8
REG
Figure 28: The computation time in two chemical environments.
32

---

Published as a conference paper at ICLR 2025
Chemical (Chain)
Chemical (Collider)
Chemical (Full)
40
40
Episode Roward Mean
Egisode Roward Mean
-35
4.0
Episode Reward
30
25
-0-
25
wi Sim
Wi Sim
20
w/l Sim
w/l
-0
Cau
Cau
wio Sha
Sha
15
Sha
10
10
200
400
600
800
1000
200
800
1000
200
400
600
800
1000
Episodes
Episodes
Epiisdes
Figure 29: Learning curves of ablation studies in three chemical environments and the shadow the standard error. w/ represents with. w/o represents without.
Cheetah Run
Walkeer Walk
Reacher Easy
1000
ECL Emp
ECL
Averago returm
600
ECL w/ Dis
MY Pon
Average retura
1000
ECL w/ Dis
1000
ECL w/ Dis
600
750
Avcrage
400
500
500
200
250
1000
2000
3000
1000
2000
3000
1000
2000
Environment steps
Environment steps
Environment steps
Figure 30: Learning curves of ablation studies in three DMC tasks and the shadow the standard eтror. w/ represents with.
D.8 ABLATION STUDIES
To further validate the effectiveness the various components comprising the proposed ECL method, we designed series of ablation experiments for verification. First, we implement the method without the first-stage model learning, simultaneously conducting causal model and task learning (w/ Sim) to verify the effectiveness of the proposed three-stage optimization framework. Second, we replace reward (w/
the curiosity reward introduced in the task learning with a causality motivation-drive KL (Peav IP&.)) and method without reward
Cau): Teau (sss,, KL (PeavllPo.... shaping (w/o Sha), respectively, to verify the effectiveness of incorporating the curiosity reward.
The results presented in Figure 29 clearly demonstrate the superior performance of the ECL all other comparative approaches. ECL achieves the highest reward scores among the evaluated methods. Moreover, when compared to the method with Sim, ECL not only attains higher cumulative rewards but also exhibits greater stability in its performance during training. Additionally, ECL significantly outperforms the methods with Cau and method without Sha, further highlighting the efficacy of our proposed curiosity- driven exploration strategy in mitigating overfitting issues. By encouraging the agent to explore novel states and gather diverse experiences, the curiosity mechanism effectively prevents the policy from becoming overly constrained.
We explore the difference between simply maximizing empowerment under the causal dynamics model (Eq. 8) versus maximizing the difference between causal and dense model empowerment (Eq. 10), Comparing ECL with empowerment (w/ Emp) against ECL with distance (w/ Dis) across three DMC tasks, our results in Figure 30 show that ECL w/ Dis achieves superior performance, and ECL w/ Emp also demonstrates strong learning capabilities.
Furthermore, conducted comparative experiments between ECL and ECL without curiosity reward. The learning curves for episodic reward and success rate, shown in Figure 31, demonstrate that the curiosity reward plays erucial role in preventing policy overfitting during the learning process. We also carried out experiments with different values of X. The success rate shown in Figure 31 shows the effectiveness of the curiosity reward.
In summary, ECL facilitates effective and controllable policy learning for agents operating in complex environments. The curiosity-driver reward enables the agent acquire comprehensive understand-
33

---

Published as a conference paper at ICLR 2025
ing of the environment while simultaneously optimizing for the desired task objectives, resulting in superior performance and improved sample efficiency.
Episode Reward Mcan
Manipulation (Reach)
Manipulation (Reach)
Manigulation (Reach)
200
w/ curiosity reward
curiosity
Success Rote
1.0
0.8
1.0
0.6
0.5
0.4
w/ curiesity
A=D
weo curicsity reward
250
500
750
1000
0.0
500 Episodes
1000
0.0
Episodes
EEiisnde
Figure 31: Learning curves ECL with and without curiosity reward manipulation reach task, and with different À settings. The shadow is the standard error. w/ represents with.
E DETAILS ON THE PROPOSED FRAMEWORK
Algorithm lists the full pipeline of ECL below.
EXPERIMENTAL PLATFORMS AND LICENSES
F.1 PLATFORMS
All experiments of this approach are implemented on 2 Intel(R) Xeon(R) Gold 6444Y and 4 NVIDIA
RTX A6000 GPUs.
F.2 LICENSES
In our code, we have utilized the following libraries, each covered by its respective license agreements:
PyTorch (BSD 3-Clause "New" or "Revised" License)
Numpy (BSD 3-Clause "New" or "Revised" License)
Tensorflow (Apache License 2.0)
Robosuite (MIT License)
CausalMBRL (MIT License)
OpenAl Gym (MIT License)
RoboDesk (Apache License 2.0)
Deep Mind Control (Apache License 2.0)
34

---

Published as a conference paper at ICLR 2025
Algorithm 1 Empowerment through causal structure learning for model-based RL
Input:
policy network Te, тæ, transition collect policy Toollect. epoch length dynamics model train- ing, causal empowerment and downstream task policy learning Hayas Hemps and Hask evaluation frequency for causal mask learning feval
Step I: Model Learning
for each environment step t do
Collect transitions {(8;; Os, ri, 8{) 1=1
[Perll with Teollect from environment
end for Add transitions to replay buffer Doolect
for epoch 1.... Hayu do
Sample transitions {(ss ai, s Dayal from Doollect
Train dynamics model Pa. with {(8;, ai x (Papll followed Eq. 4
if epoch % feval == 0 then
Sample transitions {(s,, a,, st) ) from D Doollect
Learn causal dynamics model with causal mask using different causal discovery methods followed Eq.
end if
Sample transitions {(si, ai, ri, s()) Prewl from Deollect
Train reward model Per with {(8 a, Tu s()) end for
[Prwll and del- M) followed Eq.
Step 2: Model Optimization
Collect transitions {(s;, a, rı. st) | for epoch 1.... ...., H emmp do
PPempl with policy Te
Maximize (8e4 M) Eso (8t+1)) with transitions sampled from Demp for pol-
icy learning
Add transitions sampled with मe to Demp
if epoch % feval == () then
Optimize causal mask M and reward model with transitions sampled from Demp followed Eq. and Eq.
end if end for
Step 3: Policy Learning
for epoch = 1, ... Hrask do
Collect transitions {(s, ag ri, x() [Pewkl with ma
[=1
Compute predicted rewards Ttask by learned reward predictor
Calculate curiosity reward Teur by Eq. 11
Calculate r + Ttask + Areur
Optimize policy me by the CEM planning
end for return policy n
35


Reviews:
[{"review_id": "LPsw3932fd", "reviewer": "Reviewer_YL1Y", "summary": "The paper presents a framework called Empowerment through Causal Learning (ECL), designed to integrate empowerment with causal reasoning in model-based reinforcement learning. ECL operates by training a causal dynamics model, maximizing empowerment under this structure, and updating the model through data gathered from exploration.", "questions": "1. As shown in Figure 2, what is the collection policy $\\pi_{collect}$, and how do the authors gather the initial dataset (s, a, r, ...) in the buffer?\n\n2. In line 231, the authors state, \"the dynamics encoder learned in Step 1 remains fixed, allowing for a focused optimization of both the causal structure and the empowerment in an alternating manner.\" I am wondering how the authors can fix the optimization of the encoder while still optimizing the causal structure, as shown in Eq. (5)?\n\n3. The proposed framework optimizes iteratively. How is the iteration cycle determined? Will this approach result in high computational costs and longer training times? Could the authors also provide a comparison of training times?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. **Detailed Experimental Validation**: The framework is evaluated extensively across multiple environments, including both state-based and pixel-based tasks, showcasing its adaptability and effectiveness in real-world scenarios.\n2. **Clear Presentation**: The paper is well-organized and clearly presents concepts, making it accessible and allowing readers to follow the progression of ideas and experimental setups with ease.", "weaknesses": "1. **Minor Contribution**: The current framework appears more like a combination of existing approaches rather than a novel advancement. Causal structure learning in model-based RL has been extensively studied in prior work, such as [1-2], as has empowerment in RL [3-4]. This may limit the perceived originality of the contribution, as it builds on established methodologies without significantly advancing them.\n\n2. **High Computational Cost**: The framework’s iterative process of empowerment maximization and causal model updating may result in substantial computational requirements, potentially limiting its scalability in large or dynamic environments.\n\n[1] Huang B, Feng F, Lu C, et al. Adarl: What, where, and how to adapt in transfer reinforcement learning[J]. arXiv preprint arXiv:2107.02729, 2021.\n\n[2] Huang B, Lu C, Leqi L, et al. Action-sufficient state representation learning for control with structural constraints[C]//International Conference on Machine Learning. PMLR, 2022: 9260-9279.\n\n[3] Zhang J, Wang J, Hu H, et al. Metacure: Meta reinforcement learning with empowerment-driven exploration[C]//International Conference on Machine Learning. PMLR, 2021: 12600-12610.\n\n[4] de Abril I M, Kanai R. A unified strategy for implementing curiosity and empowerment driven reinforcement learning[J]. arXiv preprint arXiv:1806.06505, 2018.", "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": "f7z1ag3I8c", "reviewer": "Reviewer_8bod", "summary": "The authors present ECL, an agent that integrates empowerment with causal reasoning to better learn to control an environment, to then perform better in downstream tasks. First, ECL learns a causal dynamics model (dense dynamics model + causal mask) and reward model to maximize the likelihood of observed trajectories, with a regularization term on the causal mask encouraging the use of as few features as possible. Second, ECL alternates between updating a policy that leverages the causal mask to maximize empowerment, and using the data generated by running the policy to improve the causal mask and reward model. Finally, the learned models can be used to learn policies for downstream tasks, mitigating overfitting of the causal model with an intrinsic reward for observing transitions where the dense model fits better than the causal model. The authors demonstrate that ECL performs well in terms of both reward and sample efficiency compared to existing methods across environments, and accurately learns the true causal structure of the environments.", "questions": "In the dynamics function f in equation (2), it’s not clear how exactly the adjacency matrices indicate the influence of current states and actions on the next state, what the output of each dot product represents and how they can be combined to get the ith dimension of the next state?\n\nWhat is the relationship between the causal mask M and the adjacency matrices?\n\nHow can the mask be updated without also having to update the dynamics encoder? If the mask was very bad at Step 1, wouldn’t the dynamics encoder also be very suboptimal and not appropriate to continue using as the mask is improved in step 2?\n\nWhy does it not work to simply maximize the empowerment given the causal dynamics model, rather than the difference between that and the empowerment under the dense model?\n\nIt is not explained in section 3 where the reward from step 1 and 2 comes from. Section 4 describes a reward function formulated to select transitions covering more state-action pairs. How sensitive is the method to the design of this reward function? Reward functions can be notoriously hard to design; a lot of the difficulty of the problem might be obfuscated in this part. \nIn Step 3 of Algorithm 1 it says the learned reward predictor predicts $r_{task}$- how can it predict that if it was learnt during step 1 and 2 in the absence of any downstream tasks? And why are the $r_i$ in the transitions collected in line 2 ignored? Step 3 of section 4 implies that the causal model is only used to generate curiosity intrinsic rewards (which does not rely on the learned reward predictor at all) so this is inconsistent with Algorithm 1.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The approach appears to be novel. The integration of causal discovery, which by itself can be quite passive, with Empowerment, to emphasize controllability, is very interesting.\n\nThe authors tested on a wide spread of environments with different types of dynamics, showing impressive causal discovery and task performance.", "weaknesses": "The presentation of the algorithm is vague and unclear, with many technical details skimmed over. For instance:\n\n - The causally-factored MDP is difficult to understand and the explanations are very brief; it would be helpful to show the causal masks and adjacency matrices for the example environment in figure 1, and provide more explanation/intuition for what each term in the equations represents. See the questions section.\n\n - How the causal models are used for downstream tasks is not specified anywhere except a single line of the Appendix “the CEM planning” \n\nThe dynamics encoder predicts an independent probability for each feature of the next state given the current state and action- how realistic is this? On a similar vein, this method seems dependent on a well-defined feature space for the states and actions.\n\nThe choice of “standard MBRL” baselines to complement the causal baselines do not seem very standard- why not compare against state of the art MBRL algorithms such as [1]? \n\nMore minor writing issues:\n\n - The use of Dense dynamics model/dynamics encoder interchangeably to describe the same thing is confusing\n\n - Section 4 is quite repetitive of section 3\n\n[1] Hafner, Danijar, et al. \"Mastering diverse domains through world models.\" arXiv preprint arXiv:2301.04104 (2023).", "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": "sVrek1LQnx", "reviewer": "Reviewer_8gmf", "summary": "This paper introduces the empowerment through Causal Learning (ECL) framework to enhance controllability and learning efficiency in model-based RL. The framework combines empowerment, a measure of an agent’s ability to control its environment through causal structure learning. ECL enables agents to understand causal relationships within the environment, improving decision-making and policy learning. ECL was evaluated with different causal discovery methods across three environments, showing improved sample efficiency, accurate causal inference, and higher episodic rewards than other causal MBRL approaches​.", "questions": "1. Curiosity-driven exploration in RL is often sensitive and can be challenging to implement effectively. Are there different settings for setting curiosity in different experimental environments? Whether ECL is sensitive to curiosity rewards?\n2. Given the complexity of the ECL structure, the ablation studies should not be omitted, i.e., ablations on reward design, basic model and other related parts should be added.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "- Overall, I think this is a good paper. The idea is interesting. The authors provide enough details and explanations in the technique section. The experiments are thorough and enough to demonstrate the advantages of ECL.\n- ECL combines causal learning with empowerment-driven exploration, which is novel. By leveraging causal structures, the model enables agents to control their environments more effectively and make informed decisions, adding depth to RL's traditional empowerment approach. Through empowerment-driven exploration, ECL enhances the agent’s ability to efficiently sample relevant experiences, reducing the data requirements compared to conventional MBRL methods. This leads to faster learning and less dependence on extensive data.\n- ECL has been tested across different environments, such as chemical, manipulation, and physical tasks, showing strong performance in sample efficiency, causal discovery accuracy, and episodic rewards.\n- By incorporating a curiosity reward during policy learning, ECL encourages exploration while reducing the risk of overfitting specific causal structures. This helps the agent generalize better to new or out-of-distribution tasks.", "weaknesses": "Refer to the questions.", "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": "omJ0Ax0qTn", "reviewer": "Reviewer_DR5q", "summary": "This paper presents a novel framework named Empowerment through Causal Learning (ECL), which integrates causal structure learning with empowerment-driven exploration in Model-Based Reinforcement Learning (MBRL) by (1) causal model learning, (2) empowerment maximization, and dynamic model updating and (3) policy learning. The proposed method is agnostic to causal discovery methods and outperforms existing causal MBRL methods across several environments.", "questions": "1. In the motivating example, why do we merely focus on improving controllablility (row 2 and 3) and do not care about whether the agent find the true target? (L79-l87) And further, how to detect the \"controllable\" trajectories? Could you provide more specific details on how controllable trajectories are identified and measured?\n2. While maximizing mutual information I(s;a), would it be helpful to also take into account the causal structure, i.e. maximizing the state and action dimensions that are dependent in the dynamics graph? Potential benefits and challenges of incorporating the related s,a dimension into the mutual information objective could be made clearer.\n3. The result in Fig. 24 in DMC environment CHeetah and Walker seem to have not converged yet. Could the author compare ECL and IFactor when they both converge? I agree that the learning curve of ECL is already going up faster and steadier than that of IFactor, and I think it could also be helpful to see the convergence point of the policies.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The writing is clear and the paper is in general easy to follow.\n2. Actively applying causal discovery to learn environment dynamics (updating through newly collected data) is a novel approach.\n3. The experiments cover a diverse set of environments, including state-based and pixel-based tasks. Both analysis on the learned causal dynamics and on the average return demonstrate substantial improvements compared to other methods and provide strong evidence of ECL across various tasks.", "weaknesses": "1. The framework aims at learning a consistent causal structure, thus cannot deal with scenarios when the causal dynamics change (change of number of objects, etc.) that might correspond to different behavior components. Maybe consider discuss potential extensions or modifications to their framework that could handle changing causal dynamics. \n2. Other issues please refer to questions.", "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": "ZRCaunTUHB", "reviewer": "Reviewer_4AH7", "summary": "The authors propose ECL (Empowerment through Causal Learning) that has two main components: (1) a causal dynamics model of the environment that is learned from data and (2) a mechanism to \"empower the causal structure for exploration, simultaneously using data gathered through exploration to update the causal dynamics model\". The objective of the causal structure for exploration is to obtain a dynamics model that \"could be more controllable than dynamics models without the causal structure\". On top of this, an intrinsic curiosity reward is developed \"to mitigate overfitting during downstream task learning\".\n\n\n--\nGiven some improvements made to the paper, I increase my score to 5.\n--\nGiven some additional improvements made to the paper, I increase my score to 6 (see comments below)", "questions": "- How can a curiosity-based reward prevent \"overfitting during task learning\" (lines 270-272)?\n- Figure 5: is it the result of only one seed? Why is there only one seed used?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The objective of learning a model with some causal structure that can be used for instance in the context of exploration is an important research topic.", "weaknesses": "Main weaknesses:\n- Notations are not clearly defined. For instance, the causal mask $M$ is introduced (i) without clear mathematical definition and (ii) it's not $M$ that is used later but $M^{s \\rightarow s'}$ in Equation 5 (see lines 204 to 215). This makes the methodology unclear.\n- The vocabulary does not relate to clearly defined concepts, e.g. \"causal understanding\", \"causal reasoning\" in lines 153-155. \n\nAdditional comments:\n- Key elements are described in the appendix instead of the main text, e.g. end of page 4: the causal loss that represents \"the objective term associated with learning the causal structure\" are given in Appendix D.2.", "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)