Add tests for the main loop in Rein

[polvalente]

We should have a simple unit test (that can use a mock agent/environment pair) showing that checkpointing works and that all functions are called in the proper order