BECAUSE: Bilinear Causal Representation for Generalizable Offline Model-based Reinforcement Learning
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 \textbf{B}ilin\textbf{E}ar \textbf{CAUS}al r\textbf{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.
Paper
Similar papers
Peer review
Summary
Distribution shifts in the sampling of MBRL will lead to the objective mismatch between model and policy learning. Therefore, this paper aims to reduce the influence of the distribution shift for MBRL. It divides the confounders into $\mu_{\pi}$ and $\mu_{c}$ and then tries to capture causal representations to reduce the influence of the distribution shift, thus mitigating the objective mismatch problem.
Strengths
1. This paper is well-written and easy to follow. 2. In RL, this paper tries to address confounders, which are key problem in causal inference. 3. There are some new theoretical results in this paper. 4. Experiment results illustrate the effectiveness of the proposed method.
Weaknesses
1. The conditions in Bilinear MDPs are difficult to achieve in real-world environments. 2. Some real-world examples can be given for confounders $\mu_{\pi}$ and $\mu_{c}$.
Questions
1. What is the world model in this paper? Does this paper use a world model to train the policy?
Rating
2
Confidence
3
Soundness
2
Presentation
2
Contribution
2
Limitations
N/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.
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).
**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.
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"?
Yes, I think how to define the model in this paper is an important problem.
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.
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.
Strengths
* The proposed methodology is very sensible and interesting. Formulating a latent bilinear MDP to enable differentiable causal discovery approaches is a brilliant idea. * The paper studies the theoretical properties of the presented approach. Although the used theoretical tools are not novel per se, they serve for their purpose excellently. The conceptual connection between causal representation learning and bilinear MDPs is indeed novel and worthwhile acknowledging. * Despite the denseness of the technical material, the paper is very easy to read and understand. * The presented empirical results are satisfactory.
Weaknesses
* The compared baselines do make sense. Some of them such as ICIL and CCIL are causal RL approaches. However, the experiments part could be even stronger if the chosen model-based ORL approaches represented the state of the art a little better. TD3+BC and MOPO are outdated. One could for instance consider to compare against a stronger baseline within the same family, such as: 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
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?
Rating
7
Confidence
4
Soundness
4
Presentation
4
Contribution
4
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.
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.
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).
Strengths
- [*Motivation and general framework*]: The idea of using causality to model the data generation process in RL (or in general, dynamic systems) gives a principle way to model RL environments. By considering and modeling potential confounders, which often lead to distribution shifts or mismatches, the paper addresses a critical challenge in RL. The decomposition of confounders into dynamics and policy-related aspects is reasonable and well-motivated in this work. - [*Technical soundness*]: The model and optimization processes are technically robust. The proofs and analyses provided in the appendix have been checked and are generally correct, offering complementary insights to the main part of the paper. - The writing is clear and easy to follow,
Weaknesses
Since here most of the weaknesses and questions are mixed. So I put both here. One of the major weaknesses, which is also a primary question, is the inherent difficulty in fully learning the causal process through the bilinear MDP optimization (details in question 1-3). **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
I listed all the questions (mixed with weaknesses) in the above section.
Rating
6
Confidence
3
Soundness
2
Presentation
3
Contribution
2
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.
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.
Summary
This paper studies challenges in Model Based rl, namely: 1. Objective mismatch between estimating the model and policy learning for value optimization 2. Confoundedness in offline RL causing this mismatch The paper proposes a theoretically sound algorithm and experiments extensively to support their claims.
Strengths
1. Well written paper - Problem presentation is very clear 2. Experiments cover a many scenarios, and ablation studies have been provided. I'm curious why multi-task RL algorithms weren't used as baselines given the setup of the experiments.
Weaknesses
1. Has missed some relevant related work particularly in theoretical offline RL which use pessimism based on uncertainty. Please cite : https://arxiv.org/pdf/2402.12570, https://arxiv.org/pdf/2403.11574 2. The result in Theorem 1 is applicable for a Tabular MDP, but this hasn't been discussed in the paper anywhere. The result becomes void in the case when some (s_h, a_h) along the optimal policy has not been seen in the offline dataset.
Questions
1. How does the Theorem 1 fare to existing works? A detailed comparison on the improvement factors in the sample complexity would be useful.
Rating
6
Confidence
5
Soundness
3
Presentation
3
Contribution
3
Limitations
There are no significant limitations of the 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.
Strengths
The paper is well-written, proposing their intuition backed by three clearly laid out steps to walk the readers through the working of the algorithm. The insight of mitigates the objective mismatch with causal awareness learned from offline data is novel. The paper also provides a thorough theoretical analysis, including proofs of error bounds and sample efficiency.
Weaknesses
The method relies on certain assumptions, such as the invariance of causal graphs, which might not always hold in real-world scenarios. Additionally, modeling the causal relationships with a linear structure may be overly simplistic, potentially limiting the accuracy and effectiveness of the algorithm in capturing more complex dependencies.
Questions
Is this method efficient in handling continuous actions? During the planning phase, computing the argmax Q can be time-consuming. Limitation: The proposed method introduces additional complexity to the MBRL framework, which may pose challenges for practical implementation and scalability.
Rating
7
Confidence
4
Soundness
4
Presentation
4
Contribution
4
Limitations
none
i thank the authors for their detailed responses, as well as adding the relevant papers in the related work. I've increased my score.
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.
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.
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.
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.
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.
[Reply to Reviewer 8enR] Regarding the "misleading" part you mentioned
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 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.
Decision
Accept (poster)