Imitation Learning fails on key error

I am trying to train a model using the mixed imitation learning and apex baseline defined in apex_il_tree_obs_25.yaml, but I getting a key error when rllib tries to read the expert demonstrations that I have converted to json using saving_experiences.py. I have successfully trained using imitation learning with rllib outside of the flatland competition, so I am fairly sure rllib correctly functions. Is it possible that saving_experience.py is outputting json files in the incorrect format?

I am posting the error here.

2021-03-28 13:35:09,847	ERROR trial_runner.py:519 -- Trial APEX_flatland_sparse_00001: Error processing event.

Traceback (most recent call last):
File “/home/dylan/miniconda3/envs/flatland/lib/python3.6/site-packages/ray/tune/trial_runner.py”, line 467, in _process_trial
result = self.trial_executor.fetch_result(trial)
File “/home/dylan/miniconda3/envs/flatland/lib/python3.6/site-packages/ray/tune/ray_trial_executor.py”, line 431, in fetch_result
result = ray.get(trial_future[0], DEFAULT_GET_TIMEOUT)
File “/home/dylan/miniconda3/envs/flatland/lib/python3.6/site-packages/ray/worker.py”, line 1515, in get
raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(KeyError): ray::APEX.train() (pid=1689497, ip=10.0.0.191)
File “python/ray/_raylet.pyx”, line 463, in ray._raylet.execute_task
File “python/ray/_raylet.pyx”, line 417, in ray._raylet.execute_task.function_executor
File “/home/dylan/miniconda3/envs/flatland/lib/python3.6/site-packages/ray/rllib/agents/trainer.py”, line 495, in train
raise e
File “/home/dylan/miniconda3/envs/flatland/lib/python3.6/site-packages/ray/rllib/agents/trainer.py”, line 484, in train
result = Trainable.train(self)
File “/home/dylan/miniconda3/envs/flatland/lib/python3.6/site-packages/ray/tune/trainable.py”, line 261, in train
result = self._train()
File “/home/dylan/miniconda3/envs/flatland/lib/python3.6/site-packages/ray/rllib/agents/trainer_template.py”, line 142, in _train
return self._train_exec_impl()
File “/home/dylan/miniconda3/envs/flatland/lib/python3.6/site-packages/ray/rllib/agents/trainer_template.py”, line 177, in _train_exec_impl
res = next(self.train_exec_impl)
File “/home/dylan/miniconda3/envs/flatland/lib/python3.6/site-packages/ray/util/iter.py”, line 634, in next
return next(self.built_iterator)
File “/home/dylan/miniconda3/envs/flatland/lib/python3.6/site-packages/ray/util/iter.py”, line 644, in apply_foreach
for item in it:
File “/home/dylan/miniconda3/envs/flatland/lib/python3.6/site-packages/ray/util/iter.py”, line 644, in apply_foreach
for item in it:
File “/home/dylan/miniconda3/envs/flatland/lib/python3.6/site-packages/ray/util/iter.py”, line 685, in apply_filter
for item in it:
File “/home/dylan/miniconda3/envs/flatland/lib/python3.6/site-packages/ray/util/iter.py”, line 685, in apply_filter
for item in it:
File “/home/dylan/miniconda3/envs/flatland/lib/python3.6/site-packages/ray/util/iter.py”, line 644, in apply_foreach
for item in it:
File “/home/dylan/miniconda3/envs/flatland/lib/python3.6/site-packages/ray/util/iter.py”, line 685, in apply_filter
for item in it:
File “/home/dylan/miniconda3/envs/flatland/lib/python3.6/site-packages/ray/util/iter.py”, line 888, in build_union
item = next(it)
File “/home/dylan/miniconda3/envs/flatland/lib/python3.6/site-packages/ray/util/iter.py”, line 634, in next
return next(self.built_iterator)
File “/home/dylan/miniconda3/envs/flatland/lib/python3.6/site-packages/ray/util/iter.py”, line 644, in apply_foreach
for item in it:
File “/home/dylan/miniconda3/envs/flatland/lib/python3.6/site-packages/ray/util/iter.py”, line 644, in apply_foreach
for item in it:
File “/home/dylan/miniconda3/envs/flatland/lib/python3.6/site-packages/ray/util/iter.py”, line 644, in apply_foreach
for item in it:
[Previous line repeated 2 more times]
File “/home/dylan/miniconda3/envs/flatland/lib/python3.6/site-packages/ray/util/iter.py”, line 470, in base_iterator
yield ray.get(obj_id)
ray.exceptions.RayTaskError(KeyError): ray::RolloutWorker.par_iter_next() (pid=1689506, ip=10.0.0.191)
File “python/ray/_raylet.pyx”, line 463, in ray._raylet.execute_task
File “python/ray/_raylet.pyx”, line 417, in ray._raylet.execute_task.function_executor
File “/home/dylan/miniconda3/envs/flatland/lib/python3.6/site-packages/ray/util/iter.py”, line 958, in par_iter_next
return next(self.local_it)
File “/home/dylan/miniconda3/envs/flatland/lib/python3.6/site-packages/ray/rllib/evaluation/rollout_worker.py”, line 260, in gen_rollouts
yield self.sample()
File “/home/dylan/miniconda3/envs/flatland/lib/python3.6/site-packages/ray/rllib/evaluation/rollout_worker.py”, line 525, in sample
batch = batches[0].concat_samples(batches)
File “/home/dylan/miniconda3/envs/flatland/lib/python3.6/site-packages/ray/rllib/policy/sample_batch.py”, line 73, in concat_samples
out[k] = concat_aligned([s[k] for s in samples])
File “/home/dylan/miniconda3/envs/flatland/lib/python3.6/site-packages/ray/rllib/policy/sample_batch.py”, line 73, in
out[k] = concat_aligned([s[k] for s in samples])
File “/home/dylan/miniconda3/envs/flatland/lib/python3.6/site-packages/ray/rllib/policy/sample_batch.py”, line 193, in getitem
return self.data[key]
KeyError: ‘action_logp’

UPDATE: To test this, I used the checkpoint from an APEX agent I trained yesterday to generate new jsons through rllib rather than saving_experiences.py. Using these logs, the mixed APEX and IL agent is successfully able to train. So, I am not sure if saving_experiences.py is outputting jsons formatted for an older version of rllib that is no longer compatible, but if anyone has any info on this, I would be interested.