README.md

metadata

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 — also published inside the live Space at https://huggingface.co/spaces/Mohith202/transformer
Colab training notebook	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, trainer_history.png, reward_log.jsonl
Held-out eval (test-subject baselines)	outputs/eval/single_m/plots_test/baseline_comparison.png, r2_curves.png, baselines_test.csv

Try the live environment

# 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)

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:

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:

{"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

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:

pip install -e .

Optional extras:

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:

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

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:

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

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

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

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:

--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.

--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.

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)

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 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

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.

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:

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:

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

Pruning rules (in prepare_parcels.py)

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.