New: ACT tuning tips
TL;DR: if your ACT policy is jerky or pauses in the middle of an episode, just train for longer! Success rate and smoothness can improve way after loss plateaus.
Project Website: https://tonyzhaozh.github.io/aloha/
This repo contains the implementation of ACT, together with 2 simulated environments: Transfer Cube and Bimanual Insertion. You can train and evaluate ACT in sim or real. For real, you would also need to install ALOHA.
You can find all scripted/human demo for simulated environments here.
imitate_episodes.py
Train and Evaluate ACTpolicy.py
An adaptor for ACT policydetr
Model definitions of ACT, modified from DETRsim_env.py
Mujoco + DM_Control environments with joint space controlee_sim_env.py
Mujoco + DM_Control environments with EE space controlscripted_policy.py
Scripted policies for sim environmentsconstants.py
Constants shared across filesutils.py
Utils such as data loading and helper functionsvisualize_episodes.py
Save videos from a .hdf5 dataset
conda create -n aloha python=3.8
conda activate aloha
pip install torchvision
pip install torch
pip install pyquaternion
pip install pyyaml
pip install rospkg
pip install pexpect
pip install mujoco
pip install dm_control
pip install opencv-python
pip install matplotlib
pip install einops
pip install packaging
pip install h5py
pip install ipython
cd act/detr && pip install -e .
To set up a new terminal, run:
conda activate aloha
cd <path to act repo>
We use sim_transfer_cube_scripted
task in the examples below. Another option is sim_insertion_scripted
.
To generated 50 episodes of scripted data, run:
python3 record_sim_episodes.py \
--task_name sim_transfer_cube_scripted \
--dataset_dir <data save dir> \
--num_episodes 50
To can add the flag --onscreen_render
to see real-time rendering.
To visualize the episode after it is collected, run
python3 visualize_episodes.py --dataset_dir <data save dir> --episode_idx 0
To train ACT:
# UPGM R2D2 data
python3 act/imitate_episodes.py \
--log_dir act/checkpoints \
--policy_class ACT --kl_weight 10 --chunk_size 15 --hidden_dim 512 --batch_size 64 --dim_feedforward 3200 \
--num_epochs 20000 --lr 1e-5 \
--seed 0 \
--data_dir=R2D2/data_1_2023-04-26_128-demos-5-objects,R2D2/data_2_2023-05-02_50-demos-red-cube-fixed-robot-pose,R2D2/data_3_2023-05-02_50-demos-blue-and-white-milk-carton-fixed-robot-pose,R2D2/data_4_2023-05-03_56-demos-2-objects,R2D2/data_5_2023-05-10_76-demos-3-objects,R2D2/data_6_2023-05-14_136-demos-5-objects,R2D2/data_7_2023-05-15_152-demos-5-objects
Add --checkpoint_dir
and --checkpoint_epoch
to load a checkpoint. Add --load_optimizer=True
if you want to load the optimizer state.
To evaluate the policy, run the same command but add --eval
. This loads the best validation checkpoint.
The success rate should be around 90% for transfer cube, and around 50% for insertion.
To enable temporal ensembling, add flag --temporal_agg
.
Videos will be saved to <log_dir>
for each rollout.
You can also add --onscreen_render
to see real-time rendering during evaluation.
For real-world data where things can be harder to model, train for at least 5000 epochs or 3-4 times the length after the loss has plateaued. Please refer to tuning tips for more info.