Compare commits

..

57 Commits

Author SHA1 Message Date
CarolinePascal 468a349f5e docs(update): updating docs with the task modification features 2026-07-30 18:50:45 +02:00
CarolinePascal 2f08bb32c1 docs(docstrings): updating docstrings 2026-07-30 17:40:52 +02:00
CarolinePascal e48b395c60 chore(fromat): formatting code 2026-07-30 17:32:15 +02:00
CarolinePascal dca9e80f03 feat(task modification precedence): Improving task modification precedence so that all modes can be used in a single run. Adapting tests accordinginly. 2026-07-30 17:30:12 +02:00
vvezre c1caa512ef Support task replacement mappings in modify_tasks
Part of #2326.

Signed-off-by: 陈伟 <woshei0a0a0a@qq.com>
2026-07-30 14:41:53 +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
107 changed files with 2529 additions and 595 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: ["*"]
+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**
```bash
pip install 'lerobot[feetech]' # SO-100/SO-101 motor stack
# pip install 'lerobot[all]' # everything
# pip install 'lerobot[aloha,pusht]' # specific features
# pip install 'lerobot[smolvla]' # add SmolVLA deps
# uv (recommended — see AGENTS.md and CLAUDE.md)
uv sync --locked --extra feetech # SO-100/SO-101 motor stack
# uv sync --locked --extra all # everything
# 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
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.
```bash
+4 -5
View File
@@ -68,17 +68,16 @@ ENV HOME=/home/user_lerobot \
# issues with MuJoCo and OpenGL drivers.
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 src/ src/
RUN uv sync --locked --extra all --no-cache
RUN uv sync --locked --extra all --no-install-project --no-cache
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
COPY --chown=user_lerobot:user_lerobot . .
RUN uv sync --locked --extra all --no-cache
# Set the default command
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 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 src/ src/
RUN uv sync --locked --extra all --no-install-project --no-cache
RUN uv sync --locked --extra all --no-cache
# Copy the rest of the application code
# Copy the application code and install the local project
# Make sure to have the git-LFS files for testing
COPY --chown=user_lerobot:user_lerobot . .
RUN uv sync --locked --extra all --no-cache
# Set the default command
CMD ["/bin/bash"]
+3 -3
View File
@@ -58,7 +58,7 @@ final_action = postprocessor(action)
## 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?
@@ -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.
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.
Then, add these same transformations on your inference script (shown here in the `record.py` script):
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 to your inference script (shown here in the `record.py` script):
```diff
action_values = predict_action(
+13
View File
@@ -136,6 +136,10 @@ config = RealSenseCameraConfig(
height=480,
color_mode=ColorMode.RGB,
use_depth=True,
# Optional fixed color controls. Omit them to leave the current sensor settings unchanged.
exposure=120,
gain=64,
white_balance=4600,
rotation=Cv2Rotation.NO_ROTATION
)
@@ -154,6 +158,15 @@ finally:
```
<!-- prettier-ignore-end -->
Manual color controls disable the corresponding automatic exposure or white-balance mode. Their
supported ranges vary by camera model; an invalid value raises an error at connection time that
includes the range reported by the sensor. Requesting an unsupported control also raises an error.
Omitted controls leave the sensor's existing automatic or manual setting unchanged. These options
require `use_rgb=True`.
On the RealSense D405, the color stream is provided by the Stereo Module, so changing manual
exposure or gain also affects the depth stream.
</hfoption>
</hfoptions>
+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:
````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
# Use SmolVLA policy with LIBERO environment
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,
)
act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
```
### 3. **Easier Experimentation**
@@ -145,7 +132,7 @@ class LiberoVelocityProcessorStep(ObservationProcessorStep):
state = torch.cat([eef_pos, eef_axisangle, eef_vel,
gripper_pos, gripper_vel], dim=-1) # 14D
return state
````
```
### 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:
1. **Select the motor** from the list by clicking on it
2. **Click on Upgrade tab**:
3. **Click on Online button**:
- If an potential firmware update is found, it will be displayed in the box
4. **Click on Upgrade button**:
2. **Click the Upgrade tab**:
3. **Click the Online button**:
- If a potential firmware update is found, it will be displayed in the box
4. **Click the Upgrade button**:
- The update progress will be displayed
## 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
- `so_leader` - SO100 / SO101 leader arm
- `bi_so_leader` - Bimanual SO100 / SO101 leader arms
> [!IMPORTANT]
> 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
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
lerobot-record \
+1 -2
View File
@@ -187,14 +187,13 @@ Teleop is optional — if omitted the robot holds its position during the reset
| `ESC` | Stop the recording session |
| Flag | Description |
| ----------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------- |
| ----------------------------------------------- | -------------------------------------------------------------------------- |
| `--dataset.num_episodes` | Number of episodes to record |
| `--dataset.episode_time_s` | Duration of each recording episode in seconds |
| `--dataset.reset_time_s` | Duration of the reset phase between episodes in seconds |
| `--teleop.type` | Optional. Teleoperator to drive the robot during resets |
| `--strategy.reset_to_initial_position` | Whether to reset the robot to its initial position between episodes |
| `--strategy.smooth_leader_to_follower_handover` | Whether to turn on or off the leader -> follower smooth handover behavior. |
| `--strategy.smooth_handover` | Smoothly hand control to the teleop at reset start (default: true). Disable for clutch-style teleops that re-reference at the current robot pose on engage |
---
+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
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):
+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]"
```
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.
# Step-by-Step Assembly Instructions
+1 -1
View File
@@ -174,7 +174,7 @@ The model takes images, text instructions, and robot state as input, and outputs
## 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:
+3 -3
View File
@@ -22,7 +22,7 @@ With processors, you choose the learning features you want to use for your polic
## Three pipelines
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)
2. Pipeline 2: Dataset action space → robot command space (EE targets → joints)
@@ -74,7 +74,7 @@ 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.
- `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.
- `transition_to_observation`: transforms the pipeline transition to a observation dict.
- `transition_to_observation`: transforms the pipeline transition to an observation dict.
Check out [src/lerobot/processor/converters.py](https://github.com/huggingface/lerobot/blob/main/src/lerobot/processor/converters.py) for more details.
@@ -82,7 +82,7 @@ Checkout [src/lerobot/processor/converters.py](https://github.com/huggingface/le
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
def transform_features(
+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")
# 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
action_queue = ActionQueue(policy_cfg.rtc_config)
@@ -100,7 +100,7 @@ Typical values: 8-12 steps
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.
+1 -1
View File
@@ -93,7 +93,7 @@ lerobot-train --help
## 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:
```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">
- 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.
- Attach the follower trigger with 4 M3x6mm screws.
+53 -5
View File
@@ -11,9 +11,10 @@ LeRobot provides several utilities for manipulating datasets:
3. **Merge Datasets** - Combine multiple datasets into one. The datasets must have identical features, and episodes are concatenated in the order specified in `repo_ids`
4. **Add Features** - Add new features to a dataset
5. **Remove Features** - Remove features from a dataset
6. **Convert to Video** - Convert image-based datasets to video format for efficient storage (RGB and depth cameras are encoded with separate encoders)
7. **Re-encode Videos** - Re-encode an existing video dataset's RGB and/or depth streams with new encoder settings
8. **Show the Info of Datasets** - Show the summary of datasets information such as number of episode etc.
6. **Modify Tasks** - Change the natural-language task descriptions associated with episodes
7. **Convert to Video** - Convert image-based datasets to video format for efficient storage (RGB and depth cameras are encoded with separate encoders)
8. **Re-encode Videos** - Re-encode an existing video dataset's RGB and/or depth streams with new encoder settings
9. **Show the Info of Datasets** - Show the summary of datasets information such as number of episode etc.
The core implementation is in `lerobot.datasets.dataset_tools`.
An example script detailing how to use the tools API is available in `examples/dataset/use_dataset_tools.py`.
@@ -50,11 +51,11 @@ lerobot-edit-dataset \
Divide a dataset into multiple subsets.
```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 \
--repo_id lerobot/pusht \
--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
lerobot-edit-dataset \
@@ -89,6 +90,53 @@ lerobot-edit-dataset \
--operation.feature_names "['observation.images.top']"
```
#### Modify Tasks
Change the natural-language task descriptions attached to episodes. This is useful for fixing typos, standardizing wording, or re-labeling episodes.
> [!WARNING]
> `modify_tasks` modifies the dataset **in-place** (updating `meta/tasks.parquet`, the `task_index` column in the data files, the `tasks` column in the episode metadata, and `total_tasks` in `meta/info.json`). The `--new_repo_id` and `--new_root` parameters are ignored for this operation.
```bash
# Set a single task for all episodes
lerobot-edit-dataset \
--repo_id lerobot/pusht \
--operation.type modify_tasks \
--operation.new_task "Pick up the cube and place it"
# Set different tasks for specific episodes
lerobot-edit-dataset \
--repo_id lerobot/pusht \
--operation.type modify_tasks \
--operation.episode_tasks '{"0": "Task A", "1": "Task B", "2": "Task A"}'
# Replace existing task strings wherever they appear
lerobot-edit-dataset \
--repo_id lerobot/pusht \
--operation.type modify_tasks \
--operation.task_replacements '{"Pick up the red cube": "Lift the red cube"}'
# Combine modes in a single run
lerobot-edit-dataset \
--repo_id lerobot/pusht \
--operation.type modify_tasks \
--operation.new_task "Default task" \
--operation.task_replacements '{"Pick up the red cube": "Lift the red cube"}' \
--operation.episode_tasks '{"5": "Special task for episode 5"}'
```
**Parameters:**
- `new_task`: A single task string used as the default for episodes not otherwise covered.
- `episode_tasks`: Mapping from episode index to task string.
- `task_replacements`: Mapping from existing task strings to their replacements, applied to episodes whose current task matches a key. Every key must be an existing task in the dataset.
The modes can be combined in a single run. Per episode, the task is resolved with the following precedence:
`episode_tasks` > `task_replacements` > `new_task` > original task
At least one of `new_task`, `episode_tasks`, or `task_replacements` must be specified. An episode that ends up with no task raises an error.
#### Convert to Video
Convert an image-based dataset to video format, creating a new LeRobotDataset where images are stored as videos. This is useful for reducing storage requirements and improving data loading performance. The new dataset will have the exact same structure as the original, but with images encoded as MP4 videos in the proper LeRobot format.
+14 -1
View File
@@ -67,7 +67,7 @@ dependencies = [
"einops>=0.8.0,<0.9.0",
# 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",
"requests>=2.32.0,<3.0.0",
@@ -494,6 +494,19 @@ ignore_errors = true
module = "lerobot.envs.*"
ignore_errors = false
[[tool.mypy.overrides]]
module = "lerobot.annotations.*"
ignore_errors = false
disallow_untyped_defs = true
disallow_incomplete_defs = true
check_untyped_defs = true
[[tool.mypy.overrides]]
module = "lerobot.transforms.*"
ignore_errors = false
disallow_untyped_defs = true
disallow_incomplete_defs = true
check_untyped_defs = true
# [[tool.mypy.overrides]]
# module = "lerobot.utils.*"
+46 -16
View File
@@ -120,14 +120,22 @@ class OpenCVCamera(Camera):
self.rotation: int | None = get_cv2_rotation(config.rotation)
self.backend: int = config.backend
if self.height and self.width:
self.capture_width, self.capture_height = self.width, self.height
if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
self.capture_width, self.capture_height = self.height, self.width
self.capture_width: int | None = None
self.capture_height: int | None = None
self._reset_connection_settings()
def __str__(self) -> str:
return f"{self.__class__.__name__}({self.index_or_path})"
def _reset_connection_settings(self) -> None:
"""Restore settings that may have been auto-detected during a failed connection."""
self.fps = self.config.fps
self.width = self.config.width
self.height = self.config.height
self.capture_width, self.capture_height = self.width, self.height
if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
self.capture_width, self.capture_height = self.height, self.width
@property
def is_connected(self) -> bool:
"""Checks if the camera is currently connected and opened."""
@@ -164,6 +172,7 @@ class OpenCVCamera(Camera):
f"Failed to open {self}.Run `lerobot-find-cameras opencv` to find available cameras."
)
try:
self._configure_capture_settings()
self._start_read_thread()
@@ -175,6 +184,13 @@ class OpenCVCamera(Camera):
with self.frame_lock:
if self.latest_frame is None:
raise ConnectionError(f"{self} failed to capture frames during warmup.")
except BaseException:
try:
self._cleanup_resources()
except Exception:
logger.exception(f"Failed to fully clean up {self} after connect() failed.")
self._reset_connection_settings()
raise
logger.info(f"{self} connected.")
@@ -312,6 +328,7 @@ class OpenCVCamera(Camera):
for target in targets_to_scan:
camera = cv2.VideoCapture(target)
try:
if camera.isOpened():
default_width = int(camera.get(cv2.CAP_PROP_FRAME_WIDTH))
default_height = int(camera.get(cv2.CAP_PROP_FRAME_HEIGHT))
@@ -321,7 +338,9 @@ class OpenCVCamera(Camera):
# Get FOURCC code and convert to string
default_fourcc_code = camera.get(cv2.CAP_PROP_FOURCC)
default_fourcc_code_int = int(default_fourcc_code)
default_fourcc = "".join([chr((default_fourcc_code_int >> 8 * i) & 0xFF) for i in range(4)])
default_fourcc = "".join(
[chr((default_fourcc_code_int >> 8 * i) & 0xFF) for i in range(4)]
)
camera_info = {
"name": f"OpenCV Camera @ {target}",
@@ -338,6 +357,7 @@ class OpenCVCamera(Camera):
}
found_cameras_info.append(camera_info)
finally:
camera.release()
return found_cameras_info
@@ -496,6 +516,26 @@ class OpenCVCamera(Camera):
self.latest_timestamp = None
self.new_frame_event.clear()
def _cleanup_resources(self) -> None:
"""Stop background reads and release the capture, including after partial setup."""
read_thread = self.thread
videocapture = self.videocapture
try:
self._stop_read_thread()
finally:
self.videocapture = None
try:
if videocapture is not None:
videocapture.release()
finally:
# Releasing the device may unblock a hardware read that outlived
# the first bounded join in _stop_read_thread().
if read_thread is not None and read_thread.is_alive():
read_thread.join(timeout=2.0)
if read_thread.is_alive(): # pragma: no cover
logger.warning(f"{self} read thread remained alive after releasing the capture.")
@check_if_not_connected
def async_read(self, timeout_ms: float = 200) -> NDArray[Any]:
"""
@@ -586,16 +626,6 @@ class OpenCVCamera(Camera):
if not self.is_connected and self.thread is None:
raise DeviceNotConnectedError(f"{self} not connected.")
if self.thread is not None:
self._stop_read_thread()
if self.videocapture is not None:
self.videocapture.release()
self.videocapture = None
with self.frame_lock:
self.latest_frame = None
self.latest_timestamp = None
self.new_frame_event.clear()
self._cleanup_resources()
logger.info(f"{self} disconnected.")
+157 -19
View File
@@ -121,6 +121,9 @@ class RealSenseCamera(Camera):
self.config = config
self.width: int | None = config.width
self.height: int | None = config.height
if config.serial_number_or_name.isdigit():
self.serial_number = config.serial_number_or_name
else:
@@ -131,6 +134,9 @@ class RealSenseCamera(Camera):
self.use_rgb = config.use_rgb
self.use_depth = config.use_depth
self.warmup_s = config.warmup_s
self.exposure: int | None = config.exposure
self.gain: int | None = config.gain
self.white_balance: int | None = config.white_balance
self.rs_pipeline: rs.pipeline | None = None
self.rs_profile: rs.pipeline_profile | None = None
@@ -145,14 +151,23 @@ class RealSenseCamera(Camera):
self.rotation: int | None = get_cv2_rotation(config.rotation)
if self.height and self.width:
self.capture_width, self.capture_height = self.width, self.height
if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
self.capture_width, self.capture_height = self.height, self.width
self.capture_width: int | None = None
self.capture_height: int | None = None
self._reset_connection_settings()
def __str__(self) -> str:
return f"{self.__class__.__name__}({self.serial_number})"
def _reset_connection_settings(self) -> None:
"""Restore settings that may have been auto-detected during a failed connection."""
self.fps = self.config.fps
self.width = self.config.width
self.height = self.config.height
self.warmup_s = self.config.warmup_s
self.capture_width, self.capture_height = self.width, self.height
if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
self.capture_width, self.capture_height = self.height, self.width
@property
def is_connected(self) -> bool:
"""Checks if the camera pipeline is started and streams are active."""
@@ -172,7 +187,8 @@ class RealSenseCamera(Camera):
Raises:
DeviceAlreadyConnectedError: If the camera is already connected.
ValueError: If the configuration is invalid (e.g., missing serial/name, name not unique).
ValueError: If the configuration is invalid, a requested sensor option is unsupported,
or a requested sensor value is invalid.
ConnectionError: If the camera is found but fails to start the pipeline or no RealSense devices are detected at all.
RuntimeError: If the pipeline starts but fails to apply requested settings.
"""
@@ -190,7 +206,9 @@ class RealSenseCamera(Camera):
f"Failed to open {self}.Run `lerobot-find-cameras realsense` to find available cameras."
) from e
try:
self._configure_capture_settings()
self._configure_sensor_options()
self._start_read_thread()
# NOTE(Steven/Caroline): Enforcing at least one second of warmup as RS cameras need a bit of time before the first read. If we don't wait, the first read from the warmup will raise.
@@ -206,6 +224,13 @@ class RealSenseCamera(Camera):
self.use_depth and self.latest_depth_frame is None
):
raise ConnectionError(f"{self} failed to capture frames during warmup.")
except BaseException:
try:
self._cleanup_resources()
except Exception:
logger.exception(f"Failed to fully clean up {self} after connect() failed.")
self._reset_connection_settings()
raise
logger.info(f"{self} connected.")
@@ -339,6 +364,111 @@ class RealSenseCamera(Camera):
self.new_frame_event.clear()
return self._async_read(timeout_ms=10000, read_depth=read_depth)
def _get_color_sensor(self) -> "rs.sensor":
"""Returns the sensor that controls the color stream.
Most RealSense cameras expose "RGB Camera" for color. The D405 has no
separate RGB module — its color stream comes from "Stereo Module".
We try RGB Camera first, then fall back to Stereo Module.
"""
if self.rs_profile is None:
raise RuntimeError(f"{self}: rs_profile must be initialized before use.")
device = self.rs_profile.get_device()
sensors = {s.get_info(rs.camera_info.name): s for s in device.query_sensors()}
for name in ("RGB Camera", "Stereo Module"):
if name in sensors:
return sensors[name]
available = list(sensors.keys())
raise RuntimeError(f"{self}: no color sensor found. Available sensors: {available}")
def _set_sensor_option(self, sensor: "rs.sensor", option: "rs.option", value: float, label: str) -> None:
"""Sets a sensor option, re-raising range errors with actionable diagnostics."""
try:
sensor.set_option(option, value)
except Exception as e:
range_info = ""
try:
option_range = sensor.get_option_range(option)
range_info = (
f" (supported range: min={option_range.min}, max={option_range.max}, "
f"step={option_range.step}, default={option_range.default})"
)
except Exception:
range_info = " (option range unavailable)"
raise ValueError(
f"{self}: failed to set {label} to {value}{range_info}. Original error: {e}"
) from e
def _configure_sensor_options(self) -> None:
"""Applies manual sensor options (exposure, gain, white balance) to the color sensor.
When exposure or gain is set, auto-exposure is disabled first. When white_balance
is set, auto white balance is disabled first. An omitted option is left unchanged,
and configuration is skipped entirely if all options are omitted.
Raises:
ValueError: If the sensor does not support a requested option or a requested
value is invalid. Invalid-value errors include the option name, requested
value, and supported range when available.
"""
if self.exposure is None and self.gain is None and self.white_balance is None:
return
color_sensor = self._get_color_sensor()
requested_options = (
(rs.option.exposure, self.exposure, "exposure"),
(rs.option.gain, self.gain, "gain"),
(rs.option.white_balance, self.white_balance, "white balance"),
)
unsupported_options = [
label
for option, value, label in requested_options
if value is not None and not color_sensor.supports(option)
]
if unsupported_options:
raise ValueError(
f"{self}: color sensor does not support requested manual options: {unsupported_options}."
)
manual_exposure_requested = self.exposure is not None or self.gain is not None
if manual_exposure_requested:
if color_sensor.supports(rs.option.enable_auto_exposure):
self._set_sensor_option(color_sensor, rs.option.enable_auto_exposure, 0, "auto-exposure")
logger.info(f"{self} auto-exposure disabled.")
else:
logger.warning(
f"{self} sensor does not support disabling auto-exposure; "
"applying manual exposure/gain directly."
)
if self.exposure is not None:
self._set_sensor_option(color_sensor, rs.option.exposure, self.exposure, "exposure")
logger.info(f"{self} exposure set to {self.exposure}.")
if self.gain is not None:
self._set_sensor_option(color_sensor, rs.option.gain, self.gain, "gain")
logger.info(f"{self} gain set to {self.gain}.")
if self.white_balance is not None:
if color_sensor.supports(rs.option.enable_auto_white_balance):
self._set_sensor_option(
color_sensor, rs.option.enable_auto_white_balance, 0, "auto white balance"
)
logger.info(f"{self} auto white balance disabled.")
else:
logger.warning(
f"{self} sensor does not support disabling auto white balance; "
"applying manual white balance directly."
)
self._set_sensor_option(
color_sensor, rs.option.white_balance, self.white_balance, "white balance"
)
logger.info(f"{self} white balance set to {self.white_balance}.")
@check_if_not_connected
def read_depth(self, timeout_ms: int = 200) -> NDArray[Any]:
"""
@@ -541,6 +671,27 @@ class RealSenseCamera(Camera):
self.latest_timestamp = None
self.new_frame_event.clear()
def _cleanup_resources(self) -> None:
"""Stop background reads and stop the pipeline, including after partial setup."""
read_thread = self.thread
rs_pipeline = self.rs_pipeline
try:
self._stop_read_thread()
finally:
self.rs_pipeline = None
self.rs_profile = None
try:
if rs_pipeline is not None:
rs_pipeline.stop()
finally:
# Stopping the pipeline may unblock a hardware read that outlived
# the first bounded join in _stop_read_thread().
if read_thread is not None and read_thread.is_alive():
read_thread.join(timeout=2.0)
if read_thread.is_alive(): # pragma: no cover
logger.warning(f"{self} read thread remained alive after stopping the pipeline.")
def _async_read(self, timeout_ms: float, read_depth: bool = False) -> NDArray[Any]:
"""Shared helper for :meth:`async_read`/:meth:`async_read_depth`: return the latest buffered frame."""
if self.thread is None or not self.thread.is_alive():
@@ -684,18 +835,5 @@ class RealSenseCamera(Camera):
f"Attempted to disconnect {self}, but it appears already disconnected."
)
if self.thread is not None:
self._stop_read_thread()
if self.rs_pipeline is not None:
self.rs_pipeline.stop()
self.rs_pipeline = None
self.rs_profile = None
with self.frame_lock:
self.latest_color_frame = None
self.latest_depth_frame = None
self.latest_timestamp = None
self.new_frame_event.clear()
self._cleanup_resources()
logger.info(f"{self} disconnected.")
@@ -46,6 +46,17 @@ class RealSenseCameraConfig(CameraConfig):
use_depth: Whether to enable depth stream. Defaults to False.
rotation: Image rotation setting (0°, 90°, 180°, or 270°). Defaults to no rotation.
warmup_s: Time reading frames before returning from connect (in seconds)
exposure: Manual exposure value for the color sensor. When set, auto-exposure is
disabled and this fixed value is used. Valid ranges are camera-model specific
and reported if the value is rejected. Defaults to None (leave unchanged).
gain: Manual gain value for the color sensor. When set, auto-exposure is disabled
and this fixed gain is used, which also freezes exposure at its current value
when no exposure is configured. Valid ranges are camera-model specific and
reported if the value is rejected. Defaults to None (leave unchanged).
white_balance: Manual white balance value for the color sensor. When set, auto
white balance is disabled and this fixed value is used. Valid ranges are
camera-model specific and reported if the value is rejected. Defaults to None
(leave unchanged).
Note:
- Either name or serial_number must be specified.
@@ -61,6 +72,9 @@ class RealSenseCameraConfig(CameraConfig):
use_depth: bool = False
rotation: Cv2Rotation = Cv2Rotation.NO_ROTATION
warmup_s: int = 1
exposure: int | None = None
gain: int | None = None
white_balance: int | None = None
def __post_init__(self) -> None:
self.color_mode = ColorMode(self.color_mode)
@@ -69,6 +83,18 @@ class RealSenseCameraConfig(CameraConfig):
if not self.use_rgb and not self.use_depth:
raise ValueError("At least one of `use_rgb` or `use_depth` must be enabled.")
manual_color_options = {
"exposure": self.exposure,
"gain": self.gain,
"white_balance": self.white_balance,
}
configured_color_options = [name for name, value in manual_color_options.items() if value is not None]
if configured_color_options and not self.use_rgb:
raise ValueError(
"Manual color sensor options require `use_rgb=True`. "
f"Configured options: {configured_color_options}."
)
values = (self.fps, self.width, self.height)
if any(v is not None for v in values) and any(v is None for v in values):
raise ValueError(
+6
View File
@@ -71,13 +71,19 @@ class DatasetRecordConfig:
# Number of threads per encoder instance. None = auto (codec default).
# Lower values reduce CPU usage, maps to 'lp' (via svtav1-params) for libsvtav1 and 'threads' for h264/hevc..
encoder_threads: int | None = None
# Skip appending the date-time tag to repo_id, keeping the user-provided name as-is
# (e.g. self-managed versioned names intended for a later `lerobot-edit-dataset merge`).
no_stamp: bool = False
def stamp_repo_id(self) -> None:
"""Append a date-time tag to ``repo_id`` so each recording session gets a unique name.
Must be called explicitly at dataset *creation* time not on resume,
where the existing ``repo_id`` (already stamped) must be preserved.
No-op when ``no_stamp`` is set, preserving a user-managed ``repo_id``.
"""
if self.no_stamp:
return
if self.repo_id:
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
self.repo_id = f"{self.repo_id}_{timestamp}"
+4 -2
View File
@@ -163,8 +163,10 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno
return None
def _save_pretrained(self, save_directory: Path) -> None:
with open(save_directory / CONFIG_NAME, "w") as f, draccus.config_type("json"):
draccus.dump(self, f, indent=4)
# Encode against the base class so draccus includes the choice "type" key,
# 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
def from_pretrained(
+4 -2
View File
@@ -103,8 +103,10 @@ class RewardModelConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
pass
def _save_pretrained(self, save_directory: Path) -> None:
with open(save_directory / CONFIG_NAME, "w") as f, draccus.config_type("json"):
draccus.dump(self, f, indent=4)
# Encode against the base class so draccus includes the choice "type" key,
# 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
def from_pretrained(
+23 -1
View File
@@ -14,6 +14,7 @@
import builtins
import datetime as dt
import json
import multiprocessing
import os
import tempfile
from dataclasses import dataclass, field
@@ -101,6 +102,12 @@ class TrainPipelineConfig(HubMixin):
batch_size: int = 8
prefetch_factor: int = 4
persistent_workers: bool = True
# DataLoader worker start method. "spawn" is safer than "fork" with
# non-fork-safe libs (PyAV / torchcodec / ffmpeg), but adds some
# worker-startup time per run since workers re-import modules instead
# of inheriting parent state. Override with `--dataloader_multiprocessing_context=fork`
# when appropriate, or set it to `null` to use Python's platform default.
dataloader_multiprocessing_context: str | None = "spawn"
steps: int = 100_000
# Run policy in the simulation environment every N steps to measure reward/success (0 = disabled).
env_eval_freq: int = 20_000
@@ -187,7 +194,11 @@ class TrainPipelineConfig(HubMixin):
)
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
elif self.job.is_remote:
return
@@ -212,6 +223,17 @@ class TrainPipelineConfig(HubMixin):
self.reward_model.pretrained_path = str(policy_dir)
def validate(self) -> None:
available_contexts = multiprocessing.get_all_start_methods()
if (
self.dataloader_multiprocessing_context is not None
and self.dataloader_multiprocessing_context not in available_contexts
):
raise ValueError(
"`dataloader_multiprocessing_context` must be None or one of "
f"{available_contexts} on this platform, got "
f"{self.dataloader_multiprocessing_context!r}."
)
self._resolve_pretrained_from_cli()
if self.policy is None and self.reward_model is None:
+96 -40
View File
@@ -19,6 +19,7 @@ import copy
import logging
import shutil
from pathlib import Path
from typing import Any, NotRequired, TypedDict
import datasets
import pandas as pd
@@ -49,8 +50,32 @@ from .utils import (
)
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.
Args:
@@ -59,14 +84,14 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
Returns:
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"]
for vk in video_keys:
video_infos = [m.features.get(vk, {}).get("info") or {} for m in all_metadata]
base_video_info = video_infos[0]
merged_encoder_info: dict = {}
merged_encoder_info: dict[str, Any] = {}
fallback_keys: list[str] = []
for info_key in VIDEO_ENCODER_INFO_KEYS:
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
if fallback_keys:
logging.warning(
logger.warning(
f"Merging heterogeneous or incomplete video encoder metadata for feature {vk}. "
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
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.
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
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.
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(
df,
dst_meta,
meta_idx,
data_idx,
videos_idx,
):
df: pd.DataFrame,
dst_meta: LeRobotDatasetMetadata,
meta_idx: IndexState,
data_idx: IndexState,
videos_idx: VideoIndexState,
) -> pd.DataFrame:
"""Updates metadata DataFrame with new chunk, file, and timestamp indices.
Adjusts all indices and timestamps to account for previously aggregated
@@ -289,7 +316,7 @@ def aggregate_datasets(
chunk_size: int | None = None,
concatenate_videos: bool = True,
concatenate_data: bool = True,
):
) -> None:
"""Aggregates multiple LeRobot datasets into a single unified dataset.
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_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:
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,
)
logging.info("Find all tasks")
logger.info("Find all tasks")
unique_tasks = pd.concat([m.tasks for m in all_metadata]).index.unique()
dst_meta.tasks = pd.DataFrame(
{"task_index": range(len(unique_tasks))}, index=pd.Index(unique_tasks, name="task")
)
meta_idx = {"chunk": 0, "file": 0}
data_idx = {"chunk": 0, "file": 0}
videos_idx = {
meta_idx: IndexState = {"chunk": 0, "file": 0}
data_idx: IndexState = {"chunk": 0, "file": 0}
videos_idx: VideoIndexState = {
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
finalize_aggregation(dst_meta, all_metadata)
logging.info("Aggregation complete.")
logger.info("Aggregation complete.")
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.
Handles video file concatenation and rotation based on file size limits.
@@ -406,7 +438,8 @@ def aggregate_videos(
videos_idx[key]["dst_file_durations"] = {}
for key, video_idx in videos_idx.items():
unique_chunk_file_pairs = {
unique_chunk_file_pairs: list[ChunkFile] = sorted(
{
(chunk, file)
for chunk, file in zip(
src_meta.episodes[f"videos/{key}/chunk_index"],
@@ -414,7 +447,7 @@ def aggregate_videos(
strict=False,
)
}
unique_chunk_file_pairs = sorted(unique_chunk_file_pairs)
)
chunk_idx = video_idx["chunk"]
file_idx = video_idx["file"]
@@ -489,7 +522,14 @@ def aggregate_videos(
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.
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:
dict: Updated data_idx with current chunk and file indices.
"""
unique_chunk_file_ids = {
unique_chunk_file_ids: list[ChunkFile] = sorted(
{
(c, f)
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
# 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
# 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:
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
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.
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:
dict: Updated meta_idx with current chunk and file indices.
"""
chunk_file_ids = {
chunk_file_ids: list[ChunkFile] = sorted(
{
(c, f)
for c, f in zip(
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,
)
}
chunk_file_ids = sorted(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)
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(
df: pd.DataFrame,
src_path: Path,
idx: dict[str, int],
idx: IndexState,
max_mb: float,
chunk_size: int,
default_path: str,
contains_images: bool = False,
aggr_root: Path = None,
aggr_root: Path | None = None,
hf_features: datasets.Features | None = None,
concatenate: bool = True,
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.
Manages file rotation when size limits are exceeded to prevent individual files
@@ -654,7 +702,13 @@ def append_or_create_parquet_file(
Returns:
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.
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_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)
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.
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.
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)
logging.info("write info")
logger.info("write info")
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_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)}"}
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])
write_stats(aggr_meta.stats, aggr_meta.root)
+2 -2
View File
@@ -188,8 +188,8 @@ class LeRobotDatasetMetadata:
def _load_metadata(self):
self.info = load_info(self.root)
check_version_compatibility(self.repo_id, self._version, CODEBASE_VERSION)
self.tasks = load_tasks(self.root)
self.episodes = load_episodes(self.root)
self.tasks = load_tasks(self.root) if self.total_tasks > 0 else None
self.episodes = load_episodes(self.root) if self.total_episodes > 0 else None
self.stats = load_stats(self.root)
def ensure_readable(self) -> None:
+38 -16
View File
@@ -1435,15 +1435,18 @@ def modify_tasks(
dataset: LeRobotDataset,
new_task: str | None = None,
episode_tasks: dict[int, str] | None = None,
task_replacements: dict[str, str] | None = None,
) -> LeRobotDataset:
"""Modify tasks in a LeRobotDataset.
This function allows you to either:
1. Set a single task for the entire dataset (using `new_task`)
2. Set specific tasks for specific episodes (using `episode_tasks`)
3. Replace existing task strings wherever they appear (using `task_replacements`)
You can combine both: `new_task` sets the default, and `episode_tasks` overrides
specific episodes.
Per episode, the task is resolved with precedence:
`episode_tasks` > `task_replacements` > `new_task` > original task. An episode that ends
up with no task (none of the above apply and it had no original task) raises an error.
The dataset is modified in-place, updating only the task-related files:
- meta/tasks.parquet
@@ -1453,11 +1456,14 @@ def modify_tasks(
Args:
dataset: The source LeRobotDataset to modify.
new_task: A single task string to apply to all episodes. If None and episode_tasks
is also None, raises an error.
episode_tasks: Optional dict mapping episode indices to their task strings.
Overrides `new_task` for specific episodes.
new_task: Default task applied to any episode not covered by `episode_tasks` or a
matching `task_replacements` entry.
episode_tasks: Optional dict mapping episode indices to task strings. Takes precedence
over both `task_replacements` and `new_task`.
task_replacements: Optional dict mapping existing task strings to new ones. Applied to
episodes whose current task matches a key. Every key must be an existing task.
At least one of `new_task`, `episode_tasks`, or `task_replacements` must be provided.
Examples:
Set a single task for all episodes:
@@ -1475,11 +1481,17 @@ def modify_tasks(
new_task="Default task",
episode_tasks={5: "Special task for episode 5"}
)
"""
if new_task is None and episode_tasks is None:
raise ValueError("Must specify at least one of new_task or episode_tasks")
if episode_tasks is not None:
Replace existing task strings in-place:
dataset = modify_tasks(
dataset,
task_replacements={"Pick up the cube": "Lift the cube"}
)
"""
if not new_task and not episode_tasks and not task_replacements:
raise ValueError("Must specify at least one of new_task, episode_tasks, or task_replacements")
if episode_tasks:
valid_indices = set(range(dataset.meta.total_episodes))
invalid = set(episode_tasks.keys()) - valid_indices
if invalid:
@@ -1489,19 +1501,29 @@ def modify_tasks(
if dataset.meta.episodes is None:
dataset.meta.episodes = load_episodes(dataset.root)
if task_replacements:
current_tasks = set(dataset.meta.tasks.index)
invalid_tasks = set(task_replacements) - current_tasks
if invalid_tasks:
raise ValueError(f"Task replacements reference unknown tasks: {sorted(invalid_tasks)}")
# Build the mapping from episode index to task string
episode_to_task: dict[int, str] = {}
for ep_idx in range(dataset.meta.total_episodes):
original_tasks = dataset.meta.episodes[ep_idx]["tasks"]
original_task = original_tasks[0] if original_tasks else None
if episode_tasks and ep_idx in episode_tasks:
episode_to_task[ep_idx] = episode_tasks[ep_idx]
elif new_task is not None:
elif task_replacements and original_task in task_replacements:
episode_to_task[ep_idx] = task_replacements[original_task]
elif new_task:
episode_to_task[ep_idx] = new_task
else:
elif original_task:
# Keep original task if not overridden and no default provided
original_tasks = dataset.meta.episodes[ep_idx]["tasks"]
if not original_tasks:
raise ValueError(f"Episode {ep_idx} has no tasks and no default task was provided")
episode_to_task[ep_idx] = original_tasks[0]
episode_to_task[ep_idx] = original_task
else:
raise ValueError(f"Episode {ep_idx} has no task; provide new_task or episode_tasks")
# Collect all unique tasks and create new task mapping
unique_tasks = sorted(set(episode_to_task.values()))
+23 -14
View File
@@ -172,6 +172,23 @@ class DatasetWriter:
def _get_image_file_dir(self, episode_index: int, image_key: str) -> Path:
return self._get_image_file_path(episode_index, image_key, frame_index=0).parent
def _get_episode_buffer_index(self) -> int:
episode_index = self.episode_buffer["episode_index"]
# episode_index is `int` when freshly created, but becomes `np.ndarray` after
# save_episode() mutates the buffer. Handle both types here.
if isinstance(episode_index, np.ndarray):
episode_index = episode_index.item() if episode_index.size == 1 else episode_index[0]
return int(episode_index)
def _delete_camera_frame_dirs(self, camera_keys: list[str]) -> None:
if self.image_writer is not None:
self._wait_image_writer()
episode_index = self._get_episode_buffer_index()
for camera_key in camera_keys:
img_dir = self._get_image_file_dir(episode_index, camera_key)
if img_dir.is_dir():
shutil.rmtree(img_dir)
def _save_image(
self, image: torch.Tensor | np.ndarray | PIL.Image.Image, fpath: Path, compress_level: int = 1
) -> None:
@@ -369,7 +386,9 @@ class DatasetWriter:
self._episodes_since_last_encoding = 0
if episode_data is None:
self.clear_episode_buffer(delete_images=len(self._meta.image_keys) > 0)
if len(self._meta.image_keys) > 0:
self._delete_camera_frame_dirs(self._meta.image_keys)
self.episode_buffer = self._create_episode_buffer()
def _batch_save_episode_video(self, start_episode: int, end_episode: int | None = None) -> None:
"""Batch save videos for multiple episodes."""
@@ -561,10 +580,10 @@ class DatasetWriter:
return metadata
def clear_episode_buffer(self, delete_images: bool = True) -> None:
"""Discard the current episode buffer and optionally delete temp images.
"""Discard the current episode buffer and optionally delete temp camera frames.
Args:
delete_images: If ``True``, remove temporary image directories
delete_images: If ``True``, remove temporary camera frame directories
written for the current episode.
"""
# Cancel streaming encoder if active
@@ -572,17 +591,7 @@ class DatasetWriter:
self._streaming_encoder.cancel_episode()
if delete_images:
if self.image_writer is not None:
self._wait_image_writer()
episode_index = self.episode_buffer["episode_index"]
# episode_index is `int` when freshly created, but becomes `np.ndarray` after
# save_episode() mutates the buffer. Handle both types here.
if isinstance(episode_index, np.ndarray):
episode_index = episode_index.item() if episode_index.size == 1 else episode_index[0]
for cam_key in self._meta.image_keys:
img_dir = self._get_image_file_dir(episode_index, cam_key)
if img_dir.is_dir():
shutil.rmtree(img_dir)
self._delete_camera_frame_dirs(self._meta.camera_keys)
self.episode_buffer = self._create_episode_buffer()
+5
View File
@@ -384,7 +384,12 @@ class LiberoEnv(gym.Env):
def close(self):
if self._env is not None:
try:
self._env.close()
finally:
# LIBERO deletes its inner env on close, so this wrapper must
# be recreated before the next reset.
self._env = None
def _make_env_fns(
+3
View File
@@ -155,6 +155,7 @@ class MetaworldEnv(gym.Env):
env.model.cam_pos[2] = [0.75, 0.075, 0.7]
env.reset()
env._freeze_rand_vec = False # otherwise no randomization
env.seeded_rand_vec = True # use seeded RNG so reset(seed=X) controls object positions
self._env = env
def render(self) -> np.ndarray:
@@ -220,6 +221,8 @@ class MetaworldEnv(gym.Env):
self._ensure_env()
super().reset(seed=seed)
if seed is not None:
self._env.seed(seed)
raw_obs, info = self._env.reset(seed=seed)
observation = self._format_raw_obs(raw_obs)
+3 -1
View File
@@ -384,7 +384,9 @@ class RoboTwinEnv(gym.Env):
self._env: Any | None = None # deferred — created on first reset() inside worker
self._step_count: int = 0
self._black_frame = np.zeros((self.observation_height, self.observation_width, 3), dtype=np.uint8)
self._black_frame: np.ndarray = np.zeros(
(self.observation_height, self.observation_width, 3), dtype=np.uint8
)
image_spaces = {
cam: spaces.Box(
+1 -1
View File
@@ -373,7 +373,7 @@ class VLABenchEnv(gym.Env):
if action.shape[0] != 7:
# Unknown layout — fall back to zero-pad so the sim doesn't crash.
padded = np.zeros(ctrl_dim, dtype=np.float64)
padded: np.ndarray = np.zeros(ctrl_dim, dtype=np.float64)
padded[: min(action.shape[0], ctrl_dim)] = action[:ctrl_dim]
return padded
+18
View File
@@ -122,6 +122,9 @@ MODEL_ENCODING_TABLE = {
"xm430-w350": X_SERIES_ENCODINGS_TABLE,
"xm540-w270": X_SERIES_ENCODINGS_TABLE,
"xc430-w150": X_SERIES_ENCODINGS_TABLE,
"xh540-w150": X_SERIES_ENCODINGS_TABLE,
"xc330-t288": X_SERIES_ENCODINGS_TABLE,
"xc330-t181": X_SERIES_ENCODINGS_TABLE,
}
# {model: model_resolution}
@@ -134,6 +137,9 @@ MODEL_RESOLUTION = {
"xm430-w350": 4096,
"xm540-w270": 4096,
"xc430-w150": 4096,
"xh540-w150": 4096,
"xc330-t288": 4096,
"xc330-t181": 4096,
}
# {model: model_number}
@@ -145,6 +151,9 @@ MODEL_NUMBER_TABLE = {
"xm430-w350": 1020,
"xm540-w270": 1120,
"xc430-w150": 1070,
"xh540-w150": 1110,
"xc330-t288": 1220,
"xc330-t181": 1210,
}
# {model: available_operating_modes}
@@ -156,6 +165,9 @@ MODEL_OPERATING_MODES = {
"xm430-w350": [0, 1, 3, 4, 5, 16],
"xm540-w270": [0, 1, 3, 4, 5, 16],
"xc430-w150": [1, 3, 4, 16],
"xh540-w150": [0, 1, 3, 4, 5, 16],
"xc330-t288": [0, 1, 3, 4, 5, 16],
"xc330-t181": [0, 1, 3, 4, 5, 16],
}
MODEL_CONTROL_TABLE = {
@@ -166,6 +178,9 @@ MODEL_CONTROL_TABLE = {
"xm430-w350": X_SERIES_CONTROL_TABLE,
"xm540-w270": X_SERIES_CONTROL_TABLE,
"xc430-w150": X_SERIES_CONTROL_TABLE,
"xh540-w150": X_SERIES_CONTROL_TABLE,
"xc330-t288": X_SERIES_CONTROL_TABLE,
"xc330-t181": X_SERIES_CONTROL_TABLE,
}
MODEL_BAUDRATE_TABLE = {
@@ -176,6 +191,9 @@ MODEL_BAUDRATE_TABLE = {
"xm430-w350": X_SERIES_BAUDRATE_TABLE,
"xm540-w270": X_SERIES_BAUDRATE_TABLE,
"xc430-w150": X_SERIES_BAUDRATE_TABLE,
"xh540-w150": X_SERIES_BAUDRATE_TABLE,
"xc330-t288": X_SERIES_BAUDRATE_TABLE,
"xc330-t181": X_SERIES_BAUDRATE_TABLE,
}
AVAILABLE_BAUDRATES = [
@@ -42,6 +42,9 @@ class Evo1Policy(PreTrainedPolicy):
config_class = Evo1Config
name = "evo1"
def supports_rtc(self) -> bool:
return True
def __init__(self, config: Evo1Config, *, vlm_hub_kwargs: dict | None = None, **kwargs):
super().__init__(config)
config.validate_features()
+35 -5
View File
@@ -302,6 +302,33 @@ def _pad_evo1_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(
config: Evo1Config,
preprocessor: PolicyProcessorPipeline,
@@ -309,16 +336,19 @@ def reconcile_evo1_processors(
) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
"""Reconcile checkpoint-loaded pipelines with the current EVO1 config.
Two 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
action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`). This
restores the converter and rebuilds the action step from the current config so those overrides
take effect.
Three things cannot be restored from a serialized pipeline alone: the EVO1 batch converter
(converters are plain functions and are never serialized), eval-time CLI overrides of the
action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`), and the
(un)normalizer stats/features when the generic override path injects raw, unpadded dataset
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
# non-observation extras (embodiment_id, state_mask, custom task fields) needed by EVO1.
preprocessor.to_transition = evo1_batch_to_transition
_refresh_evo1_normalization_steps(config, preprocessor, postprocessor)
action_step = Evo1ActionProcessorStep(
action_dim=_evo1_action_dim(config),
binarize_gripper=config.binarize_gripper,
+19 -3
View File
@@ -44,12 +44,19 @@ from lerobot.utils.constants import (
POLICY_PREPROCESSOR_DEFAULT_NAME,
)
from lerobot.utils.feature_utils import dataset_to_policy_features
from lerobot.utils.import_utils import _peft_available, require_package
from .evo1.configuration_evo1 import Evo1Config
from .groot.configuration_groot import GrootConfig
from .pretrained import PreTrainedPolicy
from .utils import validate_visual_features_consistency
if TYPE_CHECKING or _peft_available:
from peft import PeftConfig, PeftModel
else:
PeftConfig = None
PeftModel = None
def _reconnect_relative_absolute_steps(
preprocessor: PolicyProcessorPipeline, postprocessor: PolicyProcessorPipeline
@@ -177,6 +184,7 @@ def make_pre_post_processors(
return make_groot_pre_post_processors_from_pretrained(
config=policy_cfg,
pretrained_path=pretrained_path,
revision=pretrained_revision,
dataset_stats=kwargs.get("dataset_stats"),
dataset_meta=kwargs.get("dataset_meta"),
preprocessor_overrides=kwargs.get("preprocessor_overrides"),
@@ -333,12 +341,15 @@ def make_policy(
# Load a pretrained PEFT model on top of the policy. The pretrained path points to the folder/repo
# of the adapter and the adapter's config contains the path to the base policy. So we need the
# adapter config first, then load the correct policy and then apply PEFT.
from peft import PeftConfig, PeftModel
require_package("peft", extra="peft")
logging.info("Loading policy's PEFT adapter.")
peft_pretrained_path = str(cfg.pretrained_path)
peft_config = PeftConfig.from_pretrained(peft_pretrained_path)
peft_config = PeftConfig.from_pretrained(
peft_pretrained_path,
revision=cfg.pretrained_revision,
)
kwargs["pretrained_name_or_path"] = peft_config.base_model_name_or_path
if not kwargs["pretrained_name_or_path"]:
@@ -349,9 +360,14 @@ def make_policy(
"the adapter was trained."
)
kwargs["revision"] = peft_config.revision
policy = policy_cls.from_pretrained(**kwargs)
policy = PeftModel.from_pretrained(
policy, peft_pretrained_path, config=peft_config, is_trainable=True
policy,
peft_pretrained_path,
config=peft_config,
revision=cfg.pretrained_revision,
is_trainable=True,
)
else:
@@ -37,13 +37,19 @@ def is_image_feature(key: str) -> bool:
@dataclass
class ConcurrencyConfig:
"""Configuration for the concurrency of the actor and learner.
Possible values are:
- "threads": Use threads for the actor and learner.
- "processes": Use processes for the actor and learner.
``multiprocessing_context`` selects the process-wide start method when
processes are used. Set it to ``None`` to preserve Python's default or a
method already selected by the embedding application.
"""
actor: str = "threads"
learner: str = "threads"
multiprocessing_context: str | None = "spawn"
@dataclass
@@ -68,6 +68,9 @@ class GrootPolicy(PreTrainedPolicy):
name = "groot"
config_class = GrootConfig
def supports_rtc(self) -> bool:
return True
def __init__(self, config: GrootConfig, **kwargs):
"""Initialize Groot policy wrapper."""
require_package("transformers", extra="groot")
@@ -475,6 +475,7 @@ def make_groot_pre_post_processors_from_pretrained(
config: GrootConfig,
pretrained_path: str,
*,
revision: str | None = None,
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
dataset_meta: Any | None = None,
preprocessor_overrides: dict[str, Any] | None = None,
@@ -511,6 +512,7 @@ def make_groot_pre_post_processors_from_pretrained(
preprocessor, postprocessor = _load_groot_processor_pipelines(
pretrained_path,
revision=revision,
preprocessor_overrides=preprocessor_overrides,
postprocessor_overrides=postprocessor_overrides,
preprocessor_config_filename=preprocessor_config_filename,
@@ -526,6 +528,7 @@ def make_groot_pre_post_processors_from_pretrained(
def _load_groot_processor_pipelines(
pretrained_path: str,
*,
revision: str | None,
preprocessor_overrides: dict[str, Any],
postprocessor_overrides: dict[str, Any],
preprocessor_config_filename: str,
@@ -540,6 +543,7 @@ def _load_groot_processor_pipelines(
preprocessor = PolicyProcessorPipeline.from_pretrained(
pretrained_model_name_or_path=pretrained_path,
config_filename=preprocessor_config_filename,
revision=revision,
overrides=preprocessor_overrides,
to_transition=batch_to_transition,
to_output=transition_to_batch,
@@ -547,6 +551,7 @@ def _load_groot_processor_pipelines(
postprocessor = PolicyProcessorPipeline.from_pretrained(
pretrained_model_name_or_path=pretrained_path,
config_filename=postprocessor_config_filename,
revision=revision,
overrides=postprocessor_overrides,
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
@@ -43,11 +43,22 @@ from torch.distributions import Beta
from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.utils.constants import ACTION
from lerobot.utils.import_utils import _scipy_available, _transformers_available, require_package
from lerobot.utils.import_utils import (
_peft_available,
_scipy_available,
_transformers_available,
require_package,
)
from ..rtc.modeling_rtc import RTCProcessor
from .configuration_molmoact2 import MolmoAct2Config
if TYPE_CHECKING or _peft_available:
from peft import LoraConfig, get_peft_model
else:
LoraConfig = None
get_peft_model = None
logger = logging.getLogger(__name__)
@@ -509,6 +520,9 @@ class MolmoAct2Policy(PreTrainedPolicy):
config_class = MolmoAct2Config
name = "molmoact2"
def supports_rtc(self) -> bool:
return self.config.inference_action_mode == "continuous"
def __init__(
self,
config: MolmoAct2Config,
@@ -1731,13 +1745,11 @@ class MolmoAct2Policy(PreTrainedPolicy):
def _build_inner_lora_config(self):
require_package("peft", extra="molmoact2")
from peft import LoraConfig
return LoraConfig(**self._get_inner_peft_targets())
def _apply_lora_adapters(self) -> None:
require_package("peft", extra="molmoact2")
from peft import get_peft_model
peft_config = self._build_inner_lora_config()
self._validate_peft_config(peft_config)
+3
View File
@@ -749,6 +749,9 @@ class PI0Policy(PreTrainedPolicy):
config_class = PI0Config
name = "pi0"
def supports_rtc(self) -> bool:
return True
def __init__(
self,
config: PI0Config,
@@ -714,6 +714,9 @@ class PI05Policy(PreTrainedPolicy):
config_class = PI05Config
name = "pi05"
def supports_rtc(self) -> bool:
return True
def __init__(
self,
config: PI05Config,
+17 -5
View File
@@ -34,14 +34,22 @@ from lerobot.configs import PreTrainedConfig
from lerobot.configs.train import TrainPipelineConfig
from lerobot.utils.device_utils import resolve_safetensors_device
from lerobot.utils.hub import HubMixin
from lerobot.utils.import_utils import _peft_available, require_package
from .utils import log_model_loading_keys
T = TypeVar("T", bound="PreTrainedPolicy")
if TYPE_CHECKING or _peft_available:
from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType, get_peft_model
else:
PEFT_TYPE_TO_CONFIG_MAPPING = None
PeftType = None
get_peft_model = None
if TYPE_CHECKING:
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
T = TypeVar("T", bound="PreTrainedPolicy")
def _build_card_context(
cfg: TrainPipelineConfig | None,
@@ -241,6 +249,10 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
"""
raise NotImplementedError
def supports_rtc(self) -> bool:
"""Whether this policy implements Real-Time Chunking inference semantics."""
return False
# TODO(aliberts, rcadene): split into 'forward' and 'compute_loss'?
@abc.abstractmethod
def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict | None]:
@@ -384,7 +396,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
peft_cli_overrides: Optional dict of CLI overrides (method_type, target_modules, r, etc.)
These are merged with policy defaults to build the final config.
"""
from peft import get_peft_model
require_package("peft", extra="peft")
# If user provided a complete config, use it directly (with overrides)
if peft_config is not None:
@@ -455,7 +467,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
Returns:
Preprocessed dict with renamed keys and init_type mapped to method-specific key.
"""
from peft import PeftType
require_package("peft", extra="peft")
cli_overrides = cli_overrides.copy()
@@ -480,7 +492,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
def _build_peft_config(self, cli_overrides: dict):
"""Build a PEFT config from policy defaults and CLI overrides."""
from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType
require_package("peft", extra="peft")
# Determine PEFT method type (default to LORA)
method_type_str = cli_overrides.get("method_type") or "lora"
@@ -507,7 +519,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
def _apply_peft_cli_overrides(self, peft_config, cli_overrides: dict):
"""Apply CLI overrides to an existing PEFT config."""
from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType
require_package("peft", extra="peft")
# Get method type from existing config or CLI override
method_type_str = cli_overrides.get("method_type")
@@ -145,6 +145,9 @@ class SmolVLAPolicy(PreTrainedPolicy):
config_class = SmolVLAConfig
name = "smolvla"
def supports_rtc(self) -> bool:
return True
def __init__(
self,
config: SmolVLAConfig,
@@ -168,14 +168,23 @@ class SmolVLMWithExpertModel(nn.Module):
last_layers.append(self.num_vlm_layers - 2)
frozen_layers = [
"lm_head",
"text_model.model.norm.weight",
"text_model.norm.weight",
]
for layer in last_layers:
frozen_layers.append(f"text_model.model.layers.{layer}.")
frozen_layers.append(f"text_model.layers.{layer}.")
unmatched_patterns = set(frozen_layers)
for name, params in self.vlm.named_parameters():
if any(k in name for k in frozen_layers):
matched_patterns = [k for k in frozen_layers if k in name]
if matched_patterns:
params.requires_grad = False
unmatched_patterns.difference_update(matched_patterns)
if unmatched_patterns:
raise RuntimeError(
"Some frozen layer patterns matched no VLM parameters, so the corresponding layers "
"would silently remain trainable (parameter naming may have changed in transformers): "
f"{sorted(unmatched_patterns)}"
)
# To avoid unused params issue with distributed training
for name, params in self.lm_expert.named_parameters():
if "lm_head" in name:
@@ -132,10 +132,20 @@ class MapDeltaActionToRobotActionStep(RobotActionProcessorStep):
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
for axis in ["x", "y", "z", "gripper"]:
for axis in ["x", "y", "z"]:
features[PipelineFeatureType.ACTION].pop(f"delta_{axis}", None)
features[PipelineFeatureType.ACTION].pop("gripper", None)
for feat in ["enabled", "target_x", "target_y", "target_z", "target_wx", "target_wy", "target_wz"]:
for feat in [
"enabled",
"target_x",
"target_y",
"target_z",
"target_wx",
"target_wy",
"target_wz",
"gripper_vel",
]:
features[PipelineFeatureType.ACTION][f"{feat}"] = PolicyFeature(
type=FeatureType.ACTION, shape=(1,)
)
+18 -4
View File
@@ -713,6 +713,8 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
ProcessorMigrationError: If the model requires migration to processor format.
"""
model_id = str(pretrained_model_name_or_path)
model_path = Path(model_id)
is_local_source = model_path.is_dir() or model_path.is_file()
hub_download_kwargs = {
"force_download": force_download,
"resume_download": resume_download,
@@ -731,7 +733,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
# 3. Build steps with overrides
steps, validated_overrides = cls._build_steps_with_overrides(
loaded_config, overrides or {}, model_id, base_path, hub_download_kwargs
loaded_config, overrides or {}, model_id, base_path, hub_download_kwargs, is_local_source
)
# 4. Validate that all overrides were used
@@ -921,6 +923,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
model_id: str,
base_path: Path | None,
hub_download_kwargs: dict[str, Any],
is_local_source: bool = False,
) -> tuple[list[ProcessorStep], set[str]]:
"""Build all processor steps with overrides and state loading.
@@ -944,7 +947,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
3. **State Loading** (via _load_step_state):
- **If step has "state_file"**: Load tensor state from .safetensors
- **Local first**: Check base_path/state_file.safetensors
- **Hub fallback**: Download state file if not found locally
- **Hub fallback**: Download state file if the pipeline was loaded from the Hub
- **Optional**: Only load if step has load_state_dict method
4. **Override Tracking**:
@@ -962,6 +965,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
model_id: The model identifier (needed for Hub state file downloads)
base_path: Local directory path for finding state files
hub_download_kwargs: Parameters for hf_hub_download (tokens, cache, etc.)
is_local_source: Whether model_id resolved to a local directory or config file.
Returns:
Tuple of (instantiated_steps_list, unused_override_keys)
@@ -975,7 +979,9 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
steps, remaining_override_keys = cls._build_steps_from_config(loaded_config, overrides)
for step_instance, step_entry in zip(steps, loaded_config["steps"], strict=True):
cls._load_step_state(step_instance, step_entry, model_id, base_path, hub_download_kwargs)
cls._load_step_state(
step_instance, step_entry, model_id, base_path, hub_download_kwargs, is_local_source
)
return steps, remaining_override_keys
@@ -1139,6 +1145,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
model_id: str,
base_path: Path | None,
hub_download_kwargs: dict[str, Any],
is_local_source: bool = False,
) -> None:
"""Load state dictionary for a processor step if available.
@@ -1157,7 +1164,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
- **Use case**: Loading from local saved model directory
2. **Hub download fallback**: Download state file from repository
- **When triggered**: Local file not found or base_path is None
- **When triggered**: Local file not found and the pipeline source is a Hub repo
- **Process**: Use hf_hub_download with same parameters as config
- **Example**: Download "normalize_step_0.safetensors" from "user/repo"
- **Result**: Downloaded to local cache, path returned
@@ -1178,6 +1185,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
model_id: The model identifier (used for Hub downloads if needed)
base_path: Local directory path for finding state files (None for Hub-only)
hub_download_kwargs: Parameters for hf_hub_download (tokens, cache, etc.)
is_local_source: Whether model_id resolved to a local directory or config file.
Note:
This method modifies step_instance in-place and returns None.
@@ -1191,6 +1199,12 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
# Try local file first
if base_path and (base_path / state_filename).exists():
state_path = str(base_path / state_filename)
elif is_local_source:
state_path = base_path / state_filename if base_path else Path(state_filename)
raise FileNotFoundError(
f"State file '{state_filename}' was not found for local processor pipeline "
f"'{model_id}' at '{state_path}'."
)
else:
# Download from Hub
state_path = hf_hub_download(
+2 -4
View File
@@ -91,7 +91,7 @@ from lerobot.robots import so_follower # noqa: F401
from lerobot.teleoperators import gamepad, so_leader # noqa: F401
from lerobot.teleoperators.utils import TeleopEvents
from lerobot.utils.device_utils import get_safe_torch_device
from lerobot.utils.process import ProcessSignalHandler
from lerobot.utils.process import ProcessSignalHandler, ensure_multiprocessing_start_method
from lerobot.utils.random_utils import set_seed
from lerobot.utils.robot_utils import precise_sleep
from lerobot.utils.transition import (
@@ -124,9 +124,7 @@ def actor_cli(cfg: TrainRLServerPipelineConfig):
cfg.validate()
display_pid = False
if not use_threads(cfg):
import torch.multiprocessing as mp
mp.set_start_method("spawn")
ensure_multiprocessing_start_method(cfg.policy.concurrency.multiprocessing_context)
display_pid = True
# Create logs directory to ensure it exists
+5 -2
View File
@@ -16,6 +16,7 @@ from __future__ import annotations
import abc
import builtins
import json
import logging
import os
from dataclasses import dataclass, field
@@ -78,8 +79,10 @@ class RLAlgorithmConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
def _save_pretrained(self, save_directory: Path) -> None:
"""Serialize this config as ``config.json`` inside ``save_directory``."""
with open(save_directory / CONFIG_NAME, "w") as f, draccus.config_type("json"):
draccus.dump(self, f, indent=4)
# Encode against the base class so draccus includes the choice "type" key,
# which `from_pretrained` needs to resolve the concrete subclass.
with open(save_directory / CONFIG_NAME, "w") as f:
json.dump(draccus.encode(self, RLAlgorithmConfig), f, indent=4)
@classmethod
def from_pretrained(
+2 -2
View File
@@ -18,7 +18,7 @@ import functools
import threading
from collections.abc import Callable, Sequence
from contextlib import suppress
from typing import TypedDict
from typing import NotRequired, TypedDict
import torch
import torch.nn.functional as F # noqa: N812
@@ -36,7 +36,7 @@ class BatchTransition(TypedDict):
next_state: dict[str, torch.Tensor]
done: torch.Tensor
truncated: torch.Tensor
complementary_info: dict[str, torch.Tensor | float | int] | None = None
complementary_info: NotRequired[dict[str, torch.Tensor | float | int] | None]
def random_crop_vectorized(images: torch.Tensor, output_size: tuple) -> torch.Tensor:
+2 -4
View File
@@ -102,7 +102,7 @@ from lerobot.utils.constants import (
)
from lerobot.utils.device_utils import get_safe_torch_device
from lerobot.utils.io_utils import load_json, write_json
from lerobot.utils.process import ProcessSignalHandler
from lerobot.utils.process import ProcessSignalHandler, ensure_multiprocessing_start_method
from lerobot.utils.random_utils import set_seed
from lerobot.utils.utils import (
format_big_number,
@@ -123,9 +123,7 @@ def train_cli(cfg: TrainRLServerPipelineConfig):
# Fail fast with a friendly error if the optional ``hilserl`` extra is missing.
require_package("grpcio", extra="hilserl", import_name="grpc")
if not use_threads(cfg):
import torch.multiprocessing as mp
mp.set_start_method("spawn")
ensure_multiprocessing_start_method(cfg.policy.concurrency.multiprocessing_context)
# Use the job_name from the config
train(
@@ -323,6 +323,10 @@ class LeKiwiClient(Robot):
np.ndarray: the action sent to the motors, potentially clipped.
"""
# Action values may be torch tensors (e.g. replayed from a dataset) or numpy
# scalars; json.dumps only serializes Python primitives, so coerce each value to a
# plain float before sending.
action = {key: float(value) for key, value in action.items()}
self.zmq_cmd_socket.send_string(json.dumps(action)) # action is in motor space
# TODO(Steven): Remove the np conversion when it is possible to record a non-numpy array value
@@ -46,6 +46,12 @@ class SOFollowerConfig:
position_i_coefficient: int = 0
position_d_coefficient: int = 32
# Number of extra attempts when a `sync_read` of the motors fails. Feetech buses can occasionally
# return a corrupted status packet ("Incorrect status packet!"), especially when several joints move
# at once, which otherwise aborts the control loop. Retries are immediate (no sleep) and only happen on
# failure, so the steady-state read cost is unchanged.
num_read_retries: int = 2
@RobotConfig.register_subclass("so101_follower")
@RobotConfig.register_subclass("so100_follower")
@@ -510,10 +510,10 @@ class ForwardKinematicsJointsToEEAction(RobotActionProcessorStep):
# We only use the ee pose in the dataset, so we don't need the joint positions
for n in self.motor_names:
features[PipelineFeatureType.ACTION].pop(f"{n}.pos", None)
# We specify the dataset features of this step that we want to be stored in the dataset
# Store end-effector features as actions in the dataset schema
for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]:
features[PipelineFeatureType.ACTION][f"ee.{k}"] = PolicyFeature(
type=FeatureType.STATE, shape=(1,)
type=FeatureType.ACTION, shape=(1,)
)
return features
@@ -180,7 +180,7 @@ class SOFollower(Robot):
def get_observation(self) -> RobotObservation:
# Read arm position
start = time.perf_counter()
obs_dict = self.bus.sync_read("Present_Position")
obs_dict = self.bus.sync_read("Present_Position", num_retry=self.config.num_read_retries)
obs_dict = {f"{motor}.pos": val for motor, val in obs_dict.items()}
dt_ms = (time.perf_counter() - start) * 1e3
logger.debug(f"{self} read state: {dt_ms:.1f}ms")
@@ -221,7 +221,7 @@ class SOFollower(Robot):
# Cap goal position when too far away from present position.
# /!\ Slower fps expected due to reading from the follower.
if self.config.max_relative_target is not None:
present_pos = self.bus.sync_read("Present_Position")
present_pos = self.bus.sync_read("Present_Position", num_retry=self.config.num_read_retries)
goal_present_pos = {key: (g_pos, present_pos[key]) for key, g_pos in goal_pos.items()}
goal_pos = ensure_safe_goal_position(goal_present_pos, self.config.max_relative_target)
+11 -11
View File
@@ -149,15 +149,6 @@ class EpisodicStrategyConfig(RolloutStrategyConfig):
# Note that leader -> follower handover is only supported when the leader has `send_feedback` capability.
smooth_leader_to_follower_handover: bool = True
# Whether to turn on or off the smooth handover behavior at the start of the
# reset phase: the leader is driven to the follower position (actuated
# teleops, see `smooth_leader_to_follower_handover`), or the follower is
# slid to the teleop pose (non-actuated teleops). Disable for clutch-style
# teleoperators (e.g. VR controllers) that re-reference at the current robot
# pose on engage: the handover is already continuous there, and the blocking
# interpolation only delays the start of the reset phase.
smooth_handover: bool = True
@RolloutStrategyConfig.register_subclass("dagger")
@dataclass
@@ -335,8 +326,17 @@ class RolloutConfig:
policy_path = parser.get_path_arg("policy")
if policy_path:
cli_overrides = parser.get_cli_overrides("policy")
self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=cli_overrides)
yaml_overrides = parser.get_yaml_overrides("policy")
cli_overrides = parser.get_cli_overrides("policy") or []
policy_overrides = yaml_overrides + cli_overrides
pretrained_revision = parser.parse_arg("pretrained_revision", cli_overrides)
if pretrained_revision is None:
pretrained_revision = parser.parse_arg("pretrained_revision", yaml_overrides)
self.policy = PreTrainedConfig.from_pretrained(
policy_path,
revision=pretrained_revision,
cli_overrides=policy_overrides,
)
self.policy.pretrained_path = policy_path
if self.policy is None:
raise ValueError("--policy.path is required for rollout")
+59 -15
View File
@@ -24,10 +24,11 @@ from __future__ import annotations
import logging
from dataclasses import dataclass, field
from threading import Event
from typing import TYPE_CHECKING
import torch
from lerobot.configs import FeatureType
from lerobot.configs import FeatureType, PreTrainedConfig
from lerobot.datasets import (
LeRobotDataset,
aggregate_pipeline_dataset_features,
@@ -47,6 +48,7 @@ from lerobot.processor.relative_action_processor import RelativeActionsProcessor
from lerobot.robots import make_robot_from_config
from lerobot.teleoperators import Teleoperator, make_teleoperator_from_config
from lerobot.utils.feature_utils import combine_feature_dicts, hw_to_dataset_features
from lerobot.utils.import_utils import _peft_available, require_package
from .configs import BaseStrategyConfig, DAggerStrategyConfig, RolloutConfig
from .inference import (
@@ -55,8 +57,15 @@ from .inference import (
SyncInferenceConfig,
create_inference_engine,
)
from .inference.rtc import supports_rtc_inference
from .robot_wrapper import ThreadSafeRobot
if TYPE_CHECKING or _peft_available:
from peft import PeftConfig, PeftModel
else:
PeftConfig = None
PeftModel = None
logger = logging.getLogger(__name__)
@@ -159,6 +168,35 @@ class RolloutContext:
# ---------------------------------------------------------------------------
def _load_pretrained_policy(policy_config: PreTrainedConfig) -> PreTrainedPolicy:
"""Load policy weights, keeping adapter and base-model revisions independent."""
pretrained_revision = policy_config.pretrained_revision
policy_class = get_policy_class(policy_config.type)
if not policy_config.use_peft:
return policy_class.from_pretrained(
policy_config.pretrained_path,
config=policy_config,
revision=pretrained_revision,
)
require_package("peft", extra="peft")
peft_path = policy_config.pretrained_path
peft_config = PeftConfig.from_pretrained(peft_path, revision=pretrained_revision)
policy = policy_class.from_pretrained(
pretrained_name_or_path=peft_config.base_model_name_or_path,
config=policy_config,
revision=peft_config.revision,
)
return PeftModel.from_pretrained(
policy,
peft_path,
config=peft_config,
revision=pretrained_revision,
)
def build_rollout_context(
cfg: RolloutConfig,
shutdown_event: Event,
@@ -176,7 +214,6 @@ def build_rollout_context(
# --- 1. Policy (heavy I/O, but no hardware yet) -------------------
logger.info("Loading policy from '%s'...", cfg.policy.pretrained_path)
policy_config = cfg.policy
policy_class = get_policy_class(policy_config.type)
if hasattr(policy_config, "compile_model"):
policy_config.compile_model = cfg.use_torch_compile
@@ -187,19 +224,15 @@ def build_rollout_context(
"Please use `cpu` or `cuda` backend."
)
if policy_config.use_peft:
from peft import PeftConfig, PeftModel
peft_path = policy_config.pretrained_path
peft_config = PeftConfig.from_pretrained(peft_path)
policy = policy_class.from_pretrained(
pretrained_name_or_path=peft_config.base_model_name_or_path, config=policy_config
)
policy = PeftModel.from_pretrained(policy, peft_path, config=peft_config)
else:
policy = policy_class.from_pretrained(policy_config.pretrained_path, config=policy_config)
policy = _load_pretrained_policy(policy_config)
if is_rtc:
if not supports_rtc_inference(policy):
raise ValueError(
f"RTC inference is not supported by policy type '{policy_config.type}': "
"the policy must implement RTC semantics and predict_action_chunk must accept "
"inference_delay and prev_chunk_left_over. Use '--inference.type=sync' instead."
)
policy.config.rtc_config = cfg.inference.rtc
if hasattr(policy, "init_rtc_processor"):
policy.init_rtc_processor()
@@ -276,12 +309,22 @@ def build_rollout_context(
# ``observation_features`` values are either a tuple (camera shape) or the
# ``float`` type itself used as a sentinel for scalar motor features —
# see ``dict[str, type | tuple]`` annotation on ``Robot.observation_features``.
# Keep cameras (tuple) plus both joint-position (.pos) and base-velocity (.vel)
# scalar state features. LeKiwi's observation.state is 9-dim (6 arm .pos +
# x/y/theta.vel) and the policy was trained/normalized on all 9; the old .pos-only
# filter fed a 6-dim state into a 9-dim normalizer → RuntimeError (size 6 vs 9).
# Pure-arm robots have no .vel state keys, so this is a no-op for them.
observation_features_hw = {
k: v
for k, v in all_obs_features.items()
if isinstance(v, tuple) or (v is float and k.endswith(".pos"))
if isinstance(v, tuple) or (v is float and k.endswith((".pos", ".vel")))
}
action_features_hw = {k: v for k, v in robot.action_features.items() if k.endswith(".pos")}
# Keep both joint-position (.pos) and base-velocity (.vel) action features so
# mobile manipulators command the base too (e.g. LeKiwi: 6 arm .pos +
# x/y/theta.vel = 9-dim action). Pure-arm robots have no .vel keys, so this is
# a no-op for them. Without the .vel keys the base velocities are silently
# dropped from dataset_features[ACTION]/ordered_action_keys and the base never moves.
action_features_hw = {k: v for k, v in robot.action_features.items() if k.endswith((".pos", ".vel"))}
# The action side is always needed: sync inference reads action names from
# ``dataset_features[ACTION]`` to map policy tensors back to robot actions.
@@ -392,6 +435,7 @@ def build_rollout_context(
preprocessor, postprocessor = make_pre_post_processors(
policy_cfg=policy_config,
pretrained_path=cfg.policy.pretrained_path,
pretrained_revision=policy_config.pretrained_revision,
dataset_stats=dataset_stats,
preprocessor_overrides={
"device_processor": {"device": cfg.device},
+18
View File
@@ -22,6 +22,7 @@ way via ``notify_observation``.
from __future__ import annotations
import inspect
import logging
import math
import time
@@ -62,6 +63,23 @@ _RTC_JOIN_TIMEOUT_S: float = 3.0
# ---------------------------------------------------------------------------
def supports_rtc_inference(policy: PreTrainedPolicy) -> bool:
"""Whether a policy declares RTC support and accepts the RTC call shape."""
supports_rtc = getattr(policy, "supports_rtc", None)
if not callable(supports_rtc) or not supports_rtc():
return False
try:
inspect.signature(policy.predict_action_chunk).bind(
object(),
inference_delay=0,
prev_chunk_left_over=None,
)
except (TypeError, ValueError):
return False
return True
def _normalize_prev_actions_length(prev_actions: torch.Tensor, target_steps: int) -> torch.Tensor:
"""Pad or truncate RTC prefix actions to a fixed length for stable compiled inference."""
if prev_actions.ndim != 2:
@@ -143,10 +143,6 @@ class EpisodicStrategy(RolloutStrategy):
# position so the operator takes over without fighting the arm.
# For non-actuated teleops: slide the follower to the teleop's current
# pose instead, since the leader cannot be driven.
# Disabled entirely with --strategy.smooth_handover=false (useful for
# clutch-style teleops that re-reference at the current robot pose on
# engage).
if self.config.smooth_handover:
obs = robot.get_observation()
current_pos = {k: v for k, v in obs.items() if k.endswith(".pos")}
if (
@@ -36,6 +36,7 @@ python src/lerobot/scripts/augment_dataset_quantile_stats.py \
import argparse
import concurrent.futures
import logging
import os
from pathlib import Path
import numpy as np
@@ -52,6 +53,7 @@ from lerobot.datasets import (
get_feature_stats,
write_stats,
)
from lerobot.datasets.compute_stats import sample_indices
from lerobot.utils.utils import init_logging
@@ -77,12 +79,14 @@ def has_quantile_stats(stats: dict[str, dict] | None, quantile_list_keys: list[s
return False
def process_single_episode(dataset: LeRobotDataset, episode_idx: int) -> dict:
def process_single_episode(dataset: LeRobotDataset, episode_idx: int, use_sampling: bool = True) -> dict:
"""Process a single episode and return its statistics.
Args:
dataset: The LeRobot dataset
episode_idx: Index of the episode to process
use_sampling: If True, sub-sample image/video frames per episode to bound
memory. If False, use every frame (exact, higher memory).
Returns:
Dictionary containing episode statistics
@@ -92,16 +96,31 @@ def process_single_episode(dataset: LeRobotDataset, episode_idx: int) -> dict:
start_idx = dataset.meta.episodes[episode_idx]["dataset_from_index"]
end_idx = dataset.meta.episodes[episode_idx]["dataset_to_index"]
collected_data: dict[str, list] = {}
for idx in range(start_idx, end_idx):
item = dataset[idx]
for key, value in item.items():
if key not in dataset.features:
continue
episode_len = end_idx - start_idx
if key not in collected_data:
collected_data[key] = []
collected_data[key].append(value)
# Images/video are the memory hog, so sub-sample those frames per episode;
# numeric columns are cheap, so read them in full (exact).
image_keys = [k for k in dataset.features if dataset.features[k]["dtype"] in ("image", "video")]
numeric_keys = [
k for k in dataset.features if dataset.features[k]["dtype"] not in ("image", "video", "string")
]
collected_data: dict[str, list] = {}
# Numeric features: every frame, read directly from the underlying table.
if numeric_keys:
numeric_cols = dataset.hf_dataset.select_columns(numeric_keys)[start_idx:end_idx]
for key in numeric_keys:
collected_data[key] = [torch.as_tensor(v) for v in numeric_cols[key]]
# Image/video features: decode only a sampled subset of frames.
if image_keys:
sampled_offsets = sample_indices(episode_len) if use_sampling else list(range(episode_len))
for offset in sampled_offsets:
item = dataset[start_idx + offset]
for key in image_keys:
if key in item:
collected_data.setdefault(key, []).append(item[key])
ep_stats = {}
for key, data_list in collected_data.items():
@@ -131,11 +150,13 @@ def process_single_episode(dataset: LeRobotDataset, episode_idx: int) -> dict:
return ep_stats
def compute_quantile_stats_for_dataset(dataset: LeRobotDataset) -> dict[str, dict]:
def compute_quantile_stats_for_dataset(dataset: LeRobotDataset, use_sampling: bool = True) -> dict[str, dict]:
"""Compute quantile statistics for all episodes in the dataset.
Args:
dataset: The LeRobot dataset to compute statistics for
use_sampling: If True, sub-sample image/video frames per episode to bound
memory. If False, use every frame (exact, higher memory).
Returns:
Dictionary containing aggregated statistics with quantiles
@@ -153,15 +174,15 @@ def compute_quantile_stats_for_dataset(dataset: LeRobotDataset) -> dict[str, dic
if has_videos:
logging.info("Dataset contains video keys - using sequential processing for thread safety")
for episode_idx in tqdm(range(dataset.num_episodes), desc="Processing episodes"):
ep_stats = process_single_episode(dataset, episode_idx)
ep_stats = process_single_episode(dataset, episode_idx, use_sampling)
episode_stats_list.append(ep_stats)
else:
logging.info("Dataset has no video keys - using parallel processing for better performance")
max_workers = min(dataset.num_episodes, 16)
max_workers = min(dataset.num_episodes, int(os.environ.get("LEROBOT_STATS_MAX_WORKERS", 16)))
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
future_to_episode = {
executor.submit(process_single_episode, dataset, episode_idx): episode_idx
executor.submit(process_single_episode, dataset, episode_idx, use_sampling): episode_idx
for episode_idx in range(dataset.num_episodes)
}
@@ -188,6 +209,7 @@ def augment_dataset_with_quantile_stats(
repo_id: str,
root: str | Path | None = None,
overwrite: bool = False,
use_sampling: bool = True,
) -> None:
"""Augment a dataset with quantile statistics if they are missing.
@@ -195,6 +217,8 @@ def augment_dataset_with_quantile_stats(
repo_id: Repository ID of the dataset
root: Local root directory for the dataset
overwrite: Overwrite existing quantile statistics if they already exist
use_sampling: If True, sub-sample image/video frames per episode to bound
memory. If False, use every frame (exact, higher memory).
"""
logging.info(f"Loading dataset: {repo_id}")
dataset = LeRobotDataset(
@@ -208,7 +232,7 @@ def augment_dataset_with_quantile_stats(
logging.info("Dataset does not contain quantile statistics. Computing them now...")
new_stats = compute_quantile_stats_for_dataset(dataset)
new_stats = compute_quantile_stats_for_dataset(dataset, use_sampling=use_sampling)
logging.info("Updating dataset metadata with new quantile statistics")
dataset.meta.stats = new_stats
@@ -248,6 +272,14 @@ def main():
action="store_true",
help="Overwrite existing quantile statistics if they already exist",
)
parser.add_argument(
"--no-sampling",
action="store_true",
help=(
"Compute stats over every frame (exact, higher memory). By default, "
"image/video frames are sub-sampled per episode to bound memory."
),
)
args = parser.parse_args()
root = Path(args.root) if args.root else None
@@ -258,6 +290,7 @@ def main():
repo_id=args.repo_id,
root=root,
overwrite=args.overwrite,
use_sampling=not args.no_sampling,
)
@@ -61,6 +61,7 @@ import pyarrow as pa
import tqdm
from datasets import Dataset, Features, Image
from huggingface_hub import HfApi, snapshot_download
from huggingface_hub.errors import RevisionNotFoundError
from requests import HTTPError
from lerobot.datasets import CODEBASE_VERSION, LeRobotDataset, aggregate_stats
@@ -93,6 +94,8 @@ from lerobot.datasets.video_utils import concatenate_video_files, get_video_dura
from lerobot.utils.constants import HF_LEROBOT_HOME
from lerobot.utils.utils import flatten_dict, init_logging
logger = logging.getLogger(__name__)
V21 = "v2.1"
V30 = "v3.0"
@@ -475,11 +478,11 @@ def convert_dataset(
# First check if the dataset already has a v3.0 version
if root is None and not force_conversion:
try:
print("Trying to download v3.0 version of the dataset from the hub...")
logger.info("Trying to download v3.0 version of the dataset from the hub...")
snapshot_download(repo_id, repo_type="dataset", revision=V30, local_dir=HF_LEROBOT_HOME / repo_id)
return
except Exception:
print("Dataset does not have an uploaded v3.0 version. Continuing with conversion.")
logger.info("Dataset does not have an uploaded v3.0 version. Continuing with conversion.")
# Set root based on whether local dataset path is provided
use_local_dataset = False
@@ -487,7 +490,7 @@ def convert_dataset(
if root.exists():
validate_local_dataset_version(root)
use_local_dataset = True
print(f"Using local dataset at {root}")
logger.info(f"Using local dataset at {root}")
old_root = root.parent / f"{root.name}_old"
new_root = root.parent / f"{root.name}_v30"
@@ -521,8 +524,8 @@ def convert_dataset(
hub_api = HfApi()
try:
hub_api.delete_tag(repo_id, tag=CODEBASE_VERSION, repo_type="dataset")
except HTTPError as e:
print(f"tag={CODEBASE_VERSION} probably doesn't exist. Skipping exception ({e})")
except (HTTPError, RevisionNotFoundError) as e:
logger.warning(f"tag={CODEBASE_VERSION} probably doesn't exist. Skipping exception ({e})")
pass
hub_api.delete_files(
delete_patterns=["data/chunk*/episode_*", "meta/*.jsonl", "videos/chunk*"],
+6 -7
View File
@@ -154,14 +154,14 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
repo_id = cfg.new_repo_id or cfg.repo_id
commit_message = cfg.push_commit_message or "Add steerable annotations (lerobot-annotate)"
api = HfApi()
print(f"[lerobot-annotate] creating/locating dataset repo {repo_id}...", flush=True)
logger.info(f"[lerobot-annotate] creating/locating dataset repo {repo_id}...")
api.create_repo(
repo_id=repo_id,
repo_type="dataset",
private=cfg.push_private,
exist_ok=True,
)
print(f"[lerobot-annotate] uploading {root} -> {repo_id}...", flush=True)
logger.info(f"[lerobot-annotate] uploading {root} -> {repo_id}...")
commit_info = api.upload_folder(
folder_path=str(root),
repo_id=repo_id,
@@ -172,7 +172,7 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
# at the source dataset; a fresh card is generated below instead.
ignore_patterns=[".annotate_staging/**", "**/.DS_Store", "README.md"],
)
print(f"[lerobot-annotate] uploaded to https://huggingface.co/datasets/{repo_id}", flush=True)
logger.info(f"[lerobot-annotate] uploaded to https://huggingface.co/datasets/{repo_id}")
dataset_info = load_info(root)
card = create_lerobot_dataset_card(dataset_info=dataset_info, license="apache-2.0", repo_id=repo_id)
@@ -200,14 +200,13 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
with suppress(RevisionNotFoundError):
api.delete_tag(repo_id, tag=version_tag, repo_type="dataset")
api.create_tag(**tag_kwargs)
print(f"[lerobot-annotate] tagged {repo_id} as {version_tag}", flush=True)
logger.info(f"[lerobot-annotate] tagged {repo_id} as {version_tag}")
except Exception as exc: # noqa: BLE001
print(
logger.warning(
f"[lerobot-annotate] WARNING: could not create tag {version_tag!r} on {repo_id}: {exc}. "
"Dataset is uploaded but ``LeRobotDataset`` won't be able to load it until it's tagged. "
"Run: from huggingface_hub import HfApi; "
f"HfApi().create_tag({repo_id!r}, tag={version_tag!r}, repo_type='dataset', exist_ok=True)",
flush=True,
f"HfApi().create_tag({repo_id!r}, tag={version_tag!r}, repo_type='dataset', exist_ok=True)"
)
+3 -1
View File
@@ -89,6 +89,8 @@ from lerobot.datasets import LeRobotDataset
from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS
from lerobot.utils.utils import init_logging
logger = logging.getLogger(__name__)
DEFAULT_FOXGLOVE_PORT = 8765
DEFAULT_RERUN_PORT = 9090
@@ -299,7 +301,7 @@ def visualize_dataset(
while True:
time.sleep(1)
except KeyboardInterrupt:
print("Ctrl-C received. Exiting.")
logger.info("Ctrl-C received. Exiting.")
def main():
+15 -2
View File
@@ -127,6 +127,12 @@ Modify tasks - set default task with overrides for specific episodes (WARNING: m
--operation.new_task "Default task" \
--operation.episode_tasks '{"5": "Special task for episode 5"}'
Modify tasks - replace existing task strings in-place (WARNING: modifies in-place):
lerobot-edit-dataset \
--repo_id lerobot/pusht \
--operation.type modify_tasks \
--operation.task_replacements '{"Pick up the red cube": "Lift the red cube"}'
Convert image dataset to video format and save locally:
lerobot-edit-dataset \
--repo_id lerobot/pusht_image \
@@ -303,6 +309,7 @@ class RemoveFeatureConfig(OperationConfig):
class ModifyTasksConfig(OperationConfig):
new_task: str | None = None
episode_tasks: dict[str, str] | None = None
task_replacements: dict[str, str] | None = None
@OperationConfig.register_subclass("convert_image_to_video")
@@ -551,9 +558,12 @@ def handle_modify_tasks(cfg: EditDatasetConfig) -> None:
new_task = cfg.operation.new_task
episode_tasks_raw = cfg.operation.episode_tasks
task_replacements = cfg.operation.task_replacements
if new_task is None and episode_tasks_raw is None:
raise ValueError("Must specify at least one of new_task or episode_tasks for modify_tasks operation")
if new_task is None and episode_tasks_raw is None and task_replacements is None:
raise ValueError(
"Must specify at least one of new_task, episode_tasks, or task_replacements for modify_tasks operation"
)
if cfg.new_repo_id is not None or cfg.new_root is not None:
logging.warning(
@@ -573,11 +583,14 @@ def handle_modify_tasks(cfg: EditDatasetConfig) -> None:
logging.info(f" Default task: '{new_task}'")
if episode_tasks:
logging.info(f" Episode-specific tasks: {episode_tasks}")
if task_replacements:
logging.info(f" Task replacements: {task_replacements}")
modified_dataset = modify_tasks(
dataset,
new_task=new_task,
episode_tasks=episode_tasks,
task_replacements=task_replacements,
)
logging.info(f"Dataset modified at {dataset.root}")
+31 -12
View File
@@ -62,7 +62,7 @@ from dataclasses import asdict
from functools import partial
from pathlib import Path
from pprint import pformat
from typing import Any, TypedDict
from typing import TYPE_CHECKING, Any, TypedDict
import einops
import gymnasium as gym
@@ -87,7 +87,7 @@ from lerobot.processor import PolicyProcessorPipeline
from lerobot.types import PolicyAction
from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_IMAGES, OBS_STR, REWARD
from lerobot.utils.device_utils import get_safe_torch_device
from lerobot.utils.import_utils import register_third_party_plugins
from lerobot.utils.import_utils import _peft_available, register_third_party_plugins, require_package
from lerobot.utils.io_utils import write_video
from lerobot.utils.random_utils import set_seed
from lerobot.utils.utils import (
@@ -95,6 +95,14 @@ from lerobot.utils.utils import (
inside_slurm,
)
if TYPE_CHECKING or _peft_available:
from peft import PeftModel
else:
PeftModel = None
logger = logging.getLogger(__name__)
def _env_features_to_dataset_features(env_features: dict) -> dict:
"""Convert EnvConfig.features to the dict format expected by LeRobotDataset.create()."""
@@ -444,15 +452,16 @@ def eval_policy(
exc = ValueError(
f"Policy of type 'PreTrainedPolicy' is expected, but type '{type(policy)}' was provided."
)
try:
from peft import PeftModel
if not _peft_available:
raise exc
require_package("peft", extra="peft")
if not isinstance(policy, PeftModel):
raise exc
except ImportError:
raise exc from None
start = time.time()
# Preserve the mode for direct callers. eval_policy_all scopes the mode
# around all tasks so parallel evaluations cannot race with each other.
was_training = policy.training
policy.eval()
# Determine how many batched rollouts we need to get n_episodes. Note that if n_episodes is not evenly
@@ -555,7 +564,7 @@ def eval_policy(
if seeds:
all_seeds.extend(seeds)
else:
all_seeds.append(None)
all_seeds.extend([None] * env.num_envs)
# FIXME: episode_data is either None or it doesn't exist
if return_episode_data:
@@ -674,6 +683,8 @@ def eval_policy(
if save_predicted_video:
info["predicted_video_paths"] = predicted_video_paths
policy.train(was_training)
return info
@@ -791,13 +802,13 @@ def eval_main(cfg: EvalPipelineConfig):
recording_repo_id=cfg.eval.recording_repo_id,
recording_private=cfg.eval.recording_private,
)
print("Overall Aggregated Metrics:")
print(info["overall"])
logger.info("Overall Aggregated Metrics:")
logger.info(info["overall"])
# Print per-suite stats
for task_group, task_group_info in info.items():
print(f"\nAggregated Metrics for {task_group}:")
print(task_group_info)
logger.info(f"\nAggregated Metrics for {task_group}:")
logger.info(task_group_info)
# Close all vec envs
close_envs(envs)
@@ -1010,6 +1021,12 @@ def eval_policy_all(
recording_private=recording_private,
)
# Set the shared policy's mode before launching any workers. Restoring it
# inside individual tasks would let one task enable training mode while
# another task is still evaluating.
was_training = policy.training
policy.eval()
try:
if max_parallel_tasks <= 1:
prefetch_thread: threading.Thread | None = None
for i, (task_group, task_id, env) in enumerate(tasks):
@@ -1044,6 +1061,8 @@ def eval_policy_all(
per_task_infos.append({"task_group": tg, "task_id": tid, "metrics": metrics})
finally:
env.close()
finally:
policy.train(was_training)
# compute aggregated metrics helper (robust to lists/scalars)
def _agg_from_list(xs):
+31 -39
View File
@@ -28,7 +28,6 @@ lerobot-find-cameras
# NOTE(Steven): macOS cameras sometimes report different FPS at init time, not an issue here as we don't specify FPS when opening the cameras, but the information displayed might not be truthful.
import argparse
import concurrent.futures
import logging
import time
from pathlib import Path
@@ -40,6 +39,7 @@ from PIL import Image
from lerobot.cameras import ColorMode
from lerobot.cameras.opencv import OpenCVCamera, OpenCVCameraConfig
from lerobot.cameras.realsense import RealSenseCamera, RealSenseCameraConfig
from lerobot.utils.utils import init_logging
logger = logging.getLogger(__name__)
@@ -132,7 +132,7 @@ def save_image(
camera_identifier: str | int,
images_dir: Path,
camera_type: str,
):
) -> None:
"""
Saves a single image to disk using Pillow. Handles color conversion if necessary.
"""
@@ -151,7 +151,7 @@ def save_image(
logger.error(f"Failed to save image for camera {camera_identifier} (type {camera_type}): {e}")
def create_camera_instance(cam_meta: dict[str, Any]) -> dict[str, Any] | None:
def create_camera_instance(cam_meta: dict[str, Any], *, warmup_s: int = 1) -> dict[str, Any] | None:
"""Create and connect to a camera instance based on metadata."""
cam_type = cam_meta.get("type")
cam_id = cam_meta.get("id")
@@ -164,12 +164,14 @@ def create_camera_instance(cam_meta: dict[str, Any]) -> dict[str, Any] | None:
cv_config = OpenCVCameraConfig(
index_or_path=cam_id,
color_mode=ColorMode.RGB,
warmup_s=warmup_s,
)
instance = OpenCVCamera(cv_config)
elif cam_type == "RealSense":
rs_config = RealSenseCameraConfig(
serial_number_or_name=cam_id,
color_mode=ColorMode.RGB,
warmup_s=warmup_s,
)
instance = RealSenseCamera(rs_config)
else:
@@ -187,9 +189,7 @@ def create_camera_instance(cam_meta: dict[str, Any]) -> dict[str, Any] | None:
return None
def process_camera_image(
cam_dict: dict[str, Any], output_dir: Path, current_time: float
) -> concurrent.futures.Future | None:
def process_camera_image(cam_dict: dict[str, Any], output_dir: Path, current_time: float) -> None:
"""Capture and process an image from a single camera."""
cam = cam_dict["instance"]
meta = cam_dict["meta"]
@@ -199,7 +199,7 @@ def process_camera_image(
try:
image_data = cam.read()
return save_image(
save_image(
image_data,
cam_id_str,
output_dir,
@@ -214,10 +214,9 @@ def process_camera_image(
return None
def cleanup_cameras(cameras_to_use: list[dict[str, Any]]):
def cleanup_camera(cam_dict: dict[str, Any]) -> None:
"""Disconnect all cameras."""
logger.info(f"Disconnecting {len(cameras_to_use)} cameras...")
for cam_dict in cameras_to_use:
logger.info(f"Disconnecting camera with ID {cam_dict['meta'].get('id')}...")
try:
if cam_dict["instance"] and cam_dict["instance"].is_connected:
cam_dict["instance"].disconnect()
@@ -229,6 +228,7 @@ def save_images_from_all_cameras(
output_dir: Path,
record_time_s: float = 2.0,
camera_type: str | None = None,
warmup_s: int = 1,
):
"""
Connects to detected cameras (optionally filtered by type) and saves images from each.
@@ -239,6 +239,7 @@ def save_images_from_all_cameras(
record_time_s: Duration in seconds to record images.
camera_type: Optional string to filter cameras ("realsense" or "opencv").
If None, uses all detected cameras.
warmup_s: Duration in seconds to warmup camera before recording images.
"""
output_dir.mkdir(parents=True, exist_ok=True)
logger.info(f"Saving images to {output_dir}")
@@ -248,47 +249,32 @@ def save_images_from_all_cameras(
logger.warning("No cameras detected matching the criteria. Cannot save images.")
return
cameras_to_use = []
for cam_meta in all_camera_metadata:
camera_instance = create_camera_instance(cam_meta)
if camera_instance:
cameras_to_use.append(camera_instance)
logger.info(
f"Starting image capture for {record_time_s} seconds from {len(all_camera_metadata)} cameras."
)
if not cameras_to_use:
logger.warning("No cameras could be connected. Aborting image save.")
return
logger.info(f"Starting image capture for {record_time_s} seconds from {len(cameras_to_use)} cameras.")
start_time = time.perf_counter()
with concurrent.futures.ThreadPoolExecutor(max_workers=len(cameras_to_use) * 2) as executor:
try:
for cam_meta in all_camera_metadata:
cam_dict = create_camera_instance(cam_meta, warmup_s=warmup_s)
if cam_dict is None:
continue
start_time = time.perf_counter()
while time.perf_counter() - start_time < record_time_s:
futures = []
current_capture_time = time.perf_counter()
for cam_dict in cameras_to_use:
future = process_camera_image(cam_dict, output_dir, current_capture_time)
if future:
futures.append(future)
if futures:
concurrent.futures.wait(futures)
process_camera_image(cam_dict, output_dir, current_capture_time)
cleanup_camera(cam_dict)
except KeyboardInterrupt:
logger.info("Capture interrupted by user.")
finally:
print("\nFinalizing image saving...")
executor.shutdown(wait=True)
cleanup_cameras(cameras_to_use)
print(f"Image capture finished. Images saved to {output_dir}")
def main():
init_logging()
parser = argparse.ArgumentParser(
description="Unified camera utility script for listing cameras and capturing images."
)
parser.add_argument(
"camera_type",
type=str,
@@ -306,8 +292,14 @@ def main():
parser.add_argument(
"--record-time-s",
type=float,
default=6.0,
help="Time duration to attempt capturing frames. Default: 6 seconds.",
default=2.0,
help="Time duration to attempt capturing frames. Default: 2 seconds.",
)
parser.add_argument(
"--warmup-s",
type=int,
default=1,
help="Time duration to warmup camera before attempting to capture frames. Default: 1 second.",
)
args = parser.parse_args()
save_images_from_all_cameras(**vars(args))
+3 -1
View File
@@ -453,9 +453,11 @@ def record(
encoder_queue_maxsize=cfg.dataset.encoder_queue_maxsize,
)
robot.connect()
# Connect the teleoperator before the robot so the robot isn't left idle (and possibly
# tripping a firmware watchdog) during teleop init. Matches lerobot_teleoperate.py.
if teleop is not None:
teleop.connect()
robot.connect()
listener, events = init_keyboard_listener()
+1
View File
@@ -61,6 +61,7 @@ from lerobot.robots import ( # noqa: F401
earthrover_mini_plus,
hope_jr,
koch_follower,
lekiwi,
make_robot_from_config,
omx_follower,
openarm_follower,
+1
View File
@@ -165,6 +165,7 @@ from lerobot.robots import ( # noqa: F401
earthrover_mini_plus,
hope_jr,
koch_follower,
lekiwi,
omx_follower,
openarm_follower,
reachy2,
+43 -23
View File
@@ -22,7 +22,8 @@ import dataclasses
import logging
import sys
import time
from contextlib import nullcontext
from collections.abc import Iterator
from contextlib import contextmanager, nullcontext
from pprint import pformat
from typing import TYPE_CHECKING, Any
@@ -57,7 +58,7 @@ from lerobot.optim.factory import make_optimizer_and_scheduler
from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors
from lerobot.rewards import make_reward_pre_post_processors
from lerobot.utils.collate import lerobot_collate_fn
from lerobot.utils.import_utils import register_third_party_plugins
from lerobot.utils.import_utils import _peft_available, register_third_party_plugins, require_package
from lerobot.utils.logging_utils import AverageMeter, MetricsTracker
from lerobot.utils.random_utils import set_seed
from lerobot.utils.utils import (
@@ -68,9 +69,38 @@ from lerobot.utils.utils import (
inside_slurm,
)
if TYPE_CHECKING or _peft_available:
from peft import PeftModel
else:
PeftModel = None
from .lerobot_eval import eval_policy_all
@contextmanager
def _make_eval_envs(cfg: TrainPipelineConfig) -> Iterator[dict[str, dict[int, Any]]]:
"""Create evaluation environments for one run and always dispose of them."""
envs = make_env(
cfg.env,
n_envs=cfg.eval.batch_size,
use_async_envs=cfg.eval.use_async_envs,
)
try:
yield envs
finally:
close_envs(envs)
def _dataloader_worker_kwargs(cfg: TrainPipelineConfig) -> dict[str, Any]:
"""Return worker-only DataLoader options, disabling them for single-process loading."""
workers_enabled = cfg.num_workers > 0
return {
"prefetch_factor": cfg.prefetch_factor if workers_enabled else None,
"persistent_workers": cfg.persistent_workers and workers_enabled,
"multiprocessing_context": cfg.dataloader_multiprocessing_context if workers_enabled else None,
}
def update_policy(
train_metrics: MetricsTracker,
policy: PreTrainedPolicy,
@@ -197,8 +227,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if cfg.job.is_remote:
return submit_to_hf(cfg)
from lerobot.utils.import_utils import require_package
require_package("accelerate", extra="training")
from accelerate import Accelerator
from accelerate.utils import DistributedDataParallelKwargs, DistributedType
@@ -267,14 +295,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if not is_main_process:
dataset, eval_dataset = make_train_eval_datasets(cfg)
# Create environment used for evaluating checkpoints during training on simulation data.
# On real-world data, no need to create an environment as evaluations are done outside train.py,
# using the eval.py instead, with gym_dora environment and dora-rs.
eval_env = None
if cfg.env_eval_freq > 0 and cfg.env is not None and is_main_process:
logging.info("Creating env")
eval_env = make_env(cfg.env, n_envs=cfg.eval.batch_size, use_async_envs=cfg.eval.use_async_envs)
if cfg.is_reward_model_training:
if is_main_process:
logging.info("Creating reward model")
@@ -302,7 +322,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if cfg.peft is not None:
if cfg.is_reward_model_training:
raise ValueError("PEFT is only supported for policy training. ")
from peft import PeftModel
require_package("peft", extra="peft")
if isinstance(policy, PeftModel):
logging.info("PEFT adapter already loaded from checkpoint, skipping wrap_with_peft.")
@@ -328,7 +348,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
preprocessor_overrides = {
"device_processor": {"device": device.type},
"normalizer_processor": {
"stats": dataset.meta.stats,
"features": {**policy.config.input_features, **policy.config.output_features},
"norm_map": policy.config.normalization_mapping,
},
@@ -336,11 +355,17 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
}
postprocessor_overrides = {
"unnormalizer_processor": {
"stats": dataset.meta.stats,
"features": policy.config.output_features,
"norm_map": policy.config.normalization_mapping,
},
}
# On resume, the checkpoint's saved processor stats are authoritative: they may have
# been adapted by the policy (e.g. EVO1 pads state/action stats to max_state_dim),
# and force-feeding raw dataset stats over them crashes normalization (#4006).
# This mirrors the `dataset_stats` kwarg above, which is also skipped on resume.
if not cfg.resume:
preprocessor_overrides["normalizer_processor"]["stats"] = dataset.meta.stats
postprocessor_overrides["unnormalizer_processor"]["stats"] = dataset.meta.stats
if getattr(active_cfg, "use_relative_actions", False):
preprocessor_overrides["relative_actions_processor"] = {
"enabled": True,
@@ -473,8 +498,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
pin_memory=device.type == "cuda",
drop_last=False,
collate_fn=collate_fn,
prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None,
persistent_workers=cfg.persistent_workers and cfg.num_workers > 0,
**_dataloader_worker_kwargs(cfg),
)
# Build eval dataloader if a held-out split exists
@@ -500,8 +524,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
pin_memory=device.type == "cuda",
drop_last=False,
collate_fn=eval_collate_fn,
prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None,
persistent_workers=cfg.persistent_workers and cfg.num_workers > 0,
**_dataloader_worker_kwargs(cfg),
)
# Prepare everything with accelerator
@@ -684,7 +707,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if is_main_process:
step_id = get_step_identifier(step, cfg.steps)
logging.info(f"Eval policy at step {step}")
with torch.no_grad(), accelerator.autocast():
with _make_eval_envs(cfg) as eval_env, torch.no_grad(), accelerator.autocast():
eval_info = eval_policy_all(
envs=eval_env, # dict[suite][task_id] -> vec_env
policy=accelerator.unwrap_model(policy),
@@ -732,9 +755,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if is_main_process:
progbar.close()
if eval_env:
close_envs(eval_env)
is_fsdp = accelerator.distributed_type == DistributedType.FSDP
model_state_dict = accelerator.get_state_dict(policy) if is_fsdp else None
if is_main_process:
+58 -56
View File
@@ -45,6 +45,7 @@ lerobot-train-tokenizer \
"""
import json
import logging
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING
@@ -63,6 +64,9 @@ else:
from lerobot.configs import NormalizationMode, parser
from lerobot.datasets import LeRobotDataset
from lerobot.utils.constants import ACTION, OBS_STATE
from lerobot.utils.utils import init_logging
logger = logging.getLogger(__name__)
@dataclass
@@ -274,11 +278,8 @@ def process_episode(args):
return action_chunks
except Exception as e:
print(f"Error processing episode {ep_idx}: {e}")
import traceback
traceback.print_exc()
except Exception:
logger.exception("Error processing episode %s", ep_idx)
return None
@@ -300,10 +301,10 @@ def train_fast_tokenizer(
Returns:
Trained FAST tokenizer
"""
print(f"Training FAST tokenizer on {len(action_chunks)} action chunks...")
print(f"Action chunk shape: {action_chunks.shape}")
print(f"Vocab size: {vocab_size}")
print(f"DCT scale: {scale}")
logger.info(f"Training FAST tokenizer on {len(action_chunks)} action chunks...")
logger.info(f"Action chunk shape: {action_chunks.shape}")
logger.info(f"Vocab size: {vocab_size}")
logger.info(f"DCT scale: {scale}")
# download the tokenizer source code (not pretrained weights)
# we'll train a new tokenizer on our own data
@@ -314,7 +315,7 @@ def train_fast_tokenizer(
# train the new tokenizer on our action data using .fit()
# this trains the BPE tokenizer on DCT coefficients
print("Training new tokenizer (this may take a few minutes)...")
logger.info("Training new tokenizer (this may take a few minutes)...")
tokenizer = base_tokenizer.fit(
action_data_list,
scale=scale,
@@ -322,21 +323,21 @@ def train_fast_tokenizer(
time_horizon=action_chunks.shape[1], # action_horizon
action_dim=action_chunks.shape[2], # encoded dimensions
)
print("✓ Tokenizer training complete!")
logger.info("✓ Tokenizer training complete!")
# validate it works
sample_chunk = action_chunks[0]
encoded = tokenizer(sample_chunk[None])[0]
if isinstance(encoded, list):
encoded = np.array(encoded)
print(f"Sample encoding: {len(encoded)} tokens for chunk shape {sample_chunk.shape}")
logger.info(f"Sample encoding: {len(encoded)} tokens for chunk shape {sample_chunk.shape}")
return tokenizer
def compute_compression_stats(tokenizer, action_chunks: np.ndarray):
"""Compute compression statistics."""
print("\nComputing compression statistics...")
logger.info("\nComputing compression statistics...")
# sample for stats (use max 1000 chunks for speed)
sample_size = min(1000, len(action_chunks))
@@ -366,12 +367,12 @@ def compute_compression_stats(tokenizer, action_chunks: np.ndarray):
"max_token_length": float(np.max(token_lengths)),
}
print("Compression Statistics:")
print(f" Average compression ratio: {stats['compression_ratio']:.2f}x")
print(f" Mean token length: {stats['mean_token_length']:.1f}")
print(f" P99 token length: {stats['p99_token_length']:.0f}")
print(f" Min token length: {stats['min_token_length']:.0f}")
print(f" Max token length: {stats['max_token_length']:.0f}")
logger.info("Compression Statistics:")
logger.info(f" Average compression ratio: {stats['compression_ratio']:.2f}x")
logger.info(f" Mean token length: {stats['mean_token_length']:.1f}")
logger.info(f" P99 token length: {stats['p99_token_length']:.0f}")
logger.info(f" Min token length: {stats['min_token_length']:.0f}")
logger.info(f" Max token length: {stats['max_token_length']:.0f}")
return stats
@@ -385,9 +386,9 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
cfg: TokenizerTrainingConfig dataclass with all configuration parameters
"""
# load dataset
print(f"Loading dataset: {cfg.repo_id}")
logger.info(f"Loading dataset: {cfg.repo_id}")
dataset = LeRobotDataset(repo_id=cfg.repo_id, root=cfg.root)
print(f"Dataset loaded: {dataset.num_episodes} episodes, {dataset.num_frames} frames")
logger.info(f"Dataset loaded: {dataset.num_episodes} episodes, {dataset.num_frames} frames")
# parse normalization mode
try:
@@ -397,7 +398,7 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
f"Invalid normalization_mode: {cfg.normalization_mode}. "
f"Must be one of: {', '.join([m.value for m in NormalizationMode])}"
) from err
print(f"Normalization mode: {norm_mode.value}")
logger.info(f"Normalization mode: {norm_mode.value}")
# parse encoded dimensions
encoded_dim_ranges = []
@@ -406,38 +407,38 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
encoded_dim_ranges.append((start, end))
total_encoded_dims = sum(end - start for start, end in encoded_dim_ranges)
print(f"Encoding {total_encoded_dims} dimensions: {cfg.encoded_dims}")
logger.info(f"Encoding {total_encoded_dims} dimensions: {cfg.encoded_dims}")
# parse relative dimensions
relative_dim_list = None
if cfg.relative_dims is not None and cfg.relative_dims.strip():
relative_dim_list = [int(d.strip()) for d in cfg.relative_dims.split(",")]
print(f"Relative dimensions: {relative_dim_list}")
logger.info(f"Relative dimensions: {relative_dim_list}")
else:
print("No relative dimensions specified")
logger.info("No relative dimensions specified")
print(f"Use relative transform: {cfg.use_relative_transform}")
logger.info(f"Use relative transform: {cfg.use_relative_transform}")
if cfg.use_relative_transform and (relative_dim_list is None or len(relative_dim_list) == 0):
print(
logger.warning(
"Warning: use_relative_transform=True but no relative_dims specified. "
"No relative transform will be applied."
)
print(f"Action horizon: {cfg.action_horizon}")
print(f"State key: {cfg.state_key}")
logger.info(f"Action horizon: {cfg.action_horizon}")
logger.info(f"State key: {cfg.state_key}")
# determine episodes to process
num_episodes = dataset.num_episodes
if cfg.max_episodes is not None:
num_episodes = min(cfg.max_episodes, num_episodes)
print(f"Processing {num_episodes} episodes...")
logger.info(f"Processing {num_episodes} episodes...")
# process episodes sequentially (to avoid pickling issues with dataset)
all_chunks = []
for ep_idx in range(num_episodes):
if ep_idx % 10 == 0:
print(f" Processing episode {ep_idx}/{num_episodes}...")
logger.info(f" Processing episode {ep_idx}/{num_episodes}...")
chunks = process_episode(
(
@@ -455,19 +456,19 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
# concatenate all chunks
all_chunks = np.concatenate(all_chunks, axis=0)
print(f"Collected {len(all_chunks)} action chunks")
logger.info(f"Collected {len(all_chunks)} action chunks")
# extract only encoded dimensions FIRST (before normalization)
encoded_chunks = []
for start, end in encoded_dim_ranges:
encoded_chunks.append(all_chunks[:, :, start:end])
encoded_chunks = np.concatenate(encoded_chunks, axis=-1) # [N, H, D_encoded]
print(f"Extracted {encoded_chunks.shape[-1]} encoded dimensions")
logger.info(f"Extracted {encoded_chunks.shape[-1]} encoded dimensions")
# apply normalization to encoded dimensions
print("\nBefore normalization - overall stats:")
print(f" Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
print(f" Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
logger.info("\nBefore normalization - overall stats:")
logger.info(f" Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
logger.info(f" Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
# get normalization stats from dataset
norm_stats = dataset.meta.stats
@@ -489,9 +490,9 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
encoded_stats[stat_name] = stat_array[encoded_dim_indices]
if encoded_stats:
print(f"\nNormalization stats for encoded dimensions (mode: {norm_mode.value}):")
logger.info(f"\nNormalization stats for encoded dimensions (mode: {norm_mode.value}):")
for stat_name, stat_values in encoded_stats.items():
print(
logger.info(
f" {stat_name}: shape={stat_values.shape}, "
f"range=[{np.min(stat_values):.4f}, {np.max(stat_values):.4f}]"
)
@@ -499,27 +500,27 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
# apply normalization based on mode
try:
encoded_chunks = apply_normalization(encoded_chunks, encoded_stats, norm_mode, eps=1e-8)
print(f"\nApplied {norm_mode.value} normalization")
logger.info(f"\nApplied {norm_mode.value} normalization")
except ValueError as e:
print(f"Warning: {e}. Using raw actions without normalization.")
logger.warning(f"Warning: {e}. Using raw actions without normalization.")
print("\nAfter normalization - overall stats:")
print(f" Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
print(f" Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
logger.info("\nAfter normalization - overall stats:")
logger.info(f" Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
logger.info(f" Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
print("\nPer-dimension stats (after normalization):")
logger.info("\nPer-dimension stats (after normalization):")
for d in range(encoded_chunks.shape[-1]):
dim_data = encoded_chunks[:, :, d]
print(
logger.info(
f" Dim {d}: min={np.min(dim_data):7.4f}, max={np.max(dim_data):7.4f}, "
f"mean={np.mean(dim_data):7.4f}, std={np.std(dim_data):7.4f}"
)
else:
print("Warning: Could not extract stats for encoded dimensions, using raw actions")
logger.warning("Warning: Could not extract stats for encoded dimensions, using raw actions")
else:
print("Warning: No normalization stats found in dataset, using raw actions")
logger.warning("Warning: No normalization stats found in dataset, using raw actions")
print(f"Encoded chunks shape: {encoded_chunks.shape}")
logger.info(f"Encoded chunks shape: {encoded_chunks.shape}")
# train FAST tokenizer
tokenizer = train_fast_tokenizer(
@@ -561,8 +562,8 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
with open(output_path / "metadata.json", "w") as f:
json.dump(metadata, f, indent=2)
print(f"\nSaved FAST tokenizer to {output_path}")
print(f"Metadata: {json.dumps(metadata, indent=2)}")
logger.info(f"\nSaved FAST tokenizer to {output_path}")
logger.info(f"Metadata: {json.dumps(metadata, indent=2)}")
# push to Hugging Face Hub if requested
if cfg.push_to_hub:
@@ -570,10 +571,10 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
hub_repo_id = cfg.hub_repo_id
if hub_repo_id is None:
hub_repo_id = output_path.name
print(f"\nNo hub_repo_id provided, using: {hub_repo_id}")
logger.info(f"\nNo hub_repo_id provided, using: {hub_repo_id}")
print(f"\nPushing tokenizer to Hugging Face Hub: {hub_repo_id}")
print(f" Private: {cfg.hub_private}")
logger.info(f"\nPushing tokenizer to Hugging Face Hub: {hub_repo_id}")
logger.info(f" Private: {cfg.hub_private}")
try:
# use the tokenizer's push_to_hub method
@@ -593,14 +594,15 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
commit_message="Upload tokenizer metadata",
)
print(f"Successfully pushed tokenizer to: https://huggingface.co/{hub_repo_id}")
logger.info(f"Successfully pushed tokenizer to: https://huggingface.co/{hub_repo_id}")
except Exception as e:
print(f"Error pushing to hub: {e}")
print(" Make sure you're logged in with `huggingface-cli login`")
logger.error(f"Error pushing to hub: {e}")
logger.error(" Make sure you're logged in with `huggingface-cli login`")
def main():
"""CLI entry point that parses arguments and runs the tokenizer training."""
init_logging()
train_tokenizer()
@@ -67,7 +67,15 @@ class BiSOLeader(BimanualMixin, Teleoperator):
@cached_property
def feedback_features(self) -> dict[str, type]:
return {}
# Bimanual teleop has feedback (can be actuated for handover).
# Return the same structure as action_features for consistency with left/right arms.
left_arm_features = self.left_arm.feedback_features
right_arm_features = self.right_arm.feedback_features
return {
**{f"left_{k}": v for k, v in left_arm_features.items()},
**{f"right_{k}": v for k, v in right_arm_features.items()},
}
def setup_motors(self) -> None:
self.left_arm.setup_motors()
@@ -87,6 +95,43 @@ class BiSOLeader(BimanualMixin, Teleoperator):
return action_dict
def enable_torque(self) -> None:
"""Enable torque on both leader arms for smooth handover."""
self.left_arm.enable_torque()
self.right_arm.enable_torque()
def disable_torque(self) -> None:
"""Disable torque on both leader arms to allow human control."""
self.left_arm.disable_torque()
self.right_arm.disable_torque()
@check_if_not_connected
def send_feedback(self, feedback: dict[str, float]) -> None:
# TODO: Implement force feedback
raise NotImplementedError
"""Route bimanual feedback to left and right arms with proper prefix stripping.
Receives feedback dict with keys like: left_shoulder_pan.pos, right_shoulder_pan.pos, ...
Splits and routes to each arm by removing the prefix.
This enables DAgger smooth handover: when transitioning from policy control to human
intervention, both leader arms are commanded to the follower's current pose to avoid
discontinuities.
"""
# Split feedback by arm prefix
left_feedback = {}
right_feedback = {}
for key, value in feedback.items():
if key.startswith("left_"):
# Strip "left_" prefix and pass to left arm
stripped_key = key[5:] # len("left_") == 5
left_feedback[stripped_key] = value
elif key.startswith("right_"):
# Strip "right_" prefix and pass to right arm
stripped_key = key[6:] # len("right_") == 6
right_feedback[stripped_key] = value
# Send to each arm
if left_feedback:
self.left_arm.send_feedback(left_feedback)
if right_feedback:
self.right_arm.send_feedback(right_feedback)
@@ -171,7 +171,13 @@ class IOSPhone(BasePhone, Teleoperator):
# HEBI provides orientation in w, x, y, z format.
# Scipy's Rotation expects x, y, z, w.
quat_xyzw = np.concatenate((ar_quat[1:], [ar_quat[0]])) # wxyz to xyzw
# ARKit can emit zero/NaN quaternions before tracking is ready or on a
# dropped packet. Rotation.from_quat now rejects those; degrade the same
# way as a missing pose so teleop stays alive mid-session.
try:
rot = Rotation.from_quat(quat_xyzw)
except ValueError:
return False, None, None, None
pos = ar_pos - rot.apply(self.config.camera_offset)
return True, pos, rot, pose
@@ -29,6 +29,12 @@ class SOLeaderConfig:
# Whether to use degrees for angles
use_degrees: bool = True
# Number of extra attempts when a `sync_read` of the motors fails. Feetech buses can occasionally
# return a corrupted status packet ("Incorrect status packet!"), especially when several joints move
# at once, which otherwise aborts the teleoperation loop. Retries are immediate (no sleep) and only
# happen on failure, so the steady-state read cost is unchanged.
num_read_retries: int = 2
@TeleoperatorConfig.register_subclass("so101_leader")
@TeleoperatorConfig.register_subclass("so100_leader")
@@ -145,7 +145,7 @@ class SOLeader(Teleoperator):
@check_if_not_connected
def get_action(self) -> dict[str, float]:
start = time.perf_counter()
action = self.bus.sync_read("Present_Position")
action = self.bus.sync_read("Present_Position", num_retry=self.config.num_read_retries)
action = {f"{motor}.pos": val for motor, val in action.items()}
dt_ms = (time.perf_counter() - start) * 1e3
logger.debug(f"{self} read action: {dt_ms:.1f}ms")
+7 -7
View File
@@ -41,7 +41,7 @@ class RandomSubsetApply(Transform):
def __init__(
self,
transforms: Sequence[Callable],
transforms: Sequence[Callable[..., Any]],
p: list[float] | None = None,
n_subset: int | None = None,
random_order: bool = False,
@@ -50,7 +50,7 @@ class RandomSubsetApply(Transform):
if not isinstance(transforms, Sequence):
raise TypeError("Argument transforms should be a sequence of callables")
if p is None:
p = [1] * len(transforms)
p = [1.0] * len(transforms)
elif len(p) != len(transforms):
raise ValueError(
f"Length of p doesn't match the number of transforms: {len(p)} != {len(transforms)}"
@@ -69,7 +69,7 @@ class RandomSubsetApply(Transform):
self.n_subset = n_subset
self.random_order = random_order
self.selected_transforms = None
self.selected_transforms: list[Callable[..., Any]] = []
def forward(self, *inputs: Any) -> Any:
needs_unpacking = len(inputs) > 1
@@ -119,7 +119,7 @@ class SharpnessJitter(Transform):
super().__init__()
self.sharpness = self._check_input(sharpness)
def _check_input(self, sharpness):
def _check_input(self, sharpness: float | Sequence[float]) -> tuple[float, float]:
if isinstance(sharpness, (int | float)):
if sharpness < 0:
raise ValueError("If sharpness is a single number, it must be non negative.")
@@ -215,7 +215,7 @@ class ImageTransformsConfig:
)
def make_transform_from_config(cfg: ImageTransformConfig):
def make_transform_from_config(cfg: ImageTransformConfig) -> Transform:
if cfg.type == "SharpnessJitter":
return SharpnessJitter(**cfg.kwargs)
@@ -236,8 +236,8 @@ class ImageTransforms(Transform):
super().__init__()
self._cfg = cfg
self.weights = []
self.transforms = {}
self.weights: list[float] = []
self.transforms: dict[str, Transform] = {}
for tf_name, tf_cfg in cfg.tfs.items():
if tf_cfg.weight <= 0.0:
continue
+13 -4
View File
@@ -37,16 +37,25 @@ def auto_select_torch_device() -> torch.device:
# TODO(Steven): Remove log. log shouldn't be an argument, this should be handled by the logger level
def get_safe_torch_device(try_device: str, log: bool = False) -> torch.device:
"""Given a string, return a torch.device with checks on whether the device is available."""
"""Given a string, return a torch.device with checks on whether the device is available.
Raises:
ValueError: If the requested device family is known but not available on
this machine (``AssertionError`` was previously used and is easy to
mistake for a programmer bug under ``python -O`` where asserts vanish).
"""
try_device = str(try_device)
if try_device.startswith("cuda"):
assert torch.cuda.is_available()
if not torch.cuda.is_available():
raise ValueError(f"Requested device {try_device!r} but CUDA is not available.")
device = torch.device(try_device)
elif try_device == "mps":
assert torch.backends.mps.is_available()
if not torch.backends.mps.is_available():
raise ValueError("Requested device 'mps' but MPS is not available.")
device = torch.device("mps")
elif try_device == "xpu":
assert torch.xpu.is_available()
if not torch.xpu.is_available():
raise ValueError("Requested device 'xpu' but XPU is not available.")
device = torch.device("xpu")
elif try_device == "cpu":
device = torch.device("cpu")
+5 -5
View File
@@ -32,21 +32,21 @@ def load_json(fpath: Path) -> Any:
Returns:
Any: The data loaded from the JSON file.
"""
with open(fpath) as f:
with open(fpath, encoding="utf-8") as f:
return json.load(f)
def write_json(data: dict, fpath: Path) -> None:
"""Write data to a JSON file.
def write_json(data: JsonLike, fpath: Path) -> None:
"""Write JSON-serializable data to a file.
Creates parent directories if they don't exist.
Args:
data (dict): The dictionary to write.
data: JSON-serializable data to write.
fpath (Path): The path to the output JSON file.
"""
fpath.parent.mkdir(exist_ok=True, parents=True)
with open(fpath, "w") as f:
with open(fpath, "w", encoding="utf-8") as f:
json.dump(data, f, indent=4, ensure_ascii=False)
+28
View File
@@ -16,11 +16,39 @@
# limitations under the License.
import logging
import multiprocessing
import os
import signal
import sys
def ensure_multiprocessing_start_method(start_method: str | None) -> None:
"""Set a multiprocessing start method once, or verify the existing method matches.
Passing ``None`` leaves Python's process-wide default untouched. This is useful
when LeRobot is embedded in an application that owns multiprocessing setup.
"""
if start_method is None:
return
available_methods = multiprocessing.get_all_start_methods()
if start_method not in available_methods:
raise ValueError(
f"Multiprocessing start method must be one of {available_methods} on this platform, "
f"got {start_method!r}."
)
current_method = multiprocessing.get_start_method(allow_none=True)
if current_method is None:
multiprocessing.set_start_method(start_method)
elif current_method != start_method:
raise RuntimeError(
f"Multiprocessing start method is already {current_method!r}; cannot change it to "
f"{start_method!r}. Set the configured multiprocessing context to null to keep the "
"application's existing method, or launch LeRobot in a fresh process."
)
class ProcessSignalHandler:
"""Utility class to attach graceful shutdown signal handlers.
+4
View File
@@ -85,6 +85,8 @@ def serialize_torch_rng_state() -> dict[str, torch.Tensor]:
torch_rng_state_dict = {"torch_rng_state": torch.get_rng_state()}
if torch.cuda.is_available():
torch_rng_state_dict["torch_cuda_rng_state"] = torch.cuda.get_rng_state()
if torch.backends.mps.is_available():
torch_rng_state_dict["torch_mps_rng_state"] = torch.mps.get_rng_state()
return torch_rng_state_dict
@@ -95,6 +97,8 @@ def deserialize_torch_rng_state(rng_state_dict: dict[str, torch.Tensor]) -> None
torch.set_rng_state(rng_state_dict["torch_rng_state"])
if torch.cuda.is_available() and "torch_cuda_rng_state" in rng_state_dict:
torch.cuda.set_rng_state(rng_state_dict["torch_cuda_rng_state"])
if torch.backends.mps.is_available() and "torch_mps_rng_state" in rng_state_dict:
torch.mps.set_rng_state(rng_state_dict["torch_mps_rng_state"])
def serialize_rng_state() -> dict[str, torch.Tensor]:
+4
View File
@@ -30,6 +30,10 @@ def precise_sleep(seconds: float, spin_threshold: float = 0.010, sleep_margin: f
"""
if seconds <= 0:
return
if spin_threshold < 0:
raise ValueError(f"spin_threshold must be >= 0, got {spin_threshold}")
if sleep_margin < 0:
raise ValueError(f"sleep_margin must be >= 0, got {sleep_margin}")
system = platform.system()
# On macOS and Windows the scheduler / sleep granularity can make
+5 -2
View File
@@ -29,9 +29,12 @@ class Rotation:
def __init__(self, quat: np.ndarray) -> None:
"""Initialize rotation from quaternion [x, y, z, w]."""
self._quat = np.asarray(quat, dtype=float)
# Normalize quaternion
if self._quat.shape != (4,):
raise ValueError(f"Quaternion must have shape (4,), got {self._quat.shape}")
# Normalize quaternion. Reject the zero vector — it has no orientation.
norm = np.linalg.norm(self._quat)
if norm > 0:
if norm <= 0.0 or not np.isfinite(norm):
raise ValueError(f"Quaternion must be a non-zero finite vector; got {self._quat} (norm={norm})")
self._quat = self._quat / norm
@classmethod
+2 -2
View File
@@ -14,7 +14,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TypedDict
from typing import NotRequired, TypedDict
import torch
@@ -28,7 +28,7 @@ class Transition(TypedDict):
next_state: dict[str, torch.Tensor]
done: bool
truncated: bool
complementary_info: dict[str, torch.Tensor | float | int] | None = None
complementary_info: NotRequired[dict[str, torch.Tensor | float | int] | None]
def move_transition_to_device(transition: Transition, device: str = "cpu") -> Transition:
+13 -9
View File
@@ -24,7 +24,6 @@ import sys
import time
from collections.abc import Iterator
from copy import copy, deepcopy
from datetime import datetime
from pathlib import Path
from statistics import mean
from typing import TYPE_CHECKING, Any
@@ -61,14 +60,16 @@ def init_logging(
accelerator: Optional Accelerator instance (for multi-GPU detection)
"""
def custom_format(record: logging.LogRecord) -> str:
dt = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
fnameline = f"{record.pathname}:{record.lineno}"
pid_str = f"[PID: {os.getpid()}] " if display_pid else ""
return f"{record.levelname} {pid_str}{dt} {fnameline[-15:]:>15} {record.getMessage()}"
class LeRobotFormatter(logging.Formatter):
def format(self, record: logging.LogRecord) -> str:
record.lerobot_location = f"{record.pathname}:{record.lineno}"[-15:]
record.lerobot_pid = f"[PID: {os.getpid()}] " if display_pid else ""
return super().format(record)
formatter = logging.Formatter()
formatter.format = custom_format
formatter = LeRobotFormatter(
"%(levelname)s %(lerobot_pid)s%(asctime)s %(lerobot_location)15s %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger()
logger.setLevel(logging.NOTSET)
@@ -133,10 +134,13 @@ def say(text: str, blocking: bool = False):
else:
raise RuntimeError("Unsupported operating system for text-to-speech.")
try:
if blocking:
subprocess.run(cmd, check=True)
subprocess.run(cmd, check=True, timeout=5)
else:
subprocess.Popen(cmd, creationflags=subprocess.CREATE_NO_WINDOW if system == "Windows" else 0)
except (FileNotFoundError, subprocess.TimeoutExpired) as e:
logging.warning("Text-to-speech command failed: %s | Error: %s", cmd, e)
def log_say(text: str, play_sounds: bool = True, blocking: bool = False):
+68 -1
View File
@@ -20,7 +20,7 @@
# ```
from pathlib import Path
from unittest.mock import patch
from unittest.mock import MagicMock, patch
import cv2
import numpy as np
@@ -123,6 +123,73 @@ def test_invalid_width_connect():
camera.connect(warmup=False)
def test_connect_cleans_up_after_settings_failure_and_allows_retry():
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH, warmup_s=0)
camera = OpenCVCamera(config)
opened_captures = []
def fail_settings():
opened_captures.append(camera.videocapture)
raise RuntimeError("settings failed")
with (
patch.object(camera, "_configure_capture_settings", side_effect=fail_settings),
pytest.raises(RuntimeError, match="settings failed"),
):
camera.connect(warmup=False)
assert camera.videocapture is None
assert camera.thread is None
assert not camera.is_connected
assert opened_captures[0] is not None
assert not opened_captures[0].isOpened()
camera.connect(warmup=False)
assert camera.is_connected
camera.disconnect()
def test_connect_cleans_up_after_warmup_failure_and_allows_retry():
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH, warmup_s=1)
camera = OpenCVCamera(config)
read_threads = []
def fail_warmup(*_args, **_kwargs):
read_threads.append(camera.thread)
raise TimeoutError("no frame")
with (
patch.object(camera, "async_read", side_effect=fail_warmup),
pytest.raises(TimeoutError, match="no frame"),
):
camera.connect()
assert camera.videocapture is None
assert camera.thread is None
assert not camera.is_connected
assert read_threads[0] is not None
assert not read_threads[0].is_alive()
camera.connect(warmup=False)
assert camera.is_connected
camera.disconnect()
def test_find_cameras_releases_unopened_handles():
module_path = OpenCVCamera.__module__
unopened_capture = MagicMock()
unopened_capture.isOpened.return_value = False
with (
patch(f"{module_path}.platform.system", return_value="Darwin"),
patch(f"{module_path}.MAX_OPENCV_INDEX", 1),
patch(f"{module_path}.cv2.VideoCapture", return_value=unopened_capture),
):
assert OpenCVCamera.find_cameras() == []
unopened_capture.release.assert_called_once_with()
@pytest.mark.parametrize("index_or_path", TEST_IMAGE_PATHS, ids=TEST_IMAGE_SIZES)
def test_read(index_or_path):
config = OpenCVCameraConfig(index_or_path=index_or_path, warmup_s=0)
+259 -1
View File
@@ -20,7 +20,7 @@
# ```
from pathlib import Path
from unittest.mock import patch
from unittest.mock import MagicMock, patch
import numpy as np
import pytest
@@ -30,6 +30,8 @@ from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnected
pytest.importorskip("pyrealsense2")
import pyrealsense2 as rs
from lerobot.cameras.realsense import RealSenseCamera, RealSenseCameraConfig
TEST_ARTIFACTS_DIR = Path(__file__).parent.parent / "artifacts" / "cameras"
@@ -61,6 +63,17 @@ def test_abc_implementation():
_ = RealSenseCamera(config)
@pytest.mark.parametrize("option", ["exposure", "gain", "white_balance"])
def test_manual_color_option_requires_rgb(option):
with pytest.raises(ValueError, match="use_rgb=True"):
RealSenseCameraConfig(
serial_number_or_name="042",
use_rgb=False,
use_depth=True,
**{option: 100},
)
def test_connect():
config = RealSenseCameraConfig(serial_number_or_name="042", warmup_s=0)
@@ -83,6 +96,27 @@ def test_connect_invalid_camera_path(patch_realsense):
camera.connect(warmup=False)
def test_connect_cleans_up_when_sensor_configuration_fails():
config = RealSenseCameraConfig(serial_number_or_name="042", exposure=120)
camera = RealSenseCamera(config)
pipeline = MagicMock()
pipeline.start.return_value = MagicMock()
with (
patch("lerobot.cameras.realsense.camera_realsense.rs.pipeline", return_value=pipeline),
patch.object(camera, "_configure_rs_pipeline_config"),
patch.object(camera, "_configure_capture_settings"),
patch.object(camera, "_configure_sensor_options", side_effect=ValueError("invalid exposure")),
pytest.raises(ValueError, match="invalid exposure"),
):
camera.connect(warmup=False)
pipeline.stop.assert_called_once_with()
assert camera.rs_pipeline is None
assert camera.rs_profile is None
assert not camera.is_connected
def test_invalid_width_connect():
config = RealSenseCameraConfig(serial_number_or_name="042", width=99999, height=480, fps=30)
camera = RealSenseCamera(config)
@@ -91,6 +125,33 @@ def test_invalid_width_connect():
camera.connect(warmup=False)
def test_connect_cleans_up_after_warmup_failure_and_allows_retry():
config = RealSenseCameraConfig(serial_number_or_name="042", width=640, height=480, fps=30)
camera = RealSenseCamera(config)
read_threads = []
def fail_warmup(*_args, **_kwargs):
read_threads.append(camera.thread)
raise TimeoutError("no frame")
with (
patch.object(camera, "async_read", side_effect=fail_warmup),
pytest.raises(TimeoutError, match="no frame"),
):
camera.connect()
assert camera.rs_pipeline is None
assert camera.rs_profile is None
assert camera.thread is None
assert not camera.is_connected
assert read_threads[0] is not None
assert not read_threads[0].is_alive()
camera.connect(warmup=False)
assert camera.is_connected
camera.disconnect()
def test_read():
config = RealSenseCameraConfig(serial_number_or_name="042", width=640, height=480, fps=30, warmup_s=0)
with RealSenseCamera(config) as camera:
@@ -228,6 +289,203 @@ def test_read_latest_too_old():
_ = camera.read_latest(max_age_ms=0) # immediately too old
def _make_mock_sensor(name: str, supported_options: set | None = None) -> MagicMock:
"""Build a fake rs.sensor that reports a name and a configurable supported-options set."""
supported = supported_options if supported_options is not None else set()
sensor = MagicMock()
sensor.get_info.return_value = name
sensor.supports.side_effect = lambda opt: opt in supported
return sensor
def _attach_mock_color_sensor(camera: RealSenseCamera, sensor: MagicMock) -> None:
"""Wire camera.rs_profile so _get_color_sensor finds the given sensor."""
profile = MagicMock()
device = MagicMock()
device.query_sensors.return_value = [sensor]
profile.get_device.return_value = device
camera.rs_profile = profile
def test_get_color_sensor_prefers_rgb_camera():
config = RealSenseCameraConfig(serial_number_or_name="042")
camera = RealSenseCamera(config)
rgb = _make_mock_sensor("RGB Camera")
stereo = _make_mock_sensor("Stereo Module")
profile = MagicMock()
device = MagicMock()
device.query_sensors.return_value = [stereo, rgb]
profile.get_device.return_value = device
camera.rs_profile = profile
assert camera._get_color_sensor() is rgb
def test_get_color_sensor_falls_back_to_stereo_module():
"""D405 has no separate RGB module; color comes from Stereo Module."""
config = RealSenseCameraConfig(serial_number_or_name="042")
camera = RealSenseCamera(config)
stereo = _make_mock_sensor("Stereo Module")
_attach_mock_color_sensor(camera, stereo)
assert camera._get_color_sensor() is stereo
def test_get_color_sensor_raises_with_available_sensors():
config = RealSenseCameraConfig(serial_number_or_name="042")
camera = RealSenseCamera(config)
other = _make_mock_sensor("Motion Module")
_attach_mock_color_sensor(camera, other)
with pytest.raises(RuntimeError, match="Motion Module"):
camera._get_color_sensor()
def test_configure_sensor_options_skipped_when_none():
config = RealSenseCameraConfig(serial_number_or_name="042")
camera = RealSenseCamera(config)
with patch.object(RealSenseCamera, "_get_color_sensor") as mock_get:
camera._configure_sensor_options()
mock_get.assert_not_called()
def test_configure_sensor_options_applies_all_values():
config = RealSenseCameraConfig(serial_number_or_name="042", exposure=120, gain=64, white_balance=4600)
camera = RealSenseCamera(config)
sensor = _make_mock_sensor(
"RGB Camera",
supported_options={
rs.option.enable_auto_exposure,
rs.option.exposure,
rs.option.gain,
rs.option.enable_auto_white_balance,
rs.option.white_balance,
},
)
_attach_mock_color_sensor(camera, sensor)
camera._configure_sensor_options()
sensor.set_option.assert_any_call(rs.option.enable_auto_exposure, 0)
sensor.set_option.assert_any_call(rs.option.exposure, 120)
sensor.set_option.assert_any_call(rs.option.gain, 64)
sensor.set_option.assert_any_call(rs.option.enable_auto_white_balance, 0)
sensor.set_option.assert_any_call(rs.option.white_balance, 4600)
@pytest.mark.parametrize(
("config_field", "option", "label"),
[
("exposure", rs.option.exposure, "exposure"),
("gain", rs.option.gain, "gain"),
("white_balance", rs.option.white_balance, "white balance"),
],
)
def test_configure_sensor_options_raises_when_requested_option_is_unsupported(config_field, option, label):
config = RealSenseCameraConfig(serial_number_or_name="042", **{config_field: 100})
camera = RealSenseCamera(config)
sensor = _make_mock_sensor("RGB Camera", supported_options=set())
_attach_mock_color_sensor(camera, sensor)
with pytest.raises(ValueError, match=label):
camera._configure_sensor_options()
sensor.supports.assert_any_call(option)
sensor.set_option.assert_not_called()
@pytest.mark.parametrize(
("config_field", "option", "value"),
[
("exposure", rs.option.exposure, 120),
("gain", rs.option.gain, 64),
],
)
def test_configure_sensor_options_exposure_or_gain_disables_auto_exposure(config_field, option, value):
"""white_balance=None should not touch auto white balance."""
config = RealSenseCameraConfig(serial_number_or_name="042", **{config_field: value})
camera = RealSenseCamera(config)
sensor = _make_mock_sensor(
"RGB Camera",
supported_options={rs.option.enable_auto_exposure, option},
)
_attach_mock_color_sensor(camera, sensor)
camera._configure_sensor_options()
calls = [call.args for call in sensor.set_option.call_args_list]
assert (rs.option.enable_auto_exposure, 0) in calls
assert (option, value) in calls
for opt, _ in calls:
assert opt != rs.option.enable_auto_white_balance
assert opt != rs.option.white_balance
def test_configure_sensor_options_warns_when_auto_exposure_control_is_unsupported(caplog):
config = RealSenseCameraConfig(serial_number_or_name="042", exposure=120)
camera = RealSenseCamera(config)
sensor = _make_mock_sensor("RGB Camera", supported_options={rs.option.exposure})
_attach_mock_color_sensor(camera, sensor)
with caplog.at_level("WARNING"):
camera._configure_sensor_options()
sensor.set_option.assert_called_once_with(rs.option.exposure, 120)
assert "does not support disabling auto-exposure" in caplog.text
def test_configure_sensor_options_warns_when_auto_white_balance_control_is_unsupported(caplog):
config = RealSenseCameraConfig(serial_number_or_name="042", white_balance=4600)
camera = RealSenseCamera(config)
sensor = _make_mock_sensor("RGB Camera", supported_options={rs.option.white_balance})
_attach_mock_color_sensor(camera, sensor)
with caplog.at_level("WARNING"):
camera._configure_sensor_options()
sensor.set_option.assert_called_once_with(rs.option.white_balance, 4600)
assert "does not support disabling auto white balance" in caplog.text
def test_configure_sensor_options_out_of_range_raises_value_error():
"""set_option errors should be re-raised as ValueError with range diagnostics."""
config = RealSenseCameraConfig(serial_number_or_name="042", exposure=999999)
camera = RealSenseCamera(config)
sensor = _make_mock_sensor(
"RGB Camera",
supported_options={rs.option.enable_auto_exposure, rs.option.exposure},
)
def fake_set_option(option, value):
if option == rs.option.exposure:
raise RuntimeError("value out of range")
sensor.set_option.side_effect = fake_set_option
option_range = MagicMock(min=1, max=10000, step=1, default=156)
sensor.get_option_range.return_value = option_range
_attach_mock_color_sensor(camera, sensor)
with pytest.raises(ValueError, match="exposure") as exc_info:
camera._configure_sensor_options()
msg = str(exc_info.value)
assert "999999" in msg
assert "min=1" in msg
assert "max=10000" in msg
@pytest.mark.parametrize(
"rotation",
[
+24
View File
@@ -66,3 +66,27 @@ def test_from_pretrained_raises_when_no_root_config_and_no_checkpoints(monkeypat
with pytest.raises(FileNotFoundError, match="train_config.json not found"):
TrainPipelineConfig.from_pretrained("user/empty-repo")
@pytest.mark.parametrize("pass_dir", [False, True])
def test_resolve_resume_checkpoint_accepts_file_or_pretrained_model_dir(tmp_path, monkeypatch, pass_dir):
"""`--config_path` may point at the checkpoint's train_config.json or at its
pretrained_model/ directory; both must resolve `policy.pretrained_path` to the
pretrained_model/ directory (regression test for the directory case, which
previously resolved one level too high and failed on model.safetensors)."""
pretrained_dir = tmp_path / "checkpoints" / "000002" / "pretrained_model"
pretrained_dir.mkdir(parents=True)
(pretrained_dir / "train_config.json").touch()
target = pretrained_dir if pass_dir else pretrained_dir / "train_config.json"
from lerobot.policies.act.configuration_act import ACTConfig
cfg = tc.draccus.parse(TrainPipelineConfig, args=["--dataset.repo_id", "u/d"])
cfg.policy = ACTConfig()
cfg.resume = True
monkeypatch.setattr(tc.parser, "parse_arg", lambda name: str(target) if name == "config_path" else None)
cfg._resolve_resume_checkpoint()
assert cfg.policy.pretrained_path == pretrained_dir
assert cfg.checkpoint_path == pretrained_dir.parent
@@ -0,0 +1,104 @@
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import numpy as np
import pytest
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
from lerobot.scripts.augment_dataset_quantile_stats import (
compute_quantile_stats_for_dataset,
has_quantile_stats,
)
def _numeric_keys(dataset):
return [k for k, v in dataset.features.items() if v["dtype"] not in ("image", "video", "string")]
def _image_keys(dataset):
return [k for k, v in dataset.features.items() if v["dtype"] in ("image", "video")]
def test_numeric_stats_are_unaffected_by_sampling(tmp_path, lerobot_dataset_factory):
"""Sampling only touches image/video frames; numeric features are read in
full either way, so their stats must be identical with and without sampling."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=2, total_frames=400, use_videos=False
)
exact = compute_quantile_stats_for_dataset(dataset, use_sampling=False)
sampled = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
numeric_keys = _numeric_keys(dataset)
assert numeric_keys, "fixture should expose numeric features"
for key in numeric_keys:
if key not in exact:
continue
for stat in ("mean", "std", "q01", "q50", "q99"):
if stat in exact[key]:
np.testing.assert_allclose(
sampled[key][stat],
exact[key][stat],
rtol=1e-6,
atol=1e-6,
err_msg=f"numeric feature '{key}' stat '{stat}' changed under sampling",
)
def test_image_sampling_reduces_data_but_keeps_stats_close(tmp_path, lerobot_dataset_factory):
"""For images, sampling should reduce the number of samples considered while
keeping the resulting statistics close to the exact ones."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=2, total_frames=400, use_videos=False
)
exact = compute_quantile_stats_for_dataset(dataset, use_sampling=False)
sampled = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
image_keys = _image_keys(dataset)
assert image_keys, "fixture should expose at least one image feature"
for key in image_keys:
# sampling actually looked at fewer pixels
assert sampled[key]["count"][0] < exact[key]["count"][0]
# but per-channel mean stays close
np.testing.assert_allclose(
sampled[key]["mean"],
exact[key]["mean"],
rtol=0.15,
err_msg=f"image feature '{key}' mean drifted too far under sampling",
)
def test_short_episodes_use_all_frames(tmp_path, lerobot_dataset_factory):
"""With episodes shorter than the sampling floor, sampling is a no-op and
must produce exactly the same stats as the exact path."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=2, total_frames=40, use_videos=False
)
exact = compute_quantile_stats_for_dataset(dataset, use_sampling=False)
sampled = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
for key in _image_keys(dataset):
assert sampled[key]["count"][0] == exact[key]["count"][0]
def test_quantile_stats_present_after_compute(tmp_path, lerobot_dataset_factory):
"""The computed stats should contain quantile keys for the dataset."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=2, total_frames=200, use_videos=False
)
stats = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
assert has_quantile_stats(stats)
+59 -1
View File
@@ -1125,9 +1125,61 @@ def test_modify_tasks_default_with_overrides(sample_dataset):
assert ep_data["tasks"][0] == default_task
def test_modify_tasks_replacements(sample_dataset):
"""Test replacing task strings based on their current values."""
modified_dataset = modify_tasks(
sample_dataset,
task_replacements={
"task_0": "Pick the cube",
"task_1": "Place the cube",
},
)
assert len(modified_dataset.meta.tasks) == 2
assert "Pick the cube" in modified_dataset.meta.tasks.index
assert "Place the cube" in modified_dataset.meta.tasks.index
for ep_idx in range(5):
expected_task = "Pick the cube" if ep_idx % 2 == 0 else "Place the cube"
assert modified_dataset.meta.episodes[ep_idx]["tasks"][0] == expected_task
def test_modify_tasks_replacements_with_episode_overrides(sample_dataset):
"""Test that explicit episode overrides take precedence over replacements."""
modified_dataset = modify_tasks(
sample_dataset,
task_replacements={
"task_0": "Pick the cube",
"task_1": "Place the cube",
},
episode_tasks={1: "Inspect the cube"},
)
assert modified_dataset.meta.episodes[0]["tasks"][0] == "Pick the cube"
assert modified_dataset.meta.episodes[1]["tasks"][0] == "Inspect the cube"
assert modified_dataset.meta.episodes[3]["tasks"][0] == "Place the cube"
assert len(modified_dataset.meta.tasks) == 3
def test_modify_tasks_default_task_and_replacements(sample_dataset):
"""Test that new_task acts as the default for episodes not matched by task_replacements."""
modified_dataset = modify_tasks(
sample_dataset,
new_task="Default task",
task_replacements={"task_0": "Pick the cube"},
)
for ep_idx in range(5):
expected_task = "Pick the cube" if ep_idx % 2 == 0 else "Default task"
assert modified_dataset.meta.episodes[ep_idx]["tasks"][0] == expected_task
assert len(modified_dataset.meta.tasks) == 2
def test_modify_tasks_no_task_specified(sample_dataset):
"""Test error when no task is specified."""
with pytest.raises(ValueError, match="Must specify at least one of new_task or episode_tasks"):
with pytest.raises(
ValueError, match="Must specify at least one of new_task, episode_tasks, or task_replacements"
):
modify_tasks(sample_dataset)
@@ -1137,6 +1189,12 @@ def test_modify_tasks_invalid_episode_indices(sample_dataset):
modify_tasks(sample_dataset, episode_tasks={10: "Task", 20: "Task"})
def test_modify_tasks_invalid_task_replacements(sample_dataset):
"""Test error when task replacements refer to unknown task strings."""
with pytest.raises(ValueError, match="Task replacements reference unknown tasks"):
modify_tasks(sample_dataset, task_replacements={"missing_task": "New task"})
def test_modify_tasks_updates_info_json(sample_dataset):
"""Test that total_tasks is updated in info.json."""
episode_tasks = {0: "Task A", 1: "Task B", 2: "Task C", 3: "Task A", 4: "Task B"}
+32
View File
@@ -204,6 +204,38 @@ def test_clear_resets_buffer(tmp_path):
assert dataset.writer.episode_buffer["size"] == 0
def test_clear_removes_video_frame_staging_dir(tmp_path):
"""clear_episode_buffer() removes PNG staging dirs for video features."""
video_key = "observation.images.cam"
features = {
video_key: {
"dtype": "video",
"shape": (64, 96, 3),
"names": ["height", "width", "channels"],
},
"action": {"dtype": "float32", "shape": (2,), "names": None},
}
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID,
fps=DEFAULT_FPS,
features=features,
root=tmp_path / "ds",
use_videos=True,
)
dataset.add_frame(_make_frame(features))
video_staging_dir = (
dataset.root
/ Path(DEFAULT_IMAGE_PATH.format(image_key=video_key, episode_index=0, frame_index=0)).parent
)
assert video_staging_dir.is_dir()
dataset.clear_episode_buffer()
assert dataset.writer.episode_buffer["size"] == 0
assert not video_staging_dir.exists()
def test_finalize_is_idempotent(tmp_path):
"""Calling finalize() twice does not raise."""
dataset = LeRobotDataset.create(
+14
View File
@@ -482,6 +482,20 @@ def test_add_frame_works_in_write_mode(tmp_path):
# ── Resume mode ──────────────────────────────────────────────────────
def test_resume_freshly_created_empty_dataset(tmp_path):
"""resume() accepts a local dataset created before any episode was recorded."""
root = tmp_path / "resume_empty_ds"
LeRobotDataset.create(repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root)
resumed = LeRobotDataset.resume(repo_id=DUMMY_REPO_ID, root=root)
assert isinstance(resumed.writer, DatasetWriter)
assert resumed.meta.total_episodes == 0
assert resumed.meta.total_frames == 0
assert resumed.meta.tasks is None
assert resumed.meta.episodes is None
def test_resume_creates_writer(tmp_path):
"""After resume(), writer is a DatasetWriter."""
root = tmp_path / "resume_ds"
+13
View File
@@ -294,6 +294,19 @@ def test__sync_read(addr, length, ids_values, mock_motors, dummy_motors):
assert read_values == ids_values
def test__sync_read_retries_after_transient_failure(mock_motors, dummy_motors):
addr, length, ids_values = (10, 4, {1: 1337})
stub = mock_motors.build_sync_read_stub(addr, length, ids_values, num_invalid_try=1)
bus = FeetechMotorsBus(port=mock_motors.port, motors=dummy_motors)
bus.connect(handshake=False)
read_values, read_comm = bus._sync_read(addr, length, list(ids_values), num_retry=1)
assert read_comm == scs.COMM_SUCCESS
assert read_values == ids_values
assert mock_motors.stubs[stub].calls == 2
@pytest.mark.parametrize("raise_on_error", (True, False))
def test__sync_read_comm(raise_on_error, mock_motors, dummy_motors):
addr, length, ids_values = (10, 4, {1: 1337})
+54
View File
@@ -496,6 +496,60 @@ def test_evo1_processor_save_load_round_trip_applies_config_overrides(tmp_path):
assert "embodiment_id" in processed
def test_reconcile_evo1_processors_repads_overridden_stats(tmp_path):
"""Loading a checkpoint and injecting raw (unpadded) dataset stats must be re-padded.
Regression test: lerobot-train passes the raw dataset stats as normalizer/unnormalizer
overrides when resuming from a checkpoint (e.g. stage2 from a stage1 checkpoint). Those stats
are at the dataset dims (e.g. LIBERO state=8/action=7), but EVO1 pads state/action to
max_state_dim/max_action_dim before normalization, so reconcile_evo1_processors must re-pad the
stats or normalization crashes with a shape mismatch.
"""
config = make_config()
preprocessor, postprocessor = make_evo1_pre_post_processors(config, dataset_stats=make_stats())
preprocessor.save_pretrained(tmp_path)
postprocessor.save_pretrained(tmp_path)
# Reload with the generic override path injecting raw, unpadded dataset stats.
raw_stats = make_stats()
loaded_pre = PolicyProcessorPipeline.from_pretrained(
tmp_path,
config_filename=f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json",
overrides={"normalizer_processor": {"stats": raw_stats}},
to_transition=batch_to_transition,
to_output=transition_to_batch,
)
loaded_post = PolicyProcessorPipeline.from_pretrained(
tmp_path,
config_filename=f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json",
overrides={"unnormalizer_processor": {"stats": raw_stats}},
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
)
# Sanity: the override really injected unpadded stats before reconciliation.
normalizer = next(step for step in loaded_pre.steps if isinstance(step, NormalizerProcessorStep))
assert normalizer._tensor_stats[OBS_STATE]["min"].shape == (STATE_DIM,)
loaded_pre, loaded_post = reconcile_evo1_processors(config, loaded_pre, loaded_post)
normalizer = next(step for step in loaded_pre.steps if isinstance(step, NormalizerProcessorStep))
unnormalizer = next(step for step in loaded_post.steps if isinstance(step, UnnormalizerProcessorStep))
assert normalizer._tensor_stats[OBS_STATE]["min"].shape == (MAX_STATE_DIM,)
assert normalizer._tensor_stats[ACTION]["min"].shape == (MAX_ACTION_DIM,)
assert unnormalizer._tensor_stats[ACTION]["min"].shape == (MAX_ACTION_DIM,)
# Normalizing a padded state must not raise (this is the exact runtime path that crashed).
processed = loaded_pre(
{
"task": "pick the block",
OBS_STATE: torch.zeros(STATE_DIM),
f"{OBS_IMAGES}.front": torch.rand(3, 16, 16),
}
)
assert processed[OBS_STATE].shape == (1, MAX_STATE_DIM)
def test_evo1_policy_forward_and_inference_use_batched_embedding(monkeypatch):
monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
policy = modeling_evo1.Evo1Policy(make_config())
@@ -0,0 +1,83 @@
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from types import SimpleNamespace
from unittest.mock import MagicMock
import torch
import lerobot.policies.factory as policy_factory
def test_make_policy_keeps_peft_adapter_and_base_revisions_separate(monkeypatch):
cfg = SimpleNamespace(
type="mock",
device="cpu",
pretrained_path="user/adapter",
pretrained_revision="adapter-sha",
use_peft=True,
input_features={},
output_features={},
)
dataset_meta = SimpleNamespace(features={}, stats={})
base_policy = torch.nn.Linear(1, 1)
policy_from_pretrained = MagicMock(return_value=base_policy)
policy_class = SimpleNamespace(from_pretrained=policy_from_pretrained)
monkeypatch.setattr(policy_factory, "get_policy_class", lambda _: policy_class)
monkeypatch.setattr(policy_factory, "dataset_to_policy_features", lambda _: {})
monkeypatch.setattr(policy_factory, "validate_visual_features_consistency", lambda *args: None)
peft_config = SimpleNamespace(
base_model_name_or_path="user/base-policy",
revision="base-sha",
)
peft_config_from_pretrained = MagicMock(return_value=peft_config)
adapted_policy = torch.nn.Linear(1, 1)
peft_model_from_pretrained = MagicMock(return_value=adapted_policy)
require_package = MagicMock()
monkeypatch.setattr(policy_factory, "require_package", require_package)
monkeypatch.setattr(
policy_factory,
"PeftConfig",
SimpleNamespace(from_pretrained=peft_config_from_pretrained),
)
monkeypatch.setattr(
policy_factory,
"PeftModel",
SimpleNamespace(from_pretrained=peft_model_from_pretrained),
)
policy = policy_factory.make_policy(cfg, ds_meta=dataset_meta)
assert policy is adapted_policy
require_package.assert_called_once_with("peft", extra="peft")
peft_config_from_pretrained.assert_called_once_with(
"user/adapter",
revision="adapter-sha",
)
policy_from_pretrained.assert_called_once_with(
config=cfg,
dataset_stats=dataset_meta.stats,
dataset_meta=dataset_meta,
pretrained_name_or_path="user/base-policy",
revision="base-sha",
)
peft_model_from_pretrained.assert_called_once_with(
base_policy,
"user/adapter",
config=peft_config,
revision="adapter-sha",
is_trainable=True,
)
@@ -113,6 +113,7 @@ def test_gaussian_actor_config_default_initialization():
# Concurrency configuration
assert config.concurrency.actor == "threads"
assert config.concurrency.learner == "threads"
assert config.concurrency.multiprocessing_context == "spawn"
assert isinstance(config.actor_network_kwargs, ActorNetworkConfig)
assert isinstance(config.policy_kwargs, PolicyConfig)
@@ -152,6 +153,7 @@ def test_concurrency_config():
config = ConcurrencyConfig()
assert config.actor == "threads"
assert config.learner == "threads"
assert config.multiprocessing_context == "spawn"
def test_gaussian_actor_config_custom_initialization():
@@ -26,8 +26,17 @@ import tempfile
from pathlib import Path
import pytest
import torch
from safetensors.torch import save_file
from lerobot.processor.pipeline import DataProcessorPipeline, ProcessorMigrationError
from lerobot.configs import PipelineFeatureType, PolicyFeature
from lerobot.processor.pipeline import (
DataProcessorPipeline,
ProcessorMigrationError,
ProcessorStep,
ProcessorStepRegistry,
)
from lerobot.types import EnvTransition
# Simplified Config Loading Tests
@@ -98,6 +107,140 @@ def test_load_config_nonexistent_path_tries_hub():
DataProcessorPipeline._load_config("nonexistent/path", "processor.json", {})
def test_from_pretrained_local_directory_missing_state_does_not_call_hub(monkeypatch):
"""Local processor dirs must fail locally when a state file is missing."""
@ProcessorStepRegistry.register("local_missing_state_step")
class LocalMissingStateStep(ProcessorStep):
def __call__(self, transition: EnvTransition) -> EnvTransition:
return transition
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
return features
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
pass
try:
with tempfile.TemporaryDirectory() as tmp_dir:
tmp_path = Path(tmp_dir)
config = {
"name": "LocalMissingStatePipeline",
"steps": [{"registry_name": "local_missing_state_step", "state_file": "missing.safetensors"}],
}
(tmp_path / "processor.json").write_text(json.dumps(config))
def fail_hub_download(*args, **kwargs):
pytest.fail("local missing processor state should not call hf_hub_download")
monkeypatch.setattr("lerobot.processor.pipeline.hf_hub_download", fail_hub_download)
with pytest.raises(FileNotFoundError, match="missing.safetensors.*local processor pipeline"):
DataProcessorPipeline.from_pretrained(tmp_path, config_filename="processor.json")
finally:
ProcessorStepRegistry.unregister("local_missing_state_step")
def test_from_pretrained_local_config_file_missing_state_does_not_call_hub(monkeypatch):
"""Local single-file processor configs must also keep missing state resolution local."""
@ProcessorStepRegistry.register("local_file_missing_state_step")
class LocalFileMissingStateStep(ProcessorStep):
def __call__(self, transition: EnvTransition) -> EnvTransition:
return transition
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
return features
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
pass
try:
with tempfile.TemporaryDirectory() as tmp_dir:
tmp_path = Path(tmp_dir)
config_path = tmp_path / "processor.json"
config = {
"name": "LocalFileMissingStatePipeline",
"steps": [
{"registry_name": "local_file_missing_state_step", "state_file": "missing.safetensors"}
],
}
config_path.write_text(json.dumps(config))
def fail_hub_download(*args, **kwargs):
pytest.fail("local missing processor state should not call hf_hub_download")
monkeypatch.setattr("lerobot.processor.pipeline.hf_hub_download", fail_hub_download)
with pytest.raises(FileNotFoundError, match="missing.safetensors.*local processor pipeline"):
DataProcessorPipeline.from_pretrained(config_path, config_filename="ignored.json")
finally:
ProcessorStepRegistry.unregister("local_file_missing_state_step")
def test_from_pretrained_hub_source_missing_local_state_still_calls_hub(monkeypatch, tmp_path):
"""Hub sources still fall back to hf_hub_download for state files."""
@ProcessorStepRegistry.register("hub_state_step")
class HubStateStep(ProcessorStep):
def __init__(self):
self.value = torch.tensor(0)
def __call__(self, transition: EnvTransition) -> EnvTransition:
return transition
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
return features
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
self.value = state["value"]
try:
state_path = tmp_path / "downloaded.safetensors"
save_file({"value": torch.tensor(7)}, state_path)
loaded_config = {
"name": "HubStatePipeline",
"steps": [{"registry_name": "hub_state_step", "state_file": "hub_state.safetensors"}],
}
calls = []
def fake_load_config(cls, model_id, config_filename, hub_download_kwargs):
return loaded_config, tmp_path / "hub_cache"
def fake_hub_download(**kwargs):
calls.append(kwargs)
return str(state_path)
monkeypatch.setattr(DataProcessorPipeline, "_load_config", classmethod(fake_load_config))
monkeypatch.setattr("lerobot.processor.pipeline.hf_hub_download", fake_hub_download)
pipeline = DataProcessorPipeline.from_pretrained("user/repo", config_filename="processor.json")
assert calls == [
{
"repo_id": "user/repo",
"filename": "hub_state.safetensors",
"repo_type": "model",
"force_download": False,
"resume_download": None,
"proxies": None,
"token": None,
"cache_dir": None,
"local_files_only": False,
"revision": None,
}
]
assert pipeline.steps[0].value.item() == 7
finally:
ProcessorStepRegistry.unregister("hub_state_step")
# Config Validation Tests
@@ -0,0 +1,45 @@
#!/usr/bin/env python
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import pytest
from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
from lerobot.robots.so_follower.robot_kinematic_processor import (
ForwardKinematicsJointsToEEAction,
ForwardKinematicsJointsToEEObservation,
)
MOTOR_NAMES = ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper"]
EE_KEYS = {f"ee.{k}" for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]}
def _joint_bucket(feature_type: FeatureType) -> dict[str, PolicyFeature]:
return {f"{n}.pos": PolicyFeature(type=feature_type, shape=(1,)) for n in MOTOR_NAMES}
@pytest.mark.parametrize(
("step_cls", "bucket", "feature_type"),
[
(ForwardKinematicsJointsToEEAction, PipelineFeatureType.ACTION, FeatureType.ACTION),
(ForwardKinematicsJointsToEEObservation, PipelineFeatureType.OBSERVATION, FeatureType.STATE),
],
)
def test_fk_feature_schema(step_cls, bucket, feature_type):
features = {PipelineFeatureType.ACTION: {}, PipelineFeatureType.OBSERVATION: {}}
features[bucket] = _joint_bucket(feature_type)
out = step_cls(kinematics=None, motor_names=MOTOR_NAMES).transform_features(features)[bucket]
assert set(out) == EE_KEYS
assert {feature.type for feature in out.values()} == {feature_type}

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