mirror of
https://github.com/huggingface/lerobot.git
synced 2026-08-02 14:39:57 +00:00
c15f6acb7e3954a7e63924e5f67f6bddb537c68d
2 Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
|
|
c15f6acb7e |
lingbot-va: increment-encoded action KV memory under relative actions
LingBot-VA's action stream doubles as its memory: _compute_kv_cache feeds the previous chunk's predicted actions back as clean context. With use_relative_actions each chunk is anchored on the state at its own start, so the same physical motion encodes differently depending on which chunk it fell in, and the cached stream sawtooths back to zero displacement at every boundary (mean 0.22, p95 0.91 normalized units against a total output range of 2.0). Training never contains such a reset -- one sample is exactly one chunk, hence one anchor -- and the model's only cross-frame supervision is "continue from the previous clean action frame", so it continues rather than restarting and the postprocessor double-counts the displacement: ~7.3 deg/chunk on the folding data, matching the observed monotonic drift after 4-5 chunks. First-difference the conditioning instead. An increment carries no anchor, so the cached history stays consistent across chunks; measured boundary inconsistency drops 7.30 -> 1.78 deg, the latter being the controller's tracking error, which is generic to chunked control rather than an artifact of the action encoding. The conversion touches ONLY the conditioning -- at inference in _preprocess_action_state, and on the matching clean copy in _add_noise_stream so training agrees. Predicted actions keep the cumulative-offset encoding, so the processors, the checkpoint pipeline and the executed commands are untouched. That is also what makes it safe to do here: normalization is affine with a non-zero constant, so cumsum does not commute with unnormalization -- inverting an accumulated stream would drift by (n-1)*(q99+q01)/2, which is 0 on a symmetric channel but 118 deg over one chunk on right_joint_6. Nothing inverts the conditioning, so that cannot arise. Gated on use_relative_actions only: an absolute action stream is already globally consistent, so those checkpoints keep their exact behaviour. A relative checkpoint is only valid for the encoding it was trained with. Also fixes observation_delta_indices: AutoencoderKLWan._encode consumes only the first 4 * (iter_ - 1) + 1 frames of an n-frame clip, so asking for frame_chunk_size * 4 decoded 3 frames per sample that never reached the encoder (verified by ablation -- scrambling them left the latents bit-identical). Same latent frame count, less video decode. The observation window stays shorter than the action window, so sample validity is unchanged. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> |
||
|
|
e275ea3960 |
LingBot-VA: video-action world model (#3731)
* 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> |