update imitation trainer to save checkpoints and batch obs across agents

1 job for flatland-paper-baselines-2 in 6 minutes and 55 seconds (queued for 2 seconds)
Name Stage Failure
test Test There has been a runner system failure, please try again
  Downloading tensorflow-2.1.0-cp37-cp37m-manylinux2010_x86_64.whl (421.8 MB)

Pip subprocess error:
/tmp/tmpP1LbTK: line 3: 319 Killed /opt/conda/envs/flatland-baseline-cpu-env/bin/python -m pip install -U -r /builds/flatland/neurips2020-flatland-baselines/tmpwFDOwa.requirements.txt

CondaEnvException: Pip failed

ERROR: Job failed (system failure): Internal error occurred: error executing command in container: error during connect: Get http://%!F(MISSING)var%!F(MISSING)run%!F(MISSING)docker.sock/v1.38/exec/55b6b68ebae07c0a553959b562712904d26b8ec8fb91b7c9f0c58e8aec0bd0f3/json: read unix @->/var/run/docker.sock: read: connection reset by peer