Support callable entropy schedules in PPO - #660
PhysicistJohn wants to merge 3 commits into
Conversation
Adds two new parameters to `ppo.train()`:
entropy_cost_end: Optional[float] = None
entropy_schedule: str = 'linear' # or 'cosine'
When `entropy_cost_end` is set, the entropy coefficient decays from
`entropy_cost` (start) to `entropy_cost_end` over the full training
budget, following the selected schedule. Default is unchanged —
`entropy_cost_end=None` keeps the existing constant-cost behaviour.
The schedule is computed inside the existing JIT graph using
`training_state.env_steps`, so there is no per-epoch recompilation.
The `if/else` on `entropy_schedule` is a Python-level (static) branch,
so JAX traces only the selected path. `entropy_cost` is now threaded
as a keyword argument through `training_step → sgd_step →
minibatch_step → loss_and_pgrad_fn`, which is already a pure
`*args, **kwargs` pass-through, so no signature conflicts arise.
The current coefficient is logged to TensorBoard as `training/entropy_cost`.
Motivation: a fixed high entropy cost promotes exploration early in
training but prevents the policy from committing to precise actions
later. Annealing entropy over the run allows warm exploration followed
by policy consolidation without requiring separate training phases.
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
|
@PhysicistJohn and brax team, whats the status of this PR, would be a really useful feature! |
|
Hey, thanks for the ping. CLA is signed and passing now, so that part's |
PhysicistJohn
changed the title
|
Updated this to accept an Optax-compatible callable through the existing |
2 participants
Summary
PPO currently accepts a fixed entropy coefficient. Allow
entropy_costto also accept an Optax-compatible callable, supporting linear, cosine, warmup, piecewise, and custom schedules without new trainer flags. The scalar default and existing positional arguments are unchanged.For example:
Pass this as
entropy_costtoppo.train. The schedule receives the total environment-step count, includingaction_repeat, once per rollout batch. All gradient updates on that batch share the coefficient. Counts start at zero for each train call, including parameter restoration.training/entropy_costreports the mean applied coefficient.Testing
bb77b328exactly for default and explicit scalar coefficients with JAX x64 disabled and enabled, including parameters and common non-timing metrics.