mirror of
https://github.com/huggingface/lerobot.git
synced 2026-06-18 00:37:10 +00:00
improved and extended structure
This commit is contained in:
Nikodem Bartnik
@@ -29,6 +29,7 @@ from huggingface_hub.errors import HfHubHTTPError
|
||||
from safetensors.torch import load_model as load_model_as_safetensor, save_model as save_model_as_safetensor
|
||||
from torch import Tensor, nn
|
||||
|
||||
from lerobot.__version__ import __version__
|
||||
from lerobot.configs import PreTrainedConfig
|
||||
from lerobot.configs.train import TrainPipelineConfig
|
||||
from lerobot.utils.hub import HubMixin
|
||||
@@ -38,6 +39,67 @@ from .utils import log_model_loading_keys
|
||||
T = TypeVar("T", bound="PreTrainedPolicy")
|
||||
|
||||
|
||||
def _build_card_context(
|
||||
cfg: TrainPipelineConfig | None,
|
||||
dataset_repo_id: str | None,
|
||||
input_features: dict | None,
|
||||
output_features: dict | None,
|
||||
) -> dict:
|
||||
"""Collect optional data for the model-card template.
|
||||
|
||||
Returns plain values only (no Markdown) — the template in
|
||||
``lerobot/templates/lerobot_modelcard_template.md`` decides how and whether to show
|
||||
each one. Everything is best-effort: anything unavailable is left empty/None and the
|
||||
template simply skips that section, so this never breaks a Hub push.
|
||||
"""
|
||||
context = {
|
||||
"training": None,
|
||||
"input_features": input_features or {},
|
||||
"output_features": output_features or {},
|
||||
"dataset": None,
|
||||
"robot_type": None,
|
||||
"cameras": [],
|
||||
}
|
||||
|
||||
if cfg is not None:
|
||||
optimizer = getattr(cfg, "optimizer", None)
|
||||
context["training"] = {
|
||||
"steps": cfg.steps,
|
||||
"batch_size": cfg.batch_size,
|
||||
"seed": cfg.seed,
|
||||
"optimizer": getattr(optimizer, "type", None) if optimizer else None,
|
||||
"lr": getattr(optimizer, "lr", None) if optimizer else None,
|
||||
"lerobot_version": __version__,
|
||||
}
|
||||
|
||||
if dataset_repo_id:
|
||||
dataset_cfg = getattr(cfg, "dataset", None)
|
||||
try:
|
||||
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
|
||||
|
||||
meta = LeRobotDatasetMetadata(
|
||||
dataset_repo_id,
|
||||
root=getattr(dataset_cfg, "root", None),
|
||||
revision=getattr(dataset_cfg, "revision", None),
|
||||
)
|
||||
context["dataset"] = {
|
||||
"repo_id": dataset_repo_id,
|
||||
"episodes": meta.total_episodes,
|
||||
"frames": meta.total_frames,
|
||||
"fps": meta.fps,
|
||||
"tasks": [str(task) for task in meta.tasks.index],
|
||||
}
|
||||
context["robot_type"] = meta.robot_type
|
||||
context["cameras"] = [key.split(".")[-1] for key in meta.camera_keys]
|
||||
except Exception as e: # noqa: BLE001 — dataset details are optional, never fail the push
|
||||
logging.warning(
|
||||
f"Could not load dataset metadata for '{dataset_repo_id}'; those sections will be "
|
||||
f"omitted from the model card. ({e})"
|
||||
)
|
||||
|
||||
return context
|
||||
|
||||
|
||||
class ActionSelectKwargs(TypedDict, total=False):
|
||||
noise: Tensor | None
|
||||
|
||||
@@ -228,7 +290,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
|
||||
self.save_pretrained(saved_path) # Calls _save_pretrained and stores model tensors
|
||||
|
||||
card = self.generate_model_card(
|
||||
cfg.dataset.repo_id, self.config.type, self.config.license, self.config.tags
|
||||
cfg.dataset.repo_id, self.config.type, self.config.license, self.config.tags, cfg=cfg
|
||||
)
|
||||
card.save(str(saved_path / "README.md"))
|
||||
|
||||
@@ -246,7 +308,12 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
|
||||
logging.info(f"Model pushed to {commit_info.repo_url.url}")
|
||||
|
||||
def generate_model_card(
|
||||
self, dataset_repo_id: str, model_type: str, license: str | None, tags: list[str] | None
|
||||
self,
|
||||
dataset_repo_id: str,
|
||||
model_type: str,
|
||||
license: str | None,
|
||||
tags: list[str] | None,
|
||||
cfg: TrainPipelineConfig | None = None,
|
||||
) -> ModelCard:
|
||||
base_model_mapping = {
|
||||
"smolvla": "lerobot/smolvla_base",
|
||||
@@ -266,10 +333,14 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
|
||||
base_model=base_model_mapping.get(model_type),
|
||||
)
|
||||
|
||||
context = _build_card_context(
|
||||
cfg, dataset_repo_id, self.config.input_features, self.config.output_features
|
||||
)
|
||||
|
||||
template_card = (
|
||||
files("lerobot.templates").joinpath("lerobot_modelcard_template.md").read_text(encoding="utf-8")
|
||||
)
|
||||
card = ModelCard.from_template(card_data, template_str=template_card)
|
||||
card = ModelCard.from_template(card_data, template_str=template_card, **context)
|
||||
card.validate()
|
||||
return card
|
||||
|
||||
|
||||
@@ -58,14 +58,69 @@ This is a **{{ model_name }}** policy trained with [LeRobot](https://github.com/
|
||||
{% endif %}
|
||||
|
||||
This policy has been trained and pushed to the Hub using [LeRobot](https://github.com/huggingface/lerobot).
|
||||
See the full documentation at [LeRobot Docs](https://huggingface.co/docs/lerobot/index).
|
||||
{% set policy_docs = {"act": "act", "smolvla": "smolvla", "pi0": "pi0", "pi0_fast": "pi0fast", "pi05": "pi05", "eo1": "eo1", "groot": "groot"} %}
|
||||
{% if policy_docs.get(model_name) %}Learn how to train and run it in the [LeRobot {{ model_name }} guide](https://huggingface.co/docs/lerobot/main/en/{{ policy_docs[model_name] }}), or browse the [full documentation](https://huggingface.co/docs/lerobot/index).
|
||||
{% else %}See the [full LeRobot documentation](https://huggingface.co/docs/lerobot/index).
|
||||
{% endif %}
|
||||
|
||||
---
|
||||
|
||||
## Model Details
|
||||
|
||||
- **License:** {{ license | default("\[More Information Needed]", true) }}
|
||||
{% if robot_type %}- **Robot type:** `{{ robot_type }}`
|
||||
{% endif %}{% if cameras %}- **Cameras:** {% for camera in cameras %}`{{ camera }}`{% if not loop.last %}, {% endif %}{% endfor %}
|
||||
{% endif %}
|
||||
{% if input_features or output_features %}
|
||||
## Inputs & Outputs
|
||||
|
||||
The policy consumes these observation features and produces these action features.
|
||||
{% if input_features %}
|
||||
**Inputs**
|
||||
|
||||
| Feature | Type | Shape |
|
||||
| --- | --- | --- |
|
||||
{% for name, feature in input_features.items() %}| `{{ name }}` | {{ feature.type.value }} | `{{ feature.shape }}` |
|
||||
{% endfor %}{% endif %}{% if output_features %}
|
||||
**Outputs**
|
||||
|
||||
| Feature | Type | Shape |
|
||||
| --- | --- | --- |
|
||||
{% for name, feature in output_features.items() %}| `{{ name }}` | {{ feature.type.value }} | `{{ feature.shape }}` |
|
||||
{% endfor %}{% endif %}{% endif %}
|
||||
{% if dataset %}
|
||||
## Training Dataset
|
||||
|
||||
- **Repository:** [{{ dataset.repo_id }}](https://huggingface.co/datasets/{{ dataset.repo_id }})
|
||||
- **Episodes:** {{ dataset.episodes }}
|
||||
- **Frames:** {{ dataset.frames }}
|
||||
- **Frame rate:** {{ dataset.fps }} FPS
|
||||
{% if dataset.tasks %}- **Task(s):** {% for task in dataset.tasks %}"{{ task }}"{% if not loop.last %}, {% endif %}{% endfor %}
|
||||
{% endif %}{% endif %}
|
||||
{% if training %}
|
||||
## Training Configuration
|
||||
|
||||
| Setting | Value |
|
||||
| --- | --- |
|
||||
| Training steps | {{ training.steps }} |
|
||||
| Batch size | {{ training.batch_size }} |
|
||||
{% if training.optimizer %}| Optimizer | {{ training.optimizer }} |
|
||||
{% endif %}{% if training.lr %}| Learning rate | {{ training.lr }} |
|
||||
{% endif %}{% if training.seed is not none %}| Seed | {{ training.seed }} |
|
||||
{% endif %}| LeRobot version | {{ training.lerobot_version }} |
|
||||
{% endif %}
|
||||
---
|
||||
|
||||
## How to Get Started with the Model
|
||||
|
||||
For a complete walkthrough, see the [training guide](https://huggingface.co/docs/lerobot/il_robots#train-a-policy).
|
||||
Below is the short version on how to train and run inference/eval:
|
||||
New to LeRobot? These guides cover the full workflow:
|
||||
|
||||
- **[Install LeRobot](https://huggingface.co/docs/lerobot/main/en/installation)** — set up the `lerobot` package.
|
||||
- **[Hardware setup](https://huggingface.co/docs/lerobot/main/en/hardware_guide)** — assemble, wire, and calibrate your robot and cameras.
|
||||
- **[Record data & train a policy](https://huggingface.co/docs/lerobot/en/il_robots)** — the end-to-end imitation-learning walkthrough.
|
||||
- **[CLI cheat-sheet](https://huggingface.co/docs/lerobot/main/en/cheat-sheet)** — quick reference for the `lerobot-*` commands.
|
||||
|
||||
The short version to train and run this policy:
|
||||
|
||||
### Train from scratch
|
||||
|
||||
@@ -97,10 +152,27 @@ lerobot-rollout \
|
||||
|
||||
Replace every `<...>` placeholder with your own values. The `--robot.type`, `--robot.port`, and camera names/indices must match the robot and observation keys this policy was trained on, and `--task` should describe what you want the policy to do.
|
||||
|
||||
If you want to record a dataset while testing the policy use `--dataset.repo_id=<hf_user>/eval_dataset_name` it is important to use the prefix **eval\_**. For the policy path use the policy from the Hugging Face Hub or a local one. Skipping duration will make the policy run indefinitely.
|
||||
When `--strategy.type=base` is used the script doesn't record the episodes. Skipping duration will make the policy run indefinitely. For more information look at [rollout documentation](https://huggingface.co/docs/lerobot/main/en/inference).
|
||||
|
||||
---
|
||||
|
||||
## Model Details
|
||||
## Evaluation
|
||||
|
||||
- **License:** {{ license | default("\[More Information Needed]", true) }}
|
||||
<!-- Add evaluation results here: success rate, number of trials, and the conditions (robot, task, environment). -->
|
||||
|
||||
_No evaluation results have been provided for this policy yet._
|
||||
|
||||
---
|
||||
|
||||
## Citation
|
||||
|
||||
If you use this policy, please cite the method linked in the description above, along with LeRobot:
|
||||
|
||||
```bibtex
|
||||
@misc{cadene2024lerobot,
|
||||
author = {Cadene, Remi and Alibert, Simon and Soare, Alexander and Gallouedec, Quentin and Zouitine, Adil and Palma, Steven and Kooijmans, Pepijn and Aractingi, Michel and Shukor, Mustafa and Aubakirova, Dana and Russi, Martino and Capuano, Francesco and Pascal, Caroline and Choghari, Jade and Moss, Jess and Wolf, Thomas},
|
||||
title = {LeRobot: State-of-the-art Machine Learning for Real-World Robotics in Pytorch},
|
||||
howpublished = "\url{https://github.com/huggingface/lerobot}",
|
||||
year = {2024}
|
||||
}
|
||||
```
|
||||
|
||||
Reference in New Issue
Block a user