A Note on Hybrid Online Reinforcement and Imitation Learning for LLMs: Formulations and Algorithms
Yingru Li, Ziniu Li, Jiacai Liu
TL;DR
This work presents a unified framework for fine-tuning large language models by combining imitation learning through a Dense gradient and reinforcement learning through a Sparse gradient, derived from a trajectory-level KL divergence and reward objective. The key contribution is a per-step gradient decomposition that separates an analytic imitation signal from a long-horizon reward signal, enabling efficient logit-level computation and GPU-friendly implementation. The framework is shown to be mathematically equivalent to KL-regularized RLHF and unifies methods such as online imitation learning, DAPO, and standard KD under different hyperparameter choices, with a natural curriculum for balancing imitation and reward. The proposed approach offers practical pathways for scaling LLM fine-tuning with stable on-policy learning and flexible extension to multiple teachers and rewards, holding promise for more robust and aligned models in real-world applications.
Abstract
We present a unified framework for Large Language Model (LLM) fine-tuning that integrates Imitation Learning and Reinforcement Learning. By analyzing the gradient of a composite objective combining trajectory-level KL divergence with task rewards, we derive a natural decomposition into two components: (1) an analytically computable Dense Gradient for token-level imitation, and (2) a Monte Carlo estimated Sparse Gradient for long-horizon reward optimization. The Dense Gradient admits a closed-form logit-level formula, enabling efficient GPU implementation.