[Feature] Sequence/token level rejection sampling on async training

[huaiyizhao]

Checklist

• This feature will maintain backward compatibility with the current APIs in
  areal/api/. If not, please raise a refactor issue first.

Background

enhancement for handling staleness

Potential Solution

https://richardli.xyz/post/rl-collapse-part3/

Additional Information

verl implementations https://verl.readthedocs.io/en/latest/algo/rollout_corr.html
(Add any relevant context, references, or supporting data here.)

[guozhihao-224]

Hi @huaiyizhao @garrett4wade ,

I'm interested in implementing this feature.
Based on my understanding:

1. The VeRL implementation provides a good reference with multiple KL estimators
   (K1/K2/K3) and aggregation levels (token/seq_sum/seq_mean/seq_max)
2. AReaL already has the behave_imp_weight_cap mechanism, and rejection sampling
   would complement it as an earlier filtering stage
3. The implementation would involve adding new config options to PPOActorConfig
   and a compute_rejection_mask function

Before I start, I have a few questions:

1. Is there a preferred default mode you'd recommend (e.g., seq_mean_k2 for
   Geometric RS)?
2. Should this feature be independent of or integrated with the existing
   behave_imp_weight_cap mechanism?
3. Are there any specific test cases or benchmarks I should target for validation?

I'm happy to submit a draft PR for initial review once I have a basic
implementation ready.

Thanks!

[guozhihao-224]

I see that PR #1060 added a KL divergence estimator for logging. This question could build on that and actually use these estimators to filter samples during training.

[garrett4wade]

Hi @guozhihao-224 , thanks for the follow-ups.

1. AReaL has the k1/k2/k3 KL estimators in areal/utils/data.py. feat: estimators for kl divergence #1060 also provides a great complementary implementation for logging.
2. As far as I know, the core mechanism is to modify the loss_mask according to KL estimations (like here). I think we can just remove the current "behave_imp_weight_cap" parameter and rewrite it in a more general form. We can create a new config object like (just draft, I'm not familar with the actual algorithmic choices)

@dataclass
class RejectionSamplingConfig:
    filter_level: str  # token, seq
    metric: str  
    agg_method: str

@dataclass
class PPOActorConfig(...):
    ...
    rejection_sampling: RejectionSamplingConfig

3. Let's keep the default mode unchanged, i.e., using the naive importance sampling ratio w.r.t. the current policy for token-level clipping.
4. For benchmarks, I don't have a good idea. Since now people focus more on agentic tasks, how about tau2-bench? We can just train a 1.7B model on the test set for simple validation.

[huaiyizhao]

Hi @guozhihao-224 @garrett4wade , I think we can reuse #1060 for the filtering.

From my understanding, rejection sampling should happen in sequence level: the whole rollout sequence should be discarded once a choosen estimator exceeds threshold.
Token-level IS seems biased https://yingru.notion.site/When-Speed-Kills-Stability-Demystifying-RL-Collapse-from-the-Training-Inference-Mismatch-271211a558b7808d8b12d403fd15edda#27b211a558b78099ba48fa8849ab54c8