- recipes/subtask_joint.yaml: paper-style single sequence (pi0.5 §IV-B) —
the supervised subtask span gets text CE and conditions the FAST and
flow losses in the same forward.
- joint_subtask_conditioning config flag rebuilds the same layout at
inference: state on the task turn, generated subtask as a causal
assistant turn (encode_prompt_with_targets + lang_causal_marks through
sample_actions), in both the policy select_action path and the runtime
adapter.
- fast_skip_tokens default 128 -> 1152 so FAST codes land below the <loc>
range and never collide with VQA loc targets; _FAST_ACTION_VOCAB_SIZE
tightened to the universal tokenizer's 1024 codes.
- Strip the trailing space from the 'Assistant:' generation prefill —
SentencePiece folds the space into the first target token, so the
space-suffixed prefill ended in a lone '▁' never seen in training.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Align tokenizer fitting and loss reduction with the effective training dataset, and fail early when FAST supervision cannot be produced safely.
Co-authored-by: Cursor <cursoragent@cursor.com>
* feat(policies): add EVO1 policy
* fix(evo1): infer batch size after normalizing image dims
`_collect_image_batches` read `batch_size = batch[camera_keys[0]].shape[0]`
before normalizing per-camera tensors to `(B, C, H, W)`. For an unbatched
`(C, H, W)` input (which the function tries to support via the `image.dim() == 3`
branch), this picked up the channel count `C` instead of the real batch size,
making the subsequent per-sample loop iterate `C` times and indexing go
out of bounds.
Normalize each camera tensor up-front, then read `batch_size` from the
normalized batch dim. Adds `test_collect_image_batches_handles_unbatched_chw`
covering the regression.
Reported by Copilot review on huggingface/lerobot#3545.
* chore(lock): regenerate uv.lock for evo1 extra
Adds the `evo1` entry to `[package.metadata.requires-dist]` and the
`provides-extras` list so that `uv sync --locked --extra test` (used by
fast_tests.yml) no longer reports the lockfile as stale.
Generated with `uv 0.8.0` (matching `UV_VERSION` in fast_tests.yml).
The non-evo1 marker tweaks are produced by `uv lock` re-resolving the
existing dep graph and are not introduced by this PR.
* chore(evo1): align with policy contribution guide conventions
- Add `src/lerobot/policies/evo1/README.md` symlink into `docs/source/evo1.mdx`
to match the in-tree README convention (mirroring the EO-1 layout).
- Convert `transformers` import in `internvl3_embedder.py` to the standard
`TYPE_CHECKING + _transformers_available` two-step gating used by other
optional-backbone policies (e.g. diffusion). The previous lazy-in-`__init__`
import was functionally equivalent for runtime gating but didn't expose the
real symbols to type checkers.
- Add `lerobot[evo1]` to the `all` extra in `pyproject.toml` so
`pip install 'lerobot[all]'` keeps installing every optional policy.
Per the guidance in https://moon-ci-docs.huggingface.co/docs/lerobot/pr_3534/en/contributing_a_policy.
* fix(evo1): finalize policy guide alignment
* docs(evo1): format results table
* Fix EVO1 LIBERO rollout processors
* Fix EVO1 LIBERO eval action postprocessing
* Fix eval action conversion for bf16 policies
* fix(evo1): move LIBERO padding into policy processors
* refactor(evo1): use native HF InternVL3-1B-hf, drop trust_remote_code
- Switch from OpenGVLab/InternVL3-1B (requires trust_remote_code=True)
to OpenGVLab/InternVL3-1B-hf (native transformers implementation).
- Replace manual _extract_feature + _prepare_and_fuse_embeddings with
a single model.forward() call — verified bit-for-bit identical output.
- Remove ~170 lines of manual ViT/pixel-shuffle/projection logic.
- Symlink README.md to docs/source/ following repo convention.
Weights are byte-identical between both model variants; only the module
naming differs. All 12 existing unit tests pass. Local training (10 steps)
on maximellerbach/omx_pickandplace confirmed working.
* refactor(policy): evo1 GPU-batched preprocessing + vectorized attention masking + remove dead code
* fix(style): pre-commit
oops
* chore(evo1): delete added test + reduce diff
* refactor(policies): use config for evo1 + local imports
* refactor(policies): multiple improvements
* chore: update docs + remove legacy codepaths
* feat(policies): implement RTC to EVO1
---------
Co-authored-by: javadcc_mac <javadcc1@sjtu.edu.cn>
Co-authored-by: Yiming Wang <145452074+JAVAdcc@users.noreply.github.com>
Co-authored-by: Martino Russi <nopyeps@gmail.com>
* Add GR00T N1.7 support
Add GR00T N1.7 policy configuration, checkpoint compatibility, processor parity, LIBERO documentation, and focused tests.
Co-authored-by: Ryan Halabi <ryhalabi@nvidia.com>
* Move Groot processor compatibility into Groot loader
* Restore GR00T Flash Attention install guidance
* Allow Groot fake RTC chunk prefetch
* Fix GR00T N1.7 RTC action decoding
* Trim GR00T N1.7 RTC chunks to valid horizon
* Ignore padded GR00T N1.7 RTC prefix rows
* removed n1.5 dependency
* removed remaining N1.5 traces
* groot: auto-enable LIBERO gripper action transform for libero_sim
GR00T N1.7 emits gripper in [0,1] but LIBERO expects [-1,1]. The decode
transform existed but was never auto-enabled for embodiment_tag=libero_sim,
so the policy scored 0% on LIBERO eval. Auto-set it in __post_init__ (still
overridable). LIBERO Spatial eval: 0% -> 98%.
* Reconnect GR00T relative action processors
* groot: remove dead N1.5 code (eagle2_hg_model, flow_matching_action_head, action_encoder)
N1.7 backbone is nvidia/Cosmos-Reason2-2B via Qwen3VLForConditionalGeneration,
not Eagle2 — eagle2_hg_model/ had zero refs outside its own dir.
GR00TN17ActionHead (groot_n1_7.py) re-implements MultiEmbodimentActionEncoder +
CategorySpecificLinear + swish + SinusoidalPositionalEncoding locally, so
flow_matching_action_head.py (N1.5 FlowmatchingActionHead) and its sole
dependency action_encoder.py are dead. Verified: no src/ or tests/ reference.
Removed (~2037 LOC):
- eagle2_hg_model/ (4 files, ~1575 LOC)
- action_head/flow_matching_action_head.py (408 LOC)
- action_head/action_encoder.py (54 LOC)
cross_attention_dit.py KEPT (DiT/AlternateVLDiT/SelfAttentionTransformer live in N1.7).
* groot: reuse lerobot get_device_from_parameters instead of inline lookup
modeling_groot.py duplicated next(self.parameters()).device twice. LeRobot
ships get_device_from_parameters in policies/utils.py (used by diffusion,
vqbet, tdmpc, gaussian_actor). Reuse it for consistency with the framework.
* groot: fix stale Eagle VLM docstring in processor (N1.7 uses Qwen3-VL backbone)
Addresses checker nit: processor_groot.py docstring still described the N1.5
Eagle VLM path with eagle_content/eagle_* keys that no longer exist in the code.
* test(groot): add N1.7 original-vs-LeRobot output parity test
Verifies the LeRobot GR00T N1.7 integration produces equivalent raw
action_pred to NVIDIA Isaac-GR00T for the same checkpoint, inputs, seed,
precision (fp32) and attention kernel (SDPA): max|diff|=8.9e-7 on the
libero_sim embodiment (GR00T-N1.7-LIBERO/libero_10).
The two impls pin incompatible transformers majors (orig 4.57.3 vs
LeRobot 5.x) and cannot share a process, so the original outputs + exact
collated inputs are produced out-of-process and loaded from an .npz. The
test skips on CI / when the checkpoint or artifact are absent.
* test(groot): parametrize N1.7 parity across all checkpoint embodiments
Generalize the original-vs-LeRobot N1.7 output-parity test from a single
libero_sim case to every embodiment tag in the checkpoint (libero_sim, oxe_droid,
real_g1, the real_r1_pro_sharpa family, and the xdof family). Inputs are built
generically from checkpoint metadata; the test discovers per-tag .npz artifacts
and runs one parametrized case each, loading the LeRobot model once via a fixture.
All 9 embodiments match the original to fp32 epsilon (max|diff| < 3e-6), confirming
the integration is correct across the model's full embodiment space and not overfit
to libero_sim.
* test(groot): self-contained parity test + in-repo producer + docs
- Rename test_groot_n1_7_vs_original.py -> test_groot_vs_original.py
- Make the test self-contained: producer script (dump_original_n1_7.py) now lives
next to the test; default artifact dir is repo-relative
(tests/policies/groot/artifacts/), overridable via GROOT_N1_7_PARITY_DIR. The
test only reads artifacts and skips if absent -- it never creates external dirs.
- Heavy .npz artifacts (~6-9MB each) are gitignored and regenerated by the producer;
never committed.
- Drop the verbose 'MULTIPLE EMBODIMENTS' docstring block (kept a one-line note).
- Document the parity procedure in the groot policy README (docs/source/policy_groot_README.md).
- Rename test fn test_groot_n1_7_get_action_parity -> test_groot_get_action_parity.
9/9 embodiments still pass (max|diff| < 3e-6, fp32 eps).
* docs(groot): drop WHY TWO ENVIRONMENTS block from parity test docstring
* test(groot): move parity producer into utils/ package
Mirror the tests/policies/pi0_pi05/utils convention: move dump_original_n1_7.py into
a tests/policies/groot/utils/ package (with __init__.py) and update all path
references in the test docstring/skip-message and the policy README.
* test(groot): adopt test_groot_lerobot for GR00T N1.7, drop N1.5
The test loaded MODEL_PATH='aractingi/bimanual-handover-groot-10k', an N1.5
checkpoint (config base_model_path=nvidia/GR00T-N1.5-3B, no model_version). On
load, model_version defaults to n1.7 while the base path infers n1.5, so the
version-consistency guard in GrootConfig.__post_init__ raised ValueError and both
test_lerobot_groot_inference and test_lerobot_groot_forward_pass failed. N1.5 is no
longer a supported model_version.
Adopt the test for N1.7:
- MODEL_PATH -> nvidia/GR00T-N1.7-3B (root-level sharded safetensors; loads via
GrootPolicy.from_pretrained as a base N1.7 model).
- Embodiment tag 'gr1' (N1.5) -> 'gr1_unified' (valid N1.7 tag from the checkpoint
embodiment_id.json), via a single EMBODIMENT_TAG constant.
- DUMMY_ACTION_HORIZON 16 -> 40 to match N1.7's native action-chunk size.
- Docstrings/labels updated to 'GR00T N1.7'.
Both tests run and pass on CUDA; full tests/policies/groot/ suite is
73 passed / 0 failed / 0 skipped.
* docs(groot): document the N1.5 removal and the N1.7 parity test
- groot.mdx: breaking-change warning and migration path (pin lerobot==0.5.1 to
keep N1.5, or move to N1.7); the dead `huggingface-cli download` is replaced
with `hf download`.
- policy_groot_README.md: N1.5 removal note, updated paper / model-card links,
and the two-comparison (model parity + preprocessor parity) description of
the original-vs-LeRobot test, including the raw-observation artifacts and
recorded seed.
* fix(groot): N1.7 backbone loading and DiT parameter-count logging
- select_layer default tracks the N1.7-3B checkpoint value (16); real
checkpoint loads still override it from config.json.
- get_backbone_cls recognizes Cosmos-Reason2 / Qwen3-VL backbones by name and
warns (instead of silently assuming) when an unrecognized backbone is loaded
only on the strength of backbone_model_type='qwen'.
- 'revision' pins the GR00T checkpoint repo only and is no longer forwarded
into the unrelated backbone repo load; pin the backbone via
transformers_loading_kwargs instead.
- DiT / SelfAttentionTransformer parameter counts go through logging.debug
instead of print().
* fix(groot): N1.7 config defaults, N1.5 rejection, and processor/model runtime fixes
Covers the GR00T N1.7 source trio (configuration, processor, model wrapper).
Config:
- GrootConfig defaults are the N1.7 values; explicitly passed legacy N1.5-era
values (chunk_size=50, max_state_dim=64, ...) are remapped with a warning
instead of silently.
- action_decode_transform gains an 'auto' sentinel so an explicit 'none'
opt-out wins over the libero_sim default and survives save/load round-trips.
- action_delta_indices is cached on the inputs that determine it.
- Legacy N1.5 checkpoints/configs (tokenizer_assets_repo, model_type/
architectures/eagle backbone markers) are rejected with a single clear
error pointing to lerobot==0.5.1.
Processor:
- GrootN17ActionDecodeStep handles the 2-D (B, D) actions delivered by sync
select_action (relative eef/non-eef decode in eval/record flows).
- Postprocessor falls back to dataset stats when a raw checkpoint lacks the
configured embodiment tag; raw-state cache is per-instance, not
process-global; caller overrides (device, rename_map) are honored on the
raw-checkpoint branch.
- Camera/modality-key mismatches warn (including the zero-match fallback);
deprecated Qwen2VLImageProcessorFast replaced with Qwen2VLImageProcessor;
removed N1.5 processor steps are stubbed to raise the removal guidance and
the action-unpack step is re-registered as _v2.
Model:
- Flash-attention probe is diagnostic-only; forward raises on a missing loss;
print() replaced with logging; N1.5 base-path mismatch includes the
removal guidance.
* fix(groot): skip normalization overrides for training
* fix(groot): GPU/tensor N1.7 image preprocessing + resize to trained resolution
GR00T training was dataloader-bound (0->100->0 GPU-utilization sawtooth).
GrootN17VLMEncodeStep ran the Qwen3-VL image processor per frame on PIL images
on the single CPU main-loop thread, and that cost is timed inside dataloading_s
(preprocessor(batch) runs in the main process, not the dataloader workers), so
adding workers cannot hide it.
- Feed the torchvision-backed Qwen3-VL processor (C,H,W) uint8 tensors instead
of a per-frame Image.fromarray PIL roundtrip, and run resize/normalize/patchify
on config.device (GPU) when available. Bit-identical on CPU when no resize is
configured; with a resize only the PIL->torchvision bicubic backend differs
(<2/255 per pixel). The use_albumentations path stays PIL/cv2; reload on a box
without the saved device falls back to CPU.
- Default image_target_size/crop to the N1.7 backbone's training geometry
(256x256 / 230x230) when a checkpoint ships no image sizing (checkpoint_assets
is None, e.g. finetuning nvidia/GR00T-N1.7-3B via repo-id with a new
embodiment). Previously image_target_size=None disabled the resize, so
full-resolution frames were patchified into ~4.7x more vision tokens than the
model was trained on -- inflating dataloading_s (patchify) and update_s (VLM
sequence) and skewing the input distribution. Checkpoints that pin their own
sizing are honored; the default constants are shared with GR00T_N1_7_DEFAULTS.
Net: preprocessing leaves the CPU critical path and the VLM sees the resolution
it was trained on -- faster training/inference and a correct train/serve
distribution. Affects inference too (shared preprocessor); existing checkpoints
still load (backward compatible) but must be retrained to gain the benefits.
* refactor(groot): N1.7 style cleanup (utils, imports, flash-attn, config)
Mechanical refactor of the GR00T N1.7 policy to match the repo's architecture and
style standards. No change to policy algorithm/numerics; only UX/CLI and packaging
changes. Tests are intentionally left untouched (out of scope) and need updating
for the removed `model_version` field.
Cleanup & consolidation:
- Add `groot/utils.py` holding the pure, side-effect-free helpers (JSON I/O, value
coercion, stat flattening, rot6d/SE3 math, language/batch prep) shared by the
config and processor layers.
- Remove dead code: the unused `resolve_groot_n1_7_backbone_model` cache-resolver
cluster, `GR00TN17Config.to_filtered_dict/json`, and the `_copy_default` wrapper.
Imports & execution guards:
- Hoist nested imports to module top; relative imports within the package, absolute
for external modules. The version-gated Qwen3-VL classes import under the single
`_transformers_available` guard (transformers is pinned >=5.4, which ships them).
- No import-time side effects: `_register_with_transformers()` now runs in
`GR00TN17.__init__` (idempotent via `register(exist_ok=True)`), and the N1.5 step
stubs register lazily before pipeline deserialization (idempotent via the
registry, no run-once globals).
- Gate optional deps at the point of use with `require_package(..., extra="groot")`.
Dependencies & docs:
- Drop `flash-attn` (and its build-only dep `ninja`) from the `groot` extra; default
to SDPA (numerically equivalent) with opt-in via `--policy.use_flash_attention`.
Un-comment `lerobot[groot]` in the `all` extra and regenerate `uv.lock`.
- Rewrite the `groot.mdx` install section: flash-attn is a purely optional,
user-managed optimization that LeRobot neither installs nor requires.
Config & CLI:
- Surface previously-frozen knobs on `GrootConfig` (plumbed into `GR00TN17Config`;
no-ops at their defaults): inference — `num_inference_timesteps`, `rtc_ramp_rate`,
`use_flash_attention`; fine-tuning — `tune_top_llm_layers` (partial-LLM tuning)
and `tune_vlln` (previously hardwired to True).
- Convert the single-valued `model_version` and `n1_7_backbone_model` fields to
internal constants.
- Keep `base_model_path`: it is NOT equivalent to `pretrained_path` (raw NVIDIA
checkpoints have no LeRobot `type` field and load only via `base_model_path`) and
is genuinely user-tunable.
- Keep the deprecated Isaac-GR00T/N1.5 fields (and the dead LoRA fields) as a
back-compat block so a v0.5.1 N1.5 `config.json` still parses under draccus and is
rejected with the friendly N1.5 removal message instead of an opaque decode error.
* Optimize GR00T N1.7 image preprocessing
* Remove PIL fallback from GR00T preprocessing
* Fix GROOT relative action training stats
* Address GROOT relative action review feedback
* Fix GROOT N1.7 relative action stats
* Fix GROOT relative action training stats
* Fix GROOT relative action padding and RTC leftovers
* Reset rollout state after robot episode end
* Revert "Reset rollout state after robot episode end"
This reverts commit 1322f45aec.
* Move GROOT relative stats out of train script
* Guard GR00T relative action stepwise decode
* Match GR00T N1.7 OSS preprocessing and relative actions
* Apply LIBERO action decode override after loading
* Format GR00T OSS parity changes
* chore(policies): add guards, warnings and comments + recover tests n1.5 check
* fix(style): pre-commit
* fix(ci): guard dependecy checks
* chore(groot): move cv2 to the top as its in the default install tag
* chore(policies): add explicit dataset dependecy to gr00t implementation
* fix(test): add guard
* fix(groot): make N1.7 letterbox opt-in
* feat(groot): activate checkpoint-configured N1.7 raw-state dropout during training
Isaac-GR00T applies dual state regularization during fine-tuning: raw-state
zeroing driven by the processor sidecar's state_dropout_prob (0.2 for the
inspected N1.7 checkpoint) plus encoded-feature dropout. Baseline LeRobot kept
the processor in deterministic mode, so the raw-state dropout never activated
(RCA Tier-2 contributor to the LeRobot-trained SO-101 failures).
- GrootN17PackInputsStep: runtime-only 'training' flag + state_dropout_prob;
whole-sample state zeroing gated on torch.is_grad_enabled() so eval and
no_grad validation paths are unaffected
- sidecar loader reads state_dropout_prob from processor_config.json
- state_dropout_prob serializes with the step; the training flag intentionally
does not (reloaded pipelines default to eval, re-enabled only when processors
are rebuilt with dataset_meta)
- _set_groot_preprocessor_training toggles any dataclass step exposing a
'training' field on serialized-pipeline reloads
Verification: tests/policies/groot/test_groot_state_dropout.py (4 passed) on
RTX PRO 6000 / CUDA 13.3.
* fix(groot): align N1.7 fine-tuning optimizer/scheduler/precision with Isaac-GR00T
Evidence from the LeRobot-vs-OSS checkpoint comparison: the LeRobot/HF 8k
checkpoint's DiT moved only ~19% as far from base as the OSS-trained one
(0.0547 vs 0.285 relative L2) - undertrained because the scheduler decayed over
a hardcoded 10k steps regardless of --steps, on top of beta1/clip mismatches.
- AdamW betas (0.95, 0.999) -> (0.9, 0.999) and grad_clip_norm 10.0 -> 1.0
(Isaac defaults)
- scheduler: hardcoded CosineDecayWithWarmup(10k decay, floor 10% peak) ->
DiffuserSchedulerConfig HF cosine with ceil(max_steps * warmup_ratio) warmup,
deriving num_training_steps from the outer --steps at runtime
- model_params_fp32 (default true): keep master weights in FP32 and compute
under BF16 autocast like the native N1.7 recipe (fixes optimizer-update
numerics vs pure-BF16 params)
- weight-decay grouping via transformers get_parameter_names: biases and norm
parameters excluded from decay
- restore the TF4 lm_head/embedding weight tie so the unused Qwen LM head stays
frozen and deduplicated in checkpoints
- action_mask kept in native dtype for the masked flow-matching loss
- drop_n_last_frames: exclude episode tails that cannot supply a complete
action chunk (Isaac sampler behavior)
Verification: tests/policies/groot/test_groot_training_optim_contract.py
(7 passed) + remaining groot suite 11 passed/5 skipped on RTX PRO 6000 /
CUDA 13.3. Note: tests/policies/groot/test_groot_n1_7.py does not collect on
the base branch (pre-existing ImportError, fixed in PR #37).
* feat(groot): train-time random crop for N1.7 (eval keeps center crop)
Isaac-GR00T crops a random crop_fraction window during training and the
deterministic center window at eval, replaying the sampled window across all
camera views of a sample. This contract is unchanged since the N1.5 release
(gr00t/data/transform/video.py: "If mode is 'train', return a random crop
transform. If mode is 'eval', return a center crop transform.") and mirrors
LeRobot's own Diffusion/VQBeT crop_is_random pattern. The LeRobot N1.7 port
used the eval center crop for training too, so the fine-tuned projector/DiT
never sees frame borders and trains on a single fixed appearance point.
Scope: crop geometry ONLY - no color jitter, no new dependencies. The random
window is plain numpy slicing inside the existing cv2 eval transform:
- _transform_n1_7_image_for_vlm_albumentations gains crop_position=(y, x)
fractions; None keeps the center crop byte-identical to before (verified
by test)
- GrootN17VLMEncodeStep gains a runtime-only 'training' flag (never
serialized; reloaded pipelines default to eval); training samples ONE
window per sample and reuses it across (timestep, view) frames - Isaac's
cross-view consistency
- gated on torch.is_grad_enabled() so no_grad validation and frozen-eval
paths are unaffected
- wired via dataset_meta is not None in make_groot_pre_post_processors and
the existing _set_groot_preprocessor_training on serialized reloads
Verification: tests/policies/groot/test_groot_train_random_crop.py (8 passed:
center-crop bit-exactness with crop_position=None, corner/center windows,
cross-view replay, train!=eval, no_grad gating, seed reproducibility,
serialization contract) + groot suite 23 passed / 5 skipped on RTX PRO 6000 /
CUDA 13.3.
* docs(groot): update Training & hardware Evaluation commands
Replace the multi-GPU accelerate-launch Training snippet with the current
single-command 'uv run lerobot-train' N1.7 recipe (relative actions excluding
gripper, bf16, flash attention, chunk/n_action_steps=16, bs64/20k steps).
Replace the bimanual 'Evaluate in your hardware setup' rollout example with the
SO-101 follower RTC 'uv run lerobot-rollout' command (strategy.type=base,
inference.type=rtc, wrist+front cameras, place-the-vial task).
Docs-only; no source/test changes.
* docs(groot): parameterize commands with env vars + fill LIBERO results
- Introduce BASE_MODEL / DATASET_ID / REPO_ID / JOB_NAME / OUTPUT_DIR env vars
in the training command and reuse OUTPUT_DIR + BASE_MODEL in the rollout cmd.
- Fill the LIBERO benchmark table with GR00T-LeRobot success rates
(Spatial 94%, Object 98%, Goal 93%, LIBERO 10/Long 90%; avg 93.75%),
drop the OSS column and XX placeholders. LeRobot-focused.
* docs(groot): drop export block, reference env vars directly
Use $DATASET_ID / $BASE_MODEL / $REPO_ID / $OUTPUT_DIR / $JOB_NAME as
bare placeholders in the commands without concrete export assignments.
* docs(groot): keep BASE_MODEL export in training command
* docs(groot): use literal HF repo IDs for dataset/policy repo_id
Public-facing Hub references (--dataset.repo_id, --policy.repo_id) shown as
concrete IDs; local-only values ($OUTPUT_DIR, $JOB_NAME) stay as placeholders.
* docs(groot): add LIBERO training command example
* docs(groot): remove LIBERO checkpoints subdirectory section
* docs(groot): use $BASE_MODEL for base_model_path in LIBERO eval
* docs(groot): drop hf download step from LIBERO eval, fix intro
* docs(groot): restore suite checkpoint download intro sentence
* docs(groot): remove checkpoint download note above LIBERO eval
* docs(groot): update training and rollout commands with new parameters and dependencies
* Add sample so101 training command
* Remove sample so101 training command
* docs(groot): remove optional Flash Attention setup instructions and update base model path for evaluation
* docs(groot): update training command with image transformation parameters
* docs(groot): add note on inference.queue_threshold value for stable inference
* chore(style): pre-commit gr00t
* docs(groot): update
* chore(policies): minor details
* fix(groot): license headers + test guards
* chore(policies): fix tests
* docs(groot): relative actions param doc
* chore(policy): address some of the AI review items
---------
Co-authored-by: Andrew Wrenn <awrenn@nvidia.com>
Co-authored-by: Ryan Halabi <ryhalabi@nvidia.com>
Co-authored-by: nv-sachdevkartik <ksachdev@nvidia.com>
Co-authored-by: groot-validation <groot-validation@localhost>
Co-authored-by: johnnynunez <johnnynuca14@gmail.com>
Co-authored-by: lbenhorin <lbenhorin@nvidia.com>
* feat(policies): add LingBot-VA autoregressive video-action world model
Port the LingBot-VA policy (Wan2.2 dual-stream video+action world model) into
LeRobot, following the EO-1 / VLA-JEPA conventions. Covers inference, checkpoint
conversion, and predicted-video saving (training is deferred to a follow-up PR).
- Vendored Wan transformer/attention/flex/VAE/scheduler modules (key names preserved
for near-identity conversion); torch SDPA default, flashattn/flex lazy-guarded.
- LingBotVAConfig (registered "lingbot_va") + processor with fixed-quantile action
unnormalization; full dual-stream sampling loop with CFG, two flow-matching
schedulers and KV cache, mapped onto select_action with observed-keyframe feedback.
- convert_lingbot_va_checkpoints.py (libero/robotwin variants): bundles the ~5B
transformer, lazy-pulls the frozen VAE+UMT5 from the source repo.
- Predicted-video plumbing in lerobot_eval (predicted_frames_callback; opt-in via
--policy.save_predicted_video) and ConstantWithWarmupSchedulerConfig.
- pyproject: widen diffusers-dep to <0.37, add lingbot_va + imageio-dep extras,
add lingbot_va and (missing) eo1 to `all`.
- Factory + policies/__init__ wiring, docs page + toctree, and tests.
Note: the LIBERO success-rate correctness gate must be validated on a CUDA GPU
with the converted checkpoint.
* feat(lingbot_va): RoboTwin eef-pose eval, single-file model, Hub checkpoints
Make the LingBot-VA port runnable on both LIBERO and RoboTwin and clean up the
package to LeRobot conventions.
- Consolidate all vendored Wan2.2 model code (transformer, attention, VAE helpers,
flow-matching scheduler, grid utils, flex-attention) into a single
modeling_lingbot_va.py; remove the separate wan_*/schedulers modules.
- Move the fixed action (un)normalization quantiles out of the config and into the
post-processor (LIBERO 7-DoF + RoboTwin 16-d eef); remove the conversion script in
favour of ready-to-use LeRobot-format checkpoints on the Hub.
- Fixes found via on-sim validation: undo LIBERO's 180-degree image flip
(image_hflip), encode obs as a multi-frame streaming-VAE clip, reset the streaming
VAE cache between episodes, run the transformer in config.dtype, lazy-load frozen
VAE/UMT5 by subfolder with the text encoder on CPU.
- RoboTwin: add an end-effector-pose action mode to RoboTwinEnv (16-d per-arm
xyz+quat+gripper deltas composed onto the initial eef pose, executed via CuRobo IK)
and the robotwin_tshape latent layout (full-res head + half-res wrists via a second
streaming VAE) with the upstream RoboTwin action quantiles + camera mapping.
- Predicted-video saving works for both benchmarks; docs + tests updated.
* feat(lingbot_va): implement training / fine-tuning (flow-matching loss)
- Implement LingBotVAPolicy.forward(): dual-stream flow-matching training loss
(latent + action, timestep-weighted, action-masked) ported from upstream train.py;
VAE-encodes camera clips, UMT5-encodes the task, noises both streams, runs the
block-causal flex-attention training pass (forward_train).
- training_loss_from_streams() core + _build_training_streams() data prep (action
scatter into the 30-d space, multi-frame VAE encode incl. robotwin_tshape).
- get_optim_params returns only trainable transformer params (LoRA/PEFT friendly);
VAE/UMT5 stay frozen. Training needs attn_mode='flex'.
- Add a tiny-config single-training-step test (forward->loss->backward->AdamW) and a
Training/fine-tuning section in the docs.
* fix(lingbot_va): CI quality gate + fast-test collection
- Add tests/policies/lingbot_va/__init__.py so the test files don't clash by basename
with tests/policies/vla_jepa/* under pytest's default import mode (fast-test collection error).
- Fix vendored typos flagged by the typos hook (pach_scale->patch_scale, total_tolen->
total_token_len, stablized->stabilized) and a mypy union-attr in RoboTwinEnv._read_eef_pose.
- Apply Prettier formatting to docs/source/lingbot_va.mdx.
* docs(lingbot_va): document EEF action-channel schema + camera order
* Update lingbot_va.mdx
Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.com>
* Update pyproject.toml
Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.com>
* Update pyproject.toml
Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.com>
* refactor(lingbot_va): drop hardcoded action quantiles; source from checkpoint
The LIBERO/RoboTwin action (un)normalization quantiles were hardcoded as module
constants in processor_lingbot_va.py. They are already serialized into each
checkpoint's policy_postprocessor.json (via LingBotVAActionUnnormalizeStep.get_config)
and restored on load by PolicyProcessorPipeline.from_pretrained, so the constants are
dead at eval/load time for the released checkpoints (verified: libero_long/robotwin/base
all carry their quantiles on the Hub).
- Remove LIBERO_ACTION_Q01/Q99, ROBOTWIN_ACTION_Q01/Q99 and _default_action_quantiles.
- make_lingbot_va_pre_post_processors now defaults a fresh (unconverted) build to a
neutral [-1, 1] mapping (identity rescale); real per-benchmark stats come from the
saved checkpoint (or postprocessor_overrides), analogous to dataset-stats normalization.
- Update the config doc comment to point at the checkpoint as the source of truth.
- Tests: replace the LIBERO-default assertion with a neutral-default check, and add a
save_pretrained/from_pretrained round-trip guard for the quantile serialization.
* docs(lingbot_va): trim verbose comments
- configuration_lingbot_va.py: condense multi-line field comments to one-liners
(keep the ── section headers).
- processor_lingbot_va.py: shorten the action-quantile explanation block.
- modeling_lingbot_va.py: drop the bare "# ----" separator rules, keeping the
one-line section headers.
No code changes.
* docs(lingbot_va): trim provenance comments; default wan path to base repo
- configuration_lingbot_va.py: drop the "──" decorations and the
"(from transformer/config.json)" note; default wan_pretrained_path to
robbyant/lingbot-va-base (has the frozen vae/text_encoder/tokenizer subfolders).
- modeling_lingbot_va.py: remove the vendored-code banner and the
"(upstream wan_va/...)" section-header provenance/dash decorations; condense the
transformer-dtype comment to one line.
No code changes.
* refactor(lingbot_va): use built-in UnnormalizerProcessorStep for actions
Replace the bespoke LingBotVAActionUnnormalizeStep with the standard
UnnormalizerProcessorStep in QUANTILES mode, which computes the identical
(action + 1) / 2 * (q99 - q01) + q01 mapping. The per-channel q01/q99 are stored
as the step's saved state (a safetensors file) and restored on load; a fresh build
has no action stats so the step is an identity passthrough.
The 3 Hub checkpoints (lerobot/lingbot_va_{libero_long,robotwin,base}) have been
re-uploaded with the new post-processor (policy_postprocessor.json +
*_unnormalizer_processor.safetensors); reloading from the Hub round-trips q01/q99.
- processor_lingbot_va.py: drop the custom step + registry; build the post-processor
with UnnormalizerProcessorStep (explicit ACTION->QUANTILES norm_map so the
preprocessor / training path is unchanged).
- tests: assert the built-in step is used, identity-when-no-stats, correct quantile
unnormalization, and a save_pretrained/from_pretrained stats round-trip.
* docs(lingbot_va): point checkpoint paths at the lerobot org
The LeRobot-format checkpoints moved from pepijn223/* to lerobot/* (libero_long,
robotwin, base). Update the eval/train --policy.path examples accordingly.
* docs(lingbot_va): condense processor normalization comments
* fix(lingbot-va): align RoboTwin evaluation (#3784)
Thank you for the RoboTwin fix, and alignment!
* applying fixes
* updating uv lock and linting
* adjusting test to match expected values
* cleaning up deps
* cleaning up top level imports, styling, and deps guards
* cleanup
* moving wan utils and loading utils to `utils.py`
* removing ftfy by replicating the prompt_clean function without it (we don't expect to have weird chars given in the prompt anyway)
* removing unused function
* guarding for scipy dep, renaming test to avoid collision
* adding back accelerate for peak memory usage optim + justifying robotwin description dep
---------
Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.com>
Co-authored-by: pepijn223 <pepijn223@hf.co>
Co-authored-by: Gangwei XU <gwxu@hust.edu.cn>
Co-authored-by: Maxime Ellerbach <maxime.ellerbach@huggingface.co>
Make the policy adapter architecturally clean and set up a single general
entry point for any language-conditioned policy.
Adapter architecture (Template Method):
- New lerobot/runtime/adapter.py: BaseLanguageAdapter owns the generic
control loop (throttle → generate → gibberish/empty reject → subtask→memory
cascade → diagnostics) and plan_from_text/handle_interjection. A policy
supplies only select_action + generate_text + build_messages. The
subtask→memory cascade is an overridable hook (_regenerate_context).
- GenerationConfig (typed, constructor-time) replaces config smuggled through
RuntimeState.extra (temperature/top_p/min_new_tokens/chunks_per_regen).
- LanguageDiagnostics (typed, keyed by kind) replaces ~8 loose state.extra
counter keys; the panel reads it via the adapter.
- looks_like_gibberish + split_plan_and_say move to runtime (generic).
Contract:
- LanguageConditionedPolicyAdapter protocol now states the true contract
(select_action, update_language_state, handle_interjection); the runtime
drops both getattr fallbacks.
- PI052PolicyAdapter shrinks to just its primitives (132 → ~half).
General entry point:
- lerobot/runtime/registry.py maps policy type → adapter (lazy import).
- run() resolves the adapter from the registry by policy type and defaults
the panel label to it, so one CLI serves every policy.
- Rename lerobot-pi052-runtime → lerobot-language-runtime (general script);
a new policy just registers its adapter, no new script.
Tests: new tests/runtime/test_adapter.py covers throttle/reject/cascade/
interjection; adapter + runtime + CLI-smoke tests updated for the new shape.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Set up the runtime so a second language-conditioned policy reuses the
CLI/REPL/UI instead of copying pi052's. The tick loop, REPL, panel, and
interactive CLI are now policy-independent in lerobot/runtime/; a policy
plugs in only a LanguageConditionedPolicyAdapter.
- Move repl.py, ui.py, and runtime_cli.py (-> cli.py) from
pi052/inference/ into lerobot/runtime/. Generalize labels/titles
(panel_label param, [runtime] prefixes).
- lerobot.runtime.cli.run(argv, *, adapter_factory, panel_label, prog)
is the shared entry; policy loading already dispatches generically via
the factory on cfg.type.
- lerobot-pi052-runtime is now a thin entry (scripts/lerobot_pi052_runtime.py)
that passes PI052PolicyAdapter into run(). pi052/inference/ keeps only
the adapter.
- Drop PI052Runtime back-compat wrapper (no consumers).
- Drop VQA visualization: delete inference/vqa.py + test_pi052_vqa_loc.py,
remove answer_vqa/VQAResult from the Protocol + adapter, and the
/question command + overlay paths from the CLI/REPL.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* Add FastWAM policy
* Add FastWAM policy review updates
* big refactor to use models from diffusers and transformers
* changing reproducable results
* preparing for training adding some temporary debug code aswell to visualize model output
* re-parenting of some layers to enable proper zero-3 FSDP
* linting
* small fix for the preprocessor and padded images
* removing some preprocessors
* removing temporary debug code
* cleaning up
* updating uv lock after rebasing
* adding lazy imports
* linting
* fixing stale assertion
* make tokenizer/text-encoder model ids configurable + some nits
* moving and renaming files to have a cleaner file tree
* removed asserts from the model, added guard instead and completely removed useless asserts
* cleaning up imports
* removing is_main_process and custom logging logic
* removing unused / stale attention path, removing some of the stale forwards within wan/models
---------
Co-authored-by: ZibinDong <zibindong@outlook.com>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
* refactor(vla-jepa): removing gpu roundtrip for the preprocessing part
* major refactor of the forward pass and model input conversion
* linting
* adressing suggestions from reviews
* removing redundant state dtype conversion
* avoiding recreating the same tensor each foward pass
* api simplification of `_encode_qwen`
* avoiding useless video assembly during inference
* guard against video=None for the wm loss
Align the MolmoAct2 implementation with lerobot codebase conventions:
- Rename hf_model/ to molmoact2_hf_model/
- Slim config: move all I/O and runtime logic to modeling
- Remove blanket from 8 vendored files, fix 66 lint issues
- Deduplicate _hf_token() and _resolve_checkpoint_location()
- Make huggingface_hub imports lazy
- Remove custom MolmoAct2CosineDecayWithWarmupSchedulerConfig, use base class
- Extract 13 static/classmethods from MolmoAct2Policy to free functions
- Replace print() with logger in vendored action_tokenizer
- Add module docstrings, class docstring, and key method docstrings
- Add module-level loggers to modeling and processor
- Fix docs: pip to uv install, deduplicate README symlink
- Remove shebangs from all files
* feat(train): FSDP checkpoint saving
* adding docs for FSDP
* adding a test for the fsdp checkpoint path
* cleanup
* fixing final upload to hub
* refactored initial implementation to use torch fsdp api and adding new tests
Eliminate the standalone pi052/pi05_backbone.py by distributing its contents:
- Generic dual-expert transformer machinery -> lerobot/policies/pi_gemma.py
(sdpa_attention_forward, compute_layer_complete, PaliGemmaWithExpertModel,
get_gemma_config; the openpi width/depth config is renamed GemmaConfig ->
GemmaVariantConfig to avoid clashing with transformers' GemmaConfig). These
sit next to the existing PiGemma layer code they already depend on.
- pi052-specific model + helpers -> pi052/modeling_pi052.py (PI05Pytorch,
ActionSelectKwargs, make_att_2d_masks, pad_vector, resize_with_pad_torch,
create_sinusoidal_pos_embedding, sample_beta, get_safe_dtype).
DEFAULT_IMAGE_SIZE is duplicated as a plain constant in pi_gemma to avoid a
pi_gemma -> pi05 import cycle. Additive to pi_gemma; pi0/pi05 unaffected.
Verified bit-exact on pepijn223/pi052_robocasa_full (embed/predict/forward
identical) and all 34 pi052 tests pass.
Co-authored-by: Cursor <cursoragent@cursor.com>
_fast_ce/_shifted_ce were renamed to _fast_lin_ce/_shifted_lin_ce and changed
from logits-based to Liger fused-linear-CE (hidden @ lm_head_weightᵀ). Update
the tests via thin adapters that pass an identity lm_head_weight (so the
computed logits equal the provided ones), run on CUDA (Liger is GPU-only) and
skip otherwise, and loosen the allclose tolerance to absorb GPU-vs-CPU float
noise on the tiny losses.
Co-authored-by: Cursor <cursoragent@cursor.com>
The smolvla branch had modified the shared pi0/pi05 modeling + pi05 config to
support pi052 (SDPA attention, layernorm/lm_head handling, optimizer
foreach/fused/lm_head_lr_scale, embedding scaling). Decouple pi052 instead:
- Vendor the PI0.5 backbone (PaliGemmaWithExpertModel, PI05Pytorch, helpers)
into pi052/pi05_backbone.py (verbatim copy, no PI05Policy).
- Flatten PI052Policy to subclass PreTrainedPolicy directly (no longer
PI05Policy); inline the needed PI05Policy methods.
- Restore optimizer_foreach/fused + get_optimizer_preset on PI052Config.
- Revert pi0, pi0_fast, pi05 modeling and configuration_pi05 to origin/main
(byte-identical), so the shared policies carry no smolvla modifications.
Behavior verified bit-exact on pepijn223/pi052_robocasa_full: embed_language_
tokens, predict_action_chunk, and the fused flow+text+FAST training loss are
identical before/after (max_abs_diff=0). pi052 tests pass (pre-existing
stale-name collection errors unchanged).
Co-authored-by: Cursor <cursoragent@cursor.com>
* first commit
* feat(policies): add VLA-JEPA
* feat(policies): add VLA-JEPA
* support vla_jepa
* (feat)policies: add VLA-JEPA
* linting
* adding deps to pyproject.toml
* updating uv lock
* adding guards to avoid needing transformers and diffusers for type checking and basic tests
* fixing action and state dim
* fix warnings with qwen processor kwargs
* fixing wm_loss not propagating
* adjusting obs steps, tublets size to match original implementation
* some more fixes to be closer to the original implem
* adding more tests to ensure good coverage
* align VLA-JEPA architecture with original checkpoint
- Remove stale `action_num_heads` / `action_attention_head_dim` config fields;
DiT head dimensions are now always derived from the preset (DiT-B/L/test).
- Add `num_target_vision_tokens` and `action_max_seq_len` config fields required
by the action head's future-token embedding and positional embedding tables.
- Fix default `qwen_model_name` to 2B (matches all released checkpoints).
- Rename `ActionEncoder` attrs w1/w2/w3 → layer1/layer2/layer3 to match
checkpoint key names; replace `nn.Sequential` decoder/state-encoder with
`_MLP2` (layer1/layer2 naming).
- Fix `VLAJEPAActionHead` to size ActionEncoder and StateEncoder at `inner_dim`
(DiT input width) rather than `action_hidden_size` (DiT output width).
- Rename `DiT.blocks` → `transformer_blocks` and `attn` → `attn1` to match
checkpoint; add alternating cross/self attention (even blocks cross-attend to
Qwen context, odd blocks self-attend).
- Add `DiT-test` preset for unit tests.
- Rewrite `ActionConditionedVideoPredictor` with explicit ViT-style blocks
(`_PredictorBlock` with fused qkv) to match checkpoint structure; rename
`encoder`/`norm`/`proj` → `predictor_blocks`/`predictor_norm`/`predictor_proj`.
* propagate action_is_pad masking through VLA-JEPA policy pipeline
Pass the `action_is_pad` tensor from the batch through to the action head
so padded timesteps are excluded from the flow-matching loss.
* update VLA-JEPA tests for arch changes and action_is_pad
- Switch conftest to use `action_model_type="DiT-test"` now that
`action_num_heads` / `action_attention_head_dim` have been removed.
- Add action_head tests covering fully-padded loss (zero) and equivalence
of action_is_pad=None vs all-zeros mask.
- Remove obsolete `test_native_to_lerobot_wm_only` test.
* add VLA-JEPA documentation
Covers architecture overview, pretrained checkpoints, config reference,
training/eval commands for LIBERO-10, and guidance on fine-tuning for
single-camera datasets.
* add one-shot script to convert ginwind/VLA-JEPA checkpoints to safetensors (will remove once migrated)
* make default params more aligned with paper and pretrained models
- adding possibility of freezing qwen backbone and world model
- added tests for weight loading
* trying out to re-init the action head to avoid pretraining dimension mismatch
* allow different state dim and action dim
* removing missleading future_action_window_size to just use chunk_size
* lots of changes to make existing weights work, need to massively refactor the pre and post processing
* refactoring into using pre and post processor
* pre-commit cleanup
* fixing doc defaults args
Signed-off-by: Maxime Ellerbach <maxime@ellerbach.net>
* adressing dtype zeros issue
* adding guard for diffusers
* fixing training and exal examples
* trying to close success rate gap
* fix qwen norm layer output libero eval is now as expected
* adding instructions for different embodiement + fixing some tests
* smol fix to avoid having default CPU device when training
* fixing misconception about multiview / singleview handling
* removing conversion script
* adding licences
* adding .mdx docs and shortening polivy_vla_jepa_README.md
* removing useless pre-processor
* cleanup
* removing swish in favor of silu
* adding configuration gripper index and threshold
* fixing simlink
---------
Signed-off-by: Maxime Ellerbach <maxime@ellerbach.net>
Co-authored-by: ginwind <ginwind@mail.ustc.edu.cn>
Bring the authoritative annotation pipeline from the annotation branch.
The annotation surface is forced to EXACTLY match feat/language-annotation-
pipeline (the annotation branch is the source of truth for annotation
code), which also removes smolvla's stale copies:
- deleted: steerable_pipeline/vocabulary.py, tests/annotations/test_
vocabulary.py, prompts/module_0_vocabulary.txt, module_1_action_record
.txt, module_3_vqa.txt, module_1_plan.txt, and the old module_* prompt
names (now plan_*/interjections_*/vqa.txt).
- synced: all of src/lerobot/annotations/, lerobot_annotate.py,
examples/annotations/, tests/annotations/, datasets/language.py,
tests/datasets/test_language.py, docs/annotation_pipeline.mdx.
Non-annotation conflicts resolved by union (keeping both branches' intent):
- pyproject.toml: keep smolvla's pi extra (+sentencepiece) and add the
molmoact2 extra from main.
- policies/factory.py: keep both dataset_repo_id (pi052 FAST tokenizer)
and dataset_meta (both are referenced); union the policy-type docstring.
- scripts/lerobot_train.py: keep smolvla's pi052 / use_relative_actions
processor-rebuild block.
- uv.lock: regenerated from the merged pyproject.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Replaces the per-layer ``modeling_gemma.eager_attention_forward`` call
with ``torch.nn.functional.scaled_dot_product_attention`` in
``compute_layer_complete`` (pi05) and ``_compute_layer_ki`` (pi052).
PyTorch SDPA picks the memory-efficient kernel for the
block-bidirectional 4D additive mask the dual-expert model uses (FA2 /
FA3 reject it because they only accept causal / sliding-window / varlen
patterns). The shared ``sdpa_attention_forward`` helper mirrors the
eager signature so the call sites are unchanged.
Selective AC: removes the redundant outer ``_apply_checkpoint(forward_func, ...)``
wrap in ``PI05Pytorch.forward``. Per-layer checkpointing inside
``PaliGemmaWithExpertModel.forward`` already handles activation
recompute; the outer wrap was double-recomputing the whole backbone.
+14% steps/sec on its own (job 22161405 vs 22161398, 1xH100).
groot: drop ``@strict`` on ``GR00TN15Config`` — newer ``huggingface_hub``
rejects ``@strict`` on non-dataclass ``PretrainedConfig`` subclasses,
which was blocking imports of any sibling policy through
``lerobot.policies.factory``.
New ``examples/benchmark/bench_pi052_step.py`` (+ slurm sweeps v1..v8)
times PI052Policy.forward+backward (optionally with AdamW) on
synthetic inputs. Headline numbers on 1xH100 with KI=True, GC=True,
L=512, 4.14 B trainable params, AdamW state in bf16:
pre-SDPA eager BS=8 610ms 19.5 GiB -> 13.1 samples/s
sdpa BS=8 + compile=default 413ms 19.5 GiB -> 19.3 samples/s
sdpa BS=16 + compile=default 715ms 37.3 GiB -> 22.4 samples/s
sdpa BS=32 + compile=default 1325ms 44.8 GiB -> 24.2 samples/s
sdpa BS=40 + compile=default 1665ms 48.6 GiB -> 24.0 samples/s
Parity tests in ``tests/policies/pi052/test_pi052_sdpa_attention.py``
cover fp32 / bf16 / GQA / MHA forward + backward — output and grads
match the eager path within bf16 tolerance.
Also ships ``examples/benchmark/fsdp_pi052.yaml`` (FSDP2 accelerate
config wrapping GemmaDecoderLayer + SiglipEncoderLayer) for the
follow-up multi-GPU memory sharding work.
Co-authored-by: Cursor <cursoragent@cursor.com>
Parallel variant of build_robocasa_composite_seen.py modeled after the
existing slurm_port_shards.py / slurm_aggregate_shards.py pattern.
Two-phase datatrove pipeline:
* Phase 1 DOWNLOAD: tasks=16 (one per RoboCasa composite_seen task),
each worker downloads its assigned tar via RoboCasa's own
download_datasets helper. Network-bound, idempotent.
* Phase 2 AGGREGATE: tasks=1, single worker calls aggregate_datasets
over the 16 extracted directories. Submitted with depends=phase1 so
SLURM only releases it once all 16 downloads succeed.
Reuses the COMPOSITE_SEEN_TASKS list and per-task download/resolve
helpers from the single-machine script via aliased imports — single
source of truth for 'what does it mean to download a composite_seen
task'.
Local (--slurm 0) mode runs the two phases sequentially in-process for
debugging on a workstation.
Usage on SLURM:
uv run python examples/port_datasets/slurm_build_robocasa_composite_seen.py \
--output-dir=/scratch/${USER}/robocasa_composite_seen \
--hub-repo-id=${HF_USER}/robocasa_composite_seen \
--logs-dir=/scratch/${USER}/logs/robocasa \
--partition=cpu --push-to-hub
Prereq: uv sync --extra annotations (pulls datatrove)
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* chore(gr00t): sync with #3606 for fixing gr00t config crash
* fix(pi0&pi05): fix graph break caused by deepcopy of past_key_values in sample_actions
* fix(pi0&pi05): fix frequent recompile caused by compute_layer_complete
* feat(test): add compile test and benchamrk for pi0 and pi05
* feat(test): add comprehensive testing for pi0 and pi05. Including processor, forward, sample action, etc.
The trained model collapsed to spewing 40+ <loc> tokens for *every*
prompt — subtask, memory, anything — because VQA targets were supervised
to *start* with <loc>. With ~25% of all text samples beginning with a
<loc> token, the LM head learned "Assistant: → <loc>" as a strong
attractor; once one loc is emitted, autoregression chains the rest.
Flip the format so every text target — subtask, memory, speech, AND VQA
— starts with a regular word. The model still learns the <loc>
vocabulary for the spatial portion of the answer, but loc can no
longer be the first generation step out of a clean prompt.
Examples:
point : "green box <loc0162><loc0759>"
bbox : "cube <loc0082>…<loc0409>"
multi : "blue <locs> ; yellow <locs>"
The runtime parser (parse_loc_answer) strips loc tokens and uses the
remainder as label, so it's order-tolerant and works under either
format. Old loc-first checkpoints still parse cleanly at inference;
new training will use label-first.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
THE bug behind the <loc>-salad. PaliGemma's vocab reserves ids
[256000, 257023] for <locDDDD> detection / pointing tokens, but the
stock AutoTokenizer does NOT match them on raw text — it BPE-splits
<loc0162> into SEVEN pieces (<, loc, 0, 1, 6, 2, >). So a VQA target
like "<loc0162><loc0759> green box<eos>" tokenized to 16 pieces, not
5, and training the LM head supervised those generic BPE pieces
instead of one detection-vocab id. The piece logits got pumped up
across ~25% of supervised positions; at inference they dominated
every turn — even subtask prompts produced <loc>-salad followed by
the actual answer.
Register the 1024 <locDDDD> tokens via tokenizer.add_tokens once on
load, in every path the policy uses: PI052TextTokenizerStep (training
encode), _build_text_batch_pi052 (runtime encode), and
select_message's default tokenizer (runtime decode). Verified
empirically with the real PaliGemma tokenizer: VQA target now
tokenizes to 5 ids matching the loc-vocab range (256162, 256759, ...)
with correct offset_mapping.
This unlocks PaliGemma's actual detection prior; <loc>-salad cannot
recur because each <locDDDD> is a single class on the LM head, not a
character sequence the head accidentally learns to extend.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Confirmed empirically on the published dataset: VQA bbox/keypoint
coordinates are Qwen2.5-VL's 0–1000 normalized grounding output, NOT
pixels. Scanning 8207 samples showed x and y both spanning 0..1000
with ~30% of values exceeding the camera's pixel dimensions (which is
impossible if they were pixels).
_vqa_answer_to_loc was dividing by the observation image's H/W, so
e.g. point [742, 158] on a 640x480 wrist cam clamped x to <loc1023>
(the far-right edge) instead of mapping to <loc0760> (~74% across).
Fix: divide by 1000 — the actual Qwen scale. The conversion is now
camera-resolution-independent, so _camera_image_shapes and the
image_shapes plumbing through __call__ / _encode_messages /
_messages_vqa_to_loc are dropped. Tests updated to the new signature
and the 0–1000 round-trip.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Spatial VQA answers (bbox / keypoint) were trained as pixel-coordinate
JSON, which fights PaliGemma's detection prior and leaks <loc>-token
salad at inference. Convert them to PaliGemma's native <locNNNN>
vocabulary instead so the LM head reuses that prior.
Training side (text_processor_pi052.py): a target turn whose content
parses as a bbox/keypoint answer is rewritten to <loc> text, using the
camera frame's native (H, W) from the observation and the preceding
image block. Non-spatial answers, subtask/memory targets and SmolVLA2
keep their JSON form — the dataset stays backbone-agnostic.
Runtime side (smolvla2/inference/vqa.py): parse_vqa_answer detects
<loc> answers (2 locs -> keypoint, 4 -> bbox), returning normalized
[0,1] coords with a normalized flag; draw_vqa_overlay denormalizes
against the chosen camera frame's pixel size.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
PI052TextTokenizerStep masked text_labels over the assistant turn's
*content only* — the trailing newline was excluded and no EOS token was
ever a supervised label. So the LM head was never given a stop signal:
at inference select_message decoded to max_new_tokens, producing the
runaway subtask paragraphs and the "}"}"}-style VQA tails.
_format_messages now appends the tokenizer's EOS to each supervised
target turn and extends that turn's span to cover it, so the EOS lands
in text_labels. _shifted_ce then trains "<last content token> -> EOS"
and the model learns to terminate; select_message stops on it.
Inference callers (the runtime's _build_text_batch_pi052) pass no
target_indices / eos_token, so no EOS is baked into the prompt — the
model generates it. Verified end-to-end with the PaliGemma tokenizer:
the supervised span is `<content><eos>` and the trailing newline stays
unsupervised.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Select only supervised text and FAST action-code positions before cross-entropy to avoid full-vocabulary loss tensors over padded sequences.
Co-authored-by: Cursor <cursoragent@cursor.com>
Tokenize batched recipe outputs in PI052 so training batches with nested message lists do not crash before model forward.
Co-authored-by: Cursor <cursoragent@cursor.com>
Mask the FAST auxiliary loss to discrete action-code tokens so wrapper formatting tokens do not affect action co-training.
Co-authored-by: Cursor <cursoragent@cursor.com>
pi052 had the same text-CE collapse bug smolvla2 had — PaliGemma's
embed_prefix flags the language block att=0, so make_att_2d_masks makes
it fully bidirectional and the text cross-entropy degenerates into a
copy task. Ported the three model-specific fixes:
- _mark_target_span_causal: set att=1 on supervised target language
positions so the text-CE is genuine causal next-token prediction.
Applied in both _compute_all_losses_fused and _compute_text_and_fast_loss.
- flow_loss_weight 10.0 -> 5.0: the paper's a=10 swamps the LM head once
the flow-only low_level recipe fires often (matches SmolVLA2Config).
- _flatten_say_tool_calls in the text tokenizer: serialize `say` tool
calls into a <say>...</say> marker so the spoken reply is tokenized
and supervised (PaliGemma's flat prompt has no structured calls, so
they were dropped entirely).
select_message needed no change: pi052's prefix is [images, language]
with no trailing state token, so it already decodes from the last
language token.
Regression tests mirror the smolvla2 attention-masking + tool-call suite.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Three additions to the SmolVLA2 interactive runtime:
1. Startup task picker — when no --task is given, the runtime lists the
dataset's task strings as a numbered menu (plus a custom-task option)
instead of silently waiting for the first stdin line.
2. Mode toggle — /action and /vlm slash commands flip a persistent run
mode. /vlm pauses the whole action loop (HighLevelSubtaskFwd,
LowLevelForward and DispatchAction gate on state["mode"]) and clears
the action queue so the robot holds position; /action resumes it.
The mode is shown in the state panel.
3. Interactive VQA — in /vlm mode a typed line is a VQA question. The
new inference/vqa.py module asks which camera to ground on, runs the
VLM on that single camera, and when the answer is a bbox/keypoint it
draws the overlay, saves a PNG to ./vqa_overlays/ and auto-opens it.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Pepijn Kooijmans