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

TL;DR

InfiMed-ORBIT introduces rubric-based incremental training to tackle open-ended medical dialogue alignment, addressing the shortcomings of scalar RL rewards and reward models. The framework combines retrieval-augmented rubric generation with GRPO-based RL using a rubric-driven judge, enabling data-efficient, stable optimization. Empirically, ORBIT achieves state-of-the-art results for 4B-scale models on HealthBench-Hard with as few as 2k training samples and generalizes to InfoBench, implying broad applicability to open-ended alignment tasks. This approach offers a scalable, interpretable alternative to traditional reward design, reducing the need for external knowledge bases and domain-specific reward modeling while improving consultation quality across diverse medical scenarios.

Abstract

Reinforcement learning has powered many of the recent breakthroughs in large language models, especially for tasks where rewards can be computed automatically, such as code generation. However, these methods deteriorate in open-ended domains like medical consultation, where feedback is inherently ambiguous, highly context-dependent, and cannot be reduced to a reliable scalar signal. In such settings, RL must either rely on supervision-intensive reward models that often fail to generalize, or it falls into pathological behaviors such as reward hacking - an especially troubling risk for high-stakes medical dialogue. To address these limitations, we introduce ORBIT, an open-ended rubric-based incremental training framework for high-stakes medical dialogue. ORBIT integrates synthetic dialogue generation with dynamically constructed rubrics that serve as adaptive guides for incremental RL. Instead of relying on external medical knowledge bases or handcrafted rule sets, ORBIT uses rubric-driven feedback to steer the learning process. Its judge component can be instantiated with general-purpose instruction-following LLMs, removing the need for any task-specific fine-tuning. Applied to the Qwen3-4B-Instruct model, ORBIT raises the HealthBench-Hard score from 7.0 to 27.5 using only 2k training samples, achieving SOTA performance for models at this scale. With larger rubric datasets, ORBIT-trained models further compete with the strongest open-source baselines on HealthBench-Hard. Our analysis shows that rubric-guided RL consistently improves consultation quality across diverse medical scenarios. We also apply such rubric generation and training pipeline to InfoBench, where ORBIT enhances instruction-following performance, highlighting the generality of rubric-based feedback.

Figures (13)

• Figure 1: Overview of the ORBIT pipeline for rubric-guided, incremental alignment of LLMs on open-ended medical dialogue tasks. (a) Sampling medical consultation cases, truncating dialogues to form multi-turn realistic queries. (b) During the RL stage, we adopt the GRPO framework, where a dynamic, query-specific rubric is generated for each medical case. Each response is evaluated by an external judge model based on this rubric, and the resulting score serves as the reward signal for policy optimization. (c) For each medical consultation case, a Retrieval-Augmented Generation (RAG) framework constructs a tailored, multi-dimensional rubric. The retrieval module references HealthBench rubrics (no hard subset) as exemplars, while the generation model is guided to avoid direct duplication. These rubrics are used solely by the external judge model for evaluation, ensuring that the model remains unbiased and enabling stable policy exploration. The RAG framework is also extensible, supporting integration of larger and more diverse medical consultation datasets.
• Figure 2: Multi-dimensional performance comparison of ORBIT models. This figure presents a comprehensive performance analysis categorized by Theme and Axis. The results are divided into two groups: (a, c) Internal methodological comparisons among the base instructor-tuned model (Qwen3-4B-Instruct), and our proposed models (InfiMed-ORBIT-4B with 2k training samples, 8k samples and 28k samples). (b, d) Large-scale benchmark comparisons between our best-performing model, InfiMed-ORBIT-4B (28k samples), and significantly larger proprietary and open-source models, including Qwen3-30B-Instruct, GPT-4.1, Baichuan-M2-32B, and Qwen3-30B-Thinking. For brevity, "Health Data T." refers to "Health Data Tasks," and "Emerg. Referrals" denotes "Emergency Referrals."
• Figure 3: From left to right, the figure illustrates the effects of rubric-based filtering and sample-based filtering. For rubric filtering, we track the growth in output length during training, while for sample filtering, we measure the corresponding change in overall training time. Both approaches substantially enhance training efficiency, and adjusting the filtering thresholds provides fine-grained control over the trade-off between computational cost and final model performance.
• Figure 4: Distributional Analysis of Performance Metrics: Inference Scaling vs. Rubric-RL Alignment. We compare the off-the-shelf Qwen3-4B-Instruct baseline (at $K=8$ and $K=40$ rollouts) with our proposed InfiMed-ORBIT-4B (at $K=40$ rollouts). From left to right, the subplots depict Kernel Density Estimates (KDE) and histograms for: (a) Average Normalized Score per query, (b) Max Normalized Score (Best-of-$K$), (c) Rubric Pass Rate (average compliance per criterion), and (d) Rubric Coverage / Hit Rate (probability of satisfying each criterion at least once).
• Figure 5: Overall details of the designed system prompt for the rubrics generation pipeline.