BECAUSE: Bilinear Causal Representation for Generalizable Offline Model-based Reinforcement Learning
Authors: Haohong Lin
Wenhao Ding
Jian Chen
Laixi Shi
Jiacheng Zhu
Bo Li
Ding Zhao
Venue: NeurIPS 2024 poster (2024)
Keywords: Model-based RL
Causal Reasoning
Offline RL
Primary area: reinforcement_learning

Abstract: Offline model-based reinforcement learning (MBRL) enhances data efficiency by utilizing pre-collected datasets to learn models and policies, especially in scenarios where exploration is costly or infeasible. Nevertheless, its performance often suffers from the objective mismatch between model and policy learning, resulting in inferior performance despite accurate model predictions. This paper first identifies the primary source of this mismatch comes from the underlying confounders present in offline data for MBRL. Subsequently, we introduce **B**ilin**E**ar **CAUS**al r**E**presentation (BECAUSE), an algorithm to capture causal representation for both states and actions to reduce the influence of the distribution shift, thus mitigating the objective mismatch problem. Comprehensive evaluations on 18 tasks that vary in data quality and environment context demonstrate the superior performance of BECAUSE over existing offline RL algorithms. We show the generalizability and robustness of BECAUSE under fewer samples or larger numbers of confounders. Additionally, we offer theoretical analysis of BECAUSE to prove its error bound and sample efficiency when integrating causal representation into offline MBRL. See more details in our project page: [https://sites.google.com/view/be-cause](https://sites.google.com/view/be-cause).

BECAUSE:
Bilinear Causal Representation for Generalizable Offline Model-based Reinforcement Learning
Haohong Lin', Wenhao Jian Chen', Laixi Shi?, Jiacheng Zhu", Bo Li", Ding Zhao'
*CUU, "Clttec,, "MIT, "UChicago UIUC
{haohongl, wenhaod)Gandres: .mu.cmu edu
Abstract
Offline model-based reinforcement learning (MBRL) enhances data efficiency by utilizing pre-collected datasets to learn models and policies, especially in scenarios where exploration is costly or infeasible. Nevertheless, its performance often suffers from the objective mismatch between model and policy learning, resulting in inferior performance despite accurate model predictions. This paper first identifies the primary source of this mismatch comes from the underlying confounders present offline data for MBRL. Subsequently, we introduce BilinEar CAUSal rEpresentation (BECAUSE), an algorithm to capture causal representation for both states and actions to reduce the influence of the distribution shift, thus mitigating the objective mismatch problem. Comprehensive evaluations 18 tasks that vary in data quality and environment context demonstrate the superior performance BECAUSE over existing offline RI. algorithms. We show the generalizability and robustness of BECAUSE under fewer samples or larger numbers of confounders. Additionally, we offer theoretical analysis of BECAUSE prove its error bound and sample efficiency when integrating causal representation into offline MBRL. See more details in our project page: https///ttessooogll.com/iiiiiii
Introduction
Offline Reinforcement Learning (RL) has shown great promise in learning directly from pre-collected datasets, especially in scenarios where active interaction is expensive or infeasible M Specifically, offline model-based reinforcement learning (MBRL) 2 3 learning policies with an estimated world model, generally perform better than their model-free counterparts in long-horizon tasks such self-driving vehicles [.. robotics (6), and healthcare 77 However, offline RL suffers from distribution shift because the rollout data is either sampled from some suboptimal behavior policies or sampled from slightly different training environments compared to the deployment time
Although identifying distribution shift issues, many of the current offline MBRL works fail to model the shift in environment dynamics, which is ubiquitous and could cause catastrophic fail- ure of trained policy at slightly different de- ployment stage. Furthermore, since the learn ing objectives of the world models and poli- cies are isolated from each other, a significant challenge offline MBRL is objective mis- match [9 10] problem (shown in Figure[]): mod-
collection behovlor (Unknown)
Data Buffer
Model T(x' (s.))
Policy r(als)
test Environment
(s.a.etrr
Objective Mismatch
Inw model loos success X
low model loss x success V
Success
Our method
Falure
Model
Model
Figure 1: The objective mismatch problem.
els that achieve lower training loss are not nec-
essarily better for control performance. For example, dynamics model achieve relatively low
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

prediction loss, yet it may not be sufficient to guide the planner to the high-reward region. Previous works [0] have attempted reduce such objective mismatch by jointly learning the model and policy. However, the performance is suboptimal due to the lack of digging into the underlying cause of the mismatch issue [
In this work, we identify that the objective mismatch between model estimation and policy learning comes from two sources of distribution shift in offline MBRL: (1) shift between the online optimal policy and offline sub-optimal behavior policies, and (2) shift between the data collection environment and online testing environments. Unlike humans, who make decisions based on reasoning over task- relevant factors, models in offline RL memorize correlations without learning the causality. The sub-optimal behavior policies introduce spurious correlations [[]] between actions and states, making the model memorize specific actions. When online testing environments differ from the data collection environment, the model could overfit spurious correlations in the state and fail to generalize to unseen states. Based on the analysis of the above mismatch, our work differs from previous work of causal model-based RL 12 in that we model causality in both model and policy learning.. We aim to avoid spurious correlation by discovering underlying structures between abstracted states and actions.
To alleviate objective mismatch and generalize well, introduce the BilinEar CAUSal rEpresenta- tion (BECAUSE) that integrates the causal representation in both world model learning and planning of MBRL agents. Inspired by preliminary works that use bilinear MDPs to capture the structural representation in MBRL [14] we first approximate the causal representation to capture the low-rank structure the world model, then use this learned representation facilitate planning by quantifying the uncertainty of sampled transition pairs. Consequently, we factorize the spurious correlations and learn unified representation for both the world model and planner.
In summary, the contribution of this paper is threefold:
We formulate offline MBRL into the causal representation learning problem, highlighting the tight connection between structural causal models and low-rank structures in MDPs. To the best of our knowledge, this the first work that systematically reveals the connection between causal representation learning and Bilinear MDPs.
We propose BECAUSE, empirical causal representation framework, based the above formu- lation. BECAUSE first learns a causal world model, then fosters the generalizability of offline RL agents by quantifying the uncertainty of the state transition, which facilitates conservative planning to mitigate the objective mismatch.
We provide extensive empirical studies and performance analysis in tasks of multiple domains to demonstrate the superiority of BECAUSE over existing baselines, which illustrates its potential to improve the generalizability and robustness of offline MBRL algorithms.
Problem Formulation
To alleviate the objective mismatch problem and the degraded performance caused by the spurious correlation, we first provide our novel formulation of learning the underlying causal structures of Markov Decision Process (MDP) under the bilinear MDP setting. then introduce the causal discovery for MDP with confounders.
2.1 Preliminary: MDP and Bilnear MDP
We denote an episodic finite-horizon MDP by M (S which is composed of state space S, action space A, a set of transition functions T. planning horizon H and reward function associated with task preferences. Without loss of generality in many real-world practices, we assume that the reward function is bounded by r E [0, 1], wh E [H]. Specifically, we are interested in a coalicondiio reward setting, where Vg E (ss, = if and only if s 9.
Given a policy π and the state-action pair (s, a) S A. we then define the state-action value function in the timestep l as QT E, EM rs(3,, a;)(sa 8, au a] and the value func-
tion VA (s) E E ri(ss, 4s) s]. The expectation E,, here is integrated into randomness throughout the trajectory, which is essentially induced by the random action of the policy a, m(-ss:)

---

and the time- home--ooogeneuuss transition dynamics of the environment 8+41 ~ T(-ss; az), Vi E [h, H]. In the offline dataset, the data rollouts can be seen as generated by some (mixed) behavior policy Tg. resulting in dataset DD with in total samples (8i, ; d4, s4.rchissnn
Definition 1 (Bilinear MDP [44. For each (s, a) E S x A, E S, we have the corresponding feature vector :) R
R - Rd (u-) |R - R With some core matrix M E Rdxe
we can represent the transition function kernel [- as
a, s' E SxxXXxXx T(s = a(s,a) (paa)
(I)
where dbss. a) and ju(s") are embedding functions that map the original state and action to the latent space, M is the core matrix that models the transition relationship between the previous timestep and next timestep in the latent space. Such a linear decomposition in the transition dynamics allows us embed structures of the transition model without the loss of general function approximation capabilities to derive state and action representations.
2.2 Action State Confounded MDP
We consider the existence of confounders in the MDP to represent the offline data collec- tion process, and define action-state confounded MDP (ASC-MDP):
Causal Path Confounder Path
State Nodes
Confounders
Action Nodes
Definition (ASC-MDP). Besides the compo- nents in standard MDPs M = A, H.,) r). we introduce a set of unobserved confounders In ASP-MDP, confounders are factorized as u = [uz, ue MishHH H. where Un E U denotes the confounders between and Ts(s) in- duced by behavior policies, and u, denotes the confounders within the state-action pairs of the environment transition, that is, the inher- structure between (s, and s'. Here assume time-invariant confounder distribution u ~ Pu(-), Vh E [H], which is common assump-
[a] Confounded
(b) MDP
(c) ASC-MDP (Ours]
Figure 2: Comparison of our ASC-MDP with two existing formulations.
tion 15 [6 17
The resulting causal relationship of ASC-MDP is demonstrated in Figure Originating from the original MDP, ASC-MDP is different from the Confounded MDP 18 and State-Confounded MDP (SC-MDP) [9 in that it models both the spurious correlation between the current state and current action a, as well as those between the next state s' and (s, a). Yet, confounded MDP and SC-MDP only model part of the possible confounders between states and actions. The factorization of the confounder ASC-MDP aligns with the source of spurious correlation offline MBRL.
Proposed Method: BECAUSE
We propose BECAUSE, our core methodology for modeling, learning, and applying our causal representations for generalizable offline MBRL. Section 3.1 models the basic format of causal representations and analyzes their properties. Section 3.2 gives a compact way to learn the causal representation d(s, a) and yss'), as well as the core mask estimation M. Section 3.3 utilizes these learned causal representations in both world model learning and MBRL planning from offline datasets.
3.1 Causal Representation for ASC-MDP
In the presence of a hidden confounder u, we model the confounder behind the transition dynamics as linear confounded MDP [18
(( = ul
(2)
where u ~ Pa(.). Inspired by the Bilinear MDP in Definition we decompose 7(s a, u) into confounder-aware core matrix M (u) and feature mapping Φ(s, a) which factorize the influence of
3

---

the confounders. Given the factorization of confounder u fue, ux in Definition 2 we derive via d-separation the graphical model in Figure that s' IT u+ [s, a, ue] As a result, we only need to consider the confounder u. from the environment when decomposing the transition model:
T(s'|s a, u) T a, a, ue) a(s,a) M (uc)usss).
(3)
[Odx M
Definition 3 (Construction of causal graph G). In ASC-MDP G
0e'xa' for all (sparse)
0d'xd
core matrix M, the causal graph G' is bipartite, thus V G.G.G DAG.
Definition 3 reveals the connection between the core matrix M and causal graph G, as is formulated in the ASC-MDP. To reduce the influence of and estimate the unconfounded transition model T(s'|s, a), one way is to identify the causal structures induced by the confounder th for the transition dynamics [2)] Existing methods in differentiable causal discovery [21 22 23 transform causal
discovery on some causal graph G, into a maximum likelihood estimation (MLE) with regularization:
G arg max loaaxooggppppdoeeo p(D;o, A/G GEDAG
M arg max max pD;; o, n4 M) A|M (4)
M
Since our case, M is sub-matrix the causal graph G. Given the Definition 3 G automatically satisfies the formulation of ASC-MDP, thus discovering G is essentially estimating the sparse submatrix M without DAG constraints: M = arg max M log pDD; Φ. μ, M) AlM). We elaborate Definition 3 and show the relationship between core matrix M and causal graph in Appendix A.3
We also make the assumption that the sparse G' and M remain invariant with different environment confounders the offline training and online testing.
Assumption 1 (Invariant causal graph). We denote the causal graph G under confounder as G(u). The generalization problem that we aim to solve satisfies the invariance in the causal graph G(u.). where Giuu)) G(u'), M (ue) M (u(), ue and u/ are the confounders in training and testing.
Remark 1. The assumption ! can also be interpreted task independence in [24]. invariant state representation 25 or invariant action effect in 26] See detailed comparison in appendix Table
3.2 Learning Causal Representation from Offline Data
We first learn the causal world model T(s'(s,, a) in the presence confounders u in the offline datasets. As formulated in ASC-MDP there are two sets of confounders: ", and ue. To estimate unconfounded transition model and remove the effect of confounder, we first remove the impact of u, which comes from the dynam- ics shift by estimating a batch-wise transition matrix M (u.), then apply reweighting for- mula to deconfound u induced by the behavior policies and mitigate the model objective mis- match.
Uncerteinly
start Planning
O
Break corretation %
goal
(c.))
Offine
Butfor
do
M
f(''s..)
Eg(s,a) Uncertainty Quartifier
p())
Figure 3: BECAUSE learns causality aware rep- resentation from the buffer and uses it in both the world model and uncertainty quantification ob- tain a pessimistic planning policy.
As discussed in Definition 3 we only need to optimize the part of the parameters the causal graph G, i.e. M. Thus, we can remove the straints in (4). then transform the original causal discovery problem into regularized MLE prob- lem as follows:
min Lmask (M) min (- p(D; 0, H. M) + AlM() W
M
min
(E(e.a.o)) (s)) (Bl
of (os-a(((ss)) as.
M|
(5)
M
World Model Learning
Sparsity Regularization
where K,, =
== Exes u(s' (pssssss((sp ((s')T T is an invertible matrix. The derivation of Equation is elaborated
in Appendix A.4 In practice, the 2-tstt for discrete state and action space and the fast

---

Conditional Independent Test (CIT) 27 for continuous variables to estimate each entry in the core matrix M. We regularize the sparsity of M by controlling the p-value threshold in CIT and provide more detailed implementation in Appendix C.1
Estimating the core mask provides more accurate relationship between state and action represen- tations, and we further refine the state action representation function & and to help capture more accurate transition dynamics. We optimize them by solving the following problem, according the transition model loss and spectral norm regularization [28] to satisfy the regularity constraints of the feature in Assumption
min Lrep (0, 11) min (s,D,))) ll (s)) (s.a) Apllol12 + Axllalle 6tft
(6)
The world model learning process is illustrated in Figure The estimation of individual Mu,)) mitigates the spurious correlation brought by uh- To further deal with the spurious correlation in Ur induced by the behavior policy (a(a)ssu7) we utilize the conditional independence property in the ASC-MDP shown in Equation (3). The following equation shows that the true unconfounded transition T' can be rewritten in the reweighting formulas. This reweighting process in mask M serves to estimate the treatment effect in the transition function T
as the soft intervention approach [8] [29 an unconfounded way:
()) = Ep. lo(s, a) M (uc)u(ss) TT un) Ep. ae s, Urr)
E [M (ue T )
= 0088
u(s') (s) M(u)u(s)..
p. T8 (als,
The derivation of Equation 7 illustrated in Appendix A.5 Equation 7 basically shows
a re- re-weighting process given the empirical estimation of (u)) in every batch of trajectories:
(u)) Epu (all), Compared the general reweighting strategies in previous MBRL To (ala,a,))
29 which reweights the entire value function, this re-weighting process is conducted only the estimated matrix, while the representation ((s, a) and pus') are subsequently regularized by weighted estimation of M. The pipeline of causal world model learning is described in the We discuss more details of the implementation and experiment in
literatures [8
first part of the Algorithm [ Appendix
3.3 Causal Representation for Uncertainty Quantification
To avoid entering OOD states in the online deployment, we further design pessimistic planner according the uncertainty of the predicted trajecto- ries in the imagination rollout step to mitigate objective mismatch.
Algorithm 1: BECAUSE Training and Planning
Input: Offline dataset D, causal discovery frequency k Output: Causal mask M, feature function Φ 7. policy // Causal world model learning
Mo [1]d'xd
for i E [K]do do
We the feature embedding from bilinear causal representation to help quantify the uncertainty, denoted as Es, As we have access to the offline dataset, learn Energy- based Model (EBM) (30 31 based
Update &n. Ph by (d, (1) in 6) if mod k: then
Update Mn, with Cmask (M) in 5
Weighted average M, with
// Uncertainty quantifier learning
on the abstracted state representation á and core matrix M. A higher out- put of the energy function E-(--) in- dicates higher uncertainty in the cur- rent state as they are visited by the be- havior policies Tg less frequently. In practice, the energy-based model usu- ally suffers from a high-dimensional data space (32). To mitigate this over- head of training good uncertainty
Fit Eg a) with LERM (8)
Initialize 1 = 0, (s) // Pessimistic planning
while h H do Estimate the uncertainty with score Es a)
Compute Q (s, a) with 99)
Vi = max, (s,,a)
f- arg max, ((s, a)
s',r env.step (a, g)
5

---

quantifier, we first embed the state
samples through the abstract representation pus'), and the state action pair via ((s, a).
(0) = () Balu(s (o,(s, a)] Eg(s,a) Ea ju(s (l((s, a)] + 2
(8)
where s((s* refers to the positive samples from the approximated transition dynamics T(-|s, a). and s((s refers to the latent negative samples via the Langevin dynamies [30] Additionally, regularize the parameters of EBM to avoid overfitting issues. We attach more training details and
results of EBMs in Appendix C.2 The learned energy function Eo a) is used to quantify the uncertainty based on the offline data.
During the online planning stage, we use the learned EBM to adjust the reward estimation based Model Predictive Control (MPC) [33. At timestep h, we basically subtract the original step return estimation ra (8, a) by its uncertainty a))
Qn a) = Qn(s.a) Eo(s, m,(.,)) Eo(s,a)
(9)
Adjusted Retur
s'ES
3.4 Theoretical Analysis of BECAUSE
Then we move on to develop the theoretical analysis for the proposed method BECAUSE. Based on two standard Assumption 2
and on the feature's existence and regularity, we achieve the finite-sample complexity guarantee an upper bound of the suboptimality gap follows, whose proof is postponed Appendix B
Theorem (Performance guarantee). Consider any O 8 < 1 and any initial state E S. Under the Assumption 2 3 ann that the transition model T is an SCM (defined in f) for any accuracy level O E E < with probability at least - 6, the output policy π of BECAUSE (Algorithm 7 based on the historical dataset D with n L m(8, a) samples generated from behavior policy satisfies:
| VIST C.SI M 10))
log(1/5) n(0n: an)
(3) (3) S min {Cr log
Σ Ex h=
| =
where C1, Cg are some universal constants, a SCM's s noise level (see Definition and M RRdxdd is the optimal ground truth sparse transition matrix to be estimated.
The error bound shrinks as the offline sample size n over all state-action pairs increase. It also grows proportionally to the planning horizon H. SCM's 's noise level a, and the tn norm of the ground true causal mask M, which describes the intrinsic complexity of the world model.
Consequently, with Proposition in the Appendix, we can achieve e-optimal policy (V1((3) < )) as long as the historical dataset satisfies the following conditions: A 1.
min { C log²
[M
[S] MMo 1²2 log
min En [n(Sk; an) 81 = (5.x,)))SAxx(1)
4 Experiment Results
In this section, we conduct a comprehensive em- pirical evaluation of BECAUSE's generalization performance in a diverse set of environments, covering different decision-making problems in the grid world, manipulation, and autonomous driving domains, shown in Figure 4
Crash
Figure Three environments used in this paper.
4.1 Experiment Setting
Environment Design We design 18 tasks in representative RI. environments in Figure Agents need to acquire reasoning capabilities to receive higher rewards and achieve goals.

---

105
Lift: Random
Lift-Medium
Eift
Unlock-Random
Unlock Random
Suctes Rada (h0)
00
60
ab
Talk [N]
62
60
6D
60-
40
40
20
64
0.5
Sample Ratio
Sparious Level
Unlock-Random
Unlock-Medium
Unlack Expert
Unlock-
100
Unlock-Medium
Saccess tuta 192
Suttea Sate
80
60
&
83
20
20
Sample Ratio
Spartou Levet
Crash-Random
Crash-Medium
3.0
Crash-Expeett
150.
Unlock-
102
Unloct-Expert
Suctee Sate 191
Sucoues baw
8
NOT0 BECAUSE
Expert
(a)
Sample
05 10 (c)
34
Figure 5: Results of BECAUSE and baselines in different tasks. (a) Average success rate in distri- bution and out of distribution. (b) Average success rate w.r.t. ratio of offline samples. (c) Average success rate w.r.t. spurious level in the environments. We evaluate the mean and standard deviation of the best performance among 10 random seeds and report task-wise results in Appendix Table 6
Lift: Object manipulation environment in RoboSuite 34]. We designed this environment for the agent lift an object with a specific color configuration on the table to desired height. In the OOD environment Lift-O, there is an injected spurious correlation between the color of the cube and the position of the cube in the training phase. During the testing phase, the correlation between color and position is different from training.
Unlock: We designed this environment for the agent to collect a key to open doors in Minigrid [35 In the OOD environment Unlock-O, O, there will be different number of goals (doors to be opened) in the testing environments from the training environments.
Crash: Safety is critical in autonomous driving, which is reflected by the collision avoidance capability. We consider risky scenario where an AV collides with a jaywalker because its view is blocked by another car [3). We design such a crash scenario based on highway-env [37], where the goal is to create crashes between a pedestrian and AVs. In the OOD environment Crush-O, the distribution of reward (number of pedestrians) different online testing environments.
For all three different environments, we set specific subset of the state space the goal g E S, and the reward is defined as the goal-reaching reward r(s, a, g) = 1(r = g). When the episode ends in the goal state within the task horizon H, the episode is considered success. We then use the average success rate as the general evaluation metries for our BECAUSE and all basclines.
In each environment, we collect three types of offline data: random, medium, and expert based on the different levels of u, in the behavior policies. In Unlock environments, we collect 200 episodes from each level of behavior policies the offline demonstration data, and the number of episodes is 1,000 the environments Lift and Crush, which all have continuous state and action space. A more detailed view of the environment hyperparameters and behavior policy design is in Appendix
Baselines We compare our proposed BECAUSE with several offline causal RL or MBRL baselines. ICIL (38] learns a Synamic-awar invariant causal representation learning to assist generalizable
policy learning from offline datasets. CCIL (9) conducts a soft intervention in our offline setting by jointly optimizing policy parameters and masks over the state. MnM [9] unifies the objective of jointly training the model and policy, which allocates larger weights the state prediction loss the high-reward region. Delphic A0 introduces delphic uncertainty to differentiate between uncertainties caused by hidden confounders and traditional epistemic and aleatoric uncertainties. TD3+BC [41 is an offline model-free RL approach that combines the Twin Delayed Deep Deterministic Policy
7

---

Gradient (TD3) algorithm with Behavior Cloning (BC) to adopt both the actor-critic framework and supervised learning from expert demonstrations. MOPO [2] is an offline MBRL approach that uses flat latent space and count-based uncertainty quantification to maintain conservatism in online deployment. GNN 42 is a GNN-based baseline using Relational Graph Convolutional Network to model the temporal dependency of state-action pairs in the dynamic model with message passing. CDL [24) uses causal discovery to learn a task-independent world model. Denoised MDP [12] and Iaccor [13] conduct causal state abstraction based on their controllability and reward relevance. The last three methods are designed for online settings, so we only implement their model learning objectives. We attach more details of the baseline implementation in Appendix C.6
4.2 Experiment Results Analysis
We empirically answer the following research questions.
How is the generalizability of BECAUSE in the online environments (which may be unseen)? Specifically, how does BECAUSE perform under diverse qualities of demonstration data (different level of uz)), and different environment contexts (different uz)?
RQI:
RQ2: How does the design in BECAUSE contribute to the robustness of its final performance under different sample sizes or spurious levels?
RQ3: How does BECAUSE scale up to visual RL tasks with image observation input compared to other visual RL baselines?
RQ4: How does BECAUSE achieve the aforementioned generaizabbiiy by mitigating the objective mismatch problem in offline MBRL?
For RQ1, in Figure 5 a), we evaluate the success rate in the online environment against different baselines. The result shows that under different environments and different qualities of behavior policies Ty (different U,,), BECAUSE consistently achieves the best performance in 8 out of 9 for both the in-distribution (I) and out-of-distribution (O) for all the demonstration data quality (different level of u.g). Where O here indicates the tasks under unseen environment with confounder u' F ue different from offline training. Another finding is that model-based approaches generally perform better than model-free approaches at various levels of offline data, which shows the importance of world model learning for generalizable offline RL. We attach the detailed results in the Appendix and the causal masks discovered in each environment in Appendix Figure & for reference.
Table 6 7
For RQ2, we compare different aspects of BECAUSE's s robustness with MOPO without causal structures 2 We compare their performance with different ratios of the entire offline dataset and illustrate the success rates Figure The result shows that, for any selected number of samples, BECAUSE consistently outperforms MOPO with clear margin. We also evaluate BECAUSE performance at higher spurious levels in Figure We add up to 8x of the original number of confounders in the environments to test the robustness of the agent's performance. BECAUSE consistently outperforms MOPO and the margin enlarges as the spurious level grows higher.
For RQ3, we conduct experiments with visual inputs in the Unlock environments with ICIL (38 and IFactor [13]. We param- three-layer CNN with 128 dimensions hidden size for all the baselines and our methods. The re- sults in Table show that BECAUSE can significantly improve both in-distribution and out-of-distributee performance under different quality of behavior policies.
Table Comparison of visual RL performance.
Tasks
ICII.
IFactor
BECAUSE
eterize the feature encoder
Unlock-T-random
0.8
4.3
15.7=a.
Unlock-O-random
1.5
4.7
5.9409
Unlock-I-medium
5.3120
30.214.1
62.0
Unlock- Enlock-O-medie O-medium
8,6+42
15.4+24
71.6=9.1
Unlock-I-exper
8.7+3.
34.0+48
63.7=3.9
Unlock-O-expert O-expert
17.1442
16.7+3.1
73.6
For RQ4, we aim to understand whether BECAUSE achieves higher performance by resolving the objective mismatch problem. We first collect two groups of trajectories: Tpos and Tneg. each with positive reward (success) and negative reward (failure) in Unlock task with sparse goal-reaching reward. We want to have a model whose loss is informative for discriminating control results, that is, we wish Lmodet (Toas) (Tpoa) < +mmodet(aan (Tneg). According to our visualization in Figure 6 in Enlock-Experi and Unlock-Mert Medium,, the ratio of Tpos is much higher in BECAUSE than MOPO among the trajectories with low model loss. In Snlock-Random the mismatch of the model and
8

---

Unndckrandom
Unack-macium
Unlock export
Posttive Reward
Postiive Reward
Postive Reward
Negative Reward
Number of Sammpe
Negatiee Reward
Negative Reward
Negative Reward
Postiive Reward
Postiive Reward
Positive Reward
Mostive Heward
Number of Sample
Negullve Reward
Negalive Reward
Negative Rewward
BEGAUSE
0.000
0.502
0.004
0.006
0.000
0.002 0,004
0.000
0. 0.002
0.004
D 0.006
Model loss
Model loss
Model loss
Example
Figure 6: Evaluation of the difference between the distribution of episodic model loss for success and failure trajectories. The higher difference indicates a reduction in model mismatch issues. An example of failure mode is trying to open the door without having the key.
control objective is more significant, since the demonstration is poor in state coverage. MOPO cannot succeed even when the model loss is low, whereas our methods can. We perform hypothesis test with Ho modddl Topss) < ( In BECAUSE, this desired property is more significant than MOPO attributed to the causal representation we learn, indicating a reduction of objective mismatch. We attach detailed discussions for the mismatch evaluation in Appendix C.3 and Table 5
4.3 Ablation Studies
We conducted ablation studies with three variants of BECAUSE and report the aver- age success rate across nine in-distribution and nine -utoof-distribution tasks in Ta- ble 2 The Optimism variant conducts optimistic planning instead of pessimistic planning in Equation (9), which uses uni- form sampling in the planner module. The Linear variant assumes a full connection to the causal matrix M, then directly uses linear MDP to parameterize the dynamics model T, which removes the causal discovery module in BECAUSE. The Full variant learns from the full batch of data to estimate the causal mask with- out iterative update. We report the results of the task-wise ablation with confidence interval and significance in Appendix Table 8 and 9
Table 2: The ablation studies between BECAUSE and its variants. We report the overall Success rate (%) over 9 in-distribution (I) and 9 disti--distriuution (0) tasks, respectively. Bold is the best.
Variants
BECAUSE Optimism Linear
Full
Overall-I Overall-O
73.3 4.5
64.4
57.916.1
39.3 33.2+52
4..0
32.413.3
25.2:33
Related Works
Objective Mismatch in MBRL The objective mismatch in MBRL 43 4
refers to the fact that
pure MIE estimation of the world model does not align well with the control objective. Previous propose reweighting during model training to alleviate this mismatch, (46) proposes a awal prediction by redistributing model error according to their task relevance. These works essentially reweight loss for the entire model training, while our work conducts reweighting just over the estimated causal mask more efficiently. More recently, [9] 10 proposed joint training between the world model and policies. Although joint optimization improves performance, they do not address the generalizability of the learned model under the distribution shift setting. In the offline 47 employs model ensemble, pessimistic policy optimization or value iteration 48 49 and an energy-based model for planning [50] to quantify uncertainty and improve test performance, To the best of our knowledge, no previous work explored or modeled the impact of distribution shift on the objective mismatch problem in MBRL.
works 29 45)
setting, Model-based RL 2

---

Causal Discovery with Confounder Most of the existing causal discovery methods [5]]
can
be categorized into constraint-based and score-based. Constraint-based methods [52 start from
This type of method is highly data-efficient but not robust to noisy data. As a remedy, score-based methods [55] (56] use metrics such as the likelihood or BIC [57] as scores to manipulate edges in the causal graph. Recently, researchers have extended score-based methods with RL [58], order learning [59] differentiable discovery [22 60 61]. To alleviate the non-identifiability under hidden confounders, active intervention methods have been explored [62], aiming to break spurious correlations in an online fashion. With extra assumptions on confounders, some recent works detect so that models can effectively identify elusive confounders.
complete graph and iteratively remove edges with statistical hypothesis testing 53 54
such correlations
63
64
65
Causal Reinforcement Learning Recently, many RL algorithms have incorporated causality to improve reasoning capability [66] and generalizability. For instance, [67] and [68] explicitly estimate causal structures with the interventional data obtained from the environment in online setting. These structures can be used to constrain the output space
[19] or to adjust the buffer
77 71 based on causal graphs priority [69] Building dynamic models in model-based RL is widely studied. Most existing causal MBRL works focus on estimating the causal world model by predicting transition dynamics and rewards. Existing methods learn this causal world model
[24
via structural regularization 2 72 733
conditional independence test [24] 70 74. variational
inference [75 [2], counterfactual data augmentation 76 771. hierarchical skill abstraction [78] (9].
uncertainty quantification [40], reward redistribution [24 80 causal context modeling [8]] 82 and
structure-aware state abstraction [2
13
83 (84) based on the controllability and task or reward
relevance. However, the presence of confounders during data collection can skew the learned policy, making it susceptible to spurious correlations. Deconfounding solutions have been proposed either between actions and states [39 85 86 or among different dimensions of state variables [19 87
6 Conclusion
In this paper, we study how to mitigate the objective mismatch problem MBRL, especially under the offline settings where distribution shift occurs. We first propose ASC-MDP and the bilinear causal representation associated with it. Based on the formulation, we proposed how to learn this causal abstraction by alternating between causal mask learning and feature learning in fitting the world dynamics. In the planning stage, we applied the learned causal representation to an uncertainty quantification module based on EBM, which improves the robustness under uncertainty in the online planning stage. We theoretically justify BECAUSE's sub-optimality bound induced by the sparse matrix estimation problem and offline RL. Comprehensive experiments 18 different tasks show that given a diverse level of demonstration as the offline dataset, BECAUSE has better generalizability than baselines in different online environments, and it robustly outperforms baselines under different spurious levels or sample sizes. We empirically show that BECAUSE mitigates the objective mismatch with causal awareness learned from offline data. One limitation of BECAUSE lies in its simplified assumption of him--hooognneous causal structure, which may not always hold in long-horizon or non-stationary settings. Besides, the current implementation is still based on vector observations. It will be interesting to scale up the causal reasoning framework into nigh-dimensional observations to discover concept factors in long- morizoo visual RL settings.
Acknowledgement
This work is partially supported by the Defense Advanced Research Projects Agency (DARPA) under Contract No. HR00112220012 The work of L. Shi is supported in part by the Resnick Institute and Computing, Data, and Society Postdoctoral Fellowship at California Institute of Technology.
10

---

References
[1] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643 2020.
[2] Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. Advances in Neural Information Processing Systems, 33: 11119-11122,,, 2020.
[3] Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model-based offline reinforcement learning. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 21810-218 10822. Curran Associates, Inc., 2020.
[4] Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offline model-based policy optimization. Advances in neural information processing systems, 14:28994422896 2021.
[5] Zeyu Zhu and Huijing Zhao. A survey of deep and for autonomous driving policy learning. IEEE Transactions on Intelligent Transportation Systems, 2309):14033-4465, 2021.
[6] Aviral Kumar, Anikait Singh, Frederik Ebert, Mitsuhiko Nakamoto, Yanlai Yang, Chelsea Finn, and Sergey Levine. Pre-training for robots: Offline enables learning new tasks from a handful of trials. arXiv preprint arXiv:2210.05172 2022.
[7] Shengpu Tang and Jenna Wiens. Model selection for offline reinforcement learning: Practical considerations for healthcare settings. In Machine Learning for Healthcare Conference, pages 2-35. PMLR, 2021.
[8] Laixi Shi and Yuejie Chi. Distributionally robust model-based offline reinforcement learning with near-optimal sample complexity. arXiv preprint 2xXiv:2208.05767 2022.
[9] Benjamin Eysenbach, Alexander Khazatsky, Sergey Levine, and Russ R Salakhutdinov. Mis- matched no more: Joint model-policy optimization for model-based rl. Advances in Neural Information Processing Systems, 35:23230-233 222233,3-22,,...0 2022.
[10] Shentao Yang, Shujian Zhang, Yihao Feng, and Mingyuan Zhou. A unified framework for alter- nating offline model training and policy learning. Advances in Neural Information Processing Systems,
17232. 2022.
[11] Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Elements of causal inference: founda- lions and learning algorithms. The MIT Press, 2017.
[12] Tongzhou Wang, Simon S Du, Antonio Torralba, Phillip Isola, Amy Zhang, and Yuandong Tian. Denoised mdps: Learning world models better than the world itself. arXiv preprint wwXiv:2206.15477, 2022.
[13] Yu-Ren Liu, Biwei Huang, Zhengmao Zhu, Honglong Tian, Mingming Gong, Yang Yu, and Kun Zhang. Learning world models with identifiable factorization. arXiv preprint arXiv:2306.06566 2023.
[14] Lin Yang and Mengdi Wang. Reinforcement learning feature space: Matrix bandit, kernels, and regret bound. In International Conference Machine Learning, pages 10746-10756. PMLR, 2020.
[15] Weitong Zhang, Jiafan He, Dongruo Zhou, Gu, and A Zhang. Provably efficient representation selection in low-rank markov decision processes: from online offline rl. In Uncertainty in Artificial Intelligence, pages 2488 2497. PMLR, 2023.
[16] Fan Feng and Sara Magliacane. Learning dynamic attribute-factored world models for efficient multi-object reinforcement learning. arXiv preprint arXiv:2307.09205. 2023.
[17] Angela Zhou. dewarrdrelevaace--ileee linear offline reinforcement learning. arXiv preprint 1XX1:22401.12934, 2024.

---

[18] Lingxiao Wang, Zhuoran Yang, and Zhaoran Wang. Provably efficient causal reinforcement learning with confounded observational data. Advances in Neural Information Processing Systems, 34:21 (64-21155 2021.
[19] Wenhao Ding, Laixi Shi, Yuejie Chi, and Ding Zhao. Seeing is not believing: Robust reinforce- ment learning against spurious correlation. arXiv preprint arXiv:2307. 37700707777... 2023.
[20] Junzhe Zhang and Elias Bareinboim. Markov decision processes with unobserved confounders: A causal approach. Technical report, Technical report, Technical Report R-23, Purdue Al Lab,
2016.
[21] Karren Yang, Abigail Katcoff, and Caroline Uhler. Characterizing and learning equivalence classes of causal dags under interventions. In International Conference on Machine Learning, pages 5541-5550 PMLR, 2018.
[22] Philippe Brouillard, Sébastien Lachapelle, Alexandre Lacoste, Simon Lacoste-Julien, and Alexandre Drouin. Differentiable causal discovery from interventional data. Advances in Neural Information Processing Systems, 33:21865-21877 2020.
[23] Zizhao Wang, Xuesu Xiao, Yuke Zhu, and Peter Stone. Task-independent causal state abstraction. Workshop on Robot Learning: Sell-Supervised and Lifelong Leurning, NeurIPS, 2021.
[24] Zizhao Wang, Xuesu Xiao, Zifan Xu, Yuke Zhu, and Peter Stone. Causal dynamics learning for task muskinndependent state abstraction. arXiv preprint arXiv:2206. 13452, 2022.
[25] Amy Zhang, Clare Lyle, Shagun Sodhani, Angelos Filos, Marta Kwiatkowska, Joelle Pineau, Yarin Gal, and Doina Precup. Invariant causal prediction for block mdps. In International Conference on Machine Learning, pages 11214-1 1224. PMLR, 2020.
[26] Wenxuan Zhu, Chao Yu, and Qiang Zhang. Causal deep reinforcement learning using observa- tional data. arXiv preprint 2.tii:22/111533555 2022.
[27] Krzysztof Chalupka, Pietro Perona, and Frederick Eberhardt. Fast conditional independence test for vector variables with large sample sizes. arXiv preprint arXiv: 0274477.7777 2018.
[28] Yuichi Yoshida and Takeru Miyato. Spectral norm regularization for improving the generaliz- ability of deep learning. arXiv preprint arXiv: /705/ 2017.
[29] Nathan Lambert, Brandon Amos, Omry Yadan, and Roberto Calandra. Objective mismatch in model-based reinforcement learning. arXiv preprint arXiv:2002.04523, 2020.
[30] Yilun Du and Igor Mordatch. Implicit generation and modeling with energy based models. Advances in Neural Information Processing Systems, 32, 2019.
[31] Yezhen Wang, Bo Li, Tong Che, Kaiyang Zhou, Ziwei Liu, and Dongsheng Li. Energy-based open-world uncertainty modeling for confidence calibration. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 9302-933 9311, 2021.
[32] Bo Pang. Tian Han, Erik Nijkamp, Song-Chun Zhu, and Ying Nian Wu. Learning latent space energy-based prior model. Advances Neural Information Processing Systems, 13:21994- 22008, 2020.
[33] Eduardo F Camacho and Carlos Bordons Alba. Model predictive control. Springer science & business media, 2013.
[34] Yuke Zhu, Josiah Wong, Ajay Mandlekar, Roberto Martín-Martón Abhishek Joshi, Soroush Nasiriany, and Yifeng Zhu. robosuite: A modular simulation framework and benchmark for robot learning. In arXiv preprint arXiv:2009. 12299,, 2020.
[35] Maxime Chesalier-Boisvertt, Lucas Willems, and Suman Pal. Minimalistic gridworld environ- 2018.
ment for openai gym. https: /ttps://githule hwtp:///ww..iineegggmmmmniaaaaaa
[36] Zenna Tavares, James Koppel, Xin Zhang, Ria Das, and Armando Solar-Lezama. A language for counterfactual generative models. In International Conference on Machine Learning, pages 10173-10182. PMLR, 2021.
12

---

[37] Edouard Leurent. An environment for autonomous driving decision decision-making https://githtn com/ wwleurenthhighwwayeeee.e env 2018.
[38] Ioana Bica, Daniel Jarrett, and Mihaela van der Schaar. Invariant causal imitation learning for generalizable policies. Advances in Neural Information Processing Systems, 34:3952-3964, 2021.
[39] Pim De Haan, Dinesh Jayaraman, and Sergey Levine. Causal confusion in imitation learning. Advances in Neural Information Processing Systems, 32, 2019.
[40] Alizée Pace, Hugo Yèche, Bernhard Schölkopf, Gunnar Ratsch, and Guy Tennenholtz. Delphic offline reinforcement learning under nonidentifiable hidden confounding. In The Twelfth International Conference on Learning Representations, 2023.
[41] Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. Advances in neural information processing systems,
2021.
[42] Michael Schlichtkrull, Thomas N Kipf, Peter Bloem, Rianne van den Berg, Ivan Titov, and Max Welling. Modeling relational data with graph convolutional networks. In European semantic web conference, pages 593-607. Springer, 2018.
[43]
Amir-massoud Farahmand, Andre Barreto, and Daniel Nikovski. Value-aware loss function for model-based reinforcement learning. In Artificial Intelligence and Statistics, pages 1486-1494.
PMLR, 2017.
[44] Yuping Luo, Huazhe Xu, Yuanzhi Li, Yuandong Tian, Trevor Darrell, and Tengyu Ma. Algo- rithmic framework for model-based deep reinforcement learning with theoretical guarantees. arXiv preprint arXiv: 00855, 2018
[45] Pierluca D' O'Oro, Alberto Maria Metelli, Andrea Tirinzoni, Matteo Papini, and Marcello Restelli. Gradient-aware model-based policy search. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 3801-3800 2020.
[46] Suraj Nair, Silvio Savarese, and Chelsea Finn. Goal-aware prediction: Learning to model what matters. In International Conference on Machine Learning, pages 7207-7219, PMLR, 2020.
[47] Yihao Sun, Jiaji Zhang, Chengxing Jia, Haoxin Lin, Junyin Ye, and Yang Yu. Model-bellman inconsistency for model-based offline reinforcement learning. In International Conference on Machine Learning, pages 33177-33194 PMLR, 2023.
[48] Chi Jin, Zhuoran Yang. Zhaoran Wang, and Michael Jordan. Provably efficient reinforcement learning with linear function approximation. In Conference on Learning Theory, pages 2137- 2143. PMLR, 2020.
[49] Masatoshi Uehara and Wen Sun. Pessimistic model-based offline reinforcement learning under partial coverage. arXiv preprini arXiv:2107.06626. 2021.
[50] Yilun Du, Shuang Li, Joshua Tenenbaum, and Igor Mordatch. Learning iterative reasoning through energy minimization. In International Conference on Machine Learning, pages 5570- 5582. PMLR, 2022.
[51] Clark Glymour, Kun Zhang, and Peter Spirtes. Review of causal discovery methods based on graphical models. Frontiers in genetics, 10:524, 2019.
[52] Peter Spirtes, Clark N Glymour, Richard Scheines, and David Heckerman. Causation, prediction, and search. MIT press, 2000.
[53] Karl Pearson. X. on the criterion that a given system of deviations from the probable in the case of correlated system of variables is such that it can be reasonably supposed have arisen from random sampling. The London, Edinburgh, and Dublin Philosophical Magazine and Journal of Science, 50(302):157
1900.
[54] Kun Zhang, Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Kernel-based conditional independence test and application in causal discovery. arXiv preprint arXiv: 2012.
13

---

[55] David Maxwell Chickering. Optimal structure identification with greedy search, Journal of machine learning research, 3(Nov):507-554, 2002.
[56] Alain Hauser and Peter Bühlmann. Characterization and greedy learning of interventional markov equivalence classes of directed acyelic graphs. The Journal of Muchine Learning Research, 13(1):2409-2466) 2012.
[57] Andrew A Neath and Joseph E Cavanaugh. The bayesian information criterion: background, derivation, and applications. Wiley Interdisciplinary Reviews: Computational Statistics, 4(2):199-22) 2012.
[58] Shengyu Zhu, Ignavier Ng. and Zhitang Chen. Causal discovery with reinforcement learning. arXiv preprint 1eXiv:1906.04477 2019.
[59] Dezhi Yang, Guoxian Yu, Jun Wang, Zhengtian Wu, and Maozu Guo. Reinforcement causal structure learning on order graph. In Proceedings of the AAAI Conference on Artificial Intelli- gence, volume 37, pages 10737-10744, 2023.
[60] Nan Rosemary Ke, Olexa Bilaniuk, Anirudh Goyal, Stefan Bauer, Hugo Larochelle, Bernhard Schölkopf, Michael C Mozer, Chris Pal, and Yoshua Bengio. Learning neural causal models from unknown interventions. arXiv preprint arXiv:1910.01077 2019.
[61] Yunzhu Li, Antonio Torralba, Anima Anandkumar, Dieter Fox, and Animesh Garg. Causal discovery in physical systems from videos. Advances in Neural Information Processing Systems, 33:9180-9192, 2020.
[62] Nino Scherrer, Olexa Bilaniuk, Yashas Annadani, Anirudh Goyal, Patrick Schwab, Bernhard Schölkopf, Michael C Mozer, Yoshua Bengio, Stefan Bauer, and Nan Rosemary Ke. Learning neural causal models with active interventions. arXiv preprint arXiv:2109.00229, 2021.
[63] Christopher Clark, Mark Yatskar, and Luke Zettlemoyer. Don't ton't take the easy way out: Ensemble- based methods for avoiding known dataset biases. arXiv preprint 1exxi:/9099033600 2019.
[64] Divyansh Kaushik, Eduard Hovy, and Zachary C Lipton. Learning the difference that makes difference with counterfactually-aggmeenn data. arXiv preprint arXiv:1909.122 2019.
[65] Meike Nauta, Ricky Walsh, Adam Dubowski, and Christin Seifert. Uncovering and correcting shortcut learning in machine learning models for skin cancer diagnosis. Diagnostics, 12(1):40, 2021.
[66] Prashan Madumal, Tim Miller, Liz Sonenberg, and Frank Vetere. Explainable reinforcement learning through a causal lens. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 2493-2500, 2020.
[67] Suraj Nair, Yuke Zhu, Silvio Savarese, and Li Fei-Feii Causal induction from visual observations for goal directed tasks. arXiv preprint arXiv:1910.01755 2019.
[68] Sergei Volodin, Nevan Wichers, and Jeremy Nixon. Resolving spurious correlations in causal models of environments via interventions. arXiv preprint arXiv:2002.05217, 2020.
[69] Maximilian Seitzer, Bernhard Schölkopf, and Georg Martius. Causal influence detection for improving efficiency in reinforcement learning. Advances in Neural Information Processing Systems, 34, 2021.
[70] Zheng-Mao Zhu, Xiong-Hi Hui Chen, Long Tian, Kun Zhang, and Yang Yu. Offline reinforcement learning with causal structured world models. arXiv preprint 4:144:2006.01474. 2022.
[71] Mirco Mutti, Riccardo De Santi, Emanuele Rossi, Juan Felipe Calderon, Michael Bronstein, and Marcello Restelli. Provably efficient causal model- reinforcement learning for systematic generalization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pages 9251-9259) 2023.
14

---

[72] Fan Feng, Biwei Huang, Kun Zhang, and Sara Magliacane. Factored adaptation for non- stationary reinforcement learning. Advances Neural Information Processing Systems, 35:31957-31971, 2022.
[73] 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-aware entropy regularization. arXiv preprini arXiv:2402. /4528, 2024.
[74] Jiaheng Hu, Zizhao Wang, Peter Stone, and Roberto Martin-Martin. Elden: Exploration via local dependencies. arXiv preprint arXiv:2310.08700 2023.
[75] Wenhao Ding, Haohong Lin, Bo Li, and Ding Zhao. Generalizing goal-conditioned reinforce- ment learning with variational causal reasoning. Advances in Neural Information Processing Systems, 35:26532-26548, 2022.
[76] Lars Buesing, Theophane Weber, Yori Zwols, Nicolas Heess, Sebastien Racaniere, Arthur Guez, and Bean-Baptistte Lespiau. Woulda, coulda, shoulda: Counterfactually-guiie guidee policy search. In International Conference on Learning Representations, 2018.
[77] Silviu Pitis, Elliot Creager, Ajay Mandlekar, and Animesh Garg. Mocoda: Model-basee counterfactual data augmentation. arXiv preprint arXiv:2210. 11287, 2022.
[78] Tabitha E Lee, Jialiang Alan Zhao, Amrita Sawhney, Siddharth Girdhar, and Oliver Kroemer. Causal reasoning in simulation for structure and transfer learning of robot manipulation policies. In 2021 IEEE International Conference on Robotics and Automation (ICRA), pages 4776-4782. IEEE, 2021.
[79] Tabitha Edith Lee, Shivam Vats, Siddharth Girdhar, and Oliver Kroemer. Scale: Causal learning and discovery of robot manipulation skills using simulation. In Conference on Robot Learning. pages 2229 2229-2256. PMLR, 2023.
[80] Yudi Zhang, Yali Du, Biwei Huang, Ziyan Wang, Jun Wang. Meng Fang, and Mykola Pech- enizkiy. Interpretable reward redistribution in reinforcement learning: A causal approach. Advances in Neural Information Processing Systems, 36, 2024.
[81] 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, 2021.
[82] Peide Huang, Xilun Zhang, Ziang Cao, Shiqi Liu, Mengdi Xu, Wenhao Ding, Jonathan Francis, Bingging Chen, and Ding Zhao. What went wrong? closing the sim-to-real gap via differentiable causal discovery. In Conference on Robot Learning, pages 734-760. PMLR, 2023.
[83] Xiang Fu, Ge Yang, Pulkit Agrawal, and Tommi Jaakkola. Learning task informed abstractions. In International Conference on Machine Learning, pages 3480-3491. PMLR, 2021.
[84] Biwei Huang, Chaochao Lu, Liu Leqi, José Miguel Hernández-Lobate Clark Glymour, Bern- hard Schölkopf, and Kun Zhang. Action-sufficien state representation learning for control with structural constraints. In International Conference Machine Learning, pages 9260-9279 PMLR, 2022.
[85] Zhihong Deng, Zuyue Fu, Lingxiao Wang, Zhuoran Yang, Chenjia Bai, Zhaoran Wang, and Jing Jiang. Score: Spurious correlation reduction for offline reinforcement learning. arXiv preprint arXiv:2110. 12468, 2021.
[86] Maxime Gasse, Damien Grasset, Guillaume Gaudron, and Pierre-Y Oudeyer. Using con- founded data in latent model-based reinforcement learning. Transactions on Machine Learning Research, 2023.
[87] Junzhe Zhang, Daniel Kumor, and Elias Bareinboim. Causal imitation learning with unobserved confounders. Advances in neural information processing systems, 15:12263-12274 2020.
[88] Ying Jin, Zhuoran Yang, and Zhaoran Wang. pessimism provably efficient for offline rI? In International Conference Machine Learning, pages 5084-5096. PMLR, 2021.
15

---

[89] Thomas Blumensath and Mike E Davies. Iterative hard thresholding for compressed sensing. Applied and computational harmonic analysis, 27(3):265-22 274, 2009.
[90] Amir Beck and Marc Teboulle. A fast iterative shrinkage-thrrshoodinng algorithm for linear inverse problems. SIAM journal on imaging sciences, 2(1):183-202, 2009.
[91] Peter Auer, Thomas Jaksch, and Ronald Ortner. Near-optimal regret bounds for reinforcement learning. Advances in neural information processing systems, 21, 2008.
[92] Ryan Tibshirani and Larry Wasserman. Sparsity, the lasso, and friends. Lecture notes from "Statistical Machine Learning," " Carnegie Mellon University, Spring, 2017.
[93] Patrick E McKnight and Julius Najab. Mann-whitne test. The Corsini encyclopedia of psychology, pages 1-1, 2010.
16

---

A Auxiliary Details of BECAUSE Framework
A.1 Notation Summary
We illustrate all the notations used in the main paper and appendix in Table
Table 3: Notations used in this paper and their corresponding meanings.
Notation
Explanation
A.
Action space, action
S, 8
State space, state
Reward Discount factor
Transition dynamics
h, H
Timestep h, Horizon H
Optimal policy in the online environments
Behavior policy generating the offline datasets
V())
State value function Action-state value function
((:))
D
Offline datasets
n(s, a) 18,
Number of samples for (s, a) pairs in offline datasets
Bellman operator
((.-)
Feature representation of state and action
((.))
Feature state, action and confounder in Equation
((())
Feature representation of next state
d, d' Kμ
Dimensions of features for ((;--) and p(-)
Feature matrix expanded from f in Equation
Co. C a C2 K X
CB
Feature regularity
Sparsity-relater constant
Restrictive eigenvalue constant
Feature Kronecker product
the Ur M, M())
Confounders
Binary transition matrix (under certain confounders)
G
Causal graph
M. M())
Estimated causal matrix
3M BM
Optimal / estimated parameters causal matrix
PA GAG()
Parental node the causal graph G Exogenous noise in SCM by Definition Spectrum regularizer weight in Equation Standard deviation of exogenous noise in SCM
X, As. A
g
Accuracy level of the policy
K
Iterative update steps in BECAUSE
Eg
Energy-based model
8
Level of 'high probability' Uncertainty quantification function 8-uncertainty quantifier set Regularizer weight for the l norm in EBM
I(...) 3
AFRM
A.2 Equivalence of the Assumptions with Previous Works
A.3 Derivation of Definition3]
The node of this causal graph G Oddx M
contains two groups of entities: (1) The state
Odxd'
gdxd
action abstraction d(s, a), and (ii) the next state abstraction u(s'). We denote o(.,)) (0) the tth
17

---

Table 4: Summary of assumptions and equivalent forms.
Assumption
Base and Equivalent Format
T(s'ss, ((6))))(()) pe pe(u(s,,u)
Invariant State Abstraction 25
Task Independence 24
Tonline (s' a) = Σ,, (s''ss u)f(uss)
(s's
Invariant Action Effect 26
T Tottine (s 6'|s,a)
Tontine (8 (5((8,)) X Tootine ( (s|(s.0)
ul)) * (00, M (2)
Invariant Causal Graph
T(s' [ss' = T(s'|s, (2)
factor in the abstracted state action representations. and u(-) for the jth factor in the abstracted state representations.
The source node of all the nodes is o(s, a) which is the abstracted state-action representation, and the sink node of all the edges G is the to ((s))(())
Odxx M
[((s,a)] p(ss)
T(s'(s, [)) (()'7]
Odxd'
gaxd
=
(10)
Therefore, G is a bipartite graph, since there will be no edges between ((s.a) (65 or (())(),,
Consequently, show that G DAG.
A.4 Derivation of Equation
Definition (Structured Causal Model). An SCM = (8, consists of collection of functions (II.
Sj = (PA (PA)(()))))))) (s)), jE [d],
(11)
where PA C (s1,...,3a) [5;} are called parents of æy in the Directed Acyclic Graph (DAG) G, and E Jd are jointly independent. For instance, in continuous state and action space, parameterize the world model with joint Gaussian Distribution, i.c. e ~ o( Laa)).
We then use bilinear MDP to approximate the original likelihood function in Equation i.e.
p(DiO,u, & II exp(-Ila (s')K"ss - (s
(12)
(s.a
where K₁, == Eves u1s' ((s'(T an invertible matrix. Then we can apply an MLE in Equation 5
In our BECAUSE algorithm, the optimization of the causal world model is conducted by solving the regularized MLE problem in Equation (13). The biggest difference between BECAUSE and the offline version of [14 [5] is that it aims to apply to regression instead of ridge regression to estimate
matrix M:
M, arg max [log pl D: 2. A|MI
M
= arg min
Σ
1| (s) (s,a)M12
(13)
M
Sparsity Regularization
World Model Learning
18

---

A.5 Proof of Equation
The derivation depends on the following re-weighting formula in
as) Ep. ( (s'1ss, u ux
(14)
P. TT 8, "r)
Then apply equation 14 to the decomposition in Equation and Equation 3
which yields
T()ss,a) = Bpn. a ( U.x))
pp ग TT ux))
Ep.. [ō(s, a T M ( IT ( 8, uz)]
Ep.. aT 8, ut
(15)
Epm. [ (T (als, un)]
o(s, a)²
m(s')
Ep. (Ta ux)
o(s a)² M())u(d');
where the last equality holds by letting M (u)
Bp, (M(e))saa au+))
Ep. Fala
B Proof of Theorem 1
In this section, we provide the proof of the sub-optimality upper bound in Theorem [ We first show some useful definitions and lemmas in Section B.1 Armed with them, we provide the theoretical results tailored for the causal discovery setting Section B.2 Furthermore, we give a detailed proof of the uncertainty set form in our causal discovery problems in Section B.3
B.1 Preliminary
In this subsection, first define the 5-uncertainty quantifier Г, then refer the lemmas in the previous literature to construct suboptimality bound based on the defined uncertainty quantifier r.
First, we define the Bellman operator B,, for some value function V S + R, the Bellman operator can be defined as:
(Br V) a) = (s,, an) + ((s++1)(ss, = 8, a a).
(16)
Similarly, denote the approximate Bellman operator of the empirical MDP constructed from the offline dataset as B, for any [
Definition 5 (Ó- (o-Uncertainty Quantifier). We let (T I' : S A R be d-uncertainty quantifier with respect data distribution PD if the event:
E E= 8 (1 D (s, (B, (BA & Wo satisfies (E)
8.
As we consider the offline model learning and planning, we define the model evaluation error at each step h E [H] as
((, Axx:: (x a) = ()))) (s,
(17)
where i, is the error induced by the approximate Bellman operator, especially the transition ker- nel based on D. We then identify the source of sub-optimality in our offline MBRL setting by decomposing the ootimmalilt eлror in Lemma
Lemma (Decomposition of Suboptimality 888.
H
S: (s()s) (8)
E (sN anr))ss + M=h
Σ
E [ew Js:) =s)
W=
Ex (sne, 1 (-,ss) #(, sw )) Alsh =s],
(18)
19

---

where x is any learned policy, x² is the optimal policy that maximizes the cumulative return as below:
arg max ET
Σ h1=1
r(sw [=)(sn
Based on this decomposition, we will get the basic form of sub-optimality error bound for general offiine RL settings in Lemma
Lemma 2 (Suboptimality in standard MDP 88 Suppose have (Tr) δ-uncertainty
quantifier. Under & defined in Equation ( the suboptimality error bound by conservative planning satisfies:
H
E: - (2
2 E. (SW, s).
N/=h
The basic form of sub-optimality bound Lemma 2 involves an uncertainty quantifier I's, which in our case will be further replaced by an exact bound in our sparse matrix estimation problem of causal discovery algorithms.
B.2 Proof of Theorem 1
The main results hold under the following two assumptions:
Assumption (Existence of core matrix given the feature embedding). For each (s, a) S A. feature vectors (ss, a) E IR p(s) Idd' are approximated as priori. Given specific confounder set u, there exists an unknown matrix Mu))* E Idd xd such that,
T(s'ss, a, u) = o(s a) M(u)a(s').
(19)
Assumption (Feature regularity). We assume feature regularity 14 15 nents of the confounded bilinear MDP:
for the following compo-
(w))l < Cmd,
V(s,a) S x A, ())(s,a)l13 & Cad.
vs ERSS || 1 - C₂ || Io,c, |aK 12,00 S Cy.
Vs,a,s' E x A x S, ( ()) S C...
where CM, Cs. Cu. C² are some universal constants.
Here, for any matrix X, |X 12, max; E3 X represents the operator 2 + oo norm.
Proof pipeline. Armed with the above assumptions, we turn to the bilinear MDP setting, which this work focuses on. We shall develop the finite-sample analysis by specifying the main error term 5-uncertainty quantifier I (see Lemma for our hom-hhomogeneous core matrix estimation problem in the following lemma.
Lemma (Uncertainty bound for Bilinear Causal Representation). Under the Assumption 2 3 and
that is an SCM (defined in f), for the BECAUSE algorithm, for the E-optimal policy (V; (8) VT (5) < {). 4 S 1, have the b-uncertainty set as:
EBECAUSE = [BB V1+1)(s, (Bn (BAPh+))ss,a))
[ WIS,
log(1/0) n(s,a)
E min (C1 log
where C₁ is some universal constants.
20

---

Armed with the above lemma, we complete the proof of Theorem by showing that
VT - (3) 2
Σ
En (sW,
aN)ss1 =
k==1
H
[l VISI, //1MM
log(1/6) n(sh, an.
2
Ex
min {C1 log
|s
h=1
(20)
This concludes the proof Theorem
B.3 Proof of Lemma 3
The key to proving Theorem is to prove Lemma 3 The proof pipeline of Lemma is illustrated below. In Step 1, we derive the estimation of the causal transition matrix M in BECAUSE sparsity regression problem. Step 2. we decompose the error terms within 8 6-uncertainty set into two parts: (a) eπrоr due to the under-explored dataset, (b) error due optimization error in the structured causal model. Then we bound both error terms in Step and Step 4, respectively. Finally, in Step 5, sum up all the results and derive the form of 5-uncertainty quantifier which will lead to our final results Theorem
Step 1: deriving the output model of BECAUSE.
Recalling the original optimization problem in
equation 133 to estimate the core matrix:
M arg max[log p(o, n. M, - A|M
arg min
Σ Im( ((8,a)()))))))) -M +
(21)
(..0,8))
Sparsity Regularization
Model Learning
This part of derivation aims transform the above estimation problem into linear regression problem, with the regression data pairs (X.T) and some unknown parameters 3 associated with mask . be estimated. Eventually, we'll derive the representation of each part of B X, T, and eventually reach the following form:
min BM
Σ (ELOD
I1 (s/ s;,a;) Xx,3M13 3 [lo].
(22)
We define each component of this target form of lo regression follows:
For unknown parameters 3M; We first define 8M [0, 1]dd" column dimensional vector consisting of all the entries in time homogenous causal matrix M, where 3.M denotes the i-th entry of 3M, Besides, we define 3. as the true core matrix given some offline dataset and corresponding data pairs Teate. Xdata that satisfies Tdota P X c
For dataset D: Recall the transition pairs in the offline dataset D {s;, at sthsisn- Here, represents the sample size over certain state-action pairs in the rollout data by some behavior policy Tg. For simplicity, we denote n & n(ss, a) in the following derivation, which is mentioned in Section 2.1
For regression target T Then, we introduce the following transition targets Tap induced by the offline dataset D sampled with behavior policy T&.
T (s) E(M 1( = 8,0; = 0 =8)
E(maa.))e 1 s,0; = =a
(23)
Σ
1(s: s,a, a, S =s).
n(s,a)
Under the n finite samples in the offline dataset, we assume that TT ~N(BTTT,), o² The above definition specifies the regression target in the lo regression problem, and denote
21

---

Tx = (Ts,())1)1) TT, (5/ an)l² E I= as the empirical transition probabilities of certain transition pairs in the offline data D = {st, a,, &isignn
For regression data X: Next, we need to specify the data X in the regression problem. We denote the i-th row of X as the i-th sample in the offline transition pairs X; E D, which is vector of Kronecker product between d(st, ai) E Rd and normalized of generality, we assume Co = 1 and only need to normalize ()(()) by C.):
a(el)
E Rd' (without loss Cp
(s()
X; (88,,
Cp
(24)
1 Cr (0(6: (1( (
( a.) ((s))ayy
As a result, X, E Rddd, since there are in all n samples in offline dataset, X E Raxdd is the dataset-dependent matrix with all n rows of samples, and and d' are the latent dimension of & and μ., respectively. Based on the feature regularity criteria in Assumption 3 we have |x,ll2 S [| S 1, I.
The prior work regression:
14 estimate the transition kernel of a bilinear MDP using the following ridge
min (e... K o(s a) M M||
(25)
In this paper, in order to promote the sparsity of the matrix M, we introduce the to regularization term and arrive at the following optimization problem:
min gM
Σ
[Il((s)) K7 ((({) ((8;, n;)² Mu(s))1 A||3| [lo]
(EED
(26)
Σ
+ ITT (o) 0:) X,3 + M113 [1o] =: PD
min
gM
where we denote the solution associated with the offline dataset D as B... Here, we use the empirical version constructed by the finite samples in offline dataset DD
Given the goal-conditioned reward setting, for single episode ~ T, r(s, a; g) = 1 if and only if -g. otherwise r(s, a; g) O, as is specified in Section 2.1 Since are essentially predicting the
probabilities (normalized to a sum of 1) of whether the next state is the goal state, i.c. E.es ((s) = 1.
Therefore, we have (- S 1.
As denoted by Equation 23 for the specific offline dataset collected by some behavior policies nT
we have the regression target T (s'1s, Mamgerefper (8:=8,0, ma,e/=s))
LimaneeDe 11 ma) and the corresponding
features induced by the dataset X, = (ss,, a,) x mG))
We have the following equations hold:
Ca
T = X X +
(27)
where Sp is the true underlying transition mask given the offline dataset D, ~ (0 o I.) is some exogenous noise of the transition model.
Specifically, for the regression problem, we transform the original trajectory dataset D as [X, T 7 as the representation of transition pairs rolled out by the behavior policy πT. Similarly, we can define some 'well-explored dataset' 'D'' which is an infinite dataset rollout by the behavior policy TH, D* = {8i, 04, ri j similar to the definition of regression target T in Equation (3 and representation data in Equation (24p. we have EX, as the regression pairs with access to
the true transition probability distribution for all the state action pairs as, Here X E RRnxdd' is the
regression data defined in equation (22). Recalling the definition in Equation (23 we can further
build the regression target under well-explored dataset D* as:
E E[Tag ng (s'|s, a)] = E(sran 1 = 8, ai = 8 = 8')
E(m,ann()) 1(8: 8,0; 0.)
(28)
=EenT(( [1(s = 8, a, 8/ s'))
22

---

We denote B[T,y] [ (5)5), [ sample size n, we have [ (s' s, s, = ss +e= X p +e. In addition, we also introduce
(s Sn a ]]]
ER. In practice, with finite
a vector form ||S||AA)|S|| so that v(s, E x AA we have
T(+ a) = [ T(s2ss,))
[s)ss]ls,a)] E IR!S)
where the state space is denoted as S = { s are all possible states. Similar to the Kronecker product define for X in equation (26). we define X in a matrix form for any state-action pair (8, a):
(ss)) Cμ a(
p( 1" E RlSixdd
X(-Is, a) = [o(s, a) &
C₁
In addition, we let ((. R lxdd' denote the s'-th row of ((-(s, a) associated with the state s' E S. Consequently, the estimated transition kernel can be expressed follows:
T(Iss, a) o² (s, (s.a)Mm(()) X(-18, )
Step 2: decomposing the term of interest. To begin with, recalling the definition of Bellman operator B,, in Equation (16) and applying Hölder's s inequality, the term interest for any time step 1 S h < H and state-action pair (s,a) E S x A can be controlled as
(BB, (s, a)l < (T((1s, a) - T(1ss, a), Vo+i)l S [[((-s,,a) a( (
(29)
[+(ss,))) - TT--ss a))ss
where the first incquality is held given goal-conditioned reward formulation in section 2.1 To continue, we have
(s) (BB.))ss)))BBBP..))))) (1+)))))) | (s. a))ls
(30)
= =|x a)).))))
Here, we recall 33 represents the parameter vector in the estimated causal masks based on offline dataset 'D sampled by T.. Similarly, we denote By. as the estimated causal mask outputted from equation (26 based on the infinite dataset DD generated by the behavior policy nT. Then, we can further control equation (30) as
G. [X-+|s, a))3 X(-1ss, 33 1100 |X-1|s, a).b
X(-Is, GM + X(+ a)33.. - -X(- a) a) 1oo0
< X(-18, 13
Mb-ll ( a) 13 -3
X(-Is, a0o 13
BM |lce + X ()) 13 3M 11:00
3M 10 +
13M 3M 1lce
(8)
(31)
Here the last inequality comes from the fact that
I(+1s,
max
Σ
[(-1s.
max ((((s,a),ll vE|S]
sE|S]
je(da']
max ()(()) EE|s)
# C₁
=1
based on the definition of in equation 224 and assumption 3 Here (a) comes from the mismatch error between the demonstrated offline dataset and some optimal rollout datasets. And (b) comes from the error of the lo optimization of causal masks given the existence of exogenous noise defined by SCM in Definition We will control them separately in the following.
23

---

Step 3: Controlling term (a). We need to consider the optimization process in the original regression problem in equation 22 fully understand the difference between 3 and 3 where the only difference is that the latter uses a perfect dataset with infinite samples. The optimization problem target (cf. equation 2 can be solved by the iterative hard thresholding algorithm (IHT) proposed by [89] IHT offers an iterative solution for the &u regression problem, armed with a hard thresholding operator as below:
[o) max [0,3, - {} if
> X 8f Σ -=,, dd'.
(32)
0
We denote A (x) as the estimated causal mask parameters after i-th iterations with dataset DD Similarly, we denote B (i) as the estimation after i-th iterations with dataset Dr, We initialize the graph to be a full graph regardless of the datasets (D* or D) used in the optimization process, leading
to 3 (0) By (0) 1 E Rddd".
Recall that X E Raxdd TTx R., E[T E I and BD is [0, 1jdd based on the definition in the original lo optimization problem in equation (22). equation [3) and equation (28). The update rules of using either the dataset DD or "D* can be written as:
BN (i) = (3 (i - 1) x X3 1)1)
(33)
BD (1) (x (i - + + 1E(TT)) -x30- (i (6-1)) 1)])
Note that the difference between the two parameters 3 and 3M essentially relies on the difference between two pairs of transition kernel estimation datasets [X, [[T+ and [X, Tx verified that the hard theehhhhessn operator (-)( defined in equation (32 is L 1-Lipschitz.
easily
According to 190), we can set the learning rate S L = here. Using the Lipschitz property, at each iterative update step i, we can control the difference between the estimated parameters obtained by using 'D or D as
1 (i) B (i)ll
= 11ax (3 (i 1) + 7 - X3y 1)))
g2 (3p.
1) 1 [ x.M (i)1))
11 (i- - 1) + - - - X x3 (i ())) (3 1) + 1 X X 1)1) ll2
I1aa x (3 - 1) - (i - 1)] + x [T -B[Tap]lle
Lad x X 1122118p(i - -3p- (i- 1-)12 //1xX1221T E[TT
& 1i - 1) - 3D- (i 1)llz + T/ E[TT,
(34)
Here, Lad represent the identity matrix of size dd' x dd', the last inequality holds based on the fact that | I defined in equation (24). Since r7 S 1 1, result, || - nxx X||2 max(1 - 7 X)) S 1. We perform IHT for sufficient >> 0 iterations and output the last step estimation our solutions for either the offline dataset D (we use for practical optimization) or the perfect dataset D In practice, for any dataset such as D and any accuracy level < & S 1, we have the output of IHT gradually converges to the optimal solution A of the problem in [|M
equation (26 Namely, after at most K ~ log
steps 89 Corollary 1], the output satisfies
By 3/ (K) < f. Similarly, based on the perfect dataset D*, the output of IHT gradually converges to 3M at the same rate. Consequently, the term of interest (a) can be bounded recursively as:
24

---

2M 1300
PD 00
BB (K) - 3 (K) + (30
3D (k) + (30)
-B0 (K)
BB (K) - 3 (K)
+ + 5
4
BB (K) - 3D- (K)
+ $ +
4
(35)
- |1 (0) 3M (0)ll2 + Kn
E
11l
= Knl[TT mp [T(T,) 1,.ll122 +
|TT B[Tx + |M 1
S log
2"
where the second inequality holds by recursively applying equation 34
K, K 1, iterations,
the last equality due the fact that 3 (0) 3 (0), and the last inequality holds by setting
K = log( |MI22
Now the remaining of the proof will focus on controlling defined the regression target in equation (23) and equation
E nt bl Recall that we have
1 n(saa
Tap (8)18,
Σ
1(8 = 8,0; a, =s),
a(s,a)
i=1
E[TT (s' [8,a]] Ey [1{; = = o, 8 = =')].
Proposition 1 explo-ex dataset). With the offtine dataset D of in total n samples with n(s, a) samples generated independently conditioned on any (8, a). For any 5 < with probability at least 1 - 8, one has
|S|
SS|AA
Sxxa) |T a) - E[T,, (1ss I Co
log
a(s,a)
for some universal constant Cg.
The above proposition can be directly proved by applying 91 Lemma 17] over all (s,a) E x A:
14|S n(s, a)
[|T (118,a) - E[TT, (-18. a)]l S
log
(36)
8
As result, further extend the results in equation 35 to bound the term (a) as follows:
1 - 3 1loc I Knn[Ta
ET
Illa
+
|
IS
|S||AA]
5
< Cr log
log
+
E
n(ssaa a)
6
2
(37)
S
[S||AA]
CpCp log
log
E
(sssa) a)
IIMllo E
IS)
|S||AA] 8
E
A Ci
log
+ 2
n(s,a)
where Ci C3C₁ some constant that is related the feature regularity, & the level of error tolerance in the cumulative value returns, in our setting, 0 < 1.
25

---

Step 4: Controlling term (b). For (b) in Equation 31 which is || BM1 Il.0. we are interested what is the optimization error given finite well-explored offline dataset D Here the optimization error mainly originates from the Gaussian noise in the SCM formulation in Definition + Yet our causal discovery module, i.e. an ln estimator, will always encounter some estimation error induced by the exogenous noise.
Firstly, we have the bounded relationship between Eg and &oo norm:
11.%% 1331 S 11 --- 1|2
then we can analyze the error bound for lD regression in the sense of lz norm.
The derivation below generally follows the En regularized linear regression bound in [92]. Based on Assumption 2 and there exists M, we denote this optimal solution in the vector form as 3r
Besides the aforementioned optimization error in the iterative thresholding update process, according to the SCM in Definition and Assumption 2 and Proposition | we denote a finite subset of observed transition probabilities of the well-explored data E[Tx;] with size Tobe R". By definition above, we can then assume the finite-sample regression target 7 is generated by causal features of transition pairs (denoted by X = x = , X E IR" "xxdd and the ground-truth causal mask (represented by 3M Redd") with the following equation:
Tobo E = T + cc= XB +c.
(38)
Here ~ NO, II)) is the independent exogenous noise defined in Definition 4 We'll then use the above equation to bound term (b) in equation (31).
Specifically, we solved this to regression problem with its bounded form as follows:
min
Tobo - X3M122. s.1.1133 llo S 8. gar
(39)
In BECAUSE, we select the sparsity level & R ||Bx |10 = M S dd'.
For simplicity, we denote the approximate solution 3y.. in Equation (39) as 3M, By the virtue of optimality of the solution 3. in Equation (39h, we find that
|Tobe - X33 12 < |oob X3M
(40)
then by expanding and shifting the terms, above:
derive the basic incquality for the to estimator
||aab
X.3 1/ S (Tabb XBM) + (XBM - x13
38
1ee12 S ||e + -X
1.MMXX77 -33 )12
|| + 1XBM x X3313 2(e, X33 - x3M)
Since both 3M and 3M are s-sparse, the vector BM - gM is at most 2s-sparse. We denote the set of all 2s-sparse dd' dimensional vector set as T (2s), we denote e = 1 - - = 0) as an indicator vector, 24- = 0 if and only if 3.M - BMM V E [dd"].
By shifting the terms, we use the sub-Gaussian assumption and further use Hölder's inequality get the following results:
n(s.a) XX3M 1
1135
2
(x"
x,
- )B²D
n(8 a)
2
1:33 - -B 1/2
(0, X
sup
n(s,00 a)
veT
(X c)S (XTc)s
(41)
s 2 - MM
sup |s| 2a
n( a)
2s log ddd /8)
S 3 - 3M 1120 V
n(s,0) a)
26

---

where E ~~ o0 is the exogenous noise variable in the SCM in Definition! 4
Then by simply applying restricted eigenvalue (RE) condition, for some k(S)
have
x(S))l3M B"113
1 n(s,a)
x X - BM) = 211 BB 1120
2s loglod' /5)
m(s, a)
Therefore, have:
20 28 K(S)
log log(dd'/00 /8)
|M loglooo log(dd' 18)
1| 1M1
C.o
(42)
n
n(ss a)
with probability at least 8, which bounds term (b) Equation (11.
Step 5: Summing up the results Summarizing both bounds for terms (a) and (b) in Equation will get the following bounds with probability 1 6:
31
[BB Vs+1)(8, a) - (BA MiNt+11 (s, - 11 - 3M3 11 +
11.32 3b- 1l
C1 log ( [M
|S|
|S||A|]
log(dd' /0) n(s,a)
log
+ C.a |Mlloy
n(s.a)
5
2
|
log(1/0) n(88 a)
x min { log
V
(43)
Therefore, we complete the proof of Lemma 3 by showing that for all (s, E E
EBECAUSE
((B,)h+z)(ssa)) (B,, P 1).,)))))
(44)
|
log(1/8) n(s, a)
S min {Ci log
The final bound the term (b) includes a dependency 0 (+))) and logarithm dimensionality dd' in the estimated transition matrix. It also incurs a dependency of error tolerance & or SCM's noise level a square root of sparsity level v = M|
So far, prove the following bound in theorem
|| llo &
H Σ Ex
log(1/5) n(sh. an)
(*)
V (3) < min {Ci log
VIS), C..
1 81 =
(45)
In order to achieve ç g-optimal policy such that Vi - VP (3) S E. the RHS needs to satisfy:
M
log(1/5)
min {Ci log
VISI,
Σ E
|s
SE
(46)
n(sh, ax)
We first multiply /iimssa.)ssx.xxI1 E.r [n(sn, an) s1 3 and then take the square on both
sides. We have:
M E
log(1/0) (I) (sn, an)]
min
{C1 log
[M])
Σ Ex
s1
n(sh on)
$ 5
min (aa.))eSxAx(H))
Ex [n(sn, an) =
(47)
27

---

Given the definition, LHS in Equation (4
satisfies:
M
"
LHS min (Ci log
/ C. Cl
Σ Ex
[Vlog(100)
(48)
5
To ensure the satisfaction of É-optimal policy, we thus would like the RHS of Equation 47
satisfies:
min RHS = V (s.a,)))xAxx(1)
[n(Sn, an)
8]
[I
) C.oo Σ
H
2 min {Cr log
E [Vlog(1/8)
(49)
[l
min {C₂ log
(S], C. / H (og(1/8)
Consequently, can shift the terms and get the sample complexity bound for the g-optimal policy, vO E 1:
E [n(sk; an) = min [ogl ( C3) H² log
min [5xx,))
5²
(50)
C Additional Experiments Details
In this section, we provide additional experiment results and algorithm implementation details.
C.1 Implementation of Causal Discovery
We implement the causal discovery primarily based on Equation (5) However, in practice, how to control the coefficient before the sparsity regularization terms is crucial to the final performance. In practice, instead of controlling A, we use p-value threshold determine the following conditional independence:
Φ(s, () TL ((s')(
((s,a) (9)
(51)
Here ((- means that this element is the 1th factor in the abstracted state action representation, similar ()) ((s,a) -1) means all the other factors in the representation except for the ith factor.
If the valuu based the above conditional independence test is less than threshold, remove the edge by setting Mij- = 0. Please refer to the Appendix Table 13 for the selection of threshold in each environment.
C.2 Training Details of Energy-based Model
We train the EBM according the margin loss in Equation (8) In practice, attach Tanh( the output layer to clip the unnormalized score between -1 and +1. The energy networks take both conditions and samples, then concatenate them together and sent into MLP encoders. The detailed hyperparameters of EBM are listed in Table [2
In the vanilla EBM, people follow the Langevin dynamics to effectively sample the negative samples. Here, as we discover the causal mask and identify the causal representation in the model learning stage, find that we process to get some effective negative samples, which similar to the practice of augmentation of causality-guided counterfactual data [77).
these representations in both the energy networks and the sampling
The interesting trick we employ here the way to get our negative samples by mixing the latent factors from offline data. For example, for positive sample array
28

---

((s))a) ((s()
m(s) taa')
((s))
((s))
x²
(52)
[u(())))))
u(s%))
((sh)(a))
with conditions:
((81, )
((81, a1)
((11; (1)
Ó(5), a1) (1)
$(81, (2)
0(51, 01)
(53)
((8 (1)
a1
a1)
Here, 8,, 04, 5s denotes the timestep the offline samples. o(, (0 means this element is the zth
factor in the abstracted state action representation, similar to the ():))))))
as we already get the corresponding causal representation pus') that semantically meaningful,
we can mix the columns to create useful counterfactual negative samples
(Ms))
(2) 1 (-12
((s))(1)
p(s{)
(()
(54)
countefactual
(M($ (1s2) (2)
(($))(a))
These counterfactual negative samples seem to effectively use the causal representation and speed up the training, as show in Figure 7
T-o
T-100
T-200
T-300
T=600
Randimm
0.50
Causal EBM
0.00
-0.5
T=500
T-600
T=700
T-800
T-900
Handom
-0.75
Caursl ESM
Figure 7: Comparison of the convergence speed in EBM training. Compared to random negative samples, approach enjoys higher rate of convergence empirically.
C.3 Additional Mismatch Analysis
To evaluate the significance of the objective mismatch effect at three different levels of offline datasets in Unlock environments, we collect 5,000 episodes in each of the Unlock environments for each method. Then evaluate the mismatch via the following metrics.
We conduct hypothesis testing via Mann-Whitney U Test 93 with Null hypothesis
Ho : Lmodet (ooudet(Tpaal (Tpos) Cmodel (modee(TTeg)).
understand the exact difference in model loss between two groups of samples, we compute their Wasserstein-1 distance in the episodic model loss between the trajectories with positive and negative rewards, i.e. W, (T Tneg).
We report the results of p-value and the W, distance of two groups of model loss samples in the following table.
29

---

Table The comparison results of the p-value and the W1 distance (x10-4) BECAUSE. Bold means the better.
between MOPO and
Unlock Unlock-Expert
Ehlock-Medingg
Methods
p-value ())
Wi Dist (1)
p-value ())
Wi Dist (f)
p-value (4)
Wi Dist (1)
MOPO
6.5
10-3
0.7
1.0
1.1
1.0
N.A.
Ours
2 0
3.1
3.0
0.9
< 0.1
C.4 Additional Experiment Results
We report the results of the task-wise performance of all baselines in the main experiment and variants in the ablation studies Table
and
Table 6: Success rate (%) for 18 tasks three different environments. We evaluate the mean and 95% confidence interval given by the t-test of the best performance among 10 random seeds, well as the p-value between the overall performance. Bold is the best.
Eav
ICIL
CCIL
TD3-BC
MOPO
GNN
CDL
Denodsed
Factor
MnM
Delphic
Ours
Lif-1 Lin-O-R
1431999
1001114
934354
24.3=29
22.11:36
338-150 1601447
20001:10
24.02288
16.31228
20221:.11
19.2:106
8.524.7
0.000.0
13410
10.2-16
133423
155513
21.233.0
142422
1777228
21.4:39
Uabck-1-
34211
222088
439920
21.5-19
11742.1
6.220.5
69207
81111
86213
144411
327418
Ualock-ORR
116240
185237
13.3230
166213
12:115
1.22.0
7.400.8
80111
90210
12:21111
27.6:220
Crash-I-R
114443
19541005
9.1466
32.7=49
3974446
$204228
31.32444
1609:2
1744336
59.416611
5.7=3.0
5101133
41001:0088
26.6114.8
59.8:33
27.51:3
32.91:577
35.3148
41.02:57
287118
24.0::23
59.51444
Lif-O-M
4681152
20092119
13.0:1111
31.9228
25.811
26.0::422
21.03:3.2
30.22277
24.34:2.3
1833-266
$2314499
Usloc-1MM
48200
4.21133
$4::10
848=111
1755226
20.744444
12.94:55
17.755
37.61:4
14:11:
980:4499
Uolock-OM
129128
14.1220
19:2:12.5
1955
162218
20.1199
10.8:10.8
21.52:277
27.01:377
$1.7122
6.8.1
Crash-IM
3554.99
24.11222
163:2900
17
11S1111
63:51:0
638:1400
45.5
184129
58.01222
94
Crash-O-M
117335
1.244
566:3.1
17.3223
38204
2002233
20.3::1.6
1600220
63321
22.21:17
20.11:9
LitsLE
73.81700
86741577
44.3::1.22
82.4-67
63.3--000
71.05:55
74.245.
9..
33.33:
53.5-663
92841.2
LifeO-B
26.1
41.6:116.6
12.4119
60.85066
63.1
64.0:4444
91.73377
29.05:4.1
29.5-33.11
9374599
Unlock H
116433
13.71311
15.6:12.22
88.6=4.6
15.311
73241:8
50.3-133.0
29.3::2.66
60.71:22
83:24:122
97.4110
Uslock-O-E
224150
19.9-444
13.8114
413114.1
35.7:1:2.3
29.5::3.5
38.71:20
$6.119
82.1165
Crash-I-E
35.7.9.
261771
58.5-477
11:21:09
63.7.128
60.311:100
$2.0125.3
10:01:111
$7.11:22
953.11.3
Crash-O-E
113333
1233:33
6.2.1:18
148::2.1
3.5
188:113
20.7:1:2.2
13.6:322.3
3.9:2.0.6
16511:17
20.71332
Ovemli-1
27.2
21.0
73.
Dvenll-O
20.2
19.9
14.6
28.1
17.0
24.9
23.5
27.2
17.4
27.5
43.0
30

---

Table 7: p-values of different methods (each has 10 random trials) against BECAUSE in various environments. Under the significance level 0.05, we mark all the baseline results that are significantly lower than BECAUSE as green, and the rest as red. We can see that BECAUSE significantly outperforms 10 baselines in 18 tasks in 91.1 1 % of the experiments (164 out of total 180 pairs of experiments).
Env
ICII.
TD3+BC
MOPO
GNN
CDL
Denoised
IFactor
CCI.
MrM
Delphic
Lift-I-andom
0.001
0.000
0.001
0.000
0.000
0,000
0.001
0.001
0.000
0,000
LiftO-random
0.000
0.000
0.000
0.001
0.033
0.013
0.464
0.000
0.001
0.052
Unlock-Lranderr randoo
0.000
0.000
0.000
0.000
0.000
0.000
0.000
0.000
0.000
0,000
Unlack-O-mnan
0.000
0,000
0.000
0.000
0.000
0.000
0.000
0.000
0.000
0.000
Crash--randoo
0.000
0.000
0.000
0.000
0,000
0.000
0,000
0.000
0,000
Crash-O-randee
0,000
0.000
0.000
0.000
0.000
0,000
0.000
0.000
0.000
0.000
Lift-I Litt--mediim
lift---medium
0.966
0.000
0.438
0.009
0.022
0.032
0.209
0.125
0.003
0.000
Unlock-I-mediun
0.000
0.000
0.000
0.000
0.000
0.000
0.000
0.000
0.000
0.000
Unlock 0-medium
0.000
0,000
0,000
0.000
0.000
0.000
0.000
0.000
0.000
0.000
Cmas----medium
0.000
0.000
0.000
0.000
0.000
0.000
0.000
0.000
0.000
0.000
Crash-O-mnedium
0.000
0.000
0.412
0.001
0,943
Lift-I-xxerri
0.017
0.000
0,004
0.000
0.000
0,000
0.993
0.205
0,000
0,000
expert
Unlock-L-expert
0.000
0.000
0.001
0.000
0.000
0.000
0.000
0.000
0.000
0.000
Unlock Saak----eee O-expert
0.000
0,000
0,000
0.000
0.000
0.000
0.000
0.000
0.000
0.000
Cmihh Lexpert
0.000
0,000
0.000
0.000
0.000
0.000
0.000
0,000
0.000
0.000
Crash-O-cxpert
0.000
0.000
0,002
0.000
0.133
0,500
0.005
0,000
0,000
0.011
Table 8: Success rate (%) for 18 tasks in three different environments. We evaluate the mean and 95% confidence interval of the test performance among 10 random seeds. Bold means the best,
Env
BECAUSE BECAUSE-Optimish BECAUSE-Lineat BECAUSE
Lift-I-random 33.8+5.0
23.2=3.1
16.5+1.6
22.2+66
Lift- lfttOOrradoom
21.4+39
15.3
8.9448
18.245.3
Unlock-I-random
32.7428
31.2-24
20.7+1.7
10.7408
Unlock-O-random
27.6121
26.3=117
24.0+27
9.340.
Crash-I-random
59.4-162
49.9
54.3+5.
36.147.3
Crash-O-random
19.7114
14.2.0.5
14.8
10.012.
Lift-I-medium
59.5144
46.8
24.4
36.4.66.
ffft-O-medium
32.315.0
24.5
16.4
28.9
Unlock-I-medium
98.0
92.7=158
91.01188
29.9.11.9
Unleck-O-medium
68.8+1.5
58.7=2.2
60.0+20
18.7+1.3
Crash-I-medium
90.4+118
82.8=99
66.7+74
60.8
Crash-O-medium
20.3
17.5=0.0
15.9+28
24.64109
Lift-L-expert
92.8+112
68.6=4.7
75.6410.7
78.146.1
Lift-O-expert
93.7-16.0
58.3
66.1
71.946.
Unlock-I-expert
97.4+10
93.1=119
94.0+20
29,3+1.3
Unlock-O-expert
82.1.666
64.6
65.8
20.2
Crash-I-expert
95.3+14
91.0=222
77.9+28
50.3+7.6
Crash-O-expen
20.7133
12.5
26.916.1
25.0
Overall-I
73.3
64.4
57.9
39.3
Overall-O
43.0
32.4
33.2
25.2
31

---

Table 9: p-values of different methods (each has 10 random trials) against BECAUSE in various environments. Under the significance level 0.05, we mark all the baseline results that are significantly lower than BECAUSE as green, and the rest as red. We can see that BECAUSE significantly outperforms 3 variants in 18 tasks 83.3% of the experiments (45 out of total 54 pairs of experiments).
Env
BECAUSE--Opimism BECAUSE-Linear
BECAUSE-Full
Lift-I-random
0.001
0.000
0,003
Lift- Lfft-O-random
0.003
0.000
0.146
Unlock-I-random
0.189
0.000
0.000
Unlock-O-random
0.145
0.015
0.000
Crash-I-random
0,036
0.090
0,000
Crash-O-random
0,000
0.000
0,000
Lift-I-medium
0,000
0.000
0,000
Lift-O-medium
0.004
0,000
0,130
Unlock-I-mediun
0.067
0,006
0.000
Unlock-O-men O-medium
0,000
0.000
0.000
Crash-I-medium
0,061
0,000
0.000
Crash-O-medium
0.005
0.005
1.000
Lift-I-expert
0.000
0.003
0.000
Lift-O-expert
0.000
0.000
0.000
Unlock-I-eppert
0,000
0.002
0.000
Unlock-O-exper
0,000
0,000
0,000
Crash-l-expert
0.001
0.000
0.000
Crash-O-expert
0.000
0.969
0.990
C.5
Additional Environment Description
We provide a more detailed description of the environments we use in the experiment, as shown in Table 10
Lift The Lift environment based on the robosuite 34 contains 33 dimensions of state space, including the end effector pose. joint pose, joint velocity, cube pose as well as its relative position, cube color, and a contact flag. contains 4 dimensions of hybrid action space that uses Operation Space Control (OSC) to control the 3D position and the ID gripper movement. The task is counted as a success when the assigned block is lifted from the table over m. The generalization setting in the Lift environment is to use an unseen combination of position and color during online testing. This environment can be abstracted into 15 dimensions of factorizable state space and 4 dimensions of factorizable action space. The causal graph of this environment is recorded in Figure g a).
Unlock The Unlock environments based the MiniGrid world 335 contain 1 110 dimensions of discrete state space, with 3 of 36-dimensional vector inputs representing the current position of the agent, key, and door in a 6x6 grid world. The rest 2 dimensions in the state space memorize the state of whether the agent has the key in hand. The action space is also discrete (with eight dimensions) to determine the movement (up/down//eet/riihh)) and the pick-key, open-door actions. An episode will be counted as a success when the agent holds the key and uses it to open the door in the right position. The generalization setting in the Unlock environment is to change the position of the door and increase the number of total goals in the environment. The agent will only successfully finish one episode by opening all the doors. The causal graph of this environment is recorded in Figure ()).
The Crash environments are based on the Highway environment (37] which contains 22 dimensions of continuous state space, with four vector inputs representing the current position, velocity, and orientation of the surrounding vehicles and ego vehicles. There are two additional dimensions of state memorizing the collision type between the ego vehicles and surrounding vehicles or pedestrians. The 8-dimensional action space is continuous to determine the acceleration in the x y directions of the ego and surrounding agents. The generalization of the Crash environment to add different numbers of pedestrians that may cause the crash. An episode will only end when the ego vehicles have a near-miss with both of the pedestrians at the scene. We visualize the causal graph of this environment in Figure g c).
Crash
32

---

All three environments are visualized in Figure 4
We list their basic configurations in Table 10
Table 10: Environment configurations used in experiments
Environment
Parameters
Lift
Unlock
Crash
Max step size
30
15
30
State dimension
33
110
22
Action dimension
4
8
8
Action type
Hybrid
Discrete
Hybrid
Intrinsic state rank
Intrinsic action rank
(8) Lift Environment
(b) Unlock Environment
(c) Crash Environment
Figure 8: Underlying causal graph G in all environments with expert demonstration.
C.6 Additional Baseline Information
We collect data on the above 3 different environments, thus forming groups of offline datasets.
Table 11: Bahavior policies used to collect offline data in different environments,
Environment
Behavior
#Episodes
Success Rate
Additional Description
Random
1000
0.24
Random actions after few steps of initialization.
Lift
Medium
1000
0.60
Random actions before the goal roal-reaching expert.
Expert
1000
1.00
Query expert policy for all time steps.
Random
200
0.21
Random navigation with high randomness Targeted scarching in goal directions Shortest path planning via A"
Unlock
Medium
200
0.46
Expert
200
0.87
Random
1000
0.14
Fixed ego, random pedestrians Planned ego, random pedestrians Planning both ego and pedestrians
Crash
Medium
1000
0.35
Expert
1000
0.66
After collecting the data using scripted policies in different environments, we train all agents as well as BECAUSE under 10 different random seeds. Then we report the best performance of each trial and compute the mean and standard deviation over 10 seeds for each task in the Appendix 6
We refer to the following codebase to implement all the baselines we use:
Invariant Causal Imitation Learning (ICIL, 338): http:://gghhbbcommmmoonnbaaaaaaa
Causalllmitation-LLamminn MIT License.
Causal Confusion Imitation Learning (CCIL, 39
reference link to the paper
TD3 with Behavior Cloning (TD3+BC, http://gthhuh.com/ttinnn/11 BC, MIT License.
33

---

Model-based Offline Policy Otimization (MOPO, 2
https//gtthub.com//unnnng
yang/mopo.gn MIT License.
Relational Graph Neural Network (GNN, 42
MIT License.
Causal Dynamics Learning (CDL,
24
MIT License.
Denoised MDP (Denoised, 12 CC BY-NC 4.0.
Mismatch No More (MnM, reference link to the paper
World model with identifiable factorization (IFactor, [3). reference link to the paper
Delphic Offline RL (Delphic, [4)): reference link to the paper
The detailed hyperparameters use in BECAUSE and other baselines are listed in Table 12 and
Table 13
C.7
Experiment Support
Our code is available at the anonymous repo: hutp:///aoyynooosseesssseeee//es
on server with AMD EPYC 7542 32-Core Processor CPU, 2xNVIDIA RTX 3090 graphics and 2xNVIDIA RTX A6000 graphics, and 252 GB memory. For one single experiment, it takes BECAUSE and other baselines about 1.5 hours with 100, 000 iterations to train the world model and 1, 000, 000 steps to train the energy-based models.
Computing resources The experiments
Table 12: Hyper-parameters of models used in experiments of BECAUSE and baselines (Part I)
Environment Unlock
Models
Parameters
Lift
Crash
Learning rate
0.0001
0.001
0.0001
Size of data D
15000
4000
15000
Epoch per iteration Batch size Planning horizon H
20
5
10
256
256
256
15
10
20
Planning population
1500
100
1000
BECAUSE
Reward discount ?
0.99
0.99
0.99
Spectral norm regularizer X
10-4
10
10
Spectral norm regularizer A
10-1
10-4
10-1
Causal discovery Pthres
10-8
10-4
10-6 6
Encoder hiddens
256
64
128
EBM hidden
256
64
128
EBM negative buffer
5000
1000
5000
EBM training steps
1000
1000
1000
EBM regularizer AFrM
10-4
10-4
10-4
MLP hiddens
256
64
128
MOPO
MLP layers
2
2
Ensemble number
5
5
5
Initialized mask coef.
1.0
1.0
1.0
CDL
MLP hiddens
256
64
128
Sparsity regularizer
0.001
0.001
0.001
GNN hiddens
256
64
128
GNN
GNN layers
3
Use the same planning parameters as BECAUSE.
34

---

Table 13: Hyper-parameters of models used experiments baselines (Continued)
Environment
Models
Parameters
Lift
Unlock
Crash
Learning rate of MINE
0.0001
0.0001
0.0001
MINE hiddens
256
64
128
MLP hiddens
256
64
128
Learning rate of EBM
0.01
0.01
0.01
ICIL
Size of buffer of EBM
1000
1000
1000
EBM training steps
1000
1000
1000
EBM hiddens
256
64
128
K of langevin rollout
60
60
60
Avar of langevin rollout
0.01
0.01
0.01
Learning rate of Critic
0.0003
0.003
0.0003
Critic hiddens
256
64
128
Learning rate of Actor
0.0003
0.001
0.0001
Actor hiddens
256
64
128
TD3+BC
Target update rate
0.005
0.001
0.0001
Policy noise
0.2
0.2
0.2
Balance coefficient a
1.0
2.5
2.5
x belief size
256
64
128
y belief size
256
64
128
z belief size
0
0
0
state size
33
110
22
Denoised MDP
state size
33
110
22
z state size
0
0
0
embedding size
256
64
128
Learning rate
0.0001
0.001
0.0001
All hidden dim Disentangled prior output size Learning rate
256
64
128
IFactor
19
7
10
0.0001
0.001
0,0001
All hidden dim
256
64
128
MnM
Diseriminator learning rate
0.0001
0.001
0.0001
Discriminator clip norm
0.25
0.25
0.25
Reg weight Initial mask probability
0.0001
0.001
0,0001
CCIL
0.95
0.95
0.95
Learning rate
0.0001
0.001
0.0001
Ensemble model size Uncertainty penalty weight
5
5
5
0.0001
0.0001
0.00005
Delphic
KL weight
0.0001
0,0001
0.00005
35

---

C.8
Broader Impact
This work incorporates causality into reinforcement learning methods, which helps humans under- stand the underlying mechanism of algorithms and check the source of failures. However, the learned causal world model may contain reaaalreedable private information about the environment and the dataset. To mitigate this potential negative societal impact, the causal world model should only be accessible trustworthy users.
36

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility. transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and precede the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No], or [NA].
[NA] means either that the question Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will published with the paper.
The reviewers of your paper will be asked use the checklist as one of the factors their evaluation. While "[Yes] is generally preferable to "[No]" ", it is perfectly acceptable to answer "[No] providedd proper justification is given (e.g.. "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering binary way, we
"[No] or "[NA]" is not grounds for rejection. While the questions are phrased acknowledge that the true answer is often more nuanced, please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS paper checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
They are provided in the abstract and introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: They discussed in the conclusion.
37

---

Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification:
They are provided in section 3.4 and Appendix and B
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide necessary implementation details in section 3.2
3.3 and Ap-
pendix C
Guidelines:
The answer NA means that the paper does not include experiments.
38

---

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g.. with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibilityy In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We provide an open access the code and data collection scripts in the anonymous link.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
//nips.cc/public/eee.c www.deddeee.oeuuummmcecom ssi ssionyolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
39

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We provide all the experiment settings and details Appendix C.3 and C.6
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[Yes]
Justification: We report the mean and 95% confidence interval (CI) well as the significance level (p value) in Table 6
7 g and 9
Guidelines: We
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-simma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time exccution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Provided in Appendix
C.7
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
40

---

The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: Provided in the OpenReview form.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Provided in Appendix C.8
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Answer: [NA]
Justification: This paper does not contain pre-trained language models, image generators, scraped datasets, similar assets.
Guidelines:
41

---

The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
All the baselines are properly cited and introduced Appendix C.6
Guidelines:
The answer NA that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g.. website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re- ce-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and sereenshots, if applicable, as well details about compensation (if any)?
Question:
Answer: [NA]
42

---

Justification:
The paper does not involve crowdsoureing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve study participants.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
43


Reviews:
[{"review_id": "wQxSwe4wnZ", "reviewer": "Reviewer_8KHa", "summary": "The paper proposes BECAUSE, an algorithm designed to address the objective mismatch between model and policy learning in offline model-based reinforcement learning (MBRL). The algorithm first models the spurious correlations between the current state s and the current action a, as well as between the next state s' and the state-action pair (s, a), through the formulation of action-state confounded Markov Decision Processes (ASC-MDP). Based on this formulation, a compact method for learning causal representations is introduced. Once these causal representations are acquired, they are employed in both world model learning and planning to enhance the algorithm’s robustness and generalizability.", "questions": "Is this method efficient in handling continuous actions? During the planning phase, computing the argmax Q can be time-consuming.\nLimitation: The proposed method introduces additional complexity to the MBRL framework, which may pose challenges for practical implementation and scalability.", "limitations": "none", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The paper is well-written, proposing their intuition backed by three clearly laid out steps to walk the readers through the working of the algorithm. The insight of mitigates the objective mismatch with causal awareness learned from offline data is novel. The paper also provides a thorough theoretical analysis, including proofs of error bounds and sample efficiency.", "weaknesses": "The method relies on certain assumptions, such as the invariance of causal graphs, which might not always hold in real-world scenarios. Additionally, modeling the causal relationships with a linear structure may be overly simplistic, potentially limiting the accuracy and effectiveness of the algorithm in capturing more complex dependencies.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank reviewer 8KHa for their insightful and inspiring feedback. We are glad to know that the reviewer recognizes the novelty of our contributions, the clarity of our problem formulation, the theoretical contributions, and the empirical evidence that shows the advantages compared to baselines. We provide our response to the questions below.\n\n**Q1. (Assumption on Invariant Causal Graph and Bilinear Structure): The method relies on certain assumptions, such as the invariance of causal graphs, which might not always hold in real-world scenarios. Additionally, modeling the causal relationships with a linear structure may be overly simplistic, potentially limiting the accuracy and effectiveness of the algorithm in capturing more complex dependencies.**\n\nWe thank the reviewer for this insightful discussion on our assumption. We elaborate more on the generality of invariant causal graph assumption and the bilinear representation structures. \n1. Our assumption of invariant causal graph targets on the generalizable RL setting when the training and testing environments have distribution shift but share the underlying **cause-and-effect relationship**. This assumption is commonly made in many prior causal RL works, such as [1, 2, 3]. This causal invariance setting is actually applicable in many real-world decision-makin tasks like indoor navigation, autonomous driving and manipulation, as is shown in our experiment. \n2. Regarding the over-simplification concern about our bilinear structure in MDP, the nonlinearity of the transition dynamics can be captured by the feature encoder $\\phi(s, a), \\mu(s')$. In more scalable contexts such as the foundation model, linear structure is also commonly used in CLIP [4], where the inner product is applied in the latent space. \n\n\n>[1] Zhang, Amy, et al. \"Invariant causal prediction for block mdps.\" International Conference on Machine Learning. PMLR, 2020.\n> \n>[2] Wang, Zizhao, et al. \"Causal Dynamics Learning for Task-Independent State Abstraction.\" International Conference on Machine Learning. PMLR, 2022.\n>\n>[3] Zhu, Wenxuan, Chao Yu, and Qiang Zhang. \"Causal deep reinforcement learning using observational data.\" Proceedings of the Thirty-Second International Joint Conference on Artificial Intelligence. 2023.\n>\n>[4] Radford, Alec, et al. \"Learning transferable visual models from natural language supervision.\" International conference on machine learning. PMLR, 2021.\n\n**Q2. (Continuous action): Is this method efficient in handling continuous actions? During the planning phase, computing the argmax Q can be time-consuming.**\n\nWe thank the reviewer for asking this important question in implementing the planner. As the reviewer mentioned about the planning phase, our method theoretically uses the following pessimistic value iteration: \n$$\n\\overline{Q}(s, a) = r(s, a) - E_\\theta(s, a) + \\sum_{s'\\in \\mathcal{S}} \\widehat{T}(s'|s, a) \\widehat{V}(s'). \n$$\nIn practice, the above formulation is applicable to continuous actions we use model predictive control (MPC). We first uniformly sample random actions in the action space $a_{sample}\\in \\mathcal{A}$, then roll out future states by imagination using the learned dynamics model $\\hat{T}(s'|s, a_{\\text{sample}})$. Finally, we will have a set of Q functions $\\overline{Q}(s, a_{\\text{sample}})$ and take $a = \\arg\\max_{a_{\\text{sample}}\\in \\mathcal{A}} \\overline{Q}(s, a_{\\text{sample}})$ as the selected actions to rollout the next step. \n\nWe denote the number of random action samples per step as \"*planning population*\", which is a fixed hyperparameter for all the methods in experiment environments in Appendix Table 9. \n\n**Q3. (Scalability): The proposed method introduces additional complexity to the MBRL framework, which may pose challenges for practical implementation and scalability.**\n\nWe agree that our method adds additional structure on top of traditional MBRL in exchange of better OOD generalizability. To further demonstrate the scalability and computational efficiency of our method, we conduct additional experiments with **raw-pixel inputs in the Unlock environments** (see supplementary Figure (c) for details). The RGB input has a shape of (192, 192, 3) and we encode them with 3-layer CNN. The results in the **second table of the general response/Table III in the supplementary** show that BECAUSE still outperforms ICIL and IFactor baselines with visual inputs by a clear margin.\nWe also evaluate the inference speed of our method with both visual inputs and attach details in the one-page PDF. We believe the inference speed (Table IV in supplementary) is still in the acceptable range with a comparable inference speed as IFactor (>20 FPS on image data). \n\nIn conclusion, the empirical evidence verifies (a) the scalability of BECAUSE and (b) the computational efficiency with high-dimensional inputs.", "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": "7LY4gqFqyV", "reviewer": "Reviewer_aSxm", "summary": "This paper studies challenges in Model Based rl, namely:\n1. Objective mismatch between estimating the model and policy learning for value optimization\n2. Confoundedness in offline RL causing this mismatch\nThe paper proposes a theoretically sound algorithm and experiments extensively to support their claims.", "questions": "1. How does the Theorem 1 fare to existing works? A detailed comparison on the improvement factors in the sample complexity would be useful.", "limitations": "There are no significant limitations of the paper.", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Well written paper - Problem presentation is very clear\n2. Experiments cover a many scenarios, and ablation studies have been provided. I'm curious why multi-task RL algorithms weren't used as baselines given the setup of the experiments.", "weaknesses": "1. Has missed some relevant related work particularly in theoretical offline RL which use pessimism based on uncertainty. Please cite :  https://arxiv.org/pdf/2402.12570, https://arxiv.org/pdf/2403.11574\n2. The result in Theorem 1 is applicable for a Tabular MDP, but this hasn't been discussed in the paper anywhere. The result becomes void in the case when some (s_h, a_h) along the optimal policy has not been seen in the offline dataset.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to express our gratitude to reviewer aSxm for their insightful feedback and acknowledging the novelty of our contributions in practice and theory. We provide our response to the questions below.\n\n**Q1. (Additional Related works): Related works about theoretical offline RL using pessimism... multi-task RL algorithms weren't used as baselines given the setup of the experiments.**\n\nAs the reviewer suggested, we added related references to offline RL theory in the related work section. \n\n- **Comparisons with the two related works**: The reviewers raised two relevant works about multi-task RL in offline settings with linear function approximation (linear MDPs). They consider the same sampling mechanism -- offline setting, but different model assumption (linear MDPs) and problem setting (multi-tasks)  as our bilinear MDPs and action-state confounded MDPs.\n- **Differences to multi-task offline RL**: Our work does not primarily target multi-task learning, which aims to solve multiple tasks jointly in an efficient way [2] by leveraging their shared information. In our case, we focus on problems with unobserved causal effects -- the training data has a distribution mismatch with testing environments due to the unobserved confounders. We didn’t include multi-task RL as baselines since they do not target our challenges about causal effects, which makes the comparisons unfair to some extent. So we focus on more related baselines -- 10 causal RL and model-based RL baselines in Section 4.1 and Appendix C.6.\n\n>[1] Jin, Ying, Zhuoran Yang, and Zhaoran Wang. \"Is pessimism provably efficient for offline rl?.\" International Conference on Machine Learning. PMLR, 2021.\n>\n>[2] Yu, Tianhe, et al. \"Conservative data sharing for multi-task offline reinforcement learning.\" Advances in Neural Information Processing Systems 34 (2021): 11501-11516.\n\n**Q2. (Applicability of theoretical results): The result in Theorem 1 is applicable for a Tabular MDP, but this hasn't been discussed in the paper anywhere.**\n\nThe reviewer is correct that Theorem 1 is primarily applicable for the tabular case when the state and action space are finite. We added a remark under the theorem as the reviewer suggested: \n- **Our Theorem can be extended to continuous state space.** Theorem 1 provides the performance guarantees of our proposed method in tabular case with finite state and action space. It lays a solid foundation to carry out more general cases with continuous state and action space. The current results can be extended to continuous state space by using covering number argument in our key lemmas, while handling action space need more adaptation and we leave it as interesting future works. \n- **Continuous space is potentially reduced to tabular cases in practice.** Our proposed algorithm can be extended to continuous state or action space via some additional derivation on feature extraction. Prior works such as VQ-VAE [1] provide powerful encoders that can tokenize a continuous visual input into some discrete latent space with finite elements. We will add more discussion on the point in the revised manuscripts.\n\n**Q3. (Theoretical results justification): The result becomes void in the case when some (s_h, a_h) along the optimal policy has not been seen in the offline dataset.**\n\nThe reviewer is insightful that offline dataset **do** need to satisfy some properties to make offline learning possible and well-posed. To our best knowledge, the current weakest assumption for the offline dataset to ensure provably learning an optimal policy is called (clipped) single-policy concentrability assumption [2, 3] (such as Definition 1 in [3]) -- the offline dataset need to include sufficient data over the region (state-action tuples $(s_h, a_h)$) that the optimal policy can visit. Under such weakest assumption, our Theorem 1 can ensure the performance of the proposed algorithm. \n\n>[1] Van Den Oord, Aaron, and Oriol Vinyals. \"Neural discrete representation learning.\" Advances in neural information processing systems 30 (2017).\n>\n>[2] Li, Gen, et al. \"Settling the sample complexity of model-based offline reinforcement learning.\" The Annals of Statistics 52.1 (2024): 233-260.\n>\n>[3] Rashidinejad, Paria, et al. \"Bridging offline reinforcement learning and imitation learning: A tale of pessimism.\" Advances in Neural Information Processing Systems 34 (2021): 11702-11716.\n\n\n**Q4. (Comparison with existing theoretical works): How does the Theorem 1 fare to existing works? A detailed comparison on the improvement factors in the sample complexity would be useful.**\n\nWe thank the reviewer for raising this question -- which is our key theoretical technical contribution (the proof pipeline for Theorem 1). **Please refer to the general response for details due to the limited space.** We provide a shortened version here:\n\n- **Technical contributions: developed new proof pipeline.** We target a new problem --- bilinear MDPs with sparse core matrix $M$ --- brings new challenges since we need to use iterative algorithms with no closed-form solution, while prior works usually heavily depend on having a closed-form solution from the optimization problem. We develop a new proof pipeline that considers the accumulative errors during the process that we believe is of independent interest for future RL works. \n- **Sample complexity comparisons with prior works.** To our knowledge, we are the first to focus on the problem  —- bilinear MDPs with a sparse core matrix $M$ for causal deconfounding and generalization. So our sample complexity results can't be directly compared to existing works. But using prior works as references, it verifies the tightness of our results --- having a comparable dependency on all the salient parameters ($H$, $d$, $\\xi$).", "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": "kkoIy3Tku6", "reviewer": "Reviewer_XKke", "summary": "This paper aims to learn causal representations in offline RL. The authors propose a novel framework that incorporates causal graph learning through bilinear MDPs to manage mismatches in model learning and policy optimization stages. The framework models two types of confounders: ones affecting state dynamics and those influencing behavior policies. To enhance learning, the framework includes re-weighting and uncertainty quantification as well. Empirical results demonstrate the framework’s effectiveness in both causal discovery for model learning and generalization in offline RL policy learning.\n\nIn general, I hold a borderline but positive rating in this initial review. I will consider increasing my rating if the authors could give clarifications on my questions in the following sections (especially questions 1, 2, 3).", "questions": "I listed all the questions (mixed with weaknesses) in the above section.", "limitations": "Limitations have been discussed in the conclusion section. Though there is no in-depth discussion on societal impacts, I don't find any particular societal impacts in this paper compared with other RL works.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- [*Motivation and general framework*]: The idea of using causality to model the data generation process in RL (or in general, dynamic systems) gives a principle way to model RL environments. By considering and modeling potential confounders, which often lead to distribution shifts or mismatches, the paper addresses a critical challenge in RL. The decomposition of confounders into dynamics and policy-related aspects is reasonable and well-motivated in this work.\n\n- [*Technical soundness*]: The model and optimization processes are technically robust. The proofs and analyses provided in the appendix have been checked and are generally correct, offering complementary insights to the main part of the paper.\n\n- The writing is clear and easy to follow,", "weaknesses": "Since here most of the weaknesses and questions are mixed. So I put both here. One of the major weaknesses, which is also a primary question, is the inherent difficulty in fully learning the causal process through the bilinear MDP optimization (details in question 1-3).\n\n**Question 1**:  [Learning the causal graph by learning $M$] While learning M can partially capture the G graph, two main concerns: (1). Feature space entanglement: In the feature space $\\phi$ and $\\mu$, the causal variables might be entangled, potentially making the $G$ graph inaccurate (hard to identify from mixed sources). Please clarify if I have any misunderstanding on this; (2) optimization on the search space: Does the optimization search space of $M$ empirically impact the ability to learn the true $G$ graph?\n\n**Question 2**: [About assumption 1] The invariant causal graph assumption simplifies the problem but might be problematic. Specifically: (1). Mapping invariance: $M(u)$ may not be invariant since the scale of confounder effects matters, especially when some latent confounder effects are too small to model or identify in some domains (while in other domains are easy to identify); (2). Deatils on equivalence with assumptions in previous works: Further explanation, preferably with formulas, is needed to clarify the equivalence of this assumption with those in other works (e.g., invariant state representation, action effect) as mentioned in Remark 1.\n\n**Question 3**: [Multiple sources of confounders and identifiability of the causal graph]: How well can M capture different sources of confounders, especially state dynamics confounders $u_c$? How can we ensure the method separates these confounder sources, which is crucial in many RL environments with multiple unobservable latents?\n\n**Question 4**: [Further structure decomposition in world model]: Can the method incorporate algorithms like denoised MDPs to handle cases where state dynamics confounders affect only a subset of states? This could potentially make M more compact.\n\n**Question 5**: [Experimentation with raw pixel inputs]: Can the method work with raw pixel inputs, similar to experiments in works like IFactor cited in the paper? Will it still be effective when learning a visual encoder to extract states?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to express our gratitude to reviewer XKke for their comprehensive and insightful feedback. We are glad to know that the reviewer recognizes the clarity of our problem formulation, the novelty and technical soundness of the proposed method, as well as the theoretical and empirical results. We answer key clarification questions about the causal learning process through the bilinear MDP optimization below.\n\n**Q1. (Learning the causal graph by learning $M$)**: \n\nFor the first concern on the feature space entanglement, our learning process alternates between feature learning and causal mask learning, therefore, the causal variables $\\phi, \\mu$ will best align with the causal relationship. In some harder cases with poor offline data quality (like the *random* behavior policy), we do observe a performance drop and significant objective mismatch issues as is shown in Figure 6. However, our methods still outperform baselines in these cases with the help of bilinear causal representation. Besides, our method is also robust under limited sample size or various spurious level of confounders, as is demonstrated in Figure 4b and 4c. \n\nFor the second concern on the optimization on the search space, we'd like to first clarify the relationship between M and G under the formulation of bilinear MDP: \n$$\nG=\\\\left[\\\\begin{matrix} 0^{d\\\\times d} & M \\\\\\\\ 0^{d'\\\\times d} & 0^{d'\\\\times d'} \\\\end{matrix}\\\\right]\\\\in \\\\mathbb{R}^{(d+d')\\times (d+d')}.\n$$\n\nThe optimization of $M$ is a conditional independence testing (CIT) on $d\\times d'$ pairs of features. Empirically, our causal discovery process is doing CIT on all the elements in the $d'\\times d$ matrix. The optimization search space of $M$ is constrained by the size of $M\\in \\mathbb{R}^{d'\\times d}$. Since the causal discovery is based on hypothesis testing statistics like $\\chi^2$ testing, the ability to identify true causal graph $G$ will not be affected by the size of $M$. \n\n**Q2. (About assumption 1)**: \n\nWe thank the reviewer for asking insightful questions on our core assumption. \n\nRegarding the first concern on the mapping invariance, if the confounder effects are small as the reviewer mentioned, it means either (a) the confounder is not important to the success of our task in either domain or (b) the dataset we have at hand is too biased to fully learn about true causality. \n- For case (a), we empirically verify the invariance of the discovered causal graph with an experiment on the robustness to the spurious level of confounder, as is mentioned in RQ2 of Section 4.2. In this experiment, we vary the spurious levels (i.e. the scale of confounding effects) by changing the number of confounders in the environments. We show a consistent performance with fewer performance degradation compared to baselines without this bilinear causal representation, as is shown in Figure 4.c. As long as these latent confounders have enough statistical evidence to be identified with causal discovery, we can learn good causal representation. \n- For case (b), which mostly corresponds to the 'random' policy setting where the behavior policy is heavily confounded by the policy confounder $u_\\pi$, BECAUSE empirically has better performance compared to other causal RL or MBRL methods. \n\nIn conclusion, in either case (a) or (b) raised by the reviewer, BECAUSE can empirically outperform other MBRL baselines, which justifies our assumption of invariant causal mask. \n\nRegarding the second concern on the equivalence with assumptions in previous works, we add some additional discussion in **Table I of our 1-page rebuttal supplementary** with a mathematical description of the specific meaning of the invariant causal graph and how it relates to the assumptions cited in our manuscripts. \n\n**Q3: (Multiple sources of confounders and identifiability of the causal graph)**: \n\nThe learning process of $M$ is conducted with two stages: regularized world model learning to deconfound $u_c$ in Equation (5) and (6), and policy-conditioned reweighting formulas to deconfound $u_\\pi$ in (7). \n- The first stage is to use the $\\ell_0$ objective to regularize the world model loss while maintaining a sparse $M$. \n- In the second stage, we average $M$ by reweighting its expectation across different trajectory batches. \n\nWe demonstrate the effect of the above deconfounding process both empirically and theoretically:\n- **Experiment** results show that our method outperforms the existing causal RL and model-based RL baselines in both in-distribution and OOD settings in three environments. \n- **Theoretically**, we provide additional analysis on how well this de-confounding process can identify the causal matrix $M$. Under reasonable assumptions about the offline dataset and underlying dynamics model, we can bound the estimation error of $M$ and have performance guarantees for learned policy in Theorem 1.\n\n**Q4: (Further structure decomposition in world model)**: \n\nWe thank the reviewer for providing this inspiring extension of our work. As we discuss in the related works, Denoised MDPs propose a structure decomposition according to the controllability and reward relevance of the state. To ground our bilinear representation in this context, one interesting way is to decompose $\\phi, \\mu$ into the controllable/non-controllable and reward-relevant/irrelevant parts. This will be an interesting future direction to make $M$ more compact. \n\n**Q5: (Experimentation with raw pixel inputs)**: \n\nTo further demonstrate the scalability of our method, we conduct additional experiments with raw-pixel inputs in the *Unlock* environments. The RGB input has a shape of (192, 192, 3), the detailed results can be found in the general response.", "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": "zyPgCVAv28", "reviewer": "Reviewer_ne9G", "summary": "The paper introduces a method to improve model-based offline reinforcement learning by addressing the objective mismatch problem using causal discovery methods. This problem stems from the fact that in this setup the learning algorithm aims to solve the following two problems at once: i) accurate learning of the transition dynamics, ii) identifying the optimal policy on the MDP characterized by the learned transition dynamics. However in practice, an improved solution to one of these two problems does not translate to a better solution to the other one, while in theory it should. The paper's central hypothesis is that this disconnect is caused by the spurious correlations learned by the model which can be blocked by a causality based formulation.", "questions": "Why was it not possible and/or sensible to study the problem in the most well-known benchmark of offline RL research, the MuJoco environments of D4RL?", "limitations": "The paper discusses the limitations of the proposed approach in the final few sentences of the conclusion section. However, it does not discuss the potential negative societal impact of the presented work.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "* The proposed methodology is very sensible and interesting. Formulating a latent bilinear MDP to enable differentiable causal discovery approaches is a brilliant idea.\n\n* The paper studies the theoretical properties of the presented approach. Although the used theoretical tools are not novel per se, they serve for their purpose excellently. The conceptual connection between causal representation learning and bilinear MDPs is indeed novel and worthwhile acknowledging.\n\n* Despite the denseness of the technical material, the paper is very easy to read and understand.\n\n* The presented empirical results are satisfactory.", "weaknesses": "* The compared baselines do make sense. Some of them such as ICIL and CCIL are causal RL approaches. However, the experiments part could be even stronger if the chosen model-based ORL approaches represented the state of the art a little better. TD3+BC and MOPO are outdated. One could for instance consider to compare against a stronger baseline within the same family, such as:\n\nSun et al., Model-Bellman Inconsistency for Model-based Offline Reinforcement Learning, ICML, 2023.\n\n* It would be interesting to see a demonstration of how well BECAUSE can recover the underlying confounders of a simple MDP. This could be an interesting research question to be investigated in the experiments section. This would enhance the value of the proposed approach from an instrumental to an explanatory one.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for their insightful and inspiring feedback. We are glad to know that the reviewer recognizes the novelty of our contributions, the clarity of our problem formulation, and the empirical results. We answer reviewers' questions about the causal learning process through the bilinear MDP optimization below.\n\n**Q1. (Advanced Model-based Offline RL baselines): The compared baselines do make sense. Some of them such as ICIL and CCIL are causal RL approaches.... consider to compare against a stronger baseline within the same family, such as: (Sun et al., ICML, 2023.).**\n\nWe thank the reviewer for proposing this important baseline. We add additional experiments on the *Unlock* environment and illustrate the results in the first table in the general response. While MOBILE indeed demonstrates promising performance in the OOD generalization setting compared to the MOPO baselines, we still observe significant advantage on our methods that learns bilinear causal representation in the world model and use them for planning. \n\n**Q2. (Discussing recovery and explainability of confounders): It would be interesting to see a demonstration of how well BECAUSE can recover the underlying confounders... enhance the value of the proposed approach from an instrumental to an explanatory one.**\n\nAs you can see in **Figure (a) and (b) in the rebuttal supplementary**, we give a simple MDP example in the autonomous driving tasks. We have the states ['time-to-collision', 'daytime'] and actions ['throttle', 'steering']. (Time-to-collision (TTC) is the ratio of relative distance to the closest vehicle divided by the velocity, and the TTC gets closer to 0 if a collision happens). \n\nIn general, the offline dataset encounters two sources of confounders: \n\n- **Confounders from suboptimal behavior policy:** As is shown in **Figure (a)**, $u_\\pi$ is the confounder introduced by suboptimal behavior policies. For example, one aggressive driver always overtakes the front vehicles. This sub-optimal behavior policies $\\pi_\\beta(a | s)$ may have large throttle even when the 'time-to-collision' is extremely small. Such behavior policies introduce confounders that cause spurious correlations between states ('ttc') and actions ('throttle'). This spurious correlation may cause collision if we directly learn from this dataset and deploy online. In our experiment, we denote different behavior policies as 'random', 'medium', 'expert' policies, leading to different confounded levels of $u_\\pi$. \n\n- **Confounders from environment dynamics**: $u_c$ is the confounder that exists in the transition dynamics $T(s'|s, a)$, and it causes spurious correlation between different state and action components in transition dynamics. In the driving example, as shown in **Figure (b)**, the traffic is more crowded (lower ttc) during the peak hours in a day and less crowded (higher ttc) in the rest. With confounder between next 'ttc' and current 'daytime', raw model fails to predict well under crowded scenarios out of peak hours, because it will falsely predict a lower 'ttc' outside peak hours even if the actual traffic is dense. This is because the model is misled by the spurious correlation between 'time-to-collision' in $s'$ and daytime in $s$. In our experiment, we also have different confounders in different tasks in all 3 environments (see Appendix C.5.)\n\nAdditionally, we are running some numerical simulations in this simple example, BECAUSE can indeed help identify the transition dynamics. We plan to add this to our final experiment results. \n\n**Q3. (Discussing selected benchmark): Why was it not possible and/or sensible to study the problem in the most well-known benchmark of offline RL research, the MuJoco environments of D4RL?**\n\nOur experiment simulator is adopted by some published works [1, 2, 3] in the causal RL domain. The selected environments have strong complexity in causal structure that leads to distribution mismatch between different training environments. Moreover, they are related to some real-world problems, such as indoor navigation, manipulation, and autonomous vehicles. These tasks all require strong reasoning capability to understand the interaction between the decision-making agents and other static or dynamic environment entities. \nOn the other hand, in the mujoco-based D4RL environments, such cause-and-effect structures mainly focus on some physical transformation of the locomotion tasks, which have lower complexity as they do not model the interaction between different entities using the causal structure. \n\n>[1] Wang, Zizhao, et al. \"Causal Dynamics Learning for Task-Independent State Abstraction.\" International Conference on Machine Learning. PMLR, 2022.\n>\n>[2] Ding, Wenhao, et al. \"Generalizing goal-conditioned reinforcement learning with variational causal reasoning.\" Advances in Neural Information Processing Systems 35 (2022): 26532-26548.\n>\n>[3] Ding, Wenhao, et al. \"Seeing is not believing: Robust reinforcement learning against spurious correlation.\" Advances in Neural Information Processing Systems 36 (2023).\n\n**Q4. (Adding societal impact): The paper discusses the limitations of the proposed approach in the final few sentences of the conclusion section. However, it does not discuss the potential negative societal impact of the presented work.**\n\nWe thank the authors for raising the concern about potential negative societal impact. We added the following dicussion of societal impacts in our appendix:\n\n*Incorporating causality into model-based reinforcement learning framework provides the explanability of decision making, which helps humans understand the underlying mechanism of algorithms and check the source of failures. However, the learned causal graph may contain human-readable private information about the offline dataset, which could raise privacy issues. To mitigate this potential negative societal impact, the discovered causal graphs should be only accessible to trustworthy users.*", "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": "DJqNJUIZ5y", "reviewer": "Reviewer_8enR", "summary": "Distribution shifts in the sampling of MBRL will lead to the objective mismatch between model and policy learning.\nTherefore, this paper aims to reduce the influence of the distribution shift for MBRL.\nIt divides the confounders into  $\\mu_{\\pi}$ and $\\mu_{c}$ and then tries to capture causal representations to reduce the influence of the distribution shift, thus mitigating the objective mismatch problem.", "questions": "1. What is the world model in this paper? Does this paper use a world model to train the policy?", "limitations": "N/A", "rating": 2, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. This paper is well-written and easy to follow.\n2. In RL, this paper tries to address confounders, which are key problem in causal inference.\n3. There are some new theoretical results in this paper.\n4. Experiment results illustrate the effectiveness of the proposed method.", "weaknesses": "1. The conditions in Bilinear MDPs are difficult to achieve in real-world environments.\n2. Some real-world examples can be given for confounders $\\mu_{\\pi}$ and $\\mu_{c}$.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to express our gratitude to reviewer 8enR for their valuable feedback. The reviewer acknowledges the clarity of our problem formulation and presentation, as well as the theoretical and empirical contribution of this work. We will answer the main questions below: \n\n**Q1. (Generality of Bilinear MDP Formulation): The conditions in Bilinear MDPs are difficult to achieve in real-world environments. Some real-world examples can be given for confounders $u_\\pi$ and $u_c$.**\n\nWe thank the reviewer for raising these important questions about the generality of bilinear MDP formulation. We discuss two categories of confounders in our experiment, which simulate some real-world problems like indoor navigation, robotic manipulation, and self-driving cars. One real-world example about the two categories of confounders $\\mu_\\pi$ and $\\mu_c$ is in the autonomous driving environments. We have the states ['time-to-collision', 'daytime'] and actions ['throttle', 'steering']. Time-to-collision is the minimum ratio of relative distance to the closest vehicle divided by the velocity. The offline dataset may have two sources of confounders: \n- **Confounders from suboptimal behavior policy:** $u_\\pi$ is the confounder introduced by suboptimal behavior policies. For example, one aggressive driver always overtakes the front vehicles. This sub-optimal behavior policy $\\pi_\\beta(a | s)$ may have a large throttle even when the 'time-to-collision' is extremely small. Such behavior policies introduce confounders that cause spurious correlations between states ('time-to-collision') and actions (throttle). This spurious correlation may cause collision and have low reward if we directly learn from this dataset and deploy it online. In our experiment, we denote different behavior policies as 'random', 'medium', and 'expert' policies, leading to different confounded levels of $u_\\pi$. \n\n- **Confounders from environment dynamics**: $u_c$ is the confounder that exists in the transition dynamics $T(s'|s, a)$, and it causes spurious correlation between different state and action components in transition dynamics. In the real-world example, the traffic will be more crowded during the peak hours in a day and less crowded in the rest. Due to the confounder for 'crowdedness' and 'daytime', the model may fail under crowded scenarios out of peak hours, because it will falsely predict a lower 'time-to-collision' outside peak hours even if the traffic is dense. This is because the model is misled by the spurious correlation between 'time-to-collision' in $s'$ and daytime in $s$. In our experiment, we also have different confounders in different tasks in all 3 environments. We have discussed the environments in Appendix C.5. \n\nWe visualize the above confounded case in the MDP as **Figure (a) and (b) in the rebuttal supplementary**. \n\n**Q2. (Clarification of World Model): What is the world model in this paper? Does this paper use a world model to train the policy?**\n\nWe thank the authors for raising the clarification about the world model. In our setting, our world model is composed of a dynamics model $T(s' | s, a)$ and a reward model $r(s, a)$. More concretely, during the planning phase, we have the following pessimistic value iteration: \n$$\n\\overline{Q}(s, a) = r(s, a) - E_\\theta(s, a) + \\sum_{s'\\in \\mathcal{S}} \\widehat{T}(s'|s, a) \\widehat{V}(s'). \n$$\n\nTherefore, optimization of our pessimistic policy will be dependent on the learned dynamics model $\\widehat{T}(s' | s, a)$ and reward function $r(s, a)$.", "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:
## **General Response**

We thank all the reviewers for their dedicated efforts in providing valuable feedback to our work. All of the reviewers acknowledge our clarity of problem formulation and the technical soundness in bridging bilinear MDP with causality. 
In this general response, we first clarify some important technical contributions, then we attach two sets of new experiment results: (1) a comparison with an advanced MBRL baseline, and (2) another experiment with raw pixel inputs. 

### **Theoretical comparison with existing works**: for both technical contributions and final sample complexity results
Recall that we target bilinear MDP in the offline setting with the transition kernel $T = \phi^T M \mu$, where the core matrix $M$ is assumed to be sparse and unknown (represent the causal graph). 

- **New challenges in our problems: no closed-form optimization solutions.** We aim to solve bilinear MDPs with sparse core matrix $M$, which leads to a $\ell_0$ regression problem that is usually solved by iterative algorithms without a closed-form solution. Such a closed-form situation brings daunting challenges for developing provable suboptimality guarantees since we can’t explicitly write the error term to be controlled out anymore. Existing proof pipelines in prior works for (bi)-linear MDPs can’t be applied here since they heavily depend on having a closed-form optimization solution  [1, 2, 3]. 
- **Our new proof pipeline for non-closed form optimization problems.**  Prior works using ridge regression only need to control the statistical error of the final closed-form output without focusing on the optimization process. While for our $\ell_0$-regularized optimization, as no closed-form solution can be given, we consider the accumulative errors throughout the optimization process, summing errors from each iteration to determine the final error terms. This new proof pipeline potentially generalizes to a wide range of RL problems without closed-form solutions, which we believe is of independent interest to RL theory, especially for RL algorithms involving iterative methods like Lasso.
- **Sample complexity comparisons with prior works.** First, we would like to highlight that we target a different problem  —- bilinear MDPs with a sparse core matrix $M$ for causal deconfounding and generalization. To our knowledge, we are the first to focus on this problem, so our sample complexity results can't be directly compared to existing works. As suggested by the reviewer, we will include the sample complexity results of existing works targeting bilinear MDPs [1-3] as references, verifying the tightness of our results.  Specifically, we denote the size of the latent space as $d$ and the planning horizon as $H$, for any $\xi$-optimal policy, the sample complexity in [2, 3] is bounded with $\mathcal{O}(\frac{H^7 d^2 \log(H^2d)}{\xi^2} \log^2(\frac{Hd}{\xi}))$ and $\mathcal{O}(\frac{H^4d}{\xi^2})$, respectively.  In our work, this sample complexity is bounded with a rate of $\mathcal{O}(\frac{\|M\|_0 H^2 \log (d)}{\xi^2 } \log^2(\frac{\|M\|_0}{\xi}))$, where $\|M\|_0\ll d^2$ is the sparsity level of core matrix $M$. Our work has a comparable dependency on all the salient parameters (H, d,$\xi$) compared with prior works. 

>[1] Yang, Lin, and Mengdi Wang. "Reinforcement learning in feature space: Matrix bandit, kernels, and regret bound." International Conference on Machine Learning. PMLR, 2020.
>
>[2] Du, Simon, et al. "Bilinear classes: A structural framework for provable generalization in rl." International Conference on Machine Learning. PMLR, 2021.
>
>[3] Zhang, Weitong, et al. "Provably efficient representation selection in low-rank Markov decision processes: from online to offline RL." Uncertainty in Artificial Intelligence. PMLR, 2023.

### **Additional experiments**


We illustrate the success rate (%) comparison between **MOBILE**, **MOPO** and **Ours** (BECAUSE) in the *Unlock* environments with six evaluation settings. We run experiments with 10 random seeds and report the mean and 95% confidence interval with t-testing. 
We can see some improvements of MOBILE over MOPO under some OOD settings, yet the gap between MOBILE and our methods is still significant. 

|Env|MOPO|MOBILE|Ours|Env|MOPO|MOBILE|Ours|
|-|-|-|-|-|-|-|-|
|Unlock-I-R|21.5(1.9)| 15.9(1.0) | **32.7(2.8)** | Unlock-O-R | 16.6 (1.3) | 12.8(0.8) | **27.6(2.0)** |
|Unlock-I-M|84.8(5.1)| 72.4(1.7) | **98.0(4.9)** | Unlock-O-M | 39.5(4.7)| 40.7(1.8) | **68.8(1.5)** |
|Unlock-I-E|88.8(4.6)| 78.3(1.2) | **97.4(1.0)** | Unlock-O-E | 39.9(4.4) | 45.6(2.1) | **82.1(6.5)** |

----------

We experiment on **RGB images** instead of vector observation of *Unlock*, and compare our method with two of our baselines, ICIL (model-free) and IFactor (model-based). We run experiments with 10 random seeds and report the mean success rate (%) and its 95% confidence interval with t-testing.  All the methods have performance decay compared to the vector state setting, yet our method can still prevail under visual inputs: 
|Env | ICIL | IFactor | Ours | Env | ICIL | IFactor | Ours |
|-|-|-|-|-|-|-|-|
|Unlock-I-R|0.8(0.8) | 4.3(1.1) | **15.7(3.3)** | Unlock-O-R | 1.5(1.8) | 4.7(1.6) | **5.9(0.9)** |
|Unlock-I-M|5.3(2.0) | 30.2(4.1) | **62.0(4.6)** | Unlock-O-M | 8.6(4.2) | 15.4(2.4) | **71.6(9.1)** |
|Unlock-I-E|8.7(3.4) | 34.0(4.8) | **63.7(3.9)** | Unlock-O-E | 17.1(4.2) | 16.7(3.1) | **73.6(19.5)** |

In addition to the above general response, we also illustrate some tables that discuss related work, additional experiment details and new illustration examples in the one-page supplementary material.

Meta-review: This paper studies the objective mismatch in offline model-based RL between learning the transition dynamics and the optimal policy of the MDP. Improving one objective does not necessarily improve the other in practice. It hypothesizes the disconnection is caused by the spurious correlations and proposes to solve the problem with a bilinear causal representation.

**Pros**
- Most reviewers consider the paper is well written and easy to follow
- The method is well motivated and the proposed method is technically sound
- Thorough theoretical analysis
- Experiments and ablation study show the benefit of the proposed method

**Cons**
- Multiple reviewers were concerned on the assumption of invariant caual graph and bilinear causal structure. The authors provided detailed discussion on its applicability in real problems.
- Reviewer aSxm pointed out some missing related work, which authors have provided additional disucssion.
- Reviewer ne9G was concerned of a lack of experiments with stronger baselines. The authors provided more experiments accordingly.
- Reviewer 8enR is concerned of the use of term "world model" that led to a lengthy discussion. While there's still disagreement on whether its appropriate to use "world model", it does not impact the core contribution of the this submission.

I would strongly encourage the authors to incorporate their additional remarks and clarification in the rebuttal in their final response. Please also make a clear definition of the "world model" or dynamics model, if they decide to rename it to, in the revision to avoid ambiguities from the literature.
Decision: Accept (poster)