A modular, production-ready Python package for grasp generation using diffusion models.
grasp_diffuser provides a clean, maintainable implementation of diffusion-based grasp generation for robotic manipulation. Built on PyTorch and PyTorch Lightning, it offers state-of-the-art grasp synthesis with a focus on code quality, extensibility, and ease of use.
# Clone the repository
git clone https://github.com/leggedrobotics/GraspDiffuser.git
cd GraspDiffuser
# Install Python dependencies
pip install -r requirements.txt
# Install in development mode
pip install -e .See requirements.txt for the pip-installable list. Key dependencies:
- Python 3.8+
- PyTorch 2.0+ / PyTorch Lightning 2.0+
- Hydra + OmegaConf for configuration
- Additional robotics/3D libraries (see requirements.txt)
Two dependencies are not on PyPI and must be installed separately:
- GraspQP — provides the hand
models (
graspqp.hands.get_hand_model). Required for training and inference. pointopsCUDA extension — required by thePointTransformerscene encoder (the default). It ships with GraspQP / the PointTransformer sources and is compiled against your CUDA toolkit; a CUDA-capable GPU is needed to build and run it.
The unit tests (
tests/) are CPU-only and do not require a GPU, but training/inference do.
The most reliable way to get a working environment — including the pointops
CUDA extension and GraspQP with its CUDA dependencies — is the provided
Dockerfile. You need Docker and, for training/inference, a
CUDA-capable GPU with the
NVIDIA Container Toolkit.
# Build the image (first build compiles CUDA extensions; takes a while)
docker build -t grasp-diffuser:latest .
# Run the unit tests (CPU-only)
docker run --rm grasp-diffuser:latest python -m pytest tests/ -q
# Train on a dataset mounted at /data (GPU required)
docker run --rm --gpus all \
-v /path/to/your/dataset:/data:ro \
-v "$PWD/outputs:/workspace/outputs" \
grasp-diffuser:latest \
python scripts/train.py exp_name=demo dataset.paths.Ours.asset_dir=/dataOr use Docker Compose (see docker-compose.yml):
DATA_DIR=/path/to/your/dataset docker compose run --rm train
docker compose run --rm test # CPU-only unit tests
docker compose run --rm shell # interactive shellNote: the image installs GraspQP from source. GraspDiffuser targets a specific GraspQP version/API — pin the compatible branch/tag/commit with
--build-arg GRASPQP_REF=<ref>(and--build-arg GRASPQP_REPO=<git-url>for a fork). The image builds, the environment is complete, and the unit tests pass out of the box; end-to-end training additionally requires a GraspQP ref whose forward-kinematics /graspqp.handsAPI matches this codebase.
Using the provided training script with Hydra configs:
The default config uses the allegro hand and the objects_allegro task
(configs/default.yaml). Point it at a dataset (see docs/DATASET.md):
# Train with the default config on your dataset
python scripts/train.py exp_name=my_experiment \
dataset.paths.Ours.asset_dir=/path/to/your/dataset
# Override diffusion / model hyper-parameters from the command line
python scripts/train.py exp_name=my_experiment \
dataset.paths.Ours.asset_dir=/path/to/your/dataset \
diffuser.steps=200 model.d_model=512Multi-GPU is automatic: with the default gpu=auto, all visible GPUs are used
with a DDP strategy. Restrict devices with e.g. gpu=[0,1] or gpu=0.
The full pipeline (eps-model + DDPM + normalizer + hand model) is assembled by
the create_diffusion_pipeline factory from a Hydra config:
from grasp_diffuser import create_diffusion_pipeline
from graspqp.hands import get_hand_model
# cfg: a Hydra DictConfig (e.g. loaded from configs/default.yaml)
ddpm = create_diffusion_pipeline(cfg) # returns a DDPM ready for sampling
# Sample grasps for an object point cloud
entry = {"pos": point_cloud} # your object point cloud, shape (N, 3)
samples, history = ddpm.sample_ddim(
num_samples=100,
entry=entry,
steps=50,
eta=0.0, # 0.0 = deterministic (DDIM)
)Generate and visualize grasp predictions on test data:
# Basic inference with HTML visualization
python scripts/inference.py ckpt_dir="outputs/2026-02-16_13-42-28_allegro_sm"
# Enable denoising process animation (shows diffusion steps)
python scripts/inference.py ckpt_dir="outputs/2026-02-16_13-42-28_allegro_sm" task.visualizer.vis_denoising=true
# Customize number of samples and visualization settings
python scripts/inference.py \
ckpt_dir="outputs/YOUR_CHECKPOINT_DIR" \
task.visualizer.ksample=10 \
task.visualizer.n_vis_grasp=4 \
task.visualizer.vis_denoising=true# 1. Train a model
python scripts/train.py exp_name=my_grasp_model
# 2. Run inference with visualization
python scripts/inference.py \
ckpt_dir="outputs/YYYY-MM-DD_HH-MM-SS_my_grasp_model" \
task.visualizer.vis_denoising=true
# 3. View results
# Open {save_dir}/html/batch_0000_gt_vs_pred.html in browser
# Open {save_dir}/denoising/batch_0000_denoising_animation.html for animation
# 4. Load exported poses for analysis
python -c "
import torch
data = torch.load('outputs/.../batch_0000_predicted_poses.pt')
root_poses = data['parameters']['root_pose']
print(f'Generated {root_poses.shape[0]} grasps')
"GraspDiffuser/
├── grasp_diffuser/ # Main package
│ ├── core/ # Diffusion models and schedulers
│ │ ├── ddpm.py # DDPM implementation
│ │ └── schedulers.py # Noise schedules (linear, cosine, etc.)
│ ├── models/ # Neural network architectures
│ │ ├── unet.py # GraspUNet model
│ │ ├── blocks.py # Building blocks
│ │ ├── our_visualizer.py # Grasp visualization
│ │ └── scene_encoders/ # Point cloud encoders
│ │ ├── base.py # Base encoder interface
│ │ └── pointtransformer/
│ ├── data/ # Dataset and data loading
│ │ ├── dataset.py # GraspDataset
│ │ ├── noise.py # Noise augmentation
│ │ ├── collate.py # Batch collation
│ │ └── object_loader.py # Mesh/point cloud loading
│ ├── training/ # Training infrastructure
│ │ ├── lightning_module.py # PyTorch Lightning module
│ │ └── callbacks.py # Custom callbacks (viz, logging)
│ ├── utils/ # Utilities
│ │ ├── normalization.py # Pose normalization
│ │ ├── rotation.py # 6D rotation utilities
│ │ ├── visualization.py # Plotly helpers
│ │ ├── metrics.py # Loss computation
│ │ └── io.py # File I/O utilities
│ ├── factory.py # Hydra factory functions
│ └── __init__.py # Package exports
├── configs/ # Hydra configuration files
│ ├── default.yaml # Default config
│ ├── task/ # Task-specific configs
│ ├── model/ # Model configs
│ └── diffuser/ # Diffuser configs
├── scripts/ # Training and evaluation scripts
│ ├── train.py # Training entry point
│ ├── inference.py # Inference and visualization
│ ├── curobo_plan.py # Motion planning integration
│ └── vis_utils.py # Visualization utilities
├── tests/ # Test suite
├── setup.py # Package setup
├── requirements.txt # Dependencies
└── README.md # This file
grasp_diffuser uses Hydra for configuration. Example:
# configs/default.yaml
diffuser:
name: DDPM
steps: 100
schedule_cfg:
beta: [0.0001, 0.02]
model:
name: GraspUNet
d_model: 256
scene_model:
name: PointNet
num_points: 2048Override config values from command line:
python scripts/train.py diffuser.steps=200 model.d_model=512grasp_diffuser trains on per-object grasp datasets produced by
GraspQP: each object is a directory
with a remeshed.obj mesh and a grasp_predictions/…/*.pt file of grasp poses
(translation + wxyz quaternion + joint angles).
See docs/DATASET.md for:
- the exact on-disk directory layout and grasp-file schema,
- how to point training at a GraspQP dataset, and
- how to convert another data source into this format (with a runnable
reference converter,
scripts/convert_to_graspqp_format.py).
The unit tests are CPU-only and self-contained (no dataset or GPU needed), but they depend on the full scientific stack. The reproducible way to run them is inside the project Docker image:
# Run the whole suite in Docker
scripts/run_tests_docker.sh
# Forward arguments to pytest (e.g. a subset)
scripts/run_tests_docker.sh -k rotation -vTo run them directly in a suitable environment instead:
pip install -e ".[dev]"
pytest tests/scripts/run_tests_docker.sh builds the image automatically on first run (or
set BUILD=1 to force a rebuild).
| Symptom | Cause / fix |
|---|---|
ModuleNotFoundError: No module named 'torch' |
You're on a Python interpreter without the dependencies. Use the Docker image, or install into the correct environment. |
ModuleNotFoundError: No module named 'graspqp' |
GraspQP is not on PyPI. Install it from its repo, or use the Docker image (which installs it). |
PointTransformer requires pointops_cuda extension (or ImportError importing the encoder) |
The pointops CUDA extension isn't built. It needs a CUDA toolkit + GPU; the Docker image compiles it for you. |
Asset directory not found: /path/to/your/dataset |
The default dataset.paths.Ours.asset_dir is a placeholder — point it at a real dataset (see docs/DATASET.md). |
Docker: could not select device driver "nvidia" |
Install the NVIDIA Container Toolkit and pass --gpus all. |
Implements Denoising Diffusion Probabilistic Models (DDPM):
- Forward Process: Gradually adds noise to grasp poses
- Reverse Process: Learns to denoise and generate new grasps
- Scene Conditioning: Conditions on 3D object point clouds
Uses robust 6D rotation representation for continuous optimization:
- Avoids gimbal lock and discontinuities
- Gram-Schmidt orthonormalization
- Handles degenerate cases
- Position (3D): Hand translation
- Orientation (6D): Hand rotation
- Joint Angles (N): Hand configuration
## 🤝 Contributing
Contributions are welcome! Please:
1. Fork the repository
2. Create a feature branch
3. Make your changes with tests
4. Submit a pull request
## 📄 License
©2026 ETH Zurich, René Zurbrügg.
This project is licensed under the MIT License — see [LICENSE](LICENSE),
[AUTHORS](AUTHORS), and [NOTICE](NOTICE) for details.
This codebase is **derived from and adapted from**
[DexGrasp Anything](https://github.com/4DVLab/DexGrasp-Anything) (©2025 4DVLab,
MIT License). Its copyright notice is retained in [LICENSE](LICENSE) and the
derivation is documented in [NOTICE](NOTICE). External dependencies (e.g.
[GraspQP](https://github.com/leggedrobotics/graspqp), the `pointops` CUDA
extension) and any datasets you use are the property of their respective owners
and are **not** covered by this repository's license — see [NOTICE](NOTICE).
## 📧 Contact
For questions and support, please open an issue on GitHub.
## 🔗 Related Projects
- [DexGrasp Anything](https://dexgraspanything.github.io/) - Diffusion-based dexterous grasp generation (this codebase is inspired by it)
- [GraspQP](https://github.com/leggedrobotics/graspqp) - Hand models and kinematics
- [PyTorch Lightning](https://lightning.ai) - Training framework
## 🙏 Acknowledgments
This codebase is **inspired by and adapted from
[DexGrasp Anything: Towards Universal Robotic Dexterous Grasping with Physics
Awareness](https://arxiv.org/abs/2503.08257)** (Zhong et al., CVPR 2025
Highlight) — [project page](https://dexgraspanything.github.io/) ·
[paper](https://arxiv.org/abs/2503.08257). The diffusion pipeline, dataset
loading convention, and pose/normalization conventions build on ideas from that
work. We thank the authors for open-sourcing their code.
If you use this repository, please also consider citing DexGrasp Anything:
```bibtex
@article{zhong2025dexgrasp,
title={DexGrasp Anything: Towards Universal Robotic Dexterous Grasping with Physics Awareness},
author={Zhong, Yiming and Jiang, Qi and Yu, Jingyi and Ma, Yuexin},
journal={arXiv preprint arXiv:2503.08257},
year={2025}
}