Make-An-Agent: A Generalizable Policy Network Generator with Behavior-Prompted Diffusion

[Paper] [Project Website]

This repository is the official PyTorch implementation of Make-An-Agent. Make-An-Agent, a policy parameter generator that leverages the power of conditional diffusion models for behavior-to-policy generation, which demonstrates remarkable versatility and scalability on multiple tasks and has a strong generalization ability on unseen tasks to output well-performed policies with only few-shot demonstrations as inputs.

💻 Installation

1. create a virtual environment and install all required packages.

conda env create -f environment.yml 
conda activate makeagent

2. install Metaworld and mujoco_py for evaluations. Following instructions in DrM.

🛠️ Code Usage

For training autoencoder and behavior embedding, you could download the training dataset from Huggingface: train_data/training_dataset.pt or specific data for each task, e.g. train_data/door-open.pt.

Training the parameter autoencoder to encode and decode policy network parameters: Change data_root in autoencoder/config.yaml.

cd autoencoder
python train.py

Training behavior embeddings to process trajectory data: Change data_root in behavior_embedding/config_embed.yaml.

cd behavior_embedding
python train.py

Training the policy generator with conditional diffusion models:

Data processing:

Make-An-Agent uses a latent diffusion model, so the data should be processed using the autoencoder and behavior embedding.

You can directly use the pretrained models in HuggingFace.

Or use the processed training data in train_data/process_data.pt to train the policy generator.

If you want to process your own data, change the paths of data and pretrained model root in dataset/config.yaml.

cd dataset.py
python process_data.py

Ensure you now have processed data to match the latent representation dimensions, then change data_root in PolicyGenerator/config.yaml with your processed data.

cd PolicyGenerator
python train.py

Tips: We save both the best model and the last model during training.

Evaluating the synthesized parameters:

First we need to decode the latent parameters into policy parameters. Then deploy synthesized policy networks in the simulator.

We provide the processed test data in test_data/processed/. You also can process your test data using dataset/process_data.py.

Change ckpt_dir, encoder_dir and data_dir in PolicyGenerator/config.yaml.

cd PolicyGenerator
python eval.py

📗 Dataset and Pretrained Models

We release all pretrain models and data in Huggingface.

Pretrained models: Include autoencoder.pt, behavior_embedding.pt and model-best.pt.

Dataset: Training data is in train_data/ and test data is in test_data/

📝 Citation

If you find our work or code useful, please consider citing as follows:

@article{liang2024make,
title={Make-An-Agent: A Generalizable Policy Network Generator with Behavior-Prompted Diffusion},
author={Liang, Yongyuan and Xu, Tingqiang and Hu, Kaizhe and Jiang, Guangqi and Huang, Furong and Xu, Huazhe},
journal={arXiv preprint arXiv:2407.10973},
year={2024}
}

🌷 Acknowledgement

Our work is primarily based on the following projects: Diffusion Policy, pytorch-lightning, Metaworld, Robosuite, walk-these-ways. We thank these authors for their contributions to the open-source community.

For any questions or suggestions, please contact Yongyuan Liang.