Ag2x2

Code repository for Ag2x2: Robust Agent-Agnostic Visual Representations for Zero-Shot Bimanual Manipulation.

Environment Setup

1. Clone the project:

   git clone https://github.com/ziyin-xiong/Ag2x2.git
   cd Ag2x2
   

2. Create a conda environment:

   conda create -n ag2x2 python=3.6.8
   conda activate ag2x2
   

3. Install Pytorch:

   pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 -f https://download.pytorch.org/whl/torch_stable.html
   

4. Install conda dependencies:

   conda install --file dependencies.txt -y
   

5. Install pip requirements:

   pip install -r requirements.txt
   

6. Install IsaacGym following the official documentation.

Visual Representation

1. Train our visual representation model on EPIC-KITCHEN dataset:

   • cd repre_trainer

   • Run train_ddp.py to train our model on multiple GPUs in parallel, or run train.py to train on a single GPU.
2. Specify your model save path by modifying exp_name in repre_trainer/cfgs/scratch.yml.
3. Please download our checkpoint here.

Bimanual Skills

1. Change ckpt_dir according to the location you store your visual representation checkpoint.
2. Train bimanual tasks in IsaacGym with the following command:

   python train.py --task=ag2x2@close_door_outward@ag2x2 --algo=ppo --seed=42 --cfg_train=cfgs/algo/ppo/manipulation.yaml --disable_wandb --camera=default

   The best policy will be saved as model_best.pt in logs/ag2x2/close_door_outward@default/ag2x2@ppo.42/.
3. Inference and save a trajectory generated by the trained policy:

   python train.py --task=ag2x2@close_door_outward@ag2x2 --model_dir=logs/ag2x2/close_door_outward@default/ag2x2@ppo.42/model_best.pt --test --save_traj --algo=ppo --cfg_train=cfgs/algo/ppo/manipulation.yaml --camera=default --seed=0 --disable_wandb

   The trajectory will be saved as logs/ag2x2/close_door_outward@default/ag2x2@ppo.42/absres_best.pkl.
4. Plan with franka robot arms using the saved trajectory:

   python plan.py --task=ag2x2@close_door_outward@ag2x2 --traj_path=logs/ag2x2/close_door_outward@default/ag2x2@ppo.42/absres_best.pkl --pipeline=cpu --algo=ppo --cfg_train=cfgs/algo/ppo/manipulation.yaml --disable_wandb --camera=default
   

Results Visualization

1. We test our model on 13 tasks, each with 9 runs (3 seeds x 3 camera positions). Check out our 13x9 experiment videos here.
2. Additionally, please check out imitation learning training data visualization and the inference video here.