Possibility to pass a flax (from the flax.linen.activation module) activation function when creating a sbx model, through the policy_kwargs argument.
Motivation
In the current implementation of sbx, users are unable to pass custom activation functions when creating a model. This limitation restricts flexibility and may not suit all users' needs.
🚀 Feature
Possibility to pass a flax (from the
flax.linen.activation
module) activation function when creating asbx
model, through thepolicy_kwargs
argument.Motivation
In the current implementation of
sbx
, users are unable to pass custom activation functions when creating a model. This limitation restricts flexibility and may not suit all users' needs.Pitch
Example:
Idea on how to implement it
Add attribute
activation_fn
to the underlying classes that are composing the policy (like in Critic and Actor int3d/policy.py
)