Off-policy training

[Lucaweihs]

Problem

The ADVISOR code-base has support for interleaving off-policy updates (from an arbitrary pytorch dataset and with arbitrary losses) with on-policy updates. It would be great to have similar capabilities here. In particular, we should be able to:

1. Define a pipeline stage so that it performs some fixed number of off-policy updates.
2. Allow for off-policy updates to be interleaved with on-policy updates.

Solution

This requires:

• A way to specify how this training will occur (e.g. defining off-policy losses and updating the pipeline stage to allow for these types of losses + dataset).
• Updating the runner/light_engine to implement these types of updates.

Possible issues:

• Currently we log based on the number of rollout steps. This will not be reasonable if we allow a stage to do purely off-policy updates.

Dependencies

None

[Lucaweihs]

@jordis-ai2 @klemenkotar -- it might be good to have some discussion about this issue before starting on it to see what our various needs are (e.g. does ALFRED require off-policy training in a certain way?)

[jordis-ai2]

I'll try to port https://github.com/allenai/hex-embodied-rl/blob/advisor_into_master/projects/babyai_baselines/experiments/go_to_local/pure_offpolicy.py before release...

[jordis-ai2]

From Luca:
to run this code you'll want to first run

python extensions/rl_babyai/babyai_scripts/download_babyai_expert_demos.py GoToLocal # Downloads the data
python extensions/rl_babyai/babyai_scripts/truncate_expert_demos.py # Makes a small version of the IL dataset you can work with locally

and then
python go_to_local.pure_offpolicy --experiment_base projects/babyai_baselines/experiments

[Lucaweihs]

Feel free to make improvements as you see them: the current implementation is far from perfect so I would not be unhappy if you changed the API. Also of note: off-policy might be a bit tricky when considering the distributed setting (as you'd presumably need some way of partitioning the off-policy dataset into different chunks to be used by the various processes). Not sure what the best solution for this is.