Causal Information Prioritization for Efficient Reinforcement Learning
Authors: Hongye Cao
Fan Feng
Tianpei Yang
Jing Huo
Yang Gao
Venue: ICLR 2025 Poster (2025)
Keywords: causality
reinforcement learning
empowerment
sample efficiency
Primary area: reinforcement learning
TL;DR: To address limitations of blind exploration and poor sample efficiency, we introduce CIP, a novel efficient RL framework that prioritizes causal information through the lens of reward feedback.

Abstract: Current Reinforcement Learning (RL) methods often suffer from sample-inefficiency, resulting from blind exploration strategies that neglect causal relationships among states, actions, and rewards. Although recent causal approaches aim to address this problem, they lack grounded modeling of reward-guided causal understanding of states and actions for goal-orientation, thus impairing learning efficiency. To tackle this issue, we propose a novel method named Causal Information Prioritization (CIP) that improves sample efficiency by leveraging factored MDPs to infer causal relationships between different dimensions of states and actions with respect to rewards, enabling the prioritization of causal information. Specifically, CIP identifies and leverages causal relationships between states and rewards to execute counterfactual data augmentation to prioritize high-impact state features under the causal understanding of the environments. Moreover, CIP integrates a causality-aware empowerment learning objective, which significantly enhances the agent's execution of reward-guided actions for more efficient exploration in complex environments. 
To fully assess the effectiveness of CIP, we conduct extensive experiments across $39$ tasks in $5$ diverse continuous control environments, encompassing both locomotion and manipulation skills learning with pixel-based and sparse reward settings. Experimental results demonstrate that CIP consistently outperforms existing RL methods across a wide range of scenarios.

Published as a conference paper at ICLR 2025
CAUSAL INFORMATION PRIORITIZATION FOR EFFICIENT REINFORCEMENT LEARNING
Hongye Cao Fan Feng Tianpei Yang Jing Huo' Yang Gao
1 NNational Key Laboratory for Novel Software Technology, Nanjing University University of California, San Diego MBZUAI "School of Intelligence Science and Technology, Nanjing University
ABSTRACT
Reinforcement Learning (RL) methods often suffer from sample inefficiency, one of the underlying reasons is that blind exploration strategies may neglect causal lationships among states, actions, and rewards. Although recent causal approaches aim to address this problem, they lack grounded modeling of reward-guided causal understanding of states and actions for goal orientation, thus impairing learning effi- ciency. To tackle this issue, propose novel method named Causal Information Prioritization (CIP) ) that improves sample efficiency by leveraging factored MDPs to infer causal relationships between different dimensions of states and actions with respect to rewards, enabling the prioritization of causal information. Specifically, CIP identifies and leverages causal relationships between states and rewards execute counterfactual data augmentation to prioritize high-impact state features under the causal understanding of the environments. Moreover, CIP integrates a causality- aware empowerment learning objective, which significantly enhances the agent's execution of swardd-gided actions for more efficient exploration complex environments. To fully assess the effectiveness of CIP, we conduct ex- tensive experiments across 39 tasks in diverse continuous control environments, encompassing both locomotion and manipulation skills learning with based and sparse reward settings. Experimental results demonstrate that CIP consistently outperforms existing RL methods across a wide range of scenarios. The project page is https: sitess..s.///////////eee..///tte google. com/ www.wcww/////.//...oo..........t wi////iww//////////eeeee hrtpp//.
INTRODUCTION
Reinforcement Learning (RL) has emerged as a powerful paradigm for training intelligent decision- making agents learn optimal behaviors by interacting with their environments, receiving reward feedback, and iteratively optimizing their decision-maker making policies (Haarnoja et al., 2018; Sutton, 2018; Silver et al., 2017: Cao al., 2023; Hao al., 2023; Cao et al., 2022). Despite its notable successes, most RL approaches are faced with the sample-inefficiency problem, which means they typically necessitate an enormous number of interactions with the environment to learn policies, which can be impractical costly in real-world scenarios (Savva et al., 2019; Kroemer et al., 2021). Inefficient policy learning often results from blind exploration strategies that neglect causal relation- ships, leading spurious correlations and suboptimal solutions with high exploration costs (Zeng et al., 2023; Liu et al., 2024).
Causal reasoning captures essential information by analyzing causal relationships between differ- ent factors, filtering out irrelevant information, and avoiding interference from spurious correla- tions (Wang al., 2022; Pitis al., 2022; Zhang et al., 2024; Huang al., 2022b). These approaches build internal causal structural models, enabling agents to strategically focus their exploration on the most pertinent aspects of the environment. They significantly reduce the number of samples required and demonstrate remarkable performance in single-task learning, generalization, and counterfactual reasoning (Richens & Everitt, 2024; Urpí et al., 2024; Deng et al., 2023; Huang et al., 2022a; Feng Magliacane, 2023). However, most of these works overlook the reward-relevant causal relationships among different factors, or only partially consider the causal connections between states, actions, and rewards (Liu et al., 2024; Ji et al., 2024a), thus hindering efficient exploration.
"Corresponding Tianpei Yang (t (tianpei. yang@nju.nn.... edu. en).

---

Published as a conference paper at ICLR 2025
(a) Unable to reach the ball and moving chaotically
(b)
$
S
S+1
2
2
1/4
s42
Can reach the ball but cannot goal under control
3
St
1+1
++2
a
4+
a
Can reach the ball and goal under control
a
a
4
r
T T+1
time
Figure 1: (a). An example of robot manipulation soccer task with three trajectories, where the objective is to move the ball into the goal. (b). Underlying causal structure of this example in a factored MDP. Different nodes represent different dimensional states and actions.
In this work, we aim to identify and exploit task-specific causal relationships between states, actions, and rewards, enabling agents to discern relevant states and select actions that maximize rewards, ultimately facilitating precise and goal-oriented behaviors. Here we provide a motivating example in Figure 1, showing three trajectories for executing a manipulation soccer task, along with the underlying causal structure in factored Markov Decision Process (MDP) (Kearns & Koller, 1999). In the first trajectory (row 1), when the agent fails to distinguish states with more intricate causal relationships of the task, the robotic arm exhibits chaotic moving and receives no rewards. The second trajectory (row 2) shows that even without chaotic movements, uncontrollable actions unrelated to the reward lead to an inability guide the ball towards the goal. Only by filtering out irrelevant state features and executing more controllable actions can we guarantee that the ball is kicked into the goal like row Quantifying the contribution of different factors to the reward can effectively help analyzing important causal relationships.
To address the limitation of sample-inefficieng and leverage the potential of causal reasoning, we propose a novel approach named Causal Information Prioritization (CIP) for efficient RI, improving learning efficiency from the perspective of rewards. Building upon the factored MDPs, CIP infers causal relationships between states, actions, and rewards across different dimensions, respectively. CIP employs counterfactual data augmentation based on the causality between states and rewards to generate transitions, prioritizing critical state transitions. Furthermore, CIP leverages the causality between actions and rewards to reweight actions, while utilizing empowerment to maximize mutual information between causally informed actions and future states, thereby enabling better control.
Specifically, CIP leverages collected data to construct a reward- reward-guided structural model that explicitly reasons about state-reward causal influences, enabling the swapping of causally independent state features across observed trajectories to generate synthetic transitions without additional environment interactions. By swapping independent state features across different transitions (i.e., irrelevant state dimensions of chaotic movements in the soccer task), CIP accentuates causally dependent state information (i.e., relevant states to reach the ball), facilitating focused learning of critical state transitions. Subsequently. CIP constructs another structural model that incorporates actions and rewards reweight actions of dimensions. To enhance the exploration efficiency, CIP integrates a causality-aware empowerment term, quantifying the agent's capacity to exert controlled influence over its environment through the mutual information. This empowerment term, combined with causally weighted actions, is integrated into the learning objective, prioritizing actions with high causal influence. The synthesis of causal reasoning and action empowerment enables agents to focus on behaviors that causally relevant to the task, leading more efficient and effective policy learning. The main contributions of this work can be summarized as follows.
To address limitations of blind exploration and sample-inefficieney efficient RL framework that prioritizes causal information through the lens of reward. CIP bridges the gap between causal reasoning and empowerment to facilitate efficient exploration.
introduce CIP, novel
2

---

Published as a conference paper at ICLR 2025
CIP constructs reward-guided structural models to uncover causal relationships between states, actions, and rewards across dimensions. By leveraging state-reward causality, it performs coun- terfactual data augmentation. eliminating the need for additional environment interactions, and enabling learning on critical state transitions. Exploiting action-reward causality, it reweights ac- tions enhance exploration efficiency through empowerment By prioritizing causal information, CIP enables agents to focus on behaviors that have causally significant effects on their tasks. To validate the effectiveness of CIP, we conduct extensive experiments in 39 tasks across 5 diverse continuous control environments, including manipulation and locomotion. These comprehen- sive evaluations demonstrate the effectiveness of CIP in pixel-based and sparse reward settings, underscoring its versatility and reliability.
2
RELATED WORK
2.1
CAUSAL RL
The application of causal reasoning in RL has shown significant potential to improve sample effi- ciency and generalization by effectively excluding irrelevant environmental factors through causal analysis (Huang et al., 2022a; Feng & Magliacane, 2023; Mutti et al., 2023; Sun et al., 2024; Sun & Wang). Wang (Wang al., 2021) introduces novel regularization-basse method for causal dynamics learning, which explicitly identifies causal dependencies by regulating the number of variables used to predict each state variable. CDL (Wang et al., 2022) takes an innovative approach by using condi- tional mutual information to compute causal relationships between different dimensions of states and actions. IFactor (Liu et al., 2024) a general framework to model four distinct categories of latent state variables, capturing various aspects of information. ACE (Ji et al., 2024a), an off-policy actor- critic method, integrates causality-aware entropy regularization. Table 2 provides a categorization of various causal RL methods, highlighting their focus different reward-guider causal relationships. Existing approaches do not fully account for the causal relationships between both states and actions with rewards. Our goal is to explore these causal relationships from reward-guided perspective to enhance sample efficiency across a broader range of tasks.
2.2 EMPOWERMENT IN RL
Empowerment, an information theory-based concept of intrinsic motivation, has emerged as a power- ful paradigm for enhancing an agent's environmental controllability (Mohamed & Jimenez Rezende, 2015; Klyubin et al., 2005; Cao et al., 2024). This framework conceptualizes actions and future states as information transmission channels, offering a novel perspective on ugent-environment interactions. In RL, empowerment has been applied to uncover more controllable associations between states and actions, as well as to develop robust skill (Salge et al., 2014; Bharadhwaj et al., 2022; Choi et al., 2021; Eysenbach et al., 2018; Leibfried et al., 2019; Seitzer et al., 2021). Empowerment, expressed as maximizing mutual information max, I, serves as a learning objective in various RL frameworks, providing intrinsic motivation for exploration and potentially yielding more efficient and generalizable policies. Our approach extends empowerment in RL by examining the influence of state, actions, and rewards through causal lens, integrating causal understanding with empowerment to enhance exploration strategy and learning efficiency.
2.3
OBJECT-CENTRE RL AND OBJECT- ORIENTED OBJECT-ORIENTED RL
Recent advances in object-centric representation learning focus on acquiring and leveraging structured, object-wise representations from high-dimensional observations. Foundational works include Slot Attention (Locatello et al., 2020) and AIR (Eslami et al., 2016; Kosiorek et al., 2018), establishing basis for this field. Subsequent follow-ups have worked these concepts by employing state-of-the- art architectures, including DINO- based approaches Zadaianchuk et al. (2023), transformer based models (Wu et al., 2022). diffusion models (Jiang al., 2023), and state-space models (Jiang et al., 2024). Notably, learning object-centric representations can enable compositional generalization across various domains, such as video and scene generation (Wu et al., 2023; 2024). Moreover, several theoretical studies have explored the mechanisms underlying compositional generalization and the causal identifiability (Kori et al., 2024; Brady et al., 2023; Lachapelle et al., 2024).
Object-centric representations have been effectively employed in world models to capture multi- object dynamics, as demonstrated by works (Jiang et al., 2019; Lin et al., 2020; Kossen et al., 2019). Building these object-centric world models, various studies use them in RL by better modeling complex object-centric structures in partially observable MDPs (Kossen et al., 2019; Mambelli et al., 2022; Feng & Magliacane, 2023; Choi et al., 2024), identifying critical objects (Zadaianchuk et al.,
3

---

Published as a conference paper at ICLR 2025
2022; Park et al., 2021), and learning object-centric policies (Zadaianchuk et al., 2021; Yuan et al., 2022) and applications in robotic manipulation tasks (Li et al., 2020; Mitash et al., 2024; Haramati et al., 2023; Li et al., 2024), well in learning intrinsic curiosity- policies based on objects and their interactions (Watters et al., 2019; Wang et al., 2024c;b).
Another research direction explores object-oriented MDPs, with the homomorphic object-oriented world model being a notable example that leverages MDP homomorphism to model object dynamics and enable efficient planning through symmetric equivalence in MDPs (Diuk et al., 2008; Scholz et al., 2014; Wandzel et al., 2019; Van der Pol et al., 2020; Rezaei-Shoshee et al., 2022; Zhao et al., 2022), provides a powerful foundation for learning object-oriented MDPs and facilitates efficient planning (Wolfe & Barto, 2006). Our work, which focuses on uncovering general causal relationships among components in MDPs and empowerment optimization, orthogonal to object-centric RI.. However, object-centric RL could provide useful abstract object-based variables that could be useful for causal structure learning complex environments
3 PRELIMINARIES
3.1 MARKOV DECISION PROCESS
In RL, the agent-enviromment interaction is formalized as an MDP. The standard MDP is defined by the tuple M A, (4o r; 7), where S denotes the state space, represents the action space, ((s' ss a) is the transition dynamics, r(s, a) is the reward function, and is the distribution of the initial state so. The discount factor y [0, also included. The objective of RL is to
learn a policy π. : S x A 10. 1] that maximizes the expected discounted cumulative reward (T) = Eso PaInT
EE= 7'r(s), at)].
3.2 STRUCTURAL CAUSAL MODEL
Structural Causal Model (SCM) (Pearl, 2009) is defined by a distribution over random variables, defined as V (s.. s
...... ai Tt, 8:41, s4413 and Directed Acyclic Graph (DAG) G = (V, E) with a conditional distribution P P(u;(PA(e:)) for node os E V. Then the distribution can be specified as:
M
1p) II ((A)))) i=1
where (A( is the set of parents of the node o in the graph G.
Causal Structures in MDP We use a factored MDP (Kearns & Koller, 1999; Guestrin et al., 2003; 2001) to model the MDP and the underlying causal structures between states, actions, and rewards. In the factored MDP, nodes represent system variables (rewards and different dimensions of the states and actions), while the edges denote their relationships within the MDP. We employ causal discovery methods to learn the structures of g.
We can identify the graph structure in G, which can be represented the adjacency matrix M. To integrate such relationships in MDP, we explicitly encode the causal mask over variables into the reward function. Hence, the reward function in MDP with the causal structure is defined as follows:
R(M* St, St. - ag. Cr.)
(2)
where @ denotes the element-wise product. M Ris|x1 and M E R|a|x1 are the adjacency matrices indicating the influence of current states and actions on the reward, respectively, and 6,, represents i.i.d. Gaussian noise. Under the Markov condition and faithfulness assumption (Pearl, 2009; Spirtes et al., 2001), the structural vectors propositions can be found in Appendix B. In this work, our objective is to discover and leverage these two causal matrices to prioritize causal information for efficient RL.
identifiable. The detailed assumptions and
3.3 ERPENT IN RL
Empowerment quantifies an agent's capacity influence its environment and perceive the conse- quences of its actions (Klyubin et al., 2005; Bharadhwaj et al., 2022; Jung et al., 201 1). In our
We provide detailed discussion in Appendix C.I.
4

---

Published as a conference paper at ICLR 2025
(a)
(b)
$
JUS
++1
4+2
4+3
M M Cansalmatic q
Ma Weightmarii q
Empowerment
max ElE (s|
D-
Augment
Reweight (oia) oi Te....
(5), St. st. Fe......
ne
7+1
Figure 2: (a) Underlying causal structure of CIP. (b) The whole learning process of CIP includes counterfactual data augmentation, causal action reweight and causal action empowerment.
framework, the empowerment is defined as the mutual information between the agent state 8 and action az. conditioned on the present state 8t and causal mask M, as shown follows:
== max I(o;; 8++1 St, M),
where & denotes the channel capacity from actions to states. Unlike (Cao et al., 2024), which focuses on action-to-staie empowerment effects, we leverage causal understanding and more accurate entropy calculation to analyze state-to-action influences, facilitating the development of more controllable behavioral policies.
4 CAUSAL INFORMATION PRIORITIZATION
In this section, we introduce the proposed framework CIP, which implements causal informa- tion prioritization based on the causal relationships between states, actions, and rewards shown in Figure 2). First, we train a structural model based on the causal discovery method, DirectLiNGAM (Shimizu et al., 2011) using collected trajectories to obtain causal matrix M Utilizing this matrix, CIP executes the swapping of causally independent state features, generating synthetic transitions (Section 4.1). This process of swapping independent state information accen- tuates causally dependent state information, enabling focused learning on critical state transitions. Subsequently, CIP constructs another structural model to get a weight matrix M = that incorporates actions and rewards reweight actions (Section 4.2). Furthermore, CIP integrates causality aware empowerment term & (s) combined with causally weighted actions into the learning objective to promote efficient exploration. This integration encourages the agent's policy æ. prioritize actions with high causal influence, thereby enhancing its goal-achievement capabilities.
4.1 COUNTERFACTUAL DATA AUGMENTATION
To discover the causal relationships between states and rewards, we initially collect trajectories to train structural model by the DirectLiNGAM method, denoted 9h to obtain the causal matrix M Subsequently, we infer the local factorization, which is utilized to generate counterfactual transitions. For each state in the trajectories, compute the uncontrollable set, defined as the set of variables for which the agent has no causal influence on rewards:
U, {s² - (s(, 0; e [I,
where 0 is a fixed threshold and N is the dimension of the state space. The set Ut encompasses all dimensional state variables for which the causal relationship si matrix of states and rewards. Utilizing the learned causal matrix M we partition all state variables in the factored MDP into controllable and uncontrollable sets. These uncontrollable sets are then leveraged for counterfactual data augmentation, thereby prioritizing the muusall-inffrmeer state information to improve learning efficiency.
T. does not exist in the causal
To generate counterfactual samples, perform swap of variables that fall under the uncontrollable category (i.e., in set U.) sampled from the collect trajectories. Specifically, given two transitions (st, az, $141,77) 7.) and ($), āj. f3) sampled from trajectories, which share least uncontrol- lable sub- graph structure (i.e.. 2, nil + 0), we construct a counterfactual transition (5,, de, $ F. by swapping the irrelevant state variables (st, s(4)) with (s), 851) for each U, O Us- The
5

---

Published as a conference paper at ICLR 2025
augmented transitions will be added to the training data for causal reasoning during subsequent action empowerment, thus eliminating the need for additional environment interactions to prioritize causal information. Furthermore, also consider directly using controllable state sets combined with causal action empowerment to replace counterfactual data augmentation for policy learning. The comparative experimental results validating this approach presented in Appendix D.3.1.
4.2 CAUSAL ACTION PRIORITIZATION THROUGH EMPOWERMENT
Causal action reweight Having analyzed the causal relationships between states and rewards to achieve efficient data augmentation, in this section, we further discover the causal relationships be- tween actions and rewards prioritize causally-informedd pacision-making behaviors. CIP constructs a reward-guided structural model, incorporating states (including augmented states), actions, and rewards. This model forms the foundation for action prioritization in policy learning, enabling action reweighting based on causality. Leveraging this structural model to delineate relationships between policy decisions and rewards, we evaluate the causal impact of different actions on reward outcomes. In this way, the agent focuses on pivotal actions with demonstrable causal links to desired reward outcomes, potentially accelerating learning and optimizing performance in complex environments.
Specifically, in CIP, we employ DirectLiNGAM method to train causal structural model L which yields a weight matrix M delineating the relationships between actions and rewards, conditioned on the states. For given set of actions (a,, of ag. ..... we utilize the weight matrix Ma to reweight them as (wia;, aj w2a? wza ...). ), where w represents the causal weights derived from the matrix M By leveraging this causal structure, we can prioritize the most pivotal actions, potentially leading to more efficient policy exploration and targeted policy improvements.
Causal action empowerment Based the learned causal structure, propose the causal action empowerment to incorporate the reweighted actions into the learning objective for efficient exploration in controllable manner. To this end, we design causality policy optimization. We maximize the empowerment gain of the policy Te, where т. incorporates the learned causal structure. This approach allows us to quantify and maximize the empowerment that can be achieved by explicitly considering causal relationships, thereby bridging the gap between causal reasoning and empowerment.
empowerment term &., (s) for
We denote the empowerment of the causal policy as (s) = max, I (ax; 8441 s;; M). We then formulate the following objective empowerment function:
E,,, (s) max I (ot S1+1 M)
(5)
max H(Tx(aals:)) - H(Tr((allssss4))) a=s))(1ss)
where n. is the policy under the causal weighted matrix Ma-rr The first entropy term ((r(()))))) promotes action diversity within the constraints of the causal structure. It encourages the agent to explore a wide range of actions that are causally informed, while the second entropy term -H(m<(aeee [s; $141)) enhances the action predictability in state transitions. encourages the selec- tion of actions that lead to predictable outcomes, given the current and subsequent states, thereby promoting controlled and grie-tre behaviors. We train an inverse dynamics model represent the policy Te(-19 1). The detailed derivation proceeds as follows:
M or @ (o(at())) log x(aiss)
H(Te(-lss)) -EGEAA
(6)
and
H( ((936(119188441)) - EAA
Ma @ Te(a(ls6 log n (aj]osss
1=1
where d is the dimension of the action space. Hence, the learning objective of the causal action empowerment can be defined as follows:
Ex. (s) = max H((T((a)):)) - H(TT(aglss; st+1)) a--s)
max Eg (ee (olle))e))) [log Pa. (at 18t, M) log Pm. (a.)s: M)), (8)
a-1)))(-18))

---

Published as a conference paper at ICLR 2025
basketball
coffee
disasemble
press--pres--all
coffee button
dia-turn
pickstacc
stick push
sarepp
adroit hand door
sparse: door open
sparse: hand Insert
Sparse: window open
Walker
Ari
HattChectah
Hopper
HopperStand
Reacherkard
GuadrupedWalk
HumanoldRun
HumanoidWalk
ReacherEary
Wallerstand
Figure 3: The 36 experimental tasks in 5 continuous control environments
Meta-World 17 tasks
DMControl tasks
Mujoco 4 tasks
Adroit Hand task
Sparse reward 3 3tasks
100
Nomatadd Score
80-
80
80-
80
02-
60-
40
40
ACE
BAC
SAC
ACE
BAC
SAC
ACE
BAC
SAC
BAC
SAC
20
ACE
BAC
SAC
Figure 4: Experimental results with normalized score across all 36 tasks in 5 environments.
where Ps. (ag st; MM is the action distribution given current state policy Te with the causal structure, which can be denoted (a st). Pe. (ag]st41, Sti M) represents inverse dynamics model trained on the collected transitions of state variables. Hence, we update the target policy n by maximizing the empowerment objective function derived in Eq. 8.
Adhering to the maximum entropy paradigm (Haarnoja et al., 2018), we calculate (s) for maxi- mization instead of standard entropy, thus prioritizing exploration of pivotal actions that more likely to have significant causal effects on the reward. This targeted exploration strategy has the potential to accelerate learning by focusing the most influential actions in current controllable states, Based on the causality-aware empowerment, the Q-value for policy ne could be computed iteratively by applying a modified Bellman operator T with En,
(s) term as stated below:
Te Te"Q(se, 0x) r(st, at)
( 7E Eaywte (Q(s aE (s]]]
r(81, at) + +. -. [aan,, (Q(++11, + +a(((n((asss) - M(xa(ad 81 s1:++)))]]]
Hence, integrate the causality-aware empowerment term into the policy optimization objective function, ()M(TT) = Eeppppo.. Eee (r(ss,, 4)) + aEE, (s))]-
In summary, CIP harnesses empowerment to integrate the causal understanding into deciisionmmaing By maximizing the empowerment gain of the causally-informes policy, we guide the agent to prioritize actions that align with the environment's underlying causal relationships. This approach enhances the agent's 's exploration efficiency, focusing on actions with meaningful causal impacts and correlated with desired outcomes. Algorithm I illustrates the complete CIP pipeline.
7

---

Published as a conference paper at ICLR 2025
basketball
coffee push
disassemble
door open
1.00
1.00
M
1.00
1.00
M
Success
0.50
0.50
0.50
0.50
0.00
1000
0.00
1000
0.00
1000
0.00
1000
Environment steps (1e3)
Environment steps (1e3)
Environment steps (1e3)
Environment steps (1c3)
sparse: door open
sparse: hand insert
sparse: window
adroit hand door
MAN
AN
1.00
MMM W
Success rate
0.75
M
0.75
0.75
0.75
0.50
0.50
0.50
0.50
0.25
0.25
0.25
500
1000
500
1000
500
1000
1000
Environment steps (1e3)
Environment steps (1e3)
Environment steps (1e3)
Environment Steps (1e3)
Figure 5: Experimental results of 8 manipulation skill learning tasks in Meta- World and adroit hand environments including sparse reward settings. For all tasks results, please refer to Appendix D.2.
Cheetah Rua
Reacher Fasy
Walker Walk
Cartpole
1200
1020
2.2
Aversge retum
Factor
mm
Average retumn
1000
C1P
TTackor
Average return
CP
Factor
CP
Factorr
800
Average
420
420
270
250
750
1000 Enstronment steps (le2)
Envioement
Enviroement (1el)
Figure 6: Experimental results of 4 pixel-based tasks in DMControl and Cartpole environments.
EXPERIMENTS
Our experiments aim to address the following questions: (i) How does the performance of CIP compare to other RL approaches in diverse continuous control tasks, including manipulation and locomotion with sparse rewards, nigh-dimensional action spaces, and pixel-basee challenges? (ii) Can CIP, through data augmentation and empowerment, improve sample efficiency and learn reliable policies? (iii) What are the effects of the components and hyperparameters in CIP?
5.1 EXPERIMENTAL SETUP
Environments. We evaluate CIP on 5 continuous control environments, including MuJoCo (Todorov al., 2012), DMControl (Tassa al., 2018), Meta-World (Yu al., 2020), Adroit Hand (Rajeswaran Meta-Workd This comprehensive evaluation
et al., 2018), and sparse reward setting environments encompasses 36 tasks, spanning both locomotion and manipulation skill learning, illustrated in Figure: 3. We also conduct experiments in pixel- based tasks of the DMControl and Cartpole environment shown in Figure 17. Our experimental tasks incorporate wide range of challenges, including high-dimensional state and action spaces, sparse reward settings, pixel- scenarios, and locomotion. For extensive experimental settings, please refer Appendix D.I.
Baselines. We compare CIP with three popular RL baselines across all 36 tasks and against IFacoo (Liu et al., 2024) in pixel-based tasks: (1) SAC (Haarnoja al., 2018), an off-policy actor-critic algorithm featuring maximum entropy regularization. (2) ACE (Ji et al., 2024a), method employing causality-awwad aware entropy regularization. (3) BAC (Ji et al., 2024b). a method that balances sufficient exploitation of past successes with exploration optimism. (4) IFactor (Liu et al., 2024). causal framework modeling four distinct categories of latent state variables for pixel-based tasks. To ensure robustness and statistical significance, we conduct each experiment using random seeds.
5.2
MAIN RESULTS
Figure 4 presents the normalized scores CIP compare to other methods across 36 tasks in 5 environments. In 17 Meta-Worle robot-arm tasks, CIP achieves near-perfect score of 100,
8

---

Published as a conference paper at ICLR 2025
Table 1: The experimental results of average return in 8 locomotion tasks. We bold the best scores, and underline second-bes, best results, + is the standard deviation, w/o represents without. indicates CIP is statistically superior to compared method (pairwise 7-test 95% confidence interval).
Cheetah
Hopper Stand
Quadruped Reacher Hard
Method
Ant
HalfCheetah
Hopper
Walker2d
Run
Walk
CIP
6418-81
12594=210
2846-1882
5624:99
893::12
936=17
9484:5
991=111
w/o Aug
6231=81
12225+102
2308+785
5294::1
835
931
9451:3
989113
w/o Emp
6295:1210
10986:1572
2270-1904
55471991
876::21
785
924::23
971
SAC
6062::105
10888+240
2266+:981
5251
167
936#8
930
980#8
BAC
6511=30
10276:#3
2263: 1063
3316=
665#6.
932=#4e
962::24
974#16
ACE
5922+
9390++25
2312:1673
4922+
863+23
912++16
933++57
973+17
showcasing its exceptional performance in manipulation tasks. For locomotion tasks in DMControl and MuJoCo, CIP consistently attains scores exceeding 80, indicating robust performance across diverse locomotion challenges. Notably, CIP exhibits significant performance improvements in challenging scenarios, such adroit hand manipulation and tasks with the sparse reward setting. These results underscore the effectiveness in tackling complex, high-dimensional control problems. In next sections, we present a comprehensive analysis of CIP's performance across diverse tasks.
Robot-arm manipulation. Figure 5 presents the success rates across 7 Meta- World robot-arm manipulation tasks including sparse reward settings, CIP consistently outperforms all other methods across these tasks, demonstrating both faster policy learning and enhanced stability. In challenging tasks, such as disassemble, CIP achieves an impressive 100% success rate. The effectiveness of CIP can be attributed to focus on causally relevant information within the state and action spaces. In sparse reward settings, the efficient extraction causal state information and the prioritization of controllable actions enable effective task completion. By systematically eliminating noise from non-causal factors, CIP allows the agent to construct a more controllable and efficient policy.
High- ligheniimenional Adroit hand manipulation. To rigorously evaluate method's s efficacy in high-ding high-ddmensiooal tasks, we conduct comparative experiments in the Adroit Hand environment of door open task. This challenging setup involves controlling robotic hand with up to 28 actuated degrees of freedom (A E R Figure 5 illustrates the success rates achieved across all methods. Notably, while the three other comparative methods fail to demonstrate significant progress on this challenging task, CIP achieves a near 100% success rate after 700k: environment steps.
Locomotion. We further evaluate CIP in another important category: locomotion. The part experi- mental results of average return in MuJoCo and DMControl environments are presented in Table 1. Learning curves are illustrated in Figure 7. We observe that CIP achieves the best performance in six tasks and sub-optimal in other tasks, and shows statistically significant improvements in 5 out of 8 tasks. Moreover, compared to the traditional method SAC. CIP demonstrates significant perfor- mance improvements in more challenging tasks such as CheetahRun and Hopper. Compared to the causality-ban method ACE, CIP demonstrates improvements all tasks. Overall, locomotion tasks, CIP achieves superior performance and attains high sample efficiency. Detailed performance and statistical analyses are provided in Appendix D.2 and D.3.5.
Pixel-based task learning To further validate the performance in pixel-based tasks, we use 3 complex pixel-based DMControl tasks for evaluation, where video backgrounds serve as distractors. We apply the proposed counterfactual data augmentation and causal action empowerment to IFactor for comparison. As shown in Figure 6, CIP surpasses IFactor in terms of average return. These results underscore CIP's 's efficacy in pixel- pixel-based tasks and its capacity to better overcome spurious correlations arising from video backgrounds, focusing on locomotion. Moreover, the result of Cartpole task in Figure 6 demonstrate the effectiveness in discrete action space environment. For visualization trajectories in pixel-based results, please refer to Appendix D.2.4.
5.3 ANALYSIS
Ablation study. We conduct ablation experiments involving CIP, CIP without (w/o) counterfactual data augmentation (Aug), and CIP w/o Empowerment (Emp). The results in 8 locomotion tasks shown in Table 1. And the learning curves of all tasks are depicted in Appendix D.3. The

---

Published as a conference paper at ICLR 2025
Ant
HalfCheetah
Hopper
Walker2d
6000
Average return
6000
MA
10000
2000
2000
2000
5000
1000
2000
0
0
0
0
1000
Environment steps (1e3)
Environment steps (1e3)
Environment steps (1e3)
Environment sceps (1e3)
CheetahRun
HapperStand
QuadrupedWalk
ReacherHard
1000
1000
Average return
800
1000
1000
www
750
750
200
250
250
500
1000
D
500
1800
1000 Environment steps (1e3)
500
0
500
1000
Environment steps (1e3)
Environment (1e3)
Environment (1e3)
CIP
SAC
BAC
ACE
Figure 7: Experimental results with average return across tasks in locomotion tasks.
Meta-World tasks
DMControl 11 basks
MujoCo tasks
Sparse reward tasks
BAC
BAC
BMC
BAC
ACE
CIP
CIP
CI
CIP
0.45
0.15
0.45
Optimality Gap
Optimality Gap
Optimallyy Gap
Optimality
Figure 8: Experimental results of reliability evaluation by the metric Optimality Gap (lower values better) on diverse environments across 35 tasks.
experimental results reveal that the variant without the empowerment learning objective performs poorly, underscoring the critical role of empowerment maximization enhancing control capabilities. Additionally, CIP without counterfactual data augmentation is less sample efficient than CIP, highlighting the importance of augmentation.
Reliability evaluation. We evaluate CIP's reliability across 35 tasks environments, excluding the Adroit Hand door task due CIP's exceptional performance there. Figure 8 illustrates the experimental results using the Optimality Gap metric (Agarwal et al., 2021). CIP consistently achieves the lowest values across all tasks in four environments, with lower values indicating superior performance. This consistent excellence across diverse scenarios underscores the robustness and reliability of proposed method.
CONCLUSION
This study introduces an efficient RL framework, designed enhance sample efficiency. This approach begins by counterfactual data augmentation using the causality between states and rewards, effectively mitigating interference from irrelevant states without additional environmental interactions. We then develop reward- -giidd structural model that leverages causal awareness to prioritize causal actions through empowerment. We conduct extensive experiments across 39 tasks spanning 5 diverse continuous control environments which demonstrate the exceptional performance of our proposed method, showcasing its robustness and adaptability across challenging scenarios.
L Limitation and Future Work The current limitations of our work are twofold. First, CIP has not yet been extended to complex scenarios, such as real-world 3D robotics tasks. Potential approaches to address this limitation include leveraging object-centric models (Wu et al., 2023), 3D perception models (Wang et al., 2024a), and robotic foundation models (Team et al., 2024; Firoozi et al., 2023) to construct essential variables for causal world modeling. Second, CIP does not adequately consider non -on-stationarity and heterogeneity, which are critical challenges in causal discovery. Future work could integrate method designed to handle such complexities, such CD-NOD (Huang et al., 2020).
10

---

Published as a conference paper at ICLR 2025
REPRODUCIBILITY STATEMENT
We provide the core code of CIP in the supplementary material. The implementation details shown in Appendix D.I.
ACKNOWLEDGMENT
This work was supported in part by the National Natural Science Foundation of China under Grant 62276128, Grant 62192783; in part by the Jiangsu Science and Technology Major Project BG2024031; in part by the Natural Science Foundation of Jiangsu Province under Grant BK20243051; the (ereersittes(43801122 in part by the Collaborative
Fundamental Research Funds for the Central Innovation Center of Novel Software Technology and Industrialization.
REFERENCES
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Marc Bellemare. Deep reinforcement learning the edge of the statistical precipice. Advances in neural information processing systems, 34:29304-2933 293200, 2021.
Homanga Bharadhwaj. Mohammad Babaeizadeh, Dumitru Erhan, and Sergey Levine. Information prioritization through empowerment in visual model-based rl. In International Conference on Learning Representations, 2022.
Jack Brady, Roland Zimmermann, Yash Sharma, Bernhard Schölkopf, Julius Von Kügelgen, and Wieland Brendel. Provably learning object-centric representations. International Conference on Machine Learning pp. 3038 3038-3062. PMLR, 2023.
Hongye Cao, Shangdong Yang, Jing Huo, Xingguo Chen, and Yang Gao. Enhancing ood generaliza- tion in offline reinforcement learning with energy-based policy optimization. In ECAI 2023, pp. 335 342. IOS Press, 2023.
Hongye Cao, Fan Feng, Meng Fang, Shaokang Dong, Jing Huo, and Yang Gao. Towards em- powerment gain through causal structure learning in model-based rl. In ICML 2024 Workshop: Foundations of Reinforcement Learning and
and Perspectives, 2024.
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-9--19943 2022.
David Maxwell Chickering, Optimal structure identification with greedy search. Journal of machine learning research, 5075555507-554, 2002.
Jongwook Choi, Archit Sharma, Honglak Lee, Sergey Levine, and Shixiang Shane Gu. Variational empowerment representation learning for goal-based reinforcement learning. arXiv preprint arXiv:2106.01404, 2021.
Jongwook Choi, Sungtae Lee, Xinyu Wang, Sungryull Sohn, and Honglak Lee. Unsupervised object interaction learning with counterfactual dynamics models. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 11570-11578, 2024.
ZH Deng, Jiang, Long, and Zhang. Causal reinforcement learning: A survey, Transactions on Machine Learning Research, 2023.
Carlos Diuk, Andre Cohen, and Michael L Littman. An object-oriented representation for efficient reinforcement learning. In Proceedings of the 25th international conference on Machine learning, pp. 240-247, 2008
SM Eslami, Nicolas Heess, Theophane Weber. Yuval Tassa, David Szepesvari, Geoffrey E Hinton, et al. Attend, infer, repeat: Fast scene understanding with generative models. Advances in neural information processing systems, 29, 2016.
11

---

Published as a conference paper at ICLR 2025
Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. Diversity is all you need: Learning skills without a reward function. In International Conference on Learning Representations, 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.
Roya Firoozi, Johnathan Tucker, Stephen Tian, Anirudha Majumdar, Jiankai Sun, Weiyu Liu, Yuke Zhu, Shuran Song, Ashish Kapoor, Karol Hausman, et al. Foundation models in robotics: Applications, challenges, and the future. The International Journal of Robotics Research, pp- 027836492412815008 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.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference machine learning, pp. 1861-1870. PMLR, 2018.
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 on Neural Networks and Learning Systems, 2023.
Dan Haramati, Tal Daniel, and Aviv Tamar. Entity-centric reinforcement learning for object manip- ulation from pixels. In NeurlPS 2023 Workshop on Goal-Condi Reinforcement Learning, 2023.
Biwei Huang, Kun Zhang, Jiji Zhang, Joseph Ramsey, Ruben Sanchez-Romero, Clark Glymour, and Bernhard Schölkopf. Causal discovery from heterogeneou//nonsattionayy data. Journal of Machine Learning Research, 21(89): 53,, 2020.
Biwei Huang, Fan Feng, Chaochao Lu, Sara Magliacane, and Kun Zhang. Adarl: What, where, and how to adapt in transfer reinforcement learning. In International Conference on Learning Representations, 2022a.
Biwei Huang, Chaochao Lu, Liu Leqi, José Miguel Hermández-Lobato. Clark Glymour, Bernhard Schölkopf, and Kun Zhang. Action-sufficient state representation learning for control with structural constraints. In International Conference on Machine Learning, pp. 9260-9279. PMLR, 2022b.
Tianying Ji, Yongyuan Liang, Yan Zeng, Yu Luo, Guowei Xu, Jiawei Guo, Ruijie Zheng, Furong Huang, Fuchun Sun, and Huazhe Xu. Ace: Off-policy actor-critic with causality-award entropy regularization. In Forty-first International Conference on Machine Learning, 2024a.
Tianying Ji, Yu Luo, Fuchun Sun, Xianyuan Zhan, Jianwei Zhang, and Huazhe Xu. Seizing serendip- ity: Exploiting the value of past success in off off-policy actor-critic. In Forty-first International Conference on Machine Learning, 2024b.
Jindong Jiang, Sepehr Janghorbani, Gerard De Melo, and Sungjin Ahn. Sealor: Generative world mod- els with scalable object representations. In International Conference on Learning Representations, 2019.
Jindong Jiang, Fei Deng, Gautam Singh, and Sungjin Ahn. Object-centric slot diffusion. arXiv preprint arXiv:2303. 10834, 2023.
Jindong Jiang, Fei Deng, Gautam Singh. Minseung Lee, and Sungjin Ahn. Slot state space models. arXiv preprint arXiv:2406. /272 2024.
Tobias Jung, Daniel Polani, and Peter Stone. Empowerment for continuous envirommen systems. Adaptive Behavior, 19(1): 196339,,,, 2011.
12

---

Published as a conference paper at ICLR 2025
Michael Kearns and Daphne Koller. Efficient reinforcement learning in factored mdps. In IJCAI, volume 16. pp. 740-747, 1999.
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.
Avinash Kori, Francesco Locatello, Ainkaran Santhirasekaram, Francesca Toni, Ben Glocker, and Fabio De Sousa Ribeiro. Identifiable object-centric representation learning via probabilistic slot attention. XrXiv preprint 2ew::2466..1141, 2024.
Adam Kosiorek, Hyunjik Kim, Yee Whye Teh, and Ingmar Posner. Sequential attend, infer, repeat: Generative modelling of moving objects. Advances in Neural Information Processing Systems, 31, 2018.
Jannik Kossen, Karl Stelzner, Marcel Hussing, Claas Voelcker, and Kristian Kersting. Structured object-aware physics prediction for video modeling and planning. In International Conference on Learning Representations, 2019.
Oliver Kroemer, Scott Nickum, and George Konidaris. A review of robot learning for manipulation: Challenges, representations, and algorithms. Journal of machine learning research, 22(30):1-82, 2021.
Sébastien Lachapelle, Divyat Mahajan, loannis Mitliagkas, and Simon Lacoste-Julien. Additive decoders for latent variables identification and centesiaa-produce extrapolation. Advances in Neural Information Processing Systems, 36, 2024.
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.
Richard Li, Allan Jabri, Trevor Darrell, and Pulkit Agrawal. Towards practical multi-object manipu- lation using relational reinforcement learning. In 2020 ieee international conference robotics and automation (icra), pp. 4051 1058. IEEE, 2020.
Xiaoqi Li, Mingxu Zhang, Yiran Geng, Haoran Geng, Yuxing Long, Yan Shen, Renrui Zhang, Jiaming Liu, and Hao Dong. Manipllm: Embodied multimodal large language model for object-centric robotic manipulation. In Proceedings of the IEEE/CVF Conference Computer Vision and Puttern Recognition, pp. 18061-18070, 2024.
Yulong Li and Deepak Pathak. Object-aware gaussian splatting for robotic manipulation. In ICRA 2024 Workshop on 3D Visual Representations for Robot Manipulation, 2024. URL https: //openreview net wooouoriddddeeeee 43DDo.
Zhixuan Lin, Yi-Fu Wu, Skand Peri, Bofeng Fu, Jindong Jiang, and Sungjin Ahn. Improving generative imagination in object-centric world models. In International conference machine learning, pp. 6140-6149. PMLR, 2020.
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.
Francesco Locatello, Dirk Weissenbomn Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot attention. Advances in neural information processing systems, 33:1 1121112211515 1533, 2020.
Davide Mambelli, Frederik Träuble, Stefan Bauer, Bernhard Schölkopf, and Francesco Locatello. Compositional multi-object reinforcement learning with linear relation networks. In ICLR2022 Workshop the Elements of Reasoning: Objects, Structure and Causality, 2022.
Chaitanya Mitash, Mostafa Hussein, Jeroen Vanbaar, Vikedo Terhuja, and Kapil Katyal. Scaling object-centric robotic manipulation with multimodal object identification. 2024.
13

---

Published as a conference paper at ICLR 2025
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, Emanuele Rossi, Juan Felipe Calderon, Michael Bronstein, and Marcello Restelli. Provably efficient causal model-based reinforcement learning for systematic generalization. In Proceedings of the AAAI Conference Artificial Intelligence, volume 37, pp. 9251-9259, 2023.
Jongjin Park, Younggyo Seo, Chang Liu, Li Zhao, Tao Qin, Jinwoo Shin, and Tie-Yan Liu. Object- aware regularization for addressing causal confusion in imitation learning. Advances in Neural Information Processing Systems, 34:3029-3043 2021.
Judea Pearl. Causality. Cambridge university press, 2009.
Silviu Pitis, Elliot Creager, and Animesh Garg. Counterfactual data augmentation using locally factored dynamics. Advances in Neural Information Processing Systems, 33:3976-3990, 2020.
Silviu Pitis, Elliot Creager, Ajay Mandlekar, and Animesh Garg. Mocoda: Model-based counterfac- tual data augmentation. Advances in Neural Information Processing Systems, 35:18143-18156, 2022.
Aravind Rajeswaran, Vikash Kumar, Abhishek Gupta, Giulia Vezzani, John Schulman, Emanuel Todorov, and Sergey Levine. Learning complex dexterous manipulation with deep reinforcement learning and demonstrations. Robotics: Science and Systems XIV, 2018.
Sahand Rezaci-Shoshiee Rosie Zhao, Prakash Panangaden, David Meger, and Doina Precup. Con- tinuous mdp homomorphisms and homomorphic policy gradient. Advances in Neural Information Processing Systems, 35:20189-20204 2022.
Jonathan Richens and Tom Everitt. Robust agents learn causal world models. arXiv preprint arXiv:2402. 10877, 2024.
Christoph Salge, Comellus Glackin, and Daniel Polani. Empowerment-an introduction. Guided Self-Orranization: Inception, pp- 67-114, 2014.
Manolis Savva, Abhishek Kadian, Oleksandr Maksymets, Yili Zhao, Erik Wijmans, Bhavana Jain, Julian Straub, Jia Liu, Vladlen Koltun, Jitendra Malik, et al. Habitat: platform for embodied ai research. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 9339-9947, 2019.
Jonathan Scholz, Martin Levihn, Charles Isbell, and David Wingate. A physics- model prior for object-oriented mdps. In International Conference on Machine Learning, pp. 1089-1097. PMLR, 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: 22905-22918. 2021
Shohei Shimizu, Takanori Inazumi, Yasuhiro Sogawa, Aapo Hyvarinen, Yoshinobu Kawahara, Takashi Washio, Patrik Hoyer, Kenneth Bollen, and Patrik Hoyer. Directlingam: direct method for learning linear non- ganssiann structural equation model. Journal of Machine Learning Research-IMLR, 12(Apr): 1248, 2011.
David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou. Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. nature, 550(7676):354-35) 2017.
Peter Spirtes, Clark Glymour, and Richard Scheines. Causation, prediction, and search. MIT press, 2001.
Hao Sun and Taiyi Wang. Toward causal-aware rl: State-wise action-refined temporal difference. In Deep Reinforcement Learning Workshop NeurIPS 2022.
14

---

Published as a conference paper at ICLR 2025
Yuewen Sun, Erli Wang, Biwei Huang, Chaochao Lu, Lu Feng, Changyin Sun, and Kun Zhang. Acamda: Improving data efficiency in reinforeement learning through guided counterfactual data augmentation. In Proceedings of the AAAI Conference Artificial Intelligence, volume 38, pp. 15193-15201, 2024.
Richard S Sutton. Reinforcement learning: An introduction. A Bradford Book, 2018. 8.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint arXiv: 00690, 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 arXiv:2405. 12211, 2024.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ international conference on intelligent robots and systems, pp. 5026-5033.
IEEE, 2012.
Núria Armengol Urpi, Marco Bagatella, Marin Vlastelica, and Georg Martius. Causal action influence aware counterfactual data augmentation. In Forty-first International Conference on Machine Learning, 2024.
Elise Van der Pol, Thomas Kipf, Frans A Olichoek, and Max Welling. Plannable approximations to mdp homomorphisms: Equivariance under actions. arXiv preprint arXiv:2002. 11963, 2020.
Arthur Wandzel, Yoonscon Oh, Michael Fishman, Nishanth Kumar, Lawson LS Wong, and Stefanie
Tellex. Multi-object search using object-oriented pomdps. In 2019 International Conference on Robotics and Automation (ICRA), pp. 7194-7200 IEEE, 2019.
Chenxi Wang. Hongjie Fang, Hao-Shu Fang, and Cewu Lu. Rise: 3d perception makes real-world robot imitation simple and effective. arXiv preprint arXiv:2404. /2281, 2024a.
Zizhao Wang, Xuesu Xiao, Yuke Zhu, and Peter Stone. Task-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. In International Conference on Machine Learning, pp. 23151-23180 80. PMLR, 2022.
Zizhao Wang, Jiaheng Hu, Caleb Chuck, Stephen Chen, Roberto Martín-Martin, Amy Zhang, Scott Nickum, and Peter Stone. Skild: Unsupervised skill discovery guided by factor interactions. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024b.
Zizhao Wang. Jiaheng Hu, Peter Stone, and Roberto Martin-Martin. Elden: exploration via local dependencies. Advances in Neural Information Processing Systems, 36, 2024c.
Zizhao Wang, Caroline Wang, Xuesu Xiao, Yuke Zhu, and Peter Stone. Building minimal and reusable causal state abstractions for reinforcement learning. arXiv preprint arXiv:2401. 12497, 2024d.
Nicholas Watters, Loie Matthey, Matko Bosnjak, Christopher Burgess, and Alexander Lerchner. Cobra: Data-efficient model-based rl through unsupervised object discovery and curiosity-driven exploration. arXiv preprint 100.:19955000255. 2019.
Alicia Peregrin Wolfe and Andrew G Barto. Defining object types and options using mdp homomor- phisms. In Proceedings of the ICML-06 Workshop on Structural Knowledge Transfer for Machine Learning, 2006.
Ziyi Wu, Nikita Dvomik, Klaus Greff, Thomas Kipf, and Animesh Garg. Slotformer: Unsupervised visual dynamics simulation with object-centri, models. In The Eleventh International Conference Learning Representations, 2022.
15

---

Published as a conference paper at ICLR 2025
Ziyi Wu, Jingyu Hu, Wuyue Lu, Igor Gilitschenski, and Animesh Garg. Slotdiffusion: Object-centric generative modeling with diffusion models. Advances in Neural Information Processing Systems, 36:50932-50958. 2023.
Ziyi Wu, Yulia Rubanova, Rishabh Kabra, Drew A Hudson, Igor Gilitschenski, Yusuf Aytar, Sjoerd van Steenkiste, Kelsey R Allen, and Thomas Kipf. Neural assets: 3d-aware multi-object scene synthesis with image diffusion models. arXiv preprint Stti::24000092922. 2024.
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on robot learning, pp. 1094-1100. PMLR, 2020.
Wentao Yuan, Chris Paxton, Karthik Desingh, and Dieter Fox. Sornet: Spatial object-centric representations for sequential manipulation. In Conference Robot Learning, pp. 148-157,
PMLR, 2022.
Andrii Zadaianchuk, Maximilian Seitzer, and Georg Martius. Self-supervised visual reinforcement learning with object-centric representations. In International Conference on Learning Representa- tions, 2021.
Andrii Zadaianchuk, Georg Martius, and Fanny Yang. Self-supervised reinforcement learning with independently controllable subgoals. In Conference on Robot Learning, pp. 384-394. PMLR, 2022.
Andrii Zadaianchuk, Maximilian Seitzer, and Georg Martius. Object-centric learning for real- world videos by predicting temporal feature similarities. In Thirty-seventh Conference on Neural Information Processing Systems (NeurIPS 2023), 2023.
Yan Zeng, Ruichu Cai, Fuchun Sun, Libo Huang, and Zhifeng Hao. survey on causal reinforcement learning. arXiv preprint arXiv:2302.052209 2023.
Yudi Zhang, Yali Du, Biwei Huang, Ziyan Wang, Jun Wang, Meng Fang, and Mykola Pechenizkiy. Interpretable reward redistribution in reinforcement learning: a causal approach. Advances in Neural Information Processing Systems, 36, 2024.
Linfeng Zhao, Lingzhi Kong, Robin Walters, and Lawson LS Wong, Toward compositional general- ization in object-oriented world modeling. In International Conference on Machine Learning, pp. 16841-26866 PMLR, 2022.
16

---

Published as a conference paper at ICLR 2025
CONTENTS
Introduction
2 Related Work
2.1
Causal RL
2.2
Empowerment in RL
2.3
Object-centric RL and Object-Orientee RL
3
Preliminaries
3.1
Markov Decision Process
3.2
Structural Causal Model
3.3
Empowerment in RL
4
Causal Information Prioritization
4.1
Counterfactual Data Augmentation
4.2
Causal Action Prioritization Through Empowerment
5
Experiments
5.1
Experimental setup
5.2
Main Results
5.3
Analysis
6
Conclusion
10
Broader Impact
19
B Assumptions and Propositions
19
C Extended Related Work
20
C.1 Extended Discussion on object-centric RL and 3D world models
20
Details on Experimental Design and Results
D.1
Experimental setup
21
D.2 Full Results
D.2.1
Effectiveness in robot arm manipulation
21
D.2.2
Effectiveness in spare reward settings
22
D.2.3
Effectiveness in locomotion
22
................................
D.2.4
Effectiveness in pixel-based tasks
27
D.3 Property Analysis
27
D.3.1 Analysis for replacing counterfactual data augmentation
27
D.3.2 Extensive ablation study
28
17

---

Published as a conference paper at ICLR 2025
D.3.3
Hyperparameter analysis
29
D.3.4
Computation cost analysis
30
D.3.5
Statistical performance analysis
31
D.3.6
Generalization analysis
D.3.7
Causal discovery analysis
32
Details on the Proposed Framework
36
Experimental Platforms and Licenses
36
F.1
Experimental platforms
36
F2 Licenses
36
18

---

Published as a conference paper at ICLR 2025
A BROADER IMPACT
To avoid blind exploration and improve sample efficiency, we propose CIP for efficient reinforcement learning. CIP leverages the causal relationships among states, actions, and rewards prioritize causal information for efficient policy learning. CIP first learns a causal matrix between states and rewards execute counterfactual data augmentation, prioritizing important state features without additional environmental interactions. Subsequently, it learns a causal reweight matrix between actions and rewards prioritize causally-informeet behaviors. We then introduce a causal action empowerment term into the learning objective to enhance the controllability. By prioritizing the causal information, CIP enables agents to focus on behaviors that have causally significant effects on their tasks. CIP offers substantial broader impact by prioritizing causal information through individual assessment of how different factors contribute to rewards. Our novel empowerment learning objective achieves efficient policy optimization by leveraging entropy via the policy and learned inverse dynamics model. This approach shows promise for extension into research frameworks centered on maximum entropy algorithms.
Despite its strengths, CIP has limitations beyond its reliance on the method DirectLiNGAM. There's potential to explore alternative causal discovery techniques for more robust relationship mapping. Moreover, analyzing inter-entity causal connections could lead to better disentanglemest of diverse behaviors. Our future work will investigate range of causal discovery methods to refine our approach. We aim to extend CIP to model-based RL frameworks, focusing on building causal world models to enhance generalization.
B ASSUMPTIONS AND PROPOSITIONS
Assumption (d-separation (Pearl, 2009)) d-separation is graphical criterion used determine, from a given causal graph, if a set of variables X is conditionally independent of another set Y. given a third set of variables Z. In a directed acyelic graph (DAG) G, a path between nodes nı and nmm is said to be blocked by set S if there exists a node for k = 2, ... m 1, that satisfies one of the following two conditions:
(i) ny E S, and the path between nk-1 and ns+1 forms (ng-1 - 7i - n.41), (Mk-1 & Tik nk41), or (nk-1 < ni. > mk+1)).
(ūi) Neither n& nor any of its descendants is in S, and the path between n6-1 and n4++1 forms
(ni--1 + ng +- (k-1))
In DAG, we say that two nodes and ny are d-separated by a third node ne if every path between nodes Da and og is blocked by ne. denoted as na n6 n
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 (st,... ..t,...,. s1 ut,, ... ,a,, . 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, RR = p(S)A) ((R
Assumption 3 (Faithfulness Assumption (Spirtes et al., 2001; Pearl, 2009)) For set of variables V = (s), -......, s 3 ai, ... at re)T" there are no independencies between variables that are not implied by the Markovian Condition.
Assumption 4 Under the assumptions that the causal graph is Markov and faithful to the observa- tions, the edge si -> * exists for all state variables s².
Assumption No simultaneous or backward edges in time.
Proposition 1
Under the assumptions that the causal graph is Markov and faithful to the observa- tions, there exists an edge from of > T, if and only if at J.. rel0, (a, \ of 8t-
Proof: We proceed by proving both directions of the if and only if statement.
(+) Suppose there exists an edge from a; to r,. We prove that of JL rr la, \ of. s, by contradiction. Assume of Il. relar \ at, 8t. By the faithfulness assumption, this independence must be reflected in
19

---

Published as a conference paper at ICLR 2025
the graph structure. However, this implies the absence of directed path from aj to rt, contradicting the existence of the edge. Thus, of +f rela: di. St.
(+) Now, suppose of He rela, \ oi St. We prove the existence of an edge from of to r by contra- diction. Assume no such edge exists. By the Markov assumption, the absence of this edge implics a; IL ra aala \ of. St. contradicting Thus, we have shown that an edge from aj to r. exists if and only if a 14 T1 az de 8t, completing the proof.
initial supposition. Therefore, an edge from of to rt must exist.
Proposition 2 Under the assumptions that the causal graph is Markov and faithful to the observa- tions, there exists an edge from st + Tx if and only sf s Tel rel St (re..
The proof of Proposition 2 follows similar line of reasoning that of Proposition 1.
Theorem Based on above 5 assumptions and propositions, suppose 8t: 0t, 81, follow the factored MDP reward function Eq. the causal matrices M and M are identifiable.
C EXTENDED RELATED WORK
We categorize existing causal RL approaches based on problem domains and task types, providing a systematic analysis of how different methods explore causal relationships between states, actions, and rewards, as illustrated in Table 2.
In the single-task learning domain, methods such ACE (Ji al., 2024a) and IFactor (Liu et al., 2024) have shown success in learning policies for manipulation and locomotion tasks. However, both approaches are limited by focusing on single seward-guided causal relationship. Regarding generalization, AdaRL (Huang et al., 2022a) effectively leverages both state-reward and action-rewaid causal relationships. However, AdaRL focuses primarily on applying causal inference to address generalization challenges in locomotion tasks. Its application is limited to locomotion tasks, leaving more complex manipulation tasks unaddressed. Since our work focuses on the single-task problem domain, we do not provide direct comparison with AdaRL. Conversely, CBM (Wang et al., 2024d) considers the causal relationship between states and rewards but overlooks the causal link between actions and rewards. In the problem domain of counterfactual data augmentation, current causal RL methods (Urpí et al., 2024; Pitis et al., 2020; 2022) have not yet explored the inference and utilization of both causal relationships.
In summary, current research on seward-guided causal discovery remains incomplete and lacks validation across a broader spectrum of tasks. This gap underscores the need for more comprehensive investigation and application in the field of causal reinforcement learning.
C.I EXTENDED DISCUSSION ON OBJECT-CENTRIC RL AND 3D WORLD MODELS
The main similarity lies between our framework and object-centric RL is both are learning and using factored MDPs (Kearns & Koller, 1999), but they differ in granularity: our framework operates at the component level (e.g., raw state variables), whereas object-centric RL factors states based on objects.
Although our work is orthogonal to object-centric RL, we believe certain elements of object-centric RL could complement our framework specific applications, particularly in real-world robotic manipulation tasks. Potential future work include:
Using object-centric representation as input: Object-centric models can help identify object-factored variables, such as object attributes, geometry, and physical states, which are useful for planning (Jiang et al., 2019; Lin et al., 2020; Kossen et al., 2019; Mambelli et al., 2022; Feng & Magliacane, 2023; Choi et al., 2024; Zadaianchuk et al., 2022; Park et al., 2021; Zadaianchuk et al., 2021; Yuan et al., 2022; Li et al., 2020; Mitash et al., 2024; Haramati et al., 2023; Li et al., 2024). In this case, states are factored as objects, and we can learn causal graphs over these variables. This is useful in robotic environments involving numerous objects. We will leave this as future work for adapting our current framework to the applications of the object-centric robotic task.
20

---

Published as a conference paper at ICLR 2025
Table 2: Categorization of different causal RL methods with different causal relationship of state-to-reward (state-reward) and action-to-reward (action-reward).
Causal relationship
Problem domain
Task type
Method
state-reward action-reward
Single-task
manipulation; locomotion
ACE (Ji et al., 2024a)
x
manipulation; locomotion
Factor (Liu et al., 2024)
/
X
manipulation
CAI (Seitzer et al., 2021)
x
X
Generalization
manipulation
CDL (Wang et al., 2022)
x
X
locomotion
AdaRL. (Huang et al., 2022a)
/
/
manipulation; locomotion
CBM (Wang et al., 2024d)
Augmentation
manipulation
CAIAC (Urpí et al., 2024)
x
X
manipulation
CoDA (Pitis et al., 2020)
manipulation
MoCoDA (Pitis et al., 2022)
Learning more compact factored object representations with our framework: Our structure learning approach could benefit object-centric RL by disentangling the internal representations of individual objects to the reward-relevant and reward-irrelevan groups by learning the causal structures. This can enhance the compactness and interpretability of object-centric representations.
Using object-aware 3D world models for applications: In 3D environments, object- aware 3D world models (Li & Pathak, 2024) can provide essential representations of objects. Our framework could then build causal structures on top of these factored 3D-object representations.
While these directions are promising and could advance the applicability of our framework in certain domains, they are outside the primary focus of this work. We plan explore these ideas part of future work.
D DETAILS ON EXPERIMENTAL DESIGN AND RESULTS
D.1 EXPERIMENTAL SETUP
We present the detailed hyperparameter settings of the proposed method CIP across all 5 environments in Table 3. Additionally, the Q-value and V-value networks used MLP with 512 hidden size. And the policy network is the Gaussian MLP with 512 hidden size. Moreover, we set the target update interval of 2. For fair comparison, the hyperparameters of the baseline methods (SAC (Haarnoja et al., 2018), BAC (Ji et al., 2024b), ACE (Ji et al., 2024a)) follow the same settings in the experiments.
For pixel- based DMControl environments, we employ IFactor (Liu al., 2024) encode latent states and integrate the CIP framework for policy learning. We utilize the sf state features in IFactor as uncontrollable states unrelated to rewards to execute counterfactual data augmentation. Furthermore, for simplicity, we maximize the mutual information between future states and actions to facilitate empowerment. All parameter settings in these three tasks adhere those specified in IFactor. Additionally, We use the same background video for the comparison.
D.2
FULL RESULTS
D.2.1 EFFECTIVENESS IN ROBOT ARM MANIPULATION
Figure 9 presents the learning curves for all 17 manipulation skill tasks within the Meta-World environment. The CIP framework demonstrates superior learning outcomes and efficiency compared to the three baseline methods, despite exhibiting minor instabilities in the basketball and dial-turn tasks. Notably, CIP achieves a 100% success rate in more complex tasks, such as pick-place-we and assembly. The visualization results presented in Figures 11 and 12 further demonstrate CIP'S
21

---

Published as a conference paper at ICLR 2025
Table 3: Hyperparameter settings of CIP in 5 environments
Environment
Hyperparameter
Neta-World
Sparse
MuJoCo
DMControl
Adroit Hand
batch size
512
512
256
512
256
hidden size
1024
1024
256
1024
256
Q-value network hidden size
512
V-value network hidden size
512
policy network hidden size learning step
512
1000000
replay size
1000000
causal sample size
10000
gamma learning rate update interval
0.99
0.0003
2
ability to effectively and efficiently complete tasks, even in high-dimensional action spaces such as the Adroit Hand environment.
In the hammer task, CIP allows the robot arm to execute reach and pick actions with precision, enabling it to accurately identify the nail's position and successfully perform the hammering action. In the Adroit Hand door task, CIP effectively controls the complex joints to grasp the doorknob and applies the appropriate force twist thereby opening the door.
These findings affirm the effectiveness of CIP robot arm manipulation skill learning, highlighting its capacity to enhance sample efficiency while mitigating the risks associated with blind exploration.
Visualization. We employ trajectory visualization to comparatively validate the efficacy of our method. As depicted in Figure 10, the light-shaded regions delineate the policy exploration space, while the point clustering area indicates the area of frequent interaction. Our analysis reveals that CIP, leveraging counterfactual data augmentation, achieves substantially broader exploration compared to ACE and SAC. Concurrently, the causal information prioritization framework facilitates more focused execution in critical state regions. These visual findings provide robust empirical support for the effectiveness of our proposed augmentation framework.
D.2.2 EFFECTIVENESS IN SPARE REWARD SETTINGS
Figure 13 presents the learning curves for all three sparse reward setting tasks within the Woold environment, while Figure 14 showcases their corresponding visualization trajectories. These findings reveal that CIP not only achieves superior learning efficiency but also adeptly executes critical actions necessary for task completion, such as opening the door and window and maneuvering the node to the target place.
These results substantiate the effectiveness of CIP sparse reward scenarios. The counterfactual data augmentation process prioritizes salient state information, effectively filtering out irrelevant factors that could hinder learning. Mcanwhile, causal action empowerment enhances policy controllability by focusing actions that are causally linked to desired outcomes. This dual approach not only accelerates the learning process but also fosters a more robust policy capable of navigating the complexities inherent in sparse reward settings. Overall, these findings underscore CIP's potential significantly improve performance in challenging environments characterized by limited feedback.
D.2.3 EFFECTIVENESS IN LOCOMOTION
We further evaluate CIP 15 locomotion tasks in DMControl and MuJoCo environments. Figure 15 presents the learning curves, while Figure 16 showcases the corresponding visualization trajectories specific tasks. A comprehensive analysis indicates that CIP achieves faster learning efficiency and greater stability compared to ACE and SAC, while demonstrating comparable policy learning
22

---

Published as a conference paper at ICLR 2025
basketball
coffee push
disassemble
door open
1.00
1.00
M
1.00
1.00
M
Success
0.50
0.50
0.50
0.50
0.00
1000
0.00
1000
0.00
1000
0.00
1000
Environment steps (1e3)
Environment steps (1e3)
Environment steps (1e3)
Environment steps (1c3)
door unlock
drawer
hammer
pick place wall
1.00
1.00
1.00
0.75
0.75
0.75
0.75
Success
0.25
0.25
0.25
0.25
500
1000
500
1000
500
1000
1000
Environment steps (1e3)
Environment steps (1e3)
Environment steps (1e3)
Environment steps (1e3)
plate slide
stick push
sweep
0.75
0.75
0.75
0.75
Success
0.50
0.50
0.50
0.50
0.25
0.25
0.25
0.25
0.00
500
1000
0.00
500
1000
0.00
500
1000
0.00
500
1000
Environment
(1e3)
Environment steps (1e3)
Environment steps (1e3)
Environment steps (1e3)
assembly
button press wall
coffee button
dial burn
1.00
1.80
1.00
1.00
AMMM
W
Success
0.50
0.50
0.50
0.50
0.00
500
1000
0.00
1000
0.00
1000
0.00
1000
Environment steps (1e3)
Environment steps (1e3)
Environment steps (1e3)
Envirenment steps (1e3)
pick place
1.00
SUCCES rate
0.75
0.25
250
750
Environment steps (1e3)
CP
SAC
BAC
ACE
Figure 9: Experimental results across 17 manipulation skill learning tasks in Seta-World.
performance to BAC, which is known for its proficiency in control tasks. The visualization results reveal that CIP effectively executes running and walking actions in complex humanoid scenarios.
These findings collectively underscore the efficacy of CIP in locomotion tasks, highlighting its potential advance the state-of-the-art reinforcement learning for intricate motor control problems. The method's success across varied environments suggests robust framework that could generalize effectively to other challenging domains within robotics and control systems, paving the way for future research and applications in these areas.
23

---

Published as a conference paper at ICLR 2025
SAC
CIP
ACE
Figure 10: Visualization of the trajectories in soccer task.
time
baskeeball
Figure Visualization trajectories of 4 manipulation skill learning tasks in Meta- World environment.
time
Adroit hand cloor
Figure 12: Visualization trajectory of Adroit Hand door open task.
sparse: door open
sparse: insert
sparse: window open
1.00
1.00
1.00
MMM
Success rate
M
0.75
0.75
0.75
0.50
0.50
0.50
0.00
0.00
0.00
500
1000
1000
500
1000
Environment steps (1e3)
Environment steps (1e3)
Environment steps (1e3)
SAC
BAC
ACE
Figure 13: Experimental results across 3 manipulation skill learning tasks in sparse reward settings of Meta- World environment.
24

---

Published as a conference paper at ICLR 2025
time
Sparse
doxr
Spans haand lisert
Sparrer Window open
Figure 14: Visualization trajectories of manipulation skill learning tasks in sparse reward settings of Meta- environment.
CheetahRun
HopperStand
QuadrupedWaik
ReacherHard
1000
1000
800
800
750
750
400
400
200
200
250
250
0
500
1000
0
500
1000
500
1000
500
1000
Environment steps (1e3)
Environment steps (1e3)
Envirenment steps steps (1e3)
Environment steps (1e3)
WalkerRun
CartpoleSwingup
FingerSpin
ReacherEeasy
1000
1000
1000
1000
800
800
600
600
500
500
400
400
200
200
0
500
1000
1000
1000
500
1000
Environment Steps (1e3)
Environment steps (1e3)
Environment steps (1e3)
Environment steps (1e3)
WalkerStand
HumanoidRun
HumanoidWaik
1000
1000
400
Average retur
800
800
300
600
600
400
200
400
200
100
1000
2000
Environment sleps
(1e3)
Environment sleps (1e3)
Environment steps (1e3)
CIP
SAC
BAC
ACE
Figure 15: Experimental results across 11 locomotion tasks in DMControl environment.
25

---

Published as a conference paper at ICLR 2025
time
Cartpde Swingg
Chastah Run
Humaadl Run
Hunaaooi Walk
Figure 16: Visualization trajectories of 4 locomotion tasks DMControl environment.
26

---

Published as a conference paper at ICLR 2025
Figure 17: The DMControl environment of 3 pxiel-based tasks (Walker Walk, Cheetah Run, Reacher Easy) and 1 task in Cartpole environment (Liu et al., 2024).
time
Cheetah Run
Rescher Eset
Walker Walk
Figure 18: Visualization trajectories in 3 pixel-based locomotion tasks of DMControl environment with video backgrounds as distractors.
D.2.4 EFFECTIVENESS IN PIXEL- BASED TASKS
To further validate the effectiveness of our proposed framework in pixel-based environments, we evaluated CIP three DMControl pixel-based tasks. We leverage IFactor for latent state processing and differentiation of uncontrollable state features to execute counterfactual data augmentation, alongside maximizing the mutual information between future states and actions for empowerment.
Figure 6 presents the learning curves, while Figure 18 shows the visualization trajectories. The proposed framework exhibits enhanced policy learning performance and effectively mitigates in- terference from background video, facilitating efficient locomotion. These findings reinforce the effectiveness and extensibility of our causal information prioritization framework, highlighting its potential to improve learning in complex, pixel-basee environments.
D.3 PROPERTY ANALYSIS
D.3.1 ANALYSIS FOR REPLACING COUNTERFACTUAL DATA AUGMENTATION
In CIP, we exploit the causal relationship between states and rewards to perform counterfactual data augmentation on irrelevant state features, thus prioritizing critical state information. We compare this approach with alternative method: masking irrelevant state features to achieve state abstraction for subsequent causal action empowerment and policy learning. To evaluate the efficacy of both approaches, conduct experiments with CIP with counterfactual data augmentation (CIP w/i Cda) and CIP with causally-in informedd states (CIP w/i Cs) across three distinct environments.
Figure 19 illustrates comparative results for four manipulation skill learning tasks in the GoradWorld environment. Both CIP variants achieve 100% task success rates with high sample efficiency, vali- dating their effectiveness. Notably, CIP w/i Cda exhibits superior learning efficiency compared to
27

---

Published as a conference paper at ICLR 2025
plate slide
soccer
stick push
sweep into
1.00
1.00
1.00
1.00
Success rate
0.50
0.50
0.50
0.50
0.00
500
1000
0.00
500
1000
0.00
1000
0.00
1000
Environment steps (1e3)
Environment steps (1e3)
Environment steps (1e3)
Envirenment steps (1c3)
CIP w/ Cda
CP wi Cs
BAC
ACE
Figure 19: Experimental results in 4 manipulation skill learning tasks of Meta- Woold environment. w/i stands for with.
sparse: door
sparse: hand insert
sparse: window open
1.00
1.00
1.00
M MMMM
Success rate
0.75
0.75
0.50
0.50
0.50
0.25
0.25
0.25
0.00
0.00
0.00
500
1000
500
1000
500
1000
Environment steps (1e3)
Environment steps (1e3)
Environment steps (1e3)
CIP w/i Cda
CIP w/i Cs
BAC
ACE
Figure 20: Experimental results in 3 manipulation skill learning tasks of Meta- World environment with sparse reward settings.
CIP w/i Cs, underscoring the value of our counterfactual data augmentation approach in enhancing training data without additional environmental interactions. In three sparse reward setting tasks (Figure 20), CIP w/i Cda demonstrates superior policy performance. Further experiments across four locomotion environmetal tasks corroborate these findings, consistently favoring the counterfactual data augmentation approach. These comprehensive experimental results strongly support the effec- tiveness and significance incorporating counterfactual data augmentation CIP, highlighting its potential to enhance reinforcement learning across diverse task domains.
D.3.2 EXTENSIVE ABLATION STUDY
Robot arm manipulation The ablation study results in the Meta- World and Adroit Hand envi- ronments are presented in Figure 22. The findings indicate that CIP without counterfactual data augmentation exhibits reduced learning efficiency and is unable to successfully complete tasks such as pick-and-place. This underscores the importance of incorporating counterfactual data augmentation, which prioritizes causal state information, to enhance learning efficiency by mitigating the influence of irrelevant state information and preventing policy divergence.
Furthermore, CIP without causal action empowerment demonstrates a significant decline in policy performance across robot arm manipulation tasks. In complex scenarios, such as Adroit Hand door opening and assembly, fails learn effective strategies for task completion. This outcome further corroborates the efficacy of the proposed causal action empowerment mechanism, as prioritizing causally informed actions facilitates more efficient exploration of the environment, ultimately enabling successful policy learning.
Sparse reward settings Figure 22 presents the results of the ablation study conducted across three sparse reward setting tasks. These findings underscore the substantial influence of causal action empowerment on the efficacy of policy learning, demonstrating its critical role in enhancing performance in challenging environments. Additionally, the incorporation of counterfactual data augmentation proves effective in mitigating the need for additional environmental interactions, thereby significantly improving sample efficiency. This approach not only facilitates more rapid learning but
28

---

Published as a conference paper at ICLR 2025
Rhn
HopperStand
QuadrupedWalk
ReacherHard
1000
1000
1000
1000
800
800
600
600
400
400
500
500
200
200
500
1000
0
500
1000
Environment steps (1e3)
Environment steps (1e3)
Environment steps (1e3)
Environment steps (1e3)
CIP w/ Cda
CP wi Cs
BAC
ACE
Figure 21: Experimental results in 4 locomotion tasks of DMControl environment.
also ensures that the agent can effectively navigate sparse reward scenarios by focusing on the most relevant causal information.
Locomotion We further conducted ablation experiments on locomotion tasks, The experimental results the MuJoCo environment are shown in Figure 23, where evident that the performance of CIP without causal action empowerment declines significantly. Similarly, CIP without counterfactual data augmentation also exhibits reduced learning efficiency. Notably, in the 11 DMControl tasks, the decline in performance for CIP without causal action empowerment is particularly pronounced.
These experimental results further validate the effectiveness of our proposed method, which systemat- ically analyzes the causal relationships between states, actions, and rewards. This analysis enables the execution of counterfactual data augmentation to avoid interference from irrelevant factors while prioritizing important state information. Subsequently, by leveraging the causal relationships between actions and rewards, we reweight actions to prioritize causally informed actions, thereby enhancing the agent's controllability and overall learning efficacy.
D.3.3 HYPERPARAMETER ANALYSIS
We conduct a detailed analysis the hyperparameters associated with the causal update interval (I) and sample size within the CIP framework. The experimental results for four distinct tasks are illustrated in Figure 25. Across all tasks, CIP demonstrates optimal performance with a causal update interval of 2 and sample size of 10,000.
Our findings suggest that while a reduction the causal update interval can lead to improved performance, it may also result in heightened computational costs. Additionally, we observe that higher update frequencies and increased sample sizes introduce greater instability, which significantly raises computational demands. This analysis underscores the importance of carefully balancing hyperparameter settings to optimize both performance and efficiency within the CIP.
Furthermore, we analyze the performance under different settings of the temperature factor a proposed in Eq. 9. The results across 3 tasks are shown in Figure 26. Our analysis reveals that CIP manipulation tasks, while showing
demonstrates robust performance across different values some instability in locomotion tasks when a: is either too small or too large. Moreover, we observe that setting a to 0.2 yields optimal performance across all tasks, which motivated our choice of 0.2 for all experiments.
Finally, we analyze the performance under different settings of the batch size and hidden size. The results across 3 tasks shown in Figure 27. Our experimental results demonstrate that CIP exhibits robust performance across various parameter settings in coffee push and sparse hand insert tasks, while maintaining strong performance in hopper stand task. Based these experimental results, configure the hyperparameters as follows: for manipulation tasks, we set the batch size to 512 and hidden size to 1024, while for locomotion tasks, All other hyperparameters remain constant across all tasks, as detailed in Table 3.
batch size of 256 and hidden size of 256.
29

---

Published as a conference paper at ICLR 2025
basketball
coffee push
disassemble
door open
1.00
1.00
W
1.00
1.00
0.75
0.50
0.50
0.50
0.50
0.25
0.25
0.00
500
1000
0.00
500
1000
0.00
500
1000
0.00
500
1000
Environment steps
Environment steps
Environment steps
Environment steps
door unlock
drawer
hammer
pick place wall
1.00
1.00
1.00
1.00
0.75
0.75
0.75
0.75
0.50
0.50
0.50
0.50
0.25
0.25
0.25
0.25
0.00
500
1000
0.00
500
1000
0.00
Só0
1000
0.00
1000
Environment steps
Environment steps
Environment steps
Environment steps
slide
stick push
sweep
1.00
OV
0.75
0.75
0.75
0.75
0:25
0.25
0.25
0.25
0.00
0.00
500
1000
500
1000
500
1000
1000
Environment steps
Environment steps
Environment steps
Environment steps
assembly
button press wall
coffee button
dial turn
1.00
1.00
M
1.00
1.00
m
M
0.75
0.75
0.75
0.50
0.50
0.50
0.50
0.25
0.25
0.25
0.25
0.00
500
1000
0.00
500
1000
0.00
500
1000
0.00
500
1000
Environment steps
Environment steps
Environment steps
Environment steps
sparse: open
sparse:
sparse: open
1.00
1.00
1.00
1.00
M
1.00
0.75
0.75
Success
050
0.25
02
0:55
025
0.25
Environment
Environment steps
Environment steps
Environment steps
Environment steps
CIP
CIP w/o Aug
CIP w/o Emp.
Figure 22: Ablation results across manipulation skill learning tasks in Meta-World including sparse reward settings and adroit hand.
D.3.4 COMPUTATION COST ANALYSIS
We analyze the computational cost of the proposed framework. The computation time for all methods across 36 tasks is shown in Figure 28. Our experimental results demonstrate that CIP achieves its performance improvements with minimal additional computational burden specifically less than 10% increase compared to SAC, less than 5% increase compared to ACE, and actually requiring less computation time than BAC. All experiments were conducted on the same computing platform with the same computational resources detailed in Appendix F. These experimental results verify that our proposed method achieves performance improvements without incurring significant additional computational costs.
30

---

Published as a conference paper at ICLR 2025
Ant
HalfCheetah
Hopper
Walker2d
d
6000
Average return
6000
10000
3000
4000
2000
2000
5000
2000
500
1000
500
1000
500
1000
500
1000
Environment steps
Environment steps
Environment steps
Environment steps
CP
CIP w/o Aug
CPP w/o Emp
Figure 23: Ablation results across 4 locomotion tasks in MuJoCo environment.
CheetahRun
HopperStand
QuadrupedWalk
ReacherHard
1000
1000
Average return
800
800
1000
VVV
1000
600
600
400
400
500
500
200
200
500
1000
500
1000
0
500
1000
500
1000
Environment steps
Environment steps
Environment steps
Environment steps
WalkerRun
CartpoleSwingup
FingerSpin
ReacherEasy
Average return
M
1000
1000
800
750
750
600
600
400
400
500
500
200
200
500
500
1000
1000
Environment steps
Environment steps
Environment steps
Environment steps
WalkerStand
HumanoidRun
HumanoidWalk
400
Average retum
1000
200
100
1000
2000
4000
1000
2000
3000
Environment steps
Environment steps
Environment steps
CIP
CIP w/o Aug
CIP w/o Emp
Figure 24: Ablation results across 11 locomotion tasks in DMControl environment.
D.3.5 STATISTICAL PERFORMANCE ANALYSIS
To further validate the statistical significance of the performance, we select 3 statistical metrics (Agar- wal et al., 2021) IQM, Mean, and Median for analysis across 8 locomotion tasks. The results shown in Figure 29 and 30. Our findings indicate that CIP achieves notably superior performance across all tasks, with the sole exception of the ant task where performs slightly below BAC.
D.3.6 GENERALIZATION ANALYSIS
We conduct multi-task experiments in the Meta- World environment (Yu et al., 2020) to validate the generalizability. We establish MT1 and MT10 tasks for generalization validation:
Multi-Task 1 (MT1): Learning one multi-task policy that generalizes to tasks belonging to the same environment. MT1 uses single Meta- World environments, with the training "tasks" corresponding to
31

---

Published as a conference paper at ICLR 2025
basketball all
door open
Walker2d
CheetahRun
1.00
1.00
6000
1000
Success rate
0.75
A
0.75
Average: return
4000
800
600
2000
400
0.25
200
0.00
0.00
500
1000
0
500
1000
Environment steps (3e.3
Environment steps (1e3)
Environment steps (1e3)
Environment steps (1c3)
CIP: 1=1, size=5000
CP:1=1, wizc-iizzoo 10000
CI:: I- size- 5000
CIP: 1=2, size 10000
Figure 25: Hyperparameter study. Learning curves of CIP with different hyperparameter settings. The shaded regions are the standard deviation of each policy,
coffee push
sparse hand insert
HopperStand
2000
2000
1000
Average return
Average return
800
1500
1500
600
1000
1000
400
500
500
200
500
1000
500
1000
500
1000
Environment steps (1e3)
Environment steps (le3)
Environment steps (1e3)
a=0.1
a=0.2 (CIP)
a=0.5
a=1
Figure 26: Hyperparameter analysis temperature factor a across task.
5 random initial object and goal positions. The goal positions are provided in the observation and a fixed set, as to focus on the ability of algorithms in acquiring distinct skill across multiple goals, rather than generalization and robustness.
Multi-Task 10 (MT10): This task involves learning a single multi-task policy that generalizes to 50 tasks across 10 training environments, totaling 500 training tasks. A crucial step towards rapid adaptation to distinctly new tasks is the ability to train single policy capable of solving multiple distinct training tasks. The multi-task evaluation Meta-Worts tests the ability to learn multiple tasks simultaneously. without accounting for generalization to new tasks. The MT10 evaluation encompasses 10 environments: reach, push, pick and place, open door, open drawer, close drawer, press button top-down, insert peg side, open window, and close window.
We adapt our proposed CIP to multi-task learning by incorporating one-hot task as input, comparing MT-CIP with MT-SAC. The results in Figure 31 show that MT-CIP outperforms MT- SAC in both MT1 (soccer) and MT10 tasks, achieving average success rates above 50% and 40% respectively. Notably, MT-CIP exhibits strong performance in specific MT10 tasks like drawer close and window open. The superior performance of MT-CIP stems from its effective learning of causal information, enabling robust task transfer across diverse domains. While these results promising, future work will focus on causal state abstraction for enhanced generalization and sample efficiency. All experiments were conducted under the same hyperparameter settings, and the implementation will be made publicly available.
D.3.7 CAUSAL DISCOVERY ANALYSIS
In CIP, we use the linear causal discovery method DirectLiNAM for causal structure learning. To explore alternative approaches, we compare it with two other causal discovery methods: score-based GES (Chickering, 2002) and constraint-based PC (Spirtes et al., 2001). ). The experimental results in Figure 32 across three tasks demonstrate that our chosen DirectLiNAM method exhibits superior performance compared to both alternatives. During experimentation. GES and PC methods incur significant computational overhead and frequently encounter memory constraints. In contrast, our proposed method CIP, which is fundamentally reward-guided, efficiently
also observe that both
32

---

Published as a conference paper at ICLR 2025
coffee push
sparse hand insert
HopperStand
2000
1000
1500
Average return
800
1500
1000
600
1000
400
500
500
200
0
0
0
500
1000
0
500
1000
500
1000
Environment steps (1e3)
Environment steps (1e3)
Environment steps (1e3)
batch=256, 26d
batch=512, hidden=512
batch=512, hidden=1024
Figure 27: Hyperparameter analysis of batch size and hidden size across 3 task.
disassemble
open
door
open
1..55
125
10.0
10.0
10.0-
ACE
BAC
BAG
BAC
ACE
BAC
SAC
ACE
BAC
SIKE
2.5
2.5
0.0
0.0
0.0
0.01
0.0
hammer
pick place wall
plate silde
soccar
stick push
sweap Into
12.5
12.5
10.0
BAC
BAC
BAC
RAC
SAC
assembly
button press wall
coffee button
dial turn
plck place
sparse door open
12.5
10.0
100.
10.0
10.0-
10.0-
10.8
7.5
BAC
7.5
ACE
BAC
ACE
BAC
7.3
ACE
BAC
SAC
7.5
ACE
BAC
SAC
5o.
5.0
5.0
5.3-
5.00
5.0
sparse
sparse open
HopperStand
QuadrupedWalk
12.5
125
125
12.5
12.5
12.5
75
ACE
BAC
75
ACE
BAC
SAC
ACE
ACE
BAC
SAC
BAC
sa.
0.0
0.3
ReacherHard
WalkerRun
CartpoleSwingup
FingerSpin
ReacherEaasy
WalkerStand
BAC
SAC
BAC
BAG
+
HumanoidRRn
HumanoidWallk
Ant
HalfCheetah
Hopper
Walker2d
BAC
BAC
Figure 28: Computation time 36 tasks.
discovers causal relationships between dimensional factors in states and actions with respect to rewards. This approach better aligns with the requirements of policy learning while maintaining minimal computational costs.
33

---

Published as a conference paper at ICLR 2025
CheetahRun
CheetahRun
CheetahRun
SAC
SAC
SAC
BAC
BAC
BAC
ACE
ACE
640
720
800
880
640
720
800
880
640
720
800
880
IQM
Mean
Median
HopperStand
HopperStand
HopperStand
SAC
SAC
BAC
BAC
BAC
ACE
ACE
ACE
CIP
CIP
CIP
925
950
975
900
975
920
940
960
IQM
Mean
Median
QuadrupedWalk
QuadrupedWalk
QuadrupedWalk
SAC
SAC
SAC
BAC-
BAC
BAC
ACE
ACE
ACE
CIP
CIP
CIP
900
930
960
990
900
930
960
990
880
920
960
1000
IQM
Mean
Median
ReacherHard
ReacherHard
ReacherHard
SAC
SAC
BAC
ACE
ACE
CIP
CIP
960
990
900
930
960 Mean
990
930
960
990
IQM
Median
Figure 29: statistical metrics of IQM, Mean, and Median (higher values are better) on DMControl tasks.
34

---

Published as a conference paper at ICLR 2025
Ant
Ant
Ant
SAC
SAC
SAC
BAC
ACE
CIP
CIP
CIP
6000
6200
6400
6000
6200 Mean
6400
6000
6200 Median
6400
6600
IQM
HalfCheetah
HalfCheetah
HalfCheetah
SAC
SAC
SAC
BAC
BAC
BAC
ACE
ACE
ACE
CIP
CIP
CIP
10000
11000 IQM
12000
10000
11000 Mean
12000
10000
11000 Median
12000
Hopper
Hopper
Hopper
SAC
SAC
SAC
BAC
BAC
BAC
ACE
ACE
ACE
CIP
2500
3000
3500
2500
3000 Mean
3500
1800
2400
3000 Median
3600
IQM
Walker2d
Walker2d
Walker2d
SAC
BAC
BAC
BAC
ACE
ACE
ACE
CIP
CIP
CIP
3200
4800
5600
3200
5600
3200
4800
5600
IQM
Mean
Median
Figure 30: statistical metrics of IQM, Mean, and Median (higher values are better) on 4 MuJoCo tasks.
MT.1 soccer
MT-10 Average
MT-100 drawer close
MT-1D D window open
MT-CP
MI-OP
NT-CIP
1.25
Success rate
0.6
MT SAC
Success rate
1.00
MT SAC
M
0-4
0.75
0.2
0.50
0.25
150 Environment steps (1e4)
Environment
Environment steps (1o4)
steps (1e4)
Figure 31: Generalization results in MTI and MT10 tasks.
35

---

Published as a conference paper at ICLR 2025
door open
sparse window open
ReacherHard
1000
1500
Average return
1500
800
1000
1000
600
400
500
500
200
0
0
0
100
200
300
0
100
200
300
200
400
600
Environment steps (1e3)
Environment steps (1e3)
Environment steps (1e3)
DirectLiNGAM
GES
PC
Figure 32: Compared performance with different causal discovery methods across task.
E DETAILS ON THE PROPOSED FRAMEWORK
Algorithm lists the full pipeline CIP below.
F EXPERIMENTAL PLATFORMS AND LICENSES
EXPERIMENTAL PLATFORMS
All experiments of this approach are implemented on 2 Intel(R) Xeon(R) Gold 6430 and NVIDIA Tesla A800 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)
Meta- World (MIT License)
MuJoCo (Apache License 2.0)
Deep Mind Control (Apache License 2.0)
Adroit Hand (Creative Commons License 3.0)
36

---

Published as a conference paper at ICLR 2025
Algorithm 1 Causal information prioritization for efficient RL
Input: Q network Qs policy network we. inverse dynamics model de with Q network Qá replay buffer D, local causal buffer De. causal update interval causal matrix M and M n-
for each environment step / do Collect data with He from real environment Add to replay buffer D and local buffer D. end for
Step 1: Counterfactual data augmentation
every environment step then
Sample transitions D, from local buffer Dr
Learn causal mask matrix M M with {(s, a, r, s')} (Pal for causal state prioritization Compute lncontrollable set Ug followed by Eq. 4 D.
Sample (s,
for s² E U, do
Sample (8, à f') 6")~~ D. if state U4 then
Construct counterfactual transition (5, ā, F, 7,3") by swapping s²) with (3', 8*) Add (8, a, F, 8') to local huffer D:
end if end for end
Step 2: Causal weighted matrix learning
every environment step then Sample transitions Da from local buffer De Learn causal weighted matrix M with ((s, a, 8')) causal action prioritization end
Step 3: Policy optimization with causal action empowerment
for each gradient step do
Sample N transitions (s, s') from D
Compute causal action empowerment followed by Eq. 8.
Calculate the target Qa, value
Update Qo. by mins, (TQo, Q
Update &e by
Calculate the target Q+, value
Update Q by min., (ToQ Q+e)²
Update m. by max. (8,0) En, (8))
end for
37


Reviews:
[{"review_id": "dQCt2d9WX1", "reviewer": "Reviewer_wgJu", "summary": "This paper introduces Causal Information Prioritization (CIP) to address sample inefficiency in reinforcement learning by leveraging causal relationships between states, actions, and rewards. CIP uses factored Markov Decision Processes (MDPs) to model these causal relationships and employs techniques like counterfactual data augmentation to improve sample efficiency. The framework is tested across various tasks, including manipulation and locomotion, to showcase its effectiveness in different environments.", "questions": "- How does CIP address the need for task-specific causal structure input? Is it possible to reduce the engineering effort required for this, or is the approach dependent on manual structure provision?\n- Were there specific reasons for the chosen baselines (ACE, BAC)? Would other baselines, particularly those that handle causality differently, add more context to the results?\n- What are the plans for integrating 3D world understanding or extending CIP to work more naturally with object-centric or high-dimensional MDPs without predefined causal structures?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The paper presents an interesting approach by integrating causal reasoning with reinforcement learning (RL), which could potentially lead to better sample efficiency by filtering out irrelevant features.\n- The focus on causality and counterfactual augmentation aligns with a growing area in RL, aiming to make models more effective with fewer interactions by prioritizing meaningful features.\n- The experiments cover diverse environments and tasks, which is valuable for assessing the generalizability of the approach.", "weaknesses": "- The paper seems to need more references and discussion around object-centric MDPs and related works in object-centric world modeling and MDP homomorphism, which are relevant areas given the nature of CIP.\n- The proposed framework appears to require considerable task-specific engineering, as it doesn’t seem to inherently learn the causal structure but rather relies on manually provided task structures. This may limit its scalability and adaptability to new tasks without re-engineering.\n- Empirical results, especially in Table 1, show limited improvement over baselines like ACE and BAC. The gains aren’t always significant, especially given the complexity and additional engineering that CIP requires. The benefits may not justify the added complexity in some cases.\n- Some results in Table 1 indicate the CIP approach performs on par with or even below baselines for certain tasks, suggesting that the improvements might not generalize strongly across all evaluated environments.", "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": "BvI5CfZr4x", "reviewer": "Reviewer_kTcs", "summary": "This paper introduces CIP, the causal information prioritization method that leverages causal information in reinforcement learning to improve learning efficiency. Comprehensive experiments are conducted to verify the proposed method.", "questions": "Will the proposed method be suitable for offline RL or has it to be applied with online explorations?\n\nCan the authors provide more analysis on the pitfalls of the proposed method? What is the trade-off of the proposed method, and under what cases would this method face challenge? e.g. computational burden, can not handle highly stochastic environments? \n\nI appreciate the authors' efforts in their large-scale experiments, yet I have concerns about the statistical significance of the performance. On some of the tasks, the performance seems to have large variances. I wonder what would be the recommended hyper-parameter setups from the authors, and is the proposed method stable/robust to different hyper-parameter settings?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper is well-organized. The presentation is mostly clear. The experiment section is comprehensive.", "weaknesses": "1. Several related works are missed: \n\n[1] Sun, Yuewen, et al. \"ACAMDA: Improving Data Efficiency in Reinforcement Learning Through Guided Counterfactual Data Augmentation.\" Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 38. No. 14. 2024.\n\n[2] Sun, Hao, and Taiyi Wang. \"Toward Causal-Aware RL: State-Wise Action-Refined Temporal Difference.\" arXiv preprint arXiv:2201.00354 (2022).\n\n2. The limitation section in the current paper is not really a discussion of the limitations of the proposed method.", "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": "PaIj2fsA2X", "reviewer": "Reviewer_hNGa", "summary": "The paper proposes a novel reinforcement learning (RL) approach called Causal Information Prioritization (**CIP**), designed to address sample inefficiency by leveraging causal reasoning. CIP focuses on prioritizing causal relationships between states, actions, and rewards in reinforcement learning tasks, thus guiding the agent to focus on the most impactful behaviors. CIP improves sample efficiency by leveraging factored Markov Decision Processes (MDPs) to infer causal relationships and using counterfactual data augmentation. It integrates a causality-aware empowerment learning objective that enhances goal-directed action execution, leading to more efficient exploration. The approach is validated experimentally across 39 tasks in five diverse environments involving continuous control, demonstrating superior performance over existing RL techniques in terms of both sample efficiency and task success rates.", "questions": "**Questions**\n\n1. Could the authors elaborate on how CIP could be adapted or extended to discrete action spaces, where causal dependencies may be more rigid?\n\n2. How does CIP handle potential **inaccuracies in causal discovery** during the learning process? Is there a mechanism in place to correct or adapt the causal models if they are misaligned with the environment dynamics?\n\n3. **Causal Discovery**: Could the authors elaborate on the choice of DirectLiNGAM for causal discovery? Have other causal discovery techniques been tested, and if so, how do they compare in terms of performance and sample efficiency?\n\n4. Assumptions of DirectLiNGAM: How sensitive is CIP to the assumptions required for the DirectLiNGAM method (e.g., linearity and Gaussian noise)? Would non-linear extensions of causal discovery alter CIP's effectiveness?\n\n5. **Computational Overhead**: Could you provide a more detailed analysis of the computational costs associated with empowerment calculations compared to traditional entropy regularization approaches?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "Integrating the causality-aware empowerment objective into the policy optimization objective alongside strong, extensive experimental results - demonstrating its ability to outperform several state-of-the-art RL methods, particularly in challenging environments with sparse rewards and high-dimensional action spaces - makes the paper a valuable contribution to the RL community.\n\n**Originality**: The combination of causal reasoning with empowerment learning for reinforcement learning is a novel and creative approach.\n\n**Quality**: The paper is well-researched and technically rigorous, with solid theoretical foundations and thorough experimental validation.\n\n**Clarity**: The motivation and design of CIP are well-explained, with intuitive examples such as the robot-arm manipulation scenarios. The illustrations help make abstract concepts more accessible.\n\n**Significance**: CIP addresses an important issue in reinforcement learning—sample inefficiency—by leveraging causal relationships, which has the potential for broad applicability in complex environments.", "weaknesses": "**Complexity of Empowerment Calculation**: The use of empowerment as part of the learning objective, though promising, involves additional computational complexity. The paper would benefit from a more in-depth discussion of how computational costs compare to the gains in sample efficiency, especially for tasks with very high-dimensional action spaces.\n\n**Generalizability and Assumptions**: The reliance on the DirectLiNGAM method for causal discovery may limit generalizability in real-world environments where non-linear relationships are common. Exploring alternative causal discovery methods could enhance robustness.\n\n**Empirical Comparisons**: While CIP is compared to standard baselines like SAC and ACE, further empirical comparisons to model-based RL approaches or hybrid methods would strengthen claims regarding sample efficiency and generalization capabilities.", "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)