BECAUSE: Bilinear Causal Representation for Generalizable Offline Model-based Reinforcement Learning

Authors: haohong  lin
wenhao  ding
jian  chen
laixi  shi
jiacheng  zhu
bo  li
ding  zhao
  cmu
  caltech
  mit
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-09 22:48:51

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).

Main paper: 
 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 necessarily 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. 
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.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 ∀s, a, s ′ ∈ S × A × S, T (s ′ |s, a) = ϕ(s, a) T M µ(s ′ ),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. 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 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. 
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.In the presence of a hidden confounder u, we model the confounder behind the transition dynamics as a linear confounded MDP [18] :EQUATIONwhere 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:EQUATIONDefinition 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:G = arg max G∈DAG log p(D; ϕ, µ, G) -λ|G| =⇒ M = arg max M log p(D; ϕ, µ, M ) -λ|M |, (4)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 ), whereG(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 . 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:EQUATION)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:EQUATION)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:EQUATIONThe 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) = Ep u [M (uc)π β (a|s,uπ)] Ep u π β (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. M 0 ← [1] d ′ ×d for i ∈ [K] do Update ϕ n , µ n by L rep (ϕ, µ) in (6) if i mod k = 0 then Update M n with L mask (M ) in (5)Weighted average M n with (7) // Uncertainty quantifier learning Fit E θ (s, a) with L EBM (8) Initialize V H+1 (s) = 0, ∀(s, a) // Pessimistic planning while h < H do Estimate the uncertainty with score E θ (s, a)Compute Q h (s, a) with (9) V h (s, a) = max a Q(s, a) a ← arg max a Q(s, a) s ′ , r ← env.step (a, g)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).EQUATIONwhere µ(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):EQUATIONThen 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:V * 1 ( s) -V π 1 ( s) ≲ min C 1 log ∥M ∥ 0 ξ |S|, C s σ ∥M ∥ 0 ) H h=1 E π * log(1/δ) n(s h , a h ) | s 1 = s ,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, min (s,a,h)∈S×A×[H] E π ⋆ n(s h , a h ) | s 1 = s ≳ min C 2 1 log 2 ∥M ∥0 ξ |S|, C 2 s σ 2 ∥M ∥ 0 • H 2 log(1/δ) ξ 2 .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. 
experiment results: 
 Figure 4 : Three environments used in this paper.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 .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. We evaluate the mean and standard deviation of the best performance among 10 random seeds and report task-wise results in Appendix Table 6 .Lift-Random 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 AV 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 AVs. 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.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 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 . 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 .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. 
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] .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. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
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.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.  
 
 Appendix: We illustrate all the notations used in the main paper and appendix in Table 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 Task Independence [24] T (ϕ(s (1) )|ϕ(s (1) ), a; u (1) ) = T (ϕ(s (2) )|ϕ(s (2) ), a; u (2) )′ )|ϕ(s), a) = T (s ′ C |s C , s R , a)p(s ′ R |s R ) T (ϕ(sInvariant Action Effect [26] T online (s ′ |s, a) = u T (s ′ |s, a, u)p(u|s) T offline (s ′ |s, a) = u T (s ′ |s, a, u)p(u|s) T online (s ′ |s, a) ∝ T offline (s ′ |s, a) Invariant Causal Graph u (1) ̸ = u (2) , M (u (1) ) = M (u (2) )T (s ′ |s, a; u (1) ) = T (s ′ |s, a; u (2) )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) .EQUATIONTherefore, 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] ,EQUATIONwhere 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.EQUATIONwhere 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 :EQUATIONA.5 Proof of Equation ( 7)The derivation depends on the following re-weighting formula in [18] :EQUATION)Then we apply equation ( 14) to the decomposition in Equation ( 2) and Equation ( 3), which yieldsEQUATIONwhere the last equality holds by letting M (u) :=Ep u [M (uc)π β (a|s,uπ)] Ep u [π β (a|s,uπ)].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.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:EQUATION) 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:E = |( B h V h+1 )(s, a) -(B h V h+1 )(s, a)| ≤ Γ h (s, a), ∀(s, a, h) ∈ S × A × [H] satisfies P D (E) ≥ 1 -δ.As we consider the offline model learning and planning, we define the model evaluation error at each step h ∈ [H] asEQUATIONwhere ι 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] ).EQUATIONwhere π is any learned policy, π * is the optimal policy that maximizes the cumulative return as below:π * = arg max π E π H h ′ =1 γ h ′ r(s h ′ , a h ′ )|s h .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:∀s ∈ S : V * h (s) -V π h (s) ≤ 2 H h ′ =h E π * [Γ h ′ (s h ′ , a h ′ )|s 1 = s].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.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,EQUATIONAssumption 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 ′ ∈ R |S| , ∥µ T s ′ ∥ 2 ≤ C µ ∥s ′ ∥ ∞ , ∥µK -1 µ ∥ 2,∞ ≤ C ′ µ ,• ∀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:E BECAUSE = |( B h V h+1 )(s, a) -(B h V h+1 )(s, a)| ≲ min C 1 log ∥M ∥ 0 ξ |S|, C s σ ∥M ∥ 0 ) log(1/δ) n(s, a) , ∀(s, a, h) ∈ A × S × [H] ,where C 1 is some universal constants.Armed with the above lemma, we complete the proof of Theorem 1 by showing thatV * 1 ( s) -V π 1 ( s) ≤ 2 H h ′ =1 E π * [Γ h ′ (s h ′ , a h ′ )|s 1 = s] ≲ 2 H h=1 E π * min C 1 log ∥M ∥ 0 ξ |S|, C s σ ∥M ∥ 0 ) log(1/δ) n(s h , a h ) | s 1 = s (20)This concludes the proof of Theorem 1.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:EQUATION)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:EQUATIONWe 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 β M i 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 π β :EQUATION)Under the n finite samples in the offline dataset, we assume thatT π β ∼ N (E[T π β ], σ 2 I n ).The above definition specifies the regression target in the ℓ 0 regression problem, and we denoteT π β = [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 µ ):EQUATIONAs 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:minEQUATION)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:EQUATIONwhere 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 targetT π β (s ′ |s, a) = (s i ,a i ,s ′ i )∈D 1(si=s,ai=a,s ′ i =s ′ ) (s i ,a i ,s ′ i )∈D 1(si=s,ai=a), and the corresponding features induced by the datasetX i = ϕ(s i , a i ) ⊗ µ(s ′ i )Cµ . We have the following equations hold:EQUATIONwhere β 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:EQUATIONWe denoteE[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 haveT(•|s, a) = T (s ′ 1 |s, a) T (s ′ 2 |s, a) • • • T (s ′ |S| |s, a) T ∈ R |S| ,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):X(•|s, a) = ϕ(s, a) ⊗ µ(s ′ 1 ) C µ , • • • ϕ(s, a) ⊗ µ(s ′ |S| ) C µ T ∈ R |S|×dd ′ .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:T(•|s, a) = ϕ T (s, a) M µ(•) = X(•|s, a) β M D .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 asEQUATIONwhere the first inequality is held given our goal-conditioned reward formulation in section 2.1. To continue, we haveEQUATIONHere, 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) asEQUATION(31) Here the last inequality comes from the fact that∥X(•|s, a)∥ ∞ = max i∈|S| j∈[dd ′ ] |X(•|s, a) ij | = max i∈|S| ∥X(•|s, a) i ∥ 1 = max i∈|S| ∥ϕ(s, a)µ(s ′ i ) T ∥ 1 ≤ C µ C µ = 1based 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:EQUATIONWe 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:EQUATIONNote 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 * asEQUATIONHere, 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:EQUATIONwhere 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 of the proof focus on controlling ∥T π β -E[T π β ]∥ ∞ . Recall that we have defined the regression target in equation ( 23) and equation ( 28):T π β (s ′ |s, a) = 1 n(s, a) n(s,a) i=1 1(s i = s, a i = a, s ′ i = s ′ ), E[T π β (s ′ |s, a)] = E s ′ ∼T (•|s,a) [1(s i = s, a i = a, s ′ i = s ′ )].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∀(s, a) ∈ S × A : ∥T π β (•|s, a) -E[T π β (•|s, a)]∥ 2 ≤ C β |S| n(s, a) log |S||A| δ ,for some universal constant C β .The above proposition can be directly proved by applying [91, Lemma 17 ] over all (s, a) ∈ S × A:EQUATIONAs a result, we can further extend the results in equation (35) to bound the term (a) as follows:EQUATIONwhere 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:∥ β M D * -β M ∥ ∞ ≤ ∥ β M D * -β M ∥2 , 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 byX = X π * = X π β , X ∈ R n×dd ′) and the ground-truth causal mask (represented byβ M ∈ R dd ′) with the following equation:EQUATIONHere ϵ ∼ 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:EQUATIONIn 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 thatEQUATIONthen by expanding and shifting the terms, we can derive the basic inequality for the ℓ 0 estimator above:∥T obv -Xβ M ∥ 2 2 ≤ ∥(T obv -Xβ M ) + (Xβ M -X β M )∥ 2 2 (38) = = ⇒ ∥ϵ∥ 2 2 ≤ ∥ϵ + (Xβ M -X β M )∥ 2 2 = ∥ϵ∥ 2 2 + ∥X β M -Xβ M ∥ 2 2 + 2⟨ϵ, Xβ β * -X β M ⟩Since both β M and β M are s-sparse, the vector β M -β M is at most 2s-sparse. We denote the set of all 2s-sparse dd ′ -dimensional vector set as T dd ′ (2s), we denote v = I( β M -β M = 0) as an indicator vector, v i = 0 if and only ifβ M i -β M i , ∀ i ∈ [dd ′ ]. By shifting the terms, we use the sub-Gaussian assumption and further use Hölder's inequality to get the following results:EQUATIONwhere ϵ ∼ 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κ(S)∥ β M -β M ∥ 2 2 ≤ 1 n(s, a) ∥X( β M -β M )∥ 2 2 ≤ 2∥ β M -β M ∥ 2 σ2s log(dd ′ /δ) n(s, a)Therefore, we have:EQUATIONwith 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 -δ:EQUATIONTherefore, we complete the proof of Lemma 3 by showing that for all (s, a)EQUATIONThe 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:EQUATIONIn order to achieve ξ-optimal policy such that V * 1 ( s) -V π 1 ( s) ≲ ξ, the RHS needs to satisfy:EQUATIONWe 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:EQUATIONGiven the definition, LHS in Equation ( 47) satisfies:LHS ≲ min C 1 log ∥M ∥ 0 ξ |S|, C s σ ∥M ∥ 0 ) H h=1 E π * log(1/δ) | s 1 = s (48)To ensure the satisfaction of ξ-optimal policy, we thus would like the RHS of Equation ( 47) satisfies:EQUATIONConsequently, we can shift the terms and get the sample complexity bound for the ξ-optimal policy, ∀0 ≤ ξ ≤ 1:EQUATIONIn this section, we provide additional experiment results and algorithm implementation details.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:EQUATIONHere ϕ(•, •) (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.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 We report the results of the task-wise performance of all baselines in the main experiment and variants in the ablation studies in 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.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.• 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. , 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. 


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

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

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

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

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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: limitations: none
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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: limitations: There are no significant limitations of the paper.
paper_summary: paper_summary: This paper aims to learn causal representations in offline RL. The authors propose a novel framework that incorporates causal graph learning through bilinear MDPs to manage mismatches in model learning and policy optimization stages. The framework models two types of confounders: ones affecting state dynamics and those influencing behavior policies. To enhance learning, the framework includes re-weighting and uncertainty quantification as well. Empirical results demonstrate the framework’s effectiveness in both causal discovery for model learning and generalization in offline RL policy learning.

In 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).
strength_weakness: strength_weakness: Since here most of the weaknesses and questions are mixed. So I put both here. One of the major weaknesses, which is also a primary question, is the inherent difficulty in fully learning the causal process through the bilinear MDP optimization (details in question 1-3).

**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?

**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.

**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?

**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.

**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?
questions: questions: I listed all the questions (mixed with weaknesses) in the above section.
limitations: limitations: Limitations have been discussed in the conclusion section. Though there is no in-depth discussion on societal impacts, I don't find any particular societal impacts in this paper compared with other RL works.
paper_summary: paper_summary: The paper introduces a method to improve model-based offline reinforcement learning by addressing the objective mismatch problem using causal discovery methods. This problem stems from the fact that in this setup the learning algorithm aims to solve the following two problems at once: i) accurate learning of the transition dynamics, ii) identifying the optimal policy on the MDP characterized by the learned transition dynamics. However in practice, an improved solution to one of these two problems does not translate to a better solution to the other one, while in theory it should. The paper's central hypothesis is that this disconnect is caused by the spurious correlations learned by the model which can be blocked by a causality based formulation.
strength_weakness: strength_weakness: * 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:

Sun et al., Model-Bellman Inconsistency for Model-based Offline Reinforcement Learning, ICML, 2023.

* 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.
questions: questions: Why was it not possible and/or sensible to study the problem in the most well-known benchmark of offline RL research, the MuJoco environments of D4RL?
limitations: limitations: The paper discusses the limitations of the proposed approach in the final few sentences of the conclusion section. However, it does not discuss the potential negative societal impact of the presented work.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: 1. What is the world model in this paper? Does this paper use a world model to train the policy?
limitations: limitations: N/A

Comments: title: [Reply to Reviewer 8enR] Regarding the "misleading" part you mentioned
Dear Reviewer XKke
Keep score
comment: Dear Reviewer 8enR, 

Thank you for your response. To help us better address your concern, could you please kindly elaborate more on this sentence "I think how to define the model in this paper is an important problem":
1. Which 'model' do you refer to? What is the key issue in our 'definition'? 
2. How does this issue relate to our core contribution of bilinear causal representation? 

If you still feel inappropriate in that we 'redefine' the 'world model' term, we have **changed the term 'world model' to 'dynamics model'**, including the state transition and reward model (see our last response), to avoid further confusion. The term 'dynamics model' is also commonly used by some previous MBRL literature [1, 2]. 


>[1] Janner, Michael, et al. "When to trust your model: Model-based policy optimization." Advances in neural information processing systems 32 (2019).
>
>[2] Yu, Tianhe, et al. "Mopo: Model-based offline policy optimization." Advances in Neural Information Processing Systems 33 (2020): 14129-14142.
Yes, I think how to define the model in this paper is an important problem.
Thanks for the extensive discussion on the term of "World Model". I think it's useful to clarify the definition of the models learned in this paper to avoid potential ambiguity with the term in the literature.

Reviewer 8enR, do you have additional major concerns that led to your current rating apart from the use of "world model"?
**I did not mean that the current state encoded all historical information.**
In MDPs, the historical information cannot change the probability distributions of $s_{t+1}$, however, with the help of the context in historical information, we may construct a new causal function for MDPs.
I would like to keep my score.
We think Reviewer 8enR mentioned an important point: **the MDP setting assumes that the current state $s_t$ encoded all history information.** Therefore, $p(s_{t+1}|s_t, a_t)=p(s_{t+1}|s_t, a_t, s_{t-1}, s_{t-2},\cdots)$ already considers all history information prior to $t$ under the Markov assumption. In this context, the 'world model' is applicable in the MDP settings, which is the case in our paper. 

However, to avoid further confusion and address the concern of reviewer 8enR, we will change the term 'world model' in our paper to 'dynamics model' (including state transition and reward model). We believe this will address the concern of Reviewer 8enR. Please kindly let us know if there are additional questions or concerns before the discussion period ends. We would appreciate it if Reviewer 8enR would reconsider the evaluation.
All of [1, 2. 3] take action by considering encoded historcial information. 
You may change 'world model' to other terms.
We thank reviewer 8enR for engaging in the followup discussion. There are a few recent works that use 'world model' in the MDP formulation, including but not limited to [1, 2, 3]. 

Denoised MDPs [1] formulate their problem in MDPs instead of POMDPs in their problem formulation and discuss how the MDP formulation is generic enough for their world model learning problem: 
> For generality, we consider tasks in the form of Markov Decision Processes (MDPs), described in the usual manner: $\mathcal{M}\triangleq (\mathcal{S}, \mathcal{A}, R, P, p_{s_0})$ (Puterman, 1994), where $\mathcal{S}$ is the state space, $\mathcal{A}$ is the action space, $R$: $\mathcal{S}\to \Delta([0, r_{max}])$ defines the reward random variable $R(s')$ received for arriving at state $s'\in \mathcal{S}$, $P: \mathcal{S} \times \mathcal{A} \to \Delta(\mathcal{S})$ is the transition dynamics, and $p_{s_0} \in \Delta(\mathcal{S})$ defines the distribution of initial state. We use $\Delta(A)$ to denote the set of all distributions over $A$. $P$ and $R$ define the most important components of a MDP: the transition dynamics $P[s'| s, a]$ and the reward function $P[r | s']$. Usually, the objective is to find a policy $\pi: S\to \Delta(A)$ acting based on current state, that maximizes the expected cumulative (discounted) reward.
> 
> Indeed, MDPs provide a general formulation that encompasses many tasks. In fact, the entire real world may be viewed as an MDP with a rich state/observation space S that contains all possible information/signal. For an artificial agent to successfully perform real-world tasks, it must be able to process observations that are incredibly rich and high-dimensional, such as visual or audio signals.

As we mentioned in our response, our description of world model follows [1], which is formulated in the MDP setting. We also want to emphasize that **our key contribution is the bilinear causal representation learning**. If reviewer 8enR still feels using 'world model' may affect their justification on our contribution, we are happy to add more remarks in our appendix or **change 'world model' to other terms that reviewer 8enR feels more appropriate**. We believe this change won't have any influence to the core contribution of this paper. Once again, we hope the reviewer’s concern will be resolved, if they can clarify their demands during the discussion phase. 


> [1] Wang, Tongzhou, et al. "Denoised mdps: Learning world models better than the world itself." arXiv preprint arXiv:2206.15477 (2022).
> 
> [2] Zhu, Zheng-Mao, et al. "Offline reinforcement learning with causal structured world models." arXiv preprint arXiv:2206.01474 (2022).
> 
> [3] Liu, Jingbin, Xinyang Gu, and Shuai Liu. "Reinforcement learning with world model." arXiv preprint arXiv:1908.11494 (2019).
The world model consists of all the  planner's knowledge of the past, present, and future, expressed in a temporal logic [6]. 
To optimize POMDP problems,  we should consider not only present state but also the historical information to predict the future states or rewards, therefore world model can be used to optimize POMDP tasks.

However, because of the Markov property, we just need to consider present state in MDPs, which are the key problem in this paper. Can you find any literature, in which world model is for MDPs and does not consider historical information?





[6] Allen, James F., and Johannes A. Koomen. "Planning using a temporal world model." Proceedings of the Eighth international joint conference on Artificial intelligence-Volume 2. 1983.
FIrst of all, thank you for your clarification! 

However, I don't believe the vision component is the key aspect of the "world model". If we assume access to at least some of the grounded symbolic or physical states, then the world model could be primarily concerned with the state dynamics and reward model. These elements essentially characterize the underlying evolution and reward mechanisms in the observations/ measurements within RL. The vision component becomes essential only when the grounded states are not directly observable.
We thank reviewer 8enR for engaging in the rebuttal discussion. We also thank reviewer XKke for acknowledging our main contribution. First, we want to clarify that our core contribution is **NOT** proposing a new definition of the world model, but the **bilinear causal representation framework** in model-based RL and causal RL. Then, we address the concern from reviewer 8enR about the "world model" as follows.

**'World model' in our paper is not newly defined by us.** Our paper uses a notation of transition model $T(s' | s, a)$ and a reward model $r(s, a)$ to describe the conditional distribution $p(s', r | s, a)$. This "world model" description, **which is not defined by us**, directly follows previous works of model-based RL [3-5], especially two of our baselines Denoised MDP [4] and IFactor [5]. 

**Definition of 'world model' that reviewer 8enR refers to.** To better address the "redefine" concern of the reviewer, we would like to know what mathematical form of "world model" the reviewer 8enR refers to. We summarized the related work [1, 2] mentioned by the reviewer in the following table. We are glad to discuss this question further. 

**Related works of 'world model'.** Existing papers [1-9] use "world model" with different descriptions due to the broad and controversial meaning of this term. Even [1, 2] mentioned by the reviewer have somewhat different descriptions from each other. And there were works [6-9] proposing this definition that are much earlier than [1, 2]. To have a clear understanding, we summarize the following table of "world models" related to causality or RL topics:

|Paper|Description of World Model|
|-|-|
|[1] Ha, David, and Jürgen Schmidhuber. **"World models."** arXiv preprint arXiv:1803.10122 (2018). | Vision model (V): $p(z_t \| o_t)$ , RNN Memory model (M): $p(z_{t+1} \| a_t, z_t, h_t)$|
|[2] Hafner, Danijar, et al.**"Dream to control: Learning behaviors by latent imagination."** arXiv preprint arXiv:1912.01603 (2019). | Representation: $p_\theta(s_{t+1}\|s_t, a_t, o_t)$, transition $q_\theta(s_{t+1}\|s_t, a_t)$, reward: $q_\theta(r\|s)$, action $q_\phi(a_t \| s_t)$, and value $v_\psi(s_t)$| 
|[3] Li, Minne, et al. **"Causal world models by unsupervised deconfounding of physical dynamics."** arXiv preprint arXiv:2012.14228 (2020).| Conventional **world models**: observational conditional $p(s_{t+1}\|s_t, a_t)$.|
|[4] Wang, Tongzhou, et al. **"Denoised MDPs: Learning World Models Better Than the World Itself."** International Conference on Machine Learning. PMLR, 2022.|Latent posterior : $q_\psi(x, y, z\|s, a)$, prediction decoder: $p_\theta(s, r\|x, y, z, a)$.|
|[5] Liu, Yuren, et al. **"Learning world models with identifiable factorization."** Advances in Neural Information Processing Systems 36 (2023): 31831-31864. | World model with **disentangled latent dynamics**: Transition model $p(s_{t+1}\|s_t, a_t)$, reward model $p(r_t \| s_t)$, observation model $p(o_t\|s_t)$, representation model $p(s_{t+1}\|o_t, s_t, a_t)$.|
|[6] Allen, James F., and Johannes A. Koomen. **"Planning using a temporal world model."** Proceedings of the Eighth international joint conference on Artificial intelligence-Volume 2. 1983.| Temporal logic | 
|[7] Kaelbling, Leslie Pack, Michael L. Littman, and Andrew W. Moore. **"Reinforcement learning: A survey."** Journal of artificial intelligence research 4 (1996): 237-285.| Locally weighted regression, state estimator |
|[8] Wiering, Marco A., and Martijn Van Otterlo. **"Reinforcement learning."** Adaptation, learning, and optimization 12.3 (2012): 729.|Transition models embody knowledge about the environment that can be exploited in various ways.|
|[9] Doll, Bradley B., Dylan A. Simon, and Nathaniel D. Daw. **"The ubiquity of model-based reinforcement learning."** Current opinion in neurobiology 22.6 (2012): 1075-1081.| The sequential transition structure of the task (a ‘world model’) |

This table shows that prior works have different descriptions of 'world model' in different contexts of problems. In our MDP context, we follow the prior works [3-5] and define our world model with the state transition and reward model. We don't think using our world model under this context misleads readers since a lot of existing work also use the same description. 

As we mentioned, **our key contribution is the bilinear causal representation learning**. If reviewer 8enR still feels using 'world model' may affect their justification on our contribution, we are happy to add more remarks in our appendix or change 'world model' to other terms that reviewer 8enR feels more appropriate. We believe this change won't have any influence to the core contribution of this paper. We hope the reviewer’s concern will be resolved.
World model (1120 citations) was proposed in 2018 for image-based POMDPs.
It does not aim to handle MDP problems.
Many RL methods, such as Dreamer (1228 citations), are based on world models.
As you can see, the world model in this paper is not for POMDPs.
Defining such a model with the same name as an existing well-known model will mislead many readers.
Furthermore, the authors did not cite literature [1] in the paper.





[1] David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv:1803.10122, 2018.
Could you please provide further elaboration on which parts are considered 'misleading'? Thank you in advance.
Dear Reviewer XKke

I think publishing such a definition would mislead many readers, so this is a serious mistake.

Reviewer 8enR
Dear Reviewer 8enR,

While the reviewer discussions have not officially started, I would like to discuss why the definition of 'world model' presented here is inappropriate. In my opinion, the world model could be defined as the dynamics model and the reward model in the MDP, similar to what the authors have mentioned. 

I’m not saying that every aspect of this paper is flawless, but I believe rejecting it solely based on this definition might be a bit unfair. I would appreciate your thoughts on this and welcome any corrections if am wrong. 

Reviewer XKke
World models were first defined in reference [1].
Many well-known reinforcement learning algorithms, such as Dreamer [2], are based on the world model.
Redefining world models are not appropriate, so I changed my score to 3.


[1] David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv:1803.10122, 2018.
[2] Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representations, 2019.
Thank you for your further response. My concerns have been well addressed and I would like to keep my positive rating.
We appreciate the dedicated efforts of all the insightful reviews and acknowledgment from reviewer ne9G. The contributive review helps improve our paper's quality during the revision and rebuttal phase.
We appreciate the dedicated efforts of all the insightful reviews from reviewer aSxm. The contributive review helps improve our paper's quality during the revision and rebuttal phase.
We thank reviewer XKke for acknowledging our response with additional theoretical and empirical results. We will answer your last clarification question below. 

**I just have one minor follow-up question: for Q3, my initial question was actually about whether maintaining a sparse $M$ here could be the necessary and sufficient condition for estimating the multiple confounders (e.g., different confounders affect different components in the state dynamics). Could the authors provide more insights on this point?**

Given our current assumption about the Action State Confounded MDP (ASC-MDP) and Structured Causal Model (SCM), the $\ell_0$ sparsity regularization we apply in the optimization of $M$ will be a **necessary yet not sufficient** condition for estimating all the latent confounders behind the causal graph. 
- The **necessity** of sparsity regularization in $M$ is clear that we need to avoid spurious correlation by 'trimming down' as many unnecessary edges in the transition dynamics while maintaining good prediction accuracy. 
- On the **sufficiency** part, although we cannot guarantee a complete recovery of true causal matrix $M$, we can bound estimation error of $M$ by some polynomial of the following terms: 
    - SCM's noise level $\sigma$ (how strong the confounding effect could be),
    - Sample size $n$,
    - Dimension of matrix $d\times d'$,
    - Structural complexity of ground truth mask $\||M\||_0$

    We have $\text{error}\leq poly(dd', \frac{1}{n}, \sigma, \||M\||_0)$. **Our Appendix B gives a more formal proof of it**. As a result, we can have some high-probability guarantees that if we want a near-optimal ($\epsilon$ error) estimation of the causal matrix $M$, i.e. if we want $\||\widehat{M}-M\||\leq \epsilon$ with high probability $1-\delta$, we may need $n(\epsilon, \delta)$ samples.
Thanks for your satisfactory answer. I keep my score.
Thank you for your detailed response and clarification. Most of my concerns have been well addressed, especially with the newly added Table 1 and the additional results on pixel experiments. I just have one minor follow-up question: for Q3, my initial question was actually about whether maintaining a sparse $M$ here could be the necessary and sufficient condition for estimating the multiple confounders (e.g., different confounders affect different components in the state dynamics). Could the authors provide more insights on this point? 

For now, given the current quality and the rebuttal, I would increase my rating to 6.
i thank the authors for their detailed responses, as well as adding the relevant papers in the related work. I've increased my score.