Scripts heavily based on https://github.com/luchris429/purejaxrl - thank you to the authors!