Robust Offline Reinforcement Learning Through Causal Feature Disentanglement

Abstract

Offline reinforcement learning suffers from critical vulnerability to data corruption from sensor noise or adversarial attacks. Recent research has achieved a lot by downweighting corrupted samples and fixing the corrupted data, while data corruption induces feature entanglement that undermines policy robustness. Existing methods fail to identify causal features behind performance degradation caused by corruption. To analyze causal relationships in corrupted data, we propose a method, Robust Causal Feature Disentanglement(RCFD). Our method introduces a learnable causal feature disentanglement mechanism specifically designed for reinforcement learning scenarios, integrating the CausalVAE framework to disentangle causal features governing environmental dynamics from corruption-sensitive non-causal features. Theoretically, this disentanglement confers a robustness advantage under data corruption conditions. Concurrently, causality-preserving perturbation training injects Gaussian noise solely into non-causal features to generate counterfactual samples and is enhanced by dual-path feature alignment and contrastive learning for representation invariance. A dynamic graph diagnostic module further employs graph convolutional attention networks to model spatiotemporal relationships and identify corrupted edges through structural consistency analysis, enabling precise data repair. The results exhibit highly robust performance across D4rl benchmarks under diverse data corruption conditions. This confirms that causal feature invariance helps bridge distributional gaps, promoting reliable deployment in complex real-world settings.

1. Introduction

Offline reinforcement learning (RL) addresses the critical challenge of learning optimal policies solely from pre-collected, static datasets without real-time environment interaction [1]. This paradigm is essential for applications where online exploration is impractical, costly, or unsafe. However, conventional offline RL methods exhibit significant vulnerability to data corruption, which arises from sensor noise, adversarial tampering, or environmental perturbations [2]. Such corruption manifests as random noise or adversarial perturbations that optimize systematic disruption and undermine both policy robustness and generalization [3,4,5,6,7].

These vulnerabilities are particularly concerning in high-stakes domains where policy failures can have severe consequences. In autonomous robotics, adversarial attacks on perception systems could result in catastrophic navigation failures in safety-critical environments such as surgical robots or autonomous vehicles [8,9]. Financial trading systems face comparable risks, where market data corruption or adversarial manipulation could trigger substantial economic losses and systemic market instability [10,11]. The importance of corruption-robust methods extends equally to robust control and safety fault-tolerant control systems [12,13,14,15].

The criticality of these applications demands robust offline RL methods that can maintain reliable performance despite data corruption. This motivates the development of causally robust approaches that can identify and mitigate corruption-induced feature entanglement while preserving the causal relationships essential for safe and reliable decision-making.

Recent research has achieved a lot. RIQL [5] leverages Huber loss and quantile estimators to handle heavy-tailed target distributions caused by dynamics corruption, while RDT [7] integrates embedding dropout and iterative data correction to mitigate corrupted inputs in sequential modeling. TRACER [6] model corruption-induced uncertainty via variational inference, using entropy-based weighting to downweight corrupted data. While RDT employs embedding dropout to improve robustness against erroneous inputs, it operates without understanding underlying causal relationships between state variables. The method’s iterative data correction process can identify corrupted actions and rewards but lacks theoretical grounding to distinguish corruption-sensitive features from those essential for environmental dynamics.

RIQL’s fundamental limitation stems from its reliance on temporal difference learning with Huber loss mechanisms that, while effective against dynamics corruption, cannot address feature entanglement at the representation level. TRACER’s uncertainty quantification gaps manifest in its entropy-based weighting mechanism, which downweights corrupted samples based on aggregate uncertainty without understanding specific causal mechanisms underlying corruption effects.

Existing approaches primarily focus on mitigating the performance degradation caused by damaged data, failing to explore the underlying causal relationships behind data corruption. To address these fundamental limitations, this work introduces the first causal-theoretic approach to data corruption offline reinforcement learning that explicitly models and exploits the distinction between causal features governing environmental dynamics and corruption-sensitive non-causal features. Our approach recognizes that data corruption induces systematic feature entanglement that can be formally addressed through structured causal modeling and targeted interventions.

This paper proposes a robust offline reinforcement learning framework based on causal feature invariance: Robust Causal Feature Disentanglement (RCFD). Its core innovations are as follows:

• Learnable Feature Disentanglement Mechanism: Integrating the CausalVAE framework to disentangle the corrupted data into causal and non-causal features (vulnerable to interference) provides formal guarantees of identifiability through structured causal models under idealized assumptions. Theoretically, this disentanglement confers a robustness advantage under data corruption conditions;
• Causal-Preserving Perturbation Training: This is conducted by generating samples and applying gaussian perturbations to non-causal features via counterfactual interventions while designing dual-path feature alignment loss and contrastive loss to enforce perturbation-invariant causal representations;
• Dynamic Graph Diagnosis and Reconstruction: This is conducted by employing graph convolutional attention networks to model spatiotemporal causal relationships among state variables, identifying corruption-sensitive edges through graph structure consistency loss to enable precise corrupted-data repair and smooth dynamic recovery.

Our experimental results demonstrate that the RCFD method broadly outperforms other baselines under data corruption conditions. This validates that the causality-based disentanglement approach effectively discerns the underlying causal relationships behind data corruption scenarios, thereby leveraging this insight to mitigate the impact of corrupted data.

2. Related Works

2.1. Offline Reinforcement Learning

Offline reinforcement learning addresses the critical challenge of learning policies from static datasets without environment interaction. Early approaches primarily tackled distributional shifts through policy constraints or value regularization. Methods like BCQ [16] and BEAR [17] enforced explicit policy constraints via action-space restrictions or MMD divergence, while BRAC [18] introduces behavioral constraints into actor–critic frameworks. These methods often induced excessive conservatism, limiting policy improvement potential. Concurrently, value regularization methods emerged to counter overestimation of out-of-distribution (OOD) [19,20,21,22] actions. CQL [23] learned conservative Q-functions by penalizing OOD actions, while IQL [24] approximated optimal in-distribution actions by estimating state-conditional upper expectiles of Q-values. Recent advances emphasize sequence modeling and diffusion policies. The Q-Transformer [25] integrated Transformer-based trajectory modeling with Q-learning, using Q-value regularization to bias action selection toward high-return behaviors. These innovations collectively address core challenges—distribution shifts, OOD actions, and trajectory stitching—while balancing policy improvement with safety constraints. FISOR [26] enforces hard safety constraints via the HJ reachability theory. ORL-RC [27] blends policy constraints and value regularization to mitigate over-pessimism. It uses hybrid policy evaluation with behavior cloning to balance conservatism and optimality.

2.2. Data Corruption Reinforcement Learning

Offline RL faces significant challenges when training data containing corrupted transitions due to sensor errors or adversarial tampering. RIQL [5] solves the problem of heavy-tail targets of Q functions under dynamics corruption conditions via Huber loss and quantile estimation; then they propose TRACER [6], leveraging Bayesian inference to model action-value uncertainty and downweight corrupted samples through entropy-weighted learning. Ye et al. [28] achieves additive dependence on corruption under general function approximation by controlling weighted uncertainty sums via the eluder dimension and new analytical techniques. Mandal et al. [29] model corruption in RLHF via state-confounded MDPs. Under conditions of low sample efficiency, Xu et al. [7] employed the sequence modeling method and the Decision Transformer [30] (DT) to leverage contextual knowledge, thereby enhancing data utility and mitigating the impact of data corruption.

2.3. Disentangled Reinforcement Learning

Disentangled reinforcement learning approaches focus on learning representations that explicitly separate entangled features to enhance robustness and generalization, with several recent methods leveraging auxiliary tasks for disentanglement in diverse RL settings. For example, SEAR [31] uses agent masks as weak supervision to impose an auxiliary reconstruction loss, forcing the policy to disentangle robot-centric features from environmental factors via variational inference, which significantly improves sample efficiency in visual control tasks by reducing spurious correlations across simulated robotics environments. Complementing these, CMID [32] formulates disentanglement as an adversarial RL task by minimizing conditional mutual information between features given a conditioning set derived from Markov decision process dynamics, enabling resilience to correlation shifts.

3. Preliminarys

3.1. Offline Reinforcement Learning

Reinforcement learning (RL) fundamentally addresses the challenge of how an autonomous agent should learn to make optimal sequential decisions through interactions within an environment. The standard setup is formalized as a Markov decision process (MDP). An MDP is mathematically represented by a tuple as follows:

$$M=(S,A,P,R,\gamma )$$

(1)

where S denotes the set of possible states an agent can perceive; A represents the set of actions available to the agent, $P(s^{\prime }s,a)$ defines the state transition probability distribution dictating the likelihood of transitioning to state s’ after taking action a in state s; $R(s,a,s^{\prime })$ signifies the immediate reward function specifying the scalar reward obtained after transitioning from s to $s^{\prime }$ via action a; and $\gamma$ in (0, 1] is the discount factor, determining the present value of future rewards and balancing immediate versus long-term gain.

Offline reinforcement learning tackles the problem of learning effective policies ($\pi$) solely from a pre-collected, static dataset $\mathcal{D}$ of interaction trajectories, where

$$\tau =\{(s_0,a_0,r_0,s_1),(s_1,a_1,r_1,s_2),\dots \}$$

(2)

The dataset is generated by a behavioral policy, ${\pi }_{\beta }$, without any further online interaction with the actual MDP during training. This constraint introduces the critical challenge of distribution shifts, necessitating specialized approaches for policy evaluation and optimization under the offline data regime.

3.2. Causal Inference

In reinforcement learning, environmental states encountered by agents contain both causal features that directly influence decision-making and non-causal features susceptible to noise disturbances. Causal inference formally establishes its separability through interventional invariance and counterfactual invariance. Specifically, when the environment satisfies the causal sufficiency assumption (all common causes are observable) and structural stability assumption (causal relationships remain time-invariant), state features can be uniquely identified and separated.

Interventional invariance serves as the cornerstone: For any intervention $do(n_t=n^{\prime })$ (forcing non-causal feature $n_t$ to value $n^{\prime }$), the causal feature $c_t$ preserves the distribution of rewards $r_t$ and state transitions:

$$\mathbb{E}\left[r_{t}do(n_t=n^{\prime })\right]=\mathbb{E}\left[r_t\right],P(s_{t+1}{c}_t,do(n_t=n^{\prime }))=P\left(s_{t+1}{c}_t\right)$$

(3)

where $n^{\prime }\in \mathcal{N}$ denotes the intervention scope, proving that the non-causal features $n_t$ do not alter decision dynamics.

3.3. Data Corruption

Data corruption manifests in two primary forms that require distinct mathematical characterization. Random corruption models naturally occurring noise inherent in physical systems, such as sensor inaccuracies or environmental interference. This phenomenon is formally described through additive perturbation:

$$\widehat{x}=x+\lambda ·\sigma ,\lambda \sim \mathcal{U}[-ϵ,ϵ]$$

(4)

where the stochastic term $\lambda$ sampled from a uniform distribution introduces isotropic deviations scaled by the dataset’s global standard deviation ($\sigma$) and $\epsilon$ governs corruption intensity, with higher values simulating severe measurement degradation.

Conversely, adversarial corruption represents intentional manipulation designed to maximize systemic disruption. It is modeled as a constrained optimization problem as follows:

$$\widehat{x}=arg\underset{{\parallel \delta \parallel }_{\infty }\le ϵ}{min}{Q}_{\mathrm{pretrain}}(x+\delta )$$

(5)

This minimizes a pretrained agent’s Q-value ($Q_{p}retrain$). The generated corruptions systematically exploit decision boundaries, simulating sophisticated attacks.

4. Methods

This section presents a robust framework for handling data corruption in RL. The approach formalizes the separation of causality-corrupted features, designs targeted perturbation mechanisms, implements spatiotemporal fault diagnosis, and integrates multi-objective optimization. Figure 1 illustrates the overall architecture.

State observations (${\mathbf{s}}_t$) first pass through the CausalVAE encoder to obtain $[{\mathbf{c}}_t,{\mathbf{n}}_t]$; then the perturbation module operates on ${\mathbf{n}}_t$ to generate counterfactual samples (${\tilde{\mathbf{s}}}_t$), while the GCAN module processes temporal sequences to detect corrupted edges and provide structural feedback for repair. All modules are unified through a multi-objective loss function that balances feature invariance, perturbation robustness, and structural consistency.

4.1. Causal Feature Disentanglement

We fundamentally reimagine the feature disentanglement mechanism by integrating the CausalVAE [33] framework. Unlike existing causality-aware RL methods that focus on clean data scenarios, our approach specifically targets corruption-induced feature entanglement in offline datasets through causal structure modeling. The framework implements a structured causal model (SCM) [34,35] layer that transforms independent exogenous factors into causally structured endogenous representations. Unlike heuristic separation methods, this approach provides learning objectives through guarantees of identifiability under idealized conditions while enabling explicit counterfactual interventions. The architecture establishes a bijective mapping between state observations and disentangled causal/non-causal representations:

$$s_t\stackrel{\mathrm{Encoder}}{\to }ϵ\sim \mathcal{N}(0,I)\stackrel{\mathrm{SCM} \mathrm{Layer}}{\to }{c}_t={(I-A^T)}^{-1}ϵ$$

(6)

where $ϵ$ captures non-causal noise patterns and $c_t$ represents invariant causal features governing environmental dynamics. The adjacency matrix A encodes directional dependencies between variables and satisfies Directed Acyclic Graph (DAG) constraints through the acyclicity regularization $H\left(A\right)=\mathrm{tr}\left({(I+{\displaystyle \frac{c}{m}}A\circ A)}^n\right)-n$.

Building upon the original structural causal model (SCM) framework, we formalize the causal disentanglement process by introducing a variational inference framework. The generative process for the observed state $s_t$ is defined as follows:

$$s_t=g(c_t,n_t)+ϵ,ϵ\sim \mathcal{N}(0,I)$$

(7)

where $c_t$ represents the causal features and $n_t$ represents the non-causal features.

We introduce a variational lower bound (ELBO) to help the optimization objective learn the encoder parameters ($q_{\varphi }$). The derivation of the variational lower bound is defined as follows (refer to Appendix A.1):

$$\begin{array}{cc}\hfill logp\left(s_t\right)& \ge {\mathbb{E}}_{q_{\varphi }(c_t,n_t|s_t)}\left[logp(s_t,c_t,n_t)\right]\hfill \\ \hfill & -D_{\mathrm{KL}}\left[q_{\varphi }(c_t,n_t|s_t)\parallel p\left(c_t\right)p\left(n_t\right)\right]\hfill \end{array}$$

(8)

where $q_{\varphi }$ denotes the encoder, while $p\left(c_t\right)$ and $p\left(n_t\right)$ represent the prior distributions. The KL divergence term is factorized as follows:

$$D_{KL}={\mathbb{E}}_{q_{\varphi }}\left[log{\displaystyle \frac{q_{\varphi }\left(c_t{s}_t\right)q_{\varphi }\left(n_t{s}_t\right)}{p\left(c_t\right)p\left(n_t\right)}}\right]$$

(9)

Maximizing the ELBO ensures that the variational posterior $q_{\varphi }(c_t,n_t|s_t)$ closely approximates the true, intractable posterior distribution over causal and non-causal features. Crucially, the structure of the KL divergence term $D_{KL}$ (specifically, its factorization into separate terms for $c_t$ and $n_t$ as shown in Equation (9)) facilitates the disentanglement process. By minimizing this KL divergence between the variational posterior and the factorized prior $p\left(c_t\right)p\left(n_t\right)$, the model is encouraged to separate the latent space into distinct causal features ($c_t$) that govern dynamics and non-causal features ($n_t$) sensitive to corruption, thereby achieving the core goal of causal disentanglement.

We define the marginal distribution of causal features under corruption as follows:

$$p_c^{\mathrm{corrupt}}\left(c\right)={\int }_n{p}_{\mathrm{corrupt}}\left([c,n]\right)dn$$

(10)

Since corruption primarily affects non-causal features, the distributional shift in the causal space is significantly attenuated as follows:

$$\mathrm{TV}(p_c^{\mathrm{clean}},p_c^{\mathrm{corrupt}})\le \delta ·ϵ$$

(11)

where TV is total variation distance between probability distributions and $\delta \in [0,1]$ represents the attenuation factor (typically $\delta \ll 1$), reflecting how corruption in the complete state space translates to changes in causal features. Consequently, the performance loss of the causal disentanglement strategy is bounded by the following equation:

$$|J^{\mathrm{clean}}\left({\pi }_c\right)-J^{\mathrm{corrupt}}\left({\pi }_c\right)|\le {\displaystyle \frac{2R_{max}}{{(1-\gamma )}^2}}·\delta ·ϵ·\sqrt{2\mathrm{KL}(p_c^{\mathrm{clean}}\parallel p_c^{\mathrm{corrupt}})}$$

(12)

To quantify the robustness advantage of causal disentanglement, we define the robustness gain as follows:

$$\Delta = |J^{\mathrm{clean}}\left(\pi \right)-J^{\mathrm{corrupt}}\left(\pi \right)|-|J^{\mathrm{clean}}\left({\pi }_c\right)-J^{\mathrm{corrupt}}\left({\pi }_c\right)|$$

(13)

Substituting the derived bounds yields the following equation:

$$\Delta \ge {\displaystyle \frac{2R_{max}}{{(1-\gamma )}^2}}·ϵ·\left[\sqrt{2\mathrm{KL}(p^{\mathrm{clean}}\parallel p^{\mathrm{noise}})}-\delta ·\sqrt{2\mathrm{KL}(p_c^{\mathrm{clean}}\parallel p_c^{\mathrm{corrupt}})}\right]$$

(14)

Given that $\delta \ll 1$ and $\mathrm{KL}(p_c^{\mathrm{clean}}\parallel p_c^{\mathrm{corrupt}})\le \mathrm{KL}(p^{\mathrm{clean}}\parallel p^{\mathrm{noise}})$, we obtain the following equation:

$$\Delta \ge {\displaystyle \frac{2R_{max}}{{(1-\gamma )}^2}}·ϵ·(1-\delta )·\sqrt{2\mathrm{KL}(p^{\mathrm{clean}}\parallel p^{\mathrm{noise}})}$$

(15)

This theoretical result establishes that under the causal invariance assumption, for any corruption rate $ϵ>0$, there exists a constant $\kappa =(1-\delta )>0$ such that

$$J^{*}\left({\pi }^{*}\right)-J^{\mathrm{corrupt}}\left({\pi }_c\right)\le J^{*}\left({\pi }^{*}\right)-J^{\mathrm{corrupt}}\left(\pi \right)-\kappa ·{\displaystyle \frac{2R_{max}}{{(1-\gamma )}^2}}·ϵ·\sqrt{2\mathrm{KL}(p^{\mathrm{clean}}\parallel p^{\mathrm{noise}})}$$

(16)

It is important to note that the theoretical bounds presented in this analysis serve primarily as qualitative analytical tools rather than quantities requiring precise numerical computation in practice. These expressions formalize the mechanistic relationship between causal disentanglement and robustness: corruption effects are attenuated when projected into the causal subspace (attenuation factor $\delta$), with performance degradation scaling proportionally to corruption intensity ($ϵ$) and distributional shift magnitude. The KL divergence terms appearing in these bounds are analytical constructs that help understand the underlying mechanisms, not parameters requiring estimation during algorithm execution. In high-dimensional continuous state spaces without known parametric forms, these quantities cannot be reliably computed.

To empirically validate the effectiveness of our causal feature disentanglement mechanism, we visualize the learned feature representations using t-Distributed Stochastic Neighbor Embedding (t-SNE) analysis. Figure 2 demonstrates the clear separation achieved between causal and non-causal features through our RCFD framework.

As illustrated in Figure 2, the original state representations exhibit significant feature entanglement, where causal and non-causal components are intermixed and indistinguishable in the latent space. This entanglement makes the learned policy vulnerable to data corruption, as the model cannot differentiate between corruption-sensitive features and those essential for environmental dynamics. Our causal disentanglement mechanism demonstrates remarkable separation after application. The causal features (represented by blue clusters) form distinct, well-separated regions from the non-causal features (represented by red clusters).

This visualization provides compelling evidence that our causal disentanglement mechanism effectively addresses the fundamental challenge of feature entanglement in corrupted offline RL datasets. Moreover, we supplement these qualitative results with rigorous quantitative evaluation methods to ensure statistical validity of the observed separation. We quantify the degree of feature separation using mutual information between causal and non-causal feature representations. We employ kernel density estimation to approximate the probability distributions and compute mutual information as follows:

$$I(c_t,n_t)=\int \int p(c_t,n_t)log{\displaystyle \frac{p(c_t,n_t)}{p\left(c_t\right)p\left(n_t\right)}}d{c}_{t}d{n}_t$$

(17)

where $p(c_t,n_t)$ is the joint probability density, while $p\left(c_t\right)$ and $p\left(n_t\right)$ are marginal densities estimated using Gaussian kernels with the bandwidth selected via cross-validation.

Our experimental results demonstrate that mutual information between causal and non-causal features significantly decreases to 0.23 after disentanglement under the RCFD framework. Our mutual information estimation employs kernel density estimation with Gaussian kernels to approximate probability densities. The estimation uses $N=100$ state pairs sampled from trained encoder outputs, with the bandwidth selected via cross-validation. To assess estimator reliability, we conduct bootstrap resampling with 500 iterations, yielding a 95% confidence interval of $[0.19,0.28]$ with a standard deviation of $\sigma =0.031$, indicating reasonably stable estimation across data subsamples.

The t-SNE visualization in Figure 2 provides qualitative corroboration, showing clear spatial separation between causal and non-causal feature clusters with minimal overlap regions. This visual evidence is consistent with quantitative MI measurements indicating successful but imperfect disentanglement.

The visualization and quantitative assessments provide compelling evidence that our causal disentanglement mechanism effectively addresses the fundamental challenge of feature entanglement in corrupted offline RL datasets. The clear separation enables our subsequent causality-preserving perturbation training to selectively intervene in non-causal features while preserving the integrity of causal dynamics, ultimately leading to more robust policy learning under data corruption scenarios.

4.2. Causality-Preserving Perturbation Mechanism

Based on the counterfactual reasoning framework [34,35,36,37], the perturbation mechanism operates specifically on the non-causal features ${\mathbf{n}}_t$ extracted by the encoder. The module receives the factorized state representation $[{\mathbf{c}}_t,{\mathbf{n}}_t]$ and applies selective interventions to simulate corruption scenarios while maintaining causal consistency. This targeted approach leverages the clear separation achieved by the disentanglement module, ensuring that perturbations do not interfere with environmentally critical causal features (${\mathbf{c}}_t$). Based on a disentangled feature space (${\mathbf{s}}_t=[{\mathbf{c}}_t,{\mathbf{n}}_t]$), our conditional intervention generator is implemented as follows:

$$\Delta \mathbf{n}\sim \mathcal{N}(0,{\sigma }_{\mathcal{N}}^2\mathbf{I}){\tilde{\mathbf{s}}}_t=[{\mathbf{c}}_t,{\mathbf{n}}_t+\Delta \mathbf{n}]$$

(18)

where $\Delta \mathbf{n}\in {\mathbb{R}}^{d/2}$ is a Gaussian noise vector with ${\sigma }_{\mathcal{N}}=0.1$. The state intervention probability p is set to 0.2, indicating that we intervene on 20% of states per batch by perturbing their non-causal features through additive Gaussian noise ($\sigma =0.1$). This design preserves the causal features ${\mathbf{c}}_t$ to maintain environmental dynamics consistency while disturbing the non-causal features ${\mathbf{n}}_t$ to simulate data corruption.

The counterfactual samples enhance feature invariance via a dual-path framework:

$${\mathcal{L}}_{\mathrm{cause}}={\mathbb{E}}_{s\sim \mathcal{D}}\{1-cos(\varphi \left({\mathbf{n}}_t\right),\varphi \left({\tilde{\mathbf{n}}}_t\right))\}$$

(19)

where ${\varphi }_{\mathrm{causal}}:{\mathbb{R}}^d\to {\mathbb{R}}^k$ is the causal feature extractor. Feature alignment minimizes the discrepancy in causal features between original states and perturbed states.

To further enhance the robustness of the model against data corruption, we introduce a contrastive learning loss function that enforces feature invariance between clean states and their perturbed counterparts. This loss function operates within the causal feature space, minimizing the discrepancy between the representations of original states ($s_t$) and counterfactual states (${\tilde{s}}_t$) generated through our perturbation mechanism. The contrastive loss function is formally defined as follows:

$${\mathcal{L}}_{\mathrm{contrast}}=-{\mathbb{E}}_{s_t\sim \mathcal{D}}\left[log{\displaystyle \frac{exp\left(\frac{\varphi \left(s_t\right)·\varphi \left({\tilde{s}}_t\right)}{\tau }\right)}{{\sum }_{s_k\in \mathcal{N}}exp\left(\frac{\varphi \left(s_t\right)·\varphi \left(s_k\right)}{\tau }\right)}}\right]$$

(20)

where $s_t,{\tilde{s}}_t$ are positive pairs, ${\tilde{s}}_t$ is the perturbation of $s_t$, ${\left\{s_k\right\}}_{k=1}^K$ is randomly sampled from the dataset constituting the negative sample set N, and $\tau$ is the temperature hyperparameter, where $\tau =0.1$.

This design ensures that the model learns representations invariant to corruption-induced variations, which is critical for maintaining performance under adversarial or noisy conditions.

4.3. Dynamic Causal Graph Diagnosis and Reconstruction

We propose a spatiotemporal-aware fault diagnosis module whose core innovation lies in modeling data corruption as a disruption of graph-structured causality among state variables. The graph convolution attention hybrid network (GCAN) module employs a principled approach to identify and reconstruct corrupted edges through structural consistency analysis.

In offline RL scenarios, we typically have limited prior knowledge about the causal structure among state variables. Our GCAN module addresses this challenge through a hybrid initialization strategy that combines minimal domain knowledge with data-driven structure learning. The core intuition is that corrupted data will exhibit abnormal attention patterns and violate the learned causal structure, enabling targeted detection and repair. First we start prior-based initialization. We utilize only basic structural assumptions available in most RL environments. For example, in locomotion tasks, we only assume that sequential state components may have dependencies, and the position and velocity of the same joint are related. Then given the limited prior knowledge, the GCAN primarily relies on data-driven learning to discover the causal graph structure. To handle the uncertainty from limited priors, we employ probabilistic initialization as follows:

$$A_{0_{ij}}\sim Beta({\alpha }_{\mathrm{prior}},{\beta }_{\mathrm{prior}})$$

(21)

where ${\alpha }_{\mathrm{prior}}$ and ${\beta }_{\mathrm{prior}}$ are set based on the confidence level of available prior knowledge. For most connections where we have no prior knowledge, we use ${\alpha }_{\mathrm{prior}}$ = ${\beta }_{\mathrm{prior}}$ = 1.

Given the sparse initial structure ${\mathbf{A}}_0$, we learn the complete causal graph through gradient-based optimization. We first embed the temporal state–action–reward sequences into node features:

$${\mathbf{X}}^{\left(0\right)}={\mathbf{W}}_e·\mathrm{Concat}\left[{\mathbf{s}}^{\left(t\right)},{\mathbf{a}}^{\left(t\right)},r^{\left(t\right)}\right]+{\mathbf{b}}_e$$

(22)

where ${\mathbf{s}}^{\left(t\right)}\in {\mathbb{R}}^{d_s}$, ${\mathbf{a}}^{\left(t\right)}\in {\mathbb{R}}^{d_a}$, and $r^{\left(t\right)}\in \mathbb{R}$ represent the state, action, and reward at timestep t, respectively.

The adjacency matrix is then refined through attention-based learning that respects the initial structure, where

$${\mathbf{A}}_{ij}={\mathbf{A}}_{0_{ij}}\odot \sigma \left({\displaystyle \frac{{\mathbf{q}}_i^{\top }{\mathbf{k}}_j}{\sqrt{d_k}}}\right),{\mathbf{q}}_i={\mathbf{W}}_Q{\mathbf{x}}_i,{\mathbf{k}}_j={\mathbf{W}}_K{\mathbf{x}}_j$$

(23)

The complete structure learning objective combines reconstruction accuracy with structural constraints as follows:

$$L_{\mathrm{consistency}}=\parallel \mathbf{X}-f_{\theta }{(\mathbf{X},\mathbf{A})\parallel }^2+{\lambda }_{\mathrm{DAG}}h\left(\mathbf{A}\right)+{\lambda }_{\mathrm{sparse}}{\parallel \mathbf{A}\parallel }_1$$

(24)

where $h\left(\mathbf{A}\right)=tr\left({(\mathbf{I}+c\mathbf{A}\odot \mathbf{A})}^n\right)-n$ enforces DAG constraints and $D_{KL}(\mathbf{A}\parallel {\mathbf{A}}_0)$ regularizes deviation from prior knowledge.

Once the causal structure is learned, we leverage attention weights as corruption indicators. For each timestep t, we compute the attention-based edge importance scores as follows:

$$e_{ij}^{\left(t\right)}={\displaystyle \frac{1}{H}}\sum _{h=1}^H\mathrm{softmax}\left({\displaystyle \frac{{\mathbf{q}}_i^{h\top }{\mathbf{k}}_j^h}{\sqrt{d_k}}}\right)$$

(25)

where H is the number of attention heads. Corrupted edges are identified through temporal consistency analysis.

For detected corrupted edges $E_{\mathrm{corrupt}}=\left\{(i,j)\right|{\delta }_{ij}^{\left(t\right)}>\tau \}$, we apply targeted interventions as follows:

$${\mathbf{s}}_{\mathrm{repaired}}={\mathbf{s}}_{\mathrm{original}}+\sum _{(i,j)\in E_{\mathrm{corrupt}}}{\alpha }_{ij}·({\mathbf{s}}_{\mathrm{predicted}}\left[j\right]-{\mathbf{s}}_{\mathrm{original}}\left[j\right])$$

(26)

where ${\alpha }_{ij}$ represents learned repair weights and ${\mathbf{s}}_{\mathrm{predicted}}$ comes from the structural prediction using uncorrupted edges.

4.4. Multi-Objective Optimization

The RCFD framework integrates the diverse loss functions derived from causal feature disentanglement, perturbation training, and graph-based diagnostics into a unified multi-objective optimization objective, and it is formulated as follows:

$$\mathcal{J}={\mathcal{L}}_{\mathrm{BC}}+{\lambda }_{cause}{\mathcal{L}}_{\mathrm{cause}}+{\lambda }_{contrastive}{\mathcal{L}}_{\mathrm{contrastive}}+{\lambda }_{consistency}{\mathcal{L}}_{\mathrm{consistency}}$$

(27)

where ${\mathcal{L}}_{BC}$ serves as the behavioral cloning loss, which anchors the policy to the offline dataset by minimizing action distribution discrepancies, ensuring stable learning under distributional shifts.

This comprehensive objective function aims to harmonize the competing goals of policy fidelity, causal invariance, noise robustness, and structural integrity.

5. Experiment

In this section, we test our method in the D4rl dataset [38]. We set the dataset to be under different data corruption conditions, including both random and adversarial attacks on states, actions, and rewards.

5.1. Environment Setting

The evaluation environment of our method is set on D4RL, which stands as the most widely adopted offline reinforcement learning dataset and is extensively utilized under data-deficient conditions. Three core tasks were selected, HalfCheetah, Walker2d, and Hopper, each representing distinct locomotion challenges. For each task, datasets such as “Halfcheetah-medium-replay-v2” were employed, characterized by approximately 100,000–200,000 samples, as detailed in Appendix B.3. Our experiments are organized into three tasks and three diverse data corruption conditions, and both random and adversarial corruption are included. The corruption rate is set to 0.3, and the corruption scale is 1 for the main results. Random corruption involved additive Gaussian noise scaled by the dataset’s global standard deviation, while adversarial corruption is introduced through the Projected Gradient Descent attack with a pretrained EDAC [39] agent.

Following the established protocols from RIQL [5] and other robust offline RL works, we ensure a fair comparison by carefully aligning all experimental parameters and hyperparameter settings with the baseline methods. Our implementation adopts identical network architectures, learning rates ($3\times {10}^{-4}$), batch sizes (256), and training procedures (1M environment steps) across all compared methods to ensure that performance differences can be attributed to algorithmic improvements rather than implementation variations.

We combine our method with the baselines of offline reinforcement learning methods under data corruption conditions, such as IQL [24], CQL [23], RIQL [5], RDT [7], and the causal RL baseline CausalExploration [40]. All experiments are trained for 1m environment steps using four random seeds.

5.2. Evaluation Under Random Corruption

We first test the performance RCFD under random corruption conditions. RCFD demonstrated resilience against random corruption, achieving a mean normalized score of 58.6 across all tasks and attack elements, as summarized in

Table 1. Average normalized performance under random data corruption.

Environment	Attack Element	CQL	IQL	Causal	RIQL	RDT	RCFD
halfcheetah	obs	12.8 ± 2.2	16.4 ± 3.8	11.7 ± 1.4	21.6 ± 3.5	19.9 ± 3.3	29.4 ± 3.7
	acts	43.8 ± 2.7	43.7 ± 1.2	39.9 ± 2.9	41.1 ± 2.4	30.1 ± 3.2	40.0 ± 2.4
	reward	39.1 ± 6.1	39.1 ± 0.5	42.0 ± 1.1	42.5 ± 1.6	34.6 ± 3.6	40.4 ± 2.5
walker2d	obs	24.9 ± 16.3	24.9 ± 7.3	20.9 ± 10.7	33.7 ± 13.1	53.1 ± 4.8	44.7 ± 6.1
	acts	32.5 ± 27.1	77.9 ± 1.5	78.1 ± 7.4	81.3 ± 5.0	63.1 ± 4.9	71.4 ± 5.5
	reward	58.5 ± 23.5	59.6 ± 13.6	81.4 ± 4.1	56.7 ± 12.7	53.3 ± 17.2	73.1 ± 8.5
hopper	obs	40.7 ± 11.4	44.1 ± 22.4	40.1 ± 9.3	23.3 ± 8.8	65.6 ± 3.2	84.3 ± 7.3
	acts	65.3 ± 34.3	64.7 ± 18.6	56.4 ± 9.5	66.3 ± 7.3	67.3 ± 12.3	71.6 ± 4.6
	reward	60.5 ± 12.8	62.7 ± 16.5	48.1 ± 2.4	66.0 ± 21.2	66.4 ± 11.3	72.5 ± 14.8
mean		42.0	47.2	46.5	48.9	50.4	58.6

. This performance exceeded the nearest baseline RDT by approximately 16.27%, highlighting RCFD’s ability to mitigate feature entanglement through causal invariance.

A pronounced performance advantage manifests in observation corruption scenarios for HalfCheetah and Hopper environments. In HalfCheetah, RCFD attains a score of 29.4, representing a 36.2% improvement over the second-best baseline RIQL (21.58). In Hopper, RCFD achieves a score of 84.27, substantially exceeding RDT’s performance of 65.6 by 28.5%, while simultaneously demonstrating perfect stability. These empirical findings substantiate the theoretical foundation that our causal disentanglement mechanism effectively isolates corruption-sensitive non-causal features from invariant causal dynamics, particularly when corruption primarily manifests as observational noise in specific locomotion paradigms.

Conversely, the Walker2d environment under observation corruption exposes a notable limitation in our approach. RCFD achieves a score of 44.65, underperforming relative to RDT (53.05) by 15.8%, despite outperforming RIQL (33.71). The performance disparity in Walker under observation corruption warrants detailed discussion. Walker’s bipedal locomotion involves complex interdependencies between sequential joint pairs where causal relationships are highly entangled. Unlike HalfCheetah’s relatively independent joint actuations or Hopper’s simpler single-leg dynamics, Walker requires coordinated balance control where observation corruption in one joint can cascade through the kinematic chain. Our current causal disentanglement mechanism, which assumes relatively sparse causal structures, struggles to capture these dense interaction patterns. The observed performance gap indicates the necessity for incorporating environment-specific structural priors into the causal feature learning process to accommodate the diverse causal architectures inherent in different robotic morphologies and control paradigms.

The results validate that causal feature disentanglement bridges distributional gaps, as corrupted samples were downweighted while invariant features maintained environmental dynamics.

5.3. Evaluation Under Adversarial Corruption

To further evaluate RCFD, we introduce adversarial data corruption under diverse data corruption scenarios. Under adversarial corruption, RCFD maintained robust performance with a mean score of 54.6, outperforming all baselines by approximately 15.92% over the the best-performing baseline (RIQL), as evidenced in

Table 2. Average normalized performance under adversarial data corruption conditions.

Environment	Attack Element	CQL	IQL	Causal	RIQL	RDT	RCFD
halfcheetah	obs	25.3 ± 22.3	23.7 ± 2.6	12.4 ± 3.7	31.3 ± 4.8	30.2 ± 4.0	19.9 ± 3.9
	acts	33.6 ± 14.4	37.5 ± 1.6	32.8 ± 3.7	31.4 ± 4.5	24.8 ± 2.9	30.5 ± 7.2
	reward	43.7 ± 0.4	42.5 ± 2.0	44.1 ± 0.9	43.5 ± 1.5	34.2 ± 2.7	43.0 ± 0.2
walker2d	obs	69.8 ± 4.9	27.4 ± 6.1	23.6 ± 12.7	28.7 ± 18.9	49.7 ± 2.3	66.7 ± 11.7
	acts	29.5 ± 6.8	47.4 ± 8.1	39.1 ± 12.4	57.8 ± 5.6	38.6 ± 10.1	55.8 ± 11.3
	reward	47.3 ± 21.7	68.9 ± 6.0	85.1 ± 1.3	75.6 ± 8.2	63.9 ± 1.9	80.7 ± 1.1
hopper	obs	74.6 ± 11.8	41.4 ± 6.4	35.5 ± 6.6	41.2 ± 7.4	60.3 ± 12.3	46.6 ± 13.9
	acts	48.9 ± 24.9	61.7 ± 9.8	49.8 ± 15.0	53.5 ± 28.5	36.0 ± 14.5	76.2 ± 19.9
	reward	25.2 ± 3.4	59.8 ± 9.3	51.5 ± 6.1	60.5 ± 9.4	67.4 ± 11.7	71.9 ± 9.9
mean		44.2	45.6	41.6	47.1	45.0	54.6

. The causal disentanglement mechanism effectively reduces the distributional shift between clean and corrupted data.

RCFD’s robustness depends strongly on the interplay between the morphology and dynamics of the environment and the specific corruption type. In Hopper and Walker2d, the method excels when adversarial corruption targets actions or rewards, as the GCAN-based repair and contrastive representation learning effectively mitigate the impact of such perturbations. However, in HalfCheetah and Hopper, sophisticated adversarial manipulation of state variables continues to expose weaknesses in the learned disentanglement. The causal invariance framework relies on causal sufficiency and structural stability assumptions that may be violated in real-world scenarios. Considering unobserved confounders that violate causal sufficiency, suppose an unobserved confounder U affects both observed state variable $s_i$ and reward r. The true causal model includes $U\to s_i$ and $U\to r$, but our framework only observes $s_i$. This violation affects RCFD because the encoder may incorrectly attribute variations caused by U to either $c_t$ or $n_t$. If U is correlated with corruption patterns, the learned disentanglement becomes unreliable. Overall, the adversarial corruption evaluation confirms the efficacy of RCFD’s causal-feature-centric design in preserving policy integrity against structured, high-impact adversarial attack.

5.4. Ablation Experiments

5.4.1. Ablation Study on Different Modules

To analyze component contributions, we conduct ablation studies by sequentially removing core architectural elements. Ablation experiments were conducted to dissect RCFD’s core innovations, sequentially removing components to quantify their impact on corruption robustness, as shown in Figure 3. When the causal disentanglement mechanism was disabled, mean performance under random corruption dropped by 28.24%, emphasizing that feature separation is indispensable for mitigating entanglement. Similarly, omitting the causality-preserving perturbation training generates counterfactual samples via Gaussian noise in non-causal features. This reduced scores by 14.67% due to degraded invariance in causal representations. When the modification detection work of the GCAN’s damage data was ignored, the score dropped by 13.5%.

5.4.2. Ablation Study on Corruption Rates

To comprehensively evaluate RCFD’s robustness across varying corruption intensities, we conduct ablation studies on corruption rates ranging from 0.1 to 0.5 in the Walker environment. As the Figure 4 shows, This analysis examines how different corruption intensities affect the relative performance advantages of our method compared to baseline approaches.

RCFD maintains consistently strong performance under reward corruption across all tested corruption rates. This consistent performance advantage validates our theoretical framework’s effectiveness in isolating reward-related corruption from causal dynamics. At the highest tested corruption rate, RCFD maintains functional performance across all corruption types.

These results validate our theoretical framework’s core premise: by explicitly modeling and preserving causal relationships, RCFD achieves superior robustness that scales effectively with corruption intensity.

5.4.3. Ablation Study on Hyperparameters

As the Figure 5 shows, the contrastive learning coefficient ${\lambda }_{\mathrm{contrastive}}$ demonstrates a pronounced non-monotonic relationship with model performance across all corruption scenarios. At ${\lambda }_{\mathrm{contrastive}}=0.5$, the framework achieves optimal performance under corruption conditions. The dramatic performance drop at ${\lambda }_{\mathrm{contrastive}}=0.3$ suggests that insufficient contrastive regularization fails to enforce adequate feature invariance between clean and perturbed states, allowing corruption-sensitive representations to dominate the learned policy. Conversely, excessive contrastive weighting at ${\lambda }_{\mathrm{contrastive}}=1.0$ leads to over-regularization, constraining the model’s capacity to capture nuanced environmental dynamics and resulting in reduced performance scores.

The consistency regularization parameter ${\lambda }_{\mathrm{consistency}}$ exhibits a more stable performance profile, though with notable variations that illuminate the role of temporal structural coherence in corruption robustness. The optimal configuration occurs at ${\lambda }_{\mathrm{consistency}}=0.5$. This configuration strikes an effective balance between structural constraint enforcement and model flexibility, enabling the dynamic graph diagnostic module to identify corrupted edges without imposing excessive rigidity on the learned causal relationships. Interestingly, the consistency parameter shows diminishing returns at higher values, with ${\lambda }_{\mathrm{consistency}}=1.0$ producing marginally reduced performance, suggesting that overly strict temporal consistency constraints can impede the model’s ability to adapt to legitimate environmental variations.

The causal regularization coefficient ${\lambda }_{\mathrm{cause}}$ exhibits complex optimization dynamics that vary significantly across different corruption scenarios. Comprehensive analysis across all three attack vectors reveals that ${\lambda }_{\mathrm{cause}}=0.7$ delivers the most robust overall performance.

6. Conclusions

This paper introduces Robust Causal Feature Disentanglement (RCFD), a novel framework for corruption-robust offline reinforcement learning. By integrating causal disentanglement, perturbation training, and dynamic graph diagnostics, RCFD fundamentally addresses feature entanglement induced by data corruption. The experiment results confirm that isolating causal invariance is effective for bridging distributional gaps in corrupted data. While RCFD demonstrates effectiveness on D4RL locomotion benchmarks, several limitations constrain broader applicability. Our framework relies on causal sufficiency and structural stability assumptions that may be violated in complex real-world scenarios. High-dimensional visual observations present challenges for causal structure learning in pixel space, requiring integration with world models or hierarchical representations to make causal discovery tractable. Partial observability fundamentally violates causal sufficiency since hidden states constitute unobserved confounders, necessitating extensions to belief-state representations and partial identification techniques. Future work will address these limitations through hierarchical causal models for vision tasks, online structure learning for non-stationary environments, and theoretical analysis under relaxed causal assumptions.