Multi-Turn On-Policy Distillation with Prefix Replay

We study on-policy distillation (OPD) for agentic tasks, where an LLM agent interacts with an environment over multiple turns and a student imitates a teacher over these multi-turn interaction histories. Fully online OPD is costly because each update requires fresh student rollouts through the environment and teacher queries at visited histories. We propose Replayed-Prefix On-Policy Distillation (ReOPD), an off-environment alternative that reuses pre-collected teacher trajectories as replayed prefixes: the student acts at selected steps, while the teacher provides dense per-step supervision without executing new environment interactions. We show that multi-turn OPD introduces a prefix trap: making histories more student-on-policy improves relevance to the student, but can query the teacher on histories where its target is unreliable. This creates a two-sided distribution shift between student occupancy and teacher reliability. ReOPD addresses this by treating multi-turn OPD as a reliability-aware prefix distribution design and implements it with a simple step-decaying sampling schedule that emphasizes early, lower-shift prefixes. Across mathematical reasoning with Python and search environments over multiple teacher and student model scales, ReOPD preserves or improves OPD-level accuracy, uses zero tool calls during student training, and is at least 4times faster per rollout than OPD. ReOPD therefore turns expensive agent-environment interaction into a reusable offline resource, enabling scalable distillation across tools, tasks, and environments.

한 줄 요약

ReOPD는 학습 비용을 4배 이상 절감하면서도 정확도를 유지하는 오프라인 멀티턴 온폴리시 디스틸레이션 방법이다.

핵심 기여도

ReOPD는 멀티턴 온폴리시 디스틸레이션(OPD)의 비용 문제를 해결하기 위해 선호된 교사 트레이젝토리를 재사용하는 **prefix replay** 기법을 제안.
멀티턴 OPD에서 발생하는 **prefix trap** 문제를 이론적으로 분석하고, **two-sided distribution shift** 개념을 도입.
**step-decaying sampling schedule**을 통해 초기 단계의 더 신뢰할 수 있는 prefix를 강조하여, 정확도를 유지하면서도 학습 속도를 4배 이상 향상.
Python 수학 추론 및 검색 환경에서 다양한 모델 규모에 걸쳐 **zero tool calls**로 학습 가능.

핵심 아이디어

기존 멀티턴 OPD는 학습 시마다 학습자(student)가 환경과 상호작용하며 새로운 트레이젝토리를 생성해야 하므로 비용이 높다. ReOPD는 이 문제를 해결하기 위해 **선호된 교사 트레이젝토리**(pre-collected teacher trajectories)를 **prefix replay** 형태로 재사용한다. 학습자는 선택된 단계에서만 행동을 생성하고, 나머지 prefix는 교사가 제공한 기록을 사용한다. 이로 인해 환경과의 상호작용 없이도 학습이 가능하다.

하지만 이 방식은 **prefix trap**이라는 문제를 유발한다. 학습자의 행동에 맞춘 prefix가 학습자와 더 유사해질수록(occupancy shift), 교사의 신뢰도가 낮아질 수 있다(teacher reliability shift). 이는 **two-sided distribution shift**로 이어진다. ReOPD는 이 이론적 한계를 인식하고, **reliability-aware prefix distribution design**을 통해 학습자와 교사의 분포 간 균형을 맞춘다. 구체적으로, **step-decaying sampling schedule**을 사용해 초기 단계의 prefix를 더 많이 샘플링함으로써 학습 효율을 높인다.

기술적 접근법

**ReOPD**는 멀티턴 OPD를 **off-environment** 방식으로 구현.
학습자는 **replayed prefix** 상에서 특정 단계에서만 행동을 생성.
교사는 **dense per-step supervision** 제공.
**step-decaying sampling schedule**을 사용하여 초기 단계의 prefix를 더 높은 확률로 샘플링.
학습 과정에서 **zero tool calls**를 유지.
**teacher-forced prefix**와 **student-on-policy roll-in**의 극단적 경우를 포함한 **prefix distribution design**을 학습자와 교사의 신뢰도 간 균형을 고려해 구현.

주요 결과

**Python 수학 추론 및 검색 환경**에서 ReOPD는 OPD 수준의 정확도를 유지하거나 향상.
학습 시 **tool calls 0회** 사용.
**rollout당 4배 이상 빠른 학습 속도** 달성.
다양한 **teacher 및 student 모델 규모**에 걸쳐 성능 일관성 보장.

의의 및 한계

ReOPD는 기존 멀티턴 OPD의 높은 학습 비용 문제를 해결하고, **off-environment** 환경에서도 효과적인 디스틸레이션이 가능하도록 한다. 특히, **prefix trap**이라는 새로운 이론적 통찰을 제시하고, 이를 해결하기 위한 **step-decaying sampling schedule**을 구현한 점에서 학술적 의의가 크다. 또한, **teacher가 RL로 학습된 경우**, 그 트레이젝토리는 이미 수집되어 있어 추가 비용 없이 prefix pool을 구성할 수 있다.

하지만, ReOPD는 **교사의 신뢰도가 낮은 상황**에서는 성능이 저하될 수 있으며, **학습자와 교사 간의 격차**(teacher-student gap)가 클 때만 최대 효과를 얻는다는 한계가 있다. 또한, **prefix replay**는 환경과의 상호작용을 완전히 배제하므로, 일부 상황에서는 실제 환경과의 학습이 여전히 필요할 수 있다.

실용적 활용

ReOPD는 **대규모 LLM 학습**에서 환경과의 상호작용 비용을 줄이며, **다양한 도구와 환경**에서의 학습을 효율적으로 가능하게 한다. 특히, **수학 추론**, **코드 실행**, **검색 엔진 활용** 등 복잡한 멀티턴 작업에 적용 가능하다. 또한, **하나의 학습자 모델이 여러 이질적인 환경**에서 학습해야 할 때, ReOPD는 각 환경의 교사 트레이젝토리를 별도로 수집해 학습자에게 통합 제공할 수 있어 유용하다.