Compare commits

...

29 Commits

Author SHA1 Message Date
Pepijn f2b90e3ad6 feat(runtime): add interactive language rollouts 2026-07-28 10:05:27 +02:00
Pepijn 3f093d8927 feat(data): add recipe-driven language supervision 2026-07-28 10:04:28 +02:00
Steven Palma 95211b98f1 feat(config): add multiprocessing option to DataLoader context and sets spawn as default (#4139)
* Add dataloader_multiprocessing_context, default to spawn

Make the DataLoader multiprocessing start method configurable on
TrainPipelineConfig and default it to 'spawn'.

The previous default (fork on Linux) is unsafe with libraries that hold
non-fork-safe state in the parent process — common ones in this codebase
are PyAV, torchcodec, and the ffmpeg shared libs they wrap. Symptoms
reported in #2488, #2209, and observed locally include:

- multiprocessing.context.AuthenticationError: digest received was wrong
- RuntimeError: Pin memory thread exited unexpectedly
- RuntimeError: DataLoader worker exited unexpectedly
- Random SIGSEGV inside worker processes during video decode

Switching to spawn re-imports modules cleanly in each worker and
eliminates these failure modes. Added the setting as a config field
rather than hard-coding so users on platforms where fork is preferred
can opt back in via --dataloader-multiprocessing-context=fork.

* Address review: shorten config comment, note spawn startup tradeoff

Per @jashshah999, mention that spawn workers re-import modules and so
add some startup time vs fork. Also trim the failure-mode dump from
the inline comment — the linked issue covers the symptoms in detail.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* chore(scripts): add multiprocessing_context safeguards

* chore(config): add libs note

---------

Co-authored-by: 0o8o0-blip <0o8o0-blip@users.noreply.github.com>
2026-07-28 00:42:55 +02:00
Xingdong Zuo 95256d766d feat(lekiwi): support LeKiwi in lerobot-replay CLI (#3739)
Register the `lekiwi` robot module in `lerobot_replay.py` so episodes can be
replayed on a LeKiwi via `--robot.type=lekiwi_client`. The module is already
registered in `lerobot_calibrate.py` and `lerobot_setup_motors.py`; this fills
the gap so the replay CLI recognizes the same robot.

Replayed actions are loaded from the dataset as torch tensors, which
`json.dumps` cannot serialize when `LeKiwiClient.send_action` ships them over
ZMQ. Coerce each action value to a plain float before sending. This is scoped
to the LeKiwi network client and does not affect any other robot.

Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-27 19:29:44 +02:00
Thomas Landeg fd53716688 fix(envs): make metaworld seeding reproducible (#3727)
Co-authored-by: Pepijn <138571049+pkooij@users.noreply.github.com>
2026-07-27 18:42:17 +02:00
Steven Palma a96540a2c4 fix rollout policy revision loading (#4161)
Co-authored-by: RaviTeja-Kondeti <rkondet3@asu.edu>
2026-07-27 18:20:39 +02:00
WOLIKIMCHENG acd42b4d85 fix(processor): keep missing local state resolution local (#3715)
Co-authored-by: root <kinsonnee@gmail.com>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-27 15:53:31 +02:00
Steven Palma bbeacfe57d fix(record): connect teleoperator before robot to avoid watchdog jump (#4166)
* fix(record): connect teleoperator before robot to avoid watchdog jump

lerobot-record connected the robot before the teleoperator. A robot's
connect()/reset() can leave it holding a default pose under a firmware
watchdog (e.g. Unitree G1); if teleop.connect() (model loading, IK init,
network setup) then takes longer than that watchdog, the joints drop to
damping and the first send_action() makes the robot jump.

Swap the order so the teleoperator connects first, matching the ordering
already used in lerobot_teleoperate.py. Pure ordering fix, no API change.

Fixes #3684

* fix(record): trim comment and add connect-order regression test

Address review feedback on #3684:
- Trim the verbose ordering comment down to two lines.
- Add test_record_connects_teleop_before_robot to tests/test_control_robot.py,
  asserting teleop.connect() runs before robot.connect() in record().

* chore(test): remove test

---------

Co-authored-by: Jaimin Patel <jpatel@tuvalabs.com>
Co-authored-by: Martino Russi <77496684+nepyope@users.noreply.github.com>
2026-07-27 14:09:58 +02:00
Steven Palma 801346e18c fix(scripts): restore policy training mode after eval_policy() in lerobot-eval (#4162)
* fix(scripts/eval): restore policy training mode after eval_policy()

`eval_policy` calls `policy.eval()` before the rollout but never restores
the prior mode on return. When called from the training loop
(`lerobot_train.py`'s `eval_policy_all -> run_one -> eval_one ->
eval_policy` chain), the policy is left in eval mode for every subsequent
training step, which silently:

  * disables Dropout (no regularisation),
  * freezes BatchNorm running stats (no further EMA updates).

Under DDP only `is_main_process` runs eval (lerobot_train.py:527), so the
main rank ends up in eval mode while workers stay in train mode — the
all-reduced gradients then combine forward passes computed with different
dropout masks and different BN behaviour, a real DDP-correctness issue.

Scope of impact:
  * Affects every policy with Dropout in its forward path. In-tree, that
    includes the default ACT (6 Dropout layers at p=0.1), Diffusion (vision
    backbone), VQ-BeT, Multi-Task DiT, X-VLA, plus all VLA policies that
    inherit Dropout from their pretrained HF backbone (PI0/PI0.5/PI0-FAST,
    SmolVLA, GR00T-N1.5, EO1, Wall-X).
  * Triggers from the first eval onward. On the default config
    (steps=100k, eval_freq=20k) that's 80% of training; on the LIBERO /
    RoboCasa / VLABench example commands in docs/ (eval_freq=1k–5k)
    it's 95–99% of training.
  * Policies using only LayerNorm/GroupNorm and no Dropout (TDMPC, RTC)
    are unaffected. Policies using `FrozenBatchNorm2d` (ACT's ResNet
    backbone) are immune to the BN-stat half; the Dropout half still bites.

Fix:
  * Snapshot `policy.training` on entry to `eval_policy`.
  * Restore it on normal return.
  * Save-and-restore is a strict no-op for callers that pass an
    already-eval-mode policy (e.g. the standalone `lerobot-eval` script
    loading a frozen checkpoint).
  * Restoration is placed before the normal return only, not in a
    try/finally — exception paths leave the policy in eval mode, same as
    today. A try/finally upgrade would require re-indenting ~165 lines and
    can land as a separate cleanup if desired.

Tests (tests/scripts/test_eval.py, 7 tests total, ~1.6s):
  * Regression gates on the lerobot_eval fix itself: training-mode
    preservation, eval-mode preservation, dropout-active behavioural
    check, non-crash for both entry modes.
  * Quantitative mechanism demonstration
    (`test_missing_mode_restoration_hurts_generalisation`): trains a tiny
    Dropout+BatchNorm MLP under both the bug pattern and the fix pattern
    on identical data and seed, then asserts the buggy variant generalises
    at least 5% worse on a held-out val set. In repeated runs we see
    10-25% deltas on this toy problem; real policies (more layers, more
    Dropout, longer training) generally see larger gaps. Lives alongside
    the regression tests so the empirical proof is reproducible from the
    repo without adding a separate benchmarks/ directory.


* fix(scripts): keep policy train/eval

---------

Co-authored-by: ModeEric <ericjm4@illinois.edu>
2026-07-27 14:08:11 +02:00
MihaiAnca13 ab87fd9764 fix(datasets): clear video frame staging on episode reset (#3683)
* fix video frame staging cleanup on episode reset

* linting

---------

Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-27 13:44:32 +02:00
hf-dependantbot-rollout[bot] 6c57dfd2ee chore: enable Dependabot weekly GitHub Actions bumps (#3677)
Co-authored-by: hf-dependantbot-rollout[bot] <285970069+hf-dependantbot-rollout[bot]@users.noreply.github.com>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-27 13:22:04 +02:00
Kohei SENDAI d63e6e67a5 fix convverstion err (#3656)
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-27 11:54:09 +02:00
Steven Palma 0d383d09f2 feat(dataset): accept token argument for private HF Hub datasets (#4136) 2026-07-24 18:51:35 +02:00
Caroline Pascal ab2b5b04dd (depth image processing): excluding depth frames from the RGB to BGR image processing (#4135)
* (depth image processing): excluding depth frames from the RGB to BGR image processing

* test(update): updating tests to include RGB/BGR conversion checks
2026-07-24 17:43:17 +02:00
Steven Palma ac5c7b8600 chore(deps): bump diffusers to >=0.38.0,<0.40.0 (#4145)
* fix(deps): bump diffusers cap to <0.39.0 (security)

Diffusers 0.35.x is affected by GHSA-98h9-4798-4q5v (HIGH, CVSS 8.8):
'trust_remote_code bypass via custom_pipeline and local custom components'.
Fixed in diffusers 0.38.0.

Current cap 'diffusers<0.36.0' blocks downstream consumers (e.g.
strands-labs/robots) from picking up the security fix.

The lerobot diffusers surface area is narrow and stable across 0.36-0.38:
- diffusers.schedulers.scheduling_ddim.DDIMScheduler
- diffusers.schedulers.scheduling_ddpm.DDPMScheduler
- diffusers.optimization.get_scheduler
- diffusers.ConfigMixin / ModelMixin / register_to_config
- diffusers.models.attention.{Attention,FeedForward}
- diffusers.models.embeddings.*

None of these were removed, renamed, or had breaking changes in 0.36, 0.37,
or 0.38 release notes. Bumping the cap to <0.39.0 unblocks the security
fix while keeping a major-version safety bound.

* chore(dependecies): bump diffusers

* chore(deps): update uv.lock

---------

Co-authored-by: Cagatay Cali <cagataycali@users.noreply.github.com>
2026-07-24 17:13:53 +02:00
Steven Palma a6befef0ba chore(dependencies): update uv.lock (#3963) 2026-07-24 16:30:36 +02:00
Steven Palma 53843007ea feat(robot): Make SO follower P coefficient configurable (#4142)
* Make SO follower P coefficient configurable

* chore(test): minimize tests

* feat(robots): expose PID coeff in SO arms

---------

Co-authored-by: taivu1998 <46636857+taivu1998@users.noreply.github.com>
2026-07-24 16:03:04 +02:00
Maxime Ellerbach d3bed0feee chore(agents): adding additional infos to AGENTS.md and bring-your-own-policies.mdx (#3904)
* chore(agents): adding additional infos to AGENTS.md

* adding `lerobot-train` requirement inside PR checklist

* prefer using code already implemented from transformers / diffusers instead of re-implementing in tree

---------

Signed-off-by: Maxime Ellerbach <maxime.ellerbach@huggingface.co>
2026-07-24 14:58:43 +02:00
Steven Palma a0eb860d1e feat(dataset): add slice support to LeRobotDataset.__getitem__ (#4129)
* feat(dataset): add efficient slice support

* fix(dataset): handle empty dataset slices

* refactor(dataset): reuse scalar path for slices

---------

Co-authored-by: Francesco Capuano <fc.francescocapuano@gmail.com>
2026-07-23 22:05:29 +02:00
Steven Palma cfd9ff969c fix(envs): set LiberoEnvConfig.fps default to 20 to match robosuite (#4124)
* fix(envs): set LiberoEnvConfig.fps default to 20 to match robosuite

LiberoEnvConfig.fps was set to 30, but the underlying robosuite
OffScreenRenderEnv always runs at its default control_freq of 20 Hz
since fps is never passed through. This mismatch silently decouples
the dataset/eval loop rate from the actual simulation step rate.

Set the default to 20 to match the real sim rate and avoid the
footgun.

Fixes #3368

* fix(libero): apply configured control frequency

---------

Co-authored-by: xinmotlanthua <275663218+xinmotlanthua@users.noreply.github.com>
2026-07-23 19:49:19 +02:00
Steven Palma f59eae4e27 fix(robots): add retries while recording motor ranges (#4126)
* Add retries while recording motor ranges

* fix(motors): throttle calibration reads consistently

---------

Co-authored-by: tom-doerr <tomdoerr96@gmail.com>
2026-07-23 18:41:48 +02:00
Martino Russi a993af9c51 fix(openarms): stop set_zero_position()ing on connect (#4058)
* fix(damiao): make is_calibrated a plain property, not cached

`is_calibrated` was a `@cached_property`, so it froze at its first-read
value and never reflected later changes to `self.calibration` (set by
connect/calibrate/load). This caused the OpenArm teleop to re-run
calibration even when a calibration file existed, and to skip
`set_zero_position()` after a fresh calibration.

Switch to `@property` (matching the MotorsBus base contract and the
Feetech/SO-100 buses) and drop the now-unused `functools.cached_property`
import.

Co-authored-by: Cursor <cursoragent@cursor.com>

* don't set_zero_position() on connect

---------

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-23 18:34:13 +02:00
Steven Palma 392246feaf feat(diffusion): add gradient checkpointing for memory optimization (#4127)
* feat(diffusion): add gradient checkpointing for memory optimization

Add gradient_checkpointing config option to DiffusionPolicy. When
enabled, wraps the UNet encoder, mid, and decoder residual blocks
with torch.utils.checkpoint.checkpoint to trade compute for memory.

Allows training with larger batch sizes or higher-resolution inputs
on memory-constrained GPUs. Disabled by default.

Usage: --policy.gradient_checkpointing=true

Part of the 0.6.0 roadmap item 3.3 (gradient checkpointing for all
policies).

* test(diffusion): verify gradient checkpointing parity

---------

Co-authored-by: Jash Shah <jashshah.999@gmail.com>
2026-07-23 18:33:10 +02:00
Steven Palma 19dcbc19f1 fix(gamepad): Gamepad on macos often does not need fallback (#4125)
* gamepad does often work on macos

* review comments

* fix(gamepad): expose hidapi fallback in config

---------

Co-authored-by: Maxim Bonnaerens <maxim@bonnaerens.be>
2026-07-23 18:21:48 +02:00
Steven Palma 679faeaafc fix(scripts): register third-party plugins in lerobot_setup_motors (#4123)
* fix(scripts): register third-party plugins in setup-motors

* test(setup-motors): cover plugin registration

---------

Co-authored-by: Janos von Gencsy <janos.von-gencsy@tum.de>
2026-07-23 18:06:44 +02:00
YK 228cb5ddb9 Fix missing periods at end of sentences in README (#3473)
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-23 16:07:58 +02:00
Eunsung Kim ad176c6d41 Feature omx docs (#3421)
* docs(omx): add header and omx image in docs

* fix(docs):adjust image size in omx docs

---------

Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-23 15:25:46 +02:00
Duhyeon, Kim d6c605e8c5 refactor(pi05): remove unused variables in embed_suffix method (#3263)
* refactor(pi05): remove unused variables in embed_suffix method

* Refactor embed_suffix to streamline pad_masks handling

Removed unused pad_masks list and simplified its creation.

Signed-off-by: Duhyeon, Kim <49020301+dudududukim@users.noreply.github.com>

---------

Signed-off-by: Duhyeon, Kim <49020301+dudududukim@users.noreply.github.com>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-23 14:37:57 +02:00
Pepijn 9c82c39c7b feat(annotate): run lerobot-annotate on HF Jobs via --job.target (#4095)
* feat(annotate): run lerobot-annotate on HF Jobs via --job.target

Annotation needed a hand-edited launcher script (examples/annotations/run_hf_job.py)
to reach a GPU: users copied it, rewrote the embedded CMD string for their dataset,
and ran it with `python`. Fold that into the CLI instead, mirroring `lerobot-train`:
`lerobot-annotate --job.target=h200` submits the exact command you'd run locally.

- AnnotationJobConfig extends JobConfig with the annotation runtime's defaults
  (vllm/vllm-openai image, 2h cap) plus --job.lerobot_ref, so an unmerged branch
  can be exercised remotely without editing a script.
- lerobot.jobs.annotate builds the pod command by replaying the user's own CLI
  flags (minus --job.*/--root, with --repo_id re-emitted from the config) after a
  setup prelude that installs lerobot on top of the vLLM image. Job monitoring,
  log tailing and Ctrl-C-detaches reuse the training submitter's plumbing.
- Remote runs require --repo_id; a local-only dataset is pushed privately first.

The generated pod command is byte-for-byte the script's old CMD.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* fix(annotate): reject client-side config files on remote runs

draccus exposes `--config_path` plus a `--<field>` config-file arg for every
nested dataclass (`--vlm`, `--plan`, `--job`, ...). All name files on the
client's disk, so forwarding them to the pod silently dropped whatever settings
they carried. Reject them up front instead.

Bare `--job` also slipped past the `--job.` prefix filter, so a `--job=cfg.yaml`
holding `target: h200` would have reached the pod and had the job submit a job
of its own, recursively. It is dropped from the forwarded args as well.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* refactor(jobs): share the submit-and-follow loop between both submitters

`submit_annotate_to_hf` reused the leaf helpers (`_poll_until_done`, `_tail_logs`,
`_pod_forwarded_args`) but duplicated the orchestration around them: ~40 of the 50
lines that spawn the poll/log threads, install the Ctrl-C-detaches handler and
raise on a non-COMPLETED stage were identical in both files.

Extract that into `follow_job(job_id, *, detach, success_marker=None) -> bool`,
returning True when the job finished and False when we stopped watching without a
verdict (detach or Ctrl-C). Training keeps its model-pushed marker by passing it in;
annotation has no equivalent line (the CLI keeps working after the upload log to
write the card and tag) so its completion stays stage-based.

Kept in hf.py rather than a new module so every existing monkeypatch target in
test_hf.py still resolves.

Behaviour change: a training run whose job reaches COMPLETED without the marker
matching now prints its completion line instead of returning silently. The marker
was already documented as an optimisation with a stage-based fallback; the fallback
just never reported success.

Tests: adds annotate coverage for the non-detach path (completion and failure) —
previously only ever exercised with detach=true — plus a detach short-circuit test.
Both new annotate tests verified to fail under a mutation that stubs out follow_job.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-23 10:30:33 +02:00
107 changed files with 6501 additions and 1447 deletions
+11
View File
@@ -0,0 +1,11 @@
version: 2
updates:
- package-ecosystem: "github-actions"
directory: "/"
schedule:
interval: "weekly"
cooldown:
default-days: 7
groups:
actions:
patterns: ["*"]
+2 -1
View File
@@ -51,6 +51,7 @@ pre-commit run --all-files # Lint + format (ruff, typo
## Notes
- **Mypy is gradual**: strict only for `lerobot.envs`, `lerobot.configs`, `lerobot.optim`, `lerobot.model`, `lerobot.cameras`, `lerobot.motors`, `lerobot.transport`. Add type annotations when modifying these modules.
- **Optional dependencies**: many policies, envs, and robots are behind extras (e.g., `lerobot[aloha]`). New imports for optional packages must be guarded or lazy. See `pyproject.toml [project.optional-dependencies]`.
- **Imports**: prefer top-level imports; relative (`from .sibling import X`) across sibling files within a module, absolute (`from lerobot.module import X`) across modules.
- **Optional dependencies**: many policies, envs, and robots are behind extras (e.g., `lerobot[aloha]`, see `pyproject.toml`). Guard optional imports with `TYPE_CHECKING or _foo_available` at module top + a `require_package(...)` check at use time. Reuse the `_foo_available` flags in `utils/import_utils.py`; don't call `is_package_available`.
- **Video decoding**: datasets can store observations as video files. `LeRobotDataset` handles frame extraction, but tests need ffmpeg installed.
- **Prioritize use of `uv run`** to execute Python commands (not raw `python` or `pip`).
+3 -3
View File
@@ -83,7 +83,7 @@ episode_index=0
print(f"{dataset[episode_index]['action'].shape=}\n")
```
Learn more about it in the [LeRobotDataset Documentation](https://huggingface.co/docs/lerobot/lerobot-dataset-v3)
Learn more about it in the [LeRobotDataset Documentation](https://huggingface.co/docs/lerobot/lerobot-dataset-v3).
## SoTA Models
@@ -109,7 +109,7 @@ lerobot-train \
| **World Models** | [VLA-JEPA](./docs/source/vla_jepa.mdx), [LingBot-VA](./docs/source/lingbot_va.mdx), [FastWAM](./docs/source/fastwam.mdx) |
| **Reward Models** | [SARM](./docs/source/sarm.mdx), [TOPReward](./docs/source/topreward.mdx), [Robometer](./docs/source/robometer.mdx) |
Similarly to the hardware, you can easily implement your own policy & leverage LeRobot's data collection, training, and visualization tools, and share your model to the HF Hub
Similarly to the hardware, you can easily implement your own policy & leverage LeRobot's data collection, training, and visualization tools, and share your model to the HF Hub.
For detailed policy setup guides, see the [Policy Documentation](https://huggingface.co/docs/lerobot/bring_your_own_policies). For GPU/RAM requirements and expected training time per policy, see the [Compute Hardware Guide](https://huggingface.co/docs/lerobot/hardware_guide).
@@ -126,7 +126,7 @@ lerobot-eval \
--eval.n_episodes=10
```
Learn how to implement your own simulation environment or benchmark and distribute it from the HF Hub by following the [EnvHub Documentation](https://huggingface.co/docs/lerobot/envhub)
Learn how to implement your own simulation environment or benchmark and distribute it from the HF Hub by following the [EnvHub Documentation](https://huggingface.co/docs/lerobot/envhub).
## Resources
+55 -16
View File
@@ -89,8 +89,8 @@ subtask.
The resulting spans are then stitched into a gap-free, full-episode
cover, so **every frame has exactly one active subtask**. See
[`run_hf_job.py`](https://github.com/huggingface/lerobot/blob/main/examples/annotations/run_hf_job.py)
for the production settings (single camera, timestamped contact sheets,
[Running on Hugging Face Jobs](#running-on-hugging-face-jobs) for the
production settings (single camera, timestamped contact sheets,
auto-windowed subtask generation).
### Tools
@@ -110,28 +110,67 @@ not-yet-implemented.
## Running on Hugging Face Jobs
Annotation runs on [Hugging Face Jobs](https://huggingface.co/docs/hub/en/jobs).
The repo ships a launcher script you copy and tweak for your dataset:
Annotating a real dataset needs a GPU big enough to serve the VLM, so
`lerobot-annotate` can dispatch itself to
[Hugging Face Jobs](https://huggingface.co/docs/hub/en/jobs) — same as
`lerobot-train`. Add `--job.target=<flavor>` to the exact command you'd
run locally and it runs on that hardware instead:
```bash
HF_TOKEN=hf_... uv run python examples/annotations/run_hf_job.py
hf auth login # once
uv run lerobot-annotate \
--repo_id=user/my_dataset \
--new_repo_id=user/my_dataset_annotated \
--push_to_hub=true \
--vlm.model_id=Qwen/Qwen3.6-27B \
--vlm.num_gpus=1 \
--vlm.serve_command="vllm serve Qwen/Qwen3.6-27B --tensor-parallel-size 1 \
--max-model-len 32768 --gpu-memory-utilization 0.8 \
--uvicorn-log-level warning --port {port}" \
--vlm.serve_ready_timeout_s=1800 \
--vlm.chat_template_kwargs='{"enable_thinking": false}' \
--job.target=h200
```
[`run_hf_job.py`](https://github.com/huggingface/lerobot/blob/main/examples/annotations/run_hf_job.py)
starts a single-GPU `h200` job (bump it to `h200x4` for big datasets)
that:
That submits a single-GPU `h200` job that:
1. installs `lerobot` (from `main`) plus the annotation extras,
2. boots one vLLM server per GPU (using the `vllm/vllm-openai` image) and
drives it over the OpenAI-compatible API,
3. runs the `plan` / `interjections` / `vqa` modules across the dataset
with `lerobot-annotate`,
1. starts from the `vllm/vllm-openai` image and installs `lerobot` on top,
2. boots one vLLM server per GPU and drives it over the OpenAI-compatible API,
3. runs the `plan` / `interjections` / `vqa` modules across the dataset,
4. with `--push_to_hub=true`, uploads the result to `--new_repo_id` (or
back to `--repo_id` in place if you leave that unset).
To use a different dataset, model, or hub repo, edit the `CMD` block in
the script. Every flag there maps directly to a `lerobot-annotate` flag
(run `lerobot-annotate --help` for the full list).
The command streams the job's logs; `Ctrl-C` detaches without cancelling
it. List the available flavors and their pricing with `hf jobs hardware`.
<Tip warning={true}>
Qwen3.6 ships with thinking enabled, which eats the token budget the
annotator needs for its JSON answer — `--vlm.chat_template_kwargs='{"enable_thinking": false}'`
turns it off. Without `--push_to_hub=true` the annotated dataset is
discarded when the pod exits.
</Tip>
### Job options
| Flag | Default | What it does |
| ------------------- | ------------------------- | ------------------------------------------------------------------------------- |
| `--job.target` | `local` | HF Jobs flavor to run on (e.g. `h200`, `h200x4`). Omitted/`local` runs here. |
| `--job.image` | `vllm/vllm-openai:latest` | Runtime image for the pod. |
| `--job.timeout` | `2h` | Wall-clock cap. Raise it for large datasets. |
| `--job.detach` | `false` | Submit and exit instead of streaming logs. |
| `--job.lerobot_ref` | `main` | Git ref of lerobot installed on the pod — point it at a branch to test changes. |
| `--job.tags` | `[]` | Extra tags on the job and on any dataset it pushes (`lerobot` is always added). |
For a bigger dataset, scale to `h200x4` and raise
`--vlm.parallel_servers` / `--vlm.num_gpus` to match, and give the job
more headroom with e.g. `--job.timeout=8h`.
Remote runs need `--repo_id` (the pod pulls the dataset from the Hub;
`--root` names a directory only your machine has). A dataset that exists
only in your local cache is pushed to a **private** repo first.
## Key options
+162 -1
View File
@@ -165,6 +165,8 @@ Batches are flat dictionaries keyed by the constants in [`lerobot.utils.constant
LeRobot uses `PolicyProcessorPipeline`s to normalize inputs and de-normalize outputs around your policy. For a concrete reference, see [`processor_act.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/act/processor_act.py) or [`processor_diffusion.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/diffusion/processor_diffusion.py).
Pay close attention here: processors are the most common reproducibility pain point. A mismatch in normalization mode (`IDENTITY` vs `MEAN_STD` vs `MIN_MAX` vs `QUANTILES`/`QUANTILE10`) or in which features get normalized will train and eval without erroring, yet silently wreck results. Make sure the modes match how the checkpoint was trained, that the required stats exist (e.g. `QUANTILES` needs `q01`/`q99`), and that the pre- and post-processors stay consistent.
```python
# processor_my_policy.py
from typing import Any
@@ -189,6 +191,162 @@ def make_my_policy_pre_post_processors(
---
## Adding high- and low-level language control
The policy API above is sufficient for training and standard evaluation. To use a language-conditioned policy with interactive `lerobot-rollout`, also register a runtime adapter. The adapter keeps policy-specific prompting and tokenization out of the generic control loop.
The runtime supports two policy shapes:
| Policy shape | Behavior | Adapter |
| ---------------- | ----------------------------------------------------------------------- | ---------------------------------------------- |
| Low-level / flat | The operator's task or subtask directly conditions action prediction. | Reuse `DirectTaskPolicyAdapter`. |
| High + low level | The policy generates subtasks or memory, then conditions actions on it. | Subclass `BaseLanguageAdapter`, as PI052 does. |
During a rollout, `RuntimeState` stores the high-level task and the active language context:
```text
task ──> adapter.generate_text("subtask", ...) ──> state.language_context["subtask"]
observation ──> processors ──> adapter.select_action() ─┴─> action chunk ──> robot
```
The generic runtime handles generation frequency, pause/resume, prompt replacement, action queues, and dispatch. The adapter only translates between that runtime contract and your policy.
### Low-level policies
If your policy already consumes the live task through its normal preprocessor and implements `predict_action_chunk`, register the shared direct adapter. PI0.5 and MolmoAct2 use this path:
```python
# src/lerobot/runtime/registry.py
_ADAPTERS = {
# ...
"my_policy": "lerobot.runtime.adapter:DirectTaskPolicyAdapter",
}
```
Run it with direct-subtask mode so the operator supplies the instruction used by the action policy:
```bash
lerobot-rollout \
--language \
--policy.path=user/my_policy_checkpoint \
--robot.type=so101_follower \
--robot.port=/dev/ttyACM0 \
--direct_subtask
```
The rollout context builds the observation batch with the current instruction before `DirectTaskPolicyAdapter` calls `policy.predict_action_chunk(observation)`. No text-generation method is required.
### Hierarchical policies
For a policy that generates language and actions, subclass [`BaseLanguageAdapter`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/runtime/adapter.py) and implement two methods:
- `generate_text(kind, observation, state, user_text=None) -> str` generates a `subtask`, `memory`, or interjection response.
- `select_action(observation, state)` builds the low-level prompt from the active context and returns an action chunk.
This abbreviated adapter follows [`PI052PolicyAdapter`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi052/inference/pi052_adapter.py):
```python
# inference/my_policy_adapter.py
from typing import Any
from lerobot.runtime import RuntimeState
from lerobot.runtime.adapter import BaseLanguageAdapter
from lerobot.utils.constants import (
OBS_LANGUAGE_ATTENTION_MASK,
OBS_LANGUAGE_TOKENS,
)
class MyPolicyAdapter(BaseLanguageAdapter):
def select_action(self, observation: dict[str, Any], state: RuntimeState):
instruction = state.language_context.get("subtask") or state.task or ""
tokens, attention_mask = tokenize_instruction(instruction)
batch = dict(observation)
batch[OBS_LANGUAGE_TOKENS] = tokens
batch[OBS_LANGUAGE_ATTENTION_MASK] = attention_mask
return self.policy.predict_action_chunk(batch)
def generate_text(
self,
kind: str,
observation: dict[str, Any] | None,
state: RuntimeState,
user_text: str | None = None,
) -> str:
messages = self.build_messages(kind, state, user_text)
batch, tokenizer = tokenize_messages(messages, observation)
return self.policy.select_message(
batch,
tokenizer=tokenizer,
min_new_tokens=self.gen.min_new_tokens,
temperature=self.gen.temperature,
top_p=self.gen.top_p,
)
def build_messages(
self, kind: str, state: RuntimeState, user_text: str | None
) -> list[dict[str, str]]:
if kind == "subtask":
return [{"role": "user", "content": state.task or ""}]
if kind == "memory":
return [
{"role": "user", "content": state.task or ""},
{
"role": "user",
"content": f"Completed subtask: {state.extra.get('prior_subtask', '')}",
},
]
if kind == "interjection":
return [
{"role": "user", "content": state.task or ""},
{"role": "user", "content": user_text or ""},
]
raise ValueError(f"Unsupported text kind: {kind}")
```
`tokenize_instruction` and `tokenize_messages` are policy-specific helpers. They must reproduce the prompt format used during training; PI052, for example, adds the discretized robot state to its low-level subtask prompt and uses the same PaliGemma formatting for `select_message`.
`BaseLanguageAdapter` provides the default hierarchy: regenerate a subtask at action-chunk boundaries, update memory when the subtask changes, and handle user interjections. Override `_regenerate_context` only if your policy uses a different hierarchy.
Register the adapter with a lazy import so importing LeRobot does not load the model or its optional dependencies:
```python
# src/lerobot/runtime/registry.py
_ADAPTERS = {
# ...
"my_policy": "lerobot.policies.my_policy.inference.my_policy_adapter:MyPolicyAdapter",
}
```
The key must match the policy's registered type. Once registered, the same checkpoint works through the shared entry point:
```bash
lerobot-rollout \
--language \
--policy.path=user/my_hierarchical_checkpoint \
--robot.type=so101_follower \
--robot.port=/dev/ttyACM0 \
--task="put the cup in the sink"
```
For RoboCasa-compatible policies, replace the robot arguments with `--sim --sim.task=<task>`. Without `--direct_subtask`, the adapter generates the low-level subtask; with it, the operator bypasses high-level generation and supplies each subtask.
### Keep training and deployment aligned
The adapter is intentionally small, but its prompts are part of the model contract:
- Use the same tokenizer, role formatting, special tokens, image ordering, and state encoding as training.
- Condition `select_action` on `state.language_context["subtask"]`, falling back to `state.task` for direct or not-yet-generated prompts.
- Return a full action chunk from `select_action`; the runtime handles control-rate dispatch.
- Keep optional model dependencies inside lazy imports.
- Test adapter selection, generated-message routing, action-batch construction, and direct-subtask behavior with a lightweight fake policy.
PI052 is the complete in-tree reference: its [processor](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi052/processor_pi052.py) renders the training recipe, its [policy](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi052/modeling_pi052.py) exposes text and action generation, and its [adapter](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi052/inference/pi052_adapter.py) reconstructs those same prompts at deployment.
---
## Path A: Out-of-tree plugin
The fastest way to ship a policy: package it as a standalone Python distribution and install it alongside LeRobot. No PR required, you own the release cycle, and you can publish to PyPI under your own namespace.
@@ -304,7 +462,9 @@ Mirror an existing policy that's structurally similar to yours; the diff is smal
### Heavy / optional dependencies
Most policies need a heavy backbone (transformers, diffusers, a specific VLM SDK). The convention is **two-step gating**: a `TYPE_CHECKING`-guarded import at module top, and a `require_package` runtime check in the constructor. [`modeling_diffusion.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/diffusion/modeling_diffusion.py) is the canonical reference:
Most policies need a heavy backbone (transformers, diffusers, a specific VLM SDK). Wherever one exists, prefer loading it e.g from `transformers` or `diffusers` rather than re-implementing the architecture in-tree.
The convention is **two-step gating**: a `TYPE_CHECKING`-guarded import at module top, and a `require_package` runtime check in the constructor. [`modeling_diffusion.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/diffusion/modeling_diffusion.py) is the canonical reference:
```python
from typing import TYPE_CHECKING
@@ -374,6 +534,7 @@ The general expectations are in [`CONTRIBUTING.md`](https://github.com/huggingfa
- [ ] Optional deps live behind a `[project.optional-dependencies]` extra and the `TYPE_CHECKING + require_package` guard.
- [ ] `tests/policies/` updated; backward-compat artifact committed & policy-specific tests.
- [ ] `src/lerobot/policies/<name>/README.md` symlinked into `docs/source/policy_<name>_README.md`; user-facing `docs/source/<name>.mdx` written and added to `_toctree.yml`.
- [ ] `lerobot-train --policy.type my_policy ...` runs end-to-end for at least a few steps + save a checkpoint that can be loaded and run by `lerobot-eval` or `lerobot-rollout`.
- [ ] `templates/lerobot_modelcard_template.md` has a description entry and a `policy_docs` link for your policy.
- [ ] The models table in the root `README.md` lists your policy in the right category, linking to your doc page.
- [ ] At least one reproducible benchmark eval in the policy MDX with a published checkpoint (sim benchmark, or real-robot dataset + checkpoint).
+47 -1
View File
@@ -1,6 +1,6 @@
# Policy Deployment (lerobot-rollout)
`lerobot-rollout` is the single CLI for deploying trained policies on real robots. It supports multiple execution strategies and inference backends, from quick evaluation to continuous recording and human-in-the-loop data collection.
`lerobot-rollout` is the single CLI for deploying trained policies on real robots or in an interactive simulator. It supports multiple execution strategies and inference backends, from quick evaluation to continuous recording, language-driven control, and human-in-the-loop data collection.
## Quick Start
@@ -197,6 +197,52 @@ Teleop is optional — if omitted the robot holds its position during the reset
---
## Interactive language control
Language-conditioned policies can expose a high-level text head in addition to
their action head. Add `--language` to open-prompt one of these policies on a
real robot. Language-only flags such as `--direct_subtask` select this mode
automatically.
MolmoAct2 has no high-level planner, so use direct-subtask mode and type each
next low-level instruction yourself:
```bash
lerobot-rollout \
--policy.path=lerobot/MolmoAct2-SO100_101-LeRobot \
--policy.device=cuda \
--robot.type=so101_follower \
--robot.port=/dev/ttyACM1 \
--robot.cameras='{"cam0":{"type":"opencv","index_or_path":"/dev/video0","width":640,"height":480,"fps":30,"fourcc":"MJPG","backend":200},"cam1":{"type":"opencv","index_or_path":"/dev/video2","width":640,"height":480,"fps":30,"fourcc":"MJPG","backend":200}}' \
--direct_subtask \
--robot.max_relative_target='{"shoulder_pan":5,"shoulder_lift":5,"elbow_flex":5,"wrist_flex":5,"wrist_roll":5,"gripper":5}'
```
The robot starts paused. Type a subtask, then use `/resume` and `/pause` to
control action dispatch. Check the workspace and motion limits before resuming.
Without `--direct_subtask`, a policy such as PI052 generates its active subtask
from the high-level `--task` itself.
RoboCasa uses the same runtime and processor path. `--sim` selects it
automatically, so no robot configuration is needed:
```bash
MUJOCO_GL=egl lerobot-rollout \
--policy.path=lerobot/pi052_robocasa \
--sim --sim.task=CloseFridge --sim.split=pretrain \
--task="close the fridge" \
--disable_memory \
--sim.render_size=384 \
--sim.views=robot0_agentview_left,robot0_eye_in_hand,robot0_agentview_right \
--mode=action --ctrl_hz=20
```
Open `http://localhost:8010` for the live simulator view. Add
`--sim.direct_subtask` to bypass the language planner and make each typed prompt
the action policy's current subtask.
---
## Inference Backends
Select a backend with `--inference.type=<name>`. All strategies work with both backends.
+11
View File
@@ -141,6 +141,17 @@ sample["target_message_indices"]
The renderer does not apply a tokenizer chat template. Policy processors decide how to serialize the messages for their backbone, which keeps the same dataset usable across SmolVLA, Pi0.5, and any future VLM that expects OpenAI-style chat messages.
## Blends
Blend recipes select one weighted sub-recipe deterministically from the sample index.
`recipes/subtask_mem.yaml` trains the compact core blend — high-level subtask prediction, low-level execution, and memory. `recipes/subtask_mem_vqa_speech.yaml` is the fuller variant that also adds VQA and spoken interjection responses.
A message recipe with a supervised assistant turn on the `low_level` stream trains
the π0.5 paper's joint sequence instead of a blend: the target span gets text CE
while also conditioning the action losses in the same forward.
`recipes/subtask_joint.yaml` is the provided example; pair it with
`--policy.joint_subtask_conditioning=true` at inference.
## Graceful absence
If both language columns are missing, `None`, or empty, `RenderMessagesStep` is a no-op.
+8
View File
@@ -1,3 +1,11 @@
# OMX
<img
src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/omx_mainimage.png"
alt="OMX"
width=600
/>
## Order and Assemble the parts
First, assemble the OMX hardware following the official assembly guide.
+4
View File
@@ -252,6 +252,10 @@ lerobot-dataset-viz \
--episode-index 0
```
For a private or gated dataset, authenticate first with `hf auth login`, or set the
`HF_TOKEN` environment variable. The Hub client then discovers the credential
automatically; no token argument is needed.
**From a local folder:**
Add the `--root` option and set `--mode local`. For example, to search in `./my_local_data_dir/lerobot/pusht`:
-80
View File
@@ -1,80 +0,0 @@
#!/usr/bin/env python
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Launch ``lerobot-annotate`` on a Hugging Face job (vllm + Qwen3.6-27B VLM).
Spawns one single-GPU ``h200`` job that:
1. installs ``lerobot`` from ``main`` plus the annotation extras,
2. boots one vllm server with Qwen3.6-27B (dense VLM),
3. runs the plan / interjections / vqa modules across the dataset
in free-form mode (each episode generates its own subtasks +
memory),
4. uploads the annotated dataset to ``--new_repo_id`` (when set)
or back to ``--repo_id``.
Usage:
HF_TOKEN=hf_... uv run python examples/annotations/run_hf_job.py
Adjust ``CMD`` (dataset, model, hub repo) and ``flavor`` below for your
run. For larger datasets, scale to ``h200x4`` and raise
``--vlm.parallel_servers`` / ``--vlm.num_gpus`` to match.
"""
import os
from huggingface_hub import get_token, run_job
token = os.environ.get("HF_TOKEN") or get_token()
if not token:
raise RuntimeError("No HF token. Run `huggingface-cli login` or `export HF_TOKEN=hf_...`")
CMD = (
"apt-get update -qq && apt-get install -y -qq git ffmpeg && "
"pip install --no-deps "
"'lerobot @ git+https://github.com/huggingface/lerobot.git@main' && "
# Pins mirror pyproject.toml — unpinned installs pull av 18 / datasets 5 /
# draccus 0.11, which break lerobot at import time.
"pip install --upgrade-strategy only-if-needed "
"'datasets>=4.7.0,<5.0.0' 'pyarrow>=21.0.0,<30.0.0' 'av>=15.0.0,<16.0.0' 'draccus==0.10.0' "
"'pandas>=2.0.0,<3.0.0' jsonlines gymnasium torchcodec mergedeep pyyaml-include toml typing-inspect "
"openai && "
"export VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0 && "
"export VLLM_VIDEO_BACKEND=pyav && "
"lerobot-annotate "
"--repo_id=pepijn223/robocasa_pretrain_human300_v4 "
"--new_repo_id=pepijn223/robocasa_pretrain_human300_v4_annotated "
"--push_to_hub=true "
"--vlm.backend=openai "
"--vlm.model_id=Qwen/Qwen3.6-27B "
"--vlm.num_gpus=1 "
'--vlm.serve_command="vllm serve Qwen/Qwen3.6-27B '
"--tensor-parallel-size 1 --max-model-len 32768 "
'--gpu-memory-utilization 0.8 --uvicorn-log-level warning --port {port}" '
"--vlm.serve_ready_timeout_s=1800 "
# Qwen3.6 ships with thinking on; annotation wants plain JSON answers.
"--vlm.chat_template_kwargs='{\"enable_thinking\": false}'"
)
job = run_job(
image="vllm/vllm-openai:latest",
command=["bash", "-c", CMD],
flavor="h200",
secrets={"HF_TOKEN": token},
timeout="2h",
)
print(f"Job URL: {job.url}")
print(f"Job ID: {job.id}")
+1 -1
View File
@@ -155,7 +155,7 @@ accelerate-dep = ["accelerate>=1.14.0,<2.0.0"]
can-dep = ["python-can>=4.2.0,<5.0.0"]
peft-dep = ["peft>=0.18.0,<1.0.0"]
scipy-dep = ["scipy>=1.14.0,<2.0.0"]
diffusers-dep = ["diffusers>=0.27.2,<0.36.0"]
diffusers-dep = ["diffusers>=0.38.0,<0.40.0"]
qwen-vl-utils-dep = ["qwen-vl-utils>=0.0.11,<0.1.0"]
matplotlib-dep = ["matplotlib>=3.10.3,<4.0.0", "contourpy>=1.3.0,<2.0.0"] # NOTE: Explicitly listing contourpy helps the resolver converge faster.
pyserial-dep = ["pyserial>=3.5,<4.0"]
@@ -20,6 +20,29 @@ from dataclasses import dataclass, field
from pathlib import Path
from typing import Any
from lerobot.configs.default import JobConfig
# The annotation pipeline boots its own vLLM server, so the pod starts from the
# official vLLM runtime rather than the prebuilt `lerobot-gpu` training image;
# `lerobot` is pip-installed on top (see `lerobot.jobs.annotate`).
DEFAULT_ANNOTATE_JOB_IMAGE = "vllm/vllm-openai:latest"
@dataclass
class AnnotationJobConfig(JobConfig):
"""`JobConfig` with the annotation runtime's defaults.
Adds `lerobot_ref` because the vLLM image ships no lerobot: the pod installs
it from git, and the ref decides which code actually annotates. Point it at a
branch/tag/SHA to try unmerged changes remotely.
"""
image: str = DEFAULT_ANNOTATE_JOB_IMAGE
# Annotation is a bounded pass over a dataset; a tighter cap than training's
# "2d" keeps a wedged vLLM server from burning a day of GPU time.
timeout: str | None = "2h"
lerobot_ref: str = "main"
@dataclass
class PlanConfig:
@@ -207,6 +230,11 @@ class AnnotationPipelineConfig:
vlm: VlmConfig = field(default_factory=VlmConfig)
executor: ExecutorConfig = field(default_factory=ExecutorConfig)
# Where the annotation runs: omitted / "local" annotates on this machine, any
# other value is an HF Jobs flavor (e.g. "h200") and submits the run there.
# List flavors + pricing with `hf jobs hardware`.
job: AnnotationJobConfig = field(default_factory=AnnotationJobConfig)
skip_validation: bool = False
only_episodes: tuple[int, ...] | None = None
@@ -30,8 +30,8 @@ Phase 3 is why the ``plan`` module must be re-entered after the
timestamps.
Distributed execution is provided by Hugging Face Jobs (see
``examples/annotations/run_hf_job.py``); the runner inside the job
invokes ``lerobot-annotate`` which uses this in-process executor.
``lerobot.jobs.annotate``, reached via ``--job.target=<flavor>``); the pod
inside the job invokes ``lerobot-annotate`` which uses this in-process executor.
Episode-level concurrency is controlled by
``ExecutorConfig.episode_parallelism``.
"""
@@ -194,12 +194,13 @@ def make_vlm_client(config: VlmConfig) -> VlmClient:
"""Build the shared VLM client.
Only the ``openai`` backend is supported for now. The shipped workflow
is Hugging Face Jobs (``examples/annotations/run_hf_job.py``): it boots
a vLLM server inside the ``vllm/vllm-openai`` image and the pipeline
talks to it over the OpenAI-compatible API (``--vlm.backend=openai``,
optionally auto-spawning the server via ``auto_serve`` /
``serve_command``). The former in-process ``vllm`` / ``transformers``
backends were removed to keep the support surface to the HF Jobs path.
is Hugging Face Jobs (``lerobot-annotate --job.target=<flavor>``): it
boots a vLLM server inside the ``vllm/vllm-openai`` image and the
pipeline talks to it over the OpenAI-compatible API
(``--vlm.backend=openai``, optionally auto-spawning the server via
``auto_serve`` / ``serve_command``). The former in-process ``vllm`` /
``transformers`` backends were removed to keep the support surface to
the HF Jobs path.
For ``stub``, construct :class:`StubVlmClient` directly with a responder
callable; it is rejected here to make accidental misuse obvious.
@@ -213,8 +214,8 @@ def make_vlm_client(config: VlmConfig) -> VlmClient:
if config.backend in {"vllm", "transformers"}:
raise ValueError(
f"backend={config.backend!r} (in-process local model) is not supported for now — "
"only backend='openai' (the Hugging Face Jobs flow) is. Run the pipeline via "
"examples/annotations/run_hf_job.py, which serves the model with vLLM in the "
"only backend='openai' (the Hugging Face Jobs flow) is. Run the pipeline with "
"`lerobot-annotate --job.target=<flavor>`, which serves the model with vLLM in the "
"vllm/vllm-openai image and talks to it over the OpenAI-compatible API."
)
raise ValueError(f"Unknown VLM backend: {config.backend!r}")
@@ -173,7 +173,8 @@ class Reachy2Camera(Camera):
raise ValueError(
f"Invalid color mode '{self.color_mode}'. Expected {ColorMode.RGB} or {ColorMode.BGR}."
)
if self.color_mode == ColorMode.RGB:
is_depth_frame = self.config.name == "depth" and self.config.image_type == "depth"
if not is_depth_frame and self.color_mode == ColorMode.RGB:
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
self.latest_frame = frame
@@ -453,7 +453,7 @@ class RealSenseCamera(Camera):
)
processed_image = image
if self.color_mode == ColorMode.BGR:
if not depth_frame and self.color_mode == ColorMode.BGR:
processed_image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE, cv2.ROTATE_180]:
+6
View File
@@ -33,6 +33,8 @@ class DatasetConfig:
# looked up under $HF_LEROBOT_HOME/repo_id and Hub downloads use a revision-safe cache under $HF_LEROBOT_HOME/hub.
root: str | None = None
episodes: list[int] | None = None
# Episode indices to drop (e.g. corrupt or heterogeneous ones). Applied on top of `episodes`.
exclude_episodes: list[int] | None = None
image_transforms: ImageTransformsConfig = field(default_factory=ImageTransformsConfig)
revision: str | None = None
use_imagenet_stats: bool = True
@@ -62,6 +64,10 @@ class DatasetConfig:
if len(self.episodes) != len(set(self.episodes)):
duplicates = sorted({ep for ep in self.episodes if self.episodes.count(ep) > 1})
raise ValueError(f"Episode indices contain duplicates: {duplicates}")
if self.exclude_episodes is not None and any(ep < 0 for ep in self.exclude_episodes):
raise ValueError(
f"exclude_episodes must be non-negative, got: {[ep for ep in self.exclude_episodes if ep < 0]}"
)
@dataclass
+10 -4
View File
@@ -78,7 +78,7 @@ class MessageTurn:
raise ValueError(f"Unsupported message stream: {self.stream!r}")
if self.content is None and self.tool_calls_from is None:
raise ValueError("MessageTurn.content is required unless tool_calls_from is set.")
if self.content is not None and not isinstance(self.content, (str, list)):
if self.content is not None and not isinstance(self.content, str | list):
raise TypeError("MessageTurn.content must be a string, a list of HF-style blocks, or None.")
if isinstance(self.content, list):
for block in self.content:
@@ -147,7 +147,7 @@ class TrainingRecipe:
return cls.from_dict(data)
def _validate_message_recipe(self) -> None:
"""Ensure every templated binding is known and at least one turn is a target."""
"""Validate bindings and require text or low-level action supervision."""
assert self.messages is not None
known_bindings = set(DEFAULT_BINDINGS) | set(self.bindings or {}) | {"task"}
@@ -156,8 +156,14 @@ class TrainingRecipe:
if missing:
raise ValueError(f"MessageTurn references unknown binding(s): {sorted(missing)}")
if not any(turn.target for turn in self.messages):
raise ValueError("Message recipes must contain at least one target turn.")
has_target = any(turn.target for turn in self.messages)
has_low_level = any(turn.stream == "low_level" for turn in self.messages)
if not (has_target or has_low_level):
raise ValueError(
"Message recipes must contain at least one supervised turn — "
"either ``target: true`` (text CE) or ``stream: low_level`` "
"(flow/action loss)."
)
def _validate_blend_recipe(self) -> None:
"""Ensure each blend component is a non-empty, weighted message recipe."""
+16
View File
@@ -0,0 +1,16 @@
# Predicts subtasks from tasks and trains subtask-conditioned action flow without memory or plans.
# Requires `subtask` annotations; samples with missing `if_present` bindings do not render.
blend:
high_level_subtask:
weight: 0.30
messages:
- {role: user, content: "${task}", stream: high_level}
- {role: assistant, content: "${subtask}", stream: high_level, target: true, if_present: subtask}
low_level_execution:
weight: 0.70
messages:
# The low-level stream trains action flow on the generated or annotated subtask.
- {role: user, content: "${subtask}", stream: low_level, if_present: subtask}
@@ -0,0 +1,13 @@
# Paper-style joint sequence (pi0.5 §IV-B): one sample supervises the subtask
# text with CE and, because the assistant turn is part of the prefix, conditions
# the FAST and flow action losses on the same annotated subtask in one forward.
# The supervised span is attended causally; the action losses see task + subtask.
#
# Pair with `--policy.joint_subtask_conditioning=true` at inference so the flow
# prefix reproduces this layout (task turn with state + causal generated subtask).
# Samples without a `subtask` annotation fall back to a plain task-prompt
# low-level sample via `if_present`.
messages:
- {role: user, content: "${task}", stream: low_level}
- {role: assistant, content: "${subtask}", stream: low_level, target: true, if_present: subtask}
@@ -0,0 +1,30 @@
# Trains subtask prediction, subtask-conditioned action flow, and memory updates without plans.
# Requires `subtask` and `memory`; missing `if_present` bindings skip the affected sub-recipe.
blend:
high_level_subtask:
weight: 0.25
messages:
- {role: user, content: "${task}", stream: high_level}
- {role: assistant, content: "${subtask}", stream: high_level, target: true, if_present: subtask}
low_level_execution:
weight: 0.60
messages:
# The low-level stream trains action flow on the generated or annotated subtask.
- {role: user, content: "${subtask}", stream: low_level, if_present: subtask}
memory_update:
# `active_at` densifies sparse boundaries while preserving the prior-memory/subtask mapping.
# Inference controls update timing through `subtask_change` events.
weight: 0.15
bindings:
prior_memory: "nth_prev(style=memory, offset=1)"
current_memory: "active_at(t, style=memory)"
completed_subtask: "nth_prev(style=subtask, offset=1)"
messages:
- {role: user, content: "${task}", stream: high_level}
- {role: assistant, content: "Previous memory: ${prior_memory}", stream: high_level, if_present: prior_memory}
- {role: user, content: "Completed subtask: ${completed_subtask}", stream: high_level, if_present: completed_subtask}
- {role: assistant, content: "${current_memory}", stream: high_level, target: true, if_present: current_memory}
@@ -0,0 +1,70 @@
# Adds memory, spoken interjection responses, and camera-grounded VQA to subtask/action training.
# Missing optional annotations skip only their sub-recipe; `say` tool calls tokenize as `<say>...</say>`.
blend:
high_level_subtask:
weight: 0.25
messages:
- {role: user, content: "${task}", stream: high_level}
- {role: assistant, content: "${subtask}", stream: high_level, target: true, if_present: subtask}
low_level_execution:
weight: 0.40
messages:
# The low-level stream trains action flow on the generated or annotated subtask.
- {role: user, content: "${subtask}", stream: low_level, if_present: subtask}
memory_update:
# `active_at` densifies sparse boundaries while preserving the prior-memory/subtask mapping.
# Inference controls update timing through `subtask_change` events.
weight: 0.10
bindings:
prior_memory: "nth_prev(style=memory, offset=1)"
current_memory: "active_at(t, style=memory)"
completed_subtask: "nth_prev(style=subtask, offset=1)"
messages:
- {role: user, content: "${task}", stream: high_level}
- {role: assistant, content: "Previous memory: ${prior_memory}", stream: high_level, if_present: prior_memory}
- {role: user, content: "Completed subtask: ${completed_subtask}", stream: high_level, if_present: completed_subtask}
- {role: assistant, content: "${current_memory}", stream: high_level, target: true, if_present: current_memory}
user_interjection_response:
weight: 0.10
bindings:
interjection: "emitted_at(t, style=interjection)"
speech: "emitted_at(t, role=assistant, tool_name=say)"
messages:
- {role: user, content: "${task}", stream: high_level}
- {role: user, content: "${interjection}", stream: high_level, if_present: interjection}
# The assistant target is a `say` tool call flattened to a `<say>...</say>` marker.
- {role: assistant, stream: high_level, target: true, if_present: speech, tool_calls_from: speech}
# Each camera uses a separate VQA sub-recipe for view-specific binding.
ask_vqa_top:
weight: 0.075
bindings:
vqa_query: "emitted_at(t, style=vqa, role=user, camera=observation.images.front)"
vqa: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.front)"
messages:
- role: user
stream: high_level
if_present: vqa_query
content:
- {type: image, feature: observation.images.front}
- {type: text, text: "${vqa_query}"}
- {role: assistant, content: "${vqa}", stream: high_level, target: true, if_present: vqa}
ask_vqa_wrist:
weight: 0.075
bindings:
vqa_query: "emitted_at(t, style=vqa, role=user, camera=observation.images.wrist)"
vqa: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.wrist)"
messages:
- role: user
stream: high_level
if_present: vqa_query
content:
- {type: image, feature: observation.images.wrist}
- {type: text, text: "${vqa_query}"}
- {role: assistant, content: "${vqa}", stream: high_level, target: true, if_present: vqa}
+18
View File
@@ -14,6 +14,7 @@
import builtins
import datetime as dt
import json
import multiprocessing
import os
import tempfile
from dataclasses import dataclass, field
@@ -101,6 +102,12 @@ class TrainPipelineConfig(HubMixin):
batch_size: int = 8
prefetch_factor: int = 4
persistent_workers: bool = True
# DataLoader worker start method. "spawn" is safer than "fork" with
# non-fork-safe libs (PyAV / torchcodec / ffmpeg), but adds some
# worker-startup time per run since workers re-import modules instead
# of inheriting parent state. Override with `--dataloader_multiprocessing_context=fork`
# when appropriate, or set it to `null` to use Python's platform default.
dataloader_multiprocessing_context: str | None = "spawn"
steps: int = 100_000
# Run policy in the simulation environment every N steps to measure reward/success (0 = disabled).
env_eval_freq: int = 20_000
@@ -212,6 +219,17 @@ class TrainPipelineConfig(HubMixin):
self.reward_model.pretrained_path = str(policy_dir)
def validate(self) -> None:
available_contexts = multiprocessing.get_all_start_methods()
if (
self.dataloader_multiprocessing_context is not None
and self.dataloader_multiprocessing_context not in available_contexts
):
raise ValueError(
"`dataloader_multiprocessing_context` must be None or one of "
f"{available_contexts} on this platform, got "
f"{self.dataloader_multiprocessing_context!r}."
)
self._resolve_pretrained_from_cli()
if self.policy is None and self.reward_model is None:
+16 -2
View File
@@ -73,6 +73,8 @@ class LeRobotDatasetMetadata:
revision: str | None = None,
force_cache_sync: bool = False,
metadata_buffer_size: int = 10,
*,
token: str | bool | None = None,
):
"""Load or download metadata for an existing LeRobot dataset.
@@ -94,6 +96,10 @@ class LeRobotDatasetMetadata:
even when local files exist.
metadata_buffer_size: Number of episode metadata records to buffer
in memory before flushing to parquet.
token: Authentication token used for Hub requests. Pass a string
token, ``True`` to require the locally stored token, ``False``
to disable authentication, or ``None`` to use the Hugging Face
Hub default.
"""
self.repo_id = repo_id
self.revision = revision if revision else CODEBASE_VERSION
@@ -113,9 +119,12 @@ class LeRobotDatasetMetadata:
self._load_metadata()
except (FileNotFoundError, NotADirectoryError):
if is_valid_version(self.revision):
self.revision = get_safe_version(self.repo_id, self.revision)
if token is None:
self.revision = get_safe_version(self.repo_id, self.revision)
else:
self.revision = get_safe_version(self.repo_id, self.revision, token=token)
self._pull_from_repo(allow_patterns="meta/")
self._pull_from_repo(allow_patterns="meta/", token=token)
self._load_metadata()
def _flush_metadata_buffer(self) -> None:
@@ -220,7 +229,10 @@ class LeRobotDatasetMetadata:
self,
allow_patterns: list[str] | str | None = None,
ignore_patterns: list[str] | str | None = None,
*,
token: str | bool | None = None,
) -> None:
token_kwargs = {} if token is None else {"token": token}
if self._requested_root is None:
self.root = Path(
snapshot_download(
@@ -230,6 +242,7 @@ class LeRobotDatasetMetadata:
cache_dir=HF_LEROBOT_HUB_CACHE,
allow_patterns=allow_patterns,
ignore_patterns=ignore_patterns,
**token_kwargs,
)
)
return
@@ -242,6 +255,7 @@ class LeRobotDatasetMetadata:
local_dir=self._requested_root,
allow_patterns=allow_patterns,
ignore_patterns=ignore_patterns,
**token_kwargs,
)
self.root = self._requested_root
+30
View File
@@ -163,10 +163,40 @@ class DatasetReader:
def _load_hf_dataset(self) -> datasets.Dataset:
"""hf_dataset contains all the observations, states, actions, rewards, etc."""
features = get_hf_features_from_features(self._meta.features)
# Annotated datasets may have language columns absent from metadata.
# Extend the schema before the strict Parquet cast.
features = self._extend_features_with_language_columns(features)
hf_dataset = load_nested_dataset(self.root / "data", features=features, episodes=self.episodes)
hf_dataset.set_transform(hf_transform_to_torch)
return hf_dataset
def _extend_features_with_language_columns(self, features: datasets.Features) -> datasets.Features:
"""Register language columns found in Parquet but missing from metadata."""
# Leave empty datasets to fail through the normal loading path.
try:
sample = next((self.root / "data").glob("*/*.parquet"))
except StopIteration:
return features
from pyarrow import parquet as _pq # noqa: PLC0415
schema_names = set(_pq.read_schema(sample).names)
from .language import ( # noqa: PLC0415
LANGUAGE_EVENTS,
LANGUAGE_PERSISTENT,
language_events_column_feature,
language_persistent_column_feature,
)
extra: dict[str, object] = {}
if LANGUAGE_PERSISTENT in schema_names and LANGUAGE_PERSISTENT not in features:
extra[LANGUAGE_PERSISTENT] = language_persistent_column_feature()
if LANGUAGE_EVENTS in schema_names and LANGUAGE_EVENTS not in features:
extra[LANGUAGE_EVENTS] = language_events_column_feature()
if not extra:
return features
return datasets.Features({**features, **extra})
def _check_cached_episodes_sufficient(self) -> bool:
"""Check if the cached dataset contains all requested episodes and their video files."""
if self.hf_dataset is None or len(self.hf_dataset) == 0:
+23 -14
View File
@@ -172,6 +172,23 @@ class DatasetWriter:
def _get_image_file_dir(self, episode_index: int, image_key: str) -> Path:
return self._get_image_file_path(episode_index, image_key, frame_index=0).parent
def _get_episode_buffer_index(self) -> int:
episode_index = self.episode_buffer["episode_index"]
# episode_index is `int` when freshly created, but becomes `np.ndarray` after
# save_episode() mutates the buffer. Handle both types here.
if isinstance(episode_index, np.ndarray):
episode_index = episode_index.item() if episode_index.size == 1 else episode_index[0]
return int(episode_index)
def _delete_camera_frame_dirs(self, camera_keys: list[str]) -> None:
if self.image_writer is not None:
self._wait_image_writer()
episode_index = self._get_episode_buffer_index()
for camera_key in camera_keys:
img_dir = self._get_image_file_dir(episode_index, camera_key)
if img_dir.is_dir():
shutil.rmtree(img_dir)
def _save_image(
self, image: torch.Tensor | np.ndarray | PIL.Image.Image, fpath: Path, compress_level: int = 1
) -> None:
@@ -369,7 +386,9 @@ class DatasetWriter:
self._episodes_since_last_encoding = 0
if episode_data is None:
self.clear_episode_buffer(delete_images=len(self._meta.image_keys) > 0)
if len(self._meta.image_keys) > 0:
self._delete_camera_frame_dirs(self._meta.image_keys)
self.episode_buffer = self._create_episode_buffer()
def _batch_save_episode_video(self, start_episode: int, end_episode: int | None = None) -> None:
"""Batch save videos for multiple episodes."""
@@ -561,10 +580,10 @@ class DatasetWriter:
return metadata
def clear_episode_buffer(self, delete_images: bool = True) -> None:
"""Discard the current episode buffer and optionally delete temp images.
"""Discard the current episode buffer and optionally delete temp camera frames.
Args:
delete_images: If ``True``, remove temporary image directories
delete_images: If ``True``, remove temporary camera frame directories
written for the current episode.
"""
# Cancel streaming encoder if active
@@ -572,17 +591,7 @@ class DatasetWriter:
self._streaming_encoder.cancel_episode()
if delete_images:
if self.image_writer is not None:
self._wait_image_writer()
episode_index = self.episode_buffer["episode_index"]
# episode_index is `int` when freshly created, but becomes `np.ndarray` after
# save_episode() mutates the buffer. Handle both types here.
if isinstance(episode_index, np.ndarray):
episode_index = episode_index.item() if episode_index.size == 1 else episode_index[0]
for cam_key in self._meta.image_keys:
img_dir = self._get_image_file_dir(episode_index, cam_key)
if img_dir.is_dir():
shutil.rmtree(img_dir)
self._delete_camera_frame_dirs(self._meta.camera_keys)
self.episode_buffer = self._create_episode_buffer()
+16 -2
View File
@@ -66,6 +66,17 @@ def resolve_delta_timestamps(
return delta_timestamps
def _resolve_episodes(
episodes: list[int] | None, exclude_episodes: list[int] | None, total_episodes: int
) -> list[int] | None:
"""Apply an episode exclusion list on top of an optional allowlist."""
if not exclude_episodes:
return episodes
base = episodes if episodes is not None else list(range(total_episodes))
excluded = set(exclude_episodes)
return [episode for episode in base if episode not in excluded]
def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDataset:
"""Handles the logic of setting up delta timestamps and image transforms before creating a dataset.
@@ -87,11 +98,14 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
cfg.dataset.repo_id, root=cfg.dataset.root, revision=cfg.dataset.revision
)
delta_timestamps = resolve_delta_timestamps(cfg.trainable_config, ds_meta)
episodes = _resolve_episodes(
cfg.dataset.episodes, cfg.dataset.exclude_episodes, ds_meta.total_episodes
)
if not cfg.dataset.streaming:
dataset = LeRobotDataset(
cfg.dataset.repo_id,
root=cfg.dataset.root,
episodes=cfg.dataset.episodes,
episodes=episodes,
delta_timestamps=delta_timestamps,
image_transforms=image_transforms,
revision=cfg.dataset.revision,
@@ -104,7 +118,7 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
dataset = StreamingLeRobotDataset(
cfg.dataset.repo_id,
root=cfg.dataset.root,
episodes=cfg.dataset.episodes,
episodes=episodes,
delta_timestamps=delta_timestamps,
image_transforms=image_transforms,
revision=cfg.dataset.revision,
+73 -10
View File
@@ -162,14 +162,28 @@ def render_sample(
task: str | None = None,
dataset_ctx: Any | None = None,
) -> RenderedMessages | None:
"""Render the chat-style messages for a single dataset sample.
"""Resolve one sample's bindings and render its message recipe.
Resolves the recipe's bindings against ``persistent`` and ``events`` rows
at frame timestamp ``t``, then expands the recipe's message templates.
Returns ``None`` if the resolved sample contains no target message.
Returns ``None`` when no text or low-level action supervision applies.
"""
persistent_rows = _normalize_rows(persistent or [])
event_rows = _normalize_rows(events or [])
# Route sparse VQA frames to a matching view-specific component before weighted selection.
# This avoids dropping annotated frames or selecting VQA without annotations.
if recipe.blend is not None:
vqa_rendered = _render_vqa_if_present(
recipe,
persistent=persistent_rows,
events=event_rows,
t=t,
sample_idx=sample_idx,
task=task,
dataset_ctx=dataset_ctx,
)
if vqa_rendered is not None:
return vqa_rendered
selected_recipe = _select_recipe(recipe, sample_idx)
bindings = _resolve_bindings(
selected_recipe,
@@ -183,6 +197,55 @@ def render_sample(
return _render_message_recipe(selected_recipe, bindings)
def _render_vqa_if_present(
recipe: TrainingRecipe,
*,
persistent: Sequence[LanguageRow],
events: Sequence[LanguageRow],
t: float,
sample_idx: int,
task: str | None,
dataset_ctx: Any | None,
) -> RenderedMessages | None:
"""Render a matching VQA component, or return ``None`` for normal selection.
Multiple matching views are selected deterministically by relative weight.
"""
assert recipe.blend is not None
renderable: list[tuple[float, RenderedMessages]] = []
for name, component in recipe.blend.items():
if not name.startswith("ask_vqa"):
continue
bindings = _resolve_bindings(
component,
persistent=persistent,
events=events,
t=t,
sample_idx=sample_idx,
task=task,
dataset_ctx=dataset_ctx,
)
rendered = _render_message_recipe(component, bindings)
if rendered is not None:
renderable.append((float(component.weight or 0.0), rendered))
if not renderable:
return None
if len(renderable) == 1:
return renderable[0][1]
# Choose among matching cameras by relative weight, or uniformly when all weights are zero.
total = sum(w for w, _ in renderable) or float(len(renderable))
digest = hashlib.blake2b(f"vqa:{sample_idx}".encode(), digest_size=8).digest()
draw = int.from_bytes(digest, "big") / 2**64 * total
cumulative = 0.0
for w, rendered in renderable:
cumulative += w or (total / len(renderable))
if draw < cumulative:
return rendered
return renderable[-1][1]
def _select_recipe(recipe: TrainingRecipe, sample_idx: int) -> TrainingRecipe:
"""Pick a deterministic blend component for ``sample_idx`` (or return ``recipe``)."""
if recipe.blend is None:
@@ -346,7 +409,9 @@ def _render_message_recipe(
if turn.target:
target_indices.append(message_idx)
if not target_indices:
# Keep samples with either text targets or low-level action supervision.
has_low_level = any(stream == "low_level" for stream in streams)
if not target_indices and not has_low_level:
return None
rendered = {
@@ -403,14 +468,12 @@ def _validate_rendered(rendered: RenderedMessages) -> None:
if len(streams) != len(messages):
raise ValueError("message_streams must be aligned with messages.")
if not target_indices:
raise ValueError("Rendered samples must contain at least one target message.")
# Require text or low-level action supervision.
if not target_indices and not any(s == "low_level" for s in streams):
raise ValueError("Rendered samples must contain a target message or a low_level-stream message.")
for idx in target_indices:
if idx < 0 or idx >= len(messages):
raise ValueError(f"Target message index {idx} is out of bounds.")
# ``stream`` is enforced non-None at MessageTurn construction time
# (see ``MessageTurn.__post_init__``), so a missing stream here would
# mean the dataclass invariant was bypassed; no need to re-check.
def _nth_relative(
+39 -10
View File
@@ -65,6 +65,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
encoder_threads: int | None = None,
streaming_encoding: bool = False,
encoder_queue_maxsize: int = 30,
*,
token: str | bool | None = None,
):
"""
2 modes are available for instantiating this class, depending on 2 different use cases:
@@ -197,6 +199,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
instead of writing PNG images first. This makes save_episode() near-instant. Defaults to False.
encoder_queue_maxsize (int, optional): Maximum number of frames to buffer per camera when using
streaming encoding. Defaults to 30 (~1s at 30fps).
token: Authentication token used while downloading this dataset
from the Hub. Pass a string token, ``True`` to require the
locally stored token, ``False`` to disable authentication, or
``None`` to use the Hugging Face Hub default. The token is not
retained on the dataset instance after initialization.
Note:
Write-mode parameters (``streaming_encoding``, ``batch_encoding_size``) passed to
@@ -220,7 +227,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
# Load metadata (sets self.root once from the resolved metadata root)
self.meta = LeRobotDatasetMetadata(
self.repo_id, self._requested_root, self.revision, force_cache_sync=force_cache_sync
self.repo_id,
self._requested_root,
self.revision,
force_cache_sync=force_cache_sync,
token=token,
)
self.root = self.meta.root
self.revision = self.meta.revision
@@ -260,8 +271,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
# Load actual data
if force_cache_sync or not self.reader.try_load():
if is_valid_version(self.revision):
self.revision = get_safe_version(self.repo_id, self.revision)
self._download(download_videos)
if token is None:
self.revision = get_safe_version(self.repo_id, self.revision)
else:
self.revision = get_safe_version(self.repo_id, self.revision, token=token)
self._download(download_videos, token=token)
self.reader.load_and_activate()
# Detect write-mode params for backward compatibility
@@ -478,18 +492,19 @@ class LeRobotDataset(torch.utils.data.Dataset):
"""Return the number of frames in the selected episodes."""
return self.num_frames
def __getitem__(self, idx) -> dict:
"""Return a single frame by index, with all transforms applied.
def __getitem__(self, idx: int | slice) -> dict | list[dict]:
"""Return one frame or a slice of frames, with all transforms applied.
Loads the frame from the underlying HF dataset, expands delta-timestamp
windows, decodes video frames, and applies image transforms. Delegates
the core logic to :meth:`DatasetReader.get_item`.
the core logic to :class:`DatasetReader`.
Args:
idx: Index into the (possibly episode-filtered) dataset.
idx: Integer index or slice into the possibly episode-filtered dataset.
Returns:
Dict mapping feature names to their tensor values for this frame.
A frame dictionary for an integer index, or a list of frame
dictionaries for a slice.
Raises:
RuntimeError: If the dataset is currently being recorded and
@@ -499,6 +514,9 @@ class LeRobotDataset(torch.utils.data.Dataset):
raise RuntimeError(
"Cannot read from a dataset that is being recorded. Call finalize() first, then access items."
)
if isinstance(idx, slice):
return [self[item_idx] for item_idx in range(*idx.indices(len(self)))]
reader = self._ensure_reader()
if reader.hf_dataset is None:
# One-shot load after finalize()
@@ -622,10 +640,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
hub_api.delete_tag(self.repo_id, tag=CODEBASE_VERSION, repo_type="dataset")
hub_api.create_tag(self.repo_id, tag=CODEBASE_VERSION, revision=branch, repo_type="dataset")
def _download(self, download_videos: bool = True) -> None:
def _download(self, download_videos: bool = True, *, token: str | bool | None = None) -> None:
"""Downloads the dataset from the given 'repo_id' at the provided version."""
ignore_patterns = None if download_videos else "videos/"
files = None
token_kwargs = {} if token is None else {"token": token}
if self.episodes is not None:
# Reader is guaranteed to exist here (created in __init__ before _download)
files = self.reader.get_episodes_file_paths()
@@ -639,6 +658,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
cache_dir=HF_LEROBOT_HUB_CACHE,
allow_patterns=files,
ignore_patterns=ignore_patterns,
**token_kwargs,
)
)
else:
@@ -650,6 +670,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
local_dir=self._requested_root,
allow_patterns=files,
ignore_patterns=ignore_patterns,
**token_kwargs,
)
self.meta.root = self._requested_root
@@ -789,6 +810,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
image_writer_threads: int = 0,
streaming_encoding: bool = False,
encoder_queue_maxsize: int = 30,
*,
token: str | bool | None = None,
) -> "LeRobotDataset":
"""Resume recording on an existing dataset.
@@ -822,6 +845,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
streaming_encoding: If ``True``, encode video in real-time during
capture.
encoder_queue_maxsize: Max buffered frames per camera for streaming.
token: Authentication token used if metadata must be downloaded
from the Hub. The token is not retained on the dataset instance.
Returns:
A :class:`LeRobotDataset` in write mode, ready to append episodes.
@@ -850,7 +875,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
# Load metadata (revision-safe when root is not provided)
obj.meta = LeRobotDatasetMetadata(
obj.repo_id, obj._requested_root, obj.revision, force_cache_sync=force_cache_sync
obj.repo_id,
obj._requested_root,
obj.revision,
force_cache_sync=force_cache_sync,
token=token,
)
obj._encoder_threads = encoder_threads
+3
View File
@@ -48,6 +48,8 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
tolerances_s: dict | None = None,
download_videos: bool = True,
video_backend: str | None = None,
*,
token: str | bool | None = None,
):
super().__init__()
self.repo_ids = repo_ids
@@ -65,6 +67,7 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
tolerance_s=self.tolerances_s[repo_id],
download_videos=download_videos,
video_backend=video_backend,
token=token,
)
for repo_id in repo_ids
]
+14 -1
View File
@@ -256,6 +256,8 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
shuffle: bool = True,
return_uint8: bool = False,
depth_output_unit: str = DEFAULT_DEPTH_UNIT,
*,
token: str | bool | None = None,
):
"""Initialize a StreamingLeRobotDataset.
@@ -278,6 +280,11 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
shuffle (bool, optional): Whether to shuffle the dataset across exhaustions. Defaults to True.
depth_output_unit (str, optional): Physical unit depth maps are dequantized to ("m" or "mm").
Defaults to "mm".
token: Authentication token used while streaming this dataset from
the Hub. Pass a string token, ``True`` to require the locally
stored token, ``False`` to disable authentication, or ``None``
to use the Hugging Face Hub default. The token is not retained
on the dataset instance after initialization.
"""
super().__init__()
self.repo_id = repo_id
@@ -306,7 +313,11 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
# Load metadata
self.meta = LeRobotDatasetMetadata(
self.repo_id, self._requested_root, self.revision, force_cache_sync=force_cache_sync
self.repo_id,
self._requested_root,
self.revision,
force_cache_sync=force_cache_sync,
token=token,
)
self.root = self.meta.root
self.revision = self.meta.revision
@@ -334,12 +345,14 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
self.delta_timestamps = delta_timestamps
self.delta_indices = get_delta_indices(self.delta_timestamps, self.fps)
token_kwargs = {} if token is None or self.streaming_from_local else {"token": token}
self.hf_dataset: datasets.IterableDataset = load_dataset(
self.repo_id if not self.streaming_from_local else str(self.root),
split="train",
streaming=self.streaming,
data_files="data/*/*.parquet",
revision=self.revision,
**token_kwargs,
)
self.num_shards = min(self.hf_dataset.num_shards, max_num_shards)
+13 -4
View File
@@ -325,16 +325,19 @@ def check_version_compatibility(
logging.warning(FUTURE_MESSAGE.format(repo_id=repo_id, version=v_check))
def get_repo_versions(repo_id: str) -> list[packaging.version.Version]:
def get_repo_versions(repo_id: str, *, token: str | bool | None = None) -> list[packaging.version.Version]:
"""Return available valid versions (branches and tags) on a given Hub repo.
Args:
repo_id (str): The repository ID on the Hugging Face Hub.
token: Authentication token used for Hub requests. Pass a string token,
``True`` to require the locally stored token, ``False`` to disable
authentication, or ``None`` to use the Hugging Face Hub default.
Returns:
list[packaging.version.Version]: A list of valid versions found.
"""
api = HfApi()
api = HfApi() if token is None else HfApi(token=token)
repo_refs = api.list_repo_refs(repo_id, repo_type="dataset")
repo_refs = [b.name for b in repo_refs.branches + repo_refs.tags]
repo_versions = []
@@ -345,7 +348,12 @@ def get_repo_versions(repo_id: str) -> list[packaging.version.Version]:
return repo_versions
def get_safe_version(repo_id: str, version: str | packaging.version.Version) -> str:
def get_safe_version(
repo_id: str,
version: str | packaging.version.Version,
*,
token: str | bool | None = None,
) -> str:
"""Return the specified version if available on repo, or the latest compatible one.
If the exact version is not found, it looks for the latest version with the
@@ -354,6 +362,7 @@ def get_safe_version(repo_id: str, version: str | packaging.version.Version) ->
Args:
repo_id (str): The repository ID on the Hugging Face Hub.
version (str | packaging.version.Version): The target version.
token: Authentication token forwarded to the Hub version lookup.
Returns:
str: The safe version string (e.g., "v1.2.3") to use as a revision.
@@ -366,7 +375,7 @@ def get_safe_version(repo_id: str, version: str | packaging.version.Version) ->
target_version = (
packaging.version.parse(version) if not isinstance(version, packaging.version.Version) else version
)
hub_versions = get_repo_versions(repo_id)
hub_versions = get_repo_versions(repo_id) if token is None else get_repo_versions(repo_id, token=token)
if not hub_versions:
raise RevisionNotFoundError(
+14 -2
View File
@@ -322,7 +322,7 @@ class HILSerlRobotEnvConfig(EnvConfig):
class LiberoEnv(EnvConfig):
task: str = "libero_10" # can also choose libero_spatial, libero_object, etc.
task_ids: list[int] | None = None
fps: int = 30
fps: int = 20 # Must match robosuite's default control_freq (20 Hz)
episode_length: int | None = None
obs_type: str = "pixels_agent_pos"
render_mode: str = "rgb_array"
@@ -354,6 +354,9 @@ class LiberoEnv(EnvConfig):
control_mode: str = "relative" # or "absolute"
def __post_init__(self):
if self.fps <= 0:
raise ValueError(f"fps must be positive, got {self.fps}")
if self.obs_type == "pixels":
self.features[LIBERO_KEY_PIXELS_AGENTVIEW] = PolicyFeature(
type=FeatureType.VISUAL, shape=(self.observation_height, self.observation_width, 3)
@@ -412,6 +415,7 @@ class LiberoEnv(EnvConfig):
"render_mode": self.render_mode,
"observation_height": self.observation_height,
"observation_width": self.observation_width,
"control_freq": self.fps,
}
if self.task_ids is not None:
kwargs["task_ids"] = self.task_ids
@@ -556,7 +560,13 @@ class RoboCasaEnv(EnvConfig):
kwargs["split"] = self.split
return kwargs
def create_envs(self, n_envs: int, use_async_envs: bool = False):
def create_envs(
self,
n_envs: int,
use_async_envs: bool = False,
terminate_on_success: bool = True,
horizon: int | None = None,
):
from .robocasa import create_robocasa_envs
if self.task is None:
@@ -570,6 +580,8 @@ class RoboCasaEnv(EnvConfig):
env_cls=env_cls,
episode_length=self.episode_length,
obj_registries=tuple(self.obj_registries),
terminate_on_success=terminate_on_success,
horizon=horizon,
)
+5
View File
@@ -125,10 +125,13 @@ class LiberoEnv(gym.Env):
n_envs: int = 1,
camera_name_mapping: dict[str, str] | None = None,
num_steps_wait: int = 10,
control_freq: int = 20,
control_mode: str = "relative",
is_libero_plus: bool = False,
):
super().__init__()
if control_freq <= 0:
raise ValueError(f"control_freq must be positive, got {control_freq}")
self.task_id = task_id
self.is_libero_plus = is_libero_plus
self.obs_type = obs_type
@@ -154,6 +157,7 @@ class LiberoEnv(gym.Env):
}
self.camera_name_mapping = camera_name_mapping
self.num_steps_wait = num_steps_wait
self.control_freq = control_freq
self.episode_index = episode_index
self.episode_length = episode_length
# Load once and keep
@@ -260,6 +264,7 @@ class LiberoEnv(gym.Env):
bddl_file_name=self._task_bddl_file,
camera_heights=self.observation_height,
camera_widths=self.observation_width,
control_freq=self.control_freq,
)
env.reset()
self._env = env
+3
View File
@@ -155,6 +155,7 @@ class MetaworldEnv(gym.Env):
env.model.cam_pos[2] = [0.75, 0.075, 0.7]
env.reset()
env._freeze_rand_vec = False # otherwise no randomization
env.seeded_rand_vec = True # use seeded RNG so reset(seed=X) controls object positions
self._env = env
def render(self) -> np.ndarray:
@@ -220,6 +221,8 @@ class MetaworldEnv(gym.Env):
self._ensure_env()
super().reset(seed=seed)
if seed is not None:
self._env.seed(seed)
raw_obs, info = self._env.reset(seed=seed)
observation = self._format_raw_obs(raw_obs)
+33 -11
View File
@@ -33,8 +33,8 @@ logger = logging.getLogger(__name__)
# Dimensions for the flat action/state vectors used by the LeRobot wrapper.
# These correspond to the PandaOmron robot in RoboCasa365.
OBS_STATE_DIM = 16 # base_pos(3) + base_quat(4) + ee_pos_rel(3) + ee_quat_rel(4) + gripper_qpos(2)
ACTION_DIM = 12 # base_motion(4) + control_mode(1) + ee_pos(3) + ee_rot(3) + gripper(1)
OBS_STATE_DIM = 16 # ee_pos_rel(3) + ee_quat_rel(4) + base_pos(3) + base_quat(4) + gripper_qpos(2)
ACTION_DIM = 12 # ee_pos(3) + ee_rot(3) + gripper(1) + base_motion(4) + control_mode(1)
ACTION_LOW = -1.0
ACTION_HIGH = 1.0
@@ -101,14 +101,15 @@ def _resolve_tasks(task: str) -> tuple[list[str], str | None]:
def convert_action(flat_action: np.ndarray) -> dict[str, Any]:
"""Split a flat (12,) action vector into a RoboCasa action dict.
Layout: base_motion(4) + control_mode(1) + ee_pos(3) + ee_rot(3) + gripper(1)
Layout (openpi / robocasa.utils.env_utils.convert_action order):
ee_pos(3) + ee_rot(3) + gripper(1) + base_motion(4) + control_mode(1)
"""
return {
"action.base_motion": flat_action[0:4],
"action.control_mode": flat_action[4:5],
"action.end_effector_position": flat_action[5:8],
"action.end_effector_rotation": flat_action[8:11],
"action.gripper_close": flat_action[11:12],
"action.end_effector_position": flat_action[0:3],
"action.end_effector_rotation": flat_action[3:6],
"action.gripper_close": flat_action[6:7],
"action.base_motion": flat_action[7:11],
"action.control_mode": flat_action[11:12],
}
@@ -136,9 +137,16 @@ class RoboCasaEnv(gym.Env):
episode_length: int | None = None,
obj_registries: Sequence[str] = DEFAULT_OBJ_REGISTRIES,
episode_index: int = 0,
terminate_on_success: bool = True,
horizon: int | None = None,
):
super().__init__()
self.task = task
# When False, a task-success does NOT end/reset the episode — used by the
# interactive sim so one kitchen persists across sequential prompts.
self.terminate_on_success = terminate_on_success
# Underlying robosuite horizon (steps before truncation). None -> default.
self.horizon = horizon
self.obs_type = obs_type
self.render_mode = render_mode
self.observation_width = observation_width
@@ -210,12 +218,16 @@ class RoboCasaEnv(gym.Env):
# (only None/"all"/"pretrain"/"target" are valid). Always pass a
# valid value so we don't hit that default. Extra kwargs are
# forwarded to the underlying kitchen env via create_env/robosuite.make.
extra_kwargs: dict[str, Any] = {}
if self.horizon is not None:
extra_kwargs["horizon"] = int(self.horizon)
self._env = RoboCasaGymEnv(
env_name=self.task,
camera_widths=self.observation_width,
camera_heights=self.observation_height,
split=self.split if self.split is not None else "all",
obj_registries=self.obj_registries,
**extra_kwargs,
)
ep_meta = self._env.env.get_ep_meta()
@@ -230,12 +242,14 @@ class RoboCasaEnv(gym.Env):
return {"pixels": images}
# `state.*` keys come from PandaOmronKeyConverter inside the wrapper.
# openpi state order: ee first, then base, then gripper (matches the
# openpi robocasa pipeline / examples/robocasa/main.py state layout).
agent_pos = np.concatenate(
[
raw_obs.get("state.base_position", np.zeros(3)),
raw_obs.get("state.base_rotation", np.zeros(4)),
raw_obs.get("state.end_effector_position_relative", np.zeros(3)),
raw_obs.get("state.end_effector_rotation_relative", np.zeros(4)),
raw_obs.get("state.base_position", np.zeros(3)),
raw_obs.get("state.base_rotation", np.zeros(4)),
raw_obs.get("state.gripper_qpos", np.zeros(2)),
],
axis=-1,
@@ -280,7 +294,7 @@ class RoboCasaEnv(gym.Env):
raw_obs, reward, done, truncated, info = self._env.step(action_dict)
is_success = bool(info.get("success", False))
terminated = done or is_success
terminated = done or (is_success and self.terminate_on_success)
info.update({"task": self.task, "done": done, "is_success": is_success})
observation = self._format_raw_obs(raw_obs)
@@ -313,6 +327,8 @@ def _make_env_fns(
split: str | None,
episode_length: int | None,
obj_registries: Sequence[str],
terminate_on_success: bool = True,
horizon: int | None = None,
) -> list[Callable[[], RoboCasaEnv]]:
"""Build n_envs factory callables for a single task.
@@ -335,6 +351,8 @@ def _make_env_fns(
episode_length=episode_length,
obj_registries=obj_registries,
episode_index=episode_index,
terminate_on_success=terminate_on_success,
horizon=horizon,
)
return [partial(_make_env, i) for i in range(n_envs)]
@@ -348,6 +366,8 @@ def create_robocasa_envs(
env_cls: Callable[[Sequence[Callable[[], Any]]], Any] | None = None,
episode_length: int | None = None,
obj_registries: Sequence[str] = DEFAULT_OBJ_REGISTRIES,
terminate_on_success: bool = True,
horizon: int | None = None,
) -> dict[str, dict[int, Any]]:
"""Create vectorized RoboCasa365 environments with a consistent return shape.
@@ -409,6 +429,8 @@ def create_robocasa_envs(
split=split,
episode_length=episode_length,
obj_registries=obj_registries,
terminate_on_success=terminate_on_success,
horizon=horizon,
)
if is_async:
+2 -1
View File
@@ -18,6 +18,7 @@ from lerobot.utils.import_utils import require_package
# guard the optional dependency here so importing this package fails loudly if it's missing.
require_package("datasets", extra="dataset")
from .annotate import submit_annotate_to_hf
from .hf import submit_to_hf
__all__ = ["submit_to_hf"]
__all__ = ["submit_annotate_to_hf", "submit_to_hf"]
+176
View File
@@ -0,0 +1,176 @@
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Run ``lerobot-annotate`` on HF Jobs (HuggingFace GPUs).
Same shape as the training submitter in ``hf.py``, with one difference: the
annotation pipeline serves its own VLM, so the pod starts from the official
``vllm/vllm-openai`` image (which has no lerobot) instead of the prebuilt
``lerobot-gpu`` image, and installs lerobot on top before running.
Because there is no config repo to stage, the pod replays the user's own CLI
flags — everything except the client-only ``--job.*`` and the host-local
``--root``, which is replaced by ``--repo_id`` so the pod pulls the dataset
from the Hub.
"""
from __future__ import annotations
import shlex
import sys
from dataclasses import is_dataclass
from typing import TYPE_CHECKING
from huggingface_hub import HfApi, get_token, run_job
from .dataset import ensure_dataset_available
# Package-internal reuse of the training submitter's job plumbing: following a
# submitted job and forwarding argv are identical for annotation runs.
from .hf import _pod_forwarded_args, follow_job, resolve_job_tags
if TYPE_CHECKING:
from lerobot.annotations.steerable_pipeline.config import AnnotationPipelineConfig
LEROBOT_GIT_URL = "https://github.com/huggingface/lerobot.git"
# Mirrors the pins in pyproject.toml. The vLLM image resolves dependencies on its
# own otherwise, and pulls av 18 / datasets 5 / draccus 0.11 — each of which breaks
# lerobot at import time. `--upgrade-strategy only-if-needed` keeps vLLM's own
# (torch, transformers, ...) pins intact.
_RUNTIME_REQUIREMENTS = (
"'datasets>=4.7.0,<5.0.0' 'pyarrow>=21.0.0,<30.0.0' 'av>=15.0.0,<16.0.0' 'draccus==0.10.0' "
"'pandas>=2.0.0,<3.0.0' jsonlines gymnasium torchcodec mergedeep pyyaml-include toml typing-inspect "
"openai"
)
# Flags the submitter resolves itself instead of forwarding verbatim: `--root`
# names a directory only this machine has, `--repo_id` is re-emitted from the
# config, and the config-file args name local files (rejected up front by
# `submit_annotate_to_hf`). `--job.*` is dropped separately, by prefix; bare
# `--job` is not, hence its entry here — it is the one arg that could smuggle a
# remote `target` onto the pod and have the job recursively submit itself.
_SUBMITTER_OWNED_ARGS = ("--root", "--repo_id", "--config_path", "--job")
def _local_config_file_args(cfg: AnnotationPipelineConfig) -> list[str]:
"""The CLI args that name a config file on the client's disk.
draccus exposes ``--config_path`` for the whole config plus a ``--<field>``
for every nested dataclass (``--vlm``, ``--plan``, ``--job``, ...). The pod has
none of those files, so a remote run has to reject them rather than silently
drop the settings they carry.
"""
return ["--config_path", *(f"--{name}" for name in vars(cfg) if is_dataclass(getattr(cfg, name)))]
def build_pod_setup(lerobot_ref: str) -> str:
"""Shell prelude that turns the vLLM image into a ``lerobot-annotate`` runtime."""
spec = f"lerobot @ git+{LEROBOT_GIT_URL}@{lerobot_ref}"
return (
# git to install from the repo, ffmpeg to decode the dataset's videos.
"apt-get update -qq && apt-get install -y -qq git ffmpeg && "
f"pip install --no-deps {shlex.quote(spec)} && "
f"pip install --upgrade-strategy only-if-needed {_RUNTIME_REQUIREMENTS} && "
# vLLM's cudagraph memory estimate over-reserves and starves the KV cache;
# PyAV is the video backend the server can decode our frames with.
"export VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0 && "
"export VLLM_VIDEO_BACKEND=pyav"
)
def build_pod_command(repo_id: str, lerobot_ref: str, argv: list[str]) -> list[str]:
"""Build the ``bash -c`` command the pod runs: setup prelude, then annotation.
``argv`` is the user's CLI (``sys.argv[1:]``) minus the flags in
``_SUBMITTER_OWNED_ARGS``; ``--repo_id`` is re-added from the config so the pod
always annotates the dataset we just made sure is reachable on the Hub.
``--job.target=local`` stops the pod from re-dispatching to itself.
"""
forwarded = _pod_forwarded_args(argv, drop_names=_SUBMITTER_OWNED_ARGS, drop_prefixes=("--job.",))
annotate = shlex.join(["lerobot-annotate", f"--repo_id={repo_id}", *forwarded, "--job.target=local"])
return ["bash", "-c", f"{build_pod_setup(lerobot_ref)} && {annotate}"]
def submit_annotate_to_hf(cfg: AnnotationPipelineConfig) -> None:
"""Submit an annotation run to HF Jobs infrastructure.
Resolves credentials, makes sure the source dataset is reachable from the pod,
submits the job, then tails its logs until the job reaches a terminal stage —
or returns immediately with ``--job.detach``. Ctrl-C detaches without
cancelling the remote job.
"""
token = get_token()
if not token:
raise RuntimeError("Not logged in to Hugging Face. Run `hf auth login` first.")
if cfg.repo_id is None:
raise ValueError(
"Remote annotation requires --repo_id: the pod downloads the dataset from the Hub, "
"and --root only names a directory on this machine."
)
argv = sys.argv[1:]
passed = {tok.split("=", 1)[0] for tok in argv}
used_config_files = sorted(passed.intersection(_local_config_file_args(cfg)))
if used_config_files:
raise ValueError(
f"{', '.join(used_config_files)} cannot be used with a remote --job.target: the pod "
"cannot read config files from this machine. Pass the settings as CLI flags instead."
)
if not cfg.push_to_hub:
# The pod's filesystem is discarded when the job ends, so without a push the
# run produces nothing. Warn rather than fail: a smoke test over
# --only_episodes that only inspects the logs is a legitimate use.
print(
"WARNING: --push_to_hub is off. The annotated dataset lives only on the pod and is "
"discarded when the job ends. Pass --push_to_hub=true to keep the result."
)
api = HfApi(token=token)
tags = resolve_job_tags(cfg.job.tags)
ensure_dataset_available(cfg.repo_id, api=api, tags=tags)
command = build_pod_command(cfg.repo_id, cfg.job.lerobot_ref, argv)
print(f"Submitting job to HF Jobs (flavor={cfg.job.target}, image={cfg.job.image}) ...")
job_info = run_job(
image=cfg.job.image,
command=command,
flavor=cfg.job.target,
secrets={"HF_TOKEN": token},
timeout=cfg.job.timeout,
# HF Jobs labels are key/value; expose each tag as a queryable label.
labels=dict.fromkeys(tags, "true"),
)
job_id = job_info.id
job_url = getattr(job_info, "url", None)
print(f"Job submitted: {job_id}")
if job_url:
print(f" Job page: {job_url}")
target_repo_id = cfg.new_repo_id or cfg.repo_id
if cfg.push_to_hub:
print(f" Dataset repo: https://huggingface.co/datasets/{target_repo_id}")
print(f" Monitor: hf jobs logs {job_id}")
print(f" Cancel: hf jobs cancel {job_id}")
# No success marker: `lerobot-annotate` keeps working after the upload log line
# (dataset card, version tag), so completion has to be stage-based.
if not follow_job(job_id, detach=cfg.job.detach):
return
if cfg.push_to_hub:
print(f"\nAnnotation complete — dataset pushed to https://huggingface.co/datasets/{target_repo_id}")
else:
print("\nAnnotation complete. Note: --push_to_hub was off, so the result stayed on the pod.")
+69 -54
View File
@@ -223,6 +223,74 @@ def _poll_until_done(
return None
def follow_job(job_id: str, *, detach: bool = False, success_marker: str | None = None) -> bool:
"""Watch a submitted job to the end, streaming its logs to stdout.
Returns True when the job finished successfully and False when we stopped watching
without a verdict — `detach`, or the user pressing Ctrl-C, which detaches rather than
cancelling the remote job. Raises RuntimeError when the job reaches a terminal stage
other than COMPLETED.
`success_marker` finishes as soon as that string appears in the logs instead of waiting
out the platform's post-run finalization (~30s). Callers that have a log line meaning
"the artifact is on the Hub" should pass it; without one, completion is stage-based.
"""
if detach:
return False
done = threading.Event()
detached = threading.Event()
marker_seen = threading.Event()
stage_holder: dict[str, str | None] = {}
def _poll() -> None:
stage_holder["stage"] = _poll_until_done(job_id, done, status_holder=stage_holder)
poll_thread = threading.Thread(target=_poll, daemon=True)
poll_thread.start()
log_thread = threading.Thread(
target=_tail_logs, args=(job_id, done, success_marker, marker_seen), daemon=True
)
log_thread.start()
def _detach(sig, frame):
detached.set()
done.set()
print("\nDetached. Job is still running.")
print(f" Monitor: hf jobs logs {job_id}")
print(f" Cancel: hf jobs cancel {job_id}")
# signal.signal only works on the main thread; when called from a worker thread
# (e.g. an orchestration framework) skip the Ctrl-C-detaches-instead-of-cancels
# handler rather than crashing with ValueError.
install_sigint = threading.current_thread() is threading.main_thread()
original_sigint = signal.getsignal(signal.SIGINT) if install_sigint else None
if install_sigint:
signal.signal(signal.SIGINT, _detach)
try:
# Timeout-based join so SIGINT is delivered to the main thread promptly.
while poll_thread.is_alive():
poll_thread.join(timeout=0.5)
log_thread.join(timeout=5)
finally:
if install_sigint:
signal.signal(signal.SIGINT, original_sigint)
if detached.is_set():
return False
if marker_seen.is_set():
return True
stage = stage_holder.get("stage")
if stage != "COMPLETED":
message = stage_holder.get("message")
detail = f" ({message})" if message else ""
raise RuntimeError(
f"Job {job_id} ended with stage={stage}{detail}. Check logs: hf jobs logs {job_id}"
)
return True
def _pod_forwarded_args(
argv: list[str], drop_names: tuple[str, ...] = (), drop_prefixes: tuple[str, ...] = ()
) -> list[str]:
@@ -362,64 +430,11 @@ def submit_to_hf(cfg: TrainPipelineConfig) -> None:
print(f" Monitor: hf jobs logs {job_id}")
print(f" Cancel: hf jobs cancel {job_id}")
if cfg.job.detach:
return
done = threading.Event()
detached = threading.Event()
pushed_ok = threading.Event()
stage_holder: dict[str, str | None] = {}
def _poll() -> None:
stage_holder["stage"] = _poll_until_done(job_id, done, status_holder=stage_holder)
poll_thread = threading.Thread(target=_poll, daemon=True)
poll_thread.start()
# Finish as soon as the model is pushed, rather than waiting out the platform's
# post-run finalization before the job stage flips to COMPLETED. This matches the
# exact log line emitted by PreTrainedPolicy.push_model_to_hub — the two must stay
# in sync. If it ever stops matching we just fall back to stage-based completion
# (~30s slower), so the contract is an optimization, not a correctness requirement.
success_marker = f"Model pushed to https://huggingface.co/{repo_id}"
log_thread = threading.Thread(
target=_tail_logs, args=(job_id, done, success_marker, pushed_ok), daemon=True
)
log_thread.start()
def _detach(sig, frame):
detached.set()
done.set()
print("\nDetached. Job is still running.")
print(f" Monitor: hf jobs logs {job_id}")
print(f" Cancel: hf jobs cancel {job_id}")
# signal.signal only works on the main thread; when called from a worker thread
# (e.g. an orchestration framework) skip the Ctrl-C-detaches-instead-of-cancels
# handler rather than crashing with ValueError.
install_sigint = threading.current_thread() is threading.main_thread()
original_sigint = signal.getsignal(signal.SIGINT) if install_sigint else None
if install_sigint:
signal.signal(signal.SIGINT, _detach)
try:
# Timeout-based join so SIGINT is delivered to the main thread promptly.
while poll_thread.is_alive():
poll_thread.join(timeout=0.5)
log_thread.join(timeout=5)
finally:
if install_sigint:
signal.signal(signal.SIGINT, original_sigint)
if detached.is_set():
return
if pushed_ok.is_set():
if follow_job(job_id, detach=cfg.job.detach, success_marker=success_marker):
print(f"\nTraining complete — model pushed to https://huggingface.co/{repo_id}")
return
stage = stage_holder.get("stage")
if stage != "COMPLETED":
message = stage_holder.get("message")
detail = f" ({message})" if message else ""
raise RuntimeError(
f"Job {job_id} ended with stage={stage}{detail}. Check logs: hf jobs logs {job_id}"
)
+1 -2
View File
@@ -20,7 +20,6 @@ import logging
import time
from contextlib import contextmanager
from copy import deepcopy
from functools import cached_property
from typing import TYPE_CHECKING, Any, TypedDict
from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
@@ -854,7 +853,7 @@ class DamiaoMotorsBus(MotorsBusBase):
else:
raise ValueError(f"Motor {motor_obj} doesn't have a valid recv_id (None).")
@cached_property
@property
def is_calibrated(self) -> bool:
"""Check if motors are calibrated."""
return bool(self.calibration)
+9 -5
View File
@@ -23,6 +23,7 @@ from __future__ import annotations
import abc
import logging
import time
from collections.abc import Sequence
from contextlib import contextmanager
from dataclasses import dataclass
@@ -818,13 +819,13 @@ class SerialMotorsBus(MotorsBusBase):
"""
motor_names = self._get_motors_list(motors)
start_positions = self.sync_read("Present_Position", motor_names, normalize=False)
start_positions = self.sync_read("Present_Position", motor_names, normalize=False, num_retry=5)
mins = start_positions.copy()
maxes = start_positions.copy()
user_pressed_enter = False
while not user_pressed_enter:
positions = self.sync_read("Present_Position", motor_names, normalize=False)
positions = self.sync_read("Present_Position", motor_names, normalize=False, num_retry=5)
mins = {motor: min(positions[motor], min_) for motor, min_ in mins.items()}
maxes = {motor: max(positions[motor], max_) for motor, max_ in maxes.items()}
@@ -837,9 +838,12 @@ class SerialMotorsBus(MotorsBusBase):
if enter_pressed():
user_pressed_enter = True
if display_values and not user_pressed_enter:
# Move cursor up to overwrite the previous output
move_cursor_up(len(motor_names) + 3)
if not user_pressed_enter:
if display_values:
# Move cursor up to overwrite the previous output
move_cursor_up(len(motor_names) + 3)
# Throttle reads even when the live table is disabled.
time.sleep(0.02)
same_min_max = [motor for motor in motor_names if mins[motor] == maxes[motor]]
if same_min_max:
@@ -79,6 +79,8 @@ class DiffusionConfig(PreTrainedConfig):
use_film_scale_modulation: FiLM (https://huggingface.co/papers/1709.07871) is used for the Unet conditioning.
Bias modulation is used be default, while this parameter indicates whether to also use scale
modulation.
gradient_checkpointing: Whether to checkpoint the Unet residual blocks during training. This reduces
activation memory at the cost of recomputing those blocks during the backward pass.
noise_scheduler_type: Name of the noise scheduler to use. Supported options: ["DDPM", "DDIM"].
num_train_timesteps: Number of diffusion steps for the forward diffusion schedule.
beta_schedule: Name of the diffusion beta schedule as per DDPMScheduler from Hugging Face diffusers.
@@ -132,6 +134,7 @@ class DiffusionConfig(PreTrainedConfig):
n_groups: int = 8
diffusion_step_embed_dim: int = 128
use_film_scale_modulation: bool = True
gradient_checkpointing: bool = False
# Noise scheduler.
noise_scheduler_type: str = "DDPM"
num_train_timesteps: int = 100
@@ -31,6 +31,7 @@ import torch
import torch.nn.functional as F # noqa: N812
import torchvision
from torch import Tensor, nn
from torch.utils.checkpoint import checkpoint
from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGES, OBS_STATE
from lerobot.utils.import_utils import _diffusers_available, require_package
@@ -727,22 +728,35 @@ class DiffusionConditionalUnet1d(nn.Module):
else:
global_feature = timesteps_embed
use_gc = self.config.gradient_checkpointing and self.training
# Run encoder, keeping track of skip features to pass to the decoder.
encoder_skip_features: list[Tensor] = []
for resnet, resnet2, downsample in self.down_modules:
x = resnet(x, global_feature)
x = resnet2(x, global_feature)
if use_gc:
x = checkpoint(resnet, x, global_feature, use_reentrant=False)
x = checkpoint(resnet2, x, global_feature, use_reentrant=False)
else:
x = resnet(x, global_feature)
x = resnet2(x, global_feature)
encoder_skip_features.append(x)
x = downsample(x)
for mid_module in self.mid_modules:
x = mid_module(x, global_feature)
if use_gc:
x = checkpoint(mid_module, x, global_feature, use_reentrant=False)
else:
x = mid_module(x, global_feature)
# Run decoder, using the skip features from the encoder.
for resnet, resnet2, upsample in self.up_modules:
x = torch.cat((x, encoder_skip_features.pop()), dim=1)
x = resnet(x, global_feature)
x = resnet2(x, global_feature)
if use_gc:
x = checkpoint(resnet, x, global_feature, use_reentrant=False)
x = checkpoint(resnet2, x, global_feature, use_reentrant=False)
else:
x = resnet(x, global_feature)
x = resnet2(x, global_feature)
x = upsample(x)
x = self.final_conv(x)
+1
View File
@@ -177,6 +177,7 @@ def make_pre_post_processors(
return make_groot_pre_post_processors_from_pretrained(
config=policy_cfg,
pretrained_path=pretrained_path,
revision=pretrained_revision,
dataset_stats=kwargs.get("dataset_stats"),
dataset_meta=kwargs.get("dataset_meta"),
preprocessor_overrides=kwargs.get("preprocessor_overrides"),
@@ -475,6 +475,7 @@ def make_groot_pre_post_processors_from_pretrained(
config: GrootConfig,
pretrained_path: str,
*,
revision: str | None = None,
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
dataset_meta: Any | None = None,
preprocessor_overrides: dict[str, Any] | None = None,
@@ -511,6 +512,7 @@ def make_groot_pre_post_processors_from_pretrained(
preprocessor, postprocessor = _load_groot_processor_pipelines(
pretrained_path,
revision=revision,
preprocessor_overrides=preprocessor_overrides,
postprocessor_overrides=postprocessor_overrides,
preprocessor_config_filename=preprocessor_config_filename,
@@ -526,6 +528,7 @@ def make_groot_pre_post_processors_from_pretrained(
def _load_groot_processor_pipelines(
pretrained_path: str,
*,
revision: str | None,
preprocessor_overrides: dict[str, Any],
postprocessor_overrides: dict[str, Any],
preprocessor_config_filename: str,
@@ -540,6 +543,7 @@ def _load_groot_processor_pipelines(
preprocessor = PolicyProcessorPipeline.from_pretrained(
pretrained_model_name_or_path=pretrained_path,
config_filename=preprocessor_config_filename,
revision=revision,
overrides=preprocessor_overrides,
to_transition=batch_to_transition,
to_output=transition_to_batch,
@@ -547,6 +551,7 @@ def _load_groot_processor_pipelines(
postprocessor = PolicyProcessorPipeline.from_pretrained(
pretrained_model_name_or_path=pretrained_path,
config_filename=postprocessor_config_filename,
revision=revision,
overrides=postprocessor_overrides,
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
+4 -12
View File
@@ -524,8 +524,6 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
def embed_suffix(self, noisy_actions, timestep):
"""Embed noisy_actions, timestep to prepare for Expert Gemma processing."""
embs = []
pad_masks = []
att_masks = []
# Embed timestep using sine-cosine positional encoding
@@ -551,23 +549,17 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
return F.silu(x)
time_emb = self._apply_checkpoint(time_mlp_func, time_emb)
action_time_emb = action_emb
adarms_cond = time_emb
embs.append(action_time_emb)
bsize, action_time_dim = action_time_emb.shape[:2]
action_time_mask = torch.ones(bsize, action_time_dim, dtype=torch.bool, device=timestep.device)
pad_masks.append(action_time_mask)
bsize, action_time_dim = action_emb.shape[:2]
pad_masks = torch.ones(bsize, action_time_dim, dtype=torch.bool, device=timestep.device)
# Set attention masks so that image, language and state inputs do not attend to action tokens
att_masks += [1] + ([0] * (self.config.chunk_size - 1))
embs = torch.cat(embs, dim=1)
pad_masks = torch.cat(pad_masks, dim=1)
att_masks = torch.tensor(att_masks, dtype=embs.dtype, device=embs.device)
att_masks = torch.tensor(att_masks, dtype=action_emb.dtype, device=action_emb.device)
att_masks = att_masks[None, :].expand(bsize, len(att_masks))
return embs, pad_masks, att_masks, adarms_cond
return action_emb, pad_masks, att_masks, adarms_cond
def forward(self, images, img_masks, tokens, masks, actions, noise, time) -> Tensor:
"""Do a full training forward pass and compute the loss."""
-3
View File
@@ -175,9 +175,6 @@ class AddBatchDimensionComplementaryDataStep(ComplementaryDataProcessorStep):
if isinstance(task_index_value, Tensor) and task_index_value.dim() == 0:
complementary_data["task_index"] = task_index_value.unsqueeze(0)
complementary_data.pop("language_persistent", None)
complementary_data.pop("language_events", None)
if "messages" in complementary_data:
messages = complementary_data["messages"]
if isinstance(messages, list) and (not messages or isinstance(messages[0], dict)):
+101 -6
View File
@@ -41,7 +41,7 @@ from pathlib import Path
from typing import Any, TypedDict, TypeVar, cast
import torch
from huggingface_hub import hf_hub_download
from huggingface_hub import hf_hub_download, snapshot_download
from safetensors.torch import load_file, save_file
from lerobot.configs import PipelineFeatureType, PolicyFeature
@@ -205,6 +205,10 @@ class ProcessorStep(ABC):
"""
return None
def save_artifacts(self, save_directory: Path) -> dict[str, str]:
"""Save non-tensor assets and map constructor arguments to relative paths."""
return {}
def reset(self) -> None:
"""Resets the internal state of the processor step, if any."""
return None
@@ -549,6 +553,22 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
pipeline_config = self.get_config()
pipeline_state_dict = self.state_dict()
for processor_step, step_entry in zip(self.steps, pipeline_config["steps"], strict=True):
artifacts = processor_step.save_artifacts(save_directory)
if artifacts:
for config_key, relative_path in artifacts.items():
artifact_path = Path(relative_path)
if artifact_path.is_absolute() or ".." in artifact_path.parts:
raise ValueError(
f"Processor artifact path must be relative to the checkpoint: {relative_path!r}"
)
if not (save_directory / artifact_path).exists():
raise FileNotFoundError(
f"Processor step did not save declared artifact '{relative_path}'"
)
step_entry["config"][config_key] = artifact_path.as_posix()
step_entry["artifacts"] = artifacts
for state_key, step_state_dict in pipeline_state_dict.items():
state_filename = f"{state_key}.safetensors"
save_file(step_state_dict, save_directory / state_filename)
@@ -713,6 +733,8 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
ProcessorMigrationError: If the model requires migration to processor format.
"""
model_id = str(pretrained_model_name_or_path)
model_path = Path(model_id)
is_local_source = model_path.is_dir() or model_path.is_file()
hub_download_kwargs = {
"force_download": force_download,
"resume_download": resume_download,
@@ -731,7 +753,13 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
# 3. Build steps with overrides
steps, validated_overrides = cls._build_steps_with_overrides(
loaded_config, overrides or {}, model_id, base_path, hub_download_kwargs
loaded_config,
overrides or {},
model_id,
base_path,
config_filename,
hub_download_kwargs,
is_local_source,
)
# 4. Validate that all overrides were used
@@ -920,7 +948,9 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
overrides: dict[str, Any],
model_id: str,
base_path: Path | None,
config_filename: str,
hub_download_kwargs: dict[str, Any],
is_local_source: bool = False,
) -> tuple[list[ProcessorStep], set[str]]:
"""Build all processor steps with overrides and state loading.
@@ -944,7 +974,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
3. **State Loading** (via _load_step_state):
- **If step has "state_file"**: Load tensor state from .safetensors
- **Local first**: Check base_path/state_file.safetensors
- **Hub fallback**: Download state file if not found locally
- **Hub fallback**: Download state file if the pipeline was loaded from the Hub
- **Optional**: Only load if step has load_state_dict method
4. **Override Tracking**:
@@ -962,6 +992,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
model_id: The model identifier (needed for Hub state file downloads)
base_path: Local directory path for finding state files
hub_download_kwargs: Parameters for hf_hub_download (tokens, cache, etc.)
is_local_source: Whether model_id resolved to a local directory or config file.
Returns:
Tuple of (instantiated_steps_list, unused_override_keys)
@@ -972,13 +1003,68 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
ImportError: If a step class cannot be imported or found in registry
ValueError: If a step cannot be instantiated with its configuration
"""
loaded_config = deepcopy(loaded_config)
cls._resolve_artifact_paths(
loaded_config,
model_id,
base_path,
config_filename,
hub_download_kwargs,
)
steps, remaining_override_keys = cls._build_steps_from_config(loaded_config, overrides)
for step_instance, step_entry in zip(steps, loaded_config["steps"], strict=True):
cls._load_step_state(step_instance, step_entry, model_id, base_path, hub_download_kwargs)
cls._load_step_state(
step_instance,
step_entry,
model_id,
base_path,
config_filename,
hub_download_kwargs,
is_local_source,
)
return steps, remaining_override_keys
@classmethod
def _resolve_artifact_paths(
cls,
loaded_config: dict[str, Any],
model_id: str,
base_path: Path | None,
config_filename: str,
hub_download_kwargs: dict[str, Any],
) -> None:
"""Resolve declared relative processor artifacts before step construction."""
is_local = Path(model_id).is_dir() or Path(model_id).is_file()
for step_entry in loaded_config["steps"]:
artifacts = step_entry.get("artifacts", {})
for config_key, relative_path in artifacts.items():
artifact_path = Path(relative_path)
if artifact_path.is_absolute() or ".." in artifact_path.parts:
raise ValueError(
f"Processor artifact path must be relative to the checkpoint: {relative_path!r}"
)
resolved_path = base_path / artifact_path if base_path is not None else artifact_path
if not resolved_path.exists() and not is_local:
repository_path = Path(config_filename).parent / artifact_path
snapshot_download(
repo_id=model_id,
repo_type="model",
allow_patterns=f"{repository_path.as_posix()}/**",
**hub_download_kwargs,
)
if not resolved_path.exists():
step_name = step_entry.get("registry_name", step_entry.get("class", "unknown"))
raise FileNotFoundError(
f"Missing processor artifact '{relative_path}' for step '{step_name}' "
f"next to '{config_filename}'. Checkpoint artifacts are incomplete."
)
step_entry["config"][config_key] = str(resolved_path)
@classmethod
def _build_steps_from_config(
cls,
@@ -1138,7 +1224,9 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
step_entry: dict[str, Any],
model_id: str,
base_path: Path | None,
config_filename: str,
hub_download_kwargs: dict[str, Any],
is_local_source: bool = False,
) -> None:
"""Load state dictionary for a processor step if available.
@@ -1157,7 +1245,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
- **Use case**: Loading from local saved model directory
2. **Hub download fallback**: Download state file from repository
- **When triggered**: Local file not found or base_path is None
- **When triggered**: Local file not found and the pipeline source is a Hub repo
- **Process**: Use hf_hub_download with same parameters as config
- **Example**: Download "normalize_step_0.safetensors" from "user/repo"
- **Result**: Downloaded to local cache, path returned
@@ -1178,6 +1266,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
model_id: The model identifier (used for Hub downloads if needed)
base_path: Local directory path for finding state files (None for Hub-only)
hub_download_kwargs: Parameters for hf_hub_download (tokens, cache, etc.)
is_local_source: Whether model_id resolved to a local directory or config file.
Note:
This method modifies step_instance in-place and returns None.
@@ -1191,11 +1280,17 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
# Try local file first
if base_path and (base_path / state_filename).exists():
state_path = str(base_path / state_filename)
elif is_local_source:
state_path = base_path / state_filename if base_path else Path(state_filename)
raise FileNotFoundError(
f"State file '{state_filename}' was not found for local processor pipeline "
f"'{model_id}' at '{state_path}'."
)
else:
# Download from Hub
state_path = hf_hub_download(
repo_id=model_id,
filename=state_filename,
filename=(Path(config_filename).parent / state_filename).as_posix(),
repo_type="model",
**hub_download_kwargs,
)
@@ -16,7 +16,7 @@
from __future__ import annotations
from dataclasses import dataclass
from dataclasses import asdict, dataclass
from typing import Any
from lerobot.configs import PipelineFeatureType, PolicyFeature
@@ -32,17 +32,18 @@ from .pipeline import ProcessorStep, ProcessorStepRegistry
@dataclass
@ProcessorStepRegistry.register(name="render_messages_processor")
class RenderMessagesStep(ProcessorStep):
"""Processor step that turns raw language columns into rendered chat messages.
Reads ``language_persistent`` and ``language_events`` from the transition's
complementary data, renders them through ``recipe`` at the sample timestamp,
and replaces the raw columns with the resulting ``messages`` /
``message_streams`` / ``target_message_indices`` keys.
"""
"""Render language columns into recipe-defined messages and supervision metadata."""
recipe: TrainingRecipe
dataset_ctx: Any | None = None
def __post_init__(self) -> None:
if isinstance(self.recipe, dict):
self.recipe = TrainingRecipe.from_dict(self.recipe)
def get_config(self) -> dict[str, Any]:
return {"recipe": asdict(self.recipe)}
def __call__(self, transition: EnvTransition) -> EnvTransition | None:
"""Render messages for a single transition; return ``None`` to drop it."""
complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA) or {}
@@ -50,7 +51,17 @@ class RenderMessagesStep(ProcessorStep):
events = complementary_data.get(LANGUAGE_EVENTS) or []
if not persistent and not events:
return transition
rendered = _fallback_low_level_render(complementary_data.get("task"))
if rendered is None:
return transition
new_transition = transition.copy()
new_complementary_data = dict(new_transition.get(TransitionKey.COMPLEMENTARY_DATA) or {})
new_complementary_data.update(rendered)
new_transition[TransitionKey.COMPLEMENTARY_DATA] = new_complementary_data
return new_transition
if _is_batched_language(persistent) or _is_batched_language(events):
return self._call_batch(transition, complementary_data, persistent, events)
timestamp = complementary_data.get("timestamp")
if timestamp is None:
@@ -67,18 +78,147 @@ class RenderMessagesStep(ProcessorStep):
dataset_ctx=self.dataset_ctx,
)
if rendered is None:
return None
rendered = _fallback_low_level_render(complementary_data.get("task"))
if rendered is None:
return None
new_transition = transition.copy()
new_complementary_data = dict(complementary_data)
new_complementary_data = dict(new_transition.get(TransitionKey.COMPLEMENTARY_DATA) or {})
new_complementary_data.pop(LANGUAGE_PERSISTENT, None)
new_complementary_data.pop(LANGUAGE_EVENTS, None)
new_complementary_data.update(rendered)
new_transition[TransitionKey.COMPLEMENTARY_DATA] = new_complementary_data
return new_transition
def _call_batch(
self,
transition: EnvTransition,
complementary_data: dict[str, Any],
persistent_batch: list,
events_batch: list,
) -> EnvTransition | None:
timestamp = complementary_data.get("timestamp")
if timestamp is None:
raise KeyError("RenderMessagesStep requires sample timestamp in complementary data.")
batch_size = max(len(persistent_batch), len(events_batch))
messages: list[list[dict[str, Any]]] = []
message_streams: list[list[str | None]] = []
target_message_indices: list[list[int]] = []
keep_indices: list[int] = []
for i in range(batch_size):
rendered = render_sample(
recipe=self.recipe,
persistent=persistent_batch[i] if i < len(persistent_batch) else [],
events=events_batch[i] if i < len(events_batch) else [],
t=_batch_value(timestamp, i),
sample_idx=int(_batch_value(complementary_data.get("index", 0), i)),
task=_batch_value(complementary_data.get("task"), i),
dataset_ctx=self.dataset_ctx,
)
if rendered is None:
rendered = _fallback_low_level_render(_batch_value(complementary_data.get("task"), i))
if rendered is None:
continue
keep_indices.append(i)
messages.append(rendered["messages"])
message_streams.append(rendered["message_streams"])
target_message_indices.append(rendered["target_message_indices"])
if not messages:
return None
new_transition = (
_select_batch_indices(transition, keep_indices)
if len(keep_indices) != batch_size
else transition.copy()
)
new_complementary_data = dict(new_transition.get(TransitionKey.COMPLEMENTARY_DATA) or {})
new_complementary_data.pop(LANGUAGE_PERSISTENT, None)
new_complementary_data.pop(LANGUAGE_EVENTS, None)
new_complementary_data["messages"] = messages
new_complementary_data["message_streams"] = message_streams
new_complementary_data["target_message_indices"] = target_message_indices
new_transition[TransitionKey.COMPLEMENTARY_DATA] = new_complementary_data
return new_transition
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
"""Pass features through unchanged; rendering only touches complementary data."""
return features
def _scalar(value: Any) -> float | int:
"""Unwrap a tensor/array/single-element list into a Python scalar."""
if hasattr(value, "item"):
return value.item()
if isinstance(value, list):
if len(value) != 1:
raise ValueError(f"Expected a scalar, got list of length {len(value)}: {value!r}")
return _scalar(value[0])
return value
def _is_batched_language(value: Any) -> bool:
return isinstance(value, list) and bool(value) and isinstance(value[0], list)
def _batch_value(value: Any, index: int) -> Any:
if value is None:
return None
if isinstance(value, list):
return value[index]
if hasattr(value, "ndim") and value.ndim > 0:
return _scalar(value[index])
return _scalar(value)
def _select_batch_indices(transition: EnvTransition, indices: list[int]) -> EnvTransition:
selected = transition.copy()
for key in (TransitionKey.OBSERVATION, TransitionKey.COMPLEMENTARY_DATA):
data = selected.get(key)
if isinstance(data, dict):
selected[key] = {k: _select_value(v, indices) for k, v in data.items()}
action = selected.get(TransitionKey.ACTION)
if action is not None:
selected[TransitionKey.ACTION] = _select_value(action, indices)
return selected
def _select_value(value: Any, indices: list[int]) -> Any:
if isinstance(value, list) and len(value) >= len(indices):
return [value[i] for i in indices]
if hasattr(value, "index_select") and hasattr(value, "new_tensor") and getattr(value, "ndim", 0) > 0:
return value.index_select(0, value.new_tensor(indices).long())
return value
def _fallback_low_level_render(task: Any) -> dict[str, Any] | None:
"""Keep action-only samples trainable when no recipe branch matches."""
if hasattr(task, "item"):
task = task.item()
if isinstance(task, list):
messages = []
message_streams = []
target_message_indices = []
for t in task:
rendered = _fallback_low_level_render(t)
if rendered is None:
return None
messages.append(rendered["messages"])
message_streams.append(rendered["message_streams"])
target_message_indices.append(rendered["target_message_indices"])
return {
"messages": messages,
"message_streams": message_streams,
"target_message_indices": target_message_indices,
}
if not isinstance(task, str) or not task:
return None
return {
"messages": [{"role": "user", "content": task}],
"message_streams": ["low_level"],
"target_message_indices": [],
}
+50 -17
View File
@@ -25,6 +25,7 @@ from __future__ import annotations
import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Any
import torch
@@ -32,6 +33,7 @@ import torch
from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
from lerobot.types import EnvTransition, RobotObservation, TransitionKey
from lerobot.utils.constants import (
ACTION_CODE_TOKEN_MASK,
ACTION_TOKEN_MASK,
ACTION_TOKENS,
OBS_LANGUAGE_ATTENTION_MASK,
@@ -136,7 +138,7 @@ class TokenizerProcessorStep(ObservationProcessorStep):
# Standardize to a list of strings for the tokenizer
if isinstance(task, str):
return [task]
elif isinstance(task, (list, tuple)) and all(isinstance(t, str) for t in task):
elif isinstance(task, list | tuple) and all(isinstance(t, str) for t in task):
return list(task)
return None
@@ -349,6 +351,7 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
max_action_tokens: int = 256
fast_skip_tokens: int = 128
paligemma_tokenizer_name: str = "google/paligemma-3b-pt-224"
allow_truncation: bool = True
# Internal tokenizer instance (not part of the config)
action_tokenizer: Any = field(default=None, init=False, repr=False)
_paligemma_tokenizer: Any = field(default=None, init=False, repr=False)
@@ -412,14 +415,15 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
# During inference, no action is available, skip tokenization
return new_transition
# Tokenize and get both tokens and mask
tokens, mask = self._tokenize_action(action)
# Tokenize and get masks for the full formatted sequence and the discrete action codes.
tokens, mask, code_mask = self._tokenize_action(action)
# Store mask in complementary data
complementary_data = new_transition.get(TransitionKey.COMPLEMENTARY_DATA, {})
if complementary_data is None:
complementary_data = {}
complementary_data[ACTION_TOKEN_MASK] = mask
complementary_data[ACTION_CODE_TOKEN_MASK] = code_mask
complementary_data[ACTION_TOKENS] = tokens
new_transition[TransitionKey.COMPLEMENTARY_DATA] = complementary_data
return new_transition
@@ -430,7 +434,7 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
"""
return self._paligemma_tokenizer.vocab_size - 1 - self.fast_skip_tokens - tokens
def _tokenize_action(self, action: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
def _tokenize_action(self, action: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Tokenizes the action tensor and creates a mask.
@@ -459,6 +463,7 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
# The fast tokenizer expects action data and returns token IDs
tokens_list = []
masks_list = []
code_masks_list = []
for i in range(batch_size):
# Tokenize single action (move to CPU first as tokenizer uses scipy which requires numpy)
@@ -476,65 +481,82 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
if tokens.dim() > 1:
tokens = tokens.flatten()
action_code_tokens = self._act_tokens_to_paligemma_tokens(tokens)
bos_id = self._paligemma_tokenizer.bos_token_id
# add bos
prompt_tokens = torch.tensor(
self._paligemma_tokenizer.encode("Action: ", add_special_tokens=False),
device=action.device,
)
end_tokens = torch.tensor(self._paligemma_tokenizer.encode("|"), device=action.device)
code_start = 1 + len(prompt_tokens)
code_end = code_start + len(action_code_tokens)
tokens = torch.cat(
[
torch.tensor([bos_id], device=action.device),
torch.tensor(
self._paligemma_tokenizer.encode("Action: ", add_special_tokens=False),
device=action.device,
),
self._act_tokens_to_paligemma_tokens(tokens),
torch.tensor(self._paligemma_tokenizer.encode("|"), device=action.device),
prompt_tokens,
action_code_tokens,
end_tokens,
]
)
code_mask = torch.zeros(len(tokens), dtype=torch.bool, device=action.device)
code_mask[code_start:code_end] = True
# Truncate or pad to max_action_tokens
if len(tokens) > self.max_action_tokens:
if not self.allow_truncation:
raise ValueError(
f"FAST action sequence has {len(tokens)} tokens, exceeding "
f"max_action_tokens={self.max_action_tokens}."
)
logging.warning(
f"Token length ({len(tokens)}) exceeds max length ({self.max_action_tokens}), truncating. "
"Consider increasing the `max_action_tokens` in your model config if this happens frequently."
)
tokens = tokens[: self.max_action_tokens]
code_mask = code_mask[: self.max_action_tokens]
mask = torch.ones(self.max_action_tokens, dtype=torch.bool, device=action.device)
else:
pad_len = self.max_action_tokens - len(tokens)
mask = torch.cat(
[
torch.ones(len(tokens), dtype=torch.bool, device=action.device),
torch.zeros(
self.max_action_tokens - len(tokens), dtype=torch.bool, device=action.device
),
torch.zeros(pad_len, dtype=torch.bool, device=action.device),
]
)
code_mask = torch.nn.functional.pad(code_mask, (0, pad_len), value=False)
# Pad tokens with zeros
tokens = torch.nn.functional.pad(tokens, (0, self.max_action_tokens - len(tokens)), value=0)
tokens = torch.nn.functional.pad(tokens, (0, pad_len), value=0)
tokens_list.append(tokens)
masks_list.append(mask)
code_masks_list.append(code_mask)
# Stack into batched tensors
tokens_batch = torch.stack(tokens_list, dim=0) # (B, max_action_tokens)
masks_batch = torch.stack(masks_list, dim=0) # (B, max_action_tokens)
code_masks_batch = torch.stack(code_masks_list, dim=0) # (B, max_action_tokens)
# Remove batch dimension if input was single sample
if single_sample:
tokens_batch = tokens_batch.squeeze(0)
masks_batch = masks_batch.squeeze(0)
code_masks_batch = code_masks_batch.squeeze(0)
# Move to the same device as the input
if device is not None:
tokens_batch = tokens_batch.to(device)
masks_batch = masks_batch.to(device)
code_masks_batch = code_masks_batch.to(device)
return tokens_batch, masks_batch
return tokens_batch, masks_batch, code_masks_batch
def action(self, action: torch.Tensor) -> torch.Tensor:
"""
This method is not used since we override __call__.
Required by ActionProcessorStep ABC.
"""
tokens, _ = self._tokenize_action(action)
tokens, _, _ = self._tokenize_action(action)
return tokens
def get_config(self) -> dict[str, Any]:
@@ -550,6 +572,9 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
config = {
"trust_remote_code": self.trust_remote_code,
"max_action_tokens": self.max_action_tokens,
"fast_skip_tokens": self.fast_skip_tokens,
"paligemma_tokenizer_name": self.paligemma_tokenizer_name,
"allow_truncation": self.allow_truncation,
}
# Only save tokenizer_name if it was used to create the tokenizer
@@ -558,6 +583,14 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
return config
def save_artifacts(self, save_directory: Path) -> dict[str, str]:
artifact_path = Path("action_tokenizer")
save_pretrained = getattr(self.action_tokenizer, "save_pretrained", None)
if save_pretrained is None:
raise TypeError("Action tokenizer must implement save_pretrained() to save a portable pipeline.")
save_pretrained(save_directory / artifact_path)
return {"action_tokenizer_name": artifact_path.as_posix()}
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
@@ -58,6 +58,9 @@ class BiSOFollower(BimanualMixin, Robot):
port=config.left_arm_config.port,
disable_torque_on_disconnect=config.left_arm_config.disable_torque_on_disconnect,
max_relative_target=config.left_arm_config.max_relative_target,
position_p_coefficient=config.left_arm_config.position_p_coefficient,
position_i_coefficient=config.left_arm_config.position_i_coefficient,
position_d_coefficient=config.left_arm_config.position_d_coefficient,
use_degrees=config.left_arm_config.use_degrees,
cameras=left_arm_cameras,
)
@@ -68,6 +71,9 @@ class BiSOFollower(BimanualMixin, Robot):
port=config.right_arm_config.port,
disable_torque_on_disconnect=config.right_arm_config.disable_torque_on_disconnect,
max_relative_target=config.right_arm_config.max_relative_target,
position_p_coefficient=config.right_arm_config.position_p_coefficient,
position_i_coefficient=config.right_arm_config.position_i_coefficient,
position_d_coefficient=config.right_arm_config.position_d_coefficient,
use_degrees=config.right_arm_config.use_degrees,
cameras=config.right_arm_config.cameras,
)
@@ -323,6 +323,10 @@ class LeKiwiClient(Robot):
np.ndarray: the action sent to the motors, potentially clipped.
"""
# Action values may be torch tensors (e.g. replayed from a dataset) or numpy
# scalars; json.dumps only serializes Python primitives, so coerce each value to a
# plain float before sending.
action = {key: float(value) for key, value in action.items()}
self.zmq_cmd_socket.send_string(json.dumps(action)) # action is in motor space
# TODO(Steven): Remove the np conversion when it is possible to record a non-numpy array value
@@ -150,9 +150,6 @@ class OpenArmFollower(Robot):
self.configure()
if self.is_calibrated:
self.bus.set_zero_position()
self.bus.enable_torque()
logger.info(f"{self} connected.")
@@ -41,6 +41,11 @@ class SOFollowerConfig:
# Set to `True` for backward compatibility with previous policies/dataset
use_degrees: bool = True
# Position-mode PID gains written to Feetech STS3215 motors at connect time.
position_p_coefficient: int = 16
position_i_coefficient: int = 0
position_d_coefficient: int = 32
@RobotConfig.register_subclass("so101_follower")
@RobotConfig.register_subclass("so100_follower")
@@ -161,11 +161,9 @@ class SOFollower(Robot):
self.bus.configure_motors()
for motor in self.bus.motors:
self.bus.write("Operating_Mode", motor, OperatingMode.POSITION.value)
# Set P_Coefficient to lower value to avoid shakiness (Default is 32)
self.bus.write("P_Coefficient", motor, 16)
# Set I_Coefficient and D_Coefficient to default value 0 and 32
self.bus.write("I_Coefficient", motor, 0)
self.bus.write("D_Coefficient", motor, 32)
self.bus.write("P_Coefficient", motor, self.config.position_p_coefficient)
self.bus.write("I_Coefficient", motor, self.config.position_i_coefficient)
self.bus.write("D_Coefficient", motor, self.config.position_d_coefficient)
if motor == "gripper":
self.bus.write("Max_Torque_Limit", motor, 500) # 50% of max torque to avoid burnout
+3 -1
View File
@@ -21,6 +21,8 @@ from lerobot.utils.import_utils import make_device_from_device_class
from .config import RobotConfig
from .robot import Robot
logger = logging.getLogger(__name__)
def make_robot_from_config(config: RobotConfig) -> Robot:
# TODO(Steven): Consider just using the make_device_from_device_class for all types
@@ -118,7 +120,7 @@ def ensure_safe_goal_position(
}
if warnings_dict:
logging.warning(
logger.warning(
"Relative goal position magnitude had to be clamped to be safe.\n"
f"{pformat(warnings_dict, indent=4)}"
)
+11 -2
View File
@@ -326,8 +326,17 @@ class RolloutConfig:
policy_path = parser.get_path_arg("policy")
if policy_path:
cli_overrides = parser.get_cli_overrides("policy")
self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=cli_overrides)
yaml_overrides = parser.get_yaml_overrides("policy")
cli_overrides = parser.get_cli_overrides("policy") or []
policy_overrides = yaml_overrides + cli_overrides
pretrained_revision = parser.parse_arg("pretrained_revision", cli_overrides)
if pretrained_revision is None:
pretrained_revision = parser.parse_arg("pretrained_revision", yaml_overrides)
self.policy = PreTrainedConfig.from_pretrained(
policy_path,
revision=pretrained_revision,
cli_overrides=policy_overrides,
)
self.policy.pretrained_path = policy_path
if self.policy is None:
raise ValueError("--policy.path is required for rollout")
+32 -13
View File
@@ -27,7 +27,7 @@ from threading import Event
import torch
from lerobot.configs import FeatureType
from lerobot.configs import FeatureType, PreTrainedConfig
from lerobot.datasets import (
LeRobotDataset,
aggregate_pipeline_dataset_features,
@@ -159,6 +159,35 @@ class RolloutContext:
# ---------------------------------------------------------------------------
def _load_pretrained_policy(policy_config: PreTrainedConfig) -> PreTrainedPolicy:
"""Load policy weights, keeping adapter and base-model revisions independent."""
pretrained_revision = policy_config.pretrained_revision
policy_class = get_policy_class(policy_config.type)
if not policy_config.use_peft:
return policy_class.from_pretrained(
policy_config.pretrained_path,
config=policy_config,
revision=pretrained_revision,
)
from peft import PeftConfig, PeftModel
peft_path = policy_config.pretrained_path
peft_config = PeftConfig.from_pretrained(peft_path, revision=pretrained_revision)
policy = policy_class.from_pretrained(
pretrained_name_or_path=peft_config.base_model_name_or_path,
config=policy_config,
revision=peft_config.revision,
)
return PeftModel.from_pretrained(
policy,
peft_path,
config=peft_config,
revision=pretrained_revision,
)
def build_rollout_context(
cfg: RolloutConfig,
shutdown_event: Event,
@@ -176,7 +205,6 @@ def build_rollout_context(
# --- 1. Policy (heavy I/O, but no hardware yet) -------------------
logger.info("Loading policy from '%s'...", cfg.policy.pretrained_path)
policy_config = cfg.policy
policy_class = get_policy_class(policy_config.type)
if hasattr(policy_config, "compile_model"):
policy_config.compile_model = cfg.use_torch_compile
@@ -187,17 +215,7 @@ def build_rollout_context(
"Please use `cpu` or `cuda` backend."
)
if policy_config.use_peft:
from peft import PeftConfig, PeftModel
peft_path = policy_config.pretrained_path
peft_config = PeftConfig.from_pretrained(peft_path)
policy = policy_class.from_pretrained(
pretrained_name_or_path=peft_config.base_model_name_or_path, config=policy_config
)
policy = PeftModel.from_pretrained(policy, peft_path, config=peft_config)
else:
policy = policy_class.from_pretrained(policy_config.pretrained_path, config=policy_config)
policy = _load_pretrained_policy(policy_config)
if is_rtc:
policy.config.rtc_config = cfg.inference.rtc
@@ -392,6 +410,7 @@ def build_rollout_context(
preprocessor, postprocessor = make_pre_post_processors(
policy_cfg=policy_config,
pretrained_path=cfg.policy.pretrained_path,
pretrained_revision=policy_config.pretrained_revision,
dataset_stats=dataset_stats,
preprocessor_overrides={
"device_processor": {"device": cfg.device},
+38
View File
@@ -0,0 +1,38 @@
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Policy-agnostic runtime for language-conditioned policies.
Adapters registered in :mod:`lerobot.runtime.registry` are served by ``lerobot-rollout --language``.
"""
from .adapter import BaseLanguageAdapter, GenerationConfig, LanguageDiagnostics
from .language_runtime import (
LanguageConditionedPolicyAdapter,
LanguageConditionedRuntime,
RuntimeState,
Tick,
TickClock,
)
__all__ = [
"BaseLanguageAdapter",
"GenerationConfig",
"LanguageConditionedPolicyAdapter",
"LanguageConditionedRuntime",
"LanguageDiagnostics",
"RuntimeState",
"Tick",
"TickClock",
]
+165
View File
@@ -0,0 +1,165 @@
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Policy adapters for the language runtime.
The base adapter owns generation control and diagnostics while subclasses provide policy-specific actions and text.
"""
from __future__ import annotations
import re
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any
from .language_runtime import RuntimeState
_SAY_RE = re.compile(r"<\s*say\s*>(.*?)<\s*/\s*say\s*>", re.IGNORECASE | re.DOTALL)
@dataclass
class GenerationConfig:
"""Text-generation settings fixed for the adapter's lifetime."""
min_new_tokens: int = 0
temperature: float = 0.0
top_p: float = 1.0
chunks_per_regen: int = 1 # regenerate the language context every N action chunks
enable_memory: bool = True # generate a running memory note on subtask change
enable_subtask: bool = True # generate the low-level subtask (off => use the given text directly)
@dataclass
class LanguageDiagnostics:
"""Runtime-panel generation counters keyed by text kind."""
last_raw: dict[str, str] = field(default_factory=dict)
empty: dict[str, int] = field(default_factory=dict)
repeat: int = 0
def _bump(self, table: dict[str, int], kind: str) -> int:
table[kind] = table.get(kind, 0) + 1
return table[kind]
class BaseLanguageAdapter(ABC):
"""Batteries-included adapter: generic high-level control, policy primitives abstract."""
def __init__(self, policy: Any, gen: GenerationConfig | None = None) -> None:
self.policy = policy
self.gen = gen or GenerationConfig()
self.diag = LanguageDiagnostics()
self._chunks_until_regen = 0
@abstractmethod
def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any:
"""Produce an action chunk from the observation + current language context."""
@abstractmethod
def generate_text(
self,
kind: str,
observation: dict[str, Any] | None,
state: RuntimeState,
user_text: str | None = None,
) -> str:
"""Generate one text stream (``kind``) and return the decoded string."""
def update_language_state(self, observation: dict[str, Any] | None, state: RuntimeState) -> None:
"""Throttled regeneration of the language context (subtask / memory / ...)."""
if self._chunks_until_regen > 0:
self._chunks_until_regen -= 1
return
self._chunks_until_regen = max(1, self.gen.chunks_per_regen) - 1
self._regenerate_context(observation, state)
def handle_interjection(
self, user_text: str, observation: dict[str, Any] | None, state: RuntimeState
) -> None:
"""React to a mid-run user message by regenerating the plan."""
out = self.generate_text("interjection", observation, state, user_text=user_text)
plan = self.plan_from_text(out)
if plan:
state.set_context("plan", plan, label="plan")
def plan_from_text(self, text: str) -> str:
"""Strip ``<say>`` speech markers from a generated plan."""
plan, _speech = split_plan_and_say(text)
return plan
def _regenerate_context(self, observation: dict[str, Any] | None, state: RuntimeState) -> None:
"""Default hierarchy: regenerate the subtask, then memory when it changes.
Override for a policy with a different language hierarchy.
"""
if not self.gen.enable_subtask:
# Preserve operator-provided subtasks in direct mode.
return
subtask = self._generate_filtered("subtask", observation, state)
if subtask is None:
return
previous = state.language_context.get("subtask")
if not state.set_context("subtask", subtask, label="subtask"):
self.diag.repeat += 1
return
self.diag.repeat = 0
if previous:
state.extra["prior_subtask"] = previous
if not self.gen.enable_memory:
return
memory = self._generate_filtered("memory", observation, state)
if memory is not None:
state.set_context("memory", memory, label="memory")
def _generate_filtered(
self, kind: str, observation: dict[str, Any] | None, state: RuntimeState
) -> str | None:
"""Generate one ``kind``, record diagnostics, and drop empty output."""
text = self.generate_text(kind, observation, state)
self.diag.last_raw[kind] = text or ""
if not text:
count = self.diag._bump(self.diag.empty, kind)
if count == 1 or count % 5 == 0:
state.log(f" [info] {kind} gen returned empty (x{count})")
return None
return text
class DirectTaskPolicyAdapter(BaseLanguageAdapter):
"""Adapter for flat policies whose preprocessors condition actions on the operator's task."""
def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any:
return self.policy.predict_action_chunk(observation)
def generate_text(
self,
kind: str,
observation: dict[str, Any] | None,
state: RuntimeState,
user_text: str | None = None,
) -> str:
return ""
def split_plan_and_say(text: str) -> tuple[str, str]:
"""Split ``plan <say>speech</say>`` into ``(plan, speech)``."""
if not text:
return "", ""
match = _SAY_RE.search(text)
if not match:
return text.strip(), ""
speech = match.group(1).strip().strip('"').strip("'")
plan = (text[: match.start()] + text[match.end() :]).strip()
return plan, speech
File diff suppressed because it is too large Load Diff
+349
View File
@@ -0,0 +1,349 @@
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Small reusable runtime for language-conditioned robot policies."""
from __future__ import annotations
import logging
import threading
import time
from collections import deque
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any, Protocol
logger = logging.getLogger(__name__)
@dataclass
class RuntimeState:
"""Explicit state shared by the runtime and policy adapter."""
task: str = ""
language_context: dict[str, str] = field(default_factory=dict)
action_queue: deque[Any] = field(default_factory=deque)
events: set[str] = field(default_factory=set)
log_lines: list[str] = field(default_factory=list)
mode: str = "action"
stop: bool = False
tick: Tick | None = None
actions_dispatched: int = 0
action_deadline: float | None = None
extra: dict[str, Any] = field(default_factory=dict)
revision: int = 0
lock: Any = field(default_factory=threading.RLock, repr=False)
def emit(self, event_name: str) -> None:
self.events.add(event_name)
def take_event(self, event_name: str) -> bool:
if event_name not in self.events:
return False
self.events.remove(event_name)
return True
def log(self, line: str) -> None:
self.log_lines.append(line)
def set_context(self, key: str, value: str | None, *, label: str | None = None) -> bool:
with self.lock:
previous = self.language_context.get(key)
if previous == value:
return False
if value is None:
self.language_context.pop(key, None)
else:
self.language_context[key] = value
self.revision += 1
if label is not None and value:
self.log(f" {label}: {value}")
return True
def get(self, key: str, default: Any = None) -> Any:
try:
return self[key]
except KeyError:
return default
def setdefault(self, key: str, default: Any = None) -> Any:
current = self.get(key, None)
if current is not None:
return current
self[key] = default
return default
def __getitem__(self, key: str) -> Any:
if hasattr(self, key):
return getattr(self, key)
if key in self.extra:
return self.extra[key]
raise KeyError(key)
def __setitem__(self, key: str, value: Any) -> None:
with self.lock:
if hasattr(self, key):
if key == "mode" and self.mode != value:
self.revision += 1
setattr(self, key, value)
else:
self.extra[key] = value
class LanguageConditionedPolicyAdapter(Protocol):
"""Runtime policy contract, implemented directly or through ``BaseLanguageAdapter``."""
def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any: ...
def update_language_state(self, observation: dict[str, Any] | None, state: RuntimeState) -> None: ...
def handle_interjection(
self, user_text: str, observation: dict[str, Any] | None, state: RuntimeState
) -> None: ...
@dataclass
class Tick:
index: int
monotonic_seconds: float
@dataclass
class TickClock:
max_rate_hz: float = 50.0
_index: int = field(default=0, init=False)
_last_seconds: float | None = field(default=None, init=False)
def advance(self) -> Tick:
period = 1.0 / max(self.max_rate_hz, 0.1)
now = time.monotonic()
if self._last_seconds is not None:
sleep_for = (self._last_seconds + period) - now
if sleep_for > 0:
time.sleep(sleep_for)
now = time.monotonic()
self._last_seconds = now
self._index += 1
return Tick(index=self._index, monotonic_seconds=now)
@dataclass
class _RateGate:
hz: float
_last_seconds: float | None = None
def due(self, tick: Tick, *, force: bool = False) -> bool:
if force:
self._last_seconds = tick.monotonic_seconds
return True
period = 1.0 / max(self.hz, 1e-6)
if self._last_seconds is None or tick.monotonic_seconds - self._last_seconds >= period:
self._last_seconds = tick.monotonic_seconds
return True
return False
def rearm(self) -> None:
self._last_seconds = None
@dataclass
class LanguageConditionedRuntime:
"""Generic tick loop for language-conditioned robot policies."""
policy_adapter: LanguageConditionedPolicyAdapter
observation_provider: Callable[[], dict[str, Any] | None] | None = None
action_executor: Callable[[Any], None] | None = None
event_collector: Callable[[RuntimeState], None] | None = None
chunk_hz: float = 4.0
ctrl_hz: float = 50.0
high_level_hz: float = 1.0
max_rate_hz: float = 50.0
state: RuntimeState = field(default_factory=RuntimeState)
_chunk_gate: _RateGate = field(init=False)
_ctrl_gate: _RateGate = field(init=False)
_language_gate: _RateGate = field(init=False)
_stop: bool = field(default=False, init=False)
_last_dispatch_seconds: float | None = field(default=None, init=False)
def __post_init__(self) -> None:
self._chunk_gate = _RateGate(self.chunk_hz)
self._ctrl_gate = _RateGate(self.ctrl_hz)
self._language_gate = _RateGate(self.high_level_hz)
@property
def policy(self) -> Any:
return getattr(self.policy_adapter, "policy", self.policy_adapter)
def set_task(self, task: str) -> None:
with self.state.lock:
if self.state.task != task:
self.state.revision += 1
self.state.task = task
self.state.log(f"Task: {task}")
def stop(self) -> None:
self._stop = True
self.state.stop = True
def run(self, *, max_ticks: int | None = None) -> None:
clock = TickClock(max_rate_hz=self.max_rate_hz)
while not self._stop:
tick = clock.advance()
self._run_tick(tick)
self._flush_logs()
if self.state.stop:
self._stop = True
if max_ticks is not None and tick.index >= max_ticks:
break
self._on_shutdown()
def step_once(self) -> list[str]:
previous = self.state.tick.index if self.state.tick is not None else 0
tick = Tick(index=previous + 1, monotonic_seconds=time.monotonic())
self._run_tick(tick, force_rates=True)
return list(self.state.log_lines)
def _run_tick(self, tick: Tick, *, force_rates: bool = False) -> None:
self.state.tick = tick
self.state.log_lines = []
if self.event_collector is not None:
self.event_collector(self.state)
self._handle_action_deadline()
if self.state.stop:
return
self.maybe_update_language_state(force=force_rates)
self.maybe_handle_user_events()
self.maybe_enqueue_action_chunk(force=force_rates)
self.dispatch_action(force=force_rates)
self.state.events.clear()
def _current_observation(self) -> dict[str, Any] | None:
if self.observation_provider is None:
return None
try:
return self.observation_provider()
except Exception as exc: # noqa: BLE001
logger.debug("observation_provider failed: %s", exc)
return None
def maybe_update_language_state(self, *, force: bool = False) -> None:
if self.state.mode != "action" or not self.state.task:
return
if self.state.action_queue:
self._language_gate.rearm()
return
if self.state.tick is None or not self._language_gate.due(self.state.tick, force=force):
return
observation = self._current_observation()
try:
self.policy_adapter.update_language_state(observation, self.state)
except Exception as exc: # noqa: BLE001
logger.warning("language update failed: %s", exc, exc_info=logger.isEnabledFor(logging.DEBUG))
self.state.log(f" [warn] language update failed: {type(exc).__name__}: {exc}")
def maybe_handle_user_events(self) -> None:
if self.state.take_event("user_interjection"):
self._handle_user_interjection()
def _handle_user_interjection(self) -> None:
text = str(self.state.extra.get("recent_interjection") or "")
if not text:
return
observation = self._current_observation()
self.policy_adapter.handle_interjection(text, observation, self.state)
self.state.extra["recent_interjection"] = None
def maybe_enqueue_action_chunk(self, *, force: bool = False) -> None:
with self.state.lock:
if self.state.mode != "action" or not self.state.task:
return
if self.state.action_queue:
return
if self.state.tick is None or not self._chunk_gate.due(self.state.tick, force=force):
return
revision = self.state.revision
observation = self._current_observation()
if observation is None:
return
try:
chunk = self.policy_adapter.select_action(observation, self.state)
except Exception as exc: # noqa: BLE001
logger.warning("select_action failed: %s", exc, exc_info=logger.isEnabledFor(logging.DEBUG))
self.state.log(f" [warn] select_action failed: {type(exc).__name__}: {exc}")
return
with self.state.lock:
if (
self.state.revision != revision
or self.state.mode != "action"
or self.state.stop
or self._stop
):
logger.info("Discarded an action chunk invalidated during inference.")
return
self._enqueue_chunk(chunk)
def _enqueue_chunk(self, chunk: Any) -> None:
if chunk is None:
return
chunk_iter = chunk[0] if getattr(chunk, "ndim", None) == 3 else chunk
if getattr(chunk_iter, "ndim", None) == 1:
chunk_iter = chunk_iter.unsqueeze(0)
for step in chunk_iter:
self.state.action_queue.append(step.unsqueeze(0) if hasattr(step, "unsqueeze") else step)
try:
self.state.extra["last_chunk_size"] = int(chunk_iter.shape[0])
except Exception: # noqa: BLE001
self.state.extra["last_chunk_size"] = len(self.state.action_queue)
def dispatch_action(self, *, force: bool = False) -> None:
if self.state.mode != "action":
self._last_dispatch_seconds = None
return
if self.state.tick is None or not self._ctrl_gate.due(self.state.tick, force=force):
return
queue = self.state.action_queue
if not queue:
self._last_dispatch_seconds = None
return
now = time.monotonic()
if self._last_dispatch_seconds is None or self.ctrl_hz <= 0:
n_to_pop = 1
else:
n_to_pop = max(1, min(len(queue), int(round((now - self._last_dispatch_seconds) * self.ctrl_hz))))
self._last_dispatch_seconds = now
latest = None
for _ in range(n_to_pop):
if not queue:
break
latest = queue.popleft()
self.state.actions_dispatched += 1
if latest is not None and self.action_executor is not None:
self.action_executor(latest)
def _handle_action_deadline(self) -> None:
deadline = self.state.action_deadline
if self.state.mode == "action" and deadline is not None and time.monotonic() >= deadline:
self.state.mode = "paused"
self.state.action_deadline = None
self.state.action_queue.clear()
self.state.log("timed action elapsed — paused")
def _flush_logs(self) -> None:
for line in self.state.log_lines:
print(f"[runtime] {line}", flush=True)
def _on_shutdown(self) -> None:
self.state.action_queue.clear()
print("[runtime] stopped", flush=True)
+39
View File
@@ -0,0 +1,39 @@
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Lazy mapping from policy types to language-runtime adapters."""
from __future__ import annotations
import importlib
from collections.abc import Callable
from typing import Any
_ADAPTERS: dict[str, str] = {
"pi052": "lerobot.policies.pi052.inference.pi052_adapter:PI052PolicyAdapter",
"pi05": "lerobot.runtime.adapter:DirectTaskPolicyAdapter",
"molmoact2": "lerobot.runtime.adapter:DirectTaskPolicyAdapter",
}
def get_language_adapter_factory(policy_type: str) -> Callable[..., Any]:
"""Return the adapter class registered for ``policy_type``."""
spec = _ADAPTERS.get(policy_type)
if spec is None:
raise ValueError(
f"No language-runtime adapter registered for policy type {policy_type!r}. "
f"Registered: {sorted(_ADAPTERS)}. Add an entry to lerobot.runtime.registry."
)
module_path, class_name = spec.split(":")
return getattr(importlib.import_module(module_path), class_name)
+406
View File
@@ -0,0 +1,406 @@
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""RoboCasa backend for interactive language-conditioned rollouts.
It reuses the eval observation/action pipeline while prompts control a persistent selected scene.
"""
from __future__ import annotations
import logging
from collections.abc import Callable
from datetime import datetime
from pathlib import Path
from typing import Any
import numpy as np
import torch
from lerobot.utils.io_utils import StreamingVideoWriter
from lerobot.utils.video_annotation import annotate_frame
logger = logging.getLogger(__name__)
def _short_cam_name(cam: str) -> str:
"""Human-friendly view label for a RoboCasa camera name."""
c = cam.replace("robot0_", "")
return {
"agentview_left": "left",
"agentview_right": "right",
"eye_in_hand": "wrist",
}.get(c, c)
def _label_panel(img: np.ndarray, label: str) -> np.ndarray:
"""Draw a small camera-view label in the bottom-left corner of a panel."""
try:
import cv2 # noqa: PLC0415
except ImportError:
return img
y = img.shape[0] - 6
cv2.putText(img, label, (5, y), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 0, 0), 3, cv2.LINE_AA)
cv2.putText(img, label, (5, y), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (255, 255, 0), 1, cv2.LINE_AA)
return img
# Two workers avoid broken single-worker EGL rendering; only env 0 is displayed.
_SIM_N_ENVS = 2
def create_sim_env(
*,
task: str,
split: str | None,
obj_registries: list[str],
seed: int | None,
render_size: int = 384,
) -> tuple[Any, dict]:
"""Create and reset the vectorized RoboCasa environment before CUDA initializes.
Two workers keep EGL stable, while only env 0 is driven and displayed.
"""
from lerobot.envs.configs import RoboCasaEnv as RoboCasaEnvConfig # noqa: PLC0415
# The policy resizes inputs, so render_size only affects display quality and cost.
env_cfg = RoboCasaEnvConfig(
task=task,
split=split,
obj_registries=list(obj_registries),
observation_height=render_size,
observation_width=render_size,
)
# Keep one kitchen alive across sequential prompts.
envs = env_cfg.create_envs(
n_envs=_SIM_N_ENVS,
use_async_envs=True,
terminate_on_success=False,
horizon=100_000,
)
env = envs[next(iter(envs))][0]
logger.info("[sim] resetting RoboCasa scene task=%r split=%r (n_envs=%d)", task, split, _SIM_N_ENVS)
seeds = None if seed is None else [seed + i for i in range(_SIM_N_ENVS)]
obs, _ = env.reset(seed=seeds)
return env, obs
def start_mjpeg_server(port: int, get_frame: Callable[[], np.ndarray | None]) -> Any:
"""Start an MJPEG server that shows a placeholder until ``get_frame`` returns frames."""
import io # noqa: PLC0415
import threading # noqa: PLC0415
import time # noqa: PLC0415
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer # noqa: PLC0415
from PIL import Image # noqa: PLC0415
_placeholder = Image.new("RGB", (256, 256), (17, 17, 17))
class _Handler(BaseHTTPRequestHandler):
def log_message(self, *args): # silence per-request logging
pass
def do_GET(self): # noqa: N802
if self.path in ("/", "/index.html"):
self.send_response(200)
self.send_header("Content-Type", "text/html")
self.end_headers()
self.wfile.write(
b"<html><body style='margin:0;background:#111;text-align:center'>"
b"<img src='/stream' style='max-width:100vw;max-height:100vh;"
b"image-rendering:pixelated'></body></html>"
)
return
if self.path != "/stream":
self.send_response(404)
self.end_headers()
return
self.send_response(200)
self.send_header("Content-Type", "multipart/x-mixed-replace; boundary=frame")
self.end_headers()
try:
while True:
frame = get_frame()
buf = io.BytesIO()
img = Image.fromarray(frame) if frame is not None else _placeholder
img.save(buf, format="JPEG", quality=80)
data = buf.getvalue()
self.wfile.write(
b"--frame\r\nContent-Type: image/jpeg\r\nContent-Length: "
+ str(len(data)).encode()
+ b"\r\n\r\n"
+ data
+ b"\r\n"
)
time.sleep(0.05)
except (BrokenPipeError, ConnectionResetError):
pass
try:
# Bind all interfaces intentionally so the viewer remains reachable
# through the documented SSH port-forwarding workflow.
server = ThreadingHTTPServer(("0.0.0.0", port), _Handler) # nosec B104
except OSError as exc:
logger.warning("[sim] could not start live stream on port %d: %s", port, exc)
print(f"[runtime] WARNING: live stream port {port} unavailable ({exc})", flush=True)
return None
threading.Thread(target=server.serve_forever, daemon=True, name="sim-mjpeg").start()
print(
f"[runtime] live view: http://localhost:{port} "
f"(over SSH: ssh -L {port}:localhost:{port} <host>) — loading until scene is ready",
flush=True,
)
return server
class RoboCasaSimBackend:
"""Expose a RoboCasa environment through the runtime observation/action contract.
The environment must be created before the policy initializes CUDA.
"""
def __init__(
self,
*,
env: Any,
last_obs: dict,
task: str,
seed: int | None,
device: str,
preprocessor: Any,
postprocessor: Any,
record: bool = True,
output_dir: str | None = None,
view_cams: list[str] | None = None,
) -> None:
self.env = env
self._last_obs = last_obs
self._scene_task = task
self._view_cams = view_cams or [
"robot0_agentview_left",
"robot0_eye_in_hand",
"robot0_agentview_right",
]
self.device = torch.device(device) if isinstance(device, str) else device
self.preprocessor = preprocessor
self.postprocessor = postprocessor
self.seed = seed
self.record = record
self.output_dir = Path(output_dir) if output_dir else Path("outputs/runtime_sim")
self._video_writer: StreamingVideoWriter | None = None
self._video_path: Path | None = None
self._live_counter = 0
self._latest_frame: np.ndarray | None = None
self._stream_server: Any = None
self._reset_count = 0
# Bind these after runtime construction for live annotations.
self._task_getter: Callable[[], str | None] | None = None
self._subtask_getter: Callable[[], str | None] | None = None
self._memory_getter: Callable[[], str | None] | None = None
logger.info("[sim] scene ready — task_description=%r", self._scene_description())
def bind_runtime(self, runtime: Any) -> None:
"""Wire live task/subtask/memory getters from the runtime state."""
self._task_getter = lambda: runtime.state.get("task")
self._subtask_getter = lambda: runtime.state.language_context.get("subtask")
self._memory_getter = lambda: (runtime.state.get("language_context") or {}).get("memory")
def _scene_description(self) -> str:
try:
return str(self.env.get_attr("task_description")[0]) or self._scene_task
except Exception: # noqa: BLE001
return self._scene_task
def _current_task(self) -> str:
task = self._task_getter() if self._task_getter else None
return task or self._scene_description() or self._scene_task
def reset_scene(self) -> None:
"""Re-roll the kitchen: reset the env to a fresh scene (new layout/style).
Uses a new seed each call so ``/reset`` explores different kitchens.
"""
self._reset_count += 1
n = self.env.num_envs
if self.seed is None:
seeds = None
else:
base = self.seed + self._reset_count * 1000
seeds = [base + i for i in range(n)]
obs, _ = self.env.reset(seed=seeds)
self._last_obs = obs
logger.info("[sim] scene reset (#%d)", self._reset_count)
def _env0_obs(self) -> dict:
"""Slice env 0 out of the batched vec-env observation (batch of 1)."""
raw = self._last_obs or {}
pixels = raw.get("pixels")
out: dict[str, Any] = {}
if isinstance(pixels, dict):
out["pixels"] = {k: np.asarray(v)[0:1] for k, v in pixels.items()}
agent_pos = raw.get("agent_pos")
if agent_pos is not None:
out["agent_pos"] = np.asarray(agent_pos)[0:1]
return out
def observation_provider(self) -> dict | None:
from lerobot.envs.utils import preprocess_observation # noqa: PLC0415
try:
obs = preprocess_observation(self._env0_obs())
except Exception as exc: # noqa: BLE001
logger.warning("[sim] preprocess_observation failed: %s", exc)
return None
# The adapter later replaces this recipe input with its generated subtask.
obs["task"] = [self._current_task()]
if self.preprocessor is not None:
try:
obs = self.preprocessor(obs)
except Exception as exc: # noqa: BLE001
logger.warning("[sim] preprocessor failed: %s", exc)
return None
return {
k: (v.to(self.device) if isinstance(v, torch.Tensor) else v)
for k, v in obs.items()
if isinstance(k, str) and k.startswith("observation.")
}
def action_executor(self, action: Any) -> None:
try:
if self.postprocessor is not None:
action = self.postprocessor(action)
if isinstance(action, torch.Tensor):
if action.ndim > 1 and action.shape[0] == 1:
action = action.squeeze(0)
action = action.detach().to("cpu").numpy()
# Tile env 0's action because the extra workers exist only for EGL stability.
action_row = np.asarray(action, dtype=np.float32).reshape(-1)
action_np = np.tile(action_row, (self.env.num_envs, 1))
obs, _reward, terminated, truncated, _info = self.env.step(action_np)
self._last_obs = obs
self._capture_frame()
# AsyncVectorEnv resets terminated sub-environments automatically.
if bool(np.any(terminated)) or bool(np.any(truncated)):
logger.info("[sim] episode ended — scene auto-reset")
except Exception as exc: # noqa: BLE001
logger.error("[sim] env.step failed: %s", exc, exc_info=True)
def _multiview_frame(self) -> np.ndarray | None:
"""Label and compose env 0's existing observation views without extra rendering."""
pixels = (self._last_obs or {}).get("pixels")
if not isinstance(pixels, dict) or not pixels:
return None
panels: list[np.ndarray] = []
for cam in self._view_cams:
v = pixels.get(cam)
if v is None:
continue
img = np.asarray(v)
if img.ndim == 4: # (n_envs, H, W, C) -> env 0
img = img[0]
if img.ndim != 3 or img.shape[-1] != 3:
continue
panels.append(_label_panel(np.ascontiguousarray(img.astype(np.uint8)), _short_cam_name(cam)))
if not panels:
return None
h = min(p.shape[0] for p in panels)
panels = [p[:h] for p in panels]
return np.concatenate(panels, axis=1)
def _capture_frame(self) -> None:
frame = self._multiview_frame()
if frame is None: # fallback to single env.render()
try:
rendered = self.env.call("render")[0]
if isinstance(rendered, np.ndarray) and rendered.ndim == 3:
frame = rendered
except Exception as exc: # noqa: BLE001
logger.debug("[sim] render failed: %s", exc)
if frame is None:
return
subtask = self._subtask_getter() if self._subtask_getter else None
memory = self._memory_getter() if self._memory_getter else None
annotated = annotate_frame(
frame,
(("Task", self._current_task()), ("Subtask", subtask), ("Memory", memory)),
)
self._latest_frame = annotated # served by the live MJPEG stream
self._write_live_frame(annotated)
if self.record:
self._write_recording_frame(annotated)
def _write_recording_frame(self, frame: np.ndarray) -> None:
try:
if self._video_writer is None:
self.output_dir.mkdir(parents=True, exist_ok=True)
stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
self._video_path = self.output_dir / f"sim_{stamp}.mp4"
fps = int((getattr(self.env, "metadata", None) or {}).get("render_fps", 20))
self._video_writer = StreamingVideoWriter(self._video_path, fps)
self._video_writer.add_frame(frame)
except Exception as exc: # noqa: BLE001
logger.warning("[sim] video encoding failed: %s", exc)
self.record = False
def _write_live_frame(self, frame: np.ndarray) -> None:
"""Write a rolling latest.png every few frames for live viewing over SSH.
Open ``{output_dir}/latest.png`` in an editor/viewer and refresh to watch
the rollout in near-real-time without a GUI window. Written atomically
(temp + replace) so a reader never sees a half-written file.
"""
self._live_counter += 1
if self._live_counter % 3 != 0:
return
try:
import os # noqa: PLC0415
from PIL import Image # noqa: PLC0415
self.output_dir.mkdir(parents=True, exist_ok=True)
tmp = self.output_dir / ".latest.tmp.png"
Image.fromarray(frame).save(tmp)
os.replace(tmp, self.output_dir / "latest.png")
except Exception as exc: # noqa: BLE001
logger.debug("[sim] live frame write failed: %s", exc)
def _flush_video(self) -> None:
if self._video_writer is None:
return
writer = self._video_writer
self._video_writer = None
try:
writer.close()
logger.info("[sim] wrote video (%d frames) to %s", writer.frames_written, self._video_path)
print(f"[runtime] sim video saved to {self._video_path}", flush=True)
except Exception as exc: # noqa: BLE001
logger.warning("[sim] video close failed: %s", exc)
def attach_stream_server(self, server: Any) -> None:
"""Attach an already-running MJPEG server so disconnect() can stop it."""
self._stream_server = server
def disconnect(self) -> None:
"""Match the robot backend's cleanup contract."""
if self._stream_server is not None:
try:
self._stream_server.shutdown()
except Exception as exc: # noqa: BLE001
logger.debug("[sim] stream server shutdown raised %s", exc)
self._flush_video()
try:
self.env.close()
except Exception as exc: # noqa: BLE001
logger.debug("[sim] env.close raised %s", exc)
@@ -61,6 +61,7 @@ import pyarrow as pa
import tqdm
from datasets import Dataset, Features, Image
from huggingface_hub import HfApi, snapshot_download
from huggingface_hub.errors import RevisionNotFoundError
from requests import HTTPError
from lerobot.datasets import CODEBASE_VERSION, LeRobotDataset, aggregate_stats
@@ -521,7 +522,7 @@ def convert_dataset(
hub_api = HfApi()
try:
hub_api.delete_tag(repo_id, tag=CODEBASE_VERSION, repo_type="dataset")
except HTTPError as e:
except (HTTPError, RevisionNotFoundError) as e:
print(f"tag={CODEBASE_VERSION} probably doesn't exist. Skipping exception ({e})")
pass
hub_api.delete_files(
+16 -1
View File
@@ -24,7 +24,14 @@ Example:
--root=/path/to/dataset \\
--vlm.model_id=Qwen/Qwen2.5-VL-7B-Instruct
For distributed runs, see ``examples/annotations/run_hf_job.py``.
Pass ``--job.target=<flavor>`` to run the same command on a Hugging Face
Jobs GPU instead of this machine (see ``lerobot.jobs.annotate``):
uv run lerobot-annotate \\
--repo_id=user/dataset \\
--new_repo_id=user/dataset_annotated \\
--push_to_hub=true \\
--job.target=h200
"""
import logging
@@ -69,6 +76,14 @@ def _resolve_root(cfg: AnnotationPipelineConfig) -> Path:
def annotate(cfg: AnnotationPipelineConfig) -> None:
"""Run the steerable annotation pipeline against a dataset."""
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
if cfg.job.is_remote:
# Imported lazily: the submitter pulls in LeRobotDataset (the `dataset`
# extra), which a local annotation run over --root doesn't need.
from lerobot.jobs.annotate import submit_annotate_to_hf
return submit_annotate_to_hf(cfg)
root = _resolve_root(cfg)
logger.info("annotate: root=%s", root)
+81 -30
View File
@@ -94,6 +94,19 @@ from lerobot.utils.utils import (
init_logging,
inside_slurm,
)
from lerobot.utils.video_annotation import annotate_frame
def _annotate_eval_frames(frames: np.ndarray, task: str | None, subtask: str | None) -> np.ndarray:
"""Overlay the high-level task and predicted subtask onto rendered frames.
``frames`` is ``(n_envs, H, W, C)`` uint8. Best-effort: if OpenCV isn't
available the frames are returned unchanged so eval never fails over a
visualization concern.
"""
if frames.ndim != 4 or frames.shape[-1] != 3:
return frames
return np.stack([annotate_frame(frame, (("Task", task), ("Subtask", subtask))) for frame in frames])
def _env_features_to_dataset_features(env_features: dict) -> dict:
@@ -453,6 +466,9 @@ def eval_policy(
raise exc from None
start = time.time()
# Preserve the mode for direct callers. eval_policy_all scopes the mode
# around all tasks so parallel evaluations cannot race with each other.
was_training = policy.training
policy.eval()
# Determine how many batched rollouts we need to get n_episodes. Note that if n_episodes is not evenly
@@ -474,11 +490,36 @@ def eval_policy(
return
n_to_render_now = min(max_episodes_rendered - n_episodes_rendered, env.num_envs)
if isinstance(env, gym.vector.SyncVectorEnv):
ep_frames.append(np.stack([env.envs[i].render() for i in range(n_to_render_now)])) # noqa: B023
frames = np.stack([env.envs[i].render() for i in range(n_to_render_now)]) # noqa: B023
elif hasattr(env, "call"):
# Here we must render all frames and discard any we don't need.
# Covers AsyncVectorEnv and _LazyAsyncVectorEnv (which wraps one).
ep_frames.append(np.stack(env.call("render")[:n_to_render_now]))
frames = np.stack(env.call("render")[:n_to_render_now])
else:
return
# Overlay the high-level task and (for hierarchical policies like
# pi052) the predicted low-level subtask onto each frame. Both are
# best-effort: missing values just skip that line.
try:
tasks = list(env.call("task_description"))
except (AttributeError, NotImplementedError):
try:
tasks = list(env.call("task"))
except (AttributeError, NotImplementedError):
tasks = None
subtasks = getattr(policy, "last_subtasks", None)
annotated = []
for i in range(frames.shape[0]):
subtask_i = subtasks[i] if subtasks is not None and i < len(subtasks) else None
annotated.append(
_annotate_eval_frames(
frames[i : i + 1],
tasks[i] if tasks is not None and i < len(tasks) else None,
subtask_i,
)[0]
)
ep_frames.append(np.stack(annotated))
if max_episodes_rendered > 0:
video_paths: list[str] = []
@@ -674,6 +715,8 @@ def eval_policy(
if save_predicted_video:
info["predicted_video_paths"] = predicted_video_paths
policy.train(was_training)
return info
@@ -1010,40 +1053,48 @@ def eval_policy_all(
recording_private=recording_private,
)
if max_parallel_tasks <= 1:
prefetch_thread: threading.Thread | None = None
for i, (task_group, task_id, env) in enumerate(tasks):
if prefetch_thread is not None:
prefetch_thread.join()
prefetch_thread = None
# Set the shared policy's mode before launching any workers. Restoring it
# inside individual tasks would let one task enable training mode while
# another task is still evaluating.
was_training = policy.training
policy.eval()
try:
if max_parallel_tasks <= 1:
prefetch_thread: threading.Thread | None = None
for i, (task_group, task_id, env) in enumerate(tasks):
if prefetch_thread is not None:
prefetch_thread.join()
prefetch_thread = None
try:
tg, tid, metrics = task_runner(task_group, task_id, env)
_accumulate_to(tg, metrics)
per_task_infos.append({"task_group": tg, "task_id": tid, "metrics": metrics})
finally:
env.close()
# Prefetch next task's workers *after* closing current env to prevent
# GPU memory overlap between consecutive tasks.
if i + 1 < len(tasks):
next_env = tasks[i + 1][2]
if hasattr(next_env, "_ensure"):
prefetch_thread = threading.Thread(target=next_env._ensure, daemon=True)
prefetch_thread.start()
else:
with cf.ThreadPoolExecutor(max_workers=max_parallel_tasks) as executor:
fut2meta = {}
for task_group, task_id, env in tasks:
fut = executor.submit(task_runner, task_group, task_id, env)
fut2meta[fut] = (task_group, task_id, env)
for fut in cf.as_completed(fut2meta):
tg, tid, env = fut2meta[fut]
try:
tg, tid, metrics = fut.result()
tg, tid, metrics = task_runner(task_group, task_id, env)
_accumulate_to(tg, metrics)
per_task_infos.append({"task_group": tg, "task_id": tid, "metrics": metrics})
finally:
env.close()
# Prefetch next task's workers *after* closing current env to prevent
# GPU memory overlap between consecutive tasks.
if i + 1 < len(tasks):
next_env = tasks[i + 1][2]
if hasattr(next_env, "_ensure"):
prefetch_thread = threading.Thread(target=next_env._ensure, daemon=True)
prefetch_thread.start()
else:
with cf.ThreadPoolExecutor(max_workers=max_parallel_tasks) as executor:
fut2meta = {}
for task_group, task_id, env in tasks:
fut = executor.submit(task_runner, task_group, task_id, env)
fut2meta[fut] = (task_group, task_id, env)
for fut in cf.as_completed(fut2meta):
tg, tid, env = fut2meta[fut]
try:
tg, tid, metrics = fut.result()
_accumulate_to(tg, metrics)
per_task_infos.append({"task_group": tg, "task_id": tid, "metrics": metrics})
finally:
env.close()
finally:
policy.train(was_training)
# compute aggregated metrics helper (robust to lists/scalars)
def _agg_from_list(xs):
+3 -1
View File
@@ -453,9 +453,11 @@ def record(
encoder_queue_maxsize=cfg.dataset.encoder_queue_maxsize,
)
robot.connect()
# Connect the teleoperator before the robot so the robot isn't left idle (and possibly
# tripping a firmware watchdog) during teleop init. Matches lerobot_teleoperate.py.
if teleop is not None:
teleop.connect()
robot.connect()
listener, events = init_keyboard_listener()
+1
View File
@@ -61,6 +61,7 @@ from lerobot.robots import ( # noqa: F401
earthrover_mini_plus,
hope_jr,
koch_follower,
lekiwi,
make_robot_from_config,
omx_follower,
openarm_follower,
+63 -3
View File
@@ -151,6 +151,7 @@ Usage examples
"""
import logging
import sys
from lerobot.cameras.opencv import OpenCVCameraConfig # noqa: F401
from lerobot.cameras.realsense import RealSenseCameraConfig # noqa: F401
@@ -241,10 +242,69 @@ def rollout(cfg: RolloutConfig):
logger.info("Rollout finished")
def main():
"""CLI entry point for ``lerobot-rollout``."""
_LANGUAGE_RUNTIME_FLAGS = {
"--language",
"--no_robot",
"--sim",
"--direct_subtask",
"--sim.direct_subtask",
"--disable_memory",
"--fp8",
}
_LANGUAGE_RUNTIME_PREFIXES = (
"--sim.",
"--chunk_hz",
"--ctrl_hz",
"--high_level_hz",
"--subtask_chunks_per_gen",
"--text_min_new_tokens",
"--text_temperature",
"--text_top_p",
)
def _uses_language_runtime(argv: list[str]) -> bool:
"""Return whether *argv* selects the interactive language runtime.
``--language`` is the explicit selector for real-robot runs whose other
options overlap with the standard rollout CLI. Language-only options also
select it automatically, which keeps the former language-runtime examples
working after replacing their command name with ``lerobot-rollout``.
"""
return any(
arg.split("=", 1)[0] in _LANGUAGE_RUNTIME_FLAGS or arg.startswith(_LANGUAGE_RUNTIME_PREFIXES)
for arg in argv
)
def main(argv: list[str] | None = None):
"""CLI entry point for ``lerobot-rollout``.
Standard policy deployment continues through :class:`RolloutConfig`.
Interactive language-conditioned and RoboCasa runs share this entry point
and are selected with ``--language`` or any language-runtime-only option.
"""
register_third_party_plugins()
rollout()
cli_args = list(sys.argv[1:] if argv is None else argv)
if _uses_language_runtime(cli_args):
from lerobot.runtime.cli import run as run_language_runtime
# ``--language`` is a dispatcher flag, not part of the runtime's own
# argparse surface. All other arguments pass through unchanged.
runtime_args = [arg for arg in cli_args if arg != "--language"]
return run_language_runtime(runtime_args, prog="lerobot-rollout")
if argv is None:
return rollout()
# draccus reads sys.argv. Supporting an explicit argv keeps this entry
# point easy to smoke-test and mirrors the language-runtime branch above.
previous_argv = sys.argv
try:
sys.argv = [previous_argv[0], *cli_args]
return rollout()
finally:
sys.argv = previous_argv
if __name__ == "__main__":
+6 -17
View File
@@ -51,19 +51,7 @@ from lerobot.teleoperators import ( # noqa: F401
rebot_102_leader,
so_leader,
)
COMPATIBLE_DEVICES = [
"koch_follower",
"koch_leader",
"omx_follower",
"omx_leader",
"openarm_mini",
"so100_follower",
"so100_leader",
"so101_follower",
"so101_leader",
"lekiwi",
]
from lerobot.utils.import_utils import register_third_party_plugins
@dataclass
@@ -80,18 +68,19 @@ class SetupConfig:
@draccus.wrap()
def setup_motors(cfg: SetupConfig):
if cfg.device.type not in COMPATIBLE_DEVICES:
raise NotImplementedError
if isinstance(cfg.device, RobotConfig):
device = make_robot_from_config(cfg.device)
else:
device = make_teleoperator_from_config(cfg.device)
device.setup_motors()
setup = getattr(device, "setup_motors", None)
if not callable(setup):
raise NotImplementedError(f"Device type '{cfg.device.type}' does not support motor setup.")
setup()
def main():
register_third_party_plugins()
setup_motors()
+12 -4
View File
@@ -71,6 +71,16 @@ from lerobot.utils.utils import (
from .lerobot_eval import eval_policy_all
def _dataloader_worker_kwargs(cfg: TrainPipelineConfig) -> dict[str, Any]:
"""Return worker-only DataLoader options, disabling them for single-process loading."""
workers_enabled = cfg.num_workers > 0
return {
"prefetch_factor": cfg.prefetch_factor if workers_enabled else None,
"persistent_workers": cfg.persistent_workers and workers_enabled,
"multiprocessing_context": cfg.dataloader_multiprocessing_context if workers_enabled else None,
}
def update_policy(
train_metrics: MetricsTracker,
policy: PreTrainedPolicy,
@@ -473,8 +483,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
pin_memory=device.type == "cuda",
drop_last=False,
collate_fn=collate_fn,
prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None,
persistent_workers=cfg.persistent_workers and cfg.num_workers > 0,
**_dataloader_worker_kwargs(cfg),
)
# Build eval dataloader if a held-out split exists
@@ -500,8 +509,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
pin_memory=device.type == "cuda",
drop_last=False,
collate_fn=eval_collate_fn,
prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None,
persistent_workers=cfg.persistent_workers and cfg.num_workers > 0,
**_dataloader_worker_kwargs(cfg),
)
# Prepare everything with accelerator
@@ -23,3 +23,5 @@ from ..config import TeleoperatorConfig
@dataclass
class GamepadTeleopConfig(TeleoperatorConfig):
use_gripper: bool = True
# Use hidapi instead of pygame for controllers that pygame cannot detect reliably.
hidapi_fallback: bool = False
@@ -14,6 +14,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import logging
import sys
from enum import IntEnum
from typing import Any
@@ -27,6 +28,8 @@ from ..teleoperator import Teleoperator
from ..utils import TeleopEvents
from .configuration_gamepad import GamepadTeleopConfig
logger = logging.getLogger(__name__)
class GripperAction(IntEnum):
CLOSE = 0
@@ -56,6 +59,13 @@ class GamepadTeleop(Teleoperator):
self.gamepad = None
self.hidapi_fallback = config.hidapi_fallback
if sys.platform == "darwin" and not self.hidapi_fallback:
logger.warning(
"On macOS, pygame may not reliably detect input from some controllers. "
"If you experience issues, set `hidapi_fallback=true`."
)
@property
def action_features(self) -> dict:
if self.config.use_gripper:
@@ -76,9 +86,7 @@ class GamepadTeleop(Teleoperator):
return {}
def connect(self) -> None:
# use HidApi for macos
if sys.platform == "darwin":
# NOTE: On macOS, pygame doesnt reliably detect input from some controllers so we fall back to hidapi
if self.hidapi_fallback:
from .gamepad_utils import GamepadControllerHID as Gamepad
else:
from .gamepad_utils import GamepadController as Gamepad
+1 -1
View File
@@ -22,7 +22,7 @@ from torch.utils.data._utils.collate import default_collate
from lerobot.datasets.language import LANGUAGE_COLUMNS
_PYTHON_LIST_KEYS = {"messages", "message_streams", "target_message_indices"}
_PYTHON_LIST_KEYS = {"messages", "message_streams", "target_message_indices", *LANGUAGE_COLUMNS}
def lerobot_collate_fn(batch: list[dict[str, Any] | None]) -> dict[str, Any] | None:
+2
View File
@@ -26,6 +26,7 @@ OBS_IMAGES = OBS_IMAGE + "s"
OBS_LANGUAGE = OBS_STR + ".language"
OBS_LANGUAGE_TOKENS = OBS_LANGUAGE + ".tokens"
OBS_LANGUAGE_ATTENTION_MASK = OBS_LANGUAGE + ".attention_mask"
OBS_LANGUAGE_CAUSAL_MARKS = OBS_LANGUAGE + ".causal_marks"
OBS_LANGUAGE_SUBTASK = OBS_STR + ".subtask"
OBS_LANGUAGE_SUBTASK_TOKENS = OBS_LANGUAGE_SUBTASK + ".tokens"
OBS_LANGUAGE_SUBTASK_ATTENTION_MASK = OBS_LANGUAGE_SUBTASK + ".attention_mask"
@@ -34,6 +35,7 @@ ACTION = "action"
ACTION_PREFIX = ACTION + "."
ACTION_TOKENS = ACTION + ".tokens"
ACTION_TOKEN_MASK = ACTION + ".token_mask"
ACTION_CODE_TOKEN_MASK = ACTION + ".code_token_mask"
REWARD = "next.reward"
TRUNCATED = "next.truncated"
DONE = "next.done"
+45 -29
View File
@@ -23,6 +23,46 @@ logger = logging.getLogger(__name__)
JsonLike = str | int | float | bool | None | list["JsonLike"] | dict[str, "JsonLike"] | tuple["JsonLike", ...]
class StreamingVideoWriter:
"""Incrementally encode RGB frames to an MP4 without retaining them in memory."""
def __init__(self, video_path: str | Path, fps: int) -> None:
from .import_utils import require_package
require_package("av", extra="av-dep")
import av
self._av = av
self._container = av.open(str(video_path), mode="w")
self._stream = self._container.add_stream("libx264", rate=fps)
self._shape: tuple[int, int] | None = None
self.frames_written = 0
def add_frame(self, frame_array) -> None:
orig_height, orig_width = frame_array.shape[:2]
height = orig_height - orig_height % 2
width = orig_width - orig_width % 2
if self._shape is None:
self._shape = (height, width)
self._stream.width = width
self._stream.height = height
self._stream.pix_fmt = "yuv420p"
elif self._shape != (height, width):
raise ValueError(f"Video frame shape changed from {self._shape} to {(height, width)}")
frame = self._av.VideoFrame.from_ndarray(frame_array[:height, :width], format="rgb24")
for packet in self._stream.encode(frame):
self._container.mux(packet)
self.frames_written += 1
def close(self) -> None:
if self._container is None:
return
for packet in self._stream.encode():
self._container.mux(packet)
self._container.close()
self._container = None
def load_json(fpath: Path) -> Any:
"""Load data from a JSON file.
@@ -58,36 +98,12 @@ def write_video(video_path: str | Path, stacked_frames: list, fps: int) -> None:
stacked_frames: List of HWC uint8 numpy arrays (RGB).
fps: Frames per second for the output video.
"""
from .import_utils import require_package
require_package("av", extra="av-dep")
import av
with av.open(str(video_path), mode="w") as container:
orig_height, orig_width = stacked_frames[0].shape[:2]
# yuv420p requires even dimensions; crop by one pixel if needed
height = orig_height if orig_height % 2 == 0 else orig_height - 1
width = orig_width if orig_width % 2 == 0 else orig_width - 1
if height != orig_height or width != orig_width:
logger.warning(
"Frame dimensions %dx%d are not even; cropping to %dx%d for yuv420p compatibility.",
orig_width,
orig_height,
width,
height,
)
stream = container.add_stream("libx264", rate=fps)
stream.width = width
stream.height = height
stream.pix_fmt = "yuv420p"
writer = StreamingVideoWriter(video_path, fps)
try:
for frame_array in stacked_frames:
if height != orig_height or width != orig_width:
frame_array = frame_array[:height, :width]
frame = av.VideoFrame.from_ndarray(frame_array, format="rgb24")
for packet in stream.encode(frame):
container.mux(packet)
for packet in stream.encode():
container.mux(packet)
writer.add_frame(frame_array)
finally:
writer.close()
def deserialize_json_into_object[T: JsonLike](fpath: Path, obj: T) -> T:
+9 -1
View File
@@ -38,7 +38,10 @@ def _is_scalar(x):
def init_rerun(
session_name: str = "lerobot_control_loop", ip: str | None = None, port: int | None = None
session_name: str = "lerobot_control_loop",
ip: str | None = None,
port: int | None = None,
web_port: int | None = None,
) -> None:
"""
Initializes the Rerun SDK for visualizing the control loop.
@@ -47,6 +50,7 @@ def init_rerun(
session_name: Name of the Rerun session.
ip: Optional IP for connecting to a Rerun server.
port: Optional port for connecting to a Rerun server.
web_port: Serve a headless web viewer on this port, using ``port`` for gRPC.
"""
require_package("rerun-sdk", extra="viz", import_name="rerun")
@@ -60,6 +64,10 @@ def init_rerun(
memory_limit = os.getenv("LEROBOT_RERUN_MEMORY_LIMIT", "10%")
if ip and port:
rr.connect_grpc(url=f"rerun+http://{ip}:{port}/proxy")
elif web_port is not None:
grpc_port = port or 9876
url = rr.serve_grpc(grpc_port=grpc_port)
rr.serve_web_viewer(web_port=web_port, open_browser=False, connect_to=url)
else:
rr.spawn(memory_limit=memory_limit)
+71
View File
@@ -0,0 +1,71 @@
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Best-effort text overlays shared by evaluation and interactive rollouts."""
from __future__ import annotations
from collections.abc import Iterable
import numpy as np
def annotate_frame(frame: np.ndarray, fields: Iterable[tuple[str, str | None]]) -> np.ndarray:
"""Return an RGB frame annotated with the non-empty labeled ``fields``."""
if frame.ndim != 3 or frame.shape[-1] != 3:
return frame
try:
import cv2 # noqa: PLC0415
except ImportError:
return frame
text_rows = [f"{label}: {value}" for label, value in fields if value]
if not text_rows:
return frame
image = np.ascontiguousarray(frame).copy()
font, scale, thickness, margin = cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1, 6
max_width = image.shape[1] - 2 * margin
lines: list[str] = []
for text in text_rows:
current = ""
for word in text.split():
candidate = f"{current} {word}".strip()
width = cv2.getTextSize(candidate, font, scale, thickness)[0][0]
if width > max_width and current:
lines.append(current)
current = word
else:
current = candidate
if current:
lines.append(current)
line_height = 20
header_height = min(image.shape[0], len(lines) * line_height + 6)
backdrop = image.copy()
cv2.rectangle(backdrop, (0, 0), (image.shape[1], header_height), (0, 0, 0), -1)
cv2.addWeighted(backdrop, 0.55, image, 0.45, 0, dst=image)
for index, line in enumerate(lines):
cv2.putText(
image,
line,
(margin, 18 + index * line_height),
font,
scale,
(255, 255, 255),
thickness,
cv2.LINE_AA,
)
return image
+23 -1
View File
@@ -26,7 +26,7 @@ import cv2
import numpy as np
import pytest
from lerobot.cameras.configs import Cv2Rotation
from lerobot.cameras.configs import ColorMode, Cv2Rotation
from lerobot.cameras.opencv import OpenCVCamera, OpenCVCameraConfig
from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
@@ -132,6 +132,28 @@ def test_read(index_or_path):
assert isinstance(img, np.ndarray)
@pytest.mark.parametrize("index_or_path", TEST_IMAGE_PATHS, ids=TEST_IMAGE_SIZES)
def test_color_mode_conversion(index_or_path):
"""RGB and BGR reads of the same frame must differ only by a channel-axis reversal."""
rgb_config = OpenCVCameraConfig(index_or_path=index_or_path, color_mode=ColorMode.RGB, warmup_s=0)
bgr_config = OpenCVCameraConfig(index_or_path=index_or_path, color_mode=ColorMode.BGR, warmup_s=0)
with OpenCVCamera(rgb_config) as rgb_cam:
rgb = rgb_cam.read()
with OpenCVCamera(bgr_config) as bgr_cam:
bgr = bgr_cam.read()
assert rgb.shape == bgr.shape
np.testing.assert_array_equal(rgb, bgr[..., ::-1])
def test_postprocess_invalid_color_mode():
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH)
camera = OpenCVCamera(config)
camera.color_mode = "invalid"
with pytest.raises(ValueError):
camera._postprocess_image(np.zeros((120, 160, 3), dtype=np.uint8))
def test_read_before_connect():
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH)
+44 -15
View File
@@ -22,6 +22,7 @@ import pytest
pytest.importorskip("reachy2_sdk")
from lerobot.cameras.configs import ColorMode
from lerobot.cameras.reachy2_camera import Reachy2Camera, Reachy2CameraConfig
from lerobot.utils.errors import DeviceNotConnectedError
@@ -33,28 +34,19 @@ PARAMS = [
]
def _make_cam_manager_mock():
def _make_cam_manager_mock(color_frame, depth_frame=None):
c = MagicMock(name="CameraManagerMock")
teleop = MagicMock(name="TeleopCam")
teleop.width = 640
teleop.height = 480
teleop.get_frame = MagicMock(
side_effect=lambda *_, **__: (
np.zeros((480, 640, 3), dtype=np.uint8),
time.time(),
)
)
teleop.get_frame = MagicMock(side_effect=lambda *_, **__: (color_frame, time.time()))
depth = MagicMock(name="DepthCam")
depth.width = 640
depth.height = 480
depth.get_frame = MagicMock(
side_effect=lambda *_, **__: (
np.zeros((480, 640, 3), dtype=np.uint8),
time.time(),
)
)
depth.get_frame = MagicMock(side_effect=lambda *_, **__: (color_frame, time.time()))
depth.get_depth_frame = MagicMock(side_effect=lambda *_, **__: (depth_frame, time.time()))
c.is_connected.return_value = True
c.teleop = teleop
@@ -84,12 +76,14 @@ def _make_cam_manager_mock():
# ids=["teleop-left", "teleop-right", "torso-rgb", "torso-depth"],
ids=["teleop-left", "teleop-right", "torso-rgb"],
)
def camera(request):
def camera(request, img_array_factory):
name, image_type = request.param
color_frame = img_array_factory(height=480, width=640)
depth_frame = img_array_factory(height=480, width=640, channels=1, dtype=np.uint16)[..., 0]
with (
patch(
"lerobot.cameras.reachy2_camera.reachy2_camera.CameraManager",
side_effect=lambda *a, **k: _make_cam_manager_mock(),
side_effect=lambda *a, **k: _make_cam_manager_mock(color_frame, depth_frame),
),
):
config = Reachy2CameraConfig(name=name, image_type=image_type)
@@ -188,6 +182,41 @@ def test_read_latest_too_old(camera):
_ = camera.read_latest(max_age_ms=0) # immediately too old
def test_color_mode_conversion(img_array_factory):
"""teleop frames are native BGR: RGB reverses the channel axis, BGR is passed through."""
frame = img_array_factory(height=8, width=8)
outputs = {}
for color_mode in (ColorMode.RGB, ColorMode.BGR):
with patch(
"lerobot.cameras.reachy2_camera.reachy2_camera.CameraManager",
side_effect=lambda *a, **k: _make_cam_manager_mock(frame),
):
cam = Reachy2Camera(Reachy2CameraConfig(name="teleop", image_type="left", color_mode=color_mode))
cam.connect()
outputs[color_mode] = cam.read()
cam.disconnect()
np.testing.assert_array_equal(outputs[ColorMode.BGR], frame)
np.testing.assert_array_equal(outputs[ColorMode.RGB], frame[..., ::-1])
def test_depth_frame_not_color_converted(img_array_factory):
"""A depth/depth frame must be returned as-is, without BGR<->RGB conversion."""
color_frame = img_array_factory(height=8, width=8)
depth = img_array_factory(height=8, width=8, channels=1, dtype=np.uint16)[..., 0]
with patch(
"lerobot.cameras.reachy2_camera.reachy2_camera.CameraManager",
side_effect=lambda *a, **k: _make_cam_manager_mock(color_frame, depth_frame=depth),
):
cam = Reachy2Camera(Reachy2CameraConfig(name="depth", image_type="depth"))
cam.connect()
out = cam.read()
cam.disconnect()
np.testing.assert_array_equal(out, depth)
def test_wrong_camera_name():
with pytest.raises(ValueError):
_ = Reachy2CameraConfig(name="wrong-name", image_type="left")
+27 -1
View File
@@ -25,7 +25,7 @@ from unittest.mock import patch
import numpy as np
import pytest
from lerobot.cameras.configs import Cv2Rotation
from lerobot.cameras.configs import ColorMode, Cv2Rotation
from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
pytest.importorskip("pyrealsense2")
@@ -109,6 +109,32 @@ def test_read_depth():
assert isinstance(img, np.ndarray)
# These exercise _postprocess_image directly rather than read(): the bag playback returns
# non-deterministic frames we can't compare against, and the depth read() path is skipped
# (see test_read_depth) with the current pyrealsense2 version.
def test_color_mode_conversion(img_array_factory):
"""RGB (native for RealSense) is passed through; BGR reverses the channel axis."""
color = img_array_factory(height=3, width=4)
outputs = {}
for color_mode in (ColorMode.RGB, ColorMode.BGR):
camera = RealSenseCamera(RealSenseCameraConfig(serial_number_or_name="042", color_mode=color_mode))
camera.capture_height, camera.capture_width = color.shape[:2]
outputs[color_mode] = camera._postprocess_image(color)
np.testing.assert_array_equal(outputs[ColorMode.RGB], color)
np.testing.assert_array_equal(outputs[ColorMode.BGR], color[..., ::-1])
def test_depth_frame_not_color_converted(img_array_factory):
"""Depth frames must bypass color conversion, even when a BGR color_mode is set."""
camera = RealSenseCamera(RealSenseCameraConfig(serial_number_or_name="042", color_mode=ColorMode.BGR))
depth = img_array_factory(height=3, width=4, channels=1, dtype=np.uint16)[..., 0]
camera.capture_height, camera.capture_width = depth.shape
np.testing.assert_array_equal(camera._postprocess_image(depth, depth_frame=True), depth)
def test_read_before_connect():
config = RealSenseCameraConfig(serial_number_or_name="042")
camera = RealSenseCamera(config)
+7
View File
@@ -29,6 +29,13 @@ def test_message_recipe_validates_unknown_binding():
)
def test_canonical_recipe_loads():
"""The canonical PI052 blend YAML loads + validates."""
recipe = TrainingRecipe.from_yaml(Path("src/lerobot/configs/recipes/subtask_mem_vqa_speech.yaml"))
assert recipe.blend is not None
assert sum(c.weight for c in recipe.blend.values()) == pytest.approx(1.0)
def test_message_turn_requires_a_stream():
"""Every turn must declare a stream — None is rejected at construction.
+30 -1
View File
@@ -14,16 +14,21 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from types import SimpleNamespace
from unittest.mock import Mock
import pytest
import torch
from packaging.version import Version
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
from datasets import Dataset # noqa: E402
from huggingface_hub import DatasetCard
import lerobot.datasets.utils as dataset_utils
from lerobot.datasets.io_utils import hf_transform_to_torch
from lerobot.datasets.utils import create_lerobot_dataset_card
from lerobot.datasets.utils import create_lerobot_dataset_card, get_repo_versions, get_safe_version
from lerobot.utils.constants import ACTION, OBS_IMAGES
from lerobot.utils.feature_utils import combine_feature_dicts
@@ -57,6 +62,30 @@ def test_default_parameters():
]
@pytest.mark.parametrize("token", ["hf_test_token", True, False])
def test_get_repo_versions_forwards_token(monkeypatch, token):
api = Mock()
api.list_repo_refs.return_value = SimpleNamespace(
branches=[SimpleNamespace(name="v3.0")],
tags=[],
)
hf_api = Mock(return_value=api)
monkeypatch.setattr(dataset_utils, "HfApi", hf_api)
assert get_repo_versions("private/repo", token=token) == [Version("3.0")]
hf_api.assert_called_once_with(token=token)
api.list_repo_refs.assert_called_once_with("private/repo", repo_type="dataset")
@pytest.mark.parametrize("token", ["hf_test_token", True, False])
def test_get_safe_version_forwards_token(monkeypatch, token):
get_versions = Mock(return_value=[Version("3.0")])
monkeypatch.setattr(dataset_utils, "get_repo_versions", get_versions)
assert get_safe_version("private/repo", "v3.0", token=token) == "v3.0"
get_versions.assert_called_once_with("private/repo", token=token)
def test_with_tags():
tags = ["tag1", "tag2"]
card = create_lerobot_dataset_card(tags=tags)
+32
View File
@@ -204,6 +204,38 @@ def test_clear_resets_buffer(tmp_path):
assert dataset.writer.episode_buffer["size"] == 0
def test_clear_removes_video_frame_staging_dir(tmp_path):
"""clear_episode_buffer() removes PNG staging dirs for video features."""
video_key = "observation.images.cam"
features = {
video_key: {
"dtype": "video",
"shape": (64, 96, 3),
"names": ["height", "width", "channels"],
},
"action": {"dtype": "float32", "shape": (2,), "names": None},
}
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID,
fps=DEFAULT_FPS,
features=features,
root=tmp_path / "ds",
use_videos=True,
)
dataset.add_frame(_make_frame(features))
video_staging_dir = (
dataset.root
/ Path(DEFAULT_IMAGE_PATH.format(image_key=video_key, episode_index=0, frame_index=0)).parent
)
assert video_staging_dir.is_dir()
dataset.clear_episode_buffer()
assert dataset.writer.episode_buffer["size"] == 0
assert not video_staging_dir.exists()
def test_finalize_is_idempotent(tmp_path):
"""Calling finalize() twice does not raise."""
dataset = LeRobotDataset.create(
+46
View File
@@ -114,6 +114,20 @@ def test_dataset_initialization(tmp_path, lerobot_dataset_factory):
assert dataset.num_frames == len(dataset)
def test_dataset_slice(tmp_path, lerobot_dataset_factory):
dataset = lerobot_dataset_factory(
root=tmp_path / "test", total_episodes=3, total_frames=30, use_videos=False
)
assert len(dataset[:5]) == 5
assert len(dataset[::2]) == (len(dataset) + 1) // 2
assert [item["index"].item() for item in dataset[4::-1]] == [4, 3, 2, 1, 0]
assert [item["index"].item() for item in dataset[-3:]] == list(range(len(dataset) - 3, len(dataset)))
assert dataset[len(dataset) :] == []
assert isinstance(dataset[0], dict)
assert dataset[:1][0].keys() == dataset[0].keys()
# TODO(rcadene, aliberts): do not run LeRobotDataset.create, instead refactor LeRobotDatasetMetadata.create
# and test the small resulting function that validates the features
def test_dataset_feature_with_forward_slash_raises_error():
@@ -1741,6 +1755,38 @@ def test_delta_timestamps_query_returns_correct_values(tmp_path, empty_lerobot_d
assert is_pad == [True, False], f"Expected [True, False], got {is_pad}"
def test_dataset_slice_with_delta_timestamps(tmp_path, empty_lerobot_dataset_factory):
features = {
"observation.state": {"dtype": "float32", "shape": (1,), "names": ["x"]},
}
dataset = empty_lerobot_dataset_factory(
root=tmp_path / "test_slice_delta", features=features, use_videos=False, fps=10
)
for frame_idx in range(5):
dataset.add_frame(
{
"observation.state": torch.tensor([frame_idx], dtype=torch.float32),
"task": "task_0",
}
)
dataset.save_episode()
dataset.finalize()
sliced_dataset = LeRobotDataset(
dataset.repo_id,
root=dataset.root,
delta_timestamps={"observation.state": [-0.1, 0.0]},
tolerance_s=0.04,
)
items = sliced_dataset[:2]
assert items[0]["observation.state"].tolist() == [0.0, 0.0]
assert items[0]["observation.state_is_pad"].tolist() == [True, False]
assert items[1]["observation.state"].tolist() == [0.0, 1.0]
def test_episode_filter_filters_dataset(tmp_path, lerobot_dataset_factory):
"""episode_filter on LeRobotDataset narrows the loaded dataset to matching episodes."""
dataset = lerobot_dataset_factory(root=tmp_path / "test", total_episodes=8, total_frames=200)
+78
View File
@@ -343,6 +343,84 @@ def test_resolve_task_explicit_override_beats_rephrasings():
assert rendered["messages"][0]["content"] == "explicit override wins"
def test_flow_only_low_level_recipe_renders_without_target():
"""Regression: a flow-only ``low_level`` recipe has no ``target`` turn —
its supervision is the action-expert flow loss, not text-CE. It must
still render (not ``None``), otherwise every blend draw of it is dropped
and the action expert never receives a flow loss."""
recipe = TrainingRecipe(
messages=[
MessageTurn(
role="user",
content="${subtask}",
stream="low_level",
if_present="subtask",
),
],
bindings={"subtask": "active_at(t, style=subtask)"},
)
rendered = render_sample(
recipe=recipe,
persistent=PERSISTENT,
events=[],
t=0.5,
sample_idx=0,
task="clean kitchen",
)
assert rendered is not None
assert rendered["messages"] == [{"role": "user", "content": "subtask 0"}]
assert rendered["message_streams"] == ["low_level"]
assert rendered["target_message_indices"] == []
def test_vqa_frame_is_consumed_over_the_weighted_blend():
"""A frame carrying a VQA annotation renders the ``ask_vqa*`` sub-recipe
even when its blend weight is tiny VQA annotations are sparse and must
never be wasted on a subtask/action draw."""
recipe = TrainingRecipe(
blend={
"high_level_subtask": TrainingRecipe(
weight=0.99,
messages=[
MessageTurn(role="user", content="${task}", stream="high_level"),
MessageTurn(role="assistant", content="a subtask", stream="high_level", target=True),
],
),
"ask_vqa_top": TrainingRecipe(
weight=0.01,
bindings={
"vqa_query": "emitted_at(t, style=vqa, role=user, camera=observation.images.top)",
"vqa": "emitted_at(t, style=vqa, role=assistant, camera=observation.images.top)",
},
messages=[
MessageTurn(
role="user", content="${vqa_query}", stream="high_level", if_present="vqa_query"
),
MessageTurn(
role="assistant",
content="${vqa}",
stream="high_level",
target=True,
if_present="vqa",
),
],
),
}
)
# A frame WITH a vqa event renders VQA on every sample_idx, despite the
# ask_vqa weight being only 0.01.
for sample_idx in range(20):
rendered = render_sample(
recipe=recipe, persistent=PERSISTENT, events=EVENTS_AT_1, t=1.0, sample_idx=sample_idx, task="x"
)
assert rendered["messages"][-1]["content"] == '{"count": 2}', sample_idx
# A frame WITHOUT a vqa event falls back to the normal weighted blend.
rendered = render_sample(recipe=recipe, persistent=PERSISTENT, events=[], t=1.0, sample_idx=0, task="x")
assert rendered["messages"][-1]["content"] == "a subtask"
def test_emitted_at_persistent_tolerates_small_timestamp_drift():
"""Persistent ``emitted_at`` should match within EMITTED_AT_TOLERANCE_S
so callers that derive ``t`` arithmetically (``frame_idx / fps``) still
+43
View File
@@ -20,6 +20,7 @@ property delegation, and the full create-record-finalize-read lifecycle.
"""
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import Mock
import pytest
@@ -191,6 +192,48 @@ def test_metadata_without_root_uses_hub_cache_snapshot_download(
}
@pytest.mark.parametrize("token", ["hf_test_token", True, False])
def test_metadata_download_forwards_token(tmp_path, monkeypatch, token):
snapshot_root = tmp_path / "snapshot"
snapshot_download = Mock(return_value=str(snapshot_root))
get_safe_version = Mock(return_value="v3.0")
load_metadata = Mock(side_effect=[FileNotFoundError, None])
monkeypatch.setattr(dataset_metadata_module, "snapshot_download", snapshot_download)
monkeypatch.setattr(dataset_metadata_module, "get_safe_version", get_safe_version)
monkeypatch.setattr(LeRobotDatasetMetadata, "_load_metadata", load_metadata)
meta = LeRobotDatasetMetadata(
repo_id=DUMMY_REPO_ID,
revision="v3.0",
token=token,
)
assert meta.root == snapshot_root
assert not hasattr(meta, "_token")
get_safe_version.assert_called_once_with(DUMMY_REPO_ID, "v3.0", token=token)
assert snapshot_download.call_args.kwargs["token"] is token
@pytest.mark.parametrize("token", ["hf_test_token", True, False])
def test_data_download_forwards_token(tmp_path, monkeypatch, token):
snapshot_root = tmp_path / "snapshot"
snapshot_download = Mock(return_value=str(snapshot_root))
monkeypatch.setattr(lerobot_dataset_module, "snapshot_download", snapshot_download)
dataset = LeRobotDataset.__new__(LeRobotDataset)
dataset.repo_id = DUMMY_REPO_ID
dataset.revision = "main"
dataset.episodes = None
dataset._requested_root = None
dataset.meta = SimpleNamespace(root=None)
dataset.reader = SimpleNamespace(root=None)
dataset._download(token=token)
assert dataset.root == snapshot_root
assert snapshot_download.call_args.kwargs["token"] is token
def test_without_root_reads_different_revisions_from_distinct_snapshot_roots(
tmp_path,
info_factory,
+1 -3
View File
@@ -25,7 +25,7 @@ from datasets import Dataset # noqa: E402
from lerobot.datasets.io_utils import (
hf_transform_to_torch,
)
from lerobot.datasets.sampler import EpisodeAwareSampler
from lerobot.datasets.sampler import EpisodeAwareSampler, compute_sampler_state
def calculate_episode_data_index(hf_dataset: Dataset) -> dict[str, torch.Tensor]:
@@ -154,8 +154,6 @@ def test_partial_episode_drop_warns(caplog):
# --- seeded (seed, epoch) shuffling, resume, and state ---
from lerobot.datasets.sampler import compute_sampler_state # noqa: E402
EPISODE_BOUNDS = ([0, 2, 3], [2, 3, 6]) # episodes of 2, 1 and 3 frames
+38
View File
@@ -13,12 +13,16 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from types import SimpleNamespace
from unittest.mock import Mock
import numpy as np
import pytest
import torch
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
import lerobot.datasets.streaming_dataset as streaming_dataset_module
from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset
from lerobot.datasets.utils import safe_shard
from lerobot.utils.constants import ACTION
@@ -71,6 +75,40 @@ def get_frames_expected_order(streaming_ds: StreamingLeRobotDataset) -> list[int
return expected_indices
@pytest.mark.parametrize("token", ["hf_test_token", True, False])
@pytest.mark.parametrize("from_local", [False, True])
def test_streaming_dataset_forwards_hub_token_only_for_remote_data(tmp_path, monkeypatch, token, from_local):
requested_root = tmp_path / "local" if from_local else None
metadata = SimpleNamespace(
root=requested_root or tmp_path / "snapshot",
revision=streaming_dataset_module.CODEBASE_VERSION,
_version=streaming_dataset_module.CODEBASE_VERSION,
features={},
depth_keys=[],
image_keys=[],
rescale_depth_stats=Mock(),
)
metadata_cls = Mock(return_value=metadata)
load_dataset = Mock(return_value=SimpleNamespace(num_shards=1))
monkeypatch.setattr(streaming_dataset_module, "LeRobotDatasetMetadata", metadata_cls)
monkeypatch.setattr(streaming_dataset_module, "load_dataset", load_dataset)
dataset = StreamingLeRobotDataset(DUMMY_REPO_ID, root=requested_root, token=token)
metadata_cls.assert_called_once_with(
DUMMY_REPO_ID,
requested_root,
streaming_dataset_module.CODEBASE_VERSION,
force_cache_sync=False,
token=token,
)
if from_local:
assert "token" not in load_dataset.call_args.kwargs
else:
assert load_dataset.call_args.kwargs["token"] is token
assert not hasattr(dataset, "_token")
def test_single_frame_consistency(tmp_path, lerobot_dataset_factory):
"""Test if are correctly accessed"""
ds_num_frames = 400
+11
View File
@@ -35,6 +35,17 @@ def test_unknown_type():
make_env_config("nonexistent")
def test_libero_fps_controls_simulator_frequency():
cfg = LiberoEnv(fps=17)
assert cfg.gym_kwargs["control_freq"] == 17
def test_libero_rejects_nonpositive_fps():
with pytest.raises(ValueError, match="fps must be positive"):
LiberoEnv(fps=0)
def test_identity_processors():
"""Base class get_env_processors() returns identity pipelines."""
cfg = make_env_config("aloha")
+48
View File
@@ -0,0 +1,48 @@
#!/usr/bin/env python
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import numpy as np
from lerobot.envs.robocasa import RoboCasaEnv, convert_action
def test_robocasa_action_uses_openpi_checkpoint_order():
action = np.arange(12, dtype=np.float32)
converted = convert_action(action)
np.testing.assert_array_equal(converted["action.end_effector_position"], [0, 1, 2])
np.testing.assert_array_equal(converted["action.end_effector_rotation"], [3, 4, 5])
np.testing.assert_array_equal(converted["action.gripper_close"], [6])
np.testing.assert_array_equal(converted["action.base_motion"], [7, 8, 9, 10])
np.testing.assert_array_equal(converted["action.control_mode"], [11])
def test_robocasa_state_uses_openpi_checkpoint_order():
env = object.__new__(RoboCasaEnv)
env.obs_type = "pixels_agent_pos"
env.camera_name = []
raw_observation = {
"state.end_effector_position_relative": np.arange(0, 3),
"state.end_effector_rotation_relative": np.arange(3, 7),
"state.base_position": np.arange(7, 10),
"state.base_rotation": np.arange(10, 14),
"state.gripper_qpos": np.arange(14, 16),
}
observation = env._format_raw_obs(raw_observation)
np.testing.assert_array_equal(observation["agent_pos"], np.arange(16, dtype=np.float32))
+245
View File
@@ -0,0 +1,245 @@
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import shlex
import sys
from unittest.mock import MagicMock
import draccus
import pytest
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
from lerobot.annotations.steerable_pipeline.config import (
DEFAULT_ANNOTATE_JOB_IMAGE,
AnnotationJobConfig,
AnnotationPipelineConfig,
)
from lerobot.jobs.annotate import build_pod_command, build_pod_setup, submit_annotate_to_hf
def _parse(*args):
return draccus.parse(AnnotationPipelineConfig, args=list(args))
def _set_argv(monkeypatch, *args):
monkeypatch.setattr(sys, "argv", ["lerobot-annotate", *args])
# --- config ----------------------------------------------------------------
def test_annotation_job_defaults_are_local_with_vllm_image():
cfg = AnnotationJobConfig()
assert cfg.target is None
assert cfg.is_remote is False
assert cfg.image == DEFAULT_ANNOTATE_JOB_IMAGE
assert cfg.timeout == "2h"
assert cfg.lerobot_ref == "main"
def test_annotation_config_parses_job_target():
cfg = _parse("--repo_id", "u/d", "--job.target", "h200")
assert cfg.job.target == "h200"
assert cfg.job.is_remote is True
def test_annotation_config_defaults_to_local():
assert _parse("--repo_id", "u/d").job.is_remote is False
# --- pod command -----------------------------------------------------------
def test_pod_setup_installs_requested_ref():
setup = build_pod_setup("my-branch")
assert "git+https://github.com/huggingface/lerobot.git@my-branch" in setup
# The vLLM image has neither ffmpeg (video decode) nor lerobot's pinned deps.
assert "ffmpeg" in setup
assert "'draccus==0.10.0'" in setup
def _annotate_argv(command):
"""Extract the `lerobot-annotate ...` argv from a `bash -c` pod command."""
assert command[:2] == ["bash", "-c"]
_setup, _, annotate = command[2].rpartition(" && ")
return shlex.split(annotate)
def test_pod_command_forwards_user_flags_and_pins_local_target():
command = build_pod_command(
"u/d",
"main",
["--repo_id=u/d", "--new_repo_id=u/d_annotated", "--push_to_hub=true", "--job.target=h200"],
)
argv = _annotate_argv(command)
assert argv[0] == "lerobot-annotate"
# --job.* is client-side orchestration; the pod must not re-dispatch itself.
assert not any(a.startswith("--job.") for a in argv[1:-1])
assert argv[-1] == "--job.target=local"
assert "--new_repo_id=u/d_annotated" in argv
assert "--push_to_hub=true" in argv
def test_pod_command_replaces_host_local_root_with_repo_id():
"""--root points at a directory only the client has; the pod resolves by repo_id."""
command = build_pod_command("u/d", "main", ["--root", "/home/me/datasets/d", "--seed=7"])
argv = _annotate_argv(command)
assert "--root" not in argv
assert "/home/me/datasets/d" not in argv
assert argv.count("--repo_id=u/d") == 1
assert "--seed=7" in argv
def test_pod_command_does_not_duplicate_repo_id():
command = build_pod_command("u/d", "main", ["--repo_id", "u/d"])
assert _annotate_argv(command).count("--repo_id=u/d") == 1
def test_pod_command_quotes_flags_containing_spaces_and_json():
"""serve_command and chat_template_kwargs must survive the trip through `bash -c`."""
serve = "--vlm.serve_command=vllm serve Qwen/Qwen3.6-27B --max-model-len 32768 --port {port}"
kwargs = '--vlm.chat_template_kwargs={"enable_thinking": false}'
command = build_pod_command("u/d", "main", [serve, kwargs])
argv = _annotate_argv(command)
assert serve in argv
assert kwargs in argv
# --- submission ------------------------------------------------------------
def test_submit_requires_login(monkeypatch):
monkeypatch.setattr("lerobot.jobs.annotate.get_token", lambda: None)
with pytest.raises(RuntimeError, match="hf auth login"):
submit_annotate_to_hf(_parse("--repo_id", "u/d", "--job.target", "h200"))
def test_submit_requires_repo_id(monkeypatch):
"""A remote run over --root alone can't work: the pod can't see the client's disk."""
monkeypatch.setattr("lerobot.jobs.annotate.get_token", lambda: "tok")
cfg = _parse("--root", "/tmp/d", "--job.target", "h200")
with pytest.raises(ValueError, match="--repo_id"):
submit_annotate_to_hf(cfg)
@pytest.mark.parametrize("arg", ["--config_path=annotate.yaml", "--vlm=vlm.yaml", "--job=job.yaml"])
def test_submit_rejects_local_config_files(monkeypatch, arg):
"""draccus takes a config file for the whole config and for each nested one; the
pod can read none of them, so a remote run must refuse rather than drop them."""
monkeypatch.setattr("lerobot.jobs.annotate.get_token", lambda: "tok")
_set_argv(monkeypatch, arg, "--job.target=h200")
cfg = _parse("--repo_id", "u/d", "--job.target", "h200")
with pytest.raises(ValueError, match="cannot read config files"):
submit_annotate_to_hf(cfg)
def test_pod_command_drops_bare_job_config_file_arg():
"""`--job` isn't caught by the `--job.` prefix, and could carry a remote target
that would make the pod submit a job of its own recursively."""
argv = _annotate_argv(build_pod_command("u/d", "main", ["--job", "job.yaml", "--seed=7"]))
assert "--job" not in argv
assert "job.yaml" not in argv
assert argv[-1] == "--job.target=local"
def test_submit_dispatches_job(monkeypatch):
monkeypatch.setattr("lerobot.jobs.annotate.get_token", lambda: "tok")
monkeypatch.setattr("lerobot.jobs.annotate.HfApi", lambda token=None: MagicMock())
monkeypatch.setattr("lerobot.jobs.annotate.ensure_dataset_available", lambda *a, **kw: None)
run_job_calls = []
def fake_run_job(**kwargs):
run_job_calls.append(kwargs)
return MagicMock(id="job-123")
monkeypatch.setattr("lerobot.jobs.annotate.run_job", fake_run_job)
_set_argv(monkeypatch, "--repo_id=u/d", "--push_to_hub=true", "--job.target=h200", "--job.detach=true")
cfg = _parse("--repo_id", "u/d", "--push_to_hub", "true", "--job.target", "h200", "--job.detach", "true")
submit_annotate_to_hf(cfg)
assert len(run_job_calls) == 1
call = run_job_calls[0]
assert call["flavor"] == "h200"
assert call["image"] == DEFAULT_ANNOTATE_JOB_IMAGE
assert call["timeout"] == "2h"
# The Hub token is forwarded so the pod can pull a private dataset and push the result.
assert call["secrets"]["HF_TOKEN"] == "tok"
assert call["labels"].get("lerobot") == "true"
argv = _annotate_argv(call["command"])
assert argv[0] == "lerobot-annotate"
assert "--push_to_hub=true" in argv
@pytest.mark.timeout(15)
def test_submit_follows_job_to_completion(monkeypatch, capsys):
"""Non-detach path must stream logs and RETURN (not hang) once the job is terminal.
Exercises the `follow_job` helper shared with the training submitter from the
annotation side, which is why the job-state patches target `lerobot.jobs.hf`.
Asserting on the completion message and not merely on "didn't hang" is what makes
this fail if `follow_job` ever reports detached-without-a-verdict instead.
"""
monkeypatch.setattr("lerobot.jobs.annotate.get_token", lambda: "tok")
monkeypatch.setattr("lerobot.jobs.annotate.HfApi", lambda token=None: MagicMock())
monkeypatch.setattr("lerobot.jobs.annotate.ensure_dataset_available", lambda *a, **kw: None)
monkeypatch.setattr("lerobot.jobs.annotate.run_job", lambda **kw: MagicMock(id="job-1", url="http://x"))
monkeypatch.setattr(
"lerobot.jobs.hf.inspect_job",
lambda job_id: MagicMock(status=MagicMock(stage=MagicMock(value="COMPLETED"), message=None)),
)
monkeypatch.setattr("lerobot.jobs.hf.fetch_job_logs", lambda job_id, follow=True: iter(()))
_set_argv(monkeypatch, "--repo_id=u/d", "--job.target=h200")
submit_annotate_to_hf(_parse("--repo_id", "u/d", "--push_to_hub", "true", "--job.target", "h200"))
assert "Annotation complete" in capsys.readouterr().out
@pytest.mark.timeout(15)
def test_submit_raises_when_job_fails(monkeypatch):
"""A job that ends in a non-COMPLETED stage must surface as an error, not a silent return."""
monkeypatch.setattr("lerobot.jobs.annotate.get_token", lambda: "tok")
monkeypatch.setattr("lerobot.jobs.annotate.HfApi", lambda token=None: MagicMock())
monkeypatch.setattr("lerobot.jobs.annotate.ensure_dataset_available", lambda *a, **kw: None)
monkeypatch.setattr("lerobot.jobs.annotate.run_job", lambda **kw: MagicMock(id="job-1", url=None))
monkeypatch.setattr(
"lerobot.jobs.hf.inspect_job",
lambda job_id: MagicMock(status=MagicMock(stage=MagicMock(value="ERROR"), message="Job timeout")),
)
monkeypatch.setattr("lerobot.jobs.hf.fetch_job_logs", lambda job_id, follow=True: iter(()))
_set_argv(monkeypatch, "--repo_id=u/d", "--job.target=h200")
with pytest.raises(RuntimeError, match="stage=ERROR .Job timeout."):
submit_annotate_to_hf(_parse("--repo_id", "u/d", "--job.target", "h200"))
def test_submit_ensures_dataset_is_on_the_hub(monkeypatch):
"""A local-only dataset is pushed (privately) before the job can reach it by repo_id."""
monkeypatch.setattr("lerobot.jobs.annotate.get_token", lambda: "tok")
monkeypatch.setattr("lerobot.jobs.annotate.HfApi", lambda token=None: MagicMock())
monkeypatch.setattr("lerobot.jobs.annotate.run_job", lambda **kw: MagicMock(id="job-1"))
seen = []
monkeypatch.setattr(
"lerobot.jobs.annotate.ensure_dataset_available",
lambda repo_id, *, api, tags=None: seen.append((repo_id, tags)),
)
_set_argv(monkeypatch, "--repo_id=u/d", "--job.target=h200", "--job.detach=true")
submit_annotate_to_hf(
_parse("--repo_id", "u/d", "--job.target", "h200", "--job.detach", "true", "--job.tags", '["lelab"]')
)
assert seen == [("u/d", ["lerobot", "lelab"])]
+14
View File
@@ -29,12 +29,26 @@ from lerobot.jobs.hf import (
_poll_until_done,
build_remote_config_file,
build_repo_id,
follow_job,
resolve_job_tags,
resolve_wandb_api_key,
submit_to_hf,
)
def test_follow_job_detach_returns_without_watching(monkeypatch):
"""`detach` must short-circuit before any polling or log streaming starts."""
def _boom(*a, **kw):
raise AssertionError("detach must not touch the job")
monkeypatch.setattr("lerobot.jobs.hf.inspect_job", _boom)
monkeypatch.setattr("lerobot.jobs.hf.fetch_job_logs", _boom)
# False = "stopped watching without a verdict", so callers stay quiet rather than
# claiming success for a job that is still running.
assert follow_job("job-1", detach=True) is False
def test_resolve_job_tags_always_includes_lerobot_and_dedups():
assert resolve_job_tags(None) == ["lerobot"]
assert resolve_job_tags([]) == ["lerobot"]
+9 -3
View File
@@ -405,12 +405,18 @@ def test_record_ranges_of_motion(mock_motors, dummy_motors):
read_pos_stub = mock_motors.build_sequential_sync_read_stub(
*X_SERIES_CONTROL_TABLE["Present_Position"], positions
)
with patch("lerobot.motors.motors_bus.enter_pressed", side_effect=[False, True]):
bus = DynamixelMotorsBus(port=mock_motors.port, motors=dummy_motors)
bus.connect(handshake=False)
bus = DynamixelMotorsBus(port=mock_motors.port, motors=dummy_motors)
bus.connect(handshake=False)
with (
patch("lerobot.motors.motors_bus.enter_pressed", side_effect=[False, True]),
patch("lerobot.motors.motors_bus.time.sleep") as mock_sleep,
patch.object(bus, "sync_read", wraps=bus.sync_read) as mock_sync_read,
):
mins, maxes = bus.record_ranges_of_motion(display_values=False)
assert mock_motors.stubs[read_pos_stub].calls == 3
assert all(call.kwargs["num_retry"] == 5 for call in mock_sync_read.call_args_list)
mock_sleep.assert_called_once_with(0.02)
assert mins == expected_mins
assert maxes == expected_maxes
+9 -3
View File
@@ -509,12 +509,18 @@ def test_record_ranges_of_motion(mock_motors, dummy_motors):
stub = mock_motors.build_sequential_sync_read_stub(
*STS_SMS_SERIES_CONTROL_TABLE["Present_Position"], positions
)
with patch("lerobot.motors.motors_bus.enter_pressed", side_effect=[False, True]):
bus = FeetechMotorsBus(port=mock_motors.port, motors=dummy_motors)
bus.connect(handshake=False)
bus = FeetechMotorsBus(port=mock_motors.port, motors=dummy_motors)
bus.connect(handshake=False)
with (
patch("lerobot.motors.motors_bus.enter_pressed", side_effect=[False, True]),
patch("lerobot.motors.motors_bus.time.sleep") as mock_sleep,
patch.object(bus, "sync_read", wraps=bus.sync_read) as mock_sync_read,
):
mins, maxes = bus.record_ranges_of_motion(display_values=False)
assert mock_motors.stubs[stub].calls == 3
assert all(call.kwargs["num_retry"] == 5 for call in mock_sync_read.call_args_list)
mock_sleep.assert_called_once_with(0.02)
assert mins == expected_mins
assert maxes == expected_maxes
@@ -26,8 +26,17 @@ import tempfile
from pathlib import Path
import pytest
import torch
from safetensors.torch import save_file
from lerobot.processor.pipeline import DataProcessorPipeline, ProcessorMigrationError
from lerobot.configs import PipelineFeatureType, PolicyFeature
from lerobot.processor.pipeline import (
DataProcessorPipeline,
ProcessorMigrationError,
ProcessorStep,
ProcessorStepRegistry,
)
from lerobot.types import EnvTransition
# Simplified Config Loading Tests
@@ -98,6 +107,140 @@ def test_load_config_nonexistent_path_tries_hub():
DataProcessorPipeline._load_config("nonexistent/path", "processor.json", {})
def test_from_pretrained_local_directory_missing_state_does_not_call_hub(monkeypatch):
"""Local processor dirs must fail locally when a state file is missing."""
@ProcessorStepRegistry.register("local_missing_state_step")
class LocalMissingStateStep(ProcessorStep):
def __call__(self, transition: EnvTransition) -> EnvTransition:
return transition
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
return features
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
pass
try:
with tempfile.TemporaryDirectory() as tmp_dir:
tmp_path = Path(tmp_dir)
config = {
"name": "LocalMissingStatePipeline",
"steps": [{"registry_name": "local_missing_state_step", "state_file": "missing.safetensors"}],
}
(tmp_path / "processor.json").write_text(json.dumps(config))
def fail_hub_download(*args, **kwargs):
pytest.fail("local missing processor state should not call hf_hub_download")
monkeypatch.setattr("lerobot.processor.pipeline.hf_hub_download", fail_hub_download)
with pytest.raises(FileNotFoundError, match="missing.safetensors.*local processor pipeline"):
DataProcessorPipeline.from_pretrained(tmp_path, config_filename="processor.json")
finally:
ProcessorStepRegistry.unregister("local_missing_state_step")
def test_from_pretrained_local_config_file_missing_state_does_not_call_hub(monkeypatch):
"""Local single-file processor configs must also keep missing state resolution local."""
@ProcessorStepRegistry.register("local_file_missing_state_step")
class LocalFileMissingStateStep(ProcessorStep):
def __call__(self, transition: EnvTransition) -> EnvTransition:
return transition
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
return features
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
pass
try:
with tempfile.TemporaryDirectory() as tmp_dir:
tmp_path = Path(tmp_dir)
config_path = tmp_path / "processor.json"
config = {
"name": "LocalFileMissingStatePipeline",
"steps": [
{"registry_name": "local_file_missing_state_step", "state_file": "missing.safetensors"}
],
}
config_path.write_text(json.dumps(config))
def fail_hub_download(*args, **kwargs):
pytest.fail("local missing processor state should not call hf_hub_download")
monkeypatch.setattr("lerobot.processor.pipeline.hf_hub_download", fail_hub_download)
with pytest.raises(FileNotFoundError, match="missing.safetensors.*local processor pipeline"):
DataProcessorPipeline.from_pretrained(config_path, config_filename="ignored.json")
finally:
ProcessorStepRegistry.unregister("local_file_missing_state_step")
def test_from_pretrained_hub_source_missing_local_state_still_calls_hub(monkeypatch, tmp_path):
"""Hub sources still fall back to hf_hub_download for state files."""
@ProcessorStepRegistry.register("hub_state_step")
class HubStateStep(ProcessorStep):
def __init__(self):
self.value = torch.tensor(0)
def __call__(self, transition: EnvTransition) -> EnvTransition:
return transition
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
return features
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
self.value = state["value"]
try:
state_path = tmp_path / "downloaded.safetensors"
save_file({"value": torch.tensor(7)}, state_path)
loaded_config = {
"name": "HubStatePipeline",
"steps": [{"registry_name": "hub_state_step", "state_file": "hub_state.safetensors"}],
}
calls = []
def fake_load_config(cls, model_id, config_filename, hub_download_kwargs):
return loaded_config, tmp_path / "hub_cache"
def fake_hub_download(**kwargs):
calls.append(kwargs)
return str(state_path)
monkeypatch.setattr(DataProcessorPipeline, "_load_config", classmethod(fake_load_config))
monkeypatch.setattr("lerobot.processor.pipeline.hf_hub_download", fake_hub_download)
pipeline = DataProcessorPipeline.from_pretrained("user/repo", config_filename="processor.json")
assert calls == [
{
"repo_id": "user/repo",
"filename": "hub_state.safetensors",
"repo_type": "model",
"force_download": False,
"resume_download": None,
"proxies": None,
"token": None,
"cache_dir": None,
"local_files_only": False,
"revision": None,
}
]
assert pipeline.steps[0].value.item() == 7
finally:
ProcessorStepRegistry.unregister("hub_state_step")
# Config Validation Tests
@@ -12,7 +12,9 @@ from lerobot.processor.render_messages_processor import RenderMessagesStep # no
from lerobot.types import TransitionKey # noqa: E402
def test_render_messages_step_noops_without_language_columns():
def test_render_messages_step_renders_task_fallback_without_language_columns():
"""No language columns + a task string → low-level task fallback render,
matching what the policy sees at eval time on unannotated observations."""
recipe = TrainingRecipe(
messages=[
MessageTurn(role="user", content="${task}", stream="high_level"),
@@ -21,6 +23,24 @@ def test_render_messages_step_noops_without_language_columns():
)
transition = create_transition(complementary_data={"task": "do it"})
out = RenderMessagesStep(recipe)(transition)
data = out[TransitionKey.COMPLEMENTARY_DATA]
assert data["messages"] == [{"role": "user", "content": "do it"}]
assert data["message_streams"] == ["low_level"]
assert data["target_message_indices"] == []
assert data["task"] == "do it"
def test_render_messages_step_noops_without_language_columns_or_task():
recipe = TrainingRecipe(
messages=[
MessageTurn(role="user", content="${task}", stream="high_level"),
MessageTurn(role="assistant", content="${subtask}", stream="low_level", target=True),
]
)
transition = create_transition(complementary_data={})
assert RenderMessagesStep(recipe)(transition) == transition
@@ -58,3 +78,70 @@ def test_render_messages_step_renders_and_drops_raw_language():
assert data["messages"][-1]["content"] == "reach carefully"
assert data["message_streams"] == ["high_level", "low_level"]
assert data["target_message_indices"] == [1]
def test_render_messages_step_falls_back_to_low_level_task_when_recipe_misses():
recipe = TrainingRecipe(
messages=[
MessageTurn(
role="assistant",
content="${subtask}",
stream="high_level",
target=True,
if_present="subtask",
),
]
)
transition = create_transition(
complementary_data={
"task": "pick the cube",
"timestamp": torch.tensor(0.0),
"index": torch.tensor(7),
"language_persistent": [],
"language_events": [{"style": "unmatched", "timestamp": 0.0}],
}
)
out = RenderMessagesStep(recipe)(transition)
data = out[TransitionKey.COMPLEMENTARY_DATA]
assert data["messages"] == [{"role": "user", "content": "pick the cube"}]
assert data["message_streams"] == ["low_level"]
assert data["target_message_indices"] == []
def test_render_messages_step_falls_back_per_sample_in_batched_language():
recipe = TrainingRecipe(
messages=[
MessageTurn(
role="assistant",
content="${subtask}",
stream="high_level",
target=True,
if_present="subtask",
),
]
)
transition = create_transition(
action=torch.arange(4).reshape(2, 2),
complementary_data={
"task": ["pick the cube", "open the drawer"],
"timestamp": torch.tensor([0.0, 1.0]),
"index": torch.tensor([7, 8]),
"language_persistent": [[], []],
"language_events": [
[{"style": "unmatched", "timestamp": 0.0}],
[{"style": "unmatched", "timestamp": 1.0}],
],
},
)
out = RenderMessagesStep(recipe)(transition)
data = out[TransitionKey.COMPLEMENTARY_DATA]
assert data["messages"] == [
[{"role": "user", "content": "pick the cube"}],
[{"role": "user", "content": "open the drawer"}],
]
assert data["message_streams"] == [["low_level"], ["low_level"]]
assert data["target_message_indices"] == [[], []]
+41 -1
View File
@@ -25,7 +25,7 @@ import pytest
import torch
from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
from lerobot.processor import DataProcessorPipeline, TokenizerProcessorStep
from lerobot.processor import ActionTokenizerProcessorStep, DataProcessorPipeline, TokenizerProcessorStep
from lerobot.processor.converters import create_transition, identity_transition
from lerobot.types import TransitionKey
from lerobot.utils.constants import (
@@ -88,6 +88,46 @@ class MockTokenizer:
return result
def test_action_tokenizer_config_preserves_token_mapping():
processor = object.__new__(ActionTokenizerProcessorStep)
processor.trust_remote_code = True
processor.max_action_tokens = 384
processor.fast_skip_tokens = 64
processor.paligemma_tokenizer_name = "custom/paligemma"
processor.allow_truncation = False
processor.action_tokenizer_name = "custom/fast"
processor.action_tokenizer_input_object = None
assert processor.get_config() == {
"trust_remote_code": True,
"max_action_tokens": 384,
"fast_skip_tokens": 64,
"paligemma_tokenizer_name": "custom/paligemma",
"allow_truncation": False,
"action_tokenizer_name": "custom/fast",
}
def test_action_tokenizer_can_reject_truncated_sequences():
processor = object.__new__(ActionTokenizerProcessorStep)
processor.max_action_tokens = 4
processor.fast_skip_tokens = 128
processor.allow_truncation = False
processor.action_tokenizer = lambda _actions: [1, 2, 3]
processor._paligemma_tokenizer = type(
"Tokenizer",
(),
{
"vocab_size": 1000,
"bos_token_id": 2,
"encode": lambda _self, text, **_kwargs: [10, 11] if text == "Action: " else [12, 1],
},
)()
with pytest.raises(ValueError, match="max_action_tokens=4"):
processor._tokenize_action(torch.zeros(1, 2, 1))
@pytest.fixture
def mock_tokenizer():
"""Provide a mock tokenizer for testing."""
+19
View File
@@ -109,3 +109,22 @@ def test_send_action(follower):
goal_pos = {m: (i + 1) * 10 for i, m in enumerate(follower.bus.motors)}
follower.bus.sync_write.assert_called_once_with("Goal_Position", goal_pos)
def test_configure_writes_position_pid_coefficients():
bus_mock = _make_bus_mock()
bus_mock.motors = ["shoulder_pan"]
robot = MagicMock()
robot.bus = bus_mock
robot.config = SO100FollowerConfig(
port="/dev/null",
position_p_coefficient=32,
position_i_coefficient=1,
position_d_coefficient=16,
)
SO100Follower.configure(robot)
bus_mock.write.assert_any_call("P_Coefficient", "shoulder_pan", 32)
bus_mock.write.assert_any_call("I_Coefficient", "shoulder_pan", 1)
bus_mock.write.assert_any_call("D_Coefficient", "shoulder_pan", 16)

Some files were not shown because too many files have changed in this diff Show More