is it possible to add checkpoint saving/restore to this repository?

[gzadigo]

The purejaxrl code is very helpful, but after I finished the training, I hope to save the checkpoints, and found no way to do this

[luchris429]

I think enough people have asked such that I'll prioritize this feature.

For now, I'd recommend looking at this