Add IcePop Masking

[Crystalcareai]

Adds IcePop stabilization from Ring-1T (arXiv:2510.18855) to both ClippedPG and CISPO loss functions.

What it does

Masks out tokens where the engine mismatch ratio π_train/π_infer falls outside configurable bounds [alpha, beta] (default [0.5, 2.0]).

Why hard masking vs soft weighting

We discard tokens entirely rather than downweighting them because:

1. Tokens with high engine mismatch indicate that the training and inference engines disagree significantly on the action probability. The gradient signal from these tokens is fundamentally unreliable.
2. Soft weighting would still propagate some gradient from untrustworthy samples, potentially destabilizing training
3. Binary masking is simpler to implement and reason about, with no additional hyperparameters for the weighting function shape
4. The paper's results show that hard cutoffs at [0.5, 2.0] work well in practice

Configuration

loss_fn:
  icepop_enabled: true
  icepop_alpha: 0.5  # lower bound
  icepop_beta: 2.0   # upper bound

Monitoring

Logs icepop_fraction_masked metric to track what fraction of tokens are being discarded.

---

Note

Adds IcePop-based token masking to ClippedPG and CISPO losses, exposes config knobs, logs masking diagnostics, and provides an example training config.

• Loss functions (rlkit/algorithms/loss_functions.py):
  • IcePop masking: Apply token masks where engine mismatch ratio exp(prev_logprobs - generation_logprobs) is outside [icepop_alpha, icepop_beta].
    • Implemented in both ClippedPGLossFn and CISPOLossFn; affects mask and sequence-level token_mask_for_loss.
    • Logs icepop_fraction_masked metric.
  • Diagnostics/efficiency:
    • Reuse engine_mismatch_ratio in CISPO for sampling_importance_ratio.
• Config types (rlkit/config/rl/loss.py):
  • Add icepop_enabled, icepop_alpha, icepop_beta to ClippedPGLossConfig and CISPOLossConfig.
  • Add token_level_loss to CISPOLossConfig.
• Example config:
  • New examples/configs/rl/reverser/qwen3_0.6B_2gpu_icepop.yaml demonstrating CISPO + IcePop settings and training setup.

^{Written by Cursor Bugbot for commit dbb78e7. This will update automatically on new commits. Configure here.}

[cursor]

Bug: IcePop fraction metric breaks aggregation across microbatches

The icepop_fraction_masked metric is computed as 1.0 - masked_mean(...), which breaks the expected aggregation pattern. Per the code comment, metrics are designed to be summed across microbatches to get the true global value. All other metrics return masked_mean(...) directly, so summing gives the correct global mean. However, 1.0 - masked_mean(...) adds a constant 1.0 per microbatch, causing the aggregated result to be N - global_fraction_not_masked instead of the expected 1 - global_fraction_not_masked. The metric should return the fraction NOT masked directly (without the 1.0 -), letting the caller invert after aggregation if needed.

Additional Locations (1)

• rlkit/algorithms/loss_functions.py#L598-L603