CausalDream: Causal World Models by Deconfounding of Dynamics in Episodically Non-Stationary Environments

World models trained on observational data can conflate the effects of actions with those of unobserved physical factors, such as gravity or mass, that stay fixed within an episode but vary across episodes. When the latent state fails to retain these factors, they confound the learned dynamics and undermine policy robustness. Simply conditioning the model on an inferred context neither ensures that the dynamics use it nor exploits it to generate experience beyond the current policy. We propose CausalDream, which treats these factors as exogenous variables of a structural causal model. Its Causal World Model conditions latent transitions on exogenous variables that a label-supervised estimator infers from a short probe trajectory. A Causal Consistency Regularizer reconstructs these estimates from latent trajectories so that the dynamics retain them. Counterfactual imagination then intervenes on the first imagined action while holding the estimates fixed, broadening the experience available for policy learning. Across eight DeepMind Control tasks, CausalDream achieves the highest aggregate learning AUC among all compared methods, 12.9% above an architecture-matched baseline that conditions on the same estimates without the causal components. It also predicts the effects of action interventions more accurately on all four evaluated tasks, showing that its gains go beyond context conditioning alone.