Compare commits

..

2 Commits

Author SHA1 Message Date
Nikodem Bartnik 68d6335d5b fix old docs and comments 2026-07-28 14:31:30 +02:00
Nikodem Bartnik 258d521a89 add docs audit 2026-07-28 14:31:07 +02:00
37 changed files with 211 additions and 862 deletions
+113
View File
@@ -0,0 +1,113 @@
# LeRobot Docs Audit
**Status:** Fresh baseline for a new docs redesign effort · **Date:** 2026-07-27
**Supersedes:** the June 2026 `DOCS_REDESIGN.md` proposal (lived on the unmerged, now-stale `docs/complete-docs-redesign` branch — abandoned; that branch is ~67k lines behind `main` and should not be resurrected). This document re-audits `docs/source/` from scratch against current `main`, since the docs tree grew substantially in the ~6 weeks between the two audits (many new policies, robots, and benchmarks landed).
---
## Snapshot (current state, in numbers)
- **IA**: 15 top-level `_toctree.yml` sections, 81 `local:` entries, 0 broken toctree links.
- **Files**: 97 total files in `docs/source/` = 80 real `.mdx` pages + 16 orphaned `policy_*_README.md` stubs (up from 14 in June) + 1 `contributing.md` (a genuine symlink to root `CONTRIBUTING.md`).
- **Growth**: `.mdx` pages grew 49 → 80 in 6 months (+63%, ~5.2 pages/month), arriving in **bursts against a shared, hand-edited `_toctree.yml`** — e.g. 6 benchmark pages landed via 6 separate PRs within 48 hours (Apr 2026); 3 policy pages within 3 days (Jul 2026).
- **CLI docs coverage improved sharply since June**: 17 of 18 `lerobot-*` commands now documented (only `lerobot-info`, a diagnostics helper, is not) — the June baseline was 6/17 undocumented.
- **Decision content**: zero comparison/decision pages exist in `docs/source/` for Policies (15 options), Robots (11), Benchmarks (10), or Teleoperators — this equivalent content already exists, fully written, in root `AGENT_GUIDE.md` §6, but was never ported or linked in.
- **Teleoperators**: dedicated toctree section has only 2 pages vs. 16-17 registered teleoperator types in `src/`, but all types do get at least one doc mention somewhere (usually buried inside a robot's hardware page) — a discoverability gap, not an absence of content.
- **Orphan policy READMEs**: of the 16, 3 (Diffusion, TD-MPC, VQ-BeT) have **zero real published documentation anywhere** — the orphan stub is their only artifact. 13 duplicate an actively-maintained `.mdx` page under a different name. One (MolmoAct2) is proven to have silently rotted: a 39-line pre-refactor stub sits next to a 495-line maintained page that diverged 481 lines ago.
- **Root `README.md` itself has drifted**: its "Supported Hardware" line omits SO-101 (the flagship robot) entirely; its "SoTA Models" table links several policies to thin orphan stubs instead of the richer tutorial pages that exist for them.
- **4 concrete broken/stale copy-paste commands** verified in published pages (`act.mdx`, `il_robots.mdx` ×2 spots, `lerobot-dataset-v3.mdx`, `openarm.mdx` — the last added as recently as Jan 2026, so a fresh miss, not just aging drift).
- **Process/CI**: no CODEOWNERS, no `_redirects.yml`, no doc-consistency or toctree-completeness check anywhere; the docs CI workflow triggers only on `paths: docs/**`, so code-only PRs get zero automated docs signal (empirically responsible for the one real lag found — OMX's docs page, 41 days after its code). Despite this, 18 of 19 spot-checked new integrations (EVO1, FastWAM, LingBot-VA, MolmoAct2, RTC, VLA-JEPA, X-VLA, WallOSS, Damiao, Hope Jr, OpenArm, reBot B601, Reachy 2, Unitree G1, RoboCasa365, VLABench, IsaacLab Arena, etc.) shipped docs in the *same commit* as the feature — organic discipline is currently strong.
- **No content template** exists for policy or robot doc pages: heading sets/length vary 80-528 lines with almost no shared vocabulary; by contrast, benchmarks *do* have a written template (`adding_benchmarks.mdx`) and the resulting 8 pages are visibly consistent — direct proof a template is what produces consistency here.
---
## Strengths
1. **Organic contributor discipline is currently much better than the June audit's numbers implied.** 18/19 spot-checked new policies/robots/benchmarks shipped docs in the same PR as the code; 17/18 CLI commands are documented; all teleoperator types get at least one mention. The underlying practice is healthier than the structural numbers (orphans, thin nav sections) suggest — the real problem is increasingly discoverability and consistency, not absence of effort.
2. **The doc-builder CI pipeline (PR previews, main build, versioned release builds) is solid, standard infrastructure already in place** — nothing bespoke to build or maintain, and it's the same tooling Transformers uses.
3. **A working, provably-effective template pattern already exists for one catalog (benchmarks)**`adding_benchmarks.mdx`'s checklist + "at a glance" table produces 8 structurally consistent pages. This is the single best evidence in the repo for what fixes the policy/robot template gap, and it needs no new invention, just extension.
4. **Genuine single-source-of-truth patterns already exist and work**: `docs/source/contributing.md` is a real symlink to root `CONTRIBUTING.md` (verified via `ls -la`/`find -type l`). Most policy `src/README.md` files are also symlinked into `docs/source/policy_*_README.md`, eliminating drift at that layer (though the docs-side stub itself remains an unlinked orphan — see weaknesses).
5. **Zero broken toctree links**, and specific pages are genuinely high quality: `cheat-sheet.mdx` (spot-checked line-by-line against source, fully accurate), `bring_your_own_policies.mdx` (clear scope, working PR checklist), `installation.mdx`'s tables/OS tabs, and the SO-100/SO-101 hardware pages.
6. **Root `README.md` already contains a usable taxonomy for policies** (Imitation Learning / RL / VLAs / World Models / Reward Models) that the flat 15-item Policies sidebar never adopted — the raw material for a better IA already exists in-repo.
---
## Weaknesses
### High
- **No comparison/decision-guidance page for any multi-option category** (Policies=15, Robots=11, Benchmarks=10, Teleoperators) — the exact content (decision rules, profiling snapshot) already exists in `AGENT_GUIDE.md` §6 but was never ported into `docs/source`.
- **The landing page (`index.mdx`) provides zero navigation** — 23 lines of marketing copy and a Discord link, no path to installation, cheat-sheet, or audience-specific guidance.
- **Teleoperators' dedicated nav section (2 pages) badly undersells real coverage** (16-17 implementations, all mentioned somewhere) — a new user browsing that sidebar section would wrongly conclude only phone/Isaac teleop exist. Content exists; it's a pure discoverability failure.
- **Multiple concrete broken/stale CLI examples in published, frequently copy-pasted pages**: `act.mdx` tells readers to eval with `lerobot-record` while the shown code uses `lerobot-rollout`; an identical shell-breaking snippet (`\`-terminated comment swallowing a continuation) is duplicated in `il_robots.mdx` and `lerobot-dataset-v3.mdx`; `il_robots.mdx:421` uses the dead `--control.push_to_hub` flag namespace; `openarm.mdx:208-210` (added Jan 2026) gives a `lerobot-record` example with flat pre-draccus flags that don't exist on the current config classes.
- **3 shipping policies (Diffusion, TD-MPC, VQ-BeT) have zero real published documentation** — only a thin, toctree-unreferenced stub exists for each.
- **Root `README.md` has drifted as the project scaled**: its Supported Hardware line omits SO-101 entirely; its SoTA Models table links ACT/Diffusion/VQ-BeT/Multitask-DiT/TDMPC/GR00T/SmolVLA to thin orphan stubs instead of the richer tutorial pages that exist for several of them.
- **No content template for policy/robot doc pages, unlike benchmarks** — heading sets and length vary 80-528 lines with almost no shared vocabulary (`molmoact2.mdx` has no Overview section at all; `smolvla.mdx` has no Overview/Architecture/Citation/License and reads as pure tutorial).
- **Docs CI triggers only on `paths: docs/**`** — code-only PRs get zero automated docs signal; this is the mechanism directly responsible for the one measured lag (OMX's docs page shipped 41 days after its code, in a separate PR).
- **Proof that orphan-stub drift isn't hypothetical**: `policy_molmoact2_README.md` is a stale 39-line pre-refactor snapshot sitting 481 lines behind the maintained 495-line page it once mirrored — nothing caught this silently rotting.
### Medium
- 16 orphaned `policy_*_README.md` files unreferenced by the toctree; 13 are pure duplicate cruft shadowing a maintained same-topic `.mdx` page, creating a false "two files to keep in sync" impression for contributors.
- No `_redirects.yml` anywhere — future cleanup of the orphans/renames has no 404 safety net.
- `installation.mdx` ends on a dangling forward-reference ("follow the link below to use LeRobot with your robot") — no link follows.
- `cheat-sheet.mdx`'s "Policy Types" line is stale (`act, diffusion, smolvla, pi05`) against the current 15-entry Policies section, and lists `diffusion`, which has no working toctree page at all.
- "Tutorials" (10 pages) mixes beginner, contributor, and RL-researcher content in one flat, unlabeled list.
- The EnvHub feature family is split inconsistently across two unrelated sections (`envhub.mdx`/`envhub_leisaac.mdx` under Simulation, `envhub_isaaclab_arena.mdx` under Benchmarks) despite identical framing/opening text.
- Admonition syntax is split GFM `[!NOTE]` vs. doc-builder `<Tip>` across sampled pages; since the site actually builds with `hf-doc-builder`, the majority pattern may not render as a styled callout at all (plausible, not independently render-verified).
- No CODEOWNERS file anywhere — zero designated reviewer routing for `docs/source`.
- `CONTRIBUTING.md` never links the good in-repo extension guides (`adding_benchmarks.mdx`, `bring_your_own_policies.mdx`, `integrate_hardware.mdx`) and states no docs-required policy.
- The "docs required" PR checklist convention is applied inconsistently: Policies and Benchmarks have an explicit required-docs checklist; `integrate_hardware.mdx` (Robots/Teleoperators) has none — a plausible root cause of the Teleoperators gap above.
- Growth repeatedly stacks simultaneous PRs against one shared, hand-edited `_toctree.yml` (6 benchmark PRs in 48h; 3 policy PRs in 3 days) — a merge/oversight risk even though no damage from it was found yet.
- `policy_sarm_README.md` has no source of truth left to sync against at all (SARM moved `policies/``rewards/`, no README carried over) — pure abandoned cruft next to the real 593-line `sarm.mdx`.
- No automated check anywhere for toctree completeness or README/mdx sync — the only CI gate is whether the doc-builder build succeeds, which doesn't catch missing coverage.
### Low
- No dedicated "Motors" section despite `motors/` being named as a distinct hardware layer in `CLAUDE.md`; `feetech.mdx`/`damiao.mdx` sit in a catch-all "Resources" section instead.
- "Sensors" is a single-page top-level nav section (`cameras.mdx` only).
- Policies (15), Benchmarks (10), and Robots (11) are flat, ungrouped lists with no internal sub-headings, despite `README.md` already having a usable taxonomy that could be reused.
- "SO-101" vs. "SO101" naming is inconsistent across ~15 files, and even within a single file (`so100.mdx` uses both).
- `hilserl.mdx` (950 lines) and `il_robots.mdx` (638 lines) each mix tutorial, reference, and troubleshooting content in one long page.
- `reachy2_camera` has zero doc mentions anywhere; the `zmq` camera backend is documented only incidentally inside `unitree_g1.mdx` rather than centrally in `cameras.mdx`.
- The PR template's "Documentation updated" line is a self-reported, unenforced checkbox.
---
## Recommendations
### 1. Quick wins (small effort, ship this week, no maintainer proposal needed)
- Fix the 4 verified broken commands: `act.mdx` (`lerobot-record``lerobot-rollout`), the duplicated shell-breaking snippet in `il_robots.mdx` + `lerobot-dataset-v3.mdx`, `il_robots.mdx`'s `--control.push_to_hub``--dataset.push_to_hub`, and `openarm.mdx`'s invalid flat-flag record example.
- Delete the 3 fully-orphaned policy stubs (Diffusion, TD-MPC, VQ-BeT) and give each a real toctree page by promoting the existing `src/` README content.
- Delete the remaining 13 duplicate orphan stubs and the dead `policy_sarm_README.md`.
- Fix `installation.mdx`'s dangling ending and `cheat-sheet.mdx`'s stale policy-type list.
- Fix root `README.md`: add SO-101 to the Supported Hardware line; repoint the SoTA Models table's links from orphan stubs to the richer existing tutorial pages.
- Move `envhub_isaaclab_arena.mdx` into Simulation (one-line YAML change).
- Add 2-4 orientation sentences + links to the top of `index.mdx` (installation, cheat-sheet, "I have hardware" / "I don't" / "I want to contribute") without touching its marketing framing.
- Publish a single "Choosing a policy" page that ports the already-written decision rules from `AGENT_GUIDE.md` §6 — highest-leverage fix available.
- Start a `docs/source/_redirects.yml` now, before the orphan cleanup above creates the first real dead links.
- No action needed on the `contributing.md` "duplication" claim — verified it's a working symlink, not an anti-pattern.
### 2. Structural / UX changes (need a proposal + maintainer buy-in, phase as separate PRs)
- Split "Tutorials" into an explicit beginner-facing section vs. an "Advanced & Research/Contributor" section (mechanical YAML reorg, no content rewrites).
- Add a "Teleoperators" index page listing all types with one-line descriptions and links to wherever each is actually documented today.
- Re-group the flat Policies/Benchmarks/Robots lists into sub-headings reusing the taxonomy `README.md` already has.
- Treat a dedicated "Motors" section, deeper sub-grouping, and closing the remaining teleoperator/robot coverage gaps as a phased backlog of independently reviewable PRs rather than one restructure.
- Standardize SO-101/SO101 naming and (after confirming actual doc-builder rendering behavior) the admonition syntax, each as one mechanical, low-risk PR.
- Longer-term and biggest-ticket: port more of `AGENT_GUIDE.md`'s procedural content (training duration heuristics, eval targets, data-collection tips) into the published site.
### 3. Maintainability / process changes (prevent debt from reaccumulating at the current growth rate)
- Port `adding_benchmarks.mdx`'s explicit "writing a doc page" checklist/template into `bring_your_own_policies.mdx` and `integrate_hardware.mdx`.
- Add a CODEOWNERS entry for `docs/source/` (and ideally `_toctree.yml` specifically) so docs PRs get routed to a real reviewer.
- Link the extension guides from `CONTRIBUTING.md` and state a docs-required policy there.
- Add a lightweight CI/pre-commit script (not a full doc-builder run) that fails when (a) a `docs/source/*.mdx` file isn't reachable from `_toctree.yml`, or (b) a new `register_subclass` policy/robot/teleoperator/env lands with no corresponding doc file in the same diff.
- Decide the fate of the `policy_*_README.md` symlink convention going forward: it is currently self-perpetuating because `bring_your_own_policies.mdx`'s own checklist instructs new contributors to create the file that ends up orphaned. Either fold citation/paper content into the main `.mdx` tutorial, or wire the stub into the toctree as a linked citation anchor.
- Make the PR template's "Documentation updated" checkbox actionable (e.g., "(N/A if this PR only touches tests/CI/refactors)").
- Defer heavier generated-registry/support-matrix tooling (Transformers/Ultralytics-style single source of truth) until closer to 1.0.
---
## Notes on cross-checking
Five independent audit passes fed this report; two disagreements were resolved during synthesis:
- **Toctree section count** — 15 is correct (independently parsed twice from `_toctree.yml`).
- **`contributing.md`** — it's a working symlink, not a hand-duplicated anti-pattern; one pass's `diff`-based claim didn't survive checking `find -type l`.
- **Orphan-README "sync mechanism"** — mostly fixed at the `src/``docs` layer via symlinks, but that just moved the unresolved drift to the docs-side stub's absence from the toctree, and left old pre-symlink stubs (MolmoAct2) as dead leftovers.
- **Teleoperators coverage** — the nav-section framing ("~2/16") and the "mentioned somewhere" framing are both true; the real gap is discoverability, not content.
+1 -1
View File
@@ -76,7 +76,7 @@ If your local computer doesn't have a powerful GPU, you can utilize Google Colab
## Evaluating ACT
Once training is complete, you can evaluate your ACT policy using the `lerobot-record` command with your trained policy. This will run inference and record evaluation episodes:
Once training is complete, you can evaluate your ACT policy using the `lerobot-rollout` command with your trained policy. This will run inference and record evaluation episodes:
```bash
lerobot-rollout \
+1 -1
View File
@@ -194,8 +194,8 @@ lerobot-record \
--dataset.single_task="Navigate around obstacles" \
--dataset.streaming_encoding=true \
--dataset.encoder_threads=2 \
# --dataset.rgb_encoder.vcodec=auto \
--display_data=true
# Optionally, set --dataset.rgb_encoder.vcodec=auto to pick a specific video codec
```
Replace `your_username/dataset_name` with your Hugging Face username and a name for your dataset.
+2 -2
View File
@@ -232,8 +232,8 @@ lerobot-record \
--dataset.private=true \
--dataset.streaming_encoding=true \
--dataset.encoder_threads=2 \
# --dataset.rgb_encoder.vcodec=auto \
--display_data=true
# Optionally, set --dataset.rgb_encoder.vcodec=auto to pick a specific video codec
```
### Replay
@@ -278,6 +278,6 @@ lerobot-record \
--dataset.num_episodes=10 \
--dataset.streaming_encoding=true \
--dataset.encoder_threads=2 \
# --dataset.rgb_encoder.vcodec=auto \
--policy.path=outputs/train/hopejr_hand/checkpoints/last/pretrained_model
# Optionally, set --dataset.rgb_encoder.vcodec=auto to pick a specific video codec
```
+2 -2
View File
@@ -207,8 +207,8 @@ lerobot-record \
--dataset.num_episodes=5 \
--dataset.single_task="Grab the black cube" \
--dataset.streaming_encoding=true \
# --dataset.rgb_encoder.vcodec=auto \
--dataset.encoder_threads=2
# Optionally, set --dataset.rgb_encoder.vcodec=auto to pick a specific video codec
```
</hfoption>
<hfoption id="API example">
@@ -418,7 +418,7 @@ If you want to dive deeper into this important topic, you can check out the [blo
## Visualize a dataset
If you uploaded your dataset to the hub with `--control.push_to_hub=true`, you can [visualize your dataset online](https://huggingface.co/spaces/lerobot/visualize_dataset) by copy pasting your repo id given by:
If you uploaded your dataset to the hub with `--dataset.push_to_hub=true`, you can [visualize your dataset online](https://huggingface.co/spaces/lerobot/visualize_dataset) by copy pasting your repo id given by:
```bash
echo ${HF_USER}/so101_test
-11
View File
@@ -141,17 +141,6 @@ 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.
+1 -1
View File
@@ -44,8 +44,8 @@ lerobot-record \
--dataset.num_episodes=5 \
--dataset.single_task="Grab the black cube" \
--dataset.streaming_encoding=true \
# --dataset.rgb_encoder.vcodec=auto \
--dataset.encoder_threads=2
# Optionally, set --dataset.rgb_encoder.vcodec=auto to pick a specific video codec
```
See the [recording guide](./il_robots#record-a-dataset) for more details.
+4 -3
View File
@@ -205,9 +205,10 @@ lerobot-record \
--teleop.type=openarm_leader \
--teleop.port=can1 \
--teleop.id=my_leader \
--repo-id=my_hf_username/my_openarm_dataset \
--fps=30 \
--num-episodes=10
--dataset.repo_id=my_hf_username/my_openarm_dataset \
--dataset.single_task="Grab the black cube" \
--dataset.fps=30 \
--dataset.num_episodes=10
```
## Configuration Options
+2 -2
View File
@@ -161,8 +161,8 @@ lerobot-record \
--dataset.private=true \
--dataset.streaming_encoding=true \
--dataset.encoder_threads=2 \
# --dataset.rgb_encoder.vcodec=auto \
--display_data=true
# Optionally, set --dataset.rgb_encoder.vcodec=auto to pick a specific video codec
```
#### Specific Options
@@ -203,8 +203,8 @@ lerobot-record \
--dataset.private=true \
--dataset.streaming_encoding=true \
--dataset.encoder_threads=2 \
# --dataset.rgb_encoder.vcodec=auto \
--display_data=true
# Optionally, set --dataset.rgb_encoder.vcodec=auto to pick a specific video codec
```
##### `--robot.use_external_commands`
+13 -14
View File
@@ -100,20 +100,19 @@ Once you are logged in, you can run inference in your setup by doing:
lerobot-rollout \
--strategy.type=base \
--robot.type=so101_follower \
--robot.port=/dev/ttyACM0 \ # <- Use your port
--robot.id=my_blue_follower_arm \ # <- Use your robot id
--robot.cameras="{ front: {type: opencv, index_or_path: 8, width: 640, height: 480, fps: 30}}" \ # <- Use your cameras
--task="Grasp a lego block and put it in the bin." \ # <- Use the same task description you used in your dataset recording
# <- RTC optional, use when running on low power hardware \
# --inference.type=rtc \
# --inference.rtc.execution_horizon=10 \
# --inference.rtc.max_guidance_weight=10.0 \
# <- Teleop optional if you want to teleoperate in between episodes \
# --teleop.type=so100_leader \
# --teleop.port=/dev/ttyACM0 \
# --teleop.id=my_red_leader_arm \
# --display_data=true #optional use if you want to see the camera stream \
--policy.path=HF_USER/FINETUNE_MODEL_NAME # <- Use your fine-tuned model
--robot.port=/dev/ttyACM0 \
--robot.id=my_blue_follower_arm \
--robot.cameras="{ front: {type: opencv, index_or_path: 8, width: 640, height: 480, fps: 30}}" \
--task="Grasp a lego block and put it in the bin." \
--policy.path=HF_USER/FINETUNE_MODEL_NAME
```
Replace `--robot.port`, `--robot.id`, `--robot.cameras`, `--task`, and `--policy.path` with your own port, robot ID, camera setup, task description (matching what you used when recording your dataset), and fine-tuned model repo ID.
A few optional flags you can add to the command above:
- **RTC** (useful on low-power hardware): `--inference.type=rtc --inference.rtc.execution_horizon=10 --inference.rtc.max_guidance_weight=10.0`
- **Teleoperate in between episodes**: `--teleop.type=so100_leader --teleop.port=/dev/ttyACM0 --teleop.id=my_red_leader_arm`
- **See the camera stream**: `--display_data=true`
Depending on your evaluation setup, you can configure the duration and the number of episodes to record for your evaluation suite.
-6
View File
@@ -33,8 +33,6 @@ 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
@@ -64,10 +62,6 @@ 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
+4 -10
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:
"""Validate bindings and require text or low-level action supervision."""
"""Ensure every templated binding is known and at least one turn is a target."""
assert self.messages is not None
known_bindings = set(DEFAULT_BINDINGS) | set(self.bindings or {}) | {"task"}
@@ -156,14 +156,8 @@ class TrainingRecipe:
if missing:
raise ValueError(f"MessageTurn references unknown binding(s): {sorted(missing)}")
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)."
)
if not any(turn.target for turn in self.messages):
raise ValueError("Message recipes must contain at least one target turn.")
def _validate_blend_recipe(self) -> None:
"""Ensure each blend component is a non-empty, weighted message recipe."""
-16
View File
@@ -1,16 +0,0 @@
# 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}
@@ -1,13 +0,0 @@
# 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}
@@ -1,30 +0,0 @@
# 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}
@@ -1,70 +0,0 @@
# 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}
-30
View File
@@ -163,40 +163,10 @@ 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:
+2 -16
View File
@@ -66,17 +66,6 @@ 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.
@@ -98,14 +87,11 @@ 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=episodes,
episodes=cfg.dataset.episodes,
delta_timestamps=delta_timestamps,
image_transforms=image_transforms,
revision=cfg.dataset.revision,
@@ -118,7 +104,7 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
dataset = StreamingLeRobotDataset(
cfg.dataset.repo_id,
root=cfg.dataset.root,
episodes=episodes,
episodes=cfg.dataset.episodes,
delta_timestamps=delta_timestamps,
image_transforms=image_transforms,
revision=cfg.dataset.revision,
+10 -73
View File
@@ -162,28 +162,14 @@ def render_sample(
task: str | None = None,
dataset_ctx: Any | None = None,
) -> RenderedMessages | None:
"""Resolve one sample's bindings and render its message recipe.
"""Render the chat-style messages for a single dataset sample.
Returns ``None`` when no text or low-level action supervision applies.
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.
"""
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,
@@ -197,55 +183,6 @@ 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:
@@ -409,9 +346,7 @@ def _render_message_recipe(
if turn.target:
target_indices.append(message_idx)
# 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:
if not target_indices:
return None
rendered = {
@@ -468,12 +403,14 @@ def _validate_rendered(rendered: RenderedMessages) -> None:
if len(streams) != len(messages):
raise ValueError("message_streams must be aligned with messages.")
# 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.")
if not target_indices:
raise ValueError("Rendered samples must contain at least one target 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(
@@ -37,19 +37,13 @@ def is_image_feature(key: str) -> bool:
@dataclass
class ConcurrencyConfig:
"""Configuration for the concurrency of the actor and learner.
Possible values are:
- "threads": Use threads for the actor and learner.
- "processes": Use processes for the actor and learner.
``multiprocessing_context`` selects the process-wide start method when
processes are used. Set it to ``None`` to preserve Python's default or a
method already selected by the embedding application.
"""
actor: str = "threads"
learner: str = "threads"
multiprocessing_context: str | None = "spawn"
@dataclass
+3
View File
@@ -175,6 +175,9 @@ 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)):
@@ -132,20 +132,10 @@ class MapDeltaActionToRobotActionStep(RobotActionProcessorStep):
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
for axis in ["x", "y", "z"]:
for axis in ["x", "y", "z", "gripper"]:
features[PipelineFeatureType.ACTION].pop(f"delta_{axis}", None)
features[PipelineFeatureType.ACTION].pop("gripper", None)
for feat in [
"enabled",
"target_x",
"target_y",
"target_z",
"target_wx",
"target_wy",
"target_wz",
"gripper_vel",
]:
for feat in ["enabled", "target_x", "target_y", "target_z", "target_wx", "target_wy", "target_wz"]:
features[PipelineFeatureType.ACTION][f"{feat}"] = PolicyFeature(
type=FeatureType.ACTION, shape=(1,)
)
+4 -85
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, snapshot_download
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file, save_file
from lerobot.configs import PipelineFeatureType, PolicyFeature
@@ -205,10 +205,6 @@ 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
@@ -553,22 +549,6 @@ 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)
@@ -753,13 +733,7 @@ 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,
config_filename,
hub_download_kwargs,
is_local_source,
loaded_config, overrides or {}, model_id, base_path, hub_download_kwargs, is_local_source
)
# 4. Validate that all overrides were used
@@ -948,7 +922,6 @@ 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]]:
@@ -1003,68 +976,15 @@ 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,
config_filename,
hub_download_kwargs,
is_local_source,
step_instance, step_entry, model_id, base_path, 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,
@@ -1224,7 +1144,6 @@ 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:
@@ -1290,7 +1209,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
# Download from Hub
state_path = hf_hub_download(
repo_id=model_id,
filename=(Path(config_filename).parent / state_filename).as_posix(),
filename=state_filename,
repo_type="model",
**hub_download_kwargs,
)
@@ -16,7 +16,7 @@
from __future__ import annotations
from dataclasses import asdict, dataclass
from dataclasses import dataclass
from typing import Any
from lerobot.configs import PipelineFeatureType, PolicyFeature
@@ -32,18 +32,17 @@ from .pipeline import ProcessorStep, ProcessorStepRegistry
@dataclass
@ProcessorStepRegistry.register(name="render_messages_processor")
class RenderMessagesStep(ProcessorStep):
"""Render language columns into recipe-defined messages and supervision metadata."""
"""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.
"""
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 {}
@@ -51,17 +50,7 @@ class RenderMessagesStep(ProcessorStep):
events = complementary_data.get(LANGUAGE_EVENTS) or []
if not persistent and not events:
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)
return transition
timestamp = complementary_data.get("timestamp")
if timestamp is None:
@@ -78,147 +67,18 @@ class RenderMessagesStep(ProcessorStep):
dataset_ctx=self.dataset_ctx,
)
if rendered is None:
rendered = _fallback_low_level_render(complementary_data.get("task"))
if rendered is None:
return None
return None
new_transition = transition.copy()
new_complementary_data = dict(new_transition.get(TransitionKey.COMPLEMENTARY_DATA) or {})
new_complementary_data = dict(complementary_data)
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": [],
}
+17 -50
View File
@@ -25,7 +25,6 @@ from __future__ import annotations
import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Any
import torch
@@ -33,7 +32,6 @@ 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,
@@ -138,7 +136,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
@@ -351,7 +349,6 @@ 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)
@@ -415,15 +412,14 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
# During inference, no action is available, skip tokenization
return new_transition
# Tokenize and get masks for the full formatted sequence and the discrete action codes.
tokens, mask, code_mask = self._tokenize_action(action)
# Tokenize and get both tokens and mask
tokens, 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
@@ -434,7 +430,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, torch.Tensor]:
def _tokenize_action(self, action: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""
Tokenizes the action tensor and creates a mask.
@@ -463,7 +459,6 @@ 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)
@@ -481,82 +476,65 @@ 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
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)
# add bos
tokens = torch.cat(
[
torch.tensor([bos_id], device=action.device),
prompt_tokens,
action_code_tokens,
end_tokens,
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),
]
)
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(pad_len, dtype=torch.bool, device=action.device),
torch.zeros(
self.max_action_tokens - len(tokens), 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, pad_len), value=0)
tokens = torch.nn.functional.pad(tokens, (0, self.max_action_tokens - len(tokens)), 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, code_masks_batch
return tokens_batch, 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]:
@@ -572,9 +550,6 @@ 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
@@ -583,14 +558,6 @@ 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]]:
+4 -2
View File
@@ -91,7 +91,7 @@ from lerobot.robots import so_follower # noqa: F401
from lerobot.teleoperators import gamepad, so_leader # noqa: F401
from lerobot.teleoperators.utils import TeleopEvents
from lerobot.utils.device_utils import get_safe_torch_device
from lerobot.utils.process import ProcessSignalHandler, ensure_multiprocessing_start_method
from lerobot.utils.process import ProcessSignalHandler
from lerobot.utils.random_utils import set_seed
from lerobot.utils.robot_utils import precise_sleep
from lerobot.utils.transition import (
@@ -124,7 +124,9 @@ def actor_cli(cfg: TrainRLServerPipelineConfig):
cfg.validate()
display_pid = False
if not use_threads(cfg):
ensure_multiprocessing_start_method(cfg.policy.concurrency.multiprocessing_context)
import torch.multiprocessing as mp
mp.set_start_method("spawn")
display_pid = True
# Create logs directory to ensure it exists
+4 -2
View File
@@ -102,7 +102,7 @@ from lerobot.utils.constants import (
)
from lerobot.utils.device_utils import get_safe_torch_device
from lerobot.utils.io_utils import load_json, write_json
from lerobot.utils.process import ProcessSignalHandler, ensure_multiprocessing_start_method
from lerobot.utils.process import ProcessSignalHandler
from lerobot.utils.random_utils import set_seed
from lerobot.utils.utils import (
format_big_number,
@@ -123,7 +123,9 @@ def train_cli(cfg: TrainRLServerPipelineConfig):
# Fail fast with a friendly error if the optional ``hilserl`` extra is missing.
require_package("grpcio", extra="hilserl", import_name="grpc")
if not use_threads(cfg):
ensure_multiprocessing_start_method(cfg.policy.concurrency.multiprocessing_context)
import torch.multiprocessing as mp
mp.set_start_method("spawn")
# Use the job_name from the config
train(
+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", *LANGUAGE_COLUMNS}
_PYTHON_LIST_KEYS = {"messages", "message_streams", "target_message_indices"}
def lerobot_collate_fn(batch: list[dict[str, Any] | None]) -> dict[str, Any] | None:
-2
View File
@@ -26,7 +26,6 @@ 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"
@@ -35,7 +34,6 @@ 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"
-28
View File
@@ -16,39 +16,11 @@
# limitations under the License.
import logging
import multiprocessing
import os
import signal
import sys
def ensure_multiprocessing_start_method(start_method: str | None) -> None:
"""Set a multiprocessing start method once, or verify the existing method matches.
Passing ``None`` leaves Python's process-wide default untouched. This is useful
when LeRobot is embedded in an application that owns multiprocessing setup.
"""
if start_method is None:
return
available_methods = multiprocessing.get_all_start_methods()
if start_method not in available_methods:
raise ValueError(
f"Multiprocessing start method must be one of {available_methods} on this platform, "
f"got {start_method!r}."
)
current_method = multiprocessing.get_start_method(allow_none=True)
if current_method is None:
multiprocessing.set_start_method(start_method)
elif current_method != start_method:
raise RuntimeError(
f"Multiprocessing start method is already {current_method!r}; cannot change it to "
f"{start_method!r}. Set the configured multiprocessing context to null to keep the "
"application's existing method, or launch LeRobot in a fresh process."
)
class ProcessSignalHandler:
"""Utility class to attach graceful shutdown signal handlers.
-7
View File
@@ -29,13 +29,6 @@ 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.
-78
View File
@@ -343,84 +343,6 @@ 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
+3 -1
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, compute_sampler_state
from lerobot.datasets.sampler import EpisodeAwareSampler
def calculate_episode_data_index(hf_dataset: Dataset) -> dict[str, torch.Tensor]:
@@ -154,6 +154,8 @@ 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
@@ -113,7 +113,6 @@ def test_gaussian_actor_config_default_initialization():
# Concurrency configuration
assert config.concurrency.actor == "threads"
assert config.concurrency.learner == "threads"
assert config.concurrency.multiprocessing_context == "spawn"
assert isinstance(config.actor_network_kwargs, ActorNetworkConfig)
assert isinstance(config.policy_kwargs, PolicyConfig)
@@ -153,7 +152,6 @@ def test_concurrency_config():
config = ConcurrencyConfig()
assert config.actor == "threads"
assert config.learner == "threads"
assert config.multiprocessing_context == "spawn"
def test_gaussian_actor_config_custom_initialization():
@@ -12,9 +12,7 @@ from lerobot.processor.render_messages_processor import RenderMessagesStep # no
from lerobot.types import TransitionKey # noqa: E402
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."""
def test_render_messages_step_noops_without_language_columns():
recipe = TrainingRecipe(
messages=[
MessageTurn(role="user", content="${task}", stream="high_level"),
@@ -23,24 +21,6 @@ def test_render_messages_step_renders_task_fallback_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
@@ -78,70 +58,3 @@ 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"] == [[], []]
+1 -41
View File
@@ -25,7 +25,7 @@ import pytest
import torch
from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
from lerobot.processor import ActionTokenizerProcessorStep, DataProcessorPipeline, TokenizerProcessorStep
from lerobot.processor import DataProcessorPipeline, TokenizerProcessorStep
from lerobot.processor.converters import create_transition, identity_transition
from lerobot.types import TransitionKey
from lerobot.utils.constants import (
@@ -88,46 +88,6 @@ 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."""
+5 -5
View File
@@ -9,7 +9,7 @@ import torch # noqa: E402
from lerobot.utils.collate import lerobot_collate_fn # noqa: E402
def test_lerobot_collate_preserves_messages_and_raw_language():
def test_lerobot_collate_preserves_messages_and_drops_raw_language():
batch = [
{
"index": torch.tensor(0),
@@ -17,14 +17,14 @@ def test_lerobot_collate_preserves_messages_and_raw_language():
"message_streams": ["low_level"],
"target_message_indices": [0],
"language_persistent": [{"content": "raw"}],
"language_events": [{"content": "event a"}],
"language_events": [],
},
{
"index": torch.tensor(1),
"messages": [{"role": "assistant", "content": "b"}],
"message_streams": ["low_level"],
"target_message_indices": [0],
"language_persistent": [{"content": "raw b"}],
"language_persistent": [{"content": "raw"}],
"language_events": [],
},
]
@@ -36,8 +36,8 @@ def test_lerobot_collate_preserves_messages_and_raw_language():
assert out["messages"][1][0]["content"] == "b"
assert out["message_streams"] == [["low_level"], ["low_level"]]
assert out["target_message_indices"] == [[0], [0]]
assert out["language_persistent"] == [[{"content": "raw"}], [{"content": "raw b"}]]
assert out["language_events"] == [[{"content": "event a"}], []]
assert "language_persistent" not in out
assert "language_events" not in out
def test_lerobot_collate_passes_through_standard_batch():