README.md

---
title: BrainRL Region Selection Environment
emoji: 🧠
colorFrom: indigo
colorTo: green
sdk: docker
pinned: false
app_port: 8000
base_path: /web
tags:
  - openenv-0.2.3
  - openenv
  - reinforcement-learning
  - fmri
---

# BrainRL Region Selection Environment

BrainRL frames Le Petit Prince fMRI analysis as an active sensing problem: an
RL agent sequentially selects a small set of brain regions/parcels to maximize
auditory stimulus-response prediction under a fixed budget.

The defensible claim is intentionally narrow:

> We test whether semantic priors can guide efficient sequential brain region
> selection for auditory fMRI prediction, compared with non-RL baselines.

The hackathon target action space is **~200 atlas parcels** (Schaefer-2018,
7 networks). A 7-ROI easy-curriculum mode is kept for smoke tests. There is no
Gym/PPO path: the project is OpenEnv-first and trained with TRL/GRPO.

## Submission materials

All artefacts referenced by the hackathon judging criteria:

| Artefact | Link |
| --- | --- |
| OpenEnv environment on Hugging Face Spaces | <https://huggingface.co/spaces/Mohith202/transformer> |
| HF Dataset bundle (configs + stimulus annotations) | <https://huggingface.co/datasets/Mohith202/brainrl-config-data> |
| Trained GRPO checkpoint (Qwen 2.5 0.5B) | <https://huggingface.co/Mohith202/brainrl-grpo-single-m> |
| Mini-blog (long-form writeup) | [`blog.md`](blog.md) — also published inside the live Space at <https://huggingface.co/spaces/Mohith202/transformer> |
| Colab training notebook | [`BrainRL_HF_Pipeline.ipynb`](BrainRL_HF_Pipeline.ipynb) — login → deploy Space → submit HF Job → fetch plots → promote checkpoint |
| Training evidence (loss + reward plots) | [`outputs/grpo_single_m/plots/training_curve.png`](outputs/grpo_single_m/plots/training_curve.png), [`trainer_history.png`](outputs/grpo_single_m/plots/trainer_history.png), [`reward_log.jsonl`](outputs/grpo_single_m/plots/reward_log.jsonl) |
| Held-out eval (test-subject baselines) | [`outputs/eval/single_m/plots_test/baseline_comparison.png`](outputs/eval/single_m/plots_test/baseline_comparison.png), [`r2_curves.png`](outputs/eval/single_m/plots_test/r2_curves.png), [`baselines_test.csv`](outputs/eval/single_m/baselines_test.csv) |

### Try the live environment

```bash
# Hit the deployed Space directly (no local install needed)
export BRAINRL_API_URL="https://mohith202-transformer.hf.space"
curl "$BRAINRL_API_URL/health"
curl -X POST "$BRAINRL_API_URL/rollout" \
  -H 'Content-Type: application/json' \
  -d '{"seed": 42, "subject_id": "sub-21", "run_id": "run-1", "condition": "single_m"}'

# Or open the Gradio UI in a browser:
#   https://huggingface.co/spaces/Mohith202/transformer
# The "Model predictions" tab serves the trained Qwen 2.5 0.5B GRPO checkpoint
# (Mohith202/brainrl-grpo-single-m).
```

> Per hackathon rules large media and model weights live on Hugging Face Hub
> (Space / dataset / model URLs above), not inside this GitHub repo.

## Architecture (CLI-orchestrated)

```text
prepare_parcels.py  ──>  configs/parcel_candidates.json  ──>  OpenEnv  ──>  TRL/GRPO
   (atlas + pruning)        (frozen manifest, ~200 rows)      (verifier)    (LLM RL)

participant_run_info.json  ──>  data_split.py  ──>  train/test (subject, run, condition) pairs
                                                       │
                                                       ▼
                                                evaluate.py / train_grpo.py / inference.py
                                                       │
                                                       ▼
                                                outputs/.../plots/*.png + *.csv

deploy_to_hf.py  ──>  Hugging Face Space (Docker SDK, server.app:app on :8000)
```

* `prepare_parcels.py` is the *only* place atlas parcels are loaded, scored,
  and pruned. It runs once before training.
* `server/brain_environment.py` consumes the frozen manifest, never touches
  NIfTI files at episode time, and exposes the OpenEnv MCP tools.
* `prompts.py` only ever shows the compact `prompt_top_k` shortlist plus a
  per-network summary, so prompts stay small even with 200 candidates.
* `train_grpo.py` runs TRL/GRPO against the same verifier.
* A `Makefile` ties the steps together; nothing in this stack depends on
  Gymnasium or stable-baselines3.

## OpenEnv task

At each step the agent observes:

- the selected-region list,
- the remaining selection budget,
- the current cumulative R2 estimate,
- a per-network summary of unselected candidates,
- the top-K unselected candidates with semantic prior + base-R2 hint,
- a `stimulus` block describing the words the subject is currently hearing
  (a window of ~30 words from the per-condition annotation CSV) including the
  dominant POS, POS mix, mean log-frequency, speech density and the first few
  words. The reward applies a small group-level bias from this window
  (auditory_temporal benefits from acoustic density / nouns; inferior_frontal
  benefits from function words / verbs; association handles mixed content),
  giving GRPO a learnable space ↔ word signal across episodes.

The action is one parcel/region id. The OpenEnv environment is the source of
truth for state transitions, validation, and scoring. The scalar reward is
backed by independent verifier components:

```text
total_reward =
  delta_r2_reward
  - cost_penalty
  + valid_action_reward
  + duplicate_penalty
  + budget_penalty
  + success_bonus
```

The action format is one JSON object per step:

```json
{"region_id": "parcel_187"}
```

In ROI-easy mode it is the legacy seven-ROI form, e.g. `{"region_id": "pSTS"}`.

## Baselines

`evaluate.py` compares the OpenEnv-shaped policies:

- `random` – random region order
- `semantic_prior` – static prior order
- `prompt_static` – prompt-compatible deterministic policy
- `greedy` – one-step improvement maximization
- `llm_prompt` – optional remote prompt-based LLM policy

## Project structure

```text
BrainRL_hackathon/
├── README.md
├── pyproject.toml
├── requirements.txt
├── Dockerfile
├── openenv.yaml
├── Makefile
├── prepare_parcels.py        # parcel pruning orchestrator
├── data_loader.py            # reads parcel_candidates.json or region_priors.json
├── data_split.py             # condition + subject train/test split helpers
├── prompts.py                # compressed top-K prompt + parser
├── rewards.py                # independent verifier components
├── metrics.py                # R²-adjacent correlation + 2v2 eval metrics
├── baselines.py              # random/semantic/prompt/greedy
├── plotting.py               # matplotlib helpers (eval + training plots)
├── evaluate.py               # CLI: compare baselines, write CSV + plots
├── inference.py              # CLI: prompt-policy rollout (subject-aware)
├── train_grpo.py             # CLI: TRL/GRPO trainer with reward logging
├── deploy_to_hf.py           # CLI: deploy server to Hugging Face Spaces
├── client.py                 # OpenEnv MCP client wrapper
├── models.py                 # action/observation schemas
├── configs/
│   ├── subset_config.yaml
│   ├── participant_run_info.json
│   ├── parcel_candidates.json   # generated by prepare_parcels.py
│   └── region_priors.json
├── BrainRL_HF_Pipeline.ipynb # Colab launcher: deploy Space, run HF Job, fetch plots
├── blog.md                    # long-form writeup (project mini-blog)
├── hf_data.py                 # CLI: build/upload the HF Dataset config bundle
├── hf_jobs.py                 # CLI: submit BrainRL training to HF Jobs (GPU)
└── server/
    ├── __init__.py
    ├── app.py
    └── brain_environment.py
```

## Quick start

Install locally:

```bash
pip install -e .
```

Optional extras:

```bash
pip install -e ".[atlas]"     # nilearn + nibabel for real Schaefer-200 fetch
pip install -e ".[plots]"     # matplotlib for evaluation/training plots
pip install -e ".[deploy]"    # huggingface_hub for HF Spaces deploy
pip install -e ".[llm-rl]"    # trl, datasets, transformers, accelerate, peft
```

### 1. Build the parcel candidate manifest

The orchestrator tries `nilearn` first and falls back to a deterministic
synthetic Schaefer-style atlas if it isn't installed/online:

```bash
python prepare_parcels.py \
  --atlas schaefer200 \
  --max-candidates 200 \
  --min-voxels 20 \
  --prompt-top-k 30 \
  --selection-budget 20 \
  --output configs/parcel_candidates.json
```

This writes a frozen manifest with per-parcel metadata, weighted prune
scores, and a kept-by-floor flag for clearly language-relevant parcels.

### 2. Verify the OpenEnv pipeline end-to-end

```bash
python train_grpo.py --dry-run         # one prompt + verifier round-trip
python evaluate.py --episodes 8        # baseline comparison
python inference.py                    # static prompt policy rollout
```

Or use the Makefile shortcuts:

```bash
make prepare       # writes configs/parcel_candidates.json
make dryrun        # GRPO verifier dry-run
make eval          # run baselines
make inference     # prompt-policy rollout
```

### 3. Optional: prompt-based LLM rollouts

```bash
export HF_TOKEN="your_token"
export MODEL_NAME="Qwen/Qwen2.5-72B-Instruct"
python inference.py --use-llm --show-prompts
python evaluate.py --episodes 8 --use-llm
```

### 4. Optional: tiny TRL/GRPO smoke run

```bash
python train_grpo.py \
  --model-name Qwen/Qwen2.5-0.5B-Instruct \
  --max-steps 50 \
  --output-dir outputs/grpo_brainrl
```

Install Unsloth separately if your GPU/CUDA supports it, then adapt the
model-loading section in `train_grpo.py`.

### 5. Optional: serve the environment

```bash
uvicorn server.app:app --host 0.0.0.0 --port 8000
# or
make server
```

## Single-condition training with subject splits

For a defensible generalization claim, train on one condition and split
subjects so the test subjects never appear during training. The dataset
ships with `derivatives/participant_run_info.json` mapping each subject's
four runs to `single_m | single_f | mixed_m | mixed_f`.

`configs/participant_run_info.json` provides a portable copy of the
subject/run mapping for Colab and HF Spaces. `data_split.py` filters that
mapping by condition and parses ranges like
`sub-01:sub-20`. Every CLI accepts the same flags:

```text
--condition         single_m
--participant-info  configs/participant_run_info.json
--train-subjects    sub-01:sub-20
--test-subjects     sub-21:sub-26
--exclude-subjects  sub-03,sub-18         # corrupted recordings
--split             {train, test, all}
--plot-dir          outputs/.../plots
```

The OpenEnv environment uses the (subject, run, condition) tuple to apply a
deterministic per-subject reward perturbation and a per-condition boost on
the relevant redundancy groups, so training subjects and held-out subjects
genuinely have different reward landscapes.

### Skipping corrupted subjects

`sub-03` and `sub-18` in this dataset have unusable recordings, so all
subject-aware CLIs accept `--exclude-subjects`. Excluded subjects are
removed before the train/test partition, so they cannot leak into either
split. The Makefile defaults `EXCLUDE_SUBJECTS=sub-03,sub-18` and threads
it through `dryrun-single`, `eval-single-train`, `eval-single-test`,
`train-single`, and `inference-single`.

```text
--exclude-subjects sub-03,sub-18           # comma list
--exclude-subjects sub-03:sub-05           # inclusive range (no-op if missing)
EXCLUDE_SUBJECTS=none make eval-single-test  # turn the default off
```

After exclusion the `single_m` condition has 18 train-eligible subjects
(sub-01..sub-20 minus sub-03 and sub-18) and 6 test subjects
(sub-21..sub-26).

### Stimulus context (word/audio annotations)

`stimulus_loader.py` reads the per-condition word annotation CSV that ships
with the dataset
(`<data_root>/data/annotation/<condition>_word_information.csv`) and slices it
into ~30-word "stimulus windows". Every episode is bound to one window, and
that window is passed to the LLM (top words + POS mix + density) and used to
bias the reward by parcel group:

* `auditory_temporal` ↑ for noun-heavy / acoustically dense windows,
* `inferior_frontal`  ↑ for verb / function-word heavy windows,
* `association`       ↑ for mixed content windows.

This is what gives GRPO a signal to learn a space ↔ word mapping.

```text
BRAINRL_STIMULUS_DIR         /path/to/data/annotation     # default: ../data/annotation
BRAINRL_STIMULUS_WINDOW_SIZE 30                           # words per window
BRAINRL_STIMULUS_TOP_WORDS   10                           # words shown in prompt
```

These map to Makefile vars `STIMULUS_DIR`, `STIMULUS_WINDOW_SIZE`,
`STIMULUS_TOP_WORDS` and are threaded into all subject-aware targets. Mixed
conditions fall back to the male single-track annotation; if the CSV is
missing the environment falls back to no-stimulus mode automatically.

### Single-condition end-to-end (`single_m`)

```bash
make prepare ATLAS_SOURCE=auto MAX_CANDIDATES=200 SELECTION_BUDGET=20

make dryrun-single \
    CONDITION=single_m \
    TRAIN_SUBJECTS=sub-01:sub-20 \
    TEST_SUBJECTS=sub-21:sub-26 \
    EXCLUDE_SUBJECTS=sub-03,sub-18

make eval-single-train EPISODES=8     # baselines on training subjects + plots
make eval-single-test  EPISODES=8     # baselines on held-out subjects + plots

make train-single \
    CONDITION=single_m \
    TRAIN_SUBJECTS=sub-01:sub-20 \
    TEST_SUBJECTS=sub-21:sub-26 \
    EXCLUDE_SUBJECTS=sub-03,sub-18 \
    GRPO_MAX_STEPS=300 \
    GRPO_MODEL=Qwen/Qwen2.5-0.5B-Instruct
```

Plots and CSVs land in `outputs/eval/single_m/plots_test/baseline_comparison.png`,
`r2_curves.png`, `outputs/grpo_single_m/plots/training_curve.png`, and
`reward_log.jsonl`. The baseline CSV and comparison plot report final `R²`,
parcel-priority correlation, 2v2 ranking accuracy, and total reward for each
policy on the same train/test split used by GRPO.

### Colab notebook

[`BrainRL_HF_Pipeline.ipynb`](BrainRL_HF_Pipeline.ipynb) is the runnable
submission notebook. It is Colab-friendly and walks the whole HF round-trip
in one place:

1. Hugging Face login + token preflight.
2. Build and push the HF Dataset config bundle (`hf_data.py`).
3. Deploy the OpenEnv server to a HF Space (`deploy_to_hf.py`).
4. Submit a TRL/GRPO training run to HF Jobs (`hf_jobs.py`) — choose a 0.5B
   smoke run or a 7B full run, monitor logs from the notebook.
5. Pull back `training_curve.png`, `trainer_history.png`, and `reward_log.jsonl`.
6. Promote the trained checkpoint to the live Space and run an end-to-end
   `/rollout` health check.

This is the single notebook judges should re-run.

### Live prediction tab

When the Space points at a trained model (set `BRAINRL_POLICY_MODEL_REPO`),
the Space serves a Gradio UI at `/web` with two tabs:

- **OpenEnv playground** — the default action-by-action MCP UI.
- **Model predictions** — pick subject / run / condition / seed, run a 20-step
  rollout under the trained policy, and see the cumulative R² curve, per-step
  selections, network distribution, and (optionally) a glass-brain rendering.

Both surfaces share the same `/rollout` endpoint, so anything shown in the UI
is reproducible from an external client.

### One-shot CLI

```bash
make prepare ATLAS_SOURCE=auto MAX_CANDIDATES=200 SELECTION_BUDGET=20 \
  && python train_grpo.py \
       --condition single_m \
       --participant-info configs/participant_run_info.json \
       --train-subjects sub-01:sub-20 \
       --test-subjects sub-21:sub-26 \
       --exclude-subjects sub-03,sub-18 \
       --max-steps 300 \
       --output-dir outputs/grpo_single_m \
       --plot-dir outputs/grpo_single_m/plots \
       --model-name Qwen/Qwen2.5-0.5B-Instruct \
  && python evaluate.py \
       --condition single_m \
       --participant-info configs/participant_run_info.json \
       --train-subjects sub-01:sub-20 \
       --test-subjects sub-21:sub-26 \
       --exclude-subjects sub-03,sub-18 \
       --split test \
       --episodes 32 \
       --output-csv outputs/grpo_single_m/baselines_test.csv \
       --plot-dir outputs/grpo_single_m/plots
```

## Deploy the OpenEnv server to Hugging Face Spaces

The repo already has a Docker SDK Space layout (`Dockerfile`,
`openenv.yaml`, README frontmatter). `deploy_to_hf.py` automates the push.
The reference deployment for this submission lives at
<https://huggingface.co/spaces/Mohith202/transformer>.

```bash
huggingface-cli login                 # or export HF_TOKEN=...

# Preview what would be uploaded
make deploy-dry HF_REPO_ID=Mohith202/transformer

# Actually push (public Space, includes parcel manifest)
make deploy HF_REPO_ID=Mohith202/transformer HF_VISIBILITY=--public

# Or push as private:
python deploy_to_hf.py --repo-id Mohith202/transformer --include-manifest
```

After the Space finishes building, point clients/inference at it:

```bash
export BRAINRL_API_URL="https://mohith202-transformer.hf.space"
python inference.py --condition single_m --split test --episodes 1
```

Local TRL/GRPO can run against either the local environment (default) or a
deployed Space; the verifier API is the same.

## Curriculum

| Stage  | Mode             | Candidates       | Budget |
| ------ | ---------------- | ---------------- | ------ |
| Easy   | `roi_priors`     | 7 language ROIs  | 5      |
| Medium | `atlas_parcels`  | ~200 Schaefer    | 20     |
| Hard   | `atlas_parcels`  | top-N voxel/parcel hierarchy (future) | 20+ |

Switch curricula by editing `configs/subset_config.yaml`:

```yaml
candidate_mode: atlas_parcels   # or roi_priors
selection_budget: 20
prompt_top_k: 30
```

## Pruning rules (in `prepare_parcels.py`)

```text
prune_score =
  0.40 * variance/activity_proxy
+ 0.30 * semantic_prior
+ 0.20 * cached_encoding_score (or semantic-derived fallback)
+ 0.10 * voxel_count_quality
```

Parcels with `n_voxels < min_voxels_per_parcel` are dropped, and a
language-relevance floor force-keeps the top auditory/temporal and inferior
frontal parcels even if their score is borderline. The weights are CLI flags
so different prune profiles are reproducible.

## Data notes

`prepare_parcels.py` may pull the real Schaefer-2018 atlas via `nilearn` (the
file is cached locally, not committed). Optional `--cached-scores` CSVs from
prior fits can boost a parcel's encoding score. Nothing in the OpenEnv
environment touches NIfTI files, raw derivatives, or large arrays.

Keep these out of git:

- `derivatives/`
- `outputs/`
- `*.nii` / `*.nii.gz`
- `*.npy` / `*.npz`
- trained model adapters

## Demo pitch

BrainRL treats brain mapping like a constrained experiment. Given a budget of
only ~20 parcels out of ~200 Schaefer regions, an LLM agent learns which
brain areas to query, in what order, to predict auditory responses from
narrative stimuli. We compare reward-trained RL (TRL/GRPO) against random,
static semantic priors, and greedy baselines to test whether sequential
adaptation under semantic priors adds value.