From 12e7416473f9e7c577d16fd6eee73d3ae861b6f8 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Thu, 30 Jul 2026 13:00:40 +0200 Subject: [PATCH] refactor(g05): consolidate policy modules --- src/lerobot/policies/g05/action_codec_g05.py | 674 ------- src/lerobot/policies/g05/modeling_g05.py | 1869 +++++++++++++++++- src/lerobot/policies/g05/native_g05.py | 1250 ------------ src/lerobot/policies/g05/processing_g05.py | 354 ---- src/lerobot/policies/g05/processor_g05.py | 332 +++- tests/policies/g05/test_action_codec_g05.py | 2 +- tests/policies/g05/test_g05.py | 3 +- 7 files changed, 2200 insertions(+), 2284 deletions(-) delete mode 100644 src/lerobot/policies/g05/action_codec_g05.py delete mode 100644 src/lerobot/policies/g05/native_g05.py delete mode 100644 src/lerobot/policies/g05/processing_g05.py diff --git a/src/lerobot/policies/g05/action_codec_g05.py b/src/lerobot/policies/g05/action_codec_g05.py deleted file mode 100644 index d0c2331c5..000000000 --- a/src/lerobot/policies/g05/action_codec_g05.py +++ /dev/null @@ -1,674 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2026 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__() - layer_scale_init = float(config.get("layer_scale_init", 1.0)) - self.ls1 = nn.Parameter(torch.full((dimension,), layer_scale_init)) - self.ls2 = nn.Parameter(torch.full((dimension,), layer_scale_init)) - 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/modeling_g05.py b/src/lerobot/policies/g05/modeling_g05.py index f807f7a5e..9de4c362d 100644 --- a/src/lerobot/policies/g05/modeling_g05.py +++ b/src/lerobot/policies/g05/modeling_g05.py @@ -14,23 +14,39 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""First-class LeRobot wrapper around the community-licensed G0.5 implementation.""" +"""Native G0.5 policy, model, and ActionCodec implementation.""" from __future__ import annotations +import itertools import json +import math import shutil +import time from collections import deque from collections.abc import Mapping from pathlib import Path from typing import Any import torch +import torch.nn.functional as functional from huggingface_hub import snapshot_download from torch import Tensor, nn +from transformers import DynamicCache +from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig, Qwen3_5VisionConfig +from transformers.models.qwen3_5.modeling_qwen3_5 import ( + Qwen3_5Attention, + Qwen3_5DecoderLayer, + Qwen3_5MLP, + Qwen3_5RMSNorm, + Qwen3_5TextRotaryEmbedding, + Qwen3_5VisionModel, + apply_rotary_pos_emb_vision, +) from lerobot.configs.policies import PreTrainedConfig from lerobot.optim.optimizers import OptimizerParams +from lerobot.policies.pi_gemma import PiGemmaRMSNorm from lerobot.policies.pretrained import PreTrainedPolicy from lerobot.utils.constants import ACTION, OBS_STATE @@ -40,7 +56,1856 @@ from .configuration_g05 import ( make_g05_cot_prompt_template, make_g05_prompt_template, ) -from .native_g05 import G05_RUNTIME_PREDICT_COT, G05NativeBackend +from .processor_g05 import IGNORE_INDEX, G05SequenceBatch, G05Tokenizer, G05TokenType + + +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__() + layer_scale_init = float(config.get("layer_scale_init", 1.0)) + self.ls1 = nn.Parameter(torch.full((dimension,), layer_scale_init)) + self.ls2 = nn.Parameter(torch.full((dimension,), layer_scale_init)) + 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 + + +G05_RUNTIME_PREDICT_COT = "g05_runtime_predict_cot" + + +def _autoregressive_ce_loss( + logits: Tensor, + labels: Tensor, + *, + ce_weight: float, + z_loss_scale: float, +) -> Tensor: + """Apply G0.5's checkpoint-configured autoregressive objective.""" + + if ce_weight < 0: + raise ValueError("G0.5 ar.ce_weight must be non-negative.") + if z_loss_scale < 0: + raise ValueError("G0.5 ar.ce_z_loss_scale must be non-negative.") + if logits.ndim != 2 or labels.ndim != 1 or logits.shape[0] != labels.shape[0]: + raise ValueError( + "G0.5 autoregressive CE expects logits [N,V] and labels [N], " + f"got {tuple(logits.shape)} and {tuple(labels.shape)}." + ) + if logits.shape[0] == 0 or ce_weight == 0: + return logits.sum() * 0 + + token_loss = functional.cross_entropy(logits, labels, reduction="none") + if z_loss_scale: + log_z = torch.logsumexp(logits.float(), dim=-1) + token_loss = token_loss.float() + z_loss_scale * log_z.square() + return token_loss.mean() * ce_weight + + +def _qwen_text_config(values: Mapping[str, Any], *, vocab_size: int | None = None): + """Translate the serialized G0.5 Qwen config into a Transformers config.""" + + 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.""" + + 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__() + + 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]: + 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__() + + 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__() + + 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__() + + 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}") + + 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: + 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: + 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): + 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 + predict_cot = bool(batch.get(G05_RUNTIME_PREDICT_COT, self.model_config.get("predict_cot", False))) + if predict_cot: + 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] = {} + + ar_config = self.model_config.get("ar") or {} + ce_weight = float(ar_config.get("ce_weight", 1.0)) + z_loss_scale = float(ar_config.get("ce_z_loss_scale", 0.0)) + if ce_weight < 0: + raise ValueError("G0.5 ar.ce_weight must be non-negative.") + if z_loss_scale < 0: + raise ValueError("G0.5 ar.ce_z_loss_scale must be non-negative.") + 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() and ce_weight: + shift_hidden = hidden_states[:, :-1].reshape(-1, hidden_states.shape[-1]) + valid_hidden = shift_hidden[valid.reshape(-1)] + valid_labels = shift_labels.reshape(-1)[valid.reshape(-1)] + loss_dict["ce_loss"] = _autoregressive_ce_loss( + self.model.vlm.logits(valid_hidden), + valid_labels, + ce_weight=ce_weight, + z_loss_scale=z_loss_scale, + ) + 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 def _native_backend(config: G05Config) -> nn.Module: diff --git a/src/lerobot/policies/g05/native_g05.py b/src/lerobot/policies/g05/native_g05.py deleted file mode 100644 index dfed2f23c..000000000 --- a/src/lerobot/policies/g05/native_g05.py +++ /dev/null @@ -1,1250 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2026 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 itertools -import json -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 transformers import DynamicCache -from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig, Qwen3_5VisionConfig -from transformers.models.qwen3_5.modeling_qwen3_5 import ( - Qwen3_5Attention, - Qwen3_5DecoderLayer, - Qwen3_5MLP, - Qwen3_5RMSNorm, - Qwen3_5TextRotaryEmbedding, - Qwen3_5VisionModel, - apply_rotary_pos_emb_vision, -) - -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 - -G05_RUNTIME_PREDICT_COT = "g05_runtime_predict_cot" - - -def _autoregressive_ce_loss( - logits: Tensor, - labels: Tensor, - *, - ce_weight: float, - z_loss_scale: float, -) -> Tensor: - """Apply G0.5's checkpoint-configured autoregressive objective.""" - - if ce_weight < 0: - raise ValueError("G0.5 ar.ce_weight must be non-negative.") - if z_loss_scale < 0: - raise ValueError("G0.5 ar.ce_z_loss_scale must be non-negative.") - if logits.ndim != 2 or labels.ndim != 1 or logits.shape[0] != labels.shape[0]: - raise ValueError( - "G0.5 autoregressive CE expects logits [N,V] and labels [N], " - f"got {tuple(logits.shape)} and {tuple(labels.shape)}." - ) - if logits.shape[0] == 0 or ce_weight == 0: - return logits.sum() * 0 - - token_loss = functional.cross_entropy(logits, labels, reduction="none") - if z_loss_scale: - log_z = torch.logsumexp(logits.float(), dim=-1) - token_loss = token_loss.float() + z_loss_scale * log_z.square() - return token_loss.mean() * ce_weight - - -def _qwen_text_config(values: Mapping[str, Any], *, vocab_size: int | None = None): - """Translate the serialized G0.5 Qwen config into a Transformers config.""" - - 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.""" - - 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__() - - 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]: - 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__() - - 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__() - - 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__() - - 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}") - - 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: - 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: - 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): - 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 - predict_cot = bool(batch.get(G05_RUNTIME_PREDICT_COT, self.model_config.get("predict_cot", False))) - if predict_cot: - 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] = {} - - ar_config = self.model_config.get("ar") or {} - ce_weight = float(ar_config.get("ce_weight", 1.0)) - z_loss_scale = float(ar_config.get("ce_z_loss_scale", 0.0)) - if ce_weight < 0: - raise ValueError("G0.5 ar.ce_weight must be non-negative.") - if z_loss_scale < 0: - raise ValueError("G0.5 ar.ce_z_loss_scale must be non-negative.") - 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() and ce_weight: - shift_hidden = hidden_states[:, :-1].reshape(-1, hidden_states.shape[-1]) - valid_hidden = shift_hidden[valid.reshape(-1)] - valid_labels = shift_labels.reshape(-1)[valid.reshape(-1)] - loss_dict["ce_loss"] = _autoregressive_ce_loss( - self.model.vlm.logits(valid_hidden), - valid_labels, - ce_weight=ce_weight, - z_loss_scale=z_loss_scale, - ) - 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 deleted file mode 100644 index cecaefc2b..000000000 --- a/src/lerobot/policies/g05/processing_g05.py +++ /dev/null @@ -1,354 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2026 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 -from transformers import AutoTokenizer - -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: - 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/src/lerobot/policies/g05/processor_g05.py b/src/lerobot/policies/g05/processor_g05.py index ad7f6209e..853974402 100644 --- a/src/lerobot/policies/g05/processor_g05.py +++ b/src/lerobot/policies/g05/processor_g05.py @@ -14,16 +14,19 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Serializable preprocessing and inverse projection for G0.5.""" +"""G0.5 tokenization, serialization, preprocessing, and inverse projection.""" from __future__ import annotations +import re from dataclasses import dataclass, field from pathlib import Path from typing import Any import torch import torchvision.transforms.functional as vision_functional +from torch import Tensor +from transformers import AutoTokenizer from lerobot.configs import recipe as recipe_module from lerobot.configs.recipe import TrainingRecipe @@ -851,3 +854,330 @@ def make_g05_pre_post_processors( to_output=transition_to_policy_action, ) return preprocessor, postprocessor + + +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: + 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 index 7cfd523a9..3512c7912 100644 --- a/tests/policies/g05/test_action_codec_g05.py +++ b/tests/policies/g05/test_action_codec_g05.py @@ -16,7 +16,7 @@ import torch -from lerobot.policies.g05.action_codec_g05 import G05NativeActionCodec, _BinarySequenceCodec +from lerobot.policies.g05.modeling_g05 import G05NativeActionCodec, _BinarySequenceCodec def _tiny_codec_config() -> dict: diff --git a/tests/policies/g05/test_g05.py b/tests/policies/g05/test_g05.py index 08a4a2b07..b096a49f2 100644 --- a/tests/policies/g05/test_g05.py +++ b/tests/policies/g05/test_g05.py @@ -28,8 +28,7 @@ from lerobot.configs.policies import PreTrainedConfig from lerobot.configs.types import FeatureType, PolicyFeature from lerobot.policies.factory import get_policy_class, make_policy_config, make_pre_post_processors from lerobot.policies.g05.configuration_g05 import G05_CAMERA_PROFILES, G05_EMBODIMENT_MAPPINGS, G05Config -from lerobot.policies.g05.modeling_g05 import G05Policy -from lerobot.policies.g05.native_g05 import G05_RUNTIME_PREDICT_COT, G05NativeBackend +from lerobot.policies.g05.modeling_g05 import G05_RUNTIME_PREDICT_COT, G05NativeBackend, G05Policy from lerobot.processor import PolicyProcessorPipeline from lerobot.utils.constants import ACTION, OBS_STATE, POLICY_PREPROCESSOR_DEFAULT_NAME