From e275ea3960332543e2a9f441356775a53720543f Mon Sep 17 00:00:00 2001 From: Pepijn <138571049+pkooij@users.noreply.github.com> Date: Fri, 3 Jul 2026 13:32:38 +0200 Subject: [PATCH 01/11] LingBot-VA: video-action world model (#3731) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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 Co-authored-by: Gangwei XU Co-authored-by: Maxime Ellerbach --- .dockerignore | 4 + docs/source/_toctree.yml | 2 + docs/source/lingbot_va.mdx | 187 +++ pyproject.toml | 2 + src/lerobot/envs/configs.py | 8 +- src/lerobot/envs/robotwin.py | 175 ++- src/lerobot/optim/schedulers.py | 22 + src/lerobot/policies/__init__.py | 2 + src/lerobot/policies/factory.py | 15 + src/lerobot/policies/lingbot_va/README.md | 1 + src/lerobot/policies/lingbot_va/__init__.py | 21 + .../lingbot_va/configuration_lingbot_va.py | 168 +++ .../lingbot_va/modeling_lingbot_va.py | 853 ++++++++++++ .../lingbot_va/processor_lingbot_va.py | 87 ++ src/lerobot/policies/lingbot_va/utils.py | 1138 +++++++++++++++++ src/lerobot/scripts/lerobot_eval.py | 101 +- .../test_configuration_lingbot_va.py | 78 ++ tests/policies/lingbot_va/test_factory.py | 38 + tests/policies/lingbot_va/test_modules.py | 128 ++ tests/policies/lingbot_va/test_processor.py | 88 ++ uv.lock | 11 +- 21 files changed, 3110 insertions(+), 19 deletions(-) create mode 100644 docs/source/lingbot_va.mdx create mode 120000 src/lerobot/policies/lingbot_va/README.md create mode 100644 src/lerobot/policies/lingbot_va/__init__.py create mode 100644 src/lerobot/policies/lingbot_va/configuration_lingbot_va.py create mode 100644 src/lerobot/policies/lingbot_va/modeling_lingbot_va.py create mode 100644 src/lerobot/policies/lingbot_va/processor_lingbot_va.py create mode 100644 src/lerobot/policies/lingbot_va/utils.py create mode 100644 tests/policies/lingbot_va/test_configuration_lingbot_va.py create mode 100644 tests/policies/lingbot_va/test_factory.py create mode 100644 tests/policies/lingbot_va/test_modules.py create mode 100644 tests/policies/lingbot_va/test_processor.py diff --git a/.dockerignore b/.dockerignore index c0d8a84b5..3295cc1b4 100644 --- a/.dockerignore +++ b/.dockerignore @@ -22,6 +22,10 @@ outputs rl media +# Local virtualenvs (the image provides its own) +.venv +venv + # Logging logs diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml index dcd14e131..79f4bc124 100644 --- a/docs/source/_toctree.yml +++ b/docs/source/_toctree.yml @@ -69,6 +69,8 @@ title: VLA-JEPA - local: eo1 title: EO-1 + - local: lingbot_va + title: LingBot-VA - local: fastwam title: FastWAM - local: groot diff --git a/docs/source/lingbot_va.mdx b/docs/source/lingbot_va.mdx new file mode 100644 index 000000000..d33e90340 --- /dev/null +++ b/docs/source/lingbot_va.mdx @@ -0,0 +1,187 @@ +# LingBot-VA + +LingBot-VA is an **autoregressive video-action world-model policy** built on the **Wan2.2** +video-diffusion stack. It interleaves, in one autoregressive sequence, the prediction of +future **video latents** and **robot actions** ("VA" = Video-Action). The LeRobot +integration wires LingBot-VA into the standard training, evaluation and processor +interfaces. + +## Model Overview + +LingBot-VA is a **dual-stream "mixture-of-transformers"**: a video/latent stream +(`patch_embedding_mlp → blocks → proj_out`) and an action stream +(`action_embedder → blocks → action_proj_out`) share the same 30 transformer blocks and +text conditioning. + +| Component | Class | Role | +| ------------------------ | ----------------------- | ----------------------------------------------------------- | +| DiT backbone (trainable) | `WanTransformer3DModel` | ~5B-param dual-stream transformer. | +| VAE (frozen) | `AutoencoderKLWan` | Wan2.2 VAE, `z_dim=48`. Lazy-pulled from the source repo. | +| Text encoder (frozen) | `UMT5EncoderModel` | UMT5-XXL, `d_model=4096`. Lazy-pulled from the source repo. | + +At inference the policy runs an autoregressive loop per chunk: it denoises the video-latent +stream (CFG, ~20 steps) and the action stream (~50 steps) with two independent +flow-matching schedulers, maintaining a KV cache across chunks. Real observed keyframes are +fed back into the KV cache as the chunk is executed (closed-loop world modeling). + +### What the LeRobot Integration Covers + +- Standard `policy.type=lingbot_va` configuration through LeRobot. +- Ready-to-use LeRobot-format checkpoints on the Hub (converted from the released upstream ones). +- Autoregressive dual-stream inference behind the standard `select_action` interface + (single-environment eval, `--eval.batch_size=1`). +- Opt-in saving of the policy's **predicted (imagined) videos** during eval / training. +- Evaluation with `lerobot-eval` on LIBERO and RoboTwin. +- Training / fine-tuning via the dual-stream flow-matching loss (`policy.forward`), see below. + +## Installation + +1. Install LeRobot by following the [Installation Guide](./installation). +2. Install the LingBot-VA extra: + +```bash +pip install -e ".[lingbot_va]" +``` + +## Checkpoints + +The released upstream checkpoints have been converted to LeRobot format and pushed to the Hub: + +| Variant | LeRobot checkpoint | +| ---------------------- | -------------------------------- | +| LIBERO-Long post-train | `lerobot/lingbot_va_libero_long` | +| RoboTwin post-train | `lerobot/lingbot_va_robotwin` | +| Pretrained base | `lerobot/lingbot_va_base` | + +Only the trainable ~5B transformer is stored in the LeRobot +`model.safetensors`. The frozen VAE + UMT5 + tokenizer (~20 GB) are pulled from +`config.wan_pretrained_path` at load time (defaults to the source `robbyant/*` repo). The +UMT5-XXL text encoder runs on CPU by default (`config.text_encoder_device`) so the 5B +transformer + VAE fit on a single 24–32 GB GPU. + +## Evaluation (LIBERO) + +```bash +lerobot-eval \ + --policy.path=lerobot/lingbot_va_libero_long \ + --policy.device=cuda \ + --env.type=libero --env.task=libero_10 \ + --env.observation_height=128 --env.observation_width=128 \ + --eval.n_episodes=50 --eval.batch_size=1 \ + --output_dir=outputs/eval/lingbot_va_libero +``` + +LingBot-VA's streaming inference (KV cache + observed-keyframe feedback) is implemented for +single-environment eval; use `--eval.batch_size=1`. + +## Evaluation (RoboTwin) + +RoboTwin 2.0 needs the SAPIEN + CuRobo simulator stack. You can use the benchmark Docker image +(`docker/Dockerfile.benchmark.robotwin`, which also needs `warp-lang==1.3.1` and CuRobo built +with the GPU's compute capability in `TORCH_CUDA_ARCH_LIST`). RoboTwin uses **end-effector-pose +control**, so run with `--env.action_mode=ee`: the policy predicts per-arm `xyz+quaternion+gripper` +deltas (`robotwin_tshape` latent layout) that are composed onto the episode's initial eef pose and +executed via CuRobo IK. + +```bash +lerobot-eval \ + --policy.path=lerobot/lingbot_va_robotwin \ + --policy.device=cuda \ + --env.type=robotwin --env.task=beat_block_hammer --env.action_mode=ee \ + --eval.n_episodes=10 --eval.batch_size=1 \ + --output_dir=outputs/eval/lingbot_va_robotwin +``` + +### Saving predicted (imagined) videos + +Set `--policy.save_predicted_video=true` to additionally VAE-decode the predicted video +latents and write `pred_episode_*.mp4` next to the env-rendered `eval_episode_*.mp4` videos. +The same flag works for the periodic eval during `lerobot-train`. + +## Training / fine-tuning + +`LingBotVAPolicy.forward(batch)` implements the dual-stream **flow-matching** loss +(`latent_loss + action_loss`, timestep-weighted, action-masked) from the paper: it VAE-encodes +the camera clips into video latents, UMT5-encodes the task, noises both streams, runs the +transformer's block-causal training pass and returns `(loss, metrics)`. Optimizer preset is AdamW +with a linear-warmup-then-constant schedule (matching upstream). + +Requirements: + +- The block-causal masks use PyTorch **flex-attention**, so build the policy with + `--policy.attn_mode=flex` for training (the default `torch` SDPA is inference-only). +- The full 5B DiT does not fit a single 24–32 GB GPU under AdamW; fine-tune with **LoRA** + (`--policy.use_peft=true`) and/or optimizer offload. `get_optim_params` returns only the + trainable (e.g. adapter) parameters; the VAE + UMT5 text encoder stay frozen. + +```bash +lerobot-train \ + --policy.path=lerobot/lingbot_va_libero_long --policy.attn_mode=flex \ + --policy.use_peft=true \ + --dataset.repo_id= \ + --batch_size=1 --steps=... --output_dir=outputs/train/lingbot_va +``` + +The dataset must provide camera clips (a temporal window per camera, VAE-encoded to +`frame_chunk_size` latent frames) and `frame_chunk_size * action_per_frame` action steps per item. + +## Data format (action channels & camera order) + +LingBot-VA is an **end-effector (Cartesian) pose** policy, it predicts EEF poses + gripper, not +joint positions. Actions live in a fixed multi-embodiment **30-dim** layout; map your robot's +action dimensions into these channels and pad the rest with `0` (`used_action_channel_ids` selects +the channels a given checkpoint actually uses): + +| channels | meaning | +| -------- | ----------------------------------------------------- | +| 0–6 | Left-arm end-effector pose | +| 7–13 | Right-arm end-effector pose | +| 14–20 | Left-arm joints (unused by the released checkpoints) | +| 21–27 | Right-arm joints (unused by the released checkpoints) | +| 28 | Left gripper | +| 29 | Right gripper | + +- **LIBERO** uses channels `0–6`: a 6-DoF EEF delta (xyz + rotation) + gripper (single arm). +- **RoboTwin** uses channels `[0–6, 28, 7–13, 29]`: left EEF (xyz + quaternion) + left gripper + + right EEF + right gripper (16 dims). The env converts these poses to joint trajectories via + CuRobo IK — joints are never predicted. + +Joint-space datasets (or a different EEF convention) must be remapped into this schema before +fine-tuning these checkpoints. + +**Camera order is fixed and order-sensitive**, per-camera latents are concatenated spatially in +`obs_cam_keys` order, so the physical camera→slot mapping must match training: + +| benchmark | `obs_cam_keys` (in order) | `camera_layout` | +| --------- | ----------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------- | +| LIBERO | `observation.images.image` (agentview / 3rd-person), `observation.images.image2` (eye-in-hand wrist) | `width_concat` (latents concatenated on width) | +| RoboTwin | `observation.images.head_camera`, `observation.images.left_camera`, `observation.images.right_camera` | `robotwin_tshape` (full-res head below, two half-res wrists on top) | + +The first camera is the exterior/head view and the rest are wrist views. + +## Inference Hyperparameters (LIBERO) + +| Key | Value | +| -------------------------------------- | --------------------------------------------------------------------------------- | +| height × width | 128 × 128 | +| cameras | `observation.images.image` (agentview), `observation.images.image2` (eye-in-hand) | +| action channels used | 0–6 (7-DoF arm + gripper) | +| action_per_frame / frame_chunk_size | 4 / 4 | +| attn_window | 30 | +| video / action denoising steps | 20 / 50 | +| guidance_scale / action_guidance_scale | 5 / 1 | +| snr_shift / action_snr_shift | 5.0 / 0.05 | + +These are the defaults of `LingBotVAConfig`; override any of them via `--policy.=...`. + +## Notes + +- **Attention backend:** inference uses the `torch` SDPA backend (always available). The + `flashattn` and `flex` backends are optional; `flex` is only needed for training. +- **Model size:** the DiT is ~5B params and the frozen VAE+UMT5 add ~20 GB; inference needs + roughly 18–24 GB of VRAM. + +## License + +LingBot-VA is released under Apache-2.0. See the +[upstream repository](https://github.com/Robbyant/lingbot-va). diff --git a/pyproject.toml b/pyproject.toml index 03487e8c4..84fe07328 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -236,6 +236,7 @@ fastwam = [ ] hilserl = ["lerobot[transformers-dep]", "lerobot[dataset]", "gym-hil>=0.1.14,<0.2.0", "lerobot[grpcio-dep]", "lerobot[placo-dep]"] vla_jepa = ["lerobot[transformers-dep]", "lerobot[diffusers-dep]", "lerobot[qwen-vl-utils-dep]"] +lingbot_va = ["lerobot[transformers-dep]", "lerobot[diffusers-dep]", "lerobot[accelerate-dep]"] # Features async = ["lerobot[grpcio-dep]", "lerobot[matplotlib-dep]"] @@ -318,6 +319,7 @@ all = [ "lerobot[xvla]", "lerobot[hilserl]", "lerobot[vla_jepa]", + "lerobot[lingbot_va]", "lerobot[async]", "lerobot[dev]", "lerobot[test]", diff --git a/src/lerobot/envs/configs.py b/src/lerobot/envs/configs.py index 84c40472f..3624357e2 100644 --- a/src/lerobot/envs/configs.py +++ b/src/lerobot/envs/configs.py @@ -757,7 +757,7 @@ class RoboTwinEnvConfig(EnvConfig): task: str = "beat_block_hammer" # single task or comma-separated list fps: int = 25 - episode_length: int = 300 + episode_length: int = 1200 obs_type: str = "pixels_agent_pos" render_mode: str = "rgb_array" # Available cameras from RoboTwin's aloha-agilex embodiment: head_camera @@ -768,6 +768,9 @@ class RoboTwinEnvConfig(EnvConfig): # must equal what SAPIEN actually renders. observation_height: int = 240 observation_width: int = 320 + # "joint": 14-d joint-space control. "ee": 16-d end-effector-pose deltas executed via CuRobo IK + # (for world-model policies like LingBot-VA that predict per-arm xyz+quaternion+gripper poses). + action_mode: str = "joint" features: dict[str, PolicyFeature] = field( default_factory=lambda: { ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(14,)), @@ -784,6 +787,8 @@ class RoboTwinEnvConfig(EnvConfig): ) def __post_init__(self): + if self.action_mode == "ee": + self.features[ACTION] = PolicyFeature(type=FeatureType.ACTION, shape=(16,)) cam_list = [c.strip() for c in self.camera_names.split(",") if c.strip()] for cam in cam_list: self.features[f"pixels/{cam}"] = PolicyFeature( @@ -826,6 +831,7 @@ class RoboTwinEnvConfig(EnvConfig): observation_height=self.observation_height, observation_width=self.observation_width, episode_length=self.episode_length, + action_mode=self.action_mode, ) diff --git a/src/lerobot/envs/robotwin.py b/src/lerobot/envs/robotwin.py index 823f14fa0..5b03f337b 100644 --- a/src/lerobot/envs/robotwin.py +++ b/src/lerobot/envs/robotwin.py @@ -17,6 +17,7 @@ from __future__ import annotations import importlib import logging +import os from collections import defaultdict from collections.abc import Callable, Sequence from functools import partial @@ -28,9 +29,17 @@ import torch from gymnasium import spaces from lerobot.types import RobotObservation +from lerobot.utils.import_utils import _scipy_available from .utils import _LazyAsyncVectorEnv +# scipy is only used for end-effector-pose composition (``--env.action_mode=ee``); guard it so this +# module (and its base-env unit tests, which mock the RoboTwin runtime) imports without scipy installed. +if _scipy_available: + from scipy.spatial.transform import Rotation +else: + Rotation = None + logger = logging.getLogger(__name__) # Camera names as used by RoboTwin 2.0. The wrapper appends "_rgb" when looking @@ -41,10 +50,124 @@ ROBOTWIN_CAMERA_NAMES: tuple[str, ...] = ( "right_camera", ) -ACTION_DIM = 14 # 7 DOF × 2 arms +ACTION_DIM = 14 # 7 DOF × 2 arms (joint-space control mode) +# End-effector-pose control mode: per arm [x, y, z, qx, qy, qz, qw, gripper] = 8, dual-arm = 16. +# Used by world-model policies (e.g. LingBot-VA) that predict eef-pose deltas executed via CuRobo IK. +EEF_ACTION_DIM = 16 ACTION_LOW = -1.0 ACTION_HIGH = 1.0 -DEFAULT_EPISODE_LENGTH = 300 +DEFAULT_EPISODE_LENGTH = 1200 +OFFICIAL_INSTRUCTION_ENV = "LEROBOT_ROBOTWIN_OFFICIAL_INSTRUCTION" +OFFICIAL_INSTRUCTION_TYPE_ENV = "LEROBOT_ROBOTWIN_INSTRUCTION_TYPE" +OFFICIAL_INSTRUCTION_MAX_ENV = "LEROBOT_ROBOTWIN_INSTRUCTION_MAX" + + +def _compose_eef_pose(new_pose: np.ndarray, init_pose: np.ndarray) -> np.ndarray: + """Compose a single-arm predicted delta pose onto the initial pose. + + ``new_pose`` / ``init_pose`` are 8-vectors ``[x, y, z, qx, qy, qz, qw, gripper]``. Translation + is added, rotation is composed (``init_R * new_R``), and the gripper is taken from the + prediction. Mirrors ``add_eef_pose`` in the upstream LingBot-VA RoboTwin client. + """ + new_r = Rotation.from_quat(new_pose[3:7]) + init_r = Rotation.from_quat(init_pose[3:7]) + out_rot = (init_r * new_r).as_quat() + out_trans = new_pose[:3] + init_pose[:3] + return np.concatenate([out_trans, out_rot, new_pose[7:8]]) + + +def _add_init_eef_pose(delta_pose: np.ndarray, init_pose: np.ndarray) -> np.ndarray: + """Compose a dual-arm (16-d) predicted delta pose onto the initial eef pose, normalizing quats.""" + left = _compose_eef_pose(delta_pose[:8], init_pose[:8]) + right = _compose_eef_pose(delta_pose[8:], init_pose[8:]) + out = np.concatenate([left, right]) + # Normalize the two quaternions (indices 3:7 and 11:15) as the upstream client does. + out[3:7] = out[3:7] / (np.linalg.norm(out[3:7]) + 1e-8) + out[11:15] = out[11:15] / (np.linalg.norm(out[11:15]) + 1e-8) + return out + + +def _env_flag(name: str, default: bool = False) -> bool: + raw = os.environ.get(name) + if raw is None: + return default + return raw.strip().lower() in {"1", "true", "yes", "on"} + + +def _arm_for_block(block: Any) -> str: + return "left" if float(block.get_pose().p[0]) < 0 else "right" + + +def _robotwin_blocks_episode_info(task_name: str, env: Any) -> dict[str, str] | None: + """Infer the episode-info dict used by RoboTwin's official instruction generator for block ranking.""" + if task_name == "blocks_ranking_rgb": + return { + "{A}": "red block", + "{B}": "green block", + "{C}": "blue block", + "{a}": _arm_for_block(env.block1), + "{b}": _arm_for_block(env.block2), + "{c}": _arm_for_block(env.block3), + } + if task_name == "blocks_ranking_size": + return { + "{A}": "large block", + "{B}": "medium block", + "{C}": "small block", + "{a}": _arm_for_block(env.block1), + "{b}": _arm_for_block(env.block2), + "{c}": _arm_for_block(env.block3), + } + return None + + +def _generate_robotwin_official_instruction(task_name: str, env: Any) -> str: + """Generate language with RoboTwin's official task templates, matching its eval client.""" + fallback = task_name.replace("_", " ") + episode_info = _robotwin_blocks_episode_info(task_name, env) + if episode_info is None: + logger.warning( + "Official RoboTwin instruction is not implemented for task=%s; using %r.", task_name, fallback + ) + return fallback + + try: + # Part of the robotwin simulator repo, this is being pulled by the docker image running robotwin + # see https://github.com/RoboTwin-Platform/RoboTwin/tree/main/description + # Used to generate the official instructions + from description.utils.generate_episode_instructions import generate_episode_descriptions + except Exception: + logger.warning( + "Failed to import RoboTwin official instruction generator; using %r.", fallback, exc_info=True + ) + return fallback + + instruction_type = os.environ.get(OFFICIAL_INSTRUCTION_TYPE_ENV, "seen") + try: + max_descriptions = int(os.environ.get(OFFICIAL_INSTRUCTION_MAX_ENV, "1000000")) + except ValueError: + max_descriptions = 1000000 + + results = generate_episode_descriptions(task_name, [episode_info], max_descriptions=max_descriptions) + if not results: + logger.warning( + "RoboTwin generated no official instructions for task=%s; using %r.", task_name, fallback + ) + return fallback + + options = results[0].get(instruction_type) or results[0].get("seen") or results[0].get("unseen") + if not options: + logger.warning( + "RoboTwin generated no %s official instructions for task=%s; using %r.", + instruction_type, + task_name, + fallback, + ) + return fallback + + return str(np.random.choice(options)) + + # D435 dims from task_config/_camera_config.yml (what demo_clean.yml selects). DEFAULT_CAMERA_H = 240 DEFAULT_CAMERA_W = 320 @@ -234,6 +357,7 @@ class RoboTwinEnv(gym.Env): observation_width: int | None = None, episode_length: int = DEFAULT_EPISODE_LENGTH, render_mode: str = "rgb_array", + action_mode: str = "joint", ): super().__init__() self.task_name = task_name @@ -241,6 +365,13 @@ class RoboTwinEnv(gym.Env): self.task_description = task_name.replace("_", " ") self.episode_index = episode_index self._reset_stride = n_envs + # "joint": 14-d joint-space actions via take_action(action). "ee": 16-d end-effector-pose + # deltas (added onto the episode's initial eef pose) executed via take_action(.., "ee") + IK. + if action_mode not in ("joint", "ee"): + raise ValueError(f"action_mode must be 'joint' or 'ee'; got {action_mode!r}") + self.action_mode = action_mode + self._action_dim = EEF_ACTION_DIM if action_mode == "ee" else ACTION_DIM + self._init_eef_pose: np.ndarray | None = None self.camera_names = list(camera_names) # Default to D435 dims (the camera type baked into task_config/demo_clean.yml). # The YAML-driven lookup is deferred to reset() so construction doesn't @@ -271,7 +402,7 @@ class RoboTwinEnv(gym.Env): } ) self.action_space = spaces.Box( - low=ACTION_LOW, high=ACTION_HIGH, shape=(ACTION_DIM,), dtype=np.float32 + low=ACTION_LOW, high=ACTION_HIGH, shape=(self._action_dim,), dtype=np.float32 ) def _ensure_env(self) -> None: @@ -317,6 +448,18 @@ class RoboTwinEnv(gym.Env): return {"pixels": images, "agent_pos": joint_state} + def _read_eef_pose(self) -> np.ndarray: + """Read the current 16-d dual-arm eef pose [left(xyz+quat)+grip, right(xyz+quat)+grip].""" + assert self._env is not None, "_read_eef_pose called before _ensure_env()" + ep = self._env.get_obs()["endpose"] + pose = ( + list(ep["left_endpose"]) + + [ep["left_gripper"]] + + list(ep["right_endpose"]) + + [ep["right_gripper"]] + ) + return np.asarray(pose, dtype=np.float64) + def reset(self, seed: int | None = None, **kwargs) -> tuple[RobotObservation, dict]: self._ensure_env() super().reset(seed=seed) @@ -330,16 +473,32 @@ class RoboTwinEnv(gym.Env): self.episode_index += self._reset_stride self._step_count = 0 + use_official_instruction = self.task_name in {"blocks_ranking_rgb", "blocks_ranking_size"} + if _env_flag(OFFICIAL_INSTRUCTION_ENV, default=use_official_instruction): + self.task_description = _generate_robotwin_official_instruction(self.task_name, self._env) + if hasattr(self._env, "set_instruction"): + self._env.set_instruction(instruction=self.task_description) + logger.info("RoboTwin official instruction | task=%s | %s", self.task_name, self.task_description) + else: + self.task_description = self.task_name.replace("_", " ") + + # In eef mode the policy predicts pose deltas relative to the initial eef pose. + if self.action_mode == "ee": + self._init_eef_pose = self._read_eef_pose() + obs = self._get_obs() return obs, {"is_success": False, "task": self.task_name} def step(self, action: np.ndarray) -> tuple[RobotObservation, float, bool, bool, dict[str, Any]]: assert self._env is not None, "step() called before reset()" - if action.ndim != 1 or action.shape[0] != ACTION_DIM: - raise ValueError(f"Expected 1-D action of shape ({ACTION_DIM},), got {action.shape}") + if action.ndim != 1 or action.shape[0] != self._action_dim: + raise ValueError(f"Expected 1-D action of shape ({self._action_dim},), got {action.shape}") with torch.enable_grad(): - if hasattr(self._env, "take_action"): + if self.action_mode == "ee": + ee_action = _add_init_eef_pose(np.asarray(action, dtype=np.float64), self._init_eef_pose) + self._env.take_action(ee_action, action_type="ee") + elif hasattr(self._env, "take_action"): self._env.take_action(action) else: self._env.step(action) @@ -398,6 +557,7 @@ def _make_env_fns( observation_height: int, observation_width: int, episode_length: int, + action_mode: str = "joint", ) -> list[Callable[[], RoboTwinEnv]]: """Return n_envs factory callables for a single task.""" @@ -410,6 +570,7 @@ def _make_env_fns( observation_height=observation_height, observation_width=observation_width, episode_length=episode_length, + action_mode=action_mode, ) return [partial(_make_one, i) for i in range(n_envs)] @@ -423,6 +584,7 @@ def create_robotwin_envs( observation_height: int = DEFAULT_CAMERA_H, observation_width: int = DEFAULT_CAMERA_W, episode_length: int = DEFAULT_EPISODE_LENGTH, + action_mode: str = "joint", ) -> dict[str, dict[int, Any]]: """Create vectorized RoboTwin 2.0 environments. @@ -473,6 +635,7 @@ def create_robotwin_envs( observation_height=observation_height, observation_width=observation_width, episode_length=episode_length, + action_mode=action_mode, ) if is_async: lazy = _LazyAsyncVectorEnv(fns, cached_obs_space, cached_act_space, cached_metadata) diff --git a/src/lerobot/optim/schedulers.py b/src/lerobot/optim/schedulers.py index 250650089..74111e7ef 100644 --- a/src/lerobot/optim/schedulers.py +++ b/src/lerobot/optim/schedulers.py @@ -83,6 +83,28 @@ class VQBeTSchedulerConfig(LRSchedulerConfig): return LambdaLR(optimizer, lr_lambda, -1) +@LRSchedulerConfig.register_subclass("constant_with_warmup") +@dataclass +class ConstantWithWarmupSchedulerConfig(LRSchedulerConfig): + """Linear warmup followed by a constant learning rate. + + Mirrors the ``warmup_constant_lambda`` used by LingBot-VA (upstream ``wan_va/train.py``): + the LR ramps linearly from 0 to the peak over ``num_warmup_steps`` steps, then stays flat. + """ + + num_warmup_steps: int = 1000 + + def build(self, optimizer: Optimizer, num_training_steps: int) -> LambdaLR: + warmup_steps = self.num_warmup_steps or 0 + + def lr_lambda(current_step): + if current_step < warmup_steps: + return float(current_step) / float(max(1, warmup_steps)) + return 1.0 + + return LambdaLR(optimizer, lr_lambda, -1) + + @LRSchedulerConfig.register_subclass("cosine_decay_with_warmup") @dataclass class CosineDecayWithWarmupSchedulerConfig(LRSchedulerConfig): diff --git a/src/lerobot/policies/__init__.py b/src/lerobot/policies/__init__.py index 4daa6abc5..494427692 100644 --- a/src/lerobot/policies/__init__.py +++ b/src/lerobot/policies/__init__.py @@ -21,6 +21,7 @@ from .factory import get_policy_class, make_policy, make_policy_config, make_pre from .fastwam.configuration_fastwam import FastWAMConfig as FastWAMConfig from .gaussian_actor.configuration_gaussian_actor import GaussianActorConfig as GaussianActorConfig from .groot.configuration_groot import GrootConfig as GrootConfig +from .lingbot_va.configuration_lingbot_va import LingBotVAConfig as LingBotVAConfig from .molmoact2.configuration_molmoact2 import MolmoAct2Config as MolmoAct2Config from .multi_task_dit.configuration_multi_task_dit import MultiTaskDiTConfig as MultiTaskDiTConfig from .pi0.configuration_pi0 import PI0Config as PI0Config @@ -46,6 +47,7 @@ __all__ = [ "FastWAMConfig", "GaussianActorConfig", "GrootConfig", + "LingBotVAConfig", "MolmoAct2Config", "MultiTaskDiTConfig", "PI0Config", diff --git a/src/lerobot/policies/factory.py b/src/lerobot/policies/factory.py index f5acb0170..12871218e 100644 --- a/src/lerobot/policies/factory.py +++ b/src/lerobot/policies/factory.py @@ -50,6 +50,7 @@ from .eo1.configuration_eo1 import EO1Config from .fastwam.configuration_fastwam import FastWAMConfig from .gaussian_actor.configuration_gaussian_actor import GaussianActorConfig from .groot.configuration_groot import GrootConfig +from .lingbot_va.configuration_lingbot_va import LingBotVAConfig from .molmoact2.configuration_molmoact2 import MolmoAct2Config from .multi_task_dit.configuration_multi_task_dit import MultiTaskDiTConfig from .pi0.configuration_pi0 import PI0Config @@ -163,6 +164,10 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]: from .vla_jepa.modeling_vla_jepa import VLAJEPAPolicy return VLAJEPAPolicy + elif name == "lingbot_va": + from .lingbot_va.modeling_lingbot_va import LingBotVAPolicy + + return LingBotVAPolicy elif name == "fastwam": from .fastwam.modeling_fastwam import FastWAMPolicy @@ -223,6 +228,8 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig: return MolmoAct2Config(**kwargs) elif policy_type == "vla_jepa": return VLAJEPAConfig(**kwargs) + elif policy_type == "lingbot_va": + return LingBotVAConfig(**kwargs) elif policy_type == "fastwam": return FastWAMConfig(**kwargs) else: @@ -458,6 +465,14 @@ def make_pre_post_processors( dataset_stats=kwargs.get("dataset_stats"), ) + elif isinstance(policy_cfg, LingBotVAConfig): + from .lingbot_va.processor_lingbot_va import make_lingbot_va_pre_post_processors + + processors = make_lingbot_va_pre_post_processors( + config=policy_cfg, + dataset_stats=kwargs.get("dataset_stats"), + ) + elif isinstance(policy_cfg, FastWAMConfig): from .fastwam.processor_fastwam import make_fastwam_pre_post_processors diff --git a/src/lerobot/policies/lingbot_va/README.md b/src/lerobot/policies/lingbot_va/README.md new file mode 120000 index 000000000..2ec3c82af --- /dev/null +++ b/src/lerobot/policies/lingbot_va/README.md @@ -0,0 +1 @@ +../../../../docs/source/lingbot_va.mdx \ No newline at end of file diff --git a/src/lerobot/policies/lingbot_va/__init__.py b/src/lerobot/policies/lingbot_va/__init__.py new file mode 100644 index 000000000..db2ab17d1 --- /dev/null +++ b/src/lerobot/policies/lingbot_va/__init__.py @@ -0,0 +1,21 @@ +#!/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 .configuration_lingbot_va import LingBotVAConfig +from .modeling_lingbot_va import LingBotVAPolicy +from .processor_lingbot_va import make_lingbot_va_pre_post_processors + +__all__ = ["LingBotVAConfig", "LingBotVAPolicy", "make_lingbot_va_pre_post_processors"] diff --git a/src/lerobot/policies/lingbot_va/configuration_lingbot_va.py b/src/lerobot/policies/lingbot_va/configuration_lingbot_va.py new file mode 100644 index 000000000..424ea7c63 --- /dev/null +++ b/src/lerobot/policies/lingbot_va/configuration_lingbot_va.py @@ -0,0 +1,168 @@ +# 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. + +"""Configuration for the LingBot-VA policy. + +LingBot-VA is an autoregressive video-action world-model policy built on the Wan2.2 +video-diffusion stack. It interleaves prediction of future video latents and robot +actions in a single dual-stream transformer. See ``docs/source/lingbot_va.mdx`` and the +upstream repository (https://github.com/Robbyant/lingbot-va). + +Defaults below match the upstream LIBERO configuration (``wan_va/configs/va_libero_cfg.py``) +and the ``transformer/config.json`` of the released checkpoints. +""" + +from dataclasses import dataclass, field + +from lerobot.configs.policies import PreTrainedConfig +from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature +from lerobot.optim.optimizers import AdamWConfig +from lerobot.optim.schedulers import ConstantWithWarmupSchedulerConfig, LRSchedulerConfig +from lerobot.utils.constants import ACTION + + +@PreTrainedConfig.register_subclass("lingbot_va") +@dataclass +class LingBotVAConfig(PreTrainedConfig): + """Configuration for the native LingBot-VA policy integration in LeRobot.""" + + # Wan transformer architecture + patch_size: tuple[int, int, int] = (1, 2, 2) + num_attention_heads: int = 24 + attention_head_dim: int = 128 + in_channels: int = 48 + out_channels: int = 48 + action_dim: int = 30 + text_dim: int = 4096 + freq_dim: int = 256 + ffn_dim: int = 14336 + num_layers: int = 30 + cross_attn_norm: bool = True + eps: float = 1e-6 + rope_max_seq_len: int = 1024 + # "flex" = training only (needs recent torch); inference uses "torch" SDPA or "flashattn". + attn_mode: str = "torch" + + # Frozen sub-models (VAE + UMT5 text encoder + tokenizer) + # ~20 GB of frozen weights, NOT bundled in the checkpoint; lazily pulled from this HF repo / + # local dir (must hold diffusers-style ``vae/``, ``text_encoder/``, ``tokenizer/`` sub-folders). + wan_pretrained_path: str = "robbyant/lingbot-va-base" + dtype: str = "bfloat16" # transformer / VAE / text-encoder dtype: "bfloat16", "float16", "float32" + # Frozen UMT5-XXL encoder device; "cpu" frees ~11 GB VRAM (it runs once per episode). + text_encoder_device: str = "cpu" + + # Observation cameras (order matters: latents are concatenated on width; LIBERO defaults) + obs_cam_keys: list[str] = field( + default_factory=lambda: ["observation.images.image", "observation.images.image2"] + ) + # Undo the LIBERO env processor's extra horizontal flip to match the model's training orientation. + image_hflip: bool = False + # Camera latent layout: "width_concat" (cameras concatenated on width; LIBERO) or + # "robotwin_tshape" (full-res head + half-res wrists in a "T"; RoboTwin). + camera_layout: str = "width_concat" + + # Inference hyperparameters (LIBERO defaults) + n_obs_steps: int = 1 + height: int = 128 + width: int = 128 + action_per_frame: int = 4 + frame_chunk_size: int = 4 + attn_window: int = 30 + num_inference_steps: int = 20 + video_exec_step: int = -1 + action_num_inference_steps: int = 50 + guidance_scale: float = 5.0 + action_guidance_scale: float = 1.0 + snr_shift: float = 5.0 + action_snr_shift: float = 0.05 + max_sequence_length: int = 512 # UMT5 prompt length + + # Subset of the 30-d action space used by the benchmark (LIBERO = 7-DoF). The action + # (un)normalization quantiles live in the checkpoint's ``policy_postprocessor.json``, not here. + used_action_channel_ids: list[int] = field(default_factory=lambda: list(range(7))) + + # Opt-in: VAE-decode predicted video latents to ``self.last_predicted_frames`` for saving MP4s. + save_predicted_video: bool = False + + # Normalization: IDENTITY here; images are scaled + VAE-encoded and actions are + # quantile-(un)normalized inside the policy / dedicated processor steps. + normalization_mapping: dict[str, NormalizationMode] = field( + default_factory=lambda: { + "VISUAL": NormalizationMode.IDENTITY, + "STATE": NormalizationMode.IDENTITY, + "ACTION": NormalizationMode.IDENTITY, + } + ) + + # Optimizer / scheduler (training; AdamW + warmup-constant per upstream train.py) + optimizer_lr: float = 1e-5 + optimizer_betas: tuple[float, float] = (0.9, 0.95) + optimizer_eps: float = 1e-8 + optimizer_weight_decay: float = 1e-4 + optimizer_grad_clip_norm: float = 1.0 + scheduler_warmup_steps: int = 1000 + + def __post_init__(self): + super().__post_init__() + if self.attn_mode not in ("torch", "flashattn", "flex"): + raise ValueError(f"attn_mode must be one of 'torch', 'flashattn', 'flex'; got {self.attn_mode!r}") + + @property + def chunk_size(self) -> int: + """Number of single-step actions produced per autoregressive chunk.""" + return self.frame_chunk_size * self.action_per_frame + + @property + def n_action_steps(self) -> int: + """Number of actions executed before refilling (the whole chunk).""" + return self.chunk_size + + def validate_features(self) -> None: + image_features = [key for key, feat in self.input_features.items() if feat.type == FeatureType.VISUAL] + if not image_features: + raise ValueError( + "LingBot-VA requires at least one visual input feature. " + "No features of type FeatureType.VISUAL found in input_features." + ) + if ACTION not in self.output_features: + self.output_features[ACTION] = PolicyFeature( + type=FeatureType.ACTION, shape=(len(self.used_action_channel_ids),) + ) + + def get_optimizer_preset(self) -> AdamWConfig: + return AdamWConfig( + lr=self.optimizer_lr, + betas=self.optimizer_betas, + eps=self.optimizer_eps, + weight_decay=self.optimizer_weight_decay, + grad_clip_norm=self.optimizer_grad_clip_norm, + ) + + def get_scheduler_preset(self) -> LRSchedulerConfig | None: + # Upstream uses a linear warmup followed by a constant LR (warmup_constant_lambda). + return ConstantWithWarmupSchedulerConfig(num_warmup_steps=self.scheduler_warmup_steps) + + @property + def observation_delta_indices(self) -> list[int]: + temporal_downsample = 4 + stride = max(1, self.action_per_frame // temporal_downsample) + return list(range(0, self.frame_chunk_size * temporal_downsample * stride, stride)) + + @property + def action_delta_indices(self) -> list[int]: + return list(range(self.chunk_size)) + + @property + def reward_delta_indices(self) -> None: + return None diff --git a/src/lerobot/policies/lingbot_va/modeling_lingbot_va.py b/src/lerobot/policies/lingbot_va/modeling_lingbot_va.py new file mode 100644 index 000000000..0f70ad290 --- /dev/null +++ b/src/lerobot/policies/lingbot_va/modeling_lingbot_va.py @@ -0,0 +1,853 @@ +# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved. +# 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. + +"""LingBot-VA policy: an autoregressive video-action world model on the Wan2.2 stack. + +The sampling loop is a faithful re-implementation of the upstream streaming server +(``wan_va/wan_va_server.py``) and LIBERO client (``evaluation/libero/client.py``), adapted +to LeRobot's ``select_action`` interface: + + * the trainable dual-stream transformer is owned as a sub-module and round-trips in the + single ``model.safetensors`` checkpoint; + * the frozen Wan VAE + UMT5 text encoder + tokenizer are *lazily pulled* from + ``config.wan_pretrained_path`` (not bundled), so the LeRobot checkpoint stays small; + * ``predict_action_chunk`` runs one autoregressive chunk (video stream then action + stream, each with CFG and its own flow-matching scheduler) and updates the KV cache; + * ``select_action`` drains a per-step action queue and records the real observed + keyframes that are fed back into the KV cache when the queue is refilled. + +NOTE: The streaming path is written for single-environment eval (``--eval.batch_size=1``). +""" + +from collections import deque + +import torch +import torch.nn.functional as F # noqa: N812 +from einops import rearrange +from torch import Tensor + +from lerobot.policies.pretrained import PreTrainedPolicy +from lerobot.utils.constants import ACTION +from lerobot.utils.import_utils import require_package + +from .configuration_lingbot_va import LingBotVAConfig +from .utils import ( + FlowMatchScheduler, + WanTransformer3DModel, + WanVAEStreamingWrapper, + _sample_timestep_id, + _torch_dtype, + clean_prompt, + data_seq_to_patch, + denormalize_latents, + get_mesh_id, + load_text_encoder, + load_tokenizer, + load_vae, +) + + +class LingBotVAPolicy(PreTrainedPolicy): + """LeRobot wrapper for the LingBot-VA autoregressive video-action world model.""" + + config_class = LingBotVAConfig + name = "lingbot_va" + + def __init__(self, config: LingBotVAConfig, **kwargs): + require_package("diffusers", extra="lingbot_va") + require_package("transformers", extra="lingbot_va") + super().__init__(config) + config.validate_features() + self.config = config + + self.dtype = _torch_dtype(config.dtype) + + # Trainable dual-stream transformer (the only sub-module saved in the LeRobot checkpoint). + self.transformer = WanTransformer3DModel( + patch_size=tuple(config.patch_size), + num_attention_heads=config.num_attention_heads, + attention_head_dim=config.attention_head_dim, + in_channels=config.in_channels, + out_channels=config.out_channels, + action_dim=config.action_dim, + text_dim=config.text_dim, + freq_dim=config.freq_dim, + ffn_dim=config.ffn_dim, + num_layers=config.num_layers, + cross_attn_norm=config.cross_attn_norm, + eps=config.eps, + rope_max_seq_len=config.rope_max_seq_len, + attn_mode=config.attn_mode, + ) + # Run the transformer in config.dtype (bf16); norm/modulation paths upcast to fp32 internally. + self.transformer = self.transformer.to(self.dtype) + + # Frozen modules are stored OUTSIDE the nn.Module registry (plain dict) so they are + # neither saved into model.safetensors nor moved by ``.to()``. They are lazily loaded + # from ``config.wan_pretrained_path`` the first time inference runs. + self._frozen: dict = {} + + self.last_predicted_frames: Tensor | None = None + self.last_predicted_latents: Tensor | None = None + self.reset() + + # Frozen-module lazy loading (VAE + UMT5 + tokenizer) + def _ensure_frozen_modules(self): + if self._frozen: + return + path = self.config.wan_pretrained_path + device = self.config.device + + # The frozen modules always live in ``vae/``, ``text_encoder/`` and ``tokenizer/`` + # sub-folders -- both in the released diffusers-style HF repos and in the local + # ``--bundle-frozen`` output dir. ``from_pretrained(path, subfolder=...)`` resolves + # them for either a HF repo id or a local directory. + vae = load_vae(path, torch_dtype=self.dtype, torch_device=device, subfolder="vae") + # The UMT5-XXL text encoder (~11 GB) runs once per episode; keep it on its own + # (CPU by default) device so the 5B transformer + VAE fit on a single GPU. + text_encoder = load_text_encoder( + path, + torch_dtype=self.dtype, + torch_device=self.config.text_encoder_device, + subfolder="text_encoder", + ) + tokenizer = load_tokenizer(path, subfolder="tokenizer") + self._frozen = { + "vae": vae.eval(), + "streaming_vae": WanVAEStreamingWrapper(vae), + "text_encoder": text_encoder.eval(), + "tokenizer": tokenizer, + } + # RoboTwin's T-shape layout encodes the half-resolution wrist cameras through a second + # streaming VAE (separate causal cache) alongside the full-res head camera. + if self.config.camera_layout == "robotwin_tshape": + vae_half = load_vae(path, torch_dtype=self.dtype, torch_device=device, subfolder="vae") + self._frozen["streaming_vae_half"] = WanVAEStreamingWrapper(vae_half.eval()) + + @property + def _vae(self): + return self._frozen["vae"] + + @property + def _streaming_vae(self): + return self._frozen["streaming_vae"] + + # PreTrainedPolicy API + def get_optim_params(self) -> dict: + # Only the transformer is trainable; the VAE / text encoder stay frozen (kept outside the + # nn.Module registry). With PEFT/LoRA this naturally returns just the adapter params. + return [p for p in self.transformer.parameters() if p.requires_grad] + + def reset(self): + """Reset all per-episode streaming state (KV cache, queues, frame counter).""" + cfg = self.config + self._action_queue: deque = deque(maxlen=cfg.n_action_steps) + self._obs_buffer: list = [] # raw keyframe obs (one per env substep) observed this chunk + self._executed_actions: Tensor | None = ( + None # last chunk's actions (model-normalized) for KV feedback + ) + self._started = False # first select_action call uses the obs as the conditioning frame + self._exec_step = 0 # index of the action being executed within the current chunk + self._prev_j = 0 # sub-step index (within a predicted frame) of the last executed action + # Sample one keyframe every ``action_per_frame / temporal_downsample`` executed sub-steps so + # that exactly ``frame_chunk_size * temporal_downsample`` frames are VAE-encoded per chunk + # (the Wan2.2 VAE temporal downsample is 4 -> ``frame_chunk_size`` latent frames). + self._keyframe_stride = max(1, cfg.action_per_frame // 4) + self._frame_st_id = 0 + self._first_chunk = True + self._prompt: str | None = None + self._prompt_embeds = None + self._negative_prompt_embeds = None + self.last_predicted_frames = None + self.last_predicted_latents = None + self._use_cfg = (cfg.guidance_scale > 1) or (cfg.action_guidance_scale > 1) + # Two independent flow-matching schedulers (video latent + action streams). + self._scheduler = FlowMatchScheduler(shift=cfg.snr_shift, sigma_min=0.0, extra_one_step=True) + self._action_scheduler = FlowMatchScheduler( + shift=cfg.action_snr_shift, sigma_min=0.0, extra_one_step=True + ) + self._scheduler.set_timesteps(1000, training=True) + self._action_scheduler.set_timesteps(1000, training=True) + self._cache_initialised = False + # Clear KV cache on the (already-built) transformer, if present. + if hasattr(self, "transformer"): + self.transformer.clear_cache("pos") + # Reset the causal streaming-VAE feat cache between episodes (mirrors upstream ``_reset``). + # Without this the encoder carries over the previous episode's temporal state, corrupting the + # latent frame counts on the next episode's first encode. + if self._frozen: + self._frozen["streaming_vae"].clear_cache() + if "streaming_vae_half" in self._frozen: + self._frozen["streaming_vae_half"].clear_cache() + + # Training (flow-matching dual-stream loss). Requires attn_mode="flex". + def _ensure_train_schedulers(self): + if getattr(self, "_train_sched_latent", None) is None: + cfg = self.config + self._train_sched_latent = FlowMatchScheduler( + shift=cfg.snr_shift, sigma_min=0.0, extra_one_step=True + ) + self._train_sched_latent.set_timesteps(1000, training=True) + self._train_sched_action = FlowMatchScheduler( + shift=cfg.action_snr_shift, sigma_min=0.0, extra_one_step=True + ) + self._train_sched_action.set_timesteps(1000, training=True) + + @torch.no_grad() + def _add_noise_stream(self, latent, scheduler, action_mask, action_mode, noisy_cond_prob): + """Flow-matching noising of one stream (port of upstream ``Trainer._add_noise``).""" + device = latent.device + b, _c, f, _h, _w = latent.shape + p = self.config.patch_size + patch_f, patch_h, patch_w = (1, 1, 1) if action_mode else (p[0], p[1], p[2]) + + ts_ids = _sample_timestep_id(f, num_train_timesteps=scheduler.num_train_timesteps) + noise = torch.zeros_like(latent).normal_() + timesteps = scheduler.timesteps[ts_ids].to(device) + noisy_latents = scheduler.add_noise(latent, noise, timesteps, t_dim=2) + targets = scheduler.training_target(latent, noise, timesteps) + + grid_id = ( + get_mesh_id( + latent.shape[-3] // patch_f, + latent.shape[-2] // patch_h, + latent.shape[-1] // patch_w, + t=1 if action_mode else 0, + f_w=1, + f_shift=0, + action=action_mode, + ) + .to(device)[None] + .repeat(b, 1, 1) + ) + + if torch.rand(1).item() < noisy_cond_prob: + cond_ids = _sample_timestep_id( + f, min_timestep_bd=0.5, max_timestep_bd=1.0, num_train_timesteps=scheduler.num_train_timesteps + ) + cond_noise = torch.zeros_like(latent).normal_() + cond_timesteps = scheduler.timesteps[cond_ids].to(device) + latent = scheduler.add_noise(latent, cond_noise, cond_timesteps, t_dim=2) + else: + cond_timesteps = torch.zeros_like(timesteps) + + if action_mask is not None: + noisy_latents = noisy_latents * action_mask.float() + targets = targets * action_mask.float() + latent = latent * action_mask.float() + + return { + "timesteps": timesteps[None].repeat(b, 1), + "noisy_latents": noisy_latents, + "targets": targets, + "latent": latent, + "cond_timesteps": cond_timesteps[None].repeat(b, 1), + "grid_id": grid_id, + } + + def _flow_matching_loss(self, input_dict, pred): + """Dual-stream flow-matching loss (port of upstream ``Trainer.compute_loss``).""" + latent_pred, action_pred = pred + ld, ad = input_dict["latent_dict"], input_dict["action_dict"] + action_pred = rearrange(action_pred, "b (f n) c -> b c f n 1", f=ad["targets"].shape[-3]) + latent_pred = data_seq_to_patch( + self.config.patch_size, + latent_pred, + ld["targets"].shape[-3], + ld["targets"].shape[-2], + ld["targets"].shape[-1], + batch_size=latent_pred.shape[0], + ) + bn, fn = ld["timesteps"].shape + lw = self._train_sched_latent.training_weight(ld["timesteps"].flatten()).reshape(bn, fn) + aw = self._train_sched_action.training_weight(ad["timesteps"].flatten()).reshape(bn, fn) + + latent_loss = F.mse_loss(latent_pred.float(), ld["targets"].float().detach(), reduction="none") + latent_loss = ( + (latent_loss * lw[:, None, :, None, None]).permute(0, 2, 3, 4, 1).flatten(0, 1).flatten(1) + ) + latent_loss = (latent_loss.sum(dim=1) / (torch.ones_like(latent_loss).sum(dim=1) + 1e-6)).mean() + + amask = ad["actions_mask"].float() + action_loss = F.mse_loss(action_pred.float(), ad["targets"].float().detach(), reduction="none") + action_loss = ( + (action_loss * aw[:, None, :, None, None] * amask).permute(0, 2, 3, 4, 1).flatten(0, 1).flatten(1) + ) + amask_f = amask.permute(0, 2, 3, 4, 1).flatten(0, 1).flatten(1) + action_loss = (action_loss.sum(dim=1) / (amask_f.sum(dim=1) + 1e-6)).mean() + return latent_loss, action_loss + + def training_loss_from_streams(self, latents, actions, actions_mask, text_emb): + """Core dual-stream training loss given prepared latents / actions / text embeddings. + + ``latents``: ``[B, in_channels, F, h, w]`` (normalized video latents). + ``actions`` / ``actions_mask``: ``[B, action_dim, F, action_per_frame, 1]``. + ``text_emb``: ``[B, seq_len, text_dim]``. Returns ``(loss, {latent_loss, action_loss})``. + """ + if self.config.attn_mode != "flex": + raise ValueError( + "LingBot-VA training requires attn_mode='flex' (block-causal flow-matching masks). " + "Load/convert the policy with --policy.attn_mode=flex for training/fine-tuning." + ) + self._ensure_train_schedulers() + latent_dict = self._add_noise_stream( + latents, self._train_sched_latent, action_mask=None, action_mode=False, noisy_cond_prob=0.5 + ) + action_dict = self._add_noise_stream( + actions, self._train_sched_action, action_mask=actions_mask, action_mode=True, noisy_cond_prob=0.0 + ) + latent_dict["text_emb"] = text_emb + action_dict["text_emb"] = text_emb + action_dict["actions_mask"] = actions_mask + input_dict = { + "latent_dict": latent_dict, + "action_dict": action_dict, + "chunk_size": int(torch.randint(1, 5, (1,)).item()), + "window_size": int(torch.randint(4, 65, (1,)).item()), + } + pred = self.transformer(input_dict, train_mode=True) + latent_loss, action_loss = self._flow_matching_loss(input_dict, pred) + loss = latent_loss + action_loss + return loss, {"latent_loss": latent_loss.item(), "action_loss": action_loss.item()} + + def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict | None]: + """Training forward: dual-stream flow-matching loss. + + Builds the (video-latent, action, text) training streams from a LeRobot batch + (VAE-encoding the camera frames and UMT5-encoding the task), then runs the flow-matching + dual-stream loss. Requires the policy to be built with ``attn_mode='flex'``. + """ + self._ensure_frozen_modules() + latents, actions, actions_mask, text_emb = self._build_training_streams(batch) + return self.training_loss_from_streams(latents, actions, actions_mask, text_emb) + + @torch.no_grad() + def _build_training_streams(self, batch): + """Build (latents, actions, actions_mask, text_emb) from a LeRobot training batch. + + Camera frames per ``obs_cam_keys`` are expected as a temporal clip ``[B, C, T, H, W]`` (or + ``[B, T, C, H, W]``); they are VAE-encoded into ``F = T / temporal_downsample`` latent frames. + Actions ``[B, F*action_per_frame, n_used]`` are scattered into the model's ``action_dim`` space. + """ + cfg = self.config + device = cfg.device + # text embeddings + task = batch.get("task") + if isinstance(task, str): + task = [task] + text_emb = self._get_t5_prompt_embeds(list(task), cfg.max_sequence_length) + + # video latents (VAE-encode the camera clips) + latents = self._encode_training_latents(batch) + + # actions -> [B, action_dim, F, action_per_frame, 1] + act = batch[ACTION].to(device) # [B, F*apf, n_used] + b = act.shape[0] + used = cfg.used_action_channel_ids + apf, fc = cfg.action_per_frame, cfg.frame_chunk_size + act = act[:, : fc * apf].reshape(b, fc, apf, len(used)).permute(0, 3, 1, 2) # [B, n_used, F, apf] + full = act.new_zeros(b, cfg.action_dim, fc, apf) + idx = torch.as_tensor(used, device=device) + full[:, idx] = act + actions = full.unsqueeze(-1).to(self.dtype) # [B, action_dim, F, apf, 1] + mask = torch.zeros(cfg.action_dim, device=device, dtype=self.dtype) + mask[idx] = 1.0 + actions_mask = mask.view(1, -1, 1, 1, 1).expand_as(actions) + return latents, actions, actions_mask, text_emb + + @torch.no_grad() + def _encode_training_latents(self, batch) -> Tensor: + """VAE-encode the per-camera training clips into normalized video latents [B, C, F, h, w].""" + vae_device = next(self._vae.parameters()).device + + def _clip(key): + x = batch[key].to(vae_device) + if x.dim() == 4: # [B, C, H, W] -> single frame clip + x = x.unsqueeze(2) + elif x.shape[1] not in (1, 3) and x.shape[2] in (1, 3): # [B, T, C, H, W] -> [B, C, T, H, W] + x = x.permute(0, 2, 1, 3, 4) + return x.contiguous() + + def _encode(x, size): + b, c, t = x.shape[:3] + x = F.interpolate(x.flatten(0, 1).float(), size=size, mode="bilinear", align_corners=False) + x = (x.view(b, c, t, *size) * 2.0 - 1.0).to(self.dtype) + mu = self._vae.encode(x).latent_dist.mode() # [B, z_dim, F, h, w] + mean = torch.tensor(self._vae.config.latents_mean).view(1, -1, 1, 1, 1).to(mu.device) + inv_std = (1.0 / torch.tensor(self._vae.config.latents_std)).view(1, -1, 1, 1, 1).to(mu.device) + return ((mu.float() - mean) * inv_std).to(mu) + + keys = self.config.obs_cam_keys + if self.config.camera_layout == "robotwin_tshape": + h, w = self.config.height, self.config.width + head = _encode(_clip(keys[0]), (h, w)) + left = _encode(_clip(keys[1]), (h // 2, w // 2)) + right = _encode(_clip(keys[2]), (h // 2, w // 2)) + return torch.cat([torch.cat([left, right], dim=-1), head], dim=-2).to(self.config.device) + per_cam = [_encode(_clip(k), (self.config.height, self.config.width)) for k in keys] + return torch.cat(per_cam, dim=-1).to(self.config.device) + + @torch.no_grad() + def select_action(self, batch: dict[str, Tensor], **kwargs) -> Tensor: + """Return one action, refilling the chunk (and feeding back observed keyframes) as needed. + + Mirrors the upstream LIBERO client loop (``evaluation/libero/client.py``): the first obs is + the conditioning frame; every observation produced afterwards is buffered as a keyframe and, + once the chunk's actions are exhausted, the buffered frames + executed actions are fed back + into the KV cache before the next chunk is predicted. + """ + self.eval() + self._ensure_frozen_modules() + self._maybe_init_prompt(batch) + + if not self._started: + # First call: this observation conditions the first chunk (it is *not* a keyframe). + self._started = True + actions = self.predict_action_chunk(batch) # [B, chunk_size, n_used] + self._action_queue.extend(actions.transpose(0, 1)) # [chunk_size, B, n_used] + self._obs_buffer = [] + self._exec_step = 0 + else: + # This observation is the result of the previously executed action -> a candidate + # keyframe. Buffer it on the sub-step boundary the upstream client samples on. + if (self._prev_j + 1) % self._keyframe_stride == 0: + self._obs_buffer.append(self._extract_raw_obs(batch)) + if len(self._action_queue) == 0: + # All actions for the current chunk have been executed; feed the observed + # keyframes + executed actions back and predict the next chunk. + actions = self.predict_action_chunk(None) + self._action_queue.extend(actions.transpose(0, 1)) + self._exec_step = 0 + + self._prev_j = self._exec_step % self.config.action_per_frame + self._exec_step += 1 + return self._action_queue.popleft() + + @torch.no_grad() + def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs) -> Tensor: + """Run one autoregressive chunk and return actions ``[B, chunk_size, n_used]`` (normalized).""" + self.eval() + self._ensure_frozen_modules() + self._maybe_init_prompt(batch) + + is_first = self._first_chunk + if is_first: + init_latent = self._encode_frames([self._extract_raw_obs(batch)]) + self._init_latent = init_latent + self._init_streaming_cache(init_latent) + self._obs_buffer = [] # frame 0 (the init obs) conditions the chunk; it is not fed back + actions, latents = self._infer(init_latent, frame_st_id=0) + self._first_chunk = False + else: + # Feed the real observed keyframes + the executed actions back into the KV cache. + self._compute_kv_cache(self._obs_buffer, self._executed_actions) + self._obs_buffer = [] + actions, latents = self._infer(None, frame_st_id=self._frame_st_id) + + # actions: [B, action_dim, F, action_per_frame, 1] (model-normalized). Keep for KV feedback. + self._executed_actions = actions + + if self.config.save_predicted_video: + # Match upstream LingBot-VA visualization: collect chunk latents and decode the + # concatenated latent sequence once after the rollout finishes. + self.last_predicted_frames = None + self.last_predicted_latents = latents.detach().to("cpu") + + # On the first chunk, frame 0 is the conditioning frame (already "known"): the upstream + # LIBERO client skips it (start_idx=1), so we drop the first frame's actions here. + used = self.config.used_action_channel_ids + a = actions[:, used] # [B, n_used, F, action_per_frame, 1] + if is_first: + a = a[:, :, 1:] # drop frame 0 -> (F-1) frames of actions + a = a.squeeze(-1).flatten(2) # [B, n_used, n_steps] + a = a.transpose(1, 2).contiguous() # [B, n_steps, n_used] + return a.to(torch.float32) + + # Prompt / text encoding + def _maybe_init_prompt(self, batch): + if self._prompt_embeds is not None or batch is None: + return + task = batch.get("task") + prompt = task[0] if isinstance(task, list | tuple) else task + self._prompt = prompt or "" + self._prompt_embeds, self._negative_prompt_embeds = self._encode_prompt(self._prompt) + + def _get_t5_prompt_embeds(self, prompt, max_sequence_length): + tokenizer = self._frozen["tokenizer"] + text_encoder = self._frozen["text_encoder"] + device = self.config.device + + prompt = [prompt] if isinstance(prompt, str) else prompt + prompt = [clean_prompt(u) for u in prompt] + + text_inputs = tokenizer( + prompt, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask + seq_lens = mask.gt(0).sum(dim=1).long() + + te_device = next(text_encoder.parameters()).device + prompt_embeds = text_encoder(text_input_ids.to(te_device), mask.to(te_device)).last_hidden_state + prompt_embeds = prompt_embeds.to(dtype=self.dtype, device=device) + prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens, strict=False)] + prompt_embeds = torch.stack( + [torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds], + dim=0, + ) + return prompt_embeds.to(device) + + def _encode_prompt(self, prompt): + max_len = self.config.max_sequence_length + prompt_embeds = self._get_t5_prompt_embeds(prompt, max_len) + negative_prompt_embeds = None + if self._use_cfg: + negative_prompt_embeds = self._get_t5_prompt_embeds("", max_len) + return prompt_embeds, negative_prompt_embeds + + # Observation (image) encoding -> normalized video latents + def _extract_raw_obs(self, batch) -> dict[str, Tensor]: + """Snapshot the configured camera images from a batch (kept raw for later VAE encoding).""" + return {k: batch[k].detach() for k in self.config.obs_cam_keys} + + def _camera_frame(self, raw_obs, key, size=None) -> Tensor: + """Return a single-frame camera tensor [1, C, 1, H, W] resized + scaled to [-1, 1].""" + img = raw_obs[key] + if img.dim() == 3: # [C, H, W] + img = img.unsqueeze(0) + # LeRobot images arrive as float in [0, 1], shape [B, C, H, W]. + img = img.to(self.config.device, torch.float32) + if self.config.image_hflip: + img = torch.flip(img, dims=[-1]) # undo the env processor's horizontal flip + if size is None: + size = (self.config.height, self.config.width) + img = F.interpolate(img, size=size, mode="bilinear", align_corners=False) + img = img * 2.0 - 1.0 + return img.unsqueeze(2).to(self.dtype) # [1, C, F=1, H, W] + + def _normalize_vae_latent(self, enc_out: Tensor) -> Tensor: + """Take the mean of a VAE encoder output and channel-normalize it (matches upstream).""" + mu, _logvar = torch.chunk(enc_out, 2, dim=1) + latents_mean = torch.tensor(self._vae.config.latents_mean).to(mu.device) + latents_std = torch.tensor(self._vae.config.latents_std).to(mu.device) + mean = latents_mean.view(1, -1, 1, 1, 1) + inv_std = (1.0 / latents_std).view(1, -1, 1, 1, 1) + return ((mu.float() - mean) * inv_std).to(mu) + + @torch.no_grad() + def _encode_frames(self, raw_frames: list) -> Tensor: + """VAE-encode a temporal clip of observed frames and concat the per-camera latents on width. + + ``raw_frames`` is a list of per-frame obs dicts (one per env sub-step). Each configured + camera is stacked along the temporal axis into a ``[1, C, F, H, W]`` clip and encoded in a + single streaming ``encode_chunk`` call so the VAE temporal downsample (x4) collapses the F + input frames into ``F / 4`` latent frames, with the causal ``feat_cache`` carried across + chunks (mirrors upstream ``_encode_obs``). + """ + vae_device = next(self._vae.parameters()).device + if self.config.camera_layout == "robotwin_tshape": + return self._encode_frames_tshape(raw_frames, vae_device) + per_cam_videos = [] + for k in self.config.obs_cam_keys: + frames = [self._camera_frame(fb, k) for fb in raw_frames] + per_cam_videos.append(torch.cat(frames, dim=2)) # [1, C, F, H, W] + videos = torch.cat(per_cam_videos, dim=0) # [num_cam, C, F, H, W] + enc_out = self._streaming_vae.encode_chunk(videos.to(vae_device).to(self.dtype)) + mu_norm = self._normalize_vae_latent(enc_out) + # Concatenate the per-camera latents along width. + video_latent = torch.cat(mu_norm.split(1, dim=0), dim=-1) + return video_latent.to(self.config.device) + + @torch.no_grad() + def _encode_frames_tshape(self, raw_frames: list, vae_device) -> Tensor: + """RoboTwin T-shape latent assembly: full-res head + half-res wrists (second streaming VAE). + + The two wrist latents are concatenated on width and stacked (on the height axis) on top of + the head latent, mirroring upstream ``_encode_obs`` for ``env_type='robotwin_tshape'``. + """ + cfg = self.config + h, w = cfg.height, cfg.width + head_key, left_key, right_key = cfg.obs_cam_keys[0], cfg.obs_cam_keys[1], cfg.obs_cam_keys[2] + head = torch.cat([self._camera_frame(fb, head_key, size=(h, w)) for fb in raw_frames], dim=2) + left = torch.cat( + [self._camera_frame(fb, left_key, size=(h // 2, w // 2)) for fb in raw_frames], dim=2 + ) + right = torch.cat( + [self._camera_frame(fb, right_key, size=(h // 2, w // 2)) for fb in raw_frames], dim=2 + ) + wrists = torch.cat([left, right], dim=0) # [2, C, F, H/2, W/2] + enc_high = self._streaming_vae.encode_chunk(head.to(vae_device).to(self.dtype)) + enc_lr = self._frozen["streaming_vae_half"].encode_chunk(wrists.to(vae_device).to(self.dtype)) + # wrists side-by-side on width, then stacked on top of the head latent on the height axis. + enc_out = torch.cat([torch.cat(enc_lr.split(1, dim=0), dim=-1), enc_high], dim=-2) + video_latent = self._normalize_vae_latent(enc_out) + return video_latent.to(self.config.device) + + # KV cache management + @property + def _latent_hw(self): + if self.config.camera_layout == "robotwin_tshape": + # head (full) on the bottom, two half-res wrists side-by-side on top -> 1.5x height. + return ((self.config.height // 16) * 3) // 2, self.config.width // 16 + h = self.config.height // 16 + w = (self.config.width // 16) * len(self.config.obs_cam_keys) + return h, w + + def _init_streaming_cache(self, init_latent): + cfg = self.config + latent_h, latent_w = self._latent_hw + p = cfg.patch_size + latent_token_per_chunk = (cfg.frame_chunk_size * latent_h * latent_w) // (p[0] * p[1] * p[2]) + action_token_per_chunk = cfg.frame_chunk_size * cfg.action_per_frame + self.transformer.create_empty_cache( + "pos", + cfg.attn_window, + latent_token_per_chunk, + action_token_per_chunk, + device=self.config.device, + dtype=self.dtype, + batch_size=2 if self._use_cfg else 1, + ) + self._cache_initialised = True + + def _repeat_input_for_cfg(self, input_dict): + if self._use_cfg: + input_dict["noisy_latents"] = input_dict["noisy_latents"].repeat(2, 1, 1, 1, 1) + input_dict["text_emb"] = torch.cat( + [ + self._prompt_embeds.to(self.dtype).clone(), + self._negative_prompt_embeds.to(self.dtype).clone(), + ], + dim=0, + ) + input_dict["grid_id"] = input_dict["grid_id"][None].repeat(2, 1, 1) + input_dict["timesteps"] = input_dict["timesteps"][None].repeat(2, 1) + else: + input_dict["grid_id"] = input_dict["grid_id"][None] + input_dict["timesteps"] = input_dict["timesteps"][None] + return input_dict + + def _prepare_latent_input( + self, + latent_model_input, + action_model_input, + latent_t=0, + action_t=0, + latent_cond=None, + action_cond=None, + frame_st_id=0, + ): + cfg = self.config + device = self.config.device + p = cfg.patch_size + out = {} + if latent_model_input is not None: + out["latent_res_lst"] = { + "noisy_latents": latent_model_input, + "timesteps": torch.ones([latent_model_input.shape[2]], dtype=torch.float32, device=device) + * latent_t, + "grid_id": get_mesh_id( + latent_model_input.shape[-3] // p[0], + latent_model_input.shape[-2] // p[1], + latent_model_input.shape[-1] // p[2], + 0, + 1, + frame_st_id, + ).to(device), + "text_emb": self._prompt_embeds.to(self.dtype).clone(), + } + if latent_cond is not None: + out["latent_res_lst"]["noisy_latents"][:, :, 0:1] = latent_cond[:, :, 0:1] + out["latent_res_lst"]["timesteps"][0:1] *= 0 + if action_model_input is not None: + out["action_res_lst"] = { + "noisy_latents": action_model_input, + "timesteps": torch.ones([action_model_input.shape[2]], dtype=torch.float32, device=device) + * action_t, + "grid_id": get_mesh_id( + action_model_input.shape[-3], + action_model_input.shape[-2], + action_model_input.shape[-1], + 1, + 1, + frame_st_id, + action=True, + ).to(device), + "text_emb": self._prompt_embeds.to(self.dtype).clone(), + } + if action_cond is not None: + out["action_res_lst"]["noisy_latents"][:, :, 0:1] = action_cond[:, :, 0:1] + out["action_res_lst"]["timesteps"][0:1] *= 0 + out["action_res_lst"]["noisy_latents"][:, ~self._action_mask] *= 0 + return out + + @property + def _action_mask(self): + mask = torch.zeros([self.config.action_dim], dtype=torch.bool) + mask[self.config.used_action_channel_ids] = True + return mask + + # Action conditioning (executed action history) (de)normalization + def _preprocess_action_state(self, action_norm: Tensor) -> Tensor: + """Build the action-conditioning tensor from the already-normalized executed actions. + + ``action_norm`` is the model-space action chunk ``[B, action_dim, F, action_per_frame, 1]``. + Upstream re-derives the conditioning from the raw executed action via quantile norm; here + the executed actions are already in the model-normalized space, so we pass them through. + """ + return action_norm.to(self.config.device, self.dtype) + + def _compute_kv_cache(self, obs_buffer, executed_actions): + """Feed real observed keyframes + executed actions back into the KV cache.""" + if not obs_buffer or executed_actions is None: + return + self.transformer.clear_pred_cache("pos") + # Encode the buffered keyframe clip in one streaming call (carries the causal VAE cache). + latent_model_input = self._encode_frames(obs_buffer) + # On the first feedback, prepend the init latent so the latent/action frame counts align + # (upstream prepends ``init_latent`` to the observed keyframes when frame_st_id == 0). + if self._frame_st_id == 0 and getattr(self, "_init_latent", None) is not None: + latent_model_input = torch.cat([self._init_latent, latent_model_input], dim=2) + action_model_input = self._preprocess_action_state(executed_actions) + action_model_input = action_model_input.to(latent_model_input) + input_dict = self._prepare_latent_input( + latent_model_input, action_model_input, frame_st_id=self._frame_st_id + ) + with torch.no_grad(): + self.transformer( + self._repeat_input_for_cfg(input_dict["latent_res_lst"]), + update_cache=2, + cache_name="pos", + action_mode=False, + ) + self.transformer( + self._repeat_input_for_cfg(input_dict["action_res_lst"]), + update_cache=2, + cache_name="pos", + action_mode=True, + ) + self._frame_st_id += latent_model_input.shape[2] + + # The core dual-stream denoising loop (one chunk) + @torch.no_grad() + def _infer(self, init_latent, frame_st_id=0): + cfg = self.config + device = self.config.device + latent_h, latent_w = self._latent_hw + frame_chunk_size = cfg.frame_chunk_size + + latents = torch.randn(1, 48, frame_chunk_size, latent_h, latent_w, device=device, dtype=self.dtype) + actions = torch.randn( + 1, cfg.action_dim, frame_chunk_size, cfg.action_per_frame, 1, device=device, dtype=self.dtype + ) + + self._scheduler.set_timesteps(cfg.num_inference_steps) + self._action_scheduler.set_timesteps(cfg.action_num_inference_steps) + timesteps = F.pad(self._scheduler.timesteps, (0, 1), mode="constant", value=0) + if cfg.video_exec_step != -1: + timesteps = timesteps[: cfg.video_exec_step] + action_timesteps = F.pad(self._action_scheduler.timesteps, (0, 1), mode="constant", value=0) + + # 1. Video-latent denoising loop + for i, t in enumerate(timesteps): + last_step = i == len(timesteps) - 1 + latent_cond = ( + init_latent[:, :, 0:1].to(self.dtype) + if frame_st_id == 0 and init_latent is not None + else None + ) + input_dict = self._prepare_latent_input( + latents, None, t, t, latent_cond, None, frame_st_id=frame_st_id + ) + video_noise_pred = self.transformer( + self._repeat_input_for_cfg(input_dict["latent_res_lst"]), + update_cache=1 if last_step else 0, + cache_name="pos", + action_mode=False, + ) + if not last_step or cfg.video_exec_step != -1: + video_noise_pred = data_seq_to_patch( + cfg.patch_size, + video_noise_pred, + frame_chunk_size, + latent_h, + latent_w, + batch_size=2 if self._use_cfg else 1, + ) + if cfg.guidance_scale > 1: + video_noise_pred = video_noise_pred[1:] + cfg.guidance_scale * ( + video_noise_pred[:1] - video_noise_pred[1:] + ) + else: + video_noise_pred = video_noise_pred[:1] + latents = self._scheduler.step(video_noise_pred, t, latents, return_dict=False) + if frame_st_id == 0 and latent_cond is not None: + latents[:, :, 0:1] = latent_cond + + # 2. Action denoising loop + for i, t in enumerate(action_timesteps): + last_step = i == len(action_timesteps) - 1 + action_cond = ( + torch.zeros([1, cfg.action_dim, 1, cfg.action_per_frame, 1], device=device, dtype=self.dtype) + if frame_st_id == 0 + else None + ) + input_dict = self._prepare_latent_input( + None, actions, t, t, None, action_cond, frame_st_id=frame_st_id + ) + action_noise_pred = self.transformer( + self._repeat_input_for_cfg(input_dict["action_res_lst"]), + update_cache=1 if last_step else 0, + cache_name="pos", + action_mode=True, + ) + if not last_step: + action_noise_pred = rearrange(action_noise_pred, "b (f n) c -> b c f n 1", f=frame_chunk_size) + if cfg.action_guidance_scale > 1: + action_noise_pred = action_noise_pred[1:] + cfg.action_guidance_scale * ( + action_noise_pred[:1] - action_noise_pred[1:] + ) + else: + action_noise_pred = action_noise_pred[:1] + actions = self._action_scheduler.step(action_noise_pred, t, actions, return_dict=False) + if frame_st_id == 0 and action_cond is not None: + actions[:, :, 0:1] = action_cond + + actions[:, ~self._action_mask] *= 0 + return actions, latents + + # Predicted-video decoding (opt-in) + @torch.no_grad() + def decode_predicted_latents(self, latents) -> Tensor: + """Decode a concatenated predicted-latent sequence into ``[T, H, W, 3]`` uint8 frames.""" + return self._decode_predicted_video(latents) + + @torch.no_grad() + def _decode_predicted_video(self, latents) -> Tensor: + """VAE-decode predicted latents into a uint8 frame stack ``[T, H, W, 3]`` on CPU.""" + vae = self._vae + z_dim = vae.config.z_dim + vae_device = next(vae.parameters()).device + latents = latents.to(device=vae_device, dtype=vae.dtype) + latents = denormalize_latents(latents, vae.config.latents_mean, vae.config.latents_std, z_dim) + video = vae.decode(latents, return_dict=False)[0] # [B, C, F, H, W] in [-1, 1] + video = (video.float().clamp(-1, 1) + 1.0) / 2.0 + video = (video[0].permute(1, 2, 3, 0) * 255.0).round().to(torch.uint8) # [F, H, W, C] + return video.cpu() diff --git a/src/lerobot/policies/lingbot_va/processor_lingbot_va.py b/src/lerobot/policies/lingbot_va/processor_lingbot_va.py new file mode 100644 index 000000000..119f77d5b --- /dev/null +++ b/src/lerobot/policies/lingbot_va/processor_lingbot_va.py @@ -0,0 +1,87 @@ +# 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. + +"""Pre/post-processor pipelines for the LingBot-VA policy. + +The preprocessor passes inputs through (IDENTITY) and the postprocessor maps the policy's +``[-1, 1]`` actions back to physical units with the built-in ``UnnormalizerProcessorStep`` +(QUANTILES) using per-channel q01/q99 restored from the checkpoint. +""" + +from typing import Any + +import torch + +from lerobot.configs.types import FeatureType, NormalizationMode +from lerobot.processor import ( + AddBatchDimensionProcessorStep, + DeviceProcessorStep, + NormalizerProcessorStep, + PolicyAction, + PolicyProcessorPipeline, + ProcessorStep, + RenameObservationsProcessorStep, + UnnormalizerProcessorStep, +) +from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action +from lerobot.utils.constants import ( + POLICY_POSTPROCESSOR_DEFAULT_NAME, + POLICY_PREPROCESSOR_DEFAULT_NAME, +) + +from .configuration_lingbot_va import LingBotVAConfig + + +def make_lingbot_va_pre_post_processors( + config: LingBotVAConfig, + dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None, +) -> tuple[ + PolicyProcessorPipeline[dict[str, Any], dict[str, Any]], + PolicyProcessorPipeline[PolicyAction, PolicyAction], +]: + """Build the pre/post processor pipelines for LingBot-VA.""" + + input_steps: list[ProcessorStep] = [ + RenameObservationsProcessorStep(rename_map={}), + AddBatchDimensionProcessorStep(), + NormalizerProcessorStep( + features={**config.input_features, **config.output_features}, + norm_map=config.normalization_mapping, + stats=dataset_stats, + ), + DeviceProcessorStep(device=config.device), + ] + + # Unnormalize actions from [-1, 1] to physical units (QUANTILES) using q01/q99 restored from the checkpoint. + output_steps: list[ProcessorStep] = [ + UnnormalizerProcessorStep( + features=config.output_features, + norm_map={FeatureType.ACTION: NormalizationMode.QUANTILES}, + stats=dataset_stats, + ), + DeviceProcessorStep(device="cpu"), + ] + + return ( + PolicyProcessorPipeline[dict[str, Any], dict[str, Any]]( + steps=input_steps, + name=POLICY_PREPROCESSOR_DEFAULT_NAME, + ), + PolicyProcessorPipeline[PolicyAction, PolicyAction]( + steps=output_steps, + name=POLICY_POSTPROCESSOR_DEFAULT_NAME, + to_transition=policy_action_to_transition, + to_output=transition_to_policy_action, + ), + ) diff --git a/src/lerobot/policies/lingbot_va/utils.py b/src/lerobot/policies/lingbot_va/utils.py new file mode 100644 index 000000000..86f885549 --- /dev/null +++ b/src/lerobot/policies/lingbot_va/utils.py @@ -0,0 +1,1138 @@ +# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved. +# 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. + +"""Vendored Wan2.2 model code and plumbing for the LingBot-VA policy. + +Everything the policy builds on lives here: grid/patch reshaping, attention backends, +VAE (de)normalization + frozen-component loaders, the flow-matching scheduler, and the +dual-stream Wan transformer (``WanTransformer3DModel`` and its sub-modules). Only the +LeRobot-facing ``LingBotVAPolicy`` orchestrator stays in ``modeling_lingbot_va.py``; this +module imports nothing from it (one-directional dependency). +""" + +import html +import math +import re +from copy import deepcopy +from functools import partial +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +import torch.nn.functional as F # noqa: N812 +from einops import rearrange + +from lerobot.utils.import_utils import _diffusers_available, _transformers_available + +if TYPE_CHECKING or _diffusers_available: + from diffusers import AutoencoderKLWan + from diffusers.configuration_utils import ConfigMixin, register_to_config + from diffusers.models.attention import FeedForward + from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbedding, Timesteps + from diffusers.models.modeling_utils import ModelMixin + from diffusers.models.normalization import FP32LayerNorm +else: + AutoencoderKLWan = FeedForward = PixArtAlphaTextProjection = None + TimestepEmbedding = Timesteps = FP32LayerNorm = None + + class ModelMixin: + pass + + class ConfigMixin: + pass + + def register_to_config(func): + return func + + +if TYPE_CHECKING or _transformers_available: + from transformers import T5TokenizerFast, UMT5EncoderModel +else: + T5TokenizerFast = UMT5EncoderModel = None + + +# Grid-id / patch utilities +def data_seq_to_patch(patch_size, data_seq, latent_num_frames, latent_height, latent_width, batch_size=1): + """Reshape a flattened patch sequence back into a ``(B, C, F, H, W)`` latent grid.""" + p_t, p_h, p_w = patch_size + post_patch_num_frames = latent_num_frames // p_t + post_patch_height = latent_height // p_h + post_patch_width = latent_width // p_w + + data_patch = data_seq.reshape( + batch_size, post_patch_num_frames, post_patch_height, post_patch_width, p_t, p_h, p_w, -1 + ) + data_patch = data_patch.permute(0, 7, 1, 4, 2, 5, 3, 6) + data_patch = data_patch.flatten(6, 7).flatten(4, 5).flatten(2, 3) + return data_patch + + +def get_mesh_id(f, h, w, t, f_w=1, f_shift=0, action=False): + """Build the (frame, height, width, stream) grid ids used to index the rotary embedding.""" + f_idx = torch.arange(f_shift, f + f_shift) * f_w + h_idx = torch.arange(h) + w_idx = torch.arange(w) + ff, hh, ww = torch.meshgrid(f_idx, h_idx, w_idx, indexing="ij") + if action: + ff_offset = (torch.ones([h]).cumsum(0) / (h + 1)).view(1, -1, 1) + ff = ff + ff_offset + hh = torch.ones_like(hh) * -1 + ww = torch.ones_like(ww) * -1 + + grid_id = torch.cat([ff.unsqueeze(0), hh.unsqueeze(0), ww.unsqueeze(0)], dim=0).flatten(1) + grid_id = torch.cat([grid_id, torch.full_like(grid_id[:1], t)], dim=0) + return grid_id + + +# Attention backends +def custom_sdpa(q, k, v): + """Scaled-dot-product attention operating on ``(B, S, H, D)`` tensors.""" + out = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)) + return out.transpose(1, 2) + + +def _load_flash_attn_func(): + try: + from flash_attn_interface import flash_attn_func + except ImportError: + try: + from flash_attn import flash_attn_func + except ImportError as e: + raise ImportError( + "attn_mode='flashattn' requires the `flash_attn` package, which is not installed. " + "Install it, or use attn_mode='torch' (the default)." + ) from e + return flash_attn_func + + +# Wan2.2 VAE helpers (stock diffusers ``AutoencoderKLWan``) +def _vae_patchify(x, patch_size): + if patch_size is None or patch_size == 1: + return x + batch_size, channels, frames, height, width = x.shape + x = x.view( + batch_size, channels, frames, height // patch_size, patch_size, width // patch_size, patch_size + ) + x = x.permute(0, 1, 6, 4, 2, 3, 5).contiguous() + x = x.view( + batch_size, channels * patch_size * patch_size, frames, height // patch_size, width // patch_size + ) + return x + + +def denormalize_latents(latents: torch.Tensor, latents_mean, latents_std, z_dim) -> torch.Tensor: + """Inverse of the encode-time latent normalization, for VAE-decoding predicted latents.""" + mean = torch.tensor(latents_mean).view(1, z_dim, 1, 1, 1).to(latents.device, latents.dtype) + inv_std = 1.0 / torch.tensor(latents_std).view(1, z_dim, 1, 1, 1).to(latents.device, latents.dtype) + return latents / inv_std + mean + + +def load_vae(vae_path, torch_dtype, torch_device, subfolder=None): + vae = AutoencoderKLWan.from_pretrained(vae_path, subfolder=subfolder, torch_dtype=torch_dtype) + return vae.to(torch_device) + + +def load_text_encoder(text_encoder_path, torch_dtype, torch_device, subfolder=None): + text_encoder = UMT5EncoderModel.from_pretrained( + text_encoder_path, subfolder=subfolder, torch_dtype=torch_dtype + ) + return text_encoder.to(torch_device) + + +def load_tokenizer(tokenizer_path, subfolder=None): + return T5TokenizerFast.from_pretrained(tokenizer_path, subfolder=subfolder) + + +# Misc +def clean_prompt(text: str) -> str: + """Normalize a task prompt (HTML-unescape + whitespace collapse). + + Mirrors diffusers' Wan ``prompt_clean`` minus ``ftfy.fix_text``, + which is a no-op for the ASCII task strings used here, so we avoid the extra ``ftfy`` dep. + """ + text = html.unescape(html.unescape(text)).strip() + return re.sub(r"\s+", " ", text).strip() + + +def _torch_dtype(name: str) -> torch.dtype: + return {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}[name] + + +def _sample_timestep_id( + batch_size: int = 1, + min_timestep_bd: float = 0.0, + max_timestep_bd: float = 1.0, + num_train_timesteps: int = 1000, +) -> torch.Tensor: + """Sample per-frame flow-matching timestep ids (upstream ``utils.sample_timestep_id``).""" + u = torch.rand(size=[batch_size]) * (max_timestep_bd - min_timestep_bd) + min_timestep_bd + return (u * num_train_timesteps).clamp(min=0, max=num_train_timesteps - 1).to(torch.int64) + + +# Flow-matching scheduler +# LingBot-VA uses two independent instances at inference (one for the video-latent stream, +# one for the action stream), each with its own ``shift`` and number of denoising steps. +class FlowMatchScheduler: + def __init__( + self, + num_inference_steps=100, + num_train_timesteps=1000, + shift=3.0, + sigma_max=1.0, + sigma_min=0.003 / 1.002, + inverse_timesteps=False, + extra_one_step=False, + reverse_sigmas=False, + exponential_shift=False, + exponential_shift_mu=None, + shift_terminal=None, + ): + self.num_train_timesteps = num_train_timesteps + self.shift = shift + self.sigma_max = sigma_max + self.sigma_min = sigma_min + self.inverse_timesteps = inverse_timesteps + self.extra_one_step = extra_one_step + self.reverse_sigmas = reverse_sigmas + self.exponential_shift = exponential_shift + self.exponential_shift_mu = exponential_shift_mu + self.shift_terminal = shift_terminal + self.set_timesteps(num_inference_steps) + + def set_timesteps( + self, + num_inference_steps=100, + denoising_strength=1.0, + training=False, + shift=None, + dynamic_shift_len=None, + ): + if shift is not None: + self.shift = shift + sigma_start = self.sigma_min + (self.sigma_max - self.sigma_min) * denoising_strength + if self.extra_one_step: + self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps + 1)[:-1] + else: + self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps) + if self.inverse_timesteps: + self.sigmas = torch.flip(self.sigmas, dims=[0]) + if self.exponential_shift: + mu = ( + self.calculate_shift(dynamic_shift_len) + if dynamic_shift_len is not None + else self.exponential_shift_mu + ) + self.sigmas = math.exp(mu) / (math.exp(mu) + (1 / self.sigmas - 1)) + else: + self.sigmas = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas) + if self.shift_terminal is not None: + one_minus_z = 1 - self.sigmas + scale_factor = one_minus_z[-1] / (1 - self.shift_terminal) + self.sigmas = 1 - (one_minus_z / scale_factor) + if self.reverse_sigmas: + self.sigmas = 1 - self.sigmas + self.timesteps = self.sigmas * self.num_train_timesteps + if training: + x = self.timesteps + y = torch.exp(-2 * ((x - num_inference_steps / 2) / num_inference_steps) ** 2) + y_shifted = y - y.min() + bsmntw_weighing = y_shifted * (num_inference_steps / y_shifted.sum()) + self.linear_timesteps_weights = bsmntw_weighing + self.training = True + else: + self.training = False + + def step(self, model_output, timestep, sample, to_final=False, **kwargs): + if isinstance(timestep, torch.Tensor): + timestep = timestep.cpu() + timestep_id = torch.argmin((self.timesteps - timestep).abs()) + sigma = self.sigmas[timestep_id] + if to_final or timestep_id + 1 >= len(self.timesteps): + sigma_ = 1 if (self.inverse_timesteps or self.reverse_sigmas) else 0 + else: + sigma_ = self.sigmas[timestep_id + 1] + prev_sample = sample + model_output * (sigma_ - sigma) + return prev_sample + + def add_noise(self, original_samples, noise, timestep, t_dim=2): + if isinstance(timestep, torch.Tensor): + timestep = timestep.cpu() + timestep = timestep[None] + timestep_id = torch.argmin((self.timesteps[:, None] - timestep).abs(), dim=0) + shape = [1] * noise.ndim + shape[t_dim] = timestep_id.shape[0] + sigma = self.sigmas[timestep_id].to(original_samples).view(shape) + sample = (1 - sigma) * original_samples + sigma * noise + return sample + + def training_target(self, sample, noise, timestep): + target = noise - sample + return target + + def training_weight(self, timestep): + timestep_id = torch.argmin( + (self.timesteps[:, None].to(timestep.device) - timestep[None]).abs(), dim=0 + ) + weights = self.linear_timesteps_weights.to(timestep.device)[timestep_id].to(timestep.device) + return weights + + def calculate_shift( + self, + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 8192, + base_shift: float = 0.5, + max_shift: float = 0.9, + ): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + +class FlexAttnFunc(nn.Module): + """Flex-attention backend (training only; ``attn_mode='flex'``). + + Builds the block-causal / window / noise-vs-clean masks used by the dual-stream + flow-matching training. Inference uses the ``torch`` SDPA backend. The flex-attention + APIs and their ``torch.compile`` wrappers are imported/initialised lazily so importing + this module never requires a flex-attention-capable PyTorch build. + """ + + flex_attn = None + compiled_create_block_mask = None + attention_mask = None + cross_attention_mask = None + + def __init__(self, is_cross=False) -> None: + super().__init__() + self.is_cross = is_cross + + @classmethod + def _ensure_compiled(cls): + if cls.flex_attn is None: + from torch.nn.attention.flex_attention import create_block_mask, flex_attention + + cls.flex_attn = torch.compile(flex_attention, dynamic=True) + cls.compiled_create_block_mask = torch.compile(create_block_mask) + + def forward(self, query, key, value, dtype=torch.bfloat16): + self._ensure_compiled() + q_varlen = rearrange(query[0], "s n d -> 1 n s d") + k_varlen = rearrange(key[0], "s n d -> 1 n s d") + v_varlen = rearrange(value[0], "s n d -> 1 n s d") + + half_dtypes = (torch.float16, torch.bfloat16) + if dtype not in half_dtypes: + raise ValueError(f"Flex attention requires a half-precision dtype, got {dtype}.") + + def half(x): + return x if x.dtype in half_dtypes else x.to(dtype) + + q_varlen = half(q_varlen) + k_varlen = half(k_varlen) + v_varlen = half(v_varlen) + q_varlen = q_varlen.to(v_varlen.dtype) + k_varlen = k_varlen.to(v_varlen.dtype) + + block_mask = FlexAttnFunc.cross_attention_mask if self.is_cross else FlexAttnFunc.attention_mask + + x_out = FlexAttnFunc.flex_attn( + q_varlen, + k_varlen, + v_varlen, + block_mask=block_mask, + kernel_options={ + "BLOCK_M": 64, + "BLOCK_N": 64, + "BLOCK_M1": 32, + "BLOCK_N1": 64, + "BLOCK_M2": 64, + "BLOCK_N2": 32, + }, + ) + + x_out = rearrange(x_out, "b n s d -> b s n d") + return x_out + + @staticmethod + @torch.no_grad() + def init_mask(latent_shape, action_shape, padded_length, chunk_size, window_size, patch_size, device): + FlexAttnFunc._ensure_compiled() + torch._inductor.config.realize_opcount_threshold = 100 + b, _, l_f, l_h, l_w = latent_shape + _, _, a_f, a_h, a_w = action_shape + + latent_seq_id = ( + torch.arange(b)[:, None, None, None] + .expand(-1, l_f // patch_size[0], l_h // patch_size[1], l_w // patch_size[2]) + .flatten() + ) + action_seq_id = torch.arange(b)[:, None, None, None].expand(-1, a_f, a_h, a_w).flatten() + seq_ids = torch.cat([latent_seq_id] * 2 + [action_seq_id] * 2) + + latent_frame_id = ( + torch.arange(l_f)[None, :, None, None] + .expand(b, -1, l_h // patch_size[1], l_w // patch_size[2])[None] + .flatten() + ) + action_frame_id = torch.arange(a_f)[None, :, None, None].expand(b, -1, a_h, a_w)[None].flatten() + frame_ids = torch.cat( + [latent_frame_id // chunk_size * 2] * 2 + [action_frame_id // chunk_size * 2 + 1] * 2 + ) + + noise_ids = torch.cat( + [ + torch.zeros_like(latent_frame_id), + torch.ones_like(latent_frame_id), + torch.zeros_like(action_frame_id), + torch.ones_like(action_frame_id), + ] + ) + + seq_ids = F.pad(seq_ids, (0, padded_length), value=-1) + frame_ids = F.pad(frame_ids, (0, padded_length), value=-1) + noise_ids = F.pad(noise_ids, (0, padded_length), value=-1) + + mask_mod = FlexAttnFunc._get_mask_mod( + seq_ids.long().to(device), frame_ids.long().to(device), noise_ids.long().to(device), window_size + ) + block_mask = FlexAttnFunc.compiled_create_block_mask( + mask_mod, 1, 1, len(seq_ids), len(seq_ids), device=device, _compile=True + ) + FlexAttnFunc.attention_mask = block_mask + + text_seq_ids = torch.arange(b)[:, None].expand(-1, 512).flatten() + mask_mod_cross = FlexAttnFunc._get_cross_mask_mod( + seq_ids.long().to(device), text_seq_ids.long().to(device) + ) + block_mask_cross = FlexAttnFunc.compiled_create_block_mask( + mask_mod_cross, 1, 1, len(seq_ids), len(text_seq_ids), device=device, _compile=True + ) + FlexAttnFunc.cross_attention_mask = block_mask_cross + + @staticmethod + @torch.no_grad() + def _get_cross_mask_mod(seq_ids, text_seq_ids): + def seq_mask(b, h, q_idx, kv_idx): + return ( + (seq_ids[q_idx] == text_seq_ids[kv_idx]) & (seq_ids[q_idx] >= 0) & (text_seq_ids[kv_idx] >= 0) + ) + + return seq_mask + + @staticmethod + @torch.no_grad() + def _get_mask_mod(seq_ids, frame_ids, noise_ids, window_size): + from torch.nn.attention.flex_attention import and_masks, or_masks + + def seq_mask(b, h, q_idx, kv_idx): + return (seq_ids[q_idx] == seq_ids[kv_idx]) & (seq_ids[q_idx] >= 0) & (seq_ids[kv_idx] >= 0) + + def block_causal_mask(b, h, q_idx, kv_idx): + return frame_ids[kv_idx] <= frame_ids[q_idx] + + def block_causal_mask_exclude_self(b, h, q_idx, kv_idx): + return frame_ids[kv_idx] < frame_ids[q_idx] + + def block_self_mask(b, h, q_idx, kv_idx): + return frame_ids[kv_idx] == frame_ids[q_idx] + + def clean2clean_mask(b, h, q_idx, kv_idx): + return (noise_ids[q_idx] == 1) & (noise_ids[kv_idx] == 1) + + def noise2clean_mask(b, h, q_idx, kv_idx): + return (noise_ids[q_idx] == 0) & (noise_ids[kv_idx] == 1) + + def noise2noise_mask(b, h, q_idx, kv_idx): + return (noise_ids[q_idx] == 0) & (noise_ids[kv_idx] == 0) + + def block_window_mask(b, h, q_idx, kv_idx, window_size: int): + return (frame_ids[q_idx] - frame_ids[kv_idx]).abs() <= window_size + + mask_list = [] + mask_list.append(and_masks(clean2clean_mask, block_causal_mask)) + mask_list.append(and_masks(noise2clean_mask, block_causal_mask_exclude_self)) + mask_list.append(and_masks(noise2noise_mask, block_self_mask)) + mask = or_masks(*mask_list) + mask = and_masks(mask, seq_mask) + mask = and_masks(mask, partial(block_window_mask, window_size=window_size)) + return mask + + +class WanRotaryPosEmbed(nn.Module): + """Rotary position embedding with separate frequency bases for frame / height / width.""" + + def __init__(self, attention_head_dim: int, patch_size, max_seq_len: int, theta: float = 10000.0): + super().__init__() + + self.attention_head_dim = attention_head_dim + self.patch_size = patch_size + self.max_seq_len = max_seq_len + self.theta = theta + + self.f_dim = self.attention_head_dim - 2 * (self.attention_head_dim // 3) + self.h_dim = self.attention_head_dim // 3 + self.w_dim = self.attention_head_dim // 3 + + f_freqs_base, h_freqs_base, w_freqs_base = self._precompute_freqs_base() + self.f_freqs_base = f_freqs_base + self.h_freqs_base = h_freqs_base + self.w_freqs_base = w_freqs_base + + def _precompute_freqs_base(self): + f_freqs_base = 1.0 / ( + self.theta ** (torch.arange(0, self.f_dim, 2)[: (self.f_dim // 2)].double() / self.f_dim) + ) + h_freqs_base = 1.0 / ( + self.theta ** (torch.arange(0, self.h_dim, 2)[: (self.h_dim // 2)].double() / self.h_dim) + ) + w_freqs_base = 1.0 / ( + self.theta ** (torch.arange(0, self.w_dim, 2)[: (self.w_dim // 2)].double() / self.w_dim) + ) + return f_freqs_base, h_freqs_base, w_freqs_base + + def forward(self, grid_ids): + with torch.no_grad(): + f_freqs = grid_ids[:, 0, :].unsqueeze(-1) * self.f_freqs_base.to(grid_ids.device) + h_freqs = grid_ids[:, 1, :].unsqueeze(-1) * self.h_freqs_base.to(grid_ids.device) + w_freqs = grid_ids[:, 2, :].unsqueeze(-1) * self.w_freqs_base.to(grid_ids.device) + freqs = torch.cat([f_freqs, h_freqs, w_freqs], dim=-1).float() + freqs_cis = torch.polar(torch.ones_like(freqs), freqs) + + return freqs_cis + + +class WanAttention(nn.Module): + """Self/cross attention with KV-caching for autoregressive streaming inference. + + Backends: ``torch`` (default SDPA), ``flashattn`` (optional), ``flex`` (training masks). + """ + + def __init__( + self, + dim, + heads=8, + dim_head=64, + eps=1e-5, + dropout=0.0, + cross_attention_dim_head=None, + attn_mode="torch", + ): + super().__init__() + if attn_mode == "torch": + self.attn_op = custom_sdpa + elif attn_mode == "flashattn": + self.attn_op = _load_flash_attn_func() + elif attn_mode == "flex": + self.attn_op = FlexAttnFunc(cross_attention_dim_head is not None) + else: + raise ValueError( + f"Unsupported attention mode: {attn_mode}, only support 'torch', 'flashattn' and 'flex'" + ) + + self.inner_dim = dim_head * heads + self.heads = heads + self.cross_attention_dim_head = cross_attention_dim_head + self.kv_inner_dim = ( + self.inner_dim if cross_attention_dim_head is None else cross_attention_dim_head * heads + ) + + self.to_q = nn.Linear(dim, self.inner_dim, bias=True) + self.to_k = nn.Linear(dim, self.kv_inner_dim, bias=True) + self.to_v = nn.Linear(dim, self.kv_inner_dim, bias=True) + self.to_out = nn.ModuleList([nn.Linear(self.inner_dim, dim, bias=True), nn.Dropout(dropout)]) + self.norm_q = nn.RMSNorm(dim_head * heads, eps=eps, elementwise_affine=True) + self.norm_k = nn.RMSNorm(dim_head * heads, eps=eps, elementwise_affine=True) + # KV cache only lives on self-attention modules (cross_attention_dim_head is None). + self.attn_caches = {} if cross_attention_dim_head is None else None + + def clear_pred_cache(self, cache_name): + if self.attn_caches is None: + return + cache = self.attn_caches[cache_name] + is_pred = cache["is_pred"] + cache["mask"][is_pred] = False + + def clear_cache(self, cache_name): + if self.attn_caches is None: + return + self.attn_caches[cache_name] = None + + def init_kv_cache(self, cache_name, total_token_len, num_head, head_dim, device, dtype, batch_size): + if self.attn_caches is None: + return + self.attn_caches[cache_name] = { + "k": torch.empty([batch_size, total_token_len, num_head, head_dim], device=device, dtype=dtype), + "v": torch.empty([batch_size, total_token_len, num_head, head_dim], device=device, dtype=dtype), + "id": torch.full((total_token_len,), -1, device=device), + "mask": torch.zeros((total_token_len,), dtype=torch.bool, device=device), + "is_pred": torch.zeros((total_token_len,), dtype=torch.bool, device=device), + } + + def allocate_slots(self, cache_name, key_size): + cache = self.attn_caches[cache_name] + mask = cache["mask"] + ids = cache["id"] + free = (~mask).nonzero(as_tuple=False).squeeze(-1) + + if free.numel() < key_size: + used = mask.nonzero(as_tuple=False).squeeze(-1) + + used_ids = ids[used] + order = torch.argsort(used_ids) + need = key_size - free.numel() + to_free = used[order[:need]] + + mask[to_free] = False + ids[to_free] = -1 + free = (~mask).nonzero(as_tuple=False).squeeze(-1) + + if free.numel() < key_size: + raise RuntimeError(f"KV cache exhausted: need {key_size} free slots, have {free.numel()}.") + return free[:key_size] + + def _next_cache_id(self, cache_name): + ids = self.attn_caches[cache_name]["id"] + mask = self.attn_caches[cache_name]["mask"] + + if mask.any(): + return ids[mask].max() + 1 + else: + return torch.tensor(0, device=ids.device, dtype=ids.dtype) + + def update_cache(self, cache_name, key, value, is_pred): + cache = self.attn_caches[cache_name] + + key_size = key.shape[1] + slots = self.allocate_slots(cache_name, key_size) + + new_id = self._next_cache_id(cache_name) + + cache["k"][:, slots] = key + cache["v"][:, slots] = value + cache["mask"][slots] = True + cache["id"][slots] = new_id + cache["is_pred"][slots] = is_pred + return slots + + def restore_cache(self, cache_name, slots): + self.attn_caches[cache_name]["mask"][slots] = False + + def forward(self, q, k, v, rotary_emb, update_cache=0, cache_name="pos"): + kv_cache = ( + self.attn_caches[cache_name] + if (self.attn_caches is not None) and (cache_name in self.attn_caches) + else None + ) + + query, key, value = self.to_q(q), self.to_k(k), self.to_v(v) + query = self.norm_q(query) + query = query.unflatten(2, (self.heads, -1)) + key = self.norm_k(key) + key = key.unflatten(2, (self.heads, -1)) + value = value.unflatten(2, (self.heads, -1)) + if rotary_emb is not None: + + def apply_rotary_emb(x, freqs): + x_out = torch.view_as_complex( + x.to(torch.float64).reshape(x.shape[0], x.shape[1], x.shape[2], -1, 2) + ) + x_out = torch.view_as_real(x_out * freqs).flatten(3) + return x_out.to(x.dtype) + + query = apply_rotary_emb(query, rotary_emb) + key = apply_rotary_emb(key, rotary_emb) + slots = None + if kv_cache is not None and kv_cache["k"] is not None: + slots = self.update_cache(cache_name, key, value, is_pred=(update_cache == 1)) + key_pool = self.attn_caches[cache_name]["k"] + value_pool = self.attn_caches[cache_name]["v"] + mask = self.attn_caches[cache_name]["mask"] + valid = mask.nonzero(as_tuple=False).squeeze(-1) + key = key_pool[:, valid] + value = value_pool[:, valid] + + hidden_states = self.attn_op(query, key, value) + + if update_cache == 0 and kv_cache is not None and kv_cache["k"] is not None: + self.restore_cache(cache_name, slots) + + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.type_as(query) + hidden_states = self.to_out[0](hidden_states) + hidden_states = self.to_out[1](hidden_states) + return hidden_states + + +# Dual-stream Wan2.2 transformer +class WanTimeTextImageEmbedding(nn.Module): + def __init__(self, dim, time_freq_dim, time_proj_dim, text_embed_dim, pos_embed_seq_len): + super().__init__() + + self.timesteps_proj = Timesteps( + num_channels=time_freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0 + ) + self.time_embedder = TimestepEmbedding(in_channels=time_freq_dim, time_embed_dim=dim) + self.act_fn = nn.SiLU() + self.time_proj = nn.Linear(dim, time_proj_dim) + self.text_embedder = PixArtAlphaTextProjection(text_embed_dim, dim, act_fn="gelu_tanh") + + def forward(self, timestep: torch.Tensor, dtype=None): + b, seq_len = timestep.shape + timestep = timestep.reshape(-1) + timestep = self.timesteps_proj(timestep) + time_embedder_dtype = self.time_embedder.linear_1.weight.dtype + if timestep.dtype != time_embedder_dtype and time_embedder_dtype != torch.int8: + timestep = timestep.to(time_embedder_dtype) + temb = self.time_embedder(timestep).to(dtype=dtype) + timestep_proj = self.time_proj(self.act_fn(temb)) + return temb.reshape(b, seq_len, -1), timestep_proj.reshape(b, seq_len, -1) + + +class WanTransformerBlock(nn.Module): + def __init__(self, dim, ffn_dim, num_heads, cross_attn_norm=False, eps=1e-6, attn_mode: str = "torch"): + super().__init__() + self.attn_mode = attn_mode + + # 1. Self-attention + self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.attn1 = WanAttention( + dim=dim, + heads=num_heads, + dim_head=dim // num_heads, + eps=eps, + cross_attention_dim_head=None, + attn_mode=attn_mode, + ) + + # 2. Cross-attention + self.attn2 = WanAttention( + dim=dim, + heads=num_heads, + dim_head=dim // num_heads, + eps=eps, + cross_attention_dim_head=dim // num_heads, + attn_mode=attn_mode, + ) + self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() + + # 3. Feed-forward + self.ffn = FeedForward(dim, inner_dim=ffn_dim, activation_fn="gelu-approximate") + self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False) + + self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def forward( + self, hidden_states, encoder_hidden_states, temb, rotary_emb, update_cache=0, cache_name="pos" + ) -> torch.Tensor: + temb_scale_shift_table = self.scale_shift_table[None] + temb.float() + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = rearrange( + temb_scale_shift_table, "b l n c -> b n l c" + ).chunk(6, dim=1) + shift_msa = shift_msa.squeeze(1) + scale_msa = scale_msa.squeeze(1) + gate_msa = gate_msa.squeeze(1) + c_shift_msa = c_shift_msa.squeeze(1) + c_scale_msa = c_scale_msa.squeeze(1) + c_gate_msa = c_gate_msa.squeeze(1) + # 1. Self-attention + norm_hidden_states = (self.norm1(hidden_states.float()) * (1.0 + scale_msa) + shift_msa).type_as( + hidden_states + ) + attn_output = self.attn1( + norm_hidden_states, + norm_hidden_states, + norm_hidden_states, + rotary_emb, + update_cache=update_cache, + cache_name=cache_name, + ) + hidden_states = (hidden_states.float() + attn_output * gate_msa).type_as(hidden_states) + + # 2. Cross-attention + norm_hidden_states = self.norm2(hidden_states.float()).type_as(hidden_states) + attn_output = self.attn2( + norm_hidden_states, + encoder_hidden_states, + encoder_hidden_states, + None, + update_cache=0, + cache_name=cache_name, + ) + hidden_states = hidden_states + attn_output + + # 3. Feed-forward + norm_hidden_states = (self.norm3(hidden_states.float()) * (1.0 + c_scale_msa) + c_shift_msa).type_as( + hidden_states + ) + + ff_output = self.ffn(norm_hidden_states) + + hidden_states = (hidden_states.float() + ff_output.float() * c_gate_msa).type_as(hidden_states) + return hidden_states + + +class WanTransformer3DModel(ModelMixin, ConfigMixin): + """Dual-stream (video + action) Wan2.2 DiT backbone with autoregressive KV caching.""" + + _supports_gradient_checkpointing = True + _skip_layerwise_casting_patterns = [ + "patch_embedding_mlp", + "condition_embedder", + "condition_embedder_action", + "norm", + ] + _no_split_modules = ["WanTransformerBlock"] + _keep_in_fp32_modules = [ + "time_embedder", + "scale_shift_table", + "scale_shift_table_action", + "norm1", + "action_norm1", + "text_norm1", + "norm2", + "action_norm2", + "text_norm2", + "norm3", + "action_norm3", + "text_norm3", + ] + _keys_to_ignore_on_load_unexpected = ["norm_added_q"] + _repeated_blocks = ["WanTransformerBlock"] + + @register_to_config + def __init__( + self, + patch_size=(1, 2, 2), + num_attention_heads=24, + attention_head_dim=128, + in_channels=48, + out_channels=48, + action_dim=30, + text_dim=4096, + freq_dim=256, + ffn_dim=14336, + num_layers=30, + cross_attn_norm=True, + eps=1e-06, + rope_max_seq_len=1024, + pos_embed_seq_len=None, + attn_mode="torch", + ): + super().__init__() + self.patch_size = patch_size + self.num_attention_heads = num_attention_heads + self.attention_head_dim = attention_head_dim + inner_dim = num_attention_heads * attention_head_dim + self.rope = WanRotaryPosEmbed(attention_head_dim, patch_size, rope_max_seq_len) + self.patch_embedding_mlp = nn.Linear( + in_channels * patch_size[0] * patch_size[1] * patch_size[2], inner_dim + ) + self.action_embedder = nn.Linear(action_dim, inner_dim) + self.condition_embedder = WanTimeTextImageEmbedding( + dim=inner_dim, + time_freq_dim=freq_dim, + time_proj_dim=inner_dim * 6, + text_embed_dim=text_dim, + pos_embed_seq_len=pos_embed_seq_len, + ) + self.condition_embedder_action = deepcopy(self.condition_embedder) + + self.blocks = nn.ModuleList( + [ + WanTransformerBlock( + inner_dim, ffn_dim, num_attention_heads, cross_attn_norm, eps, attn_mode=attn_mode + ) + for _ in range(num_layers) + ] + ) + + self.norm_out = FP32LayerNorm(inner_dim, eps, elementwise_affine=False) + self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size)) + self.action_proj_out = nn.Linear(inner_dim, action_dim) + self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5) + + # KV-cache management for autoregressive streaming inference + def clear_cache(self, cache_name): + for block in self.blocks: + block.attn1.clear_cache(cache_name) + + def clear_pred_cache(self, cache_name): + for block in self.blocks: + block.attn1.clear_pred_cache(cache_name) + + def create_empty_cache( + self, + cache_name, + attn_window, + latent_token_per_chunk, + action_token_per_chunk, + device, + dtype, + batch_size, + ): + total_token_len = (attn_window // 2) * latent_token_per_chunk + ( + attn_window // 2 + ) * action_token_per_chunk + for block in self.blocks: + block.attn1.init_kv_cache( + cache_name, + total_token_len, + self.num_attention_heads, + self.attention_head_dim, + device, + dtype, + batch_size, + ) + + # Embedding helpers (shared by train + inference paths) + def _input_embed(self, latents, input_type="latent"): + if input_type == "latent": + hidden_states = rearrange( + latents, + "b c (f p1) (h p2) (w p3) -> b (f h w) (c p1 p2 p3)", + p1=self.patch_size[0], + p2=self.patch_size[1], + p3=self.patch_size[2], + ) + hidden_states = self.patch_embedding_mlp(hidden_states) + elif input_type == "action": + hidden_states = rearrange(latents, "b c f h w -> b (f h w) c") + hidden_states = self.action_embedder(hidden_states) + elif input_type == "text": + hidden_states = self.condition_embedder.text_embedder(latents) + else: + raise ValueError(f"Unsupported input type: {input_type}") + return hidden_states + + def _time_embed(self, timesteps, h, w, dtype, action_mode=False): + patch_scale_h, patch_scale_w = (1, 1) if action_mode else (self.patch_size[1], self.patch_size[2]) + latent_time_steps = torch.repeat_interleave( + timesteps, (h // patch_scale_h) * (w // patch_scale_w), dim=1 + ) + current_condition_embedder = ( + self.condition_embedder_action if action_mode else self.condition_embedder + ) + temb, timestep_proj = current_condition_embedder(latent_time_steps, dtype=dtype) + timestep_proj = timestep_proj.unflatten(2, (6, -1)) # B L 6 C + return temb, timestep_proj + + # Dual-stream training forward (flow matching). Requires attn_mode='flex'. + def forward_train(self, input_dict): + input_dict["latent_dict"]["noisy_latents"] = input_dict["latent_dict"]["noisy_latents"].to( + torch.bfloat16 + ) + input_dict["latent_dict"]["latent"] = input_dict["latent_dict"]["latent"].to(torch.bfloat16) + input_dict["action_dict"]["noisy_latents"] = input_dict["action_dict"]["noisy_latents"].to( + torch.bfloat16 + ) + input_dict["action_dict"]["latent"] = input_dict["action_dict"]["latent"].to(torch.bfloat16) + + latent_dict = input_dict["latent_dict"] + action_dict = input_dict["action_dict"] + batch_size = latent_dict["noisy_latents"].shape[0] + + latent_hidden_states = self._input_embed(latent_dict["noisy_latents"], input_type="latent").flatten( + 0, 1 + )[None] + action_hidden_states = self._input_embed(action_dict["noisy_latents"], input_type="action").flatten( + 0, 1 + )[None] + text_hidden_states = self._input_embed(latent_dict["text_emb"], input_type="text") + + text_hidden_states = text_hidden_states.flatten(0, 1)[None] + + condition_latent_hidden_states = self._input_embed( + latent_dict["latent"], input_type="latent" + ).flatten(0, 1)[None] + condition_action_hidden_states = self._input_embed( + action_dict["latent"], input_type="action" + ).flatten(0, 1)[None] + + hidden_states = torch.cat( + [ + latent_hidden_states, + condition_latent_hidden_states, + action_hidden_states, + condition_action_hidden_states, + ], + dim=1, + ) + + latent_grid_id = latent_dict["grid_id"].permute(1, 0, 2).flatten(1)[None] + action_grid_id = action_dict["grid_id"].permute(1, 0, 2).flatten(1)[None] + full_grid_id = torch.cat([latent_grid_id] * 2 + [action_grid_id] * 2, dim=2) + + rotary_emb = self.rope(full_grid_id)[:, :, None] + + latent_time_steps = torch.cat( + [latent_dict["timesteps"].flatten(0, 1), latent_dict["cond_timesteps"].flatten(0, 1)] + )[None] + action_time_steps = torch.cat( + [action_dict["timesteps"].flatten(0, 1), action_dict["cond_timesteps"].flatten(0, 1)] + )[None] + latent_temb, latent_timestep_proj = self._time_embed( + latent_time_steps, + latent_dict["noisy_latents"].shape[-2], + latent_dict["noisy_latents"].shape[-1], + dtype=hidden_states.dtype, + action_mode=False, + ) + action_temb, action_timestep_proj = self._time_embed( + action_time_steps, + action_dict["noisy_latents"].shape[-2], + action_dict["noisy_latents"].shape[-1], + dtype=hidden_states.dtype, + action_mode=True, + ) + temb = torch.cat([latent_temb, action_temb], dim=1) + timestep_proj = torch.cat([latent_timestep_proj, action_timestep_proj], dim=1) + + total_length = hidden_states.shape[1] + padded_length = (128 - total_length % 128) % 128 + hidden_states = F.pad(hidden_states, (0, 0, 0, padded_length)) + rotary_emb = F.pad(rotary_emb, (0, 0, 0, 0, 0, padded_length)) + temb = F.pad(temb, (0, 0, 0, padded_length)) + timestep_proj = F.pad(timestep_proj, (0, 0, 0, 0, 0, padded_length)) + + split_list = [ + latent_hidden_states.shape[1], + condition_latent_hidden_states.shape[1], + action_hidden_states.shape[1], + condition_action_hidden_states.shape[1], + padded_length, + ] + + FlexAttnFunc.init_mask( + latent_dict["noisy_latents"].shape, + action_dict["noisy_latents"].shape, + padded_length, + input_dict["chunk_size"], + window_size=input_dict["window_size"], + patch_size=self.patch_size, + device=hidden_states.device, + ) + + for block in self.blocks: + hidden_states = block( + hidden_states, text_hidden_states, timestep_proj, rotary_emb, update_cache=False + ) + temb_scale_shift_table = self.scale_shift_table[None] + temb[:, :, None, ...] + shift, scale = rearrange(temb_scale_shift_table, "b l n c -> b n l c").chunk(2, dim=1) + shift = shift.to(hidden_states.device).squeeze(1) + scale = scale.to(hidden_states.device).squeeze(1) + hidden_states = (self.norm_out(hidden_states.float()) * (1.0 + scale) + shift).type_as(hidden_states) + latent_hidden_states, _, action_hidden_states, _, _ = torch.split(hidden_states, split_list, dim=1) + latent_hidden_states = self.proj_out(latent_hidden_states) + latent_hidden_states = rearrange( + latent_hidden_states, "1 (b l) (n c) -> b (l n) c", n=math.prod(self.patch_size), b=batch_size + ) + action_hidden_states = self.action_proj_out(action_hidden_states) + action_hidden_states = rearrange(action_hidden_states, "1 (b l) c -> b l c", b=batch_size) + + return latent_hidden_states, action_hidden_states + + # Single-stream inference forward (one denoising step for one stream) + def forward(self, input_dict, update_cache=0, cache_name="pos", action_mode=False, train_mode=False): + if train_mode: + return self.forward_train(input_dict) + if action_mode: # action input emb + latent_hidden_states = rearrange(input_dict["noisy_latents"], "b c f h w -> b (f h w) c") + latent_hidden_states = self.action_embedder(latent_hidden_states) # B L1 C + else: # latent input emb + latent_hidden_states = rearrange( + input_dict["noisy_latents"], + "b c (f p1) (h p2) (w p3) -> b (f h w) (c p1 p2 p3)", + p1=self.patch_size[0], + p2=self.patch_size[1], + p3=self.patch_size[2], + ) + latent_hidden_states = self.patch_embedding_mlp(latent_hidden_states) + text_hidden_states = self.condition_embedder.text_embedder(input_dict["text_emb"]) # B L2 C + + latent_grid_id = input_dict["grid_id"] + rotary_emb = self.rope(latent_grid_id)[:, :, None] # 1 L 1 C + patch_scale_h, patch_scale_w = (1, 1) if action_mode else (self.patch_size[1], self.patch_size[2]) + + latent_time_steps = torch.repeat_interleave( + input_dict["timesteps"], + (input_dict["noisy_latents"].shape[-2] // patch_scale_h) + * (input_dict["noisy_latents"].shape[-1] // patch_scale_w), + dim=1, + ) # L + current_condition_embedder = ( + self.condition_embedder_action if action_mode else self.condition_embedder + ) + temb, timestep_proj = current_condition_embedder(latent_time_steps, dtype=latent_hidden_states.dtype) + timestep_proj = timestep_proj.unflatten(2, (6, -1)) # B L 6 C + + for block in self.blocks: + latent_hidden_states = block( + latent_hidden_states, + text_hidden_states, + timestep_proj, + rotary_emb, + update_cache=update_cache, + cache_name=cache_name, + ) + temb_scale_shift_table = self.scale_shift_table[None] + temb[:, :, None, ...] + shift, scale = rearrange(temb_scale_shift_table, "b l n c -> b n l c").chunk(2, dim=1) + shift = shift.to(latent_hidden_states.device).squeeze(1) + scale = scale.to(latent_hidden_states.device).squeeze(1) + latent_hidden_states = (self.norm_out(latent_hidden_states.float()) * (1.0 + scale) + shift).type_as( + latent_hidden_states + ) + + if action_mode: + latent_hidden_states = self.action_proj_out(latent_hidden_states) + else: + latent_hidden_states = self.proj_out(latent_hidden_states) + latent_hidden_states = rearrange( + latent_hidden_states, "b l (n c) -> b (l n) c", n=math.prod(self.patch_size) + ) + + return latent_hidden_states + + +class WanVAEStreamingWrapper: + """Wraps an ``AutoencoderKLWan`` encoder to support causal streaming encoding across chunks.""" + + def __init__(self, vae_model): + self.vae = vae_model + self.encoder = vae_model.encoder + self.quant_conv = vae_model.quant_conv + + if hasattr(self.vae, "_cached_conv_counts"): + self.enc_conv_num = self.vae._cached_conv_counts["encoder"] + else: + count = 0 + for m in self.encoder.modules(): + if m.__class__.__name__ == "WanCausalConv3d": + count += 1 + self.enc_conv_num = count + + self.clear_cache() + + def clear_cache(self): + self.feat_cache = [None] * self.enc_conv_num + + def encode_chunk(self, x_chunk): + if hasattr(self.vae.config, "patch_size") and self.vae.config.patch_size is not None: + x_chunk = _vae_patchify(x_chunk, self.vae.config.patch_size) + feat_idx = [0] + out = self.encoder(x_chunk, feat_cache=self.feat_cache, feat_idx=feat_idx) + enc = self.quant_conv(out) + return enc diff --git a/src/lerobot/scripts/lerobot_eval.py b/src/lerobot/scripts/lerobot_eval.py index 1ec4ea75f..722763d6e 100644 --- a/src/lerobot/scripts/lerobot_eval.py +++ b/src/lerobot/scripts/lerobot_eval.py @@ -169,6 +169,7 @@ def rollout( env_features: dict | None = None, recording_repo_id: str | None = None, recording_private: bool = False, + predicted_latents_callback: Callable[[PreTrainedPolicy], None] | None = None, ) -> dict: """Run a batched policy rollout once through a batch of environments. @@ -198,6 +199,9 @@ def rollout( are returned optionally because they typically take more memory to cache. Defaults to False. render_callback: Optional rendering callback to be used after the environments are reset, and after every step. + predicted_latents_callback: Optional callback invoked after every ``select_action`` with the policy + itself. World-model policies (e.g. LingBot-VA) stash predicted video latents on + ``policy.last_predicted_latents``; this lets the caller concatenate chunks and decode once. Returns: The dictionary described above. """ @@ -276,6 +280,8 @@ def rollout( observation = preprocessor(observation) with torch.inference_mode(): action = policy.select_action(observation) + if predicted_latents_callback is not None: + predicted_latents_callback(policy) action = postprocessor(action) action_transition = {ACTION: action} @@ -295,12 +301,22 @@ def rollout( # available if none of the envs finished. if "final_info" in info: final_info = info["final_info"] - if not isinstance(final_info, dict): - raise RuntimeError( - "Unsupported `final_info` format: expected dict (Gymnasium >= 1.0). " - "You're likely using an older version of gymnasium (< 1.0). Please upgrade." + if isinstance(final_info, dict): + is_success = final_info.get("is_success", [False] * env.num_envs) + successes = ( + is_success.tolist() + if hasattr(is_success, "tolist") + else [bool(is_success)] * env.num_envs ) - successes = final_info["is_success"].tolist() + else: + # Gymnasium < 1.0 returns final_info as a per-env sequence/object array, + # with entries set to a dict only for envs that just finished. + successes = [] + for item in final_info: + if isinstance(item, dict) and "is_success" in item: + successes.append(bool(item["is_success"])) + else: + successes.append(False) elif "is_success" in info: is_success = info["is_success"] successes = ( @@ -400,6 +416,7 @@ def eval_policy( env_features: dict | None = None, recording_repo_id: str | None = None, recording_private: bool = False, + save_predicted_video: bool = False, ) -> dict: """ Args: @@ -418,6 +435,11 @@ def eval_policy( if max_episodes_rendered > 0 and not videos_dir: raise ValueError("If max_episodes_rendered > 0, videos_dir must be provided.") + # World-model policies (e.g. LingBot-VA) opt into predicted-video saving via their config. + save_predicted_video = save_predicted_video or bool( + getattr(getattr(policy, "config", None), "save_predicted_video", False) + ) + if not isinstance(policy, PreTrainedPolicy): exc = ValueError( f"Policy of type 'PreTrainedPolicy' is expected, but type '{type(policy)}' was provided." @@ -461,6 +483,22 @@ def eval_policy( if max_episodes_rendered > 0: video_paths: list[str] = [] + if save_predicted_video: + if not videos_dir: + raise ValueError("If save_predicted_video is True, videos_dir must be provided.") + predicted_video_paths: list[str] = [] + n_predicted_rendered = 0 + + # Collect predicted-video latents across a rollout (world-model policies only). The latents are + # concatenated and decoded once after the rollout, matching upstream LingBot-VA's visualization path. + def collect_predicted_latents(policy: PreTrainedPolicy): + latents = getattr(policy, "last_predicted_latents", None) + if latents is not None: + pred_latents.append( + latents.detach().to("cpu") if hasattr(latents, "detach") else torch.as_tensor(latents).cpu() + ) + policy.last_predicted_latents = None + if return_episode_data: episode_data: dict | None = None @@ -472,6 +510,9 @@ def eval_policy( if max_episodes_rendered > 0: ep_frames: list[np.ndarray] = [] + if save_predicted_video: + pred_latents: list[torch.Tensor] = [] + if start_seed is None: seeds = None else: @@ -492,6 +533,7 @@ def eval_policy( env_features=env_features, recording_repo_id=recording_repo_id, recording_private=recording_private, + predicted_latents_callback=collect_predicted_latents if save_predicted_video else None, ) # Figure out where in each rollout sequence the first done condition was encountered (results after @@ -557,6 +599,35 @@ def eval_policy( threads.append(thread) n_episodes_rendered += 1 + # Maybe save the policy's predicted (imagined) video for this batch's rollout. + if save_predicted_video and len(pred_latents) > 0: + predicted_latent = torch.cat(pred_latents, dim=2) + decoder = getattr(policy, "decode_predicted_latents", None) or getattr( + policy, "_decode_predicted_video", None + ) + if decoder is None: + raise AttributeError( + "Policy config requested predicted-video saving, but the policy does not expose " + "`decode_predicted_latents` or `_decode_predicted_video`." + ) + predicted_video = decoder(predicted_latent) + if hasattr(predicted_video, "detach"): + predicted_video = predicted_video.detach().to("cpu").numpy() + videos_dir.mkdir(parents=True, exist_ok=True) + predicted_video_path = videos_dir / f"pred_episode_{n_predicted_rendered}.mp4" + predicted_video_paths.append(str(predicted_video_path)) + thread = threading.Thread( + target=write_video, + args=( + str(predicted_video_path), + predicted_video, + env.unwrapped.metadata["render_fps"], + ), + ) + thread.start() + threads.append(thread) + n_predicted_rendered += 1 + progbar.set_postfix( {"running_success_rate": f"{np.mean(all_successes[:n_episodes]).item() * 100:.1f}%"} ) @@ -600,6 +671,9 @@ def eval_policy( if max_episodes_rendered > 0: info["video_paths"] = video_paths + if save_predicted_video: + info["predicted_video_paths"] = predicted_video_paths + return info @@ -740,9 +814,10 @@ class TaskMetrics(TypedDict): max_rewards: list[float] successes: list[bool] video_paths: list[str] + predicted_video_paths: list[str] -ACC_KEYS = ("sum_rewards", "max_rewards", "successes", "video_paths") +ACC_KEYS = ("sum_rewards", "max_rewards", "successes", "video_paths", "predicted_video_paths") def eval_one( @@ -791,6 +866,7 @@ def eval_one( max_rewards=[ep["max_reward"] for ep in per_episode], successes=[ep["success"] for ep in per_episode], video_paths=task_result.get("video_paths", []), + predicted_video_paths=task_result.get("predicted_video_paths", []), ) @@ -851,6 +927,7 @@ def run_one( if max_episodes_rendered > 0: metrics.setdefault("video_paths", []) + metrics.setdefault("predicted_video_paths", []) return task_group, task_id, metrics @@ -908,11 +985,11 @@ def eval_policy_all( _append("sum_rewards", metrics.get("sum_rewards")) _append("max_rewards", metrics.get("max_rewards")) _append("successes", metrics.get("successes")) - # video_paths is list-like - paths = metrics.get("video_paths", []) - if paths: - group_acc[group]["video_paths"].extend(paths) - overall["video_paths"].extend(paths) + for key in ("video_paths", "predicted_video_paths"): + paths = metrics.get(key, []) + if paths: + group_acc[group][key].extend(paths) + overall[key].extend(paths) # Choose runner (sequential vs threaded) task_runner = partial( @@ -984,6 +1061,7 @@ def eval_policy_all( "pc_success": _agg_from_list(acc["successes"]) * 100 if acc["successes"] else float("nan"), "n_episodes": len(acc["sum_rewards"]), "video_paths": list(acc["video_paths"]), + "predicted_video_paths": list(acc["predicted_video_paths"]), } # overall aggregates @@ -995,6 +1073,7 @@ def eval_policy_all( "eval_s": time.time() - start_t, "eval_ep_s": (time.time() - start_t) / max(1, len(overall["sum_rewards"])), "video_paths": list(overall["video_paths"]), + "predicted_video_paths": list(overall["predicted_video_paths"]), } return { diff --git a/tests/policies/lingbot_va/test_configuration_lingbot_va.py b/tests/policies/lingbot_va/test_configuration_lingbot_va.py new file mode 100644 index 000000000..089d5ec02 --- /dev/null +++ b/tests/policies/lingbot_va/test_configuration_lingbot_va.py @@ -0,0 +1,78 @@ +#!/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.configs.policies import PreTrainedConfig +from lerobot.configs.types import FeatureType, PolicyFeature +from lerobot.policies.lingbot_va.configuration_lingbot_va import LingBotVAConfig +from lerobot.utils.constants import ACTION, OBS_IMAGES + + +def make_config(**overrides) -> LingBotVAConfig: + kwargs = {"device": "cpu"} + kwargs.update(overrides) + return LingBotVAConfig(**kwargs) + + +def test_registered_in_choice_registry() -> None: + assert "lingbot_va" in PreTrainedConfig.get_known_choices() + assert PreTrainedConfig.get_choice_class("lingbot_va") is LingBotVAConfig + + +def test_type_property() -> None: + assert make_config().type == "lingbot_va" + + +def test_chunk_size_and_action_steps() -> None: + cfg = make_config(frame_chunk_size=4, action_per_frame=4) + assert cfg.chunk_size == 16 + assert cfg.n_action_steps == 16 + assert cfg.action_delta_indices == list(range(16)) + assert cfg.observation_delta_indices == list(range(16)) + assert cfg.reward_delta_indices is None + + +def test_optimizer_and_scheduler_presets() -> None: + cfg = make_config() + opt = cfg.get_optimizer_preset() + assert opt.lr == cfg.optimizer_lr + sched = cfg.get_scheduler_preset() + assert sched.num_warmup_steps == cfg.scheduler_warmup_steps + + +def test_validate_features_sets_action_feature() -> None: + cfg = make_config() + cfg.input_features = {f"{OBS_IMAGES}.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 128, 128))} + cfg.output_features = {} + cfg.validate_features() + assert ACTION in cfg.output_features + assert cfg.output_features[ACTION].shape == (len(cfg.used_action_channel_ids),) + + +def test_validate_features_no_visual_raises() -> None: + cfg = make_config() + cfg.input_features = {} + cfg.output_features = {} + with pytest.raises(ValueError, match="at least one visual input feature"): + cfg.validate_features() + + +def test_invalid_attn_mode_raises() -> None: + with pytest.raises(ValueError, match="attn_mode"): + make_config(attn_mode="banana") diff --git a/tests/policies/lingbot_va/test_factory.py b/tests/policies/lingbot_va/test_factory.py new file mode 100644 index 000000000..1aec34df2 --- /dev/null +++ b/tests/policies/lingbot_va/test_factory.py @@ -0,0 +1,38 @@ +#!/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 diff --git a/tests/policies/lingbot_va/test_modules.py b/tests/policies/lingbot_va/test_modules.py new file mode 100644 index 000000000..3ca6b2ba1 --- /dev/null +++ b/tests/policies/lingbot_va/test_modules.py @@ -0,0 +1,128 @@ +#!/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. + +"""Unit tests for the vendored LingBot-VA helper code (scheduler + grid utilities).""" + +from __future__ import annotations + +import pytest +import torch + +pytest.importorskip("diffusers") # the model code lives in modeling_lingbot_va, which imports diffusers + +from lerobot.policies.lingbot_va.modeling_lingbot_va import FlowMatchScheduler +from lerobot.policies.lingbot_va.utils import data_seq_to_patch, get_mesh_id + + +def test_flow_match_scheduler_timesteps_monotone_decreasing() -> None: + sch = FlowMatchScheduler(shift=5.0, sigma_min=0.0, extra_one_step=True) + sch.set_timesteps(20) + assert sch.timesteps.shape == (20,) + diffs = sch.timesteps[1:] - sch.timesteps[:-1] + assert torch.all(diffs <= 0) # decreasing + + +def test_flow_match_scheduler_step_preserves_shape() -> None: + sch = FlowMatchScheduler(shift=5.0, sigma_min=0.0, extra_one_step=True) + sch.set_timesteps(20) + sample = torch.zeros(1, 48, 4, 8, 16) + out = sch.step(torch.ones_like(sample), sch.timesteps[0], sample) + assert out.shape == sample.shape + + +def test_flow_match_scheduler_add_noise() -> None: + sch = FlowMatchScheduler(shift=5.0, sigma_min=0.0, extra_one_step=True) + sch.set_timesteps(20) + sample = torch.randn(1, 48, 4, 8, 16) + noise = torch.randn_like(sample) + noisy = sch.add_noise(sample, noise, sch.timesteps[:4], t_dim=2) + assert noisy.shape == sample.shape + + +def test_get_mesh_id_latent_shape() -> None: + grid = get_mesh_id(4, 8, 16, 0, 1, 0) + assert grid.shape == (4, 4 * 8 * 16) # (f, h, w, stream) x tokens + + +def test_get_mesh_id_action_shape() -> None: + grid = get_mesh_id(4, 4, 1, 1, 1, 0, action=True) + assert grid.shape == (4, 4 * 4 * 1) + # Action rows for h/w are sentinel -1. + assert torch.all(grid[1] < 0) + assert torch.all(grid[2] < 0) + + +def test_data_seq_to_patch_roundtrip_shape() -> None: + b, f, h, w, c = 1, 4, 8, 16, 48 + seq = torch.arange(b * f * h * w * c, dtype=torch.float32).reshape(b, f * h * w, c) + out = data_seq_to_patch((1, 2, 2), seq, f, h, w, batch_size=b) + assert out.shape == (b, c, f, h, w) + + +def test_training_step_reduces_loss_tiny_flex() -> None: + """End-to-end single training step (flow-matching loss -> backward -> AdamW) on a tiny config. + + Exercises the flex-attention training path; requires a CUDA GPU with flex-attention support. + """ + if not torch.cuda.is_available(): + import pytest + + pytest.skip("training step test requires a CUDA GPU (flex-attention)") + + from lerobot.configs.types import FeatureType, PolicyFeature + from lerobot.policies.lingbot_va.configuration_lingbot_va import LingBotVAConfig + from lerobot.policies.lingbot_va.modeling_lingbot_va import LingBotVAPolicy + from lerobot.utils.constants import ACTION, OBS_IMAGES + + cfg = LingBotVAConfig( + attn_mode="flex", + dtype="bfloat16", + in_channels=16, + out_channels=16, + action_dim=8, + text_dim=32, + freq_dim=64, + ffn_dim=64, + num_attention_heads=2, + attention_head_dim=24, + num_layers=2, + frame_chunk_size=2, + action_per_frame=4, + used_action_channel_ids=[0, 1, 2, 3], + obs_cam_keys=[f"{OBS_IMAGES}.image"], + device="cuda", + ) + cfg.input_features = {f"{OBS_IMAGES}.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 64, 64))} + cfg.output_features = {ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(4,))} + cfg.validate_features() + + policy = LingBotVAPolicy(cfg).to("cuda") + policy.train() + opt = torch.optim.AdamW(policy.get_optim_params(), lr=1e-4) + + b, fc, apf = 1, cfg.frame_chunk_size, cfg.action_per_frame + latents = torch.randn(b, cfg.in_channels, fc, 4, 4, device="cuda", dtype=torch.bfloat16) + actions = torch.randn(b, cfg.action_dim, fc, apf, 1, device="cuda", dtype=torch.bfloat16) + amask = torch.zeros(cfg.action_dim, device="cuda") + amask[cfg.used_action_channel_ids] = 1.0 + actions_mask = amask.view(1, -1, 1, 1, 1).expand_as(actions) + text_emb = torch.randn(b, cfg.max_sequence_length, cfg.text_dim, device="cuda", dtype=torch.bfloat16) + + loss, metrics = policy.training_loss_from_streams(latents, actions, actions_mask, text_emb) + assert torch.isfinite(loss) and {"latent_loss", "action_loss"} <= set(metrics) + loss.backward() + assert any(p.grad is not None and torch.isfinite(p.grad).all() for p in policy.get_optim_params()) + opt.step() diff --git a/tests/policies/lingbot_va/test_processor.py b/tests/policies/lingbot_va/test_processor.py new file mode 100644 index 000000000..e79771600 --- /dev/null +++ b/tests/policies/lingbot_va/test_processor.py @@ -0,0 +1,88 @@ +#!/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 torch + +from lerobot.configs.types import FeatureType, PolicyFeature +from lerobot.policies.lingbot_va.configuration_lingbot_va import LingBotVAConfig +from lerobot.policies.lingbot_va.processor_lingbot_va import make_lingbot_va_pre_post_processors +from lerobot.processor import PolicyProcessorPipeline, UnnormalizerProcessorStep +from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action +from lerobot.utils.constants import ( + ACTION, + OBS_IMAGES, + POLICY_POSTPROCESSOR_DEFAULT_NAME, + POLICY_PREPROCESSOR_DEFAULT_NAME, +) + + +def _make_config() -> LingBotVAConfig: + cfg = LingBotVAConfig(device="cpu") + cfg.input_features = {f"{OBS_IMAGES}.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 128, 128))} + cfg.output_features = {} + cfg.validate_features() + return cfg + + +def test_make_pre_post_processors_names_and_steps() -> None: + cfg = _make_config() + pre, post = make_lingbot_va_pre_post_processors(cfg, dataset_stats=None) + assert pre.name == POLICY_PREPROCESSOR_DEFAULT_NAME + assert post.name == POLICY_POSTPROCESSOR_DEFAULT_NAME + # Actions are unnormalized by the standard built-in quantile unnormalizer. + assert any(isinstance(s, UnnormalizerProcessorStep) for s in post.steps) + + +def test_freshly_built_postprocessor_is_identity() -> None: + # Without action stats the quantile unnormalizer is a no-op (identity passthrough): the real + # per-benchmark q01/q99 are restored from the saved checkpoint on load, not hardcoded here. + cfg = _make_config() + _, post = make_lingbot_va_pre_post_processors(cfg, dataset_stats=None) + normed = torch.tensor([[0.3, -0.5, 1.0, -1.0, 0.0, 0.7, -0.2]]) + assert torch.allclose(post(normed), normed, atol=1e-6) + + +def test_postprocessor_quantile_unnormalization() -> None: + # QUANTILES unnormalize maps [-1, 1] -> [q01, q99]: -1 -> q01, +1 -> q99. + cfg = _make_config() + q01 = [-1.0, -0.5, 0.0, -1.0, -1.0, -1.0, -1.0] + q99 = [1.0, 0.5, 2.0, 1.0, 1.0, 1.0, 1.0] + stats = {ACTION: {"q01": q01, "q99": q99}} + _, post = make_lingbot_va_pre_post_processors(cfg, dataset_stats=stats) + out_lo = post(torch.full((1, 7), -1.0)) + out_hi = post(torch.full((1, 7), 1.0)) + assert torch.allclose(out_lo, torch.tensor(q01).unsqueeze(0), atol=1e-4) + assert torch.allclose(out_hi, torch.tensor(q99).unsqueeze(0), atol=1e-4) + + +def test_postprocessor_stats_survive_save_load(tmp_path) -> None: + # Regression guard for the Hub mechanism: the q01/q99 stats live in the saved post-processor + # state and must round-trip through save_pretrained / from_pretrained. + cfg = _make_config() + q01 = [-0.6, -0.8, -0.9, -0.1, -0.15, -0.25, -1.0] + q99 = [0.9, 0.85, 0.9, 0.17, 0.18, 0.34, 1.0] + _, post = make_lingbot_va_pre_post_processors(cfg, dataset_stats={ACTION: {"q01": q01, "q99": q99}}) + post.save_pretrained(tmp_path) + loaded = PolicyProcessorPipeline.from_pretrained( + tmp_path, + config_filename=f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json", + to_transition=policy_action_to_transition, + to_output=transition_to_policy_action, + ) + out = loaded(torch.full((1, 7), -1.0)) + assert torch.allclose(out, torch.tensor(q01).unsqueeze(0), atol=1e-4) diff --git a/uv.lock b/uv.lock index 22afec6a6..260fb07f0 100644 --- a/uv.lock +++ b/uv.lock @@ -3057,6 +3057,11 @@ libero = [ { name = "torchcodec", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'AMD64' and sys_platform == 'linux') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'arm64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'win32'" }, { name = "transformers" }, ] +lingbot-va = [ + { name = "accelerate" }, + { name = "diffusers" }, + { name = "transformers" }, +] matplotlib-dep = [ { name = "contourpy" }, { name = "matplotlib" }, @@ -3265,6 +3270,7 @@ requires-dist = [ { name = "ipykernel", marker = "extra == 'notebook'", specifier = ">=6.0.0,<7.0.0" }, { name = "jsonlines", marker = "extra == 'dataset'", specifier = ">=4.0.0,<5.0.0" }, { name = "jupyter", marker = "extra == 'notebook'", specifier = ">=1.0.0,<2.0.0" }, + { name = "lerobot", extras = ["accelerate-dep"], marker = "extra == 'lingbot-va'" }, { name = "lerobot", extras = ["accelerate-dep"], marker = "extra == 'smolvla'" }, { name = "lerobot", extras = ["accelerate-dep"], marker = "extra == 'training'" }, { name = "lerobot", extras = ["aloha"], marker = "extra == 'all'" }, @@ -3292,6 +3298,7 @@ requires-dist = [ { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'diffusion'" }, { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'fastwam'" }, { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'groot'" }, + { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'lingbot-va'" }, { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'multi-task-dit'" }, { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'vla-jepa'" }, { name = "lerobot", extras = ["diffusion"], marker = "extra == 'all'" }, @@ -3312,6 +3319,7 @@ requires-dist = [ { name = "lerobot", extras = ["kinematics"], marker = "extra == 'all'" }, { name = "lerobot", extras = ["lekiwi"], marker = "extra == 'all'" }, { name = "lerobot", extras = ["libero"], marker = "sys_platform == 'linux' and extra == 'all'" }, + { name = "lerobot", extras = ["lingbot-va"], marker = "extra == 'all'" }, { name = "lerobot", extras = ["matplotlib-dep"], marker = "extra == 'async'" }, { name = "lerobot", extras = ["matplotlib-dep"], marker = "extra == 'sarm'" }, { name = "lerobot", extras = ["matplotlib-dep"], marker = "extra == 'unitree-g1'" }, @@ -3370,6 +3378,7 @@ requires-dist = [ { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'groot'" }, { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'hilserl'" }, { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'libero'" }, + { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'lingbot-va'" }, { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'molmoact2'" }, { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'multi-task-dit'" }, { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'peft'" }, @@ -3449,7 +3458,7 @@ requires-dist = [ { name = "transformers", marker = "extra == 'transformers-dep'", specifier = ">=5.4.0,<5.6.0" }, { name = "wandb", marker = "extra == 'training'", specifier = ">=0.24.0,<0.28.0" }, ] -provides-extras = ["dataset", "training", "hardware", "viz", "core-scripts", "evaluation", "dataset-viz", "av-dep", "pygame-dep", "placo-dep", "transformers-dep", "grpcio-dep", "accelerate-dep", "can-dep", "peft-dep", "scipy-dep", "diffusers-dep", "qwen-vl-utils-dep", "matplotlib-dep", "pyserial-dep", "deepdiff-dep", "pynput-dep", "pyzmq-dep", "motorbridge-dep", "motorbridge-smart-servo-dep", "feetech", "dynamixel", "damiao", "robstride", "openarms", "gamepad", "hopejr", "lekiwi", "unitree-g1", "reachy2", "rebot", "kinematics", "intelrealsense", "phone", "diffusion", "wallx", "pi", "molmoact2", "smolvla", "multi-task-dit", "groot", "sarm", "robometer", "topreward", "xvla", "eo1", "fastwam", "hilserl", "vla-jepa", "async", "peft", "annotations", "dev", "notebook", "test", "video-benchmark", "aloha", "pusht", "libero", "metaworld", "all"] +provides-extras = ["dataset", "training", "hardware", "viz", "core-scripts", "evaluation", "dataset-viz", "av-dep", "pygame-dep", "placo-dep", "transformers-dep", "grpcio-dep", "accelerate-dep", "can-dep", "peft-dep", "scipy-dep", "diffusers-dep", "qwen-vl-utils-dep", "matplotlib-dep", "pyserial-dep", "deepdiff-dep", "pynput-dep", "pyzmq-dep", "motorbridge-dep", "motorbridge-smart-servo-dep", "feetech", "dynamixel", "damiao", "robstride", "openarms", "gamepad", "hopejr", "lekiwi", "unitree-g1", "reachy2", "rebot", "kinematics", "intelrealsense", "phone", "diffusion", "wallx", "pi", "molmoact2", "smolvla", "multi-task-dit", "groot", "sarm", "robometer", "topreward", "xvla", "eo1", "fastwam", "hilserl", "vla-jepa", "lingbot-va", "async", "peft", "annotations", "dev", "notebook", "test", "video-benchmark", "aloha", "pusht", "libero", "metaworld", "all"] [[package]] name = "librt" From 708fa1d18952fc0b0c6810926f39347c7bd721df Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Fri, 3 Jul 2026 21:15:09 +0200 Subject: [PATCH 02/11] feat(policies): add Gr00t N1.7 policy (#3922) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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 * 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 1322f45aec088d3ca346640d995d28edcf71d00f. * 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 Co-authored-by: Ryan Halabi Co-authored-by: nv-sachdevkartik Co-authored-by: groot-validation Co-authored-by: johnnynunez Co-authored-by: lbenhorin --- README.md | 2 +- docs/source/_toctree.yml | 2 +- docs/source/envhub_isaaclab_arena.mdx | 2 +- docs/source/groot.mdx | 227 +- docs/source/policy_groot_README.md | 115 +- pyproject.toml | 5 +- src/lerobot/policies/factory.py | 35 +- .../groot/action_head/action_encoder.py | 54 - .../groot/action_head/cross_attention_dit.py | 87 +- .../action_head/flow_matching_action_head.py | 408 --- .../policies/groot/configuration_groot.py | 416 ++- .../configuration_eagle2_5_vl.py | 135 - .../image_processing_eagle2_5_vl_fast.py | 503 --- .../eagle2_hg_model/modeling_eagle2_5_vl.py | 396 --- .../eagle2_hg_model/processing_eagle2_5_vl.py | 541 --- src/lerobot/policies/groot/groot_n1.py | 380 --- src/lerobot/policies/groot/groot_n1_7.py | 951 ++++++ src/lerobot/policies/groot/modeling_groot.py | 410 ++- src/lerobot/policies/groot/processor_groot.py | 2572 ++++++++++++-- src/lerobot/policies/groot/utils.py | 293 +- src/lerobot/utils/import_utils.py | 1 + tests/policies/groot/test_groot_lerobot.py | 33 +- tests/policies/groot/test_groot_n1_7.py | 3025 +++++++++++++++++ .../groot/test_groot_n1_7_oss_parity.py | 259 ++ .../groot/test_groot_state_dropout.py | 100 + .../groot/test_groot_train_random_crop.py | 169 + .../test_groot_training_optim_contract.py | 125 + .../policies/groot/test_groot_vs_original.py | 557 +-- .../groot/utils/dump_original_n1_7.py | 212 ++ uv.lock | 52 +- 30 files changed, 8593 insertions(+), 3474 deletions(-) delete mode 100644 src/lerobot/policies/groot/action_head/action_encoder.py delete mode 100644 src/lerobot/policies/groot/action_head/flow_matching_action_head.py delete mode 100755 src/lerobot/policies/groot/eagle2_hg_model/configuration_eagle2_5_vl.py delete mode 100644 src/lerobot/policies/groot/eagle2_hg_model/image_processing_eagle2_5_vl_fast.py delete mode 100755 src/lerobot/policies/groot/eagle2_hg_model/modeling_eagle2_5_vl.py delete mode 100755 src/lerobot/policies/groot/eagle2_hg_model/processing_eagle2_5_vl.py delete mode 100644 src/lerobot/policies/groot/groot_n1.py create mode 100644 src/lerobot/policies/groot/groot_n1_7.py create mode 100644 tests/policies/groot/test_groot_n1_7.py create mode 100644 tests/policies/groot/test_groot_n1_7_oss_parity.py create mode 100644 tests/policies/groot/test_groot_state_dropout.py create mode 100644 tests/policies/groot/test_groot_train_random_crop.py create mode 100644 tests/policies/groot/test_groot_training_optim_contract.py create mode 100644 tests/policies/groot/utils/dump_original_n1_7.py diff --git a/README.md b/README.md index f72952102..acfad05fb 100644 --- a/README.md +++ b/README.md @@ -105,7 +105,7 @@ lerobot-train \ | -------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | **Imitation Learning** | [ACT](./docs/source/policy_act_README.md), [Diffusion](./docs/source/policy_diffusion_README.md), [VQ-BeT](./docs/source/policy_vqbet_README.md), [Multitask DiT Policy](./docs/source/policy_multi_task_dit_README.md) | | **Reinforcement Learning** | [HIL-SERL](./docs/source/hilserl.mdx), [TDMPC](./docs/source/policy_tdmpc_README.md) & QC-FQL (coming soon) | -| **VLAs Models** | [Pi0](./docs/source/pi0.mdx), [Pi0Fast](./docs/source/pi0fast.mdx), [Pi0.5](./docs/source/pi05.mdx), [GR00T N1.5](./docs/source/policy_groot_README.md), [SmolVLA](./docs/source/policy_smolvla_README.md), [XVLA](./docs/source/xvla.mdx), [EO-1](./docs/source/eo1.mdx), [MolmoAct2](./docs/source/molmoact2.mdx), [WALL-OSS](./docs/source/walloss.mdx) | +| **VLAs Models** | [Pi0](./docs/source/pi0.mdx), [Pi0Fast](./docs/source/pi0fast.mdx), [Pi0.5](./docs/source/pi05.mdx), [GR00T N1.7](./docs/source/policy_groot_README.md), [SmolVLA](./docs/source/policy_smolvla_README.md), [XVLA](./docs/source/xvla.mdx), [EO-1](./docs/source/eo1.mdx), [MolmoAct2](./docs/source/molmoact2.mdx), [WALL-OSS](./docs/source/walloss.mdx) | | **World Models** | [VLA-JEPA](./docs/source/vla_jepa.mdx) (more coming soon) | | **Reward Models** | [SARM](./docs/source/sarm.mdx), [TOPReward](./docs/source/topreward.mdx), [Robometer](./docs/source/robometer.mdx) | diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml index 79f4bc124..d4bdf892e 100644 --- a/docs/source/_toctree.yml +++ b/docs/source/_toctree.yml @@ -74,7 +74,7 @@ - local: fastwam title: FastWAM - local: groot - title: NVIDIA GR00T N1.5 + title: NVIDIA GR00T - local: xvla title: X-VLA - local: multi_task_dit diff --git a/docs/source/envhub_isaaclab_arena.mdx b/docs/source/envhub_isaaclab_arena.mdx index b934240d6..cd077806d 100644 --- a/docs/source/envhub_isaaclab_arena.mdx +++ b/docs/source/envhub_isaaclab_arena.mdx @@ -193,7 +193,7 @@ To learn more about training policies with LeRobot, please refer to the training - [SmolVLA](./smolvla) - [Pi0.5](./pi05) -- [GR00T N1.5](./groot) +- [GR00T N1.7](./groot) Sample IsaacLab Arena datasets are available on HuggingFace Hub for experimentation: diff --git a/docs/source/groot.mdx b/docs/source/groot.mdx index 3ab202fb2..af6bfe1ab 100644 --- a/docs/source/groot.mdx +++ b/docs/source/groot.mdx @@ -1,16 +1,19 @@ -# GR00T N1.5 Policy +# GR00T Policy -GR00T N1.5 is an open foundation model from NVIDIA designed for generalized humanoid robot reasoning and skills. It is a cross-embodiment model that accepts multimodal input, including language and images, to perform manipulation tasks in diverse environments. +GR00T is an NVIDIA foundation model family for generalized humanoid robot reasoning and skills. It is a cross-embodiment policy that accepts multimodal input, including language, images, and proprioception, to perform manipulation tasks in diverse environments. -This document outlines the specifics of its integration and usage within the LeRobot framework. +LeRobot integrates GR00T N1.7 through the `groot` policy type. + +> [!WARNING] +> **Breaking change:** GR00T N1.5 support was removed from LeRobot, and current releases support GR00T N1.7 only. N1.5 checkpoints and configs are rejected with a migration note. To keep using an N1.5 checkpoint, pin the last release that supports it: `pip install 'lerobot==0.5.1'`. To use the current release, migrate to GR00T N1.7 (base model [`nvidia/GR00T-N1.7-3B`](https://huggingface.co/nvidia/GR00T-N1.7-3B)). ## Model Overview -NVIDIA Isaac GR00T N1.5 is an upgraded version of the GR00T N1 foundation model. It is built to improve generalization and language-following abilities for humanoid robots. +GR00T N1.7 uses a Cosmos-Reason2/Qwen3-VL backbone and provides checkpoints for SimplerEnv, DROID, and LIBERO. -Developers and researchers can post-train GR00T N1.5 with their own real or synthetic data to adapt it for specific humanoid robots or tasks. +Developers and researchers can post-train GR00T with their own real or synthetic data to adapt it for specific humanoid robots or tasks. -GR00T N1.5 (specifically the GR00T-N1.5-3B model) is built using pre-trained vision and language encoders. It utilizes a flow matching action transformer to model a chunk of actions, conditioned on vision, language, and proprioception. +GR00T uses pre-trained vision and language encoders with a flow matching action transformer to model a chunk of actions conditioned on vision, language, and proprioception. =2.2.1,<2.8.0" "torchvision>=0.21.0,<0.23.0" # --index-url https://download.pytorch.org/whl/cu1XX -pip install ninja "packaging>=24.2,<26.0" # flash attention dependencies -pip install "flash-attn>=2.5.9,<3.0.0" --no-build-isolation -python -c "import flash_attn; print(f'Flash Attention {flash_attn.__version__} imported successfully')" +pip install "lerobot[groot]" ``` -3. Install LeRobot by running: +For a source checkout: ```bash -pip install lerobot[groot] +pip install -e ".[groot]" ``` ## Usage -To use GR00T in your LeRobot configuration, specify the policy type as: +To use GR00T N1.7: -```python -policy.type=groot +```bash +--policy.type=groot ``` ## Training @@ -63,72 +57,171 @@ policy.type=groot Here's a complete training command for finetuning the base GR00T model on your own dataset: +This command is using the `new_embodiment` flag, which is used for the SO-101 robot, [read more about how GR00T handles different embodiments.](https://github.com/NVIDIA/Isaac-GR00T/blob/main/getting_started/policy.md#--embodiment-tag). + ```bash -# Using a multi-GPU setup -accelerate launch \ - --multi_gpu \ - --num_processes=$NUM_GPUS \ - $(which lerobot-train) \ - --output_dir=$OUTPUT_DIR \ - --save_checkpoint=true \ - --batch_size=$BATCH_SIZE \ - --steps=$NUM_STEPS \ - --save_freq=$SAVE_FREQ \ - --log_freq=$LOG_FREQ \ - --policy.push_to_hub=true \ +# install extra deps for training +pip install "lerobot[training]" + +hf auth login +wandb login + +export DATASET_NAME=your_data_set +export HF_USER=your_hf_username +export DATASET=$HF_USER/$DATASET_NAME +export REPO_ID="${DATASET}_GR00T17" #this is the model that will be uploaded to huggingface +export OUTPUT_DIR=outputs/train/$REPO_ID + +lerobot-train \ + --dataset.repo_id=$DATASET \ + --dataset.image_transforms.enable=true \ --policy.type=groot \ + --policy.device=cuda \ + --policy.base_model_path=nvidia/GR00T-N1.7-3B \ + --policy.embodiment_tag=new_embodiment \ + --policy.chunk_size=16 \ + --policy.n_action_steps=16 \ + --policy.use_relative_actions=true \ + --policy.relative_exclude_joints='["gripper"]' \ + --policy.use_bf16=true \ + --policy.push_to_hub=true \ --policy.repo_id=$REPO_ID \ - --policy.tune_diffusion_model=false \ - --dataset.repo_id=$DATASET_ID \ + --seed=42 \ + --batch_size=64 \ + --steps=20000 \ + --save_checkpoint=true \ + --save_freq=5000 \ + --use_policy_training_preset=true \ + --env_eval_freq=0 \ + --eval_steps=0 \ + --log_freq=10 \ + --output_dir=$OUTPUT_DIR \ + --job_name=$DATASET \ --wandb.enable=true \ - --wandb.disable_artifact=true \ - --job_name=$JOB_NAME + --wandb.disable_artifact=true + ``` ## Performance Results -### Libero Benchmark Results +### LIBERO Benchmark Results > [!NOTE] -> Follow our instructions for Libero usage: [Libero](./libero) +> Follow the [LIBERO](./libero) setup instructions before running `lerobot-eval`. -GR00T has demonstrated strong performance on the Libero benchmark suite. To compare and test its LeRobot implementation, we finetuned the GR00T N1.5 model for 30k steps on the Libero dataset and compared the results to the GR00T reference results. +GR00T N1.7 has demonstrated strong performance on the LIBERO benchmark suite. To reproduce LeRobot results, follow the instructions in the [LIBERO](./libero) section. -| Benchmark | LeRobot Implementation | GR00T Reference | -| ------------------ | ---------------------- | --------------- | -| **Libero Spatial** | 82.0% | 92.0% | -| **Libero Object** | 99.0% | 92.0% | -| **Libero Long** | 82.0% | 76.0% | -| **Average** | 87.0% | 87.0% | +### Train on LIBERO -These results demonstrate GR00T's strong generalization capabilities across diverse robotic manipulation tasks. To reproduce these results, you can follow the instructions in the [Libero](https://huggingface.co/docs/lerobot/libero) section. +Example training command for a LIBERO suite (here `libero_spatial`): + +```bash +IMAGE_TRANSFORMS='{ + "brightness": {"weight": 1.0, "type": "ColorJitter", "kwargs": {"brightness": [0.7, 1.3]}}, + "contrast": {"weight": 1.0, "type": "ColorJitter", "kwargs": {"contrast": [0.6, 1.4]}}, + "saturation": {"weight": 1.0, "type": "ColorJitter", "kwargs": {"saturation": [0.5, 1.5]}}, + "hue": {"weight": 1.0, "type": "ColorJitter", "kwargs": {"hue": [-0.08, 0.08]}} +}' + +lerobot-train \ + --dataset.repo_id=IPEC-COMMUNITY/libero_spatial_no_noops_1.0.0_lerobot \ + --dataset.root=/datasets/libero_spatial \ + --dataset.revision=main \ + --dataset.video_backend=pyav \ + --dataset.image_transforms.enable=true \ + --dataset.image_transforms.max_num_transforms=4 \ + --dataset.image_transforms.tfs="$IMAGE_TRANSFORMS" \ + --policy.type=groot \ + --policy.base_model_path=nvidia/GR00T-N1.7-3B \ + --policy.embodiment_tag=libero_sim \ + --policy.push_to_hub=false \ + --policy.use_relative_actions=false \ + --policy.max_steps=20000 \ + --batch_size=320 \ + --steps=20000 \ + --save_freq=2000 \ + --env_eval_freq=0 \ + --eval_steps=0 \ + --log_freq=10 \ + --wandb.enable=true \ + --wandb.project=lerobot \ + --wandb.mode=online \ + --wandb.disable_artifact=true \ + --num_workers=4 \ + --prefetch_factor=2 \ + --persistent_workers=true \ + --output_dir=$OUTPUT_DIR \ + --job_name=$JOB_NAME +``` + +This will follow the recipe found [here](https://github.com/NVIDIA/Isaac-GR00T/blob/main/examples/LIBERO/README.md). + +### GR00T N1.7 LIBERO Results + +Preliminary LeRobot integration results (GR00T-LeRobot, `eval.n_episodes >= 50` per suite): + +| Suite | Success rate | +| ---------------- | -----------: | +| LIBERO Spatial | 94% | +| LIBERO Object | 98% | +| LIBERO Goal | 93% | +| LIBERO 10 (Long) | 90% | +| **Average** | **93.75%** | + +```bash +export MODEL_ID=your_trained_model_on_huggingface + +lerobot-eval \ + --policy.type=groot \ + --policy.base_model_path=$MODEL_ID \ + --policy.embodiment_tag=libero_sim \ + --env.type=libero \ + --env.task=libero_spatial \ + --eval.n_episodes=50 +``` + +Use `eval.n_episodes >= 50` per suite when reporting success rates. ### Evaluate in your hardware setup Once you have trained your model using your parameters you can run inference in your downstream task. Follow the instructions in [Policy Deployment (lerobot-rollout)](./inference). For example: ```bash -lerobot-rollout\ - --strategy.type=sentry \ - --strategy.upload_every_n_episodes=5 \ - --robot.type=bi_so_follower \ - --robot.left_arm_port=/dev/ttyACM1 \ - --robot.right_arm_port=/dev/ttyACM0 \ - --robot.id=bimanual_follower \ - --robot.cameras='{ right: {"type": "opencv", "index_or_path": 0, "width": 640, "height": 480, "fps": 30}, - left: {"type": "opencv", "index_or_path": 2, "width": 640, "height": 480, "fps": 30}, - top: {"type": "opencv", "index_or_path": 4, "width": 640, "height": 480, "fps": 30}, - }' \ +# install extra deps for roullout and real hardware +pip install "lerobot[feetech,viz]" + +export MODEL_ID=your_trained_model_on_huggingface + +# make sure that camera index matches your setup! +# find index using `uv run lerobot-find-cameras opencv` +WRIST_CAM='wrist: {type: opencv, index_or_path: 2, width: 640, height: 480, fps: 30, fourcc: "MJPG"}' +FRONT_CAM='front: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30, fourcc: "MJPG"}' +export ROBOT_CAMERAS="{ $WRIST_CAM, $FRONT_CAM }" +export ROBOT_ID=follower_robot +export ROBOT_PORT=/dev/ttyACM0 + +uv run lerobot-rollout \ + --strategy.type=base \ + --policy.path=$MODEL_ID \ + --policy.base_model_path=nvidia/GR00T-N1.7-3B \ + --policy.n_action_steps=8 \ + --robot.type=so101_follower \ + --robot.port=$ROBOT_PORT \ + --robot.id=$ROBOT_ID \ + --robot.cameras="$ROBOT_CAMERAS" \ + --task="place the vial in the rack" \ + --duration=60 \ + --device=cuda \ --display_data=true \ - --dataset.repo_id=/eval_groot-bimanual \ - --dataset.single_task="Grab and handover the red cube to the other arm" \ - --dataset.streaming_encoding=true \ - --dataset.encoder_threads=2 \ - # --dataset.rgb_encoder.vcodec=auto \ - --policy.path=/groot-bimanual \ # your trained model - --duration=600 + --inference.type=rtc \ + --inference.rtc.enabled=True \ # set to False if it causes inference instability + --inference.rtc.execution_horizon=8 \ + --inference.queue_threshold=0 ``` +> [!NOTE] +> Value of `inference.queue_threshold` should not exceed 5 to ensure stable inference. + ## License -This model follows NVIDIA's proprietary license, consistent with the original [GR00T repository](https://github.com/NVIDIA/Isaac-GR00T). Future versions (starting from N1.7) will follow **Apache 2.0 License**. +GR00T N1.7 is released under the [NVIDIA Open Model License Agreement](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license/). diff --git a/docs/source/policy_groot_README.md b/docs/source/policy_groot_README.md index efcd76ebe..4b256fb27 100644 --- a/docs/source/policy_groot_README.md +++ b/docs/source/policy_groot_README.md @@ -1,6 +1,13 @@ ## Research Paper -Paper: https://research.nvidia.com/labs/gear/gr00t-n1_5/ +GR00T N1 technical report (covers the GR00T N1.x family, including N1.7): https://arxiv.org/abs/2503.14734 + +GR00T N1.7 model card: https://huggingface.co/nvidia/GR00T-N1.7-3B + +GR00T N1.5 research page (earlier version): https://research.nvidia.com/labs/gear/gr00t-n1_5/ + +> GR00T N1.5 support was removed from LeRobot; the last release supporting it is `lerobot==0.5.1`. +> Current releases support GR00T N1.7 only. ## Repository @@ -24,4 +31,108 @@ Code: https://github.com/NVIDIA/Isaac-GR00T Blog: https://developer.nvidia.com/isaac/gr00t -Hugging Face Model: https://huggingface.co/nvidia/GR00T-N1.5-3B +Hugging Face Models: + +- GR00T N1.7: https://huggingface.co/nvidia/GR00T-N1.7-3B +- GR00T N1.7 LIBERO checkpoints: https://huggingface.co/nvidia/GR00T-N1.7-LIBERO + +
+Original-vs-LeRobot parity test + +## Original-vs-LeRobot parity test + +`tests/policies/groot/test_groot_vs_original.py` verifies this LeRobot +reimplementation of GR00T N1.7 (Qwen3-VL backbone + flow-matching action head) +against NVIDIA's original `gr00t` package with two comparisons, each parametrized +over every embodiment tag present in the checkpoint: + +1. **Model parity** — given byte-identical pre-processed inputs and the same + flow-matching seed (recorded in each artifact), both implementations must produce + the **same raw model output** (`get_action(...)["action_pred"]`, the normalized + flow-matching prediction). Output shapes must match exactly; any action-horizon + or action-dim mismatch fails the test. +2. **Preprocessor parity** — given the identical raw observations (per-camera + frames, state vectors, language instruction), LeRobot's own preprocessor pipeline + (real Qwen3-VL chat template / tokenizer / image packing + checkpoint-driven + state normalization, no mocks) must produce the **same collated model inputs** + (`input_ids`, `attention_mask`, `pixel_values`, `image_grid_thw`, `state`, + `embodiment_id`) as the original package's processor. + +### Why two environments + +The original `gr00t` package pins `transformers==4.57.3` (Python 3.10); this +integration requires `transformers>=5.x` (Qwen3-VL). Under 5.x, `PretrainedConfig` +is itself a defaulted dataclass, so the original config dataclasses fail to import +(`non-default argument follows default argument`). The two implementations therefore +**cannot be imported in the same Python process**. + +So the test uses a **producer / consumer** split across two venvs: + +1. **Producer** — `tests/policies/groot/utils/dump_original_n1_7.py`, run in the _original_ + gr00t venv. For each embodiment it builds dummy inputs generically from the + checkpoint metadata (state dims from `statistics.json`; camera/language keys from + the processor modality configs), runs the original model, and saves to one `.npz` + per tag: the raw observations (`raw::` keys), the exact collated inputs + (`in::` keys), the seed, and the raw `action_pred`. +2. **Consumer** — the pytest above, run in the _LeRobot_ venv. It discovers every + `.npz`; the model-parity case replays the byte-identical collated inputs through + the LeRobot model with the recorded seed and asserts the outputs match, and the + preprocessor-parity case replays the raw observations through LeRobot's full + preprocessor pipeline and asserts the collated tensors match. + +> Artifacts generated by older versions of the dump script contain no `raw::` +> fields; the preprocessor-parity case then **skips** with a regeneration hint. +> Re-run the producer to refresh them. + +### Fairness controls + +- **Same pre-processed inputs (model parity)** — the original processor's `input_ids`, + `pixel_values`, `image_grid_thw`, `attention_mask`, `state`, `embodiment_id` are + fed verbatim to the LeRobot model (no re-tokenization / re-normalization), so the + model comparison isolates the model. LeRobot's own tokenization / image packing is + covered separately by the preprocessor-parity case, which compares its output + against those same collated tensors from identical raw observations. +- **Same precision + attention kernel** — both sides run **fp32 + SDPA**. The + original defaults to `use_flash_attention=True` (flash_attention_2 + bf16); the + producer forces SDPA + fp32. (With the defaults the gap is ~3e-2 — pure + kernel/rounding noise, not an implementation difference.) +- **Same flow-matching seed** — fixed right before sampling on both sides; the + producer records it in each artifact (`--seed`, default 42) and the consumer + replays the recorded value. + +### How to run + +```bash +# Resolve a local checkpoint (GR00T-N1.7-LIBERO / libero_10) +CKPT=$(python - <<'PY' +import os +from huggingface_hub import snapshot_download +print(os.path.join(snapshot_download("nvidia/GR00T-N1.7-LIBERO", + allow_patterns=["libero_10/*"]), "libero_10")) +PY +) + +# 1) Produce the original-side artifacts for all embodiments (original gr00t venv, CUDA) +CUDA_VISIBLE_DEVICES=0 /path/to/Isaac-GR00T/.venv-original/bin/python \ + tests/policies/groot/utils/dump_original_n1_7.py \ + --ckpt "$CKPT" --out-dir tests/policies/groot/artifacts --device cuda --seed 42 + +# 2) Run the parity test (LeRobot venv) — one parametrized case per embodiment +CUDA_VISIBLE_DEVICES=0 GROOT_PARITY_DEVICE=cuda \ + uv run pytest tests/policies/groot/test_groot_vs_original.py -v -s +``` + +The `.npz` artifacts are local-only (gitignored, ~6–10 MB each) and are regenerated by +the producer; they are never committed. The tests **skip** (do not fail) on CI or +when the checkpoint / artifacts are absent. + +#### Env knobs (all optional) + +| Var | Default | Purpose | +| ----------------------------------------- | -------------------------------- | ------------------------------------- | +| `GROOT_N1_7_PARITY_DIR` | `tests/policies/groot/artifacts` | directory of per-tag `.npz` artifacts | +| `GROOT_N1_7_LIBERO_CKPT` | auto (HF cache) | override checkpoint dir | +| `GROOT_PARITY_DEVICE` | `cuda` if available | `cpu` or `cuda` | +| `GROOT_PARITY_ATOL` / `GROOT_PARITY_RTOL` | `1e-3` | comparison tolerance | + +
diff --git a/pyproject.toml b/pyproject.toml index 84fe07328..882dd0b6f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -219,11 +219,10 @@ groot = [ "lerobot[transformers-dep]", "lerobot[peft-dep]", "lerobot[diffusers-dep]", + "lerobot[dataset]", # NOTE: processor_groot builds a LeRobotDataset for relative-action training stats "dm-tree>=0.1.8,<1.0.0", "timm>=1.0.0,<1.1.0", "decord>=0.6.0,<1.0.0; (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "ninja>=1.11.1,<2.0.0", - "flash-attn>=2.5.9,<3.0.0 ; sys_platform != 'darwin'" ] sarm = ["lerobot[transformers-dep]", "pydantic>=2.0.0,<3.0.0", "faker>=33.0.0,<35.0.0", "lerobot[matplotlib-dep]", "lerobot[qwen-vl-utils-dep]"] robometer = ["lerobot[transformers-dep]", "lerobot[qwen-vl-utils-dep]", "lerobot[peft-dep]"] @@ -315,7 +314,7 @@ all = [ "lerobot[molmoact2]", "lerobot[smolvla]", "lerobot[fastwam]", - # "lerobot[groot]", TODO(Steven): Gr00t requires specific installation instructions for flash-attn + "lerobot[groot]", "lerobot[xvla]", "lerobot[hilserl]", "lerobot[vla_jepa]", diff --git a/src/lerobot/policies/factory.py b/src/lerobot/policies/factory.py index 12871218e..483804a57 100644 --- a/src/lerobot/policies/factory.py +++ b/src/lerobot/policies/factory.py @@ -295,26 +295,23 @@ def make_pre_post_processors( policy configuration type. """ if pretrained_path: - # TODO(Steven): Temporary patch, implement correctly the processors for Gr00t if isinstance(policy_cfg, GrootConfig): - # GROOT handles normalization in groot_pack_inputs_v3 step - # Need to override both stats AND normalize_min_max since saved config might be empty - preprocessor_overrides = {} - postprocessor_overrides = {} - preprocessor_overrides["groot_pack_inputs_v3"] = { - "stats": kwargs.get("dataset_stats"), - "normalize_min_max": True, - } + from .groot.processor_groot import make_groot_pre_post_processors_from_pretrained - # Also ensure postprocessing slices to env action dim and unnormalizes with dataset stats - env_action_dim = policy_cfg.output_features[ACTION].shape[0] - postprocessor_overrides["groot_action_unpack_unnormalize_v1"] = { - "stats": kwargs.get("dataset_stats"), - "normalize_min_max": True, - "env_action_dim": env_action_dim, - } - kwargs["preprocessor_overrides"] = preprocessor_overrides - kwargs["postprocessor_overrides"] = postprocessor_overrides + return make_groot_pre_post_processors_from_pretrained( + config=policy_cfg, + pretrained_path=pretrained_path, + dataset_stats=kwargs.get("dataset_stats"), + dataset_meta=kwargs.get("dataset_meta"), + preprocessor_overrides=kwargs.get("preprocessor_overrides"), + postprocessor_overrides=kwargs.get("postprocessor_overrides"), + preprocessor_config_filename=kwargs.get( + "preprocessor_config_filename", f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json" + ), + postprocessor_config_filename=kwargs.get( + "postprocessor_config_filename", f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json" + ), + ) preprocessor = PolicyProcessorPipeline.from_pretrained( pretrained_model_name_or_path=pretrained_path, @@ -420,6 +417,7 @@ def make_pre_post_processors( processors = make_groot_pre_post_processors( config=policy_cfg, dataset_stats=kwargs.get("dataset_stats"), + dataset_meta=kwargs.get("dataset_meta"), ) elif isinstance(policy_cfg, XVLAConfig): @@ -570,6 +568,7 @@ def make_policy( set_dataset_feature_metadata = getattr(cfg, "set_dataset_feature_metadata", None) if callable(set_dataset_feature_metadata): set_dataset_feature_metadata(ds_meta.features) + cfg._runtime_dataset_meta = ds_meta kwargs["config"] = cfg diff --git a/src/lerobot/policies/groot/action_head/action_encoder.py b/src/lerobot/policies/groot/action_head/action_encoder.py deleted file mode 100644 index c6fa0a779..000000000 --- a/src/lerobot/policies/groot/action_head/action_encoder.py +++ /dev/null @@ -1,54 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# 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. - -import torch -import torch.nn as nn - - -def swish(x): - return x * torch.sigmoid(x) - - -class SinusoidalPositionalEncoding(nn.Module): - """ - Produces a sinusoidal encoding of shape (B, T, w) - given timesteps of shape (B, T). - """ - - def __init__(self, embedding_dim): - super().__init__() - self.embedding_dim = embedding_dim - - def forward(self, timesteps): - # timesteps: shape (B, T) - # We'll compute sin/cos frequencies across dim T - timesteps = timesteps.float() # ensure float - - b, t = timesteps.shape - device = timesteps.device - - half_dim = self.embedding_dim // 2 - # typical log space frequencies for sinusoidal encoding - exponent = -torch.arange(half_dim, dtype=torch.float, device=device) * ( - torch.log(torch.tensor(10000.0)) / half_dim - ) - # Expand timesteps to (B, T, 1) then multiply - freqs = timesteps.unsqueeze(-1) * exponent.exp() # (B, T, half_dim) - - sin = torch.sin(freqs) - cos = torch.cos(freqs) - enc = torch.cat([sin, cos], dim=-1) # (B, T, w) - - return enc diff --git a/src/lerobot/policies/groot/action_head/cross_attention_dit.py b/src/lerobot/policies/groot/action_head/cross_attention_dit.py index a4cd1a0b7..697459240 100755 --- a/src/lerobot/policies/groot/action_head/cross_attention_dit.py +++ b/src/lerobot/policies/groot/action_head/cross_attention_dit.py @@ -1,11 +1,12 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 +#!/usr/bin/env python + +# Copyright 2025 NVIDIA Corporation and 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 +# 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, @@ -14,6 +15,7 @@ # limitations under the License. +import logging from typing import TYPE_CHECKING import torch @@ -42,6 +44,9 @@ else: Timesteps = None +logger = logging.getLogger(__name__) + + class TimestepEncoder(nn.Module): def __init__(self, embedding_dim, compute_dtype=torch.float32): require_package("diffusers", extra="groot") @@ -181,8 +186,7 @@ class BasicTransformerBlock(nn.Module): attn_output = self.attn1( norm_hidden_states, encoder_hidden_states=encoder_hidden_states, - attention_mask=attention_mask, - # encoder_attention_mask=encoder_attention_mask, + attention_mask=encoder_attention_mask if encoder_hidden_states is not None else attention_mask, ) if self.final_dropout: attn_output = self.final_dropout(attn_output) @@ -266,8 +270,8 @@ class DiT(ModelMixin, ConfigMixin): self.norm_out = nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6) self.proj_out_1 = nn.Linear(self.inner_dim, 2 * self.inner_dim) self.proj_out_2 = nn.Linear(self.inner_dim, self.config.output_dim) - print( - "Total number of DiT parameters: ", + logger.debug( + "Total number of DiT parameters: %d", sum(p.numel() for p in self.parameters() if p.requires_grad), ) @@ -318,6 +322,71 @@ class DiT(ModelMixin, ConfigMixin): return self.proj_out_2(hidden_states) +class AlternateVLDiT(DiT): + """N1.7 DiT variant that alternates cross-attention over image and text tokens.""" + + def __init__(self, *args, attend_text_every_n_blocks: int = 2, **kwargs): + super().__init__(*args, **kwargs) + self.attend_text_every_n_blocks = attend_text_every_n_blocks + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.LongTensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + return_all_hidden_states: bool = False, + image_mask: torch.Tensor | None = None, + backbone_attention_mask: torch.Tensor | None = None, + ): + if image_mask is None: + raise ValueError("image_mask is required for AlternateVLDiT.") + if backbone_attention_mask is None: + raise ValueError("backbone_attention_mask is required for AlternateVLDiT.") + + temb = self.timestep_encoder(timestep) + hidden_states = hidden_states.contiguous() + encoder_hidden_states = encoder_hidden_states.contiguous() + + image_attention_mask = image_mask & backbone_attention_mask + non_image_attention_mask = (~image_mask) & backbone_attention_mask + + all_hidden_states = [hidden_states] + if not self.config.interleave_self_attention: + raise ValueError("AlternateVLDiT requires interleave_self_attention=True.") + + for idx, block in enumerate(self.transformer_blocks): + if idx % 2 == 1: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + temb=temb, + ) + else: + curr_encoder_attention_mask = ( + non_image_attention_mask + if idx % (2 * self.attend_text_every_n_blocks) == 0 + else image_attention_mask + ) + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=curr_encoder_attention_mask, + temb=temb, + ) + all_hidden_states.append(hidden_states) + + conditioning = temb + shift, scale = self.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + if return_all_hidden_states: + return self.proj_out_2(hidden_states), all_hidden_states + return self.proj_out_2(hidden_states) + + class SelfAttentionTransformer(ModelMixin, ConfigMixin): _supports_gradient_checkpointing = True @@ -362,8 +431,8 @@ class SelfAttentionTransformer(ModelMixin, ConfigMixin): for _ in range(self.config.num_layers) ] ) - print( - "Total number of SelfAttentionTransformer parameters: ", + logger.debug( + "Total number of SelfAttentionTransformer parameters: %d", sum(p.numel() for p in self.parameters() if p.requires_grad), ) diff --git a/src/lerobot/policies/groot/action_head/flow_matching_action_head.py b/src/lerobot/policies/groot/action_head/flow_matching_action_head.py deleted file mode 100644 index 986820670..000000000 --- a/src/lerobot/policies/groot/action_head/flow_matching_action_head.py +++ /dev/null @@ -1,408 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# 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 dataclasses import field -from typing import TYPE_CHECKING - -import torch -import torch.nn.functional as F # noqa: N812 -from torch import nn -from torch.distributions import Beta - -from lerobot.utils.import_utils import _transformers_available - -# Conditional import for type checking and lazy loading -if TYPE_CHECKING or _transformers_available: - from transformers import PretrainedConfig - from transformers.feature_extraction_utils import BatchFeature -else: - PretrainedConfig = object - BatchFeature = None - -from .action_encoder import ( - SinusoidalPositionalEncoding, - swish, -) -from .cross_attention_dit import DiT, SelfAttentionTransformer - - -class CategorySpecificLinear(nn.Module): - def __init__(self, num_categories, input_dim, hidden_dim): - super().__init__() - self.num_categories = num_categories - # For each category, we have separate weights and biases. - self.W = nn.Parameter(0.02 * torch.randn(num_categories, input_dim, hidden_dim)) - self.b = nn.Parameter(torch.zeros(num_categories, hidden_dim)) - - def forward(self, x, cat_ids): - selected_w = self.W[cat_ids] - selected_b = self.b[cat_ids] - return torch.bmm(x, selected_w) + selected_b.unsqueeze(1) - - -class CategorySpecificMLP(nn.Module): - def __init__(self, num_categories, input_dim, hidden_dim, output_dim): - super().__init__() - self.num_categories = num_categories - self.layer1 = CategorySpecificLinear(num_categories, input_dim, hidden_dim) - self.layer2 = CategorySpecificLinear(num_categories, hidden_dim, output_dim) - - def forward(self, x, cat_ids): - hidden = F.relu(self.layer1(x, cat_ids)) - return self.layer2(hidden, cat_ids) - - -class MultiEmbodimentActionEncoder(nn.Module): - def __init__(self, action_dim, hidden_size, num_embodiments): - super().__init__() - self.hidden_size = hidden_size - self.num_embodiments = num_embodiments - - # W1: R^{w x d}, W2: R^{w x 2w}, W3: R^{w x w} - self.W1 = CategorySpecificLinear(num_embodiments, action_dim, hidden_size) # (d -> w) - self.W2 = CategorySpecificLinear(num_embodiments, 2 * hidden_size, hidden_size) # (2w -> w) - self.W3 = CategorySpecificLinear(num_embodiments, hidden_size, hidden_size) # (w -> w) - self.pos_encoding = SinusoidalPositionalEncoding(hidden_size) - - def forward(self, actions, timesteps, cat_ids): - """ - actions: shape (B, T, action_dim) - timesteps: shape (B,) -- a single scalar per batch item - cat_ids: shape (B,) - returns: shape (B, T, hidden_size) - """ - b, t, _ = actions.shape - - # 1) Expand each batch's single scalar time 'tau' across all T steps - # so that shape => (B, T) - # e.g. if timesteps is (B,), replicate across T - if timesteps.dim() == 1 and timesteps.shape[0] == b: - # shape (B,) => (B,T) - timesteps = timesteps.unsqueeze(1).expand(-1, t) - else: - raise ValueError("Expected `timesteps` to have shape (B,) so we can replicate across T.") - - # 2) Standard action MLP step for shape => (B, T, w) - a_emb = self.W1(actions, cat_ids) - - # 3) Get the sinusoidal encoding (B, T, w) - tau_emb = self.pos_encoding(timesteps).to(dtype=a_emb.dtype) - - # 4) Concat along last dim => (B, T, 2w), then W2 => (B, T, w), swish - x = torch.cat([a_emb, tau_emb], dim=-1) - x = swish(self.W2(x, cat_ids)) - - # 5) Finally W3 => (B, T, w) - x = self.W3(x, cat_ids) - return x - - -class FlowmatchingActionHeadConfig(PretrainedConfig): - """NOTE: N1.5 uses XEmbFlowmatchingPolicyHeadConfig as action head""" - - add_pos_embed: bool = field(default=True, metadata={"help": "Whether to add positional embedding"}) - model_dtype: str = field(default="float32", metadata={"help": "Model data type."}) - diffusion_model_cfg: dict = field(default=None, metadata={"help": "Diffusion model configuration."}) - input_embedding_dim: int = field(default=1536, metadata={"help": "Input embedding channel dimension."}) - backbone_embedding_dim: int = field( - default=1536, metadata={"help": "Backbone embedding channel dimension."} - ) - - hidden_size: int = field(default=1024, metadata={"help": "Input embedding dimension."}) - max_seq_len: int = field(default=1024, metadata={"help": "Maximum Sequence Length"}) - action_dim: int = field(default=None, metadata={"help": "Action dimension."}) - action_horizon: int = field(default=None, metadata={"help": "Action horizon."}) - noise_beta_alpha: float = field(default=1.5, metadata={"help": ""}) - noise_beta_beta: float = field(default=1.0, metadata={"help": ""}) - noise_s: float = field(default=0.999, metadata={"help": "Flow matching noise Beta distribution s."}) - num_timestep_buckets: int = field( - default=1000, metadata={"help": "Number of timestep discretization buckets."} - ) - num_inference_timesteps: int = field( - default=None, - metadata={"help": "Number of inference steps for noise diffusion."}, - ) - max_num_embodiments: int = field(default=32, metadata={"help": "Number of embodiments."}) - tune_projector: bool = field(default=True, metadata={"help": "Whether to tune the projector."}) - tune_diffusion_model: bool = field( - default=True, metadata={"help": "Whether to tune the diffusion model."} - ) - load_pretrained_det_decode_layer_path: str = field( - default=None, metadata={"help": "Path to pretrained detection model."} - ) - detection_coeff: float = field(default=1.0, metadata={"help": "Detection coefficient."}) - - freeze_decode_layer: bool = field(default=False) - expand_batch: int = field(default=None) - use_vlln: bool = field(default=True) - - vl_self_attention_cfg: dict = field(default=None) - num_target_vision_tokens: int = field(default=32, metadata={"help": "Number of target vision tokens."}) - - def __init__(self, **kwargs): - super().__init__(**kwargs) - for key, value in kwargs.items(): - setattr(self, key, value) - - -class FlowmatchingActionHead(nn.Module): - config_class = FlowmatchingActionHeadConfig - supports_gradient_checkpointing = True - - def __init__( - self, - config: FlowmatchingActionHeadConfig, - ): - super().__init__() - self.hidden_size = config.hidden_size - self.input_embedding_dim = config.input_embedding_dim - - self.model = DiT(**config.diffusion_model_cfg) - self.action_dim = config.action_dim - self.action_horizon = config.action_horizon - self.num_inference_timesteps = config.num_inference_timesteps - - self.state_encoder = CategorySpecificMLP( - num_categories=config.max_num_embodiments, - input_dim=config.max_state_dim, - hidden_dim=self.hidden_size, - output_dim=self.input_embedding_dim, - ) - self.action_encoder = MultiEmbodimentActionEncoder( - action_dim=config.action_dim, - hidden_size=self.input_embedding_dim, - num_embodiments=config.max_num_embodiments, - ) - self.action_decoder = CategorySpecificMLP( - num_categories=config.max_num_embodiments, - input_dim=self.hidden_size, - hidden_dim=self.hidden_size, - output_dim=self.action_dim, - ) - self.future_tokens = nn.Embedding(config.num_target_vision_tokens, self.input_embedding_dim) - nn.init.normal_(self.future_tokens.weight, mean=0.0, std=0.02) - - self.vlln = nn.LayerNorm(config.backbone_embedding_dim) if config.use_vlln else nn.Identity() - self.vl_self_attention = ( - SelfAttentionTransformer(**config.vl_self_attention_cfg) if config.use_vlln else nn.Identity() - ) - - if config.add_pos_embed: - self.position_embedding = nn.Embedding(config.max_seq_len, self.input_embedding_dim) - nn.init.normal_(self.position_embedding.weight, mean=0.0, std=0.02) - - self._noise_beta_alpha = config.noise_beta_alpha - self._noise_beta_beta = config.noise_beta_beta - self._beta_dist = None - self.num_timestep_buckets = config.num_timestep_buckets - self.config = config - self.set_trainable_parameters(config.tune_projector, config.tune_diffusion_model) - - def set_trainable_parameters(self, tune_projector: bool, tune_diffusion_model: bool): - self.tune_projector = tune_projector - self.tune_diffusion_model = tune_diffusion_model - for p in self.parameters(): - p.requires_grad = True - if not tune_projector: - self.state_encoder.requires_grad_(False) - self.action_encoder.requires_grad_(False) - self.action_decoder.requires_grad_(False) - if self.config.add_pos_embed: - self.position_embedding.requires_grad_(False) - if not tune_diffusion_model: - self.model.requires_grad_(False) - print(f"Tune action head projector: {self.tune_projector}") - print(f"Tune action head diffusion model: {self.tune_diffusion_model}") - # Check if any parameters are still trainable. If not, print a warning. - if not tune_projector and not tune_diffusion_model: - for name, p in self.named_parameters(): - if p.requires_grad: - print(f"Action head trainable parameter: {name}") - if not any(p.requires_grad for p in self.parameters()): - print("Warning: No action head trainable parameters found.") - - def set_frozen_modules_to_eval_mode(self): - """ - Huggingface will call model.train() at each training_step. To ensure - the expected behaviors for modules like dropout, batchnorm, etc., we - need to call model.eval() for the frozen modules. - """ - if self.training: - if not self.tune_projector: - self.state_encoder.eval() - self.action_encoder.eval() - self.action_decoder.eval() - if self.config.add_pos_embed: - self.position_embedding.eval() - if not self.tune_diffusion_model: - self.model.eval() - - def sample_time(self, batch_size, device, dtype): - if self._beta_dist is None: - self._beta_dist = Beta(self._noise_beta_alpha, self._noise_beta_beta, validate_args=False) - sample = self._beta_dist.sample([batch_size]).to(device, dtype=dtype) - return (self.config.noise_s - sample) / self.config.noise_s - - def prepare_input(self, batch: dict) -> BatchFeature: - return BatchFeature(data=batch) - - def process_backbone_output(self, backbone_output: BatchFeature) -> BatchFeature: - backbone_features = backbone_output["backbone_features"] - backbone_features = self.vlln(backbone_features) - backbone_features = self.vl_self_attention(backbone_features) - backbone_output["backbone_features"] = backbone_features - return backbone_output - - def forward(self, backbone_output: BatchFeature, action_input: BatchFeature) -> BatchFeature: - # Set frozen modules to eval - self.set_frozen_modules_to_eval_mode() - - backbone_output = self.process_backbone_output(backbone_output) - - if self.config.expand_batch is not None: - for k, v in backbone_output.items(): - ndim = len(v.shape) - factors = [self.config.expand_batch] - while len(factors) < ndim: - factors.append(1) - factors = tuple(factors) - expanded = v.repeat(*factors) - backbone_output[k] = expanded - - for k, v in action_input.items(): - ndim = len(v.shape) - factors = [self.config.expand_batch] - while len(factors) < ndim: - factors.append(1) - factors = tuple(factors) - expanded = v.repeat(*factors) - action_input[k] = expanded - - # Get vision and language embeddings. - vl_embs = backbone_output.backbone_features - device = vl_embs.device - - # Get embodiment ID. - embodiment_id = action_input.embodiment_id - - # Embed state. - state_features = self.state_encoder(action_input.state, embodiment_id) - - # Embed noised action trajectory. - actions = action_input.action - noise = torch.randn(actions.shape, device=actions.device, dtype=actions.dtype) - t = self.sample_time(actions.shape[0], device=actions.device, dtype=actions.dtype) - t = t[:, None, None] # shape (B,1,1) for broadcast - - noisy_trajectory = (1 - t) * noise + t * actions - velocity = actions - noise - - # Convert (continuous) t -> discrete if needed - t_discretized = (t[:, 0, 0] * self.num_timestep_buckets).long() - action_features = self.action_encoder(noisy_trajectory, t_discretized, embodiment_id) - - # Maybe add position embedding. - if self.config.add_pos_embed: - pos_ids = torch.arange(action_features.shape[1], dtype=torch.long, device=device) - pos_embs = self.position_embedding(pos_ids).unsqueeze(0) - action_features = action_features + pos_embs - - # Join vision, language, state and action embedding along sequence dimension. - future_tokens = self.future_tokens.weight.unsqueeze(0).expand(vl_embs.shape[0], -1, -1) - sa_embs = torch.cat((state_features, future_tokens, action_features), dim=1) - - vl_attn_mask = backbone_output.backbone_attention_mask - - model_output = self.model( - hidden_states=sa_embs, - encoder_hidden_states=vl_embs, - encoder_attention_mask=vl_attn_mask, - timestep=t_discretized, - return_all_hidden_states=False, # NOTE (YL): not using flare now - ) - pred = self.action_decoder(model_output, embodiment_id) - pred_actions = pred[:, -actions.shape[1] :] - - # Slice out only the action portion of pred and target. - action_mask = action_input.action_mask - loss = F.mse_loss(pred_actions, velocity, reduction="none") * action_mask - loss = loss.sum() / action_mask.sum() - output_dict = { - "loss": loss, - } - return BatchFeature(data=output_dict) - - @torch.no_grad() - def get_action(self, backbone_output: BatchFeature, action_input: BatchFeature) -> BatchFeature: - backbone_output = self.process_backbone_output(backbone_output) - - # Get vision and language embeddings. - vl_embs = backbone_output.backbone_features - embodiment_id = action_input.embodiment_id - - # Embed state. - state_features = self.state_encoder(action_input.state, embodiment_id) - - # Set initial actions as the sampled noise. - batch_size = vl_embs.shape[0] - device = vl_embs.device - actions = torch.randn( - size=(batch_size, self.config.action_horizon, self.config.action_dim), - dtype=vl_embs.dtype, - device=device, - ) - - num_steps = self.num_inference_timesteps - dt = 1.0 / num_steps - - # Run denoising steps. - for t in range(num_steps): - t_cont = t / float(num_steps) # e.g. goes 0, 1/N, 2/N, ... - t_discretized = int(t_cont * self.num_timestep_buckets) - - # Embed noised action trajectory. - timesteps_tensor = torch.full(size=(batch_size,), fill_value=t_discretized, device=device) - action_features = self.action_encoder(actions, timesteps_tensor, embodiment_id) - # Maybe add position embedding. - if self.config.add_pos_embed: - pos_ids = torch.arange(action_features.shape[1], dtype=torch.long, device=device) - pos_embs = self.position_embedding(pos_ids).unsqueeze(0) - action_features = action_features + pos_embs - - # Join vision, language, state and action embedding along sequence dimension. - future_tokens = self.future_tokens.weight.unsqueeze(0).expand(vl_embs.shape[0], -1, -1) - sa_embs = torch.cat((state_features, future_tokens, action_features), dim=1) - - # Run model forward. - model_output = self.model( - hidden_states=sa_embs, - encoder_hidden_states=vl_embs, - timestep=timesteps_tensor, - ) - pred = self.action_decoder(model_output, embodiment_id) - - pred_velocity = pred[:, -self.action_horizon :] - - # Update actions using euler integration. - actions = actions + dt * pred_velocity - return BatchFeature(data={"action_pred": actions}) - - @property - def device(self): - return next(iter(self.parameters())).device - - @property - def dtype(self): - return next(iter(self.parameters())).dtype diff --git a/src/lerobot/policies/groot/configuration_groot.py b/src/lerobot/policies/groot/configuration_groot.py index 17cb631d7..97e08bb76 100644 --- a/src/lerobot/policies/groot/configuration_groot.py +++ b/src/lerobot/policies/groot/configuration_groot.py @@ -14,12 +14,229 @@ # See the License for the specific language governing permissions and # limitations under the License. +import logging +import math from dataclasses import dataclass, field +from pathlib import Path from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature, PreTrainedConfig -from lerobot.optim import AdamWConfig, CosineDecayWithWarmupSchedulerConfig +from lerobot.optim import AdamWConfig, DiffuserSchedulerConfig from lerobot.utils.constants import ACTION, OBS_STATE +from .utils import read_json + +logger = logging.getLogger(__name__) + +GROOT_N1_7 = "n1.7" +# Legacy GR00T N1.5 identifier. N1.5 is NOT a supported model_version (it is +# intentionally absent from _GROOT_MODEL_VERSION_ALIASES so normalize_groot_model_version +# still rejects it). It is retained only so that infer_groot_model_version can recognise +# an N1.5 base path/checkpoint and the N1.7 config/loader can reject the mismatch. +GROOT_N1_5 = "n1.5" +# Canonical guidance appended to every error raised when an N1.5 checkpoint, config, +# or processor pipeline is detected. Keep this message in sync with docs/source/groot.mdx. +GROOT_N1_5_REMOVAL_GUIDANCE = ( + "GR00T N1.5 support was removed from LeRobot. " + "To keep using an N1.5 checkpoint, pin the last release that supports it: " + "`pip install 'lerobot==0.5.1'`. To use the current release, migrate to GR00T N1.7 " + "(model_version='n1.7', base model nvidia/GR00T-N1.7-3B)." +) +GROOT_N1_7_BASE_MODEL = "nvidia/GR00T-N1.7-3B" +GROOT_N1_7_BACKBONE_MODEL = "nvidia/Cosmos-Reason2-2B" +# Default GR00T N1.7 training resolution. Fallback if processor_config lacks sizing. Prevents mismatched +# full-res patchification by forcing a resize. Mirrored by GR00T_N1_7_DEFAULTS in groot_n1_7.py. +N1_7_DEFAULT_IMAGE_TARGET_SIZE = (256, 256) +N1_7_DEFAULT_IMAGE_CROP_SIZE = (230, 230) +GROOT_ACTION_DECODE_TRANSFORM_LIBERO = "libero" +# Sentinel meaning "the user did not pick an action decode transform": __post_init__ resolves it +# to the embodiment default ('libero' for 'libero_sim', otherwise None). It is distinct from an +# explicit 'none' (resolved to None) so an opt-out survives a draccus save/load round-trip. +GROOT_ACTION_DECODE_TRANSFORM_AUTO = "auto" + +_GROOT_MODEL_VERSION_ALIASES = { + "n1.7": GROOT_N1_7, + "n1_7": GROOT_N1_7, + "n1d7": GROOT_N1_7, + "n17": GROOT_N1_7, + "1.7": GROOT_N1_7, +} + +# Legacy N1.5 spellings, kept ONLY so they can be detected and rejected with +# GROOT_N1_5_REMOVAL_GUIDANCE (see GROOT_N1_5 above). Never map these to a supported version. +_GROOT_N1_5_VERSION_ALIASES = {"n1.5", "n1_5", "n1d5", "n15", "1.5"} + +_GROOT_ACTION_DECODE_TRANSFORM_ALIASES = { + GROOT_ACTION_DECODE_TRANSFORM_AUTO: GROOT_ACTION_DECODE_TRANSFORM_AUTO, + "none": None, + "": None, + GROOT_ACTION_DECODE_TRANSFORM_LIBERO: GROOT_ACTION_DECODE_TRANSFORM_LIBERO, +} + + +def normalize_groot_model_version(model_version: str) -> str: + normalized = _GROOT_MODEL_VERSION_ALIASES.get(model_version.lower()) + if normalized is None: + supported = GROOT_N1_7 + message = f"Unsupported GR00T model_version '{model_version}'. Supported versions: {supported}." + if model_version.lower() in _GROOT_N1_5_VERSION_ALIASES: + message = f"{message} {GROOT_N1_5_REMOVAL_GUIDANCE}" + raise ValueError(message) + return normalized + + +def normalize_groot_action_decode_transform(transform: str | None) -> str | None: + if transform is None: + return None + normalized = _GROOT_ACTION_DECODE_TRANSFORM_ALIASES.get(transform.lower()) + if normalized is None and transform.lower() not in _GROOT_ACTION_DECODE_TRANSFORM_ALIASES: + supported = ", ".join( + sorted(key for key, value in _GROOT_ACTION_DECODE_TRANSFORM_ALIASES.items() if value is not None) + ) + raise ValueError( + f"Unsupported GR00T N1.7 action decode transform '{transform}'. " + f"Supported transforms: none, {supported}." + ) + return normalized + + +def infer_groot_model_version(model_path: str | None) -> str | None: + if not model_path: + return None + model_path_lower = model_path.lower() + if "gr00t-n1.7" in model_path_lower or "gr00t_n1.7" in model_path_lower: + return GROOT_N1_7 + # Detect legacy N1.5 paths so the N1.7 config/loader can reject the mismatch. + # N1.5 is unsupported, but it must still be recognised here to fail loudly + # rather than silently treating an N1.5 checkpoint as N1.7. + if "gr00t-n1.5" in model_path_lower or "gr00t_n1.5" in model_path_lower: + return GROOT_N1_5 + config_version = _infer_groot_model_version_from_local_config(model_path) + if config_version is not None: + return config_version + return None + + +def is_raw_groot_n1_7_checkpoint(model_path: str | Path | None) -> bool: + if model_path is None: + return False + + path = Path(model_path).expanduser() + if path.is_dir(): + config_path = path / "config.json" + elif path.name == "config.json": + config_path = path + else: + return False + + config = read_json(config_path) + return "type" not in config and _infer_groot_model_version_from_config(config) == GROOT_N1_7 + + +def infer_groot_n1_7_embodiment_tag(model_path: str | Path | None) -> str | None: + if model_path is None: + return None + + processor_config_path = Path(model_path).expanduser() / "processor_config.json" + processor_config = read_json(processor_config_path) + + modality_configs = processor_config.get("processor_kwargs", {}).get("modality_configs", {}) + if not isinstance(modality_configs, dict): + return None + if "libero_sim" in modality_configs: + return "libero_sim" + if len(modality_configs) == 1: + return next(iter(modality_configs)) + return None + + +def infer_groot_n1_7_action_horizon( + model_path: str | Path | None, embodiment_tag: str | None = None +) -> int | None: + if model_path is None: + return None + + processor_config_path = Path(model_path).expanduser() / "processor_config.json" + processor_config = read_json(processor_config_path) + + processor_kwargs = processor_config.get("processor_kwargs", {}) + if not isinstance(processor_kwargs, dict): + return None + modality_configs = processor_kwargs.get("modality_configs", {}) + if not isinstance(modality_configs, dict): + return None + + if embodiment_tag is None: + embodiment_tag = infer_groot_n1_7_embodiment_tag(model_path) + if embodiment_tag is None: + return None + + embodiment_config = modality_configs.get(embodiment_tag, {}) + if not isinstance(embodiment_config, dict): + return None + action_config = embodiment_config.get("action", {}) + if not isinstance(action_config, dict): + return None + delta_indices = action_config.get("delta_indices", []) + if not isinstance(delta_indices, list): + return None + return len(delta_indices) or None + + +def infer_groot_n1_7_action_execution_horizon( + model_path: str | Path | None, embodiment_tag: str | None = None +) -> int | None: + action_horizon = infer_groot_n1_7_action_horizon(model_path, embodiment_tag) + if action_horizon is None: + return None + + if embodiment_tag is None: + embodiment_tag = infer_groot_n1_7_embodiment_tag(model_path) + if embodiment_tag == "libero_sim": + # NVIDIA's N1.7 LIBERO rollout wrapper replans after 8 of the 16 decoded + # actions. Keeping that execution cadence avoids stale open-loop chunks. + return min(action_horizon, 8) + return action_horizon + + +def _infer_groot_model_version_from_local_config(model_path: str) -> str | None: + path = Path(model_path).expanduser() + if path.is_dir(): + config_path = path / "config.json" + elif path.name == "config.json": + config_path = path + else: + return None + + return _infer_groot_model_version_from_config(read_json(config_path)) + + +def _infer_groot_model_version_from_config(config: dict) -> str | None: + model_version = config.get("model_version") + if isinstance(model_version, str): + if model_version.lower() in _GROOT_N1_5_VERSION_ALIASES: + return GROOT_N1_5 + try: + return normalize_groot_model_version(model_version) + except ValueError: + return None + + candidates = [config.get("model_type"), *(config.get("architectures") or [])] + for candidate in candidates: + if not isinstance(candidate, str): + continue + normalized = candidate.lower().replace("-", "_") + if normalized in {"gr00tn1d7", "gr00t_n1d7", "gr00t_n1_7"}: + return GROOT_N1_7 + if normalized in {"gr00t_n1_5", "gr00tn1_5", "gr00t_n15", "gr00t_n1d5", "gr00tn1d5"}: + return GROOT_N1_5 + if config.get("model_name") == GROOT_N1_7_BACKBONE_MODEL: + return GROOT_N1_7 + # The Eagle VLM backbone is specific to pre-N1.7 GR00T checkpoints (N1.7 uses Cosmos/Qwen3-VL). + backbone_cfg = config.get("backbone_cfg") + if isinstance(backbone_cfg, dict) and "eagle_path" in backbone_cfg: + return GROOT_N1_5 + return None + @PreTrainedConfig.register_subclass("groot") @dataclass @@ -28,35 +245,44 @@ class GrootConfig(PreTrainedConfig): # Basic policy settings n_obs_steps: int = 1 - chunk_size: int = 50 - n_action_steps: int = 50 + chunk_size: int = 40 + n_action_steps: int = 40 # Dimension settings (must match pretrained GR00T model expectations) # Maximum state dimension. Shorter states will be zero-padded. - max_state_dim: int = 64 + max_state_dim: int = 132 # Maximum action dimension. Shorter actions will be zero-padded. - max_action_dim: int = 32 + max_action_dim: int = 132 - # Normalization (start with identity, adjust as needed) + # GR00T normalizes state/action internally in its processor steps (min/max with + # q01/q99 percentiles, per embodiment), and the Qwen3-VL backbone's image processor + # handles image normalization. The policy therefore does NOT use LeRobot's + # NormalizerProcessorStep/UnnormalizerProcessorStep, so this mapping is intentionally + # IDENTITY for every feature and is not consulted by make_groot_pre_post_processors. normalization_mapping: dict[str, NormalizationMode] = field( default_factory=lambda: { "VISUAL": NormalizationMode.IDENTITY, - "STATE": NormalizationMode.MEAN_STD, - "ACTION": NormalizationMode.MEAN_STD, + "STATE": NormalizationMode.IDENTITY, + "ACTION": NormalizationMode.IDENTITY, } ) - # Image preprocessing (adjust to match Groot's expected input) - image_size: tuple[int, int] = (224, 224) + # Groot-specific model parameters - # Groot-specific model parameters (from groot_finetune_script.py) + # Path or HuggingFace model ID for the base GR00T N1.7 model whose backbone weights and + # checkpoint sidecars (statistics.json, processor_config.json, ...) are loaded. This is the + # model *source*, and is intentionally distinct from the inherited `pretrained_path`: + # `pretrained_path` (`--policy.path`) points at a saved LeRobot checkpoint directory whose + # `config.json` carries a `type` field, whereas a raw NVIDIA GR00T checkpoint has no such + # field and so can only be loaded through `base_model_path` (`--policy.base_model_path`). + # Defaults to GROOT_N1_7_BASE_MODEL when unset (resolved in __post_init__). + base_model_path: str | None = None - # Path or HuggingFace model ID for the base Groot model - base_model_path: str = "nvidia/GR00T-N1.5-3B" - - # HF repo ID (or local path) that hosts vocab.json and merges.txt for Eagle tokenizer. - tokenizer_assets_repo: str = "lerobot/eagle2hg-processor-groot-n1p5" + # Optional named action transform applied after raw N1.7 checkpoint decoding and before env.step(). + # 'auto' (default) resolves to the embodiment default ('libero' for 'libero_sim', otherwise no + # transform). Pass 'none' to explicitly disable the transform, including for 'libero_sim'. + action_decode_transform: str | None = GROOT_ACTION_DECODE_TRANSFORM_AUTO # Embodiment tag to use for training (e.g. 'new_embodiment', 'gr1') embodiment_tag: str = "new_embodiment" @@ -75,38 +301,67 @@ class GrootConfig(PreTrainedConfig): # Whether to fine-tune the diffusion model tune_diffusion_model: bool = True - # LoRA parameters (from groot_finetune_script.py) - # Rank for the LORA model. If 0, no LORA will be used. - lora_rank: int = 0 + # Whether to fine-tune the VL LayerNorm + VL self-attention projector in the action head. + tune_vlln: bool = True - # Alpha value for the LORA model - lora_alpha: int = 16 + # Number of top LLM backbone layers to fine-tune (0 = none). Lets you adapt just the final + # language layers without unfreezing the whole backbone; independent of `tune_llm`, which tunes + # the entire LLM. + tune_top_llm_layers: int = 0 - # Dropout rate for the LORA model - lora_dropout: float = 0.1 + # Inference-time knob: Number of flow-matching denoising steps used to decode an action chunk. + # Trades inference latency for action quality. + # None keeps the checkpoint value (GR00T N1.7 default: 4). + num_inference_timesteps: int | None = None - # Whether to use the full model for LORA - lora_full_model: bool = False + # Inference-time knob: Real-Time Chunking (RTC) overlap-blend ramp rate, used when the RTC engine + # supplies a previous-chunk prefix. Higher values blend the overlapping prefix more aggressively. + # None keeps the checkpoint value (GR00T N1.7 default: 6.0). + rtc_ramp_rate: float | None = None - # Training parameters (matching groot_finetune_script.py) + # Inference-time knob: Whether to request the flash-attention-2 kernel for the Qwen3-VL backbone. + # flash-attn is an optional, user-managed optimization; when it is absent (the default), + # the backbone transparently falls back to SDPA, which is numerically equivalent. + # Set to True only after installing a flash-attn build matching your torch/CUDA env. + use_flash_attention: bool = False + + # Enable GR00T-style state-relative action chunks (action chunk expressed relative to the current + # observation state). + use_relative_actions: bool = False + + # relative_exclude_joints names the action dimensions that stay absolute; the + # match is substring/case-insensitive against the dataset action feature names. With the empty + # default every dimension is treated as relative, including the gripper -- set e.g. ["gripper"] to + # keep the gripper absolute, matching the Isaac-GR00T single-arm + absolute-gripper convention. + relative_exclude_joints: list[str] = field(default_factory=list) + + # Training parameters optimizer_lr: float = 1e-4 - optimizer_betas: tuple[float, float] = (0.95, 0.999) + # Isaac-GR00T N1.7 fine-tunes with AdamW betas (0.9, 0.999). + optimizer_betas: tuple[float, float] = (0.9, 0.999) optimizer_eps: float = 1e-8 optimizer_weight_decay: float = 1e-5 warmup_ratio: float = 0.05 use_bf16: bool = True + # The native N1.7 fine-tuning recipe keeps model parameters in FP32 and computes under BF16 autocast. + model_params_fp32: bool = True - # Dataset parameters - # Video backend to use for training ('decord' or 'torchvision_av') + # TODO(Steven): Remove these deprecated fields in a future release. + # Deprecated Isaac-GR00T runner / GR00T N1.5 fields, plus the (never-wired) LoRA fields — all + # unused by the LeRobot N1.7 implementation except the `tokenizer_assets_repo` N1.5 tripwire and + # the `image_size` legacy remap in __post_init__. They are kept ONLY so a config.json saved by an + # earlier lerobot release (notably a GR00T N1.5 checkpoint) still parses under draccus — which + # rejects unknown fields — and is then rejected with a clear N1.5 removal message rather than an + # opaque draccus decoding error. + image_size: tuple[int, int] = (256, 256) # image sizing is handled by the backbone's image processor. + tokenizer_assets_repo: str | None = None + lora_rank: int = 0 + lora_alpha: int = 16 + lora_dropout: float = 0.1 + lora_full_model: bool = False video_backend: str = "decord" - - # Whether to balance dataset weights in mixture datasets balance_dataset_weights: bool = True - - # Whether to sample trajectories weighted by their length balance_trajectory_weights: bool = True - - # Optional dataset paths for delegating training to Isaac-GR00T runner dataset_paths: list[str] | None = None output_dir: str = "./tmp/gr00t" save_steps: int = 1000 @@ -117,6 +372,65 @@ class GrootConfig(PreTrainedConfig): resume: bool = False def __post_init__(self): + if self.tokenizer_assets_repo is not None: + raise ValueError( + "Config sets 'tokenizer_assets_repo', which only existed for GR00T N1.5; this looks " + f"like a legacy GR00T N1.5 checkpoint or config. {GROOT_N1_5_REMOVAL_GUIDANCE}" + ) + + self.action_decode_transform = normalize_groot_action_decode_transform(self.action_decode_transform) + if self.base_model_path is None: + self.base_model_path = GROOT_N1_7_BASE_MODEL + + # The N1.7 LIBERO checkpoints emit a [0, 1] gripper action, but the LIBERO + # simulator expects the OpenVLA/[-1, 1] sign convention. NVIDIA's rollout + # wrapper applies this conversion; mirror it here so eval on the + # 'libero_sim' embodiment grasps correctly instead of scoring 0% success. + # This matches the embodiment-specific handling already done for the + # action execution horizon (see infer_groot_n1_7_action_execution_horizon). + # Only the 'auto' sentinel resolves to the embodiment default; an explicit + # 'none' (normalized to None above) keeps the transform disabled. + if self.action_decode_transform == GROOT_ACTION_DECODE_TRANSFORM_AUTO: + self.action_decode_transform = ( + GROOT_ACTION_DECODE_TRANSFORM_LIBERO if self.embodiment_tag == "libero_sim" else None + ) + + # GR00T N1.5-era default values (e.g. --policy.chunk_size=50 from old commands or + # stale configs) are migrated to the values the N1.7 checkpoints expect, with a + # warning. The dataclass defaults are already the N1.7 values, so a plain + # GrootConfig() never triggers this. + legacy_default_remaps = ( + ("max_state_dim", 64, 132), + ("max_action_dim", 32, 132), + ("chunk_size", 50, 40), + ("n_action_steps", 50, 40), + ("image_size", (224, 224), (256, 256)), + ) + for field_name, legacy_value, n1_7_value in legacy_default_remaps: + current_value = getattr(self, field_name) + if isinstance(legacy_value, tuple): + current_value = tuple(current_value) + if current_value == legacy_value: + logger.warning( + "GrootConfig.%s=%s matches a legacy GR00T N1.5-era default; remapping it to %s, " + "the value expected by GR00T N1.7 checkpoints. Set a different value explicitly " + "if this is not what you want.", + field_name, + legacy_value, + n1_7_value, + ) + setattr(self, field_name, n1_7_value) + + inferred_version = infer_groot_model_version(self.base_model_path) + if inferred_version is not None and inferred_version != GROOT_N1_7: + message = ( + f"GR00T model_version '{GROOT_N1_7}' does not match base_model_path " + f"'{self.base_model_path}', which looks like '{inferred_version}'." + ) + if inferred_version == GROOT_N1_5: + message = f"{message} {GROOT_N1_5_REMOVAL_GUIDANCE}" + raise ValueError(message) + super().__post_init__() if self.n_action_steps > self.chunk_size: @@ -124,9 +438,6 @@ class GrootConfig(PreTrainedConfig): f"n_action_steps ({self.n_action_steps}) cannot exceed chunk_size ({self.chunk_size})" ) - # groot_repo_path is now optional since we ported the components - # No validation needed - def validate_features(self) -> None: """Validate and set up input/output features for Groot.""" image_features = [key for key, feat in self.input_features.items() if feat.type == FeatureType.VISUAL] @@ -173,15 +484,20 @@ class GrootConfig(PreTrainedConfig): betas=self.optimizer_betas, eps=self.optimizer_eps, weight_decay=self.optimizer_weight_decay, + grad_clip_norm=1.0, ) - def get_scheduler_preset(self) -> CosineDecayWithWarmupSchedulerConfig: - """Return scheduler configuration.""" - return CosineDecayWithWarmupSchedulerConfig( - num_warmup_steps=int(10000 * self.warmup_ratio), # 5% warmup by default - num_decay_steps=10000, # Adjust based on training steps - peak_lr=self.optimizer_lr, - decay_lr=self.optimizer_lr * 0.1, + def get_scheduler_preset(self) -> DiffuserSchedulerConfig: + """Return scheduler configuration. + + Isaac-GR00T uses the HF Trainer cosine schedule with ~5% warmup over the + actual training update count; DiffuserSchedulerConfig wraps the same + diffusers/transformers `get_scheduler("cosine")` implementation and + derives num_training_steps from the outer --steps value at runtime. + """ + return DiffuserSchedulerConfig( + name="cosine", + num_warmup_steps=math.ceil(self.max_steps * self.warmup_ratio), ) @property @@ -192,7 +508,15 @@ class GrootConfig(PreTrainedConfig): @property def action_delta_indices(self) -> list[int]: """Return indices for delta actions.""" - return list(range(min(self.chunk_size, 16))) + model_action_horizon = ( + infer_groot_n1_7_action_horizon(self.base_model_path, self.embodiment_tag) or 40 + ) + return list(range(min(self.chunk_size, model_action_horizon))) + + @property + def drop_n_last_frames(self) -> int: + """Exclude episode tails that cannot supply a complete N1.7 action chunk.""" + return max(0, len(self.action_delta_indices) - 1) @property def reward_delta_indices(self) -> None: diff --git a/src/lerobot/policies/groot/eagle2_hg_model/configuration_eagle2_5_vl.py b/src/lerobot/policies/groot/eagle2_hg_model/configuration_eagle2_5_vl.py deleted file mode 100755 index 526b4f7a2..000000000 --- a/src/lerobot/policies/groot/eagle2_hg_model/configuration_eagle2_5_vl.py +++ /dev/null @@ -1,135 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# 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. - - -import copy - -from transformers.configuration_utils import PretrainedConfig -from transformers.models.llama.configuration_llama import LlamaConfig -from transformers.models.qwen2.configuration_qwen2 import Qwen2Config -from transformers.models.qwen3.configuration_qwen3 import Qwen3Config -from transformers.models.siglip.configuration_siglip import SiglipVisionConfig -from transformers.utils import logging - -logger = logging.get_logger(__name__) - - -class Eagle25VLConfig(PretrainedConfig): - model_type = "eagle_2_5_vl" - is_composition = True - sub_configs = {"vision_config": SiglipVisionConfig, "text_config": Qwen2Config} - - def __init__( - self, - vision_config=None, - text_config=None, - use_backbone_lora=0, - use_llm_lora=0, - pad2square=False, - select_layer=-4, - force_image_size=None, - downsample_ratio=0.5, - template=None, - dynamic_image_size=False, - use_thumbnail=False, - loss_version="v1", - min_dynamic_tiles=1, - max_dynamic_tiles=6, - mlp_checkpoint=False, - initializer_range=0.02, - _attn_implementation="flash_attention_2", - _attn_implementation_autoset=False, - llm_config=None, - image_token_index=None, - use_pixel_shuffle=True, - mlp_connector_layers=2, - **kwargs, - ): - super().__init__(**kwargs) - - if vision_config is None: - vision_config = {"model_type": "siglip_vision_model"} - logger.info("vision_config is None. Initializing the InternVisionConfig with default values.") - - if text_config is None: - text_config = {"architectures": ["Qwen2ForCausalLM"]} - logger.info( - "text_config is None. Initializing the LlamaConfig config with default values (`LlamaConfig`)." - ) - - if vision_config["model_type"] == "siglip_vision_model": - self.vision_config = SiglipVisionConfig(**vision_config) - else: - raise ValueError("Unsupported model_type: {}".format(vision_config["model_type"])) - - if text_config["architectures"][0] == "LlamaForCausalLM": - self.text_config = LlamaConfig(**text_config) - elif text_config["architectures"][0] == "Qwen2ForCausalLM": - self.text_config = Qwen2Config(**text_config) - elif text_config["architectures"][0] == "Qwen3ForCausalLM": - self.text_config = Qwen3Config(**text_config) - else: - raise ValueError("Unsupported architecture: {}".format(text_config["architectures"][0])) - self.use_backbone_lora = use_backbone_lora - self.use_llm_lora = use_llm_lora - self.mlp_checkpoint = mlp_checkpoint - self.pad2square = pad2square - self.select_layer = select_layer - self.force_image_size = force_image_size - self.downsample_ratio = downsample_ratio - self.template = template - self.dynamic_image_size = dynamic_image_size - self.use_thumbnail = use_thumbnail - self.loss_version = loss_version - self.initializer_range = initializer_range - self.min_dynamic_tiles = min_dynamic_tiles - self.max_dynamic_tiles = max_dynamic_tiles - self.tie_word_embeddings = self.text_config.tie_word_embeddings - self._attn_implementation = _attn_implementation - self._attn_implementation_autoset = _attn_implementation_autoset - self.image_token_index = image_token_index - self.use_pixel_shuffle = use_pixel_shuffle - self.mlp_connector_layers = mlp_connector_layers - logger.info(f"min_dynamic_tiles: {self.min_dynamic_tiles}") - logger.info(f"max_dynamic_tiles: {self.max_dynamic_tiles}") - - def to_dict(self): - """ - Serializes this instance to a Python dictionary. Override the default [`~PretrainedConfig.to_dict`]. - - Returns: - `Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance, - """ - output = copy.deepcopy(self.__dict__) - output["vision_config"] = self.vision_config.to_dict() - output["text_config"] = self.text_config.to_dict() - output["model_type"] = self.__class__.model_type - output["use_backbone_lora"] = self.use_backbone_lora - output["use_llm_lora"] = self.use_llm_lora - output["pad2square"] = self.pad2square - output["select_layer"] = self.select_layer - output["force_image_size"] = self.force_image_size - output["downsample_ratio"] = self.downsample_ratio - output["template"] = self.template - output["dynamic_image_size"] = self.dynamic_image_size - output["use_thumbnail"] = self.use_thumbnail - output["min_dynamic_tiles"] = self.min_dynamic_tiles - output["max_dynamic_tiles"] = self.max_dynamic_tiles - output["tie_word_embeddings"] = self.tie_word_embeddings - output["_attn_implementation"] = self._attn_implementation - output["_attn_implementation_autoset"] = self._attn_implementation_autoset - output["use_pixel_shuffle"] = self.use_pixel_shuffle - output["mlp_connector_layers"] = self.mlp_connector_layers - return output diff --git a/src/lerobot/policies/groot/eagle2_hg_model/image_processing_eagle2_5_vl_fast.py b/src/lerobot/policies/groot/eagle2_hg_model/image_processing_eagle2_5_vl_fast.py deleted file mode 100644 index 90e9dcecc..000000000 --- a/src/lerobot/policies/groot/eagle2_hg_model/image_processing_eagle2_5_vl_fast.py +++ /dev/null @@ -1,503 +0,0 @@ -# -------------------------------------------------------- -# NVIDIA -# Copyright (c) 2025 NVIDIA -# Licensed under The MIT License [see LICENSE for details] -# -------------------------------------------------------- - -from __future__ import annotations - -# copy from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llava_onevision/image_processing_llava_onevision_fast.py -from transformers.image_processing_utils import ( - BatchFeature, - get_patch_output_size, -) -from transformers.image_processing_utils_fast import ( - BaseImageProcessorFast, - ImagesKwargs, - group_images_by_shape, - reorder_images, -) -from transformers.image_utils import ( - IMAGENET_STANDARD_MEAN, # 0.5, 0.5, 0.5 - IMAGENET_STANDARD_STD, # 0.5, 0.5, 0.5 - ChannelDimension, - ImageInput, - PILImageResampling, - SizeDict, - get_image_size, - make_flat_list_of_images, - validate_kwargs, -) -from transformers.processing_utils import Unpack -from transformers.utils import ( - TensorType, - add_start_docstrings, - is_torch_available, - is_torchvision_v2_available, -) -from transformers.video_utils import VideoInput - -if is_torch_available(): - import torch -if is_torchvision_v2_available(): - from torchvision.transforms.v2 import functional as F # noqa: N812 - from transformers.image_utils import pil_torch_interpolation_mapping -else: - from torchvision.transforms import functional as F # noqa: N812 - - -def crop(img: torch.Tensor, left: int, top: int, right: int, bottom: int) -> torch.Tensor: - """Crop the given numpy array. - - Args: - img (torch.Tensor): Image to be cropped. Format should be (C, H, W). - left (int): The left coordinate of the crop box. - top (int): The top coordinate of the crop box. - right (int): The right coordinate of the crop box. - bottom (int): The bottom coordinate of the crop box. - - Returns: - torch.Tensor: Cropped image. - """ - if not isinstance(img, torch.Tensor): - raise TypeError(f"img should be torch.Tensor. Got {type(img)}") - - if img.ndim not in [2, 3]: - raise ValueError(f"Image should have 2 or 3 dimensions. Got {img.ndim}") - - img_height = img.shape[1] - img_width = img.shape[2] - if top < 0 or left < 0 or bottom > img_height or right > img_width: - raise ValueError("Crop coordinates out of bounds") - - if top >= bottom or left >= right: - raise ValueError("Invalid crop coordinates") - - return img[:, top:bottom, left:right] - - -class Eagle25VLFastImageProcessorKwargs(ImagesKwargs): - max_dynamic_tiles: int | None - min_dynamic_tiles: int | None - use_thumbnail: bool | None - pad_during_tiling: bool | None - do_pad: bool | None - - -@add_start_docstrings( - "Constructs a fast ConvNeXT image processor. Based on [`SiglipImageProcessor`] with incorporation of processing each video frame.", - # BASE_IMAGE_PROCESSOR_FAST_DOCSTRING, TODO: this was depreciated from transformers remove! - """ - image_grid_pinpoints (`List[List[int]]`, *optional*): - A list of possible resolutions to use for processing high resolution images. The best resolution is selected - based on the original size of the image. Can be overridden by `image_grid_pinpoints` in the `preprocess` - method. Not used for processing videos. - do_pad (`bool`, *optional*): - Whether to pad the image. If `True`, will pad the patch dimension of the images in the batch to the largest - number of patches in the batch. Padding will be applied to the bottom and right with zeros. - """, -) -class Eagle25VLImageProcessorFast(BaseImageProcessorFast): - resample = PILImageResampling.BICUBIC - image_mean = IMAGENET_STANDARD_MEAN - image_std = IMAGENET_STANDARD_STD - size = {"height": 448, "width": 448} - default_to_square = False - crop_size = None - do_resize = True - do_center_crop = None - do_rescale = True - do_normalize = True - do_convert_rgb = True - do_pad = True - max_dynamic_tiles = 12 - min_dynamic_tiles = 1 - use_thumbnail = True - pad_during_tiling = False - valid_kwargs = Eagle25VLFastImageProcessorKwargs - model_input_names = ["pixel_values_videos"] - - def __init__(self, **kwargs: Unpack[Eagle25VLFastImageProcessorKwargs]): - super().__init__(**kwargs) - - @add_start_docstrings( - # BASE_IMAGE_PROCESSOR_FAST_DOCSTRING_PREPROCESS, TODO: this was depreciated from transformers remove! - """ - max_dynamic_tiles (`int`, *optional*): - The maximum number of dynamic tiles to use for processing high resolution images. - min_dynamic_tiles (`int`, *optional*): - The minimum number of dynamic tiles to use for processing high resolution images. - use_thumbnail (`bool`, *optional*): - Whether to use a thumbnail for processing high resolution images. - pad_during_tiling (`bool`, *optional*): - Whether to pad the image during tiling. - do_pad (`bool`, *optional*): - Whether to pad the image. If `True`, will pad the patch dimension of the images in the batch to the largest - number of patches in the batch. Padding will be applied to the bottom and right with zeros. - """, - ) - - # NOTE(YL): we will overload the preprocess method to add the image_flags - # def preprocess( - # self, images: ImageInput, **kwargs: Unpack[Eagle25VLFastImageProcessorKwargs] - # ) -> BatchFeature: - # return super().preprocess(images, **kwargs) - - def _prepare_images_structure( - self, - images: ImageInput, - expected_ndims: int = 3, - ) -> ImageInput: - """ - Prepare the images structure for processing. - - Args: - images (`ImageInput`): - The input images to process. - expected_ndims (`int`, *optional*, defaults to 3): - Expected number of dimensions for the images (added for transformers >=4.53.0 compatibility). - - Returns: - `ImageInput`: The images with a valid nesting. - """ - return make_flat_list_of_images(images) - - def _resize_for_patching( - self, - image: torch.Tensor, - target_resolution: tuple, - interpolation: F.InterpolationMode, - input_data_format: ChannelDimension, - ) -> torch.Tensor: - """ - Resizes an image to a target resolution while maintaining aspect ratio. - - Args: - image ("torch.Tensor"): - The input image. - target_resolution (tuple): - The target resolution (height, width) of the image. - interpolation (`InterpolationMode`): - Resampling filter to use if resizing the image. - input_data_format (`ChannelDimension` or `str`): - The channel dimension format of the input image. - - Returns: - "torch.Tensor": The resized and padded image. - """ - new_height, new_width = get_patch_output_size(image, target_resolution, input_data_format) - - # Resize the image - resized_image = F.resize(image, (new_height, new_width), interpolation=interpolation) - - return resized_image - - def find_closest_aspect_ratio(self, aspect_ratio, target_ratios, width, height, image_size): - """ - previous version mainly focus on ratio. - We also consider area ratio here. - """ - best_factor = float("-inf") - best_ratio = (1, 1) - area = width * height - for ratio in target_ratios: - target_aspect_ratio = ratio[0] / ratio[1] - # ratio_diff = abs(aspect_ratio - target_aspect_ratio) - # area_ratio = (ratio[0] * ratio[1] * image_size * image_size) / area - """ - new area > 60% of original image area is enough. - """ - factor_based_on_area_n_ratio = min( - (ratio[0] * ratio[1] * image_size * image_size) / area, 0.6 - ) * min(target_aspect_ratio / aspect_ratio, aspect_ratio / target_aspect_ratio) - - if factor_based_on_area_n_ratio > best_factor: - best_factor = factor_based_on_area_n_ratio - best_ratio = ratio - - return best_ratio - - def _pad_for_patching( - self, image: torch.Tensor, target_resolution: tuple, input_data_format: ChannelDimension - ) -> torch.Tensor: - """ - Pad an image to a target resolution while maintaining aspect ratio. - """ - target_height, target_width = target_resolution - new_height, new_width = get_patch_output_size(image, target_resolution, input_data_format) - - paste_x = (target_width - new_width) // 2 - paste_y = (target_height - new_height) // 2 - - padded_image = F.pad(image, padding=[paste_x, paste_y, paste_x, paste_y]) - - return padded_image - - def _get_image_patches( - self, - image: torch.Tensor, - min_num: int, - max_num: int, - size: tuple, - tile_size: int, - use_thumbnail: bool, - interpolation: F.InterpolationMode, - pad_during_tiling: bool, - ) -> list[torch.Tensor]: - image_size = get_image_size(image, channel_dim=ChannelDimension.FIRST) - orig_height, orig_width = image_size - aspect_ratio = orig_width / orig_height - - # calculate the existing image aspect ratio - target_ratios = { - (i, j) - for n in range(min_num, max_num + 1) - for i in range(1, n + 1) - for j in range(1, n + 1) - if i * j <= max_num and i * j >= min_num - } - target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1]) - - # find the closest aspect ratio to the target - target_aspect_ratio = self.find_closest_aspect_ratio( - aspect_ratio, target_ratios, orig_width, orig_height, tile_size - ) - - # calculate the target width and height - target_width = tile_size * target_aspect_ratio[0] - target_height = tile_size * target_aspect_ratio[1] - blocks = target_aspect_ratio[0] * target_aspect_ratio[1] - if pad_during_tiling: - resized_image = self._resize_for_patching( - image, - (target_height, target_width), - interpolation=interpolation, - input_data_format=ChannelDimension.FIRST, - ) - padded_image = self._pad_for_patching( - resized_image, - (target_height, target_width), - input_data_format=ChannelDimension.FIRST, - ) - image_used_to_split = padded_image - else: - image_used_to_split = F.resize(image, (target_height, target_width), interpolation=interpolation) - - processed_tiles = [] - for i in range(blocks): - box = ( - (i % (target_width // tile_size)) * tile_size, - (i // (target_width // tile_size)) * tile_size, - ((i % (target_width // tile_size)) + 1) * tile_size, - ((i // (target_width // tile_size)) + 1) * tile_size, - ) - # split the image - split_img = crop(image_used_to_split, box[0], box[1], box[2], box[3]) - processed_tiles.append(split_img) - assert len(processed_tiles) == blocks - - if use_thumbnail and len(processed_tiles) != 1: - thumbnail_img = F.resize(image, (tile_size, tile_size), interpolation=interpolation) - processed_tiles.append(thumbnail_img) - - return processed_tiles - - def _pad_for_batching( - self, - pixel_values: list[torch.Tensor], - ) -> list[torch.Tensor]: - """ - Pads images on the `num_of_patches` dimension with zeros to form a batch of same number of patches. - - Args: - pixel_values (`List[torch.Tensor]`): - An array of pixel values of each images of shape (`batch_size`, `num_patches`, `image_in_3D`) - - Returns: - List[`torch.Tensor`]: The padded images. - """ - max_patch = max(len(x) for x in pixel_values) - pixel_values = [ - torch.nn.functional.pad(image, pad=[0, 0, 0, 0, 0, 0, 0, max_patch - image.shape[0]]) - for image in pixel_values - ] - - return pixel_values - - def _preprocess( - self, - images: list[torch.Tensor], - do_resize: bool, - size: SizeDict, - max_dynamic_tiles: int, - min_dynamic_tiles: int, - use_thumbnail: bool, - pad_during_tiling: bool, - interpolation: F.InterpolationMode | None, - do_center_crop: bool, - crop_size: SizeDict, - do_rescale: bool, - rescale_factor: float, - do_normalize: bool, - image_mean: float | list[float] | None, - image_std: float | list[float] | None, - do_pad: bool, - return_tensors: str | TensorType | None, - pad_size: SizeDict | None = None, # Added for transformers >=4.53.0 compatibility - disable_grouping: bool | None = None, # Added for transformers >=4.53.0 compatibility - ) -> BatchFeature: - processed_images = [] - image_sizes = [] - # Determine the size tuple - if size and size.height and size.width: - size_tuple = (size.height, size.width) - else: - size_tuple = (size.shortest_edge, size.shortest_edge) - - # Determine the patch size - if crop_size and crop_size.height: - tile_size = crop_size.height - elif size and size.height: - tile_size = size.height - else: - tile_size = size.shortest_edge - - for image in images: - image_patches = self._get_image_patches( - image, - min_num=min_dynamic_tiles, - max_num=max_dynamic_tiles, - size=size_tuple, - tile_size=tile_size, - use_thumbnail=use_thumbnail, - interpolation=interpolation, - pad_during_tiling=pad_during_tiling, - ) - - # Group images by size for batched processing - processed_image_patches_grouped = {} - # Added for transformers >=4.53.0 compatibility - grouped_image_patches, grouped_image_patches_index = group_images_by_shape( - image_patches, - disable_grouping=disable_grouping, - ) - - for shape, stacked_image_patches in grouped_image_patches.items(): - if do_resize: - stacked_image_patches = self.resize( - image=stacked_image_patches, - size=size, - interpolation=interpolation, - ) - if do_center_crop: - stacked_image_patches = self.center_crop(stacked_image_patches, crop_size) - # Fused rescale and normalize - stacked_image_patches = self.rescale_and_normalize( - stacked_image_patches, - do_rescale, - rescale_factor, - do_normalize, - image_mean, - image_std, - ) - processed_image_patches_grouped[shape] = stacked_image_patches - processed_image_patches = reorder_images( - processed_image_patches_grouped, grouped_image_patches_index - ) - processed_image_patches = ( - torch.stack(processed_image_patches, dim=0) if return_tensors else processed_image_patches - ) - processed_images.append(processed_image_patches) - image_sizes.append(get_image_size(image, ChannelDimension.FIRST)) - - if do_pad: - processed_images = self._pad_for_batching(processed_images) - - # processed_images = torch.stack(processed_images, dim=0) if return_tensors else processed_images - processed_images = torch.cat(processed_images, dim=0) if return_tensors else processed_images - return BatchFeature( - data={"pixel_values": processed_images, "image_sizes": image_sizes}, - tensor_type=return_tensors, - ) - - def preprocess( - self, - images: ImageInput, - videos: VideoInput = None, - **kwargs: Unpack[Eagle25VLFastImageProcessorKwargs], - ) -> BatchFeature: - validate_kwargs( - captured_kwargs=kwargs.keys(), - valid_processor_keys=self.valid_kwargs.__annotations__.keys(), - ) - # Set default kwargs from self. This ensures that if a kwarg is not provided - # by the user, it gets its default value from the instance, or is set to None. - for kwarg_name in self.valid_kwargs.__annotations__: - kwargs.setdefault(kwarg_name, getattr(self, kwarg_name, None)) - - # Extract parameters that are only used for preparing the input images - do_convert_rgb = kwargs.pop("do_convert_rgb") - input_data_format = kwargs.pop("input_data_format") - device = kwargs.pop("device") - # Prepare input images - # transformers >= 4.53.0: uses _prepare_image_like_inputs instead of _prepare_input_images - if images is not None: - images = self._prepare_image_like_inputs( - images=images, - do_convert_rgb=do_convert_rgb, - input_data_format=input_data_format, - device=device, - ) - - if videos is not None: - videos = self._prepare_image_like_inputs( - images=videos, - do_convert_rgb=do_convert_rgb, - input_data_format=input_data_format, - device=device, - ) - - # Update kwargs that need further processing before being validated - kwargs = self._further_process_kwargs(**kwargs) - - # Validate kwargs - self._validate_preprocess_kwargs(**kwargs) - - # torch resize uses interpolation instead of resample - # Added for transformers >=4.53.0 compatibility - resample = kwargs.pop("resample", self.resample) - kwargs["interpolation"] = ( - pil_torch_interpolation_mapping[resample] - if isinstance(resample, PILImageResampling | int) - else resample - ) - - # Filter kwargs to only include those accepted by _preprocess - valid_preprocess_kwargs = { - "do_resize", - "size", - "max_dynamic_tiles", - "min_dynamic_tiles", - "use_thumbnail", - "pad_during_tiling", - "interpolation", - "do_center_crop", - "crop_size", - "do_rescale", - "rescale_factor", - "do_normalize", - "image_mean", - "image_std", - "do_pad", - "return_tensors", - "pad_size", - "disable_grouping", - } - filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_preprocess_kwargs} - if images is not None: - return self._preprocess(images, **filtered_kwargs) - elif videos is not None: - return self._preprocess(videos, **filtered_kwargs) - - -__all__ = ["Eagle25VLImageProcessorFast"] diff --git a/src/lerobot/policies/groot/eagle2_hg_model/modeling_eagle2_5_vl.py b/src/lerobot/policies/groot/eagle2_hg_model/modeling_eagle2_5_vl.py deleted file mode 100755 index 6e5532ea4..000000000 --- a/src/lerobot/policies/groot/eagle2_hg_model/modeling_eagle2_5_vl.py +++ /dev/null @@ -1,396 +0,0 @@ -# -------------------------------------------------------- -# NVIDIA -# Copyright (c) 2025 NVIDIA -# Licensed under The MIT License [see LICENSE for details] -# -------------------------------------------------------- - -import inspect - -import torch -import torch.utils.checkpoint as cp -from peft import LoraConfig, get_peft_model -from torch import nn -from torch.nn import CrossEntropyLoss -from transformers import GenerationConfig -from transformers.generation import GenerationMixin -from transformers.modeling_outputs import CausalLMOutputWithPast -from transformers.modeling_utils import PreTrainedModel -from transformers.models.llama.modeling_llama import LlamaForCausalLM -from transformers.models.qwen2.modeling_qwen2 import Qwen2ForCausalLM -from transformers.models.qwen3.modeling_qwen3 import Qwen3ForCausalLM -from transformers.models.siglip.modeling_siglip import SiglipVisionModel -from transformers.utils import add_start_docstrings, logging - -from .configuration_eagle2_5_vl import Eagle25VLConfig - -logger = logging.get_logger(__name__) - - -# copy from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llava_onevision/modeling_llava_onevision.py#L241C1-L280C1 -EAGLE2_5_VL_START_DOCSTRING = r""" - This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the - library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads - etc.) - - This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. - Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage - and behavior. - - Parameters: - config ([`Eagle25VLConfig`]): - Model configuration class with all the parameters of the model. Initializing with a config file does not - load the weights associated with the model, only the configuration. Check out the - [`~PreTrainedModel.from_pretrained`] method to load the model weights. -""" - - -@add_start_docstrings( - "The bare Eagle2_5_VL Model outputting raw hidden-states without any specific head on top.", - EAGLE2_5_VL_START_DOCSTRING, -) -class Eagle25VLPreTrainedModel(PreTrainedModel): - config_class = Eagle25VLConfig - base_model_prefix = "model" - main_input_name = "input_ids" - supports_gradient_checkpointing = True - _no_split_modules = [ - "Qwen2DecoderLayer", - "LlamaDecoderLayer", - "Siglip2EncoderLayer", - "SiglipEncoderLayer", - ] - _skip_keys_device_placement = "past_key_values" - _supports_flash_attn = True - _supports_flash_attn_2 = True - _supports_cache_class = True - _supports_static_cache = True - _supports_quantized_cache = True - _supports_sdpa = True - - def _init_weights(self, module): - std = self.config.initializer_range - if isinstance(module, nn.Linear | nn.Conv2d): - module.weight.data.normal_(mean=0.0, std=std) - if module.bias is not None: - module.bias.data.zero_() - elif isinstance(module, nn.Embedding): - module.weight.data.normal_(mean=0.0, std=std) - if module.padding_idx is not None: - module.weight.data[module.padding_idx].zero_() - - -class Eagle25VLForConditionalGeneration(Eagle25VLPreTrainedModel, GenerationMixin): - config_class = Eagle25VLConfig - - def __init__(self, config: Eagle25VLConfig, vision_model=None, language_model=None): - super().__init__(config) - - image_size = config.force_image_size or config.vision_config.image_size - patch_size = config.vision_config.patch_size - self.patch_size = patch_size - if config.use_pixel_shuffle: - self.num_image_token = int((image_size // patch_size) ** 2 * (config.downsample_ratio**2)) - else: - self.num_image_token = int((image_size // patch_size) ** 2) - - self.select_layer = config.select_layer - self.downsample_ratio = config.downsample_ratio - self.loss_version = config.loss_version - self.mlp_checkpoint = config.mlp_checkpoint - self.use_pixel_shuffle = config.use_pixel_shuffle - self.mlp_connector_layers = config.mlp_connector_layers - logger.info(f"num_image_token: {self.num_image_token}") - logger.info(f"mlp_checkpoint: {self.mlp_checkpoint}") - if vision_model is not None: - self.vision_model = vision_model - else: - if config.vision_config.model_type == "siglip_vision_model": - config.vision_config._attn_implementation = "flash_attention_2" - self.vision_model = SiglipVisionModel(config.vision_config) - else: - raise NotImplementedError(f"{config.vision_config.model_type} is not implemented.") - - if language_model is not None: - self.language_model = language_model - else: - if config.text_config.architectures[0] == "LlamaForCausalLM": - self.language_model = LlamaForCausalLM(config.text_config) - elif config.text_config.architectures[0] == "Phi3ForCausalLM": - raise NotImplementedError("Phi3 is not implemented.") - # self.language_model = Phi3ForCausalLM(config.text_config) - elif config.text_config.architectures[0] == "Qwen2ForCausalLM": - assert config.text_config._attn_implementation == "flash_attention_2", ( - f"Qwen2 must use flash_attention_2 but got {config.text_config._attn_implementation}" - ) - self.language_model = Qwen2ForCausalLM(config.text_config) - elif config.text_config.architectures[0] == "Qwen3ForCausalLM": - self.language_model = Qwen3ForCausalLM(config.text_config) - else: - raise NotImplementedError(f"{config.text_config.architectures[0]} is not implemented.") - - vit_hidden_size = config.vision_config.hidden_size - llm_hidden_size = config.text_config.hidden_size - - if config.mlp_connector_layers == 2: - self.mlp1 = nn.Sequential( - nn.LayerNorm(vit_hidden_size * int(1 / self.downsample_ratio) ** 2), - nn.Linear(vit_hidden_size * int(1 / self.downsample_ratio) ** 2, llm_hidden_size), - nn.GELU(), - nn.Linear(llm_hidden_size, llm_hidden_size), - ) - elif config.mlp_connector_layers == 1 and config.use_pixel_shuffle: - self.mlp1 = nn.Sequential( - nn.Linear(vit_hidden_size * int(1 / self.downsample_ratio) ** 2, llm_hidden_size), - ) - elif config.mlp_connector_layers == 1 and not config.use_pixel_shuffle: - self.mlp1 = nn.Sequential( - nn.Linear(vit_hidden_size, llm_hidden_size), - ) - else: - raise NotImplementedError(f"{config.mlp_connector_layers} is not implemented.") - - self.image_token_index = config.image_token_index - self.neftune_alpha = None - - if config.use_backbone_lora: - self.wrap_backbone_lora(r=config.use_backbone_lora, lora_alpha=2 * config.use_backbone_lora) - - self.use_llm_lora = config.use_llm_lora - if config.use_llm_lora: - self.wrap_llm_lora(r=config.use_llm_lora, lora_alpha=2 * config.use_llm_lora) - - self.check_forward_kwargs() - - def check_forward_kwargs(self): - # We intentionally avoid using **kwargs in forward because Hugging Face Transformers - # has special handling for functions with **kwargs parameters that would affect - # how our model is processed during training and inference. - forward_params = inspect.signature(self.forward).parameters - assert not any(k.kind == inspect.Parameter.VAR_KEYWORD for k in forward_params.values()) - - def wrap_backbone_lora(self, r=128, lora_alpha=256, lora_dropout=0.05): - lora_config = LoraConfig( - r=r, - target_modules=[ - "self_attn.q_proj", - "self_attn.k_proj", - "self_attn.v_proj", - "self_attn.out_proj", - "mlp.fc1", - "mlp.fc2", - ], - lora_alpha=lora_alpha, - lora_dropout=lora_dropout, - ) - self.vision_model = get_peft_model(self.vision_model, lora_config) - self.vision_model.print_trainable_parameters() - - def wrap_llm_lora(self, r=128, lora_alpha=256, lora_dropout=0.05): - lora_config = LoraConfig( - r=r, - target_modules=[ - "self_attn.q_proj", - "self_attn.k_proj", - "self_attn.v_proj", - "self_attn.o_proj", - "mlp.gate_proj", - "mlp.down_proj", - "mlp.up_proj", - ], - lora_alpha=lora_alpha, - lora_dropout=lora_dropout, - task_type="CAUSAL_LM", - ) - self.language_model = get_peft_model(self.language_model, lora_config) - self.language_model.enable_input_require_grads() - self.language_model.print_trainable_parameters() - self.use_llm_lora = True - - def forward( - self, - pixel_values: torch.FloatTensor, - input_ids: torch.LongTensor = None, - attention_mask: torch.Tensor | None = None, - position_ids: torch.LongTensor | None = None, - image_flags: torch.LongTensor | None = None, - past_key_values: list[torch.FloatTensor] | None = None, - labels: torch.LongTensor | None = None, - use_cache: bool | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - return_dict: bool | None = None, - num_tiles_list: list[torch.Tensor] | None = None, - ) -> tuple | CausalLMOutputWithPast: - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - input_embeds = self.language_model.get_input_embeddings()(input_ids) - - vit_embeds = self.extract_feature(pixel_values) - - if image_flags is not None: - image_flags = image_flags.view(-1) - vit_embeds = vit_embeds[image_flags == 1] - - b, n, c = input_embeds.shape - input_embeds = input_embeds.reshape(b * n, c) - - input_ids = input_ids.reshape(b * n) - selected = input_ids == self.image_token_index - try: - input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds.reshape(-1, c) - except Exception as e: - vit_embeds = vit_embeds.reshape(-1, c) - print( - f"warning: {e}, input_embeds[selected].shape={input_embeds[selected].shape}, " - f"vit_embeds.shape={vit_embeds.shape}" - ) - n_token = selected.sum() - input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds[:n_token] - - input_embeds = input_embeds.reshape(b, n, c) - - outputs = self.language_model( - inputs_embeds=input_embeds, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - ) - logits = outputs.logits - - loss = None - if labels is not None: - # Shift so that tokens < n predict n - shift_logits = logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - # Flatten the tokens - loss_fct = CrossEntropyLoss() - shift_logits = shift_logits.view(-1, self.language_model.config.vocab_size) - shift_labels = shift_labels.view(-1) - # Enable model parallelism - shift_labels = shift_labels.to(shift_logits.device) - loss = loss_fct(shift_logits, shift_labels) - - if not return_dict: - output = (logits,) + outputs[1:] - return (loss,) + output if loss is not None else output - - return CausalLMOutputWithPast( - loss=loss, - logits=logits, - past_key_values=outputs.past_key_values, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) - - def pixel_shuffle(self, x, scale_factor=0.5): - n, w, h, c = x.size() - # N, W, H, C --> N, W, H * scale, C // scale - x = x.view(n, w, int(h * scale_factor), int(c / scale_factor)) - # N, W, H * scale, C // scale --> N, H * scale, W, C // scale - x = x.permute(0, 2, 1, 3).contiguous() - # N, H * scale, W, C // scale --> N, H * scale, W * scale, C // (scale ** 2) - x = x.view(n, int(h * scale_factor), int(w * scale_factor), int(c / (scale_factor * scale_factor))) - - x = x.permute(0, 2, 1, 3).contiguous() - return x - - def extract_feature(self, pixel_values): - if self.select_layer == -1: - vit_embeds = self.vision_model( - pixel_values=pixel_values, output_hidden_states=False, return_dict=True - ) - if hasattr(vit_embeds, "last_hidden_state"): - vit_embeds = vit_embeds.last_hidden_state - - else: - vit_embeds = self.vision_model( - pixel_values=pixel_values, output_hidden_states=True, return_dict=True - ).hidden_states[self.select_layer] - - if self.use_pixel_shuffle: - h = w = int(vit_embeds.shape[1] ** 0.5) - vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1) - vit_embeds = self.pixel_shuffle( - vit_embeds, scale_factor=self.downsample_ratio - ) # torch.Size([B, 1024, 1024]) -> torch.Size([B, 16, 16, 4096]) - vit_embeds = vit_embeds.reshape( - vit_embeds.shape[0], -1, vit_embeds.shape[-1] - ) # torch.Size([B, 16, 16, 4096]) -> torch.Size([B, 256, 4096]) - - if self.mlp_checkpoint and vit_embeds.requires_grad: - vit_embeds = cp.checkpoint(self.mlp1, vit_embeds) - else: - vit_embeds = self.mlp1(vit_embeds) - - return vit_embeds - - @torch.no_grad() - def generate( - self, - pixel_values: torch.FloatTensor | None = None, - input_ids: torch.FloatTensor | None = None, - attention_mask: torch.LongTensor | None = None, - visual_features: torch.FloatTensor | None = None, - generation_config: GenerationConfig | None = None, - output_hidden_states: bool | None = None, - image_sizes: list[tuple[int, int]] | None = None, - **generate_kwargs, - ) -> torch.LongTensor: - if pixel_values is not None: - if visual_features is not None: - vit_embeds = visual_features - else: - vit_embeds = self.extract_feature(pixel_values) - - input_embeds = self.language_model.get_input_embeddings()(input_ids) - b, n, c = input_embeds.shape - input_embeds = input_embeds.reshape(b * n, c) - - input_ids = input_ids.reshape(b * n) - selected = input_ids == self.config.image_token_index - assert selected.sum() != 0 - input_embeds[selected] = vit_embeds.reshape(-1, c).to(input_embeds.device) - - input_embeds = input_embeds.reshape(b, n, c) - else: - input_embeds = self.language_model.get_input_embeddings()(input_ids) - - if "use_cache" not in generate_kwargs: - generate_kwargs["use_cache"] = True - - outputs = self.language_model.generate( - inputs_embeds=input_embeds, - attention_mask=attention_mask, - generation_config=generation_config, - output_hidden_states=output_hidden_states, - **generate_kwargs, - ) - - return outputs - - # Copied from transformers.models.llava_next.modeling_llava_next.LlavaNextForConditionalGeneration.get_input_embeddings - def get_input_embeddings(self): - return self.language_model.get_input_embeddings() - - # Copied from transformers.models.llava_next.modeling_llava_next.LlavaNextForConditionalGeneration.set_input_embeddings - def set_input_embeddings(self, value): - self.language_model.set_input_embeddings(value) - - # Copied from transformers.models.llava_next.modeling_llava_next.LlavaNextForConditionalGeneration.get_output_embeddings - def get_output_embeddings(self): - return self.language_model.get_output_embeddings() - - # Copied from transformers.models.llava_next.modeling_llava_next.LlavaNextForConditionalGeneration.set_output_embeddings - def set_output_embeddings(self, new_embeddings): - self.language_model.set_output_embeddings(new_embeddings) - - # Copied from transformers.models.llava_next.modeling_llava_next.LlavaNextForConditionalGeneration.set_decoder - def set_decoder(self, decoder): - self.language_model.set_decoder(decoder) - - # Copied from transformers.models.llava_next.modeling_llava_next.LlavaNextForConditionalGeneration.get_decoder - def get_decoder(self): - return self.language_model.get_decoder() diff --git a/src/lerobot/policies/groot/eagle2_hg_model/processing_eagle2_5_vl.py b/src/lerobot/policies/groot/eagle2_hg_model/processing_eagle2_5_vl.py deleted file mode 100755 index b36e70c47..000000000 --- a/src/lerobot/policies/groot/eagle2_hg_model/processing_eagle2_5_vl.py +++ /dev/null @@ -1,541 +0,0 @@ -# Copyright 2024 The HuggingFace Inc. team. -# -# 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. -""" -Processor class for Eagle25VL. -copy from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llava_onevision/processing_llava_onevision.py -""" - -import base64 -import os -import re -from io import BytesIO - -import requests -import torch -from PIL import Image -from transformers.feature_extraction_utils import BatchFeature -from transformers.image_utils import ImageInput -from transformers.processing_utils import ProcessingKwargs, ProcessorMixin, Unpack -from transformers.tokenization_utils_base import PreTokenizedInput, TextInput -from transformers.utils import logging -from transformers.video_utils import VideoInput - -logger = logging.get_logger(__name__) - - -FRAME_FACTOR = 2 -FPS = 2.0 -FPS_MIN_FRAMES = 4 -FPS_MAX_FRAMES = 256 - - -def to_rgb(pil_image: Image.Image) -> Image.Image: - if pil_image.mode == "RGBA": - white_background = Image.new("RGB", pil_image.size, (255, 255, 255)) - white_background.paste(pil_image, mask=pil_image.split()[3]) # Use alpha channel as mask - return white_background - else: - return pil_image.convert("RGB") - - -def fetch_image(ele: dict[str, str | Image.Image]) -> Image.Image: - image = ele["image"] if "image" in ele else ele["image_url"] - image_obj = None - if isinstance(image, Image.Image): - image_obj = image - elif image.startswith("http://") or image.startswith("https://"): - response = requests.get(image, stream=True, timeout=10) - image_obj = Image.open(BytesIO(response.content)) - elif image.startswith("file://"): - image_obj = Image.open(image[7:]) - elif image.startswith("data:image"): - if "base64," in image: - _, base64_data = image.split("base64,", 1) - data = base64.b64decode(base64_data) - image_obj = Image.open(BytesIO(data)) - else: - image_obj = Image.open(image) - if image_obj is None: - raise ValueError( - f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}" - ) - image = to_rgb(image_obj) - if "scale_factor" in ele: - scale_factor = ele["scale_factor"] - image = image.resize((image.width * scale_factor, image.height * scale_factor), Image.BILINEAR) - return image - - -class Eagle25VLProcessorKwargs(ProcessingKwargs, total=False): - # see processing_utils.ProcessingKwargs documentation for usage. - _defaults = { - "text_kwargs": { - "padding": False, - }, - "images_kwargs": {}, - "videos_kwargs": {"max_dynamic_tiles": 1}, - } - - -class Eagle25VLProcessor(ProcessorMixin): - r""" - Constructs a Eagle25VL processor which wraps a Eagle25VL video processor, Eagle25VL image processor and a Eagle25VL tokenizer into a single processor. - - [`Eagle25VLProcessor`] offers all the functionalities of [`Eagle25VLVideoProcessor`], [`Eagle25VLImageProcessor`] and [`Eagle25VLTokenizer`]. See the - [`~Eagle25VLVideoProcessor.__call__`], [`~Eagle25VLProcessor.__call__`] and [`~Eagle25VLProcessor.decode`] for more information. - - Args: - image_processor ([`LlavaOnevisionImageProcessor`], *optional*): - The image processor is a required input. - tokenizer ([`LlamaTokenizerFast`], *optional*): - The tokenizer is a required input. - num_image_tokens (`int`, *optional*): - Number of image tokens for one imagethat will be returned by vision tower. - vision_feature_select_strategy (`str`, *optional*): - The feature selection strategy used to select the vision feature from the vision backbone. - Should be same as in model's config - chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages - in a chat into a tokenizable string. - image_token (`str`, *optional*, defaults to `""`): - Special token used to denote image location. - video_token (`str`, *optional*, defaults to `"