Installation | Prepare Dataset | Train️ | Evaluate | Project Website
git clone https://github.com/VT-Collab/TransMASK.git
cd TransMASK
You can install the vitual conda environment using the following command:
conda create -n transmask python=3.10
conda activate transmask
pip install -r requirements.txt
The following command will install the panda-gym package in your virtual conda environment:
cd panda_gym
pip install -e .
cd ..
Run the following command to download the panda-gym dataset for three tasks - Pick and Place, Push, and Rotate:
python download_datasets.py
This will create a datasets/ directory in the main folder.
The details about the dataset structure can be found here: https://huggingface.co/datasets/SagarParekh/TransMASK
To train a Behavior Cloning (BC) policy on a dataset with privileged state information, run the following command:
python train.py --dataset_path [dataset_path] --device cuda --env_name [env_name] --policy [policy_type] --savename [save_name] --epochs 500 --num_distractors 4 --method [method]
To train a Behavior Cloning (BC) policy on a dataset with image observations, run the following command:
python train.py --dataset_path [dataset_path] --device cuda --env_name [env_name] --policy [policy_type] --savename [save_name] --epochs 500 --num_distractors 4 --method [method] --use_image_obs
The --env_name argument has the following options:
- PandaPickAndPlace-v3
- PandaPush-v3
- PandaFlipJoints-v3
The --policy argument has two options:
- mlp (to train a multi layer perceptron action head)
- diffusion (to train a diffusion policy)
For diffusion policy, you can also specify the length of input sequence as well as the length of predicted actions using the arguments --obs_horizon and --pred_horizon respectively.
The --method argument allows you to choose whether to train a vanilla BC, VAE BC, or TransMASK BC policy
- default
- vae
- our
The trained models are stored in ./results/[task_name]/[save_name].
To evaluate the trained policy in the simulated environment, run the following command:
python eval.py --num_evals 100 --device cuda --render_mode rgb_array --loadloc [path_to_trained_model] --n_rollout_actions [n_rollout_actions]
The --n_rollout_actions argument allows you to specify the length of the predicted action sequence that will be executed.
The results are saved in the same directory as the trained models.