Self-Supervised Skill Abstraction and Decision Making

Model specific installations

0. conda with python 3.8
1. pip install vector-quantize-pytorch==1.8.1
2. pip install positional-encodings==6.0.3
3. pip install einops==0.3.2
4. pip install hydra-core==1.2.0
5. Other dependencies like torch, torchvision, transformers, etc
6. Libero specific dependencies are in the requirements.txt file.

Model Files

Model files are in the libero/lifelong/models directory. Description:

1. skill_vae.py: stage 1 policy file, uses base_policy class, imports stage 1 model and observation encoder.
2. modules/skill_vae_modules.py: SkillVAE model - stage 1.
3. skill_GPT.py: stage 2 policy file, uses base_policy class, imports stage 1 model and stage 2 GPT model both.
4. modules/skill_utils.py: contains stage 2 GPT model and other stage 1 modules.
5. modules/rgb_modules.py: contains visual observation encoder.

Config Files

Config files are in the libero/configs directory. Description:

1. pretrain.yaml: stage 1 training.
2. skillGPT.yaml: stage 2 training.
3. few_shot.yaml: few-shot training on libero-10. All of these files use default files in the libero/configs directory.
   Following are model specific config files:
4. policy/skill_vae.yaml: stage 1 model.
5. policy/skill_GPT.yaml: stage 2 model.

Read comments in the config files for more details.

Training Scripts

These are libero specific scripts that loads the model and data and environment (for evals) and trains the model.
Main scripts for training are in the libero/lifelong directory. The main scripts are:

1. pretrain.py: Pretrain the SkillVAE model.

2. skill_policy.py: Train the skill policy. These scripts load algo/multitask.py algo that loads multitask data and trains the model. (libero specific)

3. Command to train stage 1:
   export CUDA_VISIBLE_DEVICES=0 && export MUJOCO_EGL_DEVICE_ID=0 && python libero/lifelong/pretrain.py

4. Command to train stage 2: export CUDA_VISIBLE_DEVICES=0 && export MUJOCO_EGL_DEVICE_ID=0 && python libero/lifelong/skill_policy.py

5. Command to finetune: export CUDA_VISIBLE_DEVICES=0 && export MUJOCO_EGL_DEVICE_ID=0 && python libero/lifelong/few_shot.py

6. Command to evaluate: export CUDA_VISIBLE_DEVICES=0 && export MUJOCO_EGL_DEVICE_ID=0 && python libero/lifelong/skill_policy_eval.py

Running this on PACE

1. Start a free interactive session with a RTX6000 GPU: salloc -A gts-agarg35 -N1 --mem-per-gpu=12G -q embers -t8:00:00 --gres=gpu:RTX_6000:1
2. Load in Conda: module load anaconda3/2022.05.0.1
3. sbatch slurm/run.sbatch python libero/lifelong/pretrain.py