---
title: RL-BioAug Framework for EEG Augmentation
url: https://www.emergentmind.com/topics/rl-bioaug-framework
type: topic
---

# RL-BioAug Framework for EEG Augmentation

RL-BioAug is a label-efficient reinforcement learning (RL) framework that establishes an autonomous paradigm for data augmentation in self-supervised electroencephalography (EEG) representation learning. RL-BioAug employs an RL agent, guided by only a minimal fraction (10%) of labeled data, to determine context-appropriate augmentation strategies for each data sample. Its primary aim is to improve representation quality for downstream tasks under the inherent non-stationarity of EEG signals, outperforming heuristic or random composite augmentation techniques [2601.13964].

## 1. System Architecture

RL-BioAug integrates three major modules: a self-supervised encoder, a data-augmentation module, and a transformer-based RL agent. The encoder $f_\theta$ is implemented as a 1D ResNet-18 mapping EEG segment $x$ to an embedding $z = f_\theta(x)$. The data-augmentation module generates a "weak" view ($x_{\text{weak}}$) using fixed small jitter and scaling, and a "strong" view ($x_{\text{strong}}$) by applying augmentations chosen by the RL agent. The RL agent $\pi_\phi$ uses both the current embedding (state $s_t$) and a history of past actions and rewards to generate a probability distribution across augmentation actions, sampling one per instance.

The training loop, executed per batch and sample, consists of:
1. Computing $s_t$ using the frozen encoder $f_\theta$.
2. Having $\pi_\phi$ select $a_t$ (augmentation) given $(s_t, a_{t-K:t-1}, r_{t-K:t-1})$.
3. Augmenting $x$ into $x_{\text{strong}}$ (using $a_t$) and $x_{\text{weak}}$ (using fixed jitter+scaling).
4. Updating $f_\theta$ via InfoNCE loss between the two views.
5. Computing RL agent reward $r_t$ (using Soft-KNN metric and 10% labeled reference set), then updating $\phi$ via a policy-gradient step with entropy regularization.

## 2. Reinforcement-Learning Components

The RL agent operates with:
- **State space**: $s_t = f_\theta(x_t) \in \mathbb{R}^d$, representing the encoder's output for each EEG segment.
- **Action space**: Discrete, $A = \{1,2,3,4,5\}$, each corresponding to a distinct "strong" augmentation:
    1. Time Masking: random interval zeroed out, length $L \sim \text{Uniform}(L_\text{min}, L_\text{max})$.
    2. Time Permutation: segment sequence shuffled.
    3. Crop & Resize: contiguous window cropped (size $\alpha T$), resized back to length $T$.
    4. Time Flip: reverse time axis.
    5. Time Warp: random subinterval sped up or slowed down by $v \sim \text{Uniform}(v_\text{min}, v_\text{max})$.
- **Reward**: Soft-KNN consistency score calculated on a labeled reference set $R$, comparing whether the embedding $z_i$ of the current sample consistently shares labels with its $K$ nearest neighbors:
  $$
  r_t = P(y_i|z_i) = \frac{\sum_{j\in\mathcal{N}_i} \exp(\text{sim}(z_i,z_j)/\tau) \cdot \mathbb{I}[y_j = y_i]}{\sum_{k\in\mathcal{N}_i} \exp(\text{sim}(z_i,z_k)/\tau)}
  $$
  with $\mathcal{N}_i$ denoting $K$ nearest neighbors, $\text{sim}$ as cosine similarity, and $r_t \in [0,1]$.

## 3. Policy Network Details

The RL agent’s policy network is a transformer, receiving:
- The current sample’s state embedding $s_t$.
- Embedded representations of the last $K$ actions and rewards.
- All inputs fused via fully connected layers to a shared dimension.
- Learned positional encodings and $L$ self-attention blocks model the temporal policy-reward dynamics.
- Final logits over action space $A$ produced by an output head followed by softmax; actions sampled via Top-K strategy.

Policy gradients use the REINFORCE++ update:
$$
A_t = r_t - b, \quad b = \frac{1}{|B|}\sum_{i \in \text{batch}} r_i
$$
$$
L_{\text{policy}} = -\mathbb{E}[ \log \pi_\phi(a_t|s_t) \cdot A_t ] - \beta \gamma H(\pi_\phi)
$$
where entropy $H(\pi_\phi)$ fosters exploration, and $\phi$ is updated via gradient descent.

## 4. Contrastive Learning Objective

RL-BioAug employs SimCLR's InfoNCE objective:
$$
L_{\text{contrastive}} = -\mathbb{E}_{i \in B} \log \left( \frac{\exp(\text{sim}(z_i^{\text{weak}}, z_i^{\text{strong}})/\tau)}{\sum_{j=1}^{2|B|} \exp(\text{sim}(z_i^{\text{weak}}, z_j)/\tau)} \right)
$$
Each sample produces two augmented views:
- Weak augmentation: fixed jitter + scaling.
- Strong augmentation: transformation selected by agent $\pi_\phi$.

## 5. Training Strategy and Label Utilization

Training proceeds in two phases:
**Phase 1 (agent pre-training):**
- 10% of data with labels forms both the reference set $R$ and the unlabeled pool.
- RL agent $\pi_\phi$ learns augmentation policy using only reward signals, without propagating label supervision into the encoder.

**Phase 2 (self-supervised encoder training):**
- $\pi_\phi$ frozen.
- For each sample, agent suggests strongest augmentation; encoder $f_\theta$ trained using contrastive loss over all (unlabeled) data.

Pseudocode for the loop is:

```python
Initialize encoder θ, policy φ
Split 10% labeled → Reference set R; rest unlabeled for sampling
for epoch in 1…N1 (Phase 1):
  for each batch of unlabeled x:
    for each sample x_i in batch:
      s_t = f_θ(frozen)(x_i)
      a_t ∼ TopK_sample( π_φ(·|s_t) )
      x_strong = Aug_strong(x_i; a_t)
      x_weak   = Aug_weak(x_i)
      # self-supervised update
      z_w = f_θ(x_weak);  z_s = f_θ(x_strong)
      L_c = InfoNCE({z_w,z_s})
      θ ← θ – η_enc ∇_θ L_c
      # compute reward on z_w (or z_s)
      r_t = SoftKNN(z_w; R)
      collect (s_t,a_t,r_t)
    # after batch:
    compute A_t = r_t – mean(r) over batch
    L_p = –∑ log π_φ(a_t|s_t)·A_t  – βγ H(π_φ)
    φ ← φ – η_agent ∇_φ L_p
Freeze π_φ
for epoch in 1…N2 (Phase 2):
  for each batch x:
    for x_i in batch:
      a_i = argmax π_φ(·| f_θ(frozen)(x_i) )
      x_strong = Aug_strong(x_i; a_i)
      x_weak   = Aug_weak(x_i)
    compute InfoNCE, update θ only
```
This approach minimizes the need for expert intervention and labels, leveraging context-aware augmentation strategies to maximize discrimination in self-supervised spaces.

## 6. Experimental Validation

RL-BioAug was evaluated on Sleep-EDFX (5-class sleep staging) and CHB-MIT (binary seizure detection) datasets. Results reveal:

| Dataset    | Random Composite (MF1) | RL-BioAug (MF1) | Absolute Gain |
|------------|------------------------|-----------------|--------------|
| Sleep-EDFX | ~59.86%                | 69.55%          | +9.69%       |
| CHB-MIT    | ~62.70%                | 71.50%          | +8.80%       |

Analysis of learned augmentation strategies:
- On Sleep-EDFX, Time Masking was selected with ~62% probability.
- On CHB-MIT, Crop & Resize dominated at ~77% probability.

## 7. Context, Rationales, and Implications

EEG’s pronounced non-stationarity necessitates adaptive data augmentation: different neural states (e.g., REM, seizure) exhibit variable tolerance to distortions; fixed or random policies risk under- or over-distortion. RL-BioAug frames augmentation selection as a reinforcement learning problem, enabling per-sample policy adaptation that maximizes downstream representation quality (as measured by Soft-KNN consistency). Entropy regularization and Top-K sampling balance exploration and exploitation in policy selection.

Label-efficient RL guidance (only 10% of labels needed) allows the self-supervised encoder to function without direct label supervision, suggesting wider applicability to unlabeled biomedical contexts. The framework’s capability to supplant heuristic-based augmentation approaches and its demonstrated improvements in Macro-F1 scores indicate its significance for autonomous EEG data augmentation [2601.13964]. *A plausible implication is that similar RL-driven augmentation policies may benefit other non-stationary time-series domains with limited labeled data.*

Source: https://www.emergentmind.com/topics/rl-bioaug-framework