Skip to content

Latest commit

 

History

History
106 lines (89 loc) · 5.59 KB

README.md

File metadata and controls

106 lines (89 loc) · 5.59 KB

PoC of ACT imitation learning

Example

  • We train the policy that pick the small stapler box and place that in the other black box.
  • The position of the box can be random but within the range of robot motion.
  • The following video shows the ACT policy employed to controlling koch-v1-1.
  • The box is placed at the random position in each pick-and-place trials.
  • It takes about 30 minutes to complete the 500 epochs training of this policy.
camera1.mp4
camera2.mp4

  • The ACT policy in this video are trained in the following conditions
    GPU Intel Core i7
    GPU NVIDIA GeForce RTX 2070
    Size of memory 32 GB
    Number of epochs 500
    Training sample size 60
    Batch size 8
    Number of VAE encoder layers 4
    Number of encoder layers 4
    Number of decoder layers 5
    Backbone ResNet18
    Number of hidenn dimension 512
    Number of feedforward MLP dimension 3027
    Episode length 1000
    Image size 640 × 480
    Number of camera 2

About this repository

  • If you only want to know usage or installation, please read Usage or Installation.
  • In this section, we describe details of the implementations of this repository.

What contents contained

  • Re-implementation of action chunking transformer (ACT) which is originally developed by Tony Z. Zhao.
  • The python scripts for model training, teleoperation, and model evaluation for real the robot arm.
  • Robot client class for low cost robot arm koch-v1-1.
  • Dynamixel client for DIY robot arms which is compatible with XL430-W250 and XL330-M288-T. You can use this client for any low cost robot arms which is composed of these two types of dynamixel motors.

Re-implementation of ACT

  • We re-implemented the original action chunking transformer from scratch.
  • The original ACT are using some utils and transformer architectures which is partialy dapted from detr.
  • We use pytorch official implementation of transformer modules.
  • We don't use the positional encoding of original ACT. Alternately, simple 1-dimensional sinusoidal positional encoding is employed. This is same positional encoding as the one orignally propose in "Attention is all you need".

About python scripts

  • train.py
    • This script is responsible for training ACT.
    • By default, train_dataset and test_dataset are expected to include all train and test dataset.
  • teleoperation_and_data_collection.py
    • This script is responsible for teloperation.
    • The follower arm are controlled to synchronize with the joint angles of the leader arm.
    • So, by directly controlling the leader arm, you can see the follower arm follow the human demonstrations.
  • evaluate_policy.py
    • This script is responsible for evaluation of trained policy.
    • By giving specific checkpoint path, you can deploy the policy you trained to the real robot.

Robot client for koch-v1-1

  • We prepared robot client base which is abstract class for the client controlling robot arms.
  • Because this client prepares some basic methods that is common for any robot arms, you can customize this class even for the robot arms which is composed of actuators other than DYNAMIXEL.
  • However, this repository contains only the specific implementaion koch-v1-1.
  • If want know more details, please refer to robot_client.py.
  • And, we prepare a few examples of usage of the dynamixel robot client class. Please refer to examples/robot_client

Dynamixel client for DIY robot arms

Installation

  • Make sure you have already installed poetry for package manager.
  • Run the followeing commands to install all dependencies in the root of the repository:
poetry install

Usage

How we get the robot arms ?

  • This repository assumes that you already have follower and leader arms of koch-v1-1. Please refer to koch-v1-1 to prepare your own robot arms.

Collection of the data by teleoperation

  • To collect the data by teloperation, run
python teleoperation_and_data_collection.py --dataset_dir ./train_dataset --initial_episode_id 0

ACT training

  • To train ACT from scratch, run
python train.py --num_epochs 10000 --train_dataset_dir ./train_dataset --test_dataset_dir ./test_dataset --num_episodes_train 30

Evaluate your policy

  • To evaluate trained policym run
python evaluate_policy.py --checkpoint <checkpoints_path>

Acknowledgement

I appreciate these refenrence implementations!

  • We used the repository ACT as a reference implementation for teleoperation of koch-v1-1.
  • We adapted model arthictectures and hyper parameters for training and inferecne from original ACT repository.
  • We used the low cost robot arm koch-v1-1 as an actual robot for evaluating policy performacne.