Commit Graph

2 Commits

Author SHA1 Message Date
Maxime Ellerbach 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>
2026-07-31 11:22:00 +00:00
Pepijn 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>
2026-07-03 13:32:38 +02:00