This is the official repo for the paper:

RoboDreamer: Learning Compositional World Models for Robot Imagination
Siyuan Zhou, Yilun Du, Jiaben Chen, Yandong Li, Dit-Yan Yeung, Chuang Gan
website | arxiv

Installation

conda create -n rtx python==3.9
conda activate rtx
pip install -r requirement.txt

Dataset

This repo contains example dataset in datasets/

You can download dataset from open-x

Training

torchrun train_rtx.py

You can set your own hyperparameters by using config.py

Citation

@article{zhou2024robodreamer,
  title={RoboDreamer: Learning Compositional World Models for Robot Imagination},
  author={Zhou, Siyuan and Du, Yilun and Chen, Jiaben and Li, Yandong and Yeung, Dit-Yan and Gan, Chuang},
  journal={arXiv preprint arXiv:2404.12377},
  year={2024}
}

Acknowledgements

This codebase is modified from the following repositories: imagen-pytorch and AVDC