Title: BECAUSE: Bilinear Causal Representation for Generalizable Offline Model-based Reinforcement Learning
Venue: NeurIPS 2024
Accepted: yes
Keywords: Model-based RL

Causal Reasoning

Offline RL
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

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.

1 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 [1]. Specifically, offline model-based reinforcement learning (MBRL) [2, 3, 4], learning policies with an estimated world model, generally perform better than their model-free counterparts in long-horizon tasks such as self-driving vehicles [5], robotics [6], and healthcare [7]. 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 [8].

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 failure of trained policy at a slightly different deployment stage. Furthermore, since the learning objectives of the world models and policies are isolated from each other, a significant challenge in offline MBRL is objective mismatch [9, 10] problem (shown in Figure 1): models that achieve a lower training loss are not nec-

Policy

𝜋(𝑎|𝑠)

Model

𝑇 ( (𝑠

!

|𝑠, 𝑎)

Environment

𝑇(𝑠′|𝑠, 𝑎)

{𝑠, 𝑎, 𝑠

!

, 𝑟}

Data collection with behavior policy

Model Loss

low model loss

≠

success

low model loss

≈

success

test

Our method

(Unknown)

Success

Failure

essarily better for control performance. For example, a dynamics model achieve relatively low

prediction loss, yet it may not be sufficient to guide the planner to the high-reward region. Previous works [9, 10] have attempted to 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 [8].

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 taskrelevant factors, models in offline RL memorize correlations without learning the causality. The sub-optimal behavior policies introduce spurious correlations [11] 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, 13] 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, we introduce the BilinEar CAUSal rEpresentation (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 in the world model, then use this learned representation to facilitate planning by quantifying the uncertainty of sampled transition pairs. Consequently, we factorize the spurious correlations and learn a 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 is the first work that systematically reveals the connection between causal representation learning and Bilinear MDPs.

We propose BECAUSE, an empirical causal representation framework, based on the above formulation. 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.

2 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 , A , T , H, r } , which is composed of state space S , action space A , a set of transition functions T , planning horizon H and reward function r associated with task preferences. Without loss of generality in many real-world practices, we assume that the reward function is bounded by r h ∈ [0 , 1] , ∀ h ∈ [ H ] . Specifically, we are interested in a goal-conditioned reward setting, where ∀ g ∈ S , r ( s, a ; g ) = 1 if and only if s = g .

Given a policy π and the state-action pair ( s, a ) ∈ S × A , we then define the state-action value function in the timestep h as Q π h ( s, a ) = E π [ ∑ H i = h r i ( s i , a i ) | s h = s, a h = a ] , and the value function V π h ( s ) = E π [ ∑ H i = h r i ( s i , a i ) | s h = s ] . The expectation E π here is integrated into randomness throughout the trajectory, which is essentially induced by the random action of the policy a i ∼ π ( ·| s i )

and the time-homogeneous transition dynamics of the environment s i +1 ∼ T ( ·| s i , a i ) , ∀ i ∈ [ h, H ] . In the offline dataset, the data rollouts can be seen as generated by some (mixed) behavior policy π β , resulting in a dataset D with in total n samples { s i , a i , s ′ i , r i } 1 ≤ i ≤ n .

Definition 1 (Bilinear MDP [14]) . For each ( s, a ) ∈ S × A , s ′ ∈ S , we have the corresponding feature vector ϕ ( · , · ) : R | S | × R | A | → R d , µ ( · ) : R |S| → R d ′ . With some core matrix M ∈ R d × d ′ , we can represent the transition function kernel T ( ·|· , · ) as

where ϕ ( s, a ) and µ ( 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 to 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 collection process, and define action-state confounded MDP (ASC-MDP):

Definition 2 (ASC-MDP) . Besides the components in standard MDPs M = { S , A , T, H, r } , we introduce a set of unobserved confounders u . In ASP-MDP, confounders are factorized as u = { u π , u c } 1 ≤ h ≤ H , where u π ∈ U denotes the confounders between s and a ∼ π β ( s ) induced by behavior policies, and u c ∈ U denotes the confounders within the state-action pairs of the environment transition, that is, the inherent structure between ( s, a ) and s ′ . Here we

𝑠

𝑠′

𝑎

𝑢

!

𝑢

!

′

𝑢

"

𝑠

𝑠′

𝑎

𝑢

𝑠

𝑠′

𝑎

𝑢

State Nodes

Action Nodes

Confounders

(a) Confounded MDP

(b) SC-MDP

(c) ASC-MDP (Ours)

assume a time-invariant confounder distribution u ∼ P u ( · ) , ∀ h ∈ [ H ] , which is a common assumption [15, 16, 17]

The resulting causal relationship of ASC-MDP is demonstrated in Figure 2. Originating from the original MDP, ASC-MDP is different from the Confounded MDP [18] and State-Confounded MDP (SC-MDP) [19] in that it models both the spurious correlation between the current state s and the 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 in ASC-MDP aligns with the source of spurious correlation in offline MBRL.

3 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 ϕ ( s, a ) and µ ( s ′ ) , 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 a linear confounded MDP [18]:

where u ∼ P u ( · ) . Inspired by the Bilinear MDP in Definition 1, we decompose ˜ ϕ ( s, a, u ) into a confounder-aware core matrix M ( u ) and a feature mapping ϕ ( s, a ) , which factorize the influence of

the confounders. Given the factorization of confounder u = { u c , u π } in Definition 2, we derive via d-separation in the graphical model in Figure 2 that s ′ ⊥ ⊥ u π |{ s, a, u c } . As a result, we only need to consider the confounder u c from the environment when decomposing the transition model:

Definition 3 (Construction of causal graph G ) . In ASC-MDP, G = [ 0 d × d M 0 d ′ × d 0 d ′ × d ′ ] . for all (sparse) core matrix M , the causal graph G is bipartite, thus ∀ 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 u and estimate the unconfounded transition model T ( s ′ | s, a ) , one way is to identify the causal structures induced by the confounder u c for the transition dynamics [20]. 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:

Since in our case, M is a sub-matrix of 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 p ( D ; ϕ, µ, M ) -λ | M | . We elaborate Definition 3 and show the relationship between core matrix M and causal graph G in Appendix A.3.

We also make the assumption that the sparse G and M remain invariant with different environment confounders in the offline training and online testing.

Assumption 1 (Invariant causal graph) . We denote the causal graph G under confounder u as G ( u ) . The generalization problem that we aim to solve satisfies the invariance in the causal graph G ( u c ) , where G ( u c ) = G ( u ′ c ) , M ( u c ) = M ( u ′ c ) , u c and u ′ c are the confounders in training and testing.

Remark 1 . The assumption 1 can also be interpreted as task independence in [24], invariant state representation [25], or invariant action effect in [26]. See detailed comparison in appendix Table 4.

3.2 Learning Causal Representation from Offline Data

We first learn the causal world model T ( s ′ | s, a ) in the presence of confounders u in the offline datasets. As formulated in ASC-MDP 2, there are two sets of confounders: u π and u c . To estimate an unconfounded transition model and remove the effect of confounder, we first remove the impact of u c which comes from the dynamics shift by estimating a batch-wise transition matrix M ( u c ) , then we apply a reweighting formula to deconfound u π induced by the behavior policies and mitigate the model objective mismatch.

As discussed in Definition 3, we only need to optimize the part of the parameters of the causal graph G , i.e. M . Thus, we can remove the constraints in (4), then transform the original causal discovery problem into a regularized MLE problem as follows:

Planning

World Model

…

···

···

···

Uncertainty

start

goal

𝑠, 𝑎

𝑠′

Offline

Buffer

𝑇 " (𝑠

!

|𝑠, 𝑎)

𝑢

𝜙(⋅,⋅)

𝜇(⋅)

𝑀 )

Uncertainty Quantifier

𝐸

!

(s, a)

Break correlation

where K µ := ∑ s ′ ∈S µ ( s ′ ) µ ( s ′ ) T is an invertible matrix. The derivation of Equation (5) is elaborated in Appendix A.4. In practice, we use the χ 2 -test 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 a more detailed implementation in Appendix C.1.

Estimating the core mask provides a more accurate relationship between state and action representations, 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 to the transition model loss and spectral norm regularization [28] to satisfy the regularity constraints of the feature in Assumption 3:

The world model learning process is illustrated in Figure 3. The estimation of individual M ( u c ) mitigates the spurious correlation brought by u c . To further deal with the spurious correlation in u π induced by the behavior policy π β ( a | s, u π ) , 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 as the soft intervention approach [18, 29] to estimate the treatment effect in the transition function T in an unconfounded way:

The derivation of Equation (7) is illustrated in Appendix A.5. Equation (7) basically shows a re-weighting process given the empirical estimation of M ( u c ) in every batch of trajectories: M ( u ) = E pu [ M ( u c ) π β ( a | s,u π )] E pu π β ( a | s,u π ) . Compared to the general reweighting strategies in previous MBRL literatures [18, 29] which reweights the entire value function, this re-weighting process is conducted only on the estimated matrix, while the representation ϕ ( s, a ) and µ ( s ′ ) are subsequently regularized by weighted estimation of M . The pipeline of causal world model learning is described in the first part of the Algorithm 1. We discuss more details of the implementation and experiment in Appendix C.

3.3 Causal Representation for Uncertainty Quantification

To avoid entering OOD states in the online deployment, we further design a pessimistic planner according to the uncertainty of the predicted trajectories in the imagination rollout step to mitigate objective mismatch.

We use the feature embedding from bilinear causal representation to help quantify the uncertainty, denoted as E θ ( s, a ) . As we have access to the offline dataset, we learn an Energybased Model (EBM) [30, 31] based on the abstracted state representation ϕ and core matrix M . A higher output of the energy function E θ ( · , · ) indicates a higher uncertainty in the current state as they are visited by the behavior policies π β less frequently. In practice, the energy-based model usually suffers from a high-dimensional data space [32]. To mitigate this overhead of training a good uncertainty

quantifier, we first embed the state samples through the abstract representation µ ( s ′ ) , and the state action pair via ϕ ( s, a ) .

where µ ( s ) + refers to the positive samples from the approximated transition dynamics ̂ T ( ·| s, a ) , and µ ( s -) refers to the latent negative samples via the Langevin dynamics [30]. Additionally, we 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 E θ ( s, 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 on Model Predictive Control (MPC) [33]. At timestep h , we basically subtract the original step return estimation r h ( s, a ) by its uncertainty E θ ( s, a ) :

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 3 on the feature's existence and regularity, we achieve the finite-sample complexity guarantee - an upper bound of the suboptimality gap as follows, whose proof is postponed to Appendix B.

Theorem 1 (Performance guarantee) . Consider any 0 < δ < 1 and any initial state ˜ s ∈ S . Under the Assumption 2, 3 and that the transition model T is an SCM (defined in 4), for any accuracy level 0 ≤ ξ ≤ 1 , with probability at least 1 -δ , the output policy π of BECAUSE (Algorithm 1) based on the historical dataset D with n = ∑ ( s,a ) ∈S×A n ( s, a ) samples generated from a behavior policy π β satisfies:

where C 1 , C s are some universal constants, σ is SCM's noise level (see Definition 4), and M ∈ R d × d ′ 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 noise level σ , and the ℓ 0 norm of the ground true causal mask M , which describes the intrinsic complexity of the world model.

Consequently, with Proposition 1 in the Appendix, we can achieve ξ -optimal policy ( V ∗ 1 ( ˜ s ) -V π 1 ( ˜ s ) ≤ ξ ) as long as the historical dataset satisfies the following conditions: ∀ 0 ≤ ξ ≤ 1 ,

4 Experiment Results

In this section, we conduct a comprehensive empirical 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.

4.1 Experiment Setting

Environment Design We design 18 tasks in 3 representative RL environments in Figure 4. Agents need to acquire reasoning capabilities to receive higher rewards and achieve goals.

Lift-Medium

Lift-Expert

Unlock-Random

Unlock-Medium

Unlock-Expert

Crash-Random

Crash-Medium

Crash-Expert

Unlock-Random

Unlock-Random

Unlock-Medium

Unlock-Medium

Unlock-Expert

Unlock-Expert

(a)

(b)

(c)

Lift : Object manipulation environment in RoboSuite [34]. We designed this environment for the agent to lift an object with a specific color configuration on the table to a 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 , there will be a 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 a risky scenario where an A V collides with a jaywalker because its view is blocked by another car [36]. We design such a crash scenario based on highway-env [37], where the goal is to create crashes between a pedestrian and A Vs. In the OOD environment Crash-O , the distribution of reward (number of pedestrians) is different in online testing environments.

For all three different environments, we set a specific subset of the state space as the goal g ∈ S , and the reward is defined as the goal-reaching reward r ( s, a, g ) = I ( r = g ) . When the episode ends in the goal state within the task horizon H , the episode is considered a success. We then use the average success rate as the general evaluation metrics for our BECAUSE and all baselines.

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 as the offline demonstration data, and the number of episodes is 1,000 in the environments Lift and Crash , which all have continuous state and action space. A more detailed view of the environment hyperparameters and behavior policy design is in Appendix C.5.

Baselines We compare our proposed BECAUSE with several offline causal RL or MBRL baselines. ICIL [38] learns a dynamic-aware invariant causal representation learning to assist a generalizable policy learning from offline datasets. CCIL [39] 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 in the state prediction loss in the high-reward region. Delphic [40] 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

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 a 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 IFactor [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.

RQ1 : 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 u π ), and different environment contexts (different u c )?

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 generalizability 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 π β (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 π ). Where O here indicates the tasks under unseen environment with confounder u ′ c = u c 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 Table 6, 7 and the causal masks discovered in each environment in Appendix Figure 8 for reference.

For RQ2 , we compare different aspects of BECAUSE'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 in Figure 5(b). The result shows that, for any selected number of samples, BECAUSE consistently outperforms MOPO with a clear margin. We also evaluate BECAUSE performance at higher spurious levels in Figure 5(c). We add up to 8 × 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 parameterize the feature encoder as a three-layer CNN with 128 dimensions hidden size for all the baselines and our methods. The results in Table 1 show that BECAUSE can significantly improve both in-distribution and out-of-distribution performance under different quality of behavior policies.

For RQ4 , we aim to understand whether BECAUSE achieves higher performance by resolving the objective mismatch problem. We first collect two groups of trajectories: τ pos and τ neg , 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 L model ( τ pos ) < L model ( τ neg ) . According to our visualization in Figure 6, in Unlock-Expert and Unlock-Medium , the ratio of τ pos is much higher in BECAUSE than MOPO among the trajectories with low model loss. In Unlock-Random , the mismatch of the model and

MOPO

BECAUSE

Example Policies

Positive Reward

Negative Reward

Unlock-random

Unlock-medium

Unlock-expert

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 a hypothesis test with H 0 : L model ( τ pos ) < L model ( τ neg ) . 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 average success rate across nine in-distribution and nine out-of-distribution tasks in Table 2. The Optimism variant conducts optimistic planning instead of pessimistic planning in Equation (9), which uses uniform 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 without iterative update. We report the results of the task-wise ablation with confidence interval and significance in Appendix Table 8 and 9.

5 Related Works

Objective Mismatch in MBRL The objective mismatch in MBRL [43, 44] refers to the fact that pure MLE estimation of the world model does not align well with the control objective. Previous works [29, 45] propose reweighting during model training to alleviate this mismatch, [46] proposes a goal-aware 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 a 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 setting, Model-based RL [2, 3, 4, 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.

Causal Discovery with Confounder Most of the existing causal discovery methods [51] can be categorized into constraint-based and score-based. Constraint-based methods [52] start from a complete graph and iteratively remove edges with statistical hypothesis testing [53, 54]. 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] or 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 such correlations [63, 64, 65] so that models can effectively identify elusive confounders.

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 an online setting. These structures can be used to constrain the output space [19] or to adjust the buffer priority [69]. Building dynamic models in model-based RL [24, 70, 71] based on causal graphs 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 via structural regularization [23, 72, 73], conditional independence test [24, 70, 74], variational inference [75, 12], counterfactual data augmentation [76, 77], hierarchical skill abstraction [78, 79], uncertainty quantification [40], reward redistribution [24, 80], causal context modeling [81, 82] and structure-aware state abstraction [12, 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 in 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 on 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 time-homogeneous 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 high-dimensional observations to discover concept factors in long-horizon visual RL settings.

References

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.

Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. Advances in Neural Information Processing Systems , 33:14129-14142, 2020.

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-21823. Curran Associates, Inc., 2020.

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 , 34:28954-28967, 2021.

Zeyu Zhu and Huijing Zhao. A survey of deep rl and il for autonomous driving policy learning. IEEE Transactions on Intelligent Transportation Systems , 23(9):14043-14065, 2021.

Aviral Kumar, Anikait Singh, Frederik Ebert, Mitsuhiko Nakamoto, Yanlai Yang, Chelsea Finn, and Sergey Levine. Pre-training for robots: Offline rl enables learning new tasks from a handful of trials. arXiv preprint arXiv:2210.05178 , 2022.

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.

Laixi Shi and Yuejie Chi. Distributionally robust model-based offline reinforcement learning with near-optimal sample complexity. arXiv preprint arXiv:2208.05767 , 2022.

Benjamin Eysenbach, Alexander Khazatsky, Sergey Levine, and Russ R Salakhutdinov. Mismatched no more: Joint model-policy optimization for model-based rl. Advances in Neural Information Processing Systems , 35:23230-23243, 2022.

Shentao Yang, Shujian Zhang, Yihao Feng, and Mingyuan Zhou. A unified framework for alternating offline model training and policy learning. Advances in Neural Information Processing Systems , 35:17216-17232, 2022.

Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Elements of causal inference: foundations and learning algorithms . The MIT Press, 2017.

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 arXiv:2206.15477 , 2022.

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.06561 , 2023.

Lin Yang and Mengdi Wang. Reinforcement learning in feature space: Matrix bandit, kernels, and regret bound. In International Conference on Machine Learning , pages 10746-10756. PMLR, 2020.

Weitong Zhang, Jiafan He, Dongruo Zhou, Q Gu, and A Zhang. Provably efficient representation selection in low-rank markov decision processes: from online to offline rl. In Uncertainty in Artificial Intelligence , pages 2488-2497. PMLR, 2023.

Fan Feng and Sara Magliacane. Learning dynamic attribute-factored world models for efficient multi-object reinforcement learning. arXiv preprint arXiv:2307.09205 , 2023.

Angela Zhou. Reward-relevance-filtered linear offline reinforcement learning. arXiv preprint arXiv:2401.12934 , 2024.

Lingxiao Wang, Zhuoran Yang, and Zhaoran Wang. Provably efficient causal reinforcement learning with confounded observational data. Advances in Neural Information Processing Systems , 34:21164-21175, 2021.

Wenhao Ding, Laixi Shi, Yuejie Chi, and Ding Zhao. Seeing is not believing: Robust reinforcement learning against spurious correlation. arXiv preprint arXiv:2307.07907 , 2023.

Junzhe Zhang and Elias Bareinboim. Markov decision processes with unobserved confounders: A causal approach. Technical report, Technical report, Technical Report R-23, Purdue AI Lab, 2016.

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.

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.

Zizhao Wang, Xuesu Xiao, Yuke Zhu, and Peter Stone. Task-independent causal state abstraction. Workshop on Robot Learning: Self-Supervised and Lifelong Learning, NeurIPS , 2021.

Zizhao Wang, Xuesu Xiao, Zifan Xu, Yuke Zhu, and Peter Stone. Causal dynamics learning for task-independent state abstraction. arXiv preprint arXiv:2206.13452 , 2022.

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-11224. PMLR, 2020.

Wenxuan Zhu, Chao Yu, and Qiang Zhang. Causal deep reinforcement learning using observational data. arXiv preprint arXiv:2211.15355 , 2022.

Krzysztof Chalupka, Pietro Perona, and Frederick Eberhardt. Fast conditional independence test for vector variables with large sample sizes. arXiv preprint arXiv:1804.02747 , 2018.

Yuichi Yoshida and Takeru Miyato. Spectral norm regularization for improving the generalizability of deep learning. arXiv preprint arXiv:1705.10941 , 2017.

Nathan Lambert, Brandon Amos, Omry Yadan, and Roberto Calandra. Objective mismatch in model-based reinforcement learning. arXiv preprint arXiv:2002.04523 , 2020.

Yilun Du and Igor Mordatch. Implicit generation and modeling with energy based models. Advances in Neural Information Processing Systems , 32, 2019.

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-9311, 2021.

Bo Pang, Tian Han, Erik Nijkamp, Song-Chun Zhu, and Ying Nian Wu. Learning latent space energy-based prior model. Advances in Neural Information Processing Systems , 33:2199422008, 2020.

Eduardo F Camacho and Carlos Bordons Alba. Model predictive control . Springer science & business media, 2013.

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.12293 , 2020.

Maxime Chevalier-Boisvert, Lucas Willems, and Suman Pal. Minimalistic gridworld environment for openai gym. https://github.com/maximecb/gym-minigrid , 2018.

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.

Edouard Leurent. An environment for autonomous driving decision-making. https://github. com/eleurent/highway-env , 2018.

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.

Pim De Haan, Dinesh Jayaraman, and Sergey Levine. Causal confusion in imitation learning. Advances in Neural Information Processing Systems , 32, 2019.

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.

Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. Advances in neural information processing systems , 34:20132-20145, 2021.

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.

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.

Yuping Luo, Huazhe Xu, Yuanzhi Li, Yuandong Tian, Trevor Darrell, and Tengyu Ma. Algorithmic framework for model-based deep reinforcement learning with theoretical guarantees. arXiv preprint arXiv:1807.03858 , 2018.

Pierluca D'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-3808, 2020.

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.

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.

Chi Jin, Zhuoran Yang, Zhaoran Wang, and Michael I Jordan. Provably efficient reinforcement learning with linear function approximation. In Conference on Learning Theory , pages 21372143. PMLR, 2020.

Masatoshi Uehara and Wen Sun. Pessimistic model-based offline reinforcement learning under partial coverage. arXiv preprint arXiv:2107.06226 , 2021.

Yilun Du, Shuang Li, Joshua Tenenbaum, and Igor Mordatch. Learning iterative reasoning through energy minimization. In International Conference on Machine Learning , pages 55705582. PMLR, 2022.

Clark Glymour, Kun Zhang, and Peter Spirtes. Review of causal discovery methods based on graphical models. Frontiers in genetics , 10:524, 2019.

Peter Spirtes, Clark N Glymour, Richard Scheines, and David Heckerman. Causation, prediction, and search . MIT press, 2000.

Karl Pearson. X. on the criterion that a given system of deviations from the probable in the case of a correlated system of variables is such that it can be reasonably supposed to have arisen from random sampling. The London, Edinburgh, and Dublin Philosophical Magazine and Journal of Science , 50(302):157-175, 1900.

Kun Zhang, Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Kernel-based conditional independence test and application in causal discovery. arXiv preprint arXiv:1202.3775 , 2012.

David Maxwell Chickering. Optimal structure identification with greedy search. Journal of machine learning research , 3(Nov):507-554, 2002.

Alain Hauser and Peter Bühlmann. Characterization and greedy learning of interventional markov equivalence classes of directed acyclic graphs. The Journal of Machine Learning Research , 13(1):2409-2464, 2012.

Andrew A Neath and Joseph E Cavanaugh. The bayesian information criterion: background, derivation, and applications. Wiley Interdisciplinary Reviews: Computational Statistics , 4(2):199-203, 2012.

Shengyu Zhu, Ignavier Ng, and Zhitang Chen. Causal discovery with reinforcement learning. arXiv preprint arXiv:1906.04477 , 2019.

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 Intelligence , volume 37, pages 10737-10744, 2023.

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.01075 , 2019.

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.

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.02429 , 2021.

Christopher Clark, Mark Yatskar, and Luke Zettlemoyer. Don't take the easy way out: Ensemblebased methods for avoiding known dataset biases. arXiv preprint arXiv:1909.03683 , 2019.

Divyansh Kaushik, Eduard Hovy, and Zachary C Lipton. Learning the difference that makes a difference with counterfactually-augmented data. arXiv preprint arXiv:1909.12434 , 2019.

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.

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.

Suraj Nair, Yuke Zhu, Silvio Savarese, and Li Fei-Fei. Causal induction from visual observations for goal directed tasks. arXiv preprint arXiv:1910.01751 , 2019.

Sergei Volodin, Nevan Wichers, and Jeremy Nixon. Resolving spurious correlations in causal models of environments via interventions. arXiv preprint arXiv:2002.05217 , 2020.

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.

Zheng-Mao Zhu, Xiong-Hui Chen, Hong-Long Tian, Kun Zhang, and Yang Yu. Offline reinforcement learning with causal structured world models. arXiv preprint arXiv:2206.01474 , 2022.

Mirco Mutti, Riccardo De Santi, Emanuele Rossi, Juan Felipe Calderon, Michael Bronstein, and Marcello Restelli. Provably efficient causal model-based reinforcement learning for systematic generalization. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 37, pages 9251-9259, 2023.

Fan Feng, Biwei Huang, Kun Zhang, and Sara Magliacane. Factored adaptation for nonstationary reinforcement learning. Advances in Neural Information Processing Systems , 35:31957-31971, 2022.

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 preprint arXiv:2402.14528 , 2024.

Jiaheng Hu, Zizhao Wang, Peter Stone, and Roberto Martin-Martin. Elden: Exploration via local dependencies. arXiv preprint arXiv:2310.08702 , 2023.

Wenhao Ding, Haohong Lin, Bo Li, and Ding Zhao. Generalizing goal-conditioned reinforcement learning with variational causal reasoning. Advances in Neural Information Processing Systems , 35:26532-26548, 2022.

Lars Buesing, Theophane Weber, Yori Zwols, Nicolas Heess, Sebastien Racaniere, Arthur Guez, and Jean-Baptiste Lespiau. Woulda, coulda, shoulda: Counterfactually-guided policy search. In International Conference on Learning Representations , 2018.

Silviu Pitis, Elliot Creager, Ajay Mandlekar, and Animesh Garg. Mocoda: Model-based counterfactual data augmentation. arXiv preprint arXiv:2210.11287 , 2022.

Tabitha E Lee, Jialiang Alan Zhao, Amrita S Sawhney, Siddharth Girdhar, and Oliver Kroemer. Causal reasoning in simulation for structure and transfer learning of robot manipulation policies. In 2021 IEEE International Conference on Robotics and Automation (ICRA) , pages 4776-4782. IEEE, 2021.

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-2256. PMLR, 2023.

Yudi Zhang, Yali Du, Biwei Huang, Ziyan Wang, Jun Wang, Meng Fang, and Mykola Pechenizkiy. Interpretable reward redistribution in reinforcement learning: A causal approach. Advances in Neural Information Processing Systems , 36, 2024.

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.

Peide Huang, Xilun Zhang, Ziang Cao, Shiqi Liu, Mengdi Xu, Wenhao Ding, Jonathan Francis, Bingqing 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.

Xiang Fu, Ge Yang, Pulkit Agrawal, and Tommi Jaakkola. Learning task informed abstractions. In International Conference on Machine Learning , pages 3480-3491. PMLR, 2021.

Biwei Huang, Chaochao Lu, Liu Leqi, José Miguel Hernández-Lobato, Clark Glymour, Bernhard Schölkopf, and Kun Zhang. Action-sufficient state representation learning for control with structural constraints. In International Conference on Machine Learning , pages 9260-9279. PMLR, 2022.

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.

Maxime Gasse, Damien Grasset, Guillaume Gaudron, and Pierre-Yves Oudeyer. Using confounded data in latent model-based reinforcement learning. Transactions on Machine Learning Research , 2023.

Junzhe Zhang, Daniel Kumor, and Elias Bareinboim. Causal imitation learning with unobserved confounders. Advances in neural information processing systems , 33:12263-12274, 2020.

Ying Jin, Zhuoran Yang, and Zhaoran Wang. Is pessimism provably efficient for offline rl? In International Conference on Machine Learning , pages 5084-5096. PMLR, 2021.

Thomas Blumensath and Mike E Davies. Iterative hard thresholding for compressed sensing. Applied and computational harmonic analysis , 27(3):265-274, 2009.

Amir Beck and Marc Teboulle. A fast iterative shrinkage-thresholding algorithm for linear inverse problems. SIAM journal on imaging sciences , 2(1):183-202, 2009.

Peter Auer, Thomas Jaksch, and Ronald Ortner. Near-optimal regret bounds for reinforcement learning. Advances in neural information processing systems , 21, 2008.

Ryan Tibshirani and Larry Wasserman. Sparsity, the lasso, and friends. Lecture notes from 'Statistical Machine Learning, ' Carnegie Mellon University, Spring , 2017.

Patrick E McKnight and Julius Najab. Mann-whitney u test. The Corsini encyclopedia of psychology , pages 1-1, 2010.

A Auxiliary Details of BECAUSE Framework

A.1 Notation Summary

We illustrate all the notations used in the main paper and appendix in Table 3.

A.2 Equivalence of the Assumptions with Previous Works

A.3 Derivation of Definition 3

The node of this causal graph G = [ 0 d × d M 0 d × d ′ 0 d × d ] contains two groups of entities: (1) The state action abstraction ϕ ( s, a ) , and (ii) the next state abstraction µ ( s ′ ) . We denote ϕ ( · , · ) ( i ) as the i th

̸

factor in the abstracted state action representations, and µ ( · ) ( j ) for the j th factor in the abstracted state representations.

The source node of all the nodes G is ϕ ( s, a ) ( i ) , which is the abstracted state-action representation, and the sink node of all the edges in G is the to µ ( s ) ′ ( i ) .

Therefore, G is a bipartite graph, since there will be no edges between ϕ ( s, a ) ( i ) , ϕ ( s, a ) ( j ) , or µ ( s ) ( i ) , µ ( s ) ( j ) .

Consequently, we show that G ∈ DAG.

A.4 Derivation of Equation (5)

Definition 4 (Structured Causal Model) . An SCM θ := ( S , E ) consists of a collection S of d functions [11],

where PA G j ⊂ { s 1 , . . . , s d }\{ s j } are called parents of x j in the Directed Acyclic Graph (DAG) G , and E = { ϵ i } d i =1 are jointly independent. For instance, in continuous state and action space, we parameterize the world model with joint Gaussian Distribution, i.e. ϵ ∼ N (0 , σI dd ′ ) .

We then use bilinear MDP to approximate the original likelihood function in Equation (4), i.e.

where K µ := ∑ s ′ ∈S µ ( s ′ ) µ ( s ′ ) T is 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, 15] is that it aims to apply ℓ 0 regression instead of ridge regression to estimate matrix M :

A.5 Proof of Equation (7)

The derivation depends on the following re-weighting formula in [18]:

Then we apply equation (14) to the decomposition in Equation (2) and Equation (3), which yields

where the last equality holds by letting M ( u ) := E pu [ M ( u c ) π β ( a | s,u π )] E pu [ π β ( a | s,u π )] .

B Proof of Theorem 1

In this section, we provide the proof of the sub-optimality upper bound in Theorem 1. 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 in 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, we first define the δ -uncertainty quantifier Γ , then we refer to the lemmas in the previous literature to construct a suboptimality bound based on the defined uncertainty quantifier Γ .

First, we define the Bellman operator B h , for some value function V : S ↦→ R , the Bellman operator can be defined as:

Similarly, we denote the approximate Bellman operator of the empirical MDP constructed from the offline dataset D as ̂ B h for any h ∈ [ H ] .

Definition 5 ( δ -Uncertainty Quantifier) . We let { Γ h } H h =1 , Γ h : S × A ↦→ R to be a δ -uncertainty quantifier with respect to data distribution P D if the event:

satisfies P ( E ) ≥ 1 -δ .

As we consider the offline model learning and planning, we define the model evaluation error at each step h ∈ [ H ] as

where ι h is the error induced by the approximate Bellman operator, especially the transition kernel based on D . We then identify the source of sub-optimality in our offline MBRL setting by decomposing the sub-optimality error in Lemma 1.

Lemma 1 (Decomposition of Suboptimality [88]) .

where π is any learned policy, π ∗ is the optimal policy that maximizes the cumulative return as below:

Based on this decomposition, we will get the basic form of sub-optimality error bound for general offline RL settings in Lemma 2:

Lemma 2 (Suboptimality in standard MDP [88]) . Suppose we have { Γ h } H h =1 as δ -uncertainty quantifier. Under E defined in Equation (5), the suboptimality error bound by conservative planning satisfies:

The basic form of sub-optimality bound in Lemma 2 involves an uncertainty quantifier Γ h , 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 2 (Existence of a core matrix given the feature embedding) . For each ( s, a ) ∈ S × A , feature vectors ϕ ( s, a ) ∈ R d , µ ( s ) ∈ R d ′ are approximated as a priori. Given a specific confounder set u , there exists an unknown matrix M ( u ) ∗ ∈ R d ′ × d such that,

Assumption 3 (Feature regularity) . We assume feature regularity [14, 15] for the following components of the confounded bilinear MDP:

∀ u, ∥ M ( u ) ∥ 2 F ≤ C M d ,

∀ ( s, a ) ∈ S × A , ∥ ϕ ( s, a ) ∥ 2 2 ≤ C ϕ d ,

∀ s, a, s ′ ∈ S × A × S , ∥ ϕ ( s, a ) µ ( s ′ ) T ∥ 1 ≤ C µ .

where C M , C ϕ , C µ , C ′ µ are some universal constants.

Here, for any matrix X , ∥ X ∥ 2 , ∞ := max i √ ∑ j X 2 ij represents the operator 2 ↦→∞ 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 -δ -uncertainty quantifier Γ (see Lemma 2) for our time-homogeneous core matrix estimation problem in the following lemma.

Lemma 3 (Uncertainty bound for Bilinear Causal Representation) . Under the Assumption 2, 3 and that T is an SCM (defined in 4), for the BECAUSE algorithm, for the ξ -optimal policy ( V ∗ 1 ( ˜ s ) -V π 1 ( ˜ s ) ≤ ξ ), ∀ 0 ≤ ξ ≤ 1 , we have the δ -uncertainty set as:

where C 1 is some universal constants.

Armed with the above lemma, we complete the proof of Theorem 1 by showing that

This concludes the proof of Theorem 1.

B.3 Proof of Lemma 3

The key to proving Theorem 1 is to prove Lemma 3. The proof pipeline of Lemma 3 is illustrated below. In Step 1 , we derive the estimation of the causal transition matrix M in BECAUSE as a sparsity regression problem. In Step 2 , we decompose the error terms within δ -uncertainty set into two parts: (a) error due to the under-explored dataset, (b) error due to optimization error in the structured causal model. Then we bound both error terms in Step 3 and Step 4 , respectively. Finally, in Step 5 , we sum up all the results and derive the form of δ -uncertainty quantifier which will lead to our final results in Theorem 1.

Step 1: deriving the output model of BECAUSE. Recalling the original optimization problem in equation (13) to estimate the core matrix:

This part of derivation aims to transform the above estimation problem into a linear regression problem, with the regression data pairs ( X,T ) and some unknown parameters β associated with mask M to be estimated. Eventually, we'll derive the representation of each part of β M , X, T , and eventually reach the following form:

We define each component of this target form of ℓ 0 regression as follows:

For unknown parameters β M : We first define β M ∈ [0 , 1] dd ′ as a column dimensional vector consisting of all the entries in time-homogenous causal matrix M , where β i M denotes the i -th entry of β M . Besides, we define β M D as the true core matrix given some offline dataset D and corresponding data pairs T data , X data that satisfies T data = β M D X data + ϵ .

For dataset D : Recall the transition pairs in the offline dataset D = { s i , a i , s ′ i } 1 ≤ i ≤ n . Here, n represents the sample size over certain state-action pairs in the rollout data by some behavior policy π β . For simplicity, we denote n ≜ n ( s, a ) in the following derivation, which is mentioned in Section 2.1.

For regression target T π β : Then, we introduce the following transition targets T π β induced by the offline dataset D sampled with behavior policy π β :

Under the n finite samples in the offline dataset, we assume that T π β ∼ N ( E [ T π β ] , σ 2 I n ) . The above definition specifies the regression target in the ℓ 0 regression problem, and we denote

T π β = [ T π β ( s ′ 1 | s 1 , a 1 ) , · · · , T π β ( s ′ n | s n , a n )] T ∈ R n as the empirical transition probabilities of certain transition pairs in the offline data D = { s i , a i , s ′ i } 1 ≤ i ≤ n .

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 i ∈ D , which is a vector of Kronecker product between ϕ ( s i , a i ) ∈ R d and normalized µ ( s ′ i ) C µ ∈ R d ′ (without loss of generality, we assume C ϕ = 1 and only need to normalize µ ( s ′ i ) by C µ ):

As a result, X i ∈ R dd ′ , since there are in all n samples in offline dataset, X ∈ R n × dd ′ is the dataset-dependent matrix with all n rows of samples, and d and d ′ are the latent dimension of ϕ and µ , respectively. Based on the feature regularity criteria in Assumption 3, we have ∥ X i ∥ 2 ≤ ∥ X i ∥ 1 ≤ 1 , ∥ X ∥ ∞ ≤ 1 .

The prior work [14] estimate the transition kernel of a bilinear MDP using the following ridge regression:

In this paper, in order to promote the sparsity of the matrix M , we introduce the ℓ 0 regularization term and arrive at the following optimization problem:

where we denote the solution associated with the offline dataset D as ̂ β M D . Here, we use the empirical version constructed by the finite samples in offline dataset D .

Given the goal-conditioned reward setting, for a single episode s ∼ τ , r ( s, a ; g ) = 1 if and only if s = g , otherwise r ( s, a ; g ) = 0 , as is specified in Section 2.1. Since we are essentially predicting the probabilities (normalized to a sum of 1) of whether the next state is the goal state, i.e. ∑ s ∈S ̂ V ( s ) = 1 . Therefore, we have ∥ ̂ V ( · ) ∥ 1 ≤ 1 .

As is denoted by Equation (23), for the specific offline dataset collected by some behavior policies π β , we have the regression target T π β ( s ′ | s, a ) = ∑ ( s i ,a i ,s ′ i ) ∈D 1 ( s i = s,a i = a,s ′ i = s ′ ) ∑ ( s i ,a i ,s ′ i ) ∈D 1 ( s i = s,a i = a ) , and the corresponding ′

features induced by the dataset X i = ϕ ( s i , a i ) ⊗ µ ( s i ) C µ . We have the following equations hold:

where β D is the true underlying transition mask given the offline dataset D , ϵ ∼ N (0 , σ · I n ) is some exogenous noise of the transition model.

Specifically, for the regression problem, we transform the original trajectory dataset D as [ X,T π β ] , as the representation of transition pairs rolled out by the behavior policy π β . Similarly, we can define some 'well-explored dataset' D ∗ , which is an infinite dataset rollout by the behavior policy π β , D ∗ = { s i , a i , r i } ∞ i =0 , similar to the definition of regression target T in Equation (23) and representation data X in Equation (24), we have [ X, E [ T π β ]] as the regression pairs with access to the true transition probability distribution for all the state action pairs ( s, a ) . Here X ∈ R n × dd ′ 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:

We denote E [ T π β ] = [ E [ T π β ( s ′ 1 | s 1 , a 1 )] , · · · E [ T π β ( s ′ n | s n , a n )] ] T ∈ R n . In practice, with finite sample size n , we have E [ T π β ( s ′ | s, a )] = T ( s ′ | s, a ) + ϵ = Xβ M D ∗ + ϵ . In addition, we also introduce a vector form T ∈ R |S||A|×|S| so that ∀ ( s, a ) ∈ S ×A , we have

where the state space is denoted as S = { s ′ 1 , · · · , s |S| } are all possible states. Similar to the Kronecker product we define for X in equation (26), we define X in a matrix form for any state-action pair ( s, a ) :

In addition, we let X ( s ′ | s, a ) ∈ R 1 × dd ′ denote the s ′ -th row of X ( ·| s, a ) associated with the state s ′ ∈ S . Consequently, the estimated transition kernel can be expressed as follows:

Step 2: decomposing the term of interest. To begin with, recalling the definition of Bellman operator B h in Equation (16) and applying Hölder's inequality, the term of interest for any time step 1 ≤ h ≤ H and state-action pair ( s, a ) ∈ S × A can be controlled as

where the first inequality is held given our goal-conditioned reward formulation in section 2.1. To continue, we have

Here, we recall ̂ β M D represents the parameter vector in the estimated causal masks based on the offline dataset D sampled by π β . Similarly, we denote ̂ β M D ∗ as the estimated causal mask outputted from equation (26) based on the infinite dataset D ∗ generated by the behavior policy π β . Then, we can further control equation (30) as

Here the last inequality comes from the fact that

based on the definition of X in equation (24) 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 ℓ 0 optimization of causal masks given the existence of exogenous noise σ defined by SCM in Definition 4. We will control them separately in the following.

Step 3: Controlling term (a). We need to consider the optimization process in the original regression problem in equation (22) to fully understand the difference between ̂ β M D and ̂ β M D ∗ , where the only difference is that the latter uses a perfect dataset with infinite samples. The optimization problem we target (cf. equation (26)) can be solved by the iterative hard thresholding algorithm (IHT) proposed by [89] IHT offers an iterative solution for the ℓ 0 regression problem, armed with a hard thresholding operator as below:

We denote ̂ β M D ( i ) as the estimated causal mask parameters after i -th iterations with dataset D . Similarly, we denote ̂ β M D ∗ ( i ) as the estimation after i -th iterations with dataset D ∗ . We initialize the graph to be a full graph regardless of the datasets ( D ∗ or D ) used in the optimization process, leading to ̂ β M D ∗ (0) = ̂ β M D (0) = 1 ∈ R dd ′ .

Recall that X ∈ R n × dd ′ , T π β ∈ R n , E [ T π β ] ∈ R n and ̂ β M D is [0 , 1] dd ′ based on the definition in the original ℓ 0 optimization problem in equation (22), equation (23) and equation (28). The update rules of using either the dataset D or D ∗ can be written as:

Note that the difference between the two parameters ̂ β M D and ̂ β M D ∗ essentially relies on the difference between two pairs of transition kernel estimation datasets [ X, E [ T π β ]] and [ X,T π β ] . It is easily verified that the hard-thresholding operator [ g λ ( · )] i defined in equation (32) is L = 1 -Lipschitz. According to [90], we can set the learning rate η ≤ 1 L = 1 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

Here, I dd ′ represent the identity matrix of size dd ′ × dd ′ , the last inequality holds based on the fact that ∥ X ∥ 2 ≤ 1 defined in equation (24). Since η ≤ 1 L = 1 , as a result, ∥ I dd ′ -ηX T X ∥ 2 = max(1 -ηλ ( X T X )) ≤ 1 . We perform IHT for sufficient K > 0 iterations and output the last step estimation as 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 0 ≤ ξ ≤ 1 , we have the output of IHT gradually converges to the optimal solution ̂ β M D of the problem in equation (26). Namely, after at most K ≃ log ( ∥ M ∥ 0 ξ ) steps [89, Corollary 1], the output satisfies

∥ ∥ ∥ ̂ β M D -̂ β M D ( K ) ∥ ∥ ∥ 2 ≤ ξ 4 . Similarly, based on the perfect dataset D ∗ , the output of IHT gradually converges to ̂ β M D ∗ at the same rate. Consequently, the term of interest (a) can be bounded recursively as:

where the second inequality holds by recursively applying equation (34) to K,K -1 , · · · , 0 iterations, the last equality is due to the fact that ̂ β M D (0) = ̂ β M D ∗ (0) , and the last inequality holds by setting K = log( ∥ M ∥ 2 ξ ) .

Now the remaining of the proof will focus on controlling ∥ T π β -E [ T π β ] ∥ ∞ . Recall that we have defined the regression target in equation (23) and equation (28):

Proposition 1 (Well-explored dataset) . With the offline dataset D of in total n samples with n ( s, a ) samples generated independently conditioned on any ( s, a ) . For any 0 < δ < 1 , with probability at least 1 -δ , one has

for some universal constant C β .

The above proposition can be directly proved by applying [91, Lemma 17] over all ( s, a ) ∈ S × A :

As a result, we can further extend the results in equation (35) to bound the term (a) as follows:

where C 1 = C β C µ is some constant that is related to the feature regularity, ξ is the level of error tolerance in the cumulative value returns, in our setting, 0 ≤ ξ ≤ 1 .

Step 4: Controlling term (b). For (b) in Equation (31), which is ∥ ̂ β M D ∗ -β M ∥ ∞ , we are interested in 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 4. Yet our causal discovery module, i.e. an ℓ 0 estimator, will always encounter some estimation error induced by the exogenous noise.

Firstly, we have the bounded relationship between ℓ 2 and ℓ ∞ norm:

then we can analyze the error bound for ℓ 0 regression in the sense of ℓ 2 norm.

The derivation below generally follows the ℓ 0 regularized linear regression bound in [92]. Based on Assumption 2 and 3, there exists M , we denote this optimal solution in the vector form as β M .

Besides the aforementioned optimization error in the iterative thresholding update process, according to the SCM in Definition 4 and Assumption 2 and Proposition 1, we denote a finite subset of observed transition probabilities of the well-explored data E [ T π β ] with size n : T obv ∈ R n . By definition above, we can then assume the finite-sample regression target T is generated by causal features of transition pairs (denoted by X = X π ∗ = X π β , X ∈ R n × dd ′ ) and the ground-truth causal mask (represented by β M ∈ R dd ′ ) with the following equation:

Here ϵ ∼ N (0 , σI n ) 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 ℓ 0 regression problem with its bounded form as follows:

In BECAUSE, we select the sparsity level s ≈ ∥ β M ∥ 0 = ∥ M ∥ 0 ≤ dd ′ .

For simplicity, we denote the approximate solution ̂ β M D ∗ in Equation (39) as ̂ β M . By the virtue of optimality of the solution β M in Equation (39), we find that

then by expanding and shifting the terms, we can derive the basic inequality for the ℓ 0 estimator above:

Since both ̂ β M and β M are s -sparse, the vector ̂ β M -β M is at most 2 s -sparse. We denote the set of all 2 s -sparse dd ′ -dimensional vector set as T dd ′ (2 s ) , we denote v = I ( ̂ β M -β M = 0) as an indicator vector, v i = 0 if and only if ̂ β i M -β i M , ∀ i ∈ [ dd ′ ] .

By shifting the terms, we use the sub-Gaussian assumption and further use Hölder's inequality to get the following results:

where ϵ ∼ N (0 , σ 2 I n ) is the exogenous noise variable in the SCM in Definition 4.

Then by simply applying restricted eigenvalue (RE) condition, for some κ ( S ) > 0 , we have

Therefore, we have:

with probability at least 1 -δ , which bounds term (b) in Equation (31).

Step 5: Summing up the results Summarizing both bounds for terms (a) and (b) in Equation (31), we will get the following bounds with probability 1 -δ :

Therefore, we complete the proof of Lemma 3 by showing that for all ( s, a ) ∈ A × A , h ∈ [ H ] ,

The final bound of the term (b) includes a dependency of O ( 1 √ n ) and logarithm of dimensionality dd ′ in the estimated transition matrix. It also incurs a dependency of error tolerance ξ or SCM's noise level σ square root of sparsity level √ s = √ ∥ M ∥ 0 .

So far, we prove the following bound in theorem 1:

In order to achieve ξ -optimal policy such that V ∗ 1 ( ˜ s ) -V π 1 ( ˜ s ) ≲ ξ , the RHS needs to satisfy:

We first multiply √ min ( s,a,h ) ∈S×A× [ H ] E π ⋆ [ n ( s h , a h ) | s 1 = ˜ s ] , and then take the square on both sides. We have:

Given the definition, LHS in Equation (47) satisfies:

To ensure the satisfaction of ξ -optimal policy, we thus would like the RHS of Equation (47) satisfies:

Consequently, we can shift the terms and get the sample complexity bound for the ξ -optimal policy, ∀ 0 ≤ ξ ≤ 1 :

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 λ , we use p -value as a threshold to determine the following conditional independence:

Here ϕ ( · , · ) ( i ) means that this element is the i th factor in the abstracted state action representation, similar to µ ( · , · ) ( j ) , ϕ ( s, a ) -( i ) means all the other factors in the representation except for the i th factor.

If the p-value based on the above conditional independence test is less than a threshold, we can remove the edge by setting M ij = 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 to the margin loss in Equation (8). In practice, we attach Tanh() to the output layer to clip the unnormalized score between -1 and +1. The energy networks take in both conditions and samples, then concatenate them together and sent it into MLP encoders. The detailed hyperparameters of EBM are listed in Table 12.

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, we find that we can use these representations in both the energy networks and the sampling process to get some effective negative samples, which is similar to the practice of augmentation of causality-guided counterfactual data [77].

The interesting trick we employ here is the way to get our negative samples by mixing the latent factors from offline data. For example, for a positive sample array

with conditions:

Here, s i , a i , s ′ i denotes the timestep of the offline samples. ϕ ( · , · ) ( i ) means this element is the i th factor in the abstracted state action representation, similar to the µ ( · , · ) ( i )

′

as we already get the corresponding causal representation µ ( s ) that is semantically meaningful, we can mix the columns to create useful counterfactual negative samples

These counterfactual negative samples seem to effectively use the causal representation and speed up the training, as we show in Figure 7.

Random

Causal EBM

T=0                    T=100                  T=200                   T=300                  T=400

T=500                  T=600                   T=700                  T=800                  T=900

Random

Causal EBM

Next State

GT

0.75

0.50

0.25

0.00

-0.25

-0.50

-0.75

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 we evaluate the mismatch via the following two metrics.

We conduct hypothesis testing via Mann-Whitney U Test [93], with Null hypothesis

To 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 1 ( τ pos ∥ τ neg ) .

We report the results of p -value and the W 1 distance of two groups of model loss samples in the following table.

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 in Table 6, 7, 8, and 9.

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. It contains 4 dimensions of hybrid action space that uses Operation Space Control (OSC) to control the 3D position and the 1D gripper movement. The task is counted as a success when the assigned block is lifted from the table over 0.1m. 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 8(a).

Unlock The Unlock environments based on the MiniGrid world [35] contain 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/left/right) 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 8(b).

Crash 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 is 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 8(c).

All three environments are visualized in Figure 4. We list their basic configurations in Table 10.

(a) Lift Environment

(b) Unlock Environment

(c) Crash Environment

pick key

Controll car2

open door

Control car3

ped(h)

car1(h)

car2(h)

car3(h)

ego(h)

collision(h)

(t+ugsod juate

ped(h+1)

carl(h+1)

key pos(h+1)

car2(h+1)

door pos(h+1)

car3(h+1)

ego(h+1)

have key(h+1)

collision(h+1)

C.6 Additional Baseline Information

We collect data on the above 3 different environments, thus forming 9 groups of offline datasets.

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 , [38]): https://github.com/ioanabica/InvariantCausal-Imitation-Learning, MIT License.

Causal Confusion Imitation Learning ( CCIL , [39]): reference link to the paper.

TD3 with Behavior Cloning ( TD3+BC , [41]): https://github.com/sfujim/TD3_BC, MIT License.

Model-based Offline Policy Otimization ( MOPO , [2]): https://github.com/junmingyang/mopo.git, MIT License.

Relational Graph Neural Network ( GNN , [42]): https://github.com/MichSchli/RelationPrediction.git, MIT License.

Causal Dynamics Learning ( CDL , [24]): https://github.com/wangzizhao/robosuite/tree/cdl, MIT License.

Denoised MDP ( Denoised , [12]): https://github.com/facebookresearch/denoised_mdp.git, CC BY-NC 4.0.

Mismatch No More ( MnM , [9]): reference link to the paper.

World model with identifiable factorization ( IFactor , [13]), reference link to the paper.

Delphic Offline RL ( Delphic , [40]): reference link to the paper.

The detailed hyperparameters we 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: https://anonymous.4open.science/r/BECAUSE-NeurIPS

Computing resources The experiments are run on a server with 2 × AMDEPYC 7542 32-Core Processor CPU, 2 × NVIDIA RTX 3090 graphics and 2 × NVIDIA 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.

C.8 Broader Impact

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

NeurIPS Paper Checklist

The checklist is designed to 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 is Not Applicable for that particular paper or the relevant information is Not Available.

Please provide a 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 it (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.

The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation. While "[Yes] " is generally preferable to "[No] ", it is perfectly acceptable to answer "[No] " provided a 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 "[No] " or "[NA] " is not grounds for rejection. While the questions are phrased in a binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write a 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 be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it 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 are discussed in the conclusion.

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 in 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 a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text 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 fairness.

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 important 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 a complete (and correct) proof?

Answer: [Yes]

Justification: They are provided in section 3.4 and Appendix A 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 crossreferenced.

All assumptions should be clearly stated or referenced in 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 a 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 experimental results of the paper to the extent that it 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 Appendix C.

Guidelines:

The answer NA means that the paper does not include experiments.

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 a 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 a specific model and empirical evaluation, it may be necessary to either make it 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 submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We provide an open access to 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 NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) 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 (e.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/guides/CodeSubmissionPolicy ) 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 a 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).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We provide all the experiment settings and details in Appendix C.5 and C.6. Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a 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, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We report the mean and 95% confidence interval (CI) as well as the significance level (p value) in Table 6, 7, 8 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, confidence intervals, or statistical significance tests, at 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.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 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 of 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 in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: Provided in Appendix C.7.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as 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 it 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 https://neurips.cc/public/EthicsGuidelines ?

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 a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or 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]

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or 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 a 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 as 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 to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of 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, or scraped datasets)?

Answer: [NA]

Justification: This paper does not contain pre-trained language models, image generators, scraped datasets, or similar assets.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a 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 to 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: 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?

Answer: [Yes]

Justification: All the baselines are properly cited and introduced in Appendix C.6.

Guidelines:

The answer NA means that the paper does not use 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 is used and, if possible, include a 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 (e.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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-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 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 of 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 is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing or research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in 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 crowdsourcing nor research with human subjects.

Depending on 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 as the institution conducting the review.

Peer reviews:
comments: SUMMARY:
Distribution shifts in the sampling of MBRL will lead to the objective mismatch between model and policy learning.
Therefore, this paper aims to reduce the influence of the distribution shift for MBRL.
It 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.

STRENGTHS:
1. This paper is well-written and easy to follow.
2. In RL, this paper tries to address confounders, which are key problem in causal inference.
3. There are some new theoretical results in this paper.
4. Experiment results illustrate the effectiveness of the proposed method.

WEAKNESSES:
1. The conditions in Bilinear MDPs are difficult to achieve in real-world environments.
2. Some real-world examples can be given for confounders $\mu_{\pi}$ and $\mu_{c}$.

QUESTIONS:
1. What is the world model in this paper? Does this paper use a world model to train the policy?

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-06-22

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\nWhy 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\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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).\n\nSTRENGTHS:\n- [*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,\n\nWEAKNESSES:\nSince 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?\n\nQUESTIONS:\nI listed all the questions (mixed with weaknesses) in the above section.\n\nLIMITATIONS:\nLimitations 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-12"}

comments: SUMMARY:
This paper studies challenges in Model Based rl, namely:
1. Objective mismatch between estimating the model and policy learning for value optimization
2. Confoundedness in offline RL causing this mismatch
The paper proposes a theoretically sound algorithm and experiments extensively to support their claims.

STRENGTHS:
1. Well written paper - Problem presentation is very clear
2. 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
2. 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.

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12

comments: 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.

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.

QUESTIONS:
Is this method efficient in handling continuous actions? During the planning phase, computing the argmax Q can be time-consuming.
Limitation: The proposed method introduces additional complexity to the MBRL framework, which may pose challenges for practical implementation and scalability.

LIMITATIONS:
none
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-20