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(annotate): WGO-tuned subtask prompt (atomic completed-events + duration prior)
Rework the plan-module subtask segmentation prompt toward the WGO-Bench
atomic annotation protocol: segment by completed world-state changes
(grasp/place/open/close/pour/insert), fold approach+retreat into their
event, keep separate events separate, and add a 2-10s duration prior.
Drops the pi0.7 "fewer larger composites preferred" bias that drove
under-segmentation on the benchmark. Output JSON shape unchanged.
Co-authored-by: Cursor <cursoragent@cursor.com>
* feat(annotate): seeded-relabeling second pass for subtasks
Add an opt-in relabel pass (plan.subtask_seeded_relabel) that, after
segmentation, re-labels each span using previous/current/next segment
contact sheets and the seed label as a strong prior, minimally correcting
it. Mirrors macrodata's best end-to-end labeling step. Boundaries are
untouched; one extra VLM call per span. Off by default.
Co-authored-by: Cursor <cursoragent@cursor.com>
* feat(annotate): robust OpenAI-compat client for hosted VLMs
Guard against a choice with no message (safety filter or a thinking model
that spends its whole budget before emitting content) so one empty reply
no longer crashes the whole annotation run; treat it as an empty response
and let the existing JSON-retry path handle it.
Add an optional `reasoning_effort` knob on VlmConfig, forwarded to the
server when set, to cap a thinking model's reasoning (needed for Gemini
via its OpenAI-compatible endpoint).
Co-authored-by: Cursor <cursoragent@cursor.com>
* feat(annotate): legible tile-scaled timestamp on contact sheets
The burned-in timestamp used the ~10px bitmap default font, which blurs
once the model downsamples a full contact sheet into 768px tiles, so the
VLM can no longer read the exact source time a boundary depends on. Scale
the timestamp to the tile height (with a graceful fallback on older
Pillow) so the visual time cue stays readable at sheet resolution.
Co-authored-by: Cursor <cursoragent@cursor.com>
* feat(annotate): lean GEPA-aligned subtask segmentation prompt
Replace the verbose, label-heavy segmentation prompt with a lean
adaptation of the blog's GEPA-found completed_events_duration_prior
recipe: focus on completed manipulation events, explicit no-split /
no-merge rules, a 2-10s duration prior, and an instruction to prioritize
temporally correct boundaries over label wording. The previous prompt
over-weighted label guidance, which traded away boundary precision.
Co-authored-by: Cursor <cursoragent@cursor.com>
* revert: restore original subtask segmentation prompt
The lean GEPA-aligned paraphrase (dd4b0110d) regressed Gemini on the
30-ep subset: Seg F1 0.259 -> 0.189 and E2E 0.184 -> 0.135, driven by
worse under-segmentation (224 -> 188 preds). The blog's 0.306 came from
the actual GEPA-search artifact, which a hand paraphrase does not
reproduce. Restore the original prompt, which remains our best config.
Co-authored-by: Cursor <cursoragent@cursor.com>
* feat(annotate): env-var override for prompt templates
Allow LEROBOT_PROMPT_OVERRIDE_<name> to supersede the packaged prompt
file at load time. Enables prompt search (GEPA) to inject candidate
segmentation prompts into a remote annotate job via an env secret,
without committing a branch per candidate.
Co-authored-by: Cursor <cursoragent@cursor.com>
* docs(annotate): genericize hosted-VLM comments (no model name)
Co-authored-by: Cursor <cursoragent@cursor.com>
* docs(annotate): document seeded-relabel and reasoning_effort flags
Co-authored-by: Cursor <cursoragent@cursor.com>
* test(annotate): update subtask-prompt marker to match WGO-tuned prompt
The three plan-module tests keyed the canned VLM responder on the
literal 'atomic subtasks', which the WGO-tuned segmentation prompt no
longer contains (it now segments 'COMPLETED manipulation events'). Point
the fixture markers at the current wording so the subtask call is matched
again.
Co-authored-by: Cursor <cursoragent@cursor.com>
---------
Co-authored-by: Cursor <cursoragent@cursor.com>
Some checkpoints ship config.device=cpu (e.g. MolmoAct2 SO100/101). The
language runtime had no device override, so it always ran on the config
device. --policy.device=cuda (or cpu) now overrides cfg.device at load.
Co-authored-by: Cursor <cursoragent@cursor.com>
Load normalization stats from the checkpoint (norm_tag) and derive the
observation/action feature schema from the connected robot when no dataset
is given, mirroring lerobot-rollout. A dataset is still honoured when
supplied and its stats take precedence.
Co-authored-by: Cursor <cursoragent@cursor.com>
Enable running MolmoAct2 policies (e.g. on an SO101) in the interactive
language runtime with direct-subtask prompting.
- policies/molmoact2/molmoact2_adapter.py: MolmoAct2PolicyAdapter — flat VLA
bridge; select_action predicts an action chunk from the packed observation,
generate_text is a no-op (no text head; use --direct_subtask).
- runtime/registry.py: register "molmoact2" -> MolmoAct2PolicyAdapter.
- runtime/cli.py:
- Preserve model-input keys emitted outside observation.* (MolmoAct2 packs
the prompt+images into input_ids/pixel_values/...) through the robot
observation filter; no-op for PI0-family policies.
- Robot observation provider now reads the live task/subtask each frame via a
get_task callback, so a typed command re-packs the instruction (also fixes
stale-task for other flat VLAs). Bound to runtime state after creation.
Co-authored-by: Cursor <cursoragent@cursor.com>
Newer PI052 training runs serialize use_flex_attention into config.json.
This branch's attention path is SDPA/eager (mathematically equivalent), so
the field is accepted as an inert no-op (mirrors the existing use_hf_kernels
compat field) — otherwise loading those checkpoints raises DecodingError.
Co-authored-by: Cursor <cursoragent@cursor.com>
Wire the existing (but previously unreachable from the runtime) PI052
FlashRT FP8 MLP swap into the language runtime. --fp8 sets
config.use_flashrt_fp8_mlp before load; the policy calibrates and swaps
every Gemma + SigLIP MLP to fused FP8 on its first predict_action_chunk.
Ignored with a warning for policies without the flag (PI052 only).
Measured ~1.12x faster action-chunk inference (124 -> 111 ms) on an
RTX 5090; needs the `kernels` package (pin <0.13 for transformers) and
CUDA SM>=8.9, else it degrades to BF16.
Co-authored-by: Cursor <cursoragent@cursor.com>
Add physical-robot support to the language runtime, plus a live rerun viewer.
- runtime/rerun_viz.py: headless rerun gRPC + web viewer; logs camera frames
(every control tick) and joint state. Prints an auto-connect ?url= view URL.
- runtime/cli.py:
- _run_robot_interactive: real-time control loop (background thread) with a
clean chat prompt — a typed command switches task/subtask immediately and
regenerates. Starts running as soon as a task is set (via --task or the
picker); otherwise paused until the first command. No flag needed.
- --rerun (+ --rerun.web_port / --rerun.grpc_port): live camera view; the
robot obs provider and action executor log frames to rerun.
- --direct_subtask (general, sim or robot): the typed text is the subtask fed
to the action expert; the LM subtask generator is disabled.
- Inference overrides: force compile_model=False and gradient_checkpointing
=False (torch.compile recompiled on every prompt-length change -> >1min per
chunk; grad checkpointing only slows the forward pass).
Co-authored-by: Cursor <cursoragent@cursor.com>
Drive a persistent RoboCasa kitchen with open-ended prompts and watch it live.
- runtime/sim_robocasa.py: single-scene RoboCasa backend (n_envs=2 for stable
EGL rendering — single-worker rendering is broken), high-res multi-view
compositing incl. wrist cam, annotated MP4 + rolling latest.png + MJPEG live
viewer, and /reset scene re-roll.
- runtime/cli.py: --sim mode with a main-thread control loop (background-thread
rendering corrupts EGL), clean chat-style prompt (a new command switches the
task and regenerates the subtask immediately), plus --sim.render_size,
--sim.views, --sim.stream_port, --sim.direct_subtask and --disable_memory.
- runtime/adapter.py: GenerationConfig.enable_memory / enable_subtask toggles.
- runtime/registry.py + policies/pi05/pi05_adapter.py: register pi05 (flat VLA,
direct task-text conditioning; no subtask/memory head).
- policies/pi052/inference/pi052_adapter.py: condition the action expert on
"{subtask}, State: {..}" to match eval/training.
- envs/robocasa.py + envs/configs.py: terminate_on_success + horizon options so
the interactive kitchen persists across tasks (defaults preserve eval).
Co-authored-by: Cursor <cursoragent@cursor.com>
Drop the opt-in EMA-shadow feature entirely: EMAConfig, the `ema` field on
TrainPipelineConfig, all EMA logic in lerobot_train.py (setup/resume, per-step
update, W&B observability, checkpoint save, EMA-model eval, and the sibling
`<repo_id>-ema` hub push), and the ema-pytorch dependency.
Co-authored-by: Cursor <cursoragent@cursor.com>
Keep PI052Policy.forward's loss components as detached tensors and only
materialize loss/grad_norm/update_s to python floats on logging steps
(1-in-log_freq) via a new update_policy(log_metrics=...) gate. Also dedupe
the predict_actions .any().item() control-flow sync (2 -> 1 per step).
Keeps the training step fully async on non-logging steps so the next batch's
dataloading/enqueue overlaps GPU compute instead of stalling on a per-step
CUDA sync.
Co-authored-by: Cursor <cursoragent@cursor.com>
pepijn