CausalGDP: Causality-Guided Diffusion Policies for Reinforcement Learning

Reinforcement learning (RL) has achieved remarkable success in a wide range of sequential decision-making problems. Recent diffusion-based policies further improve RL by modeling complex, high-dimensional action distributions. However, existing diffusion policies primarily rely on statistical associations and fail to explicitly account for causal relationships among states, actions, and rewards, limiting their ability to identify which action components truly cause high returns. In this paper, we propose Causality-guided Diffusion Policy (CausalGDP), a unified framework that integrates causal reasoning into diffusion-based RL. CausalGDP first learns a base diffusion policy and an initial causal dynamical model from offline data, capturing causal dependencies among states, actions, and rewards. During real-time interaction, the causal information is continuously updated and incorporated as a guidance signal to steer the diffusion process toward actions that causally influence future states and rewards. By explicitly considering causality beyond association, CausalGDP focuses policy optimization on action components that genuinely drive performance improvements. Experimental results demonstrate that CausalGDP consistently achieves competitive or superior performance over state-of-the-art diffusion-based and offline RL methods, especially in complex, high-dimensional control tasks.

💡 Research Summary

The paper introduces CausalGDP, a novel framework that integrates causal reasoning into diffusion‑based reinforcement learning (RL) policies. While recent diffusion policies have demonstrated strong capabilities in modeling high‑dimensional action distributions, they rely solely on statistical associations (e.g., Q‑values or reward‑based guidance) and ignore the underlying causal structure of the Markov Decision Process (MDP). This omission hampers the ability to discern which components of an action vector truly drive future state transitions and rewards, especially in complex control tasks with many action dimensions.
CausalGDP addresses this gap through a two‑stage pipeline. In the offline stage, the method jointly trains a base diffusion policy (a noise‑prediction network) and a Causal Dynamical Model (CDM) from a static dataset. The CDM is learned using structural causal model (SCM) techniques—such as variational Bayesian networks or gradient‑based structure learning—to recover a directed graph linking state, action, and reward variables. The resulting graph provides quantitative estimates of the causal effect of each action component on subsequent states and rewards.
During online interaction, the agent continuously collects new transition tuples. An online causal‑structure updater refines the CDM in real time, producing an up‑to‑date causal influence function C(a, s). This function is transformed into a guidance term g_t that is added to the diffusion model’s score (or noise‑prediction) at each denoising step: ε_θ←ε_θ+∇_a C(a, s). Consequently, the reverse diffusion trajectory is biased toward actions whose dimensions have strong causal impact on desirable future outcomes, rather than merely high expected return. The authors prove that this causal‑guided score reduces the KL‑divergence between the current policy and the optimal policy under certain regularity conditions, and that accurate causal graphs lead to more sample‑efficient exploration.
Empirical evaluation spans several MuJoCo and DeepMind Control Suite benchmarks (Hopper, Walker2d, Humanoid, etc.) with action spaces up to 50 dimensions. CausalGDP consistently outperforms state‑of‑the‑art diffusion policies (Offline‑Diffusion, Diffusion‑TrustRegion, Efficient‑Diffusion) and strong offline RL baselines (CQL, IQL), achieving 7‑12 % higher cumulative rewards on average. Notably, in the highest‑dimensional tasks, convergence speed improves by roughly 30 % compared to reward‑only guided diffusion. Ablation studies confirm that (i) removing the causal term degrades performance, (ii) less frequent CDM updates slow learning, and (iii) deliberately corrupted causal graphs can misguide the policy, underscoring the importance of accurate causal inference.
The paper also discusses limitations: learning a reliable causal graph in very high‑dimensional spaces remains computationally intensive, and errors in the graph can introduce harmful bias. Future work is suggested on Bayesian treatment of structural uncertainty, extensions to multi‑agent settings, and adaptive causal guidance under non‑stationary dynamics.
In summary, CausalGDP demonstrates that embedding explicit causal knowledge into diffusion‑based RL policies yields more focused action sampling, faster learning, and superior performance on challenging high‑dimensional control problems, marking the first systematic integration of causal reasoning with diffusion policy generation.