Conversation
- 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
These files are not part of entropy regularization and already landed on main via areal-project#1568.
Le8r0nJames
left a comment
There was a problem hiding this comment.
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.
| 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 |
There was a problem hiding this comment.
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.
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:
PPOActorConfiglacked anentropy_coeffoption.grpo_loss_fn(areal/trainer/ppo/actor.py),entropywas 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 withentropy_requires_grad=True) compute differentiable entropy.This PR adds configurable entropy regularization:
entropy_coeff(default0.0) toPPOActorConfigwith non-negative validation in__post_init__.grpo_loss_fn, whenentropy_coeff > 0.0,entropyis 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 byentropy_coeff.entropy_coeff == 0.0,entropyis immediately detached, preserving the zero-autograd-overhead default.entropy_lossinstats_trackerfor logging to WandB/SwanLab/TensorBoard.docs/en/cli_reference.mdanddocs/zh/cli_reference.md.Verification Plan
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 whenentropy_coeff > 0.0, gradients flow through entropy to model logits.test_grpo_loss_fn_entropy_zero_is_detached: verifies that whenentropy_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.tests/test_ppo_stats.py(15 passed).pre-commithooks (ruff check,ruff format,generate-cli-docs, etc.).