CausalGDP: Causality-Guided Diffusion Policies for Reinforcement Learning

TL;DR

CausalGDP introduces a causality-aware diffusion policy framework for reinforcement learning, enabling better identification of causally effective actions and improving performance in complex control tasks.

Contribution

It integrates causal reasoning into diffusion-based RL, learning causal dependencies from offline data and guiding actions based on causal influence.

Findings

Outperforms state-of-the-art diffusion-based RL methods.

Effectively identifies causally impactful actions.

Achieves superior results in high-dimensional control tasks.

Abstract

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.