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

@@ -38,7 +38,13 @@ from lerobot.common.policies.vqbet.vqbet_utils import GPT, ResidualVQ
# ruff: noqa: N806
class VQBeTPolicy(nn.Module, PyTorchModelHubMixin):
class VQBeTPolicy(
nn.Module,
PyTorchModelHubMixin,
library_name="lerobot",
repo_url="https://github.com/huggingface/lerobot",
tags=["robotics", "vqbet"],
):
"""
VQ-BeT Policy as per "Behavior Generation with Latent Actions"
"""