mirror of
https://github.com/huggingface/lerobot.git
synced 2026-08-08 17:39:44 +00:00
feat(g05): make policy runtime native
This commit is contained in:
+27
-40
@@ -9,15 +9,16 @@ conditioned on the same post-reasoning KV state.
|
||||
Transformers includes the native multimodal Qwen3.5 backbone, vision tower, and
|
||||
processor. G0.5 is not a stock `Qwen3_5ForConditionalGeneration` checkpoint,
|
||||
however: it adds the proprioception/action path, action expert, flow-matching
|
||||
head, ActionCodec, and unified CoT/action decode. The current integration
|
||||
therefore keeps the pinned G0.5 author package as the model backend.
|
||||
head, ActionCodec, and unified CoT/action decode. LeRobot implements those G0.5
|
||||
components natively and loads converted checkpoints without the OpenGalaxea
|
||||
Python package, Hydra, or OmegaConf.
|
||||
|
||||
> [!WARNING]
|
||||
> G0.5 code and checkpoints use the
|
||||
> [G0.5 Community License](https://huggingface.co/OpenGalaxea/G05/blob/main/licenses/LICENSE-G0.5),
|
||||
> including non-commercial restrictions. LeRobot does not vendor the author model,
|
||||
> download gated files, or imply that Apache-2.0 applies to those materials. Accept
|
||||
> the license yourself and use a private or local checkpoint.
|
||||
> including non-commercial restrictions. LeRobot does not redistribute checkpoint
|
||||
> weights, download gated files, or imply that Apache-2.0 applies to those materials.
|
||||
> Accept the license yourself and use a private or local checkpoint.
|
||||
|
||||
## Supported checkpoint contracts
|
||||
|
||||
@@ -29,8 +30,8 @@ therefore keeps the pinned G0.5 author package as the model backend.
|
||||
| `g05-so101` | Flow or AR ActionCodec + native CoT | right arm joints 6 → 20D grouped layout | exterior, optional left + right wrist | 32 | 16 | stepwise q01/q99 |
|
||||
|
||||
Each packaged checkpoint stores the resolved model and processor configuration,
|
||||
ActionCodec metadata, statistics, exact prompt template, source revision, and
|
||||
license. Loading rejects a different head, horizon, processor mode, or
|
||||
ActionCodec metadata, statistics, exact prompt template, and license. Loading
|
||||
rejects a different head, horizon, processor mode, or
|
||||
normalization contract.
|
||||
|
||||
The converted checkpoints are private under the LeRobot organization:
|
||||
@@ -51,23 +52,12 @@ Atomic-4 dataset; reusing the R1 Lite statistics would be incorrect.
|
||||
|
||||
## Install
|
||||
|
||||
Install LeRobot's small config dependency:
|
||||
Install LeRobot with the G0.5 Transformers dependency:
|
||||
|
||||
```bash
|
||||
uv sync --extra g05 --extra test
|
||||
```
|
||||
|
||||
Then clone the audited author source and install it only after reviewing and
|
||||
accepting its license. The author package currently declares Python 3.10 while
|
||||
LeRobot uses Python 3.12, so a compatible deployment environment or an upstream
|
||||
Python-support update is required for real-model execution.
|
||||
|
||||
```bash
|
||||
git clone https://github.com/OpenGalaxea/GalaxeaVLA.git
|
||||
git -C GalaxeaVLA checkout b34966f387dd2ae0f003143b81494afd9213e613
|
||||
export PYTHONPATH="/path/to/GalaxeaVLA/src:${PYTHONPATH}"
|
||||
```
|
||||
|
||||
The LeRobot organization hosts the prepared base, LIBERO, RoboTwin, and SO-101
|
||||
checkpoints privately. Authenticate with `hf auth login` before loading them.
|
||||
SO-100 and SO-101 share the released `so100` embodiment token and six-joint
|
||||
@@ -124,9 +114,10 @@ the checkpoint's native System 2 CoT telemetry.
|
||||
|
||||
## Fine-tune with `lerobot-train`
|
||||
|
||||
G0.5 implements LeRobot's training surface: `forward` runs the author training
|
||||
backend, the policy exposes the author VLM/vision/action optimizer groups, and
|
||||
the checkpoint can be saved, resumed, and loaded by the normal LeRobot scripts.
|
||||
G0.5 implements LeRobot's training surface natively: `forward` computes
|
||||
assistant-token cross entropy and flow-matching loss, the policy exposes
|
||||
VLM/vision/action optimizer groups, and the checkpoint can be saved, resumed,
|
||||
and loaded by the normal LeRobot scripts.
|
||||
For example, fine-tune the private SO-101 checkpoint on a LeRobot dataset:
|
||||
|
||||
```bash
|
||||
@@ -157,7 +148,7 @@ Grounded VQA boxes are converted from pixel-space `xyxy` JSON using the source
|
||||
camera dimensions captured before image resizing, then serialized as G0.5
|
||||
`<locXXXX>` tokens. Joint samples preserve the released checkpoint's
|
||||
`BBox → Subtask → Action` order. The user/task conditioning tokens remain
|
||||
masked; the author backend applies its language/action objective to the
|
||||
masked; the native backend applies its language/action objective to the
|
||||
assistant sequence.
|
||||
|
||||
Generate the required `subtask` and grounded `vqa` language columns with
|
||||
@@ -197,30 +188,26 @@ lerobot-train \
|
||||
|
||||
CPU unit tests cover factory loading, config incompatibilities, prompt pass-through,
|
||||
LIBERO and `atomic_4` mappings, padding masks, inverse action projection, a finite
|
||||
forward/backward/update, author optimizer-group wiring, and save/reload parity:
|
||||
forward/backward/update, optimizer-group wiring, and save/reload parity:
|
||||
|
||||
```bash
|
||||
uv run pytest tests/policies/g05 tests/runtime/test_g05_adapter.py -q
|
||||
uv run ruff check src/lerobot/policies/g05 tests/policies/g05
|
||||
```
|
||||
|
||||
An RTX 5090 smoke test loaded the real private `g05_libero` checkpoint and
|
||||
completed a batch-size-one BF16 forward, backward, gradient clip, and AdamW
|
||||
step. It produced finite loss `2.77356`, finite pre-clip gradient norm `54.38`,
|
||||
all six author optimizer groups, and 23.59 GiB peak allocated CUDA memory.
|
||||
The converted private `g05_base` checkpoint strict-loaded all 945 native model
|
||||
tensors, and its ActionCodec sidecar strict-loaded all 208 tensors, with no
|
||||
missing, unexpected, or shape-mismatched keys. Prompt token IDs and masks match
|
||||
the released runtime exactly. ActionCodec code IDs match exactly; decoded values
|
||||
differ only by normal floating-point noise. A batch-size-one System 1 flow smoke
|
||||
produced a finite `[1, 32, 27]` action chunk on an RTX 5090.
|
||||
|
||||
The same RTX 5090 also loaded the private `g05_so101` checkpoint and completed
|
||||
a real joint `BBox → Subtask → Action` forward and backward through the
|
||||
recipe-driven path. It produced finite total loss `4.38332`, including non-zero
|
||||
`ce_loss=4.04244` and `fm_loss=0.340881`, with 15.58 GiB peak allocated CUDA
|
||||
memory.
|
||||
|
||||
The converted private `g05_base` checkpoint strict-loaded all 946 mapped tensors
|
||||
with no missing, duplicate, unexpected, or shape-mismatched keys. Author/LeRobot
|
||||
parity was exact for images, tokens, and masks; normalized state/action inputs
|
||||
matched within `1.2e-7`, and postprocessing within `1.8e-6`. A System 2 smoke
|
||||
generated a native `Subtask:` trace and a finite `[1, 32, 27]` action chunk in
|
||||
the same pass, using 11.02 GiB peak allocated memory.
|
||||
The same native base checkpoint completed a real joint language/action forward
|
||||
and backward on the RTX 5090. It produced finite `ce_loss=13.8141` and
|
||||
`fm_loss=1.32647`, finite gradients for all 945 model tensors, all six optimizer
|
||||
groups, and 19.32 GiB peak allocated CUDA memory. A System 2 smoke generated
|
||||
`Subtask: grasp and lift the red cup with the right gripper` and a finite
|
||||
same-pass `[1, 32, 27]` action in 1.07 seconds at 11.0 GiB peak.
|
||||
|
||||
A 50-episode LIBERO/RoboTwin success-rate comparison additionally requires the
|
||||
matching simulator, task assets, reset seeds, and author evaluator; no task-level
|
||||
|
||||
+3
-4
@@ -238,10 +238,9 @@ evo1 = ["lerobot[transformers-dep]"]
|
||||
hilserl = ["lerobot[transformers-dep]", "lerobot[dataset]", "gym-hil>=0.1.14,<0.2.0", "lerobot[grpcio-dep]", "lerobot[placo-dep]"]
|
||||
vla_jepa = ["lerobot[transformers-dep]", "lerobot[diffusers-dep]", "lerobot[qwen-vl-utils-dep]"]
|
||||
lingbot_va = ["lerobot[transformers-dep]", "lerobot[diffusers-dep]", "lerobot[accelerate-dep]"]
|
||||
# The G0.5 model implementation and weights carry Galaxea's separate community
|
||||
# license and are intentionally not redistributed as an Apache LeRobot dependency.
|
||||
# This extra supplies only the config runtime used by converted local checkpoints.
|
||||
g05 = ["omegaconf>=2.3.0,<3.0.0"]
|
||||
# Checkpoint weights retain their upstream licenses. The runtime itself is native
|
||||
# LeRobot code built on Transformers' Qwen3.5 implementation.
|
||||
g05 = ["transformers>=5.4.0,<5.6.0"]
|
||||
|
||||
# Features
|
||||
async = ["lerobot[grpcio-dep]", "lerobot[matplotlib-dep]"]
|
||||
|
||||
@@ -0,0 +1,671 @@
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Native inference implementation of G0.5's ActionCodec sidecar."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as functional
|
||||
from torch import Tensor, nn
|
||||
|
||||
|
||||
class _BlockDCT(nn.Module):
|
||||
def __init__(self, block_size: int) -> None:
|
||||
super().__init__()
|
||||
self.block_size = block_size
|
||||
frequency = torch.arange(block_size, dtype=torch.float32)
|
||||
time = torch.arange(block_size, dtype=torch.float32)
|
||||
basis = torch.cos(math.pi / block_size * (time + 0.5).unsqueeze(0) * frequency.unsqueeze(1))
|
||||
basis[0] *= math.sqrt(1 / block_size)
|
||||
basis[1:] *= math.sqrt(2 / block_size)
|
||||
self.register_buffer("basis", basis, persistent=False)
|
||||
|
||||
def dct(self, values: Tensor) -> Tensor:
|
||||
batch, horizon, dimension = values.shape
|
||||
pad = (-horizon) % self.block_size
|
||||
if pad:
|
||||
values = functional.pad(values, (0, 0, 0, pad))
|
||||
blocks = values.shape[1] // self.block_size
|
||||
values = values.reshape(batch * blocks, self.block_size, dimension)
|
||||
transformed = torch.einsum("kn,bnd->bkd", self.basis.to(values), values)
|
||||
return transformed.reshape(batch, blocks * self.block_size, dimension)
|
||||
|
||||
def idct(self, values: Tensor, horizon: int) -> Tensor:
|
||||
batch, padded_horizon, dimension = values.shape
|
||||
blocks = padded_horizon // self.block_size
|
||||
values = values.reshape(batch * blocks, self.block_size, dimension)
|
||||
restored = torch.einsum("nk,bkd->bnd", self.basis.to(values), values)
|
||||
return restored.reshape(batch, padded_horizon, dimension)[:, :horizon]
|
||||
|
||||
|
||||
def _rotate_half(values: Tensor) -> Tensor:
|
||||
first, second = values.chunk(2, dim=-1)
|
||||
return torch.cat((-second, first), dim=-1)
|
||||
|
||||
|
||||
class _CodecAttention(nn.Module):
|
||||
def __init__(self, dimension: int, num_heads: int, head_dim: int, rope_base: int) -> None:
|
||||
super().__init__()
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = head_dim
|
||||
inner_dim = num_heads * head_dim
|
||||
self.to_qkv = nn.Linear(dimension, inner_dim * 3, bias=False)
|
||||
self.to_out = nn.Linear(inner_dim, dimension, bias=False)
|
||||
self.q_norm = nn.LayerNorm(head_dim, eps=1e-6)
|
||||
self.k_norm = nn.LayerNorm(head_dim, eps=1e-6)
|
||||
rope_dim = max(head_dim // 2, 32)
|
||||
inverse = 1 / (rope_base ** (torch.arange(0, rope_dim, 2, dtype=torch.float32) / rope_dim))
|
||||
self.register_buffer("_inverse_frequency", inverse, persistent=False)
|
||||
|
||||
def forward(self, hidden_states: Tensor) -> Tensor:
|
||||
batch, sequence_length, _ = hidden_states.shape
|
||||
query, key, value = self.to_qkv(hidden_states).chunk(3, dim=-1)
|
||||
|
||||
def heads(values: Tensor) -> Tensor:
|
||||
return values.view(batch, sequence_length, self.num_heads, self.head_dim).transpose(1, 2)
|
||||
|
||||
query, key, value = (heads(values) for values in (query, key, value))
|
||||
query, key = self.q_norm(query), self.k_norm(key)
|
||||
time = torch.arange(sequence_length, device=hidden_states.device, dtype=torch.float32)
|
||||
phase = torch.outer(time, self._inverse_frequency.to(hidden_states.device))
|
||||
phase = torch.cat((phase, phase), dim=-1).to(hidden_states.dtype)[None, None]
|
||||
cosine, sine = phase.cos(), phase.sin()
|
||||
rotary_dim = cosine.shape[-1]
|
||||
query_rotary, query_pass = query[..., :rotary_dim], query[..., rotary_dim:]
|
||||
key_rotary, key_pass = key[..., :rotary_dim], key[..., rotary_dim:]
|
||||
query = torch.cat((query_rotary * cosine + _rotate_half(query_rotary) * sine, query_pass), dim=-1)
|
||||
key = torch.cat((key_rotary * cosine + _rotate_half(key_rotary) * sine, key_pass), dim=-1)
|
||||
attended = functional.scaled_dot_product_attention(query, key, value)
|
||||
attended = attended.transpose(1, 2).reshape(batch, sequence_length, -1)
|
||||
return self.to_out(attended)
|
||||
|
||||
|
||||
class _CodecFFN(nn.Module):
|
||||
def __init__(self, dimension: int, multiplier: float) -> None:
|
||||
super().__init__()
|
||||
inner_dim = int(dimension * multiplier)
|
||||
self.w_up = nn.Linear(dimension, inner_dim * 2, bias=False)
|
||||
self.w_down = nn.Linear(inner_dim, dimension, bias=False)
|
||||
|
||||
def forward(self, hidden_states: Tensor) -> Tensor:
|
||||
value, gate = self.w_up(hidden_states).chunk(2, dim=-1)
|
||||
return self.w_down(value * functional.gelu(gate))
|
||||
|
||||
|
||||
class _CodecTransformerLayer(nn.Module):
|
||||
def __init__(self, dimension: int, config: Mapping[str, Any]) -> None:
|
||||
super().__init__()
|
||||
self.ls1 = nn.Parameter(torch.empty(dimension))
|
||||
self.ls2 = nn.Parameter(torch.empty(dimension))
|
||||
self.norm1 = nn.LayerNorm(dimension, eps=1e-6)
|
||||
self.attn = _CodecAttention(
|
||||
dimension,
|
||||
int(config["num_heads"]),
|
||||
int(config["dim_heads"]),
|
||||
int(config["rope_base"]),
|
||||
)
|
||||
self.norm2 = nn.LayerNorm(dimension, eps=1e-6)
|
||||
self.ffn = _CodecFFN(dimension, float(config["ffn_mult"]))
|
||||
|
||||
def forward(self, hidden_states: Tensor) -> Tensor:
|
||||
hidden_states = hidden_states + self.attn(self.norm1(hidden_states)) * self.ls1
|
||||
return hidden_states + self.ffn(self.norm2(hidden_states)) * self.ls2
|
||||
|
||||
|
||||
class _CodecDownBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
input_channels: int,
|
||||
output_channels: int,
|
||||
stride: tuple[int, int],
|
||||
depth: int,
|
||||
config: Mapping[str, Any],
|
||||
) -> None:
|
||||
super().__init__()
|
||||
stride_h, stride_a = stride
|
||||
if stride_h > 1 or input_channels != output_channels:
|
||||
kernel_h = 2 * stride_h if stride_h > 1 else 1
|
||||
self.conv = nn.Conv2d(
|
||||
input_channels,
|
||||
output_channels,
|
||||
kernel_size=(kernel_h, 1),
|
||||
stride=(stride_h, stride_a),
|
||||
padding=(kernel_h // 2 - int(stride_h > 1), 0),
|
||||
)
|
||||
else:
|
||||
self.conv = nn.Identity()
|
||||
self.transformer_layers = nn.ModuleList(
|
||||
[_CodecTransformerLayer(output_channels, config) for _ in range(depth)]
|
||||
)
|
||||
|
||||
def forward(self, hidden_states: Tensor) -> Tensor:
|
||||
hidden_states = self.conv(hidden_states)
|
||||
batch, channels, height, action_dim = hidden_states.shape
|
||||
sequence = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * action_dim, channels)
|
||||
for layer in self.transformer_layers:
|
||||
sequence = layer(sequence)
|
||||
return sequence.reshape(batch, height, action_dim, channels).permute(0, 3, 1, 2)
|
||||
|
||||
|
||||
class _CodecUpBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
input_channels: int,
|
||||
output_channels: int,
|
||||
stride: tuple[int, int],
|
||||
depth: int,
|
||||
config: Mapping[str, Any],
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.transformer_layers = nn.ModuleList(
|
||||
[_CodecTransformerLayer(input_channels, config) for _ in range(depth)]
|
||||
)
|
||||
stride_h, stride_a = stride
|
||||
if stride_h > 1 or input_channels != output_channels:
|
||||
kernel_h = 2 * stride_h if stride_h > 1 else 1
|
||||
self.conv = nn.ConvTranspose2d(
|
||||
input_channels,
|
||||
output_channels,
|
||||
kernel_size=(kernel_h, 1),
|
||||
stride=(stride_h, stride_a),
|
||||
padding=(kernel_h // 2 - int(stride_h > 1), 0),
|
||||
)
|
||||
else:
|
||||
self.conv = nn.Identity()
|
||||
|
||||
def forward(self, hidden_states: Tensor) -> Tensor:
|
||||
batch, channels, height, action_dim = hidden_states.shape
|
||||
sequence = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * action_dim, channels)
|
||||
for layer in self.transformer_layers:
|
||||
sequence = layer(sequence)
|
||||
hidden_states = sequence.reshape(batch, height, action_dim, channels).permute(0, 3, 1, 2)
|
||||
return self.conv(hidden_states)
|
||||
|
||||
|
||||
class _CodecEncoder(nn.Module):
|
||||
def __init__(self, config: Mapping[str, Any]) -> None:
|
||||
super().__init__()
|
||||
base = int(config["encoder_channels"])
|
||||
channel_dims = [base * int(multiplier) for multiplier in config["c_mults"]]
|
||||
dims = [base] + channel_dims
|
||||
self.blocks = nn.ModuleList(
|
||||
[
|
||||
_CodecDownBlock(
|
||||
dims[index],
|
||||
dims[index + 1],
|
||||
tuple(stride),
|
||||
int(config["transformer_depths"][index]),
|
||||
config,
|
||||
)
|
||||
for index, stride in enumerate(config["strides"])
|
||||
]
|
||||
)
|
||||
self.out_proj = nn.Conv2d(dims[-1], int(config["latent_dim"]), kernel_size=1)
|
||||
|
||||
def forward(self, hidden_states: Tensor) -> Tensor:
|
||||
for block in self.blocks:
|
||||
hidden_states = block(hidden_states)
|
||||
return self.out_proj(hidden_states)
|
||||
|
||||
|
||||
class _CodecDecoder(nn.Module):
|
||||
def __init__(self, config: Mapping[str, Any]) -> None:
|
||||
super().__init__()
|
||||
base = int(config["encoder_channels"])
|
||||
channel_dims = [base * int(multiplier) for multiplier in config["c_mults"]]
|
||||
dims = [base] + channel_dims
|
||||
self.in_proj = nn.Conv2d(int(config["latent_dim"]), dims[-1], kernel_size=1)
|
||||
self.blocks = nn.ModuleList(
|
||||
[
|
||||
_CodecUpBlock(
|
||||
input_channels,
|
||||
output_channels,
|
||||
tuple(stride),
|
||||
int(depth),
|
||||
config,
|
||||
)
|
||||
for stride, depth, input_channels, output_channels in zip(
|
||||
reversed(config["strides"]),
|
||||
reversed(config["transformer_depths"]),
|
||||
reversed(dims[1:]),
|
||||
reversed(dims[:-1]),
|
||||
strict=True,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
def forward(self, hidden_states: Tensor) -> Tensor:
|
||||
hidden_states = self.in_proj(hidden_states)
|
||||
for block in self.blocks:
|
||||
hidden_states = block(hidden_states)
|
||||
return hidden_states
|
||||
|
||||
|
||||
class _CodecQuantizer(nn.Module):
|
||||
def __init__(self, config: Mapping[str, Any]) -> None:
|
||||
super().__init__()
|
||||
input_dim = int(config["latent_dim"])
|
||||
codebook_dim = int(config["codebook_dim"])
|
||||
codebook_size = int(config["codebook_size"])
|
||||
self.input_dim = input_dim
|
||||
self.in_proj = nn.Linear(input_dim, codebook_dim, bias=False)
|
||||
self.out_proj = nn.Linear(codebook_dim, input_dim, bias=False)
|
||||
self.register_buffer("codebook", torch.zeros(codebook_size, codebook_dim))
|
||||
self.register_buffer("embed_avg", torch.zeros(codebook_size, codebook_dim))
|
||||
self.register_buffer("cluster_size", torch.zeros(codebook_size))
|
||||
self.register_buffer("inited", torch.tensor(False))
|
||||
|
||||
def encode(self, values: Tensor) -> tuple[Tensor, Tensor]:
|
||||
projected = self.in_proj(values.transpose(1, 2))
|
||||
flat = projected.reshape(-1, projected.shape[-1]).float()
|
||||
codebook = self.codebook.float()
|
||||
distances = (
|
||||
flat.square().sum(dim=1, keepdim=True)
|
||||
- 2 * flat @ codebook.t()
|
||||
+ codebook.square().sum(dim=1)[None]
|
||||
)
|
||||
codes = distances.argmin(dim=-1).reshape(values.shape[0], values.shape[2])
|
||||
quantized = functional.embedding(codes, self.codebook)
|
||||
quantized = self.out_proj(quantized).transpose(1, 2)
|
||||
return quantized.to(values.dtype), codes
|
||||
|
||||
def decode_codes(self, codes: Tensor) -> Tensor:
|
||||
return self.out_proj(functional.embedding(codes, self.codebook)).transpose(1, 2)
|
||||
|
||||
|
||||
class _ResidualCodecQuantizer(nn.Module):
|
||||
def __init__(self, config: Mapping[str, Any]) -> None:
|
||||
super().__init__()
|
||||
self.quantizers = nn.ModuleList([_CodecQuantizer(config) for _ in range(int(config["n_codebooks"]))])
|
||||
|
||||
def encode(self, values: Tensor) -> Tensor:
|
||||
residual = values
|
||||
codes = []
|
||||
for quantizer in self.quantizers:
|
||||
quantized, level_codes = quantizer.encode(residual)
|
||||
residual = residual - quantized
|
||||
codes.append(level_codes)
|
||||
return torch.stack(codes, dim=1)
|
||||
|
||||
def from_codes(self, codes: Tensor) -> Tensor:
|
||||
quantized = torch.zeros(
|
||||
codes.shape[0],
|
||||
self.quantizers[0].input_dim,
|
||||
codes.shape[-1],
|
||||
dtype=self.quantizers[0].codebook.dtype,
|
||||
device=codes.device,
|
||||
)
|
||||
for level, quantizer in enumerate(self.quantizers[: codes.shape[1]]):
|
||||
quantized = quantized + quantizer.decode_codes(codes[:, level])
|
||||
return quantized
|
||||
|
||||
|
||||
class _ActionCodecModel(nn.Module):
|
||||
def __init__(self, config: Mapping[str, Any]) -> None:
|
||||
super().__init__()
|
||||
self.config = dict(config)
|
||||
self.block_dct = (
|
||||
_BlockDCT(int(config["block_dct_block_size"]))
|
||||
if bool(config.get("use_block_dct", False))
|
||||
else None
|
||||
)
|
||||
self.conv_in = nn.Conv2d(
|
||||
int(config["horizon_patch_size"]),
|
||||
int(config["encoder_channels"]),
|
||||
kernel_size=(1, int(config["conv_in_action_kernel"])),
|
||||
)
|
||||
self.encoder = _CodecEncoder(config)
|
||||
self.rvq = _ResidualCodecQuantizer(config)
|
||||
self.decoder = _CodecDecoder(config)
|
||||
self.conv_out = nn.ConvTranspose2d(
|
||||
int(config["encoder_channels"]),
|
||||
int(config["horizon_patch_size"]),
|
||||
kernel_size=(1, int(config["conv_in_action_kernel"])),
|
||||
)
|
||||
|
||||
@property
|
||||
def code_h(self) -> int:
|
||||
height = int(self.config["horizon"]) // int(self.config["horizon_patch_size"])
|
||||
for stride_h, _ in self.config["strides"]:
|
||||
height //= int(stride_h)
|
||||
return height
|
||||
|
||||
@property
|
||||
def code_a(self) -> int:
|
||||
return int(self.config["max_component_dim"]) - int(self.config["conv_in_action_kernel"]) + 1
|
||||
|
||||
def _pad(self, values: Tensor) -> Tensor:
|
||||
maximum = int(self.config["max_component_dim"])
|
||||
if values.shape[-1] < maximum:
|
||||
return functional.pad(values, (0, maximum - values.shape[-1]))
|
||||
return values[..., :maximum]
|
||||
|
||||
def encode(self, components: dict[str, Tensor]) -> dict[str, Tensor]:
|
||||
names = list(components)
|
||||
batch_size = next(iter(components.values())).shape[0]
|
||||
values = torch.cat([self._pad(components[name].float()) for name in names], dim=0)
|
||||
if self.block_dct is not None:
|
||||
values = self.block_dct.dct(values)
|
||||
patch = int(self.config["horizon_patch_size"])
|
||||
values = values.reshape(values.shape[0], -1, patch, values.shape[-1]).transpose(1, 2)
|
||||
latent = self.encoder(self.conv_in(values)).flatten(2)
|
||||
codes = self.rvq.encode(latent)
|
||||
return {
|
||||
name: codes[index * batch_size : (index + 1) * batch_size] for index, name in enumerate(names)
|
||||
}
|
||||
|
||||
def decode(self, components: dict[str, Tensor], dimensions: Mapping[str, int]) -> dict[str, Tensor]:
|
||||
names = list(components)
|
||||
batch_size = next(iter(components.values())).shape[0]
|
||||
codes = torch.cat([components[name] for name in names], dim=0)
|
||||
quantized = self.rvq.from_codes(codes)
|
||||
quantized = quantized.reshape(
|
||||
quantized.shape[0],
|
||||
quantized.shape[1],
|
||||
self.code_h,
|
||||
self.code_a,
|
||||
)
|
||||
decoded = self.conv_out(self.decoder(quantized))
|
||||
decoded = decoded.transpose(1, 2).reshape(decoded.shape[0], -1, decoded.shape[-1])
|
||||
if self.block_dct is not None:
|
||||
decoded = self.block_dct.idct(decoded, int(self.config["horizon"]))
|
||||
return {
|
||||
name: decoded[index * batch_size : (index + 1) * batch_size, :, : dimensions[name]]
|
||||
for index, name in enumerate(names)
|
||||
}
|
||||
|
||||
|
||||
class _NativeCodecModule(nn.Module):
|
||||
"""Module hierarchy matching ``action_tokenizer.pt`` exactly."""
|
||||
|
||||
def __init__(self, config: Mapping[str, Any]) -> None:
|
||||
super().__init__()
|
||||
self.model = _ActionCodecModel(config)
|
||||
|
||||
|
||||
class _BinarySequenceCodec:
|
||||
def __init__(self, sequence_length: int, min_block_length: int, vocab_size: int) -> None:
|
||||
self.sequence_length = sequence_length
|
||||
self.min_block_length = min_block_length
|
||||
self.vocab_size = vocab_size
|
||||
self._count_cache: dict[tuple[int, int, int, bool], int] = {}
|
||||
self.num_sequences = self._count(sequence_length, -1, 0, True)
|
||||
self.num_tokens = max(1, math.ceil(math.log(self.num_sequences, vocab_size)))
|
||||
|
||||
def _count(self, remaining: int, last: int, run_length: int, first: bool) -> int:
|
||||
cache_key = (remaining, last, run_length, first)
|
||||
if cache_key in self._count_cache:
|
||||
return self._count_cache[cache_key]
|
||||
if remaining == 0:
|
||||
return 1
|
||||
total = 0
|
||||
for bit in (0, 1):
|
||||
if last == -1 or bit == last:
|
||||
total += self._count(
|
||||
remaining - 1,
|
||||
bit,
|
||||
min(run_length + 1, self.min_block_length + 1),
|
||||
first,
|
||||
)
|
||||
elif first or run_length > self.min_block_length:
|
||||
total += self._count(remaining - 1, bit, 1, False)
|
||||
self._count_cache[cache_key] = total
|
||||
return total
|
||||
|
||||
def _repair(self, bits: list[int]) -> list[int]:
|
||||
bits = bits.copy()
|
||||
while True:
|
||||
runs = []
|
||||
start = 0
|
||||
for index in range(1, len(bits)):
|
||||
if bits[index] != bits[index - 1]:
|
||||
runs.append((bits[start], start, index))
|
||||
start = index
|
||||
runs.append((bits[start], start, len(bits)))
|
||||
invalid = next(
|
||||
(
|
||||
(start, stop, runs[index - 1][0])
|
||||
for index, (_, start, stop) in enumerate(runs[1:-1], start=1)
|
||||
if stop - start <= self.min_block_length
|
||||
),
|
||||
None,
|
||||
)
|
||||
if invalid is None:
|
||||
return bits
|
||||
start, stop, value = invalid
|
||||
bits[start:stop] = [value] * (stop - start)
|
||||
|
||||
def _zero_completions(self, remaining: int, last: int, run: int, first: bool) -> int:
|
||||
if last in (-1, 0):
|
||||
return self._count(
|
||||
remaining,
|
||||
0,
|
||||
1 if last == -1 else min(run + 1, self.min_block_length + 1),
|
||||
first,
|
||||
)
|
||||
return self._count(remaining, 0, 1, False) if first or run > self.min_block_length else 0
|
||||
|
||||
def encode(self, values: Tensor, threshold: float) -> Tensor:
|
||||
output = []
|
||||
for row in values:
|
||||
bits = self._repair([int(value >= threshold) for value in row.tolist()])
|
||||
rank, last, run, first = 0, -1, 0, True
|
||||
for position, bit in enumerate(bits):
|
||||
remaining = len(bits) - position - 1
|
||||
if bit:
|
||||
rank += self._zero_completions(remaining, last, run, first)
|
||||
if last == -1:
|
||||
last, run = bit, 1
|
||||
elif bit == last:
|
||||
run = min(run + 1, self.min_block_length + 1)
|
||||
else:
|
||||
last, run, first = bit, 1, False
|
||||
tokens = []
|
||||
for _ in range(self.num_tokens):
|
||||
tokens.append(rank % self.vocab_size)
|
||||
rank //= self.vocab_size
|
||||
output.append(list(reversed(tokens)))
|
||||
return torch.tensor(output, dtype=torch.long, device=values.device)
|
||||
|
||||
def decode(self, tokens: Tensor) -> Tensor:
|
||||
rows = []
|
||||
for row in tokens.tolist():
|
||||
rank = 0
|
||||
for token in row:
|
||||
rank = rank * self.vocab_size + max(0, min(int(token), self.vocab_size - 1))
|
||||
rank = min(rank, self.num_sequences - 1)
|
||||
bits, last, run, first = [], -1, 0, True
|
||||
for position in range(self.sequence_length):
|
||||
remaining = self.sequence_length - position - 1
|
||||
zeros = self._zero_completions(remaining, last, run, first)
|
||||
if rank < zeros:
|
||||
bit = 0
|
||||
else:
|
||||
rank -= zeros
|
||||
bit = 1
|
||||
bits.append(bit)
|
||||
if last == -1:
|
||||
last, run = bit, 1
|
||||
elif bit == last:
|
||||
run = min(run + 1, self.min_block_length + 1)
|
||||
else:
|
||||
last, run, first = bit, 1, False
|
||||
rows.append(bits)
|
||||
return torch.tensor(rows, dtype=torch.float32, device=tokens.device)
|
||||
|
||||
|
||||
class G05NativeActionCodec:
|
||||
"""Non-registered sidecar wrapper for native ActionCodec encode/decode."""
|
||||
|
||||
def __init__(self, config: Mapping[str, Any], *, action_token_begin: int) -> None:
|
||||
self.config = dict(config)
|
||||
architecture = self.config["model_arch"]
|
||||
self.module = _NativeCodecModule(architecture)
|
||||
self.model = self.module.model
|
||||
self.action_token_begin = action_token_begin
|
||||
self.parts = {
|
||||
key: int(value) for key, value in self.config["parts_meta"].items() if value is not None
|
||||
}
|
||||
patterns = tuple(self.config.get("rule_based_key_patterns") or ())
|
||||
self.rule_parts = [key for key in self.parts if any(pattern in key for pattern in patterns)]
|
||||
self.neural_parts = [key for key in self.parts if key not in self.rule_parts]
|
||||
self.codebook_size = int(architecture["codebook_size"])
|
||||
self.max_residuals = int(architecture["n_codebooks"])
|
||||
self.num_residuals = int(self.config.get("num_residuals") or self.max_residuals)
|
||||
self.code_length = self.model.code_h * self.model.code_a
|
||||
marker_names = [
|
||||
f"<{part}_{level}>" for level in range(self.max_residuals) for part in self.neural_parts
|
||||
] + [f"<{part}>" for part in self.rule_parts]
|
||||
self.marker_indices = {name: self.codebook_size + index for index, name in enumerate(marker_names)}
|
||||
self.rule_codec = _BinarySequenceCodec(
|
||||
int(architecture["horizon"]),
|
||||
int(self.config.get("rule_based_min_block_len", 1)),
|
||||
self.codebook_size,
|
||||
)
|
||||
|
||||
@property
|
||||
def action_token_length(self) -> int:
|
||||
neural = len(self.neural_parts) * self.num_residuals * (self.code_length + 1)
|
||||
rules = len(self.rule_parts) * (self.rule_codec.num_tokens + 1)
|
||||
return neural + rules
|
||||
|
||||
@classmethod
|
||||
def load(
|
||||
cls,
|
||||
config: Mapping[str, Any],
|
||||
*,
|
||||
action_token_begin: int,
|
||||
) -> G05NativeActionCodec:
|
||||
codec = cls(config, action_token_begin=action_token_begin)
|
||||
checkpoint = torch.load(
|
||||
Path(str(config["ckpt_dir"])),
|
||||
map_location="cpu",
|
||||
mmap=True,
|
||||
weights_only=True,
|
||||
)
|
||||
state_dict = checkpoint.get("model_state_dict", checkpoint)
|
||||
codec.module.load_state_dict(state_dict, strict=True)
|
||||
codec.module.eval()
|
||||
return codec
|
||||
|
||||
def to(self, device: torch.device | str) -> G05NativeActionCodec:
|
||||
self.module.to(device=device, dtype=torch.float32)
|
||||
return self
|
||||
|
||||
def _split(self, actions: Tensor) -> dict[str, Tensor]:
|
||||
splits = torch.split(actions[..., : sum(self.parts.values())], list(self.parts.values()), dim=-1)
|
||||
return dict(zip(self.parts, splits, strict=True))
|
||||
|
||||
@torch.no_grad()
|
||||
def encode_for_language(self, payload: Mapping[str, Any]) -> list[int]:
|
||||
actions = torch.as_tensor(payload["value"])
|
||||
if actions.ndim == 2:
|
||||
actions = actions.unsqueeze(0)
|
||||
components = self._split(actions)
|
||||
neural = {key: components[key] for key in self.neural_parts}
|
||||
codes = self.model.encode(neural)
|
||||
rule_codes = {
|
||||
key: self.rule_codec.encode(
|
||||
components[key][..., 0],
|
||||
float(self.config.get("rule_based_binarize_threshold", 0)),
|
||||
)
|
||||
for key in self.rule_parts
|
||||
}
|
||||
indices = []
|
||||
for level in range(self.num_residuals):
|
||||
for key in self.neural_parts:
|
||||
indices.append(self.marker_indices[f"<{key}_{level}>"])
|
||||
indices.extend(codes[key][0, level].tolist())
|
||||
for key in self.rule_parts:
|
||||
indices.append(self.marker_indices[f"<{key}>"])
|
||||
indices.extend(rule_codes[key][0].tolist())
|
||||
return [self.action_token_begin + int(index) for index in indices]
|
||||
|
||||
@torch.no_grad()
|
||||
def decode_language_tokens(
|
||||
self,
|
||||
token_ids: Tensor,
|
||||
*,
|
||||
horizon: int,
|
||||
action_dim: int,
|
||||
) -> tuple[Tensor, set[str]]:
|
||||
indices = (token_ids.long() - self.action_token_begin).tolist()
|
||||
marker_to_name = {value: name for name, value in self.marker_indices.items()}
|
||||
neural: dict[str, list[list[int] | None]] = {
|
||||
key: [None] * self.num_residuals for key in self.neural_parts
|
||||
}
|
||||
rules: dict[str, list[int]] = {}
|
||||
cursor = 0
|
||||
while cursor < len(indices):
|
||||
marker = marker_to_name.get(indices[cursor])
|
||||
if marker is None:
|
||||
cursor += 1
|
||||
continue
|
||||
marker = marker[1:-1]
|
||||
if marker in self.rule_parts:
|
||||
length = self.rule_codec.num_tokens
|
||||
values = indices[cursor + 1 : cursor + 1 + length]
|
||||
if len(values) == length and all(0 <= value < self.codebook_size for value in values):
|
||||
rules[marker] = values
|
||||
cursor += length + 1
|
||||
continue
|
||||
part, level_text = marker.rsplit("_", 1)
|
||||
level = int(level_text)
|
||||
if part in neural and level < self.num_residuals:
|
||||
values = indices[cursor + 1 : cursor + 1 + self.code_length]
|
||||
if len(values) == self.code_length and all(
|
||||
0 <= value < self.codebook_size for value in values
|
||||
):
|
||||
neural[part][level] = values
|
||||
cursor += self.code_length + 1
|
||||
|
||||
absent = {
|
||||
key
|
||||
for key in self.parts
|
||||
if (key in neural and not any(level is not None for level in neural[key]))
|
||||
or (key in self.rule_parts and key not in rules)
|
||||
}
|
||||
device = next(self.module.parameters()).device
|
||||
code_tensors = {}
|
||||
for key, levels in neural.items():
|
||||
if not any(level is not None for level in levels):
|
||||
continue
|
||||
filled = [level if level is not None else [0] * self.code_length for level in levels]
|
||||
code_tensors[key] = torch.tensor([filled], dtype=torch.long, device=device)
|
||||
decoded = (
|
||||
self.model.decode(code_tensors, {key: self.parts[key] for key in code_tensors})
|
||||
if code_tensors
|
||||
else {}
|
||||
)
|
||||
for key in self.rule_parts:
|
||||
if key in rules:
|
||||
tokens = torch.tensor([rules[key]], dtype=torch.long, device=device)
|
||||
binary = self.rule_codec.decode(tokens)
|
||||
decoded[key] = binary[:, :, None] * 2 - 1
|
||||
# ``absent_key_fill_value`` is an internal partitioner sentinel. The
|
||||
# released marker-aware final decoder converts absent/no-op body parts
|
||||
# to zero motion before returning an action.
|
||||
batch = torch.zeros((1, horizon, action_dim), dtype=torch.float32, device=device)
|
||||
offset = 0
|
||||
for key, dimension in self.parts.items():
|
||||
if key in decoded:
|
||||
batch[..., offset : offset + dimension] = decoded[key][..., :dimension]
|
||||
offset += dimension
|
||||
return batch[0], absent
|
||||
@@ -188,9 +188,9 @@ _PROFILE_DEFAULTS = {
|
||||
class G05Config(PreTrainedConfig):
|
||||
"""LeRobot-side, checkpoint-auditable configuration for G0.5.
|
||||
|
||||
``author_model_config`` comes from the packaged checkpoint's resolved Hydra
|
||||
config. It is intentionally checkpoint state rather than a collection of
|
||||
guessed LeRobot defaults.
|
||||
``author_model_config`` is the backward-compatible serialized field holding
|
||||
the packaged checkpoint's resolved native architecture. It is checkpoint
|
||||
state rather than a collection of guessed LeRobot defaults.
|
||||
"""
|
||||
|
||||
checkpoint_profile: str = "g05-base"
|
||||
|
||||
@@ -10,7 +10,6 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import json
|
||||
import shutil
|
||||
from collections import deque
|
||||
@@ -33,27 +32,25 @@ from .configuration_g05 import (
|
||||
make_g05_cot_prompt_template,
|
||||
make_g05_prompt_template,
|
||||
)
|
||||
from .native_g05 import G05NativeBackend
|
||||
|
||||
|
||||
def _author_backend(config: G05Config) -> nn.Module:
|
||||
def _native_backend(config: G05Config) -> nn.Module:
|
||||
if not config.author_model_config:
|
||||
raise ValueError(
|
||||
"G0.5 author_model_config is empty. Load a packaged checkpoint, or "
|
||||
"inject a backend explicitly for testing."
|
||||
)
|
||||
try:
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
module = importlib.import_module("g05.models.g05.g05_policy_qwen35")
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"The OpenGalaxea G0.5 author package is required for real model execution. "
|
||||
"Clone the pinned GalaxeaVLA source, accept LICENSE-G0.5, and install its "
|
||||
"runtime dependencies in a compatible environment. LeRobot does not vendor "
|
||||
"or silently download that non-commercial code."
|
||||
) from exc
|
||||
backend_cls = module.G05PolicyQwen35
|
||||
return backend_cls(**OmegaConf.to_container(OmegaConf.create(config.author_model_config)))
|
||||
model_config = dict(config.author_model_config)
|
||||
model_config.update(
|
||||
{
|
||||
"predict_cot": config.predict_cot,
|
||||
"discrete_action": config.discrete_action,
|
||||
"continuous_action": config.continuous_action,
|
||||
"return_continuous_action": config.return_continuous_action,
|
||||
}
|
||||
)
|
||||
return G05NativeBackend.from_config(model_config)
|
||||
|
||||
|
||||
class G05Policy(PreTrainedPolicy):
|
||||
@@ -65,7 +62,7 @@ class G05Policy(PreTrainedPolicy):
|
||||
def __init__(self, config: G05Config, backend: nn.Module | None = None):
|
||||
super().__init__(config)
|
||||
config.validate_features()
|
||||
self.backend = backend if backend is not None else _author_backend(config)
|
||||
self.backend = backend if backend is not None else _native_backend(config)
|
||||
if not isinstance(self.backend, nn.Module):
|
||||
raise TypeError(f"G0.5 backend must be an nn.Module, got {type(self.backend)}.")
|
||||
self._action_queue: deque[Tensor] = deque()
|
||||
@@ -150,7 +147,7 @@ class G05Policy(PreTrainedPolicy):
|
||||
shutil.copy2(source, save_directory / name)
|
||||
|
||||
# Serialized paths are portable sidecar names. Local/Hub loading resolves them
|
||||
# against the downloaded checkpoint directory before constructing the author model.
|
||||
# against the downloaded checkpoint directory before constructing the native model.
|
||||
if (processor_path is not None and processor_path.exists()) or (
|
||||
tokenizer_path is not None and tokenizer_path.exists()
|
||||
):
|
||||
@@ -191,7 +188,7 @@ class G05Policy(PreTrainedPolicy):
|
||||
weight.data = weight.data.float()
|
||||
|
||||
def to(self, *args, **kwargs) -> G05Policy:
|
||||
"""Apply author inference precision and move the ActionCodec sidecar."""
|
||||
"""Apply the released inference precision and move the ActionCodec sidecar."""
|
||||
|
||||
result = super().to(*args, **kwargs)
|
||||
explicit_dtype = "dtype" in kwargs or any(isinstance(arg, torch.dtype | Tensor) for arg in args)
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,353 @@
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Native G0.5 prompt serialization and tokenizer registration."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
IGNORE_INDEX = -100
|
||||
|
||||
|
||||
class G05TokenType:
|
||||
"""Token categories stored in G0.5's attention-mask tensor."""
|
||||
|
||||
PADDING = 0
|
||||
IMAGE = 1
|
||||
PROPRIO = 2
|
||||
ACTION = 3
|
||||
TEXT = 4
|
||||
COT = 5
|
||||
PRED_TEXT = 6
|
||||
|
||||
|
||||
@dataclass
|
||||
class G05SequenceBatch:
|
||||
input_ids: Tensor
|
||||
labels: Tensor
|
||||
token_types: Tensor
|
||||
split_index: int | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Segment:
|
||||
kind: str
|
||||
content: str = ""
|
||||
sample_key: str = ""
|
||||
processor: str = ""
|
||||
masked: bool = False
|
||||
max_tokens: int | None = None
|
||||
|
||||
|
||||
class G05Tokenizer:
|
||||
"""Checkpoint-compatible G0.5 tokenizer and template serializer.
|
||||
|
||||
G0.5 extends Qwen3.5's tokenizer with ActionCodec codes, per-group
|
||||
residual markers, ``<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:
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
self.processor_path = Path(processor_path)
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(
|
||||
self.processor_path,
|
||||
trust_remote_code=False,
|
||||
local_files_only=True,
|
||||
)
|
||||
self.model_config = model_config
|
||||
at_config = model_config["AT_CONFIG"]
|
||||
architecture = at_config["model_arch"]
|
||||
|
||||
self.action_tokens = [f"<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],
|
||||
)
|
||||
@@ -0,0 +1,83 @@
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import torch
|
||||
|
||||
from lerobot.policies.g05.action_codec_g05 import G05NativeActionCodec, _BinarySequenceCodec
|
||||
|
||||
|
||||
def _tiny_codec_config() -> dict:
|
||||
return {
|
||||
"parts_meta": {
|
||||
"left_control": 3,
|
||||
"left_gripper": 1,
|
||||
"right_control": 3,
|
||||
"right_gripper": 1,
|
||||
},
|
||||
"rule_based_key_patterns": ["gripper"],
|
||||
"rule_based_min_block_len": 1,
|
||||
"rule_based_binarize_threshold": 0.0,
|
||||
"num_residuals": 2,
|
||||
"model_arch": {
|
||||
"horizon": 8,
|
||||
"horizon_patch_size": 2,
|
||||
"max_component_dim": 3,
|
||||
"conv_in_action_kernel": 2,
|
||||
"encoder_channels": 64,
|
||||
"c_mults": [1],
|
||||
"strides": [[1, 1]],
|
||||
"transformer_depths": [1],
|
||||
"latent_dim": 16,
|
||||
"num_heads": 1,
|
||||
"dim_heads": 64,
|
||||
"rope_base": 10_000,
|
||||
"ffn_mult": 2,
|
||||
"n_codebooks": 2,
|
||||
"codebook_size": 16,
|
||||
"codebook_dim": 4,
|
||||
"use_block_dct": False,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def test_binary_sequence_codec_roundtrip_repairs_short_middle_runs() -> None:
|
||||
codec = _BinarySequenceCodec(sequence_length=8, min_block_length=1, vocab_size=16)
|
||||
values = torch.tensor([[0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0]])
|
||||
|
||||
tokens = codec.encode(values, threshold=0.5)
|
||||
decoded = codec.decode(tokens)
|
||||
|
||||
assert tokens.shape == (1, codec.num_tokens)
|
||||
assert decoded.tolist() == [[0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0]]
|
||||
|
||||
|
||||
def test_native_action_codec_language_roundtrip_and_absent_groups() -> None:
|
||||
codec = G05NativeActionCodec(_tiny_codec_config(), action_token_begin=100)
|
||||
actions = torch.linspace(-1, 1, 8 * 8).reshape(8, 8)
|
||||
|
||||
token_ids = codec.encode_for_language({"value": actions})
|
||||
decoded, absent = codec.decode_language_tokens(
|
||||
torch.tensor(token_ids),
|
||||
horizon=8,
|
||||
action_dim=8,
|
||||
)
|
||||
|
||||
assert len(token_ids) == codec.action_token_length
|
||||
assert decoded.shape == (8, 8)
|
||||
assert torch.isfinite(decoded).all()
|
||||
assert absent == set()
|
||||
assert all(key.startswith("model.") for key in codec.module.state_dict())
|
||||
|
||||
empty, absent = codec.decode_language_tokens(torch.empty(0, dtype=torch.long), horizon=8, action_dim=8)
|
||||
assert absent == set(codec.parts)
|
||||
assert torch.equal(empty, torch.zeros_like(empty))
|
||||
@@ -2879,7 +2879,6 @@ all = [
|
||||
{ name = "motorbridge-smart-servo" },
|
||||
{ name = "mypy" },
|
||||
{ name = "num2words" },
|
||||
{ name = "omegaconf" },
|
||||
{ name = "pandas" },
|
||||
{ name = "peft" },
|
||||
{ name = "placo" },
|
||||
@@ -3022,7 +3021,7 @@ feetech = [
|
||||
{ name = "pyserial" },
|
||||
]
|
||||
g05 = [
|
||||
{ name = "omegaconf" },
|
||||
{ name = "transformers" },
|
||||
]
|
||||
gamepad = [
|
||||
{ name = "hidapi" },
|
||||
@@ -3456,7 +3455,6 @@ requires-dist = [
|
||||
{ name = "mypy", marker = "extra == 'dev'", specifier = ">=1.19.1" },
|
||||
{ name = "num2words", marker = "extra == 'smolvla'", specifier = ">=0.5.14,<0.6.0" },
|
||||
{ name = "numpy", specifier = ">=2.0.0,<2.3.0" },
|
||||
{ name = "omegaconf", marker = "extra == 'g05'", specifier = ">=2.3.0,<3.0.0" },
|
||||
{ name = "onnx", marker = "extra == 'unitree-g1'", specifier = ">=1.16.0,<2.0.0" },
|
||||
{ name = "onnxruntime", marker = "extra == 'unitree-g1'", specifier = ">=1.16.0,<2.0.0" },
|
||||
{ name = "openai", marker = "extra == 'annotations'", specifier = ">=1.40,<2.0" },
|
||||
@@ -3505,6 +3503,7 @@ requires-dist = [
|
||||
{ name = "torchvision", marker = "sys_platform != 'linux'", specifier = ">=0.22.0,<0.27.0" },
|
||||
{ name = "torchvision", marker = "sys_platform == 'linux'", specifier = ">=0.22.0,<0.27.0", index = "https://download.pytorch.org/whl/cu128" },
|
||||
{ name = "tqdm", specifier = ">=4.66.0,<5.0.0" },
|
||||
{ name = "transformers", marker = "extra == 'g05'", specifier = ">=5.4.0,<5.6.0" },
|
||||
{ name = "transformers", marker = "extra == 'transformers-dep'", specifier = ">=5.4.0,<5.6.0" },
|
||||
{ name = "wandb", marker = "extra == 'training'", specifier = ">=0.24.0,<0.28.0" },
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user