RL-BioAug: Label-Efficient Reinforcement Learning for Self-Supervised EEG Representation Learning

TL;DR

RL-BioAug introduces a reinforcement learning-based framework that autonomously optimizes data augmentation policies for self-supervised EEG representation learning, significantly improving performance with minimal labeled data.

Contribution

It presents a novel RL-based approach to automatically select augmentation strategies, outperforming random methods and reducing reliance on heuristic augmentations in EEG tasks.

Findings

Achieves 9.69% and 8.80% improvements in Macro-F1 on two EEG datasets.

Uses only 10% labeled data to guide augmentation policy.

Demonstrates the potential to replace heuristic augmentation methods.

Abstract

The quality of data augmentation serves as a critical determinant for the performance of contrastive learning in EEG tasks. Although this paradigm is promising for utilizing unlabeled data, static or random augmentation strategies often fail to preserve intrinsic information due to the non-stationarity of EEG signals where statistical properties change over time. To address this, we propose RL-BioAug, a framework that leverages a label-efficient reinforcement learning (RL) agent to autonomously determine optimal augmentation policies. While utilizing only a minimal fraction (10%) of labeled data to guide the agent's policy, our method enables the encoder to learn robust representations in a strictly self-supervised manner. Experimental results demonstrate that RL-BioAug significantly outperforms the random selection strategy, achieving substantial improvements of 9.69% and 8.80% in Macro-F1 score on the Sleep-EDFX and CHB-MIT datasets, respectively. Notably, this agent mainly chose optimal strategies for each task--for example, Time Masking with a 62% probability for sleep stage classification and Crop & Resize with a 77% probability for seizure detection. Our framework suggests its potential to replace conventional heuristic-based augmentations and establish a new autonomous paradigm for data augmentation. The source code is available at https://github.com/dlcjfgmlnasa/RL-BioAug.