GitHub - ziyin-xiong/Ag2x2 (original) (raw)

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:
    • 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.