EAG-RL: Reinforcement Learning with Expert Attention Guidance for EHR Reasoning in LLMs

Table of Contents

• Overview
• Paper Abstract
• Installation
• Usage
• Code Structure
• Datasets
• Citation
• License
• Acknowledgments

Overview

This repository implements EAG-RL, a novel framework introduced in the paper "Toward Better EHR Reasoning in LLMs: Reinforcement Learning with Expert Attention Guidance". EAG-RL enhances Large Language Models (LLMs) for clinical predictions based on Electronic Health Records (EHRs) by integrating expert-guided trajectory distillation and attention-aligned reinforcement learning.

Key Contributions

• Expert-Guided Trajectory Distillation: Uses Monte Carlo Tree Search (MCTS) guided by expert EHR models to generate high-quality multi-step reasoning trajectories for supervised fine-tuning (SFT).
• Attention-Aligned Policy Optimization: Employs reinforcement learning with joint rewards from classification accuracy and attention alignment to expert models, plus entropy-aware adaptive upper clipping to encourage exploration.
• Superior Performance: Achieves an average improvement of 14.62% over baselines in AUROC and AUPRC on real-world EHR datasets (MIMIC-IV and TJH) for tasks like in-hospital mortality and 30-day readmission prediction.
• Robustness and Generalization: Demonstrates better robustness to feature perturbations and out-of-distribution generalization compared to traditional expert models and other LLM baselines.

The framework addresses challenges in EHR reasoning, such as high-dimensional time-series data, by teaching LLMs to decompose problems into subquestions, integrate evidence, and align attention with clinically important features.

Code is open-sourced for reproducibility and further research.

Paper Abstract

Large Language Models (LLMs) excel in unstructured medical text tasks but underperform specialized deep learning models ("expert EHR models") in clinical predictions based on Electronic Health Records (EHRs), which are time-series and high-dimensional. To address this, we propose EAG-RL, a two-stage framework: (1) Expert-guided MCTS constructs high-quality step-by-step reasoning trajectories for trajectory-level SFT to warm-start the LLM policy; (2) Reinforcement learning aligns the LLM's attention with expert-identified key clinical features. Experiments on two real EHR datasets show EAG-RL improves by an average of 14.62%, with enhanced robustness to perturbations and better cross-domain generalization. Code is available at this repository.

Installation

Prerequisites

• Python 3.8+
• Conda (recommended for environment management)

Setup Environment

Create and activate the Conda environment using the provided environment.yml:

conda env create -f environment.yml
conda activate verl

For logging with SwanLab (optional but recommended for tracking experiments):

1. Copy the example environment file:

   cp .env.example .env

2. Edit .env to add your SwanLab API key:

   SWANLAB_API_KEY = "your-api-key"
   SWANLAB_LOG_DIR = "./logs"
   

Dependencies

Key libraries include:

• PyTorch
• Transformers (Hugging Face)
• Verl (for reinforcement learning components)
• Other dependencies listed in environment.yml

Usage

EAG-RL training consists of two stages: Supervised Fine-Tuning (SFT) for trajectory distillation and Gradient-based Reinforcement Learning from Human Feedback (GRPO) for policy optimization.

Stage 1: Supervised Fine-Tuning (SFT)

Run the SFT script to distill expert-guided trajectories:

bash scripts/run_sft_baseline.sh

This initializes the LLM policy with high-quality reasoning trajectories generated via expert-guided MCTS.

Stage 2: Reinforcement Learning (GRPO)

After SFT, run the GRPO script for attention-aligned policy optimization:

bash scripts/train_grpo_lora_numerical.sh

Evaluation

Evaluate the trained model using vLLM for efficient inference:

bash scripts/eval_llm_lora_adapter_vllm.sh

Metrics include AUROC and AUPRC, computed with 100 bootstrap samples for statistical robustness.

Code Structure

The repository is organized as follows:

• data/: Raw and processed EHR datasets (e.g., MIMIC-IV, TJH).
• logs/: Training logs (SwanLab integration).
• models/: Saved model checkpoints.
• outputs/: Generated outputs from runs.
• results/: Evaluation results and metrics.
• scripts/: Bash scripts for training, evaluation, and utilities.
• src/:
  • args/: Argument parsers (e.g., ehrpo_args.py).
  • data/: Data preparation (e.g., prepare_ehr_data.py, numerical variants).
  • eval/: Metric calculation (e.g., calculate_metrics_bootstrap.py).
  • models/:
    • ehr_worker/: Expert EHR models (e.g., Concare).
    • llm_worker/: LLM inference methods (e.g., MCTS, baseline direct).
  • prompt/: Prompt templates and wrappers (e.g., prompt_mcts.py, numerical variants).
  • run/: Experiment runners (e.g., runexp.py).
  • trainer/: Training modules (e.g., eagrl_trainer.py, grpo_trainer.py, SFT and STAR baselines).
  • utils/: Utilities (e.g., bootstrap.py, numerical aggregation).
  • verl/: VERL framework integration for RL (includes trainers, workers, utils).
• README.md: This file.
• environment.yml: Conda environment specification.

For detailed implementation, refer to src/main.py as the entry point.

Datasets

• MIMIC-IV: ICU data for in-hospital mortality and readmission prediction.
• TJH: Annotated inpatient data.

Citation

If you use this code or find the paper useful, please cite:

@misc{fang2025betterehrreasoningllms,
      title={Toward Better EHR Reasoning in LLMs: Reinforcement Learning with Expert Attention Guidance}, 
      author={Yue Fang and Yuxin Guo and Jiaran Gao and Hongxin Ding and Xinke Jiang and Weibin Liao and Yongxin Xu and Yinghao Zhu and Zhibang Yang and Liantao Ma and Junfeng Zhao and Yasha Wang},
      year={2025},
      eprint={2508.13579},
      archivePrefix={arXiv},
      primaryClass={cs.AI},
      url={https://arxiv.org/abs/2508.13579}, 
}

License

This project is licensed under the Apache License. See the LICENSE file for details.

Acknowledgments

• Built on VERL for RL components.
• Thanks to the authors of expert EHR models like Concare.

For issues or contributions, please open a GitHub issue or pull request.