Mask replay for GRPO

[vx120]

This PR adds sampling replay for selective log-probability computation.

During rollout, the sampler records the token support used at each generation step. During training, sampling replay reconstructs log-probabilities over the same support instead of recomputing them over a potentially different distribution. This better aligns the rollout and training policies.

Results

Reward is comparable with and without replay, with a small improvement when replay is enabled:

Setting	mean train/total_reward
replay off	1.4297
replay on	1.4547

The reward curves show similar total_reward, accuracy_reward, and format_reward trajectories in both settings. Replay preserves reward quality and is slightly better at the final rollout.

Setting	mean importance_ratio_mean	mean logp_diff_mean
replay off	0.99752	-0.00279
replay on	1.00000	-0.00031

importance_ratio_mean measures the average ratio between the training-time and rollout-time probabilities of sampled tokens. The ideal value is 1.0.

logp_diff_mean measures the average difference between training-time and rollout-time token log-probabilities. The ideal value is 0.0.

With replay enabled, the importance ratio stays tightly centered around 1.0, while the log-probability difference remains much closer to 0.0. This reduces rollout/training mismatch and provides a more stable optimization signal, which should be beneficial for longer training runs.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.