Policy Adaptation via Language Optimization: Decomposing Tasks for Few-Shot Imitation

Vivek Myers, Bill Chunyuan Zheng, Oier Mees, Sergey Levine, Kuan Fang

---

This repository contains the code for Policy Adaptation via Language Optimization (PALO), which combines a handful of demonstrations of a task with proposed language decompositions sampled from a VLM to quickly enable rapid nonparametric adaptation, avoiding the need for a larger fine-tuning dataset.

Environment

conda create -n palo python=3.10
conda activate palo
pip install -e . 
pip install -r requirements.txt

For GPU:

pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

For TPU

pip install --upgrade "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

See the Jax Github page for more details on installing Jax.

Using PALO

Running Optimization

To get the best language decomposition from PALO, you can run the following commands:

python palo/optimize.py --instruction [Your Instruction Here] --trajectory_path [Your data here] \
 --checkpoint_path "./agent/checkpoint/" --im_size 224 --config_dir "./agent/config.pkl"

Citation

PLease consider citing our work if you find it useful:

@inproceedings{myers2024policy,
  title={Policy Adaptation via Language Optimization: Decomposing Tasks for Few-Shot Imitation},
  author={Vivek Myers and Bill Chunyuan Zheng and Oier Mees and Sergey Levine and Kuan Fang},
  booktitle={Conference on Robot Learning},
  year={2024}
}