refactor(g05): consolidate policy modules

src/lerobot/policies/g05/action_codec_g05.py

@@ -1,674 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Native inference implementation of G0.5's ActionCodec sidecar."""
-
-from __future__ import annotations
-
-import math
-from collections.abc import Mapping
-from pathlib import Path
-from typing import Any
-
-import torch
-import torch.nn.functional as functional
-from torch import Tensor, nn
-
-
-class _BlockDCT(nn.Module):
-    def __init__(self, block_size: int) -> None:
-        super().__init__()
-        self.block_size = block_size
-        frequency = torch.arange(block_size, dtype=torch.float32)
-        time = torch.arange(block_size, dtype=torch.float32)
-        basis = torch.cos(math.pi / block_size * (time + 0.5).unsqueeze(0) * frequency.unsqueeze(1))
-        basis[0] *= math.sqrt(1 / block_size)
-        basis[1:] *= math.sqrt(2 / block_size)
-        self.register_buffer("basis", basis, persistent=False)
-
-    def dct(self, values: Tensor) -> Tensor:
-        batch, horizon, dimension = values.shape
-        pad = (-horizon) % self.block_size
-        if pad:
-            values = functional.pad(values, (0, 0, 0, pad))
-        blocks = values.shape[1] // self.block_size
-        values = values.reshape(batch * blocks, self.block_size, dimension)
-        transformed = torch.einsum("kn,bnd->bkd", self.basis.to(values), values)
-        return transformed.reshape(batch, blocks * self.block_size, dimension)
-
-    def idct(self, values: Tensor, horizon: int) -> Tensor:
-        batch, padded_horizon, dimension = values.shape
-        blocks = padded_horizon // self.block_size
-        values = values.reshape(batch * blocks, self.block_size, dimension)
-        restored = torch.einsum("nk,bkd->bnd", self.basis.to(values), values)
-        return restored.reshape(batch, padded_horizon, dimension)[:, :horizon]
-
-
-def _rotate_half(values: Tensor) -> Tensor:
-    first, second = values.chunk(2, dim=-1)
-    return torch.cat((-second, first), dim=-1)
-
-
-class _CodecAttention(nn.Module):
-    def __init__(self, dimension: int, num_heads: int, head_dim: int, rope_base: int) -> None:
-        super().__init__()
-        self.num_heads = num_heads
-        self.head_dim = head_dim
-        inner_dim = num_heads * head_dim
-        self.to_qkv = nn.Linear(dimension, inner_dim * 3, bias=False)
-        self.to_out = nn.Linear(inner_dim, dimension, bias=False)
-        self.q_norm = nn.LayerNorm(head_dim, eps=1e-6)
-        self.k_norm = nn.LayerNorm(head_dim, eps=1e-6)
-        rope_dim = max(head_dim // 2, 32)
-        inverse = 1 / (rope_base ** (torch.arange(0, rope_dim, 2, dtype=torch.float32) / rope_dim))
-        self.register_buffer("_inverse_frequency", inverse, persistent=False)
-
-    def forward(self, hidden_states: Tensor) -> Tensor:
-        batch, sequence_length, _ = hidden_states.shape
-        query, key, value = self.to_qkv(hidden_states).chunk(3, dim=-1)
-
-        def heads(values: Tensor) -> Tensor:
-            return values.view(batch, sequence_length, self.num_heads, self.head_dim).transpose(1, 2)
-
-        query, key, value = (heads(values) for values in (query, key, value))
-        query, key = self.q_norm(query), self.k_norm(key)
-        time = torch.arange(sequence_length, device=hidden_states.device, dtype=torch.float32)
-        phase = torch.outer(time, self._inverse_frequency.to(hidden_states.device))
-        phase = torch.cat((phase, phase), dim=-1).to(hidden_states.dtype)[None, None]
-        cosine, sine = phase.cos(), phase.sin()
-        rotary_dim = cosine.shape[-1]
-        query_rotary, query_pass = query[..., :rotary_dim], query[..., rotary_dim:]
-        key_rotary, key_pass = key[..., :rotary_dim], key[..., rotary_dim:]
-        query = torch.cat((query_rotary * cosine + _rotate_half(query_rotary) * sine, query_pass), dim=-1)
-        key = torch.cat((key_rotary * cosine + _rotate_half(key_rotary) * sine, key_pass), dim=-1)
-        attended = functional.scaled_dot_product_attention(query, key, value)
-        attended = attended.transpose(1, 2).reshape(batch, sequence_length, -1)
-        return self.to_out(attended)
-
-
-class _CodecFFN(nn.Module):
-    def __init__(self, dimension: int, multiplier: float) -> None:
-        super().__init__()
-        inner_dim = int(dimension * multiplier)
-        self.w_up = nn.Linear(dimension, inner_dim * 2, bias=False)
-        self.w_down = nn.Linear(inner_dim, dimension, bias=False)
-
-    def forward(self, hidden_states: Tensor) -> Tensor:
-        value, gate = self.w_up(hidden_states).chunk(2, dim=-1)
-        return self.w_down(value * functional.gelu(gate))
-
-
-class _CodecTransformerLayer(nn.Module):
-    def __init__(self, dimension: int, config: Mapping[str, Any]) -> None:
-        super().__init__()
-        layer_scale_init = float(config.get("layer_scale_init", 1.0))
-        self.ls1 = nn.Parameter(torch.full((dimension,), layer_scale_init))
-        self.ls2 = nn.Parameter(torch.full((dimension,), layer_scale_init))
-        self.norm1 = nn.LayerNorm(dimension, eps=1e-6)
-        self.attn = _CodecAttention(
-            dimension,
-            int(config["num_heads"]),
-            int(config["dim_heads"]),
-            int(config["rope_base"]),
-        )
-        self.norm2 = nn.LayerNorm(dimension, eps=1e-6)
-        self.ffn = _CodecFFN(dimension, float(config["ffn_mult"]))
-
-    def forward(self, hidden_states: Tensor) -> Tensor:
-        hidden_states = hidden_states + self.attn(self.norm1(hidden_states)) * self.ls1
-        return hidden_states + self.ffn(self.norm2(hidden_states)) * self.ls2
-
-
-class _CodecDownBlock(nn.Module):
-    def __init__(
-        self,
-        input_channels: int,
-        output_channels: int,
-        stride: tuple[int, int],
-        depth: int,
-        config: Mapping[str, Any],
-    ) -> None:
-        super().__init__()
-        stride_h, stride_a = stride
-        if stride_h > 1 or input_channels != output_channels:
-            kernel_h = 2 * stride_h if stride_h > 1 else 1
-            self.conv = nn.Conv2d(
-                input_channels,
-                output_channels,
-                kernel_size=(kernel_h, 1),
-                stride=(stride_h, stride_a),
-                padding=(kernel_h // 2 - int(stride_h > 1), 0),
-            )
-        else:
-            self.conv = nn.Identity()
-        self.transformer_layers = nn.ModuleList(
-            [_CodecTransformerLayer(output_channels, config) for _ in range(depth)]
-        )
-
-    def forward(self, hidden_states: Tensor) -> Tensor:
-        hidden_states = self.conv(hidden_states)
-        batch, channels, height, action_dim = hidden_states.shape
-        sequence = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * action_dim, channels)
-        for layer in self.transformer_layers:
-            sequence = layer(sequence)
-        return sequence.reshape(batch, height, action_dim, channels).permute(0, 3, 1, 2)
-
-
-class _CodecUpBlock(nn.Module):
-    def __init__(
-        self,
-        input_channels: int,
-        output_channels: int,
-        stride: tuple[int, int],
-        depth: int,
-        config: Mapping[str, Any],
-    ) -> None:
-        super().__init__()
-        self.transformer_layers = nn.ModuleList(
-            [_CodecTransformerLayer(input_channels, config) for _ in range(depth)]
-        )
-        stride_h, stride_a = stride
-        if stride_h > 1 or input_channels != output_channels:
-            kernel_h = 2 * stride_h if stride_h > 1 else 1
-            self.conv = nn.ConvTranspose2d(
-                input_channels,
-                output_channels,
-                kernel_size=(kernel_h, 1),
-                stride=(stride_h, stride_a),
-                padding=(kernel_h // 2 - int(stride_h > 1), 0),
-            )
-        else:
-            self.conv = nn.Identity()
-
-    def forward(self, hidden_states: Tensor) -> Tensor:
-        batch, channels, height, action_dim = hidden_states.shape
-        sequence = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * action_dim, channels)
-        for layer in self.transformer_layers:
-            sequence = layer(sequence)
-        hidden_states = sequence.reshape(batch, height, action_dim, channels).permute(0, 3, 1, 2)
-        return self.conv(hidden_states)
-
-
-class _CodecEncoder(nn.Module):
-    def __init__(self, config: Mapping[str, Any]) -> None:
-        super().__init__()
-        base = int(config["encoder_channels"])
-        channel_dims = [base * int(multiplier) for multiplier in config["c_mults"]]
-        dims = [base] + channel_dims
-        self.blocks = nn.ModuleList(
-            [
-                _CodecDownBlock(
-                    dims[index],
-                    dims[index + 1],
-                    tuple(stride),
-                    int(config["transformer_depths"][index]),
-                    config,
-                )
-                for index, stride in enumerate(config["strides"])
-            ]
-        )
-        self.out_proj = nn.Conv2d(dims[-1], int(config["latent_dim"]), kernel_size=1)
-
-    def forward(self, hidden_states: Tensor) -> Tensor:
-        for block in self.blocks:
-            hidden_states = block(hidden_states)
-        return self.out_proj(hidden_states)
-
-
-class _CodecDecoder(nn.Module):
-    def __init__(self, config: Mapping[str, Any]) -> None:
-        super().__init__()
-        base = int(config["encoder_channels"])
-        channel_dims = [base * int(multiplier) for multiplier in config["c_mults"]]
-        dims = [base] + channel_dims
-        self.in_proj = nn.Conv2d(int(config["latent_dim"]), dims[-1], kernel_size=1)
-        self.blocks = nn.ModuleList(
-            [
-                _CodecUpBlock(
-                    input_channels,
-                    output_channels,
-                    tuple(stride),
-                    int(depth),
-                    config,
-                )
-                for stride, depth, input_channels, output_channels in zip(
-                    reversed(config["strides"]),
-                    reversed(config["transformer_depths"]),
-                    reversed(dims[1:]),
-                    reversed(dims[:-1]),
-                    strict=True,
-                )
-            ]
-        )
-
-    def forward(self, hidden_states: Tensor) -> Tensor:
-        hidden_states = self.in_proj(hidden_states)
-        for block in self.blocks:
-            hidden_states = block(hidden_states)
-        return hidden_states
-
-
-class _CodecQuantizer(nn.Module):
-    def __init__(self, config: Mapping[str, Any]) -> None:
-        super().__init__()
-        input_dim = int(config["latent_dim"])
-        codebook_dim = int(config["codebook_dim"])
-        codebook_size = int(config["codebook_size"])
-        self.input_dim = input_dim
-        self.in_proj = nn.Linear(input_dim, codebook_dim, bias=False)
-        self.out_proj = nn.Linear(codebook_dim, input_dim, bias=False)
-        self.register_buffer("codebook", torch.zeros(codebook_size, codebook_dim))
-        self.register_buffer("embed_avg", torch.zeros(codebook_size, codebook_dim))
-        self.register_buffer("cluster_size", torch.zeros(codebook_size))
-        self.register_buffer("inited", torch.tensor(False))
-
-    def encode(self, values: Tensor) -> tuple[Tensor, Tensor]:
-        projected = self.in_proj(values.transpose(1, 2))
-        flat = projected.reshape(-1, projected.shape[-1]).float()
-        codebook = self.codebook.float()
-        distances = (
-            flat.square().sum(dim=1, keepdim=True)
-            - 2 * flat @ codebook.t()
-            + codebook.square().sum(dim=1)[None]
-        )
-        codes = distances.argmin(dim=-1).reshape(values.shape[0], values.shape[2])
-        quantized = functional.embedding(codes, self.codebook)
-        quantized = self.out_proj(quantized).transpose(1, 2)
-        return quantized.to(values.dtype), codes
-
-    def decode_codes(self, codes: Tensor) -> Tensor:
-        return self.out_proj(functional.embedding(codes, self.codebook)).transpose(1, 2)
-
-
-class _ResidualCodecQuantizer(nn.Module):
-    def __init__(self, config: Mapping[str, Any]) -> None:
-        super().__init__()
-        self.quantizers = nn.ModuleList([_CodecQuantizer(config) for _ in range(int(config["n_codebooks"]))])
-
-    def encode(self, values: Tensor) -> Tensor:
-        residual = values
-        codes = []
-        for quantizer in self.quantizers:
-            quantized, level_codes = quantizer.encode(residual)
-            residual = residual - quantized
-            codes.append(level_codes)
-        return torch.stack(codes, dim=1)
-
-    def from_codes(self, codes: Tensor) -> Tensor:
-        quantized = torch.zeros(
-            codes.shape[0],
-            self.quantizers[0].input_dim,
-            codes.shape[-1],
-            dtype=self.quantizers[0].codebook.dtype,
-            device=codes.device,
-        )
-        for level, quantizer in enumerate(self.quantizers[: codes.shape[1]]):
-            quantized = quantized + quantizer.decode_codes(codes[:, level])
-        return quantized
-
-
-class _ActionCodecModel(nn.Module):
-    def __init__(self, config: Mapping[str, Any]) -> None:
-        super().__init__()
-        self.config = dict(config)
-        self.block_dct = (
-            _BlockDCT(int(config["block_dct_block_size"]))
-            if bool(config.get("use_block_dct", False))
-            else None
-        )
-        self.conv_in = nn.Conv2d(
-            int(config["horizon_patch_size"]),
-            int(config["encoder_channels"]),
-            kernel_size=(1, int(config["conv_in_action_kernel"])),
-        )
-        self.encoder = _CodecEncoder(config)
-        self.rvq = _ResidualCodecQuantizer(config)
-        self.decoder = _CodecDecoder(config)
-        self.conv_out = nn.ConvTranspose2d(
-            int(config["encoder_channels"]),
-            int(config["horizon_patch_size"]),
-            kernel_size=(1, int(config["conv_in_action_kernel"])),
-        )
-
-    @property
-    def code_h(self) -> int:
-        height = int(self.config["horizon"]) // int(self.config["horizon_patch_size"])
-        for stride_h, _ in self.config["strides"]:
-            height //= int(stride_h)
-        return height
-
-    @property
-    def code_a(self) -> int:
-        return int(self.config["max_component_dim"]) - int(self.config["conv_in_action_kernel"]) + 1
-
-    def _pad(self, values: Tensor) -> Tensor:
-        maximum = int(self.config["max_component_dim"])
-        if values.shape[-1] < maximum:
-            return functional.pad(values, (0, maximum - values.shape[-1]))
-        return values[..., :maximum]
-
-    def encode(self, components: dict[str, Tensor]) -> dict[str, Tensor]:
-        names = list(components)
-        batch_size = next(iter(components.values())).shape[0]
-        values = torch.cat([self._pad(components[name].float()) for name in names], dim=0)
-        if self.block_dct is not None:
-            values = self.block_dct.dct(values)
-        patch = int(self.config["horizon_patch_size"])
-        values = values.reshape(values.shape[0], -1, patch, values.shape[-1]).transpose(1, 2)
-        latent = self.encoder(self.conv_in(values)).flatten(2)
-        codes = self.rvq.encode(latent)
-        return {
-            name: codes[index * batch_size : (index + 1) * batch_size] for index, name in enumerate(names)
-        }
-
-    def decode(self, components: dict[str, Tensor], dimensions: Mapping[str, int]) -> dict[str, Tensor]:
-        names = list(components)
-        batch_size = next(iter(components.values())).shape[0]
-        codes = torch.cat([components[name] for name in names], dim=0)
-        quantized = self.rvq.from_codes(codes)
-        quantized = quantized.reshape(
-            quantized.shape[0],
-            quantized.shape[1],
-            self.code_h,
-            self.code_a,
-        )
-        decoded = self.conv_out(self.decoder(quantized))
-        decoded = decoded.transpose(1, 2).reshape(decoded.shape[0], -1, decoded.shape[-1])
-        if self.block_dct is not None:
-            decoded = self.block_dct.idct(decoded, int(self.config["horizon"]))
-        return {
-            name: decoded[index * batch_size : (index + 1) * batch_size, :, : dimensions[name]]
-            for index, name in enumerate(names)
-        }
-
-
-class _NativeCodecModule(nn.Module):
-    """Module hierarchy matching ``action_tokenizer.pt`` exactly."""
-
-    def __init__(self, config: Mapping[str, Any]) -> None:
-        super().__init__()
-        self.model = _ActionCodecModel(config)
-
-
-class _BinarySequenceCodec:
-    def __init__(self, sequence_length: int, min_block_length: int, vocab_size: int) -> None:
-        self.sequence_length = sequence_length
-        self.min_block_length = min_block_length
-        self.vocab_size = vocab_size
-        self._count_cache: dict[tuple[int, int, int, bool], int] = {}
-        self.num_sequences = self._count(sequence_length, -1, 0, True)
-        self.num_tokens = max(1, math.ceil(math.log(self.num_sequences, vocab_size)))
-
-    def _count(self, remaining: int, last: int, run_length: int, first: bool) -> int:
-        cache_key = (remaining, last, run_length, first)
-        if cache_key in self._count_cache:
-            return self._count_cache[cache_key]
-        if remaining == 0:
-            return 1
-        total = 0
-        for bit in (0, 1):
-            if last == -1 or bit == last:
-                total += self._count(
-                    remaining - 1,
-                    bit,
-                    min(run_length + 1, self.min_block_length + 1),
-                    first,
-                )
-            elif first or run_length > self.min_block_length:
-                total += self._count(remaining - 1, bit, 1, False)
-        self._count_cache[cache_key] = total
-        return total
-
-    def _repair(self, bits: list[int]) -> list[int]:
-        bits = bits.copy()
-        while True:
-            runs = []
-            start = 0
-            for index in range(1, len(bits)):
-                if bits[index] != bits[index - 1]:
-                    runs.append((bits[start], start, index))
-                    start = index
-            runs.append((bits[start], start, len(bits)))
-            invalid = next(
-                (
-                    (start, stop, runs[index - 1][0])
-                    for index, (_, start, stop) in enumerate(runs[1:-1], start=1)
-                    if stop - start <= self.min_block_length
-                ),
-                None,
-            )
-            if invalid is None:
-                return bits
-            start, stop, value = invalid
-            bits[start:stop] = [value] * (stop - start)
-
-    def _zero_completions(self, remaining: int, last: int, run: int, first: bool) -> int:
-        if last in (-1, 0):
-            return self._count(
-                remaining,
-                0,
-                1 if last == -1 else min(run + 1, self.min_block_length + 1),
-                first,
-            )
-        return self._count(remaining, 0, 1, False) if first or run > self.min_block_length else 0
-
-    def encode(self, values: Tensor, threshold: float) -> Tensor:
-        output = []
-        for row in values:
-            bits = self._repair([int(value >= threshold) for value in row.tolist()])
-            rank, last, run, first = 0, -1, 0, True
-            for position, bit in enumerate(bits):
-                remaining = len(bits) - position - 1
-                if bit:
-                    rank += self._zero_completions(remaining, last, run, first)
-                if last == -1:
-                    last, run = bit, 1
-                elif bit == last:
-                    run = min(run + 1, self.min_block_length + 1)
-                else:
-                    last, run, first = bit, 1, False
-            tokens = []
-            for _ in range(self.num_tokens):
-                tokens.append(rank % self.vocab_size)
-                rank //= self.vocab_size
-            output.append(list(reversed(tokens)))
-        return torch.tensor(output, dtype=torch.long, device=values.device)
-
-    def decode(self, tokens: Tensor) -> Tensor:
-        rows = []
-        for row in tokens.tolist():
-            rank = 0
-            for token in row:
-                rank = rank * self.vocab_size + max(0, min(int(token), self.vocab_size - 1))
-            rank = min(rank, self.num_sequences - 1)
-            bits, last, run, first = [], -1, 0, True
-            for position in range(self.sequence_length):
-                remaining = self.sequence_length - position - 1
-                zeros = self._zero_completions(remaining, last, run, first)
-                if rank < zeros:
-                    bit = 0
-                else:
-                    rank -= zeros
-                    bit = 1
-                bits.append(bit)
-                if last == -1:
-                    last, run = bit, 1
-                elif bit == last:
-                    run = min(run + 1, self.min_block_length + 1)
-                else:
-                    last, run, first = bit, 1, False
-            rows.append(bits)
-        return torch.tensor(rows, dtype=torch.float32, device=tokens.device)
-
-
-class G05NativeActionCodec:
-    """Non-registered sidecar wrapper for native ActionCodec encode/decode."""
-
-    def __init__(self, config: Mapping[str, Any], *, action_token_begin: int) -> None:
-        self.config = dict(config)
-        architecture = self.config["model_arch"]
-        self.module = _NativeCodecModule(architecture)
-        self.model = self.module.model
-        self.action_token_begin = action_token_begin
-        self.parts = {
-            key: int(value) for key, value in self.config["parts_meta"].items() if value is not None
-        }
-        patterns = tuple(self.config.get("rule_based_key_patterns") or ())
-        self.rule_parts = [key for key in self.parts if any(pattern in key for pattern in patterns)]
-        self.neural_parts = [key for key in self.parts if key not in self.rule_parts]
-        self.codebook_size = int(architecture["codebook_size"])
-        self.max_residuals = int(architecture["n_codebooks"])
-        self.num_residuals = int(self.config.get("num_residuals") or self.max_residuals)
-        self.code_length = self.model.code_h * self.model.code_a
-        marker_names = [
-            f"<{part}_{level}>" for level in range(self.max_residuals) for part in self.neural_parts
-        ] + [f"<{part}>" for part in self.rule_parts]
-        self.marker_indices = {name: self.codebook_size + index for index, name in enumerate(marker_names)}
-        self.rule_codec = _BinarySequenceCodec(
-            int(architecture["horizon"]),
-            int(self.config.get("rule_based_min_block_len", 1)),
-            self.codebook_size,
-        )
-
-    @property
-    def action_token_length(self) -> int:
-        neural = len(self.neural_parts) * self.num_residuals * (self.code_length + 1)
-        rules = len(self.rule_parts) * (self.rule_codec.num_tokens + 1)
-        return neural + rules
-
-    @classmethod
-    def load(
-        cls,
-        config: Mapping[str, Any],
-        *,
-        action_token_begin: int,
-    ) -> G05NativeActionCodec:
-        codec = cls(config, action_token_begin=action_token_begin)
-        checkpoint = torch.load(
-            Path(str(config["ckpt_dir"])),
-            map_location="cpu",
-            mmap=True,
-            weights_only=True,
-        )
-        state_dict = checkpoint.get("model_state_dict", checkpoint)
-        codec.module.load_state_dict(state_dict, strict=True)
-        codec.module.eval()
-        return codec
-
-    def to(self, device: torch.device | str) -> G05NativeActionCodec:
-        self.module.to(device=device, dtype=torch.float32)
-        return self
-
-    def _split(self, actions: Tensor) -> dict[str, Tensor]:
-        splits = torch.split(actions[..., : sum(self.parts.values())], list(self.parts.values()), dim=-1)
-        return dict(zip(self.parts, splits, strict=True))
-
-    @torch.no_grad()
-    def encode_for_language(self, payload: Mapping[str, Any]) -> list[int]:
-        actions = torch.as_tensor(payload["value"])
-        if actions.ndim == 2:
-            actions = actions.unsqueeze(0)
-        components = self._split(actions)
-        neural = {key: components[key] for key in self.neural_parts}
-        codes = self.model.encode(neural)
-        rule_codes = {
-            key: self.rule_codec.encode(
-                components[key][..., 0],
-                float(self.config.get("rule_based_binarize_threshold", 0)),
-            )
-            for key in self.rule_parts
-        }
-        indices = []
-        for level in range(self.num_residuals):
-            for key in self.neural_parts:
-                indices.append(self.marker_indices[f"<{key}_{level}>"])
-                indices.extend(codes[key][0, level].tolist())
-        for key in self.rule_parts:
-            indices.append(self.marker_indices[f"<{key}>"])
-            indices.extend(rule_codes[key][0].tolist())
-        return [self.action_token_begin + int(index) for index in indices]
-
-    @torch.no_grad()
-    def decode_language_tokens(
-        self,
-        token_ids: Tensor,
-        *,
-        horizon: int,
-        action_dim: int,
-    ) -> tuple[Tensor, set[str]]:
-        indices = (token_ids.long() - self.action_token_begin).tolist()
-        marker_to_name = {value: name for name, value in self.marker_indices.items()}
-        neural: dict[str, list[list[int] | None]] = {
-            key: [None] * self.num_residuals for key in self.neural_parts
-        }
-        rules: dict[str, list[int]] = {}
-        cursor = 0
-        while cursor < len(indices):
-            marker = marker_to_name.get(indices[cursor])
-            if marker is None:
-                cursor += 1
-                continue
-            marker = marker[1:-1]
-            if marker in self.rule_parts:
-                length = self.rule_codec.num_tokens
-                values = indices[cursor + 1 : cursor + 1 + length]
-                if len(values) == length and all(0 <= value < self.codebook_size for value in values):
-                    rules[marker] = values
-                cursor += length + 1
-                continue
-            part, level_text = marker.rsplit("_", 1)
-            level = int(level_text)
-            if part in neural and level < self.num_residuals:
-                values = indices[cursor + 1 : cursor + 1 + self.code_length]
-                if len(values) == self.code_length and all(
-                    0 <= value < self.codebook_size for value in values
-                ):
-                    neural[part][level] = values
-            cursor += self.code_length + 1
-
-        absent = {
-            key
-            for key in self.parts
-            if (key in neural and not any(level is not None for level in neural[key]))
-            or (key in self.rule_parts and key not in rules)
-        }
-        device = next(self.module.parameters()).device
-        code_tensors = {}
-        for key, levels in neural.items():
-            if not any(level is not None for level in levels):
-                continue
-            filled = [level if level is not None else [0] * self.code_length for level in levels]
-            code_tensors[key] = torch.tensor([filled], dtype=torch.long, device=device)
-        decoded = (
-            self.model.decode(code_tensors, {key: self.parts[key] for key in code_tensors})
-            if code_tensors
-            else {}
-        )
-        for key in self.rule_parts:
-            if key in rules:
-                tokens = torch.tensor([rules[key]], dtype=torch.long, device=device)
-                binary = self.rule_codec.decode(tokens)
-                decoded[key] = binary[:, :, None] * 2 - 1
-        # ``absent_key_fill_value`` is an internal partitioner sentinel. The
-        # released marker-aware final decoder converts absent/no-op body parts
-        # to zero motion before returning an action.
-        batch = torch.zeros((1, horizon, action_dim), dtype=torch.float32, device=device)
-        offset = 0
-        for key, dimension in self.parts.items():
-            if key in decoded:
-                batch[..., offset : offset + dimension] = decoded[key][..., :dimension]
-            offset += dimension
-        return batch[0], absent

src/lerobot/policies/g05/processing_g05.py

@@ -1,354 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Native G0.5 prompt serialization and tokenizer registration."""
-
-from __future__ import annotations
-
-import re
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Any
-
-import torch
-from torch import Tensor
-from transformers import AutoTokenizer
-
-IGNORE_INDEX = -100
-
-
-class G05TokenType:
-    """Token categories stored in G0.5's attention-mask tensor."""
-
-    PADDING = 0
-    IMAGE = 1
-    PROPRIO = 2
-    ACTION = 3
-    TEXT = 4
-    COT = 5
-    PRED_TEXT = 6
-
-
-@dataclass
-class G05SequenceBatch:
-    input_ids: Tensor
-    labels: Tensor
-    token_types: Tensor
-    split_index: int | None = None
-
-
-@dataclass
-class _Segment:
-    kind: str
-    content: str = ""
-    sample_key: str = ""
-    processor: str = ""
-    masked: bool = False
-    max_tokens: int | None = None
-
-
-class G05Tokenizer:
-    """Checkpoint-compatible G0.5 tokenizer and template serializer.
-
-    G0.5 extends Qwen3.5's tokenizer with ActionCodec codes, per-group
-    residual markers, ``<EOV>``, and ``<state>``. Registration order is model
-    state: changing it changes the rows used by the tied language head.
-    """
-
-    _PLACEHOLDER = re.compile(r"<([^<>|]+)>")
-
-    def __init__(self, processor_path: str | Path, model_config: dict[str, Any]) -> None:
-        self.processor_path = Path(processor_path)
-        self.tokenizer = AutoTokenizer.from_pretrained(
-            self.processor_path,
-            trust_remote_code=False,
-            local_files_only=True,
-        )
-        self.model_config = model_config
-        at_config = model_config["AT_CONFIG"]
-        architecture = at_config["model_arch"]
-
-        self.action_tokens = [f"<action{index:04d}>" for index in range(int(architecture["codebook_size"]))]
-        parts = list(at_config["parts_meta"])
-        rule_patterns = tuple(at_config.get("rule_based_key_patterns") or ())
-        self.rule_parts = [name for name in parts if any(pattern in name for pattern in rule_patterns)]
-        self.neural_parts = [name for name in parts if name not in self.rule_parts]
-        self.group_tokens = [
-            f"<{part}_{residual}>"
-            for residual in range(int(architecture["n_codebooks"]))
-            for part in self.neural_parts
-        ] + [f"<{part}>" for part in self.rule_parts]
-        self.tokenizer.add_tokens(self.action_tokens + self.group_tokens + ["<EOV>", "<state>"])
-
-        self.pad_token_id = int(model_config["pad_token_id"])
-        self.eos_token_id = int(model_config["eos_token_id"])
-        self.image_token_id = int(model_config["image_token_index"])
-        self.vision_start_token_id = int(self.tokenizer.convert_tokens_to_ids("<|vision_start|>"))
-        self.vision_end_token_id = int(self.tokenizer.convert_tokens_to_ids("<|vision_end|>"))
-        self.eov_token_id = int(self.tokenizer.convert_tokens_to_ids("<EOV>"))
-        self.state_token_id = int(self.tokenizer.convert_tokens_to_ids("<state>"))
-        self.action_token_begin = int(self.tokenizer.convert_tokens_to_ids(self.action_tokens[0]))
-        self.action_token_end = self.action_token_begin + len(self.action_tokens)
-        self.action_token_end_with_markers = self.action_token_end + len(self.group_tokens)
-
-    def __len__(self) -> int:
-        return len(self.tokenizer)
-
-    def encode_text(self, text: str) -> list[int]:
-        return self.tokenizer(text, add_special_tokens=False)["input_ids"]
-
-    def decode(self, ids: Tensor | list[int]) -> str:
-        if isinstance(ids, Tensor):
-            ids = ids.detach().cpu().tolist()
-        return self.tokenizer.decode(ids, skip_special_tokens=False)
-
-    @staticmethod
-    def _resolve_template(template: str) -> str:
-        replacements = {
-            "<bos>": "",
-            "<eos>": "<|endoftext|>",
-            "<chat_user_prefix>": "",
-            "<chat_user_suffix>": "",
-            "<chat_assistant_prefix>": "",
-        }
-        for placeholder, value in replacements.items():
-            template = template.replace(placeholder, value)
-        return template
-
-    def _parse(self, template: str) -> list[_Segment]:
-        template = self._resolve_template(template)
-        segments: list[_Segment] = []
-        last = 0
-        for match in self._PLACEHOLDER.finditer(template):
-            if match.start() > last:
-                segments.append(_Segment("static", template[last : match.start()]))
-            raw = match.group(1).strip()
-            if raw in {"EOC", "EOV"}:
-                segments.append(_Segment("control", raw))
-                last = match.end()
-                continue
-
-            max_tokens = None
-            limit = re.match(r"^(.+)_(\d+)$", raw)
-            if limit:
-                raw, max_tokens = limit.group(1), int(limit.group(2))
-            masked = raw.endswith("_!")
-            key = raw[:-2] if masked else raw
-            if "_" not in key:
-                token = f"<{raw}>"
-                token_id = self.tokenizer.convert_tokens_to_ids(token)
-                if token_id is None:
-                    raise ValueError(f"Unknown G0.5 template token {token!r}.")
-                segments.append(_Segment("static", token))
-            else:
-                sample_key, processor = key.rsplit("_", 1)
-                segments.append(
-                    _Segment(
-                        "dynamic",
-                        sample_key=sample_key,
-                        processor=processor,
-                        masked=masked,
-                        max_tokens=max_tokens,
-                    )
-                )
-            last = match.end()
-        if last < len(template):
-            segments.append(_Segment("static", template[last:]))
-        return segments
-
-    @staticmethod
-    def _slice_segments(segments: list[_Segment], mode: str | None, *, pred_eov: bool) -> list[_Segment]:
-        eoc = next(
-            (
-                index
-                for index, segment in enumerate(segments)
-                if segment.kind == "control" and segment.content == "EOC"
-            ),
-            None,
-        )
-        eov = next(
-            (
-                index
-                for index, segment in enumerate(segments)
-                if segment.kind == "control" and segment.content == "EOV"
-            ),
-            None,
-        )
-
-        def strip(values: list[_Segment], keep_eov: bool) -> list[_Segment]:
-            output: list[_Segment] = []
-            after_eoc = False
-            for segment in values:
-                if segment.kind == "control":
-                    if segment.content == "EOC":
-                        after_eoc = True
-                    elif segment.content == "EOV" and keep_eov:
-                        output.append(
-                            _Segment(
-                                "dynamic",
-                                content="<EOV>",
-                                processor="text",
-                                masked=not pred_eov,
-                            )
-                        )
-                    continue
-                if after_eoc and segment.kind == "static":
-                    output.append(
-                        _Segment("dynamic", content=segment.content, processor="text", masked=False)
-                    )
-                else:
-                    output.append(segment)
-            return output
-
-        if mode == "context":
-            return strip(segments if eoc is None else segments[:eoc], keep_eov=True)
-        if mode == "prefix":
-            return strip(segments if eov is None else segments[: eov + 1], keep_eov=True)
-        if mode == "suffix":
-            return [] if eov is None else strip(segments[eov + 1 :], keep_eov=False)
-        return strip(segments, keep_eov=True)
-
-    def _serialize_segment(
-        self,
-        segment: _Segment,
-        sample: dict[str, Any],
-        *,
-        action_codec: Any | None,
-    ) -> tuple[list[int], list[int], list[float]]:
-        if segment.kind == "static":
-            ids = self.encode_text(segment.content)
-            return ids, [IGNORE_INDEX] * len(ids), [float(G05TokenType.TEXT)] * len(ids)
-
-        if segment.sample_key:
-            if segment.sample_key not in sample:
-                raise KeyError(f"G0.5 prompt is missing sample field {segment.sample_key!r}.")
-            value = sample[segment.sample_key]
-        else:
-            value = segment.content
-
-        if segment.processor == "text":
-            ids = self.encode_text(value if isinstance(value, str) else str(value))
-            token_type = G05TokenType.TEXT if segment.masked else G05TokenType.PRED_TEXT
-            labels = [IGNORE_INDEX] * len(ids) if segment.masked else ids.copy()
-        elif segment.processor == "image":
-            if not isinstance(value, (tuple, list)) or len(value) != 2:
-                raise ValueError("G0.5 image placeholders require an (height, width) pair.")
-            height, width = (int(item) for item in value)
-            vision = self.model_config["vision"]
-            count = (height // int(vision["patch_size"]) // int(vision["spatial_merge_size"])) * (
-                width // int(vision["patch_size"]) // int(vision["spatial_merge_size"])
-            )
-            ids = [self.vision_start_token_id] + [self.image_token_id] * count + [self.vision_end_token_id]
-            labels = [IGNORE_INDEX] * len(ids)
-            types = (
-                [float(G05TokenType.TEXT)] + [float(G05TokenType.IMAGE)] * count + [float(G05TokenType.TEXT)]
-            )
-            return ids, labels, types
-        elif segment.processor == "proprio":
-            state = value["value"] if isinstance(value, dict) else value
-            count = 1 if torch.as_tensor(state).ndim <= 1 else int(torch.as_tensor(state).shape[0])
-            ids = [self.state_token_id] * count
-            labels = [IGNORE_INDEX] * count
-            token_type = G05TokenType.PROPRIO
-        elif segment.processor == "action":
-            if action_codec is None:
-                raise RuntimeError(
-                    "This G0.5 training template includes ActionCodec targets, but the "
-                    "checkpoint has no native ActionCodec sidecar loaded."
-                )
-            ids = action_codec.encode_for_language(value)
-            labels = ids.copy()
-            token_type = G05TokenType.ACTION
-        else:
-            ids = self.encode_text(value if isinstance(value, str) else str(value))
-            labels = [IGNORE_INDEX] * len(ids) if segment.masked else ids.copy()
-            token_type = G05TokenType.COT
-
-        if segment.max_tokens is not None:
-            ids = ids[: segment.max_tokens]
-            labels = labels[: segment.max_tokens]
-        return ids, labels, [float(token_type)] * len(ids)
-
-    def _serialize(
-        self,
-        sample: dict[str, Any],
-        *,
-        mode: str | None,
-        action_codec: Any | None,
-    ) -> tuple[list[int], list[int], list[float]]:
-        pred_eov = bool(self.model_config.get("input_preprocessor", {}).get("pred_eov", False))
-        segments = self._slice_segments(self._parse(sample["template"]), mode, pred_eov=pred_eov)
-        ids: list[int] = []
-        labels: list[int] = []
-        types: list[float] = []
-        for segment in segments:
-            segment_ids, segment_labels, segment_types = self._serialize_segment(
-                segment, sample, action_codec=action_codec
-            )
-            ids.extend(segment_ids)
-            labels.extend(segment_labels)
-            types.extend(segment_types)
-        return ids, labels, types
-
-    def _pad(
-        self,
-        rows: list[tuple[list[int], list[int], list[float]]],
-        *,
-        right_align: bool,
-        device: torch.device,
-    ) -> G05SequenceBatch:
-        length = max(len(row[0]) for row in rows)
-        input_ids = torch.full((len(rows), length), self.pad_token_id, dtype=torch.long, device=device)
-        labels = torch.full((len(rows), length), IGNORE_INDEX, dtype=torch.long, device=device)
-        token_types = torch.zeros((len(rows), length), dtype=torch.float32, device=device)
-        for index, (ids, row_labels, types) in enumerate(rows):
-            start = length - len(ids) if right_align else 0
-            stop = start + len(ids)
-            input_ids[index, start:stop] = torch.tensor(ids, dtype=torch.long, device=device)
-            labels[index, start:stop] = torch.tensor(row_labels, dtype=torch.long, device=device)
-            token_types[index, start:stop] = torch.tensor(types, dtype=torch.float32, device=device)
-        return G05SequenceBatch(input_ids, labels, token_types)
-
-    def encode_inference(
-        self,
-        samples: list[dict[str, Any]],
-        *,
-        device: torch.device,
-    ) -> G05SequenceBatch:
-        rows = [self._serialize(sample, mode="context", action_codec=None) for sample in samples]
-        return self._pad(rows, right_align=True, device=device)
-
-    def encode_train(
-        self,
-        samples: list[dict[str, Any]],
-        *,
-        device: torch.device,
-        action_codec: Any | None,
-    ) -> G05SequenceBatch:
-        prefix_rows = [
-            self._serialize(sample, mode="prefix", action_codec=action_codec) for sample in samples
-        ]
-        suffix_rows = [
-            self._serialize(sample, mode="suffix", action_codec=action_codec) for sample in samples
-        ]
-        prefix = self._pad(prefix_rows, right_align=True, device=device)
-        suffix = self._pad(suffix_rows, right_align=False, device=device)
-        return G05SequenceBatch(
-            input_ids=torch.cat((prefix.input_ids, suffix.input_ids), dim=1),
-            labels=torch.cat((prefix.labels, suffix.labels), dim=1),
-            token_types=torch.cat((prefix.token_types, suffix.token_types), dim=1),
-            split_index=prefix.input_ids.shape[1],
-        )

src/lerobot/policies/g05/processor_g05.py

@@ -14,16 +14,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Serializable preprocessing and inverse projection for G0.5."""
+"""G0.5 tokenization, serialization, preprocessing, and inverse projection."""
 
 from __future__ import annotations
 
+import re
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any
 
 import torch
 import torchvision.transforms.functional as vision_functional
+from torch import Tensor
+from transformers import AutoTokenizer
 
 from lerobot.configs import recipe as recipe_module
 from lerobot.configs.recipe import TrainingRecipe
@@ -851,3 +854,330 @@ def make_g05_pre_post_processors(
         to_output=transition_to_policy_action,
     )
     return preprocessor, postprocessor
+
+
+IGNORE_INDEX = -100
+
+
+class G05TokenType:
+    """Token categories stored in G0.5's attention-mask tensor."""
+
+    PADDING = 0
+    IMAGE = 1
+    PROPRIO = 2
+    ACTION = 3
+    TEXT = 4
+    COT = 5
+    PRED_TEXT = 6
+
+
+@dataclass
+class G05SequenceBatch:
+    input_ids: Tensor
+    labels: Tensor
+    token_types: Tensor
+    split_index: int | None = None
+
+
+@dataclass
+class _Segment:
+    kind: str
+    content: str = ""
+    sample_key: str = ""
+    processor: str = ""
+    masked: bool = False
+    max_tokens: int | None = None
+
+
+class G05Tokenizer:
+    """Checkpoint-compatible G0.5 tokenizer and template serializer.
+
+    G0.5 extends Qwen3.5's tokenizer with ActionCodec codes, per-group
+    residual markers, ``<EOV>``, and ``<state>``. Registration order is model
+    state: changing it changes the rows used by the tied language head.
+    """
+
+    _PLACEHOLDER = re.compile(r"<([^<>|]+)>")
+
+    def __init__(self, processor_path: str | Path, model_config: dict[str, Any]) -> None:
+        self.processor_path = Path(processor_path)
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            self.processor_path,
+            trust_remote_code=False,
+            local_files_only=True,
+        )
+        self.model_config = model_config
+        at_config = model_config["AT_CONFIG"]
+        architecture = at_config["model_arch"]
+
+        self.action_tokens = [f"<action{index:04d}>" for index in range(int(architecture["codebook_size"]))]
+        parts = list(at_config["parts_meta"])
+        rule_patterns = tuple(at_config.get("rule_based_key_patterns") or ())
+        self.rule_parts = [name for name in parts if any(pattern in name for pattern in rule_patterns)]
+        self.neural_parts = [name for name in parts if name not in self.rule_parts]
+        self.group_tokens = [
+            f"<{part}_{residual}>"
+            for residual in range(int(architecture["n_codebooks"]))
+            for part in self.neural_parts
+        ] + [f"<{part}>" for part in self.rule_parts]
+        self.tokenizer.add_tokens(self.action_tokens + self.group_tokens + ["<EOV>", "<state>"])
+
+        self.pad_token_id = int(model_config["pad_token_id"])
+        self.eos_token_id = int(model_config["eos_token_id"])
+        self.image_token_id = int(model_config["image_token_index"])
+        self.vision_start_token_id = int(self.tokenizer.convert_tokens_to_ids("<|vision_start|>"))
+        self.vision_end_token_id = int(self.tokenizer.convert_tokens_to_ids("<|vision_end|>"))
+        self.eov_token_id = int(self.tokenizer.convert_tokens_to_ids("<EOV>"))
+        self.state_token_id = int(self.tokenizer.convert_tokens_to_ids("<state>"))
+        self.action_token_begin = int(self.tokenizer.convert_tokens_to_ids(self.action_tokens[0]))
+        self.action_token_end = self.action_token_begin + len(self.action_tokens)
+        self.action_token_end_with_markers = self.action_token_end + len(self.group_tokens)
+
+    def __len__(self) -> int:
+        return len(self.tokenizer)
+
+    def encode_text(self, text: str) -> list[int]:
+        return self.tokenizer(text, add_special_tokens=False)["input_ids"]
+
+    def decode(self, ids: Tensor | list[int]) -> str:
+        if isinstance(ids, Tensor):
+            ids = ids.detach().cpu().tolist()
+        return self.tokenizer.decode(ids, skip_special_tokens=False)
+
+    @staticmethod
+    def _resolve_template(template: str) -> str:
+        replacements = {
+            "<bos>": "",
+            "<eos>": "<|endoftext|>",
+            "<chat_user_prefix>": "",
+            "<chat_user_suffix>": "",
+            "<chat_assistant_prefix>": "",
+        }
+        for placeholder, value in replacements.items():
+            template = template.replace(placeholder, value)
+        return template
+
+    def _parse(self, template: str) -> list[_Segment]:
+        template = self._resolve_template(template)
+        segments: list[_Segment] = []
+        last = 0
+        for match in self._PLACEHOLDER.finditer(template):
+            if match.start() > last:
+                segments.append(_Segment("static", template[last : match.start()]))
+            raw = match.group(1).strip()
+            if raw in {"EOC", "EOV"}:
+                segments.append(_Segment("control", raw))
+                last = match.end()
+                continue
+
+            max_tokens = None
+            limit = re.match(r"^(.+)_(\d+)$", raw)
+            if limit:
+                raw, max_tokens = limit.group(1), int(limit.group(2))
+            masked = raw.endswith("_!")
+            key = raw[:-2] if masked else raw
+            if "_" not in key:
+                token = f"<{raw}>"
+                token_id = self.tokenizer.convert_tokens_to_ids(token)
+                if token_id is None:
+                    raise ValueError(f"Unknown G0.5 template token {token!r}.")
+                segments.append(_Segment("static", token))
+            else:
+                sample_key, processor = key.rsplit("_", 1)
+                segments.append(
+                    _Segment(
+                        "dynamic",
+                        sample_key=sample_key,
+                        processor=processor,
+                        masked=masked,
+                        max_tokens=max_tokens,
+                    )
+                )
+            last = match.end()
+        if last < len(template):
+            segments.append(_Segment("static", template[last:]))
+        return segments
+
+    @staticmethod
+    def _slice_segments(segments: list[_Segment], mode: str | None, *, pred_eov: bool) -> list[_Segment]:
+        eoc = next(
+            (
+                index
+                for index, segment in enumerate(segments)
+                if segment.kind == "control" and segment.content == "EOC"
+            ),
+            None,
+        )
+        eov = next(
+            (
+                index
+                for index, segment in enumerate(segments)
+                if segment.kind == "control" and segment.content == "EOV"
+            ),
+            None,
+        )
+
+        def strip(values: list[_Segment], keep_eov: bool) -> list[_Segment]:
+            output: list[_Segment] = []
+            after_eoc = False
+            for segment in values:
+                if segment.kind == "control":
+                    if segment.content == "EOC":
+                        after_eoc = True
+                    elif segment.content == "EOV" and keep_eov:
+                        output.append(
+                            _Segment(
+                                "dynamic",
+                                content="<EOV>",
+                                processor="text",
+                                masked=not pred_eov,
+                            )
+                        )
+                    continue
+                if after_eoc and segment.kind == "static":
+                    output.append(
+                        _Segment("dynamic", content=segment.content, processor="text", masked=False)
+                    )
+                else:
+                    output.append(segment)
+            return output
+
+        if mode == "context":
+            return strip(segments if eoc is None else segments[:eoc], keep_eov=True)
+        if mode == "prefix":
+            return strip(segments if eov is None else segments[: eov + 1], keep_eov=True)
+        if mode == "suffix":
+            return [] if eov is None else strip(segments[eov + 1 :], keep_eov=False)
+        return strip(segments, keep_eov=True)
+
+    def _serialize_segment(
+        self,
+        segment: _Segment,
+        sample: dict[str, Any],
+        *,
+        action_codec: Any | None,
+    ) -> tuple[list[int], list[int], list[float]]:
+        if segment.kind == "static":
+            ids = self.encode_text(segment.content)
+            return ids, [IGNORE_INDEX] * len(ids), [float(G05TokenType.TEXT)] * len(ids)
+
+        if segment.sample_key:
+            if segment.sample_key not in sample:
+                raise KeyError(f"G0.5 prompt is missing sample field {segment.sample_key!r}.")
+            value = sample[segment.sample_key]
+        else:
+            value = segment.content
+
+        if segment.processor == "text":
+            ids = self.encode_text(value if isinstance(value, str) else str(value))
+            token_type = G05TokenType.TEXT if segment.masked else G05TokenType.PRED_TEXT
+            labels = [IGNORE_INDEX] * len(ids) if segment.masked else ids.copy()
+        elif segment.processor == "image":
+            if not isinstance(value, (tuple, list)) or len(value) != 2:
+                raise ValueError("G0.5 image placeholders require an (height, width) pair.")
+            height, width = (int(item) for item in value)
+            vision = self.model_config["vision"]
+            count = (height // int(vision["patch_size"]) // int(vision["spatial_merge_size"])) * (
+                width // int(vision["patch_size"]) // int(vision["spatial_merge_size"])
+            )
+            ids = [self.vision_start_token_id] + [self.image_token_id] * count + [self.vision_end_token_id]
+            labels = [IGNORE_INDEX] * len(ids)
+            types = (
+                [float(G05TokenType.TEXT)] + [float(G05TokenType.IMAGE)] * count + [float(G05TokenType.TEXT)]
+            )
+            return ids, labels, types
+        elif segment.processor == "proprio":
+            state = value["value"] if isinstance(value, dict) else value
+            count = 1 if torch.as_tensor(state).ndim <= 1 else int(torch.as_tensor(state).shape[0])
+            ids = [self.state_token_id] * count
+            labels = [IGNORE_INDEX] * count
+            token_type = G05TokenType.PROPRIO
+        elif segment.processor == "action":
+            if action_codec is None:
+                raise RuntimeError(
+                    "This G0.5 training template includes ActionCodec targets, but the "
+                    "checkpoint has no native ActionCodec sidecar loaded."
+                )
+            ids = action_codec.encode_for_language(value)
+            labels = ids.copy()
+            token_type = G05TokenType.ACTION
+        else:
+            ids = self.encode_text(value if isinstance(value, str) else str(value))
+            labels = [IGNORE_INDEX] * len(ids) if segment.masked else ids.copy()
+            token_type = G05TokenType.COT
+
+        if segment.max_tokens is not None:
+            ids = ids[: segment.max_tokens]
+            labels = labels[: segment.max_tokens]
+        return ids, labels, [float(token_type)] * len(ids)
+
+    def _serialize(
+        self,
+        sample: dict[str, Any],
+        *,
+        mode: str | None,
+        action_codec: Any | None,
+    ) -> tuple[list[int], list[int], list[float]]:
+        pred_eov = bool(self.model_config.get("input_preprocessor", {}).get("pred_eov", False))
+        segments = self._slice_segments(self._parse(sample["template"]), mode, pred_eov=pred_eov)
+        ids: list[int] = []
+        labels: list[int] = []
+        types: list[float] = []
+        for segment in segments:
+            segment_ids, segment_labels, segment_types = self._serialize_segment(
+                segment, sample, action_codec=action_codec
+            )
+            ids.extend(segment_ids)
+            labels.extend(segment_labels)
+            types.extend(segment_types)
+        return ids, labels, types
+
+    def _pad(
+        self,
+        rows: list[tuple[list[int], list[int], list[float]]],
+        *,
+        right_align: bool,
+        device: torch.device,
+    ) -> G05SequenceBatch:
+        length = max(len(row[0]) for row in rows)
+        input_ids = torch.full((len(rows), length), self.pad_token_id, dtype=torch.long, device=device)
+        labels = torch.full((len(rows), length), IGNORE_INDEX, dtype=torch.long, device=device)
+        token_types = torch.zeros((len(rows), length), dtype=torch.float32, device=device)
+        for index, (ids, row_labels, types) in enumerate(rows):
+            start = length - len(ids) if right_align else 0
+            stop = start + len(ids)
+            input_ids[index, start:stop] = torch.tensor(ids, dtype=torch.long, device=device)
+            labels[index, start:stop] = torch.tensor(row_labels, dtype=torch.long, device=device)
+            token_types[index, start:stop] = torch.tensor(types, dtype=torch.float32, device=device)
+        return G05SequenceBatch(input_ids, labels, token_types)
+
+    def encode_inference(
+        self,
+        samples: list[dict[str, Any]],
+        *,
+        device: torch.device,
+    ) -> G05SequenceBatch:
+        rows = [self._serialize(sample, mode="context", action_codec=None) for sample in samples]
+        return self._pad(rows, right_align=True, device=device)
+
+    def encode_train(
+        self,
+        samples: list[dict[str, Any]],
+        *,
+        device: torch.device,
+        action_codec: Any | None,
+    ) -> G05SequenceBatch:
+        prefix_rows = [
+            self._serialize(sample, mode="prefix", action_codec=action_codec) for sample in samples
+        ]
+        suffix_rows = [
+            self._serialize(sample, mode="suffix", action_codec=action_codec) for sample in samples
+        ]
+        prefix = self._pad(prefix_rows, right_align=True, device=device)
+        suffix = self._pad(suffix_rows, right_align=False, device=device)
+        return G05SequenceBatch(
+            input_ids=torch.cat((prefix.input_ids, suffix.input_ids), dim=1),
+            labels=torch.cat((prefix.labels, suffix.labels), dim=1),
+            token_types=torch.cat((prefix.token_types, suffix.token_types), dim=1),
+            split_index=prefix.input_ids.shape[1],
+        )

tests/policies/g05/test_action_codec_g05.py

@@ -16,7 +16,7 @@
 
 import torch
 
-from lerobot.policies.g05.action_codec_g05 import G05NativeActionCodec, _BinarySequenceCodec
+from lerobot.policies.g05.modeling_g05 import G05NativeActionCodec, _BinarySequenceCodec
 
 
 def _tiny_codec_config() -> dict:

tests/policies/g05/test_g05.py

@@ -28,8 +28,7 @@ from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import FeatureType, PolicyFeature
 from lerobot.policies.factory import get_policy_class, make_policy_config, make_pre_post_processors
 from lerobot.policies.g05.configuration_g05 import G05_CAMERA_PROFILES, G05_EMBODIMENT_MAPPINGS, G05Config
-from lerobot.policies.g05.modeling_g05 import G05Policy
+from lerobot.policies.g05.modeling_g05 import G05_RUNTIME_PREDICT_COT, G05NativeBackend, G05Policy
-from lerobot.policies.g05.native_g05 import G05_RUNTIME_PREDICT_COT, G05NativeBackend
 from lerobot.processor import PolicyProcessorPipeline
 from lerobot.utils.constants import ACTION, OBS_STATE, POLICY_PREPROCESSOR_DEFAULT_NAME