Optimizer state isn't preserved across runs

[corbt]

I often have to restart a run, either to fix something in my reward function, in response to an OOM or crash that broke training, etc. When I do, by restarting the training process the optimizer state is thrown away. I’m worried that this might lead to worse performance than just letting a run go all the way through. Is it easy to save the optimizer state along with the weights so we can truly resume as if nothing happened?