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

We would like to express our gratitude to reviewer 8enR for their valuable feedback. The reviewer acknowledges the clarity of our problem formulation and presentation, as well as the theoretical and empirical contribution of this work. We will answer the main questions below:

Q1. (Generality of Bilinear MDP Formulation): The conditions in Bilinear MDPs are difficult to achieve in real-world environments. Some real-world examples can be given for confounders $u_\pi$ and $u_c$.

We thank the reviewer for raising these important questions about the generality of bilinear MDP formulation. We discuss two categories of confounders in our experiment, which simulate some real-world problems like indoor navigation, robotic manipulation, and self-driving cars. One real-world example about the two categories of confounders $\mu_\pi$ and $\mu_c$ is in the autonomous driving environments. We have the states ['time-to-collision', 'daytime'] and actions ['throttle', 'steering']. Time-to-collision is the minimum ratio of relative distance to the closest vehicle divided by the velocity. The offline dataset may have two sources of confounders:

• Confounders from suboptimal behavior policy: $u_\pi$ is the confounder introduced by suboptimal behavior policies. For example, one aggressive driver always overtakes the front vehicles. This sub-optimal behavior policy $\pi_\beta(a | s)$ may have a large throttle even when the 'time-to-collision' is extremely small. Such behavior policies introduce confounders that cause spurious correlations between states ('time-to-collision') and actions (throttle). This spurious correlation may cause collision and have low reward if we directly learn from this dataset and deploy it online. In our experiment, we denote different behavior policies as 'random', 'medium', and 'expert' policies, leading to different confounded levels of $u_\pi$.

• Confounders from environment dynamics: $u_c$ is the confounder that exists in the transition dynamics $T(s'|s, a)$, and it causes spurious correlation between different state and action components in transition dynamics. In the real-world example, the traffic will be more crowded during the peak hours in a day and less crowded in the rest. Due to the confounder for 'crowdedness' and 'daytime', the model may fail under crowded scenarios out of peak hours, because it will falsely predict a lower 'time-to-collision' outside peak hours even if the traffic is dense. This is because the model is misled by the spurious correlation between 'time-to-collision' in $s'$ and daytime in $s$. In our experiment, we also have different confounders in different tasks in all 3 environments. We have discussed the environments in Appendix C.5.

We visualize the above confounded case in the MDP as Figure (a) and (b) in the rebuttal supplementary.

Q2. (Clarification of World Model): What is the world model in this paper? Does this paper use a world model to train the policy?

We thank the authors for raising the clarification about the world model. In our setting, our world model is composed of a dynamics model $T(s' | s, a)$ and a reward model $r(s, a)$. More concretely, during the planning phase, we have the following pessimistic value iteration:

Therefore, optimization of our pessimistic policy will be dependent on the learned dynamics model $\widehat{T}(s' | s, a)$ and reward function $r(s, a)$.

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

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.

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 sincerely thank the reviewer for their insightful and inspiring feedback. We are glad to know that the reviewer recognizes the novelty of our contributions, the clarity of our problem formulation, and the empirical results. We answer reviewers' questions about the causal learning process through the bilinear MDP optimization below.

Q1. (Advanced Model-based Offline RL baselines): The compared baselines do make sense. Some of them such as ICIL and CCIL are causal RL approaches.... consider to compare against a stronger baseline within the same family, such as: (Sun et al., ICML, 2023.).

We thank the reviewer for proposing this important baseline. We add additional experiments on the Unlock environment and illustrate the results in the first table in the general response. While MOBILE indeed demonstrates promising performance in the OOD generalization setting compared to the MOPO baselines, we still observe significant advantage on our methods that learns bilinear causal representation in the world model and use them for planning.

Q2. (Discussing recovery and explainability of confounders): It would be interesting to see a demonstration of how well BECAUSE can recover the underlying confounders... enhance the value of the proposed approach from an instrumental to an explanatory one.

As you can see in Figure (a) and (b) in the rebuttal supplementary, we give a simple MDP example in the autonomous driving tasks. We have the states ['time-to-collision', 'daytime'] and actions ['throttle', 'steering']. (Time-to-collision (TTC) is the ratio of relative distance to the closest vehicle divided by the velocity, and the TTC gets closer to 0 if a collision happens).

In general, the offline dataset encounters two sources of confounders:

• Confounders from suboptimal behavior policy: As is shown in Figure (a), $u_\pi$ is the confounder introduced by suboptimal behavior policies. For example, one aggressive driver always overtakes the front vehicles. This sub-optimal behavior policies $\pi_\beta(a | s)$ may have large throttle even when the 'time-to-collision' is extremely small. Such behavior policies introduce confounders that cause spurious correlations between states ('ttc') and actions ('throttle'). This spurious correlation may cause collision if we directly learn from this dataset and deploy online. In our experiment, we denote different behavior policies as 'random', 'medium', 'expert' policies, leading to different confounded levels of $u_\pi$.

• Confounders from environment dynamics: $u_c$ is the confounder that exists in the transition dynamics $T(s'|s, a)$, and it causes spurious correlation between different state and action components in transition dynamics. In the driving example, as shown in Figure (b), the traffic is more crowded (lower ttc) during the peak hours in a day and less crowded (higher ttc) in the rest. With confounder between next 'ttc' and current 'daytime', raw model fails to predict well under crowded scenarios out of peak hours, because it will falsely predict a lower 'ttc' outside peak hours even if the actual traffic is dense. This is because the model is misled by the spurious correlation between 'time-to-collision' in $s'$ and daytime in $s$. In our experiment, we also have different confounders in different tasks in all 3 environments (see Appendix C.5.)

Additionally, we are running some numerical simulations in this simple example, BECAUSE can indeed help identify the transition dynamics. We plan to add this to our final experiment results.

Q3. (Discussing selected benchmark): Why was it not possible and/or sensible to study the problem in the most well-known benchmark of offline RL research, the MuJoco environments of D4RL?

Our experiment simulator is adopted by some published works [1, 2, 3] in the causal RL domain. The selected environments have strong complexity in causal structure that leads to distribution mismatch between different training environments. Moreover, they are related to some real-world problems, such as indoor navigation, manipulation, and autonomous vehicles. These tasks all require strong reasoning capability to understand the interaction between the decision-making agents and other static or dynamic environment entities. On the other hand, in the mujoco-based D4RL environments, such cause-and-effect structures mainly focus on some physical transformation of the locomotion tasks, which have lower complexity as they do not model the interaction between different entities using the causal structure.

[1] Wang, Zizhao, et al. "Causal Dynamics Learning for Task-Independent State Abstraction." International Conference on Machine Learning. PMLR, 2022.

[2] Ding, Wenhao, et al. "Generalizing goal-conditioned reinforcement learning with variational causal reasoning." Advances in Neural Information Processing Systems 35 (2022): 26532-26548.

[3] Ding, Wenhao, et al. "Seeing is not believing: Robust reinforcement learning against spurious correlation." Advances in Neural Information Processing Systems 36 (2023).

Q4. (Adding societal impact): The paper discusses the limitations of the proposed approach in the final few sentences of the conclusion section. However, it does not discuss the potential negative societal impact of the presented work.

We thank the authors for raising the concern about potential negative societal impact. We added the following dicussion of societal impacts in our appendix:

Incorporating causality into model-based reinforcement learning framework provides the explanability of decision making, which helps humans understand the underlying mechanism of algorithms and check the source of failures. However, the learned causal graph may contain human-readable private information about the offline dataset, which could raise privacy issues. To mitigate this potential negative societal impact, the discovered causal graphs should be only accessible to trustworthy users.

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 would like to express our gratitude to reviewer XKke for their comprehensive and insightful feedback. We are glad to know that the reviewer recognizes the clarity of our problem formulation, the novelty and technical soundness of the proposed method, as well as the theoretical and empirical results. We answer key clarification questions about the causal learning process through the bilinear MDP optimization below.

Q1. (Learning the causal graph by learning $M$):

For the first concern on the feature space entanglement, our learning process alternates between feature learning and causal mask learning, therefore, the causal variables $\phi, \mu$ will best align with the causal relationship. In some harder cases with poor offline data quality (like the random behavior policy), we do observe a performance drop and significant objective mismatch issues as is shown in Figure 6. However, our methods still outperform baselines in these cases with the help of bilinear causal representation. Besides, our method is also robust under limited sample size or various spurious level of confounders, as is demonstrated in Figure 4b and 4c.

For the second concern on the optimization on the search space, we'd like to first clarify the relationship between M and G under the formulation of bilinear MDP:

The optimization of $M$ is a conditional independence testing (CIT) on $d\times d'$ pairs of features. Empirically, our causal discovery process is doing CIT on all the elements in the $d'\times d$ matrix. The optimization search space of $M$ is constrained by the size of $M\in \mathbb{R}^{d'\times d}$. Since the causal discovery is based on hypothesis testing statistics like $\chi^2$ testing, the ability to identify true causal graph $G$ will not be affected by the size of $M$.

Q2. (About assumption 1):

We thank the reviewer for asking insightful questions on our core assumption.

Regarding the first concern on the mapping invariance, if the confounder effects are small as the reviewer mentioned, it means either (a) the confounder is not important to the success of our task in either domain or (b) the dataset we have at hand is too biased to fully learn about true causality.

• For case (a), we empirically verify the invariance of the discovered causal graph with an experiment on the robustness to the spurious level of confounder, as is mentioned in RQ2 of Section 4.2. In this experiment, we vary the spurious levels (i.e. the scale of confounding effects) by changing the number of confounders in the environments. We show a consistent performance with fewer performance degradation compared to baselines without this bilinear causal representation, as is shown in Figure 4.c. As long as these latent confounders have enough statistical evidence to be identified with causal discovery, we can learn good causal representation.
• For case (b), which mostly corresponds to the 'random' policy setting where the behavior policy is heavily confounded by the policy confounder $u_\pi$, BECAUSE empirically has better performance compared to other causal RL or MBRL methods.

In conclusion, in either case (a) or (b) raised by the reviewer, BECAUSE can empirically outperform other MBRL baselines, which justifies our assumption of invariant causal mask.

Regarding the second concern on the equivalence with assumptions in previous works, we add some additional discussion in Table I of our 1-page rebuttal supplementary with a mathematical description of the specific meaning of the invariant causal graph and how it relates to the assumptions cited in our manuscripts.

Q3: (Multiple sources of confounders and identifiability of the causal graph):

The learning process of $M$ is conducted with two stages: regularized world model learning to deconfound $u_c$ in Equation (5) and (6), and policy-conditioned reweighting formulas to deconfound $u_\pi$ in (7).

• The first stage is to use the $\ell_0$ objective to regularize the world model loss while maintaining a sparse $M$.
• In the second stage, we average $M$ by reweighting its expectation across different trajectory batches.

We demonstrate the effect of the above deconfounding process both empirically and theoretically:

• Experiment results show that our method outperforms the existing causal RL and model-based RL baselines in both in-distribution and OOD settings in three environments.
• Theoretically, we provide additional analysis on how well this de-confounding process can identify the causal matrix $M$. Under reasonable assumptions about the offline dataset and underlying dynamics model, we can bound the estimation error of $M$ and have performance guarantees for learned policy in Theorem 1.

Q4: (Further structure decomposition in world model):

We thank the reviewer for providing this inspiring extension of our work. As we discuss in the related works, Denoised MDPs propose a structure decomposition according to the controllability and reward relevance of the state. To ground our bilinear representation in this context, one interesting way is to decompose $\phi, \mu$ into the controllable/non-controllable and reward-relevant/irrelevant parts. This will be an interesting future direction to make $M$ more compact.

Q5: (Experimentation with raw pixel inputs):

To further demonstrate the scalability of our method, we conduct additional experiments with raw-pixel inputs in the Unlock environments. The RGB input has a shape of (192, 192, 3), the detailed results can be found in the general response.

Thank you for your further response. My concerns have been well addressed and I would like to keep my positive rating.

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.

We would like to express our gratitude to reviewer aSxm for their insightful feedback and acknowledging the novelty of our contributions in practice and theory. We provide our response to the questions below.

Q1. (Additional Related works): Related works about theoretical offline RL using pessimism... multi-task RL algorithms weren't used as baselines given the setup of the experiments.

As the reviewer suggested, we added related references to offline RL theory in the related work section.

• Comparisons with the two related works: The reviewers raised two relevant works about multi-task RL in offline settings with linear function approximation (linear MDPs). They consider the same sampling mechanism -- offline setting, but different model assumption (linear MDPs) and problem setting (multi-tasks) as our bilinear MDPs and action-state confounded MDPs.
• Differences to multi-task offline RL: Our work does not primarily target multi-task learning, which aims to solve multiple tasks jointly in an efficient way [2] by leveraging their shared information. In our case, we focus on problems with unobserved causal effects -- the training data has a distribution mismatch with testing environments due to the unobserved confounders. We didn’t include multi-task RL as baselines since they do not target our challenges about causal effects, which makes the comparisons unfair to some extent. So we focus on more related baselines -- 10 causal RL and model-based RL baselines in Section 4.1 and Appendix C.6.

[1] Jin, Ying, Zhuoran Yang, and Zhaoran Wang. "Is pessimism provably efficient for offline rl?." International Conference on Machine Learning. PMLR, 2021.

[2] Yu, Tianhe, et al. "Conservative data sharing for multi-task offline reinforcement learning." Advances in Neural Information Processing Systems 34 (2021): 11501-11516.

Q2. (Applicability of theoretical results): The result in Theorem 1 is applicable for a Tabular MDP, but this hasn't been discussed in the paper anywhere.

The reviewer is correct that Theorem 1 is primarily applicable for the tabular case when the state and action space are finite. We added a remark under the theorem as the reviewer suggested:

• Our Theorem can be extended to continuous state space. Theorem 1 provides the performance guarantees of our proposed method in tabular case with finite state and action space. It lays a solid foundation to carry out more general cases with continuous state and action space. The current results can be extended to continuous state space by using covering number argument in our key lemmas, while handling action space need more adaptation and we leave it as interesting future works.
• Continuous space is potentially reduced to tabular cases in practice. Our proposed algorithm can be extended to continuous state or action space via some additional derivation on feature extraction. Prior works such as VQ-VAE [1] provide powerful encoders that can tokenize a continuous visual input into some discrete latent space with finite elements. We will add more discussion on the point in the revised manuscripts.

Q3. (Theoretical results justification): The result becomes void in the case when some (s_h, a_h) along the optimal policy has not been seen in the offline dataset.

The reviewer is insightful that offline dataset do need to satisfy some properties to make offline learning possible and well-posed. To our best knowledge, the current weakest assumption for the offline dataset to ensure provably learning an optimal policy is called (clipped) single-policy concentrability assumption [2, 3] (such as Definition 1 in [3]) -- the offline dataset need to include sufficient data over the region (state-action tuples $(s_h, a_h)$) that the optimal policy can visit. Under such weakest assumption, our Theorem 1 can ensure the performance of the proposed algorithm.

[1] Van Den Oord, Aaron, and Oriol Vinyals. "Neural discrete representation learning." Advances in neural information processing systems 30 (2017).

[2] Li, Gen, et al. "Settling the sample complexity of model-based offline reinforcement learning." The Annals of Statistics 52.1 (2024): 233-260.

[3] Rashidinejad, Paria, et al. "Bridging offline reinforcement learning and imitation learning: A tale of pessimism." Advances in Neural Information Processing Systems 34 (2021): 11702-11716.

Q4. (Comparison with existing theoretical works): How does the Theorem 1 fare to existing works? A detailed comparison on the improvement factors in the sample complexity would be useful.

We thank the reviewer for raising this question -- which is our key theoretical technical contribution (the proof pipeline for Theorem 1). Please refer to the general response for details due to the limited space. We provide a shortened version here:

• Technical contributions: developed new proof pipeline. We target a new problem --- bilinear MDPs with sparse core matrix $M$ --- brings new challenges since we need to use iterative algorithms with no closed-form solution, while prior works usually heavily depend on having a closed-form solution from the optimization problem. We develop a new proof pipeline that considers the accumulative errors during the process that we believe is of independent interest for future RL works.
• Sample complexity comparisons with prior works. To our knowledge, we are the first to focus on the problem —- bilinear MDPs with a sparse core matrix $M$ for causal deconfounding and generalization. So our sample complexity results can't be directly compared to existing works. But using prior works as references, it verifies the tightness of our results --- having a comparable dependency on all the salient parameters ($H$, $d$, $\xi$).

We sincerely thank reviewer 8KHa for their insightful and inspiring feedback. We are glad to know that the reviewer recognizes the novelty of our contributions, the clarity of our problem formulation, the theoretical contributions, and the empirical evidence that shows the advantages compared to baselines. We provide our response to the questions below.

Q1. (Assumption on Invariant Causal Graph and Bilinear Structure): The method relies on certain assumptions, such as the invariance of causal graphs, which might not always hold in real-world scenarios. Additionally, modeling the causal relationships with a linear structure may be overly simplistic, potentially limiting the accuracy and effectiveness of the algorithm in capturing more complex dependencies.

We thank the reviewer for this insightful discussion on our assumption. We elaborate more on the generality of invariant causal graph assumption and the bilinear representation structures.

1. Our assumption of invariant causal graph targets on the generalizable RL setting when the training and testing environments have distribution shift but share the underlying cause-and-effect relationship. This assumption is commonly made in many prior causal RL works, such as [1, 2, 3]. This causal invariance setting is actually applicable in many real-world decision-makin tasks like indoor navigation, autonomous driving and manipulation, as is shown in our experiment.
2. Regarding the over-simplification concern about our bilinear structure in MDP, the nonlinearity of the transition dynamics can be captured by the feature encoder $\phi(s, a), \mu(s')$. In more scalable contexts such as the foundation model, linear structure is also commonly used in CLIP [4], where the inner product is applied in the latent space.

[1] Zhang, Amy, et al. "Invariant causal prediction for block mdps." International Conference on Machine Learning. PMLR, 2020.

[2] Wang, Zizhao, et al. "Causal Dynamics Learning for Task-Independent State Abstraction." International Conference on Machine Learning. PMLR, 2022.

[3] Zhu, Wenxuan, Chao Yu, and Qiang Zhang. "Causal deep reinforcement learning using observational data." Proceedings of the Thirty-Second International Joint Conference on Artificial Intelligence. 2023.

[4] Radford, Alec, et al. "Learning transferable visual models from natural language supervision." International conference on machine learning. PMLR, 2021.

Q2. (Continuous action): Is this method efficient in handling continuous actions? During the planning phase, computing the argmax Q can be time-consuming.

We thank the reviewer for asking this important question in implementing the planner. As the reviewer mentioned about the planning phase, our method theoretically uses the following pessimistic value iteration:

In practice, the above formulation is applicable to continuous actions we use model predictive control (MPC). We first uniformly sample random actions in the action space $a_{sample}\in \mathcal{A}$, then roll out future states by imagination using the learned dynamics model $\hat{T}(s'|s, a_{\text{sample}})$. Finally, we will have a set of Q functions $\overline{Q}(s, a_{\text{sample}})$ and take $a = \arg\max_{a_{\text{sample}}\in \mathcal{A}} \overline{Q}(s, a_{\text{sample}})$ as the selected actions to rollout the next step.

We denote the number of random action samples per step as "planning population", which is a fixed hyperparameter for all the methods in experiment environments in Appendix Table 9.

Q3. (Scalability): The proposed method introduces additional complexity to the MBRL framework, which may pose challenges for practical implementation and scalability.

We agree that our method adds additional structure on top of traditional MBRL in exchange of better OOD generalizability. To further demonstrate the scalability and computational efficiency of our method, we conduct additional experiments with raw-pixel inputs in the Unlock environments (see supplementary Figure (c) for details). The RGB input has a shape of (192, 192, 3) and we encode them with 3-layer CNN. The results in the second table of the general response/Table III in the supplementary show that BECAUSE still outperforms ICIL and IFactor baselines with visual inputs by a clear margin. We also evaluate the inference speed of our method with both visual inputs and attach details in the one-page PDF. We believe the inference speed (Table IV in supplementary) is still in the acceptable range with a comparable inference speed as IFactor (>20 FPS on image data).

In conclusion, the empirical evidence verifies (a) the scalability of BECAUSE and (b) the computational efficiency with high-dimensional inputs.