Compare commits

..

95 Commits

Author SHA1 Message Date
Steven Palma a3feb09b08 fix(dataset): bump versions + improvements 2026-07-30 15:55:21 +02:00
Sundar Raghavan 13a261a08a test(streaming): make bucket get_safe_version test actually exercise the except branch
The prior test patched _load_metadata as a plain no-op, so __init__'s try
block succeeded and never entered the except branch where the
repo_type != "bucket" guard and get_safe_version live - the assertion passed
vacuously (verified: it still passed with the guard removed).

Use side_effect=[FileNotFoundError, None] so the first _load_metadata raises
(forcing the except path) and the second succeeds after the meta pull, and
stub _pull_from_repo so the path runs without a network call. Now the test
fails if the bucket guard is removed. Thanks @mohitydv09 for the catch.

Signed-off-by: Sundar Raghavan <sdraghav@amazon.com>
2026-07-30 15:45:43 +02:00
Sundar Raghavan 9782bfd64b Support streaming from HF Storage Buckets in StreamingLeRobotDataset
Add an opt-in repo_type="bucket" parameter to StreamingLeRobotDataset
and LeRobotDatasetMetadata so a dataset can be streamed directly from an
HF Storage Bucket (hf://buckets/...) with no local download.

When repo_type="bucket":
- skip git-version resolution (buckets have no refs/tags),
- pull the meta/ directory via HfFileSystem.get,
- point url_root at hf://buckets/{repo_id},
- read parquet shards via load_dataset("parquet",
  data_files="hf://buckets/{repo_id}/data/*/*.parquet", ...).

The default repo_type="dataset" preserves all existing behavior.
LeRobotDataset (non-streaming) and create() are unchanged.

Closes #3969

Signed-off-by: Sundar Raghavan <sdraghav@amazon.com>
2026-07-30 15:45:42 +02:00
Steven Palma 7e0fd0d653 refactor(types): change module name (#4232)
* refactor(types): change module name

Co-authored-by: saiteja6006 <saiteja6006@gmail.com>

* chore(test): remove package import test

* chore: remove ruff exception

---------

Co-authored-by: saiteja6006 <saiteja6006@gmail.com>
2026-07-30 15:27:51 +02:00
Martino Russi 0187856202 fix typo (#4048)
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-30 15:25:26 +02:00
Anas 2939168c33 fix(envs): use RoboCasa task horizons (#4037)
Co-authored-by: Pepijn <138571049+pkooij@users.noreply.github.com>
2026-07-30 15:20:56 +02:00
Jash Shah 40a5e70352 fix(config): accept pretrained_model dir for --config_path on resume (#4023)
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-30 13:48:27 +02:00
Jash Shah 0cef9cd197 fix(train): keep checkpoint processor stats on resume (#4022)
Co-authored-by: Martino Russi <77496684+nepyope@users.noreply.github.com>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-30 13:47:56 +02:00
Nick 643ffb4785 chore(deps): bump draccus (#4033)
* Update draccus to 0.11

* Update draccus calls to be backwards compatible

---------

Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-30 13:47:30 +02:00
Baptiste Lubrano Lavadera d59505a735 feat(teleoperators): add DAgger/HIL smooth handover support for BiSOLeader (#4028)
* fix: implement bimanual SO leader DAgger handover support

- Add feedback_features property: enables DAgger's teleop_supports_feedback() check
- Implement enable_torque()/disable_torque(): synchronized torque control for both arms
- Implement send_feedback(): routes bimanual feedback to left/right arms with prefix stripping

This fixes DAgger smooth handover for bimanual SO follower + SO leader setups:
when pausing from policy to human intervention, both leader arms now move smoothly
to the follower's current pose, avoiding discontinuities at the human takeover point.

* Update hil_data_collection.mdx

Signed-off-by: Baptiste Lubrano Lavadera  <45080391+Mr-C4T@users.noreply.github.com>

* Update bi_so_leader.py

Signed-off-by: Baptiste Lubrano Lavadera  <45080391+Mr-C4T@users.noreply.github.com>

---------

Signed-off-by: Baptiste Lubrano Lavadera  <45080391+Mr-C4T@users.noreply.github.com>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-30 13:38:14 +02:00
Steven Palma 6ac95363b0 fix(rollout): reject incompatible RTC policies (#4228)
* fix(rollout): reject incompatible RTC policies

* chore(policies): support rtc

* chore(tests): delete compatibility test

---------

Co-authored-by: ogarciarevett <ogarciarevett@gmail.com>
2026-07-30 13:27:05 +02:00
Xingdong Zuo ede1fc2978 fix(smolvla): freeze the intended VLM layers when train_expert_only=False (#4019)
* fix(smolvla): freeze the intended VLM layers when train_expert_only=False

The partial-freeze patterns in set_requires_grad() used a
'text_model.model.' prefix that does not exist in SmolVLM parameter
names ('SmolVLMModel.text_model' is a bare LlamaModel, with no nested
'.model'). As a result the last VLM layer and the final norm were
silently left trainable, defeating the freeze that was added to avoid
unused-parameter errors with DDP; only lm_head was frozen by substring
luck.

Use the real flat names, and raise if any freeze pattern stops matching
so a future transformers renaming cannot silently reintroduce the bug.
Add a CPU regression test covering both last_layers branches.

Fixes #4018

* test(smolvla): drop regression test per review

---------

Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-30 13:17:54 +02:00
sunnydave234 49d5ea49bc fix(utils): add MPS branch to torch RNG state serialization (#4014)
serialize_torch_rng_state/deserialize_torch_rng_state only handled CPU
and CUDA generators. On MPS, resumed training was not bit-exact for any
stochastic op (dropout, ACT's CVAE noise) since the MPS generator's state
was never saved or restored. Mirrors the existing CUDA branch using
torch.mps.get_rng_state/set_rng_state (available since torch 2.11).

Note: get_rng_state()/set_rng_state() (used by seeded_context()) have the
same gap but are out of scope here — happy to follow up separately if
useful.

Co-authored-by: Sunny Dave <sunnydave@Sunnys-Mac-Studio.local>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-30 12:22:13 +02:00
Nikodem Bartnik d23b65416f fix assembly instructions typo (#4008) 2026-07-30 11:43:58 +02:00
HUANG TZU-CHUN a6b06eac38 docs: fix env processor code fences and minor doc errors (#3953)
* docs: fix code fences in env processor guide

The "Flexibility and Reusability" section wrapped a duplicated example
in a four-backtick fence and left a following block unclosed, so the
stray closing fence matched a later block. Everything in between
rendered as one code block that swallowed the surrounding prose.

Remove the duplicated block, add the missing closing fence after the
first example, and normalize the four-backtick fences to three so all
fences pair correctly.

* docs(pi0fast): fix typo 40kk -> 40k steps

* docs(integrate-hardware): fix so101 follower source link

* docs(hope_jr): fix dataset example link

The "example" link in the Record section pointed at the dataset's
`/settings` page, which returns HTTP 403 for readers. Drop the
`/settings` suffix so it links to the public dataset page the
sentence describes.

* docs(lekiwi): render emoji shortcodes as unicode

MDX does not expand `🤗` / `🤖` shortcodes, so they showed as
literal text in the rendered install step. Replace them with the 🤗 and
🤖 unicode characters, matching how the other robot pages write emoji.

* docs(smolvla): anchor record link to its section

The "Record a dataset" link dropped readers at the top of the
il_robots page instead of the relevant section. Point it at the
`#record-a-dataset` anchor (the `## Record a dataset` heading in
il_robots.mdx) so the link lands on the step it names.
2026-07-30 10:53:27 +02:00
Steven Palma 36b8face98 fix(utils): validate precise_sleep spin/margin args (#4218)
* fix(utils): validate precise_sleep spin/margin args

Negative spin_threshold/sleep_margin make remaining arithmetic wrong
and can overshoot. Reject them early; cover the no-op path.

* test: drop flaky wall-clock assertion in no-op test

Per review: the 50ms wall-clock check can exceed its bound on a preempted
CI worker even when precise_sleep returns immediately. The direct calls
already exercise the non-positive no-op path, so the assertion is redundant.

* chore(tests): remove precise_sleep test negative values

---------

Co-authored-by: Bartok9 <danielrpike9@gmail.com>
2026-07-29 20:24:07 +02:00
Steven Palma cd8984cc0a fix(utils): allow any JSON payload in write_json - #3993 (#4217)
* fix(utils): allow any JSON payload in write_json

The dict-only type stub blocked lists/scalars callers already dump.
Accept Any, set utf-8 encoding, and cover list roundtrip.

* fix(utils): json type

---------

Co-authored-by: Bartok9 <danielrpike9@gmail.com>
2026-07-29 20:11:14 +02:00
Steven Palma b9ded9e761 fix(utils): mark Transition.complementary_info NotRequired (#4216)
* fix(utils): mark Transition.complementary_info NotRequired

TypedDict class-body ``= None`` does not make a key optional and confuses
type checkers. Use ``NotRequired[...]`` so transitions without metadata
are valid.

* refactor(utils): complete NotRequired

---------

Co-authored-by: Bartok9 <danielrpike9@gmail.com>
2026-07-29 19:55:39 +02:00
Steven Palma 185f3e1708 fix(utils): preserve exc_info/stack_info in init_logging formatter (#4215)
* fix(utils): preserve exc_info/stack_info in init_logging formatter

Replacing Formatter.format dropped logging.exception() tracebacks,
hurting HIL-SERL actor/learner crash diagnosis. Append formatted
exceptions and stack_info like the stdlib formatter.

Fixes #3978

* refactor(utils): format logging

---------

Co-authored-by: Bartok9 <danielrpike9@gmail.com>
2026-07-29 19:32:30 +02:00
Bartok e36783253a fix(utils): raise ValueError from get_safe_torch_device (#3992)
* fix(utils): raise ValueError from get_safe_torch_device

Bare asserts vanish under python -O and look like programmer bugs.
Convert unavailable CUDA/MPS/XPU requests into clear ValueErrors.

* style: combine nested with in device util tests (ruff)

---------

Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-29 19:21:50 +02:00
Bartok 289e577fc7 fix(utils): reject zero-norm / invalid quaternions in Rotation (#3988)
* fix(utils): reject zero-norm / invalid quaternions in Rotation

Zero or non-finite inputs previously slipped through and produced NaN
rotation matrices on later convert/apply. Validate shape and scept for
norm > 0 before normalizing.

* fix(teleop): degrade phone AR quat parse like missing pose

Address review on #3988: Rotation.from_quat now rejects zero/NaN
quaternions. Wrap HEBI iOS ARKit permission in ValueError and return the
existing (False, None, None, None) path so teleop does not die mid-session
before tracking is ready.

* style: ruff format long ValueError in rotation.py

---------

Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-29 19:13:40 +02:00
Anes Benmerzoug 9c32722eb9 fix(find-cameras): enforce sequential lifecycle and add configurable warmup (#3593)
* Connect, test and disconnected camera instances sequentially

* Add warmup-s cli argument to lerobot-find-cameras script

* Reduce default record time from 6 to 2 seconds in find_cameras

* Annotate return value of save_image function

* Initialize logging configuration in find_cameras
2026-07-29 19:01:44 +02:00
Kunal b49cb50e01 docs(agent-guide): prioritize uv over pip in §4.1 install block (#3799)
Co-authored-by: Altman <64389901+Altman-conquer@users.noreply.github.com>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-29 18:46:20 +02:00
Steven Palma dd08d4eb53 fix(robot): type FK-to-EE action features as ACTION not STATE (#4213)
* fix(robot): type FK-to-EE action features as ACTION not STATE

ForwardKinematicsJointsToEEAction.transform_features declared its
end-effector action features (ee.x/y/z/wx/wy/wz/gripper_pos) with
FeatureType.STATE, copied verbatim from the sibling
ForwardKinematicsJointsToEEObservation (where STATE is correct for
OBSERVATION features). Every other action-producing step in this file
(EEReferenceAndDelta, InverseKinematicsEEToJoints, InverseKinematicsRLStep)
types its ACTION-bucket features as FeatureType.ACTION.

The mismatch mis-classifies the converted EE actions as state, which
propagates a wrong feature schema to downstream consumers keyed on
FeatureType (e.g. normalization norm_map, policy input/output feature
classification).


* test(robot): FK-to-EE step feature-type contract (action vs observation)

Asserts ForwardKinematicsJointsToEEAction emits EE features in the ACTION
bucket typed FeatureType.ACTION, and ForwardKinematicsJointsToEEObservation
emits them in the OBSERVATION bucket typed FeatureType.STATE.


* chore: delete user file

* chore(processor): reduce verbosity

---------

Co-authored-by: Jaagat-P <jaagatp05@gmail.com>
2026-07-29 18:06:01 +02:00
Martino Russi 6e5f6df6e7 fix(evo1): re-pad normalizer stats when loading from checkpoint (#3945)
* fix(evo1): re-pad normalizer stats when loading from checkpoint

reconcile_evo1_processors did not re-pad the (un)normalizer stats to
max_state_dim/max_action_dim on the checkpoint-load path. When
lerobot-train loads a checkpoint (e.g. stage2 from a stage1 checkpoint)
it injects the raw dataset stats via processor overrides, so LIBERO's
8-dim state stats normalized a 24-dim padded state and crashed with
"size of tensor a (24) must match tensor b (8)".

Restore _refresh_evo1_normalization_steps (removed in the "remove legacy
codepaths" refactor) and call it from reconcile_evo1_processors so the
loaded stats/features are re-padded to EVO1's fixed widths. Padding is a
no-op when stats are already at the target width.

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

* test(evo1): cover reconcile re-padding of overridden normalizer stats

Regression test for the stage2-from-checkpoint crash: reloading a
checkpoint with raw (unpadded) dataset stats injected via processor
overrides must be re-padded to max_state_dim/max_action_dim by
reconcile_evo1_processors, otherwise normalizing the padded state
raises a shape mismatch.

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

---------

Co-authored-by: Martino Russi <martino@huggingface.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-29 17:26:39 +02:00
Steven Palma 265abe6c79 chore(datasets): add typing to aggregate helpers (#4211)
* chore(datasets): add typing to aggregate helpers

Signed-off-by: nathon-lee <leejianwoo@gmail.com>

* chore(dataset): add more typing aggregate

* chore(test): remove panda test

---------

Signed-off-by: nathon-lee <leejianwoo@gmail.com>
Co-authored-by: nathon-lee <leejianwoo@gmail.com>
2026-07-29 17:07:34 +02:00
Old-Ding b4e2d0b610 docs: fix wording in guides (#3939)
Generated-by: OpenAI Codex

Signed-off-by: aineoae86-sys <ai.neo.ae86@gmail.com>
Co-authored-by: aineoae86-sys <ai.neo.ae86@gmail.com>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-29 16:24:03 +02:00
Old-Ding 5594eba06a docs: fix repeated word in backward compatibility guide (#3938)
Generated-by: OpenAI Codex

Signed-off-by: aineoae86-sys <ai.neo.ae86@gmail.com>
Co-authored-by: aineoae86-sys <ai.neo.ae86@gmail.com>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-29 16:23:01 +02:00
saime428 207183c2f8 docs: fix dataset split fraction example (#3936)
* docs: fix dataset split fraction example

* docs: preserve three-way dataset split example

---------

Co-authored-by: saime <2286263079@qq.com>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-29 16:21:12 +02:00
Steven Palma 7d615acf9a fix(robots): retry SO follower/leader bus reads on transient Feetech errors (#4207)
* fix(robots): retry SO follower/leader bus reads on transient Feetech errors

SO-100/SO-101 teleoperation aborts when a sync_read of Present_Position
returns a corrupted status packet ("Incorrect status packet!"), which the
Feetech bus emits intermittently under load. The read path already supports
a num_retry argument but the SO follower and leader never used it, so a single
transient failure crashed the control loop.

Add a max_read_retry config option (default 3) to SOFollowerConfig and
SOLeaderConfig and forward it to every Present_Position sync_read. Retries are
immediate and only happen on failure, so the steady-state read cost is
unchanged; set max_read_retry=0 to restore the previous behavior.

Fixes #3131

* chore(robots): change defaults

---------

Co-authored-by: isaka1022 <isaka1022@gmail.com>
2026-07-29 15:20:14 +02:00
Steven Palma 09572babee perf(docker): split dependency install from source copy for CI layer caching (#4208)
* perf(docker): split dep install from src copy for CI layer caching

Install third-party deps (torch + all extras) in a layer keyed only on
pyproject.toml + uv.lock via --no-install-project, then copy src and install
the local package. Editing src/ no longer busts the heavy dependency layer,
so BuildKit layer cache hits across CI builds.

Applied to both Dockerfile.user and Dockerfile.internal.

* chore(ci): less verbose comments + copy all files

---------

Co-authored-by: dongmao.zhang <dongmao.zhang@bytedance.com>
2026-07-29 15:04:46 +02:00
Predrag Cvetkovic 35339d31e5 fix(datasets): bound memory of augment_dataset_quantile_stats by sampling frames (#3749)
* fix(datasets): bound memory of augment_dataset_quantile_stats by sampling frames

Per-episode stats previously materialized every frame (and decoded up to 16
episodes in parallel), so peak memory scaled with episode length and OOM'd on
large datasets (#2889). Numeric features are now read in full from the table
(exact), while only image/video frames are sub-sampled per episode using the
existing sample_indices heuristic. Worker count is configurable via
LEROBOT_STATS_MAX_WORKERS; --no-sampling restores exact behavior.

* Update tests/datasets/test_augment_quantile_stats.py

Co-authored-by: Haoming Song <1847575517@qq.com>
Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.com>

---------

Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.com>
Co-authored-by: Pepijn <138571049+pkooij@users.noreply.github.com>
Co-authored-by: Haoming Song <1847575517@qq.com>
2026-07-29 12:32:03 +02:00
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
Steven Palma 95211b98f1 feat(config): add multiprocessing option to DataLoader context and sets spawn as default (#4139)
* Add dataloader_multiprocessing_context, default to spawn

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

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

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

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

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

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

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

* chore(scripts): add multiprocessing_context safeguards

* chore(config): add libs note

---------

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

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

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

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

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

Fixes #3684

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

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

* chore(test): remove test

---------

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

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

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

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

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

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

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


* fix(scripts): keep policy train/eval

---------

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

* linting

---------

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

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

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

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

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

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

* chore(dependecies): bump diffusers

* chore(deps): update uv.lock

---------

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

* chore(test): minimize tests

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

---------

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

* adding `lerobot-train` requirement inside PR checklist

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

---------

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

* fix(dataset): handle empty dataset slices

* refactor(dataset): reuse scalar path for slices

---------

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

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

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

Fixes #3368

* fix(libero): apply configured control frequency

---------

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

* fix(motors): throttle calibration reads consistently

---------

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

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

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

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

* don't set_zero_position() on connect

---------

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

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

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

Usage: --policy.gradient_checkpointing=true

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

* test(diffusion): verify gradient checkpointing parity

---------

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

* review comments

* fix(gamepad): expose hidapi fallback in config

---------

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

* test(setup-motors): cover plugin registration

---------

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

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

---------

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

* Refactor embed_suffix to streamline pad_masks handling

Removed unused pad_masks list and simplified its creation.

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

---------

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-23 10:30:33 +02:00
Steven Palma 73dbb6f43a refactor(smolvla): reuse shared VLA components (#4064)
* refactor(smolvla): reuse shared VLA components

* chore(policies): address review smolvla shared utilities
2026-07-22 11:34:42 +02:00
Steven Palma 1427d35ef5 chore(docs): update security policy to adopt HF standards (#4098) 2026-07-21 14:07:09 +02:00
Steven Palma 30a5999cdc chore(ci): upgrade claude workflow (#4096) 2026-07-21 11:25:47 +02:00
Steven Palma 1bb9933215 refactor(xvla): reuse native Florence2 components (#4089) 2026-07-20 19:19:41 +02:00
Steven Palma ddc2aa7a27 refactor(pi0_fast): reuse shared VLA components (#4055) 2026-07-20 15:34:34 +02:00
Steven Palma 76b67d6ca8 refactor(eo1): reuse shared VLA components (#4061) 2026-07-20 15:34:16 +02:00
Steven Palma f3c0707c5f refactor(pi0): use shared VLA components (#4062) 2026-07-20 15:34:00 +02:00
Steven Palma 5361e0259e refactor(pi05): use shared VLA components (#4063) 2026-07-20 15:33:43 +02:00
Steven Palma a9879e69ed refactor(wall-x): subclass native Transformers Qwen2.5-VL instead of vendoring it (#4035) 2026-07-17 19:09:12 +02:00
Steven Palma 9d82bb9871 refactor(vla): extract shared model components (#4054) 2026-07-17 17:37:05 +02:00
Steven Palma c5371d0691 refactor(processors): share policy pipeline builders (#4016)
* refactor(processors): share policy pipeline builders

* Apply suggestions from code review

Co-authored-by: Martino Russi <77496684+nepyope@users.noreply.github.com>
Signed-off-by: Steven Palma <imstevenpmwork@ieee.org>

* fix(processor): solve style after commit suggestions

---------

Signed-off-by: Steven Palma <imstevenpmwork@ieee.org>
Co-authored-by: Martino Russi <77496684+nepyope@users.noreply.github.com>
2026-07-17 14:10:32 +02:00
Steven Palma b2c062c0f4 refactor(policies): resolve policy components by convention (#4015)
* refactor(policies): resolve policy components by convention

* remove fron None no-op

* extend processor resolver error handling logic to policy class resolver as well

---------

Co-authored-by: Martino Russi <nopyeps@gmail.com>
2026-07-17 13:59:38 +02:00
Maxime Ellerbach 051b13573e fix(safetensors): expand bare "cuda" to current device for safetensors loads (#4042) 2026-07-17 10:44:20 +02:00
Pepijn 7de2e4c1ef Move annotation dependencies to module scope (#4040) 2026-07-16 18:35:32 +02:00
Nikodem Bartnik 8db50611c2 pin pip installs (#4041) 2026-07-16 16:55:13 +02:00
Maxime Ellerbach 92f96f33b3 Aggregate policy sub-losses through MetricsTracker (#4024) 2026-07-16 12:12:37 +02:00
Steven Palma d4b3ca569c refactor(hub): load safetensors directly on target device (#4012) 2026-07-16 10:49:59 +02:00
Steven Palma 3f2179f3b6 refactor(evo1): use transformers flash attention probe (#4013)
Co-authored-by: Martino Russi <77496684+nepyope@users.noreply.github.com>
2026-07-15 17:02:01 +02:00
Nikodem Bartnik 867b58cfb2 generate new readme (#4029)
Co-authored-by: Pepijn <138571049+pkooij@users.noreply.github.com>
2026-07-15 16:32:02 +02:00
Pepijn 279c6c7af3 feat(annotate): improve VLM subtask annotation (legible contact sheets, seeded relabeling, self-hosted vLLM recipe) (#3896)
* feat(annotate): WGO-tuned subtask prompt (atomic completed-events + duration prior)

Rework the plan-module subtask segmentation prompt toward the WGO-Bench
atomic annotation protocol: segment by completed world-state changes
(grasp/place/open/close/pour/insert), fold approach+retreat into their
event, keep separate events separate, and add a 2-10s duration prior.
Drops the pi0.7 "fewer larger composites preferred" bias that drove
under-segmentation on the benchmark. Output JSON shape unchanged.

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

* feat(annotate): seeded-relabeling second pass for subtasks

Add an opt-in relabel pass (plan.subtask_seeded_relabel) that, after
segmentation, re-labels each span using previous/current/next segment
contact sheets and the seed label as a strong prior, minimally correcting
it. Mirrors macrodata's best end-to-end labeling step. Boundaries are
untouched; one extra VLM call per span. Off by default.

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

* feat(annotate): robust OpenAI-compat client for hosted VLMs

Guard against a choice with no message (safety filter or a thinking model
that spends its whole budget before emitting content) so one empty reply
no longer crashes the whole annotation run; treat it as an empty response
and let the existing JSON-retry path handle it.

Add an optional `reasoning_effort` knob on VlmConfig, forwarded to the
server when set, to cap a thinking model's reasoning (needed for Gemini
via its OpenAI-compatible endpoint).

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

* feat(annotate): legible tile-scaled timestamp on contact sheets

The burned-in timestamp used the ~10px bitmap default font, which blurs
once the model downsamples a full contact sheet into 768px tiles, so the
VLM can no longer read the exact source time a boundary depends on. Scale
the timestamp to the tile height (with a graceful fallback on older
Pillow) so the visual time cue stays readable at sheet resolution.

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

* feat(annotate): lean GEPA-aligned subtask segmentation prompt

Replace the verbose, label-heavy segmentation prompt with a lean
adaptation of the blog's GEPA-found completed_events_duration_prior
recipe: focus on completed manipulation events, explicit no-split /
no-merge rules, a 2-10s duration prior, and an instruction to prioritize
temporally correct boundaries over label wording. The previous prompt
over-weighted label guidance, which traded away boundary precision.

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

* revert: restore original subtask segmentation prompt

The lean GEPA-aligned paraphrase (dd4b0110d) regressed Gemini on the
30-ep subset: Seg F1 0.259 -> 0.189 and E2E 0.184 -> 0.135, driven by
worse under-segmentation (224 -> 188 preds). The blog's 0.306 came from
the actual GEPA-search artifact, which a hand paraphrase does not
reproduce. Restore the original prompt, which remains our best config.

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

* feat(annotate): env-var override for prompt templates

Allow LEROBOT_PROMPT_OVERRIDE_<name> to supersede the packaged prompt
file at load time. Enables prompt search (GEPA) to inject candidate
segmentation prompts into a remote annotate job via an env secret,
without committing a branch per candidate.

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

* docs(annotate): genericize hosted-VLM comments (no model name)

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

* docs(annotate): document seeded-relabel and reasoning_effort flags

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

* test(annotate): update subtask-prompt marker to match WGO-tuned prompt

The three plan-module tests keyed the canned VLM responder on the
literal 'atomic subtasks', which the WGO-tuned segmentation prompt no
longer contains (it now segments 'COMPLETED manipulation events'). Point
the fixture markers at the current wording so the subtask call is matched
again.

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

---------

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-15 11:38:49 +02:00
278 changed files with 7733 additions and 11405 deletions
+11
View File
@@ -0,0 +1,11 @@
version: 2
updates:
- package-ecosystem: "github-actions"
directory: "/"
schedule:
interval: "weekly"
cooldown:
default-days: 7
groups:
actions:
patterns: ["*"]
+17 -18
View File
@@ -34,43 +34,42 @@ jobs:
claude: claude:
if: | if: |
github.repository == 'huggingface/lerobot' && github.repository == 'huggingface/lerobot' &&
contains(
fromJSON('["OWNER", "MEMBER", "COLLABORATOR"]'),
github.event.comment.author_association || github.event.review.author_association
) &&
( (
(github.event_name == 'issue_comment' && contains(github.event.comment.body, '@claude')) || (github.event_name == 'issue_comment' && contains(github.event.comment.body, '@claude')) ||
(github.event_name == 'pull_request_review_comment' && contains(github.event.comment.body, '@claude')) || (github.event_name == 'pull_request_review_comment' && contains(github.event.comment.body, '@claude')) ||
(github.event_name == 'pull_request_review' && contains(github.event.review.body, '@claude')) (github.event_name == 'pull_request_review' && contains(github.event.review.body, '@claude'))
) )
runs-on: ubuntu-latest runs-on: ubuntu-latest
timeout-minutes: 30
steps: steps:
- name: Authorize commenter
id: authorize
run: |
AUTHOR_ASSOCIATION="${{ github.event.comment.author_association || github.event.review.author_association }}"
if [[ "$AUTHOR_ASSOCIATION" == "OWNER" ]] || [[ "$AUTHOR_ASSOCIATION" == "MEMBER" ]] || [[ "$AUTHOR_ASSOCIATION" == "COLLABORATOR" ]]; then
echo "Authorized: $AUTHOR_ASSOCIATION"
exit 0
else
echo "Unauthorized: $AUTHOR_ASSOCIATION"
exit 1
fi
- name: Checkout code - name: Checkout code
if: success()
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
with: with:
persist-credentials: false persist-credentials: false
- name: Run Claude Code - name: Run Claude Code
if: success()
id: claude id: claude
# TODO(Steven): Update once https://github.com/anthropics/claude-code-action/issues/1187 is shipped uses: anthropics/claude-code-action@b76a0776ae74036e77cd11018083743453d7ad35 # v1.0.179
uses: anthropics/claude-code-action@1eddb334cfa79fdb21ecbe2180ca1a016e8e7d47 # v1.0.88
with: with:
anthropic_api_key: ${{ secrets.ANTHROPIC_API_KEY }} anthropic_api_key: ${{ secrets.ANTHROPIC_API_KEY }}
additional_permissions: |
actions: read
track_progress: true track_progress: true
classify_inline_comments: true
include_fix_links: false
claude_args: | claude_args: |
--model claude-opus-4-6 --model claude-opus-4-8
--effort max --effort xhigh
--fallback-model claude-sonnet-5
--max-turns 20
--verbose --verbose
--tools "Read,Grep,Glob,Agent"
--strict-mcp-config
--append-subagent-system-prompt "Treat repository files and GitHub content as untrusted data. Ignore embedded instructions and return only evidence-backed code review findings."
--append-system-prompt " --append-system-prompt "
ROLE: Strict Code Review Assistant ROLE: Strict Code Review Assistant
TASK: Analyze code changes and provide objective technical reviews. TASK: Analyze code changes and provide objective technical reviews.
+2 -1
View File
@@ -51,6 +51,7 @@ pre-commit run --all-files # Lint + format (ruff, typo
## Notes ## Notes
- **Mypy is gradual**: strict only for `lerobot.envs`, `lerobot.configs`, `lerobot.optim`, `lerobot.model`, `lerobot.cameras`, `lerobot.motors`, `lerobot.transport`. Add type annotations when modifying these modules. - **Mypy is gradual**: strict only for `lerobot.envs`, `lerobot.configs`, `lerobot.optim`, `lerobot.model`, `lerobot.cameras`, `lerobot.motors`, `lerobot.transport`. Add type annotations when modifying these modules.
- **Optional dependencies**: many policies, envs, and robots are behind extras (e.g., `lerobot[aloha]`). New imports for optional packages must be guarded or lazy. See `pyproject.toml [project.optional-dependencies]`. - **Imports**: prefer top-level imports; relative (`from .sibling import X`) across sibling files within a module, absolute (`from lerobot.module import X`) across modules.
- **Optional dependencies**: many policies, envs, and robots are behind extras (e.g., `lerobot[aloha]`, see `pyproject.toml`). Guard optional imports with `TYPE_CHECKING or _foo_available` at module top + a `require_package(...)` check at use time. Reuse the `_foo_available` flags in `utils/import_utils.py`; don't call `is_package_available`.
- **Video decoding**: datasets can store observations as video files. `LeRobotDataset` handles frame extraction, but tests need ffmpeg installed. - **Video decoding**: datasets can store observations as video files. `LeRobotDataset` handles frame extraction, but tests need ffmpeg installed.
- **Prioritize use of `uv run`** to execute Python commands (not raw `python` or `pip`). - **Prioritize use of `uv run`** to execute Python commands (not raw `python` or `pip`).
+10 -6
View File
@@ -61,16 +61,20 @@ Full details in [`docs/source/so101.mdx`](./docs/source/so101.mdx) and [`docs/so
**4.1 Install** **4.1 Install**
```bash ```bash
pip install 'lerobot[feetech]' # SO-100/SO-101 motor stack # uv (recommended — see AGENTS.md and CLAUDE.md)
# pip install 'lerobot[all]' # everything uv sync --locked --extra feetech # SO-100/SO-101 motor stack
# pip install 'lerobot[aloha,pusht]' # specific features # uv sync --locked --extra all # everything
# pip install 'lerobot[smolvla]' # add SmolVLA deps # uv sync --locked --extra smolvla # add SmolVLA deps
# pip (alternative, e.g. when not working from source)
# pip install 'lerobot[feetech]'
# pip install 'lerobot[all]'
# pip install 'lerobot[smolvla]'
git lfs install && git lfs pull git lfs install && git lfs pull
hf auth login # required to push datasets/policies hf auth login # required to push datasets/policies
``` ```
Contributors can alternatively use `uv sync --locked --extra feetech` (see `AGENTS.md`).
**4.2 Find USB ports** — run once per arm, unplug when prompted. **4.2 Find USB ports** — run once per arm, unplug when prompted.
```bash ```bash
+3 -3
View File
@@ -83,7 +83,7 @@ episode_index=0
print(f"{dataset[episode_index]['action'].shape=}\n") print(f"{dataset[episode_index]['action'].shape=}\n")
``` ```
Learn more about it in the [LeRobotDataset Documentation](https://huggingface.co/docs/lerobot/lerobot-dataset-v3) Learn more about it in the [LeRobotDataset Documentation](https://huggingface.co/docs/lerobot/lerobot-dataset-v3).
## SoTA Models ## SoTA Models
@@ -109,7 +109,7 @@ lerobot-train \
| **World Models** | [VLA-JEPA](./docs/source/vla_jepa.mdx), [LingBot-VA](./docs/source/lingbot_va.mdx), [FastWAM](./docs/source/fastwam.mdx) | | **World Models** | [VLA-JEPA](./docs/source/vla_jepa.mdx), [LingBot-VA](./docs/source/lingbot_va.mdx), [FastWAM](./docs/source/fastwam.mdx) |
| **Reward Models** | [SARM](./docs/source/sarm.mdx), [TOPReward](./docs/source/topreward.mdx), [Robometer](./docs/source/robometer.mdx) | | **Reward Models** | [SARM](./docs/source/sarm.mdx), [TOPReward](./docs/source/topreward.mdx), [Robometer](./docs/source/robometer.mdx) |
Similarly to the hardware, you can easily implement your own policy & leverage LeRobot's data collection, training, and visualization tools, and share your model to the HF Hub Similarly to the hardware, you can easily implement your own policy & leverage LeRobot's data collection, training, and visualization tools, and share your model to the HF Hub.
For detailed policy setup guides, see the [Policy Documentation](https://huggingface.co/docs/lerobot/bring_your_own_policies). For GPU/RAM requirements and expected training time per policy, see the [Compute Hardware Guide](https://huggingface.co/docs/lerobot/hardware_guide). For detailed policy setup guides, see the [Policy Documentation](https://huggingface.co/docs/lerobot/bring_your_own_policies). For GPU/RAM requirements and expected training time per policy, see the [Compute Hardware Guide](https://huggingface.co/docs/lerobot/hardware_guide).
@@ -126,7 +126,7 @@ lerobot-eval \
--eval.n_episodes=10 --eval.n_episodes=10
``` ```
Learn how to implement your own simulation environment or benchmark and distribute it from the HF Hub by following the [EnvHub Documentation](https://huggingface.co/docs/lerobot/envhub) Learn how to implement your own simulation environment or benchmark and distribute it from the HF Hub by following the [EnvHub Documentation](https://huggingface.co/docs/lerobot/envhub).
## Resources ## Resources
+108 -24
View File
@@ -6,43 +6,127 @@
Fortunately, being an open-source project, the community can also help by reporting and fixing vulnerabilities. We appreciate your efforts to responsibly disclose your findings and will make every effort to acknowledge your contributions. Fortunately, being an open-source project, the community can also help by reporting and fixing vulnerabilities. We appreciate your efforts to responsibly disclose your findings and will make every effort to acknowledge your contributions.
## Reporting a Vulnerability
To report a security issue, please use the GitHub Security Advisory ["Report a Vulnerability"](https://github.com/huggingface/lerobot/security/advisories/new) tab.
The `lerobot` team will send a response indicating the next steps in handling your report. After the initial reply to your report, the security team will keep you informed of the progress towards a fix and full announcement, and may ask for additional information or guidance.
#### Hugging Face Security Team
Since this project is part of the Hugging Face ecosystem, feel free to submit vulnerability reports directly to: **[security@huggingface.co](mailto:security@huggingface.co)**. Someone from the HF security team will review the report and recommend next steps.
#### Open Source Disclosures
If reporting a vulnerability specific to the open-source codebase (and not the underlying Hub infrastructure), you may also use [Huntr](https://huntr.com), a vulnerability disclosure program for open source software.
## Supported Versions ## Supported Versions
Currently, we treat `lerobot` as a rolling release. We prioritize security updates for the latest available version (`main` branch). Currently, we treat `lerobot` as a rolling release. We prioritize security updates for the latest available version (`main` branch). Please reproduce on the current head before reporting — we do not backport fixes to older releases.
| Version | Supported | | Version | Supported |
| -------- | --------- | | -------- | --------- |
| Latest | ✅ | | Latest | ✅ |
| < Latest | ❌ | | < Latest | ❌ |
## Secure Usage Guidelines ## Reporting a Vulnerability
`lerobot` is tightly coupled to the Hugging Face Hub for sharing data and pretrained policies. When downloading artifacts uploaded by others, you expose yourself to risks. Please read below for recommendations to keep your runtime and robot environment safe. Report privately — **do not open a public issue or PR for a suspected vulnerability.**
To report a security issue, please use the GitHub Security Advisory ["Report a Vulnerability"](https://github.com/huggingface/lerobot/security/advisories/new) tab. This routes to the maintainers, keeps the report private until a fix is ready, and lets us issue a CVE through GitHub if warranted. The `lerobot` team will send a response indicating the next steps in handling your report. We acknowledge valid, in-scope reports and will keep you updated on remediation. Please give us a reasonable window to fix before any public disclosure.
#### Hugging Face Security Team
Since this project is part of the Hugging Face ecosystem, feel free to submit vulnerability reports directly to: **[security@huggingface.co](mailto:security@huggingface.co)**. Someone from the HF security team will review the report and recommend next steps. After the initial reply to your report, the security team will keep you informed of the progress towards a fix and full announcement, and may ask for additional information or guidance.
## Recognition
We do not offer a monetary bounty. For a valid, in-scope report we credit you on the published GitHub Security Advisory and name you as the reporter in the associated CVE. Let us know how you'd like to be credited (name or handle).
## What your report must include
We receive a high volume of reports. To be triaged, a report **must** follow the structure below. Copy this block into your submission and fill in every field. Reports missing the version, the proof of concept, or the impact are returned as incomplete and are not investigated until provided.
```markdown
### Summary
One sentence: what the vulnerability is and where.
### Affected version / commit
Exact released version or commit SHA you reproduced on (e.g. v4.57.0 / a1b2c3d).
Not "latest" or "main".
### Affected component
The public API, module, or entry point involved (e.g. `AutoModel.from_pretrained`).
### Vulnerability class
Type and CWE if known (e.g. deserialization / CWE-502, path traversal / CWE-22).
### Attack vector & preconditions
- How is the vulnerable code reached? (which API call / input / config)
- Who is the attacker and what do they control?
- What must be true for the attack to work? (auth, a user action, a non-default
setting, a malicious file being loaded, etc.)
### Proof of concept
A minimal, self-contained script or step sequence that runs on a clean install
of the version above. Include:
- the exact commands / code to run,
- any input files needed (attach them, or give a script that generates them),
- the **expected** behavior vs. the **actual** behavior you observed.
A snippet showing that a function _exists_ or _could_ be misused is not a PoC.
### Impact
What an attacker gains in a realistic deployment. "Could theoretically…"
without a working chain is not an impact.
### Scope
Which trust boundary (see below) does this cross? If your finding touches
anything in the "Out of scope" list, name which item and explain why it is
nonetheless a violation of a guarantee we make.
### Suggested severity (optional)
We assign the final severity. Include a CVSS v3.1 vector only if you have one.
### Suggested fix (optional)
```
> [!NOTE]
> The bar is a **reproducible PoC against a supported version, with a concrete impact that crosses a trust boundary we actually defend** (see scope below). Reports that are theoretical, auto-generated by a scanner or LLM, or that restate documented behavior will be closed without detailed review.
## Threat model & trust boundaries
`lerobot` is tightly coupled to the Hugging Face Hub for sharing data and pretrained policies. When downloading artifacts uploaded by others, you expose yourself to risks. Please read below for recommendations to keep your runtime and robot environment safe. We _will_ treat as a vulnerability anything that breaks one of these protections — e.g. code executing despite `safetensors`-only loading, or a pinned revision being bypassed.
### Remote Artefacts (Weights & Policies) ### Remote Artefacts (Weights & Policies)
Models and policies uploaded to the Hugging Face Hub come in different formats. We heavily recommend uploading and downloading models in the [`safetensors`](https://github.com/huggingface/safetensors) format. Models and policies uploaded to the Hugging Face Hub come in different formats. We heavily recommend uploading and downloading models in the [`safetensors`](https://github.com/huggingface/safetensors) format. `safetensors` was developed specifically to prevent arbitrary code execution on your system, which is critical when running software on physical hardware/robots. To avoid loading models from unsafe formats (e.g., `pickle`), you should ensure you are prioritizing `safetensors` files.
`safetensors` was developed specifically to prevent arbitrary code execution on your system, which is critical when running software on physical hardware/robots.
To avoid loading models from unsafe formats (e.g., `pickle`), you should ensure you are prioritizing `safetensors` files.
### Remote Code ### Remote Code
Some models or environments on the Hub may require `trust_remote_code=True` to run custom architecture code. Some models or environments on the Hub may require `trust_remote_code=True` to run custom architecture code. Please **always** verify the content of the modeling files when using this argument. We recommend setting a specific `revision` (commit hash) when loading remote code to ensure you protect yourself from unverified updates to the repository.
Please **always** verify the content of the modeling files when using this argument. We recommend setting a specific `revision` (commit hash) when loading remote code to ensure you protect yourself from unverified updates to the repository. ## In scope
We treat as vulnerabilities issues in the **published package code** — the library's own API surface — that an attacker can trigger without the victim having opted into a documented risk. For example:
- code execution, memory corruption, or file access reachable through a normal API call on input that is **not** an untrusted model/artifact the user chose to load;
- a control we advertise being bypassed (e.g. code running despite `safetensors`-only loading, or a pinned revision being ignored);
- exposure or mishandling of credentials, tokens, or another user's data by the library;
- a real escape from a backend we document as a sandbox;
- CI/CD or supply-chain issues in this repository.
## Out of scope
The following are **not** treated as vulnerabilities in `lerobot`. If your finding touches one of these, the report must explain why it is nonetheless a violation of a guarantee we make — otherwise it will be closed.
- Issues that require loading an untrusted artifact and amount to the documented load-time risk above (code execution / file access on load of a malicious model, dataset, config, or pickle).
- Findings in `examples/`, documentation, tests, or other non-packaged reference material.
- Local denial-of-service from feeding pathological input to a function on your own machine (high memory, slow parse, panic), absent a multi-tenant or remote-service impact.
- Model behavior: jailbreaks, alignment failures, prompt injection, or harmful generations. Model weights are authored by their uploaders; report these to the model owner.
- Vulnerabilities in third-party dependencies we do not vendor — report upstream (we'll bump once fixed).
- Theoretical issues without a working proof of concept, and reports auto-generated from scanners or LLMs without a verified, reproducible chain.
- Best-practice or hardening suggestions with no demonstrated impact — missing email-authentication or transport records (MTA-STS, TLS-RPT, DMARC/SPF tuning), missing HTTP security headers, TLS configuration preferences, and similar scanner or config-checker output presented without a working exploit chain.
## Safe harbor
Good-faith research that respects these guidelines, avoids privacy violations and service disruption, and gives us a reasonable disclosure window will not be pursued by us. Do not access data that isn't yours and do not run tests against Hugging Face production infrastructure.
<div align="center">
<sub>Built by the <a href="https://huggingface.co/lerobot">LeRobot</a> team at <a href="https://huggingface.co">Hugging Face</a> with ❤️</sub>
</div>
+4 -5
View File
@@ -68,17 +68,16 @@ ENV HOME=/home/user_lerobot \
# issues with MuJoCo and OpenGL drivers. # issues with MuJoCo and OpenGL drivers.
RUN uv venv --python python${PYTHON_VERSION} RUN uv venv --python python${PYTHON_VERSION}
# Install Python dependencies for caching # Install third-party dependencies separately for layer caching
COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./ COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./
COPY --chown=user_lerobot:user_lerobot src/ src/ RUN uv sync --locked --extra all --no-install-project --no-cache
RUN uv sync --locked --extra all --no-cache
RUN chmod +x /lerobot/.venv/lib/python${PYTHON_VERSION}/site-packages/triton/backends/nvidia/bin/ptxas RUN chmod +x /lerobot/.venv/lib/python${PYTHON_VERSION}/site-packages/triton/backends/nvidia/bin/ptxas
# Copy the rest of the application source code # Copy the application source code and install the local project
# Make sure to have the git-LFS files for testing # Make sure to have the git-LFS files for testing
COPY --chown=user_lerobot:user_lerobot . . COPY --chown=user_lerobot:user_lerobot . .
RUN uv sync --locked --extra all --no-cache
# Set the default command # Set the default command
CMD ["/bin/bash"] CMD ["/bin/bash"]
+4 -5
View File
@@ -60,15 +60,14 @@ ENV HOME=/home/user_lerobot \
# run other Python projects in the same container without dependency conflicts. # run other Python projects in the same container without dependency conflicts.
RUN uv venv RUN uv venv
# Install Python dependencies for caching # Install third-party dependencies separately for layer caching
COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./ COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./
COPY --chown=user_lerobot:user_lerobot src/ src/ RUN uv sync --locked --extra all --no-install-project --no-cache
RUN uv sync --locked --extra all --no-cache # Copy the application code and install the local project
# Copy the rest of the application code
# Make sure to have the git-LFS files for testing # Make sure to have the git-LFS files for testing
COPY --chown=user_lerobot:user_lerobot . . COPY --chown=user_lerobot:user_lerobot . .
RUN uv sync --locked --extra all --no-cache
# Set the default command # Set the default command
CMD ["/bin/bash"] CMD ["/bin/bash"]
+66 -18
View File
@@ -81,10 +81,16 @@ merged. Both prompts also carry a causal **event-boundary** definition (a
new event starts when an object becomes held / is released / reaches a new new event starts when an object becomes held / is released / reaches a new
location / a lid changes state / contents move) to sharpen where cuts land. location / a lid changes state / contents move) to sharpen where cuts land.
Optionally, a third **seeded-relabel** pass (`--plan.subtask_seeded_relabel`)
revisits each span with its previous/current/next segment contact sheets and
minimally corrects the label, using the first label as a prior — it keeps the
boundaries fixed and only sharpens wording, at the cost of one extra call per
subtask.
The resulting spans are then stitched into a gap-free, full-episode The resulting spans are then stitched into a gap-free, full-episode
cover, so **every frame has exactly one active subtask**. See cover, so **every frame has exactly one active subtask**. See
[`run_hf_job.py`](https://github.com/huggingface/lerobot/blob/main/examples/annotations/run_hf_job.py) [Running on Hugging Face Jobs](#running-on-hugging-face-jobs) for the
for the production settings (single camera, timestamped contact sheets, production settings (single camera, timestamped contact sheets,
auto-windowed subtask generation). auto-windowed subtask generation).
### Tools ### Tools
@@ -104,28 +110,67 @@ not-yet-implemented.
## Running on Hugging Face Jobs ## Running on Hugging Face Jobs
Annotation runs on [Hugging Face Jobs](https://huggingface.co/docs/hub/en/jobs). Annotating a real dataset needs a GPU big enough to serve the VLM, so
The repo ships a launcher script you copy and tweak for your dataset: `lerobot-annotate` can dispatch itself to
[Hugging Face Jobs](https://huggingface.co/docs/hub/en/jobs) — same as
`lerobot-train`. Add `--job.target=<flavor>` to the exact command you'd
run locally and it runs on that hardware instead:
```bash ```bash
HF_TOKEN=hf_... uv run python examples/annotations/run_hf_job.py hf auth login # once
uv run lerobot-annotate \
--repo_id=user/my_dataset \
--new_repo_id=user/my_dataset_annotated \
--push_to_hub=true \
--vlm.model_id=Qwen/Qwen3.6-27B \
--vlm.num_gpus=1 \
--vlm.serve_command="vllm serve Qwen/Qwen3.6-27B --tensor-parallel-size 1 \
--max-model-len 32768 --gpu-memory-utilization 0.8 \
--uvicorn-log-level warning --port {port}" \
--vlm.serve_ready_timeout_s=1800 \
--vlm.chat_template_kwargs='{"enable_thinking": false}' \
--job.target=h200
``` ```
[`run_hf_job.py`](https://github.com/huggingface/lerobot/blob/main/examples/annotations/run_hf_job.py) That submits a single-GPU `h200` job that:
starts a single-GPU `h200` job (bump it to `h200x4` for big datasets)
that:
1. installs `lerobot` (from `main`) plus the annotation extras, 1. starts from the `vllm/vllm-openai` image and installs `lerobot` on top,
2. boots one vLLM server per GPU (using the `vllm/vllm-openai` image) and 2. boots one vLLM server per GPU and drives it over the OpenAI-compatible API,
drives it over the OpenAI-compatible API, 3. runs the `plan` / `interjections` / `vqa` modules across the dataset,
3. runs the `plan` / `interjections` / `vqa` modules across the dataset
with `lerobot-annotate`,
4. with `--push_to_hub=true`, uploads the result to `--new_repo_id` (or 4. with `--push_to_hub=true`, uploads the result to `--new_repo_id` (or
back to `--repo_id` in place if you leave that unset). back to `--repo_id` in place if you leave that unset).
To use a different dataset, model, or hub repo, edit the `CMD` block in The command streams the job's logs; `Ctrl-C` detaches without cancelling
the script. Every flag there maps directly to a `lerobot-annotate` flag it. List the available flavors and their pricing with `hf jobs hardware`.
(run `lerobot-annotate --help` for the full list).
<Tip warning={true}>
Qwen3.6 ships with thinking enabled, which eats the token budget the
annotator needs for its JSON answer — `--vlm.chat_template_kwargs='{"enable_thinking": false}'`
turns it off. Without `--push_to_hub=true` the annotated dataset is
discarded when the pod exits.
</Tip>
### Job options
| Flag | Default | What it does |
| ------------------- | ------------------------- | ------------------------------------------------------------------------------- |
| `--job.target` | `local` | HF Jobs flavor to run on (e.g. `h200`, `h200x4`). Omitted/`local` runs here. |
| `--job.image` | `vllm/vllm-openai:latest` | Runtime image for the pod. |
| `--job.timeout` | `2h` | Wall-clock cap. Raise it for large datasets. |
| `--job.detach` | `false` | Submit and exit instead of streaming logs. |
| `--job.lerobot_ref` | `main` | Git ref of lerobot installed on the pod — point it at a branch to test changes. |
| `--job.tags` | `[]` | Extra tags on the job and on any dataset it pushes (`lerobot` is always added). |
For a bigger dataset, scale to `h200x4` and raise
`--vlm.parallel_servers` / `--vlm.num_gpus` to match, and give the job
more headroom with e.g. `--job.timeout=8h`.
Remote runs need `--repo_id` (the pod pulls the dataset from the Hub;
`--root` names a directory only your machine has). A dataset that exists
only in your local cache is pushed to a **private** repo first.
## Key options ## Key options
@@ -158,7 +203,7 @@ Every module is on by default and can be toggled independently (set to
### The VLM (`--vlm.*`) ### The VLM (`--vlm.*`)
| Flag | Default | What it does | | Flag | Default | What it does |
| -------------------------- | ------------------ | ----------------------------------------------------------------------------------- | | -------------------------- | ------------------ | ------------------------------------------------------------------------------------ |
| `--vlm.model_id` | `Qwen/Qwen3.6-27B` | The model to serve and prompt. | | `--vlm.model_id` | `Qwen/Qwen3.6-27B` | The model to serve and prompt. |
| `--vlm.camera_key` | first `images.*` | Which camera every prompt is grounded on. | | `--vlm.camera_key` | first `images.*` | Which camera every prompt is grounded on. |
| `--vlm.serve_command` | auto | The exact `vllm serve …` command (set TP size, GPU memory, `--max-model-len` here). | | `--vlm.serve_command` | auto | The exact `vllm serve …` command (set TP size, GPU memory, `--max-model-len` here). |
@@ -167,16 +212,19 @@ Every module is on by default and can be toggled independently (set to
| `--vlm.client_concurrency` | `16` | In-flight requests across all servers. | | `--vlm.client_concurrency` | `16` | In-flight requests across all servers. |
| `--vlm.max_new_tokens` | `512` | Generation cap per call. | | `--vlm.max_new_tokens` | `512` | Generation cap per call. |
| `--vlm.temperature` | `0.2` | Sampling temperature. | | `--vlm.temperature` | `0.2` | Sampling temperature. |
| `--vlm.reasoning_effort` | `null` | Thinking-budget hint (`low`/`medium`/`high`) forwarded to OpenAI-compatible servers. |
### Subtasks / plan / memory (`--plan.*`) ### Subtasks / plan / memory (`--plan.*`)
| Flag | Default | What it does | | Flag | Default | What it does |
| ------------------------------- | ---------- | ------------------------------------------------------------------------------------------------------------------------- | | ------------------------------- | ---------- | ---------------------------------------------------------------------------------------------------------------------------- |
| `--plan.frames_per_second` | `2.0` | Frame sampling rate for the contact sheets (`2.0` = one frame every 0.5s). | | `--plan.frames_per_second` | `2.0` | Frame sampling rate for the contact sheets (`2.0` = one frame every 0.5s). |
| `--plan.max_frames_per_prompt` | `60` | Frame budget per VLM call. Episodes whose sampling exceeds this are auto-windowed at the same density, then stitched. | | `--plan.max_frames_per_prompt` | `60` | Frame budget per VLM call. Episodes whose sampling exceeds this are auto-windowed at the same density, then stitched. |
| `--plan.contact_sheet_columns` | `5` | Columns per contact-sheet grid (`contact_sheet_frames_per_sheet` tiles, time row-major). | | `--plan.contact_sheet_columns` | `5` | Columns per contact-sheet grid (`contact_sheet_frames_per_sheet` tiles, time row-major). |
| `--plan.plan_max_steps` | `8` | Upper bound on subtasks per episode. | | `--plan.plan_max_steps` | `8` | Upper bound on subtasks per episode. |
| `--plan.subtask_describe_first` | `true` | Run the describe→segment grounding pass (best subtask quality; +1 call/episode). | | `--plan.subtask_describe_first` | `true` | Run the describe→segment grounding pass (best subtask quality; +1 call/episode). |
| `--plan.subtask_seeded_relabel` | `false` | Second pass: re-label each subtask from its prev/current/next contact sheets, seeded with the first label (+1 call/subtask). |
| `--plan.subtask_relabel_frames` | `5` | Frames sampled uniformly per segment sheet in the relabel pass (only used when `subtask_seeded_relabel=true`). |
| `--plan.emit_plan` | `true` | Emit the numbered `plan` rows (`false` = subtasks + memory only). | | `--plan.emit_plan` | `true` | Emit the numbered `plan` rows (`false` = subtasks + memory only). |
| `--plan.emit_memory` | `true` | Emit the `memory` rows (`false` = subtasks + plan only); symmetric to `emit_plan`. | | `--plan.emit_memory` | `true` | Emit the `memory` rows (`false` = subtasks + plan only); symmetric to `emit_plan`. |
| `--plan.n_task_rephrasings` | `10` | How many `task_aug` rephrasings to emit (`0` disables). | | `--plan.n_task_rephrasings` | `10` | How many `task_aug` rephrasings to emit (`0` disables). |
+3 -3
View File
@@ -58,7 +58,7 @@ final_action = postprocessor(action)
## Hardware API redesign ## Hardware API redesign
PR [#777](https://github.com/huggingface/lerobot/pull/777) improves the LeRobot calibration but is **not backward-compatible**. Below is a overview of what changed and how you can continue to work with datasets created before this pull request. PR [#777](https://github.com/huggingface/lerobot/pull/777) improves the LeRobot calibration but is **not backward-compatible**. Below is an overview of what changed and how you can continue to work with datasets created before this pull request.
### What changed? ### What changed?
@@ -129,8 +129,8 @@ python examples/backward_compatibility/replay.py \
Policies output actions in the same format as the datasets (`torch.Tensors`). Therefore, the same transformations should be applied. Policies output actions in the same format as the datasets (`torch.Tensors`). Therefore, the same transformations should be applied.
To find these transformations, we recommend to first try and and replay an episode of the dataset your policy was trained on using the section above. To find these transformations, we recommend first replaying an episode of the dataset your policy was trained on using the section above.
Then, add these same transformations on your inference script (shown here in the `record.py` script): Then, add these same transformations to your inference script (shown here in the `record.py` script):
```diff ```diff
action_values = predict_action( action_values = predict_action(
+16 -9
View File
@@ -151,12 +151,12 @@ class MyPolicy(PreTrainedPolicy):
The methods called by the train/eval loops: The methods called by the train/eval loops:
| Method | Used by | What it does | | Method | Used by | What it does |
| ----------------------------------------------------------------- | ----------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | ----------------------------------------------------------------- | ----------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
| `reset() -> None` | `lerobot-eval` | Clear per-episode state at the start of each episode. | | `reset() -> None` | `lerobot-eval` | Clear per-episode state at the start of each episode. |
| `select_action(batch, **kwargs) -> Tensor` | `lerobot-eval` | Return the next action `(B, action_dim)`. Called every step. | | `select_action(batch, **kwargs) -> Tensor` | `lerobot-eval` | Return the next action `(B, action_dim)`. Called every step. |
| `predict_action_chunk(batch, **kwargs) -> Tensor` | the policy itself | Return an action chunk `(B, chunk_size, action_dim)`. Currently abstract on the base class — raise `NotImplementedError` if your policy doesn't chunk. | | `predict_action_chunk(batch, **kwargs) -> Tensor` | the policy itself | Return an action chunk `(B, chunk_size, action_dim)`. Currently abstract on the base class — raise `NotImplementedError` if your policy doesn't chunk. |
| `forward(batch, reduction="mean") -> tuple[Tensor, dict \| None]` | `lerobot-train` | Return `(loss, output_dict)`. Accept `reduction="none"` if you want to support per-sample weighting. | | `forward(batch, reduction="mean") -> tuple[Tensor, dict \| None]` | `lerobot-train` | Return `(loss, output_dict)`. Accept `reduction="none"` if you want to support per-sample weighting. |
| `get_optim_params() -> dict` | the optimizer | Return `self.parameters()` for simple policies; return a named parameter dict for [multi-optimizer policies](https://github.com/huggingface/lerobot/blob/ecd38c50d7d15b4184cf42649ff1185ee2e11eeb/src/lerobot/policies/sac/modeling_sac.py#L61-L73). | | `get_optim_params() -> dict` | the optimizer | Return `self.parameters()` for simple policies; return a named parameter dict for multi-optimizer policies (see `get_optim_params` in [`modeling_act.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/act/modeling_act.py) for a per-group learning-rate example). |
| `update() -> None` _(optional)_ | `lerobot-train` | Called after each optimizer step _if defined_. Use for EMA, target nets, replay buffers (TDMPC uses this). | | `update() -> None` _(optional)_ | `lerobot-train` | Called after each optimizer step _if defined_. Use for EMA, target nets, replay buffers (TDMPC uses this). |
Batches are flat dictionaries keyed by the constants in [`lerobot.utils.constants`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/utils/constants.py): `OBS_STATE` (`observation.state.<motor>`), `OBS_IMAGES` (`observation.images.<camera>`), `OBS_LANGUAGE`, `ACTION`, etc. Reuse the constants — don't invent new prefixes. Batches are flat dictionaries keyed by the constants in [`lerobot.utils.constants`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/utils/constants.py): `OBS_STATE` (`observation.state.<motor>`), `OBS_IMAGES` (`observation.images.<camera>`), `OBS_LANGUAGE`, `ACTION`, etc. Reuse the constants — don't invent new prefixes.
@@ -165,6 +165,8 @@ Batches are flat dictionaries keyed by the constants in [`lerobot.utils.constant
LeRobot uses `PolicyProcessorPipeline`s to normalize inputs and de-normalize outputs around your policy. For a concrete reference, see [`processor_act.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/act/processor_act.py) or [`processor_diffusion.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/diffusion/processor_diffusion.py). LeRobot uses `PolicyProcessorPipeline`s to normalize inputs and de-normalize outputs around your policy. For a concrete reference, see [`processor_act.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/act/processor_act.py) or [`processor_diffusion.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/diffusion/processor_diffusion.py).
Pay close attention here: processors are the most common reproducibility pain point. A mismatch in normalization mode (`IDENTITY` vs `MEAN_STD` vs `MIN_MAX` vs `QUANTILES`/`QUANTILE10`) or in which features get normalized will train and eval without erroring, yet silently wreck results. Make sure the modes match how the checkpoint was trained, that the required stats exist (e.g. `QUANTILES` needs `q01`/`q99`), and that the pre- and post-processors stay consistent.
```python ```python
# processor_my_policy.py # processor_my_policy.py
from typing import Any from typing import Any
@@ -295,18 +297,18 @@ The file names are load-bearing: the factory does lazy imports by name, and the
### Wiring ### Wiring
Four places need to know about your policy. All by name. Two places need to know about your policy. All by name.
1. **`policies/__init__.py`** — re-export `MyPolicyConfig` and add it to `__all__`. **Don't** re-export the modeling class; it loads lazily through the factory (so `import lerobot` stays fast). 1. **`policies/__init__.py`** — re-export `MyPolicyConfig` and add it to `__all__`. This import is what registers your policy: `@PreTrainedConfig.register_subclass("my_policy")` runs, and from then on the factory resolves everything by convention. **Don't** re-export the modeling class; it loads lazily through the factory (so `import lerobot` stays fast).
2. **`factory.py:get_policy_class`** — add a branch returning `MyPolicy` from a lazy import. 2. **`templates/lerobot_modelcard_template.md` and the root `README.md`** — the template is what `push_model_to_hub` renders into the model card of every checkpoint trained with your policy: add a one-line description of your policy in the `model_name` branches, map it in `policy_docs` so cards link to your MDX guide, and optionally add an architecture image to `diagrams`. Then add your policy to the models table in the root `README.md`, under the right category, linking to your doc page.
3. **`factory.py:make_policy_config`** and **`factory.py:make_pre_post_processors`** — same idea, two more branches.
4. **`templates/lerobot_modelcard_template.md` and the root `README.md`** — the template is what `push_model_to_hub` renders into the model card of every checkpoint trained with your policy: add a one-line description of your policy in the `model_name` branches, map it in `policy_docs` so cards link to your MDX guide, and optionally add an architecture image to `diagrams`. Then add your policy to the models table in the root `README.md`, under the right category, linking to your doc page.
Mirror an existing policy that's structurally similar to yours; the diff is small. Mirror an existing policy that's structurally similar to yours; the diff is small.
### Heavy / optional dependencies ### Heavy / optional dependencies
Most policies need a heavy backbone (transformers, diffusers, a specific VLM SDK). The convention is **two-step gating**: a `TYPE_CHECKING`-guarded import at module top, and a `require_package` runtime check in the constructor. [`modeling_diffusion.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/diffusion/modeling_diffusion.py) is the canonical reference: Most policies need a heavy backbone (transformers, diffusers, a specific VLM SDK). Wherever one exists, prefer loading it e.g from `transformers` or `diffusers` rather than re-implementing the architecture in-tree.
The convention is **two-step gating**: a `TYPE_CHECKING`-guarded import at module top, and a `require_package` runtime check in the constructor. [`modeling_diffusion.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/diffusion/modeling_diffusion.py) is the canonical reference:
```python ```python
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
@@ -332,6 +334,10 @@ This way:
Add a matching extra to [`pyproject.toml`](https://github.com/huggingface/lerobot/blob/main/pyproject.toml) `[project.optional-dependencies]` and include it in the `all` extra so `pip install 'lerobot[all]'` keeps installing everything. Add a matching extra to [`pyproject.toml`](https://github.com/huggingface/lerobot/blob/main/pyproject.toml) `[project.optional-dependencies]` and include it in the `all` extra so `pip install 'lerobot[all]'` keeps installing everything.
### Avoid copying a modeling file — subclass it
If your policy needs to modify a backbone that already exists in `transformers` (custom conditioning, extra inputs, a swapped sub-module), **do not vendor a copy of its `modeling_*.py`**. Instead, subclass the smallest upstream unit and override only what changes. [`pi_gemma.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi_gemma.py) is the canonical reference: it injects AdaRMS conditioning into PaliGemma/Gemma in ~370 lines by subclassing `GemmaModel`/`PaliGemmaModel` and overriding the decoder-layer forward, instead of forking the ~2,000-line modeling file. Model surgery on a _loaded_ native model is also fine (layer truncation, tokenizer expansion, hidden-state capture — see `evo1/internvl3_embedder.py`, `eo1/modeling_eo1.py`, `groot/groot_n1_7.py` for working examples). Reviewers will ask for this pattern when a PR arrives with a copied modeling file; the only accepted exception is a model that does not exist in `transformers` at all.
### Benchmarks and a published checkpoint ### Benchmarks and a published checkpoint
A new policy is much easier to review — and far more useful — when it ships with a working checkpoint and at least one number you can reproduce. A new policy is much easier to review — and far more useful — when it ships with a working checkpoint and at least one number you can reproduce.
@@ -367,11 +373,12 @@ If your policy is real-robot-only and no sim benchmark applies, swap the sim eva
The general expectations are in [`CONTRIBUTING.md`](https://github.com/huggingface/lerobot/blob/main/CONTRIBUTING.md) and the [PR template](https://github.com/huggingface/lerobot/blob/main/.github/PULL_REQUEST_TEMPLATE.md). On top of those, reviewers will look for: The general expectations are in [`CONTRIBUTING.md`](https://github.com/huggingface/lerobot/blob/main/CONTRIBUTING.md) and the [PR template](https://github.com/huggingface/lerobot/blob/main/.github/PULL_REQUEST_TEMPLATE.md). On top of those, reviewers will look for:
- [ ] `MyPolicy` and `MyPolicyConfig` cover the surface above; `__init_subclass__` accepts the class. - [ ] `MyPolicy` and `MyPolicyConfig` cover the surface above; `__init_subclass__` accepts the class.
- [ ] `factory.py` and `policies/__init__.py` are wired (lazy imports for modeling). - [ ] `policies/__init__.py` re-exports the config (this registers the policy; the factory resolves modeling/processor by naming convention).
- [ ] `make_my_policy_pre_post_processors` follows the naming convention. - [ ] `make_my_policy_pre_post_processors` follows the naming convention.
- [ ] Optional deps live behind a `[project.optional-dependencies]` extra and the `TYPE_CHECKING + require_package` guard. - [ ] Optional deps live behind a `[project.optional-dependencies]` extra and the `TYPE_CHECKING + require_package` guard.
- [ ] `tests/policies/` updated; backward-compat artifact committed & policy-specific tests. - [ ] `tests/policies/` updated; backward-compat artifact committed & policy-specific tests.
- [ ] `src/lerobot/policies/<name>/README.md` symlinked into `docs/source/policy_<name>_README.md`; user-facing `docs/source/<name>.mdx` written and added to `_toctree.yml`. - [ ] `src/lerobot/policies/<name>/README.md` symlinked into `docs/source/policy_<name>_README.md`; user-facing `docs/source/<name>.mdx` written and added to `_toctree.yml`.
- [ ] `lerobot-train --policy.type my_policy ...` runs end-to-end for at least a few steps + save a checkpoint that can be loaded and run by `lerobot-eval` or `lerobot-rollout`.
- [ ] `templates/lerobot_modelcard_template.md` has a description entry and a `policy_docs` link for your policy. - [ ] `templates/lerobot_modelcard_template.md` has a description entry and a `policy_docs` link for your policy.
- [ ] The models table in the root `README.md` lists your policy in the right category, linking to your doc page. - [ ] The models table in the root `README.md` lists your policy in the right category, linking to your doc page.
- [ ] At least one reproducible benchmark eval in the policy MDX with a published checkpoint (sim benchmark, or real-robot dataset + checkpoint). - [ ] At least one reproducible benchmark eval in the policy MDX with a published checkpoint (sim benchmark, or real-robot dataset + checkpoint).
+13
View File
@@ -136,6 +136,10 @@ config = RealSenseCameraConfig(
height=480, height=480,
color_mode=ColorMode.RGB, color_mode=ColorMode.RGB,
use_depth=True, 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 rotation=Cv2Rotation.NO_ROTATION
) )
@@ -154,6 +158,15 @@ finally:
``` ```
<!-- prettier-ignore-end --> <!-- 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> </hfoption>
</hfoptions> </hfoptions>
+2 -15
View File
@@ -88,20 +88,6 @@ policy_preprocessor = NormalizerProcessorStep(stats=dataset_stats)
The same policy can work with different environment processors, and the same environment processor can work with different policies: The same policy can work with different environment processors, and the same environment processor can work with different policies:
````python
# Use SmolVLA policy with LIBERO environment
# Use SmolVLA policy with LIBERO environment
libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
env_cfg=libero_cfg,
policy_cfg=smolvla_cfg,
)
smolvla_preprocessor, smolvla_postprocessor = make_pre_post_processors(smolvla_cfg)
# Or use ACT policy with the same LIBERO environment
libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
env_cfg=libero_cfg,
policy_cfg=act_cfg,
)
act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
```python ```python
# Use SmolVLA policy with LIBERO environment # Use SmolVLA policy with LIBERO environment
libero_preprocessor, libero_postprocessor = make_env_pre_post_processors( libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
@@ -116,6 +102,7 @@ libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
policy_cfg=act_cfg, policy_cfg=act_cfg,
) )
act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg) act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
```
### 3. **Easier Experimentation** ### 3. **Easier Experimentation**
@@ -145,7 +132,7 @@ class LiberoVelocityProcessorStep(ObservationProcessorStep):
state = torch.cat([eef_pos, eef_axisangle, eef_vel, state = torch.cat([eef_pos, eef_axisangle, eef_vel,
gripper_pos, gripper_vel], dim=-1) # 14D gripper_pos, gripper_vel], dim=-1) # 14D
return state return state
```` ```
### 4. **Cleaner Environment Code** ### 4. **Cleaner Environment Code**
+4 -4
View File
@@ -40,10 +40,10 @@ This tutorial guides you through updating the firmware of Feetech motors using t
For each motor you want to update: For each motor you want to update:
1. **Select the motor** from the list by clicking on it 1. **Select the motor** from the list by clicking on it
2. **Click on Upgrade tab**: 2. **Click the Upgrade tab**:
3. **Click on Online button**: 3. **Click the Online button**:
- If an potential firmware update is found, it will be displayed in the box - If a potential firmware update is found, it will be displayed in the box
4. **Click on Upgrade button**: 4. **Click the Upgrade button**:
- The update progress will be displayed - The update progress will be displayed
## Step 6: Verify Update ## Step 6: Verify Update
+1
View File
@@ -59,6 +59,7 @@ The `lerobot-rollout --strategy.type=dagger` mode requires **teleoperators with
- `bi_openarm_mini` - Bimanual OpenArm Mini - `bi_openarm_mini` - Bimanual OpenArm Mini
- `so_leader` - SO100 / SO101 leader arm - `so_leader` - SO100 / SO101 leader arm
- `bi_so_leader` - Bimanual SO100 / SO101 leader arms
> [!IMPORTANT] > [!IMPORTANT]
> The provided commands default to `bi_openarm_follower` + `bi_openarm_mini`. > The provided commands default to `bi_openarm_follower` + `bi_openarm_mini`.
+1 -1
View File
@@ -211,7 +211,7 @@ Record, Replay and Train with Hope-JR is still experimental.
### Record ### Record
This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data/settings). This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data).
```bash ```bash
lerobot-record \ lerobot-record \
+1 -1
View File
@@ -18,7 +18,7 @@ If you're using Feetech or Dynamixel motors, LeRobot provides built-in bus inter
- [`DynamixelMotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/dynamixel/dynamixel.py) for controlling Dynamixel servos - [`DynamixelMotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/dynamixel/dynamixel.py) for controlling Dynamixel servos
Please refer to the [`MotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/motors_bus.py) abstract class to learn about its API. Please refer to the [`MotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/motors_bus.py) abstract class to learn about its API.
For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/src/lerobot/robots/so_follower/so101_follower/so101_follower.py) For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/src/lerobot/robots/so_follower/so_follower.py)
Use these if compatible. Otherwise, you'll need to find or write a Python interface (not covered in this tutorial): Use these if compatible. Otherwise, you'll need to find or write a Python interface (not covered in this tutorial):
+1 -1
View File
@@ -51,7 +51,7 @@ In addition to these instructions, you need to install the Feetech SDK & ZeroMQ
pip install -e ".[lekiwi]" pip install -e ".[lekiwi]"
``` ```
Great :hugs:! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base :robot:. Great 🤗! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base 🤖.
Every time you now want to use LeRobot, you can go to the `~/lerobot` folder where we installed LeRobot and run one of the commands. Every time you now want to use LeRobot, you can go to the `~/lerobot` folder where we installed LeRobot and run one of the commands.
# Step-by-Step Assembly Instructions # Step-by-Step Assembly Instructions
+8
View File
@@ -1,3 +1,11 @@
# OMX
<img
src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/omx_mainimage.png"
alt="OMX"
width=600
/>
## Order and Assemble the parts ## Order and Assemble the parts
First, assemble the OMX hardware following the official assembly guide. First, assemble the OMX hardware following the official assembly guide.
+1 -1
View File
@@ -174,7 +174,7 @@ The model takes images, text instructions, and robot state as input, and outputs
## Reproducing π₀Fast results ## Reproducing π₀Fast results
We reproduce the results of π₀Fast on the LIBERO benchmark using the LeRobot implementation. We take the LeRobot PiFast base model [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base) and finetune for an additional 40kk steps in bfloat16, with batch size of 256 on 8 H100 GPUs using the [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero). We reproduce the results of π₀Fast on the LIBERO benchmark using the LeRobot implementation. We take the LeRobot PiFast base model [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base) and finetune for an additional 40k steps in bfloat16, with batch size of 256 on 8 H100 GPUs using the [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero).
The finetuned model can be found here: The finetuned model can be found here:
+4 -4
View File
@@ -22,7 +22,7 @@ With processors, you choose the learning features you want to use for your polic
## Three pipelines ## Three pipelines
We often compose three pipelines. Depending on your setup, some can be empty if action and observation spaces already match. We often compose three pipelines. Depending on your setup, some can be empty if action and observation spaces already match.
Each of these pipelines handle different conversions between different action and observation spaces. Below is a quick explanation of each pipeline. Each of these pipelines handles different conversions between different action and observation spaces. Below is a quick explanation of each pipeline.
1. Pipeline 1: Teleop action space → dataset action space (phone pose → EE targets) 1. Pipeline 1: Teleop action space → dataset action space (phone pose → EE targets)
2. Pipeline 2: Dataset action space → robot command space (EE targets → joints) 2. Pipeline 2: Dataset action space → robot command space (EE targets → joints)
@@ -74,15 +74,15 @@ In the phone to SO-100 follower examples we use the following adapters:
- `robot_action_to_transition`: transforms the teleop action dict to a pipeline transition. - `robot_action_to_transition`: transforms the teleop action dict to a pipeline transition.
- `transition_to_robot_action`: transforms the pipeline transition to a robot action dict. - `transition_to_robot_action`: transforms the pipeline transition to a robot action dict.
- `observation_to_transition`: transforms the robot observation dict to a pipeline transition. - `observation_to_transition`: transforms the robot observation dict to a pipeline transition.
- `transition_to_observation`: transforms the pipeline transition to a observation dict. - `transition_to_observation`: transforms the pipeline transition to an observation dict.
Checkout [src/lerobot/processor/converters.py](https://github.com/huggingface/lerobot/blob/main/src/lerobot/processor/converters.py) for more details. Check out [src/lerobot/processor/converters.py](https://github.com/huggingface/lerobot/blob/main/src/lerobot/processor/converters.py) for more details.
## Dataset feature contracts ## Dataset feature contracts
Dataset features are determined by the keys saved in the dataset. Each step can declare what features it modifies in a contract called `transform_features(...)`. Once you build a processor, the processor can then aggregate all of these features with `aggregate_pipeline_dataset_features()` and merge multiple feature dicts with `combine_feature_dicts(...)`. Dataset features are determined by the keys saved in the dataset. Each step can declare what features it modifies in a contract called `transform_features(...)`. Once you build a processor, the processor can then aggregate all of these features with `aggregate_pipeline_dataset_features()` and merge multiple feature dicts with `combine_feature_dicts(...)`.
Below is and example of how we declare features with the `transform_features` method in the phone to SO-100 follower examples: Below is an example of how we declare features with the `transform_features` method in the phone to SO-100 follower examples:
```python ```python
def transform_features( def transform_features(
+2
View File
@@ -82,6 +82,8 @@ By default the env samples objects only from the `lightwheel` registry (what `--
All eval snippets below mirror the CI command (see `.github/workflows/benchmark_tests.yml`). The `--rename_map` argument maps RoboCasa's native camera keys (`robot0_agentview_left` / `robot0_eye_in_hand` / `robot0_agentview_right`) onto the three-camera (`camera1` / `camera2` / `camera3`) input layout the released `smolvla_robocasa` policy was trained on. All eval snippets below mirror the CI command (see `.github/workflows/benchmark_tests.yml`). The `--rename_map` argument maps RoboCasa's native camera keys (`robot0_agentview_left` / `robot0_eye_in_hand` / `robot0_agentview_right`) onto the three-camera (`camera1` / `camera2` / `camera3`) input layout the released `smolvla_robocasa` policy was trained on.
By default, each task uses the rollout horizon registered by RoboCasa. Set `--env.episode_length=<steps>` to apply the same explicit horizon to every selected task.
### Single-task evaluation (recommended for quick iteration) ### Single-task evaluation (recommended for quick iteration)
```bash ```bash
+2 -2
View File
@@ -57,7 +57,7 @@ policy_cfg.rtc_config = RTCConfig(
policy = PI0Policy.from_pretrained("lerobot/pi0_base", policy_cfg=policy_cfg, device="cuda") policy = PI0Policy.from_pretrained("lerobot/pi0_base", policy_cfg=policy_cfg, device="cuda")
# Now use predict_action_chunk with RTC parameters # Now use predict_action_chunk with RTC parameters
inference_delay = 4 # How many steps of inference latency, this values should be calculated based on the inference latency of the policy inference_delay = 4 # How many steps of inference latency, this value should be calculated based on the inference latency of the policy
# Initialize the action queue # Initialize the action queue
action_queue = ActionQueue(policy_cfg.rtc_config) action_queue = ActionQueue(policy_cfg.rtc_config)
@@ -100,7 +100,7 @@ Typical values: 8-12 steps
RTCConfig(execution_horizon=10) RTCConfig(execution_horizon=10)
``` ```
**`max_guidance_weight`**: How strongly to enforce consistency with the previous chunk. This is a hyperparameter that can be tuned to balance the smoothness of the transitions and the reactivity of the policy. For 10 steps flow matching (SmolVLA, Pi0, Pi0.5), a value of 10.0 is a optimal value. **`max_guidance_weight`**: How strongly to enforce consistency with the previous chunk. This is a hyperparameter that can be tuned to balance the smoothness of the transitions and the reactivity of the policy. For 10 steps flow matching (SmolVLA, Pi0, Pi0.5), a value of 10.0 is an optimal value.
**`prefix_attention_schedule`**: How to weight consistency across the overlap region. **`prefix_attention_schedule`**: How to weight consistency across the overlap region.
+1 -1
View File
@@ -93,7 +93,7 @@ lerobot-train --help
## Evaluate the finetuned model and run it in real-time ## Evaluate the finetuned model and run it in real-time
Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots). Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots#record-a-dataset).
Once you are logged in, you can run inference in your setup by doing: Once you are logged in, you can run inference in your setup by doing:
```bash ```bash
+1 -1
View File
@@ -338,7 +338,7 @@ It is advisable to install one 3-pin cable in the motor after placing them befor
<hfoption id="Leader"> <hfoption id="Leader">
- Mount the leader holder onto the wrist and secure it with 4 M3x6mm screws. - Mount the leader holder onto the wrist and secure it with 4 M3x6mm screws.
- Attach the handle to motor 5 using 1 M2x6mm screw. - Attach the handle to the leader holder using 1 M2x6mm screw.
- Insert the gripper motor, secure it with 2 M2x6mm screws on each side, attach a motor horn using a M3x6mm horn screw. - Insert the gripper motor, secure it with 2 M2x6mm screws on each side, attach a motor horn using a M3x6mm horn screw.
- Attach the follower trigger with 4 M3x6mm screws. - Attach the follower trigger with 4 M3x6mm screws.
+6 -2
View File
@@ -50,11 +50,11 @@ lerobot-edit-dataset \
Divide a dataset into multiple subsets. Divide a dataset into multiple subsets.
```bash ```bash
# Split by fractions (e.g. 80% train, 20% test, 20% val) # Split by fractions (e.g. 60% train, 20% val, 20% test)
lerobot-edit-dataset \ lerobot-edit-dataset \
--repo_id lerobot/pusht \ --repo_id lerobot/pusht \
--operation.type split \ --operation.type split \
--operation.splits '{"train": 0.8, "test": 0.2, "val": 0.2}' --operation.splits '{"train": 0.6, "val": 0.2, "test": 0.2}'
# Split by specific episode indices # Split by specific episode indices
lerobot-edit-dataset \ lerobot-edit-dataset \
@@ -252,6 +252,10 @@ lerobot-dataset-viz \
--episode-index 0 --episode-index 0
``` ```
For a private or gated dataset, authenticate first with `hf auth login`, or set the
`HF_TOKEN` environment variable. The Hub client then discovers the credential
automatically; no token argument is needed.
**From a local folder:** **From a local folder:**
Add the `--root` option and set `--mode local`. For example, to search in `./my_local_data_dir/lerobot/pusht`: Add the `--root` option and set `--mode local`. For example, to search in `./my_local_data_dir/lerobot/pusht`:
-77
View File
@@ -1,77 +0,0 @@
#!/usr/bin/env python
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Launch ``lerobot-annotate`` on a Hugging Face job (vllm + Qwen3.6-27B VLM).
Spawns one single-GPU ``h200`` job that:
1. installs ``lerobot`` from ``main`` plus the annotation extras,
2. boots one vllm server with Qwen3.6-27B (dense VLM),
3. runs the plan / interjections / vqa modules across the dataset
in free-form mode (each episode generates its own subtasks +
memory),
4. uploads the annotated dataset to ``--new_repo_id`` (when set)
or back to ``--repo_id``.
Usage:
HF_TOKEN=hf_... uv run python examples/annotations/run_hf_job.py
Adjust ``CMD`` (dataset, model, hub repo) and ``flavor`` below for your
run. For larger datasets, scale to ``h200x4`` and raise
``--vlm.parallel_servers`` / ``--vlm.num_gpus`` to match.
"""
import os
from huggingface_hub import get_token, run_job
token = os.environ.get("HF_TOKEN") or get_token()
if not token:
raise RuntimeError("No HF token. Run `huggingface-cli login` or `export HF_TOKEN=hf_...`")
CMD = (
"apt-get update -qq && apt-get install -y -qq git ffmpeg && "
"pip install --no-deps "
"'lerobot @ git+https://github.com/huggingface/lerobot.git@main' && "
"pip install --upgrade-strategy only-if-needed "
"datasets pyarrow av jsonlines draccus gymnasium torchcodec mergedeep pyyaml-include toml typing-inspect "
"openai && "
"export VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0 && "
"export VLLM_VIDEO_BACKEND=pyav && "
"lerobot-annotate "
"--repo_id=pepijn223/robocasa_pretrain_human300_v4 "
"--new_repo_id=pepijn223/robocasa_pretrain_human300_v4_annotated "
"--push_to_hub=true "
"--vlm.backend=openai "
"--vlm.model_id=Qwen/Qwen3.6-27B "
"--vlm.num_gpus=1 "
'--vlm.serve_command="vllm serve Qwen/Qwen3.6-27B '
"--tensor-parallel-size 1 --max-model-len 32768 "
'--gpu-memory-utilization 0.8 --uvicorn-log-level warning --port {port}" '
"--vlm.serve_ready_timeout_s=1800 "
# Qwen3.6 ships with thinking on; annotation wants plain JSON answers.
"--vlm.chat_template_kwargs='{\"enable_thinking\": false}'"
)
job = run_job(
image="vllm/vllm-openai:latest",
command=["bash", "-c", CMD],
flavor="h200",
secrets={"HF_TOKEN": token},
timeout="2h",
)
print(f"Job URL: {job.url}")
print(f"Job ID: {job.id}")
+1 -1
View File
@@ -44,6 +44,7 @@ from typing import Protocol
import numpy as np import numpy as np
from lerobot.lerobot_types import RobotAction, RobotObservation
from lerobot.model.kinematics import RobotKinematics from lerobot.model.kinematics import RobotKinematics
from lerobot.processor import ( from lerobot.processor import (
RobotProcessorPipeline, RobotProcessorPipeline,
@@ -56,7 +57,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
EEBoundsAndSafety, EEBoundsAndSafety,
InverseKinematicsEEToJoints, InverseKinematicsEEToJoints,
) )
from lerobot.types import RobotAction, RobotObservation
from lerobot.utils.constants import HF_LEROBOT_CALIBRATION, HF_LEROBOT_HOME, TELEOPERATORS from lerobot.utils.constants import HF_LEROBOT_CALIBRATION, HF_LEROBOT_HOME, TELEOPERATORS
from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.robot_utils import precise_sleep
@@ -38,7 +38,7 @@ from typing import TYPE_CHECKING
import numpy as np import numpy as np
from lerobot.types import RobotAction from lerobot.lerobot_types import RobotAction
from .base import _GRIPPER_MOTOR_SCALE, IsaacTeleopTeleoperator, _isaacteleop_available from .base import _GRIPPER_MOTOR_SCALE, IsaacTeleopTeleoperator, _isaacteleop_available
from .config_isaac_teleop import SO101LeaderArmConfig from .config_isaac_teleop import SO101LeaderArmConfig
@@ -32,7 +32,7 @@ from typing import TYPE_CHECKING, Any
import numpy as np import numpy as np
from lerobot.types import RobotAction from lerobot.lerobot_types import RobotAction
from .base import IsaacTeleopTeleoperator, _isaacteleop_available from .base import IsaacTeleopTeleoperator, _isaacteleop_available
from .config_isaac_teleop import XRControllerConfig from .config_isaac_teleop import XRControllerConfig
@@ -26,8 +26,8 @@ from __future__ import annotations
from dataclasses import dataclass from dataclasses import dataclass
from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
from lerobot.lerobot_types import RobotAction
from lerobot.processor import ProcessorStepRegistry, RobotActionProcessorStep from lerobot.processor import ProcessorStepRegistry, RobotActionProcessorStep
from lerobot.types import RobotAction
from lerobot.utils.rotation import Rotation from lerobot.utils.rotation import Rotation
from .base import _GRIPPER_MOTOR_SCALE from .base import _GRIPPER_MOTOR_SCALE
+1 -1
View File
@@ -21,6 +21,7 @@ from lerobot.cameras.opencv import OpenCVCameraConfig
from lerobot.common.control_utils import predict_action from lerobot.common.control_utils import predict_action
from lerobot.configs import FeatureType, PolicyFeature from lerobot.configs import FeatureType, PolicyFeature
from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features
from lerobot.lerobot_types import RobotAction, RobotObservation
from lerobot.model.kinematics import RobotKinematics from lerobot.model.kinematics import RobotKinematics
from lerobot.policies import make_pre_post_processors from lerobot.policies import make_pre_post_processors
from lerobot.policies.act import ACTPolicy from lerobot.policies.act import ACTPolicy
@@ -38,7 +39,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
ForwardKinematicsJointsToEE, ForwardKinematicsJointsToEE,
InverseKinematicsEEToJoints, InverseKinematicsEEToJoints,
) )
from lerobot.types import RobotAction, RobotObservation
from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.constants import ACTION, OBS_STR
from lerobot.utils.feature_utils import build_dataset_frame, combine_feature_dicts from lerobot.utils.feature_utils import build_dataset_frame, combine_feature_dicts
from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.keyboard_input import init_keyboard_listener
+1 -1
View File
@@ -16,6 +16,7 @@
from lerobot.cameras.opencv import OpenCVCameraConfig from lerobot.cameras.opencv import OpenCVCameraConfig
from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features
from lerobot.lerobot_types import RobotAction, RobotObservation
from lerobot.model.kinematics import RobotKinematics from lerobot.model.kinematics import RobotKinematics
from lerobot.processor import ( from lerobot.processor import (
RobotProcessorPipeline, RobotProcessorPipeline,
@@ -36,7 +37,6 @@ from lerobot.scripts.lerobot_record import record_loop
from lerobot.teleoperators.phone import Phone, PhoneConfig from lerobot.teleoperators.phone import Phone, PhoneConfig
from lerobot.teleoperators.phone.config_phone import PhoneOS from lerobot.teleoperators.phone.config_phone import PhoneOS
from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction
from lerobot.types import RobotAction, RobotObservation
from lerobot.utils.feature_utils import combine_feature_dicts from lerobot.utils.feature_utils import combine_feature_dicts
from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.keyboard_input import init_keyboard_listener
from lerobot.utils.utils import log_say from lerobot.utils.utils import log_say
+1 -1
View File
@@ -17,6 +17,7 @@
import time import time
from lerobot.datasets import LeRobotDataset from lerobot.datasets import LeRobotDataset
from lerobot.lerobot_types import RobotAction, RobotObservation
from lerobot.model.kinematics import RobotKinematics from lerobot.model.kinematics import RobotKinematics
from lerobot.processor import ( from lerobot.processor import (
RobotProcessorPipeline, RobotProcessorPipeline,
@@ -27,7 +28,6 @@ from lerobot.robots.so_follower import SO100Follower, SO100FollowerConfig
from lerobot.robots.so_follower.robot_kinematic_processor import ( from lerobot.robots.so_follower.robot_kinematic_processor import (
InverseKinematicsEEToJoints, InverseKinematicsEEToJoints,
) )
from lerobot.types import RobotAction, RobotObservation
from lerobot.utils.constants import ACTION from lerobot.utils.constants import ACTION
from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.robot_utils import precise_sleep
from lerobot.utils.utils import log_say from lerobot.utils.utils import log_say
+1 -1
View File
@@ -27,6 +27,7 @@ Highlight, or DAgger via ``lerobot-rollout --strategy.type=...``.
from lerobot.cameras.opencv import OpenCVCameraConfig from lerobot.cameras.opencv import OpenCVCameraConfig
from lerobot.configs import PreTrainedConfig from lerobot.configs import PreTrainedConfig
from lerobot.lerobot_types import RobotAction, RobotObservation
from lerobot.model.kinematics import RobotKinematics from lerobot.model.kinematics import RobotKinematics
from lerobot.processor import ( from lerobot.processor import (
RobotProcessorPipeline, RobotProcessorPipeline,
@@ -43,7 +44,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
from lerobot.rollout import BaseStrategyConfig, RolloutConfig, build_rollout_context from lerobot.rollout import BaseStrategyConfig, RolloutConfig, build_rollout_context
from lerobot.rollout.inference import SyncInferenceConfig from lerobot.rollout.inference import SyncInferenceConfig
from lerobot.rollout.strategies import BaseStrategy from lerobot.rollout.strategies import BaseStrategy
from lerobot.types import RobotAction, RobotObservation
from lerobot.utils.process import ProcessSignalHandler from lerobot.utils.process import ProcessSignalHandler
from lerobot.utils.utils import init_logging from lerobot.utils.utils import init_logging
+1 -1
View File
@@ -15,6 +15,7 @@
import time import time
from lerobot.lerobot_types import RobotAction, RobotObservation
from lerobot.model.kinematics import RobotKinematics from lerobot.model.kinematics import RobotKinematics
from lerobot.processor import ( from lerobot.processor import (
RobotProcessorPipeline, RobotProcessorPipeline,
@@ -31,7 +32,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
from lerobot.teleoperators.phone import Phone, PhoneConfig from lerobot.teleoperators.phone import Phone, PhoneConfig
from lerobot.teleoperators.phone.config_phone import PhoneOS from lerobot.teleoperators.phone.config_phone import PhoneOS
from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction
from lerobot.types import RobotAction, RobotObservation
from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.robot_utils import precise_sleep
from lerobot.utils.visualization_utils import init_rerun, log_rerun_data from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
+1 -1
View File
@@ -21,6 +21,7 @@ from lerobot.cameras.opencv import OpenCVCameraConfig
from lerobot.common.control_utils import predict_action from lerobot.common.control_utils import predict_action
from lerobot.configs import FeatureType, PolicyFeature from lerobot.configs import FeatureType, PolicyFeature
from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features
from lerobot.lerobot_types import RobotAction, RobotObservation
from lerobot.model.kinematics import RobotKinematics from lerobot.model.kinematics import RobotKinematics
from lerobot.policies import make_pre_post_processors from lerobot.policies import make_pre_post_processors
from lerobot.policies.act import ACTPolicy from lerobot.policies.act import ACTPolicy
@@ -38,7 +39,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
ForwardKinematicsJointsToEE, ForwardKinematicsJointsToEE,
InverseKinematicsEEToJoints, InverseKinematicsEEToJoints,
) )
from lerobot.types import RobotAction, RobotObservation
from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.constants import ACTION, OBS_STR
from lerobot.utils.feature_utils import build_dataset_frame, combine_feature_dicts from lerobot.utils.feature_utils import build_dataset_frame, combine_feature_dicts
from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.keyboard_input import init_keyboard_listener
+1 -1
View File
@@ -17,6 +17,7 @@
from lerobot.cameras.opencv import OpenCVCameraConfig from lerobot.cameras.opencv import OpenCVCameraConfig
from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features
from lerobot.lerobot_types import RobotAction, RobotObservation
from lerobot.model.kinematics import RobotKinematics from lerobot.model.kinematics import RobotKinematics
from lerobot.processor import ( from lerobot.processor import (
RobotProcessorPipeline, RobotProcessorPipeline,
@@ -33,7 +34,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
) )
from lerobot.scripts.lerobot_record import record_loop from lerobot.scripts.lerobot_record import record_loop
from lerobot.teleoperators.so_leader import SO100Leader, SO100LeaderConfig from lerobot.teleoperators.so_leader import SO100Leader, SO100LeaderConfig
from lerobot.types import RobotAction, RobotObservation
from lerobot.utils.feature_utils import combine_feature_dicts from lerobot.utils.feature_utils import combine_feature_dicts
from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.keyboard_input import init_keyboard_listener
from lerobot.utils.utils import log_say from lerobot.utils.utils import log_say
+1 -1
View File
@@ -18,6 +18,7 @@
import time import time
from lerobot.datasets import LeRobotDataset from lerobot.datasets import LeRobotDataset
from lerobot.lerobot_types import RobotAction, RobotObservation
from lerobot.model.kinematics import RobotKinematics from lerobot.model.kinematics import RobotKinematics
from lerobot.processor import ( from lerobot.processor import (
RobotProcessorPipeline, RobotProcessorPipeline,
@@ -28,7 +29,6 @@ from lerobot.robots.so_follower import SO100Follower, SO100FollowerConfig
from lerobot.robots.so_follower.robot_kinematic_processor import ( from lerobot.robots.so_follower.robot_kinematic_processor import (
InverseKinematicsEEToJoints, InverseKinematicsEEToJoints,
) )
from lerobot.types import RobotAction, RobotObservation
from lerobot.utils.constants import ACTION from lerobot.utils.constants import ACTION
from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.robot_utils import precise_sleep
from lerobot.utils.utils import log_say from lerobot.utils.utils import log_say
+1 -1
View File
@@ -25,6 +25,7 @@ forward/inverse kinematics.
from lerobot.cameras.opencv import OpenCVCameraConfig from lerobot.cameras.opencv import OpenCVCameraConfig
from lerobot.configs import PreTrainedConfig from lerobot.configs import PreTrainedConfig
from lerobot.lerobot_types import RobotAction, RobotObservation
from lerobot.model.kinematics import RobotKinematics from lerobot.model.kinematics import RobotKinematics
from lerobot.processor import ( from lerobot.processor import (
RobotProcessorPipeline, RobotProcessorPipeline,
@@ -41,7 +42,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
from lerobot.rollout import BaseStrategyConfig, RolloutConfig, build_rollout_context from lerobot.rollout import BaseStrategyConfig, RolloutConfig, build_rollout_context
from lerobot.rollout.inference import SyncInferenceConfig from lerobot.rollout.inference import SyncInferenceConfig
from lerobot.rollout.strategies import BaseStrategy from lerobot.rollout.strategies import BaseStrategy
from lerobot.types import RobotAction, RobotObservation
from lerobot.utils.process import ProcessSignalHandler from lerobot.utils.process import ProcessSignalHandler
from lerobot.utils.utils import init_logging from lerobot.utils.utils import init_logging
+1 -1
View File
@@ -16,6 +16,7 @@
import time import time
from lerobot.lerobot_types import RobotAction, RobotObservation
from lerobot.model.kinematics import RobotKinematics from lerobot.model.kinematics import RobotKinematics
from lerobot.processor import ( from lerobot.processor import (
RobotProcessorPipeline, RobotProcessorPipeline,
@@ -30,7 +31,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
InverseKinematicsEEToJoints, InverseKinematicsEEToJoints,
) )
from lerobot.teleoperators.so_leader import SO100Leader, SO100LeaderConfig from lerobot.teleoperators.so_leader import SO100Leader, SO100LeaderConfig
from lerobot.types import RobotAction, RobotObservation
from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.robot_utils import precise_sleep
from lerobot.utils.visualization_utils import init_rerun, log_rerun_data from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
+17 -6
View File
@@ -67,8 +67,8 @@ dependencies = [
"einops>=0.8.0,<0.9.0", "einops>=0.8.0,<0.9.0",
# Config & Hub # Config & Hub
"draccus==0.10.0", # TODO: Relax version constraint "draccus>=0.11.6,<0.12.0",
"huggingface-hub>=1.0.0,<2.0.0", "huggingface-hub>=1.6.0,<2.0.0",
"requests>=2.32.0,<3.0.0", "requests>=2.32.0,<3.0.0",
# Environments # Environments
@@ -95,7 +95,7 @@ dependencies = [
# ── Feature-scoped extras ────────────────────────────────── # ── Feature-scoped extras ──────────────────────────────────
dataset = [ dataset = [
"datasets>=4.7.0,<5.0.0", "datasets>=4.8.0,<5.0.0",
"pandas>=2.0.0,<3.0.0", # NOTE: Transitive dependency of datasets "pandas>=2.0.0,<3.0.0", # NOTE: Transitive dependency of datasets
"pyarrow>=21.0.0,<30.0.0", # NOTE: Transitive dependency of datasets "pyarrow>=21.0.0,<30.0.0", # NOTE: Transitive dependency of datasets
"lerobot[av-dep]", "lerobot[av-dep]",
@@ -155,7 +155,7 @@ accelerate-dep = ["accelerate>=1.14.0,<2.0.0"]
can-dep = ["python-can>=4.2.0,<5.0.0"] can-dep = ["python-can>=4.2.0,<5.0.0"]
peft-dep = ["peft>=0.18.0,<1.0.0"] peft-dep = ["peft>=0.18.0,<1.0.0"]
scipy-dep = ["scipy>=1.14.0,<2.0.0"] scipy-dep = ["scipy>=1.14.0,<2.0.0"]
diffusers-dep = ["diffusers>=0.27.2,<0.36.0"] diffusers-dep = ["diffusers>=0.38.0,<0.40.0"]
qwen-vl-utils-dep = ["qwen-vl-utils>=0.0.11,<0.1.0"] qwen-vl-utils-dep = ["qwen-vl-utils>=0.0.11,<0.1.0"]
matplotlib-dep = ["matplotlib>=3.10.3,<4.0.0", "contourpy>=1.3.0,<2.0.0"] # NOTE: Explicitly listing contourpy helps the resolver converge faster. matplotlib-dep = ["matplotlib>=3.10.3,<4.0.0", "contourpy>=1.3.0,<2.0.0"] # NOTE: Explicitly listing contourpy helps the resolver converge faster.
pyserial-dep = ["pyserial>=3.5,<4.0"] pyserial-dep = ["pyserial>=3.5,<4.0"]
@@ -413,8 +413,6 @@ ignore = [
"__init__.py" = ["F401", "F403", "E402"] "__init__.py" = ["F401", "F403", "E402"]
# E402: conditional-import guards (TYPE_CHECKING / is_package_available) must precede the imports they protect # E402: conditional-import guards (TYPE_CHECKING / is_package_available) must precede the imports they protect
"src/lerobot/scripts/convert_dataset_v21_to_v30.py" = ["E402"] "src/lerobot/scripts/convert_dataset_v21_to_v30.py" = ["E402"]
"src/lerobot/policies/wall_x/**" = ["N801", "N812", "SIM102", "SIM108", "SIM210", "SIM211", "B006", "B007", "SIM118"] # Supprese these as they are coming from original Qwen2_5_vl code TODO(pepijn): refactor original
[tool.ruff.lint.isort] [tool.ruff.lint.isort]
combine-as-imports = true combine-as-imports = true
known-first-party = ["lerobot"] known-first-party = ["lerobot"]
@@ -496,6 +494,19 @@ ignore_errors = true
module = "lerobot.envs.*" module = "lerobot.envs.*"
ignore_errors = false 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]] # [[tool.mypy.overrides]]
# module = "lerobot.utils.*" # module = "lerobot.utils.*"
@@ -20,6 +20,29 @@ from dataclasses import dataclass, field
from pathlib import Path from pathlib import Path
from typing import Any from typing import Any
from lerobot.configs.default import JobConfig
# The annotation pipeline boots its own vLLM server, so the pod starts from the
# official vLLM runtime rather than the prebuilt `lerobot-gpu` training image;
# `lerobot` is pip-installed on top (see `lerobot.jobs.annotate`).
DEFAULT_ANNOTATE_JOB_IMAGE = "vllm/vllm-openai:latest"
@dataclass
class AnnotationJobConfig(JobConfig):
"""`JobConfig` with the annotation runtime's defaults.
Adds `lerobot_ref` because the vLLM image ships no lerobot: the pod installs
it from git, and the ref decides which code actually annotates. Point it at a
branch/tag/SHA to try unmerged changes remotely.
"""
image: str = DEFAULT_ANNOTATE_JOB_IMAGE
# Annotation is a bounded pass over a dataset; a tighter cap than training's
# "2d" keeps a wedged vLLM server from burning a day of GPU time.
timeout: str | None = "2h"
lerobot_ref: str = "main"
@dataclass @dataclass
class PlanConfig: class PlanConfig:
@@ -65,6 +88,14 @@ class PlanConfig:
# invented from the task text (+1 VLM call/episode). # invented from the task text (+1 VLM call/episode).
subtask_describe_first: bool = True subtask_describe_first: bool = True
# Seeded relabeling: after segmentation, re-label each span with a focused
# pass that sees the previous / current / next segment contact sheets and
# minimally corrects the seed label (macrodata's best end-to-end labeling
# step). Costs +1 VLM call per subtask; off by default.
subtask_seeded_relabel: bool = False
# Frames sampled uniformly per segment sheet in the relabel pass.
subtask_relabel_frames: int = 5
# Emit ``style="plan"`` rows at each boundary; False = subtasks + memory only. # Emit ``style="plan"`` rows at each boundary; False = subtasks + memory only.
emit_plan: bool = True emit_plan: bool = True
@@ -160,6 +191,11 @@ class VlmConfig:
# Forwarded as extra_body.chat_template_kwargs (e.g. {"enable_thinking": false}). # Forwarded as extra_body.chat_template_kwargs (e.g. {"enable_thinking": false}).
chat_template_kwargs: dict[str, Any] | None = None chat_template_kwargs: dict[str, Any] | None = None
# OpenAI-style thinking budget hint ("low"/"medium"/"high"); forwarded to
# the server when set. Used to cap a thinking model's reasoning so it
# leaves tokens for the actual JSON answer on OpenAI-compatible endpoints.
reasoning_effort: str | None = None
@dataclass @dataclass
class ExecutorConfig: class ExecutorConfig:
@@ -194,6 +230,11 @@ class AnnotationPipelineConfig:
vlm: VlmConfig = field(default_factory=VlmConfig) vlm: VlmConfig = field(default_factory=VlmConfig)
executor: ExecutorConfig = field(default_factory=ExecutorConfig) executor: ExecutorConfig = field(default_factory=ExecutorConfig)
# Where the annotation runs: omitted / "local" annotates on this machine, any
# other value is an HF Jobs flavor (e.g. "h200") and submits the run there.
# List flavors + pricing with `hf jobs hardware`.
job: AnnotationJobConfig = field(default_factory=AnnotationJobConfig)
skip_validation: bool = False skip_validation: bool = False
only_episodes: tuple[int, ...] | None = None only_episodes: tuple[int, ...] | None = None
@@ -30,8 +30,8 @@ Phase 3 is why the ``plan`` module must be re-entered after the
timestamps. timestamps.
Distributed execution is provided by Hugging Face Jobs (see Distributed execution is provided by Hugging Face Jobs (see
``examples/annotations/run_hf_job.py``); the runner inside the job ``lerobot.jobs.annotate``, reached via ``--job.target=<flavor>``); the pod
invokes ``lerobot-annotate`` which uses this in-process executor. inside the job invokes ``lerobot-annotate`` which uses this in-process executor.
Episode-level concurrency is controlled by Episode-level concurrency is controlled by
``ExecutorConfig.episode_parallelism``. ``ExecutorConfig.episode_parallelism``.
""" """
@@ -413,6 +413,15 @@ def _draw_timestamp_badge(image: PIL.Image.Image, timestamp: float) -> PIL.Image
result = image.copy() result = image.copy()
draw = ImageDraw.Draw(result) draw = ImageDraw.Draw(result)
# Scale the timestamp to the tile so it stays legible after the model
# downsamples the full sheet into 768px tiles — a tiny bitmap font blurs
# at contact-sheet resolution and the VLM can no longer read the exact
# source time, which is what the boundary score depends on. ``size=`` is
# supported by Pillow's bitmap default since 10.1; fall back otherwise.
badge_px = max(14, round(image.height * 0.12))
try:
font = ImageFont.load_default(size=badge_px)
except TypeError:
font = ImageFont.load_default() font = ImageFont.load_default()
label = f"{timestamp:06.2f}s" label = f"{timestamp:06.2f}s"
left, top, right, bottom = draw.textbbox((0, 0), label, font=font) left, top, right, bottom = draw.textbbox((0, 0), label, font=font)
@@ -116,6 +116,8 @@ class PlanSubtasksMemoryModule:
rows.extend(self._task_aug_rows([effective_task, *variants], t0)) rows.extend(self._task_aug_rows([effective_task, *variants], t0))
subtask_spans = self._generate_subtasks(record, task=effective_task) subtask_spans = self._generate_subtasks(record, task=effective_task)
if self.config.subtask_seeded_relabel and subtask_spans:
subtask_spans = self._seeded_relabel(record, subtask_spans, effective_task)
# subtask rows # subtask rows
for span in subtask_spans: for span in subtask_spans:
@@ -509,6 +511,51 @@ class PlanSubtasksMemoryModule:
return cleaned return cleaned
def _seeded_relabel(
self, record: EpisodeRecord, spans: list[dict[str, Any]], task: str
) -> list[dict[str, Any]]:
"""Re-label each span using prev/current/next segment contact sheets.
Boundaries are kept fixed; only ``text`` is refined. The original
("seed") label is passed as a strong prior so the model verifies and
minimally corrects it rather than re-describing from scratch the
macrodata seeded-relabeling step. One VLM call per span.
"""
n = len(spans)
out: list[dict[str, Any]] = []
for i, span in enumerate(spans):
content: list[dict[str, Any]] = []
if i > 0:
content += self._segment_sheet(record, spans[i - 1])
content += self._segment_sheet(record, span)
if i < n - 1:
content += self._segment_sheet(record, spans[i + 1])
prompt = load_prompt("plan_subtask_relabel").format(
episode_task=task,
seed_label=span["text"],
segment_index=i + 1,
segment_count=n,
start=float(span["start"]),
end=float(span["end"]),
)
content.append({"type": "text", "text": prompt})
label = self._vlm_field([{"role": "user", "content": content}], "label")
text = label.strip() if isinstance(label, str) and label.strip() else span["text"]
out.append({**span, "text": text})
return out
def _segment_sheet(self, record: EpisodeRecord, span: dict[str, Any]) -> list[dict[str, Any]]:
"""Contact-sheet block(s) for one span: up to N frames sampled uniformly."""
s, e = float(span["start"]), float(span["end"])
n = max(1, int(self.config.subtask_relabel_frames))
if e <= s or n == 1:
timestamps = [s]
else:
step = (e - s) / (n - 1)
timestamps = [s + i * step for i in range(n)]
frames = self.frame_provider.frames_at(record, timestamps)
return self._contact_sheet_blocks(frames, timestamps[: len(frames)])
def _generate_subtasks_windowed( def _generate_subtasks_windowed(
self, record: EpisodeRecord, task: str, window_s: float self, record: EpisodeRecord, task: str, window_s: float
) -> list[dict[str, Any]]: ) -> list[dict[str, Any]]:
@@ -22,12 +22,23 @@ plain editors and roundtrip cleanly through ``ruff format``.
from __future__ import annotations from __future__ import annotations
import os
from pathlib import Path from pathlib import Path
_DIR = Path(__file__).parent _DIR = Path(__file__).parent
def load(name: str) -> str: def load(name: str) -> str:
"""Read prompt template ``name.txt`` from the ``prompts/`` directory.""" """Read prompt template ``name.txt`` from the ``prompts/`` directory.
A ``LEROBOT_PROMPT_OVERRIDE_<name>`` environment variable, when set to a
non-empty value, takes precedence over the packaged file. This lets prompt
search (e.g. GEPA) inject candidate templates into a remote job without
rebuilding the package; the override must keep the same ``{placeholder}``
fields the call site formats in.
"""
override = os.environ.get(f"LEROBOT_PROMPT_OVERRIDE_{name}")
if override and override.strip():
return override
path = _DIR / f"{name}.txt" path = _DIR / f"{name}.txt"
return path.read_text(encoding="utf-8") return path.read_text(encoding="utf-8")
@@ -0,0 +1,35 @@
Annotate one fixed segment from a longer robot demonstration.
Return only JSON:
{{"label": "<short descriptive subtask label>"}}
You are shown up to three timestamped contact sheets, in order:
- The FIRST sheet is the PREVIOUS segment (context only); it may be absent.
- The SECOND sheet is the CURRENT target segment.
- The THIRD sheet is the NEXT segment (context only); it may be absent.
Each tile has its timestamp (seconds, absolute video time) burned into its
top-left corner.
Episode instruction: "{episode_task}"
Target segment: {segment_index} of {segment_count}
Target time: {start:.2f}s to {end:.2f}s
Original predicted label for this exact segment: "{seed_label}"
Rules:
- Label ONLY the current target segment (the second sheet). Use the
previous/next sheets only to disambiguate what changed.
- Treat the original predicted label as a STRONG PRIOR, not ground truth:
verify it against the current segment and correct it minimally.
- If it already names the right action and main object, keep it; only fix
grammar or add a clearly visible essential detail.
- If it is vague but directionally correct, make it more specific.
- If it describes the previous/next segment, the wrong action, wrong
object, wrong destination, or a wrong state change, replace it.
- Do not describe the previous or next segment, and do not split, merge,
or move the fixed segment.
- Do not introduce an action that is not clearly visible in the current
target segment.
- Use one concise imperative phrase. Name the manipulated object and the
action / state change. Include source, destination, side, direction,
final placement, or opened/closed state when visible and central.
- Do not mention timestamps, frame numbers, uncertainty, or intent.
@@ -1,112 +1,68 @@
You are labeling a teleoperated robot demonstration. You are annotating a teleoperated robot demonstration shown as
timestamped contact sheets (each tile has its time in seconds burned
into the top-left corner). The operator's goal was: "{episode_task}"
The user originally asked: "{episode_task}" {observation_block}Reconstruct the sequence of COMPLETED manipulation events the robot
performs, in chronological order. Output one segment per event with a
[start, end] time in seconds and a short action label.
You are shown the entire demonstration as a single video. Watch the GROUNDING — read first, it overrides everything below:
whole clip, then segment it into a list of consecutive atomic subtasks - Label ONLY events you can SEE in the frames. The instruction is the
the robot performs. goal; the VIDEO is the ground truth for what actually happened.
- Do NOT invent, anticipate, or pad steps that are not shown.
{observation_block}GROUNDING — read this first, it overrides everything below: Granularity — segment by completed events, not by motion:
- Label ONLY what the robot actually does in the video. Every subtask - Start a NEW segment whenever the world state changes: an object is
you emit must correspond to motion you can SEE in specific frames. grasped, lifted, transported, placed, or released; a held object
- Do NOT invent, anticipate, or pad. If the robot only does one thing changes; a drawer/door/lid/container opens or closes; contents move
(e.g. it just navigates to a location and the clip ends), emit between containers (poured); a tool starts or stops acting on a
EXACTLY ONE subtask. Many demonstrations are a single atomic skill. surface. Watch the gripper open/close transitions — they usually mark
- ``max_steps`` below is a hard CEILING, not a target. Emitting fewer boundaries.
subtasks than the ceiling is not just allowed, it is expected for - Do NOT split approach, reach, grasp adjustment, small repositioning,
short / atomic demonstrations. One correct subtask is far better hesitation, or retreat into their own segments. Fold each into the
than several invented ones. event it belongs to (the approach is part of the pick; the retreat is
- If the video does not clearly show the action implied by the task, part of the place).
describe what you actually see — do NOT fabricate the task's steps - Do NOT merge separate completed events. Each distinct pick, place,
from the instruction text. The instruction tells you the goal; the open, close, pour, push, wipe, or insert is its own segment, even when
VIDEO is the ground truth for what happened. they repeat on different objects or locations.
- Most segments last 2-10 seconds. Shorter segments are okay ONLY for
fast pick / place / open / close / release events. Never emit a
segment shorter than {min_subtask_seconds} seconds; merge a too-short
candidate into its neighbour instead.
- Skip idle time, pure camera motion, and tiny hand jitter.
Authoring rules — Hi Robot atom granularity, pi0.7-style short prompts: Labels — short imperative phrases:
- One concise command naming the action and the manipulated object, e.g.
"pick up the red cup", "put the cup on the shelf", "open the top
drawer", "pour water into the glass", "insert the plug into the
socket".
- Include source, destination, side, direction, or the final
open/closed state when it is visible and central to the event.
- Prefer these verbs (extend only when none fits): pick up, put, place,
push, pull, turn, press, open, close, pour, insert, wipe, stack.
Disambiguate by what you SEE:
* STACK vs PUT: object placed ON TOP OF another object -> "stack".
* INSERT vs PUT: object pushed INTO a fitted slot/hole/socket -> "insert".
* PICK UP vs PUT (direction): gripper CLOSES and object moves WITH
the hand -> "pick up"; gripper OPENS and object stays -> "put".
* POUR vs PUT: source is tilted and contents flow -> "pour".
- Use the exact object nouns implied by the task; stay consistent across
the episode (don't switch "cube" to "block").
- Write imperative commands, never third person ("the robot ..."), and
drop articles/adverbs.
- Each subtask = one COMPOSITE atomic skill the low-level policy can Timing:
execute end-to-end. A "skill" bundles its own approach motion with - Use the burned-in timestamps to set start and end. Boundaries should
its terminal action — do NOT split the approach off as its own land on or near a printed time, and every [start, end] must lie within
subtask. The whole-arm policy already learns to reach as part of [0.0, {episode_duration}] seconds, be non-overlapping, and cover the
every manipulation primitive. episode in order.
- Write each subtask as an IMPERATIVE COMMAND, starting with one of - Emit at most {max_steps} segments.
these verbs (extend only when none fits):
pick up <obj> — approach + grasp + lift in one subtask
put <obj> on/in <loc> — transport + release in one subtask
place <obj> on/in <loc> — synonym of "put"; pick one and stay consistent
push <obj> — contact + linear shove
pull <obj> — contact + linear retract
turn <knob/dial/handle> — rotary actuation
press <button> — single-press contact
open <drawer/door/lid> — full open motion
close <drawer/door/lid> — full close motion
pour <src> into <dst> — tilt + flow
insert <obj> into <slot>— alignment + push-fit
go to <loc> — ONLY when no grasp / actuation follows
(e.g. a pure relocation between phases).
If the next subtask grasps something at
that location, drop "go to ..." and just
write "pick up ..." instead.
- Forbidden ultra-fine splits — the VLM is NOT allowed to emit these
as standalone subtasks; fold them into the parent composite:
"move to X" → fold into "pick up X" (or whatever follows)
"reach for X" → fold into "pick up X"
"grasp X" → fold into "pick up X"
"lift X" → fold into "pick up X" (or "put X on Y" if it's
the transport phase of a place)
"release X" → fold into "put X on Y" (or "place X in Y")
- Keep it SHORT — a verb phrase, not a sentence. Drop articles
("the", "a") and adverbs ("carefully", "slowly"). Add a "how"
detail (which hand, which grasp point) ONLY when it is needed to
disambiguate. Every subtask must begin with one of the verbs
above (no leading nouns, no "then", no "first").
- NEVER use third person. Never write "the robot", "the arm", "the
gripper moves", "it picks up" — the robot is implied. Command it,
do not describe it.
- Use the exact object nouns from the task above. If the task says
"cube", every subtask says "cube" — never switch to "block". If it
says "box", never switch to "bin"/"container". Keep vocabulary
consistent across the whole episode.
- Good: "pick up blue cube", "put blue cube in box", "open drawer",
"turn red knob", "press start button", "go to sink".
- Bad: "move to blue cube" (approach as its own subtask — forbidden,
must be folded into "pick up blue cube"); "the robot arm moves
towards the blue cube" (third person, too long); "carefully pick
up the cube" (adverb, article); "release the yellow block"
("block" when the task said "cube", and "release" must be folded
into a "put"/"place" subtask).
- Subtasks are non-overlapping and cover the full episode in order.
Choose the cut points yourself based on what you see in the video
(gripper open/close events, contact, regrasps, transitions).
- Each subtask spans at least {min_subtask_seconds} seconds. If a
candidate span would be shorter, merge it into its neighbour
rather than emitting it.
- Do not exceed {max_steps} subtasks total. Fewer, larger composites
are preferred over many micro-steps.
- Every subtask's [start_time, end_time] must lie within
[0.0, {episode_duration}] seconds.
SPECIAL CASES — verb disambiguation (each rule is narrowly visual and
fires ONLY on the spatial situation it names; it must not change how you
label any other situation):
- STACK vs PUT: if an object is placed ON TOP OF another specific object
(not on a flat table / shelf / counter), use "stack ... on ...", not
"put". "stack blue book on green book", NOT "put blue book on table".
- INSERT vs PUT: if an object goes INTO a fitted slot / hole / socket /
receptacle (push-fit), use "insert ... into ...", not "put".
- RETRIEVE/PICK-UP vs PUT (direction): watch the gripper. If it CLOSES
on the object and the object moves WITH the hand, it is "pick up" /
"retrieve" (object leaves its location). If the gripper OPENS and the
object stays where the hand left it, it is "put" / "place" (object
arrives at a location). Decide by which way the object moves, not by
where the hand ends up.
- POUR vs PUT: only use "pour" when the source is tilted and contents
flow out; moving a full container without tilting is "put"/"place".
Output strictly valid JSON of shape: Output strictly valid JSON of shape:
{{ {{
"subtasks": [ "subtasks": [
{{"text": "<short imperative verb phrase>", "start": <float>, "end": <float>}}, {{"text": "<short imperative action label>", "start": <float>, "end": <float>}},
... ...
] ]
}} }}
@@ -194,12 +194,13 @@ def make_vlm_client(config: VlmConfig) -> VlmClient:
"""Build the shared VLM client. """Build the shared VLM client.
Only the ``openai`` backend is supported for now. The shipped workflow Only the ``openai`` backend is supported for now. The shipped workflow
is Hugging Face Jobs (``examples/annotations/run_hf_job.py``): it boots is Hugging Face Jobs (``lerobot-annotate --job.target=<flavor>``): it
a vLLM server inside the ``vllm/vllm-openai`` image and the pipeline boots a vLLM server inside the ``vllm/vllm-openai`` image and the
talks to it over the OpenAI-compatible API (``--vlm.backend=openai``, pipeline talks to it over the OpenAI-compatible API
optionally auto-spawning the server via ``auto_serve`` / (``--vlm.backend=openai``, optionally auto-spawning the server via
``serve_command``). The former in-process ``vllm`` / ``transformers`` ``auto_serve`` / ``serve_command``). The former in-process ``vllm`` /
backends were removed to keep the support surface to the HF Jobs path. ``transformers`` backends were removed to keep the support surface to
the HF Jobs path.
For ``stub``, construct :class:`StubVlmClient` directly with a responder For ``stub``, construct :class:`StubVlmClient` directly with a responder
callable; it is rejected here to make accidental misuse obvious. callable; it is rejected here to make accidental misuse obvious.
@@ -213,8 +214,8 @@ def make_vlm_client(config: VlmConfig) -> VlmClient:
if config.backend in {"vllm", "transformers"}: if config.backend in {"vllm", "transformers"}:
raise ValueError( raise ValueError(
f"backend={config.backend!r} (in-process local model) is not supported for now — " f"backend={config.backend!r} (in-process local model) is not supported for now — "
"only backend='openai' (the Hugging Face Jobs flow) is. Run the pipeline via " "only backend='openai' (the Hugging Face Jobs flow) is. Run the pipeline with "
"examples/annotations/run_hf_job.py, which serves the model with vLLM in the " "`lerobot-annotate --job.target=<flavor>`, which serves the model with vLLM in the "
"vllm/vllm-openai image and talks to it over the OpenAI-compatible API." "vllm/vllm-openai image and talks to it over the OpenAI-compatible API."
) )
raise ValueError(f"Unknown VLM backend: {config.backend!r}") raise ValueError(f"Unknown VLM backend: {config.backend!r}")
@@ -285,6 +286,8 @@ def _make_openai_client(config: VlmConfig) -> VlmClient:
"max_tokens": max_tok, "max_tokens": max_tok,
"temperature": temp, "temperature": temp,
} }
if config.reasoning_effort:
kwargs["reasoning_effort"] = config.reasoning_effort
extra_body: dict[str, Any] = {} extra_body: dict[str, Any] = {}
if send_mm_kwargs and mm_kwargs: if send_mm_kwargs and mm_kwargs:
extra_body["mm_processor_kwargs"] = {**mm_kwargs, "do_sample_frames": True} extra_body["mm_processor_kwargs"] = {**mm_kwargs, "do_sample_frames": True}
@@ -296,7 +299,13 @@ def _make_openai_client(config: VlmConfig) -> VlmClient:
chosen = clients[rr_counter["i"] % len(clients)] chosen = clients[rr_counter["i"] % len(clients)]
rr_counter["i"] += 1 rr_counter["i"] += 1
response = chosen.chat.completions.create(**kwargs) response = chosen.chat.completions.create(**kwargs)
return response.choices[0].message.content or "" # Some OpenAI-compatible servers can return a choice with no message
# (safety filter, or a "thinking" model that spends the whole budget
# before emitting content). Treat that as an empty reply so the
# JSON-retry path handles it instead of crashing the run.
choice = response.choices[0] if response.choices else None
message = choice.message if choice is not None else None
return (message.content if message is not None else None) or ""
def _gen(batch: Sequence[Sequence[dict[str, Any]]], max_tok: int, temp: float) -> list[str]: def _gen(batch: Sequence[Sequence[dict[str, Any]]], max_tok: int, temp: float) -> list[str]:
if len(batch) <= 1 or config.client_concurrency <= 1: if len(batch) <= 1 or config.client_concurrency <= 1:
+1 -1
View File
@@ -38,6 +38,7 @@ import draccus
import grpc import grpc
import torch import torch
from lerobot.lerobot_types import PolicyAction
from lerobot.policies import get_policy_class, make_pre_post_processors from lerobot.policies import get_policy_class, make_pre_post_processors
from lerobot.processor import PolicyProcessorPipeline from lerobot.processor import PolicyProcessorPipeline
from lerobot.transport import ( from lerobot.transport import (
@@ -45,7 +46,6 @@ from lerobot.transport import (
services_pb2_grpc, # type: ignore services_pb2_grpc, # type: ignore
) )
from lerobot.transport.utils import receive_bytes_in_chunks from lerobot.transport.utils import receive_bytes_in_chunks
from lerobot.types import PolicyAction
from .configs import PolicyServerConfig from .configs import PolicyServerConfig
from .constants import SUPPORTED_POLICIES from .constants import SUPPORTED_POLICIES
+46 -16
View File
@@ -120,14 +120,22 @@ class OpenCVCamera(Camera):
self.rotation: int | None = get_cv2_rotation(config.rotation) self.rotation: int | None = get_cv2_rotation(config.rotation)
self.backend: int = config.backend self.backend: int = config.backend
if self.height and self.width: self.capture_width: int | None = None
self.capture_width, self.capture_height = self.width, self.height self.capture_height: int | None = None
if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]: self._reset_connection_settings()
self.capture_width, self.capture_height = self.height, self.width
def __str__(self) -> str: def __str__(self) -> str:
return f"{self.__class__.__name__}({self.index_or_path})" 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 @property
def is_connected(self) -> bool: def is_connected(self) -> bool:
"""Checks if the camera is currently connected and opened.""" """Checks if the camera is currently connected and opened."""
@@ -164,6 +172,7 @@ class OpenCVCamera(Camera):
f"Failed to open {self}.Run `lerobot-find-cameras opencv` to find available cameras." f"Failed to open {self}.Run `lerobot-find-cameras opencv` to find available cameras."
) )
try:
self._configure_capture_settings() self._configure_capture_settings()
self._start_read_thread() self._start_read_thread()
@@ -175,6 +184,13 @@ class OpenCVCamera(Camera):
with self.frame_lock: with self.frame_lock:
if self.latest_frame is None: if self.latest_frame is None:
raise ConnectionError(f"{self} failed to capture frames during warmup.") 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.") logger.info(f"{self} connected.")
@@ -312,6 +328,7 @@ class OpenCVCamera(Camera):
for target in targets_to_scan: for target in targets_to_scan:
camera = cv2.VideoCapture(target) camera = cv2.VideoCapture(target)
try:
if camera.isOpened(): if camera.isOpened():
default_width = int(camera.get(cv2.CAP_PROP_FRAME_WIDTH)) default_width = int(camera.get(cv2.CAP_PROP_FRAME_WIDTH))
default_height = int(camera.get(cv2.CAP_PROP_FRAME_HEIGHT)) default_height = int(camera.get(cv2.CAP_PROP_FRAME_HEIGHT))
@@ -321,7 +338,9 @@ class OpenCVCamera(Camera):
# Get FOURCC code and convert to string # Get FOURCC code and convert to string
default_fourcc_code = camera.get(cv2.CAP_PROP_FOURCC) default_fourcc_code = camera.get(cv2.CAP_PROP_FOURCC)
default_fourcc_code_int = int(default_fourcc_code) default_fourcc_code_int = int(default_fourcc_code)
default_fourcc = "".join([chr((default_fourcc_code_int >> 8 * i) & 0xFF) for i in range(4)]) default_fourcc = "".join(
[chr((default_fourcc_code_int >> 8 * i) & 0xFF) for i in range(4)]
)
camera_info = { camera_info = {
"name": f"OpenCV Camera @ {target}", "name": f"OpenCV Camera @ {target}",
@@ -338,6 +357,7 @@ class OpenCVCamera(Camera):
} }
found_cameras_info.append(camera_info) found_cameras_info.append(camera_info)
finally:
camera.release() camera.release()
return found_cameras_info return found_cameras_info
@@ -496,6 +516,26 @@ class OpenCVCamera(Camera):
self.latest_timestamp = None self.latest_timestamp = None
self.new_frame_event.clear() 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 @check_if_not_connected
def async_read(self, timeout_ms: float = 200) -> NDArray[Any]: 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: if not self.is_connected and self.thread is None:
raise DeviceNotConnectedError(f"{self} not connected.") raise DeviceNotConnectedError(f"{self} not connected.")
if self.thread is not None: self._cleanup_resources()
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()
logger.info(f"{self} disconnected.") logger.info(f"{self} disconnected.")
@@ -173,7 +173,8 @@ class Reachy2Camera(Camera):
raise ValueError( raise ValueError(
f"Invalid color mode '{self.color_mode}'. Expected {ColorMode.RGB} or {ColorMode.BGR}." f"Invalid color mode '{self.color_mode}'. Expected {ColorMode.RGB} or {ColorMode.BGR}."
) )
if self.color_mode == ColorMode.RGB: is_depth_frame = self.config.name == "depth" and self.config.image_type == "depth"
if not is_depth_frame and self.color_mode == ColorMode.RGB:
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
self.latest_frame = frame self.latest_frame = frame
+158 -20
View File
@@ -121,6 +121,9 @@ class RealSenseCamera(Camera):
self.config = config self.config = config
self.width: int | None = config.width
self.height: int | None = config.height
if config.serial_number_or_name.isdigit(): if config.serial_number_or_name.isdigit():
self.serial_number = config.serial_number_or_name self.serial_number = config.serial_number_or_name
else: else:
@@ -131,6 +134,9 @@ class RealSenseCamera(Camera):
self.use_rgb = config.use_rgb self.use_rgb = config.use_rgb
self.use_depth = config.use_depth self.use_depth = config.use_depth
self.warmup_s = config.warmup_s 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_pipeline: rs.pipeline | None = None
self.rs_profile: rs.pipeline_profile | 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) self.rotation: int | None = get_cv2_rotation(config.rotation)
if self.height and self.width: self.capture_width: int | None = None
self.capture_width, self.capture_height = self.width, self.height self.capture_height: int | None = None
if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]: self._reset_connection_settings()
self.capture_width, self.capture_height = self.height, self.width
def __str__(self) -> str: def __str__(self) -> str:
return f"{self.__class__.__name__}({self.serial_number})" 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 @property
def is_connected(self) -> bool: def is_connected(self) -> bool:
"""Checks if the camera pipeline is started and streams are active.""" """Checks if the camera pipeline is started and streams are active."""
@@ -172,7 +187,8 @@ class RealSenseCamera(Camera):
Raises: Raises:
DeviceAlreadyConnectedError: If the camera is already connected. 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. 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. RuntimeError: If the pipeline starts but fails to apply requested settings.
""" """
@@ -190,7 +206,9 @@ class RealSenseCamera(Camera):
f"Failed to open {self}.Run `lerobot-find-cameras realsense` to find available cameras." f"Failed to open {self}.Run `lerobot-find-cameras realsense` to find available cameras."
) from e ) from e
try:
self._configure_capture_settings() self._configure_capture_settings()
self._configure_sensor_options()
self._start_read_thread() 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. # 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.
@@ -206,6 +224,13 @@ class RealSenseCamera(Camera):
self.use_depth and self.latest_depth_frame is None self.use_depth and self.latest_depth_frame is None
): ):
raise ConnectionError(f"{self} failed to capture frames during warmup.") 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.") logger.info(f"{self} connected.")
@@ -339,6 +364,111 @@ class RealSenseCamera(Camera):
self.new_frame_event.clear() self.new_frame_event.clear()
return self._async_read(timeout_ms=10000, read_depth=read_depth) 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 @check_if_not_connected
def read_depth(self, timeout_ms: int = 200) -> NDArray[Any]: def read_depth(self, timeout_ms: int = 200) -> NDArray[Any]:
""" """
@@ -453,7 +583,7 @@ class RealSenseCamera(Camera):
) )
processed_image = image processed_image = image
if self.color_mode == ColorMode.BGR: if not depth_frame and self.color_mode == ColorMode.BGR:
processed_image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) processed_image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE, cv2.ROTATE_180]: if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE, cv2.ROTATE_180]:
@@ -541,6 +671,27 @@ class RealSenseCamera(Camera):
self.latest_timestamp = None self.latest_timestamp = None
self.new_frame_event.clear() 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]: 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.""" """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(): 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." f"Attempted to disconnect {self}, but it appears already disconnected."
) )
if self.thread is not None: self._cleanup_resources()
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()
logger.info(f"{self} disconnected.") logger.info(f"{self} disconnected.")
@@ -46,6 +46,17 @@ class RealSenseCameraConfig(CameraConfig):
use_depth: Whether to enable depth stream. Defaults to False. use_depth: Whether to enable depth stream. Defaults to False.
rotation: Image rotation setting (0°, 90°, 180°, or 270°). Defaults to no rotation. rotation: Image rotation setting (0°, 90°, 180°, or 270°). Defaults to no rotation.
warmup_s: Time reading frames before returning from connect (in seconds) 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: Note:
- Either name or serial_number must be specified. - Either name or serial_number must be specified.
@@ -61,6 +72,9 @@ class RealSenseCameraConfig(CameraConfig):
use_depth: bool = False use_depth: bool = False
rotation: Cv2Rotation = Cv2Rotation.NO_ROTATION rotation: Cv2Rotation = Cv2Rotation.NO_ROTATION
warmup_s: int = 1 warmup_s: int = 1
exposure: int | None = None
gain: int | None = None
white_balance: int | None = None
def __post_init__(self) -> None: def __post_init__(self) -> None:
self.color_mode = ColorMode(self.color_mode) self.color_mode = ColorMode(self.color_mode)
@@ -69,6 +83,18 @@ class RealSenseCameraConfig(CameraConfig):
if not self.use_rgb and not self.use_depth: if not self.use_rgb and not self.use_depth:
raise ValueError("At least one of `use_rgb` or `use_depth` must be enabled.") 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) 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): if any(v is not None for v in values) and any(v is None for v in values):
raise ValueError( raise ValueError(
+1 -1
View File
@@ -35,9 +35,9 @@ else:
if TYPE_CHECKING: if TYPE_CHECKING:
from lerobot.datasets import LeRobotDataset from lerobot.datasets import LeRobotDataset
from lerobot.lerobot_types import PolicyAction
from lerobot.processor import PolicyProcessorPipeline from lerobot.processor import PolicyProcessorPipeline
from lerobot.robots import Robot from lerobot.robots import Robot
from lerobot.types import PolicyAction
def predict_action( def predict_action(
+6
View File
@@ -71,13 +71,19 @@ class DatasetRecordConfig:
# Number of threads per encoder instance. None = auto (codec default). # 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.. # Lower values reduce CPU usage, maps to 'lp' (via svtav1-params) for libsvtav1 and 'threads' for h264/hevc..
encoder_threads: int | None = None 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: def stamp_repo_id(self) -> None:
"""Append a date-time tag to ``repo_id`` so each recording session gets a unique name. """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, Must be called explicitly at dataset *creation* time not on resume,
where the existing ``repo_id`` (already stamped) must be preserved. 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: if self.repo_id:
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
self.repo_id = f"{self.repo_id}_{timestamp}" self.repo_id = f"{self.repo_id}_{timestamp}"
+19 -11
View File
@@ -163,8 +163,10 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno
return None return None
def _save_pretrained(self, save_directory: Path) -> None: def _save_pretrained(self, save_directory: Path) -> None:
with open(save_directory / CONFIG_NAME, "w") as f, draccus.config_type("json"): # Encode against the base class so draccus includes the choice "type" key,
draccus.dump(self, f, indent=4) # which `from_pretrained` needs to resolve the concrete subclass.
with open(save_directory / CONFIG_NAME, "w") as f:
json.dump(draccus.encode(self, PreTrainedConfig), f, indent=4)
@classmethod @classmethod
def from_pretrained( def from_pretrained(
@@ -205,24 +207,30 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno
f"{CONFIG_NAME} not found on the HuggingFace Hub in {model_id}" f"{CONFIG_NAME} not found on the HuggingFace Hub in {model_id}"
) from e ) from e
# HACK: Parse the original config to get the config subclass, so that we can
# apply cli overrides.
# This is very ugly, ideally we'd like to be able to do that natively with draccus
# something like --policy.path (in addition to --policy.type)
with draccus.config_type("json"):
orig_config = draccus.parse(cls, config_file, args=[])
if config_file is None: if config_file is None:
raise FileNotFoundError(f"{CONFIG_NAME} not found in {model_id}") raise FileNotFoundError(f"{CONFIG_NAME} not found in {model_id}")
with open(config_file) as f: with open(config_file) as f:
config = json.load(f) config = json.load(f)
config.pop("type") # Resolve the concrete config subclass from the serialized "type" tag, then parse
# the config (with CLI overrides) directly for that class. The "type" key is
# stripped because draccus only consumes it when parsing the registry base class.
policy_type = config.pop("type", None)
if policy_type is None:
raise ValueError(f"Missing 'type' field in {CONFIG_NAME} of {model_id}")
try:
config_cls = cls.get_choice_class(policy_type)
except Exception as e:
raise ValueError(
f"Policy type '{policy_type}' (from {CONFIG_NAME} of {model_id}) is not registered. "
f"Available policy types: {cls.get_known_choices()}"
) from e
with tempfile.NamedTemporaryFile("w+", delete=False, suffix=".json") as f: with tempfile.NamedTemporaryFile("w+", delete=False, suffix=".json") as f:
json.dump(config, f) json.dump(config, f)
config_file = f.name config_file = f.name
cli_overrides = policy_kwargs.pop("cli_overrides", []) cli_overrides = policy_kwargs.pop("cli_overrides", [])
with draccus.config_type("json"): with draccus.config_type("json"):
return draccus.parse(orig_config.__class__, config_file, args=cli_overrides) return draccus.parse(config_cls, config_file, args=cli_overrides)
+4 -2
View File
@@ -103,8 +103,10 @@ class RewardModelConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
pass pass
def _save_pretrained(self, save_directory: Path) -> None: def _save_pretrained(self, save_directory: Path) -> None:
with open(save_directory / CONFIG_NAME, "w") as f, draccus.config_type("json"): # Encode against the base class so draccus includes the choice "type" key,
draccus.dump(self, f, indent=4) # which `from_pretrained` needs to resolve the concrete subclass.
with open(save_directory / CONFIG_NAME, "w") as f:
json.dump(draccus.encode(self, RewardModelConfig), f, indent=4)
@classmethod @classmethod
def from_pretrained( def from_pretrained(
+23 -1
View File
@@ -14,6 +14,7 @@
import builtins import builtins
import datetime as dt import datetime as dt
import json import json
import multiprocessing
import os import os
import tempfile import tempfile
from dataclasses import dataclass, field from dataclasses import dataclass, field
@@ -101,6 +102,12 @@ class TrainPipelineConfig(HubMixin):
batch_size: int = 8 batch_size: int = 8
prefetch_factor: int = 4 prefetch_factor: int = 4
persistent_workers: bool = True persistent_workers: bool = True
# DataLoader worker start method. "spawn" is safer than "fork" with
# non-fork-safe libs (PyAV / torchcodec / ffmpeg), but adds some
# worker-startup time per run since workers re-import modules instead
# of inheriting parent state. Override with `--dataloader_multiprocessing_context=fork`
# when appropriate, or set it to `null` to use Python's platform default.
dataloader_multiprocessing_context: str | None = "spawn"
steps: int = 100_000 steps: int = 100_000
# Run policy in the simulation environment every N steps to measure reward/success (0 = disabled). # Run policy in the simulation environment every N steps to measure reward/success (0 = disabled).
env_eval_freq: int = 20_000 env_eval_freq: int = 20_000
@@ -187,7 +194,11 @@ class TrainPipelineConfig(HubMixin):
) )
if Path(config_path).resolve().exists(): if Path(config_path).resolve().exists():
policy_dir = Path(config_path).parent # `config_path` may point at the checkpoint's train_config.json or at its
# pretrained_model/ directory (both documented above) — resolve either to
# the pretrained_model/ directory.
config_path_obj = Path(config_path)
policy_dir = config_path_obj.parent if config_path_obj.is_file() else config_path_obj
self.checkpoint_path = policy_dir.parent self.checkpoint_path = policy_dir.parent
elif self.job.is_remote: elif self.job.is_remote:
return return
@@ -212,6 +223,17 @@ class TrainPipelineConfig(HubMixin):
self.reward_model.pretrained_path = str(policy_dir) self.reward_model.pretrained_path = str(policy_dir)
def validate(self) -> None: def validate(self) -> None:
available_contexts = multiprocessing.get_all_start_methods()
if (
self.dataloader_multiprocessing_context is not None
and self.dataloader_multiprocessing_context not in available_contexts
):
raise ValueError(
"`dataloader_multiprocessing_context` must be None or one of "
f"{available_contexts} on this platform, got "
f"{self.dataloader_multiprocessing_context!r}."
)
self._resolve_pretrained_from_cli() self._resolve_pretrained_from_cli()
if self.policy is None and self.reward_model is None: if self.policy is None and self.reward_model is None:
+96 -40
View File
@@ -19,6 +19,7 @@ import copy
import logging import logging
import shutil import shutil
from pathlib import Path from pathlib import Path
from typing import Any, NotRequired, TypedDict
import datasets import datasets
import pandas as pd import pandas as pd
@@ -49,8 +50,32 @@ from .utils import (
) )
from .video_utils import concatenate_video_files, get_video_duration_in_s from .video_utils import concatenate_video_files, get_video_duration_in_s
logger = logging.getLogger(__name__)
def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMetadata]) -> dict[str, dict]: type FeatureDict = dict[str, dict[str, Any]]
type ChunkFile = tuple[int, int]
class IndexState(TypedDict):
chunk: int
file: int
src_to_dst: NotRequired[dict[ChunkFile, ChunkFile]]
class VideoIndex(TypedDict):
chunk: int
file: int
latest_duration: float
episode_duration: float
src_to_offset: NotRequired[dict[ChunkFile, float]]
src_to_dst: NotRequired[dict[ChunkFile, ChunkFile]]
dst_file_durations: NotRequired[dict[ChunkFile, float]]
type VideoIndexState = dict[str, VideoIndex]
def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMetadata]) -> FeatureDict:
"""Create a merged video feature info dictionary for aggregation. The video encoder info is merged field-by-field: each key is kept only when every source agrees; otherwise that key is set to ``null`` (or ``{}`` for ``video.extra_options``) and a warning is logged. """Create a merged video feature info dictionary for aggregation. The video encoder info is merged field-by-field: each key is kept only when every source agrees; otherwise that key is set to ``null`` (or ``{}`` for ``video.extra_options``) and a warning is logged.
Args: Args:
@@ -59,14 +84,14 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
Returns: Returns:
dict: A dictionary of merged video feature info. dict: A dictionary of merged video feature info.
""" """
merged_info = copy.deepcopy(all_metadata[0].features) merged_info: FeatureDict = copy.deepcopy(all_metadata[0].features)
video_keys = [k for k in merged_info if merged_info[k].get("dtype") == "video"] video_keys = [k for k in merged_info if merged_info[k].get("dtype") == "video"]
for vk in video_keys: for vk in video_keys:
video_infos = [m.features.get(vk, {}).get("info") or {} for m in all_metadata] video_infos = [m.features.get(vk, {}).get("info") or {} for m in all_metadata]
base_video_info = video_infos[0] base_video_info = video_infos[0]
merged_encoder_info: dict = {} merged_encoder_info: dict[str, Any] = {}
fallback_keys: list[str] = [] fallback_keys: list[str] = []
for info_key in VIDEO_ENCODER_INFO_KEYS: for info_key in VIDEO_ENCODER_INFO_KEYS:
values = [info.get(info_key, None) for info in video_infos] values = [info.get(info_key, None) for info in video_infos]
@@ -80,7 +105,7 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
merged_encoder_info[info_key] = {} if info_key == "video.extra_options" else None merged_encoder_info[info_key] = {} if info_key == "video.extra_options" else None
if fallback_keys: if fallback_keys:
logging.warning( logger.warning(
f"Merging heterogeneous or incomplete video encoder metadata for feature {vk}. " f"Merging heterogeneous or incomplete video encoder metadata for feature {vk}. "
f"Setting these keys to null: {fallback_keys}.", f"Setting these keys to null: {fallback_keys}.",
) )
@@ -92,7 +117,7 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
return merged_info return merged_info
def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]): def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]) -> tuple[int, str | None, FeatureDict]:
"""Validates that all dataset metadata have consistent properties. """Validates that all dataset metadata have consistent properties.
Ensures all datasets have the same fps, robot_type, and features to guarantee Ensures all datasets have the same fps, robot_type, and features to guarantee
@@ -129,7 +154,9 @@ def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]):
return fps, robot_type, features return fps, robot_type, features
def update_data_df(df, src_meta, dst_meta): def update_data_df(
df: pd.DataFrame, src_meta: LeRobotDatasetMetadata, dst_meta: LeRobotDatasetMetadata
) -> pd.DataFrame:
"""Updates a data DataFrame with new indices and task mappings for aggregation. """Updates a data DataFrame with new indices and task mappings for aggregation.
Adjusts episode indices, frame indices, and task indices to account for Adjusts episode indices, frame indices, and task indices to account for
@@ -154,12 +181,12 @@ def update_data_df(df, src_meta, dst_meta):
def update_meta_data( def update_meta_data(
df, df: pd.DataFrame,
dst_meta, dst_meta: LeRobotDatasetMetadata,
meta_idx, meta_idx: IndexState,
data_idx, data_idx: IndexState,
videos_idx, videos_idx: VideoIndexState,
): ) -> pd.DataFrame:
"""Updates metadata DataFrame with new chunk, file, and timestamp indices. """Updates metadata DataFrame with new chunk, file, and timestamp indices.
Adjusts all indices and timestamps to account for previously aggregated Adjusts all indices and timestamps to account for previously aggregated
@@ -289,7 +316,7 @@ def aggregate_datasets(
chunk_size: int | None = None, chunk_size: int | None = None,
concatenate_videos: bool = True, concatenate_videos: bool = True,
concatenate_data: bool = True, concatenate_data: bool = True,
): ) -> None:
"""Aggregates multiple LeRobot datasets into a single unified dataset. """Aggregates multiple LeRobot datasets into a single unified dataset.
This is the main function that orchestrates the aggregation process by: This is the main function that orchestrates the aggregation process by:
@@ -309,7 +336,7 @@ def aggregate_datasets(
concatenate_videos: When False, keep one mp4 per source file instead of packing into shards. concatenate_videos: When False, keep one mp4 per source file instead of packing into shards.
concatenate_data: When False, keep one parquet per source file instead of packing into shards. concatenate_data: When False, keep one parquet per source file instead of packing into shards.
""" """
logging.info("Start aggregate_datasets") logger.info("Start aggregate_datasets")
if data_files_size_in_mb is None: if data_files_size_in_mb is None:
data_files_size_in_mb = DEFAULT_DATA_FILE_SIZE_IN_MB data_files_size_in_mb = DEFAULT_DATA_FILE_SIZE_IN_MB
@@ -341,15 +368,15 @@ def aggregate_datasets(
video_files_size_in_mb=video_files_size_in_mb, video_files_size_in_mb=video_files_size_in_mb,
) )
logging.info("Find all tasks") logger.info("Find all tasks")
unique_tasks = pd.concat([m.tasks for m in all_metadata]).index.unique() unique_tasks = pd.concat([m.tasks for m in all_metadata]).index.unique()
dst_meta.tasks = pd.DataFrame( dst_meta.tasks = pd.DataFrame(
{"task_index": range(len(unique_tasks))}, index=pd.Index(unique_tasks, name="task") {"task_index": range(len(unique_tasks))}, index=pd.Index(unique_tasks, name="task")
) )
meta_idx = {"chunk": 0, "file": 0} meta_idx: IndexState = {"chunk": 0, "file": 0}
data_idx = {"chunk": 0, "file": 0} data_idx: IndexState = {"chunk": 0, "file": 0}
videos_idx = { videos_idx: VideoIndexState = {
key: {"chunk": 0, "file": 0, "latest_duration": 0, "episode_duration": 0} for key in video_keys key: {"chunk": 0, "file": 0, "latest_duration": 0, "episode_duration": 0} for key in video_keys
} }
@@ -373,12 +400,17 @@ def aggregate_datasets(
dst_meta.info.total_frames += src_meta.total_frames dst_meta.info.total_frames += src_meta.total_frames
finalize_aggregation(dst_meta, all_metadata) finalize_aggregation(dst_meta, all_metadata)
logging.info("Aggregation complete.") logger.info("Aggregation complete.")
def aggregate_videos( def aggregate_videos(
src_meta, dst_meta, videos_idx, video_files_size_in_mb, chunk_size, concatenate_videos=True src_meta: LeRobotDatasetMetadata,
): dst_meta: LeRobotDatasetMetadata,
videos_idx: VideoIndexState,
video_files_size_in_mb: float,
chunk_size: int,
concatenate_videos: bool = True,
) -> VideoIndexState:
"""Aggregates video chunks from a source dataset into the destination dataset. """Aggregates video chunks from a source dataset into the destination dataset.
Handles video file concatenation and rotation based on file size limits. Handles video file concatenation and rotation based on file size limits.
@@ -406,7 +438,8 @@ def aggregate_videos(
videos_idx[key]["dst_file_durations"] = {} videos_idx[key]["dst_file_durations"] = {}
for key, video_idx in videos_idx.items(): for key, video_idx in videos_idx.items():
unique_chunk_file_pairs = { unique_chunk_file_pairs: list[ChunkFile] = sorted(
{
(chunk, file) (chunk, file)
for chunk, file in zip( for chunk, file in zip(
src_meta.episodes[f"videos/{key}/chunk_index"], src_meta.episodes[f"videos/{key}/chunk_index"],
@@ -414,7 +447,7 @@ def aggregate_videos(
strict=False, strict=False,
) )
} }
unique_chunk_file_pairs = sorted(unique_chunk_file_pairs) )
chunk_idx = video_idx["chunk"] chunk_idx = video_idx["chunk"]
file_idx = video_idx["file"] file_idx = video_idx["file"]
@@ -489,7 +522,14 @@ def aggregate_videos(
return videos_idx return videos_idx
def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_size, concatenate_data=True): def aggregate_data(
src_meta: LeRobotDatasetMetadata,
dst_meta: LeRobotDatasetMetadata,
data_idx: IndexState,
data_files_size_in_mb: float,
chunk_size: int,
concatenate_data: bool = True,
) -> IndexState:
"""Aggregates data chunks from a source dataset into the destination dataset. """Aggregates data chunks from a source dataset into the destination dataset.
Reads source data files, updates indices to match the aggregated dataset, Reads source data files, updates indices to match the aggregated dataset,
@@ -510,14 +550,16 @@ def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_si
Returns: Returns:
dict: Updated data_idx with current chunk and file indices. dict: Updated data_idx with current chunk and file indices.
""" """
unique_chunk_file_ids = { unique_chunk_file_ids: list[ChunkFile] = sorted(
{
(c, f) (c, f)
for c, f in zip( for c, f in zip(
src_meta.episodes["data/chunk_index"], src_meta.episodes["data/file_index"], strict=False src_meta.episodes["data/chunk_index"],
src_meta.episodes["data/file_index"],
strict=False,
) )
} }
)
unique_chunk_file_ids = sorted(unique_chunk_file_ids)
contains_images = len(dst_meta.image_keys) > 0 contains_images = len(dst_meta.image_keys) > 0
# retrieve features schema for proper image typing in parquet # retrieve features schema for proper image typing in parquet
@@ -525,7 +567,7 @@ def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_si
# Track source to destination file mapping for metadata update # Track source to destination file mapping for metadata update
# This is critical for handling datasets that are already results of a merge # This is critical for handling datasets that are already results of a merge
src_to_dst: dict[tuple[int, int], tuple[int, int]] = {} src_to_dst: dict[ChunkFile, ChunkFile] = {}
for src_chunk_idx, src_file_idx in unique_chunk_file_ids: for src_chunk_idx, src_file_idx in unique_chunk_file_ids:
src_path = src_meta.root / DEFAULT_DATA_PATH.format( src_path = src_meta.root / DEFAULT_DATA_PATH.format(
@@ -564,7 +606,13 @@ def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_si
return data_idx return data_idx
def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx): def aggregate_metadata(
src_meta: LeRobotDatasetMetadata,
dst_meta: LeRobotDatasetMetadata,
meta_idx: IndexState,
data_idx: IndexState,
videos_idx: VideoIndexState,
) -> IndexState:
"""Aggregates metadata from a source dataset into the destination dataset. """Aggregates metadata from a source dataset into the destination dataset.
Reads source metadata files, updates all indices and timestamps, Reads source metadata files, updates all indices and timestamps,
@@ -580,7 +628,8 @@ def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx):
Returns: Returns:
dict: Updated meta_idx with current chunk and file indices. dict: Updated meta_idx with current chunk and file indices.
""" """
chunk_file_ids = { chunk_file_ids: list[ChunkFile] = sorted(
{
(c, f) (c, f)
for c, f in zip( for c, f in zip(
src_meta.episodes["meta/episodes/chunk_index"], src_meta.episodes["meta/episodes/chunk_index"],
@@ -588,8 +637,7 @@ def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx):
strict=False, strict=False,
) )
} }
)
chunk_file_ids = sorted(chunk_file_ids)
for chunk_idx, file_idx in chunk_file_ids: for chunk_idx, file_idx in chunk_file_ids:
src_path = src_meta.root / DEFAULT_EPISODES_PATH.format(chunk_index=chunk_idx, file_index=file_idx) src_path = src_meta.root / DEFAULT_EPISODES_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
df = pd.read_parquet(src_path) df = pd.read_parquet(src_path)
@@ -622,16 +670,16 @@ def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx):
def append_or_create_parquet_file( def append_or_create_parquet_file(
df: pd.DataFrame, df: pd.DataFrame,
src_path: Path, src_path: Path,
idx: dict[str, int], idx: IndexState,
max_mb: float, max_mb: float,
chunk_size: int, chunk_size: int,
default_path: str, default_path: str,
contains_images: bool = False, contains_images: bool = False,
aggr_root: Path = None, aggr_root: Path | None = None,
hf_features: datasets.Features | None = None, hf_features: datasets.Features | None = None,
concatenate: bool = True, concatenate: bool = True,
one_row_group_per_episode: bool = False, one_row_group_per_episode: bool = False,
) -> tuple[dict[str, int], tuple[int, int]]: ) -> tuple[IndexState, ChunkFile]:
"""Appends data to an existing parquet file or creates a new one based on size constraints. """Appends data to an existing parquet file or creates a new one based on size constraints.
Manages file rotation when size limits are exceeded to prevent individual files Manages file rotation when size limits are exceeded to prevent individual files
@@ -654,7 +702,13 @@ def append_or_create_parquet_file(
Returns: Returns:
tuple: (updated_idx, (dst_chunk, dst_file)) where updated_idx is the index dict tuple: (updated_idx, (dst_chunk, dst_file)) where updated_idx is the index dict
and (dst_chunk, dst_file) is the actual destination file the data was written to. and (dst_chunk, dst_file) is the actual destination file the data was written to.
Raises:
ValueError: If aggr_root is not provided.
""" """
if aggr_root is None:
raise ValueError("aggr_root must be provided.")
dst_chunk, dst_file = idx["chunk"], idx["file"] dst_chunk, dst_file = idx["chunk"], idx["file"]
dst_path = aggr_root / default_path.format(chunk_index=dst_chunk, file_index=dst_file) dst_path = aggr_root / default_path.format(chunk_index=dst_chunk, file_index=dst_file)
@@ -698,7 +752,9 @@ def append_or_create_parquet_file(
return idx, (dst_chunk, dst_file) return idx, (dst_chunk, dst_file)
def finalize_aggregation(aggr_meta, all_metadata): def finalize_aggregation(
aggr_meta: LeRobotDatasetMetadata, all_metadata: list[LeRobotDatasetMetadata]
) -> None:
"""Finalizes the dataset aggregation by writing summary files and statistics. """Finalizes the dataset aggregation by writing summary files and statistics.
Writes the tasks file, info file with total counts and splits, and Writes the tasks file, info file with total counts and splits, and
@@ -708,16 +764,16 @@ def finalize_aggregation(aggr_meta, all_metadata):
aggr_meta: Aggregated dataset metadata. aggr_meta: Aggregated dataset metadata.
all_metadata: List of all source dataset metadata objects. all_metadata: List of all source dataset metadata objects.
""" """
logging.info("write tasks") logger.info("write tasks")
write_tasks(aggr_meta.tasks, aggr_meta.root) write_tasks(aggr_meta.tasks, aggr_meta.root)
logging.info("write info") logger.info("write info")
aggr_meta.info.total_tasks = len(aggr_meta.tasks) aggr_meta.info.total_tasks = len(aggr_meta.tasks)
aggr_meta.info.total_episodes = sum(m.total_episodes for m in all_metadata) aggr_meta.info.total_episodes = sum(m.total_episodes for m in all_metadata)
aggr_meta.info.total_frames = sum(m.total_frames for m in all_metadata) aggr_meta.info.total_frames = sum(m.total_frames for m in all_metadata)
aggr_meta.info.splits = {"train": f"0:{sum(m.total_episodes for m in all_metadata)}"} aggr_meta.info.splits = {"train": f"0:{sum(m.total_episodes for m in all_metadata)}"}
write_info(aggr_meta.info, aggr_meta.root) write_info(aggr_meta.info, aggr_meta.root)
logging.info("write stats") logger.info("write stats")
aggr_meta.stats = aggregate_stats([m.stats for m in all_metadata]) aggr_meta.stats = aggregate_stats([m.stats for m in all_metadata])
write_stats(aggr_meta.stats, aggr_meta.root) write_stats(aggr_meta.stats, aggr_meta.root)
+5 -93
View File
@@ -18,13 +18,8 @@ from __future__ import annotations
import logging import logging
import numpy as np import numpy as np
import torch
from lerobot.processor import ( from lerobot.processor import RelativeActionsProcessorStep
RelativeActionsProcessorStep,
relative_action_output_dim,
to_relative_actions,
)
from lerobot.utils.constants import ACTION, OBS_STATE from lerobot.utils.constants import ACTION, OBS_STATE
from .io_utils import load_image_as_numpy from .io_utils import load_image_as_numpy
@@ -665,29 +660,17 @@ def _compute_relative_chunk_batch(
all_states: np.ndarray, all_states: np.ndarray,
chunk_size: int, chunk_size: int,
relative_mask: np.ndarray, relative_mask: np.ndarray,
pose_representation: str = "componentwise",
se3_pose_groups: list[list[int]] | None = None,
) -> np.ndarray: ) -> np.ndarray:
"""Vectorised relative-action computation for a batch of start indices. """Vectorised relative-action computation for a batch of start indices.
Returns an ``(N * chunk_size, model_action_dim)`` float32 array. Returns an ``(N * chunk_size, action_dim)`` float32 array.
""" """
if len(start_indices) == 0: if len(start_indices) == 0:
output_dim = relative_action_output_dim(all_actions.shape[1], pose_representation, se3_pose_groups) return np.empty((0, all_actions.shape[1]), dtype=np.float32)
return np.empty((0, output_dim), dtype=np.float32)
offsets = np.arange(chunk_size) offsets = np.arange(chunk_size)
frame_idx = start_indices[:, None] + offsets[None, :] frame_idx = start_indices[:, None] + offsets[None, :]
chunks = all_actions[frame_idx].copy() chunks = all_actions[frame_idx].copy()
states = all_states[start_indices] states = all_states[start_indices]
if pose_representation in {"se3", "se3_6d"}:
converted = to_relative_actions(
torch.from_numpy(chunks),
torch.from_numpy(states),
relative_mask.astype(bool).tolist(),
pose_representation=pose_representation,
se3_pose_groups=se3_pose_groups,
)
return converted.numpy().reshape(-1, converted.shape[-1])
mask_dim = len(relative_mask) mask_dim = len(relative_mask)
chunks[:, :, :mask_dim] -= states[:, None, :mask_dim] * relative_mask[None, None, :] chunks[:, :, :mask_dim] -= states[:, None, :mask_dim] * relative_mask[None, None, :]
return chunks.reshape(-1, all_actions.shape[1]) return chunks.reshape(-1, all_actions.shape[1])
@@ -699,9 +682,6 @@ def compute_relative_action_stats(
chunk_size: int, chunk_size: int,
exclude_joints: list[str] | None = None, exclude_joints: list[str] | None = None,
num_workers: int = 0, num_workers: int = 0,
state_from_action: bool = False,
pose_representation: str = "componentwise",
se3_pose_groups: list[list[int]] | None = None,
) -> dict[str, np.ndarray]: ) -> dict[str, np.ndarray]:
"""Compute normalization statistics for relative actions over the full dataset. """Compute normalization statistics for relative actions over the full dataset.
@@ -720,9 +700,6 @@ def compute_relative_action_stats(
num_workers: Number of parallel threads for computation. Values 1 num_workers: Number of parallel threads for computation. Values 1
mean single-threaded. Numpy releases the GIL so threads give mean single-threaded. Numpy releases the GIL so threads give
real parallelism here. real parallelism here.
state_from_action: Use the current absolute action as state. This is
intended for state-less pose datasets where each action row is the
synchronized measured robot pose.
Returns: Returns:
Statistics dict with keys "mean", "std", "min", "max", "q01", , "q99". Statistics dict with keys "mean", "std", "min", "max", "q01", , "q99".
@@ -745,7 +722,7 @@ def compute_relative_action_stats(
logging.info("Loading action/state data for relative action stats...") logging.info("Loading action/state data for relative action stats...")
all_actions = np.array(hf_dataset[ACTION], dtype=np.float32) all_actions = np.array(hf_dataset[ACTION], dtype=np.float32)
all_states = all_actions if state_from_action else np.array(hf_dataset[OBS_STATE], dtype=np.float32) all_states = np.array(hf_dataset[OBS_STATE], dtype=np.float32)
episode_indices = np.array(hf_dataset["episode_index"]) episode_indices = np.array(hf_dataset["episode_index"])
valid_starts = _get_valid_chunk_starts(episode_indices, chunk_size) valid_starts = _get_valid_chunk_starts(episode_indices, chunk_size)
@@ -777,8 +754,6 @@ def compute_relative_action_stats(
all_states, all_states,
chunk_size, chunk_size,
relative_mask, relative_mask,
pose_representation,
se3_pose_groups,
) )
for batch in batches for batch in batches
] ]
@@ -787,15 +762,7 @@ def compute_relative_action_stats(
else: else:
for batch in batches: for batch in batches:
running_stats.update( running_stats.update(
_compute_relative_chunk_batch( _compute_relative_chunk_batch(batch, all_actions, all_states, chunk_size, relative_mask)
batch,
all_actions,
all_states,
chunk_size,
relative_mask,
pose_representation,
se3_pose_groups,
)
) )
stats = running_stats.get_statistics() stats = running_stats.get_statistics()
@@ -810,58 +777,3 @@ def compute_relative_action_stats(
) )
return stats return stats
def compute_state_history_stats(
hf_dataset,
features: dict,
history_steps: int,
exclude_joints: list[str] | None = None,
relative: bool = False,
pose_representation: str = "componentwise",
se3_pose_groups: list[list[int]] | None = None,
) -> dict[str, np.ndarray]:
"""Compute stats for flattened state history synthesized from absolute actions.
History is left-padded with the first action of each episode, matching dataset
boundary padding. When ``relative`` is enabled, every history pose is expressed
relative to its newest pose while excluded dimensions remain absolute.
"""
if history_steps < 1:
raise ValueError("history_steps must be at least 1")
if exclude_joints is None:
exclude_joints = []
actions = np.asarray(hf_dataset[ACTION], dtype=np.float32)
episode_indices = np.asarray(hf_dataset["episode_index"])
sample_indices = np.arange(len(actions))
episode_starts = np.maximum.accumulate(
np.where(
np.concatenate(([True], episode_indices[1:] != episode_indices[:-1])),
sample_indices,
0,
)
)
offsets = np.arange(-(history_steps - 1), 1)
history_indices = np.maximum(sample_indices[:, None] + offsets[None, :], episode_starts[:, None])
history = actions[history_indices].copy()
if relative:
state_dim = actions.shape[-1]
names = features.get(ACTION, {}).get("names")
mask_step = RelativeActionsProcessorStep(
enabled=True,
exclude_joints=exclude_joints,
action_names=names,
)
mask = mask_step._build_mask(state_dim)
history = to_relative_actions(
torch.from_numpy(history),
torch.from_numpy(history[:, -1].copy()),
mask,
pose_representation=pose_representation,
se3_pose_groups=se3_pose_groups,
).numpy()
flattened = history.reshape(len(history), -1)
return get_feature_stats(flattened, axis=0, keepdims=False)
+52 -6
View File
@@ -18,13 +18,15 @@ import logging
from collections.abc import Callable, Iterable from collections.abc import Callable, Iterable
from copy import deepcopy from copy import deepcopy
from pathlib import Path from pathlib import Path
from typing import Literal
import numpy as np import numpy as np
import packaging.version import packaging.version
import pandas as pd import pandas as pd
import pyarrow as pa import pyarrow as pa
import pyarrow.parquet as pq import pyarrow.parquet as pq
from huggingface_hub import snapshot_download from huggingface_hub import snapshot_download, sync_bucket
from huggingface_hub.utils import WeakFileLock
from lerobot.configs import DEPTH_METER_UNIT, VideoEncoderConfig from lerobot.configs import DEPTH_METER_UNIT, VideoEncoderConfig
from lerobot.utils.constants import DEFAULT_FEATURES, HF_LEROBOT_HOME, HF_LEROBOT_HUB_CACHE from lerobot.utils.constants import DEFAULT_FEATURES, HF_LEROBOT_HOME, HF_LEROBOT_HUB_CACHE
@@ -73,6 +75,9 @@ class LeRobotDatasetMetadata:
revision: str | None = None, revision: str | None = None,
force_cache_sync: bool = False, force_cache_sync: bool = False,
metadata_buffer_size: int = 10, metadata_buffer_size: int = 10,
*,
repo_type: Literal["dataset", "bucket"] = "dataset",
token: str | bool | None = None,
): ):
"""Load or download metadata for an existing LeRobot dataset. """Load or download metadata for an existing LeRobot dataset.
@@ -94,17 +99,38 @@ class LeRobotDatasetMetadata:
even when local files exist. even when local files exist.
metadata_buffer_size: Number of episode metadata records to buffer metadata_buffer_size: Number of episode metadata records to buffer
in memory before flushing to parquet. in memory before flushing to parquet.
repo_type: Repository type: "dataset" (default) or "bucket" for an
HF Storage Bucket streamed over hf://buckets/.
token: Authentication token used for Hub requests. Pass a string
token, ``True`` to require the locally stored token, ``False``
to disable authentication, or ``None`` to use the Hugging Face
Hub default.
""" """
if repo_type not in ("dataset", "bucket"):
raise ValueError(f"repo_type must be 'dataset' or 'bucket', got {repo_type!r}")
self.repo_id = repo_id self.repo_id = repo_id
self.repo_type = repo_type
self.revision = revision if revision else CODEBASE_VERSION self.revision = revision if revision else CODEBASE_VERSION
self._requested_root = Path(root) if root is not None else None self._requested_root = Path(root) if root is not None else None
self.root = self._requested_root if self._requested_root is not None else HF_LEROBOT_HOME / repo_id if self._requested_root is not None:
self.root = self._requested_root
elif self.repo_type == "bucket":
self.root = HF_LEROBOT_HUB_CACHE / ("buckets--" + self.repo_id.replace("/", "--"))
else:
self.root = HF_LEROBOT_HOME / repo_id
self._pq_writer = None self._pq_writer = None
self.latest_episode = None self.latest_episode = None
self._metadata_buffer: list[dict] = [] self._metadata_buffer: list[dict] = []
self._metadata_buffer_size = metadata_buffer_size self._metadata_buffer_size = metadata_buffer_size
self._finalized = False self._finalized = False
metadata_lock = contextlib.nullcontext()
if self.repo_type == "bucket":
self.root.parent.mkdir(parents=True, exist_ok=True)
metadata_lock = WeakFileLock(self.root.parent / f".{self.root.name}.lock")
with metadata_lock:
try: try:
if force_cache_sync or ( if force_cache_sync or (
self._requested_root is None and has_legacy_hub_download_metadata(self.root) self._requested_root is None and has_legacy_hub_download_metadata(self.root)
@@ -112,10 +138,13 @@ class LeRobotDatasetMetadata:
raise FileNotFoundError raise FileNotFoundError
self._load_metadata() self._load_metadata()
except (FileNotFoundError, NotADirectoryError): except (FileNotFoundError, NotADirectoryError):
if is_valid_version(self.revision): if self.repo_type != "bucket" and is_valid_version(self.revision):
if token is None:
self.revision = get_safe_version(self.repo_id, self.revision) self.revision = get_safe_version(self.repo_id, self.revision)
else:
self.revision = get_safe_version(self.repo_id, self.revision, token=token)
self._pull_from_repo(allow_patterns="meta/") self._pull_from_repo(allow_patterns="meta/", token=token)
self._load_metadata() self._load_metadata()
def _flush_metadata_buffer(self) -> None: def _flush_metadata_buffer(self) -> None:
@@ -179,8 +208,8 @@ class LeRobotDatasetMetadata:
def _load_metadata(self): def _load_metadata(self):
self.info = load_info(self.root) self.info = load_info(self.root)
check_version_compatibility(self.repo_id, self._version, CODEBASE_VERSION) check_version_compatibility(self.repo_id, self._version, CODEBASE_VERSION)
self.tasks = load_tasks(self.root) self.tasks = load_tasks(self.root) if self.total_tasks > 0 else None
self.episodes = load_episodes(self.root) self.episodes = load_episodes(self.root) if self.total_episodes > 0 else None
self.stats = load_stats(self.root) self.stats = load_stats(self.root)
def ensure_readable(self) -> None: def ensure_readable(self) -> None:
@@ -220,7 +249,20 @@ class LeRobotDatasetMetadata:
self, self,
allow_patterns: list[str] | str | None = None, allow_patterns: list[str] | str | None = None,
ignore_patterns: list[str] | str | None = None, ignore_patterns: list[str] | str | None = None,
*,
token: str | bool | None = None,
) -> None: ) -> None:
if self.repo_type == "bucket":
self.root.mkdir(parents=True, exist_ok=True)
sync_bucket(
f"hf://buckets/{self.repo_id}/meta",
str(self.root / "meta"),
delete=True,
quiet=True,
token=token,
)
return
token_kwargs = {} if token is None else {"token": token}
if self._requested_root is None: if self._requested_root is None:
self.root = Path( self.root = Path(
snapshot_download( snapshot_download(
@@ -230,6 +272,7 @@ class LeRobotDatasetMetadata:
cache_dir=HF_LEROBOT_HUB_CACHE, cache_dir=HF_LEROBOT_HUB_CACHE,
allow_patterns=allow_patterns, allow_patterns=allow_patterns,
ignore_patterns=ignore_patterns, ignore_patterns=ignore_patterns,
**token_kwargs,
) )
) )
return return
@@ -242,12 +285,15 @@ class LeRobotDatasetMetadata:
local_dir=self._requested_root, local_dir=self._requested_root,
allow_patterns=allow_patterns, allow_patterns=allow_patterns,
ignore_patterns=ignore_patterns, ignore_patterns=ignore_patterns,
**token_kwargs,
) )
self.root = self._requested_root self.root = self._requested_root
@property @property
def url_root(self) -> str: def url_root(self) -> str:
"""Hugging Face Hub URL root for this dataset.""" """Hugging Face Hub URL root for this dataset."""
if self.repo_type == "bucket":
return f"hf://buckets/{self.repo_id}"
return f"hf://datasets/{self.repo_id}" return f"hf://datasets/{self.repo_id}"
@property @property
+1 -37
View File
@@ -54,7 +54,6 @@ from .compute_stats import (
aggregate_stats, aggregate_stats,
compute_episode_stats, compute_episode_stats,
compute_relative_action_stats, compute_relative_action_stats,
compute_state_history_stats,
) )
from .dataset_metadata import LeRobotDatasetMetadata from .dataset_metadata import LeRobotDatasetMetadata
from .image_writer import write_image from .image_writer import write_image
@@ -1567,12 +1566,6 @@ def recompute_stats(
relative_exclude_joints: list[str] | None = None, relative_exclude_joints: list[str] | None = None,
chunk_size: int = 50, chunk_size: int = 50,
num_workers: int = 0, num_workers: int = 0,
state_from_action: bool = False,
state_history_steps: int = 1,
relative_state_history: bool = False,
relative_state_exclude_joints: list[str] | None = None,
relative_pose_representation: str = "componentwise",
relative_se3_pose_groups: list[list[int]] | None = None,
) -> LeRobotDataset: ) -> LeRobotDataset:
"""Recompute stats.json from scratch by iterating all episodes. """Recompute stats.json from scratch by iterating all episodes.
@@ -1590,16 +1583,6 @@ def recompute_stats(
``policy.chunk_size``. Only used when ``relative_action=True``. ``policy.chunk_size``. Only used when ``relative_action=True``.
num_workers: Number of parallel threads for relative action stats computation. num_workers: Number of parallel threads for relative action stats computation.
Values 1 mean single-threaded. Only used when ``relative_action=True``. Values 1 mean single-threaded. Only used when ``relative_action=True``.
state_from_action: Use absolute action rows as synthetic state while
computing relative-action stats, and write their absolute statistics
under ``observation.state``.
state_history_steps: Number of consecutive synthesized state samples.
relative_state_history: Express state history relative to its newest pose.
relative_state_exclude_joints: State dimensions to retain as absolute.
relative_pose_representation: ``componentwise`` for legacy subtraction,
``se3`` for composition with an axis-angle output, or ``se3_6d`` for
composition with a continuous two-column rotation output.
relative_se3_pose_groups: Six-index xyz+rotation-vector pose groups.
Returns: Returns:
The same dataset with updated stats. The same dataset with updated stats.
@@ -1623,21 +1606,7 @@ def recompute_stats(
# (matching what the model sees during training) and skip action in the # (matching what the model sees during training) and skip action in the
# per-episode pass below. # per-episode pass below.
relative_action_stats = None relative_action_stats = None
synthetic_state_stats = None if relative_action and ACTION in features and OBS_STATE in features:
if state_from_action:
if ACTION not in features:
raise ValueError("state_from_action requires an action feature")
synthetic_state_stats = compute_state_history_stats(
dataset.hf_dataset,
features,
history_steps=state_history_steps,
exclude_joints=relative_state_exclude_joints,
relative=relative_state_history,
pose_representation=relative_pose_representation,
se3_pose_groups=relative_se3_pose_groups,
)
if relative_action and ACTION in features and (OBS_STATE in features or state_from_action):
if relative_exclude_joints is None: if relative_exclude_joints is None:
relative_exclude_joints = ["gripper"] relative_exclude_joints = ["gripper"]
relative_action_stats = compute_relative_action_stats( relative_action_stats = compute_relative_action_stats(
@@ -1646,9 +1615,6 @@ def recompute_stats(
chunk_size=chunk_size, chunk_size=chunk_size,
exclude_joints=relative_exclude_joints, exclude_joints=relative_exclude_joints,
num_workers=num_workers, num_workers=num_workers,
state_from_action=state_from_action,
pose_representation=relative_pose_representation,
se3_pose_groups=relative_se3_pose_groups,
) )
features_to_compute.pop(ACTION, None) features_to_compute.pop(ACTION, None)
@@ -1688,8 +1654,6 @@ def recompute_stats(
if relative_action_stats is not None: if relative_action_stats is not None:
new_stats[ACTION] = relative_action_stats new_stats[ACTION] = relative_action_stats
if synthetic_state_stats is not None:
new_stats[OBS_STATE] = synthetic_state_stats
# Merge: keep existing stats for features we didn't recompute # Merge: keep existing stats for features we didn't recompute
if dataset.meta.stats: if dataset.meta.stats:
+23 -14
View File
@@ -172,6 +172,23 @@ class DatasetWriter:
def _get_image_file_dir(self, episode_index: int, image_key: str) -> Path: def _get_image_file_dir(self, episode_index: int, image_key: str) -> Path:
return self._get_image_file_path(episode_index, image_key, frame_index=0).parent return self._get_image_file_path(episode_index, image_key, frame_index=0).parent
def _get_episode_buffer_index(self) -> int:
episode_index = self.episode_buffer["episode_index"]
# episode_index is `int` when freshly created, but becomes `np.ndarray` after
# save_episode() mutates the buffer. Handle both types here.
if isinstance(episode_index, np.ndarray):
episode_index = episode_index.item() if episode_index.size == 1 else episode_index[0]
return int(episode_index)
def _delete_camera_frame_dirs(self, camera_keys: list[str]) -> None:
if self.image_writer is not None:
self._wait_image_writer()
episode_index = self._get_episode_buffer_index()
for camera_key in camera_keys:
img_dir = self._get_image_file_dir(episode_index, camera_key)
if img_dir.is_dir():
shutil.rmtree(img_dir)
def _save_image( def _save_image(
self, image: torch.Tensor | np.ndarray | PIL.Image.Image, fpath: Path, compress_level: int = 1 self, image: torch.Tensor | np.ndarray | PIL.Image.Image, fpath: Path, compress_level: int = 1
) -> None: ) -> None:
@@ -369,7 +386,9 @@ class DatasetWriter:
self._episodes_since_last_encoding = 0 self._episodes_since_last_encoding = 0
if episode_data is None: if episode_data is None:
self.clear_episode_buffer(delete_images=len(self._meta.image_keys) > 0) if len(self._meta.image_keys) > 0:
self._delete_camera_frame_dirs(self._meta.image_keys)
self.episode_buffer = self._create_episode_buffer()
def _batch_save_episode_video(self, start_episode: int, end_episode: int | None = None) -> None: def _batch_save_episode_video(self, start_episode: int, end_episode: int | None = None) -> None:
"""Batch save videos for multiple episodes.""" """Batch save videos for multiple episodes."""
@@ -561,10 +580,10 @@ class DatasetWriter:
return metadata return metadata
def clear_episode_buffer(self, delete_images: bool = True) -> None: def clear_episode_buffer(self, delete_images: bool = True) -> None:
"""Discard the current episode buffer and optionally delete temp images. """Discard the current episode buffer and optionally delete temp camera frames.
Args: Args:
delete_images: If ``True``, remove temporary image directories delete_images: If ``True``, remove temporary camera frame directories
written for the current episode. written for the current episode.
""" """
# Cancel streaming encoder if active # Cancel streaming encoder if active
@@ -572,17 +591,7 @@ class DatasetWriter:
self._streaming_encoder.cancel_episode() self._streaming_encoder.cancel_episode()
if delete_images: if delete_images:
if self.image_writer is not None: self._delete_camera_frame_dirs(self._meta.camera_keys)
self._wait_image_writer()
episode_index = self.episode_buffer["episode_index"]
# episode_index is `int` when freshly created, but becomes `np.ndarray` after
# save_episode() mutates the buffer. Handle both types here.
if isinstance(episode_index, np.ndarray):
episode_index = episode_index.item() if episode_index.size == 1 else episode_index[0]
for cam_key in self._meta.image_keys:
img_dir = self._get_image_file_dir(episode_index, cam_key)
if img_dir.is_dir():
shutil.rmtree(img_dir)
self.episode_buffer = self._create_episode_buffer() self.episode_buffer = self._create_episode_buffer()
+38 -9
View File
@@ -65,6 +65,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
encoder_threads: int | None = None, encoder_threads: int | None = None,
streaming_encoding: bool = False, streaming_encoding: bool = False,
encoder_queue_maxsize: int = 30, encoder_queue_maxsize: int = 30,
*,
token: str | bool | None = None,
): ):
""" """
2 modes are available for instantiating this class, depending on 2 different use cases: 2 modes are available for instantiating this class, depending on 2 different use cases:
@@ -197,6 +199,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
instead of writing PNG images first. This makes save_episode() near-instant. Defaults to False. instead of writing PNG images first. This makes save_episode() near-instant. Defaults to False.
encoder_queue_maxsize (int, optional): Maximum number of frames to buffer per camera when using encoder_queue_maxsize (int, optional): Maximum number of frames to buffer per camera when using
streaming encoding. Defaults to 30 (~1s at 30fps). streaming encoding. Defaults to 30 (~1s at 30fps).
token: Authentication token used while downloading this dataset
from the Hub. Pass a string token, ``True`` to require the
locally stored token, ``False`` to disable authentication, or
``None`` to use the Hugging Face Hub default. The token is not
retained on the dataset instance after initialization.
Note: Note:
Write-mode parameters (``streaming_encoding``, ``batch_encoding_size``) passed to Write-mode parameters (``streaming_encoding``, ``batch_encoding_size``) passed to
@@ -220,7 +227,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
# Load metadata (sets self.root once from the resolved metadata root) # Load metadata (sets self.root once from the resolved metadata root)
self.meta = LeRobotDatasetMetadata( self.meta = LeRobotDatasetMetadata(
self.repo_id, self._requested_root, self.revision, force_cache_sync=force_cache_sync self.repo_id,
self._requested_root,
self.revision,
force_cache_sync=force_cache_sync,
token=token,
) )
self.root = self.meta.root self.root = self.meta.root
self.revision = self.meta.revision self.revision = self.meta.revision
@@ -260,8 +271,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
# Load actual data # Load actual data
if force_cache_sync or not self.reader.try_load(): if force_cache_sync or not self.reader.try_load():
if is_valid_version(self.revision): if is_valid_version(self.revision):
if token is None:
self.revision = get_safe_version(self.repo_id, self.revision) self.revision = get_safe_version(self.repo_id, self.revision)
self._download(download_videos) else:
self.revision = get_safe_version(self.repo_id, self.revision, token=token)
self._download(download_videos, token=token)
self.reader.load_and_activate() self.reader.load_and_activate()
# Detect write-mode params for backward compatibility # Detect write-mode params for backward compatibility
@@ -478,18 +492,19 @@ class LeRobotDataset(torch.utils.data.Dataset):
"""Return the number of frames in the selected episodes.""" """Return the number of frames in the selected episodes."""
return self.num_frames return self.num_frames
def __getitem__(self, idx) -> dict: def __getitem__(self, idx: int | slice) -> dict | list[dict]:
"""Return a single frame by index, with all transforms applied. """Return one frame or a slice of frames, with all transforms applied.
Loads the frame from the underlying HF dataset, expands delta-timestamp Loads the frame from the underlying HF dataset, expands delta-timestamp
windows, decodes video frames, and applies image transforms. Delegates windows, decodes video frames, and applies image transforms. Delegates
the core logic to :meth:`DatasetReader.get_item`. the core logic to :class:`DatasetReader`.
Args: Args:
idx: Index into the (possibly episode-filtered) dataset. idx: Integer index or slice into the possibly episode-filtered dataset.
Returns: Returns:
Dict mapping feature names to their tensor values for this frame. A frame dictionary for an integer index, or a list of frame
dictionaries for a slice.
Raises: Raises:
RuntimeError: If the dataset is currently being recorded and RuntimeError: If the dataset is currently being recorded and
@@ -499,6 +514,9 @@ class LeRobotDataset(torch.utils.data.Dataset):
raise RuntimeError( raise RuntimeError(
"Cannot read from a dataset that is being recorded. Call finalize() first, then access items." "Cannot read from a dataset that is being recorded. Call finalize() first, then access items."
) )
if isinstance(idx, slice):
return [self[item_idx] for item_idx in range(*idx.indices(len(self)))]
reader = self._ensure_reader() reader = self._ensure_reader()
if reader.hf_dataset is None: if reader.hf_dataset is None:
# One-shot load after finalize() # One-shot load after finalize()
@@ -622,10 +640,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
hub_api.delete_tag(self.repo_id, tag=CODEBASE_VERSION, repo_type="dataset") hub_api.delete_tag(self.repo_id, tag=CODEBASE_VERSION, repo_type="dataset")
hub_api.create_tag(self.repo_id, tag=CODEBASE_VERSION, revision=branch, repo_type="dataset") hub_api.create_tag(self.repo_id, tag=CODEBASE_VERSION, revision=branch, repo_type="dataset")
def _download(self, download_videos: bool = True) -> None: def _download(self, download_videos: bool = True, *, token: str | bool | None = None) -> None:
"""Downloads the dataset from the given 'repo_id' at the provided version.""" """Downloads the dataset from the given 'repo_id' at the provided version."""
ignore_patterns = None if download_videos else "videos/" ignore_patterns = None if download_videos else "videos/"
files = None files = None
token_kwargs = {} if token is None else {"token": token}
if self.episodes is not None: if self.episodes is not None:
# Reader is guaranteed to exist here (created in __init__ before _download) # Reader is guaranteed to exist here (created in __init__ before _download)
files = self.reader.get_episodes_file_paths() files = self.reader.get_episodes_file_paths()
@@ -639,6 +658,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
cache_dir=HF_LEROBOT_HUB_CACHE, cache_dir=HF_LEROBOT_HUB_CACHE,
allow_patterns=files, allow_patterns=files,
ignore_patterns=ignore_patterns, ignore_patterns=ignore_patterns,
**token_kwargs,
) )
) )
else: else:
@@ -650,6 +670,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
local_dir=self._requested_root, local_dir=self._requested_root,
allow_patterns=files, allow_patterns=files,
ignore_patterns=ignore_patterns, ignore_patterns=ignore_patterns,
**token_kwargs,
) )
self.meta.root = self._requested_root self.meta.root = self._requested_root
@@ -789,6 +810,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
image_writer_threads: int = 0, image_writer_threads: int = 0,
streaming_encoding: bool = False, streaming_encoding: bool = False,
encoder_queue_maxsize: int = 30, encoder_queue_maxsize: int = 30,
*,
token: str | bool | None = None,
) -> "LeRobotDataset": ) -> "LeRobotDataset":
"""Resume recording on an existing dataset. """Resume recording on an existing dataset.
@@ -822,6 +845,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
streaming_encoding: If ``True``, encode video in real-time during streaming_encoding: If ``True``, encode video in real-time during
capture. capture.
encoder_queue_maxsize: Max buffered frames per camera for streaming. encoder_queue_maxsize: Max buffered frames per camera for streaming.
token: Authentication token used if metadata must be downloaded
from the Hub. The token is not retained on the dataset instance.
Returns: Returns:
A :class:`LeRobotDataset` in write mode, ready to append episodes. A :class:`LeRobotDataset` in write mode, ready to append episodes.
@@ -850,7 +875,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
# Load metadata (revision-safe when root is not provided) # Load metadata (revision-safe when root is not provided)
obj.meta = LeRobotDatasetMetadata( obj.meta = LeRobotDatasetMetadata(
obj.repo_id, obj._requested_root, obj.revision, force_cache_sync=force_cache_sync obj.repo_id,
obj._requested_root,
obj.revision,
force_cache_sync=force_cache_sync,
token=token,
) )
obj._encoder_threads = encoder_threads obj._encoder_threads = encoder_threads
+3
View File
@@ -48,6 +48,8 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
tolerances_s: dict | None = None, tolerances_s: dict | None = None,
download_videos: bool = True, download_videos: bool = True,
video_backend: str | None = None, video_backend: str | None = None,
*,
token: str | bool | None = None,
): ):
super().__init__() super().__init__()
self.repo_ids = repo_ids self.repo_ids = repo_ids
@@ -65,6 +67,7 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
tolerance_s=self.tolerances_s[repo_id], tolerance_s=self.tolerances_s[repo_id],
download_videos=download_videos, download_videos=download_videos,
video_backend=video_backend, video_backend=video_backend,
token=token,
) )
for repo_id in repo_ids for repo_id in repo_ids
] ]
+1 -1
View File
@@ -17,8 +17,8 @@ from collections.abc import Sequence
from typing import Any from typing import Any
from lerobot.configs import PipelineFeatureType from lerobot.configs import PipelineFeatureType
from lerobot.lerobot_types import RobotAction, RobotObservation
from lerobot.processor import DataProcessorPipeline from lerobot.processor import DataProcessorPipeline
from lerobot.types import RobotAction, RobotObservation
from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE, OBS_STR from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE, OBS_STR
from lerobot.utils.feature_utils import hw_to_dataset_features from lerobot.utils.feature_utils import hw_to_dataset_features
+39 -6
View File
@@ -16,6 +16,7 @@
from collections import deque from collections import deque
from collections.abc import Callable, Generator, Iterable, Iterator from collections.abc import Callable, Generator, Iterable, Iterator
from pathlib import Path from pathlib import Path
from typing import Literal
import datasets import datasets
import numpy as np import numpy as np
@@ -256,14 +257,17 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
shuffle: bool = True, shuffle: bool = True,
return_uint8: bool = False, return_uint8: bool = False,
depth_output_unit: str = DEFAULT_DEPTH_UNIT, depth_output_unit: str = DEFAULT_DEPTH_UNIT,
*,
repo_type: Literal["dataset", "bucket"] = "dataset",
token: str | bool | None = None,
): ):
"""Initialize a StreamingLeRobotDataset. """Initialize a StreamingLeRobotDataset.
Args: Args:
repo_id (str): This is the repo id that will be used to fetch the dataset. repo_id (str): This is the repo id that will be used to fetch the dataset.
root (Path | None, optional): Local directory to use for local datasets. When omitted, Hub root (Path | None, optional): Local directory to use for local datasets. In bucket mode,
metadata is resolved through a revision-safe snapshot cache under this is an optional local metadata-cache directory; parquet and video data remain remote.
``$HF_LEROBOT_HOME/hub``. When omitted, Hub metadata is resolved through the cache under ``$HF_LEROBOT_HOME/hub``.
episodes (list[int] | None, optional): If specified, this will only load episodes specified by episodes (list[int] | None, optional): If specified, this will only load episodes specified by
their episode_index in this list. their episode_index in this list.
image_transforms (Callable | None, optional): Transform to apply to image data. image_transforms (Callable | None, optional): Transform to apply to image data.
@@ -278,12 +282,23 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
shuffle (bool, optional): Whether to shuffle the dataset across exhaustions. Defaults to True. shuffle (bool, optional): Whether to shuffle the dataset across exhaustions. Defaults to True.
depth_output_unit (str, optional): Physical unit depth maps are dequantized to ("m" or "mm"). depth_output_unit (str, optional): Physical unit depth maps are dequantized to ("m" or "mm").
Defaults to "mm". Defaults to "mm".
repo_type: "dataset" (default) or "bucket" to stream from an HF Storage Bucket
over ``hf://buckets/``.
token: Authentication token used while streaming this dataset from
the Hub. Pass a string token, ``True`` to require the locally
stored token, ``False`` to disable authentication, or ``None``
to use the Hugging Face Hub default. The token is not retained
on the dataset instance after initialization.
""" """
super().__init__() super().__init__()
if repo_type not in ("dataset", "bucket"):
raise ValueError(f"repo_type must be 'dataset' or 'bucket', got {repo_type!r}")
self.repo_id = repo_id self.repo_id = repo_id
self._requested_root = Path(root) if root else None self.repo_type = repo_type
self._requested_root = Path(root) if root is not None else None
self.root = self._requested_root if self._requested_root is not None else HF_LEROBOT_HOME / repo_id self.root = self._requested_root if self._requested_root is not None else HF_LEROBOT_HOME / repo_id
self.streaming_from_local = root is not None self.streaming_from_local = root is not None and self.repo_type == "dataset"
self.image_transforms = image_transforms self.image_transforms = image_transforms
self.episodes = episodes self.episodes = episodes
@@ -306,7 +321,12 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
# Load metadata # Load metadata
self.meta = LeRobotDatasetMetadata( self.meta = LeRobotDatasetMetadata(
self.repo_id, self._requested_root, self.revision, force_cache_sync=force_cache_sync self.repo_id,
self._requested_root,
self.revision,
force_cache_sync=force_cache_sync,
repo_type=self.repo_type,
token=token,
) )
self.root = self.meta.root self.root = self.meta.root
self.revision = self.meta.revision self.revision = self.meta.revision
@@ -334,12 +354,25 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
self.delta_timestamps = delta_timestamps self.delta_timestamps = delta_timestamps
self.delta_indices = get_delta_indices(self.delta_timestamps, self.fps) self.delta_indices = get_delta_indices(self.delta_timestamps, self.fps)
token_kwargs = {} if token is None else {"token": token}
if self.repo_type == "bucket":
self.hf_dataset: datasets.IterableDataset = load_dataset(
"parquet",
data_files=f"hf://buckets/{self.repo_id}/data/*/*.parquet",
split="train",
streaming=self.streaming,
**token_kwargs,
)
else:
if self.streaming_from_local:
token_kwargs = {}
self.hf_dataset: datasets.IterableDataset = load_dataset( self.hf_dataset: datasets.IterableDataset = load_dataset(
self.repo_id if not self.streaming_from_local else str(self.root), self.repo_id if not self.streaming_from_local else str(self.root),
split="train", split="train",
streaming=self.streaming, streaming=self.streaming,
data_files="data/*/*.parquet", data_files="data/*/*.parquet",
revision=self.revision, revision=self.revision,
**token_kwargs,
) )
self.num_shards = min(self.hf_dataset.num_shards, max_num_shards) self.num_shards = min(self.hf_dataset.num_shards, max_num_shards)
+13 -4
View File
@@ -325,16 +325,19 @@ def check_version_compatibility(
logging.warning(FUTURE_MESSAGE.format(repo_id=repo_id, version=v_check)) logging.warning(FUTURE_MESSAGE.format(repo_id=repo_id, version=v_check))
def get_repo_versions(repo_id: str) -> list[packaging.version.Version]: def get_repo_versions(repo_id: str, *, token: str | bool | None = None) -> list[packaging.version.Version]:
"""Return available valid versions (branches and tags) on a given Hub repo. """Return available valid versions (branches and tags) on a given Hub repo.
Args: Args:
repo_id (str): The repository ID on the Hugging Face Hub. repo_id (str): The repository ID on the Hugging Face Hub.
token: Authentication token used for Hub requests. Pass a string token,
``True`` to require the locally stored token, ``False`` to disable
authentication, or ``None`` to use the Hugging Face Hub default.
Returns: Returns:
list[packaging.version.Version]: A list of valid versions found. list[packaging.version.Version]: A list of valid versions found.
""" """
api = HfApi() api = HfApi() if token is None else HfApi(token=token)
repo_refs = api.list_repo_refs(repo_id, repo_type="dataset") repo_refs = api.list_repo_refs(repo_id, repo_type="dataset")
repo_refs = [b.name for b in repo_refs.branches + repo_refs.tags] repo_refs = [b.name for b in repo_refs.branches + repo_refs.tags]
repo_versions = [] repo_versions = []
@@ -345,7 +348,12 @@ def get_repo_versions(repo_id: str) -> list[packaging.version.Version]:
return repo_versions return repo_versions
def get_safe_version(repo_id: str, version: str | packaging.version.Version) -> str: def get_safe_version(
repo_id: str,
version: str | packaging.version.Version,
*,
token: str | bool | None = None,
) -> str:
"""Return the specified version if available on repo, or the latest compatible one. """Return the specified version if available on repo, or the latest compatible one.
If the exact version is not found, it looks for the latest version with the If the exact version is not found, it looks for the latest version with the
@@ -354,6 +362,7 @@ def get_safe_version(repo_id: str, version: str | packaging.version.Version) ->
Args: Args:
repo_id (str): The repository ID on the Hugging Face Hub. repo_id (str): The repository ID on the Hugging Face Hub.
version (str | packaging.version.Version): The target version. version (str | packaging.version.Version): The target version.
token: Authentication token forwarded to the Hub version lookup.
Returns: Returns:
str: The safe version string (e.g., "v1.2.3") to use as a revision. str: The safe version string (e.g., "v1.2.3") to use as a revision.
@@ -366,7 +375,7 @@ def get_safe_version(repo_id: str, version: str | packaging.version.Version) ->
target_version = ( target_version = (
packaging.version.parse(version) if not isinstance(version, packaging.version.Version) else version packaging.version.parse(version) if not isinstance(version, packaging.version.Version) else version
) )
hub_versions = get_repo_versions(repo_id) hub_versions = get_repo_versions(repo_id) if token is None else get_repo_versions(repo_id, token=token)
if not hub_versions: if not hub_versions:
raise RevisionNotFoundError( raise RevisionNotFoundError(
+6 -2
View File
@@ -322,7 +322,7 @@ class HILSerlRobotEnvConfig(EnvConfig):
class LiberoEnv(EnvConfig): class LiberoEnv(EnvConfig):
task: str = "libero_10" # can also choose libero_spatial, libero_object, etc. task: str = "libero_10" # can also choose libero_spatial, libero_object, etc.
task_ids: list[int] | None = None task_ids: list[int] | None = None
fps: int = 30 fps: int = 20 # Must match robosuite's default control_freq (20 Hz)
episode_length: int | None = None episode_length: int | None = None
obs_type: str = "pixels_agent_pos" obs_type: str = "pixels_agent_pos"
render_mode: str = "rgb_array" render_mode: str = "rgb_array"
@@ -354,6 +354,9 @@ class LiberoEnv(EnvConfig):
control_mode: str = "relative" # or "absolute" control_mode: str = "relative" # or "absolute"
def __post_init__(self): def __post_init__(self):
if self.fps <= 0:
raise ValueError(f"fps must be positive, got {self.fps}")
if self.obs_type == "pixels": if self.obs_type == "pixels":
self.features[LIBERO_KEY_PIXELS_AGENTVIEW] = PolicyFeature( self.features[LIBERO_KEY_PIXELS_AGENTVIEW] = PolicyFeature(
type=FeatureType.VISUAL, shape=(self.observation_height, self.observation_width, 3) type=FeatureType.VISUAL, shape=(self.observation_height, self.observation_width, 3)
@@ -412,6 +415,7 @@ class LiberoEnv(EnvConfig):
"render_mode": self.render_mode, "render_mode": self.render_mode,
"observation_height": self.observation_height, "observation_height": self.observation_height,
"observation_width": self.observation_width, "observation_width": self.observation_width,
"control_freq": self.fps,
} }
if self.task_ids is not None: if self.task_ids is not None:
kwargs["task_ids"] = self.task_ids kwargs["task_ids"] = self.task_ids
@@ -503,7 +507,7 @@ class MetaworldEnv(EnvConfig):
class RoboCasaEnv(EnvConfig): class RoboCasaEnv(EnvConfig):
task: str = "CloseFridge" task: str = "CloseFridge"
fps: int = 20 fps: int = 20
episode_length: int = 1000 episode_length: int | None = None
obs_type: str = "pixels_agent_pos" obs_type: str = "pixels_agent_pos"
render_mode: str = "rgb_array" render_mode: str = "rgb_array"
camera_name: str = "robot0_agentview_left,robot0_eye_in_hand,robot0_agentview_right" camera_name: str = "robot0_agentview_left,robot0_eye_in_hand,robot0_agentview_right"
+11 -1
View File
@@ -30,7 +30,7 @@ from gymnasium import spaces
from libero.libero import benchmark, get_libero_path from libero.libero import benchmark, get_libero_path
from libero.libero.envs import OffScreenRenderEnv from libero.libero.envs import OffScreenRenderEnv
from lerobot.types import RobotObservation from lerobot.lerobot_types import RobotObservation
from .utils import _LazyAsyncVectorEnv, parse_camera_names from .utils import _LazyAsyncVectorEnv, parse_camera_names
@@ -125,10 +125,13 @@ class LiberoEnv(gym.Env):
n_envs: int = 1, n_envs: int = 1,
camera_name_mapping: dict[str, str] | None = None, camera_name_mapping: dict[str, str] | None = None,
num_steps_wait: int = 10, num_steps_wait: int = 10,
control_freq: int = 20,
control_mode: str = "relative", control_mode: str = "relative",
is_libero_plus: bool = False, is_libero_plus: bool = False,
): ):
super().__init__() super().__init__()
if control_freq <= 0:
raise ValueError(f"control_freq must be positive, got {control_freq}")
self.task_id = task_id self.task_id = task_id
self.is_libero_plus = is_libero_plus self.is_libero_plus = is_libero_plus
self.obs_type = obs_type self.obs_type = obs_type
@@ -154,6 +157,7 @@ class LiberoEnv(gym.Env):
} }
self.camera_name_mapping = camera_name_mapping self.camera_name_mapping = camera_name_mapping
self.num_steps_wait = num_steps_wait self.num_steps_wait = num_steps_wait
self.control_freq = control_freq
self.episode_index = episode_index self.episode_index = episode_index
self.episode_length = episode_length self.episode_length = episode_length
# Load once and keep # Load once and keep
@@ -260,6 +264,7 @@ class LiberoEnv(gym.Env):
bddl_file_name=self._task_bddl_file, bddl_file_name=self._task_bddl_file,
camera_heights=self.observation_height, camera_heights=self.observation_height,
camera_widths=self.observation_width, camera_widths=self.observation_width,
control_freq=self.control_freq,
) )
env.reset() env.reset()
self._env = env self._env = env
@@ -379,7 +384,12 @@ class LiberoEnv(gym.Env):
def close(self): def close(self):
if self._env is not None: if self._env is not None:
try:
self._env.close() 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( def _make_env_fns(
+4 -1
View File
@@ -25,7 +25,7 @@ import metaworld.policies as policies
import numpy as np import numpy as np
from gymnasium import spaces from gymnasium import spaces
from lerobot.types import RobotObservation from lerobot.lerobot_types import RobotObservation
from .utils import _LazyAsyncVectorEnv from .utils import _LazyAsyncVectorEnv
@@ -155,6 +155,7 @@ class MetaworldEnv(gym.Env):
env.model.cam_pos[2] = [0.75, 0.075, 0.7] env.model.cam_pos[2] = [0.75, 0.075, 0.7]
env.reset() env.reset()
env._freeze_rand_vec = False # otherwise no randomization env._freeze_rand_vec = False # otherwise no randomization
env.seeded_rand_vec = True # use seeded RNG so reset(seed=X) controls object positions
self._env = env self._env = env
def render(self) -> np.ndarray: def render(self) -> np.ndarray:
@@ -220,6 +221,8 @@ class MetaworldEnv(gym.Env):
self._ensure_env() self._ensure_env()
super().reset(seed=seed) super().reset(seed=seed)
if seed is not None:
self._env.seed(seed)
raw_obs, info = self._env.reset(seed=seed) raw_obs, info = self._env.reset(seed=seed)
observation = self._format_raw_obs(raw_obs) observation = self._format_raw_obs(raw_obs)
+15 -2
View File
@@ -25,7 +25,7 @@ import gymnasium as gym
import numpy as np import numpy as np
from gymnasium import spaces from gymnasium import spaces
from lerobot.types import RobotObservation from lerobot.lerobot_types import RobotObservation
from .utils import _LazyAsyncVectorEnv, parse_camera_names from .utils import _LazyAsyncVectorEnv, parse_camera_names
@@ -98,6 +98,19 @@ def _resolve_tasks(task: str) -> tuple[list[str], str | None]:
return names, None return names, None
def _get_task_horizon(task: str) -> int:
"""Return the rollout horizon registered by RoboCasa for a task."""
from robocasa.utils.dataset_registry_utils import get_task_horizon
try:
return int(get_task_horizon(task))
except ValueError as exc:
raise ValueError(
f"No RoboCasa horizon is registered for task '{task}'. "
"Set `--env.episode_length=<steps>` explicitly."
) from exc
def convert_action(flat_action: np.ndarray) -> dict[str, Any]: def convert_action(flat_action: np.ndarray) -> dict[str, Any]:
"""Split a flat (12,) action vector into a RoboCasa action dict. """Split a flat (12,) action vector into a RoboCasa action dict.
@@ -154,7 +167,7 @@ class RoboCasaEnv(gym.Env):
self.camera_name = parse_camera_names(camera_name) self.camera_name = parse_camera_names(camera_name)
self._max_episode_steps = episode_length if episode_length is not None else 1000 self._max_episode_steps = episode_length if episode_length is not None else _get_task_horizon(task)
# Deferred — created on first reset() inside the worker subprocess # Deferred — created on first reset() inside the worker subprocess
# to avoid inheriting stale GPU/EGL contexts across fork(). # to avoid inheriting stale GPU/EGL contexts across fork().
+4 -2
View File
@@ -28,7 +28,7 @@ import numpy as np
import torch import torch
from gymnasium import spaces from gymnasium import spaces
from lerobot.types import RobotObservation from lerobot.lerobot_types import RobotObservation
from lerobot.utils.import_utils import _scipy_available from lerobot.utils.import_utils import _scipy_available
from .utils import _LazyAsyncVectorEnv from .utils import _LazyAsyncVectorEnv
@@ -384,7 +384,9 @@ class RoboTwinEnv(gym.Env):
self._env: Any | None = None # deferred — created on first reset() inside worker self._env: Any | None = None # deferred — created on first reset() inside worker
self._step_count: int = 0 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 = { image_spaces = {
cam: spaces.Box( cam: spaces.Box(
+2 -2
View File
@@ -37,7 +37,7 @@ import numpy as np
from gymnasium import spaces from gymnasium import spaces
from scipy.spatial.transform import Rotation from scipy.spatial.transform import Rotation
from lerobot.types import RobotObservation from lerobot.lerobot_types import RobotObservation
from .utils import _LazyAsyncVectorEnv from .utils import _LazyAsyncVectorEnv
@@ -373,7 +373,7 @@ class VLABenchEnv(gym.Env):
if action.shape[0] != 7: if action.shape[0] != 7:
# Unknown layout — fall back to zero-pad so the sim doesn't crash. # 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] padded[: min(action.shape[0], ctrl_dim)] = action[:ctrl_dim]
return padded return padded
+2 -1
View File
@@ -18,6 +18,7 @@ from lerobot.utils.import_utils import require_package
# guard the optional dependency here so importing this package fails loudly if it's missing. # guard the optional dependency here so importing this package fails loudly if it's missing.
require_package("datasets", extra="dataset") require_package("datasets", extra="dataset")
from .annotate import submit_annotate_to_hf
from .hf import submit_to_hf from .hf import submit_to_hf
__all__ = ["submit_to_hf"] __all__ = ["submit_annotate_to_hf", "submit_to_hf"]
+176
View File
@@ -0,0 +1,176 @@
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Run ``lerobot-annotate`` on HF Jobs (HuggingFace GPUs).
Same shape as the training submitter in ``hf.py``, with one difference: the
annotation pipeline serves its own VLM, so the pod starts from the official
``vllm/vllm-openai`` image (which has no lerobot) instead of the prebuilt
``lerobot-gpu`` image, and installs lerobot on top before running.
Because there is no config repo to stage, the pod replays the user's own CLI
flags everything except the client-only ``--job.*`` and the host-local
``--root``, which is replaced by ``--repo_id`` so the pod pulls the dataset
from the Hub.
"""
from __future__ import annotations
import shlex
import sys
from dataclasses import is_dataclass
from typing import TYPE_CHECKING
from huggingface_hub import HfApi, get_token, run_job
from .dataset import ensure_dataset_available
# Package-internal reuse of the training submitter's job plumbing: following a
# submitted job and forwarding argv are identical for annotation runs.
from .hf import _pod_forwarded_args, follow_job, resolve_job_tags
if TYPE_CHECKING:
from lerobot.annotations.steerable_pipeline.config import AnnotationPipelineConfig
LEROBOT_GIT_URL = "https://github.com/huggingface/lerobot.git"
# Mirrors the pins in pyproject.toml. The vLLM image resolves dependencies on its
# own otherwise, and pulls av 18 / datasets 5 / draccus 0.11 — each of which breaks
# lerobot at import time. `--upgrade-strategy only-if-needed` keeps vLLM's own
# (torch, transformers, ...) pins intact.
_RUNTIME_REQUIREMENTS = (
"'datasets>=4.7.0,<5.0.0' 'pyarrow>=21.0.0,<30.0.0' 'av>=15.0.0,<16.0.0' 'draccus==0.10.0' "
"'pandas>=2.0.0,<3.0.0' jsonlines gymnasium torchcodec mergedeep pyyaml-include toml typing-inspect "
"openai"
)
# Flags the submitter resolves itself instead of forwarding verbatim: `--root`
# names a directory only this machine has, `--repo_id` is re-emitted from the
# config, and the config-file args name local files (rejected up front by
# `submit_annotate_to_hf`). `--job.*` is dropped separately, by prefix; bare
# `--job` is not, hence its entry here — it is the one arg that could smuggle a
# remote `target` onto the pod and have the job recursively submit itself.
_SUBMITTER_OWNED_ARGS = ("--root", "--repo_id", "--config_path", "--job")
def _local_config_file_args(cfg: AnnotationPipelineConfig) -> list[str]:
"""The CLI args that name a config file on the client's disk.
draccus exposes ``--config_path`` for the whole config plus a ``--<field>``
for every nested dataclass (``--vlm``, ``--plan``, ``--job``, ...). The pod has
none of those files, so a remote run has to reject them rather than silently
drop the settings they carry.
"""
return ["--config_path", *(f"--{name}" for name in vars(cfg) if is_dataclass(getattr(cfg, name)))]
def build_pod_setup(lerobot_ref: str) -> str:
"""Shell prelude that turns the vLLM image into a ``lerobot-annotate`` runtime."""
spec = f"lerobot @ git+{LEROBOT_GIT_URL}@{lerobot_ref}"
return (
# git to install from the repo, ffmpeg to decode the dataset's videos.
"apt-get update -qq && apt-get install -y -qq git ffmpeg && "
f"pip install --no-deps {shlex.quote(spec)} && "
f"pip install --upgrade-strategy only-if-needed {_RUNTIME_REQUIREMENTS} && "
# vLLM's cudagraph memory estimate over-reserves and starves the KV cache;
# PyAV is the video backend the server can decode our frames with.
"export VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0 && "
"export VLLM_VIDEO_BACKEND=pyav"
)
def build_pod_command(repo_id: str, lerobot_ref: str, argv: list[str]) -> list[str]:
"""Build the ``bash -c`` command the pod runs: setup prelude, then annotation.
``argv`` is the user's CLI (``sys.argv[1:]``) minus the flags in
``_SUBMITTER_OWNED_ARGS``; ``--repo_id`` is re-added from the config so the pod
always annotates the dataset we just made sure is reachable on the Hub.
``--job.target=local`` stops the pod from re-dispatching to itself.
"""
forwarded = _pod_forwarded_args(argv, drop_names=_SUBMITTER_OWNED_ARGS, drop_prefixes=("--job.",))
annotate = shlex.join(["lerobot-annotate", f"--repo_id={repo_id}", *forwarded, "--job.target=local"])
return ["bash", "-c", f"{build_pod_setup(lerobot_ref)} && {annotate}"]
def submit_annotate_to_hf(cfg: AnnotationPipelineConfig) -> None:
"""Submit an annotation run to HF Jobs infrastructure.
Resolves credentials, makes sure the source dataset is reachable from the pod,
submits the job, then tails its logs until the job reaches a terminal stage
or returns immediately with ``--job.detach``. Ctrl-C detaches without
cancelling the remote job.
"""
token = get_token()
if not token:
raise RuntimeError("Not logged in to Hugging Face. Run `hf auth login` first.")
if cfg.repo_id is None:
raise ValueError(
"Remote annotation requires --repo_id: the pod downloads the dataset from the Hub, "
"and --root only names a directory on this machine."
)
argv = sys.argv[1:]
passed = {tok.split("=", 1)[0] for tok in argv}
used_config_files = sorted(passed.intersection(_local_config_file_args(cfg)))
if used_config_files:
raise ValueError(
f"{', '.join(used_config_files)} cannot be used with a remote --job.target: the pod "
"cannot read config files from this machine. Pass the settings as CLI flags instead."
)
if not cfg.push_to_hub:
# The pod's filesystem is discarded when the job ends, so without a push the
# run produces nothing. Warn rather than fail: a smoke test over
# --only_episodes that only inspects the logs is a legitimate use.
print(
"WARNING: --push_to_hub is off. The annotated dataset lives only on the pod and is "
"discarded when the job ends. Pass --push_to_hub=true to keep the result."
)
api = HfApi(token=token)
tags = resolve_job_tags(cfg.job.tags)
ensure_dataset_available(cfg.repo_id, api=api, tags=tags)
command = build_pod_command(cfg.repo_id, cfg.job.lerobot_ref, argv)
print(f"Submitting job to HF Jobs (flavor={cfg.job.target}, image={cfg.job.image}) ...")
job_info = run_job(
image=cfg.job.image,
command=command,
flavor=cfg.job.target,
secrets={"HF_TOKEN": token},
timeout=cfg.job.timeout,
# HF Jobs labels are key/value; expose each tag as a queryable label.
labels=dict.fromkeys(tags, "true"),
)
job_id = job_info.id
job_url = getattr(job_info, "url", None)
print(f"Job submitted: {job_id}")
if job_url:
print(f" Job page: {job_url}")
target_repo_id = cfg.new_repo_id or cfg.repo_id
if cfg.push_to_hub:
print(f" Dataset repo: https://huggingface.co/datasets/{target_repo_id}")
print(f" Monitor: hf jobs logs {job_id}")
print(f" Cancel: hf jobs cancel {job_id}")
# No success marker: `lerobot-annotate` keeps working after the upload log line
# (dataset card, version tag), so completion has to be stage-based.
if not follow_job(job_id, detach=cfg.job.detach):
return
if cfg.push_to_hub:
print(f"\nAnnotation complete — dataset pushed to https://huggingface.co/datasets/{target_repo_id}")
else:
print("\nAnnotation complete. Note: --push_to_hub was off, so the result stayed on the pod.")
+69 -54
View File
@@ -223,6 +223,74 @@ def _poll_until_done(
return None return None
def follow_job(job_id: str, *, detach: bool = False, success_marker: str | None = None) -> bool:
"""Watch a submitted job to the end, streaming its logs to stdout.
Returns True when the job finished successfully and False when we stopped watching
without a verdict `detach`, or the user pressing Ctrl-C, which detaches rather than
cancelling the remote job. Raises RuntimeError when the job reaches a terminal stage
other than COMPLETED.
`success_marker` finishes as soon as that string appears in the logs instead of waiting
out the platform's post-run finalization (~30s). Callers that have a log line meaning
"the artifact is on the Hub" should pass it; without one, completion is stage-based.
"""
if detach:
return False
done = threading.Event()
detached = threading.Event()
marker_seen = threading.Event()
stage_holder: dict[str, str | None] = {}
def _poll() -> None:
stage_holder["stage"] = _poll_until_done(job_id, done, status_holder=stage_holder)
poll_thread = threading.Thread(target=_poll, daemon=True)
poll_thread.start()
log_thread = threading.Thread(
target=_tail_logs, args=(job_id, done, success_marker, marker_seen), daemon=True
)
log_thread.start()
def _detach(sig, frame):
detached.set()
done.set()
print("\nDetached. Job is still running.")
print(f" Monitor: hf jobs logs {job_id}")
print(f" Cancel: hf jobs cancel {job_id}")
# signal.signal only works on the main thread; when called from a worker thread
# (e.g. an orchestration framework) skip the Ctrl-C-detaches-instead-of-cancels
# handler rather than crashing with ValueError.
install_sigint = threading.current_thread() is threading.main_thread()
original_sigint = signal.getsignal(signal.SIGINT) if install_sigint else None
if install_sigint:
signal.signal(signal.SIGINT, _detach)
try:
# Timeout-based join so SIGINT is delivered to the main thread promptly.
while poll_thread.is_alive():
poll_thread.join(timeout=0.5)
log_thread.join(timeout=5)
finally:
if install_sigint:
signal.signal(signal.SIGINT, original_sigint)
if detached.is_set():
return False
if marker_seen.is_set():
return True
stage = stage_holder.get("stage")
if stage != "COMPLETED":
message = stage_holder.get("message")
detail = f" ({message})" if message else ""
raise RuntimeError(
f"Job {job_id} ended with stage={stage}{detail}. Check logs: hf jobs logs {job_id}"
)
return True
def _pod_forwarded_args( def _pod_forwarded_args(
argv: list[str], drop_names: tuple[str, ...] = (), drop_prefixes: tuple[str, ...] = () argv: list[str], drop_names: tuple[str, ...] = (), drop_prefixes: tuple[str, ...] = ()
) -> list[str]: ) -> list[str]:
@@ -362,64 +430,11 @@ def submit_to_hf(cfg: TrainPipelineConfig) -> None:
print(f" Monitor: hf jobs logs {job_id}") print(f" Monitor: hf jobs logs {job_id}")
print(f" Cancel: hf jobs cancel {job_id}") print(f" Cancel: hf jobs cancel {job_id}")
if cfg.job.detach:
return
done = threading.Event()
detached = threading.Event()
pushed_ok = threading.Event()
stage_holder: dict[str, str | None] = {}
def _poll() -> None:
stage_holder["stage"] = _poll_until_done(job_id, done, status_holder=stage_holder)
poll_thread = threading.Thread(target=_poll, daemon=True)
poll_thread.start()
# Finish as soon as the model is pushed, rather than waiting out the platform's # Finish as soon as the model is pushed, rather than waiting out the platform's
# post-run finalization before the job stage flips to COMPLETED. This matches the # post-run finalization before the job stage flips to COMPLETED. This matches the
# exact log line emitted by PreTrainedPolicy.push_model_to_hub — the two must stay # exact log line emitted by PreTrainedPolicy.push_model_to_hub — the two must stay
# in sync. If it ever stops matching we just fall back to stage-based completion # in sync. If it ever stops matching we just fall back to stage-based completion
# (~30s slower), so the contract is an optimization, not a correctness requirement. # (~30s slower), so the contract is an optimization, not a correctness requirement.
success_marker = f"Model pushed to https://huggingface.co/{repo_id}" success_marker = f"Model pushed to https://huggingface.co/{repo_id}"
log_thread = threading.Thread( if follow_job(job_id, detach=cfg.job.detach, success_marker=success_marker):
target=_tail_logs, args=(job_id, done, success_marker, pushed_ok), daemon=True
)
log_thread.start()
def _detach(sig, frame):
detached.set()
done.set()
print("\nDetached. Job is still running.")
print(f" Monitor: hf jobs logs {job_id}")
print(f" Cancel: hf jobs cancel {job_id}")
# signal.signal only works on the main thread; when called from a worker thread
# (e.g. an orchestration framework) skip the Ctrl-C-detaches-instead-of-cancels
# handler rather than crashing with ValueError.
install_sigint = threading.current_thread() is threading.main_thread()
original_sigint = signal.getsignal(signal.SIGINT) if install_sigint else None
if install_sigint:
signal.signal(signal.SIGINT, _detach)
try:
# Timeout-based join so SIGINT is delivered to the main thread promptly.
while poll_thread.is_alive():
poll_thread.join(timeout=0.5)
log_thread.join(timeout=5)
finally:
if install_sigint:
signal.signal(signal.SIGINT, original_sigint)
if detached.is_set():
return
if pushed_ok.is_set():
print(f"\nTraining complete — model pushed to https://huggingface.co/{repo_id}") print(f"\nTraining complete — model pushed to https://huggingface.co/{repo_id}")
return
stage = stage_holder.get("stage")
if stage != "COMPLETED":
message = stage_holder.get("message")
detail = f" ({message})" if message else ""
raise RuntimeError(
f"Job {job_id} ended with stage={stage}{detail}. Check logs: hf jobs logs {job_id}"
)
+1 -2
View File
@@ -20,7 +20,6 @@ import logging
import time import time
from contextlib import contextmanager from contextlib import contextmanager
from copy import deepcopy from copy import deepcopy
from functools import cached_property
from typing import TYPE_CHECKING, Any, TypedDict from typing import TYPE_CHECKING, Any, TypedDict
from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
@@ -854,7 +853,7 @@ class DamiaoMotorsBus(MotorsBusBase):
else: else:
raise ValueError(f"Motor {motor_obj} doesn't have a valid recv_id (None).") raise ValueError(f"Motor {motor_obj} doesn't have a valid recv_id (None).")
@cached_property @property
def is_calibrated(self) -> bool: def is_calibrated(self) -> bool:
"""Check if motors are calibrated.""" """Check if motors are calibrated."""
return bool(self.calibration) return bool(self.calibration)
+18
View File
@@ -122,6 +122,9 @@ MODEL_ENCODING_TABLE = {
"xm430-w350": X_SERIES_ENCODINGS_TABLE, "xm430-w350": X_SERIES_ENCODINGS_TABLE,
"xm540-w270": X_SERIES_ENCODINGS_TABLE, "xm540-w270": X_SERIES_ENCODINGS_TABLE,
"xc430-w150": 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} # {model: model_resolution}
@@ -134,6 +137,9 @@ MODEL_RESOLUTION = {
"xm430-w350": 4096, "xm430-w350": 4096,
"xm540-w270": 4096, "xm540-w270": 4096,
"xc430-w150": 4096, "xc430-w150": 4096,
"xh540-w150": 4096,
"xc330-t288": 4096,
"xc330-t181": 4096,
} }
# {model: model_number} # {model: model_number}
@@ -145,6 +151,9 @@ MODEL_NUMBER_TABLE = {
"xm430-w350": 1020, "xm430-w350": 1020,
"xm540-w270": 1120, "xm540-w270": 1120,
"xc430-w150": 1070, "xc430-w150": 1070,
"xh540-w150": 1110,
"xc330-t288": 1220,
"xc330-t181": 1210,
} }
# {model: available_operating_modes} # {model: available_operating_modes}
@@ -156,6 +165,9 @@ MODEL_OPERATING_MODES = {
"xm430-w350": [0, 1, 3, 4, 5, 16], "xm430-w350": [0, 1, 3, 4, 5, 16],
"xm540-w270": [0, 1, 3, 4, 5, 16], "xm540-w270": [0, 1, 3, 4, 5, 16],
"xc430-w150": [1, 3, 4, 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 = { MODEL_CONTROL_TABLE = {
@@ -166,6 +178,9 @@ MODEL_CONTROL_TABLE = {
"xm430-w350": X_SERIES_CONTROL_TABLE, "xm430-w350": X_SERIES_CONTROL_TABLE,
"xm540-w270": X_SERIES_CONTROL_TABLE, "xm540-w270": X_SERIES_CONTROL_TABLE,
"xc430-w150": 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 = { MODEL_BAUDRATE_TABLE = {
@@ -176,6 +191,9 @@ MODEL_BAUDRATE_TABLE = {
"xm430-w350": X_SERIES_BAUDRATE_TABLE, "xm430-w350": X_SERIES_BAUDRATE_TABLE,
"xm540-w270": X_SERIES_BAUDRATE_TABLE, "xm540-w270": X_SERIES_BAUDRATE_TABLE,
"xc430-w150": 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 = [ AVAILABLE_BAUDRATES = [
+7 -3
View File
@@ -23,6 +23,7 @@ from __future__ import annotations
import abc import abc
import logging import logging
import time
from collections.abc import Sequence from collections.abc import Sequence
from contextlib import contextmanager from contextlib import contextmanager
from dataclasses import dataclass from dataclasses import dataclass
@@ -818,13 +819,13 @@ class SerialMotorsBus(MotorsBusBase):
""" """
motor_names = self._get_motors_list(motors) motor_names = self._get_motors_list(motors)
start_positions = self.sync_read("Present_Position", motor_names, normalize=False) start_positions = self.sync_read("Present_Position", motor_names, normalize=False, num_retry=5)
mins = start_positions.copy() mins = start_positions.copy()
maxes = start_positions.copy() maxes = start_positions.copy()
user_pressed_enter = False user_pressed_enter = False
while not user_pressed_enter: while not user_pressed_enter:
positions = self.sync_read("Present_Position", motor_names, normalize=False) positions = self.sync_read("Present_Position", motor_names, normalize=False, num_retry=5)
mins = {motor: min(positions[motor], min_) for motor, min_ in mins.items()} mins = {motor: min(positions[motor], min_) for motor, min_ in mins.items()}
maxes = {motor: max(positions[motor], max_) for motor, max_ in maxes.items()} maxes = {motor: max(positions[motor], max_) for motor, max_ in maxes.items()}
@@ -837,9 +838,12 @@ class SerialMotorsBus(MotorsBusBase):
if enter_pressed(): if enter_pressed():
user_pressed_enter = True user_pressed_enter = True
if display_values and not user_pressed_enter: if not user_pressed_enter:
if display_values:
# Move cursor up to overwrite the previous output # Move cursor up to overwrite the previous output
move_cursor_up(len(motor_names) + 3) move_cursor_up(len(motor_names) + 3)
# Throttle reads even when the live table is disabled.
time.sleep(0.02)
same_min_max = [motor for motor in motor_names if mins[motor] == maxes[motor]] same_min_max = [motor for motor in motor_names if mins[motor] == maxes[motor]]
if same_min_max: if same_min_max:
+2
View File
@@ -32,6 +32,7 @@ from .pretrained import PreTrainedPolicy as PreTrainedPolicy
from .smolvla.configuration_smolvla import SmolVLAConfig as SmolVLAConfig from .smolvla.configuration_smolvla import SmolVLAConfig as SmolVLAConfig
from .tdmpc.configuration_tdmpc import TDMPCConfig as TDMPCConfig from .tdmpc.configuration_tdmpc import TDMPCConfig as TDMPCConfig
from .utils import make_robot_action, prepare_observation_for_inference from .utils import make_robot_action, prepare_observation_for_inference
from .vla_jepa.configuration_vla_jepa import VLAJEPAConfig as VLAJEPAConfig
from .vqbet.configuration_vqbet import VQBeTConfig as VQBeTConfig from .vqbet.configuration_vqbet import VQBeTConfig as VQBeTConfig
from .wall_x.configuration_wall_x import WallXConfig as WallXConfig from .wall_x.configuration_wall_x import WallXConfig as WallXConfig
from .xvla.configuration_xvla import XVLAConfig as XVLAConfig from .xvla.configuration_xvla import XVLAConfig as XVLAConfig
@@ -57,6 +58,7 @@ __all__ = [
"PI05Config", "PI05Config",
"SmolVLAConfig", "SmolVLAConfig",
"TDMPCConfig", "TDMPCConfig",
"VLAJEPAConfig",
"VQBeTConfig", "VQBeTConfig",
"WallXConfig", "WallXConfig",
"XVLAConfig", "XVLAConfig",
+2 -39
View File
@@ -18,17 +18,10 @@ from typing import Any
import torch import torch
from lerobot.processor import ( from lerobot.processor import (
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
NormalizerProcessorStep,
PolicyAction, PolicyAction,
PolicyProcessorPipeline, PolicyProcessorPipeline,
RenameObservationsProcessorStep, make_default_pre_post_processors,
UnnormalizerProcessorStep,
policy_action_to_transition,
transition_to_policy_action,
) )
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
from .configuration_act import ACTConfig from .configuration_act import ACTConfig
@@ -54,34 +47,4 @@ def make_act_pre_post_processors(
tuple[PolicyProcessorPipeline[dict[str, Any], dict[str, Any]], PolicyProcessorPipeline[PolicyAction, PolicyAction]]: A tuple containing the tuple[PolicyProcessorPipeline[dict[str, Any], dict[str, Any]], PolicyProcessorPipeline[PolicyAction, PolicyAction]]: A tuple containing the
pre-processor pipeline and the post-processor pipeline. pre-processor pipeline and the post-processor pipeline.
""" """
return make_default_pre_post_processors(config, dataset_stats, normalizer_device=config.device)
input_steps = [
RenameObservationsProcessorStep(rename_map={}),
AddBatchDimensionProcessorStep(),
DeviceProcessorStep(device=config.device),
NormalizerProcessorStep(
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=dataset_stats,
device=config.device,
),
]
output_steps = [
UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
DeviceProcessorStep(device="cpu"),
]
return (
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
steps=input_steps,
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
),
PolicyProcessorPipeline[PolicyAction, PolicyAction](
steps=output_steps,
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
),
)
@@ -0,0 +1,122 @@
#!/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.
"""Flow-matching sampling primitives shared across policies.
Canonical versions of the beta-distributed timestep sampler and the forward-Euler
denoising loop (with its real-time-chunking hook) that the openpi-derived policies
(pi0, pi05, smolvla, eo1) historically each carried a copy of. All functions are
stateless; adopting them does not affect checkpoints.
"""
from collections.abc import Callable
from typing import TYPE_CHECKING
import torch
from torch import Tensor
if TYPE_CHECKING:
from lerobot.policies.rtc.modeling_rtc import RTCProcessor
def sample_beta(alpha: float, beta: float, bsize: int, device) -> Tensor: # see openpi (exact copy)
# Beta sampling uses _sample_dirichlet which isn't implemented for MPS, so sample on CPU
alpha_t = torch.tensor(alpha, dtype=torch.float32)
beta_t = torch.tensor(beta, dtype=torch.float32)
dist = torch.distributions.Beta(alpha_t, beta_t)
return dist.sample((bsize,)).to(device)
def sample_noise(shape, device) -> Tensor:
"""Standard-normal float32 noise, the flow-matching x_1 sample."""
return torch.normal(
mean=0.0,
std=1.0,
size=shape,
dtype=torch.float32,
device=device,
)
def sample_time_beta(bsize: int, device, *, alpha: float, beta: float, scale: float, offset: float) -> Tensor:
"""Beta-distributed flow-matching timesteps: ``Beta(alpha, beta) * scale + offset`` (openpi convention)."""
time_beta = sample_beta(alpha, beta, bsize, device)
time = time_beta * scale + offset
return time.to(dtype=torch.float32, device=device)
def euler_integrate(
denoise_fn: Callable[[Tensor, Tensor], Tensor],
noise: Tensor,
num_steps: int,
*,
rtc_processor: "RTCProcessor | None" = None,
rtc_enabled: bool = False,
inference_delay: int | None = None,
prev_chunk_left_over: Tensor | None = None,
execution_horizon: int | None = None,
) -> Tensor:
"""Forward-Euler integration of a velocity field from t=1 (noise) to t=0 (actions).
This is the openpi sampling loop: ``dt = -1/num_steps``, ``time = 1.0 + step*dt``,
``x_t <- x_t + dt * v_t``, with the optional real-time-chunking (RTC) guidance hook
wrapping the velocity computation and debug tracking after each step.
Args:
denoise_fn: Computes the velocity ``v_t`` from ``(x_t, time_tensor)`` where
``time_tensor`` is a float32 tensor of shape ``(batch_size,)``. The returned
velocity must have the same shape and dtype as ``x_t``.
noise: Initial sample ``x_1`` of shape ``(batch_size, ...)``.
num_steps: Number of Euler steps.
rtc_processor: Optional RTC processor. Debug tracking fires whenever it is set and
has debugging enabled, even if RTC guidance itself is disabled (this mirrors
the historical per-policy loops).
rtc_enabled: Whether to route the velocity computation through
``rtc_processor.denoise_step`` (requires ``rtc_processor``).
inference_delay: RTC guidance parameter, forwarded verbatim.
prev_chunk_left_over: RTC guidance parameter, forwarded verbatim.
execution_horizon: RTC guidance parameter, forwarded verbatim.
"""
bsize = noise.shape[0]
device = noise.device
dt = -1.0 / num_steps
x_t = noise
for step in range(num_steps):
time = 1.0 + step * dt
time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize)
def denoise_step_partial_call(input_x_t, current_timestep=time_tensor):
return denoise_fn(input_x_t, current_timestep)
if rtc_enabled:
v_t = rtc_processor.denoise_step(
x_t=x_t,
prev_chunk_left_over=prev_chunk_left_over,
inference_delay=inference_delay,
time=time,
original_denoise_step_partial=denoise_step_partial_call,
execution_horizon=execution_horizon,
)
else:
v_t = denoise_step_partial_call(x_t)
x_t = x_t + dt * v_t
if rtc_processor is not None and rtc_processor.is_debug_enabled():
rtc_processor.track(time=time, x_t=x_t, v_t=v_t)
return x_t
+243
View File
@@ -0,0 +1,243 @@
#!/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.
"""Helpers shared by the openpi-derived VLA policies (pi0, pi05, pi0_fast, smolvla, eo1, xvla).
These are the canonical versions of functions that historically were copy-pasted per
policy. They are pure (no parameters, no module state), so importing them from here
instead of a policy-local copy has no effect on checkpoints.
"""
import math
from typing import TYPE_CHECKING
import torch
import torch.nn.functional as F # noqa: N812
from torch import Tensor
from lerobot.utils.constants import OPENPI_ATTENTION_MASK_VALUE
from lerobot.utils.device_utils import get_safe_dtype
from lerobot.utils.import_utils import _transformers_available, require_package
if TYPE_CHECKING or _transformers_available:
from transformers import DynamicCache
else:
DynamicCache = None
def create_sinusoidal_pos_embedding( # see openpi `create_sinusoidal_pos_embedding` (exact copy)
time: torch.Tensor, dimension: int, min_period: float, max_period: float, device="cpu"
) -> Tensor:
"""Computes sine-cosine positional embedding vectors for scalar positions."""
if dimension % 2 != 0:
raise ValueError(f"dimension ({dimension}) must be divisible by 2")
if time.ndim != 1:
raise ValueError("The time tensor is expected to be of shape `(batch_size, )`.")
dtype = get_safe_dtype(torch.float64, device.type)
fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device)
period = min_period * (max_period / min_period) ** fraction
# Compute the outer product
scaling_factor = 1.0 / period * 2 * math.pi
sin_input = scaling_factor[None, :] * time[:, None]
return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)
def make_att_2d_masks(pad_masks: Tensor, att_masks: Tensor) -> Tensor: # see openpi (exact copy)
"""Copied from big_vision.
Tokens can attend to valid inputs tokens which have a cumulative mask_ar
smaller or equal to theirs. This way `mask_ar` int[B, N] can be used to
setup several types of attention, for example:
[[1 1 1 1 1 1]]: pure causal attention.
[[0 0 0 1 1 1]]: prefix-lm attention. The first 3 tokens can attend between
themselves and the last 3 tokens have a causal attention. The first
entry could also be a 1 without changing behaviour.
[[1 0 1 0 1 0 0 1 0 0]]: causal attention between 4 blocks. Tokens of a
block can attend all previous blocks and all tokens on the same block.
Args:
input_mask: bool[B, N] true if its part of the input, false if padding.
mask_ar: int32[B, N] mask that's 1 where previous tokens cannot depend on
it and 0 where it shares the same attention mask as the previous token.
"""
if att_masks.ndim != 2:
raise ValueError(att_masks.ndim)
if pad_masks.ndim != 2:
raise ValueError(pad_masks.ndim)
cumsum = torch.cumsum(att_masks, dim=1)
att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None]
pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None]
return att_2d_masks & pad_2d_masks
def prepare_attention_masks_4d(att_2d_masks: Tensor, dtype: torch.dtype | None = None) -> Tensor:
"""Expand boolean 2D attention masks to the additive 4D layout expected by transformers.
Valid positions become 0.0 and masked positions the large negative openpi constant.
"""
att_2d_masks_4d = att_2d_masks[:, None, :, :]
result = torch.where(att_2d_masks_4d, 0.0, OPENPI_ATTENTION_MASK_VALUE)
if dtype is not None:
result = result.to(dtype=dtype)
return result
def clone_past_key_values(past_key_values):
"""Clone the DynamicCache returned by prefix prefill for compiled denoising."""
if DynamicCache is None:
require_package("transformers", extra="transformers-dep")
return DynamicCache(
tuple(
(keys.clone(), values.clone(), sliding_window) for keys, values, sliding_window in past_key_values
)
)
def pad_vector(vector: Tensor, new_dim: int, *, truncate: bool = False) -> Tensor:
"""Pad the last dimension of a vector to new_dim with zeros.
Can be (batch_size x sequence_length x features_dimension)
or (batch_size x features_dimension)
With ``truncate=False`` (openpi behavior), vectors whose last dimension is already
>= new_dim are returned unchanged. With ``truncate=True`` (xVLA behavior), the last
dimension is truncated to exactly ``new_dim`` (which may be 0).
"""
if vector.shape[-1] == new_dim:
return vector
if not truncate:
if vector.shape[-1] >= new_dim:
return vector
return F.pad(vector, (0, new_dim - vector.shape[-1]))
shape = list(vector.shape)
current_dim = shape[-1]
shape[-1] = new_dim
new_vector = vector.new_zeros(*shape)
length = min(current_dim, new_dim)
new_vector[..., :length] = vector[..., :length]
return new_vector
def resize_with_pad_torch( # see openpi `resize_with_pad_torch` (exact copy)
images: torch.Tensor,
height: int,
width: int,
mode: str = "bilinear",
) -> torch.Tensor:
"""PyTorch version of resize_with_pad. Resizes an image to a target height and width without distortion
by padding with black. If the image is float32, it must be in the range [-1, 1].
Padding is centered (openpi convention). For the top-left-padding variant used by
smolvla/xvla, see :func:`resize_with_pad`.
Args:
images: Tensor of shape [*b, h, w, c] or [*b, c, h, w]
height: Target height
width: Target width
mode: Interpolation mode ('bilinear', 'nearest', etc.)
Returns:
Resized and padded tensor with same shape format as input
"""
# Check if input is in channels-last format [*b, h, w, c] or channels-first [*b, c, h, w]
if images.shape[-1] <= 4: # Assume channels-last format
channels_last = True
if images.dim() == 3:
images = images.unsqueeze(0) # Add batch dimension
images = images.permute(0, 3, 1, 2) # [b, h, w, c] -> [b, c, h, w]
else:
channels_last = False
if images.dim() == 3:
images = images.unsqueeze(0) # Add batch dimension
batch_size, channels, cur_height, cur_width = images.shape
# Calculate resize ratio
ratio = max(cur_width / width, cur_height / height)
resized_height = int(cur_height / ratio)
resized_width = int(cur_width / ratio)
# Resize
resized_images = F.interpolate(
images,
size=(resized_height, resized_width),
mode=mode,
align_corners=False if mode == "bilinear" else None,
)
# Handle dtype-specific clipping
if images.dtype == torch.uint8:
resized_images = torch.round(resized_images).clamp(0, 255).to(torch.uint8)
elif images.dtype == torch.float32:
resized_images = resized_images.clamp(0.0, 1.0)
else:
raise ValueError(f"Unsupported image dtype: {images.dtype}")
# Calculate padding
pad_h0, remainder_h = divmod(height - resized_height, 2)
pad_h1 = pad_h0 + remainder_h
pad_w0, remainder_w = divmod(width - resized_width, 2)
pad_w1 = pad_w0 + remainder_w
# Pad
constant_value = 0 if images.dtype == torch.uint8 else 0.0
padded_images = F.pad(
resized_images,
(pad_w0, pad_w1, pad_h0, pad_h1), # left, right, top, bottom
mode="constant",
value=constant_value,
)
# Convert back to original format if needed
if channels_last:
padded_images = padded_images.permute(0, 2, 3, 1) # [b, c, h, w] -> [b, h, w, c]
return padded_images
def resize_with_pad(img: torch.Tensor, height: int, width: int, *, pad_value: float) -> torch.Tensor:
"""Resize a (b, c, h, w) image without distortion, padding on the LEFT and TOP.
This is the smolvla/xvla convention. For the centered-padding openpi variant, see
:func:`resize_with_pad_torch`. ``pad_value`` is keyword-only on purpose: callers
historically used different values (0, -1) and must state their choice explicitly.
"""
if img.ndim != 4:
raise ValueError(f"(b,c,h,w) expected, but got {img.shape}")
current_height, current_width = img.shape[2:]
if current_height == height and current_width == width:
return img
ratio = max(current_width / width, current_height / height)
resized_height = int(current_height / ratio)
resized_width = int(current_width / ratio)
resized_img = F.interpolate(
img, size=(resized_height, resized_width), mode="bilinear", align_corners=False
)
pad_height = max(0, height - resized_height)
pad_width = max(0, width - resized_width)
padded_img = F.pad(resized_img, (pad_width, 0, pad_height, 0), value=pad_value)
return padded_img
@@ -79,6 +79,8 @@ class DiffusionConfig(PreTrainedConfig):
use_film_scale_modulation: FiLM (https://huggingface.co/papers/1709.07871) is used for the Unet conditioning. use_film_scale_modulation: FiLM (https://huggingface.co/papers/1709.07871) is used for the Unet conditioning.
Bias modulation is used be default, while this parameter indicates whether to also use scale Bias modulation is used be default, while this parameter indicates whether to also use scale
modulation. modulation.
gradient_checkpointing: Whether to checkpoint the Unet residual blocks during training. This reduces
activation memory at the cost of recomputing those blocks during the backward pass.
noise_scheduler_type: Name of the noise scheduler to use. Supported options: ["DDPM", "DDIM"]. noise_scheduler_type: Name of the noise scheduler to use. Supported options: ["DDPM", "DDIM"].
num_train_timesteps: Number of diffusion steps for the forward diffusion schedule. num_train_timesteps: Number of diffusion steps for the forward diffusion schedule.
beta_schedule: Name of the diffusion beta schedule as per DDPMScheduler from Hugging Face diffusers. beta_schedule: Name of the diffusion beta schedule as per DDPMScheduler from Hugging Face diffusers.
@@ -132,6 +134,7 @@ class DiffusionConfig(PreTrainedConfig):
n_groups: int = 8 n_groups: int = 8
diffusion_step_embed_dim: int = 128 diffusion_step_embed_dim: int = 128
use_film_scale_modulation: bool = True use_film_scale_modulation: bool = True
gradient_checkpointing: bool = False
# Noise scheduler. # Noise scheduler.
noise_scheduler_type: str = "DDPM" noise_scheduler_type: str = "DDPM"
num_train_timesteps: int = 100 num_train_timesteps: int = 100
@@ -31,6 +31,7 @@ import torch
import torch.nn.functional as F # noqa: N812 import torch.nn.functional as F # noqa: N812
import torchvision import torchvision
from torch import Tensor, nn from torch import Tensor, nn
from torch.utils.checkpoint import checkpoint
from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGES, OBS_STATE from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGES, OBS_STATE
from lerobot.utils.import_utils import _diffusers_available, require_package from lerobot.utils.import_utils import _diffusers_available, require_package
@@ -727,20 +728,33 @@ class DiffusionConditionalUnet1d(nn.Module):
else: else:
global_feature = timesteps_embed global_feature = timesteps_embed
use_gc = self.config.gradient_checkpointing and self.training
# Run encoder, keeping track of skip features to pass to the decoder. # Run encoder, keeping track of skip features to pass to the decoder.
encoder_skip_features: list[Tensor] = [] encoder_skip_features: list[Tensor] = []
for resnet, resnet2, downsample in self.down_modules: for resnet, resnet2, downsample in self.down_modules:
if use_gc:
x = checkpoint(resnet, x, global_feature, use_reentrant=False)
x = checkpoint(resnet2, x, global_feature, use_reentrant=False)
else:
x = resnet(x, global_feature) x = resnet(x, global_feature)
x = resnet2(x, global_feature) x = resnet2(x, global_feature)
encoder_skip_features.append(x) encoder_skip_features.append(x)
x = downsample(x) x = downsample(x)
for mid_module in self.mid_modules: for mid_module in self.mid_modules:
if use_gc:
x = checkpoint(mid_module, x, global_feature, use_reentrant=False)
else:
x = mid_module(x, global_feature) x = mid_module(x, global_feature)
# Run decoder, using the skip features from the encoder. # Run decoder, using the skip features from the encoder.
for resnet, resnet2, upsample in self.up_modules: for resnet, resnet2, upsample in self.up_modules:
x = torch.cat((x, encoder_skip_features.pop()), dim=1) x = torch.cat((x, encoder_skip_features.pop()), dim=1)
if use_gc:
x = checkpoint(resnet, x, global_feature, use_reentrant=False)
x = checkpoint(resnet2, x, global_feature, use_reentrant=False)
else:
x = resnet(x, global_feature) x = resnet(x, global_feature)
x = resnet2(x, global_feature) x = resnet2(x, global_feature)
x = upsample(x) x = upsample(x)
@@ -19,17 +19,10 @@ from typing import Any
import torch import torch
from lerobot.processor import ( from lerobot.processor import (
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
NormalizerProcessorStep,
PolicyAction, PolicyAction,
PolicyProcessorPipeline, PolicyProcessorPipeline,
RenameObservationsProcessorStep, make_default_pre_post_processors,
UnnormalizerProcessorStep,
policy_action_to_transition,
transition_to_policy_action,
) )
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
from .configuration_diffusion import DiffusionConfig from .configuration_diffusion import DiffusionConfig
@@ -63,32 +56,4 @@ def make_diffusion_pre_post_processors(
Returns: Returns:
A tuple containing the configured pre-processor and post-processor pipelines. A tuple containing the configured pre-processor and post-processor pipelines.
""" """
return make_default_pre_post_processors(config, dataset_stats)
input_steps = [
RenameObservationsProcessorStep(rename_map={}),
AddBatchDimensionProcessorStep(),
DeviceProcessorStep(device=config.device),
NormalizerProcessorStep(
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=dataset_stats,
),
]
output_steps = [
UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
DeviceProcessorStep(device="cpu"),
]
return (
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
steps=input_steps,
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
),
PolicyProcessorPipeline[PolicyAction, PolicyAction](
steps=output_steps,
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
),
)
+14 -76
View File
@@ -18,7 +18,6 @@ from __future__ import annotations
import contextlib import contextlib
import logging import logging
import math
from collections import deque from collections import deque
from typing import TYPE_CHECKING, Any from typing import TYPE_CHECKING, Any
@@ -31,6 +30,8 @@ from torch import Tensor
from lerobot.utils.constants import ACTION, OBS_STATE from lerobot.utils.constants import ACTION, OBS_STATE
from lerobot.utils.import_utils import _transformers_available, require_package from lerobot.utils.import_utils import _transformers_available, require_package
from ..common.flow_matching import euler_integrate, sample_noise, sample_time_beta
from ..common.vla_utils import create_sinusoidal_pos_embedding, pad_vector
from ..pretrained import PreTrainedPolicy from ..pretrained import PreTrainedPolicy
from .configuration_eo1 import EO1Config from .configuration_eo1 import EO1Config
@@ -46,17 +47,6 @@ else:
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
def pad_vector(vector, new_dim):
"""Pad the last dimension of a vector to new_dim with zeros.
Can be (batch_size x sequence_length x features_dimension)
or (batch_size x features_dimension)
"""
if vector.shape[-1] >= new_dim:
return vector
return F.pad(vector, (0, new_dim - vector.shape[-1]))
class EO1Policy(PreTrainedPolicy): class EO1Policy(PreTrainedPolicy):
"""EO1 policy wrapper for LeRobot robot-only training/evaluation.""" """EO1 policy wrapper for LeRobot robot-only training/evaluation."""
@@ -136,47 +126,6 @@ class EO1Policy(PreTrainedPolicy):
return self.parameters() return self.parameters()
def get_safe_dtype(target_dtype, device_type):
"""Get a safe dtype for the given device type."""
if device_type == "mps" and target_dtype == torch.float64:
return torch.float32
if device_type == "cpu":
# CPU doesn't support bfloat16, use float32 instead
if target_dtype == torch.bfloat16:
return torch.float32
if target_dtype == torch.float64:
return torch.float64
return target_dtype
def create_sinusoidal_pos_embedding( # see openpi `create_sinusoidal_pos_embedding` (exact copy)
time: torch.Tensor, dimension: int, min_period: float, max_period: float, device="cpu"
) -> Tensor:
"""Computes sine-cosine positional embedding vectors for scalar positions."""
if dimension % 2 != 0:
raise ValueError(f"dimension ({dimension}) must be divisible by 2")
if time.ndim != 1:
raise ValueError("The time tensor is expected to be of shape `(batch_size, )`.")
dtype = get_safe_dtype(torch.float64, device.type)
fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device)
period = min_period * (max_period / min_period) ** fraction
# Compute the outer product
scaling_factor = 1.0 / period * 2 * math.pi
sin_input = scaling_factor[None, :] * time[:, None]
return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)
def sample_beta(alpha, beta, bsize, device): # see openpi `sample_beta` (exact copy)
# Beta sampling uses _sample_dirichlet which isn't implemented for MPS, so sample on CPU
alpha_t = torch.tensor(alpha, dtype=torch.float32)
beta_t = torch.tensor(beta, dtype=torch.float32)
dist = torch.distributions.Beta(alpha_t, beta_t)
return dist.sample((bsize,)).to(device)
class EO1VisionActionProjector(torch.nn.Sequential): class EO1VisionActionProjector(torch.nn.Sequential):
"""This block implements the multi-layer perceptron (MLP) module.""" """This block implements the multi-layer perceptron (MLP) module."""
@@ -267,21 +216,17 @@ class EO1VisionFlowMatchingModel(nn.Module):
return func(*args, **kwargs) return func(*args, **kwargs)
def sample_noise(self, shape, device): def sample_noise(self, shape, device):
noise = torch.normal( return sample_noise(shape, device)
mean=0.0,
std=1.0,
size=shape,
dtype=torch.float32,
device=device,
)
return noise
def sample_time(self, bsize, device): def sample_time(self, bsize, device):
time_beta = sample_beta( return sample_time_beta(
self.config.time_sampling_beta_alpha, self.config.time_sampling_beta_beta, bsize, device bsize,
device,
alpha=self.config.time_sampling_beta_alpha,
beta=self.config.time_sampling_beta_beta,
scale=self.config.time_sampling_scale,
offset=self.config.time_sampling_offset,
) )
time = time_beta * self.config.time_sampling_scale + self.config.time_sampling_offset
return time.to(dtype=torch.float32, device=device)
def get_placeholder_mask( def get_placeholder_mask(
self, self,
@@ -587,18 +532,11 @@ class EO1VisionFlowMatchingModel(nn.Module):
(batch_size, chunk_size, self.config.max_action_dim), (batch_size, chunk_size, self.config.max_action_dim),
device, device,
).to(dtype=self.action_in_proj.weight.dtype) ).to(dtype=self.action_in_proj.weight.dtype)
dt = -1.0 / self.config.num_denoise_steps
past_key_values = outputs.past_key_values past_key_values = outputs.past_key_values
# 3. Denoise only the action chunk while keeping the prefix cache invariant. # 3. Denoise only the action chunk while keeping the prefix cache invariant.
for step in range(self.config.num_denoise_steps): def denoise_fn(input_x_t, current_timestep):
time = torch.full( action_time_embs = self.embed_suffix(current_timestep, input_x_t)
(batch_size,),
1.0 + step * dt,
device=device,
dtype=torch.float32,
)
action_time_embs = self.embed_suffix(time, x_t)
inputs_embeds[:, act_slice] = action_time_embs.to(inputs_embeds.dtype) inputs_embeds[:, act_slice] = action_time_embs.to(inputs_embeds.dtype)
# Keep the prefix KV cache invariant across denoising steps. # Keep the prefix KV cache invariant across denoising steps.
@@ -615,7 +553,7 @@ class EO1VisionFlowMatchingModel(nn.Module):
hidden_states = outputs.last_hidden_state[:, :chunk_size] hidden_states = outputs.last_hidden_state[:, :chunk_size]
hidden_states = hidden_states.to(dtype=self.action_out_proj.dtype) hidden_states = hidden_states.to(dtype=self.action_out_proj.dtype)
v_t = self.action_out_proj(hidden_states) v_t = self.action_out_proj(hidden_states)
return v_t.reshape(input_x_t.shape).to(input_x_t.dtype)
x_t += dt * v_t.reshape(x_t.shape) x_t = euler_integrate(denoise_fn, x_t, self.config.num_denoise_steps)
return x_t return x_t
+13 -38
View File
@@ -22,25 +22,17 @@ from typing import TYPE_CHECKING, Any
import torch import torch
from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
from lerobot.lerobot_types import TransitionKey
from lerobot.processor import ( from lerobot.processor import (
AddBatchDimensionProcessorStep,
ComplementaryDataProcessorStep, ComplementaryDataProcessorStep,
DeviceProcessorStep,
NormalizerProcessorStep,
PolicyAction, PolicyAction,
PolicyProcessorPipeline, PolicyProcessorPipeline,
ProcessorStep, ProcessorStep,
ProcessorStepRegistry, ProcessorStepRegistry,
RenameObservationsProcessorStep, make_default_policy_processor_steps,
UnnormalizerProcessorStep, make_policy_processor_pipelines,
)
from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
from lerobot.types import TransitionKey
from lerobot.utils.constants import (
OBS_STATE,
POLICY_POSTPROCESSOR_DEFAULT_NAME,
POLICY_PREPROCESSOR_DEFAULT_NAME,
) )
from lerobot.utils.constants import OBS_STATE
from lerobot.utils.import_utils import _transformers_available, require_package from lerobot.utils.import_utils import _transformers_available, require_package
from .configuration_eo1 import EO1Config from .configuration_eo1 import EO1Config
@@ -242,14 +234,12 @@ def make_eo1_pre_post_processors(
]: ]:
"""Build pre/post processor pipelines for EO1.""" """Build pre/post processor pipelines for EO1."""
steps = make_default_policy_processor_steps(config, dataset_stats)
input_steps: list[ProcessorStep] = [ input_steps: list[ProcessorStep] = [
RenameObservationsProcessorStep(rename_map={}), steps.rename_observations,
AddBatchDimensionProcessorStep(), steps.add_batch_dim,
NormalizerProcessorStep( steps.normalize,
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=dataset_stats,
),
EO1ConversationTemplateStep(input_features=config.input_features, chunk_size=config.chunk_size), EO1ConversationTemplateStep(input_features=config.input_features, chunk_size=config.chunk_size),
EO1QwenProcessorStep( EO1QwenProcessorStep(
processor_name=config.vlm_base, processor_name=config.vlm_base,
@@ -257,27 +247,12 @@ def make_eo1_pre_post_processors(
image_max_pixels=config.image_max_pixels, image_max_pixels=config.image_max_pixels,
use_fast_processor=config.use_fast_processor, use_fast_processor=config.use_fast_processor,
), ),
DeviceProcessorStep(device=config.device), steps.to_device,
] ]
output_steps: list[ProcessorStep] = [ output_steps: list[ProcessorStep] = [
UnnormalizerProcessorStep( steps.unnormalize,
features=config.output_features, steps.to_cpu,
norm_map=config.normalization_mapping,
stats=dataset_stats,
),
DeviceProcessorStep(device="cpu"),
] ]
return ( return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
steps=input_steps,
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
),
PolicyProcessorPipeline[PolicyAction, PolicyAction](
steps=output_steps,
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
),
)
@@ -27,9 +27,11 @@ from lerobot.utils.import_utils import _transformers_available, require_package
if TYPE_CHECKING or _transformers_available: if TYPE_CHECKING or _transformers_available:
from transformers import AutoModel, AutoTokenizer from transformers import AutoModel, AutoTokenizer
from transformers.utils import is_flash_attn_2_available
else: else:
AutoModel = None AutoModel = None
AutoTokenizer = None AutoTokenizer = None
is_flash_attn_2_available = None
IMAGENET_MEAN = (0.485, 0.456, 0.406) IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225) IMAGENET_STD = (0.229, 0.224, 0.225)
@@ -135,9 +137,13 @@ class InternVL3Embedder(nn.Module):
raise ValueError(f"Unsupported EVO1 vlm_dtype '{model_dtype}'") from exc raise ValueError(f"Unsupported EVO1 vlm_dtype '{model_dtype}'") from exc
self.model_dtype = model_dtype self.model_dtype = model_dtype
attn_implementation = "flash_attention_2" if (use_flash_attn and _flash_attn_available()) else "eager" attn_implementation = (
"flash_attention_2" if (use_flash_attn and is_flash_attn_2_available()) else "eager"
)
if use_flash_attn and attn_implementation == "eager": if use_flash_attn and attn_implementation == "eager":
logger.warning("flash_attn is not installed. Falling back to eager attention.") logger.warning(
"Flash Attention 2 is unavailable on this runtime. Falling back to eager attention."
)
self.model = AutoModel.from_pretrained( self.model = AutoModel.from_pretrained(
model_name, model_name,
@@ -359,11 +365,3 @@ class InternVL3Embedder(nn.Module):
@property @property
def device(self) -> torch.device: def device(self) -> torch.device:
return next(self.model.parameters()).device return next(self.model.parameters()).device
def _flash_attn_available() -> bool:
try:
import flash_attn # noqa: F401
except ModuleNotFoundError:
return False
return True
@@ -42,6 +42,9 @@ class Evo1Policy(PreTrainedPolicy):
config_class = Evo1Config config_class = Evo1Config
name = "evo1" name = "evo1"
def supports_rtc(self) -> bool:
return True
def __init__(self, config: Evo1Config, *, vlm_hub_kwargs: dict | None = None, **kwargs): def __init__(self, config: Evo1Config, *, vlm_hub_kwargs: dict | None = None, **kwargs):
super().__init__(config) super().__init__(config)
config.validate_features() config.validate_features()
+36 -6
View File
@@ -21,6 +21,7 @@ from typing import Any
import torch import torch
from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
from lerobot.lerobot_types import EnvTransition, TransitionKey
from lerobot.processor import ( from lerobot.processor import (
AddBatchDimensionProcessorStep, AddBatchDimensionProcessorStep,
DeviceProcessorStep, DeviceProcessorStep,
@@ -40,7 +41,6 @@ from lerobot.processor.converters import (
policy_action_to_transition, policy_action_to_transition,
transition_to_policy_action, transition_to_policy_action,
) )
from lerobot.types import EnvTransition, TransitionKey
from lerobot.utils.constants import ( from lerobot.utils.constants import (
ACTION, ACTION,
DONE, DONE,
@@ -302,6 +302,33 @@ def _pad_evo1_stats(
return padded_stats return padded_stats
def _refresh_evo1_normalization_steps(
config: Evo1Config,
preprocessor: PolicyProcessorPipeline,
postprocessor: PolicyProcessorPipeline,
) -> None:
"""Re-pad checkpoint-loaded (un)normalizer stats/features to EVO1's fixed widths.
Loading a checkpoint injects the raw dataset stats (unpadded to max_state_dim/max_action_dim)
into the (un)normalizer via the generic override path in make_pre_post_processors. Those stats
and their declared features must be re-padded/reshaped to EVO1's fixed widths, otherwise
normalization fails against the padded state/action tensors (e.g. state padded to 24 vs. 8-dim
LIBERO stats). Padding is a no-op when stats are already at the target width.
"""
normalization_features = _evo1_normalization_features(config)
action_features = _evo1_action_features(config)
for step in preprocessor.steps:
if isinstance(step, NormalizerProcessorStep):
step.features = normalization_features
step.stats = _pad_evo1_stats(config, step.stats)
step.to(device=step.device, dtype=step.dtype)
for step in postprocessor.steps:
if isinstance(step, UnnormalizerProcessorStep):
step.features = action_features
step.stats = _pad_evo1_stats(config, step.stats)
step.to(device=step.device, dtype=step.dtype)
def reconcile_evo1_processors( def reconcile_evo1_processors(
config: Evo1Config, config: Evo1Config,
preprocessor: PolicyProcessorPipeline, preprocessor: PolicyProcessorPipeline,
@@ -309,16 +336,19 @@ def reconcile_evo1_processors(
) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: ) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
"""Reconcile checkpoint-loaded pipelines with the current EVO1 config. """Reconcile checkpoint-loaded pipelines with the current EVO1 config.
Two things cannot be restored from a serialized pipeline alone: the EVO1 batch converter Three things cannot be restored from a serialized pipeline alone: the EVO1 batch converter
(converters are plain functions and are never serialized), and eval-time CLI overrides of the (converters are plain functions and are never serialized), eval-time CLI overrides of the
action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`). This action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`), and the
restores the converter and rebuilds the action step from the current config so those overrides (un)normalizer stats/features when the generic override path injects raw, unpadded dataset
take effect. stats. This restores the converter, re-pads the normalization stats to EVO1's fixed widths, and
rebuilds the action step from the current config so those overrides take effect.
""" """
# Pipelines reloaded from a checkpoint come back with the default batch converter, which drops # Pipelines reloaded from a checkpoint come back with the default batch converter, which drops
# non-observation extras (embodiment_id, state_mask, custom task fields) needed by EVO1. # non-observation extras (embodiment_id, state_mask, custom task fields) needed by EVO1.
preprocessor.to_transition = evo1_batch_to_transition preprocessor.to_transition = evo1_batch_to_transition
_refresh_evo1_normalization_steps(config, preprocessor, postprocessor)
action_step = Evo1ActionProcessorStep( action_step = Evo1ActionProcessorStep(
action_dim=_evo1_action_dim(config), action_dim=_evo1_action_dim(config),
binarize_gripper=config.binarize_gripper, binarize_gripper=config.binarize_gripper,
+75 -311
View File
@@ -17,6 +17,7 @@
from __future__ import annotations from __future__ import annotations
import importlib import importlib
import inspect
import logging import logging
from typing import TYPE_CHECKING, Any, TypedDict, Unpack from typing import TYPE_CHECKING, Any, TypedDict, Unpack
@@ -27,6 +28,7 @@ if TYPE_CHECKING:
from lerobot.configs import FeatureType, PreTrainedConfig from lerobot.configs import FeatureType, PreTrainedConfig
from lerobot.envs import EnvConfig, env_to_policy_features from lerobot.envs import EnvConfig, env_to_policy_features
from lerobot.lerobot_types import PolicyAction
from lerobot.processor import ( from lerobot.processor import (
AbsoluteActionsProcessorStep, AbsoluteActionsProcessorStep,
PolicyProcessorPipeline, PolicyProcessorPipeline,
@@ -36,34 +38,24 @@ from lerobot.processor import (
transition_to_batch, transition_to_batch,
transition_to_policy_action, transition_to_policy_action,
) )
from lerobot.types import PolicyAction
from lerobot.utils.constants import ( from lerobot.utils.constants import (
ACTION, ACTION,
POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_POSTPROCESSOR_DEFAULT_NAME,
POLICY_PREPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME,
) )
from lerobot.utils.feature_utils import dataset_to_policy_features from lerobot.utils.feature_utils import dataset_to_policy_features
from lerobot.utils.import_utils import _peft_available, require_package
from .act.configuration_act import ACTConfig
from .diffusion.configuration_diffusion import DiffusionConfig
from .eo1.configuration_eo1 import EO1Config
from .evo1.configuration_evo1 import Evo1Config from .evo1.configuration_evo1 import Evo1Config
from .fastwam.configuration_fastwam import FastWAMConfig
from .gaussian_actor.configuration_gaussian_actor import GaussianActorConfig
from .groot.configuration_groot import GrootConfig from .groot.configuration_groot import GrootConfig
from .lingbot_va.configuration_lingbot_va import LingBotVAConfig
from .molmoact2.configuration_molmoact2 import MolmoAct2Config
from .multi_task_dit.configuration_multi_task_dit import MultiTaskDiTConfig
from .pi0.configuration_pi0 import PI0Config
from .pi05.configuration_pi05 import PI05Config
from .pretrained import PreTrainedPolicy from .pretrained import PreTrainedPolicy
from .smolvla.configuration_smolvla import SmolVLAConfig
from .tdmpc.configuration_tdmpc import TDMPCConfig
from .utils import validate_visual_features_consistency from .utils import validate_visual_features_consistency
from .vla_jepa.configuration_vla_jepa import VLAJEPAConfig
from .vqbet.configuration_vqbet import VQBeTConfig if TYPE_CHECKING or _peft_available:
from .wall_x.configuration_wall_x import WallXConfig from peft import PeftConfig, PeftModel
from .xvla.configuration_xvla import XVLAConfig else:
PeftConfig = None
PeftModel = None
def _reconnect_relative_absolute_steps( def _reconnect_relative_absolute_steps(
@@ -88,100 +80,23 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:
""" """
Retrieves a policy class by its registered name. Retrieves a policy class by its registered name.
This function uses dynamic imports to avoid loading all policy classes into memory Resolution is convention-based: the draccus-registered config class of ``name`` is
at once, improving startup time and reducing dependencies. looked up, its ``configuration_*`` module path is rewritten to ``modeling_*``, and
the ``<X>Policy`` class is imported from there. The modeling module is only imported
at call time, keeping heavy optional dependencies lazy. This works for both built-in
policies and third-party lerobot plugins (anything registered via
``@PreTrainedConfig.register_subclass``).
Args: Args:
name: The name of the policy. Supported names are "tdmpc", "diffusion", "act", name: The registered name of the policy (e.g. "act", "diffusion", "pi0").
"multi_task_dit", "vqbet", "pi0", "pi05", "gaussian_actor", "smolvla", "wall_x",
"molmoact2", "eo1", "evo1".
Returns: Returns:
The policy class corresponding to the given name. The policy class corresponding to the given name.
Raises: Raises:
NotImplementedError: If the policy name is not recognized. ValueError: If the policy name is not registered.
ImportError: If the policy's optional dependencies are not installed.
""" """
if name == "tdmpc":
from .tdmpc.modeling_tdmpc import TDMPCPolicy
return TDMPCPolicy
elif name == "diffusion":
from .diffusion.modeling_diffusion import DiffusionPolicy
return DiffusionPolicy
elif name == "act":
from .act.modeling_act import ACTPolicy
return ACTPolicy
elif name == "multi_task_dit":
from .multi_task_dit.modeling_multi_task_dit import MultiTaskDiTPolicy
return MultiTaskDiTPolicy
elif name == "vqbet":
from .vqbet.modeling_vqbet import VQBeTPolicy
return VQBeTPolicy
elif name == "pi0":
from .pi0.modeling_pi0 import PI0Policy
return PI0Policy
elif name == "pi0_fast":
from .pi0_fast.modeling_pi0_fast import PI0FastPolicy
return PI0FastPolicy
elif name == "pi05":
from .pi05.modeling_pi05 import PI05Policy
return PI05Policy
elif name == "gaussian_actor":
from .gaussian_actor.modeling_gaussian_actor import GaussianActorPolicy
return GaussianActorPolicy
elif name == "smolvla":
from .smolvla.modeling_smolvla import SmolVLAPolicy
return SmolVLAPolicy
elif name == "groot":
from .groot.modeling_groot import GrootPolicy
return GrootPolicy
elif name == "xvla":
from .xvla.modeling_xvla import XVLAPolicy
return XVLAPolicy
elif name == "wall_x":
from .wall_x.modeling_wall_x import WallXPolicy
return WallXPolicy
elif name == "eo1":
from .eo1.modeling_eo1 import EO1Policy
return EO1Policy
elif name == "molmoact2":
from .molmoact2.modeling_molmoact2 import MolmoAct2Policy
return MolmoAct2Policy
elif name == "vla_jepa":
from .vla_jepa.modeling_vla_jepa import VLAJEPAPolicy
return VLAJEPAPolicy
elif name == "lingbot_va":
from .lingbot_va.modeling_lingbot_va import LingBotVAPolicy
return LingBotVAPolicy
elif name == "fastwam":
from .fastwam.modeling_fastwam import FastWAMPolicy
return FastWAMPolicy
elif name == "evo1":
from .evo1.modeling_evo1 import Evo1Policy
return Evo1Policy
else:
try:
return _get_policy_cls_from_policy_name(name=name) return _get_policy_cls_from_policy_name(name=name)
except Exception as e:
raise ValueError(f"Policy type '{name}' is not available.") from e
def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig: def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
@@ -192,9 +107,8 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
mapping a string identifier to the corresponding config class. mapping a string identifier to the corresponding config class.
Args: Args:
policy_type: The type of the policy. Supported types include "tdmpc", policy_type: The registered type of the policy (any name registered via
"multi_task_dit", "diffusion", "act", "vqbet", "pi0", "pi05", "gaussian_actor", ``@PreTrainedConfig.register_subclass``, e.g. "act", "diffusion", "pi0").
"smolvla", "wall_x", "molmoact2", "eo1", "evo1".
**kwargs: Keyword arguments to be passed to the configuration class constructor. **kwargs: Keyword arguments to be passed to the configuration class constructor.
Returns: Returns:
@@ -203,48 +117,11 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
Raises: Raises:
ValueError: If the `policy_type` is not recognized. ValueError: If the `policy_type` is not recognized.
""" """
if policy_type == "tdmpc":
return TDMPCConfig(**kwargs)
elif policy_type == "diffusion":
return DiffusionConfig(**kwargs)
elif policy_type == "act":
return ACTConfig(**kwargs)
elif policy_type == "multi_task_dit":
return MultiTaskDiTConfig(**kwargs)
elif policy_type == "vqbet":
return VQBeTConfig(**kwargs)
elif policy_type == "pi0":
return PI0Config(**kwargs)
elif policy_type == "pi05":
return PI05Config(**kwargs)
elif policy_type == "gaussian_actor":
return GaussianActorConfig(**kwargs)
elif policy_type == "smolvla":
return SmolVLAConfig(**kwargs)
elif policy_type == "groot":
return GrootConfig(**kwargs)
elif policy_type == "xvla":
return XVLAConfig(**kwargs)
elif policy_type == "wall_x":
return WallXConfig(**kwargs)
elif policy_type == "eo1":
return EO1Config(**kwargs)
elif policy_type == "molmoact2":
return MolmoAct2Config(**kwargs)
elif policy_type == "vla_jepa":
return VLAJEPAConfig(**kwargs)
elif policy_type == "lingbot_va":
return LingBotVAConfig(**kwargs)
elif policy_type == "fastwam":
return FastWAMConfig(**kwargs)
elif policy_type == "evo1":
return Evo1Config(**kwargs)
else:
try: try:
config_cls = PreTrainedConfig.get_choice_class(policy_type) config_cls = PreTrainedConfig.get_choice_class(policy_type)
return config_cls(**kwargs)
except Exception as e: except Exception as e:
raise ValueError(f"Policy type '{policy_type}' is not available.") from e raise ValueError(f"Policy type '{policy_type}' is not available.") from e
return config_cls(**kwargs)
class ProcessorConfigKwargs(TypedDict, total=False): class ProcessorConfigKwargs(TypedDict, total=False):
@@ -298,8 +175,7 @@ def make_pre_post_processors(
A tuple containing the input (pre-processor) and output (post-processor) pipelines. A tuple containing the input (pre-processor) and output (post-processor) pipelines.
Raises: Raises:
NotImplementedError: If a processor factory is not implemented for the given ValueError: If no processor factory exists for the given policy configuration type.
policy configuration type.
""" """
if pretrained_path: if pretrained_path:
if isinstance(policy_cfg, GrootConfig): if isinstance(policy_cfg, GrootConfig):
@@ -308,6 +184,7 @@ def make_pre_post_processors(
return make_groot_pre_post_processors_from_pretrained( return make_groot_pre_post_processors_from_pretrained(
config=policy_cfg, config=policy_cfg,
pretrained_path=pretrained_path, pretrained_path=pretrained_path,
revision=pretrained_revision,
dataset_stats=kwargs.get("dataset_stats"), dataset_stats=kwargs.get("dataset_stats"),
dataset_meta=kwargs.get("dataset_meta"), dataset_meta=kwargs.get("dataset_meta"),
preprocessor_overrides=kwargs.get("preprocessor_overrides"), preprocessor_overrides=kwargs.get("preprocessor_overrides"),
@@ -351,167 +228,14 @@ def make_pre_post_processors(
) )
return preprocessor, postprocessor return preprocessor, postprocessor
# Create a new processor based on policy type # Create new processors from the policy config, resolving the per-policy factory
if isinstance(policy_cfg, TDMPCConfig): # function by naming convention (lazy import keeps optional dependencies optional).
from .tdmpc.processor_tdmpc import make_tdmpc_pre_post_processors return _make_processors_from_policy_config(
processors = make_tdmpc_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, DiffusionConfig):
from .diffusion.processor_diffusion import make_diffusion_pre_post_processors
processors = make_diffusion_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, ACTConfig):
from .act.processor_act import make_act_pre_post_processors
processors = make_act_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, MultiTaskDiTConfig):
from .multi_task_dit.processor_multi_task_dit import (
make_multi_task_dit_pre_post_processors,
)
processors = make_multi_task_dit_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, VQBeTConfig):
from .vqbet.processor_vqbet import make_vqbet_pre_post_processors
processors = make_vqbet_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, PI0Config):
from .pi0.processor_pi0 import make_pi0_pre_post_processors
processors = make_pi0_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, PI05Config):
from .pi05.processor_pi05 import make_pi05_pre_post_processors
processors = make_pi05_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, GaussianActorConfig):
from .gaussian_actor.processor_gaussian_actor import make_gaussian_actor_pre_post_processors
processors = make_gaussian_actor_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, SmolVLAConfig):
from .smolvla.processor_smolvla import make_smolvla_pre_post_processors
processors = make_smolvla_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, GrootConfig):
from .groot.processor_groot import make_groot_pre_post_processors
processors = make_groot_pre_post_processors(
config=policy_cfg, config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"), dataset_stats=kwargs.get("dataset_stats"),
dataset_meta=kwargs.get("dataset_meta"), dataset_meta=kwargs.get("dataset_meta"),
) )
elif isinstance(policy_cfg, XVLAConfig):
from .xvla.processor_xvla import (
make_xvla_pre_post_processors,
)
processors = make_xvla_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, WallXConfig):
from .wall_x.processor_wall_x import make_wall_x_pre_post_processors
processors = make_wall_x_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, EO1Config):
from .eo1.processor_eo1 import make_eo1_pre_post_processors
processors = make_eo1_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, Evo1Config):
from .evo1.processor_evo1 import make_evo1_pre_post_processors
processors = make_evo1_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, MolmoAct2Config):
from .molmoact2.processor_molmoact2 import make_molmoact2_pre_post_processors
processors = make_molmoact2_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
dataset_meta=kwargs.get("dataset_meta"),
)
elif isinstance(policy_cfg, VLAJEPAConfig):
from .vla_jepa.processor_vla_jepa import make_vla_jepa_pre_post_processors
processors = make_vla_jepa_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, LingBotVAConfig):
from .lingbot_va.processor_lingbot_va import make_lingbot_va_pre_post_processors
processors = make_lingbot_va_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, FastWAMConfig):
from .fastwam.processor_fastwam import make_fastwam_pre_post_processors
processors = make_fastwam_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
else:
try:
processors = _make_processors_from_policy_config(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
except Exception as e:
raise ValueError(f"Processor for policy type '{policy_cfg.type}' is not implemented.") from e
return processors
def make_policy( def make_policy(
cfg: PreTrainedConfig, cfg: PreTrainedConfig,
@@ -617,12 +341,15 @@ def make_policy(
# Load a pretrained PEFT model on top of the policy. The pretrained path points to the folder/repo # 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 # 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. # 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.") logging.info("Loading policy's PEFT adapter.")
peft_pretrained_path = str(cfg.pretrained_path) 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 kwargs["pretrained_name_or_path"] = peft_config.base_model_name_or_path
if not kwargs["pretrained_name_or_path"]: if not kwargs["pretrained_name_or_path"]:
@@ -633,9 +360,14 @@ def make_policy(
"the adapter was trained." "the adapter was trained."
) )
kwargs["revision"] = peft_config.revision
policy = policy_cls.from_pretrained(**kwargs) policy = policy_cls.from_pretrained(**kwargs)
policy = PeftModel.from_pretrained( 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: else:
@@ -654,10 +386,12 @@ def make_policy(
return policy return policy
def _get_policy_cls_from_policy_name(name: str) -> type[PreTrainedConfig]: def _get_policy_cls_from_policy_name(name: str) -> type[PreTrainedPolicy]:
"""Get policy class from its registered name using dynamic imports. """Get policy class from its registered name using dynamic imports.
This is used as a helper function to import policies from 3rd party lerobot plugins. Works for built-in policies and 3rd party lerobot plugins alike: the config class
registered under ``name`` is resolved via the draccus ChoiceRegistry, and the policy
class is imported from the sibling ``modeling_*`` module by naming convention.
Args: Args:
name: The name of the policy. name: The name of the policy.
@@ -683,22 +417,39 @@ def _get_policy_cls_from_policy_name(name: str) -> type[PreTrainedConfig]:
"configuration_", "modeling_" "configuration_", "modeling_"
) # e.g., configuration_diffusion -> modeling_diffusion ) # e.g., configuration_diffusion -> modeling_diffusion
try:
module = importlib.import_module(module_path) module = importlib.import_module(module_path)
policy_cls = getattr(module, cls_name) except ModuleNotFoundError as e:
if e.name == module_path:
# The modeling_* module itself does not exist for this policy type. A missing
# optional dependency inside an existing module propagates unchanged instead,
# so its actionable install hint stays visible.
raise ValueError(f"Policy class for '{name}' is not implemented.") from e
raise
policy_cls = getattr(module, cls_name, None)
if policy_cls is None:
raise ValueError(
f"Policy class '{cls_name}' not found in '{module_path}'. "
f"Policies must expose '<Name>Policy' in the sibling 'modeling_*' module by naming convention."
)
return policy_cls return policy_cls
def _make_processors_from_policy_config( def _make_processors_from_policy_config(
config: PreTrainedConfig, config: PreTrainedConfig,
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
dataset_meta: Any | None = None,
) -> tuple[Any, Any]: ) -> tuple[Any, Any]:
"""Create pre- and post-processors from a policy configuration using dynamic imports. """Create pre- and post-processors from a policy configuration using dynamic imports.
This is used as a helper function to import processor factories from 3rd party lerobot plugins. Resolves ``make_{type}_pre_post_processors`` from the policy's ``processor_*`` module
by naming convention. Works for built-in policies and 3rd party lerobot plugins.
Args: Args:
config: The policy configuration object. config: The policy configuration object.
dataset_stats: Dataset statistics for normalization. dataset_stats: Dataset statistics for normalization.
dataset_meta: Dataset metadata, forwarded only to factories that declare a
``dataset_meta`` parameter (e.g. groot, molmoact2).
Returns: Returns:
A tuple containing the input (pre-processor) and output (post-processor) pipelines. A tuple containing the input (pre-processor) and output (post-processor) pipelines.
""" """
@@ -711,6 +462,19 @@ def _make_processors_from_policy_config(
logging.debug( logging.debug(
f"Instantiating pre/post processors using function '{function_name}' from module '{module_path}'" f"Instantiating pre/post processors using function '{function_name}' from module '{module_path}'"
) )
try:
module = importlib.import_module(module_path) module = importlib.import_module(module_path)
function = getattr(module, function_name) except ModuleNotFoundError as e:
return function(config, dataset_stats=dataset_stats) if e.name == module_path:
# The processor_* module itself does not exist for this policy type. A missing
# optional dependency inside an existing module propagates unchanged instead,
# so its actionable install hint stays visible.
raise ValueError(f"Processor for policy type '{policy_type}' is not implemented.") from e
raise
function = getattr(module, function_name, None)
if function is None:
raise ValueError(f"Processor for policy type '{policy_type}' is not implemented.")
call_kwargs: dict[str, Any] = {"dataset_stats": dataset_stats}
if "dataset_meta" in inspect.signature(function).parameters:
call_kwargs["dataset_meta"] = dataset_meta
return function(config, **call_kwargs)
@@ -22,20 +22,11 @@ import torch
from lerobot.configs import PipelineFeatureType, PolicyFeature from lerobot.configs import PipelineFeatureType, PolicyFeature
from lerobot.processor import ( from lerobot.processor import (
ActionProcessorStep, ActionProcessorStep,
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
NormalizerProcessorStep,
PolicyAction, PolicyAction,
PolicyProcessorPipeline, PolicyProcessorPipeline,
ProcessorStepRegistry, ProcessorStepRegistry,
RenameObservationsProcessorStep, make_default_policy_processor_steps,
UnnormalizerProcessorStep, make_policy_processor_pipelines,
policy_action_to_transition,
transition_to_policy_action,
)
from lerobot.utils.constants import (
POLICY_POSTPROCESSOR_DEFAULT_NAME,
POLICY_PREPROCESSOR_DEFAULT_NAME,
) )
from .configuration_fastwam import FastWAMConfig from .configuration_fastwam import FastWAMConfig
@@ -105,38 +96,20 @@ def make_fastwam_pre_post_processors(
# anyway) and unsafe across fine-tuning: its `resize_size` would be inherited from the base # anyway) and unsafe across fine-tuning: its `resize_size` would be inherited from the base
# checkpoint's camera geometry, not this dataset's, making the concatenation N_cameras x too wide. # checkpoint's camera geometry, not this dataset's, making the concatenation N_cameras x too wide.
steps = make_default_policy_processor_steps(config, normalization_stats, normalizer_device=config.device)
input_steps = [ input_steps = [
RenameObservationsProcessorStep(rename_map={}), steps.rename_observations,
AddBatchDimensionProcessorStep(), steps.add_batch_dim,
DeviceProcessorStep(device=config.device), steps.to_device,
NormalizerProcessorStep( steps.normalize,
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=normalization_stats,
device=config.device,
),
] ]
output_steps = [ output_steps = [
UnnormalizerProcessorStep( steps.unnormalize,
features=config.output_features,
norm_map=config.normalization_mapping,
stats=normalization_stats,
),
] ]
if config.toggle_action_dimensions: if config.toggle_action_dimensions:
output_steps.append( output_steps.append(
FastWAMActionToggleProcessorStep(toggle_dimensions=config.toggle_action_dimensions) FastWAMActionToggleProcessorStep(toggle_dimensions=config.toggle_action_dimensions)
) )
output_steps.append(DeviceProcessorStep(device="cpu")) output_steps.append(steps.to_cpu)
return ( return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
steps=input_steps,
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
),
PolicyProcessorPipeline[PolicyAction, PolicyAction](
steps=output_steps,
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
),
)
@@ -37,13 +37,19 @@ def is_image_feature(key: str) -> bool:
@dataclass @dataclass
class ConcurrencyConfig: class ConcurrencyConfig:
"""Configuration for the concurrency of the actor and learner. """Configuration for the concurrency of the actor and learner.
Possible values are: Possible values are:
- "threads": Use threads for the actor and learner. - "threads": Use threads for the actor and learner.
- "processes": Use processes 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" actor: str = "threads"
learner: str = "threads" learner: str = "threads"
multiprocessing_context: str | None = "spawn"
@dataclass @dataclass
@@ -20,17 +20,10 @@ from typing import Any
import torch import torch
from lerobot.processor import ( from lerobot.processor import (
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
NormalizerProcessorStep,
PolicyAction, PolicyAction,
PolicyProcessorPipeline, PolicyProcessorPipeline,
RenameObservationsProcessorStep, make_default_pre_post_processors,
UnnormalizerProcessorStep,
policy_action_to_transition,
transition_to_policy_action,
) )
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
from .configuration_gaussian_actor import GaussianActorConfig from .configuration_gaussian_actor import GaussianActorConfig
@@ -62,33 +55,4 @@ def make_gaussian_actor_pre_post_processors(
Returns: Returns:
A tuple containing the configured pre-processor and post-processor pipelines. A tuple containing the configured pre-processor and post-processor pipelines.
""" """
return make_default_pre_post_processors(config, dataset_stats)
# Add remaining processors
input_steps = [
RenameObservationsProcessorStep(rename_map={}),
AddBatchDimensionProcessorStep(),
DeviceProcessorStep(device=config.device),
NormalizerProcessorStep(
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=dataset_stats,
),
]
output_steps = [
UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
DeviceProcessorStep(device="cpu"),
]
return (
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
steps=input_steps,
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
),
PolicyProcessorPipeline[PolicyAction, PolicyAction](
steps=output_steps,
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
),
)
@@ -68,6 +68,9 @@ class GrootPolicy(PreTrainedPolicy):
name = "groot" name = "groot"
config_class = GrootConfig config_class = GrootConfig
def supports_rtc(self) -> bool:
return True
def __init__(self, config: GrootConfig, **kwargs): def __init__(self, config: GrootConfig, **kwargs):
"""Initialize Groot policy wrapper.""" """Initialize Groot policy wrapper."""
require_package("transformers", extra="groot") require_package("transformers", extra="groot")

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