On-policy のはずが Off-policy になる：LLM 強化学習の rollout mismatchと対策(rollout correction)

ABEJAでデータサイエンティストをしている服部です。

今回はGRPOなどのLLMの強化学習の中で実際に起こりうる少しマイナーだけど重要かもしれない強化学習のRollout時に発生するトピックについての記事です。

過去にもAgenticRLに関しての記事も書いてますので、こちらも興味ある方いましたら読んでいただけると幸いです。

tech-blog.abeja.asia

本記事は何となくでもGRPOなどの手法で何をしているか知っている前提の記事になっています。

上記記事をベースにした内容をW&B様のホワイトペーパーにも寄稿しており、こちらには強化学習の基礎から書かれているのでこちらももし良ければご確認いただければです。

wandb.ai

学習とRolloutでのズレ
実際にズレが起こりうるかを見てみる。
分布のmismatch
mismatchへの対策
実ライブラリでの実装
- NeMo-RL
- verl
まとめ
We Are Hiring!

学習とRolloutでのズレ

GRPOを始めとする最近のLLMでの強化学習では、モデルのPolicyに従ってRolloutを行い、そのRollout結果を元に与えられた報酬を用いてモデルを更新します。

このRolloutというのは学習データ（指示などのインプット）に対しての推論を行うわけですが

GRPO系では同一sampleに対してGroup内での出力数分、推論を行う
一つのRollout出力が、Reasoningも含めると時間がかかる
Tool Useなどの処理時間、multi-step化により時間がかかる

といった背景から、学習時間よりRollout(推論)の時間がボトルネックになりがちです。

そのため実際のシステムでは、

学習: Megatron / FSDP / DeepSpeed
Rollout: vLLM / SG-Lang

のように、学習backendとRollout(推論)backendを分離する構成が一般的です。

しかし、Rollout時のPolicyと、学習時に仮定しているPolicyが一致しないことがあるという課題があります。

このズレは例えば、

FP16 / BF16 / FP8 の違い
attention実装の差
backend実装の差
sampling実装の差
CUDA kernelの差

などによって、同じweightを使っていても tokenの出力確率が完全には一致しないことがあります。

GRPO/PPO系の学習では、Rolloutで得られたtoken列に対して、次のような確率比を使ってPolicyを更新します。

$r_t(\theta)=\frac{\pi_{\theta}(a_t|s_t)}{\pi_{old}(a_t|s_t)}$

ここで重要なのは、分母の $\pi_{old}$ が単なる正規化項ではなく、そのtoken列を実際に生成したPolicyであるという前提を持っている点です。

つまりPPO/GRPOでは、Rolloutで得られたaction $a_t$ は

$a_t \sim \pi_{old}(\cdot|s_t)$

からサンプリングされたものとして扱われます。

しかし、実際のシステムではRollout backendと学習backendが分かれていることが多く、Rolloutを生成したPolicyを $\mu$ 、学習側でold policyとして再計算したPolicyを $\pi_{old}$ とすると、

$\mu(a_t|s_t) \neq \pi_{old}(a_t|s_t)$

となることがあります。

このとき、サンプルは $\mu$ から来ているのに、loss上では $\pi_{old}$ から来たものとして扱ってしまうことになってしまいます。

例えば、あるtokenについてRollout backendでは

$\mu(a_t|s_t)=0.6$

だったのに、学習backendで再計算すると

$\pi_{old}(a_t|s_t)=0.3$

だったとします。

このtokenは、Rollout backend上では比較的出やすいtokenとしてサンプリングされています。しかし学習backend側で見ると、同じtokenが「old policyではかなり出にくかったtoken」として扱われます。

例えば現在のPolicyがそのtokenに $\pi_{\theta}(a_t|s_t)=0.6$ を割り当てていた場合、本来Rollout時の確率 $\mu(a_t|s_t)=0.6$ を基準にすればratioは

$\frac{0.6}{0.6}=1.0$

ですが、学習側の $\pi_{old}(a_t|s_t)=0.3$ を分母に使うと

$\frac{0.6}{0.3}=2.0$

になります。

つまり、PolicyがRollout時から変わっていないようなケースでも、backend mismatchによって「確率が2倍に増えた」ように見えてしまいます。このように、backend間の差がPPO/GRPOのratioに入ると、loss上の更新量やclipのかかり方に影響が出ます。

実際にズレが起こりうるかを見てみる。

ここまで「Rollout時のPolicyと学習時のPolicyがズレることがある」という話をしてきましたが、実際どのくらいズレるのか、手元で測ってみます。

今回は Qwen3-4B を使って、次のような実験をしました。

まず vLLM でいくつかのpromptに対して生成（Rollout）を行い、各tokenについて「vLLMがそのtokenを選んだ確率」を記録します。これが先ほどの $\mu$ にあたります。
次に、生成されたまったく同じtoken列を transformers (HuggingFace) 側にもう一度通して、同じtokenの確率を測り直します。これが学習側が見る $\pi_{old}$ にあたります。
この2つを token ごとに並べて比較します。

weightは同じものを使い、dtypeも両方 bf16 に揃えています。つまりここで出てくる差は、モデルや精度の違いではなく、純粋に vLLM と transformers の実装（attentionやkernelなど）の違いから来るもの、ということになります。実験はA6000上で行っています。

同じtokenでも、確率は一致しない

まず、vLLMが出した確率と transformers が出した確率を、tokenごとに散布図にしてみます。

対角線（赤い点線）の上に乗っていれば「2つのbackendが完全に一致している」状態です。見てのとおり、多くの点は対角線の近くに集まっていますが、そこそこの数の点が対角線から外れています。同じweightで同じtokenを見ているはずなのに、vLLM と transformers で「このtokenを出す確率」がそれなりに違う、というのがまず確認できます。

具体的にどれくらい違うのか、ズレの大きかったtokenをいくつか挙げてみます。

token	vLLM が出した確率	transformers が出した確率
`Sydney`	0.59	0.38
`the`	0.67	0.50
`is`	0.57	0.44
`We`	0.38	0.50

上記の例は「オーストラリアの首都はどこか、なぜSydneyではないのか（原文: What is the capital of Australia, and why is it not Sydney?）」という質問に対して、

The capital of Australia is **Canberra**.

  **Why is it not Sydney?**

  While Sydney is the largest and most well-known city in Australia, it is not
  the capital. The capital was chosen to be **Canberra** for several reasons:

  1. **Balanced Geographical Position**: Sydney
....

といった文章を生成しており、"1. Balanced Geographical Position: Sydney" の Sydney というtokenに対して、vLLMは 0.59 の確率を割り当てているのに、transformers では 0.38 になっています。0.2 以上の差があり、この1 tokenの出力だけでも計算上影響がでそうに思えます。

ズレの大きさの分布を見てみる

とはいえ、全部のtokenがこんなに大きくズレているわけではありません。むしろ逆で、大半のtokenはほとんど一致しています。

ズレの大きさ（確率の差 |Δp|）で分布を見ると、今回のデータ（約3900 token）では次のようになりました。

確率の差が 0.001 未満（ほぼ完全一致）: 81.1%
確率の差が 0.01 以上: 10.9%
確率の差が 0.05 以上: 2.0%
確率の差が 0.10 以上: 0.3%
最大のズレ: 0.215

ヒストグラム（縦軸は対数目盛にしています）を見ると、ゼロ付近に鋭いピークがあって、そこから両側に細く裾が伸びている形になっています。8割以上のtokenはほぼほぼ一致していて、ズレているのは一部のtokenだけです。ただしその一部は、上記の例のように 0.2 近くズレることもある、という分布です。

また、このズレが系列の中で均一に散らばっているわけではなく、特定の位置に固まって現れるのも興味深いですね。1つの出力について、tokenごとのズレを生成順に並べてみると、こんな感じになりました。

ほとんどの位置はゼロ付近に張り付いていて、ところどころ大きく飛び出している箇所があるように見えてます。

今回の範囲で言えそうなこと

今回のデータで見えたのは次のような傾向でした。

vLLM と transformers のズレは、確かに存在する
ただし全tokenが一律に少しずつズレるわけではなく、大半は一致していて、一部のtokenだけが大きくズレる
平均的なズレ（log-ratio の平均）はほぼ 0 に近く、今回の範囲では効いてきそうなのはむしろ「たまに大きくズレるtoken」の方だった

冒頭で触れたような「各tokenが少しずつズレて、それが sequence 全体で積み重なっていく」というイメージとは少し違って、今回見た限りでは「ほとんどのtokenは一致するけれど、たまに大きく外れるtokenが混じる」という印象でした。

ただし、これはあくまで今回試した条件（Qwen3-4B、temperature=1.0、bf16、vLLMとtransformersの組み合わせ）での結果です。モデルやサンプリング設定、backendのバージョンによって傾向は変わりうるので、あくまで一例と思ってもらえるとです。（ハイパラなどのデフォルト値などで予期せぬ差分が発生してる可能性なども捨てきれません）

分布のmismatch

実際にはrolloutとtrainのbackendの差分以外でもこういった分布差は問題として存在しています。

主なところでいうと、今回のもの含めて以下の3つがあります。

推論ライブラリ（今回）
非同期RL の stale policy
classical off-policy RL

2つ目の非同期RLというのは、RLの推論に時間がかかりすぎる問題への対策として最近ホットになってきているものです。通常の同期RLでは、Rollout→Policy更新→ Rolloutという順で更新後の最新のPolicyでRolloutを行います。

これだとPolicy更新中Rolloutを止めることになりGPU効率が落ちるため、非同期RLでは推論効率重視でRolloutするPolicyが最新である保証を捨ててRolloutを行います。（今回の記事の本題ではないためざっくりとした説明です）

この場合も、実際にRolloutを生成したPolicyと学習側が仮定しているPolicyの間に違いが生じます。

3つ目のclassical off-policy RLというのは、従来のOff Policyを前提としたRL全般のことを指します。 LLMでいうと、DPOやその派生手法のように、別モデルや過去モデルから収集したデータを用いて学習するアプローチがこれに該当します。off-policy RLは古くから研究されており、補正手法が提案されてきました。こうした古典的なoff-policyの枠組みも切り離せないものでしょう。

mismatchへの対策

この補正のために使われる代表的な手法が Importance Sampling (IS) です。

ISは、ある分布から得られたサンプルを、別の分布から得られたサンプルであるかのように扱うための補正です。今回の文脈では、Rollout backendで実際にサンプリングに使ったPolicyを $\mu$ 、学習backend側でold policyとして扱うPolicyを $\pi_{old}$ とすると、 $\mu$ から得られたtoken列を $\pi_{old}$ から得られたものとして補正したい、ということになります。

sequence全体で見ると、Rolloutで得られるtoken列 $a_{1:T}$ は $\mu$ によって生成されています。そのため、理論上の重みはsequence全体の生成確率の比になります。

$w_{seq}=\frac{\pi_{old}(a_{1:T}|s)}{\mu(a_{1:T}|s)}$

自己回帰モデルではsequenceの生成確率はtokenごとの条件付き確率の積で書けるため、これは次のように分解できます。

$w_{seq}=\prod_{t=1}^{T}\frac{\pi_{old}(a_t|s_t)}{\mu(a_t|s_t)}$

ここで、各tokenに対応する確率比

$w_t=\frac{\pi_{old}(a_t|s_t)}{\mu(a_t|s_t)}$

が、sequence全体の重みを構成する因子になっています。対数で見ると、sequence全体の重みはtoken単位のlog-ratioの和になります。

$\log w_{seq}=\sum_{t=1}^{T}(\log \pi_{old}(a_t|s_t)-\log \mu(a_t|s_t))$

ここで問題になるのが、このsequence全体の重み $w_{seq}$ をそのまま使うと非常に不安定になりやすいという点です。

先ほどの実験では、token単位のlog-ratioの平均はほぼ0に近い値でした。つまり今回の結果に限れば、「各tokenが少しずつ同じ方向にズレて、それが $1.01^{100} \approx 2.7$ のように膨らむ」というより、ほとんどのtokenは一致しているが、一部のtokenだけ大きくズレるという形でした。

このような分布では、sequence全体でlog-ratioを足し合わせたとき、平均よりも分散や外れ値の影響が問題になります。平均が0付近であっても、足し合わせるtoken数が増えるほどsequenceごとのばらつきは大きくなります。また、log-ratioが大きいtokenが数個混ざるだけで、 $w_{seq}$ は数倍、あるいはその逆方向に大きく振れる可能性があります。

その結果、sequence全体の重み $w_{seq}$ はサンプルごとに大きく上下する高分散な量になり、lossにそのまま掛けると学習が不安定になりやすくなります。

そのため実際には、素朴なsequence-level ISをそのまま使うのではなく、

token-level IS
truncated IS
幾何平均
rejection sampling
IcePOP

などの安定化手法が使われます。

token-level IS

最もシンプルかつ広く使われているアプローチです。 sequence全体でまとめて一つの重みを掛けるのではなく、token単位で重みを計算してloss に掛けます。

$w_t=\frac{\pi_{old}(a_t|s_t)}{\mu(a_t|s_t)}$

ここで $\pi_{old}$ は学習backend側で計算し直したlogprob、 $\mu$ はRollout backendが実際にサンプリングに使ったlogprobです。

実際の実装では、Rolloutで得られたサンプルに対して、学習backend側でforwardを一度走らせて $\pi_{old}$ のlogprobを取得し、Rollout backendが返したlogprob $\mu$ との差分から重みを作ります。これをPPO/GRPOのloss項に掛ける形で、 $L(\theta)=E[w_t r_t(\theta) A_t$ ] のように補正します。

sequence全体の積を取らないため、先ほどの $1.01^{100} \approx 2.7$ のような指数的な発散は起こりません。verlやNeMo-RLなどの主要なRLライブラリで、この形のtoken-level補正が標準的に実装されています。

truncated IS

token-level ISであっても、特定のtokenで $\mu$ の確率が極端に小さい場合、重み $w_t$ が巨大になり学習が不安定になることがあります。これを防ぐために、重みに上限を設けるのが truncated IS です。

$w_t=\min(\frac{\pi_{old}(a_t|s_t)}{\mu(a_t|s_t)},C)$

ここで $C$ は事前に決めた閾値です（例えば $C=10$ など）。 truncationを入れると厳密にはbiasの入った推定量になりますが、varianceが大幅に下がるため実用上は安定して学習が進むケースが多いです。 backend間のズレは通常そこまで大きくないため、ほとんどの token では truncation は発動せず、外れ値的に大きく乖離した token だけ抑え込むイメージになります。

なお、これはPPO/GRPOで元々ある $r_t(\theta)$ へのclip（policy drift補正側のclip）とは別物で、あくまでbackend mismatch補正用の重み $w_t$ に対する処理です。実装上は両方が同時に効いている形になります。 verl の Rollout Correction や、TIS (Truncated Importance Sampling) として論文化されているのもこのアプローチです。

幾何平均

token-level ISは安定性が高い一方で、token単位で独立に重みを掛けるため「sequenceとしての一貫した重み付け」にはなっていません。

本来、sequenceとして生成されたサンプルに対する重みは

$w_{seq}=\prod_{t=1}^{T}\frac{\pi_{old}(a_t|s_t)}{\mu(a_t|s_t)}$

ですが、これをそのまま使うとsequence長 $T$ に対して指数的に発散します。

そこで、積の代わりに幾何平均を取って一つのsequence-levelの重みにするアプローチがあります。

$w_{seq}=\left(\prod_{t=1}^{T}\frac{\pi_{old}(a_t|s_t)}{\mu(a_t|s_t)}\right)^{1/T}$

対数の世界で見ると、

$\log w_{seq}=\frac{1}{T}\sum_{t=1}^{T}(\log \pi_{old}(a_t|s_t)-\log \mu(a_t|s_t))$

つまり、token単位のlog-ratioの平均を取っていることになります。

こうすることで、sequence長 $T$ による指数的なスケーリングを抑えつつ、sequence全体としての $\pi_{old}$ と $\mu$ のズレを一つの重みに集約できます。

なお、類似の幾何平均をsequence-level ratioに使うアプローチとして GSPO (Group Sequence Policy Optimization) や GMPO (Geometric-Mean Policy Optimization) が知られていますが、これらはここで議論しているbackend mismatchの補正ではなく、学習中の $\pi_{\theta}$ とold policy $\pi_{\theta_{old}}$ の比をsequence-levelで取るための安定化を目的としたものです。問題設定は別ですが、「token単位のlog-ratioを平均する」という手法のアイデアは共通しています。

rejection sampling

これは重みで補正するのではなく、「 $\mu$ と $\pi_{old}$ のズレが大きすぎるサンプルは学習から外す」という割り切ったアプローチです。具体的には、sample（あるいはtoken）ごとに

$w_t=\frac{\pi_{old}(a_t|s_t)}{\mu(a_t|s_t)}$

を計算し、これが事前に決めた範囲 $[1/C,C$ ] から外れる場合には、そのサンプルをloss計算から除外します。

メリットとしては、

実装が単純
巨大な重みによる勾配の暴走が原理的に起きない
残ったサンプルだけで見ると、Rollout backendと学習backendがほぼ一致している状況で学習できる

といった点があります。

一方で、

一定割合のサンプルを捨てることになるためデータ効率が落ちる
特定のtoken（例えば低確率token、長いsequenceの後半など）でズレが系統的に大きくなる場合、それらが選択的に捨てられることでデータ分布に偏りが生じる可能性がある

というデメリットもあります。 verlの Rollout Correction でも、Importance Sampling weight と並んで Rejection Sampling が用意されており、両者を組み合わせて使うことができます。

実ライブラリでの実装

実際にRLを行えるライブラリで、このあたりどう対応されているかを見ていきます。

NeMo-RL

NeMo-RLでは、loss関数（ClippedPGLossConfig）の設定で対応がされています。（ブログ執筆時点）

github.com

    # --- Importance sampling correction ---
    # Async GRPO requires importance sampling correction enabled
    # Set to true when async_grpo.enabled is true
    use_importance_sampling_correction: bool = False
    # --- Truncated importance sampling ---
    # Type of truncated importance sampling:
    #   "tis"          – clamp IS weights to max
    #   "icepop"       – zero out tokens with IS weight outside [min, max]
    #   "seq-mask-tis" – zero out sequences by geometric-mean IS ratio, non-truncated token IS correction
    truncated_importance_sampling_type: Optional[str] = None
    truncated_importance_sampling_ratio: Optional[float] = None
    # Lower bound for ICE-POP / seq-mask-tis filtering
    truncated_importance_sampling_ratio_min: Optional[float] = None

use_importance_sampling_correction というフラグが用意されています。

また、 truncated_importance_sampling_type や truncated_importance_sampling_ratio といったパラメータも存在しており、 truncated_importance_sampling_type は tis(Truncated Importance Sampling), icepop, seq-mask-tis(幾何平均でのsequnce単位でのIS) が用意されているようです。

verl

verlでは、この一連の補正が Rollout Correction という機能としてまとめられており、importance sampling (IS) と rejection sampling (RS) を組み合わせて使えるようになっています。（ブログ執筆時点）

verl.readthedocs.io

主な設定パラメータは以下の通りです。

rollout_is : IS補正のレベル（null / token / sequence）
rollout_is_threshold : IS重みの閾値（デフォルト 2.0。"0.5_5.0" のように指定するとIcePop）
rollout_rs : rejection sampling のモード（token_k1 / seq_sum_k1 / seq_mean_k1 など）
rollout_rs_threshold : rejection sampling の閾値
bypass_mode : πold の計算を省略し πold = π_rollout とする

rollout_is で token-level か sequence-level かを選び、rollout_rs を併用するかどうかで truncated IS だけにするか rejection sampling も効かせるかを切り替えます。よく使う組み合わせは decoupled_token_is（token-level IS）や decoupled_seq_is（log-ratio平均で集約する sequence-level IS）などの preset として用意されており、ここまで紹介した「token-level IS」「truncated IS」「幾何平均的なsequence-level補正」「rejection sampling」が設定の組み合わせで一通り表現できます。

まとめ

というわけで今回はrolloutでのズレによる問題をとりあげました。バックエンドの違いによる差分という直近のLLM学習の課題のところから、On-policy/Off-policyという従来の課題まで繋がっている感じがして個人的には調べながら面白かったです。

これ自体がどこまで学習に影響を与えるかまでは調べられていませんが、今後気にしていきたいと思っています。

We Are Hiring!

ABEJAは、テクノロジーの社会実装に取り組んでいます。技術はもちろん、技術をどのようにして社会やビジネスに組み込んでいくかを考えるのが好きな方は、下記採用ページからエントリーください！（新卒の方やインターンシップのエントリーもお待ちしております！） careers.abejainc.com

特に下記ポジションの募集を強化しています！ぜひ御覧ください！

トランスフォーメーション領域：データサイエンティスト

トランスフォーメーション領域：データサイエンティスト（ミドル）

トランスフォーメーション領域：データサイエンティスト（シニア）

ABEJA Tech Blog

中の人の興味のある情報を発信していきます

On-policy のはずが Off-policy になる：LLM 強化学習の rollout mismatchと対策(rollout correction)

学習とRolloutでのズレ