This repository contains a re-adapatation of Action Chunking Transformer that works for this low-cost robot design (250$).
We are sharing the repo so anyone (non-experts included) can train a robot policy after a few teleoperated demonstraions.
The sorting task in the video was trained with less than 30 demonstrations on an RTX 3080 and took less than 30min.
trained.mp4
Create conda environment
conda create --name act python=3.9
conda activate act
Install torch (for reference we add the versions we are using)
conda install pytorch==1.13.1 torchvision==0.14.1
You can now install the requirements:
pip install -r requirements.txt
Go to TASK_CONFIG in config/config.py and change the paths of the ports that connect leader and follower robots to your computer.
You will also need to connect a camera to your computer and point it towards the robot while collecting the data via teleoperation. You can change the camera port in the config (set to 0 by default). It's important the camera doesn't move otherwise evaluation of the policy is likely to fail.
In order to collect data simply run:
python record_episodes.py --task sort
You can define the name of the task you are doing and the episodes will be stored at data/<task>. You can also select how many episodes to collect when running the script by passing the argument --num_episodes 1 (set to 1 by default). After getting a hold of it you can easily do 20 tasks in a row.
Turn on the volume of your pc-- data for each episode will be recorded after you hear "Go" and it will stop when you hear "Stop".
We slightly re-adapt Action Chunking Tranfosrmer to account for our setup. To start training simply run:
python train.py --task sort
The policy will be saved in checkpoints/<task>.
Make sure to keep the same setup while you were collecting the data. To evaluate the policy simply run:
python evaluate.py --task sort