From 2e3be32361c82ef58028d7b0aaed1768baf5ed41 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Tue, 28 Jul 2026 11:36:21 +0200 Subject: [PATCH] feat(g05): add OpenGalaxea G0.5 policy integration --- docs/source/_toctree.yml | 2 + docs/source/g05.mdx | 110 +++++ pyproject.toml | 5 + src/lerobot/policies/__init__.py | 2 + src/lerobot/policies/g05/README.md | 8 + src/lerobot/policies/g05/__init__.py | 6 + src/lerobot/policies/g05/configuration_g05.py | 282 ++++++++++++ .../policies/g05/convert_g05_checkpoint.py | 357 +++++++++++++++ .../policies/g05/inference/g05_adapter.py | 11 +- src/lerobot/policies/g05/modeling_g05.py | 334 ++++++++++++++ src/lerobot/policies/g05/processor_g05.py | 358 +++++++++++++++ .../templates/lerobot_modelcard_template.md | 5 +- tests/policies/g05/test_g05.py | 417 ++++++++++++++++++ tests/runtime/test_g05_adapter.py | 6 +- uv.lock | 6 + 15 files changed, 1902 insertions(+), 7 deletions(-) create mode 100644 docs/source/g05.mdx create mode 100644 src/lerobot/policies/g05/README.md create mode 100644 src/lerobot/policies/g05/__init__.py create mode 100644 src/lerobot/policies/g05/configuration_g05.py create mode 100644 src/lerobot/policies/g05/convert_g05_checkpoint.py create mode 100644 src/lerobot/policies/g05/modeling_g05.py create mode 100644 src/lerobot/policies/g05/processor_g05.py create mode 100644 tests/policies/g05/test_g05.py diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml index 7f7a34e6a..6ae00127d 100644 --- a/docs/source/_toctree.yml +++ b/docs/source/_toctree.yml @@ -71,6 +71,8 @@ title: EO-1 - local: lingbot_va title: LingBot-VA + - local: g05 + title: OpenGalaxea G0.5 - local: fastwam title: FastWAM - local: evo1 diff --git a/docs/source/g05.mdx b/docs/source/g05.mdx new file mode 100644 index 000000000..2dea16ee1 --- /dev/null +++ b/docs/source/g05.mdx @@ -0,0 +1,110 @@ +# OpenGalaxea G0.5 + +G0.5 is a Qwen3.5-2B vision-language-action model that can generate embodied +reasoning and actions in one stream. LeRobot exposes the action phase as System 1 +and the optional native chain-of-thought phase as System 2. They are not separate +models: the runtime obtains both from one inference call and the action stays +conditioned on the same post-reasoning KV state. + +> [!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 local checkpoint. + +## Supported checkpoint contracts + +| Profile | Released action path | Raw → policy layout | Cameras | Chunk | Normalization | +| ---------------- | ------------------------------------------ | -------------------------------------------- | ------------------------------ | -----------------: | ------------------------ | +| `g05-base` | Explicitly selected AR ActionCodec or flow | selected named embodiment | selected named camera contract | checkpoint-defined | resolved checkpoint mode | +| `g05-libero` | Continuous flow | right EEF 6 + gripper 1 → 20D grouped layout | exterior, right wrist | 32 | stepwise q01/q99 | +| `g05-robotwin20` | Continuous flow | two arms 6+gripper → 20D grouped layout | high, left wrist, right wrist | 32 | stepwise q01/q99 | + +The converter stores the resolved Hydra model, processor, ActionCodec metadata, +statistics, exact prompt template, source revision, and license with the converted +checkpoint. Loading rejects a different head, horizon, processor mode, or +normalization contract. + +The named `atomic_4` adapter is intentionally separate from LIBERO. Its raw state +is EEF relative xyz+quaternion, base xyz+quaternion, and two gripper positions +(16D). Its action is EEF delta xyz+rpy, gripper, four mobile-base commands, and a +control-mode flag (12D). This needs G0.5's 27D whole-body layout; a released 20D +LIBERO checkpoint is rejected. + +## Install + +Install LeRobot's small config 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 +``` + +## Convert a local checkpoint + +The command never contacts the Hub. Point it at a complete local bundle containing +the checkpoint Hydra config, weights, `dataset_stats.json`, +`action_tokenizer.pt`, and `hf_processor/`. + +```bash +uv run python -m lerobot.policies.g05.convert_g05_checkpoint \ + --source-dir /path/to/checkpoints/g05-libero \ + --output-dir outputs/g05-libero-lerobot \ + --profile g05-libero \ + --license-file /path/to/GalaxeaVLA/LICENSE-G0.5 +``` + +Use `--profile g05-base` for the base checkpoint or `--profile +g05-robotwin20` for RoboTwin. `conversion_report.json` records every mapped, +missing, unexpected, duplicate, and shape-mismatched tensor; conversion fails +when strict required-state validation fails. + +## Interactive System 1 and System 2 runtime + +System 1 executes the selected ActionCodec or flow chunk directly: + +```bash +lerobot-rollout \ + --policy.path=outputs/g05-base-lerobot \ + --language --direct_subtask \ + --task="pick up the cup" \ + --mode=action +``` + +System 2 is available only when the converted checkpoint metadata has +`predict_cot=true`. The adapter forwards the operator task byte-for-byte and +returns CoT telemetry and the matching action chunk atomically. It never samples +`task_aug` text, launches a second planner, or feeds generated CoT back as a +replacement task. + +```bash +lerobot-rollout \ + --policy.path=outputs/g05-system2-lerobot \ + --language \ + --task="clear the table" \ + --mode=action +``` + +## Validation status + +CPU unit tests cover factory loading, config incompatibilities, prompt pass-through, +LIBERO and `atomic_4` mappings, padding masks, inverse action projection, strict +conversion diagnostics, a finite forward/backward/update, and save/reload parity. +Numerical author-oracle parity and 50-episode LIBERO/RoboTwin evaluation require +licensed checkpoint access and suitable CUDA hardware; no benchmark number is +claimed by this integration until those gates run. + +```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 +``` diff --git a/pyproject.toml b/pyproject.toml index 9e88e8eca..308156656 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -238,6 +238,10 @@ 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"] # Features async = ["lerobot[grpcio-dep]", "lerobot[matplotlib-dep]"] @@ -322,6 +326,7 @@ all = [ "lerobot[hilserl]", "lerobot[vla_jepa]", "lerobot[lingbot_va]", + "lerobot[g05]", "lerobot[async]", "lerobot[dev]", "lerobot[test]", diff --git a/src/lerobot/policies/__init__.py b/src/lerobot/policies/__init__.py index a95d23b91..7f21d6680 100644 --- a/src/lerobot/policies/__init__.py +++ b/src/lerobot/policies/__init__.py @@ -20,6 +20,7 @@ from .eo1.configuration_eo1 import EO1Config as EO1Config from .evo1.configuration_evo1 import Evo1Config as Evo1Config from .factory import get_policy_class, make_policy, make_policy_config, make_pre_post_processors from .fastwam.configuration_fastwam import FastWAMConfig as FastWAMConfig +from .g05.configuration_g05 import G05Config as G05Config from .gaussian_actor.configuration_gaussian_actor import GaussianActorConfig as GaussianActorConfig from .groot.configuration_groot import GrootConfig as GrootConfig from .lingbot_va.configuration_lingbot_va import LingBotVAConfig as LingBotVAConfig @@ -48,6 +49,7 @@ __all__ = [ "EO1Config", "FastWAMConfig", "GaussianActorConfig", + "G05Config", "Evo1Config", "GrootConfig", "LingBotVAConfig", diff --git a/src/lerobot/policies/g05/README.md b/src/lerobot/policies/g05/README.md new file mode 100644 index 000000000..2a7c89920 --- /dev/null +++ b/src/lerobot/policies/g05/README.md @@ -0,0 +1,8 @@ +# OpenGalaxea G0.5 + +LeRobot wrapper and runtime adapter for +[OpenGalaxea G0.5](https://opengalaxea.github.io/G05/). + +The author implementation and checkpoints are licensed under the separate +[G0.5 Community License](https://huggingface.co/OpenGalaxea/G05/blob/main/licenses/LICENSE-G0.5), +not LeRobot's Apache-2.0 license. See the [LeRobot G0.5 guide](../../../../docs/source/g05.mdx). diff --git a/src/lerobot/policies/g05/__init__.py b/src/lerobot/policies/g05/__init__.py new file mode 100644 index 000000000..5d8c61e29 --- /dev/null +++ b/src/lerobot/policies/g05/__init__.py @@ -0,0 +1,6 @@ +"""OpenGalaxea G0.5 policy integration.""" + +from .configuration_g05 import G05Config +from .processor_g05 import make_g05_pre_post_processors + +__all__ = ["G05Config", "make_g05_pre_post_processors"] diff --git a/src/lerobot/policies/g05/configuration_g05.py b/src/lerobot/policies/g05/configuration_g05.py new file mode 100644 index 000000000..c346a6f59 --- /dev/null +++ b/src/lerobot/policies/g05/configuration_g05.py @@ -0,0 +1,282 @@ +# 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. + +"""Configuration for the OpenGalaxea G0.5 policy adapter.""" + +from dataclasses import dataclass, field +from typing import Any + +from lerobot.configs.policies import PreTrainedConfig +from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature +from lerobot.optim.optimizers import AdamWConfig +from lerobot.optim.schedulers import ConstantWithWarmupSchedulerConfig, LRSchedulerConfig +from lerobot.utils.constants import ACTION, OBS_STATE + +G05_SOURCE_REVISION = "b34966f387dd2ae0f003143b81494afd9213e613" +G05_HUB_REVISION = "e312be81e90c56a55bcb26b57429bd39a335b449" + +G05_CAMERA_PROFILES: dict[str, tuple[str, ...]] = { + "libero": ( + "observation.images.image", + "observation.images.wrist_image", + ), + "robotwin20": ( + "observation.images.cam_high", + "observation.images.cam_left_wrist", + "observation.images.cam_right_wrist", + ), + "atomic_4": ( + "observation.images.robot0_agentview_left", + "observation.images.robot0_eye_in_hand", + "observation.images.robot0_agentview_right", + ), +} + +G05_CAMERA_SIZE_PROFILES: dict[str, dict[str, tuple[int, int]]] = { + "libero": dict.fromkeys(G05_CAMERA_PROFILES["libero"], (224, 224)), + "robotwin20": dict.fromkeys(G05_CAMERA_PROFILES["robotwin20"], (256, 256)), + "atomic_4": dict.fromkeys(G05_CAMERA_PROFILES["atomic_4"], (256, 256)), +} + + +def make_g05_prompt_template(num_images: int, *, predict_cot: bool, flow_only: bool) -> str: + """Reproduce the selected author SamplesBuilder template exactly.""" + + images = "".join(f"" for index in range(num_images)) + prefix = ( + f"{images}" + "Embodiment: ; Task: " + "State: ;" + "" + ) + if predict_cot: + action = "Action: " if flow_only else "Action: |" + return f"{prefix}\n|{action}" + if flow_only: + # BaseActionSamplesBuilderFMOnly intentionally has no chat wrapper. + return ( + f"{images}Embodiment: ; " + "Task: State: ;\n" + "Action: " + ) + return f"{prefix}Action: |" + + +# Raw dimensions are inserted in these exact policy slots. The G0.5 shared layout is: +# left_control[9] | left_gripper[1] | right_control[9] | right_gripper[1] | lower_body[7]. +# LIBERO uses only the right EEF delta and right gripper. atomic_4 is a single-arm mobile +# manipulator and therefore has a deliberately separate state/action map. +G05_EMBODIMENT_MAPPINGS: dict[str, dict[str, tuple[int, ...]]] = { + "libero": { + "state": (10, 11, 12, 13, 14, 15, 19), + "action": (10, 11, 12, 13, 14, 15, 19), + }, + "robotwin20": { + "state": (0, 1, 2, 3, 4, 5, 9, 10, 11, 12, 13, 14, 15, 19), + "action": (0, 1, 2, 3, 4, 5, 9, 10, 11, 12, 13, 14, 15, 19), + }, + "atomic_4": { + # EEF relative xyz+quat -> right_control[0:7], base xyz+quat -> lower_body[0:7], + # the two parallel-jaw qpos values -> the two one-dimensional gripper slots. + "state": (10, 11, 12, 13, 14, 15, 16, 20, 21, 22, 23, 24, 25, 26, 9, 19), + # EEF delta xyz+rpy -> right_control[0:6], gripper -> right_gripper, + # base motion[4] -> lower_body[0:4], control mode -> lower_body[4]. + "action": (10, 11, 12, 13, 14, 15, 19, 20, 21, 22, 23, 24), + }, +} + +_PROFILE_DEFAULTS = { + "g05-base": ("checkpoint", 20, 16), + "g05-libero": ("q01_q99", 20, 32), + "g05-robotwin20": ("q01_q99", 20, 32), +} + + +@PreTrainedConfig.register_subclass("g05") +@dataclass +class G05Config(PreTrainedConfig): + """LeRobot-side, checkpoint-auditable configuration for G0.5. + + ``author_model_config`` is populated by the conversion script from the selected + checkpoint's Hydra config. It is intentionally checkpoint state rather than a + collection of guessed LeRobot defaults. + """ + + checkpoint_profile: str = "g05-base" + embodiment: str = "libero" + action_head: str = "actioncodec" # actioncodec (AR) or flow (continuous) + runtime_system: str = "system1" # system1 actions, or unified system2 CoT+actions + predict_cot: bool = False + discrete_action: bool = True + continuous_action: bool = False + return_continuous_action: bool = False + + policy_action_dim: int = 20 + policy_state_dim: int = 20 + raw_action_dim: int = 7 + raw_state_dim: int = 7 + chunk_size: int = 16 + normalization_mode: str = "checkpoint" + use_stepwise_action_norm: bool = False + gripper_indices: tuple[int, ...] = (6,) + camera_order: tuple[str, ...] = field(default_factory=lambda: G05_CAMERA_PROFILES["libero"]) + camera_sizes: dict[str, tuple[int, int]] = field(default_factory=dict) + image_mean: tuple[float, float, float] = (0.5, 0.5, 0.5) + image_std: tuple[float, float, float] = (0.5, 0.5, 0.5) + num_input_images: int = 0 + + author_source_revision: str = G05_SOURCE_REVISION + source_checkpoint_revision: str = G05_HUB_REVISION + author_model_config: dict[str, Any] = field(default_factory=dict) + processor_metadata: dict[str, Any] = field(default_factory=dict) + action_codec_metadata: dict[str, Any] = field(default_factory=dict) + prompt_template: str = "" + + normalization_mapping: dict[str, NormalizationMode] = field( + default_factory=lambda: { + "VISUAL": NormalizationMode.IDENTITY, + "STATE": NormalizationMode.IDENTITY, + "ACTION": NormalizationMode.IDENTITY, + } + ) + optimizer_lr: float = 8e-5 + optimizer_betas: tuple[float, float] = (0.9, 0.95) + optimizer_weight_decay: float = 0.01 + optimizer_grad_clip_norm: float = 1.0 + scheduler_warmup_steps: int = 500 + + def __post_init__(self) -> None: + super().__post_init__() + self.camera_order = tuple(self.camera_order) + self.camera_sizes = {key: tuple(size) for key, size in (self.camera_sizes or {}).items()} + if not self.camera_sizes and self.embodiment in G05_CAMERA_SIZE_PROFILES: + self.camera_sizes = G05_CAMERA_SIZE_PROFILES[self.embodiment].copy() + if self.num_input_images == 0: + self.num_input_images = len(self.camera_order) * self.n_obs_steps + if not self.prompt_template: + self.prompt_template = make_g05_prompt_template( + self.num_input_images, + predict_cot=self.predict_cot, + flow_only=self.continuous_action and not self.discrete_action, + ) + if self.checkpoint_profile not in _PROFILE_DEFAULTS and self.checkpoint_profile != "custom": + raise ValueError( + f"Unknown G0.5 checkpoint_profile={self.checkpoint_profile!r}; " + f"expected one of {sorted(_PROFILE_DEFAULTS)} or 'custom'." + ) + if self.action_head not in {"actioncodec", "flow"}: + raise ValueError("action_head must be 'actioncodec' or 'flow'.") + if self.runtime_system not in {"system1", "system2"}: + raise ValueError("runtime_system must be 'system1' or 'system2'.") + if self.runtime_system == "system2" and not self.predict_cot: + raise ValueError("G0.5 System 2 requires predict_cot=True in the converted checkpoint.") + if self.action_head == "actioncodec" and not self.discrete_action: + raise ValueError("The ActionCodec runtime requires discrete_action=True.") + if self.action_head == "flow" and not self.continuous_action: + raise ValueError("The flow runtime requires continuous_action=True.") + if self.action_head == "flow" and not self.return_continuous_action: + raise ValueError("The flow runtime requires return_continuous_action=True.") + if not (self.discrete_action or self.continuous_action): + raise ValueError("At least one G0.5 action path must be enabled.") + if self.embodiment not in G05_EMBODIMENT_MAPPINGS: + raise ValueError(f"No named G0.5 embodiment mapping for {self.embodiment!r}.") + if self.embodiment == "atomic_4": + if self.policy_action_dim < 27 or self.policy_state_dim < 27: + raise ValueError( + "atomic_4 includes mobile-base/control-mode semantics and requires the 27D " + "G0.5 shared layout; a 20D LIBERO checkpoint is incompatible." + ) + if self.raw_action_dim != 12 or self.raw_state_dim != 16: + raise ValueError("atomic_4 requires raw_state_dim=16 and raw_action_dim=12.") + mapping = G05_EMBODIMENT_MAPPINGS.get(self.embodiment) + if mapping is not None: + if len(mapping["state"]) != self.raw_state_dim: + raise ValueError("raw_state_dim does not match the selected embodiment mapping.") + if len(mapping["action"]) != self.raw_action_dim: + raise ValueError("raw_action_dim does not match the selected embodiment mapping.") + if max(mapping["state"]) >= self.policy_state_dim: + raise ValueError("Selected state mapping exceeds policy_state_dim.") + if max(mapping["action"]) >= self.policy_action_dim: + raise ValueError("Selected action mapping exceeds policy_action_dim.") + if self.normalization_mode not in {"checkpoint", "q01_q99", "z_score", "identity"}: + raise ValueError("normalization_mode must be checkpoint, q01_q99, z_score, or identity.") + if self.checkpoint_profile == "g05-libero": + if self.chunk_size != 32 or self.normalization_mode != "q01_q99": + raise ValueError("g05-libero requires a 32-step chunk and q01/q99 normalization.") + if self.action_head != "flow": + raise ValueError("The released g05-libero config enables only the continuous flow path.") + expected_cameras = G05_CAMERA_PROFILES.get(self.embodiment) + if expected_cameras is not None and tuple(self.camera_order) != expected_cameras: + raise ValueError( + f"{self.embodiment} camera_order must be {expected_cameras}, got {self.camera_order}." + ) + if set(self.camera_sizes) != set(self.camera_order): + raise ValueError("camera_sizes must contain exactly the ordered checkpoint camera keys.") + if self.num_input_images != len(self.camera_order) * self.n_obs_steps: + raise ValueError( + "num_input_images must equal len(camera_order) * n_obs_steps for the selected checkpoint." + ) + if any(len(size) != 2 or min(size) <= 0 for size in self.camera_sizes.values()): + raise ValueError("Every G0.5 camera size must be a positive (height, width) pair.") + if len(self.image_mean) != 3 or len(self.image_std) != 3 or min(self.image_std) <= 0: + raise ValueError("G0.5 image_mean/image_std must be three channels with positive std.") + + def validate_features(self) -> None: + if self.input_features is None: + self.input_features = {} + if self.output_features is None: + self.output_features = {} + state = self.input_features.get(OBS_STATE) + if state is not None and state.shape[-1] != self.raw_state_dim: + raise ValueError( + f"G0.5 {self.embodiment} expects {self.raw_state_dim} raw state dimensions, " + f"got {state.shape[-1]}." + ) + action = self.output_features.get(ACTION) + if action is not None and action.shape[-1] != self.raw_action_dim: + raise ValueError( + f"G0.5 {self.embodiment} expects {self.raw_action_dim} raw action dimensions, " + f"got {action.shape[-1]}." + ) + if OBS_STATE not in self.input_features: + self.input_features[OBS_STATE] = PolicyFeature( + type=FeatureType.STATE, shape=(self.raw_state_dim,) + ) + if ACTION not in self.output_features: + self.output_features[ACTION] = PolicyFeature( + type=FeatureType.ACTION, shape=(self.raw_action_dim,) + ) + + def get_optimizer_preset(self) -> AdamWConfig: + return AdamWConfig( + lr=self.optimizer_lr, + betas=self.optimizer_betas, + weight_decay=self.optimizer_weight_decay, + grad_clip_norm=self.optimizer_grad_clip_norm, + ) + + def get_scheduler_preset(self) -> LRSchedulerConfig | None: + return ConstantWithWarmupSchedulerConfig(num_warmup_steps=self.scheduler_warmup_steps) + + @property + def observation_delta_indices(self) -> list[int]: + return list(range(-(self.n_obs_steps - 1), 1)) + + @property + def action_delta_indices(self) -> list[int]: + return list(range(self.chunk_size)) + + @property + def reward_delta_indices(self) -> None: + return None diff --git a/src/lerobot/policies/g05/convert_g05_checkpoint.py b/src/lerobot/policies/g05/convert_g05_checkpoint.py new file mode 100644 index 000000000..fe6eb4a56 --- /dev/null +++ b/src/lerobot/policies/g05/convert_g05_checkpoint.py @@ -0,0 +1,357 @@ +#!/usr/bin/env python +# 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 + +"""Deterministically package a user-authorized G0.5 checkpoint for LeRobot. + +This command never downloads from the gated Hub. The user supplies a local checkpoint +after accepting Galaxea's license. +""" + +from __future__ import annotations + +import argparse +import json +import shutil +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any + +import torch +import yaml +from safetensors.torch import save_file + +from lerobot.policies.g05.configuration_g05 import ( + G05_CAMERA_PROFILES, + G05_HUB_REVISION, + G05_SOURCE_REVISION, + G05Config, +) +from lerobot.policies.g05.processor_g05 import make_g05_pre_post_processors +from lerobot.utils.constants import ACTION, OBS_STATE + +_EXACT_RENAMES = { + "model.embed_tokens.weight": "model.vlm.input_proj.weight", +} +_PREFIX_RENAMES = ( + ("model.joint_model.mixtures.vlm.", "model.vlm."), + ("model.joint_model.mixtures.action.", "model.action_expert."), + ("model.action_encoder.", "model.action_expert.input_proj."), + ("model.action_decoder.", "model.action_expert.output_proj."), +) +_REQUIRED_PREFIXES = ( + "backend.model.vlm.", + "backend.model.vision_tower.", + "backend.model.action_expert.", +) + + +@dataclass +class ConversionReport: + mapped: dict[str, str] = field(default_factory=dict) + missing: list[str] = field(default_factory=list) + unexpected: list[str] = field(default_factory=list) + duplicate: list[str] = field(default_factory=list) + shape_mismatched: dict[str, dict[str, list[int]]] = field(default_factory=dict) + + def fail_if_invalid(self) -> None: + if self.missing or self.unexpected or self.duplicate or self.shape_mismatched: + raise ValueError( + "G0.5 conversion failed strict validation: " + f"missing={len(self.missing)}, unexpected={len(self.unexpected)}, " + f"duplicate={len(self.duplicate)}, shape_mismatched={len(self.shape_mismatched)}" + ) + + +def _mapped_key(key: str) -> str: + key = _EXACT_RENAMES.get(key, key) + for old, new in _PREFIX_RENAMES: + if key.startswith(old): + key = new + key.removeprefix(old) + break + return key if key.startswith("backend.") else f"backend.{key}" + + +def convert_state_dict( + source: dict[str, torch.Tensor], + expected: dict[str, torch.Tensor] | None = None, +) -> tuple[dict[str, torch.Tensor], ConversionReport]: + """Map every tensor exactly once and optionally validate a target state dict.""" + + converted: dict[str, torch.Tensor] = {} + report = ConversionReport() + for old_key in sorted(source): + value = source[old_key] + if not isinstance(value, torch.Tensor): + report.unexpected.append(old_key) + continue + new_key = _mapped_key(old_key) + if new_key in converted: + report.duplicate.append(new_key) + continue + converted[new_key] = value.detach().cpu().contiguous() + report.mapped[old_key] = new_key + + if expected is not None: + report.missing = sorted(set(expected) - set(converted)) + report.unexpected.extend(sorted(set(converted) - set(expected))) + for key in sorted(set(expected) & set(converted)): + if expected[key].shape != converted[key].shape: + report.shape_mismatched[key] = { + "source": list(converted[key].shape), + "expected": list(expected[key].shape), + } + else: + for prefix in _REQUIRED_PREFIXES: + if not any(key.startswith(prefix) for key in converted): + report.missing.append(f"{prefix}*") + return converted, report + + +def _load_checkpoint(path: Path) -> dict[str, torch.Tensor]: + payload = torch.load(path, map_location="cpu", weights_only=True) + if isinstance(payload, dict) and isinstance(payload.get("model_state_dict"), dict): + payload = payload["model_state_dict"] + if not isinstance(payload, dict): + raise TypeError(f"Expected a state-dict mapping in {path}.") + return payload + + +def _profile_config(profile: str, hydra: dict[str, Any]) -> G05Config: + model = hydra.get("model", {}) + arch = model.get("model_arch", {}) + data = hydra.get("data", {}) + horizon = int(data.get("action_size", arch.get("horizon_steps", 16))) + predict_cot = bool(arch.get("predict_cot", False)) + discrete = bool(arch.get("discrete_action", True)) + continuous = bool(arch.get("continuous_action", False)) + norm_name = str(model.get("processor", {}).get("norm_default_mode", "")).lower() + checkpoint_normalization = { + "q01/q99": "q01_q99", + "z-score": "z_score", + "identity": "identity", + "dummy": "identity", + }.get(norm_name) + arch = dict(arch) + arch.pop("_target_", None) + num_input_images = int(arch.get("num_input_images", len(G05_CAMERA_PROFILES["libero"]))) + if profile == "g05-libero": + return G05Config( + checkpoint_profile=profile, + embodiment="libero", + action_head="flow", + runtime_system="system1", + predict_cot=predict_cot, + discrete_action=discrete, + continuous_action=continuous, + return_continuous_action=True, + chunk_size=horizon, + normalization_mode="q01_q99", + use_stepwise_action_norm=True, + camera_order=G05_CAMERA_PROFILES["libero"], + num_input_images=num_input_images, + author_model_config=arch, + processor_metadata=model.get("processor", {}), + action_codec_metadata=model.get("tokenizer", hydra.get("tokenizer", {})), + author_source_revision=G05_SOURCE_REVISION, + source_checkpoint_revision=G05_HUB_REVISION, + license="other", + tags=["g05", "robotics", "non-commercial"], + ) + if profile == "g05-robotwin20": + return G05Config( + checkpoint_profile=profile, + embodiment="robotwin20", + action_head="flow", + runtime_system="system1", + predict_cot=predict_cot, + discrete_action=discrete, + continuous_action=continuous, + return_continuous_action=True, + raw_state_dim=14, + raw_action_dim=14, + chunk_size=horizon, + normalization_mode="q01_q99", + use_stepwise_action_norm=True, + camera_order=G05_CAMERA_PROFILES["robotwin20"], + num_input_images=num_input_images, + author_model_config=arch, + processor_metadata=model.get("processor", {}), + action_codec_metadata=model.get("tokenizer", hydra.get("tokenizer", {})), + author_source_revision=G05_SOURCE_REVISION, + source_checkpoint_revision=G05_HUB_REVISION, + license="other", + tags=["g05", "robotics", "non-commercial"], + ) + action_head = "actioncodec" if discrete else "flow" + if checkpoint_normalization is None: + raise ValueError("g05-base conversion requires the resolved checkpoint processor.norm_default_mode.") + exceptions = model.get("processor", {}).get("norm_exception_mode") + if exceptions: + raise ValueError( + "g05-base has per-part normalization exceptions; select a concrete benchmark " + "profile or add a named LeRobot normalization contract instead of flattening it." + ) + return G05Config( + checkpoint_profile="g05-base", + embodiment="libero", + action_head=action_head, + runtime_system="system2" if predict_cot else "system1", + predict_cot=predict_cot, + discrete_action=discrete, + continuous_action=continuous, + return_continuous_action=action_head == "flow", + chunk_size=horizon, + normalization_mode=checkpoint_normalization, + use_stepwise_action_norm=bool(model.get("processor", {}).get("use_stepwise_action_norm", False)), + camera_order=G05_CAMERA_PROFILES["libero"], + num_input_images=num_input_images, + author_model_config=arch, + processor_metadata=model.get("processor", {}), + action_codec_metadata=model.get("tokenizer", hydra.get("tokenizer", {})), + author_source_revision=G05_SOURCE_REVISION, + source_checkpoint_revision=G05_HUB_REVISION, + license="other", + tags=["g05", "robotics", "non-commercial"], + ) + + +def _camera_sizes( + processor_metadata: dict[str, Any], camera_order: tuple[str, ...] +) -> dict[str, tuple[int, int]]: + images = (processor_metadata.get("shape_meta") or {}).get("images") or [] + by_lerobot_key = { + item.get("lerobot_key"): tuple(item["shape"][-2:]) + for item in images + if item.get("lerobot_key") and len(item.get("shape") or ()) >= 3 + } + if by_lerobot_key: + missing = [key for key in camera_order if key not in by_lerobot_key] + if missing: + raise ValueError(f"Checkpoint processor shape_meta is missing cameras {missing}.") + return {key: by_lerobot_key[key] for key in camera_order} + return {} + + +def convert_dataset_stats(payload: dict[str, Any], config: G05Config) -> dict[str, dict[str, torch.Tensor]]: + """Flatten author per-part stats into LeRobot raw state/action feature stats.""" + + if config.embodiment in payload: + payload = payload[config.embodiment] + shape_meta = config.processor_metadata.get("shape_meta") or {} + stat_names = ( + ("q01", "q99") + if config.normalization_mode == "q01_q99" + else ("mean", "std") + if config.normalization_mode == "z_score" + else () + ) + result: dict[str, dict[str, torch.Tensor]] = {} + for source_category, feature_name in (("state", OBS_STATE), ("action", ACTION)): + if not stat_names: + continue + category_stats = payload.get(source_category) + component_meta = shape_meta.get(source_category) + if not isinstance(category_stats, dict) or not isinstance(component_meta, list): + raise ValueError( + f"dataset_stats.json and processor shape_meta must define {source_category!r} components." + ) + converted_feature: dict[str, torch.Tensor] = {} + for short_name in stat_names: + source_name = ( + f"stepwise_{short_name}" + if source_category == "action" and config.use_stepwise_action_norm + else f"global_{short_name}" + ) + parts = [] + for component in component_meta: + key = component["key"] + if key not in category_stats or source_name not in category_stats[key]: + raise ValueError(f"Missing checkpoint statistic {source_category}.{key}.{source_name}.") + parts.append(torch.as_tensor(category_stats[key][source_name], dtype=torch.float32)) + converted_feature[short_name] = torch.cat(parts, dim=-1) + expected_width = config.raw_state_dim if feature_name == OBS_STATE else config.raw_action_dim + if converted_feature[stat_names[0]].shape[-1] != expected_width: + raise ValueError( + f"Flattened {feature_name} statistics have width " + f"{converted_feature[stat_names[0]].shape[-1]}, expected {expected_width}." + ) + result[feature_name] = converted_feature + return result + + +def convert_checkpoint( + source_dir: Path, + output_dir: Path, + profile: str, + *, + license_file: Path, + expected_state: dict[str, torch.Tensor] | None = None, +) -> ConversionReport: + hydra_path = source_dir / ".hydra" / "config.yaml" + stats_path = source_dir / "dataset_stats.json" + candidates = ( + source_dir / "model.pt", + source_dir / "checkpoints" / "model_state_dict.pt", + ) + checkpoint_path = next((path for path in candidates if path.is_file()), None) + tokenizer_path = source_dir / "action_tokenizer.pt" + processor_path = source_dir / "hf_processor" + required = [hydra_path, stats_path, tokenizer_path, processor_path, license_file] + missing_files = [str(path) for path in required if not path.exists()] + if checkpoint_path is None: + missing_files.append("model.pt or checkpoints/model_state_dict.pt") + if missing_files: + raise FileNotFoundError(f"Incomplete G0.5 checkpoint bundle: {missing_files}") + + hydra = yaml.safe_load(hydra_path.read_text()) + config = _profile_config(profile, hydra) + camera_sizes = _camera_sizes(config.processor_metadata, config.camera_order) + if camera_sizes: + config.camera_sizes = camera_sizes + stats_payload = json.loads(stats_path.read_text()) + lerobot_stats = convert_dataset_stats(stats_payload, config) + converted, report = convert_state_dict(_load_checkpoint(checkpoint_path), expected_state) + report.fail_if_invalid() + + output_dir.mkdir(parents=True, exist_ok=True) + save_file(converted, output_dir / "model.safetensors") + config._save_pretrained(output_dir) + preprocessor, postprocessor = make_g05_pre_post_processors(config, dataset_stats=lerobot_stats) + preprocessor.save_pretrained(output_dir) + postprocessor.save_pretrained(output_dir) + shutil.copy2(stats_path, output_dir / "g05_dataset_stats.json") + shutil.copy2(tokenizer_path, output_dir / "action_tokenizer.pt") + shutil.copytree(processor_path, output_dir / "hf_processor", dirs_exist_ok=True) + shutil.copy2(hydra_path, output_dir / "author_config.yaml") + shutil.copy2(license_file, output_dir / "LICENSE-G0.5") + (output_dir / "NOTICE").write_text( + "G0.5 is licensed under the G0.5 Community License Agreement " + "(Non-Commercial + Limited Patent License), not sold, Copyright © 2026 " + "Galaxea. All rights reserved by Galaxea. “Galaxea” and related marks are " + "trademarks of Galaxea or its affiliates.\n" + ) + (output_dir / "conversion_report.json").write_text(json.dumps(asdict(report), indent=2, sort_keys=True)) + return report + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--source-dir", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--profile", choices=("g05-base", "g05-libero", "g05-robotwin20"), required=True) + parser.add_argument("--license-file", type=Path, required=True) + args = parser.parse_args() + report = convert_checkpoint( + args.source_dir, args.output_dir, args.profile, license_file=args.license_file + ) + print(json.dumps(asdict(report), indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/src/lerobot/policies/g05/inference/g05_adapter.py b/src/lerobot/policies/g05/inference/g05_adapter.py index b3ed5ab33..27176427d 100644 --- a/src/lerobot/policies/g05/inference/g05_adapter.py +++ b/src/lerobot/policies/g05/inference/g05_adapter.py @@ -80,6 +80,8 @@ class G05PolicyAdapter(BaseLanguageAdapter): raw_mode = requested if raw_mode is None: raw_mode = _read_config(config, "runtime_system_mode") + if raw_mode is None: + raw_mode = _read_config(config, "runtime_system") if raw_mode is None: raw_mode = "auto" normalized = _MODE_ALIASES.get(str(raw_mode).strip().lower()) @@ -166,9 +168,12 @@ class G05PolicyAdapter(BaseLanguageAdapter): explicit_subtask = _as_text(metadata.get("subtask")) subtask = explicit_subtask or _extract_labeled_text(_SUBTASK_RE, text) if subtask: - previous = state.language_context.get("subtask") - if _set_generated_context(state, "subtask", subtask, label="subtask") and previous: - state.extra["prior_subtask"] = previous + # PR #4183's rollout observation provider treats language_context["subtask"] + # as the next policy command. G0.5 CoT must not replace the operator's task: + # it is telemetry from the same unified stream, not a separately supervised + # low-level prompt. Keep the parsed value visible without changing prompt routing. + state.extra["g05_subtask"] = subtask + state.log(f" subtask: {subtask}") explicit_memory = _as_text(metadata.get("memory")) memory = explicit_memory or _extract_labeled_text(_MEMORY_RE, text) diff --git a/src/lerobot/policies/g05/modeling_g05.py b/src/lerobot/policies/g05/modeling_g05.py new file mode 100644 index 000000000..f5b74ed6a --- /dev/null +++ b/src/lerobot/policies/g05/modeling_g05.py @@ -0,0 +1,334 @@ +# 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 + +"""First-class LeRobot wrapper around the community-licensed G0.5 implementation.""" + +from __future__ import annotations + +import importlib +import shutil +from collections import deque +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import torch +from huggingface_hub import snapshot_download +from torch import Tensor, nn + +from lerobot.configs.policies import PreTrainedConfig +from lerobot.policies.pretrained import PreTrainedPolicy +from lerobot.utils.constants import ACTION, OBS_STATE + +from .configuration_g05 import G05Config + + +def _author_backend(config: G05Config) -> nn.Module: + if not config.author_model_config: + raise ValueError( + "G0.5 author_model_config is empty. Convert an official checkpoint first, 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))) + + +class G05Policy(PreTrainedPolicy): + """LeRobot policy surface for G0.5's unified CoT and action stream.""" + + config_class = G05Config + name = "g05" + + 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) + 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() + + @classmethod + def from_pretrained( + cls, + pretrained_name_or_path: str | Path, + *, + config: G05Config | None = None, + **kwargs, + ) -> G05Policy: + resolved_path = Path(pretrained_name_or_path) + if not resolved_path.is_dir(): + resolved_path = Path( + snapshot_download( + repo_id=str(pretrained_name_or_path), + token=kwargs.get("token"), + cache_dir=kwargs.get("cache_dir"), + local_files_only=kwargs.get("local_files_only", False), + revision=kwargs.get("revision"), + ) + ) + if config is None: + config = PreTrainedConfig.from_pretrained( + resolved_path, + token=kwargs.get("token"), + cache_dir=kwargs.get("cache_dir"), + local_files_only=kwargs.get("local_files_only", False), + revision=kwargs.get("revision"), + ) + if not isinstance(config, G05Config): + raise TypeError(f"Expected a G05Config, got {type(config).__name__}.") + author_config = dict(config.author_model_config) + author_config["hf_processor_path"] = str(resolved_path / "hf_processor") + at_config = dict(author_config.get("AT_CONFIG") or {}) + at_config["ckpt_dir"] = str(resolved_path / "action_tokenizer.pt") + author_config["AT_CONFIG"] = at_config + author_config["pretrained_model_path"] = None + config.author_model_config = author_config + return super().from_pretrained( + resolved_path, + config=config, + **kwargs, + ) + + def _save_pretrained(self, save_directory: Path, state_dict: dict[str, Tensor] | None = None) -> None: + super()._save_pretrained(save_directory, state_dict=state_dict) + author_config = dict(self.config.author_model_config) + processor_value = author_config.get("hf_processor_path") + processor_path = Path(str(processor_value)) if processor_value else None + at_config = dict(author_config.get("AT_CONFIG") or {}) + tokenizer_value = at_config.get("ckpt_dir") + tokenizer_path = Path(str(tokenizer_value)) if tokenizer_value else None + roots = [ + path.parent for path in (processor_path, tokenizer_path) if path is not None and path.exists() + ] + + if ( + processor_path is not None + and processor_path.is_dir() + and processor_path.resolve() != (save_directory / "hf_processor").resolve() + ): + shutil.copytree(processor_path, save_directory / "hf_processor", dirs_exist_ok=True) + if ( + tokenizer_path is not None + and tokenizer_path.is_file() + and tokenizer_path.resolve() != (save_directory / "action_tokenizer.pt").resolve() + ): + shutil.copy2(tokenizer_path, save_directory / "action_tokenizer.pt") + for name in ( + "g05_dataset_stats.json", + "author_config.yaml", + "LICENSE-G0.5", + "NOTICE", + "conversion_report.json", + ): + source = next((root / name for root in roots if (root / name).is_file()), None) + if source is not None and source.resolve() != (save_directory / name).resolve(): + 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. + if (processor_path is not None and processor_path.exists()) or ( + tokenizer_path is not None and tokenizer_path.exists() + ): + portable = dict(author_config) + portable["hf_processor_path"] = "hf_processor" + portable_at = dict(portable.get("AT_CONFIG") or {}) + portable_at["ckpt_dir"] = "action_tokenizer.pt" + portable["AT_CONFIG"] = portable_at + portable["pretrained_model_path"] = None + runtime_config = self.config.author_model_config + self.config.author_model_config = portable + self.config._save_pretrained(save_directory) + self.config.author_model_config = runtime_config + + def reset(self) -> None: + self._action_queue.clear() + reset = getattr(self.backend, "reset", None) + if callable(reset): + reset() + + def get_optim_params(self) -> dict: + get_params = getattr(self.backend, "get_optim_params", None) + if callable(get_params): + return get_params() + return {"params": [parameter for parameter in self.parameters() if parameter.requires_grad]} + + @staticmethod + def _task_values(batch: Mapping[str, Any], task: str | None, batch_size: int) -> list[str]: + if task is not None: + return [task] * batch_size + value = batch.get("task") + if isinstance(value, str): + return [value] * batch_size + if isinstance(value, list | tuple) and len(value) == batch_size: + return [str(item) for item in value] + raise ValueError( + "G0.5 requires the already-selected LeRobot task string; no task augmentation " + "or model-local sampling is performed." + ) + + def _prepare_author_batch(self, batch: Mapping[str, Any], task: str | None = None) -> dict[str, Any]: + prepare = getattr(self.backend, "prepare_lerobot_batch", None) + if callable(prepare): + return prepare(batch, task=task, config=self.config) + + state = batch.get(OBS_STATE) + if not isinstance(state, Tensor): + raise ValueError(f"G0.5 requires tensor {OBS_STATE!r}.") + if state.ndim == 1: + state = state.unsqueeze(0) + batch_size = state.shape[0] + tasks = self._task_values(batch, task, batch_size) + state_mask = batch.get("proprio_dim_is_pad") + if state_mask is None: + state_mask = torch.zeros( + batch_size, self.config.policy_state_dim, dtype=torch.bool, device=state.device + ) + elif isinstance(state_mask, Tensor) and state_mask.ndim == 1: + state_mask = state_mask.unsqueeze(0).expand(batch_size, -1) + + pixel_values: dict[str, Tensor] = {} + for key in self.config.camera_order: + image = batch.get(key) + if not isinstance(image, Tensor): + raise ValueError(f"G0.5 requires camera {key!r}; camera order is checkpoint state.") + if image.ndim == 4: + image = image.unsqueeze(1) + pixel_values[key] = image + image_count = sum(image.shape[1] for image in pixel_values.values()) + if image_count != self.config.num_input_images: + raise ValueError( + f"G0.5 received {image_count} camera/history frames, but the checkpoint " + f"template requires {self.config.num_input_images}." + ) + + samples = [] + for index, raw_task in enumerate(tasks): + proprio = state[index] + if proprio.ndim == 1: + proprio = proprio.unsqueeze(0) + sample = { + "template": self.config.prompt_template, + # This is the author InputPreprocessor command slot. Keep it byte-for-byte + # unchanged; checkpoint-specific chat formatting occurs downstream. + "command": raw_task, + "embodiment": self.config.embodiment, + "proprio": { + "value": proprio, + "proprio_dim_is_pad": state_mask[index], + }, + } + frequency = self.config.processor_metadata.get("frequency") + if frequency is not None: + sample["frequency"] = frequency + if self.config.predict_cot: + sample["prompt"] = "predict subtask" + for image_index in range(self.config.num_input_images): + camera = self.config.camera_order[image_index % len(self.config.camera_order)] + sample[f"image{image_index}"] = self.config.camera_sizes[camera] + action = batch.get(ACTION) + if " tuple[Tensor, dict[str, Any]]: + prepared = self._prepare_author_batch(batch, task=task) + predict = getattr(self.backend, "predict_action", None) + result = predict(prepared) if callable(predict) else self.backend(prepared) + if isinstance(result, Tensor): + result = {ACTION: result} + if not isinstance(result, Mapping): + raise TypeError("G0.5 backend inference must return a tensor or mapping.") + + if self.config.action_head == "actioncodec": + action = result.get("ar_action", result.get(ACTION)) + else: + action = result.get(ACTION) + if not isinstance(action, Tensor): + raise ValueError(f"G0.5 {self.config.action_head} output is missing its action tensor.") + metadata = { + key: result[key] + for key in ("cot_text", "generated_ids", "decoded_action_tokens", "ar_absent_keys", "_timing") + if key in result + } + return action, metadata + + def predict_action_chunk_with_runtime( + self, batch: dict[str, Any], *, task: str + ) -> tuple[Tensor, dict[str, Any]]: + """Return System 1 actions and same-pass System 2 telemetry atomically.""" + + return self._run_inference(batch, task=task) + + @torch.no_grad() + def predict_action_chunk(self, batch: dict[str, Any], **kwargs) -> Tensor: + action, _ = self._run_inference(batch) + return action + + @torch.no_grad() + def select_action(self, batch: dict[str, Any], **kwargs) -> Tensor: + if not self._action_queue: + chunk = self.predict_action_chunk(batch, **kwargs) + if chunk.ndim != 3: + raise ValueError(f"G0.5 action chunk must be [B,T,D], got {tuple(chunk.shape)}.") + # LeRobot's synchronous select_action queue is intentionally batch-size one. + if chunk.shape[0] != 1: + raise ValueError( + "G0.5 select_action requires batch size 1; use predict_action_chunk for B>1." + ) + self._action_queue.extend(chunk[0]) + return self._action_queue.popleft().unsqueeze(0) + + def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict[str, Any] | None]: + prepared = self._prepare_author_batch(batch) + result = self.backend(prepared) + if isinstance(result, tuple) and len(result) == 2: + loss, loss_dict = result + elif isinstance(result, Mapping) and "loss" in result: + loss = result["loss"] + loss_dict = {key: value for key, value in result.items() if key != "loss"} + else: + raise TypeError("G0.5 training backend must return (loss, loss_dict) or {'loss': ...}.") + if not isinstance(loss, Tensor): + raise TypeError("G0.5 training loss must be a torch.Tensor.") + logging_values = { + key: value.detach().item() if isinstance(value, Tensor) and value.numel() == 1 else value + for key, value in (loss_dict or {}).items() + } + return loss, logging_values diff --git a/src/lerobot/policies/g05/processor_g05.py b/src/lerobot/policies/g05/processor_g05.py new file mode 100644 index 000000000..1a7701d6d --- /dev/null +++ b/src/lerobot/policies/g05/processor_g05.py @@ -0,0 +1,358 @@ +# 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 + +"""Serializable preprocessing and inverse projection for G0.5.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn.functional as functional + +from lerobot.configs.types import FeatureType, NormalizationMode, PipelineFeatureType, PolicyFeature +from lerobot.processor import ( + AddBatchDimensionProcessorStep, + DeviceProcessorStep, + NormalizerProcessorStep, + PolicyAction, + PolicyProcessorPipeline, + ProcessorStep, + ProcessorStepRegistry, + UnnormalizerProcessorStep, +) +from lerobot.processor.converters import ( + batch_to_transition, + policy_action_to_transition, + transition_to_batch, + transition_to_policy_action, +) +from lerobot.types import EnvTransition, TransitionKey +from lerobot.utils.constants import ( + ACTION, + OBS_STATE, + POLICY_POSTPROCESSOR_DEFAULT_NAME, + POLICY_PREPROCESSOR_DEFAULT_NAME, +) + +from .configuration_g05 import G05_EMBODIMENT_MAPPINGS, G05Config + + +def _copy_feature_tree( + features: dict[PipelineFeatureType, dict[str, PolicyFeature]], +) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]: + return {kind: values.copy() for kind, values in features.items()} + + +@dataclass +@ProcessorStepRegistry.register(name="g05_image_transform") +class G05ImageTransformStep(ProcessorStep): + """Apply the checkpoint's per-camera resize and ``[0,1]`` to ``[-1,1]`` transform.""" + + camera_order: tuple[str, ...] + camera_sizes: dict[str, tuple[int, int]] + mean: tuple[float, float, float] + std: tuple[float, float, float] + + def __post_init__(self) -> None: + self.camera_order = tuple(self.camera_order) + self.camera_sizes = {key: tuple(value) for key, value in self.camera_sizes.items()} + self.mean = tuple(self.mean) + self.std = tuple(self.std) + + def __call__(self, transition: EnvTransition) -> EnvTransition: + observation = transition.get(TransitionKey.OBSERVATION) + if observation is None: + return transition + missing = [key for key in self.camera_order if key not in observation] + if missing: + raise ValueError(f"G0.5 is missing camera(s) {missing}; required order is {self.camera_order}.") + transition = transition.copy() + observation = dict(observation) + for key in self.camera_order: + image = torch.as_tensor(observation[key]) + if image.ndim < 3 or image.shape[-3] != 3: + raise ValueError(f"G0.5 camera {key!r} must end in [3,H,W], got {image.shape}.") + was_floating_point = torch.is_floating_point(image) + image = image.float() + if not was_floating_point: + image = image / 255.0 + flat = image.reshape(-1, *image.shape[-3:]) + target_size = self.camera_sizes[key] + if tuple(flat.shape[-2:]) != target_size: + flat = functional.interpolate(flat, size=target_size, mode="bilinear", align_corners=False) + mean = flat.new_tensor(self.mean).view(1, 3, 1, 1) + std = flat.new_tensor(self.std).view(1, 3, 1, 1) + observation[key] = ((flat - mean) / std).reshape(*image.shape[:-3], 3, *target_size) + transition[TransitionKey.OBSERVATION] = observation + return transition + + def transform_features( + self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]] + ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]: + result = _copy_feature_tree(features) + observations = result.setdefault(PipelineFeatureType.OBSERVATION, {}) + for key in self.camera_order: + if key in observations: + height, width = self.camera_sizes[key] + observations[key] = PolicyFeature(type=FeatureType.VISUAL, shape=(3, height, width)) + return result + + def get_config(self) -> dict[str, Any]: + return { + "camera_order": list(self.camera_order), + "camera_sizes": {key: list(value) for key, value in self.camera_sizes.items()}, + "mean": list(self.mean), + "std": list(self.std), + } + + +@dataclass +@ProcessorStepRegistry.register(name="g05_embodiment_projection") +class G05EmbodimentProjectionStep(ProcessorStep): + """Map raw embodiment coordinates into the checkpoint's padded policy layout.""" + + embodiment: str + policy_state_dim: int + policy_action_dim: int + camera_order: tuple[str, ...] + + def __post_init__(self) -> None: + self.camera_order = tuple(self.camera_order) + if self.embodiment not in G05_EMBODIMENT_MAPPINGS: + raise ValueError(f"No projection is defined for G0.5 embodiment {self.embodiment!r}.") + + @property + def mapping(self) -> dict[str, tuple[int, ...]]: + return G05_EMBODIMENT_MAPPINGS[self.embodiment] + + @staticmethod + def _project(value: torch.Tensor, indices: tuple[int, ...], width: int) -> torch.Tensor: + if value.shape[-1] != len(indices): + raise ValueError(f"Raw G0.5 tensor has {value.shape[-1]} dimensions, expected {len(indices)}.") + projected = value.new_zeros(*value.shape[:-1], width) + projected[..., list(indices)] = value + return projected + + def __call__(self, transition: EnvTransition) -> EnvTransition: + transition = transition.copy() + observation = dict(transition.get(TransitionKey.OBSERVATION) or {}) + missing_cameras = [key for key in self.camera_order if key not in observation] + if missing_cameras and any(key.startswith("observation.images.") for key in observation): + raise ValueError( + f"G0.5 {self.embodiment} is missing camera(s) {missing_cameras}; " + f"required order is {self.camera_order}." + ) + if OBS_STATE in observation: + raw_state = observation[OBS_STATE] + observation[OBS_STATE] = self._project(raw_state, self.mapping["state"], self.policy_state_dim) + state_mask = torch.ones( + *raw_state.shape[:-1], + self.policy_state_dim, + dtype=torch.bool, + device=observation[OBS_STATE].device, + ) + state_mask[..., list(self.mapping["state"])] = False + complementary = dict(transition.get(TransitionKey.COMPLEMENTARY_DATA) or {}) + complementary["proprio_dim_is_pad"] = state_mask + complementary["g05_camera_order"] = self.camera_order + transition[TransitionKey.COMPLEMENTARY_DATA] = complementary + action = transition.get(TransitionKey.ACTION) + if isinstance(action, torch.Tensor): + transition[TransitionKey.ACTION] = self._project( + action, self.mapping["action"], self.policy_action_dim + ) + batch_shape = action.shape[:1] if action.ndim >= 3 else () + action_mask = torch.ones( + *batch_shape, self.policy_action_dim, dtype=torch.bool, device=action.device + ) + action_mask[..., list(self.mapping["action"])] = False + complementary = dict(transition.get(TransitionKey.COMPLEMENTARY_DATA) or {}) + complementary["action_dim_is_pad"] = action_mask + if "action_is_pad" not in complementary: + complementary["action_is_pad"] = torch.zeros( + *action.shape[:-1], dtype=torch.bool, device=action.device + ) + transition[TransitionKey.COMPLEMENTARY_DATA] = complementary + transition[TransitionKey.OBSERVATION] = observation + return transition + + def transform_features( + self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]] + ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]: + result = _copy_feature_tree(features) + observations = result.setdefault(PipelineFeatureType.OBSERVATION, {}) + if OBS_STATE in observations: + observations[OBS_STATE] = PolicyFeature(type=FeatureType.STATE, shape=(self.policy_state_dim,)) + actions = result.setdefault(PipelineFeatureType.ACTION, {}) + if ACTION in actions: + actions[ACTION] = PolicyFeature(type=FeatureType.ACTION, shape=(self.policy_action_dim,)) + return result + + def get_config(self) -> dict[str, Any]: + return { + "embodiment": self.embodiment, + "policy_state_dim": self.policy_state_dim, + "policy_action_dim": self.policy_action_dim, + "camera_order": list(self.camera_order), + } + + +@dataclass +@ProcessorStepRegistry.register(name="g05_inverse_action_projection") +class G05InverseActionProjectionStep(ProcessorStep): + """Project policy-layout actions back to the environment's exact raw layout.""" + + embodiment: str + policy_action_dim: int + + @property + def indices(self) -> tuple[int, ...]: + return G05_EMBODIMENT_MAPPINGS[self.embodiment]["action"] + + def __call__(self, transition: EnvTransition) -> EnvTransition: + action = transition.get(TransitionKey.ACTION) + if not isinstance(action, torch.Tensor): + return transition + if action.shape[-1] != self.policy_action_dim: + raise ValueError( + f"G0.5 policy action has {action.shape[-1]} dimensions, expected {self.policy_action_dim}." + ) + transition = transition.copy() + transition[TransitionKey.ACTION] = action[..., list(self.indices)] + return transition + + def transform_features( + self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]] + ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]: + result = _copy_feature_tree(features) + actions = result.setdefault(PipelineFeatureType.ACTION, {}) + actions[ACTION] = PolicyFeature(type=FeatureType.ACTION, shape=(len(self.indices),)) + return result + + def get_config(self) -> dict[str, Any]: + return {"embodiment": self.embodiment, "policy_action_dim": self.policy_action_dim} + + +def _normalization_mode(config: G05Config) -> NormalizationMode: + if config.normalization_mode == "q01_q99": + return NormalizationMode.QUANTILES + if config.normalization_mode == "z_score": + return NormalizationMode.MEAN_STD + return NormalizationMode.IDENTITY + + +def _project_stats( + config: G05Config, + dataset_stats: dict[str, dict[str, torch.Tensor]] | None, +) -> dict[str, dict[str, torch.Tensor]] | None: + if not dataset_stats: + return dataset_stats + result: dict[str, dict[str, torch.Tensor]] = {} + mapping = G05_EMBODIMENT_MAPPINGS[config.embodiment] + widths = {OBS_STATE: config.policy_state_dim, ACTION: config.policy_action_dim} + index_maps = {OBS_STATE: mapping["state"], ACTION: mapping["action"]} + for feature_name, stats in dataset_stats.items(): + if feature_name not in widths: + result[feature_name] = stats + continue + projected_stats: dict[str, torch.Tensor] = {} + for stat_name, raw_value in stats.items(): + value = torch.as_tensor(raw_value) + if value.shape[-1] != len(index_maps[feature_name]): + raise ValueError( + f"{feature_name}.{stat_name} has width {value.shape[-1]}, " + f"expected {len(index_maps[feature_name])} for {config.embodiment}." + ) + fill = { + "std": 1.0, + "q01": -1.0, + "q99": 1.0, + "min": -1.0, + "max": 1.0, + }.get(stat_name, 0.0) + projected = value.new_full((*value.shape[:-1], widths[feature_name]), fill) + projected[..., list(index_maps[feature_name])] = value + projected_stats[stat_name] = projected + result[feature_name] = projected_stats + return result + + +def make_g05_pre_post_processors( + config: G05Config, + dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None, +) -> tuple[ + PolicyProcessorPipeline[dict[str, Any], dict[str, Any]], + PolicyProcessorPipeline[PolicyAction, PolicyAction], +]: + """Build serializable G0.5 pipelines from checkpoint-authoritative metadata.""" + + if config.normalization_mode == "checkpoint" and not config.processor_metadata: + raise ValueError( + "normalization_mode='checkpoint' requires processor_metadata from the converted checkpoint." + ) + mode = _normalization_mode(config) + if mode is NormalizationMode.QUANTILES and dataset_stats: + for key in (OBS_STATE, ACTION): + if key in dataset_stats and not {"q01", "q99"} <= set(dataset_stats[key]): + raise ValueError(f"{key} requires real q01/q99 statistics; min/max must not be substituted.") + policy_features = dict(config.input_features or {}) + policy_features.update(config.output_features or {}) + policy_features[OBS_STATE] = PolicyFeature(type=FeatureType.STATE, shape=(config.policy_state_dim,)) + policy_features[ACTION] = PolicyFeature(type=FeatureType.ACTION, shape=(config.policy_action_dim,)) + projected_stats = _project_stats(config, dataset_stats) + norm_map = { + FeatureType.STATE: mode, + FeatureType.ACTION: mode, + FeatureType.VISUAL: NormalizationMode.IDENTITY, + } + + preprocessor = PolicyProcessorPipeline[dict[str, Any], dict[str, Any]]( + steps=[ + AddBatchDimensionProcessorStep(), + G05ImageTransformStep( + camera_order=config.camera_order, + camera_sizes=config.camera_sizes, + mean=config.image_mean, + std=config.image_std, + ), + G05EmbodimentProjectionStep( + embodiment=config.embodiment, + policy_state_dim=config.policy_state_dim, + policy_action_dim=config.policy_action_dim, + camera_order=config.camera_order, + ), + NormalizerProcessorStep( + features=policy_features, + norm_map=norm_map, + stats=projected_stats, + ), + DeviceProcessorStep(device=config.device), + ], + name=POLICY_PREPROCESSOR_DEFAULT_NAME, + to_transition=batch_to_transition, + to_output=transition_to_batch, + ) + postprocessor = PolicyProcessorPipeline[PolicyAction, PolicyAction]( + steps=[ + UnnormalizerProcessorStep( + features={ACTION: policy_features[ACTION]}, + norm_map={FeatureType.ACTION: mode}, + stats=projected_stats, + ), + G05InverseActionProjectionStep( + embodiment=config.embodiment, policy_action_dim=config.policy_action_dim + ), + DeviceProcessorStep(device="cpu"), + ], + name=POLICY_POSTPROCESSOR_DEFAULT_NAME, + to_transition=policy_action_to_transition, + to_output=transition_to_policy_action, + ) + return preprocessor, postprocessor diff --git a/src/lerobot/templates/lerobot_modelcard_template.md b/src/lerobot/templates/lerobot_modelcard_template.md index c0eb9893e..c450c32df 100644 --- a/src/lerobot/templates/lerobot_modelcard_template.md +++ b/src/lerobot/templates/lerobot_modelcard_template.md @@ -43,6 +43,8 @@ This is a Gaussian Actor policy (Gaussian policy with a tanh squash) — the pol [FastWAM](https://arxiv.org/abs/2603.16666) is a World Action Model policy that keeps video world-modeling during training but predicts actions directly at inference time, initializing its visual world-model components from the Wan2.2 video-diffusion stack. {% elif model_name == "lingbot_va" %} [LingBot-VA](https://github.com/Robbyant/lingbot-va) is an autoregressive video-action world-model policy built on the Wan2.2 video-diffusion stack. It interleaves the prediction of future video latents and robot actions in a single autoregressive sequence, feeding observed keyframes back into its KV cache for closed-loop world modeling. +{% elif model_name == "g05" %} +[OpenGalaxea G0.5](https://opengalaxea.github.io/G05/) is a Qwen3.5-based VLA that emits optional embodied reasoning and actions in one unified stream. The author model and weights use the separate non-commercial G0.5 Community License; inspect the checkpoint's license before use. {% else %} This is a **{{ model_name }}** policy trained with [LeRobot](https://github.com/huggingface/lerobot). {% endif %} @@ -84,7 +86,8 @@ This policy has been trained and pushed to the Hub using [LeRobot](https://githu "wall_x": "walloss", "evo1": "evo1", "fastwam": "fastwam", - "lingbot_va": "lingbot_va" + "lingbot_va": "lingbot_va", + "g05": "g05" } %} {% if policy_docs.get(model_name) %}Learn how to train and run it in the [LeRobot {{ model_name }} guide](https://huggingface.co/docs/lerobot/main/en/{{ policy_docs[model_name] }}), or browse the [full documentation](https://huggingface.co/docs/lerobot/index). {% else %}See the [full LeRobot documentation](https://huggingface.co/docs/lerobot/index). diff --git a/tests/policies/g05/test_g05.py b/tests/policies/g05/test_g05.py new file mode 100644 index 000000000..12f6930b3 --- /dev/null +++ b/tests/policies/g05/test_g05.py @@ -0,0 +1,417 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. + +from __future__ import annotations + +import json +import os +from pathlib import Path + +import pytest +import torch +from torch import nn + +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_EMBODIMENT_MAPPINGS, G05Config +from lerobot.policies.g05.convert_g05_checkpoint import convert_checkpoint, convert_state_dict +from lerobot.policies.g05.modeling_g05 import G05Policy +from lerobot.processor import PolicyProcessorPipeline +from lerobot.utils.constants import ACTION, OBS_STATE, POLICY_PREPROCESSOR_DEFAULT_NAME + + +class TinyG05Backend(nn.Module): + def __init__(self): + super().__init__() + self.proj = nn.Linear(20, 20) + self.last_samples = None + + def predict_action(self, batch): + self.last_samples = batch["samples"] + state = batch[OBS_STATE] + if state.ndim == 2: + state = state.unsqueeze(1) + step = self.proj(state[:, -1]) + return { + ACTION: step.unsqueeze(1).expand(-1, 4, -1), + "ar_action": (step + 1).unsqueeze(1).expand(-1, 4, -1), + "cot_text": ["Subtask: move carefully"] * step.shape[0], + } + + def forward(self, batch): + prediction = self.proj(batch[OBS_STATE][:, -1]) + target = batch[ACTION][:, 0] + loss = torch.nn.functional.mse_loss(prediction, target) + return loss, {"fm_loss": loss.detach()} + + +def _features(): + return { + OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(7,)), + "observation.images.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 8, 8)), + "observation.images.wrist_image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 8, 8)), + } + + +def _config(**kwargs): + normalization_mode = kwargs.pop("normalization_mode", "identity") + return G05Config( + checkpoint_profile="custom", + normalization_mode=normalization_mode, + input_features=_features(), + output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(7,))}, + chunk_size=4, + device="cpu", + **kwargs, + ) + + +def _policy_batch(task: str = " Pick café cup\nverbatim "): + return { + OBS_STATE: torch.zeros(1, 1, 20), + ACTION: torch.zeros(1, 4, 20), + "observation.images.image": torch.zeros(1, 3, 8, 8), + "observation.images.wrist_image": torch.zeros(1, 3, 8, 8), + "task": [task], + "proprio_dim_is_pad": torch.zeros(20, dtype=torch.bool), + } + + +def test_factory_wiring_is_lazy(): + assert make_policy_config("g05", checkpoint_profile="custom").type == "g05" + assert get_policy_class("g05") is G05Policy + + +def test_libero_and_atomic4_are_distinct_validated_mappings(): + with pytest.raises(ValueError, match="27D"): + G05Config( + checkpoint_profile="custom", + embodiment="atomic_4", + raw_state_dim=16, + raw_action_dim=12, + camera_order=( + "observation.images.robot0_agentview_left", + "observation.images.robot0_eye_in_hand", + "observation.images.robot0_agentview_right", + ), + ) + + cfg = G05Config( + checkpoint_profile="custom", + embodiment="atomic_4", + raw_state_dim=16, + raw_action_dim=12, + policy_state_dim=27, + policy_action_dim=27, + camera_order=( + "observation.images.robot0_agentview_left", + "observation.images.robot0_eye_in_hand", + "observation.images.robot0_agentview_right", + ), + ) + assert cfg.embodiment == "atomic_4" + + +def test_libero_projection_mask_and_inverse_roundtrip(): + config = _config() + preprocessor, postprocessor = make_pre_post_processors(config) + raw_action = torch.arange(7, dtype=torch.float32).repeat(4, 1) + batch = { + OBS_STATE: torch.arange(7, dtype=torch.float32), + ACTION: raw_action, + "observation.images.image": torch.zeros(3, 8, 8), + "observation.images.wrist_image": torch.zeros(3, 8, 8), + "task": "test", + } + + processed = preprocessor(batch) + assert processed[OBS_STATE].shape == (1, 20) + assert processed[ACTION].shape == (4, 20) + assert processed["action_dim_is_pad"].sum() == 13 + assert torch.equal(processed[ACTION][:, [10, 11, 12, 13, 14, 15, 19]], raw_action) + restored = postprocessor(processed[ACTION]) + assert torch.equal(restored, raw_action) + + +def test_atomic4_projection_has_mobile_base_control_mode_and_exact_inverse(): + config = G05Config( + checkpoint_profile="custom", + embodiment="atomic_4", + raw_state_dim=16, + raw_action_dim=12, + policy_state_dim=27, + policy_action_dim=27, + normalization_mode="identity", + camera_order=( + "observation.images.robot0_agentview_left", + "observation.images.robot0_eye_in_hand", + "observation.images.robot0_agentview_right", + ), + input_features={ + OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(16,)), + "observation.images.robot0_agentview_left": PolicyFeature( + type=FeatureType.VISUAL, shape=(3, 8, 8) + ), + "observation.images.robot0_eye_in_hand": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 8, 8)), + "observation.images.robot0_agentview_right": PolicyFeature( + type=FeatureType.VISUAL, shape=(3, 8, 8) + ), + }, + output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(12,))}, + device="cpu", + ) + preprocessor, postprocessor = make_pre_post_processors(config) + raw_action = torch.arange(12, dtype=torch.float32).repeat(3, 1) + batch = { + OBS_STATE: torch.arange(16, dtype=torch.float32), + ACTION: raw_action, + **{camera: torch.zeros(3, 8, 8) for camera in config.camera_order}, + "task": "atomic", + } + + processed = preprocessor(batch) + indices = G05_EMBODIMENT_MAPPINGS["atomic_4"]["action"] + assert torch.equal(processed[ACTION][..., list(indices)], raw_action) + assert torch.equal(postprocessor(processed[ACTION]), raw_action) + # Last five raw dimensions are base motion[4] and control mode. + assert indices[-5:] == (20, 21, 22, 23, 24) + + +def test_quantile_mode_refuses_minmax_substitution(): + config = _config(normalization_mode="q01_q99") + stats = { + OBS_STATE: {"min": torch.zeros(7), "max": torch.ones(7)}, + ACTION: {"min": torch.zeros(7), "max": torch.ones(7)}, + } + with pytest.raises(ValueError, match="real q01/q99"): + make_pre_post_processors(config, dataset_stats=stats) + + +def test_stepwise_quantiles_constant_dimension_are_finite_and_serializable(tmp_path: Path): + config = _config(normalization_mode="q01_q99") + q01_action = torch.zeros(4, 7) + q99_action = torch.ones(4, 7) + q99_action[:, 2] = 0 + stats = { + OBS_STATE: {"q01": torch.zeros(7), "q99": torch.ones(7)}, + ACTION: {"q01": q01_action, "q99": q99_action}, + } + preprocessor, postprocessor = make_pre_post_processors(config, dataset_stats=stats) + processed = preprocessor( + { + OBS_STATE: torch.zeros(7), + ACTION: torch.zeros(4, 7), + "observation.images.image": torch.zeros(3, 8, 8), + "observation.images.wrist_image": torch.zeros(3, 8, 8), + "task": "constant", + } + ) + assert torch.isfinite(processed[ACTION]).all() + torch.testing.assert_close(postprocessor(processed[ACTION]), torch.zeros(4, 7)) + + preprocessor.save_pretrained(tmp_path) + loaded = PolicyProcessorPipeline.from_pretrained( + tmp_path, config_filename=f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json" + ) + assert [step.__class__.__name__ for step in loaded.steps] == [ + step.__class__.__name__ for step in preprocessor.steps + ] + + +def test_exact_raw_task_reaches_author_command_and_head_selection(): + backend = TinyG05Backend() + policy = G05Policy(_config(), backend=backend) + raw_task = " 把 red cup 放到左边\nexactly as written " + + action, metadata = policy.predict_action_chunk_with_runtime(_policy_batch(), task=raw_task) + + assert backend.last_samples[0]["command"] == raw_task + assert action.shape == (1, 4, 20) + assert metadata["cot_text"] == ["Subtask: move carefully"] + + +def test_batch_two_preserves_each_raw_task_and_every_camera_slot(): + backend = TinyG05Backend() + policy = G05Policy(_config(), backend=backend) + batch = _policy_batch() + batch[OBS_STATE] = batch[OBS_STATE].expand(2, -1, -1) + batch[ACTION] = batch[ACTION].expand(2, -1, -1) + batch["observation.images.image"] = batch["observation.images.image"].expand(2, -1, -1, -1) + batch["observation.images.wrist_image"] = batch["observation.images.wrist_image"].expand(2, -1, -1, -1) + batch["proprio_dim_is_pad"] = torch.zeros(2, 20, dtype=torch.bool) + batch["task"] = [" first\n", "第二个 task"] + + action = policy.predict_action_chunk(batch) + + assert action.shape == (2, 4, 20) + assert [sample["command"] for sample in backend.last_samples] == batch["task"] + assert all(sample["image0"] == (224, 224) for sample in backend.last_samples) + assert all(sample["image1"] == (224, 224) for sample in backend.last_samples) + + +def test_forward_backward_update_and_save_reload(tmp_path: Path): + policy = G05Policy(_config(), backend=TinyG05Backend()) + optimizer = torch.optim.AdamW(policy.get_optim_params()["params"], lr=1e-3) + loss, metrics = policy(_policy_batch("train")) + loss.backward() + grad_norm = torch.stack( + [parameter.grad.norm() for parameter in policy.parameters() if parameter.grad is not None] + ).sum() + assert torch.isfinite(loss) + assert grad_norm > 0 and torch.isfinite(grad_norm) + optimizer.step() + assert metrics is not None and metrics["fm_loss"] >= 0 + + policy.save_pretrained(tmp_path) + reloaded = G05Policy.from_pretrained( + tmp_path, backend=TinyG05Backend(), local_files_only=True, strict=True + ) + expected = policy.predict_action_chunk(_policy_batch("save")) + actual = reloaded.predict_action_chunk(_policy_batch("save")) + torch.testing.assert_close(actual, expected) + + +def test_save_pretrained_copies_required_gated_sidecars_portably(tmp_path: Path): + source = tmp_path / "converted" + processor = source / "hf_processor" + processor.mkdir(parents=True) + (processor / "tokenizer.json").write_text("{}") + tokenizer = source / "action_tokenizer.pt" + torch.save({"codec": "ActionCodec"}, tokenizer) + for name in ("LICENSE-G0.5", "NOTICE", "conversion_report.json"): + (source / name).write_text("{}") + config = _config( + author_model_config={ + "hf_processor_path": str(processor), + "AT_CONFIG": {"ckpt_dir": str(tokenizer)}, + } + ) + output = tmp_path / "saved" + + G05Policy(config, backend=TinyG05Backend()).save_pretrained(output) + + assert (output / "hf_processor" / "tokenizer.json").is_file() + assert (output / "action_tokenizer.pt").is_file() + assert (output / "LICENSE-G0.5").is_file() + loaded_config = PreTrainedConfig.from_pretrained(output) + assert isinstance(loaded_config, G05Config) + assert loaded_config.author_model_config["hf_processor_path"] == "hf_processor" + assert loaded_config.author_model_config["AT_CONFIG"]["ckpt_dir"] == "action_tokenizer.pt" + + +def test_tiny_fixed_batch_overfit_reduces_loss(): + policy = G05Policy(_config(), backend=TinyG05Backend()) + optimizer = torch.optim.AdamW(policy.get_optim_params()["params"], lr=5e-2) + batch = _policy_batch("overfit") + initial = policy(batch)[0].item() + for _ in range(20): + optimizer.zero_grad() + loss, _ = policy(batch) + loss.backward() + optimizer.step() + final = policy(batch)[0].item() + assert final < initial * 0.25 + + +def test_conversion_reports_mapping_duplicates_shapes_and_required_prefixes(): + source = { + "model.embed_tokens.weight": torch.zeros(2, 3), + "model.vision_tower.block.weight": torch.ones(2, 2), + "model.action_expert.block.weight": torch.ones(2, 2), + } + converted, report = convert_state_dict(source) + assert "backend.model.vlm.input_proj.weight" in converted + assert report.missing == [] + + expected = { + "backend.model.vlm.input_proj.weight": torch.zeros(3, 3), + "backend.model.vision_tower.block.weight": torch.zeros(2, 2), + "backend.model.action_expert.block.weight": torch.ones(2, 2), + } + _, strict_report = convert_state_dict(source, expected) + assert strict_report.shape_mismatched["backend.model.vlm.input_proj.weight"]["source"] == [2, 3] + + +def test_libero_conversion_packages_model_processors_and_provenance(tmp_path: Path): + source = tmp_path / "author" + output = tmp_path / "lerobot" + (source / ".hydra").mkdir(parents=True) + (source / "hf_processor").mkdir() + (source / "hf_processor" / "tokenizer.json").write_text("{}") + (source / ".hydra" / "config.yaml").write_text( + """ +model: + model_arch: + num_input_images: 2 + horizon_steps: 32 + predict_cot: false + discrete_action: false + continuous_action: true + processor: + use_stepwise_action_norm: true + shape_meta: + state: + - {key: right_ee_pose, shape: 6} + - {key: right_gripper, shape: 1} + action: + - {key: right_ee_pose, shape: 6} + - {key: right_gripper, shape: 1} + images: + - {key: image, lerobot_key: observation.images.image, shape: [3, 224, 224]} + - {key: wrist_image, lerobot_key: observation.images.wrist_image, shape: [3, 224, 224]} +data: + action_size: 32 +tokenizer: + vq_config: {block_wise_autoregressive: false} +""" + ) + action_stats = { + "right_ee_pose": { + "stepwise_q01": torch.zeros(32, 6).tolist(), + "stepwise_q99": torch.ones(32, 6).tolist(), + }, + "right_gripper": { + "stepwise_q01": torch.zeros(32, 1).tolist(), + "stepwise_q99": torch.ones(32, 1).tolist(), + }, + } + state_stats = { + "right_ee_pose": {"global_q01": [0.0] * 6, "global_q99": [1.0] * 6}, + "right_gripper": {"global_q01": [0.0], "global_q99": [1.0]}, + } + (source / "dataset_stats.json").write_text( + json.dumps({"libero": {"state": state_stats, "action": action_stats}}) + ) + torch.save( + { + "model.vlm.block.weight": torch.zeros(2, 2), + "model.vision_tower.block.weight": torch.zeros(2, 2), + "model.action_expert.block.weight": torch.zeros(2, 2), + }, + source / "model.pt", + ) + torch.save({"tokenizer_meta": {"codec": "ActionCodec"}}, source / "action_tokenizer.pt") + license_file = source / "LICENSE-G0.5" + license_file.write_text("test fixture license") + + report = convert_checkpoint(source, output, "g05-libero", license_file=license_file) + + assert len(report.mapped) == 3 + assert (output / "model.safetensors").is_file() + assert (output / "policy_preprocessor.json").is_file() + assert (output / "policy_postprocessor.json").is_file() + assert (output / "conversion_report.json").is_file() + config = PreTrainedConfig.from_pretrained(output) + assert isinstance(config, G05Config) + assert config.source_checkpoint_revision + assert config.prompt_template.startswith("") + + +@pytest.mark.skipif( + not os.environ.get("LEROBOT_G05_CHECKPOINT"), + reason="requires an accepted gated OpenGalaxea/G05 checkpoint and author CUDA environment", +) +def test_gated_checkpoint_loads_strictly(): + checkpoint = Path(os.environ["LEROBOT_G05_CHECKPOINT"]) + policy = G05Policy.from_pretrained(checkpoint, local_files_only=True, strict=True) + assert policy.config.source_checkpoint_revision diff --git a/tests/runtime/test_g05_adapter.py b/tests/runtime/test_g05_adapter.py index 2f6c8078b..2125e7a62 100644 --- a/tests/runtime/test_g05_adapter.py +++ b/tests/runtime/test_g05_adapter.py @@ -82,7 +82,7 @@ def test_system2_surfaces_same_pass_cot_and_action(): state.language_context["cot_text"] == "BBox: cup [1,2,3,4]|\nSubtask: grasp the cup|Updated Memory: cup located" ) - assert state.language_context["subtask"] == "grasp the cup" + assert state.extra["g05_subtask"] == "grasp the cup" assert state.language_context["memory"] == "cup located" @@ -100,7 +100,7 @@ def test_system2_accepts_batch_safe_tuple_metadata(): chunk = G05PolicyAdapter(ReasoningPolicy()).select_action({}, state) assert chunk == "chunk" - assert state.language_context["subtask"] == "move left" + assert state.extra["g05_subtask"] == "move left" assert state.language_context["plan"] == "first move left" @@ -126,7 +126,7 @@ def test_system2_reasoning_does_not_invalidate_same_pass_action_chunk(): assert executed == ["a0"] assert list(runtime.state.action_queue) == ["a1"] - assert runtime.state.language_context["subtask"] == "pick cup" + assert runtime.state.extra["g05_subtask"] == "pick cup" def test_system2_rejects_checkpoint_without_predict_cot(): diff --git a/uv.lock b/uv.lock index a7055011e..59627efce 100644 --- a/uv.lock +++ b/uv.lock @@ -2879,6 +2879,7 @@ all = [ { name = "motorbridge-smart-servo" }, { name = "mypy" }, { name = "num2words" }, + { name = "omegaconf" }, { name = "pandas" }, { name = "peft" }, { name = "placo" }, @@ -3020,6 +3021,9 @@ feetech = [ { name = "feetech-servo-sdk" }, { name = "pyserial" }, ] +g05 = [ + { name = "omegaconf" }, +] gamepad = [ { name = "hidapi" }, { name = "pygame" }, @@ -3349,6 +3353,7 @@ requires-dist = [ { name = "lerobot", extras = ["feetech"], marker = "extra == 'all'" }, { name = "lerobot", extras = ["feetech"], marker = "extra == 'hopejr'" }, { name = "lerobot", extras = ["feetech"], marker = "extra == 'lekiwi'" }, + { name = "lerobot", extras = ["g05"], marker = "extra == 'all'" }, { name = "lerobot", extras = ["gamepad"], marker = "extra == 'all'" }, { name = "lerobot", extras = ["groot"], marker = "extra == 'all'" }, { name = "lerobot", extras = ["grpcio-dep"], marker = "extra == 'async'" }, @@ -3451,6 +3456,7 @@ 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" },