brax icon indicating copy to clipboard operation
brax copied to clipboard

policy_params_fn

Open jamesheald opened this issue 4 months ago • 0 comments

Brax PPO can be passed a policy_params_fn:

policy_params_fn: a user-defined callback function that can be used for saving custom policy checkpoints or creating policy rollouts and videos

I personally find this feature useful for rendering throughout training.

I've added this functionality to BRAX SAC.

jamesheald avatar Sep 09 '25 13:09 jamesheald