allow saving env status into training checkpoints

[le-horizon]

Currently environment statuses like teacher's internal status are not saved in model checkpoints during training. In settings like curriculum training, if we stopped training in the middle, and resume, everything starts from simpliest level again.

If we can save and load env status into model checkpoints, we can resume training at the appropriate level of the curriculum.

Good to have as we introduce teacher tasks in different levels.