mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-21 08:51:47 +00:00
Pepijn
4dfa8cea65
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. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
53 lines
1.9 KiB
Python
53 lines
1.9 KiB
Python
#!/usr/bin/env python
|
|
|
|
# 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 __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from lerobot.policies.factory import make_policy_config
|
|
from lerobot.policies.lingbot_va.configuration_lingbot_va import LingBotVAConfig
|
|
|
|
|
|
def test_make_policy_config_returns_lingbot_va() -> None:
|
|
cfg = make_policy_config("lingbot_va", device="cpu")
|
|
assert isinstance(cfg, LingBotVAConfig)
|
|
|
|
|
|
def test_get_policy_class_resolves_lazily() -> None:
|
|
# Importing the policy class pulls in diffusers (Wan2.2 stack); skip if unavailable.
|
|
pytest.importorskip("diffusers")
|
|
pytest.importorskip("transformers")
|
|
from lerobot.policies.factory import get_policy_class
|
|
|
|
cls = get_policy_class("lingbot_va")
|
|
assert cls.name == "lingbot_va"
|
|
assert cls.config_class is LingBotVAConfig
|
|
|
|
|
|
def test_convert_build_config_libero() -> None:
|
|
pytest.importorskip("diffusers")
|
|
from lerobot.policies.lingbot_va.convert_lingbot_va_checkpoints import build_config
|
|
|
|
cfg = build_config("libero", wan_pretrained_path="dummy/path", dtype="float32")
|
|
assert cfg.height == 128 and cfg.width == 128
|
|
assert cfg.used_action_channel_ids == list(range(7))
|
|
# validate_features (called inside build_config) must have populated the action feature.
|
|
from lerobot.utils.constants import ACTION
|
|
|
|
assert cfg.output_features[ACTION].shape == (7,)
|
|
assert len(cfg.obs_cam_keys) == 2
|