Aggregated masked autoencoding for offline reinforcement learning

Abstract

Viewing offline reinforcement learning (RL) as a sequence modeling problem has emerged as a new research trend. Recent approaches leverage self-supervised learning to improve sequence representations, yet most rely on state sequences for pretraining, thereby disrupting the intrinsic state–action coupling, which complicates the distinction of trajectory bifurcations caused by action quality differences. Moreover, actions from stochastic policies in offline datasets may cause low-quality state transitions to be mistakenly identified as salient information, hindering representation learning and degrading policy performance. To mitigate these issues, we propose aggregated masked future prediction (AMFP), a self-supervised learning framework for offline RL. AMFP introduces a new pretext task that combines weighted aggregation and masked autoencoding through global fusion tokens to perform aggregated masked reconstruction. The weighted aggregation mechanism is to assign higher weights to samples that are semantically similar to the anchor in the representation space, enabling the model to emphasize reliable state transitions and suppress misleading transitions from stochastic or low-quality actions. Meanwhile, the global fusion tokens serve a dual purpose: they facilitate the integration of weighted aggregation and masked autoencoding, and, after encoding, function as compressed representations of the state trajectory and implicit action-state coupling. The encoded representations are then utilized as the latent contextual factor to guide policy learning and improve robustness. Experimental evaluation on D4RL benchmarks demonstrates the effectiveness of our method in improving policy learning.