From e6d7eea88a8ac0737ca457c5db35cf21e0dfc569 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Wed, 29 Jul 2026 16:58:02 +0200 Subject: [PATCH] feat(g05): make policy runtime native --- docs/source/g05.mdx | 67 +- pyproject.toml | 7 +- src/lerobot/policies/g05/action_codec_g05.py | 671 +++++++++ src/lerobot/policies/g05/configuration_g05.py | 6 +- src/lerobot/policies/g05/modeling_g05.py | 33 +- src/lerobot/policies/g05/native_g05.py | 1215 +++++++++++++++++ src/lerobot/policies/g05/processing_g05.py | 353 +++++ tests/policies/g05/test_action_codec_g05.py | 83 ++ uv.lock | 5 +- 9 files changed, 2372 insertions(+), 68 deletions(-) create mode 100644 src/lerobot/policies/g05/action_codec_g05.py create mode 100644 src/lerobot/policies/g05/native_g05.py create mode 100644 src/lerobot/policies/g05/processing_g05.py create mode 100644 tests/policies/g05/test_action_codec_g05.py diff --git a/docs/source/g05.mdx b/docs/source/g05.mdx index d0b9bc85a..24fa3b358 100644 --- a/docs/source/g05.mdx +++ b/docs/source/g05.mdx @@ -9,15 +9,16 @@ conditioned on the same post-reasoning KV state. Transformers includes the native multimodal Qwen3.5 backbone, vision tower, and processor. G0.5 is not a stock `Qwen3_5ForConditionalGeneration` checkpoint, however: it adds the proprioception/action path, action expert, flow-matching -head, ActionCodec, and unified CoT/action decode. The current integration -therefore keeps the pinned G0.5 author package as the model backend. +head, ActionCodec, and unified CoT/action decode. LeRobot implements those G0.5 +components natively and loads converted checkpoints without the OpenGalaxea +Python package, Hydra, or OmegaConf. > [!WARNING] > G0.5 code and checkpoints use the > [G0.5 Community License](https://huggingface.co/OpenGalaxea/G05/blob/main/licenses/LICENSE-G0.5), -> including non-commercial restrictions. LeRobot does not vendor the author model, -> download gated files, or imply that Apache-2.0 applies to those materials. Accept -> the license yourself and use a private or local checkpoint. +> including non-commercial restrictions. LeRobot does not redistribute checkpoint +> weights, download gated files, or imply that Apache-2.0 applies to those materials. +> Accept the license yourself and use a private or local checkpoint. ## Supported checkpoint contracts @@ -29,8 +30,8 @@ therefore keeps the pinned G0.5 author package as the model backend. | `g05-so101` | Flow or AR ActionCodec + native CoT | right arm joints 6 → 20D grouped layout | exterior, optional left + right wrist | 32 | 16 | stepwise q01/q99 | Each packaged checkpoint stores the resolved model and processor configuration, -ActionCodec metadata, statistics, exact prompt template, source revision, and -license. Loading rejects a different head, horizon, processor mode, or +ActionCodec metadata, statistics, exact prompt template, and license. Loading +rejects a different head, horizon, processor mode, or normalization contract. The converted checkpoints are private under the LeRobot organization: @@ -51,23 +52,12 @@ Atomic-4 dataset; reusing the R1 Lite statistics would be incorrect. ## Install -Install LeRobot's small config dependency: +Install LeRobot with the G0.5 Transformers dependency: ```bash uv sync --extra g05 --extra test ``` -Then clone the audited author source and install it only after reviewing and -accepting its license. The author package currently declares Python 3.10 while -LeRobot uses Python 3.12, so a compatible deployment environment or an upstream -Python-support update is required for real-model execution. - -```bash -git clone https://github.com/OpenGalaxea/GalaxeaVLA.git -git -C GalaxeaVLA checkout b34966f387dd2ae0f003143b81494afd9213e613 -export PYTHONPATH="/path/to/GalaxeaVLA/src:${PYTHONPATH}" -``` - The LeRobot organization hosts the prepared base, LIBERO, RoboTwin, and SO-101 checkpoints privately. Authenticate with `hf auth login` before loading them. SO-100 and SO-101 share the released `so100` embodiment token and six-joint @@ -124,9 +114,10 @@ the checkpoint's native System 2 CoT telemetry. ## Fine-tune with `lerobot-train` -G0.5 implements LeRobot's training surface: `forward` runs the author training -backend, the policy exposes the author VLM/vision/action optimizer groups, and -the checkpoint can be saved, resumed, and loaded by the normal LeRobot scripts. +G0.5 implements LeRobot's training surface natively: `forward` computes +assistant-token cross entropy and flow-matching loss, the policy exposes +VLM/vision/action optimizer groups, and the checkpoint can be saved, resumed, +and loaded by the normal LeRobot scripts. For example, fine-tune the private SO-101 checkpoint on a LeRobot dataset: ```bash @@ -157,7 +148,7 @@ Grounded VQA boxes are converted from pixel-space `xyxy` JSON using the source camera dimensions captured before image resizing, then serialized as G0.5 `` tokens. Joint samples preserve the released checkpoint's `BBox → Subtask → Action` order. The user/task conditioning tokens remain -masked; the author backend applies its language/action objective to the +masked; the native backend applies its language/action objective to the assistant sequence. Generate the required `subtask` and grounded `vqa` language columns with @@ -197,30 +188,26 @@ lerobot-train \ CPU unit tests cover factory loading, config incompatibilities, prompt pass-through, LIBERO and `atomic_4` mappings, padding masks, inverse action projection, a finite -forward/backward/update, author optimizer-group wiring, and save/reload parity: +forward/backward/update, optimizer-group wiring, and save/reload parity: ```bash uv run pytest tests/policies/g05 tests/runtime/test_g05_adapter.py -q uv run ruff check src/lerobot/policies/g05 tests/policies/g05 ``` -An RTX 5090 smoke test loaded the real private `g05_libero` checkpoint and -completed a batch-size-one BF16 forward, backward, gradient clip, and AdamW -step. It produced finite loss `2.77356`, finite pre-clip gradient norm `54.38`, -all six author optimizer groups, and 23.59 GiB peak allocated CUDA memory. +The converted private `g05_base` checkpoint strict-loaded all 945 native model +tensors, and its ActionCodec sidecar strict-loaded all 208 tensors, with no +missing, unexpected, or shape-mismatched keys. Prompt token IDs and masks match +the released runtime exactly. ActionCodec code IDs match exactly; decoded values +differ only by normal floating-point noise. A batch-size-one System 1 flow smoke +produced a finite `[1, 32, 27]` action chunk on an RTX 5090. -The same RTX 5090 also loaded the private `g05_so101` checkpoint and completed -a real joint `BBox → Subtask → Action` forward and backward through the -recipe-driven path. It produced finite total loss `4.38332`, including non-zero -`ce_loss=4.04244` and `fm_loss=0.340881`, with 15.58 GiB peak allocated CUDA -memory. - -The converted private `g05_base` checkpoint strict-loaded all 946 mapped tensors -with no missing, duplicate, unexpected, or shape-mismatched keys. Author/LeRobot -parity was exact for images, tokens, and masks; normalized state/action inputs -matched within `1.2e-7`, and postprocessing within `1.8e-6`. A System 2 smoke -generated a native `Subtask:` trace and a finite `[1, 32, 27]` action chunk in -the same pass, using 11.02 GiB peak allocated memory. +The same native base checkpoint completed a real joint language/action forward +and backward on the RTX 5090. It produced finite `ce_loss=13.8141` and +`fm_loss=1.32647`, finite gradients for all 945 model tensors, all six optimizer +groups, and 19.32 GiB peak allocated CUDA memory. A System 2 smoke generated +`Subtask: grasp and lift the red cup with the right gripper` and a finite +same-pass `[1, 32, 27]` action in 1.07 seconds at 11.0 GiB peak. A 50-episode LIBERO/RoboTwin success-rate comparison additionally requires the matching simulator, task assets, reset seeds, and author evaluator; no task-level diff --git a/pyproject.toml b/pyproject.toml index 308156656..ab5f31a51 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -238,10 +238,9 @@ evo1 = ["lerobot[transformers-dep]"] 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]"] -# The G0.5 model implementation and weights carry Galaxea's separate community -# license and are intentionally not redistributed as an Apache LeRobot dependency. -# This extra supplies only the config runtime used by converted local checkpoints. -g05 = ["omegaconf>=2.3.0,<3.0.0"] +# Checkpoint weights retain their upstream licenses. The runtime itself is native +# LeRobot code built on Transformers' Qwen3.5 implementation. +g05 = ["transformers>=5.4.0,<5.6.0"] # Features async = ["lerobot[grpcio-dep]", "lerobot[matplotlib-dep]"] diff --git a/src/lerobot/policies/g05/action_codec_g05.py b/src/lerobot/policies/g05/action_codec_g05.py new file mode 100644 index 000000000..8b565aff5 --- /dev/null +++ b/src/lerobot/policies/g05/action_codec_g05.py @@ -0,0 +1,671 @@ +# 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. + +"""Native inference implementation of G0.5's ActionCodec sidecar.""" + +from __future__ import annotations + +import math +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import torch +import torch.nn.functional as functional +from torch import Tensor, nn + + +class _BlockDCT(nn.Module): + def __init__(self, block_size: int) -> None: + super().__init__() + self.block_size = block_size + frequency = torch.arange(block_size, dtype=torch.float32) + time = torch.arange(block_size, dtype=torch.float32) + basis = torch.cos(math.pi / block_size * (time + 0.5).unsqueeze(0) * frequency.unsqueeze(1)) + basis[0] *= math.sqrt(1 / block_size) + basis[1:] *= math.sqrt(2 / block_size) + self.register_buffer("basis", basis, persistent=False) + + def dct(self, values: Tensor) -> Tensor: + batch, horizon, dimension = values.shape + pad = (-horizon) % self.block_size + if pad: + values = functional.pad(values, (0, 0, 0, pad)) + blocks = values.shape[1] // self.block_size + values = values.reshape(batch * blocks, self.block_size, dimension) + transformed = torch.einsum("kn,bnd->bkd", self.basis.to(values), values) + return transformed.reshape(batch, blocks * self.block_size, dimension) + + def idct(self, values: Tensor, horizon: int) -> Tensor: + batch, padded_horizon, dimension = values.shape + blocks = padded_horizon // self.block_size + values = values.reshape(batch * blocks, self.block_size, dimension) + restored = torch.einsum("nk,bkd->bnd", self.basis.to(values), values) + return restored.reshape(batch, padded_horizon, dimension)[:, :horizon] + + +def _rotate_half(values: Tensor) -> Tensor: + first, second = values.chunk(2, dim=-1) + return torch.cat((-second, first), dim=-1) + + +class _CodecAttention(nn.Module): + def __init__(self, dimension: int, num_heads: int, head_dim: int, rope_base: int) -> None: + super().__init__() + self.num_heads = num_heads + self.head_dim = head_dim + inner_dim = num_heads * head_dim + self.to_qkv = nn.Linear(dimension, inner_dim * 3, bias=False) + self.to_out = nn.Linear(inner_dim, dimension, bias=False) + self.q_norm = nn.LayerNorm(head_dim, eps=1e-6) + self.k_norm = nn.LayerNorm(head_dim, eps=1e-6) + rope_dim = max(head_dim // 2, 32) + inverse = 1 / (rope_base ** (torch.arange(0, rope_dim, 2, dtype=torch.float32) / rope_dim)) + self.register_buffer("_inverse_frequency", inverse, persistent=False) + + def forward(self, hidden_states: Tensor) -> Tensor: + batch, sequence_length, _ = hidden_states.shape + query, key, value = self.to_qkv(hidden_states).chunk(3, dim=-1) + + def heads(values: Tensor) -> Tensor: + return values.view(batch, sequence_length, self.num_heads, self.head_dim).transpose(1, 2) + + query, key, value = (heads(values) for values in (query, key, value)) + query, key = self.q_norm(query), self.k_norm(key) + time = torch.arange(sequence_length, device=hidden_states.device, dtype=torch.float32) + phase = torch.outer(time, self._inverse_frequency.to(hidden_states.device)) + phase = torch.cat((phase, phase), dim=-1).to(hidden_states.dtype)[None, None] + cosine, sine = phase.cos(), phase.sin() + rotary_dim = cosine.shape[-1] + query_rotary, query_pass = query[..., :rotary_dim], query[..., rotary_dim:] + key_rotary, key_pass = key[..., :rotary_dim], key[..., rotary_dim:] + query = torch.cat((query_rotary * cosine + _rotate_half(query_rotary) * sine, query_pass), dim=-1) + key = torch.cat((key_rotary * cosine + _rotate_half(key_rotary) * sine, key_pass), dim=-1) + attended = functional.scaled_dot_product_attention(query, key, value) + attended = attended.transpose(1, 2).reshape(batch, sequence_length, -1) + return self.to_out(attended) + + +class _CodecFFN(nn.Module): + def __init__(self, dimension: int, multiplier: float) -> None: + super().__init__() + inner_dim = int(dimension * multiplier) + self.w_up = nn.Linear(dimension, inner_dim * 2, bias=False) + self.w_down = nn.Linear(inner_dim, dimension, bias=False) + + def forward(self, hidden_states: Tensor) -> Tensor: + value, gate = self.w_up(hidden_states).chunk(2, dim=-1) + return self.w_down(value * functional.gelu(gate)) + + +class _CodecTransformerLayer(nn.Module): + def __init__(self, dimension: int, config: Mapping[str, Any]) -> None: + super().__init__() + self.ls1 = nn.Parameter(torch.empty(dimension)) + self.ls2 = nn.Parameter(torch.empty(dimension)) + self.norm1 = nn.LayerNorm(dimension, eps=1e-6) + self.attn = _CodecAttention( + dimension, + int(config["num_heads"]), + int(config["dim_heads"]), + int(config["rope_base"]), + ) + self.norm2 = nn.LayerNorm(dimension, eps=1e-6) + self.ffn = _CodecFFN(dimension, float(config["ffn_mult"])) + + def forward(self, hidden_states: Tensor) -> Tensor: + hidden_states = hidden_states + self.attn(self.norm1(hidden_states)) * self.ls1 + return hidden_states + self.ffn(self.norm2(hidden_states)) * self.ls2 + + +class _CodecDownBlock(nn.Module): + def __init__( + self, + input_channels: int, + output_channels: int, + stride: tuple[int, int], + depth: int, + config: Mapping[str, Any], + ) -> None: + super().__init__() + stride_h, stride_a = stride + if stride_h > 1 or input_channels != output_channels: + kernel_h = 2 * stride_h if stride_h > 1 else 1 + self.conv = nn.Conv2d( + input_channels, + output_channels, + kernel_size=(kernel_h, 1), + stride=(stride_h, stride_a), + padding=(kernel_h // 2 - int(stride_h > 1), 0), + ) + else: + self.conv = nn.Identity() + self.transformer_layers = nn.ModuleList( + [_CodecTransformerLayer(output_channels, config) for _ in range(depth)] + ) + + def forward(self, hidden_states: Tensor) -> Tensor: + hidden_states = self.conv(hidden_states) + batch, channels, height, action_dim = hidden_states.shape + sequence = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * action_dim, channels) + for layer in self.transformer_layers: + sequence = layer(sequence) + return sequence.reshape(batch, height, action_dim, channels).permute(0, 3, 1, 2) + + +class _CodecUpBlock(nn.Module): + def __init__( + self, + input_channels: int, + output_channels: int, + stride: tuple[int, int], + depth: int, + config: Mapping[str, Any], + ) -> None: + super().__init__() + self.transformer_layers = nn.ModuleList( + [_CodecTransformerLayer(input_channels, config) for _ in range(depth)] + ) + stride_h, stride_a = stride + if stride_h > 1 or input_channels != output_channels: + kernel_h = 2 * stride_h if stride_h > 1 else 1 + self.conv = nn.ConvTranspose2d( + input_channels, + output_channels, + kernel_size=(kernel_h, 1), + stride=(stride_h, stride_a), + padding=(kernel_h // 2 - int(stride_h > 1), 0), + ) + else: + self.conv = nn.Identity() + + def forward(self, hidden_states: Tensor) -> Tensor: + batch, channels, height, action_dim = hidden_states.shape + sequence = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * action_dim, channels) + for layer in self.transformer_layers: + sequence = layer(sequence) + hidden_states = sequence.reshape(batch, height, action_dim, channels).permute(0, 3, 1, 2) + return self.conv(hidden_states) + + +class _CodecEncoder(nn.Module): + def __init__(self, config: Mapping[str, Any]) -> None: + super().__init__() + base = int(config["encoder_channels"]) + channel_dims = [base * int(multiplier) for multiplier in config["c_mults"]] + dims = [base] + channel_dims + self.blocks = nn.ModuleList( + [ + _CodecDownBlock( + dims[index], + dims[index + 1], + tuple(stride), + int(config["transformer_depths"][index]), + config, + ) + for index, stride in enumerate(config["strides"]) + ] + ) + self.out_proj = nn.Conv2d(dims[-1], int(config["latent_dim"]), kernel_size=1) + + def forward(self, hidden_states: Tensor) -> Tensor: + for block in self.blocks: + hidden_states = block(hidden_states) + return self.out_proj(hidden_states) + + +class _CodecDecoder(nn.Module): + def __init__(self, config: Mapping[str, Any]) -> None: + super().__init__() + base = int(config["encoder_channels"]) + channel_dims = [base * int(multiplier) for multiplier in config["c_mults"]] + dims = [base] + channel_dims + self.in_proj = nn.Conv2d(int(config["latent_dim"]), dims[-1], kernel_size=1) + self.blocks = nn.ModuleList( + [ + _CodecUpBlock( + input_channels, + output_channels, + tuple(stride), + int(depth), + config, + ) + for stride, depth, input_channels, output_channels in zip( + reversed(config["strides"]), + reversed(config["transformer_depths"]), + reversed(dims[1:]), + reversed(dims[:-1]), + strict=True, + ) + ] + ) + + def forward(self, hidden_states: Tensor) -> Tensor: + hidden_states = self.in_proj(hidden_states) + for block in self.blocks: + hidden_states = block(hidden_states) + return hidden_states + + +class _CodecQuantizer(nn.Module): + def __init__(self, config: Mapping[str, Any]) -> None: + super().__init__() + input_dim = int(config["latent_dim"]) + codebook_dim = int(config["codebook_dim"]) + codebook_size = int(config["codebook_size"]) + self.input_dim = input_dim + self.in_proj = nn.Linear(input_dim, codebook_dim, bias=False) + self.out_proj = nn.Linear(codebook_dim, input_dim, bias=False) + self.register_buffer("codebook", torch.zeros(codebook_size, codebook_dim)) + self.register_buffer("embed_avg", torch.zeros(codebook_size, codebook_dim)) + self.register_buffer("cluster_size", torch.zeros(codebook_size)) + self.register_buffer("inited", torch.tensor(False)) + + def encode(self, values: Tensor) -> tuple[Tensor, Tensor]: + projected = self.in_proj(values.transpose(1, 2)) + flat = projected.reshape(-1, projected.shape[-1]).float() + codebook = self.codebook.float() + distances = ( + flat.square().sum(dim=1, keepdim=True) + - 2 * flat @ codebook.t() + + codebook.square().sum(dim=1)[None] + ) + codes = distances.argmin(dim=-1).reshape(values.shape[0], values.shape[2]) + quantized = functional.embedding(codes, self.codebook) + quantized = self.out_proj(quantized).transpose(1, 2) + return quantized.to(values.dtype), codes + + def decode_codes(self, codes: Tensor) -> Tensor: + return self.out_proj(functional.embedding(codes, self.codebook)).transpose(1, 2) + + +class _ResidualCodecQuantizer(nn.Module): + def __init__(self, config: Mapping[str, Any]) -> None: + super().__init__() + self.quantizers = nn.ModuleList([_CodecQuantizer(config) for _ in range(int(config["n_codebooks"]))]) + + def encode(self, values: Tensor) -> Tensor: + residual = values + codes = [] + for quantizer in self.quantizers: + quantized, level_codes = quantizer.encode(residual) + residual = residual - quantized + codes.append(level_codes) + return torch.stack(codes, dim=1) + + def from_codes(self, codes: Tensor) -> Tensor: + quantized = torch.zeros( + codes.shape[0], + self.quantizers[0].input_dim, + codes.shape[-1], + dtype=self.quantizers[0].codebook.dtype, + device=codes.device, + ) + for level, quantizer in enumerate(self.quantizers[: codes.shape[1]]): + quantized = quantized + quantizer.decode_codes(codes[:, level]) + return quantized + + +class _ActionCodecModel(nn.Module): + def __init__(self, config: Mapping[str, Any]) -> None: + super().__init__() + self.config = dict(config) + self.block_dct = ( + _BlockDCT(int(config["block_dct_block_size"])) + if bool(config.get("use_block_dct", False)) + else None + ) + self.conv_in = nn.Conv2d( + int(config["horizon_patch_size"]), + int(config["encoder_channels"]), + kernel_size=(1, int(config["conv_in_action_kernel"])), + ) + self.encoder = _CodecEncoder(config) + self.rvq = _ResidualCodecQuantizer(config) + self.decoder = _CodecDecoder(config) + self.conv_out = nn.ConvTranspose2d( + int(config["encoder_channels"]), + int(config["horizon_patch_size"]), + kernel_size=(1, int(config["conv_in_action_kernel"])), + ) + + @property + def code_h(self) -> int: + height = int(self.config["horizon"]) // int(self.config["horizon_patch_size"]) + for stride_h, _ in self.config["strides"]: + height //= int(stride_h) + return height + + @property + def code_a(self) -> int: + return int(self.config["max_component_dim"]) - int(self.config["conv_in_action_kernel"]) + 1 + + def _pad(self, values: Tensor) -> Tensor: + maximum = int(self.config["max_component_dim"]) + if values.shape[-1] < maximum: + return functional.pad(values, (0, maximum - values.shape[-1])) + return values[..., :maximum] + + def encode(self, components: dict[str, Tensor]) -> dict[str, Tensor]: + names = list(components) + batch_size = next(iter(components.values())).shape[0] + values = torch.cat([self._pad(components[name].float()) for name in names], dim=0) + if self.block_dct is not None: + values = self.block_dct.dct(values) + patch = int(self.config["horizon_patch_size"]) + values = values.reshape(values.shape[0], -1, patch, values.shape[-1]).transpose(1, 2) + latent = self.encoder(self.conv_in(values)).flatten(2) + codes = self.rvq.encode(latent) + return { + name: codes[index * batch_size : (index + 1) * batch_size] for index, name in enumerate(names) + } + + def decode(self, components: dict[str, Tensor], dimensions: Mapping[str, int]) -> dict[str, Tensor]: + names = list(components) + batch_size = next(iter(components.values())).shape[0] + codes = torch.cat([components[name] for name in names], dim=0) + quantized = self.rvq.from_codes(codes) + quantized = quantized.reshape( + quantized.shape[0], + quantized.shape[1], + self.code_h, + self.code_a, + ) + decoded = self.conv_out(self.decoder(quantized)) + decoded = decoded.transpose(1, 2).reshape(decoded.shape[0], -1, decoded.shape[-1]) + if self.block_dct is not None: + decoded = self.block_dct.idct(decoded, int(self.config["horizon"])) + return { + name: decoded[index * batch_size : (index + 1) * batch_size, :, : dimensions[name]] + for index, name in enumerate(names) + } + + +class _NativeCodecModule(nn.Module): + """Module hierarchy matching ``action_tokenizer.pt`` exactly.""" + + def __init__(self, config: Mapping[str, Any]) -> None: + super().__init__() + self.model = _ActionCodecModel(config) + + +class _BinarySequenceCodec: + def __init__(self, sequence_length: int, min_block_length: int, vocab_size: int) -> None: + self.sequence_length = sequence_length + self.min_block_length = min_block_length + self.vocab_size = vocab_size + self._count_cache: dict[tuple[int, int, int, bool], int] = {} + self.num_sequences = self._count(sequence_length, -1, 0, True) + self.num_tokens = max(1, math.ceil(math.log(self.num_sequences, vocab_size))) + + def _count(self, remaining: int, last: int, run_length: int, first: bool) -> int: + cache_key = (remaining, last, run_length, first) + if cache_key in self._count_cache: + return self._count_cache[cache_key] + if remaining == 0: + return 1 + total = 0 + for bit in (0, 1): + if last == -1 or bit == last: + total += self._count( + remaining - 1, + bit, + min(run_length + 1, self.min_block_length + 1), + first, + ) + elif first or run_length > self.min_block_length: + total += self._count(remaining - 1, bit, 1, False) + self._count_cache[cache_key] = total + return total + + def _repair(self, bits: list[int]) -> list[int]: + bits = bits.copy() + while True: + runs = [] + start = 0 + for index in range(1, len(bits)): + if bits[index] != bits[index - 1]: + runs.append((bits[start], start, index)) + start = index + runs.append((bits[start], start, len(bits))) + invalid = next( + ( + (start, stop, runs[index - 1][0]) + for index, (_, start, stop) in enumerate(runs[1:-1], start=1) + if stop - start <= self.min_block_length + ), + None, + ) + if invalid is None: + return bits + start, stop, value = invalid + bits[start:stop] = [value] * (stop - start) + + def _zero_completions(self, remaining: int, last: int, run: int, first: bool) -> int: + if last in (-1, 0): + return self._count( + remaining, + 0, + 1 if last == -1 else min(run + 1, self.min_block_length + 1), + first, + ) + return self._count(remaining, 0, 1, False) if first or run > self.min_block_length else 0 + + def encode(self, values: Tensor, threshold: float) -> Tensor: + output = [] + for row in values: + bits = self._repair([int(value >= threshold) for value in row.tolist()]) + rank, last, run, first = 0, -1, 0, True + for position, bit in enumerate(bits): + remaining = len(bits) - position - 1 + if bit: + rank += self._zero_completions(remaining, last, run, first) + if last == -1: + last, run = bit, 1 + elif bit == last: + run = min(run + 1, self.min_block_length + 1) + else: + last, run, first = bit, 1, False + tokens = [] + for _ in range(self.num_tokens): + tokens.append(rank % self.vocab_size) + rank //= self.vocab_size + output.append(list(reversed(tokens))) + return torch.tensor(output, dtype=torch.long, device=values.device) + + def decode(self, tokens: Tensor) -> Tensor: + rows = [] + for row in tokens.tolist(): + rank = 0 + for token in row: + rank = rank * self.vocab_size + max(0, min(int(token), self.vocab_size - 1)) + rank = min(rank, self.num_sequences - 1) + bits, last, run, first = [], -1, 0, True + for position in range(self.sequence_length): + remaining = self.sequence_length - position - 1 + zeros = self._zero_completions(remaining, last, run, first) + if rank < zeros: + bit = 0 + else: + rank -= zeros + bit = 1 + bits.append(bit) + if last == -1: + last, run = bit, 1 + elif bit == last: + run = min(run + 1, self.min_block_length + 1) + else: + last, run, first = bit, 1, False + rows.append(bits) + return torch.tensor(rows, dtype=torch.float32, device=tokens.device) + + +class G05NativeActionCodec: + """Non-registered sidecar wrapper for native ActionCodec encode/decode.""" + + def __init__(self, config: Mapping[str, Any], *, action_token_begin: int) -> None: + self.config = dict(config) + architecture = self.config["model_arch"] + self.module = _NativeCodecModule(architecture) + self.model = self.module.model + self.action_token_begin = action_token_begin + self.parts = { + key: int(value) for key, value in self.config["parts_meta"].items() if value is not None + } + patterns = tuple(self.config.get("rule_based_key_patterns") or ()) + self.rule_parts = [key for key in self.parts if any(pattern in key for pattern in patterns)] + self.neural_parts = [key for key in self.parts if key not in self.rule_parts] + self.codebook_size = int(architecture["codebook_size"]) + self.max_residuals = int(architecture["n_codebooks"]) + self.num_residuals = int(self.config.get("num_residuals") or self.max_residuals) + self.code_length = self.model.code_h * self.model.code_a + marker_names = [ + f"<{part}_{level}>" for level in range(self.max_residuals) for part in self.neural_parts + ] + [f"<{part}>" for part in self.rule_parts] + self.marker_indices = {name: self.codebook_size + index for index, name in enumerate(marker_names)} + self.rule_codec = _BinarySequenceCodec( + int(architecture["horizon"]), + int(self.config.get("rule_based_min_block_len", 1)), + self.codebook_size, + ) + + @property + def action_token_length(self) -> int: + neural = len(self.neural_parts) * self.num_residuals * (self.code_length + 1) + rules = len(self.rule_parts) * (self.rule_codec.num_tokens + 1) + return neural + rules + + @classmethod + def load( + cls, + config: Mapping[str, Any], + *, + action_token_begin: int, + ) -> G05NativeActionCodec: + codec = cls(config, action_token_begin=action_token_begin) + checkpoint = torch.load( + Path(str(config["ckpt_dir"])), + map_location="cpu", + mmap=True, + weights_only=True, + ) + state_dict = checkpoint.get("model_state_dict", checkpoint) + codec.module.load_state_dict(state_dict, strict=True) + codec.module.eval() + return codec + + def to(self, device: torch.device | str) -> G05NativeActionCodec: + self.module.to(device=device, dtype=torch.float32) + return self + + def _split(self, actions: Tensor) -> dict[str, Tensor]: + splits = torch.split(actions[..., : sum(self.parts.values())], list(self.parts.values()), dim=-1) + return dict(zip(self.parts, splits, strict=True)) + + @torch.no_grad() + def encode_for_language(self, payload: Mapping[str, Any]) -> list[int]: + actions = torch.as_tensor(payload["value"]) + if actions.ndim == 2: + actions = actions.unsqueeze(0) + components = self._split(actions) + neural = {key: components[key] for key in self.neural_parts} + codes = self.model.encode(neural) + rule_codes = { + key: self.rule_codec.encode( + components[key][..., 0], + float(self.config.get("rule_based_binarize_threshold", 0)), + ) + for key in self.rule_parts + } + indices = [] + for level in range(self.num_residuals): + for key in self.neural_parts: + indices.append(self.marker_indices[f"<{key}_{level}>"]) + indices.extend(codes[key][0, level].tolist()) + for key in self.rule_parts: + indices.append(self.marker_indices[f"<{key}>"]) + indices.extend(rule_codes[key][0].tolist()) + return [self.action_token_begin + int(index) for index in indices] + + @torch.no_grad() + def decode_language_tokens( + self, + token_ids: Tensor, + *, + horizon: int, + action_dim: int, + ) -> tuple[Tensor, set[str]]: + indices = (token_ids.long() - self.action_token_begin).tolist() + marker_to_name = {value: name for name, value in self.marker_indices.items()} + neural: dict[str, list[list[int] | None]] = { + key: [None] * self.num_residuals for key in self.neural_parts + } + rules: dict[str, list[int]] = {} + cursor = 0 + while cursor < len(indices): + marker = marker_to_name.get(indices[cursor]) + if marker is None: + cursor += 1 + continue + marker = marker[1:-1] + if marker in self.rule_parts: + length = self.rule_codec.num_tokens + values = indices[cursor + 1 : cursor + 1 + length] + if len(values) == length and all(0 <= value < self.codebook_size for value in values): + rules[marker] = values + cursor += length + 1 + continue + part, level_text = marker.rsplit("_", 1) + level = int(level_text) + if part in neural and level < self.num_residuals: + values = indices[cursor + 1 : cursor + 1 + self.code_length] + if len(values) == self.code_length and all( + 0 <= value < self.codebook_size for value in values + ): + neural[part][level] = values + cursor += self.code_length + 1 + + absent = { + key + for key in self.parts + if (key in neural and not any(level is not None for level in neural[key])) + or (key in self.rule_parts and key not in rules) + } + device = next(self.module.parameters()).device + code_tensors = {} + for key, levels in neural.items(): + if not any(level is not None for level in levels): + continue + filled = [level if level is not None else [0] * self.code_length for level in levels] + code_tensors[key] = torch.tensor([filled], dtype=torch.long, device=device) + decoded = ( + self.model.decode(code_tensors, {key: self.parts[key] for key in code_tensors}) + if code_tensors + else {} + ) + for key in self.rule_parts: + if key in rules: + tokens = torch.tensor([rules[key]], dtype=torch.long, device=device) + binary = self.rule_codec.decode(tokens) + decoded[key] = binary[:, :, None] * 2 - 1 + # ``absent_key_fill_value`` is an internal partitioner sentinel. The + # released marker-aware final decoder converts absent/no-op body parts + # to zero motion before returning an action. + batch = torch.zeros((1, horizon, action_dim), dtype=torch.float32, device=device) + offset = 0 + for key, dimension in self.parts.items(): + if key in decoded: + batch[..., offset : offset + dimension] = decoded[key][..., :dimension] + offset += dimension + return batch[0], absent diff --git a/src/lerobot/policies/g05/configuration_g05.py b/src/lerobot/policies/g05/configuration_g05.py index e21441a93..e029fbf05 100644 --- a/src/lerobot/policies/g05/configuration_g05.py +++ b/src/lerobot/policies/g05/configuration_g05.py @@ -188,9 +188,9 @@ _PROFILE_DEFAULTS = { class G05Config(PreTrainedConfig): """LeRobot-side, checkpoint-auditable configuration for G0.5. - ``author_model_config`` comes from the packaged checkpoint's resolved Hydra - config. It is intentionally checkpoint state rather than a collection of - guessed LeRobot defaults. + ``author_model_config`` is the backward-compatible serialized field holding + the packaged checkpoint's resolved native architecture. It is checkpoint + state rather than a collection of guessed LeRobot defaults. """ checkpoint_profile: str = "g05-base" diff --git a/src/lerobot/policies/g05/modeling_g05.py b/src/lerobot/policies/g05/modeling_g05.py index b53d6020b..912149b62 100644 --- a/src/lerobot/policies/g05/modeling_g05.py +++ b/src/lerobot/policies/g05/modeling_g05.py @@ -10,7 +10,6 @@ from __future__ import annotations -import importlib import json import shutil from collections import deque @@ -33,27 +32,25 @@ from .configuration_g05 import ( make_g05_cot_prompt_template, make_g05_prompt_template, ) +from .native_g05 import G05NativeBackend -def _author_backend(config: G05Config) -> nn.Module: +def _native_backend(config: G05Config) -> nn.Module: if not config.author_model_config: raise ValueError( "G0.5 author_model_config is empty. Load a packaged checkpoint, or " "inject a backend explicitly for testing." ) - try: - from omegaconf import OmegaConf - - module = importlib.import_module("g05.models.g05.g05_policy_qwen35") - except ImportError as exc: - raise ImportError( - "The OpenGalaxea G0.5 author package is required for real model execution. " - "Clone the pinned GalaxeaVLA source, accept LICENSE-G0.5, and install its " - "runtime dependencies in a compatible environment. LeRobot does not vendor " - "or silently download that non-commercial code." - ) from exc - backend_cls = module.G05PolicyQwen35 - return backend_cls(**OmegaConf.to_container(OmegaConf.create(config.author_model_config))) + model_config = dict(config.author_model_config) + model_config.update( + { + "predict_cot": config.predict_cot, + "discrete_action": config.discrete_action, + "continuous_action": config.continuous_action, + "return_continuous_action": config.return_continuous_action, + } + ) + return G05NativeBackend.from_config(model_config) class G05Policy(PreTrainedPolicy): @@ -65,7 +62,7 @@ class G05Policy(PreTrainedPolicy): def __init__(self, config: G05Config, backend: nn.Module | None = None): super().__init__(config) config.validate_features() - self.backend = backend if backend is not None else _author_backend(config) + self.backend = backend if backend is not None else _native_backend(config) if not isinstance(self.backend, nn.Module): raise TypeError(f"G0.5 backend must be an nn.Module, got {type(self.backend)}.") self._action_queue: deque[Tensor] = deque() @@ -150,7 +147,7 @@ class G05Policy(PreTrainedPolicy): shutil.copy2(source, save_directory / name) # Serialized paths are portable sidecar names. Local/Hub loading resolves them - # against the downloaded checkpoint directory before constructing the author model. + # against the downloaded checkpoint directory before constructing the native model. if (processor_path is not None and processor_path.exists()) or ( tokenizer_path is not None and tokenizer_path.exists() ): @@ -191,7 +188,7 @@ class G05Policy(PreTrainedPolicy): weight.data = weight.data.float() def to(self, *args, **kwargs) -> G05Policy: - """Apply author inference precision and move the ActionCodec sidecar.""" + """Apply the released inference precision and move the ActionCodec sidecar.""" result = super().to(*args, **kwargs) explicit_dtype = "dtype" in kwargs or any(isinstance(arg, torch.dtype | Tensor) for arg in args) diff --git a/src/lerobot/policies/g05/native_g05.py b/src/lerobot/policies/g05/native_g05.py new file mode 100644 index 000000000..ac28d1213 --- /dev/null +++ b/src/lerobot/policies/g05/native_g05.py @@ -0,0 +1,1215 @@ +# 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. + +"""Native G0.5 model components built on Transformers' Qwen3.5 implementation.""" + +from __future__ import annotations + +import math +import time +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import torch +import torch.nn.functional as functional +from torch import Tensor, nn + +from lerobot.policies.pi_gemma import PiGemmaRMSNorm +from lerobot.utils.constants import ACTION + +from .action_codec_g05 import G05NativeActionCodec +from .processing_g05 import IGNORE_INDEX, G05SequenceBatch, G05Tokenizer, G05TokenType + + +def _qwen_text_config(values: Mapping[str, Any], *, vocab_size: int | None = None): + """Translate the serialized G0.5 Qwen config into a Transformers config.""" + + from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig + + return Qwen3_5TextConfig( + vocab_size=int(vocab_size if vocab_size is not None else values.get("vocab_size", 1)), + hidden_size=int(values["hidden_size"]), + intermediate_size=int(values["intermediate_size"]), + num_hidden_layers=int(values["num_hidden_layers"]), + num_attention_heads=int(values["num_attention_heads"]), + num_key_value_heads=int(values["num_key_value_heads"]), + head_dim=int(values["head_dim"]), + rms_norm_eps=float(values["rms_norm_eps"]), + max_position_embeddings=int(values["max_position_embeddings"]), + attention_bias=bool(values.get("attention_bias", False)), + hidden_act=str(values.get("hidden_act", "silu")), + rope_parameters=dict(values["rope_parameters"]), + linear_conv_kernel_dim=int(values.get("linear_conv_kernel_dim", 4)), + linear_key_head_dim=int(values.get("linear_key_head_dim", 128)), + linear_value_head_dim=int(values.get("linear_value_head_dim", 128)), + linear_num_key_heads=int(values.get("linear_num_key_heads", 16)), + linear_num_value_heads=int(values.get("linear_num_value_heads", 16)), + layer_types=list(values["layer_types"]), + pad_token_id=values.get("pad_token_id"), + tie_word_embeddings=True, + ) + + +def _qwen_vision_config(values: Mapping[str, Any]): + """Translate the serialized G0.5 vision config into Transformers.""" + + from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5VisionConfig + + config = Qwen3_5VisionConfig( + depth=int(values["depth"]), + hidden_size=int(values["hidden_size"]), + num_heads=int(values["num_heads"]), + patch_size=int(values["patch_size"]), + temporal_patch_size=int(values["temporal_patch_size"]), + spatial_merge_size=int(values["spatial_merge_size"]), + in_channels=int(values.get("in_channels", 3)), + intermediate_size=int(values["intermediate_size"]), + out_hidden_size=int(values["out_hidden_size"]), + num_position_embeddings=int(values["num_position_embeddings"]), + hidden_act=str(values.get("hidden_act", "gelu_pytorch_tanh")), + ) + config.temporal_freq = int(values.get("temporal_freq", 0)) + config.spacetime_mode = str(values.get("spacetime_mode", "factorized")) + config.token_drop_layer = values.get("token_drop_layer") + config.temporal_pe_pretrain_frames = values.get("temporal_pe_pretrain_frames") + config.batch_all_cameras = bool(values.get("batch_all_cameras", False)) + return config + + +class G05ProprioEmbedder(nn.Module): + """Project the padded G0.5 proprioception vector into the VLM hidden size.""" + + def __init__(self, proprio_dim: int, hidden_size: int) -> None: + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(proprio_dim, hidden_size), + nn.GELU(), + nn.LayerNorm(hidden_size), + nn.Linear(hidden_size, hidden_size), + ) + + def forward(self, proprio: Tensor) -> Tensor: + with torch.autocast(proprio.device.type, enabled=False): + return self.mlp(proprio.float()) + + +class G05QwenTextModel(nn.Module): + """Qwen3.5 text stack with G0.5's checkpoint-compatible module names.""" + + def __init__(self, values: Mapping[str, Any], *, vocab_size: int) -> None: + super().__init__() + from transformers.models.qwen3_5.modeling_qwen3_5 import ( + Qwen3_5DecoderLayer, + Qwen3_5RMSNorm, + Qwen3_5TextRotaryEmbedding, + ) + + self.config = _qwen_text_config(values, vocab_size=vocab_size) + self.input_proj = nn.Embedding(vocab_size, self.config.hidden_size, self.config.pad_token_id) + self.layers = nn.ModuleList( + [ + Qwen3_5DecoderLayer(self.config, layer_idx) + for layer_idx in range(self.config.num_hidden_layers) + ] + ) + self.norm = Qwen3_5RMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps) + self.rotary_emb = Qwen3_5TextRotaryEmbedding(self.config) + + def embed(self, input_ids: Tensor) -> Tensor: + return self.input_proj(input_ids) + + def logits(self, hidden_states: Tensor) -> Tensor: + return torch.nn.functional.linear(hidden_states, self.input_proj.weight) + + def forward( + self, + inputs_embeds: Tensor, + *, + full_attention_mask: Tensor, + linear_attention_mask: Tensor, + position_ids: Tensor, + cache=None, + ) -> tuple[Tensor, Any]: + from transformers import DynamicCache + + if cache is None: + cache = DynamicCache(config=self.config) + position_embeddings = self.rotary_emb(inputs_embeds, position_ids) + hidden_states = inputs_embeds + for layer_index, layer in enumerate(self.layers): + attention_mask = ( + linear_attention_mask + if self.config.layer_types[layer_index] == "linear_attention" + else full_attention_mask + ) + hidden_states = layer( + hidden_states, + position_embeddings=position_embeddings, + attention_mask=attention_mask, + position_ids=position_ids[0], + past_key_values=cache, + use_cache=True, + ) + return self.norm(hidden_states), cache + + +class G05ActionDecoderLayer(nn.Module): + """Qwen3.5 decoder layer with G0.5 adaptive RMSNorm conditioning.""" + + def __init__(self, config, layer_idx: int) -> None: + super().__init__() + from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5Attention, Qwen3_5MLP + + if config.layer_types[layer_idx] != "full_attention": + raise ValueError("The released G0.5 action expert requires full-attention layers.") + self.layer_idx = layer_idx + self.self_attn = Qwen3_5Attention(config, layer_idx) + self.mlp = Qwen3_5MLP(config, config.intermediate_size) + self.input_layernorm = PiGemmaRMSNorm( + config.hidden_size, + eps=config.rms_norm_eps, + cond_dim=config.hidden_size, + ) + self.post_attention_layernorm = PiGemmaRMSNorm( + config.hidden_size, + eps=config.rms_norm_eps, + cond_dim=config.hidden_size, + ) + + def forward( + self, + hidden_states: Tensor, + *, + attention_mask: Tensor, + position_embeddings: tuple[Tensor, Tensor], + position_ids: Tensor, + cache, + time_cond: Tensor, + ) -> Tensor: + residual = hidden_states + hidden_states, gate = self.input_layernorm(hidden_states, cond=time_cond) + key_length = cache.layers[self.layer_idx].get_seq_length() + hidden_states.shape[1] + layer_attention_mask = ( + attention_mask[..., -key_length:] if attention_mask.shape[-1] != key_length else attention_mask + ) + hidden_states, _ = self.self_attn( + hidden_states, + attention_mask=layer_attention_mask, + position_ids=position_ids[0], + past_key_values=cache, + position_embeddings=position_embeddings, + use_cache=True, + ) + hidden_states = residual + hidden_states if gate is None else residual + hidden_states * gate + + residual = hidden_states + hidden_states, gate = self.post_attention_layernorm(hidden_states, cond=time_cond) + hidden_states = self.mlp(hidden_states) + return residual + hidden_states if gate is None else residual + hidden_states * gate + + +class G05ActionExpert(nn.Module): + """Continuous G0.5 action expert with checkpoint-compatible parameter names.""" + + def __init__(self, values: Mapping[str, Any]) -> None: + super().__init__() + from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5TextRotaryEmbedding + + self.config = _qwen_text_config(values) + input_dim = int(values["input_dim"]) + output_dim = int(values["output_dim"]) + hidden_size = int(values["hidden_size"]) + self.input_proj = nn.Linear(input_dim, hidden_size) + self.layers = nn.ModuleList( + [ + G05ActionDecoderLayer(self.config, layer_idx) + for layer_idx in range(self.config.num_hidden_layers) + ] + ) + self.norm = PiGemmaRMSNorm( + hidden_size, + eps=self.config.rms_norm_eps, + cond_dim=hidden_size, + ) + self.output_proj = nn.Linear(hidden_size, output_dim) + self.time_mlp_in = nn.Linear(hidden_size, hidden_size) + self.time_mlp_out = nn.Linear(hidden_size, hidden_size) + self.rotary_emb = Qwen3_5TextRotaryEmbedding(self.config) + + def embed(self, actions: Tensor) -> Tensor: + return self.input_proj(actions) + + def encode_time(self, timesteps: Tensor) -> Tensor: + half = self.config.hidden_size // 2 + fraction = torch.linspace(0.0, 1.0, half, device=timesteps.device, dtype=torch.float32) + periods = 4e-3 * (4.0 / 4e-3) ** fraction + phase = timesteps.float().unsqueeze(-1) * (2 * math.pi / periods) + embedding = torch.cat((phase.sin(), phase.cos()), dim=-1) + with torch.autocast(timesteps.device.type, enabled=False): + return torch.nn.functional.silu( + self.time_mlp_out(torch.nn.functional.silu(self.time_mlp_in(embedding))) + ) + + def forward( + self, + inputs_embeds: Tensor, + *, + attention_mask: Tensor, + position_ids: Tensor, + cache, + time_cond: Tensor, + ) -> Tensor: + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids) + for layer in self.layers: + hidden_states = layer( + hidden_states, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + position_ids=position_ids, + cache=cache, + time_cond=time_cond, + ) + hidden_states, _ = self.norm(hidden_states, cond=time_cond) + return hidden_states + + def decode(self, hidden_states: Tensor) -> Tensor: + with torch.autocast(hidden_states.device.type, enabled=False): + return self.output_proj(hidden_states.float()) + + +class G05NativeModel(nn.Module): + """Weight-owning native G0.5 model assembled from serialized checkpoint config.""" + + def __init__(self, model_config: Mapping[str, Any], *, vocab_size: int) -> None: + super().__init__() + from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5VisionModel + + self.vision_tower = Qwen3_5VisionModel(_qwen_vision_config(model_config["vision"])) + self.vlm = G05QwenTextModel(model_config["vlm"], vocab_size=vocab_size) + self.action_expert = G05ActionExpert(model_config["action_expert"]) + self.proprio_embedder = G05ProprioEmbedder( + int(model_config["proprio_dim"]), + int(model_config["vlm"]["hidden_size"]), + ) + + +def _temporal_embedding(timesteps: Tensor, dimension: int) -> Tensor: + half = dimension // 2 + frequencies = torch.exp( + -math.log(10000.0) * torch.arange(half, device=timesteps.device, dtype=torch.float32) / max(half, 1) + ) + phase = timesteps.float().unsqueeze(1) * frequencies.unsqueeze(0) + return torch.stack((phase.sin(), phase.cos() - 1), dim=-1).reshape(len(timesteps), dimension) + + +class G05NativeBackend(nn.Module): + """Native LeRobot backend for G0.5. + + Inference and training routing are added around this checkpoint-compatible + model core; no OpenGalaxea Python package is imported. + """ + + def __init__( + self, + model_config: Mapping[str, Any], + *, + vocab_size: int, + processor_path: str | Path, + ) -> None: + super().__init__() + self.model_config = dict(model_config) + self.model = G05NativeModel(self.model_config, vocab_size=vocab_size) + attention_implementation = str(self.model_config.get("attn_implementation", "eager")) + self.model.vlm.config._attn_implementation = attention_implementation + self.model.action_expert.config._attn_implementation = attention_implementation + self.model.vision_tower.config._attn_implementation = attention_implementation + self.processor = G05Tokenizer(processor_path, self.model_config) + if len(self.processor) != vocab_size: + raise ValueError( + f"G0.5 tokenizer has {len(self.processor)} rows, but model expects {vocab_size}." + ) + self.action_tokenizer = None + action_config = self.model_config.get("AT_CONFIG") + if isinstance(action_config, Mapping): + checkpoint = Path(str(action_config.get("ckpt_dir", ""))) + if checkpoint.is_file(): + self.action_tokenizer = G05NativeActionCodec.load( + action_config, + action_token_begin=self.processor.action_token_begin, + ) + self._last_vision_grids: list[tuple[int, int, int]] = [] + + def apply_fp32_params(self) -> None: + """Restore the FP32 islands used by the released mixed-precision runtime.""" + + patterns = ( + "vision_tower.patch_embed", + "vision_tower.pos_embed", + "vision_tower.merger", + "norm1", + "norm2", + "input_layernorm", + "post_attention_layernorm", + "q_norm", + "k_norm", + "linear_attn.norm", + "linear_attn.A_log", + "linear_attn.dt_bias", + "vlm.norm", + "action_expert.norm", + "action_expert.input_proj", + "action_expert.output_proj", + "action_expert.time_mlp", + "proprio_embedder", + ) + for name, parameter in self.named_parameters(): + if any(pattern in name for pattern in patterns): + parameter.data = parameter.data.float() + + @staticmethod + def _should_apply_weight_decay( + owner_module: nn.Module | None, + leaf_name: str, + parameter: nn.Parameter, + ) -> bool: + return leaf_name != "bias" and parameter.ndim > 1 and not isinstance(owner_module, nn.Embedding) + + def get_optim_param_groups( + self, + *, + lr: float, + weight_decay: float, + apply_decay_on_norm_and_bias: bool = False, + backbone_lr_multiplier: float = 1.0, + vision_lr_multiplier: float = 1.0, + ) -> list[dict[str, Any]]: + """Build the released six native backbone/action/vision parameter groups.""" + + action_parameters = {id(parameter) for parameter in self.model.action_expert.parameters()} + vision_parameters = {id(parameter) for parameter in self.model.vision_tower.parameters()} + modules = dict(self.model.named_modules()) + grouped: dict[str, list[nn.Parameter]] = { + "backbone_decay": [], + "action_decay": [], + "vision_decay": [], + "backbone_no_decay": [], + "action_no_decay": [], + "vision_no_decay": [], + } + for name, parameter in self.model.named_parameters(): + if not parameter.requires_grad: + continue + owner_name, _, leaf_name = name.rpartition(".") + decay = apply_decay_on_norm_and_bias or self._should_apply_weight_decay( + modules.get(owner_name), + leaf_name, + parameter, + ) + if id(parameter) in action_parameters: + family = "action" + elif id(parameter) in vision_parameters: + family = "vision" + else: + family = "backbone" + grouped[f"{family}_{'decay' if decay else 'no_decay'}"].append(parameter) + + learning_rates = { + "backbone": lr * backbone_lr_multiplier, + "action": lr, + "vision": lr * backbone_lr_multiplier * vision_lr_multiplier, + } + parameter_groups = [ + { + "params": grouped[name], + "lr": learning_rates[name.split("_", 1)[0]], + "weight_decay": weight_decay if not name.endswith("_no_decay") else 0.0, + "name": name, + } + for name in ( + "backbone_decay", + "action_decay", + "vision_decay", + "backbone_no_decay", + "action_no_decay", + "vision_no_decay", + ) + ] + expected = sum(parameter.requires_grad for parameter in self.model.parameters()) + actual = sum(len(group["params"]) for group in parameter_groups) + if actual != expected: + raise RuntimeError( + f"G0.5 optimizer grouping lost parameters: grouped {actual}, expected {expected}." + ) + return parameter_groups + + @classmethod + def from_config(cls, model_config: Mapping[str, Any]) -> G05NativeBackend: + processor_path = Path(str(model_config["hf_processor_path"])) + tokenizer_config = processor_path / "tokenizer_config.json" + if not tokenizer_config.is_file(): + raise FileNotFoundError(f"G0.5 tokenizer config not found: {tokenizer_config}") + import json + + tokenizer_metadata = json.loads(tokenizer_config.read_text()) + added = tokenizer_metadata.get("added_tokens_decoder") or {} + base_vocab_size = max((int(token_id) for token_id in added), default=-1) + 1 + at_config = model_config["AT_CONFIG"] + codebook_size = int(at_config["model_arch"]["codebook_size"]) + parts = at_config["parts_meta"] + rule_patterns = tuple(at_config.get("rule_based_key_patterns") or ()) + rule_parts = [name for name in parts if any(pattern in name for pattern in rule_patterns)] + neural_parts = [name for name in parts if name not in rule_parts] + residuals = int(at_config["model_arch"]["n_codebooks"]) + marker_count = len(neural_parts) * residuals + len(rule_parts) + # Action-code tokens, group markers, , and the MLP token. + vocab_size = base_vocab_size + codebook_size + marker_count + 2 + return cls(model_config, vocab_size=vocab_size, processor_path=processor_path) + + @staticmethod + def _patchify(images: Tensor, patch_size: int, temporal_patch_size: int, merge_size: int) -> Tensor: + batch_frames, channels, height, width = images.shape + grid_h, grid_w = height // patch_size, width // patch_size + temporal = images.unsqueeze(2).expand(-1, -1, temporal_patch_size, -1, -1) + return ( + temporal.reshape( + batch_frames, + temporal_patch_size, + channels, + grid_h // merge_size, + merge_size, + patch_size, + grid_w // merge_size, + merge_size, + patch_size, + ) + .permute(0, 3, 6, 4, 7, 2, 1, 5, 8) + .reshape(batch_frames * grid_h * grid_w, -1) + ) + + def _vision_temporal_block( + self, + block, + hidden_states: Tensor, + *, + position_embeddings: tuple[Tensor, Tensor], + batch_size: int, + num_frames: int, + patches_per_frame: int, + temporal_pe: Tensor, + temporal_mask: Tensor, + ) -> Tensor: + from transformers.models.qwen3_5.modeling_qwen3_5 import apply_rotary_pos_emb_vision + + total, hidden_size = hidden_states.shape + num_heads = block.attn.num_heads + head_dim = block.attn.head_dim + residual = hidden_states + conditioned = ( + hidden_states.view(batch_size, num_frames, patches_per_frame, hidden_size) + + temporal_pe[None, :, None, :] + ).reshape(total, hidden_size) + normed = block.norm1(conditioned) + query, key, value = ( + block.attn.qkv(normed).reshape(total, 3, num_heads, head_dim).permute(1, 0, 2, 3).unbind(0) + ) + + def temporal_view(tensor: Tensor) -> Tensor: + return ( + tensor.view(batch_size, num_frames, patches_per_frame, num_heads, head_dim) + .permute(0, 2, 3, 1, 4) + .reshape(batch_size * patches_per_frame, num_heads, num_frames, head_dim) + ) + + query_t, key_t, value_t = (temporal_view(tensor) for tensor in (query, key, value)) + weights = torch.matmul(query_t, key_t.transpose(-2, -1)) * block.attn.scaling + weights = functional.softmax(weights + temporal_mask[None, None], dim=-1, dtype=torch.float32).to( + query_t.dtype + ) + mixed_value = torch.matmul(weights, value_t) + mixed_value = ( + mixed_value.view(batch_size, patches_per_frame, num_heads, num_frames, head_dim) + .permute(0, 3, 1, 2, 4) + .reshape(total, num_heads, head_dim) + ) + + cosine, sine = position_embeddings + query, key = apply_rotary_pos_emb_vision(query, key, cosine, sine) + spatial_outputs = [] + for start in range(0, total, patches_per_frame): + stop = start + patches_per_frame + spatial_outputs.append( + functional.scaled_dot_product_attention( + query[start:stop].transpose(0, 1).unsqueeze(0), + key[start:stop].transpose(0, 1).unsqueeze(0), + mixed_value[start:stop].transpose(0, 1).unsqueeze(0), + scale=block.attn.scaling, + ) + ) + spatial = torch.cat(spatial_outputs, dim=2).squeeze(0).transpose(0, 1).reshape(total, hidden_size) + hidden_states = residual + block.attn.proj(spatial) + return hidden_states + block.mlp(block.norm2(hidden_states)) + + def _encode_camera(self, frames: Tensor) -> tuple[Tensor, tuple[int, int, int]]: + """Encode one camera, including G0.5's causal temporal-memory mixing.""" + + tower = self.model.vision_tower + batch_size, num_frames, _, height, width = frames.shape + patch_size = int(tower.config.patch_size) + merge_size = int(tower.config.spatial_merge_size) + temporal_patch_size = int(tower.config.temporal_patch_size) + grid_h, grid_w = height // patch_size, width // patch_size + patches_per_frame = grid_h * grid_w + flattened = frames.reshape(batch_size * num_frames, *frames.shape[2:]) + patches = self._patchify(flattened, patch_size, temporal_patch_size, merge_size) + grid = torch.tensor( + [[1, grid_h, grid_w]] * (batch_size * num_frames), + dtype=torch.long, + device=frames.device, + ) + + with torch.autocast(frames.device.type, enabled=False): + hidden_states = tower.patch_embed(patches) + hidden_states = hidden_states + tower.fast_pos_embed_interpolate(grid) + rotary = tower.rot_pos_emb(grid).reshape(hidden_states.shape[0], -1) + rotary = torch.cat((rotary, rotary), dim=-1) + position_embeddings = (rotary.cos(), rotary.sin()) + cu_seqlens = torch.arange( + 0, + (batch_size * num_frames + 1) * patches_per_frame, + patches_per_frame, + dtype=torch.int32, + device=frames.device, + ) + + temporal_frequency = int(getattr(tower.config, "temporal_freq", 0)) + if num_frames > 1 and temporal_frequency > 0: + timesteps = torch.arange(-(num_frames - 1), 1, device=frames.device) + temporal_pe = _temporal_embedding(timesteps, hidden_states.shape[-1]).to(hidden_states.dtype) + temporal_mask = torch.triu( + torch.full( + (num_frames, num_frames), + float("-inf"), + device=frames.device, + dtype=hidden_states.dtype, + ), + diagonal=1, + ) + drop_layer = int(getattr(tower.config, "token_drop_layer", None) or len(tower.blocks)) - 1 + else: + temporal_pe = temporal_mask = None + drop_layer = -1 + + for layer_index, block in enumerate(tower.blocks): + use_temporal = ( + temporal_pe is not None + and layer_index <= drop_layer + and (drop_layer - layer_index) % temporal_frequency == 0 + ) + if use_temporal: + hidden_states = self._vision_temporal_block( + block, + hidden_states, + position_embeddings=position_embeddings, + batch_size=batch_size, + num_frames=num_frames, + patches_per_frame=patches_per_frame, + temporal_pe=temporal_pe, + temporal_mask=temporal_mask, + ) + else: + hidden_states = block( + hidden_states, + cu_seqlens=cu_seqlens, + position_embeddings=position_embeddings, + ) + if layer_index == drop_layer and num_frames > 1: + hidden_states = hidden_states.view(batch_size, num_frames, patches_per_frame, -1)[ + :, -1 + ].reshape(batch_size * patches_per_frame, -1) + num_frames = 1 + cu_seqlens = torch.arange( + 0, + (batch_size + 1) * patches_per_frame, + patches_per_frame, + dtype=torch.int32, + device=frames.device, + ) + cosine, sine = position_embeddings + position_embeddings = ( + cosine[: batch_size * patches_per_frame], + sine[: batch_size * patches_per_frame], + ) + + with torch.autocast(frames.device.type, enabled=False): + merged = tower.merger(hidden_states) + tokens_per_frame = (grid_h // merge_size) * (grid_w // merge_size) + return merged.reshape(batch_size, tokens_per_frame, -1), (1, grid_h, grid_w) + + def _encode_vision(self, pixel_values: Mapping[str, Tensor]) -> Tensor: + features = [] + grids = [] + for frames in pixel_values.values(): + feature, grid = self._encode_camera(frames) + features.append(feature) + grids.append(grid) + self._last_vision_grids = grids + return torch.cat(features, dim=1) + + def _embed( + self, + sequence: G05SequenceBatch, + pixel_values: Mapping[str, Tensor], + proprio: Tensor, + ) -> Tensor: + image_features = self._encode_vision(pixel_values) + text_features = self.model.vlm.embed(sequence.input_ids).to(image_features.dtype) + embeddings = text_features.clone() + + image_mask = sequence.token_types == G05TokenType.IMAGE + image_indices = (image_mask.long().cumsum(dim=1) - 1).clamp(min=0) + if image_mask.any() and int(image_indices[image_mask].max()) >= image_features.shape[1]: + raise ValueError("G0.5 prompt image-token count does not match the native vision encoder output.") + gathered_images = torch.gather( + image_features, + 1, + image_indices.unsqueeze(-1).expand(-1, -1, image_features.shape[-1]), + ) + embeddings[image_mask] = gathered_images[image_mask] + + state_mask = sequence.token_types == G05TokenType.PROPRIO + state_features = self.model.proprio_embedder(proprio).to(embeddings.dtype) + state_indices = (state_mask.long().cumsum(dim=1) - 1).clamp(min=0) + gathered_state = torch.gather( + state_features, + 1, + state_indices.unsqueeze(-1).expand(-1, -1, state_features.shape[-1]), + ) + embeddings[state_mask] = gathered_state[state_mask] + return embeddings + + def _mrope_positions(self, token_types: Tensor) -> Tensor: + import itertools + + batch_size, sequence_length = token_types.shape + positions = torch.zeros( + 3, + batch_size, + sequence_length, + dtype=torch.long, + device=token_types.device, + ) + position_mode = str(self.model_config.get("position_ids_type", "pi0fast")) + for batch_index in range(batch_size): + cursor = 0 + grid_index = 0 + values = token_types[batch_index].detach().cpu().tolist() + for token_type, entries in itertools.groupby(enumerate(values), key=lambda item: item[1]): + entries = list(entries) + start, stop = entries[0][0], entries[-1][0] + 1 + if int(token_type) == G05TokenType.PADDING: + continue + length = stop - start + if int(token_type) == G05TokenType.IMAGE: + if grid_index >= len(self._last_vision_grids): + raise ValueError("G0.5 MRoPE received more image segments than vision grids.") + _, raw_h, raw_w = self._last_vision_grids[grid_index] + grid_index += 1 + merge = int(self.model_config["vision"]["spatial_merge_size"]) + grid_h, grid_w = raw_h // merge, raw_w // merge + height = ( + torch.arange(grid_h, device=token_types.device).repeat_interleave(grid_w)[:length] + + cursor + ) + width = torch.arange(grid_w, device=token_types.device).repeat(grid_h)[:length] + cursor + positions[0, batch_index, start:stop] = cursor + positions[1, batch_index, start:stop] = height + positions[2, batch_index, start:stop] = width + cursor += max(grid_h, grid_w) + continue + + if position_mode == "gaussian": + if self.training: + steps = ( + torch.normal( + mean=2.0, + std=0.5, + size=(length,), + device=token_types.device, + ) + .round() + .clamp(1, 3) + .long() + ) + else: + steps = torch.full((length,), 2, dtype=torch.long, device=token_types.device) + else: + steps = torch.ones(length, dtype=torch.long, device=token_types.device) + text_positions = cursor + steps.cumsum(0) - steps[0] + positions[:, batch_index, start:stop] = text_positions + cursor = int(text_positions[-1]) + int(steps[-1]) + return positions + + @staticmethod + def _causal_mask(token_types: Tensor, dtype: torch.dtype) -> tuple[Tensor, Tensor]: + valid = token_types != G05TokenType.PADDING + sequence_length = token_types.shape[1] + causal = torch.ones( + sequence_length, + sequence_length, + dtype=torch.bool, + device=token_types.device, + ).tril() + allowed = causal[None] & valid[:, None, :] & valid[:, :, None] + full = torch.zeros( + token_types.shape[0], + 1, + sequence_length, + sequence_length, + dtype=dtype, + device=token_types.device, + ) + full.masked_fill_(~allowed[:, None], torch.finfo(dtype).min) + return full, valid.to(dtype) + + @staticmethod + def _proprio(samples: list[dict[str, Any]], device: torch.device) -> Tensor: + rows = [] + for sample in samples: + value = sample["proprio"] + value = value["value"] if isinstance(value, Mapping) else value + value = torch.as_tensor(value, dtype=torch.float32, device=device) + rows.append(value.unsqueeze(0) if value.ndim == 1 else value) + return torch.stack(rows) + + def _prefill( + self, + sequence: G05SequenceBatch, + pixel_values: Mapping[str, Tensor], + proprio: Tensor, + ) -> tuple[Tensor, Any, Tensor]: + embeddings = self._embed(sequence, pixel_values, proprio) + positions = self._mrope_positions(sequence.token_types) + full_mask, linear_mask = self._causal_mask(sequence.token_types, embeddings.dtype) + hidden_states, cache = self.model.vlm( + embeddings, + full_attention_mask=full_mask, + linear_attention_mask=linear_mask, + position_ids=positions, + ) + return hidden_states, cache, positions + + def _decode_token( + self, + token_ids: Tensor, + *, + token_types: Tensor, + positions: Tensor, + cache, + ) -> tuple[Tensor, Tensor, Tensor]: + embeddings = self.model.vlm.embed(token_ids[:, None]) + batch_size = token_ids.shape[0] + next_positions = positions.amax(dim=-1, keepdim=True) + 1 + prefix_length = token_types.shape[1] + prefix_mask = (token_types == G05TokenType.PADDING).to(embeddings.dtype) + full_mask = torch.zeros( + batch_size, + 1, + 1, + prefix_length + 1, + dtype=embeddings.dtype, + device=embeddings.device, + ) + full_mask[..., :prefix_length].masked_fill_( + prefix_mask[:, None, None].bool(), torch.finfo(embeddings.dtype).min + ) + hidden_states, cache = self.model.vlm( + embeddings, + full_attention_mask=full_mask, + linear_attention_mask=torch.ones(batch_size, 1, dtype=embeddings.dtype, device=embeddings.device), + position_ids=next_positions, + cache=cache, + ) + next_types = torch.full( + (batch_size, 1), + float(G05TokenType.PRED_TEXT), + dtype=token_types.dtype, + device=token_types.device, + ) + return ( + hidden_states[:, -1], + torch.cat((token_types, next_types), dim=1), + torch.cat((positions, next_positions), dim=-1), + ) + + def _generate_text( + self, + last_hidden: Tensor, + *, + token_types: Tensor, + positions: Tensor, + cache, + max_new_tokens: int, + stop_token_id: int, + ) -> tuple[Tensor, Any, Tensor, Tensor, Tensor]: + generated = [] + finished = torch.zeros(last_hidden.shape[0], dtype=torch.bool, device=last_hidden.device) + for _ in range(max_new_tokens): + logits = self.model.vlm.logits(last_hidden) + next_token = logits.argmax(dim=-1) + next_token = torch.where( + finished, + torch.full_like(next_token, stop_token_id), + next_token, + ) + generated.append(next_token) + last_hidden, token_types, positions = self._decode_token( + next_token, + token_types=token_types, + positions=positions, + cache=cache, + ) + finished |= next_token == stop_token_id + if bool(finished.all()): + break + generated_ids = ( + torch.stack(generated, dim=1) + if generated + else torch.empty(last_hidden.shape[0], 0, dtype=torch.long, device=last_hidden.device) + ) + return generated_ids, cache, last_hidden, token_types, positions + + def _action_cache(self, vlm_cache, prefix_length: int, *, repeats: int = 1): + from transformers import DynamicCache + + cache = DynamicCache(config=self.model.action_expert.config) + layer_types = self.model.vlm.config.layer_types + for layer_index, layer_type in enumerate(layer_types): + if layer_type != "full_attention": + continue + source = vlm_cache.layers[layer_index] + if not source.is_initialized: + continue + key = source.keys[..., :prefix_length, :].detach() + value = source.values[..., :prefix_length, :].detach() + if repeats > 1: + key = key.repeat_interleave(repeats, dim=0) + value = value.repeat_interleave(repeats, dim=0) + cache.layers[layer_index].update(key, value) + return cache + + def _action_mask_and_positions( + self, + token_types: Tensor, + positions: Tensor, + horizon: int, + dtype: torch.dtype, + ) -> tuple[Tensor, Tensor]: + batch_size, prefix_length = token_types.shape + prefix_mask = (token_types == G05TokenType.PADDING).to(dtype) * torch.finfo(dtype).min + action_mask = torch.zeros( + batch_size, + horizon, + horizon, + dtype=dtype, + device=token_types.device, + ) + if bool(self.model_config["fm"].get("action_causal", False)): + action_mask = torch.triu(torch.full_like(action_mask, torch.finfo(dtype).min), diagonal=1) + mask = torch.cat((prefix_mask[:, None].expand(-1, horizon, -1), action_mask), dim=-1).unsqueeze(1) + offset = positions.amax(dim=-1, keepdim=True) + action_positions = torch.arange(1, horizon + 1, device=token_types.device)[None, None] + offset + return mask, action_positions + + def _velocity( + self, + actions: Tensor, + timesteps: Tensor, + *, + vlm_cache, + token_types: Tensor, + positions: Tensor, + ) -> Tensor: + action_embeddings = self.model.action_expert.embed(actions) + time_cond = self.model.action_expert.encode_time(timesteps) + mask, action_positions = self._action_mask_and_positions( + token_types, positions, actions.shape[1], action_embeddings.dtype + ) + cache = self._action_cache(vlm_cache, token_types.shape[1]) + hidden_states = self.model.action_expert( + action_embeddings, + attention_mask=mask, + position_ids=action_positions, + cache=cache, + time_cond=time_cond, + ) + return self.model.action_expert.decode(hidden_states) + + def _infer_flow( + self, + *, + vlm_cache, + token_types: Tensor, + positions: Tensor, + action_dim_is_pad: Tensor | None, + dtype: torch.dtype, + ) -> Tensor: + fm = self.model_config["fm"] + batch_size = token_types.shape[0] + horizon = int(fm["horizon_steps"]) + action_dim = int(fm["action_dim"]) + action = torch.randn( + batch_size, + horizon, + action_dim, + device=token_types.device, + dtype=dtype, + ) + dim_mask = ( + action_dim_is_pad.bool().unsqueeze(1) + if action_dim_is_pad is not None and not bool(fm["zero_pad_action_target"]) + else None + ) + if dim_mask is not None: + action.masked_fill_(dim_mask, 0) + steps = int(fm["num_inference_steps"]) + delta = 1.0 / steps + pi_convention = fm["time_convention"] == "pi_convention" + time_value = 1.0 if pi_convention else 0.0 + timesteps = torch.full((batch_size,), time_value, dtype=dtype, device=token_types.device) + for _ in range(steps): + velocity = self._velocity( + action, + timesteps, + vlm_cache=vlm_cache, + token_types=token_types, + positions=positions, + ) + action = action - delta * velocity if pi_convention else action + delta * velocity + timesteps = timesteps - delta if pi_convention else timesteps + delta + if dim_mask is not None: + action.masked_fill_(dim_mask, 0) + clip = fm.get("final_action_clip_value") + return action.clamp(-float(clip), float(clip)) if clip is not None else action + + def _flow_loss( + self, + actions: Tensor, + *, + action_is_pad: Tensor, + action_dim_is_pad: Tensor | None, + vlm_cache, + token_types: Tensor, + positions: Tensor, + ) -> Tensor: + fm = self.model_config["fm"] + samples = int(fm.get("num_flow_samples", 1)) + batch_size = actions.shape[0] + beta = torch.distributions.Beta(1.5, 1.0) + z = beta.sample((samples, batch_size)).to(actions.device, actions.dtype) + if fm["time_convention"] == "pi_convention": + timesteps = 1 - (1 - float(fm["flow_sig_min"])) * (1 - z) + else: + timesteps = (1 - float(fm["flow_sig_min"])) * (1 - z) + timesteps = timesteps.reshape(-1) + noise = torch.randn( + samples, + *actions.shape, + device=actions.device, + dtype=actions.dtype, + ).flatten(0, 1) + target_actions = actions.repeat(samples, 1, 1) + t = timesteps[:, None, None] + if fm["time_convention"] == "pi_convention": + interpolated = (1 - t) * target_actions + t * noise + target_velocity = noise - target_actions + else: + interpolated = t * target_actions + (1 - t) * noise + target_velocity = target_actions - noise + + repeated_dim_mask = None + if action_dim_is_pad is not None: + repeated_dim_mask = action_dim_is_pad.repeat(samples, 1) + if not bool(fm["zero_pad_action_target"]): + interpolated = interpolated.masked_fill(repeated_dim_mask[:, None], 0) + + repeated_types = token_types.repeat(samples, 1) + repeated_positions = positions.repeat(1, samples, 1) + action_embeddings = self.model.action_expert.embed(interpolated) + time_cond = self.model.action_expert.encode_time(timesteps) + mask, action_positions = self._action_mask_and_positions( + repeated_types, repeated_positions, actions.shape[1], action_embeddings.dtype + ) + cache = self._action_cache(vlm_cache, token_types.shape[1], repeats=samples) + predicted = self.model.action_expert.decode( + self.model.action_expert( + action_embeddings, + attention_mask=mask, + position_ids=action_positions, + cache=cache, + time_cond=time_cond, + ) + ) + weights = torch.ones_like(predicted) + weights[action_is_pad.repeat(samples, 1)] = float(fm["padding_action_weight"]) + if repeated_dim_mask is not None and not bool(fm["zero_pad_action_target"]): + weights.masked_fill_( + repeated_dim_mask[:, None], + float(fm["padding_action_weight"]), + ) + loss = (weights * (predicted - target_velocity).square()).sum() / weights.sum().clamp_min(1) + return loss * float(fm["fm_weight"]) + + def predict_action(self, batch: Mapping[str, Any]) -> dict[str, Any]: + start = time.monotonic() + samples = list(batch["samples"]) + pixel_values = batch["pixel_values"] + first_image = next(iter(pixel_values.values())) + sequence = self.processor.encode_inference(samples, device=first_image.device) + proprio = self._proprio(samples, first_image.device) + hidden_states, cache, positions = self._prefill(sequence, pixel_values, proprio) + result: dict[str, Any] = {} + last_hidden = hidden_states[:, -1] + token_types = sequence.token_types + if bool(self.model_config.get("predict_cot", False)): + generated, cache, last_hidden, token_types, positions = self._generate_text( + last_hidden, + token_types=sequence.token_types, + positions=positions, + cache=cache, + max_new_tokens=int(self.model_config["ar"].get("max_new_tokens", 300)), + stop_token_id=self.processor.eov_token_id, + ) + sequence.token_types = token_types + result["generated_ids"] = generated + result["cot_text"] = [ + self.processor.decode( + ids[ + : next( + ( + index + for index, token_id in enumerate(ids.tolist()) + if token_id == self.processor.eov_token_id + ), + len(ids), + ) + ] + ) + for ids in generated + ] + sequence.token_types = token_types + if bool(self.model_config.get("continuous_action", False)): + dim_mask = batch.get("action_dim_is_pad") + if not isinstance(dim_mask, Tensor): + dim_mask = None + result[ACTION] = self._infer_flow( + vlm_cache=cache, + token_types=sequence.token_types, + positions=positions, + action_dim_is_pad=dim_mask, + dtype=first_image.dtype, + ) + if bool(self.model_config.get("discrete_action", False)): + if self.action_tokenizer is None: + if ACTION not in result: + raise RuntimeError( + "The native G0.5 ActionCodec checkpoint has not been loaded; " + "select the continuous flow head for this checkpoint." + ) + else: + generated_action, _, _, _, _ = self._generate_text( + last_hidden, + token_types=sequence.token_types, + positions=positions, + cache=cache, + max_new_tokens=self.action_tokenizer.action_token_length + 32, + stop_token_id=self.processor.eos_token_id, + ) + decoded_actions = [] + decoded_tokens = [] + absent_keys = [] + for token_row in generated_action: + is_action = (token_row >= self.processor.action_token_begin) & ( + token_row < self.processor.action_token_end_with_markers + ) + action_tokens = token_row[is_action] + decoded, absent = self.action_tokenizer.decode_language_tokens( + action_tokens, + horizon=int(self.model_config["fm"]["horizon_steps"]), + action_dim=int(self.model_config["fm"]["action_dim"]), + ) + decoded_actions.append(decoded) + decoded_tokens.append(action_tokens) + absent_keys.append(absent) + result["ar_action"] = torch.stack(decoded_actions) + result["decoded_action_tokens"] = decoded_tokens + result["ar_absent_keys"] = absent_keys + if ACTION not in result: + result[ACTION] = result["ar_action"] + result["_timing"] = {"forward_inference_total_ms": (time.monotonic() - start) * 1000} + return result + + def forward(self, batch: Mapping[str, Any]) -> tuple[Tensor, dict[str, Tensor]]: + samples = list(batch["samples"]) + pixel_values = batch["pixel_values"] + first_image = next(iter(pixel_values.values())) + sequence = self.processor.encode_train( + samples, + device=first_image.device, + action_codec=self.action_tokenizer, + ) + proprio = self._proprio(samples, first_image.device) + hidden_states, cache, positions = self._prefill(sequence, pixel_values, proprio) + loss_dict: dict[str, Tensor] = {} + + skip_ce = ( + bool(self.model_config.get("continuous_action", False)) + and not bool(self.model_config.get("discrete_action", False)) + and not bool(self.model_config.get("predict_cot", False)) + ) + if not skip_ce: + shift_labels = sequence.labels[:, 1:] + valid = shift_labels != IGNORE_INDEX + if valid.any(): + logits = self.model.vlm.logits(hidden_states[:, :-1]) + loss_dict["ce_loss"] = functional.cross_entropy( + logits.reshape(-1, logits.shape[-1]), + shift_labels.reshape(-1), + ignore_index=IGNORE_INDEX, + ) + else: + loss_dict["ce_loss"] = hidden_states.sum() * 0 + + if bool(self.model_config.get("continuous_action", False)): + actions = batch.get(ACTION) + if not isinstance(actions, Tensor): + raise ValueError("G0.5 flow training requires an action tensor.") + action_is_pad = batch.get("action_is_pad") + if not isinstance(action_is_pad, Tensor): + action_is_pad = torch.zeros(actions.shape[:2], dtype=torch.bool, device=actions.device) + action_dim_is_pad = batch.get("action_dim_is_pad") + if not isinstance(action_dim_is_pad, Tensor): + action_dim_is_pad = None + prefix = int(sequence.split_index) + loss_dict["fm_loss"] = self._flow_loss( + actions, + action_is_pad=action_is_pad, + action_dim_is_pad=action_dim_is_pad, + vlm_cache=cache, + token_types=sequence.token_types[:, :prefix], + positions=positions[..., :prefix], + ) + loss = sum(loss_dict.values()) + return loss, loss_dict diff --git a/src/lerobot/policies/g05/processing_g05.py b/src/lerobot/policies/g05/processing_g05.py new file mode 100644 index 000000000..57dec5e7f --- /dev/null +++ b/src/lerobot/policies/g05/processing_g05.py @@ -0,0 +1,353 @@ +# 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. + +"""Native G0.5 prompt serialization and tokenizer registration.""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch +from torch import Tensor + +IGNORE_INDEX = -100 + + +class G05TokenType: + """Token categories stored in G0.5's attention-mask tensor.""" + + PADDING = 0 + IMAGE = 1 + PROPRIO = 2 + ACTION = 3 + TEXT = 4 + COT = 5 + PRED_TEXT = 6 + + +@dataclass +class G05SequenceBatch: + input_ids: Tensor + labels: Tensor + token_types: Tensor + split_index: int | None = None + + +@dataclass +class _Segment: + kind: str + content: str = "" + sample_key: str = "" + processor: str = "" + masked: bool = False + max_tokens: int | None = None + + +class G05Tokenizer: + """Checkpoint-compatible G0.5 tokenizer and template serializer. + + G0.5 extends Qwen3.5's tokenizer with ActionCodec codes, per-group + residual markers, ````, and ````. Registration order is model + state: changing it changes the rows used by the tied language head. + """ + + _PLACEHOLDER = re.compile(r"<([^<>|]+)>") + + def __init__(self, processor_path: str | Path, model_config: dict[str, Any]) -> None: + from transformers import AutoTokenizer + + self.processor_path = Path(processor_path) + self.tokenizer = AutoTokenizer.from_pretrained( + self.processor_path, + trust_remote_code=False, + local_files_only=True, + ) + self.model_config = model_config + at_config = model_config["AT_CONFIG"] + architecture = at_config["model_arch"] + + self.action_tokens = [f"" for index in range(int(architecture["codebook_size"]))] + parts = list(at_config["parts_meta"]) + rule_patterns = tuple(at_config.get("rule_based_key_patterns") or ()) + self.rule_parts = [name for name in parts if any(pattern in name for pattern in rule_patterns)] + self.neural_parts = [name for name in parts if name not in self.rule_parts] + self.group_tokens = [ + f"<{part}_{residual}>" + for residual in range(int(architecture["n_codebooks"])) + for part in self.neural_parts + ] + [f"<{part}>" for part in self.rule_parts] + self.tokenizer.add_tokens(self.action_tokens + self.group_tokens + ["", ""]) + + self.pad_token_id = int(model_config["pad_token_id"]) + self.eos_token_id = int(model_config["eos_token_id"]) + self.image_token_id = int(model_config["image_token_index"]) + self.vision_start_token_id = int(self.tokenizer.convert_tokens_to_ids("<|vision_start|>")) + self.vision_end_token_id = int(self.tokenizer.convert_tokens_to_ids("<|vision_end|>")) + self.eov_token_id = int(self.tokenizer.convert_tokens_to_ids("")) + self.state_token_id = int(self.tokenizer.convert_tokens_to_ids("")) + self.action_token_begin = int(self.tokenizer.convert_tokens_to_ids(self.action_tokens[0])) + self.action_token_end = self.action_token_begin + len(self.action_tokens) + self.action_token_end_with_markers = self.action_token_end + len(self.group_tokens) + + def __len__(self) -> int: + return len(self.tokenizer) + + def encode_text(self, text: str) -> list[int]: + return self.tokenizer(text, add_special_tokens=False)["input_ids"] + + def decode(self, ids: Tensor | list[int]) -> str: + if isinstance(ids, Tensor): + ids = ids.detach().cpu().tolist() + return self.tokenizer.decode(ids, skip_special_tokens=False) + + @staticmethod + def _resolve_template(template: str) -> str: + replacements = { + "": "", + "": "<|endoftext|>", + "": "", + "": "", + "": "", + } + for placeholder, value in replacements.items(): + template = template.replace(placeholder, value) + return template + + def _parse(self, template: str) -> list[_Segment]: + template = self._resolve_template(template) + segments: list[_Segment] = [] + last = 0 + for match in self._PLACEHOLDER.finditer(template): + if match.start() > last: + segments.append(_Segment("static", template[last : match.start()])) + raw = match.group(1).strip() + if raw in {"EOC", "EOV"}: + segments.append(_Segment("control", raw)) + last = match.end() + continue + + max_tokens = None + limit = re.match(r"^(.+)_(\d+)$", raw) + if limit: + raw, max_tokens = limit.group(1), int(limit.group(2)) + masked = raw.endswith("_!") + key = raw[:-2] if masked else raw + if "_" not in key: + token = f"<{raw}>" + token_id = self.tokenizer.convert_tokens_to_ids(token) + if token_id is None: + raise ValueError(f"Unknown G0.5 template token {token!r}.") + segments.append(_Segment("static", token)) + else: + sample_key, processor = key.rsplit("_", 1) + segments.append( + _Segment( + "dynamic", + sample_key=sample_key, + processor=processor, + masked=masked, + max_tokens=max_tokens, + ) + ) + last = match.end() + if last < len(template): + segments.append(_Segment("static", template[last:])) + return segments + + @staticmethod + def _slice_segments(segments: list[_Segment], mode: str | None, *, pred_eov: bool) -> list[_Segment]: + eoc = next( + ( + index + for index, segment in enumerate(segments) + if segment.kind == "control" and segment.content == "EOC" + ), + None, + ) + eov = next( + ( + index + for index, segment in enumerate(segments) + if segment.kind == "control" and segment.content == "EOV" + ), + None, + ) + + def strip(values: list[_Segment], keep_eov: bool) -> list[_Segment]: + output: list[_Segment] = [] + after_eoc = False + for segment in values: + if segment.kind == "control": + if segment.content == "EOC": + after_eoc = True + elif segment.content == "EOV" and keep_eov: + output.append( + _Segment( + "dynamic", + content="", + processor="text", + masked=not pred_eov, + ) + ) + continue + if after_eoc and segment.kind == "static": + output.append( + _Segment("dynamic", content=segment.content, processor="text", masked=False) + ) + else: + output.append(segment) + return output + + if mode == "context": + return strip(segments if eoc is None else segments[:eoc], keep_eov=True) + if mode == "prefix": + return strip(segments if eov is None else segments[: eov + 1], keep_eov=True) + if mode == "suffix": + return [] if eov is None else strip(segments[eov + 1 :], keep_eov=False) + return strip(segments, keep_eov=True) + + def _serialize_segment( + self, + segment: _Segment, + sample: dict[str, Any], + *, + action_codec: Any | None, + ) -> tuple[list[int], list[int], list[float]]: + if segment.kind == "static": + ids = self.encode_text(segment.content) + return ids, [IGNORE_INDEX] * len(ids), [float(G05TokenType.TEXT)] * len(ids) + + if segment.sample_key: + if segment.sample_key not in sample: + raise KeyError(f"G0.5 prompt is missing sample field {segment.sample_key!r}.") + value = sample[segment.sample_key] + else: + value = segment.content + + if segment.processor == "text": + ids = self.encode_text(value if isinstance(value, str) else str(value)) + token_type = G05TokenType.TEXT if segment.masked else G05TokenType.PRED_TEXT + labels = [IGNORE_INDEX] * len(ids) if segment.masked else ids.copy() + elif segment.processor == "image": + if not isinstance(value, (tuple, list)) or len(value) != 2: + raise ValueError("G0.5 image placeholders require an (height, width) pair.") + height, width = (int(item) for item in value) + vision = self.model_config["vision"] + count = (height // int(vision["patch_size"]) // int(vision["spatial_merge_size"])) * ( + width // int(vision["patch_size"]) // int(vision["spatial_merge_size"]) + ) + ids = [self.vision_start_token_id] + [self.image_token_id] * count + [self.vision_end_token_id] + labels = [IGNORE_INDEX] * len(ids) + types = ( + [float(G05TokenType.TEXT)] + [float(G05TokenType.IMAGE)] * count + [float(G05TokenType.TEXT)] + ) + return ids, labels, types + elif segment.processor == "proprio": + state = value["value"] if isinstance(value, dict) else value + count = 1 if torch.as_tensor(state).ndim <= 1 else int(torch.as_tensor(state).shape[0]) + ids = [self.state_token_id] * count + labels = [IGNORE_INDEX] * count + token_type = G05TokenType.PROPRIO + elif segment.processor == "action": + if action_codec is None: + raise RuntimeError( + "This G0.5 training template includes ActionCodec targets, but the " + "checkpoint has no native ActionCodec sidecar loaded." + ) + ids = action_codec.encode_for_language(value) + labels = ids.copy() + token_type = G05TokenType.ACTION + else: + ids = self.encode_text(value if isinstance(value, str) else str(value)) + labels = [IGNORE_INDEX] * len(ids) if segment.masked else ids.copy() + token_type = G05TokenType.COT + + if segment.max_tokens is not None: + ids = ids[: segment.max_tokens] + labels = labels[: segment.max_tokens] + return ids, labels, [float(token_type)] * len(ids) + + def _serialize( + self, + sample: dict[str, Any], + *, + mode: str | None, + action_codec: Any | None, + ) -> tuple[list[int], list[int], list[float]]: + pred_eov = bool(self.model_config.get("input_preprocessor", {}).get("pred_eov", False)) + segments = self._slice_segments(self._parse(sample["template"]), mode, pred_eov=pred_eov) + ids: list[int] = [] + labels: list[int] = [] + types: list[float] = [] + for segment in segments: + segment_ids, segment_labels, segment_types = self._serialize_segment( + segment, sample, action_codec=action_codec + ) + ids.extend(segment_ids) + labels.extend(segment_labels) + types.extend(segment_types) + return ids, labels, types + + def _pad( + self, + rows: list[tuple[list[int], list[int], list[float]]], + *, + right_align: bool, + device: torch.device, + ) -> G05SequenceBatch: + length = max(len(row[0]) for row in rows) + input_ids = torch.full((len(rows), length), self.pad_token_id, dtype=torch.long, device=device) + labels = torch.full((len(rows), length), IGNORE_INDEX, dtype=torch.long, device=device) + token_types = torch.zeros((len(rows), length), dtype=torch.float32, device=device) + for index, (ids, row_labels, types) in enumerate(rows): + start = length - len(ids) if right_align else 0 + stop = start + len(ids) + input_ids[index, start:stop] = torch.tensor(ids, dtype=torch.long, device=device) + labels[index, start:stop] = torch.tensor(row_labels, dtype=torch.long, device=device) + token_types[index, start:stop] = torch.tensor(types, dtype=torch.float32, device=device) + return G05SequenceBatch(input_ids, labels, token_types) + + def encode_inference( + self, + samples: list[dict[str, Any]], + *, + device: torch.device, + ) -> G05SequenceBatch: + rows = [self._serialize(sample, mode="context", action_codec=None) for sample in samples] + return self._pad(rows, right_align=True, device=device) + + def encode_train( + self, + samples: list[dict[str, Any]], + *, + device: torch.device, + action_codec: Any | None, + ) -> G05SequenceBatch: + prefix_rows = [ + self._serialize(sample, mode="prefix", action_codec=action_codec) for sample in samples + ] + suffix_rows = [ + self._serialize(sample, mode="suffix", action_codec=action_codec) for sample in samples + ] + prefix = self._pad(prefix_rows, right_align=True, device=device) + suffix = self._pad(suffix_rows, right_align=False, device=device) + return G05SequenceBatch( + input_ids=torch.cat((prefix.input_ids, suffix.input_ids), dim=1), + labels=torch.cat((prefix.labels, suffix.labels), dim=1), + token_types=torch.cat((prefix.token_types, suffix.token_types), dim=1), + split_index=prefix.input_ids.shape[1], + ) diff --git a/tests/policies/g05/test_action_codec_g05.py b/tests/policies/g05/test_action_codec_g05.py new file mode 100644 index 000000000..2858a7031 --- /dev/null +++ b/tests/policies/g05/test_action_codec_g05.py @@ -0,0 +1,83 @@ +# 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. + +import torch + +from lerobot.policies.g05.action_codec_g05 import G05NativeActionCodec, _BinarySequenceCodec + + +def _tiny_codec_config() -> dict: + return { + "parts_meta": { + "left_control": 3, + "left_gripper": 1, + "right_control": 3, + "right_gripper": 1, + }, + "rule_based_key_patterns": ["gripper"], + "rule_based_min_block_len": 1, + "rule_based_binarize_threshold": 0.0, + "num_residuals": 2, + "model_arch": { + "horizon": 8, + "horizon_patch_size": 2, + "max_component_dim": 3, + "conv_in_action_kernel": 2, + "encoder_channels": 64, + "c_mults": [1], + "strides": [[1, 1]], + "transformer_depths": [1], + "latent_dim": 16, + "num_heads": 1, + "dim_heads": 64, + "rope_base": 10_000, + "ffn_mult": 2, + "n_codebooks": 2, + "codebook_size": 16, + "codebook_dim": 4, + "use_block_dct": False, + }, + } + + +def test_binary_sequence_codec_roundtrip_repairs_short_middle_runs() -> None: + codec = _BinarySequenceCodec(sequence_length=8, min_block_length=1, vocab_size=16) + values = torch.tensor([[0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0]]) + + tokens = codec.encode(values, threshold=0.5) + decoded = codec.decode(tokens) + + assert tokens.shape == (1, codec.num_tokens) + assert decoded.tolist() == [[0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0]] + + +def test_native_action_codec_language_roundtrip_and_absent_groups() -> None: + codec = G05NativeActionCodec(_tiny_codec_config(), action_token_begin=100) + actions = torch.linspace(-1, 1, 8 * 8).reshape(8, 8) + + token_ids = codec.encode_for_language({"value": actions}) + decoded, absent = codec.decode_language_tokens( + torch.tensor(token_ids), + horizon=8, + action_dim=8, + ) + + assert len(token_ids) == codec.action_token_length + assert decoded.shape == (8, 8) + assert torch.isfinite(decoded).all() + assert absent == set() + assert all(key.startswith("model.") for key in codec.module.state_dict()) + + empty, absent = codec.decode_language_tokens(torch.empty(0, dtype=torch.long), horizon=8, action_dim=8) + assert absent == set(codec.parts) + assert torch.equal(empty, torch.zeros_like(empty)) diff --git a/uv.lock b/uv.lock index d8a5494e5..32376a64b 100644 --- a/uv.lock +++ b/uv.lock @@ -2879,7 +2879,6 @@ all = [ { name = "motorbridge-smart-servo" }, { name = "mypy" }, { name = "num2words" }, - { name = "omegaconf" }, { name = "pandas" }, { name = "peft" }, { name = "placo" }, @@ -3022,7 +3021,7 @@ feetech = [ { name = "pyserial" }, ] g05 = [ - { name = "omegaconf" }, + { name = "transformers" }, ] gamepad = [ { name = "hidapi" }, @@ -3456,7 +3455,6 @@ requires-dist = [ { name = "mypy", marker = "extra == 'dev'", specifier = ">=1.19.1" }, { name = "num2words", marker = "extra == 'smolvla'", specifier = ">=0.5.14,<0.6.0" }, { name = "numpy", specifier = ">=2.0.0,<2.3.0" }, - { name = "omegaconf", marker = "extra == 'g05'", specifier = ">=2.3.0,<3.0.0" }, { name = "onnx", marker = "extra == 'unitree-g1'", specifier = ">=1.16.0,<2.0.0" }, { name = "onnxruntime", marker = "extra == 'unitree-g1'", specifier = ">=1.16.0,<2.0.0" }, { name = "openai", marker = "extra == 'annotations'", specifier = ">=1.40,<2.0" }, @@ -3505,6 +3503,7 @@ requires-dist = [ { name = "torchvision", marker = "sys_platform != 'linux'", specifier = ">=0.22.0,<0.27.0" }, { name = "torchvision", marker = "sys_platform == 'linux'", specifier = ">=0.22.0,<0.27.0", index = "https://download.pytorch.org/whl/cu128" }, { name = "tqdm", specifier = ">=4.66.0,<5.0.0" }, + { name = "transformers", marker = "extra == 'g05'", specifier = ">=5.4.0,<5.6.0" }, { 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" }, ]