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 Bilinear Causal Representation~(BECAUSE), an algorithm to capture causal representation for both states and actions to reduce the influence of the distribution shift, thus mitigating the objective mismatch problem. Comprehensive evaluations 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. 

In the presence of a hidden confounder u, we model the confounder behind the transition dynamics as a linear confounded MDP, and model the policy confounder induced by the sub-optimal behavior policies between the current state and action. According to the formulation in the overall Action-State Confounded MDP (ASC-MDP), the future state s'  will be conditionally independent given the dynamics confounder and the past state and action pairs (s, a). We parameterize the transition function into a bilinear structure, with two feature encoders $\mu, \phi$ and a core matrix M. 

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

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_\theta(s,a)$.  As we have access to the offline dataset, we learn an Energy-based Model~(EBM) to quantify the causal uncertainty of the state action pairs along the planning trajectories and take this into account during the online planning phase. We will plan a rewarding and most causally confident path to the goal in the online deployment. 

Lift: Object manipulation environment in RoboSuite. 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, there is an injected spurious correlation between the color of the cube and the position of the cube in the training phase. \diff{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. In the OOD environment \textit{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. We design such a crash scenario based on highway-env, where the goal is to create crashes between a pedestrian and AVs. In the OOD environment, the distribution of reward~(number of pedestrians) is different in online testing environments.