RL-BioAug: RL for Bio Data Augmentation

• RL-BioAug is a reinforcement learning framework that replaces static, heuristic methods with autonomous policies for data augmentation and gene selection in biological data analysis.
• It employs specialized architectures—using a Transformer for EEG and MLP-based agents for scRNA-seq—to optimize performance via reward-guided selection of augmentations and biomarkers.
• Empirical results show significant improvements in classification and clustering metrics, reducing manual tuning and enhancing scalability in bioinformatics applications.

RL-BioAug encompasses a class of reinforcement learning (RL)-driven frameworks for optimizing data augmentation or biomarker selection strategies in biological data analysis. Distinct instantiations have been proposed for electroencephalography (EEG) self-supervised representation learning and for label-free single-cell RNA-seq (scRNA-seq) biomarker identification. The common principle underlying RL-BioAug methods is the replacement of static, heuristic, or random strategies with autonomous, learned policies that maximize task-relevant objectives and leverage partial supervision or ensemble priors (Lee et al., 20 Jan 2026, Xiao et al., 2 Jan 2025).

1. Core Concepts and Motivation

RL-BioAug addresses the inadequacy of static or randomly chosen strategies in biological settings characterized by high dynamism or heterogeneity. In EEG analysis, non-stationary dynamics render random augmentation susceptible to information loss or distortion; for scRNA-seq, manual or heuristic gene selection embeds user bias and suboptimality. RL-BioAug frameworks recast augmentation or panel selection as a Markov Decision Process (MDP) or multi-agent RL problem, in which agents adaptively select transformation or selection actions guided by reward signals tightly coupled to downstream self-supervised, unsupervised, or biologically meaningful metrics. RL-BioAug frameworks typically require minimal label supervision—e.g., only 10% of EEG data labels for augmentation policy learning—thereby supporting label efficiency and scalability across bioinformatic domains (Lee et al., 20 Jan 2026, Xiao et al., 2 Jan 2025).

2. RL-BioAug for Self-Supervised EEG Representation Learning

In the RL-BioAug framework for EEG, a reinforcement learning agent autonomously selects strong augmentation operations to optimize self-supervised contrastive representation learning, specifically with a SimCLR-style InfoNCE objective. An encoder network $f_\theta$ generates continuous state embeddings $s_t = f_\theta(x_t) \in \mathbb{R}^d$ from each sample $x_t$, and the RL agent $\pi_\varphi(a_t\mid s_t)$ outputs a distribution over a fixed action space $\mathbb{A}$ of five strong augmentations: Time Masking, Time Permutation, Crop & Resize, Time Flip, and Time Warp. Weak augmentations (Gaussian jitter, amplitude scaling) anchor the contrastive views.

The reward $r_t$ is based on a Soft-KNN consistency score evaluating whether augmented views preserve class similarity, computed with 10% of labeled data as reference. Augmentation selection is formalized as a one-step non-causal MDP. The learning objective for the augmentation policy involves advantage-weighted policy gradients with entropy regularization:

$$\mathcal{L}_\mathrm{policy}(\varphi) = -\mathbb{E}_{(s,a)\sim\pi_\varphi} \left[\log\pi_\varphi(a\mid s)\,A\right] - \beta\gamma H(\pi_\varphi(\cdot\mid s))$$

with $A_t = r_t - b_t$, and $b_t$ as the batch baseline. The encoder is jointly trained with SimCLR's InfoNCE loss (Lee et al., 20 Jan 2026).

The policy agent employs a Transformer-based architecture that conditions on the current state and past $K$ action–reward pairs with self-attention, employing top-$s_t = f_\theta(x_t) \in \mathbb{R}^d$0 sampling at each step. Training proceeds in two phases: SSL+RL (agent and encoder updated jointly with 10% labeled data-guided rewards), followed by SSL-only (frozen agent, full data). Empirical results on Sleep-EDFX and CHB-MIT datasets indicate substantial gains over both random and best-single-augmentation baselines, with Macro-F1 improvement of 9.72 and 7.25 percentage points, respectively. Notably, the agent reliably discovers task-specific augmentation preferences—e.g., consistently selecting Time Masking ($s_t = f_\theta(x_t) \in \mathbb{R}^d$162%) for sleep-stage classification and Crop & Resize ($s_t = f_\theta(x_t) \in \mathbb{R}^d$277%) for seizure detection. Ablation analysis highlights the superiority of Soft-KNN rewards and the necessity of top-$s_t = f_\theta(x_t) \in \mathbb{R}^d$3 policy sampling (Lee et al., 20 Jan 2026).

3. RL-BioAug for Knowledge-Guided Biomarker Identification

In single-cell transcriptomics, RL-BioAug (termed RiGPS) is applied as a multi-agent RL paradigm for label-free gene panel selection. The framework consists of two stages:

• Ensemble Knowledge Guidance: A suite of classical gene selection methods (e.g., HVG, Random Forest, SVM, geneBasis, KBest) produces initial scoring and candidate gene sets. Meta-voting weights methods by downstream clustering performance, forming a refined candidate pool $s_t = f_\theta(x_t) \in \mathbb{R}^d$4—genes with meta-votes more than two standard deviations above the mean.
• Multi-Agent RL Optimization: Each candidate gene in $s_t = f_\theta(x_t) \in \mathbb{R}^d$5 is controlled by a dedicated agent, which selects at each step whether to retain or discard its gene. The unified reward $s_t = f_\theta(x_t) \in \mathbb{R}^d$6 combines (a) unsupervised cluster separability (normalized mutual information on clustering with currently selected panel), and (b) panel compactness (favoring minimal size), with the trade-off parameter $s_t = f_\theta(x_t) \in \mathbb{R}^d$7. The agent state is a summary embedding (flattened descriptive statistics over $s_t = f_\theta(x_t) \in \mathbb{R}^d$8, encoded via a bottleneck autoencoder). Policy and value functions are updated using actor-critic with prioritized experience replay. Prior gene selections are injected into agent replay buffers as initial expert experiences.

This iterative procedure refines gene panels over multiple episodes, and achieves superior clustering (highest average NMI/ARI/SI on 21/24 scRNA-seq datasets tested), more compact panels (often 20–50% smaller than the next-best), and improved resistance to batch effects. Visualization analyses (volcano plots, t-SNE/UMAP, cluster-wise heatmaps) confirm both quantitative and biological interpretability advantages. Ablations show that removing RL optimization, knowledge injection, or pre-filtering consistently degrades panel quality (Xiao et al., 2 Jan 2025).

4. Architectural and Training Design

Domain	Policy Network	State	Action Space	Reward	Optimization
EEG	Transformer	Encoder output	5 strong augs	Soft-KNN consistency (SSL)	REINFORCE++
scRNA-seq	MLP (per gene)	AE bottleneck	Select/Discard	Cluster NMI & panel size	Actor-critic multi-agent

The EEG policy leverages a Transformer over the current state and action–reward log, allowing task-specific adaptation for non-stationary temporal data. The scRNA-seq approach, in contrast, manages high-dimensional discrete actions via a multi-agent architecture, with each gene represented as an agent and leveraging expert-initialized experience replay to accelerate convergence and reduce suboptimal exploration.

Both frameworks employ two-phase training schedules; in EEG, RL and self-supervised learning are co-trained, then the policy is frozen. In scRNA-seq, RL is iterative until both reward and panel size stabilize.

5. Empirical Performance and Comparative Results

On benchmark EEG datasets, RL-BioAug delivers substantial gains in discriminative performance: on Sleep-EDFX (5-class sleep stage classification), Macro-F1 improves from 61.83% (best single augmentation) and 59.86% (random composite) to 69.55%. On CHB-MIT (seizure detection), Macro-F1 increases from 64.25% and 62.70% to 71.50%. Policy analysis reveals interpretable augmentation specialization per downstream task.

In scRNA-seq, RL-BioAug achieves top clustering and supervised annotation metrics over 24 datasets (e.g., NMI, ARI, SI), selecting smaller gene panels than all tested baselines without loss—and often with improvement—of separability and biological interpretability. Functional markers chosen by RL-BioAug coincide with literature-validated cell-type markers (e.g., INS, S100A8/A9). On batch-effect-sensitive tasks, RL-BioAug preserves clean separation of biologically meaningful clusters where competing methods fail (Lee et al., 20 Jan 2026, Xiao et al., 2 Jan 2025).

6. Limitations and Prospective Developments

Reported limitations for RL-BioAug include the constraint of discrete augmentation action spaces (EEG), computational cost of reinforcement learning, and reliance on partial supervision (for labeled rewards or expert priors). Extensions under consideration include continuous-action policy architectures, generative augmentation operators, and broader application to alternative biosignal modalities such as ECG and EMG for EEG, or expansion to multi-omic biomarker selection for scRNA-seq.

A plausible implication is that RL-BioAug’s agent-based approach, highly adaptive to the local structure of biological signals and data distributions, may generalize across a range of biomedical data augmentation and selection tasks, especially where full supervision and stable signal properties are unavailable.

7. Significance and Future Impact

RL-BioAug represents a paradigm shift from fixed or empirically chosen augmentation/selection pipelines to a fully data-driven, autonomous optimization regime. In both self-supervised representation learning (e.g., EEG) and label-free biomarker identification (scRNA-seq), RL-BioAug demonstrates marked improvements in both performance and efficiency, while systematically reducing reliance on exhaustive manual parameter or heuristic design. Its empirical successes suggest applicability to a broad spectrum of biomedical domains where robust, label-efficient adaptation to non-stationary or heterogeneous data is essential (Lee et al., 20 Jan 2026, Xiao et al., 2 Jan 2025).