Brax is imported when we import qdax.types

[Lookatator]

As suggested by @manon-but-yes, it would be great to make QDax types independent from external libraries as much as possible.

https://github.com/adaptive-intelligent-robotics/QDax/blob/54fae67977f6cb300cee8d5b18e965573f16b913/qdax/types.py#L20C1-L20C38