--- 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 | | | HF Dataset bundle (configs + stimulus annotations) | | | Trained GRPO checkpoint (Qwen 2.5 0.5B) | | | Mini-blog (long-form writeup) | [`blog.md`](blog.md) — also published inside the live Space at | | 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/annotation/_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 . ```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.