🧠 Closed-Loop Neuroadaptive Stimulation

A reinforcement learning system that learns when and how to stimulate a simulated pathological brain to suppress abnormal neural oscillations.

Inspired by next-generation BCIs for Parkinson's, epilepsy, and depression therapy.

---

Overview

This project builds a closed-loop neurostimulation controller using deep reinforcement learning. The system:

1. Observes simulated neural activity (Kuramoto oscillator network)
2. Decides a stimulation action (none / low / high intensity)
3. Receives reward for suppressing pathological beta-band oscillations
4. Closes the loop — stimulation changes brain dynamics, which changes future observations

The key distinction from passive decoding projects: the agent actively intervenes and learns adaptive timing, not just classification.

---

Architecture

┌────────────────────────────────────────────┐
│             BrainEnv (Gymnasium)           │
│                                            │
│  Kuramoto oscillators (N=50)               │
│       ↓                                    │
│  Population signal → FFT → PSD obs        │
│       ↑                                    │
│  Stimulation term applied to phases        │
└────────────────────────────────────────────┘
           ↕  (obs, reward, done)
┌────────────────────────────────────────────┐
│              PPO Agent (SB3)               │
│   MLP policy [256, 256]                    │
│   Action: {0=none, 1=low, 2=high stim}     │
└────────────────────────────────────────────┘

---

Project Structure

neuroadaptive-stim/
├── env/
│   └── brain_env.py          # Gymnasium environment (Kuramoto model)
├── agents/
│   ├── ppo_agent.py          # PPO agent (stable-baselines3 wrapper)
│   └── baselines.py          # NoStim, FixedStim, ThresholdAgent
├── models/
│   └── decoder.py            # Optional LSTM brain state decoder
├── utils/
│   ├── viz.py                # Visualization functions
│   └── metrics.py            # Evaluation metrics + comparison table
├── experiments/              # Experiment result analysis notebooks/scripts
├── scripts/
│   └── run_experiment.py     # Main training + evaluation script
├── tests/
│   └── test_env.py           # Sanity checks for the environment
├── configs/
│   ├── default.yaml          # Default hyperparameters
│   └── high_disorder.yaml    # Ablation: severe disorder
├── results/
│   ├── plots/                # Generated figures
│   └── checkpoints/          # Saved model checkpoints
└── requirements.txt

---

Quickstart

# 1. Install dependencies
pip install -r requirements.txt

# 2. Run environment tests
python tests/test_env.py

# 3. Train PPO agent + evaluate all agents
python scripts/run_experiment.py

# 4. Eval only (load saved model)
python scripts/run_experiment.py --eval-only --model-path results/ppo_model

# 5. Run with custom config (e.g., high disorder ablation)
python scripts/run_experiment.py --config configs/high_disorder.yaml

---

The Brain Model

The simulation uses a Kuramoto oscillator network — the standard model for synchrony-based neurological disorders:

dθᵢ/dt = ωᵢ + (K/N) Σⱼ sin(θⱼ - θᵢ) − stim · sin(θᵢ) + noise

• ωᵢ: Natural frequency of oscillator i (drawn from a Lorentzian centered at 20 Hz)
• K: Coupling strength — high K → pathological synchrony (Parkinson's-like)
• Stimulation term stim · sin(θᵢ): desynchronizing input (mimics DBS)
• Order parameter r = |⟨e^{iθ}⟩|: r≈1 = full synchrony, r≈0 = desynchronized (healthy)

---

Reward Function

r = (P_baseline − P_current) / P_baseline  −  λ · stim_intensity

The agent is rewarded for suppressing beta power but penalized for unnecessary stimulation (λ=0.1). This encourages the agent to learn energy-efficient, adaptive stimulation timing.

---

Agents Compared

Agent	Strategy
NoStimAgent	Never stimulates (pathological baseline)
FixedStimAgent	Always applies the same stimulation
ThresholdAgent	Rule-based: stimulate if beta power > threshold
PPO	Learned adaptive policy via deep RL ← our model

---

Key Results (What to Expect)

• PPO learns to stimulate selectively — only when beta power is elevated
• PPO achieves higher beta suppression than fixed/threshold agents
• Policy heatmap reveals a learned threshold structure in action selection
• Adaptive stimulation uses ~30-40% less total stimulation than fixed strategies

---

Extensions (Grad School Portfolio Ideas)

• Add an LSTM decoder and use latent state as RL observation
• Multi-frequency stimulation (alpha + beta bands)
• Transfer learning: train on one disorder severity, test on another
• Safety constraints: penalize overstimulation (safe RL)
• Real EEG data from PhysioNet as pre-training signal

---

References

• Tass et al. (1998). Detection of n:m Phase Locking. PRL
• Little et al. (2013). Adaptive deep brain stimulation. Annals of Neurology
• Schulman et al. (2017). Proximal Policy Optimization. arXiv:1707.06347
• Kuramoto (1984). Chemical Oscillations, Waves, and Turbulence