This is the codebase of the paper Learning Task Decomposition with Order-Memory Policy Network It contains a version of the craft environment with gym wrappers. Dependency
python3.6
torch==1.5.1
Install locally pip install -e .
If you find the environment or the paper to be useful, please cite
@inproceedings{
lu2021learning,
title={Learning Task Decomposition with Order-Memory Policy Network},
author={Yuchen Lu and Yikang Shen and Siyuan Zhou and Aaron Courville and Joshua B. Tenenbaum and Chuang Gan},
booktitle={International Conference on Learning Representations},
year={2021},
url={https://openreview.net/forum?id=vcopnwZ7bC}
}
import gym_psketch
print(gym_psketch.env_list)
Keyboard interactive mode. Use arrow keys to move and u
to use. Other key triggers done
action.
python scripts/enjoy.py -mode keyboard -env [ENV_NAME]
See rule-based bot.
python scripts/enjoy.py -mode demo -env [ENV_NAME]
More scripts see in scripts
Use main.py
as main entry for both IL and RL
Generate demo
python main.py --mode demo \
--envs <ENVS_NAME> \
--demo_episodes 1500
Run OMPN on unsupervised task information
python main.py --mode IL --arch omstack \
--flagfile ilflagfile \
--nb_slots 3 \
--cuda \
--envs <ENVS_NAME> \
--env_arch noenv
Run OMPN on unsupervised with sketch information
python main.py --mode IL --arch omstack \
--flagfile ilflagfile \
--nb_slots 3 \
--cuda \
--envs <ENVS_NAME> \
--env_arch sketch
Visualize the learned expanding position by
python scripts/analysis.py --model_ckpt PATH_TO_PKL \
--envs makebedfull-v0 --use_demo --episodes 20