Imitation Learning algorithms and Co-training for Mobile ALOHA (velocity added)

This repo contains the implementation of ACT, Diffusion Policy and VINN, together with 2 simulated environments: Transfer Cube and Bimanual Insertion. You can train and evaluate them in sim or real. For real, you would also need to install Mobile ALOHA. This repo is forked from the ACT repo. The main updates/differences to the original repo are as follow;

• added all joint velocities to the state that improves successful rate
• added some homebrewed performance evaluations (plots and vides)
• fixed some dependendy issues and possible troubleshooting guide.

You can find all scripted/human demo for simulated environments here.

git clone --recursive git@github.com:inkyusa/act_private.git

• imitate_episodes.py Train and Evaluate ACT
• policy.py An adaptor for ACT policy
• detr Model definitions of ACT, modified from DETR
• sim_env.py Mujoco + DM_Control environments with joint space control
• ee_sim_env.py Mujoco + DM_Control environments with EE space control
• scripted_policy.py Scripted policies for sim environments
• constants.py Constants shared across files
• utils.py Utils such as data loading and helper functions
• visualize_episodes.py Save videos from a .hdf5 dataset

It is important to exactly match the below software package version in order to correctly make sure simulation dataset generation.

    conda create -n aloha python=3.8.10
    conda activate aloha
    pip install torch==2.0.1+cu118 torchvision==0.15.2+cu118 torchaudio==2.0.2+cu118 --extra-index-url https://download.pytorch.org/whl/cu118
    pip install mujoco==2.3.7 dm_env==1.6 dm_control==1.0.14
    pip install pyquaternion pyyaml rospkg pexpect opencv-python matplotlib einops packaging h5py ipython wandb diffusers
    
    cd act_private/detr
    pip install -e .
    cd act_private/robomimic
    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. After running these scripts, you should be able to check Success: 50 / 50. Otherwise, check the installed python packages (e.g., Mujoco, dm_env, and dm_control)

To generated 50 episodes of transfer_cube_scripted data, run:

    python3 record_sim_episodes.py --task_name sim_transfer_cube_scripted --dataset_dir data/sim/sim_transfer_cube_scripted_top_angle_cams --num_episodes 50

This will create data/sim/sim_transfer_cube_scripted_top_angle_cams folder under the root repository.

python3 record_sim_episodes.py --task_name sim_insertion_scripted --dataset_dir data/sim/sim_insertion_scripted_top_angle_cams --num_episodes 50

Adding the flag --onscreen_render to see real-time rendering. To visualize i episode after it is collected, run

    python3 visualize_sim_episodes.py --dataset_dir data/sim/sim_transfer_cube_scripted_top_angle_cams/ --episode_idx 0

or this will visualise all episodes

for i in {0..50}; do python3 visualize_sim_episodes.py --dataset_dir data/sim/sim_transfer_cube_scripted_top_angle_cams --episode_idx $i; done

    # Transfer Cube task
    python3 imitate_episodes.py --task_name sim_transfer_cube_scripted --ckpt_dir ckpt/sim_transfer_cube_scripted_top_angle_cams --policy_class ACT --kl_weight 10 --chunk_size 100 --hidden_dim 512 --batch_size 16 --dim_feedforward 3200 --num_steps 20000  --lr 1e-5 --seed 0 --save_every 1000

The model should perform above 90% success after 15k steps as shown below

Below demonstrates the evaluation of the best performed policy that we trained from above in various views.

    # Insertion task
    python3 imitate_episodes.py --task_name sim_insertion_scripted --ckpt_dir ckpt/sim_insertion_scripted_top_angle_cams --policy_class ACT --kl_weight 10 --chunk_size 100 --hidden_dim 512 --batch_size 16 --dim_feedforward 3200 --num_steps 20000  --lr 1e-5 --seed 0 --save_every 1000

For this task, the model achieved 80% success after 17k iterations.

Below demonstrates the evaluation of the best performed policy that we trained from above.

    # Transfer Cube task
    python3 imitate_episodes.py --task_name sim_transfer_cube_human --ckpt_dir ckpt/sim_transfer_cube_human --policy_class ACT --kl_weight 10 --chunk_size 50 --hidden_dim 512 --batch_size 16 --dim_feedforward 3200 --num_steps 26500 --lr 1e-5 --seed 0 --save_every 1000

For this task, the model achieved 90% success after 26k iterations.

Below demonstrates the evaluation of the best performed policy that we trained from above.

(Success)

(Fail)

    # Transfer Cube task
    python3 imitate_episodes.py --task_name sim_insertion_human --ckpt_dir ckpt/sim_insertion_human --policy_class ACT --kl_weight 10 --chunk_size 100 --hidden_dim 512 --batch_size 16 --dim_feedforward 3200 --num_steps 26500 --lr 1e-5 --seed 0 --save_every 1000

For this task, the model achieved 40% best success after 30k iterations.

Below demonstrates the evaluation of the best performed policy that we trained from above.

(Success)

(Fail)

The model trained with this dataset performs the worst among above 4 tasks. There may exist some room to improve the performance.

Note that you downloaded simulated environments here and used that dataset for training (only has top view camera so that need to change this in constants.py file).

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 <ckpt_dir> for each rollout. You can also add --onscreen_render to see real-time rendering during evaluation.

One can download and test policies that we trained above for the cube transfer and insertion tasks.

• Sim cube transfer task (scripted) ckpt
• Sim insertion task (scripted) ckpt
• Sim cube transfer task (human) ckpt
• Sim insertion task (human) ckpt

After downloading these ckpts, if you want to evaluate them, just place them under ckpt folder and specify the path for --ckpt_dir argument. Below shows an example command to do this.

python3 imitate_episodes.py --task_name sim_transfer_cube_scripted --ckpt_dir ckpt/sim_transfer_cube_scripted_top_angle_cams --policy_class ACT --kl_weight 10 --chunk_size 100 --hidden_dim 512 --batch_size 16 --dim_feedforward 3200 --num_steps 20000  --lr 1e-5 --seed 0 --save_every 1000 --eval

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.

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.

Note: Mobile-aloha github https://github.com/MarkFzp/mobile-aloha

• Error loading /home/user/workspace/act_private/data/sim/sim_insertion_scripted/episode_32.hdf5 in __getitem__ -> check SIM_TASK_CONFIGS in constants.py file. dataset_dir, camera_names set properly as of the recorded dataset?
• Poor successful rate. check the installed python packages (e.g., Mujoco, dm_env, and dm_control). see here MarkFzp/act-plus-plus#12
• It appears to be challenging to learn from images taken by cameras on the left and right wrists. For instance, model training was unsuccessful (with a low success rate of about 10%) using 3 cameras [top, left wrist, right wrist], whereas it achieved a 100% success rate with two cameras [top, angle]. The positions and orientations of the cameras on the left and right wrists change with every iteration, whereas the 'top' and 'angle' cameras are ego-centric and static.