feat(g05): make policy runtime native

This commit is contained in:
Pepijn
2026-07-29 16:58:02 +02:00
parent 916fc866f5
commit e6d7eea88a
9 changed files with 2372 additions and 68 deletions
+27 -40
View File
@@ -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
View File
@@ -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"
+15 -18
View File
@@ -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
+353
View File
@@ -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))
Generated
+2 -3
View File
@@ -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" },
]