Save and load policies from a file

[vincekurtz]

The MJX tutorial has a partial example, but it's incomplete, as it relies on a make_inference_fn that was created during training and never saved. We should have some utilities to make saving and loading easy, particularly for non-default policy architectures.

[alberthli]

Whoever picks this up should start from here: https://flax.readthedocs.io/en/latest/api_reference/flax.serialization.html