KEPO: Knowledge-Enhanced Preference Optimization for Reinforcement Learning with Reasoning

TL;DR

Reasoning-based reinforcement learning with sparse rewards suffers from credit assignment and exploration failures, leading to a learning cliff. The authors propose KEPO, which integrates a quality-gated on-policy distillation signal with a knowledge-enhanced exploration strategy that uses teacher hints to inject reward-bearing trajectories when naive rollout fails. The method is instantiated in a medical vision-language VQA setting (OmniMedVQA) with single-domain MRI training and cross-modality evaluation, showing improved training stability, more coherent reasoning traces, and stronger out-of-distribution generalization than RL and uniform distillation baselines. This work offers a general framework for injecting structured teacher guidance into reinforcement-based post-training to enable robust reasoning emergence under sparse rewards.

Abstract

Reinforcement learning (RL) has emerged as a promising paradigm for inducing explicit reasoning behaviors in large language and vision-language models. However, reasoning-oriented RL post-training remains fundamentally challenging due to sparse trajectory-level rewards, leading to ambiguous credit assignment and severe exploration failures that can trap the policy in a ``learning cliff.'' Recent on-policy distillation methods introduce dense teacher supervision to stabilize optimization, but apply it uniformly across all generated trajectories. We argue that such uniform distillation is ill-suited for reasoning-intensive tasks, as low-quality on-policy trajectories often originate from early logical errors, and distillation under flawed contexts injects noisy and misaligned gradients. To address these challenges, we propose Knowledge-Enhanced Preference Optimization (KEPO), a unified post-training framework that integrates: (i) a quality-gated on-policy distillation objective that selectively applies dense teacher guidance only to high-quality trajectories, and (ii) a knowledge-enhanced exploration strategy that leverages hints learned from a teacher model to rejectively sample reward-positive on-policy trajectories for RL, thereby mitigating exploration collapse. Evaluated on a challenging medical visual question answering benchmark under single-source generalization, KEPO demonstrates improved training stability, more coherent reasoning behaviors, and superior out-of-distribution performance over reinforcement learning and on-policy distillation baselines.

Figures (2)

• Figure 1: Overview of the KEPO framework. KEPO augments reinforcement-based post-training by jointly integrating knowledge-enhanced exploration and quality-gated distillation. Given an input $(x, y)$ during training, the student model adaptively triggers knowledge-enhanced exploration when all sampled trajectories receive zero rewards. During exploration, a teacher model first generates auxiliary reasoning hints. Then, conditioned on the hint and the ground-truth answer, the student model generates reward-positive candidates $(y_{n+1}, \dots, y_{G=n+m})$ via a rejection sampling procedure. By supplementing the standard on-policy rollouts $(y_1, \dots, y_n)$, the candidate pool contains a total $G$ responses. Finally, the reinforcement learning objective is optimized over the full pool, while an auxiliary on-policy distillation loss is selectively applied only to high-quality samples satisfying $r_i \geq \tau$. The student policy is updated by jointly optimizing the resulting objective $\mathcal{J}_{\text{KEPO}}$.
• Figure 2: Training dynamics under different post-training strategies. All experiments are trained for 5 epochs, with checkpoints evaluated every 10 steps on two metrics: Left: accuracy on the in-domain (ID) MRI modality. Right: average accuracy across out-of-distribution (OOD) modalities. Compared to GRPO, KEPO-KE achieves stronger ID and OOD performance in the early training stage, while the full KEPO framework consistently outperforms all baselines throughout training and converges faster.