update train script with arguments

1 job for flatland-paper-baselines-2 in 2 minutes and 41 seconds (queued for 25 seconds)
Name Stage Failure
failed
test Test
  Downloading tensorflow-2.1.0-cp37-cp37m-manylinux2010_x86_64.whl (421.8 MB)

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


CondaEnvException: Pip failed

ERROR: Job failed: command terminated with exit code 1