Technical Summary: SHAP-Guided Kernel Actor-Critic for Explainable Reinforcement Learning

1. Problem Statement

Actor-Critic (AC) methods are foundational to Reinforcement Learning (RL) but suffer from limited interpretability. Standard AC architectures treat state features equally, failing to reveal the heterogeneous impacts of individual state dimensions on the reward. While existing Explainable RL (XRL) methods exist, they are predominantly post-hoc (analyzing a trained policy without altering training) or intrinsic but coarse-grained (e.g., decision trees or symbolic policies). Crucially, few methods provide dimension-level attributions that can be integrated directly into the learning loop to guide policy updates.

Furthermore, integrating attribution mechanisms into AC presents two fundamental challenges:

1. Attribution Reliability: Using attributions to modify the functional form of the Critic can violate compatible function approximation and disrupt the contraction properties of the Bellman operator, especially when the underlying kernel dictionary evolves.
2. Optimization Stability: RL states are often perturbed by noise, stochastic dynamics, and distributional shifts. Attribution-sensitive methods are vulnerable to spurious correlations unless carefully regularized, as instability in attribution values can amplify approximation errors and lead to brittle learning dynamics.

The core research question is how to design an efficient, provable, attribution-aware AC algorithm that computes state attributions online while maintaining stability under state perturbations.

2. Methodology: RSA2C

The authors propose RKHS-SHAP-based Advanced Actor-Critic (RSA2C), an attribution-aware, two-timescale AC algorithm. RSA2C integrates state attributions directly into the learning process via a Mahalanobis-weighted operator-valued kernel.

2.1. Architecture

RSA2C consists of three components instantiated in Reproducing Kernel Hilbert Spaces (RKHS), controlled by sparse dictionaries maintained via the Approximate Linear Dependence (ALD) method:

• Advanced Actor: Instantiated in a vector-valued RKHS using an Operator-Valued Kernel (OVK). The kernel is a Mahalanobis-weighted Gaussian kernel, $K(s, s_j) = \kappa_\phi(s, s_j)\Sigma_K$, where the weight matrix $W = \text{diag}(\tilde{\phi})$ is dynamically adjusted based on state feature attributions ($\phi_i$). This allows the policy to adaptively reweight state dimensions based on their importance.
• Value Critic: Resides in a scalar RKHS with its own dictionary. It estimates the value function $V(s)$ using Temporal Difference (TD) learning. Crucially, the Value Critic is not reweighted by SHAP scores to prevent circular dependencies and ensure a stable baseline for attribution computation.
• Advantage Critic: Resides in a scalar RKHS sharing the Actor's dictionary. It uses "compatible features" (aligned with $\nabla \log \pi$) to estimate the advantage function $A(s, a)$, ensuring low-variance policy gradients.

2.2. RKHS-SHAP Integration

Instead of using Monte Carlo sampling (which is computationally expensive), RSA2C computes SHAP values analytically using Kernel Mean Embeddings (KME) and Conditional Mean Embeddings (CME) within the RKHS framework:

• On-manifold (Observational): Uses KME to average over the empirical state distribution, ensuring attributions reflect environment-consistent configurations.
• Off-manifold (Interventional): Uses CME to condition on subsets of dimensions while respecting data geometry, probing feature coalitions without generating invalid states.

These computed SHAP attributions are converted into Mahalanobis-gated weights. These weights modulate the Actor's gradient updates and the Advantage Critic's targets, effectively injecting dimension-level importance signals directly into the optimization dynamics.

2.3. Two-Timescale Mechanism

RSA2C operates on two timescales:

• Fast timescale: The Value Critic updates rapidly to provide stable value estimates and SHAP attributions.
• Slow timescale: The Actor and Advantage Critic update more slowly, utilizing the stable attributions from the Critic to guide policy improvements.

3. Key Contributions

1. Kernelized Two-Timescale Actor-Critic: The introduction of RSA2C, which embeds an Actor and two Critics in RKHS. It uniquely injects adaptive state feature importance via RKHS-SHAP into a Mahalanobis-weighted OVK, modeling vector-valued policies while encoding correlations among state dimensions. The use of sparse dictionaries ensures linear computational complexity relative to dictionary size, making it suitable for online RL.
2. State Attribution to Learning Signal: The method computes SHAP from the Value Critic via two routes (KME for on-manifold and CME for off-manifold expectations). These signals are gated and injected into the Actor and Advantage Critic, achieving both efficiency and intrinsic interpretability without relying solely on post-hoc explanation.
3. Global Non-Asymptotic Convergence: The authors derive a global convergence bound for RSA2C under state perturbations. The learning gap is decomposed into a perturbation error (quantifying stability against noise) and a convergence error (quantifying efficiency). The perturbation error is further split into attribution-induced and policy-learning terms, demonstrating the algorithm's robustness.
4. Empirical Validation: Extensive simulations on three continuous-control environments (Pendulum-v1, BipedalWalker-v3, Ant-v5) demonstrate that RSA2C achieves competitive returns with significantly lower computational overhead (FLOPs) compared to deep RL baselines (SAC, PPO), while maintaining superior stability under state perturbations and providing interpretable feature attributions.

4. Results

• Efficiency: In continuous control tasks, RSA2C variants (KME and CME) converge to higher returns than ablations (Advanced AC without SHAP, Uniform SHAP) and traditional RKHS-AC. While deep RL methods like SAC may reach higher asymptotic returns in high-dimensional tasks (e.g., Ant-v5), RSA2C achieves comparable or better performance in lower-dimensional tasks with significantly fewer FLOPs and runtime, particularly on CPU.
• Stability: Under state perturbations (zero-mean noise with varying variance), RSA2C variants exhibit the most stable performance among RKHS-based methods. Their mean returns degrade only mildly as noise increases, whereas methods without learned attributions (Advanced AC, RKHS-AC) suffer significant performance drops. RSA2C also outperforms deep RL baselines in stability; for instance, PPO's performance degrades severely under high noise, while RSA2C remains robust.
• Interpretability: Visualization via beeswarm plots and heatmaps confirms that RSA2C captures meaningful feature contributions aligned with physical dynamics (e.g., distinguishing between angle, angular velocity, and torque in the Pendulum task). RSA2C-CME, which models conditional dependencies, produces more balanced and temporally adaptive attributions compared to RSA2C-KME.

5. Significance and Claims

The paper claims that RSA2C provides a principled mechanism to transform state feature importance into a Mahalanobis distance on an OVK-based policy, thereby enhancing sample-efficient learning with intrinsic interpretability.

• Theoretical Guarantee: The work establishes a global, non-asymptotic convergence guarantee under state perturbations, a significant contribution as non-asymptotic guarantees for two-timescale RKHS-enhanced AC were previously an open gap.
• Robustness: The method demonstrates that learned feature attributions are critical for stabilizing network updates against stochastic disturbances, outperforming fixed-weight or non-attribution baselines.
• Limitations: The authors modestly acknowledge that the scalability of RKHS-based representations is limited in high-dimensional continuous-control tasks (e.g., Ant-v5 with 105D states), where the representational capacity of linear RKHS with RBF kernels may become a bottleneck compared to deep neural networks. Additionally, memory footprints grow with task complexity and training horizons.

In summary, RSA2C bridges the gap between high-performance RL and explainability by integrating dimension-level SHAP attributions directly into the kernelized learning loop, offering a stable, efficient, and theoretically grounded approach for continuous control.