Compare commits

..

16 Commits

Author SHA1 Message Date
Steven Palma f37be3edbe fix(eval): prevent eval_policy crash when start_seed is None and num_envs>1 (#4203)
* fix(eval): align seed list length with num_envs when unseeded

eval_policy appended a single None per batch to all_seeds on the unseeded path while the reward and success lists grew by num_envs. The per-episode zip(..., strict=True) then raised ValueError for num_envs > 1. Extend all_seeds by num_envs so the lists stay aligned.

* chore(tests): delete lerobot_eval test

---------

Co-authored-by: Devin Lai <markauto75@gmail.com>
2026-07-28 18:41:28 +02:00
Khalil Meftah 4d076845ac fix peft factory test mocking (#4201) 2026-07-28 17:54:58 +02:00
Steven Palma 413972c812 fix(env): eval env lifecycle (#4194)
Co-authored-by: itxaiohanglover <1531137510@qq.com>
Co-authored-by: nickndeng <nickndeng@gmail.com>
Co-authored-by: nickndeng <107904079+nickndeng@users.noreply.github.com>
Co-authored-by: Pepijn <138571049+pkooij@users.noreply.github.com>
2026-07-28 16:45:48 +02:00
Steven Palma 0449aa02f6 fix(utils): handle missing/unresponsive TTS on Linux (#4199)
* fix: handle missing/unresponsive TTS on Linux

spd-say may be installed but hang indefinitely when speech-dispatcher
is not running. Add a 5s timeout and catch TimeoutExpired alongside
FileNotFoundError so recording continues without audio.

* chore(utils): add log warning for say

---------

Co-authored-by: Jiwen Cai <jiwenc@nvidia.com>
2026-07-28 16:45:32 +02:00
Alexandre Edmond a05c0833e1 chore(mypy): cover annotations and transforms (#3860)
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-28 16:25:45 +02:00
Alexandre Edmond 7b76d94c5b Handle resuming empty local datasets (#3859)
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-28 16:25:42 +02:00
Khalil Meftah ec2dbc1c98 fix(policy): honor revisions when loading PEFT checkpoints (#4189) 2026-07-28 15:41:47 +02:00
Steven Palma d526785e47 fix(dependencies): protect peft import (#4188)
Signed-off-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-28 14:57:50 +02:00
Steven Palma 4af7c70664 refactor(logging): standardize logging with getLogger(__name__) in scripts (#4192)
* refactor(logging): replace print with logger in lerobot_info

* refactor(logging): replace print with logger in convert_dataset_v21_to_v30

* refactor(logging): replace print with logger in lerobot_annotate

* refactor(logging): replace print with logger in lerobot_dataset_viz

* refactor(logging): replace print with logger in lerobot_eval

* refactor(logging): replace print with logger in lerobot_find_cameras

* refactor(logging): replace print with logger in lerobot_find_joint_limits

* refactor(logging): replace print with logger in lerobot_find_port

* refactor(logging): replace print with logger in lerobot_imgtransform_viz

* refactor(logging): replace print with logger in lerobot_setup_can

* refactor(logging): replace print with logger in lerobot_teleoperate

* refactor(logging): replace print with logger in lerobot_train_tokenizer

* fix(logging): preserve CLI output semantics

---------

Co-authored-by: ailisilob <2248345706@qq.com>
2026-07-28 14:42:38 +02:00
charlie8612 a855570097 feat(motors): add XH540-W150, XC330-T288, XC330-T181 to Dynamixel tables (#3815)
Register three X-series Dynamixel models so they can be driven by
DynamixelMotorsBus: XH540-W150 (model 1110), XC330-T288 (1220) and
XC330-T181 (1210). All are standard Protocol 2.0 X-series motors that
share the existing X_SERIES control/baudrate/encoding tables and 4096
resolution; only the model number and operating-mode list are
model-specific. Values verified against the ROBOTIS e-manual.

These motors are used by the ROBOTIS OMY-L100 arm, among others.

Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-28 14:27:48 +02:00
Steven Palma 167e22ba51 feat(record): add --dataset.no_stamp to opt out of repo_id timestamping (#4190)
* feat(record): add `--dataset.no_stamp` to opt out of repo_id timestamping

stamp_repo_id() unconditionally appended a date-time tag to repo_id for every new (non-resume) dataset, so users managing their own versioned repo names (e.g. for a later lerobot-edit-dataset merge) could not opt out. Add a no_stamp field to DatasetRecordConfig and make stamp_repo_id() a no-op when it is set. The flag covers both lerobot-record and lerobot-rollout since they share this config, and no call-site changes are needed. Fixes #3722.

* chore(tests): delete dataset config test

---------

Co-authored-by: Philipp Sinitsin <ph.sinitsin@gmail.com>
2026-07-28 14:06:28 +02:00
Steven Palma 00c25c65c2 feat(camera): add manual exposure, gain, and white balance options for RealSense cameras (#4130)
* feat(camera): add manual exposure, gain, and white balance options for RealSense cameras

The RealSense camera integration lacked sensor-level controls, causing
issues like unstable lighting from auto-exposure hunting. This adds
optional `exposure`, `gain`, and `white_balance` fields to
RealSenseCameraConfig that disable the corresponding auto modes and
apply fixed values when set.

* fix: support D405 stereo module for sensor options

D405 exposes color stream via "Stereo Module", not "RGB Camera".
Fall back to Stereo Module when RGB Camera is not found.

* test(camera): add unit tests + range-aware errors for RealSense sensor options

Address PR #3220 review:
- Wrap set_option calls; re-raise ValueError with option name, value,
  and sensor.get_option_range() diagnostics on out-of-range values.
- Add unit tests for _get_color_sensor (RGB Camera, D405 Stereo Module
  fallback, diagnostic error) and _configure_sensor_options (no-op,
  all values, unsupported warns, partial config, out-of-range raise).

* fix(realsense): validate manual color controls

* refactor(camera): apply feedback

---------

Co-authored-by: Lev Kozlov <kozlov.l.a10@gmail.com>
2026-07-28 13:41:09 +02:00
Steven Palma 23f6d5dabd fix(cameras): release device handle when connect() setup fails (#4187)
Co-authored-by: Ryan Rana <39924576+RyanRana@users.noreply.github.com>
2026-07-28 13:21:06 +02:00
Xingdong Zuo 9b25b7fe0a feat(lekiwi): support LeKiwi in the rollout/eval CLI (#3742)
* feat(lekiwi): support LeKiwi in the rollout/eval CLI

Register the lekiwi robot in lerobot_rollout.py so policies can be rolled out
on a LeKiwi, and keep base-velocity (.vel) features in build_rollout_context.

LeKiwi's observation.state and action are 9-dim (6 arm .pos + x/y/theta.vel)
and the policy is normalized on all 9. The old filter kept only .pos features,
so it fed a 6-dim vector into a 9-dim normalizer (RuntimeError, size 6 vs 9) and
silently dropped the base velocities from the action, leaving the base unable to
move. Keeping both .pos and .vel fixes both. Pure-arm robots have no .vel keys,
so this is a no-op for them.

* style: format LeKiwi rollout action features

---------

Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
Co-authored-by: Steven Palma <steven.palma@huggingface.co>
2026-07-28 12:52:02 +02:00
Steven Palma c1b6ea85d6 feat(rl): add multiprocessing option to training pipeline and sets spawn as default + guard (#4140) 2026-07-28 12:19:53 +02:00
Steven Palma ffe25afb8f fix(processors): wrong feature key dropped in delta-action transform_features (#4165) 2026-07-28 11:18:23 +02:00
76 changed files with 1130 additions and 4080 deletions
-156
View File
@@ -191,162 +191,6 @@ 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.
+13
View File
@@ -136,6 +136,10 @@ config = RealSenseCameraConfig(
height=480,
color_mode=ColorMode.RGB,
use_depth=True,
# Optional fixed color controls. Omit them to leave the current sensor settings unchanged.
exposure=120,
gain=64,
white_balance=4600,
rotation=Cv2Rotation.NO_ROTATION
)
@@ -154,6 +158,15 @@ finally:
```
<!-- prettier-ignore-end -->
Manual color controls disable the corresponding automatic exposure or white-balance mode. Their
supported ranges vary by camera model; an invalid value raises an error at connection time that
includes the range reported by the sensor. Requesting an unsupported control also raises an error.
Omitted controls leave the sensor's existing automatic or manual setting unchanged. These options
require `use_rgb=True`.
On the RealSense D405, the color stream is provided by the Stereo Module, so changing manual
exposure or gain also affects the depth stream.
</hfoption>
</hfoptions>
+1 -47
View File
@@ -1,6 +1,6 @@
# Policy Deployment (lerobot-rollout)
`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.
`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.
## Quick Start
@@ -197,52 +197,6 @@ 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,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.
+13
View File
@@ -494,6 +494,19 @@ ignore_errors = true
module = "lerobot.envs.*"
ignore_errors = false
[[tool.mypy.overrides]]
module = "lerobot.annotations.*"
ignore_errors = false
disallow_untyped_defs = true
disallow_incomplete_defs = true
check_untyped_defs = true
[[tool.mypy.overrides]]
module = "lerobot.transforms.*"
ignore_errors = false
disallow_untyped_defs = true
disallow_incomplete_defs = true
check_untyped_defs = true
# [[tool.mypy.overrides]]
# module = "lerobot.utils.*"
+78 -48
View File
@@ -120,14 +120,22 @@ class OpenCVCamera(Camera):
self.rotation: int | None = get_cv2_rotation(config.rotation)
self.backend: int = config.backend
if self.height and self.width:
self.capture_width, self.capture_height = self.width, self.height
if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
self.capture_width, self.capture_height = self.height, self.width
self.capture_width: int | None = None
self.capture_height: int | None = None
self._reset_connection_settings()
def __str__(self) -> str:
return f"{self.__class__.__name__}({self.index_or_path})"
def _reset_connection_settings(self) -> None:
"""Restore settings that may have been auto-detected during a failed connection."""
self.fps = self.config.fps
self.width = self.config.width
self.height = self.config.height
self.capture_width, self.capture_height = self.width, self.height
if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
self.capture_width, self.capture_height = self.height, self.width
@property
def is_connected(self) -> bool:
"""Checks if the camera is currently connected and opened."""
@@ -164,17 +172,25 @@ class OpenCVCamera(Camera):
f"Failed to open {self}.Run `lerobot-find-cameras opencv` to find available cameras."
)
self._configure_capture_settings()
self._start_read_thread()
try:
self._configure_capture_settings()
self._start_read_thread()
if warmup and self.warmup_s > 0:
start_time = time.time()
while time.time() - start_time < self.warmup_s:
self.async_read(timeout_ms=self.warmup_s * 1000)
time.sleep(0.1)
with self.frame_lock:
if self.latest_frame is None:
raise ConnectionError(f"{self} failed to capture frames during warmup.")
if warmup and self.warmup_s > 0:
start_time = time.time()
while time.time() - start_time < self.warmup_s:
self.async_read(timeout_ms=self.warmup_s * 1000)
time.sleep(0.1)
with self.frame_lock:
if self.latest_frame is None:
raise ConnectionError(f"{self} failed to capture frames during warmup.")
except BaseException:
try:
self._cleanup_resources()
except Exception:
logger.exception(f"Failed to fully clean up {self} after connect() failed.")
self._reset_connection_settings()
raise
logger.info(f"{self} connected.")
@@ -312,32 +328,36 @@ class OpenCVCamera(Camera):
for target in targets_to_scan:
camera = cv2.VideoCapture(target)
if camera.isOpened():
default_width = int(camera.get(cv2.CAP_PROP_FRAME_WIDTH))
default_height = int(camera.get(cv2.CAP_PROP_FRAME_HEIGHT))
default_fps = camera.get(cv2.CAP_PROP_FPS)
default_format = camera.get(cv2.CAP_PROP_FORMAT)
try:
if camera.isOpened():
default_width = int(camera.get(cv2.CAP_PROP_FRAME_WIDTH))
default_height = int(camera.get(cv2.CAP_PROP_FRAME_HEIGHT))
default_fps = camera.get(cv2.CAP_PROP_FPS)
default_format = camera.get(cv2.CAP_PROP_FORMAT)
# Get FOURCC code and convert to string
default_fourcc_code = camera.get(cv2.CAP_PROP_FOURCC)
default_fourcc_code_int = int(default_fourcc_code)
default_fourcc = "".join([chr((default_fourcc_code_int >> 8 * i) & 0xFF) for i in range(4)])
# Get FOURCC code and convert to string
default_fourcc_code = camera.get(cv2.CAP_PROP_FOURCC)
default_fourcc_code_int = int(default_fourcc_code)
default_fourcc = "".join(
[chr((default_fourcc_code_int >> 8 * i) & 0xFF) for i in range(4)]
)
camera_info = {
"name": f"OpenCV Camera @ {target}",
"type": "OpenCV",
"id": target,
"backend_api": camera.getBackendName(),
"default_stream_profile": {
"format": default_format,
"fourcc": default_fourcc,
"width": default_width,
"height": default_height,
"fps": default_fps,
},
}
camera_info = {
"name": f"OpenCV Camera @ {target}",
"type": "OpenCV",
"id": target,
"backend_api": camera.getBackendName(),
"default_stream_profile": {
"format": default_format,
"fourcc": default_fourcc,
"width": default_width,
"height": default_height,
"fps": default_fps,
},
}
found_cameras_info.append(camera_info)
found_cameras_info.append(camera_info)
finally:
camera.release()
return found_cameras_info
@@ -496,6 +516,26 @@ class OpenCVCamera(Camera):
self.latest_timestamp = None
self.new_frame_event.clear()
def _cleanup_resources(self) -> None:
"""Stop background reads and release the capture, including after partial setup."""
read_thread = self.thread
videocapture = self.videocapture
try:
self._stop_read_thread()
finally:
self.videocapture = None
try:
if videocapture is not None:
videocapture.release()
finally:
# Releasing the device may unblock a hardware read that outlived
# the first bounded join in _stop_read_thread().
if read_thread is not None and read_thread.is_alive():
read_thread.join(timeout=2.0)
if read_thread.is_alive(): # pragma: no cover
logger.warning(f"{self} read thread remained alive after releasing the capture.")
@check_if_not_connected
def async_read(self, timeout_ms: float = 200) -> NDArray[Any]:
"""
@@ -586,16 +626,6 @@ class OpenCVCamera(Camera):
if not self.is_connected and self.thread is None:
raise DeviceNotConnectedError(f"{self} not connected.")
if self.thread is not None:
self._stop_read_thread()
if self.videocapture is not None:
self.videocapture.release()
self.videocapture = None
with self.frame_lock:
self.latest_frame = None
self.latest_timestamp = None
self.new_frame_event.clear()
self._cleanup_resources()
logger.info(f"{self} disconnected.")
+171 -33
View File
@@ -121,6 +121,9 @@ class RealSenseCamera(Camera):
self.config = config
self.width: int | None = config.width
self.height: int | None = config.height
if config.serial_number_or_name.isdigit():
self.serial_number = config.serial_number_or_name
else:
@@ -131,6 +134,9 @@ class RealSenseCamera(Camera):
self.use_rgb = config.use_rgb
self.use_depth = config.use_depth
self.warmup_s = config.warmup_s
self.exposure: int | None = config.exposure
self.gain: int | None = config.gain
self.white_balance: int | None = config.white_balance
self.rs_pipeline: rs.pipeline | None = None
self.rs_profile: rs.pipeline_profile | None = None
@@ -145,14 +151,23 @@ class RealSenseCamera(Camera):
self.rotation: int | None = get_cv2_rotation(config.rotation)
if self.height and self.width:
self.capture_width, self.capture_height = self.width, self.height
if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
self.capture_width, self.capture_height = self.height, self.width
self.capture_width: int | None = None
self.capture_height: int | None = None
self._reset_connection_settings()
def __str__(self) -> str:
return f"{self.__class__.__name__}({self.serial_number})"
def _reset_connection_settings(self) -> None:
"""Restore settings that may have been auto-detected during a failed connection."""
self.fps = self.config.fps
self.width = self.config.width
self.height = self.config.height
self.warmup_s = self.config.warmup_s
self.capture_width, self.capture_height = self.width, self.height
if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
self.capture_width, self.capture_height = self.height, self.width
@property
def is_connected(self) -> bool:
"""Checks if the camera pipeline is started and streams are active."""
@@ -172,7 +187,8 @@ class RealSenseCamera(Camera):
Raises:
DeviceAlreadyConnectedError: If the camera is already connected.
ValueError: If the configuration is invalid (e.g., missing serial/name, name not unique).
ValueError: If the configuration is invalid, a requested sensor option is unsupported,
or a requested sensor value is invalid.
ConnectionError: If the camera is found but fails to start the pipeline or no RealSense devices are detected at all.
RuntimeError: If the pipeline starts but fails to apply requested settings.
"""
@@ -190,22 +206,31 @@ class RealSenseCamera(Camera):
f"Failed to open {self}.Run `lerobot-find-cameras realsense` to find available cameras."
) from e
self._configure_capture_settings()
self._start_read_thread()
try:
self._configure_capture_settings()
self._configure_sensor_options()
self._start_read_thread()
# NOTE(Steven/Caroline): Enforcing at least one second of warmup as RS cameras need a bit of time before the first read. If we don't wait, the first read from the warmup will raise.
self.warmup_s = max(self.warmup_s, 1)
# NOTE(Steven/Caroline): Enforcing at least one second of warmup as RS cameras need a bit of time before the first read. If we don't wait, the first read from the warmup will raise.
self.warmup_s = max(self.warmup_s, 1)
warmup_read = self.async_read if self.use_rgb else self.async_read_depth
start_time = time.time()
while time.time() - start_time < self.warmup_s:
warmup_read(timeout_ms=self.warmup_s * 1000)
time.sleep(0.1)
with self.frame_lock:
if (self.use_rgb and self.latest_color_frame is None) or (
self.use_depth and self.latest_depth_frame is None
):
raise ConnectionError(f"{self} failed to capture frames during warmup.")
warmup_read = self.async_read if self.use_rgb else self.async_read_depth
start_time = time.time()
while time.time() - start_time < self.warmup_s:
warmup_read(timeout_ms=self.warmup_s * 1000)
time.sleep(0.1)
with self.frame_lock:
if (self.use_rgb and self.latest_color_frame is None) or (
self.use_depth and self.latest_depth_frame is None
):
raise ConnectionError(f"{self} failed to capture frames during warmup.")
except BaseException:
try:
self._cleanup_resources()
except Exception:
logger.exception(f"Failed to fully clean up {self} after connect() failed.")
self._reset_connection_settings()
raise
logger.info(f"{self} connected.")
@@ -339,6 +364,111 @@ class RealSenseCamera(Camera):
self.new_frame_event.clear()
return self._async_read(timeout_ms=10000, read_depth=read_depth)
def _get_color_sensor(self) -> "rs.sensor":
"""Returns the sensor that controls the color stream.
Most RealSense cameras expose "RGB Camera" for color. The D405 has no
separate RGB module — its color stream comes from "Stereo Module".
We try RGB Camera first, then fall back to Stereo Module.
"""
if self.rs_profile is None:
raise RuntimeError(f"{self}: rs_profile must be initialized before use.")
device = self.rs_profile.get_device()
sensors = {s.get_info(rs.camera_info.name): s for s in device.query_sensors()}
for name in ("RGB Camera", "Stereo Module"):
if name in sensors:
return sensors[name]
available = list(sensors.keys())
raise RuntimeError(f"{self}: no color sensor found. Available sensors: {available}")
def _set_sensor_option(self, sensor: "rs.sensor", option: "rs.option", value: float, label: str) -> None:
"""Sets a sensor option, re-raising range errors with actionable diagnostics."""
try:
sensor.set_option(option, value)
except Exception as e:
range_info = ""
try:
option_range = sensor.get_option_range(option)
range_info = (
f" (supported range: min={option_range.min}, max={option_range.max}, "
f"step={option_range.step}, default={option_range.default})"
)
except Exception:
range_info = " (option range unavailable)"
raise ValueError(
f"{self}: failed to set {label} to {value}{range_info}. Original error: {e}"
) from e
def _configure_sensor_options(self) -> None:
"""Applies manual sensor options (exposure, gain, white balance) to the color sensor.
When exposure or gain is set, auto-exposure is disabled first. When white_balance
is set, auto white balance is disabled first. An omitted option is left unchanged,
and configuration is skipped entirely if all options are omitted.
Raises:
ValueError: If the sensor does not support a requested option or a requested
value is invalid. Invalid-value errors include the option name, requested
value, and supported range when available.
"""
if self.exposure is None and self.gain is None and self.white_balance is None:
return
color_sensor = self._get_color_sensor()
requested_options = (
(rs.option.exposure, self.exposure, "exposure"),
(rs.option.gain, self.gain, "gain"),
(rs.option.white_balance, self.white_balance, "white balance"),
)
unsupported_options = [
label
for option, value, label in requested_options
if value is not None and not color_sensor.supports(option)
]
if unsupported_options:
raise ValueError(
f"{self}: color sensor does not support requested manual options: {unsupported_options}."
)
manual_exposure_requested = self.exposure is not None or self.gain is not None
if manual_exposure_requested:
if color_sensor.supports(rs.option.enable_auto_exposure):
self._set_sensor_option(color_sensor, rs.option.enable_auto_exposure, 0, "auto-exposure")
logger.info(f"{self} auto-exposure disabled.")
else:
logger.warning(
f"{self} sensor does not support disabling auto-exposure; "
"applying manual exposure/gain directly."
)
if self.exposure is not None:
self._set_sensor_option(color_sensor, rs.option.exposure, self.exposure, "exposure")
logger.info(f"{self} exposure set to {self.exposure}.")
if self.gain is not None:
self._set_sensor_option(color_sensor, rs.option.gain, self.gain, "gain")
logger.info(f"{self} gain set to {self.gain}.")
if self.white_balance is not None:
if color_sensor.supports(rs.option.enable_auto_white_balance):
self._set_sensor_option(
color_sensor, rs.option.enable_auto_white_balance, 0, "auto white balance"
)
logger.info(f"{self} auto white balance disabled.")
else:
logger.warning(
f"{self} sensor does not support disabling auto white balance; "
"applying manual white balance directly."
)
self._set_sensor_option(
color_sensor, rs.option.white_balance, self.white_balance, "white balance"
)
logger.info(f"{self} white balance set to {self.white_balance}.")
@check_if_not_connected
def read_depth(self, timeout_ms: int = 200) -> NDArray[Any]:
"""
@@ -541,6 +671,27 @@ class RealSenseCamera(Camera):
self.latest_timestamp = None
self.new_frame_event.clear()
def _cleanup_resources(self) -> None:
"""Stop background reads and stop the pipeline, including after partial setup."""
read_thread = self.thread
rs_pipeline = self.rs_pipeline
try:
self._stop_read_thread()
finally:
self.rs_pipeline = None
self.rs_profile = None
try:
if rs_pipeline is not None:
rs_pipeline.stop()
finally:
# Stopping the pipeline may unblock a hardware read that outlived
# the first bounded join in _stop_read_thread().
if read_thread is not None and read_thread.is_alive():
read_thread.join(timeout=2.0)
if read_thread.is_alive(): # pragma: no cover
logger.warning(f"{self} read thread remained alive after stopping the pipeline.")
def _async_read(self, timeout_ms: float, read_depth: bool = False) -> NDArray[Any]:
"""Shared helper for :meth:`async_read`/:meth:`async_read_depth`: return the latest buffered frame."""
if self.thread is None or not self.thread.is_alive():
@@ -684,18 +835,5 @@ class RealSenseCamera(Camera):
f"Attempted to disconnect {self}, but it appears already disconnected."
)
if self.thread is not None:
self._stop_read_thread()
if self.rs_pipeline is not None:
self.rs_pipeline.stop()
self.rs_pipeline = None
self.rs_profile = None
with self.frame_lock:
self.latest_color_frame = None
self.latest_depth_frame = None
self.latest_timestamp = None
self.new_frame_event.clear()
self._cleanup_resources()
logger.info(f"{self} disconnected.")
@@ -46,6 +46,17 @@ class RealSenseCameraConfig(CameraConfig):
use_depth: Whether to enable depth stream. Defaults to False.
rotation: Image rotation setting (0°, 90°, 180°, or 270°). Defaults to no rotation.
warmup_s: Time reading frames before returning from connect (in seconds)
exposure: Manual exposure value for the color sensor. When set, auto-exposure is
disabled and this fixed value is used. Valid ranges are camera-model specific
and reported if the value is rejected. Defaults to None (leave unchanged).
gain: Manual gain value for the color sensor. When set, auto-exposure is disabled
and this fixed gain is used, which also freezes exposure at its current value
when no exposure is configured. Valid ranges are camera-model specific and
reported if the value is rejected. Defaults to None (leave unchanged).
white_balance: Manual white balance value for the color sensor. When set, auto
white balance is disabled and this fixed value is used. Valid ranges are
camera-model specific and reported if the value is rejected. Defaults to None
(leave unchanged).
Note:
- Either name or serial_number must be specified.
@@ -61,6 +72,9 @@ class RealSenseCameraConfig(CameraConfig):
use_depth: bool = False
rotation: Cv2Rotation = Cv2Rotation.NO_ROTATION
warmup_s: int = 1
exposure: int | None = None
gain: int | None = None
white_balance: int | None = None
def __post_init__(self) -> None:
self.color_mode = ColorMode(self.color_mode)
@@ -69,6 +83,18 @@ class RealSenseCameraConfig(CameraConfig):
if not self.use_rgb and not self.use_depth:
raise ValueError("At least one of `use_rgb` or `use_depth` must be enabled.")
manual_color_options = {
"exposure": self.exposure,
"gain": self.gain,
"white_balance": self.white_balance,
}
configured_color_options = [name for name, value in manual_color_options.items() if value is not None]
if configured_color_options and not self.use_rgb:
raise ValueError(
"Manual color sensor options require `use_rgb=True`. "
f"Configured options: {configured_color_options}."
)
values = (self.fps, self.width, self.height)
if any(v is not None for v in values) and any(v is None for v in values):
raise ValueError(
+6
View File
@@ -71,13 +71,19 @@ class DatasetRecordConfig:
# Number of threads per encoder instance. None = auto (codec default).
# Lower values reduce CPU usage, maps to 'lp' (via svtav1-params) for libsvtav1 and 'threads' for h264/hevc..
encoder_threads: int | None = None
# Skip appending the date-time tag to repo_id, keeping the user-provided name as-is
# (e.g. self-managed versioned names intended for a later `lerobot-edit-dataset merge`).
no_stamp: bool = False
def stamp_repo_id(self) -> None:
"""Append a date-time tag to ``repo_id`` so each recording session gets a unique name.
Must be called explicitly at dataset *creation* time — not on resume,
where the existing ``repo_id`` (already stamped) must be preserved.
No-op when ``no_stamp`` is set, preserving a user-managed ``repo_id``.
"""
if self.no_stamp:
return
if self.repo_id:
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
self.repo_id = f"{self.repo_id}_{timestamp}"
-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}
+2 -2
View File
@@ -188,8 +188,8 @@ class LeRobotDatasetMetadata:
def _load_metadata(self):
self.info = load_info(self.root)
check_version_compatibility(self.repo_id, self._version, CODEBASE_VERSION)
self.tasks = load_tasks(self.root)
self.episodes = load_episodes(self.root)
self.tasks = load_tasks(self.root) if self.total_tasks > 0 else None
self.episodes = load_episodes(self.root) if self.total_episodes > 0 else None
self.stats = load_stats(self.root)
def ensure_readable(self) -> None:
-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(
+1 -9
View File
@@ -560,13 +560,7 @@ class RoboCasaEnv(EnvConfig):
kwargs["split"] = self.split
return kwargs
def create_envs(
self,
n_envs: int,
use_async_envs: bool = False,
terminate_on_success: bool = True,
horizon: int | None = None,
):
def create_envs(self, n_envs: int, use_async_envs: bool = False):
from .robocasa import create_robocasa_envs
if self.task is None:
@@ -580,8 +574,6 @@ 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,
)
+6 -1
View File
@@ -384,7 +384,12 @@ class LiberoEnv(gym.Env):
def close(self):
if self._env is not None:
self._env.close()
try:
self._env.close()
finally:
# LIBERO deletes its inner env on close, so this wrapper must
# be recreated before the next reset.
self._env = None
def _make_env_fns(
+11 -33
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 # 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)
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)
ACTION_LOW = -1.0
ACTION_HIGH = 1.0
@@ -101,15 +101,14 @@ 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 (openpi / robocasa.utils.env_utils.convert_action order):
ee_pos(3) + ee_rot(3) + gripper(1) + base_motion(4) + control_mode(1)
Layout: base_motion(4) + control_mode(1) + ee_pos(3) + ee_rot(3) + gripper(1)
"""
return {
"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],
"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],
}
@@ -137,16 +136,9 @@ 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
@@ -218,16 +210,12 @@ 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()
@@ -242,14 +230,12 @@ 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.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.end_effector_position_relative", np.zeros(3)),
raw_obs.get("state.end_effector_rotation_relative", np.zeros(4)),
raw_obs.get("state.gripper_qpos", np.zeros(2)),
],
axis=-1,
@@ -294,7 +280,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 and self.terminate_on_success)
terminated = done or is_success
info.update({"task": self.task, "done": done, "is_success": is_success})
observation = self._format_raw_obs(raw_obs)
@@ -327,8 +313,6 @@ 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.
@@ -351,8 +335,6 @@ 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)]
@@ -366,8 +348,6 @@ 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.
@@ -429,8 +409,6 @@ 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:
+3 -1
View File
@@ -384,7 +384,9 @@ class RoboTwinEnv(gym.Env):
self._env: Any | None = None # deferred — created on first reset() inside worker
self._step_count: int = 0
self._black_frame = np.zeros((self.observation_height, self.observation_width, 3), dtype=np.uint8)
self._black_frame: np.ndarray = np.zeros(
(self.observation_height, self.observation_width, 3), dtype=np.uint8
)
image_spaces = {
cam: spaces.Box(
+1 -1
View File
@@ -373,7 +373,7 @@ class VLABenchEnv(gym.Env):
if action.shape[0] != 7:
# Unknown layout — fall back to zero-pad so the sim doesn't crash.
padded = np.zeros(ctrl_dim, dtype=np.float64)
padded: np.ndarray = np.zeros(ctrl_dim, dtype=np.float64)
padded[: min(action.shape[0], ctrl_dim)] = action[:ctrl_dim]
return padded
+18
View File
@@ -122,6 +122,9 @@ MODEL_ENCODING_TABLE = {
"xm430-w350": X_SERIES_ENCODINGS_TABLE,
"xm540-w270": X_SERIES_ENCODINGS_TABLE,
"xc430-w150": X_SERIES_ENCODINGS_TABLE,
"xh540-w150": X_SERIES_ENCODINGS_TABLE,
"xc330-t288": X_SERIES_ENCODINGS_TABLE,
"xc330-t181": X_SERIES_ENCODINGS_TABLE,
}
# {model: model_resolution}
@@ -134,6 +137,9 @@ MODEL_RESOLUTION = {
"xm430-w350": 4096,
"xm540-w270": 4096,
"xc430-w150": 4096,
"xh540-w150": 4096,
"xc330-t288": 4096,
"xc330-t181": 4096,
}
# {model: model_number}
@@ -145,6 +151,9 @@ MODEL_NUMBER_TABLE = {
"xm430-w350": 1020,
"xm540-w270": 1120,
"xc430-w150": 1070,
"xh540-w150": 1110,
"xc330-t288": 1220,
"xc330-t181": 1210,
}
# {model: available_operating_modes}
@@ -156,6 +165,9 @@ MODEL_OPERATING_MODES = {
"xm430-w350": [0, 1, 3, 4, 5, 16],
"xm540-w270": [0, 1, 3, 4, 5, 16],
"xc430-w150": [1, 3, 4, 16],
"xh540-w150": [0, 1, 3, 4, 5, 16],
"xc330-t288": [0, 1, 3, 4, 5, 16],
"xc330-t181": [0, 1, 3, 4, 5, 16],
}
MODEL_CONTROL_TABLE = {
@@ -166,6 +178,9 @@ MODEL_CONTROL_TABLE = {
"xm430-w350": X_SERIES_CONTROL_TABLE,
"xm540-w270": X_SERIES_CONTROL_TABLE,
"xc430-w150": X_SERIES_CONTROL_TABLE,
"xh540-w150": X_SERIES_CONTROL_TABLE,
"xc330-t288": X_SERIES_CONTROL_TABLE,
"xc330-t181": X_SERIES_CONTROL_TABLE,
}
MODEL_BAUDRATE_TABLE = {
@@ -176,6 +191,9 @@ MODEL_BAUDRATE_TABLE = {
"xm430-w350": X_SERIES_BAUDRATE_TABLE,
"xm540-w270": X_SERIES_BAUDRATE_TABLE,
"xc430-w150": X_SERIES_BAUDRATE_TABLE,
"xh540-w150": X_SERIES_BAUDRATE_TABLE,
"xc330-t288": X_SERIES_BAUDRATE_TABLE,
"xc330-t181": X_SERIES_BAUDRATE_TABLE,
}
AVAILABLE_BAUDRATES = [
+18 -3
View File
@@ -44,12 +44,19 @@ from lerobot.utils.constants import (
POLICY_PREPROCESSOR_DEFAULT_NAME,
)
from lerobot.utils.feature_utils import dataset_to_policy_features
from lerobot.utils.import_utils import _peft_available, require_package
from .evo1.configuration_evo1 import Evo1Config
from .groot.configuration_groot import GrootConfig
from .pretrained import PreTrainedPolicy
from .utils import validate_visual_features_consistency
if TYPE_CHECKING or _peft_available:
from peft import PeftConfig, PeftModel
else:
PeftConfig = None
PeftModel = None
def _reconnect_relative_absolute_steps(
preprocessor: PolicyProcessorPipeline, postprocessor: PolicyProcessorPipeline
@@ -334,12 +341,15 @@ def make_policy(
# Load a pretrained PEFT model on top of the policy. The pretrained path points to the folder/repo
# of the adapter and the adapter's config contains the path to the base policy. So we need the
# adapter config first, then load the correct policy and then apply PEFT.
from peft import PeftConfig, PeftModel
require_package("peft", extra="peft")
logging.info("Loading policy's PEFT adapter.")
peft_pretrained_path = str(cfg.pretrained_path)
peft_config = PeftConfig.from_pretrained(peft_pretrained_path)
peft_config = PeftConfig.from_pretrained(
peft_pretrained_path,
revision=cfg.pretrained_revision,
)
kwargs["pretrained_name_or_path"] = peft_config.base_model_name_or_path
if not kwargs["pretrained_name_or_path"]:
@@ -350,9 +360,14 @@ def make_policy(
"the adapter was trained."
)
kwargs["revision"] = peft_config.revision
policy = policy_cls.from_pretrained(**kwargs)
policy = PeftModel.from_pretrained(
policy, peft_pretrained_path, config=peft_config, is_trainable=True
policy,
peft_pretrained_path,
config=peft_config,
revision=cfg.pretrained_revision,
is_trainable=True,
)
else:
@@ -37,13 +37,19 @@ 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
@@ -43,11 +43,22 @@ from torch.distributions import Beta
from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.utils.constants import ACTION
from lerobot.utils.import_utils import _scipy_available, _transformers_available, require_package
from lerobot.utils.import_utils import (
_peft_available,
_scipy_available,
_transformers_available,
require_package,
)
from ..rtc.modeling_rtc import RTCProcessor
from .configuration_molmoact2 import MolmoAct2Config
if TYPE_CHECKING or _peft_available:
from peft import LoraConfig, get_peft_model
else:
LoraConfig = None
get_peft_model = None
logger = logging.getLogger(__name__)
@@ -1731,13 +1742,11 @@ class MolmoAct2Policy(PreTrainedPolicy):
def _build_inner_lora_config(self):
require_package("peft", extra="molmoact2")
from peft import LoraConfig
return LoraConfig(**self._get_inner_peft_targets())
def _apply_lora_adapters(self) -> None:
require_package("peft", extra="molmoact2")
from peft import get_peft_model
peft_config = self._build_inner_lora_config()
self._validate_peft_config(peft_config)
+13 -5
View File
@@ -34,14 +34,22 @@ from lerobot.configs import PreTrainedConfig
from lerobot.configs.train import TrainPipelineConfig
from lerobot.utils.device_utils import resolve_safetensors_device
from lerobot.utils.hub import HubMixin
from lerobot.utils.import_utils import _peft_available, require_package
from .utils import log_model_loading_keys
T = TypeVar("T", bound="PreTrainedPolicy")
if TYPE_CHECKING or _peft_available:
from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType, get_peft_model
else:
PEFT_TYPE_TO_CONFIG_MAPPING = None
PeftType = None
get_peft_model = None
if TYPE_CHECKING:
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
T = TypeVar("T", bound="PreTrainedPolicy")
def _build_card_context(
cfg: TrainPipelineConfig | None,
@@ -384,7 +392,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
peft_cli_overrides: Optional dict of CLI overrides (method_type, target_modules, r, etc.)
These are merged with policy defaults to build the final config.
"""
from peft import get_peft_model
require_package("peft", extra="peft")
# If user provided a complete config, use it directly (with overrides)
if peft_config is not None:
@@ -455,7 +463,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
Returns:
Preprocessed dict with renamed keys and init_type mapped to method-specific key.
"""
from peft import PeftType
require_package("peft", extra="peft")
cli_overrides = cli_overrides.copy()
@@ -480,7 +488,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
def _build_peft_config(self, cli_overrides: dict):
"""Build a PEFT config from policy defaults and CLI overrides."""
from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType
require_package("peft", extra="peft")
# Determine PEFT method type (default to LORA)
method_type_str = cli_overrides.get("method_type") or "lora"
@@ -507,7 +515,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
def _apply_peft_cli_overrides(self, peft_config, cli_overrides: dict):
"""Apply CLI overrides to an existing PEFT config."""
from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType
require_package("peft", extra="peft")
# Get method type from existing config or CLI override
method_type_str = cli_overrides.get("method_type")
+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,10 +132,20 @@ class MapDeltaActionToRobotActionStep(RobotActionProcessorStep):
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
for axis in ["x", "y", "z", "gripper"]:
for axis in ["x", "y", "z"]:
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"]:
for feat in [
"enabled",
"target_x",
"target_y",
"target_z",
"target_wx",
"target_wy",
"target_wz",
"gripper_vel",
]:
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]]:
+2 -4
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
from lerobot.utils.process import ProcessSignalHandler, ensure_multiprocessing_start_method
from lerobot.utils.random_utils import set_seed
from lerobot.utils.robot_utils import precise_sleep
from lerobot.utils.transition import (
@@ -124,9 +124,7 @@ def actor_cli(cfg: TrainRLServerPipelineConfig):
cfg.validate()
display_pid = False
if not use_threads(cfg):
import torch.multiprocessing as mp
mp.set_start_method("spawn")
ensure_multiprocessing_start_method(cfg.policy.concurrency.multiprocessing_context)
display_pid = True
# Create logs directory to ensure it exists
+2 -4
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
from lerobot.utils.process import ProcessSignalHandler, ensure_multiprocessing_start_method
from lerobot.utils.random_utils import set_seed
from lerobot.utils.utils import (
format_big_number,
@@ -123,9 +123,7 @@ 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):
import torch.multiprocessing as mp
mp.set_start_method("spawn")
ensure_multiprocessing_start_method(cfg.policy.concurrency.multiprocessing_context)
# Use the job_name from the config
train(
+1 -3
View File
@@ -21,8 +21,6 @@ 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
@@ -120,7 +118,7 @@ def ensure_safe_goal_position(
}
if warnings_dict:
logger.warning(
logging.warning(
"Relative goal position magnitude had to be clamped to be safe.\n"
f"{pformat(warnings_dict, indent=4)}"
)
+21 -3
View File
@@ -24,6 +24,7 @@ from __future__ import annotations
import logging
from dataclasses import dataclass, field
from threading import Event
from typing import TYPE_CHECKING
import torch
@@ -47,6 +48,7 @@ from lerobot.processor.relative_action_processor import RelativeActionsProcessor
from lerobot.robots import make_robot_from_config
from lerobot.teleoperators import Teleoperator, make_teleoperator_from_config
from lerobot.utils.feature_utils import combine_feature_dicts, hw_to_dataset_features
from lerobot.utils.import_utils import _peft_available, require_package
from .configs import BaseStrategyConfig, DAggerStrategyConfig, RolloutConfig
from .inference import (
@@ -57,6 +59,12 @@ from .inference import (
)
from .robot_wrapper import ThreadSafeRobot
if TYPE_CHECKING or _peft_available:
from peft import PeftConfig, PeftModel
else:
PeftConfig = None
PeftModel = None
logger = logging.getLogger(__name__)
@@ -171,7 +179,7 @@ def _load_pretrained_policy(policy_config: PreTrainedConfig) -> PreTrainedPolicy
revision=pretrained_revision,
)
from peft import PeftConfig, PeftModel
require_package("peft", extra="peft")
peft_path = policy_config.pretrained_path
peft_config = PeftConfig.from_pretrained(peft_path, revision=pretrained_revision)
@@ -294,12 +302,22 @@ def build_rollout_context(
# ``observation_features`` values are either a tuple (camera shape) or the
# ``float`` type itself used as a sentinel for scalar motor features —
# see ``dict[str, type | tuple]`` annotation on ``Robot.observation_features``.
# Keep cameras (tuple) plus both joint-position (.pos) and base-velocity (.vel)
# scalar state features. LeKiwi's observation.state is 9-dim (6 arm .pos +
# x/y/theta.vel) and the policy was trained/normalized on all 9; the old .pos-only
# filter fed a 6-dim state into a 9-dim normalizer → RuntimeError (size 6 vs 9).
# Pure-arm robots have no .vel state keys, so this is a no-op for them.
observation_features_hw = {
k: v
for k, v in all_obs_features.items()
if isinstance(v, tuple) or (v is float and k.endswith(".pos"))
if isinstance(v, tuple) or (v is float and k.endswith((".pos", ".vel")))
}
action_features_hw = {k: v for k, v in robot.action_features.items() if k.endswith(".pos")}
# Keep both joint-position (.pos) and base-velocity (.vel) action features so
# mobile manipulators command the base too (e.g. LeKiwi: 6 arm .pos +
# x/y/theta.vel = 9-dim action). Pure-arm robots have no .vel keys, so this is
# a no-op for them. Without the .vel keys the base velocities are silently
# dropped from dataset_features[ACTION]/ordered_action_keys and the base never moves.
action_features_hw = {k: v for k, v in robot.action_features.items() if k.endswith((".pos", ".vel"))}
# The action side is always needed: sync inference reads action names from
# ``dataset_features[ACTION]`` to map policy tensors back to robot actions.
-38
View File
@@ -1,38 +0,0 @@
# 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
@@ -1,165 +0,0 @@
# 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
@@ -1,349 +0,0 @@
# 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
@@ -1,39 +0,0 @@
# 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
@@ -1,406 +0,0 @@
# 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)
@@ -94,6 +94,8 @@ from lerobot.datasets.video_utils import concatenate_video_files, get_video_dura
from lerobot.utils.constants import HF_LEROBOT_HOME
from lerobot.utils.utils import flatten_dict, init_logging
logger = logging.getLogger(__name__)
V21 = "v2.1"
V30 = "v3.0"
@@ -476,11 +478,11 @@ def convert_dataset(
# First check if the dataset already has a v3.0 version
if root is None and not force_conversion:
try:
print("Trying to download v3.0 version of the dataset from the hub...")
logger.info("Trying to download v3.0 version of the dataset from the hub...")
snapshot_download(repo_id, repo_type="dataset", revision=V30, local_dir=HF_LEROBOT_HOME / repo_id)
return
except Exception:
print("Dataset does not have an uploaded v3.0 version. Continuing with conversion.")
logger.info("Dataset does not have an uploaded v3.0 version. Continuing with conversion.")
# Set root based on whether local dataset path is provided
use_local_dataset = False
@@ -488,7 +490,7 @@ def convert_dataset(
if root.exists():
validate_local_dataset_version(root)
use_local_dataset = True
print(f"Using local dataset at {root}")
logger.info(f"Using local dataset at {root}")
old_root = root.parent / f"{root.name}_old"
new_root = root.parent / f"{root.name}_v30"
@@ -523,7 +525,7 @@ def convert_dataset(
try:
hub_api.delete_tag(repo_id, tag=CODEBASE_VERSION, repo_type="dataset")
except (HTTPError, RevisionNotFoundError) as e:
print(f"tag={CODEBASE_VERSION} probably doesn't exist. Skipping exception ({e})")
logger.warning(f"tag={CODEBASE_VERSION} probably doesn't exist. Skipping exception ({e})")
pass
hub_api.delete_files(
delete_patterns=["data/chunk*/episode_*", "meta/*.jsonl", "videos/chunk*"],
+6 -7
View File
@@ -154,14 +154,14 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
repo_id = cfg.new_repo_id or cfg.repo_id
commit_message = cfg.push_commit_message or "Add steerable annotations (lerobot-annotate)"
api = HfApi()
print(f"[lerobot-annotate] creating/locating dataset repo {repo_id}...", flush=True)
logger.info(f"[lerobot-annotate] creating/locating dataset repo {repo_id}...")
api.create_repo(
repo_id=repo_id,
repo_type="dataset",
private=cfg.push_private,
exist_ok=True,
)
print(f"[lerobot-annotate] uploading {root} -> {repo_id}...", flush=True)
logger.info(f"[lerobot-annotate] uploading {root} -> {repo_id}...")
commit_info = api.upload_folder(
folder_path=str(root),
repo_id=repo_id,
@@ -172,7 +172,7 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
# at the source dataset; a fresh card is generated below instead.
ignore_patterns=[".annotate_staging/**", "**/.DS_Store", "README.md"],
)
print(f"[lerobot-annotate] uploaded to https://huggingface.co/datasets/{repo_id}", flush=True)
logger.info(f"[lerobot-annotate] uploaded to https://huggingface.co/datasets/{repo_id}")
dataset_info = load_info(root)
card = create_lerobot_dataset_card(dataset_info=dataset_info, license="apache-2.0", repo_id=repo_id)
@@ -200,14 +200,13 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
with suppress(RevisionNotFoundError):
api.delete_tag(repo_id, tag=version_tag, repo_type="dataset")
api.create_tag(**tag_kwargs)
print(f"[lerobot-annotate] tagged {repo_id} as {version_tag}", flush=True)
logger.info(f"[lerobot-annotate] tagged {repo_id} as {version_tag}")
except Exception as exc: # noqa: BLE001
print(
logger.warning(
f"[lerobot-annotate] WARNING: could not create tag {version_tag!r} on {repo_id}: {exc}. "
"Dataset is uploaded but ``LeRobotDataset`` won't be able to load it until it's tagged. "
"Run: from huggingface_hub import HfApi; "
f"HfApi().create_tag({repo_id!r}, tag={version_tag!r}, repo_type='dataset', exist_ok=True)",
flush=True,
f"HfApi().create_tag({repo_id!r}, tag={version_tag!r}, repo_type='dataset', exist_ok=True)"
)
+3 -1
View File
@@ -89,6 +89,8 @@ from lerobot.datasets import LeRobotDataset
from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS
from lerobot.utils.utils import init_logging
logger = logging.getLogger(__name__)
DEFAULT_FOXGLOVE_PORT = 8765
DEFAULT_RERUN_PORT = 9090
@@ -299,7 +301,7 @@ def visualize_dataset(
while True:
time.sleep(1)
except KeyboardInterrupt:
print("Ctrl-C received. Exiting.")
logger.info("Ctrl-C received. Exiting.")
def main():
+20 -52
View File
@@ -62,7 +62,7 @@ from dataclasses import asdict
from functools import partial
from pathlib import Path
from pprint import pformat
from typing import Any, TypedDict
from typing import TYPE_CHECKING, Any, TypedDict
import einops
import gymnasium as gym
@@ -87,26 +87,21 @@ from lerobot.processor import PolicyProcessorPipeline
from lerobot.types import PolicyAction
from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_IMAGES, OBS_STR, REWARD
from lerobot.utils.device_utils import get_safe_torch_device
from lerobot.utils.import_utils import register_third_party_plugins
from lerobot.utils.import_utils import _peft_available, register_third_party_plugins, require_package
from lerobot.utils.io_utils import write_video
from lerobot.utils.random_utils import set_seed
from lerobot.utils.utils import (
init_logging,
inside_slurm,
)
from lerobot.utils.video_annotation import annotate_frame
if TYPE_CHECKING or _peft_available:
from peft import PeftModel
else:
PeftModel = None
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])
logger = logging.getLogger(__name__)
def _env_features_to_dataset_features(env_features: dict) -> dict:
@@ -457,13 +452,11 @@ def eval_policy(
exc = ValueError(
f"Policy of type 'PreTrainedPolicy' is expected, but type '{type(policy)}' was provided."
)
try:
from peft import PeftModel
if not isinstance(policy, PeftModel):
raise exc
except ImportError:
raise exc from None
if not _peft_available:
raise exc
require_package("peft", extra="peft")
if not isinstance(policy, PeftModel):
raise exc
start = time.time()
# Preserve the mode for direct callers. eval_policy_all scopes the mode
@@ -490,36 +483,11 @@ def eval_policy(
return
n_to_render_now = min(max_episodes_rendered - n_episodes_rendered, env.num_envs)
if isinstance(env, gym.vector.SyncVectorEnv):
frames = np.stack([env.envs[i].render() for i in range(n_to_render_now)]) # noqa: B023
ep_frames.append(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).
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))
ep_frames.append(np.stack(env.call("render")[:n_to_render_now]))
if max_episodes_rendered > 0:
video_paths: list[str] = []
@@ -596,7 +564,7 @@ def eval_policy(
if seeds:
all_seeds.extend(seeds)
else:
all_seeds.append(None)
all_seeds.extend([None] * env.num_envs)
# FIXME: episode_data is either None or it doesn't exist
if return_episode_data:
@@ -834,13 +802,13 @@ def eval_main(cfg: EvalPipelineConfig):
recording_repo_id=cfg.eval.recording_repo_id,
recording_private=cfg.eval.recording_private,
)
print("Overall Aggregated Metrics:")
print(info["overall"])
logger.info("Overall Aggregated Metrics:")
logger.info(info["overall"])
# Print per-suite stats
for task_group, task_group_info in info.items():
print(f"\nAggregated Metrics for {task_group}:")
print(task_group_info)
logger.info(f"\nAggregated Metrics for {task_group}:")
logger.info(task_group_info)
# Close all vec envs
close_envs(envs)
@@ -40,6 +40,7 @@ from PIL import Image
from lerobot.cameras import ColorMode
from lerobot.cameras.opencv import OpenCVCamera, OpenCVCameraConfig
from lerobot.cameras.realsense import RealSenseCamera, RealSenseCameraConfig
from lerobot.utils.utils import init_logging
logger = logging.getLogger(__name__)
@@ -285,6 +286,8 @@ def save_images_from_all_cameras(
def main():
init_logging()
parser = argparse.ArgumentParser(
description="Unified camera utility script for listing cameras and capturing images."
)
+4 -63
View File
@@ -151,7 +151,6 @@ Usage examples
"""
import logging
import sys
from lerobot.cameras.opencv import OpenCVCameraConfig # noqa: F401
from lerobot.cameras.realsense import RealSenseCameraConfig # noqa: F401
@@ -166,6 +165,7 @@ from lerobot.robots import ( # noqa: F401
earthrover_mini_plus,
hope_jr,
koch_follower,
lekiwi,
omx_follower,
openarm_follower,
reachy2,
@@ -242,69 +242,10 @@ def rollout(cfg: RolloutConfig):
logger.info("Rollout finished")
_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.
"""
def main():
"""CLI entry point for ``lerobot-rollout``."""
register_third_party_plugins()
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
rollout()
if __name__ == "__main__":
+24 -17
View File
@@ -22,7 +22,8 @@ import dataclasses
import logging
import sys
import time
from contextlib import nullcontext
from collections.abc import Iterator
from contextlib import contextmanager, nullcontext
from pprint import pformat
from typing import TYPE_CHECKING, Any
@@ -57,7 +58,7 @@ from lerobot.optim.factory import make_optimizer_and_scheduler
from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors
from lerobot.rewards import make_reward_pre_post_processors
from lerobot.utils.collate import lerobot_collate_fn
from lerobot.utils.import_utils import register_third_party_plugins
from lerobot.utils.import_utils import _peft_available, register_third_party_plugins, require_package
from lerobot.utils.logging_utils import AverageMeter, MetricsTracker
from lerobot.utils.random_utils import set_seed
from lerobot.utils.utils import (
@@ -68,9 +69,28 @@ from lerobot.utils.utils import (
inside_slurm,
)
if TYPE_CHECKING or _peft_available:
from peft import PeftModel
else:
PeftModel = None
from .lerobot_eval import eval_policy_all
@contextmanager
def _make_eval_envs(cfg: TrainPipelineConfig) -> Iterator[dict[str, dict[int, Any]]]:
"""Create evaluation environments for one run and always dispose of them."""
envs = make_env(
cfg.env,
n_envs=cfg.eval.batch_size,
use_async_envs=cfg.eval.use_async_envs,
)
try:
yield envs
finally:
close_envs(envs)
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
@@ -207,8 +227,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if cfg.job.is_remote:
return submit_to_hf(cfg)
from lerobot.utils.import_utils import require_package
require_package("accelerate", extra="training")
from accelerate import Accelerator
from accelerate.utils import DistributedDataParallelKwargs, DistributedType
@@ -277,14 +295,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if not is_main_process:
dataset, eval_dataset = make_train_eval_datasets(cfg)
# Create environment used for evaluating checkpoints during training on simulation data.
# On real-world data, no need to create an environment as evaluations are done outside train.py,
# using the eval.py instead, with gym_dora environment and dora-rs.
eval_env = None
if cfg.env_eval_freq > 0 and cfg.env is not None and is_main_process:
logging.info("Creating env")
eval_env = make_env(cfg.env, n_envs=cfg.eval.batch_size, use_async_envs=cfg.eval.use_async_envs)
if cfg.is_reward_model_training:
if is_main_process:
logging.info("Creating reward model")
@@ -312,7 +322,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if cfg.peft is not None:
if cfg.is_reward_model_training:
raise ValueError("PEFT is only supported for policy training. ")
from peft import PeftModel
require_package("peft", extra="peft")
if isinstance(policy, PeftModel):
logging.info("PEFT adapter already loaded from checkpoint, skipping wrap_with_peft.")
@@ -692,7 +702,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if is_main_process:
step_id = get_step_identifier(step, cfg.steps)
logging.info(f"Eval policy at step {step}")
with torch.no_grad(), accelerator.autocast():
with _make_eval_envs(cfg) as eval_env, torch.no_grad(), accelerator.autocast():
eval_info = eval_policy_all(
envs=eval_env, # dict[suite][task_id] -> vec_env
policy=accelerator.unwrap_model(policy),
@@ -740,9 +750,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if is_main_process:
progbar.close()
if eval_env:
close_envs(eval_env)
is_fsdp = accelerator.distributed_type == DistributedType.FSDP
model_state_dict = accelerator.get_state_dict(policy) if is_fsdp else None
if is_main_process:
+58 -56
View File
@@ -45,6 +45,7 @@ lerobot-train-tokenizer \
"""
import json
import logging
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING
@@ -63,6 +64,9 @@ else:
from lerobot.configs import NormalizationMode, parser
from lerobot.datasets import LeRobotDataset
from lerobot.utils.constants import ACTION, OBS_STATE
from lerobot.utils.utils import init_logging
logger = logging.getLogger(__name__)
@dataclass
@@ -274,11 +278,8 @@ def process_episode(args):
return action_chunks
except Exception as e:
print(f"Error processing episode {ep_idx}: {e}")
import traceback
traceback.print_exc()
except Exception:
logger.exception("Error processing episode %s", ep_idx)
return None
@@ -300,10 +301,10 @@ def train_fast_tokenizer(
Returns:
Trained FAST tokenizer
"""
print(f"Training FAST tokenizer on {len(action_chunks)} action chunks...")
print(f"Action chunk shape: {action_chunks.shape}")
print(f"Vocab size: {vocab_size}")
print(f"DCT scale: {scale}")
logger.info(f"Training FAST tokenizer on {len(action_chunks)} action chunks...")
logger.info(f"Action chunk shape: {action_chunks.shape}")
logger.info(f"Vocab size: {vocab_size}")
logger.info(f"DCT scale: {scale}")
# download the tokenizer source code (not pretrained weights)
# we'll train a new tokenizer on our own data
@@ -314,7 +315,7 @@ def train_fast_tokenizer(
# train the new tokenizer on our action data using .fit()
# this trains the BPE tokenizer on DCT coefficients
print("Training new tokenizer (this may take a few minutes)...")
logger.info("Training new tokenizer (this may take a few minutes)...")
tokenizer = base_tokenizer.fit(
action_data_list,
scale=scale,
@@ -322,21 +323,21 @@ def train_fast_tokenizer(
time_horizon=action_chunks.shape[1], # action_horizon
action_dim=action_chunks.shape[2], # encoded dimensions
)
print("✓ Tokenizer training complete!")
logger.info("✓ Tokenizer training complete!")
# validate it works
sample_chunk = action_chunks[0]
encoded = tokenizer(sample_chunk[None])[0]
if isinstance(encoded, list):
encoded = np.array(encoded)
print(f"Sample encoding: {len(encoded)} tokens for chunk shape {sample_chunk.shape}")
logger.info(f"Sample encoding: {len(encoded)} tokens for chunk shape {sample_chunk.shape}")
return tokenizer
def compute_compression_stats(tokenizer, action_chunks: np.ndarray):
"""Compute compression statistics."""
print("\nComputing compression statistics...")
logger.info("\nComputing compression statistics...")
# sample for stats (use max 1000 chunks for speed)
sample_size = min(1000, len(action_chunks))
@@ -366,12 +367,12 @@ def compute_compression_stats(tokenizer, action_chunks: np.ndarray):
"max_token_length": float(np.max(token_lengths)),
}
print("Compression Statistics:")
print(f" Average compression ratio: {stats['compression_ratio']:.2f}x")
print(f" Mean token length: {stats['mean_token_length']:.1f}")
print(f" P99 token length: {stats['p99_token_length']:.0f}")
print(f" Min token length: {stats['min_token_length']:.0f}")
print(f" Max token length: {stats['max_token_length']:.0f}")
logger.info("Compression Statistics:")
logger.info(f" Average compression ratio: {stats['compression_ratio']:.2f}x")
logger.info(f" Mean token length: {stats['mean_token_length']:.1f}")
logger.info(f" P99 token length: {stats['p99_token_length']:.0f}")
logger.info(f" Min token length: {stats['min_token_length']:.0f}")
logger.info(f" Max token length: {stats['max_token_length']:.0f}")
return stats
@@ -385,9 +386,9 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
cfg: TokenizerTrainingConfig dataclass with all configuration parameters
"""
# load dataset
print(f"Loading dataset: {cfg.repo_id}")
logger.info(f"Loading dataset: {cfg.repo_id}")
dataset = LeRobotDataset(repo_id=cfg.repo_id, root=cfg.root)
print(f"Dataset loaded: {dataset.num_episodes} episodes, {dataset.num_frames} frames")
logger.info(f"Dataset loaded: {dataset.num_episodes} episodes, {dataset.num_frames} frames")
# parse normalization mode
try:
@@ -397,7 +398,7 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
f"Invalid normalization_mode: {cfg.normalization_mode}. "
f"Must be one of: {', '.join([m.value for m in NormalizationMode])}"
) from err
print(f"Normalization mode: {norm_mode.value}")
logger.info(f"Normalization mode: {norm_mode.value}")
# parse encoded dimensions
encoded_dim_ranges = []
@@ -406,38 +407,38 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
encoded_dim_ranges.append((start, end))
total_encoded_dims = sum(end - start for start, end in encoded_dim_ranges)
print(f"Encoding {total_encoded_dims} dimensions: {cfg.encoded_dims}")
logger.info(f"Encoding {total_encoded_dims} dimensions: {cfg.encoded_dims}")
# parse relative dimensions
relative_dim_list = None
if cfg.relative_dims is not None and cfg.relative_dims.strip():
relative_dim_list = [int(d.strip()) for d in cfg.relative_dims.split(",")]
print(f"Relative dimensions: {relative_dim_list}")
logger.info(f"Relative dimensions: {relative_dim_list}")
else:
print("No relative dimensions specified")
logger.info("No relative dimensions specified")
print(f"Use relative transform: {cfg.use_relative_transform}")
logger.info(f"Use relative transform: {cfg.use_relative_transform}")
if cfg.use_relative_transform and (relative_dim_list is None or len(relative_dim_list) == 0):
print(
logger.warning(
"Warning: use_relative_transform=True but no relative_dims specified. "
"No relative transform will be applied."
)
print(f"Action horizon: {cfg.action_horizon}")
print(f"State key: {cfg.state_key}")
logger.info(f"Action horizon: {cfg.action_horizon}")
logger.info(f"State key: {cfg.state_key}")
# determine episodes to process
num_episodes = dataset.num_episodes
if cfg.max_episodes is not None:
num_episodes = min(cfg.max_episodes, num_episodes)
print(f"Processing {num_episodes} episodes...")
logger.info(f"Processing {num_episodes} episodes...")
# process episodes sequentially (to avoid pickling issues with dataset)
all_chunks = []
for ep_idx in range(num_episodes):
if ep_idx % 10 == 0:
print(f" Processing episode {ep_idx}/{num_episodes}...")
logger.info(f" Processing episode {ep_idx}/{num_episodes}...")
chunks = process_episode(
(
@@ -455,19 +456,19 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
# concatenate all chunks
all_chunks = np.concatenate(all_chunks, axis=0)
print(f"Collected {len(all_chunks)} action chunks")
logger.info(f"Collected {len(all_chunks)} action chunks")
# extract only encoded dimensions FIRST (before normalization)
encoded_chunks = []
for start, end in encoded_dim_ranges:
encoded_chunks.append(all_chunks[:, :, start:end])
encoded_chunks = np.concatenate(encoded_chunks, axis=-1) # [N, H, D_encoded]
print(f"Extracted {encoded_chunks.shape[-1]} encoded dimensions")
logger.info(f"Extracted {encoded_chunks.shape[-1]} encoded dimensions")
# apply normalization to encoded dimensions
print("\nBefore normalization - overall stats:")
print(f" Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
print(f" Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
logger.info("\nBefore normalization - overall stats:")
logger.info(f" Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
logger.info(f" Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
# get normalization stats from dataset
norm_stats = dataset.meta.stats
@@ -489,9 +490,9 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
encoded_stats[stat_name] = stat_array[encoded_dim_indices]
if encoded_stats:
print(f"\nNormalization stats for encoded dimensions (mode: {norm_mode.value}):")
logger.info(f"\nNormalization stats for encoded dimensions (mode: {norm_mode.value}):")
for stat_name, stat_values in encoded_stats.items():
print(
logger.info(
f" {stat_name}: shape={stat_values.shape}, "
f"range=[{np.min(stat_values):.4f}, {np.max(stat_values):.4f}]"
)
@@ -499,27 +500,27 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
# apply normalization based on mode
try:
encoded_chunks = apply_normalization(encoded_chunks, encoded_stats, norm_mode, eps=1e-8)
print(f"\nApplied {norm_mode.value} normalization")
logger.info(f"\nApplied {norm_mode.value} normalization")
except ValueError as e:
print(f"Warning: {e}. Using raw actions without normalization.")
logger.warning(f"Warning: {e}. Using raw actions without normalization.")
print("\nAfter normalization - overall stats:")
print(f" Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
print(f" Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
logger.info("\nAfter normalization - overall stats:")
logger.info(f" Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
logger.info(f" Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
print("\nPer-dimension stats (after normalization):")
logger.info("\nPer-dimension stats (after normalization):")
for d in range(encoded_chunks.shape[-1]):
dim_data = encoded_chunks[:, :, d]
print(
logger.info(
f" Dim {d}: min={np.min(dim_data):7.4f}, max={np.max(dim_data):7.4f}, "
f"mean={np.mean(dim_data):7.4f}, std={np.std(dim_data):7.4f}"
)
else:
print("Warning: Could not extract stats for encoded dimensions, using raw actions")
logger.warning("Warning: Could not extract stats for encoded dimensions, using raw actions")
else:
print("Warning: No normalization stats found in dataset, using raw actions")
logger.warning("Warning: No normalization stats found in dataset, using raw actions")
print(f"Encoded chunks shape: {encoded_chunks.shape}")
logger.info(f"Encoded chunks shape: {encoded_chunks.shape}")
# train FAST tokenizer
tokenizer = train_fast_tokenizer(
@@ -561,8 +562,8 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
with open(output_path / "metadata.json", "w") as f:
json.dump(metadata, f, indent=2)
print(f"\nSaved FAST tokenizer to {output_path}")
print(f"Metadata: {json.dumps(metadata, indent=2)}")
logger.info(f"\nSaved FAST tokenizer to {output_path}")
logger.info(f"Metadata: {json.dumps(metadata, indent=2)}")
# push to Hugging Face Hub if requested
if cfg.push_to_hub:
@@ -570,10 +571,10 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
hub_repo_id = cfg.hub_repo_id
if hub_repo_id is None:
hub_repo_id = output_path.name
print(f"\nNo hub_repo_id provided, using: {hub_repo_id}")
logger.info(f"\nNo hub_repo_id provided, using: {hub_repo_id}")
print(f"\nPushing tokenizer to Hugging Face Hub: {hub_repo_id}")
print(f" Private: {cfg.hub_private}")
logger.info(f"\nPushing tokenizer to Hugging Face Hub: {hub_repo_id}")
logger.info(f" Private: {cfg.hub_private}")
try:
# use the tokenizer's push_to_hub method
@@ -593,14 +594,15 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
commit_message="Upload tokenizer metadata",
)
print(f"Successfully pushed tokenizer to: https://huggingface.co/{hub_repo_id}")
logger.info(f"Successfully pushed tokenizer to: https://huggingface.co/{hub_repo_id}")
except Exception as e:
print(f"Error pushing to hub: {e}")
print(" Make sure you're logged in with `huggingface-cli login`")
logger.error(f"Error pushing to hub: {e}")
logger.error(" Make sure you're logged in with `huggingface-cli login`")
def main():
"""CLI entry point that parses arguments and runs the tokenizer training."""
init_logging()
train_tokenizer()
+7 -7
View File
@@ -41,7 +41,7 @@ class RandomSubsetApply(Transform):
def __init__(
self,
transforms: Sequence[Callable],
transforms: Sequence[Callable[..., Any]],
p: list[float] | None = None,
n_subset: int | None = None,
random_order: bool = False,
@@ -50,7 +50,7 @@ class RandomSubsetApply(Transform):
if not isinstance(transforms, Sequence):
raise TypeError("Argument transforms should be a sequence of callables")
if p is None:
p = [1] * len(transforms)
p = [1.0] * len(transforms)
elif len(p) != len(transforms):
raise ValueError(
f"Length of p doesn't match the number of transforms: {len(p)} != {len(transforms)}"
@@ -69,7 +69,7 @@ class RandomSubsetApply(Transform):
self.n_subset = n_subset
self.random_order = random_order
self.selected_transforms = None
self.selected_transforms: list[Callable[..., Any]] = []
def forward(self, *inputs: Any) -> Any:
needs_unpacking = len(inputs) > 1
@@ -119,7 +119,7 @@ class SharpnessJitter(Transform):
super().__init__()
self.sharpness = self._check_input(sharpness)
def _check_input(self, sharpness):
def _check_input(self, sharpness: float | Sequence[float]) -> tuple[float, float]:
if isinstance(sharpness, (int | float)):
if sharpness < 0:
raise ValueError("If sharpness is a single number, it must be non negative.")
@@ -215,7 +215,7 @@ class ImageTransformsConfig:
)
def make_transform_from_config(cfg: ImageTransformConfig):
def make_transform_from_config(cfg: ImageTransformConfig) -> Transform:
if cfg.type == "SharpnessJitter":
return SharpnessJitter(**cfg.kwargs)
@@ -236,8 +236,8 @@ class ImageTransforms(Transform):
super().__init__()
self._cfg = cfg
self.weights = []
self.transforms = {}
self.weights: list[float] = []
self.transforms: dict[str, Transform] = {}
for tf_name, tf_cfg in cfg.tfs.items():
if tf_cfg.weight <= 0.0:
continue
+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"
+29 -45
View File
@@ -23,46 +23,6 @@ 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.
@@ -98,12 +58,36 @@ 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.
"""
writer = StreamingVideoWriter(video_path, fps)
try:
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"
for frame_array in stacked_frames:
writer.add_frame(frame_array)
finally:
writer.close()
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)
def deserialize_json_into_object[T: JsonLike](fpath: Path, obj: T) -> T:
+28
View File
@@ -16,11 +16,39 @@
# 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.
+1 -9
View File
@@ -38,10 +38,7 @@ def _is_scalar(x):
def init_rerun(
session_name: str = "lerobot_control_loop",
ip: str | None = None,
port: int | None = None,
web_port: int | None = None,
session_name: str = "lerobot_control_loop", ip: str | None = None, port: int | None = None
) -> None:
"""
Initializes the Rerun SDK for visualizing the control loop.
@@ -50,7 +47,6 @@ 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")
@@ -64,10 +60,6 @@ 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)
+7 -4
View File
@@ -133,10 +133,13 @@ def say(text: str, blocking: bool = False):
else:
raise RuntimeError("Unsupported operating system for text-to-speech.")
if blocking:
subprocess.run(cmd, check=True)
else:
subprocess.Popen(cmd, creationflags=subprocess.CREATE_NO_WINDOW if system == "Windows" else 0)
try:
if blocking:
subprocess.run(cmd, check=True, timeout=5)
else:
subprocess.Popen(cmd, creationflags=subprocess.CREATE_NO_WINDOW if system == "Windows" else 0)
except (FileNotFoundError, subprocess.TimeoutExpired) as e:
logging.warning("Text-to-speech command failed: %s | Error: %s", cmd, e)
def log_say(text: str, play_sounds: bool = True, blocking: bool = False):
-71
View File
@@ -1,71 +0,0 @@
# 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
+68 -1
View File
@@ -20,7 +20,7 @@
# ```
from pathlib import Path
from unittest.mock import patch
from unittest.mock import MagicMock, patch
import cv2
import numpy as np
@@ -123,6 +123,73 @@ def test_invalid_width_connect():
camera.connect(warmup=False)
def test_connect_cleans_up_after_settings_failure_and_allows_retry():
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH, warmup_s=0)
camera = OpenCVCamera(config)
opened_captures = []
def fail_settings():
opened_captures.append(camera.videocapture)
raise RuntimeError("settings failed")
with (
patch.object(camera, "_configure_capture_settings", side_effect=fail_settings),
pytest.raises(RuntimeError, match="settings failed"),
):
camera.connect(warmup=False)
assert camera.videocapture is None
assert camera.thread is None
assert not camera.is_connected
assert opened_captures[0] is not None
assert not opened_captures[0].isOpened()
camera.connect(warmup=False)
assert camera.is_connected
camera.disconnect()
def test_connect_cleans_up_after_warmup_failure_and_allows_retry():
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH, warmup_s=1)
camera = OpenCVCamera(config)
read_threads = []
def fail_warmup(*_args, **_kwargs):
read_threads.append(camera.thread)
raise TimeoutError("no frame")
with (
patch.object(camera, "async_read", side_effect=fail_warmup),
pytest.raises(TimeoutError, match="no frame"),
):
camera.connect()
assert camera.videocapture is None
assert camera.thread is None
assert not camera.is_connected
assert read_threads[0] is not None
assert not read_threads[0].is_alive()
camera.connect(warmup=False)
assert camera.is_connected
camera.disconnect()
def test_find_cameras_releases_unopened_handles():
module_path = OpenCVCamera.__module__
unopened_capture = MagicMock()
unopened_capture.isOpened.return_value = False
with (
patch(f"{module_path}.platform.system", return_value="Darwin"),
patch(f"{module_path}.MAX_OPENCV_INDEX", 1),
patch(f"{module_path}.cv2.VideoCapture", return_value=unopened_capture),
):
assert OpenCVCamera.find_cameras() == []
unopened_capture.release.assert_called_once_with()
@pytest.mark.parametrize("index_or_path", TEST_IMAGE_PATHS, ids=TEST_IMAGE_SIZES)
def test_read(index_or_path):
config = OpenCVCameraConfig(index_or_path=index_or_path, warmup_s=0)
+259 -1
View File
@@ -20,7 +20,7 @@
# ```
from pathlib import Path
from unittest.mock import patch
from unittest.mock import MagicMock, patch
import numpy as np
import pytest
@@ -30,6 +30,8 @@ from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnected
pytest.importorskip("pyrealsense2")
import pyrealsense2 as rs
from lerobot.cameras.realsense import RealSenseCamera, RealSenseCameraConfig
TEST_ARTIFACTS_DIR = Path(__file__).parent.parent / "artifacts" / "cameras"
@@ -61,6 +63,17 @@ def test_abc_implementation():
_ = RealSenseCamera(config)
@pytest.mark.parametrize("option", ["exposure", "gain", "white_balance"])
def test_manual_color_option_requires_rgb(option):
with pytest.raises(ValueError, match="use_rgb=True"):
RealSenseCameraConfig(
serial_number_or_name="042",
use_rgb=False,
use_depth=True,
**{option: 100},
)
def test_connect():
config = RealSenseCameraConfig(serial_number_or_name="042", warmup_s=0)
@@ -83,6 +96,27 @@ def test_connect_invalid_camera_path(patch_realsense):
camera.connect(warmup=False)
def test_connect_cleans_up_when_sensor_configuration_fails():
config = RealSenseCameraConfig(serial_number_or_name="042", exposure=120)
camera = RealSenseCamera(config)
pipeline = MagicMock()
pipeline.start.return_value = MagicMock()
with (
patch("lerobot.cameras.realsense.camera_realsense.rs.pipeline", return_value=pipeline),
patch.object(camera, "_configure_rs_pipeline_config"),
patch.object(camera, "_configure_capture_settings"),
patch.object(camera, "_configure_sensor_options", side_effect=ValueError("invalid exposure")),
pytest.raises(ValueError, match="invalid exposure"),
):
camera.connect(warmup=False)
pipeline.stop.assert_called_once_with()
assert camera.rs_pipeline is None
assert camera.rs_profile is None
assert not camera.is_connected
def test_invalid_width_connect():
config = RealSenseCameraConfig(serial_number_or_name="042", width=99999, height=480, fps=30)
camera = RealSenseCamera(config)
@@ -91,6 +125,33 @@ def test_invalid_width_connect():
camera.connect(warmup=False)
def test_connect_cleans_up_after_warmup_failure_and_allows_retry():
config = RealSenseCameraConfig(serial_number_or_name="042", width=640, height=480, fps=30)
camera = RealSenseCamera(config)
read_threads = []
def fail_warmup(*_args, **_kwargs):
read_threads.append(camera.thread)
raise TimeoutError("no frame")
with (
patch.object(camera, "async_read", side_effect=fail_warmup),
pytest.raises(TimeoutError, match="no frame"),
):
camera.connect()
assert camera.rs_pipeline is None
assert camera.rs_profile is None
assert camera.thread is None
assert not camera.is_connected
assert read_threads[0] is not None
assert not read_threads[0].is_alive()
camera.connect(warmup=False)
assert camera.is_connected
camera.disconnect()
def test_read():
config = RealSenseCameraConfig(serial_number_or_name="042", width=640, height=480, fps=30, warmup_s=0)
with RealSenseCamera(config) as camera:
@@ -228,6 +289,203 @@ def test_read_latest_too_old():
_ = camera.read_latest(max_age_ms=0) # immediately too old
def _make_mock_sensor(name: str, supported_options: set | None = None) -> MagicMock:
"""Build a fake rs.sensor that reports a name and a configurable supported-options set."""
supported = supported_options if supported_options is not None else set()
sensor = MagicMock()
sensor.get_info.return_value = name
sensor.supports.side_effect = lambda opt: opt in supported
return sensor
def _attach_mock_color_sensor(camera: RealSenseCamera, sensor: MagicMock) -> None:
"""Wire camera.rs_profile so _get_color_sensor finds the given sensor."""
profile = MagicMock()
device = MagicMock()
device.query_sensors.return_value = [sensor]
profile.get_device.return_value = device
camera.rs_profile = profile
def test_get_color_sensor_prefers_rgb_camera():
config = RealSenseCameraConfig(serial_number_or_name="042")
camera = RealSenseCamera(config)
rgb = _make_mock_sensor("RGB Camera")
stereo = _make_mock_sensor("Stereo Module")
profile = MagicMock()
device = MagicMock()
device.query_sensors.return_value = [stereo, rgb]
profile.get_device.return_value = device
camera.rs_profile = profile
assert camera._get_color_sensor() is rgb
def test_get_color_sensor_falls_back_to_stereo_module():
"""D405 has no separate RGB module; color comes from Stereo Module."""
config = RealSenseCameraConfig(serial_number_or_name="042")
camera = RealSenseCamera(config)
stereo = _make_mock_sensor("Stereo Module")
_attach_mock_color_sensor(camera, stereo)
assert camera._get_color_sensor() is stereo
def test_get_color_sensor_raises_with_available_sensors():
config = RealSenseCameraConfig(serial_number_or_name="042")
camera = RealSenseCamera(config)
other = _make_mock_sensor("Motion Module")
_attach_mock_color_sensor(camera, other)
with pytest.raises(RuntimeError, match="Motion Module"):
camera._get_color_sensor()
def test_configure_sensor_options_skipped_when_none():
config = RealSenseCameraConfig(serial_number_or_name="042")
camera = RealSenseCamera(config)
with patch.object(RealSenseCamera, "_get_color_sensor") as mock_get:
camera._configure_sensor_options()
mock_get.assert_not_called()
def test_configure_sensor_options_applies_all_values():
config = RealSenseCameraConfig(serial_number_or_name="042", exposure=120, gain=64, white_balance=4600)
camera = RealSenseCamera(config)
sensor = _make_mock_sensor(
"RGB Camera",
supported_options={
rs.option.enable_auto_exposure,
rs.option.exposure,
rs.option.gain,
rs.option.enable_auto_white_balance,
rs.option.white_balance,
},
)
_attach_mock_color_sensor(camera, sensor)
camera._configure_sensor_options()
sensor.set_option.assert_any_call(rs.option.enable_auto_exposure, 0)
sensor.set_option.assert_any_call(rs.option.exposure, 120)
sensor.set_option.assert_any_call(rs.option.gain, 64)
sensor.set_option.assert_any_call(rs.option.enable_auto_white_balance, 0)
sensor.set_option.assert_any_call(rs.option.white_balance, 4600)
@pytest.mark.parametrize(
("config_field", "option", "label"),
[
("exposure", rs.option.exposure, "exposure"),
("gain", rs.option.gain, "gain"),
("white_balance", rs.option.white_balance, "white balance"),
],
)
def test_configure_sensor_options_raises_when_requested_option_is_unsupported(config_field, option, label):
config = RealSenseCameraConfig(serial_number_or_name="042", **{config_field: 100})
camera = RealSenseCamera(config)
sensor = _make_mock_sensor("RGB Camera", supported_options=set())
_attach_mock_color_sensor(camera, sensor)
with pytest.raises(ValueError, match=label):
camera._configure_sensor_options()
sensor.supports.assert_any_call(option)
sensor.set_option.assert_not_called()
@pytest.mark.parametrize(
("config_field", "option", "value"),
[
("exposure", rs.option.exposure, 120),
("gain", rs.option.gain, 64),
],
)
def test_configure_sensor_options_exposure_or_gain_disables_auto_exposure(config_field, option, value):
"""white_balance=None should not touch auto white balance."""
config = RealSenseCameraConfig(serial_number_or_name="042", **{config_field: value})
camera = RealSenseCamera(config)
sensor = _make_mock_sensor(
"RGB Camera",
supported_options={rs.option.enable_auto_exposure, option},
)
_attach_mock_color_sensor(camera, sensor)
camera._configure_sensor_options()
calls = [call.args for call in sensor.set_option.call_args_list]
assert (rs.option.enable_auto_exposure, 0) in calls
assert (option, value) in calls
for opt, _ in calls:
assert opt != rs.option.enable_auto_white_balance
assert opt != rs.option.white_balance
def test_configure_sensor_options_warns_when_auto_exposure_control_is_unsupported(caplog):
config = RealSenseCameraConfig(serial_number_or_name="042", exposure=120)
camera = RealSenseCamera(config)
sensor = _make_mock_sensor("RGB Camera", supported_options={rs.option.exposure})
_attach_mock_color_sensor(camera, sensor)
with caplog.at_level("WARNING"):
camera._configure_sensor_options()
sensor.set_option.assert_called_once_with(rs.option.exposure, 120)
assert "does not support disabling auto-exposure" in caplog.text
def test_configure_sensor_options_warns_when_auto_white_balance_control_is_unsupported(caplog):
config = RealSenseCameraConfig(serial_number_or_name="042", white_balance=4600)
camera = RealSenseCamera(config)
sensor = _make_mock_sensor("RGB Camera", supported_options={rs.option.white_balance})
_attach_mock_color_sensor(camera, sensor)
with caplog.at_level("WARNING"):
camera._configure_sensor_options()
sensor.set_option.assert_called_once_with(rs.option.white_balance, 4600)
assert "does not support disabling auto white balance" in caplog.text
def test_configure_sensor_options_out_of_range_raises_value_error():
"""set_option errors should be re-raised as ValueError with range diagnostics."""
config = RealSenseCameraConfig(serial_number_or_name="042", exposure=999999)
camera = RealSenseCamera(config)
sensor = _make_mock_sensor(
"RGB Camera",
supported_options={rs.option.enable_auto_exposure, rs.option.exposure},
)
def fake_set_option(option, value):
if option == rs.option.exposure:
raise RuntimeError("value out of range")
sensor.set_option.side_effect = fake_set_option
option_range = MagicMock(min=1, max=10000, step=1, default=156)
sensor.get_option_range.return_value = option_range
_attach_mock_color_sensor(camera, sensor)
with pytest.raises(ValueError, match="exposure") as exc_info:
camera._configure_sensor_options()
msg = str(exc_info.value)
assert "999999" in msg
assert "min=1" in msg
assert "max=10000" in msg
@pytest.mark.parametrize(
"rotation",
[
-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
+14
View File
@@ -482,6 +482,20 @@ def test_add_frame_works_in_write_mode(tmp_path):
# ── Resume mode ──────────────────────────────────────────────────────
def test_resume_freshly_created_empty_dataset(tmp_path):
"""resume() accepts a local dataset created before any episode was recorded."""
root = tmp_path / "resume_empty_ds"
LeRobotDataset.create(repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root)
resumed = LeRobotDataset.resume(repo_id=DUMMY_REPO_ID, root=root)
assert isinstance(resumed.writer, DatasetWriter)
assert resumed.meta.total_episodes == 0
assert resumed.meta.total_frames == 0
assert resumed.meta.tasks is None
assert resumed.meta.episodes is None
def test_resume_creates_writer(tmp_path):
"""After resume(), writer is a DatasetWriter."""
root = tmp_path / "resume_ds"
+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
-48
View File
@@ -1,48 +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.
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))
@@ -0,0 +1,83 @@
# 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.
from types import SimpleNamespace
from unittest.mock import MagicMock
import torch
import lerobot.policies.factory as policy_factory
def test_make_policy_keeps_peft_adapter_and_base_revisions_separate(monkeypatch):
cfg = SimpleNamespace(
type="mock",
device="cpu",
pretrained_path="user/adapter",
pretrained_revision="adapter-sha",
use_peft=True,
input_features={},
output_features={},
)
dataset_meta = SimpleNamespace(features={}, stats={})
base_policy = torch.nn.Linear(1, 1)
policy_from_pretrained = MagicMock(return_value=base_policy)
policy_class = SimpleNamespace(from_pretrained=policy_from_pretrained)
monkeypatch.setattr(policy_factory, "get_policy_class", lambda _: policy_class)
monkeypatch.setattr(policy_factory, "dataset_to_policy_features", lambda _: {})
monkeypatch.setattr(policy_factory, "validate_visual_features_consistency", lambda *args: None)
peft_config = SimpleNamespace(
base_model_name_or_path="user/base-policy",
revision="base-sha",
)
peft_config_from_pretrained = MagicMock(return_value=peft_config)
adapted_policy = torch.nn.Linear(1, 1)
peft_model_from_pretrained = MagicMock(return_value=adapted_policy)
require_package = MagicMock()
monkeypatch.setattr(policy_factory, "require_package", require_package)
monkeypatch.setattr(
policy_factory,
"PeftConfig",
SimpleNamespace(from_pretrained=peft_config_from_pretrained),
)
monkeypatch.setattr(
policy_factory,
"PeftModel",
SimpleNamespace(from_pretrained=peft_model_from_pretrained),
)
policy = policy_factory.make_policy(cfg, ds_meta=dataset_meta)
assert policy is adapted_policy
require_package.assert_called_once_with("peft", extra="peft")
peft_config_from_pretrained.assert_called_once_with(
"user/adapter",
revision="adapter-sha",
)
policy_from_pretrained.assert_called_once_with(
config=cfg,
dataset_stats=dataset_meta.stats,
dataset_meta=dataset_meta,
pretrained_name_or_path="user/base-policy",
revision="base-sha",
)
peft_model_from_pretrained.assert_called_once_with(
base_policy,
"user/adapter",
config=peft_config,
revision="adapter-sha",
is_trainable=True,
)
@@ -113,6 +113,7 @@ 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)
@@ -152,6 +153,7 @@ 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."""
-105
View File
@@ -1,105 +0,0 @@
# 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.
from lerobot.runtime import RuntimeState
from lerobot.runtime.adapter import (
BaseLanguageAdapter,
DirectTaskPolicyAdapter,
GenerationConfig,
)
class ScriptedAdapter(BaseLanguageAdapter):
"""Base adapter whose text generation returns queued strings per kind."""
def __init__(self, scripts, gen=None):
super().__init__(policy=object(), gen=gen)
self.scripts = {k: list(v) for k, v in scripts.items()}
self.calls = []
def select_action(self, observation, state):
return None
def generate_text(self, kind, observation, state, user_text=None):
self.calls.append(kind)
queue = self.scripts.get(kind, [])
return queue.pop(0) if queue else ""
def test_cascade_sets_subtask_then_memory():
adapter = ScriptedAdapter({"subtask": ["pick the red cup"], "memory": ["the cup is grasped"]})
state = RuntimeState(task="clean")
adapter.update_language_state(None, state)
assert state.language_context["subtask"] == "pick the red cup"
assert state.language_context["memory"] == "the cup is grasped"
assert adapter.calls == ["subtask", "memory"]
def test_nonempty_generation_is_used_verbatim():
adapter = ScriptedAdapter({"subtask": [":::: ::"], "memory": ["memory"]})
state = RuntimeState(task="clean")
adapter.update_language_state(None, state)
assert state.language_context["subtask"] == ":::: ::"
assert state.language_context["memory"] == "memory"
assert adapter.calls == ["subtask", "memory"]
def test_throttle_regenerates_every_n_chunks():
adapter = ScriptedAdapter(
{
"subtask": ["pick the first cup", "pick the second cup"],
"memory": ["memory one two three", "memory four five six"],
},
gen=GenerationConfig(chunks_per_regen=2),
)
state = RuntimeState(task="clean")
adapter.update_language_state(None, state) # generates
assert state.language_context["subtask"] == "pick the first cup"
adapter.update_language_state(None, state) # throttled — no generation
assert state.language_context["subtask"] == "pick the first cup"
adapter.update_language_state(None, state) # generates again
assert state.language_context["subtask"] == "pick the second cup"
def test_handle_interjection_sets_plan_and_strips_say():
adapter = ScriptedAdapter({"interjection": ["turn to the left now <say>heading left</say>"]})
state = RuntimeState(task="clean")
adapter.handle_interjection("turn", None, state)
assert state.language_context["plan"] == "turn to the left now"
def test_direct_task_adapter_delegates_action_chunk():
class Policy:
def predict_action_chunk(self, observation):
return ("chunk", observation)
observation = {"task": "pick up the cube"}
adapter = DirectTaskPolicyAdapter(Policy())
assert adapter.select_action(observation, RuntimeState()) == ("chunk", observation)
assert adapter.generate_text("subtask", observation, RuntimeState()) == ""
def test_flat_policy_registry_reuses_direct_task_adapter():
from lerobot.runtime.registry import get_language_adapter_factory
assert get_language_adapter_factory("pi05") is DirectTaskPolicyAdapter
assert get_language_adapter_factory("molmoact2") is DirectTaskPolicyAdapter
-75
View File
@@ -1,75 +0,0 @@
# 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.
from types import SimpleNamespace
from unittest.mock import MagicMock
import pytest
import torch
from lerobot.runtime.cli import _build_rollout_runtime_io, _parse_args
def test_parse_args_preserves_rollout_robot_overrides():
args = _parse_args(
[
"--policy.path=checkpoint",
"--robot.type=so101_follower",
"--robot.calibration_dir=/tmp/calibration",
]
)
assert args.robot_type == "so101_follower"
assert "--robot.calibration_dir=/tmp/calibration" in args.raw_argv
def test_parse_args_rejects_removed_dataset_replay_flags():
with pytest.raises(SystemExit):
_parse_args(["--policy.path=checkpoint", "--dataset.repo_id=dataset"])
def test_rollout_runtime_io_uses_context_processors():
robot = MagicMock()
robot.robot_type = "mock_robot"
robot.cameras = {}
robot.get_observation.return_value = {"joint.pos": 1.5}
ctx = SimpleNamespace(
hardware=SimpleNamespace(robot_wrapper=robot),
runtime=SimpleNamespace(cfg=SimpleNamespace(device="cpu")),
processors=SimpleNamespace(
robot_observation_processor=lambda observation: observation,
robot_action_processor=lambda pair: pair[0],
),
policy=SimpleNamespace(
preprocessor=lambda observation: observation,
postprocessor=lambda action: action,
),
data=SimpleNamespace(
dataset_features={
"observation.state": {
"dtype": "float32",
"shape": (1,),
"names": ["joint.pos"],
},
"action": {"dtype": "float32", "shape": (1,), "names": ["joint.pos"]},
}
),
)
provider, executor = _build_rollout_runtime_io(ctx, rerun_log=False, get_task=lambda: "move")
observation = provider()
executor(torch.tensor([[2.0]]))
assert observation["observation.state"].shape == (1, 1)
robot.send_action.assert_called_once_with({"joint.pos": 2.0})
-100
View File
@@ -1,100 +0,0 @@
# 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 threading
import time
from lerobot.runtime import LanguageConditionedRuntime, Tick
class FakeAdapter:
def __init__(self):
self.updated = False
self.interjections = []
def select_action(self, observation, state):
assert observation == {"observation.state": 1}
assert state.task == "clean"
return ["a0", "a1"]
def update_language_state(self, observation, state):
self.updated = True
state.set_context("subtask", "pick cup", label="subtask")
def handle_interjection(self, user_text, observation, state):
self.interjections.append(user_text)
state.set_context("plan", "new plan", label="plan")
def test_runtime_tick_updates_language_enqueues_and_dispatches_action():
adapter = FakeAdapter()
executed = []
runtime = LanguageConditionedRuntime(
policy_adapter=adapter,
observation_provider=lambda: {"observation.state": 1},
action_executor=executed.append,
)
runtime.set_task("clean")
logs = runtime.step_once()
assert adapter.updated
assert runtime.state.language_context["subtask"] == "pick cup"
assert executed == ["a0"]
assert list(runtime.state.action_queue) == ["a1"]
assert " subtask: pick cup" in logs
def test_runtime_handles_user_interjection():
adapter = FakeAdapter()
runtime = LanguageConditionedRuntime(
policy_adapter=adapter,
observation_provider=lambda: {"observation.state": 1},
)
runtime.set_task("clean")
runtime.state.extra["recent_interjection"] = "please say ok"
runtime.state.emit("user_interjection")
runtime.step_once()
assert "please say ok" in adapter.interjections
assert runtime.state.language_context["plan"] == "new plan"
def test_prompt_change_discards_in_flight_action_chunk():
started = threading.Event()
release = threading.Event()
class BlockingAdapter(FakeAdapter):
def select_action(self, observation, state):
started.set()
assert release.wait(timeout=2)
return ["stale"]
runtime = LanguageConditionedRuntime(
policy_adapter=BlockingAdapter(),
observation_provider=lambda: {"observation.state": 1},
)
runtime.set_task("old task")
runtime.state.tick = Tick(index=1, monotonic_seconds=time.monotonic())
inference = threading.Thread(target=runtime.maybe_enqueue_action_chunk, kwargs={"force": True})
inference.start()
assert started.wait(timeout=2)
runtime.set_task("new task")
release.set()
inference.join(timeout=2)
assert not inference.is_alive()
assert list(runtime.state.action_queue) == []
-82
View File
@@ -1,82 +0,0 @@
# 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 sys
from types import SimpleNamespace
import numpy as np
from lerobot.runtime.sim_robocasa import RoboCasaSimBackend
from lerobot.utils.video_annotation import annotate_frame
def test_overlay_draws_each_label_once(monkeypatch):
put_text_calls = []
rectangle_calls = []
def put_text(image, text, origin, font, scale, color, thickness, line_type):
put_text_calls.append((text, color, thickness))
return image
def rectangle(image, start, end, color, thickness):
rectangle_calls.append((start, end, color, thickness))
return image
def add_weighted(src1, alpha, src2, beta, gamma, *, dst):
dst[:] = src1 * alpha + src2 * beta + gamma
return dst
fake_cv2 = SimpleNamespace(
FONT_HERSHEY_SIMPLEX=0,
LINE_AA=16,
getTextSize=lambda text, font, scale, thickness: ((len(text) * 7, 10), 0),
putText=put_text,
rectangle=rectangle,
addWeighted=add_weighted,
)
monkeypatch.setitem(sys.modules, "cv2", fake_cv2)
frame = np.full((120, 480, 3), 200, dtype=np.uint8)
annotated = annotate_frame(
frame,
(("Task", "close the fridge"), ("Subtask", "reach for the handle"), ("Memory", None)),
)
assert [call[0] for call in put_text_calls] == [
"Task: close the fridge",
"Subtask: reach for the handle",
]
assert all(color == (255, 255, 255) and thickness == 1 for _, color, thickness in put_text_calls)
assert len(rectangle_calls) == 1
assert not np.shares_memory(annotated, frame)
def test_capture_updates_live_frame_when_recording_is_disabled(monkeypatch):
backend = object.__new__(RoboCasaSimBackend)
frame = np.full((8, 8, 3), 42, dtype=np.uint8)
written = []
backend.record = False
backend.runtime_state = None
backend._multiview_frame = lambda: frame
backend._current_task = lambda: "task"
backend._subtask_getter = None
backend._memory_getter = None
backend._latest_frame = None
backend._write_live_frame = written.append
monkeypatch.setattr("lerobot.runtime.sim_robocasa.annotate_frame", lambda image, labels: image)
backend._capture_frame()
assert backend._latest_frame is frame
assert written == [frame]
+14 -7
View File
@@ -185,18 +185,25 @@ def test_load_pretrained_peft_policy_keeps_adapter_and_base_revisions_separate(m
peft_config_from_pretrained = MagicMock(return_value=peft_config)
adapted_policy = MagicMock()
peft_model_from_pretrained = MagicMock(return_value=adapted_policy)
monkeypatch.setitem(
sys.modules,
"peft",
SimpleNamespace(
PeftConfig=SimpleNamespace(from_pretrained=peft_config_from_pretrained),
PeftModel=SimpleNamespace(from_pretrained=peft_model_from_pretrained),
),
require_package = MagicMock()
monkeypatch.setattr(rollout_context, "require_package", require_package)
monkeypatch.setattr(
rollout_context,
"PeftConfig",
SimpleNamespace(from_pretrained=peft_config_from_pretrained),
raising=False,
)
monkeypatch.setattr(
rollout_context,
"PeftModel",
SimpleNamespace(from_pretrained=peft_model_from_pretrained),
raising=False,
)
policy = rollout_context._load_pretrained_policy(policy_config)
assert policy is adapted_policy
require_package.assert_called_once_with("peft", extra="peft")
peft_config_from_pretrained.assert_called_once_with("user/adapter", revision="adapter-sha")
policy_class.from_pretrained.assert_called_once_with(
pretrained_name_or_path="user/base-policy",