Aprendizaje por Refuerzo a partir de Retroalimentación Rica con DAgger Distribucional

Resumen

Los modelos de razonamiento han avanzado rápidamente, pero la receta dominante de aprendizaje por refuerzo a partir de recompensas verificables (RLVR) sigue siendo sorprendentemente limitada: muestrear muchas respuestas y recompensar cada una con un solo bit que indique si la respuesta final es correcta. Sin embargo, muchos entornos proporcionan retroalimentación enriquecida, incluyendo trazas de ejecución, salidas de herramientas, correcciones de expertos y autoevaluaciones del modelo. Estudiamos cómo utilizar dicha retroalimentación mediante una variante distribucional del clásico algoritmo de aprendizaje por imitación DAgger, donde el aprendiz tiene acceso local a una distribución experta sobre los estados visitados por la política actual. Esto produce un objetivo simple de entropía cruzada directa que admite un experto en caja negra y cuyo gradiente a nivel de secuencia realiza una rica asignación de crédito al propagar el desacuerdo futuro entre experto y estudiante hacia decisiones anteriores. Mostramos que el RL previo con objetivos de autodestilación basados en KL inversa o Jensen-Shannon no logra garantizar una mejora monótona de la política: incluso cuando el experto tiene una recompensa más alta, sus actualizaciones pueden aumentar la probabilidad de acciones peores. Por el contrario, mostramos que la entropía cruzada directa permite una mejora monótona de la política y disfruta de garantías sobre el arrepentimiento. Además, mostramos que nuestro objetivo optimiza un límite inferior en la probabilidad ponderada por el profesor del éxito, lo que conduce a una mejora en Pass@N. Empíricamente, nuestro enfoque, DistIL, mejora con respecto a RLVR y RL con líneas base de autodestilación en una variedad de dominios: razonamiento científico, codificación y resolución de problemas matemáticos difíciles.

English

Reasoning models have advanced rapidly, but the dominant reinforcement learning from verifiable rewards (RLVR) recipe remains surprisingly narrow: sample many responses and reward each with a single bit indicating whether the final answer is correct. Yet many settings provide rich feedback, including execution traces, tool outputs, expert corrections, and model self-evaluations. We study how to use such feedback through a distributional variant of the classic imitation learning algorithm DAgger, where the learner has local access to an expert distribution on states visited by the current policy. This yields a simple forward cross-entropy objective that admits a blackbox expert and whose sequence-level gradient {conduct rich credit assignment by propagating} future expert-student disagreement back to earlier decisions. We show that prior RL with self-distillation objectives based on reverse KL or Jensen-Shannon fail to guarantee monotonic policy improvement: even when the expert has higher reward, their updates may increase probability on worse actions. In contrast, we show that forward cross-entropy admits monotonic policy improvement and enjoys guarantees on regret. We further show that our objective optimizes a lower bound on teacher-weighted likelihood of success, leading to improved Pass@N. Empirically, our approach, DistIL, improves over RLVR and RL with self-distillation baselines across a variety of domains: scientific reasoning, coding, and solving hard mathematical problems.