GRAFT — Guided Reinforcement Learning with Agentic Fused Teachers

GRAFT is a reinforcement learning–guided knowledge distillation framework for clinically significant prostate cancer (csPCa) detection on bi-parametric MRI. It trains a single lightweight 3D U-Net student model by adaptively distilling knowledge from multiple teacher models, using a policy-gradient RL agent that dynamically weights teacher contributions based on their predictive accuracy and mutual disagreement. The result is a compact, deployable detector that matches or surpasses full teacher ensembles while requiring only a single forward pass at inference — achieving >1,000× speedup over the Teacher Average.

This release contains the training and evaluation code for the 5-fold nnUNet ensemble configuration (foldGRAFT) used in the GRAFT paper. The teachers are five fold-specific nnUNet checkpoints; the student is a full-resolution 3D U-Net trained jointly with a small policy network that learns how to weight each teacher's intermediate features at every training step.

Dataset / image / teacher-feature paths in this release have been replaced with <PLACEHOLDER> strings. Before running anything, edit the launch script and the dataset YAML to point at your own data.

Paper: Li, C.X., et al. GRAFT: Guided Reinforcement Learning with Agentic Fused Teachers. (2026). [GitHub]

For each training batch the student forwards the 3D MR volume and produces intermediate decoder features. The teachers' features for that same case are loaded from disk (pre-extracted once and cached). A tiny policy network observes per-teacher state — cosine similarity, BCE accuracy, and (optionally) inter-teacher disagreement — and outputs a softmax weight vector w ∈ Δ^{N_teachers}. The distillation loss is a weighted feature MSE

L_student = L_seg + w * Σᵢ wᵢ · D(F_student, F_teacherᵢ)

combined with cross-entropy / Dice / focal supervised losses on the labels. The policy is updated with REINFORCE-style returns where the reward is the per-episode improvement in validation Dice (and optionally training loss reduction). See meta_teacher_optimizer.py and the train_agent_episode function in train_loops_new_policy.py.

• Automated detection and segmentation of clinically significant prostate cancer (ISUP Grade Group ≥ 2) on bi-parametric MRI as a decision-support tool for radiologists.
• Efficient deployment alternative to full teacher model ensembles — a single forward pass replaces 3–15 independent inferences.
• Multi-institutional and privacy-preserving deployment via fedGRAFT, where sites share only intermediate features, not patient images or model weights.

Out of scope: This model is not intended for autonomous clinical diagnosis without radiologist review. Not validated for prostate cancer staging, Gleason grading, non-prostate anatomical sites, or DCE sequences.

Lesion-level detection using a sextant biopsy–inspired protocol (true positive: 90th percentile of predicted voxels within ground-truth lesion boundary classified as cancer).

multiGRAFT detected 7.5% more clinically significant cancers than radiologists on the independent radical prostatectomy cohort (93 patients, GE scanners).

• All experiments were conducted on retrospective datasets; prospective validation has not yet been performed.
• Teacher pool is currently limited to three architectural families; behavior with larger or more diverse pools is unexplored.
• A formal privacy analysis for fedGRAFT has not yet been conducted.
• Missed lesions tend to be smaller (median 226.5 mm³); small-lesion sensitivity may be a limitation in screening contexts.

GRAFT/
├── run_5fold_201_binary_focal_neg1.sh   # Example launch script (5-fold, binary, focal)
├── train_student_rl_new_policy.py       # Main DDP training entry point
├── train_loops_new_policy.py            # Train / val / episode-RL inner loops
├── meta_teacher_optimizer.py            # Policy-gradient teacher-weight optimizer
├── validation_visualizer.py             # Sample qualitative validation panels
├── utils.py                             # Logging, AMP helpers, distillation losses
├── metrics.py                           # Dice / IoU / lesion-level metrics
├── setting.py                           # Optional 2D-teacher checkpoint registry
├── dataset/
│   ├── pimed_dataset_configs_local_region_5fold_binary_201.yaml  # Edit me
│   ├── pimedloader_3d.py                # Base 3D prostate-MR dataloader
│   ├── pimedloader_3d_with_features.py  # Adds teacher feature loading + LRU cache
│   ├── fixed_subset_sampler.py          # Per-rank fixed-size sampler for DDP
│   ├── region_converter.py              # 3-class ↔ region-based label encoding
│   └── base.py
├── distiller_zoo/
│   ├── contrastive_kd.py                # Contrastive feature KD loss
│   ├── feature_mse_mtkd_rl.py           # Per-teacher weighted feature MSE
│   ├── normalized_feature_losses.py     # Channel-normalized feature distance
│   └── segmentation_3d_losses.py        # Dice / focal / region-based seg losses
└── models/
    ├── fullres_3d_unet.py               # nnUNet-style full-resolution 3D U-Net (student)
    ├── unet3d.py                        # Lightweight 3D U-Net variants
    ├── swin_unetr3d.py                  # 3D Swin-UNETR backbones
    ├── policy.py                        # Policy network producing per-teacher weights
    ├── util.py / util_learned_resize.py # Building blocks
    └── (2D classification models kept for compatibility)

1. Multi-cohort 3D prostate-MR dataset in nnUNet-style layout:

   • imagesTr/ — <case>_0000.nii.gz, <case>_0001.nii.gz, <case>_0002.nii.gz (T2w / ADC / DWI stacked channels)
   • masksTr/ — prostate masks
   • labelsTr/ — lesion / cancer labels

2. Per-cohort 5-fold split JSONs — case-name lists for each fold.

3. Per-case z-score normalization stats — a JSON { "<case_id>": { "mean": <float>, "std": <float> }, ... }.

4. Pre-trained nnUNet 5-fold teachers, plus their extracted decoder features saved per case as .pt (one directory per fold). The extractor used to produce these files is not included in this release — any compatible per-case feature dump that keys on the same case IDs as imagesTr/ will work.

All five locations are configured in dataset/pimed_dataset_configs_local_region_5fold_binary_201.yaml.

Input volumes are expected to be preprocessed as follows before feature extraction:

The code targets PyTorch ≥ 2.0 with CUDA. Recommended:

python -m venv venv
source venv/bin/activate
pip install --upgrade pip
pip install \
    "torch>=2.0" torchvision \
    numpy scipy scikit-image scikit-learn pandas \
    SimpleITK torchio nibabel matplotlib tqdm pyyaml

Multi-GPU training uses torch.distributed.run with NCCL.

1. Open run_5fold_201_binary_focal_neg1.sh and edit the activation line for your Python environment and CUDA_VISIBLE_DEVICES / NUM_GPUS.

2. Open dataset/pimed_dataset_configs_local_region_5fold_binary_201.yaml and replace every <PATH_TO_...> placeholder with a real path on your system.

3. Launch:

   bash run_5fold_201_binary_focal_neg1.sh

The script defaults to 2 GPUs, batch size 2 × 2-step gradient accumulation, 200 epochs, BF16 mixed precision, full-resolution 3D U-Net student, deep supervision, region-based training, episode-based RL rewards (γ = 0.95, α = 0.7), and feature distillation from the second-to-last decoder stage (--distill-features neg1).

For the full list run:

python train_student_rl_new_policy.py --help

• Teacher feature directories must contain one .pt (or .npz) per training/validation case, keyed by the same case ID used in the nnUNet imagesTr/ filenames. Feature shapes must match the student's corresponding decoder stage.
• pimedloader_3d_with_features.py uses an LRU cache (--cache-size, default 300 cases/GPU) so feature loading does not blow up RAM when shuffling is enabled.
• If you train without the RL agent (--agent-warmup-epochs ≥ total epochs), GRAFT reduces to a uniformly-weighted multi-teacher feature distillation baseline.
• setting.py exists for compatibility with legacy 2D classification teachers; it is not needed for the 3D segmentation pipeline.

If you use this code or model, please cite:

@article{li2026graft,
  title   = {{GRAFT}: Guided Reinforcement Learning with Agentic Fused Teachers},
  author  = {Li, Cynthia X. and others},
  journal = {arXiv preprint},
  year    = {2026}
}

This work was supported by Stanford University (Departments of Radiology and Urology) and by the National Cancer Institute, National Institutes of Health (R37CA260346).

Released for research use.