Just a list of things I want to change in the code for the sake of simplicity or performance or both:

• Remove hidden file saving for the RL updates in favour of having a data class recorder setup
• JIT compile all loss updates
• Better way to save RL rewards during / after training
• Make the colloid object JIT compilable so that we can JIT or VMAP more functions