Improve discoverability on the hub (#325)

Co-authored-by: Lucain <lucainp@gmail.com>
Co-authored-by: Simon Alibert <75076266+aliberts@users.noreply.github.com>
This commit is contained in:
NielsRogge
2024-08-19 15:16:46 +02:00
committed by GitHub
parent 0f6e0f6d74
commit 86bbd16d43
5 changed files with 29 additions and 4 deletions

View File

@@ -43,7 +43,13 @@ from lerobot.common.policies.utils import (
)
class DiffusionPolicy(nn.Module, PyTorchModelHubMixin):
class DiffusionPolicy(
nn.Module,
PyTorchModelHubMixin,
library_name="lerobot",
repo_url="https://github.com/huggingface/lerobot",
tags=["robotics", "diffusion-policy"],
):
"""
Diffusion Policy as per "Diffusion Policy: Visuomotor Policy Learning via Action Diffusion"
(paper: https://arxiv.org/abs/2303.04137, code: https://github.com/real-stanford/diffusion_policy).