InfiMed-ORBIT: Aligning LLMs on Open-Ended Complex Tasks via Rubric-Based Incremental Training

ABSTRACT

ORBIT is an open-ended rubric-based incremental training framework for high-stakes medical dialogue. It combines synthetic dialogue generation, retrieval-augmented rubric construction, difficulty filtering, and rubric-guided reinforcement learning to align LLMs on context-dependent consultation tasks. Applied to Qwen3-4B-Instruct, ORBIT raises the HealthBench-Hard score from 7.0 to 27.5 using only 2k training samples, with larger rubric sets improving performance to 33.6 and 37.3.

Medial LLMOpen-Ended LLM AlignmentScaling Rubrics-RL
Apr 30, 2026
Pengkai Wang, Pengwei Liu, Qi Zuo, Zhijie Sang, Congkai Xie, Hongxia Yang
Overview of the ORBIT pipeline

ORBIT is a rubric-guided reinforcement learning framework for aligning LLMs on open-ended medical dialogue.


Abstract

Reinforcement learning has powered many recent breakthroughs in large language models, especially for tasks where rewards can be computed automatically, such as code generation. However, RL is less reliable in open-ended medical dialogue, where feedback is ambiguous, context-dependent, and difficult to collapse into a single scalar signal. This often requires heavily supervised reward models and risks reward hacking.

ORBIT is an open-ended rubric-based incremental training framework for high-stakes medical dialogue. It integrates synthetic dialogue generation with dynamically constructed rubrics that serve as adaptive guides for incremental RL. Unlike approaches that rely on external medical knowledge bases or handcrafted rules, ORBIT uses rubric-guided evaluation and can be implemented with general-purpose instruction-following LLMs, avoiding task-specific judge fine-tuning.

Applied to Qwen3-4B-Instruct, ORBIT raises the HealthBench-Hard score from 7.0 to 27.5 using only 2k training samples. With larger rubric sets, ORBIT remains competitive with the strongest open-source baselines and consistently improves consultation quality across diverse medical scenarios. The same rubric pipeline also improves InfoBench instruction-following performance.

Introduction

As high-quality pretraining data become harder to expand and model scaling brings smaller marginal gains, progress increasingly comes from post-training. Supervised fine-tuning aligns models to target behaviors with token-level demonstrations, while reinforcement learning optimizes policies against preference objectives. In domains with explicit, verifiable ground truth, RL with verifiable rewards can deliver reliable gains.

Medical consultation is different. Response quality depends on multiple context-sensitive axes, including safety, empathy, clinical appropriateness, context awareness, and completeness. Reducing such judgments to a single opaque scalar reward is fragile, expensive, and hard to transfer. In high-stakes dialogue, this can also amplify reward hacking risks.

Rubric-based evaluation provides a more interpretable alternative. Rubrics decompose response quality into explicit criteria and support transparent credit assignment. HealthBench shows the value of expert-designed rubrics for clinical reasoning, but existing medical LLMs still perform poorly on HealthBench-Hard, exposing a gap between QA-style optimization and realistic consultation.

ORBIT addresses this gap by using a small collection of rubric examples to guide evaluation. Through retrieval-augmented in-context prompting, it directs a general-purpose judge model to assess responses against query-specific criteria inside an incremental RL framework.

The contributions are:

  • Rubric-based open-ended alignment: ORBIT replaces opaque scalar rewards and costly reward models with interpretable, rubric-driven feedback for high-stakes medical dialogue.
  • Context-aware rubric generation and training: the pipeline combines retrieval-augmented in-context prompting, multi-stage filtering, dynamic sampling, and incremental optimization.
  • Empirical validation: ORBIT substantially improves HealthBench-Hard and transfers to InfoBench, suggesting that rubric-guided RL is useful beyond medical dialogue.

Open-Ended Benchmarks

LLM evaluation for open-ended generation is shifting from short-form automatic metrics toward holistic rubric-based frameworks. Benchmarks such as HealthBench, PaperBench, WildBench, AMEGA, and MultiChallenge define fine-grained criteria across diverse scenarios. HealthBench is especially relevant because it tests realistic medical consultation, where response quality is multidimensional and context-dependent.

Reward Models in LLMs

RLHF and related post-training methods usually depend on reward models or rule-based reward signals. Rule-based rewards work when criteria are explicit and verifiable, while learned reward models can approximate human judgment but are expensive to construct, domain-dependent, and sensitive to distribution shift. ORBIT instead uses automatically derived, case-specific rubrics to provide structured and transparent feedback.

LLMs for Health

Medical LLM research spans differential diagnosis, clinical documentation, mental health assistance, radiology reporting, and agentic clinical reasoning. Many systems remain specialized to narrow tasks and struggle with heterogeneous, context-dependent reasoning. ORBIT targets this open-ended setting by aligning models with clinical consultation rubrics rather than only training on question-answer style supervision.

ORBIT Framework

Given a clinical case in chat format or outpatient chart format, ORBIT turns it into practical RL training data. The system produces final <dialogue, rubrics> pairs through three stages: dialogue QA simulation, rubric generation with in-context learning, and rubric-based reinforcement learning.

Dialogue QA Simulation

ORBIT uses patient query dialogues from DoctorAgent-RL as seeds for simulated consultations and supplements them with additional open-source medical dialogue datasets. These sources provide realistic multi-turn consultation contexts. After dialogue QA data are prepared, advanced LLMs generate rubrics step by step through in-context learning.

Rubric Generation With In-Context Learning

ORBIT builds a diagnostic database from seed HealthBench rubrics. Each consultation dialogue and rubric is embedded and stored in vector pools:

  • a case-rubric pair pool that keeps each case together with its rubric set;
  • a rubric pool that collects unique rubric criteria and their embeddings.

For a new consultation query, ORBIT retrieves similar cases and semantically aligned rubric candidates, reranks them, and prompts a generation model to synthesize a query-specific checklist. The generated rubric set is used only by the judge model for response evaluation.

Difficulty Filtering With Pass@k

ORBIT filters both queries and rubrics to focus RL on useful learning signals. For each query, the current model generates multiple rollouts and the judge scores response-rubric pairs.

The average query score is:

sˉq=1nrolloutRqi=1nrolloutrRqS(yi,r)\bar{s}_q = \frac{1}{n_{\text{rollout}}\cdot|\mathcal{R}_q|} \sum_{i=1}^{n_{\text{rollout}}}\sum_{r\in\mathcal{R}_q}S(y_i,r)

High-scoring cases are too easy, while very low-scoring cases may be unsolved. ORBIT keeps cases in an intermediate difficulty band.

At the rubric level, the pass rate is:

P(r,q)=1nrollouti=1nrollout1{S(yi,r)τs}P(r,q)= \frac{1}{n_{\text{rollout}}} \sum_{i=1}^{n_{\text{rollout}}} \mathbf{1}\{S(y_i,r)\ge\tau_s\}

Rubrics with very high pass rates are removed because they provide little learning signal.

Rubric-Based Reinforcement Learning

ORBIT optimizes the policy with Group Relative Policy Optimization (GRPO). For each query, it samples a group of outputs and computes advantages from the group mean and standard deviation of rubric rewards.

The rubric-aware reward is:

R(q,oi)=j=1RqI[Mjudge(oi,critj)True]wjR(q,o_i)= \sum_{j=1}^{|\mathcal{R}_q|} \mathbb{I}\left[ \mathcal{M}_{\text{judge}}(o_i,\text{crit}_j)\to\text{True} \right]\cdot w_j

This gives dense, semantic feedback over clinical reasoning steps rather than only final answers.

ORBIT adds two stability strategies:

  • Variance-aware dynamic sampling: batches with near-zero reward variance are filtered so updates come from discriminative samples.
  • Staged entropic restarts: at stage transitions, the policy starts from the best previous checkpoint while sampling temperature is reset to reintroduce exploration.

Experiments

Settings

The primary experiments use Qwen3-4B-Instruct-2507 as the base model and Qwen3-30B-Instruct-2507 as the rubric evaluation judge during training. The main benchmark is HealthBench-Hard, a challenging 1k-case subset of HealthBench designed to test open-ended medical consultation.

The data are organized into:

  • Core Experimental Dataset: 2,082 multi-turn consultations from IMCS21, CHIP-MDCFNPC, and MedDG.
  • Scalability Dataset: roughly 8k curated samples from DoctorAgent-RL.
  • Large-Scale Validation Dataset: 20k additional samples from ReMeDi.

To avoid contamination, HealthBench-Hard samples are excluded from seed data. Only non-Hard HealthBench-4k rubrics are used for RAG seed examples, with additional lexical and semantic filtering.

HealthBench-Hard Results

ORBIT substantially improves smaller open-source models and allows a 4B model to compete with much larger systems.

ModelEmergency referralsContext seekingGlobal healthHealth data tasksCommunicationHedgingResponse depthAccuracyCompletenessCommunication qualityContext awarenessInstruction followingTotal Score
GPT-4.120.512.312.19.714.912.317.530.5070.6060.513.2
GPT-5 thinking------------46.2
Qwen3-4B-Instruct base9.38.57.108.612.25.124.10.857.5045.07.0
Qwen3-4B-Thinking14.412.52.403.58.5023.2042.5039.65.2
InfiMed-ORBIT-4B (2k)39.937.830.26.226.632.26.631.838.145.316.843.727.5
InfiMed-ORBIT-4B (8k)44.649.134.69.328.042.610.433.642.046.732.350.533.6
InfiMed-ORBIT-4B (28k)51.650.541.910.830.943.813.438.148.942.131.349.137.3
Qwen3-30B-Instruct18.312.914.717.919.49.528.528.5045.2033.713.1
Baichuan-M2-32B45.639.535.621.332.040.919.941.344.651.619.348.034.5

With 2k samples, ORBIT improves the total score from 7.0 to 27.5, a 293% relative improvement. Scaling to 8k and 28k samples further raises the score to 33.6 and 37.3.

Multi-dimensional performance comparison of ORBIT models

Multi-dimensional performance comparison by clinical theme and evaluation axis.

Ablations and Data Efficiency

ORBIT studies the effects of rubric model selection, judge selection, pass@k filtering, dynamic sampling, and multi-stage restart training. The core finding is that not all samples and rubrics contribute equally: easy samples and redundant rubrics can add computation without useful gradients.

SettingTotal Score
Qwen3-4B-Instruct base7.2
InfiMed-ORBIT-4B, no filter20.2
Rubric pass@k filter 0-0.7519.9
Rubric pass@k filter 0-0.5017.9
Rubric pass@k filter 0-0.2518.7
Sample pass@k filter 0-0.7519.7
Sample pass@k filter 0-0.5014.5
InfiMed-ORBIT-4B, 8k data25.9
InfiMed-ORBIT-4B, restart27.3

Selective filtering improves the compute-performance tradeoff. Multi-stage restart training further improves the 8k setting from 25.9 to 27.3 under the GPT-OSS-120B-middle evaluation setup.

Training dynamics under filtering strategies

Filtering controls output length and training runtime while preserving useful reward signals.

Rubric-RL Versus Inference Scaling

The paper compares inference scaling on the baseline Qwen3-4B-Instruct model with ORBIT-trained models. Increasing baseline rollouts from K=8 to K=40 only gives marginal improvements, suggesting a capability ceiling. ORBIT instead reshapes the model distribution: high-scoring, rubric-compliant responses become much more likely, rather than appearing only through sampling luck.

Distributional comparison of inference scaling versus rubric-guided RL

Rubric-guided RL shifts the policy distribution toward high rubric compliance beyond inference scaling alone.

Generalization to InfoBench

ORBIT is also evaluated on InfoBench, an open-ended instruction-following benchmark. Using the easy split as seed data and retrieving 2k WildChat samples for rubric-augmented RAG, ORBIT improves Qwen3-4B-Base from 42.0 to 82.9 on the hard split. This suggests that rubric-based reinforcement learning can generalize beyond medical dialogue.

Supplementary Analysis

Rubric Generator

The supplementary material describes a retrieval-augmented rubric generator. It retrieves reference cases relevant to a medical scenario and uses them as in-context examples for rubric generation. The prompt explicitly asks for both positive criteria and negative criteria, so the judge can reward desirable behaviors and penalize unsafe or incomplete behaviors.

Judge and Rubric Model Selection

The paper compares several open-source judge models against GPT-4.1. GPT-OSS-120B-middle is used during algorithm design and data construction because it aligns relatively well with GPT-4.1 while being more computationally practical. Final HealthBench-Hard results use GPT-4.1 for comparability with the official protocol.

For rubric generation, DeepSeek-R1 and Gemini-2.5-Pro produce the most consistent improvements under GPT-OSS-120B-based evaluation. The default pipeline adopts DeepSeek-R1 as the rubric generator.

Data Integrity Diagnostics

The appendix examines semantic differences between HealthBench Consensus and HealthBench-Hard rubrics. Consensus rubrics form tighter clusters, while Hard rubrics are sparse and fragmented, reflecting greater constraint variability and reasoning difficulty.

t-SNE visualization of HealthBench Consensus rubrics

HealthBench Consensus rubrics form dense, cohesive clusters.

t-SNE visualization of HealthBench Hard rubrics

HealthBench-Hard rubrics show a sparse and fragmented topology.

The paper also performs prompt-level and rubric-level contamination diagnostics between NoHard and Hard data. High rubric similarity can reflect benign template reuse, so the joint prompt-rubric test is used to screen for instance leakage.

Contamination diagnostics between NoHard and Hard data

Prompt-rubric contamination diagnostics suggest high rubric similarity rarely co-occurs with near-duplicate prompts.

Impact Statement

Open-ended, high-stakes domains such as medical consultation require alignment methods that are reliable and interpretable. ORBIT translates qualitative clinical judgments into explicit training signals, improving data efficiency and controllability. It is not a substitute for clinical expertise, but it provides a practical step toward safer and more controllable alignment of language models in realistic medical dialogue settings.

Limitations

ORBIT still relies on a seed set of human-crafted rubrics. Generating seed rubrics from clinical guidelines and best practices is an important future direction. The experiments also fix Qwen3-30B-Instruct as the judge model during training; judge calibration and reasoning ability can vary across model families and scales, so judge-robust training remains important for deployment.

BibTeX

@misc{infimed-orbit-2026,
      title={InfiMed-ORBIT: Aligning LLMs on Open-Ended Complex Tasks via Rubric-Based Incremental Training},
      author={Anonymous},
      year={2026},
      url={},
}