Mini-Diffuse-Actor

Highlights

• Achieves 95% of the performance of 3D-Diffuser-Actor on the RLBench-18 multi-task benchmark, with just 13 hours of training on a single RTX 4090 GPU, or 1 day on an A100. Check our Wandb Reports.
• Complete codebase for RLBench training and headless testing, as well as real-world training and evaluation.
• Includes checkpoints, training logs, and test logs.
• Training and testing can be run on cloud platforms. For example, on Vast.ai it costs less than $5 to train your own Mini-Diffuser.

This codebase builds upon 3D-LUTOS — many thanks to the original authors for open-sourcing their container and providing a robust headless training-evaluation-testing pipeline. The idea was inspired by comparing 3D-Diffuser-Actor and Act3D.

🚧 Project in progress

Installation

See INSTALL.md.
For training-only usage, you don’t need to install the simulator.

Training

Run:

python train/train_diffusion_policy.py

To reproduce results, ensure arguments match those used in the wandb logs. Example:

python train/train_diffusion_policy.py MODEL.mini_batches=256 TRAIN.learning_rate=3e-4 wandb_name=rlbench18_256 TRAIN_DATASET.num_points=4000 VAL_DATASET.num_points=4000 MODEL.diffusion.total_timesteps=100 TRAIN.num_epochs=800 SEED=2024 TRAIN.lr_sched=cosine TRAIN.num_cosine_cycles=0.6

Testing

You can test your own model or download pretrained checkpoints here.

Ensure the RLBench simulator is working properly according to INSTALL.md.

Run:

python minidiffuser/evaluation/eval_simple_policy_parrallel.py

For headless testing:

xvfb-run -a python minidiffuser/evaluation/eval_simple_policy_parrallel.py

To reproduce benchmark results:

bash scripts/locals_policy_peract.sh 

⚠️ Remember to delete results in the seed folder first — otherwise, runs with existing seeds will be skipped.

Real-World Experiments

See our real-world setup in this repo. In short, we seperate the ROS and CUDA learning environment and use websocket to communicate, to aviod python version conflict.

In that repo we will update with guides on:

Data Collection

Preprocessing Collected Data

Training

Real-World Testing