[Feature Request] Custom action distribution in networks?

Hi! Currently in the brax agents, there's the ability to pass a `parametric_action_distribution` to the network but in the `make_network` function the `distribution_type` arg is fixed to only `normal` or `tanh_normal`. It would be nice to be able to pass the constructor for a custom action distribution directly.