Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

medR: Reward Engineering for Clinical Offline Reinforcement Learning via Tri-Drive Potential Functions

Qianyi Xu, Gousia Habib, Feng Wu, Yanrui Du, Zhihui Chen, Swapnil Mishra, Dilruk Perera, Mengling Feng

TL;DR

medR addresses the central challenge of reward engineering in offline clinical RL by introducing a Tri-Drive potential framework that combines Survival, Confidence, and Competence signals with offline verification. It leverages LLM-driven interpretable feature selection to tailor rewards to specific diseases, and uses a potential-based shaping approach to provide dense, clinically meaningful guidance without altering the optimal policy. Offline tri-drive fitness metrics (Survival, Confidence, Competence) are optimized with NSGA-II to balance outcomes, uncertainty, and intervention cost, yielding policies that outperform clinician baselines on three high-stakes tasks as measured by WIS and action agreement. The approach emphasizes safety, interpretability, and disease-specific customization, offering a scalable pathway toward safer, more effective AI-assisted decision support in ICU settings, while noting the need for prospective validation and careful bias management.

Abstract

Reinforcement Learning (RL) offers a powerful framework for optimizing dynamic treatment regimes (DTRs). However, clinical RL is fundamentally bottlenecked by reward engineering: the challenge of defining signals that safely and effectively guide policy learning in complex, sparse offline environments. Existing approaches often rely on manual heuristics that fail to generalize across diverse pathologies. To address this, we propose an automated pipeline leveraging Large Language Models (LLMs) for offline reward design and verification. We formulate the reward function using potential functions consisted of three core components: survival, confidence, and competence. We further introduce quantitative metrics to rigorously evaluate and select the optimal reward structure prior to deployment. By integrating LLM-driven domain knowledge, our framework automates the design of reward functions for specific diseases while significantly enhancing the performance of the resulting policies.

medR: Reward Engineering for Clinical Offline Reinforcement Learning via Tri-Drive Potential Functions

TL;DR

medR addresses the central challenge of reward engineering in offline clinical RL by introducing a Tri-Drive potential framework that combines Survival, Confidence, and Competence signals with offline verification. It leverages LLM-driven interpretable feature selection to tailor rewards to specific diseases, and uses a potential-based shaping approach to provide dense, clinically meaningful guidance without altering the optimal policy. Offline tri-drive fitness metrics (Survival, Confidence, Competence) are optimized with NSGA-II to balance outcomes, uncertainty, and intervention cost, yielding policies that outperform clinician baselines on three high-stakes tasks as measured by WIS and action agreement. The approach emphasizes safety, interpretability, and disease-specific customization, offering a scalable pathway toward safer, more effective AI-assisted decision support in ICU settings, while noting the need for prospective validation and careful bias management.

Abstract

Reinforcement Learning (RL) offers a powerful framework for optimizing dynamic treatment regimes (DTRs). However, clinical RL is fundamentally bottlenecked by reward engineering: the challenge of defining signals that safely and effectively guide policy learning in complex, sparse offline environments. Existing approaches often rely on manual heuristics that fail to generalize across diverse pathologies. To address this, we propose an automated pipeline leveraging Large Language Models (LLMs) for offline reward design and verification. We formulate the reward function using potential functions consisted of three core components: survival, confidence, and competence. We further introduce quantitative metrics to rigorously evaluate and select the optimal reward structure prior to deployment. By integrating LLM-driven domain knowledge, our framework automates the design of reward functions for specific diseases while significantly enhancing the performance of the resulting policies.
Paper Structure (55 sections, 2 theorems, 16 equations, 7 figures, 8 tables, 1 algorithm)

This paper contains 55 sections, 2 theorems, 16 equations, 7 figures, 8 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Intrinsic Motivation
  4. Potential-Based Reward
  5. Methodology
  6. Interpretable Feature Selection
  7. Potential Based Reward Function Generation
  8. Tri-Drive Reward Selection
  9. Experiments
  10. Tasks
  11. Baselines
  12. Heuristic Reward Design
  13. LLM-based Design
  14. Evaluation Metrics
  15. RQ1: Should we use mortality as fitness or reward?
  16. ...and 40 more sections

Key Result

Proposition 4.1

Let the reward function be defined as $R(s_t, a_t, s_{t+1}) = F(s_t, s_{t+1}) - \lambda \mathcal{C}(a_t)$, where $F(s_t, s_{t+1}) = \gamma \Phi(s_{t+1}) - \Phi(s_t)$ is a potential-based shaping function and $\mathcal{C}(a_t) > 0$ is a strictly positive action cost. While the shaping term $F$ forms

Figures (7)

  • Figure 1: Framework pipeline of medR.
  • Figure 2: The relationship between cumulative process reward and patient mortality probability for different tasks. A valid process should demonstrate clear downward trend with mortality, baselines that include mortality as part of the reward design are excluded for fair comparison.
  • Figure 3: WIS training plot.
  • Figure 4: Decomposition of Tri-Drive fitness across tasks. The stacked bars represent the contribution of three distinct evaluation metrics: Survival ($J_{surv}$), Confidence ($J_{conf}$), and Competence ($J_{comp}$).
  • Figure 5: Sepsis Policy Agreement. Confusion matrices comparing AI vs. Clinician actions for IV Fluids (left) and Vasopressors (right) across different reward functions.
  • ...and 2 more figures

Theorems & Definitions (4)

  • Proposition 4.1: Cost Regularization Breaks Invariance and Enforces Efficiency
  • proof
  • Proposition 4.2: Equivalence to Lagrangian Relaxation of CMDP
  • proof