Less is More: Clustered Cross-Covariance Control for Offline RL
Nan Qiao, Sheng Yue, Shuning Wang, Yongheng Deng, Ju Ren
TL;DR
This work identifies harmful cross-covariance in the TD second moment as a primary instability source in offline RL under weak data coverage. It introduces Clustered Cross-Covariance Control for TD (C^4), an EM-style, gradient-space clustering framework that partitions the replay buffer into local clusters and performs single-cluster updates while applying a gradient-based penalty to bound within-cluster cross-covariance; it also proves that these modifications preserve a computable lower bound on the standard improvement targets. Theoretical analysis decomposes the TD variance into beneficial supervised-like terms and a harmful TD cross term, motivating the cluster-based and mixture-regularized objective. Empirically, C^4 yields up to about 30% improvements in returns on small data regimes across D4RL benchmarks, with modest computational overhead and strong plug-and-play compatibility with multiple offline RL backbones. Collectively, the approach offers a robust, scalable path to stabilizing offline RL under limited coverage and distributional shift.
Abstract
A fundamental challenge in offline reinforcement learning is distributional shift. Scarce data or datasets dominated by out-of-distribution (OOD) areas exacerbate this issue. Our theoretical analysis and experiments show that the standard squared error objective induces a harmful TD cross covariance. This effect amplifies in OOD areas, biasing optimization and degrading policy learning. To counteract this mechanism, we develop two complementary strategies: partitioned buffer sampling that restricts updates to localized replay partitions, attenuates irregular covariance effects, and aligns update directions, yielding a scheme that is easy to integrate with existing implementations, namely Clustered Cross-Covariance Control for TD (C^4). We also introduce an explicit gradient-based corrective penalty that cancels the covariance induced bias within each update. We prove that buffer partitioning preserves the lower bound property of the maximization objective, and that these constraints mitigate excessive conservatism in extreme OOD areas without altering the core behavior of policy constrained offline reinforcement learning. Empirically, our method showcases higher stability and up to 30% improvement in returns over prior methods, especially with small datasets and splits that emphasize OOD areas.