Entity Factored RL

This contains code for running the experiments in Policy Architectures for Compositional Generalization in Control.

Imitation Learning

Data Collection

Download the weights for generating BC data from this link. Unzip the file in ./weights. Then generate behavior cloning data by running:

./script/make_data.sh

BC Training

Sweep over different environments and architectures.

python launch_bc.py -m +bc_experiment=big_transformer,deepset,mlp +bc_setup=3p,3s,2s2p

Reinforcement Learning

In general, setup specifies the environment and exploration schedule and experiment specifies the architecture. Some examples:

# 3 push for transformer and MLP (padded for extrapolation eval)
python launch.py -m +experiment=3pdense_fastexp +setup=transformer,padded_mlp seed="range(5)"

# 3 push for deepset uses a faster exploration schedule
python launch.py -m +experiment=3pdense_fastexp +setup=deepset seed="range(5)"

License

The majority of the code for "Entity Factored RL" is licensed under CC-BY-NC, however portions of the project are available under separate license terms.

• Portions based on Hindsight Experience Replay, Stable Baselines 3, and Gym are licensed under the MIT license.