Attention-driven Robotic Manipulation (ARM)

Codebase of Q-attention, coarse-to-fine Q-attention, and other variants. Code from the following papers:

• Q-attention: Enabling Efficient Learning for Vision-based Robotic Manipulation (ARM system)
• Coarse-to-Fine Q-attention: Efficient Learning for Visual Robotic Manipulation via Discretisation (C2F-ARM system)
• Coarse-to-Fine Q-attention with Learned Path Ranking (C2F-ARM+LPR system)
• Coarse-to-Fine Q-attention with Tree Expansion

Installation

ARM is trained using the YARR framework. Head to the YARR github page and follow installation instructions.

ARM is evaluated on RLBench 1.1.0. Head to the RLBench github page and follow installation instructions.

Now install project requirements:

pip install -r requirements.txt

Running experiments

Be sure to have RLBench demos saved on your machine before proceeding. To generate demos for a task, go to the tools directory in RLBench (rlbench/tools), and run:

python dataset_generator.py --save_path=/mnt/my/save/dir --tasks=take_lid_off_saucepan --image_size=128,128 \
--renderer=opengl --episodes_per_task=100 --variations=1 --processes=1

Experiments are launched via Hydra. To start training C2F-ARM on the take_lid_off_saucepan task with the default parameters on gpu 0:

python launch.py method=C2FARM rlbench.task=take_lid_off_saucepan rlbench.demo_path=/mnt/my/save/dir framework.gpu=0

To launch C2F-ARM+LPR:

python launch.py method=LPR rlbench.task=take_lid_off_saucepan rlbench.demo_path=/mnt/my/save/dir framework.gpu=0

To launch C2F-ARM+QTE:

python launch.py method=QTE rlbench.task=take_lid_off_saucepan rlbench.demo_path=/mnt/my/save/dir framework.gpu=0