brax
brax copied to clipboard
policy_params_fn
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.