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.
Hi! Currently in the brax agents, there's the ability to pass a
parametric_action_distributionto the network but in themake_networkfunction thedistribution_typearg is fixed to onlynormalortanh_normal. It would be nice to be able to pass the constructor for a custom action distribution directly.