Less is More: Clustered Cross-Covariance Control for Offline RL

TL;DR

This paper introduces Clustered Cross-Covariance Control for TD (C^4), a novel offline RL method that mitigates harmful cross-covariance effects in OOD data, leading to more stable learning and up to 30% higher returns.

Contribution

It proposes buffer partitioning and gradient penalties to reduce covariance bias, improving offline RL performance especially with limited or OOD-dominated datasets.

Findings

Achieves up to 30% improvement in returns over prior methods.

Enhances stability in offline RL with small datasets.

Effectively mitigates distributional shift issues.

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.