mirror of
https://github.com/huggingface/lerobot.git
synced 2026-08-08 17:39:44 +00:00
refactor(g05): consolidate policy modules
This commit is contained in:
Pepijn
@@ -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
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -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, ``<EOV>``, and ``<state>``. 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"<action{index:04d}>" 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 + ["<EOV>", "<state>"])
|
||||
|
||||
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("<EOV>"))
|
||||
self.state_token_id = int(self.tokenizer.convert_tokens_to_ids("<state>"))
|
||||
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 = {
|
||||
"<bos>": "",
|
||||
"<eos>": "<|endoftext|>",
|
||||
"<chat_user_prefix>": "",
|
||||
"<chat_user_suffix>": "",
|
||||
"<chat_assistant_prefix>": "",
|
||||
}
|
||||
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="<EOV>",
|
||||
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],
|
||||
)
|
||||
@@ -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, ``<EOV>``, and ``<state>``. 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"<action{index:04d}>" 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 + ["<EOV>", "<state>"])
|
||||
|
||||
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("<EOV>"))
|
||||
self.state_token_id = int(self.tokenizer.convert_tokens_to_ids("<state>"))
|
||||
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 = {
|
||||
"<bos>": "",
|
||||
"<eos>": "<|endoftext|>",
|
||||
"<chat_user_prefix>": "",
|
||||
"<chat_user_suffix>": "",
|
||||
"<chat_assistant_prefix>": "",
|
||||
}
|
||||
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="<EOV>",
|
||||
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],
|
||||
)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user