Skip to content

feat(trainer): support configurable entropy regularization in PPO and GRPO actor - #1680

Open
hsusul wants to merge 4 commits into
areal-project:mainfrom
hsusul:feat/entropy-regularization
Open

hsusul wants to merge 4 commits into
areal-project:mainfrom
hsusul:feat/entropy-regularization

Conversation

@hsusul

@hsusul hsusul commented Sep 5, 2026

Copy link
Copy Markdown
Contributor

Description

In large-scale reinforcement learning for reasoning models (e.g. DeepSeek-R1, QwQ, and mathematical reasoning / tool use), policy optimization often suffers from premature entropy collapse, where the policy rapidly concentrates on repetitive degenerate sequences or narrow sub-optimal generation paths. In classical and modern RL (TRL, CleanRL, OpenRLHF, verl), an entropy regularization bonus (-\beta \mathcal{H}(\pi)) is standard practice to encourage token exploration.

Previously in AReaL:

  1. PPOActorConfig lacked an entropy_coeff option.
  2. In grpo_loss_fn (areal/trainer/ppo/actor.py), entropy was unconditionally detached (entropy = entropy.detach()) purely for metric logging, making it impossible to propagate gradients back to model weights even when engines (such as FSDP or Megatron with entropy_requires_grad=True) compute differentiable entropy.

This PR adds configurable entropy regularization:

  • Configuration: Adds entropy_coeff (default 0.0) to PPOActorConfig with non-negative validation in __post_init__.
  • Differentiable Entropy Loss: In grpo_loss_fn, when entropy_coeff > 0.0, entropy is kept attached to the autograd graph, and a masked token-level entropy loss term (\mathcal{L}{\text{entropy}} = -\frac{1}{|\mathcal{M}|} \sum{i \in \mathcal{M}} \mathcal{H}(\pi_\theta(s_i))) is added to the objective scaled by entropy_coeff.
  • Zero-Overhead Detached Default: When entropy_coeff == 0.0, entropy is immediately detached, preserving the zero-autograd-overhead default.
  • Observability: Records entropy_loss in stats_tracker for logging to WandB/SwanLab/TensorBoard.
  • CLI Docs: Regenerated documentation in docs/en/cli_reference.md and docs/zh/cli_reference.md.

Verification Plan

  • Unit test suite in tests/test_ppo_entropy.py:
    • test_ppo_actor_config_entropy_coeff_validation: verifies default, positive acceptance, and rejection of negative values.
    • test_grpo_loss_fn_entropy_bonus_gradient: verifies that when entropy_coeff > 0.0, gradients flow through entropy to model logits.
    • test_grpo_loss_fn_entropy_zero_is_detached: verifies that when entropy_coeff == 0.0, entropy receives no gradients.
    • test_grpo_loss_fn_entropy_bonus_modifies_loss: verifies mathematical correctness of the entropy bonus in the overall loss.
  • Verified existing PPO tests in tests/test_ppo_stats.py (15 passed).
  • Passed all 16 pre-commit hooks (ruff check, ruff format, generate-cli-docs, etc.).

- Use lazy evaluation in testing model path dictionaries to avoid eager downloads on test collection
- Safely handle PackageNotFoundError in version check helper functions
… GRPO actor

- Add entropy_coeff configuration in PPOActorConfig with non-negative validation
- Pass entropy_coeff to grpo_loss_fn in PPOActor._train_batch
- Compute masked mean entropy loss when entropy_coeff > 0.0 and maintain gradient tracking through entropy
- Maintain detached entropy behavior when entropy_coeff == 0.0 for zero autograd overhead
- Log entropy loss metrics to stats_tracker
- Update CLI reference documentation
- Add unit tests in tests/test_ppo_entropy.py verifying entropy gradient flow, detached behavior, and config validation

@Le8r0nJames Le8r0nJames left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks useful to support as an opt-in feature. Have you run any end-to-end training experiments with a positive entropy_coeff? It would be helpful to share the model/backend, coefficient, and how entropy and validation performance compare with the zero-coefficient baseline.

Comment on lines +739 to +742
if entropy_coeff > 0.0:
entropy_loss = -(entropy * loss_mask).sum() / loss_mask.sum().clamp(min=1)
loss = loss + entropy_coeff * entropy_loss
entropy_stat = -entropy.detach() * loss_mask

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we validate the backend requirements when entropy_coeff > 0? With Megatron's enable_chunked_logits=True and the default entropy_requires_grad=False, entropy is non-differentiable upstream, so this term changes the loss value without contributing gradients.

Please reject incompatible configurations or enable a supported differentiable path, and add a test using engine-produced entropy.

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.

2 participants