Skip to content

feat: add QAE and adaptive PPO clipping - #89

Open
nightlessbaron wants to merge 1 commit into
prodfrom
feature/qae-adaptive-clipping
Open

nightlessbaron wants to merge 1 commit into
prodfrom
feature/qae-adaptive-clipping

Conversation

@nightlessbaron

Copy link
Copy Markdown

Summary

  • add Quantile Advantage Estimation for GRPO/GSPO via --qae-quantile
  • use the right-continuous empirical quantile (order statistic ceil(K * group_size)) while preserving MILES' existing standard-deviation normalization
  • add the MAI-Thinking-1 entropy-feedback controller for the PPO upper clipping bound
  • update the controller once per optimizer step using importance-weighted token entropy
  • support both Megatron and experimental FSDP backends
  • persist adaptive-clipping state alongside checkpoints
  • add argument validation, focused tests, and an example recipe

The implemented clipping interval is:

[1 - eps_clip, 1 / (1 - eps_clip) + k]

with:

k <- clip(k + step_size * sign(target_entropy - estimated_entropy), 0, k_max)

Example

python train.py \
  ... \
  --advantage-estimator grpo \
  --qae-quantile 0.4 \
  --eps-clip 0.6 \
  --use-adaptive-clip \
  --adaptive-clip-target-entropy 0.3 \
  --adaptive-clip-step-size 0.25 \
  --adaptive-clip-max-relaxation 2.5 \
  --entropy-coef 0

Metrics

  • train/adaptive_clip_entropy
  • train/adaptive_clip_relaxation
  • train/adaptive_clip_high

Validation

  • python -m compileall -q miles tests
  • git diff --check
  • clean patch application against prod

Focused unit tests are included. They were not executed in the construction environment because PyTorch and pytest were unavailable.

References

Add quantile reward baselines for GRPO/GSPO and the MAI-Thinking-1 entropy-feedback controller for the PPO upper clipping bound. Persist controller state across checkpoints and support Megatron and FSDP training.
@nightlessbaron
nightlessbaron requested a review from a team September 16, 2026 21:32
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant