Оптимизация целевой политики

Аннотация

В обучении с подкреплением (RL) при заданном промпте мы сэмплируем набор завершений из модели и оцениваем их. Возникают два вопроса: какие завершения должны получить большую вероятностную массу и как должны измениться параметры, чтобы реализовать это изменение? Стандартные методы policy-gradient отвечают на оба вопроса одновременно, поэтому обновление может привести к перерегулированию или недорегулированию в зависимости от скорости обучения, отсечения (clipping) и других настроек оптимизатора. Мы представляем Target Policy Optimization (TPO), который разделяет эти два вопроса. Для заданных оцененных завершений TPO строит целевое распределение q_i ∝ p_i^{старое} exp(u_i) и подгоняет политику к нему с помощью перекрестной энтропии. Градиент потерь по логитам сэмплированных завершений равен p^θ - q, который обращается в ноль, когда политика совпадает с целевой. На табличных бандитах, задачах с трансформерными последовательностями и RLVR для LLM с миллиардами параметров TPO показывает результаты, сопоставимые с PG, PPO, GRPO и DG на простых задачах, и существенно превосходит их в условиях разреженного вознаграждения. Код доступен по адресу https://github.com/JeanKaddour/tpo.

English

In RL, given a prompt, we sample a group of completions from a model and score them. Two questions follow: which completions should gain probability mass, and how should the parameters move to realize that change? Standard policy-gradient methods answer both at once, so the update can overshoot or undershoot depending on the learning rate, clipping, and other optimizer choices. We introduce Target Policy Optimization (TPO), which separates the two questions. Given scored completions, TPO constructs a target distribution q_i propto p_i^{,old} exp(u_i) and fits the policy to it by cross-entropy. The loss gradient on sampled-completion logits is p^θ- q, which vanishes once the policy matches the target. On tabular bandits, transformer sequence tasks, and billion-parameter LLM RLVR, TPO matches PG, PPO, GRPO, and DG on easy tasks and substantially outperforms them under sparse reward. Code is available at https://github.com/JeanKaddour/tpo.