Scaling Manipulation Learning with Visual Kinematic Chain Prediction

Arixv Paper Website

This repo contains code for the VKT implementation, download datasets from box drive, extract and change the dataset path in configs/vkt/default.py.

Then, install torch and torchvision, and other dependencies.

pip install -r ./requirements.txt

Next, run the training command:

accelerate launch  --config_file ./accelerate.config.yaml  train.py  --config  ./configs/vkt/default.py   --wandb YOUR_WANDB_PROJECT_ID   --module.data.num_workers 10

Note:

1. The above dataset does not include UR5, sawyer and language table, even thought language table is shown in the above gif.
2. This code includes the visual kinematics chain forecasting part, the remaining code will be uploaded later.