Skip to content

leggedrobotics/GraspDiffuser

BranchesTags

Repository files navigation

grasp_diffuser

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.

Python 3.8+ PyTorch License: MIT

📦 Installation

From Source

# 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 .

Dependencies

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.
  • pointops CUDA extension — required by the PointTransformer scene 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.

Docker (recommended)

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=/data

Or 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 shell

Note: 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.hands API matches this codebase.

🚀 Quick Start

Training a Model

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=512

Multi-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.

Using the Python API

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)
)

🎨 Visualization

Running Inference with Visualization

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

Example Workflow

# 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')
"

🏗️ Package Structure

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

🔧 Configuration

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: 2048

Override config values from command line:

python scripts/train.py diffuser.steps=200 model.d_model=512

📊 Dataset

grasp_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).

🧪 Testing

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 -v

To 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).

🛠️ Troubleshooting

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.

🎓 Core Concepts

Diffusion Models

Implements Denoising Diffusion Probabilistic Models (DDPM):

  1. Forward Process: Gradually adds noise to grasp poses
  2. Reverse Process: Learns to denoise and generate new grasps
  3. Scene Conditioning: Conditions on 3D object point clouds

6D Rotation Representation

Uses robust 6D rotation representation for continuous optimization:

  • Avoids gimbal lock and discontinuities
  • Gram-Schmidt orthonormalization
  • Handles degenerate cases

Grasp Representation

  • 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}
}

About

No description, website, or topics provided.

Resources

Readme

License

View license
Activity
Custom properties

Stars

0 stars

Watchers

0 watching

Forks

0 forks
Report repository

Releases

No releases published

Packages

 
 
 

Contributors

Languages