Skip to content

AgnesAI-Labs/Evo-PI

 
 

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

2 Commits
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Evo-PI

Introduction

This repository contains the code for our ACL 2026 main conference paper: Evo-PI: Aligning Medical Reasoning via Evolving Principle-Guided Supervision.

The project studies medical visual reasoning alignment with evolving principle-guided supervision, where modality-specific reasoning principles are used by the reward function and updated during iterative RL training.

@inproceedings{Evo_PI,
  title = {Evo-PI: Aligning Medical Reasoning via Evolving Principle-Guided Supervision},
  author = {Xianda Zheng, Huan Gao, Meng-Fen Chiang, Michael Witbrock, Kaiqi Zhao and Shangyang Li},
  booktitle = {Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics},
  year = {2026}
}

Overview

This repository contains the training code and scripts for Evo-PI. It is built on verl and provides modality-specific training entry points for:

  • CT
  • Dermoscopy
  • FP
  • MI
  • MR
  • OCT
  • US
  • X-Ray

Each modality folder under src/ contains training scripts, a reward function, an evolving principles.json file, and an update.py script for updating the principles.

Setup

  1. Install verl.

    Evo-PI is built on the verl training framework. Set up a Python environment with CUDA/PyTorch support, install verl, and make sure the local verl/ package in this repository is available when running the training scripts.

  2. Install the required training dependencies.

    The helper script below may be used as a reference for installing the expected vllm, sglang, transformers, datasets, ray, wandb, and related packages:

bash scripts/install_vllm_sglang_mcore.sh
  1. Run commands from the repository root.

    Add the repository to PYTHONPATH so that python -m verl.trainer.main_ppo can import the local code:

export PYTHONPATH=$PWD:$PYTHONPATH

Before Training

The scripts contain local paths from the original experiment environment. Update them before running.

  1. Set the model path.

    Each training script loads a base or SFT model through MODEL_ID, for example:

    MODEL_ID="/path/to/your/model"

    This path is passed to actor_rollout_ref.model.path=${MODEL_ID}. Some scripts hard-code MODEL_ID; edit those lines directly or change them to read from an environment variable.

  2. Set the data paths.

    Update data.train_files and data.val_files to your local parquet files:

    data.train_files=/path/to/med_CT/train.parquet
    data.val_files=/path/to/med_CT/test.parquet

    The expected parquet columns include prompt, images, reward_model, and extra_info. See src/med_parquet.py for an example conversion script.

  3. Fix script paths.

    The original launch scripts refer to run/<MODALITY>/..., while this repository stores the scripts under src/<MODALITY>/.... Before running, either replace paths such as:

    run/CT/med_reward.py
    run/CT/update.py

    with:

    src/CT/med_reward.py
    src/CT/update.py

    or create a compatible run/ directory/symlink in your local workspace.

  4. Update the OpenAI key in update.py.

    Each src/<MODALITY>/update.py calls the OpenAI API to update principles.json. Replace:

    api_key="sk-your_key"

    with your own key, and update base_url if you are not using the proxy URL in the script. Also make sure the principles.json path in update.py and med_reward.py points to the correct modality folder.

  5. Adjust GPU settings.

    The scripts usually assume four GPUs, for example:

    CUDA_VISIBLE_DEVICES=0,1,2,3
    trainer.n_gpus_per_node=4
    actor_rollout_ref.rollout.tensor_model_parallel_size=2

    Change these values for your hardware.

Running Training

src/train.sh is the main reference for how to launch experiments. It shows two common patterns: running individual modality scripts with nohup, and merging FSDP checkpoints after training.

To run one experiment directly:

bash src/CT/med_r1_grpo.sh

Other available script variants follow the same naming pattern:

bash src/CT/med_r1_gspo.sh
bash src/CT/huatuo_grpo.sh
bash src/CT/huatuo_gspo.sh

To run in the background, follow the style in src/train.sh:

nohup bash src/CT/med_r1_grpo.sh > src/CT/med_r1_grpo.log 2>&1 &
tail -f src/CT/med_r1_grpo.log

Only launch multiple lines from src/train.sh when you have enough GPUs. The file starts many jobs in parallel, so it is best to edit it and keep only the experiments you want to run.

Merging Checkpoints

src/train.sh also includes examples for merging FSDP checkpoints into Hugging Face format:

python -m verl.model_merger merge \
    --local_dir /path/to/checkpoints/<experiment>/global_step_<step>/actor \
    --target_dir /path/to/output_model \
    --backend fsdp

Change --local_dir to the saved checkpoint directory and --target_dir to the desired output model directory.

About

No description, website, or topics provided.

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages