refactor(g05): keep checkpoint tooling out of runtime PR

This commit is contained in:
Pepijn
2026-07-28 16:55:46 +02:00
parent 85a24cebc6
commit 99e16174e3
8 changed files with 76 additions and 1524 deletions
-328
View File
@@ -1,328 +0,0 @@
#!/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.
"""Compare a converted G0.5 checkpoint with the pinned author implementation.
This is intentionally an opt-in checkpoint test: it requires an accepted gated
checkpoint, the pinned GalaxeaVLA source checkout, and its CUDA dependencies.
"""
from __future__ import annotations
import argparse
import copy
import json
import sys
from pathlib import Path
from typing import Any
import torch
from lerobot.configs.policies import PreTrainedConfig
from lerobot.policies.factory import make_pre_post_processors
from lerobot.policies.g05.configuration_g05 import G05Config
from lerobot.policies.g05.modeling_g05 import G05Policy
from lerobot.utils.constants import ACTION, OBS_STATE
def _tensor_error(reference: torch.Tensor, actual: torch.Tensor) -> dict[str, float]:
reference = reference.detach().float().cpu()
actual = actual.detach().float().cpu()
if reference.shape != actual.shape:
raise AssertionError(f"shape mismatch: {tuple(reference.shape)} != {tuple(actual.shape)}")
error = (reference - actual).abs().flatten()
return {
"max": error.max().item() if error.numel() else 0.0,
"p99": torch.quantile(error, 0.99).item() if error.numel() else 0.0,
"mean": error.mean().item() if error.numel() else 0.0,
}
def _assert_bool_equal(name: str, reference: torch.Tensor, actual: torch.Tensor) -> None:
if not torch.equal(reference.detach().cpu().bool(), actual.detach().cpu().bool()):
raise AssertionError(f"{name} differs")
def _move_to(value: Any, device: torch.device) -> Any:
if isinstance(value, torch.Tensor):
return value.to(device)
if isinstance(value, dict):
return {key: _move_to(item, device) for key, item in value.items()}
if isinstance(value, list):
return [_move_to(item, device) for item in value]
if isinstance(value, tuple):
return tuple(_move_to(item, device) for item in value)
return value
def _raw_sample(processor: Any, index: int, task: str) -> dict[str, Any]:
images: dict[str, torch.Tensor] = {}
for camera_index, meta in enumerate(processor.shape_meta["images"]):
channels, height, width = meta["raw_shape"]
values = torch.arange(channels * height * width, dtype=torch.int64)
images[meta["key"]] = (
(values.reshape(channels, height, width) + index * 37 + camera_index * 71)
.remainder(256)
.to(torch.uint8)
)
state = {}
for part_index, meta in enumerate(processor.shape_meta["state"]):
width = int(meta["raw_shape"])
state[meta["key"]] = torch.linspace(
-0.2 + 0.03 * index + 0.01 * part_index,
0.2 + 0.03 * index + 0.01 * part_index,
width,
)
action = {}
horizon = int(processor.action_horizon)
for part_index, meta in enumerate(processor.shape_meta["action"]):
width = int(meta["raw_shape"])
action[meta["key"]] = torch.linspace(
-0.1 + 0.02 * index + 0.01 * part_index,
0.1 + 0.02 * index + 0.01 * part_index,
horizon * width,
).reshape(horizon, width)
return {
"images": {
key: value.unsqueeze(0).expand(processor.num_obs_steps, -1, -1, -1)
for key, value in images.items()
},
"state": {
key: value.unsqueeze(0).expand(processor.num_obs_steps, -1) for key, value in state.items()
},
"action": action,
"action_is_pad": torch.zeros(horizon, dtype=torch.bool),
"state_is_pad": torch.zeros(processor.num_obs_steps, dtype=torch.bool),
"image_is_pad": torch.zeros(processor.num_obs_steps, dtype=torch.bool),
"task": task,
"frequency": 15.0,
"idx": index,
}
def _lerobot_input(raw: dict[str, Any], processor: Any, config: G05Config) -> dict[str, Any]:
image_pairs = zip(processor.shape_meta["images"], config.camera_order, strict=True)
return {
OBS_STATE: torch.cat(
[raw["state"][meta["key"]][-1] for meta in processor.shape_meta["state"]], dim=-1
),
ACTION: torch.cat([raw["action"][meta["key"]] for meta in processor.shape_meta["action"]], dim=-1),
**{lerobot_key: raw["images"][meta["key"]][-1] for meta, lerobot_key in image_pairs},
"action_is_pad": raw["action_is_pad"],
"task": raw["task"],
}
def _collate_lerobot(samples: list[dict[str, Any]], config: G05Config) -> dict[str, Any]:
result: dict[str, Any] = {
OBS_STATE: torch.cat([sample[OBS_STATE] for sample in samples], dim=0),
ACTION: torch.stack([sample[ACTION] for sample in samples], dim=0),
"task": [
sample["task"][0] if isinstance(sample["task"], list) else sample["task"] for sample in samples
],
}
for key in (*config.camera_order, "proprio_dim_is_pad", "action_dim_is_pad"):
result[key] = torch.cat([sample[key] for sample in samples], dim=0)
action_pad = [sample["action_is_pad"] for sample in samples]
result["action_is_pad"] = torch.stack(
[value.squeeze(0) if value.ndim == 2 else value for value in action_pad], dim=0
)
return result
def _flatten_author_action(action: dict[str, torch.Tensor], processor: Any) -> torch.Tensor:
return torch.cat([action[meta["key"]] for meta in processor.shape_meta["action"]], dim=-1)
def run(args: argparse.Namespace) -> dict[str, Any]:
sys.path.insert(0, str(args.author_source / "src"))
from g05.utils.data.data_utils import collate_fn_pad_sequences
from g05.utils.data.normalizer import load_dataset_stats_from_json
from g05.utils.data.processor_utils import build_processors
from omegaconf import OmegaConf
device = torch.device(args.device)
checkpoint = args.checkpoint.resolve()
if not OmegaConf.has_resolver("oc.load"):
def _oc_load(path: str, key: str | None = None) -> Any:
loaded = OmegaConf.load(args.author_source / path)
return OmegaConf.select(loaded, key) if key is not None else loaded
OmegaConf.register_new_resolver(
"oc.load",
_oc_load,
)
author_cfg = OmegaConf.load(checkpoint / "author_config.yaml")
author_processors = build_processors(author_cfg)
author_processors.set_normalizer_from_stats(
load_dataset_stats_from_json(checkpoint / "g05_dataset_stats.json")
)
author_processors.eval()
author_processor = author_processors.processors[args.embodiment]
author_processor.action_horizon = int(author_cfg.data.action_size)
config = PreTrainedConfig.from_pretrained(checkpoint)
if not isinstance(config, G05Config):
raise TypeError(f"Expected G05Config, got {type(config).__name__}")
preprocessor, postprocessor = make_pre_post_processors(config, pretrained_path=checkpoint)
tasks = [
" Pick café cup\nverbatim ",
"第二个 task — keep Unicode and whitespace\t",
][: args.batch_size]
raw_samples = [_raw_sample(author_processor, index, task) for index, task in enumerate(tasks)]
author_samples = [author_processor.preprocess(copy.deepcopy(sample)) for sample in raw_samples]
author_batch = collate_fn_pad_sequences(copy.deepcopy(author_samples))
lerobot_samples = [
preprocessor(_lerobot_input(sample, author_processor, config)) for sample in raw_samples
]
lerobot_batch = _collate_lerobot(lerobot_samples, config)
policy = G05Policy.from_pretrained(
checkpoint,
local_files_only=True,
strict=True,
).to(device)
policy.eval()
port_author_batch = policy._prepare_author_batch(lerobot_batch)
report: dict[str, Any] = {
"batch_size": args.batch_size,
"device": str(device),
"dtype": str(next(policy.parameters()).dtype),
"prompt_exact": all(
left["template"] == right["template"]
and left["command"] == right["command"]
and left["embodiment"] == right["embodiment"]
for left, right in zip(author_batch["samples"], port_author_batch["samples"], strict=True)
),
}
if not report["prompt_exact"]:
raise AssertionError("author and LeRobot prompt payloads differ")
image_errors = {}
for (author_key, author_images), (port_key, port_images) in zip(
author_batch["pixel_values"].items(),
port_author_batch["pixel_values"].items(),
strict=True,
):
image_errors[f"{author_key}->{port_key}"] = _tensor_error(author_images, port_images)
report["images"] = image_errors
author_proprio = torch.stack([sample["proprio"]["value"] for sample in author_batch["samples"]], dim=0)
port_proprio = torch.stack([sample["proprio"]["value"] for sample in port_author_batch["samples"]], dim=0)
report["proprio"] = _tensor_error(author_proprio, port_proprio)
report["normalized_input_action"] = _tensor_error(author_batch[ACTION], lerobot_batch[ACTION])
for index, (author_sample, port_sample) in enumerate(
zip(author_batch["samples"], port_author_batch["samples"], strict=True)
):
_assert_bool_equal(
f"proprio_dim_is_pad[{index}]",
author_sample["proprio"]["proprio_dim_is_pad"],
port_sample["proprio"]["proprio_dim_is_pad"],
)
_assert_bool_equal(
"action_dim_is_pad", author_batch["action_dim_is_pad"], port_author_batch["action_dim_is_pad"]
)
report["masks_exact"] = True
author_ids, author_attention = policy.backend.processor.encode_inference(
copy.deepcopy(author_batch["samples"]), device=device, mode="fm"
)
port_ids, port_attention = policy.backend.processor.encode_inference(
copy.deepcopy(port_author_batch["samples"]), device=device, mode="fm"
)
_assert_bool_equal("input_ids", author_ids, port_ids)
_assert_bool_equal("attention_mask", author_attention, port_attention)
report["tokens_exact"] = True
report["token_shape"] = list(author_ids.shape)
author_cuda = _move_to(copy.deepcopy(author_batch), device)
port_cuda = _move_to(lerobot_batch, device)
torch.manual_seed(args.seed)
with (
torch.inference_mode(),
torch.autocast(
device_type=device.type,
dtype=torch.bfloat16,
enabled=config.model_weights_to_bf16 and device.type == "cuda",
),
):
author_output = policy.backend.predict_action(author_cuda)
torch.manual_seed(args.seed)
with torch.inference_mode():
port_action = policy.predict_action_chunk(port_cuda)
report["normalized_action"] = _tensor_error(author_output[ACTION], port_action)
author_post = author_processor.postprocess(
{
ACTION: author_output[ACTION].detach().cpu(),
"proprio": author_cuda["proprio"].detach().cpu(),
"action_dim_is_pad": author_cuda.get("action_dim_is_pad"),
"proprio_dim_is_pad": author_cuda.get("proprio_dim_is_pad"),
}
)
author_env_action = _flatten_author_action(author_post[ACTION], author_processor)
# Isolate processor parity from small repeated BF16 sampling drift by feeding
# both postprocessors the same normalized author action chunk.
port_env_action = postprocessor(author_output[ACTION])
report["environment_action"] = _tensor_error(author_env_action, port_env_action)
if args.compare_training_loss:
policy.train()
torch.manual_seed(args.seed)
with torch.no_grad():
author_loss, _ = policy.backend(_move_to(copy.deepcopy(author_batch), device))
torch.manual_seed(args.seed)
with torch.no_grad():
port_loss, _ = policy(_move_to(lerobot_batch, device))
report["training_loss"] = {
"author": author_loss.detach().float().item(),
"lerobot": port_loss.detach().float().item(),
"absolute_error": abs(author_loss.detach().float().item() - port_loss.detach().float().item()),
}
numeric_sections = (
*report["images"].values(),
report["proprio"],
report["normalized_input_action"],
report["normalized_action"],
report["environment_action"],
)
report["tolerance"] = args.atol
report["passed"] = all(section["max"] <= args.atol for section in numeric_sections)
if "training_loss" in report:
report["passed"] &= report["training_loss"]["absolute_error"] <= args.atol
return report
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--author-source", type=Path, required=True)
parser.add_argument("--checkpoint", type=Path, required=True)
parser.add_argument("--embodiment", default="libero")
parser.add_argument("--batch-size", type=int, choices=(1, 2), default=1)
parser.add_argument("--device", default="cuda")
parser.add_argument("--seed", type=int, default=1234)
parser.add_argument("--atol", type=float, default=5e-5)
parser.add_argument("--compare-training-loss", action="store_true")
parser.add_argument("--output", type=Path)
args = parser.parse_args()
report = run(args)
payload = json.dumps(report, indent=2, sort_keys=True)
if args.output is not None:
args.output.write_text(f"{payload}\n")
print(payload)
if not report["passed"]:
raise SystemExit(1)
if __name__ == "__main__":
main()
+18 -85
View File
@@ -6,6 +6,12 @@ and the optional native chain-of-thought phase as System 2. They are not separat
models: the runtime obtains both from one inference call and the action stays
conditioned on the same post-reasoning KV state.
Transformers includes the native multimodal Qwen3.5 backbone, vision tower, and
processor. G0.5 is not a stock `Qwen3_5ForConditionalGeneration` checkpoint,
however: it adds the proprioception/action path, action expert, flow-matching
head, ActionCodec, and unified CoT/action decode. The current integration
therefore keeps the pinned G0.5 author package as the model backend.
> [!WARNING]
> G0.5 code and checkpoints use the
> [G0.5 Community License](https://huggingface.co/OpenGalaxea/G05/blob/main/licenses/LICENSE-G0.5),
@@ -22,9 +28,9 @@ conditioned on the same post-reasoning KV state.
| `g05-robotwin20` | Continuous flow | two arms 6+gripper → 20D grouped layout | high, left wrist, right wrist | 32 | 8 | stepwise q01/q99 |
| `g05-so101` | Flow or AR ActionCodec + native CoT | right arm joints 6 → 20D grouped layout | exterior, optional left + right wrist | 32 | 16 | stepwise q01/q99 |
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
Each packaged checkpoint stores the resolved model and processor configuration,
ActionCodec metadata, statistics, exact prompt template, source revision, and
license. Loading rejects a different head, horizon, processor mode, or
normalization contract.
The named `atomic_4` adapter is intentionally separate from LIBERO. Its raw state
@@ -52,67 +58,11 @@ git -C GalaxeaVLA checkout b34966f387dd2ae0f003143b81494afd9213e613
export PYTHONPATH="/path/to/GalaxeaVLA/src:${PYTHONPATH}"
```
## 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/`. For released bundles, the converter
also discovers the tokenizer and `qwen3_5_2b_base_processor/` in the parent
directory.
```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
```
For the base checkpoint, select both its concrete embodiment and one of its
enabled output heads:
```bash
uv run python -m lerobot.policies.g05.convert_g05_checkpoint \
--source-dir /path/to/checkpoints/g05-base \
--output-dir outputs/g05-base-r1lite-flow-lerobot \
--profile g05-base \
--embodiment galaxea_r1lite \
--action-head flow \
--license-file /path/to/GalaxeaVLA/LICENSE-G0.5
```
Use `--action-head actioncodec` for the autoregressive action path,
`--embodiment galaxea_r1pro` for R1 Pro, or `--profile g05-robotwin20` for
RoboTwin. The base model's 32-step head contains five history-alignment steps;
postprocessing returns the 27 executable actions. `conversion_report.json`
records every mapped, missing, unexpected, duplicate, and shape-mismatched
tensor; conversion fails when strict required-state validation fails.
The official RoboTwin config contains a pinned data include. Resolve and package
it from the audited author checkout:
```bash
uv run python -m lerobot.policies.g05.convert_g05_checkpoint \
--source-dir /path/to/G05/g05-robotwin20 \
--output-dir outputs/g05-robotwin20-lerobot \
--profile g05-robotwin20 \
--author-source /path/to/GalaxeaVLA \
--license-file /path/to/GalaxeaVLA/LICENSE-G0.5
```
The LeRobot organization hosts the prepared LIBERO, RoboTwin, and SO-101
checkpoints privately. Authenticate with `hf auth login` before loading them.
SO-100 and SO-101 share the released `so100` embodiment token and six-joint
right-arm contract. A missing left-wrist camera is zero-padded exactly as in the
author deployment client:
```bash
uv run python -m lerobot.policies.g05.convert_g05_checkpoint \
--source-dir /path/to/G05/g05-so101 \
--output-dir outputs/g05-so101-lerobot \
--profile g05-so101 \
--action-head flow \
--author-source /path/to/GalaxeaVLA \
--license-file /path/to/GalaxeaVLA/LICENSE-G0.5
```
author deployment client.
## Interactive System 1 and System 2 runtime
@@ -120,13 +70,13 @@ System 1 executes the selected ActionCodec or flow chunk directly:
```bash
lerobot-rollout \
--policy.path=outputs/g05-base-lerobot \
--policy.path=lerobot/g05_so101 \
--language --direct_subtask \
--task="pick up the cup" \
--mode=action
```
System 2 is available only when the converted checkpoint metadata has
System 2 is available only when the packaged 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
@@ -134,7 +84,7 @@ replacement task.
```bash
lerobot-rollout \
--policy.path=outputs/g05-system2-lerobot \
--policy.path=lerobot/g05_so101 \
--language \
--task="clear the table" \
--mode=action
@@ -165,31 +115,14 @@ the checkpoint's native System 2 CoT telemetry.
## 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.
The opt-in author-oracle gate compares the exact prompt tokens, every image/state/action
tensor and mask, seeded flow loss and samples, and de-normalized environment actions
for batch sizes one and two:
LIBERO and `atomic_4` mappings, padding masks, inverse action projection, a finite
forward/backward/update, and save/reload parity:
```bash
uv run pytest tests/policies/g05 tests/runtime/test_g05_adapter.py -q
uv run ruff check src/lerobot/policies/g05 tests/policies/g05
PYTHONPATH=/path/to/GalaxeaVLA/src uv run python benchmarks/g05_checkpoint_parity.py \
--author-source /path/to/GalaxeaVLA \
--checkpoint outputs/g05-libero-lerobot \
--batch-size 1 \
--compare-training-loss \
--output outputs/g05-libero-parity-b1.json
PYTHONPATH=/path/to/GalaxeaVLA/src uv run python benchmarks/g05_checkpoint_parity.py \
--author-source /path/to/GalaxeaVLA \
--checkpoint outputs/g05-libero-lerobot \
--batch-size 2 \
--output outputs/g05-libero-parity-b2.json
```
This command still requires licensed checkpoint access and suitable CUDA hardware.
A 50-episode LIBERO/RoboTwin success-rate comparison additionally requires the
matching simulator, task assets, reset seeds, and author evaluator; no task-level
benchmark number is claimed until that separate gate runs.
@@ -199,7 +132,7 @@ LeRobot rollout with the author camera names and relative control:
```bash
lerobot-eval \
--policy.path=outputs/g05-libero-lerobot \
--policy.path=lerobot/g05_libero \
--policy.device=cuda \
--env.type=libero \
--env.task=libero_goal \
@@ -155,9 +155,9 @@ _PROFILE_DEFAULTS = {
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.
``author_model_config`` comes from the packaged checkpoint's resolved Hydra
config. It is intentionally checkpoint state rather than a collection of
guessed LeRobot defaults.
"""
checkpoint_profile: str = "g05-base"
@@ -250,7 +250,7 @@ class G05Config(PreTrainedConfig):
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.")
raise ValueError("G0.5 System 2 requires predict_cot=True in the packaged checkpoint.")
if not 1 <= self.n_action_steps <= self.chunk_size:
raise ValueError("n_action_steps must be between 1 and chunk_size.")
if self.action_head == "actioncodec" and not self.discrete_action:
@@ -1,691 +0,0 @@
#!/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)
shared_aliases: 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 save_converted_state_dict(state_dict: dict[str, torch.Tensor], path: Path) -> dict[str, str]:
"""Save exact tensor aliases once, matching safetensors' strict model loader."""
aliases: dict[str, str] = {}
unique: dict[str, torch.Tensor] = {}
seen: dict[tuple[int, int, tuple[int, ...], tuple[int, ...]], str] = {}
for key in sorted(state_dict):
tensor = state_dict[key]
identity = (
tensor.untyped_storage().data_ptr(),
tensor.storage_offset(),
tuple(tensor.shape),
tuple(tensor.stride()),
)
if identity in seen:
aliases[key] = seen[identity]
else:
seen[identity] = key
unique[key] = tensor
save_file(unique, path, metadata=aliases or None)
return aliases
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 _resolve_model_arch(hydra: dict[str, Any], data_override: dict[str, Any] | None) -> dict[str, Any]:
"""Resolve the official Hydra snapshot without importing the author package."""
from omegaconf import OmegaConf
def multiply(expression: str) -> int:
factors = [part.strip() for part in str(expression).split("*")]
if not factors or any(not part.isdigit() for part in factors):
raise ValueError(f"Unsupported G0.5 eval expression {expression!r}.")
result = 1
for factor in factors:
result *= int(factor)
return result
if not OmegaConf.has_resolver("obs_image_steps"):
OmegaConf.register_new_resolver("obs_image_steps", lambda value: int(value))
if not OmegaConf.has_resolver("eval"):
OmegaConf.register_new_resolver("eval", multiply)
payload = dict(hydra)
if data_override is not None:
payload["data"] = data_override
resolved = OmegaConf.to_container(OmegaConf.create(payload)["model"]["model_arch"], resolve=True)
if not isinstance(resolved, dict):
raise TypeError("Expected model.model_arch to resolve to a mapping.")
return resolved
def _profile_config(
profile: str,
hydra: dict[str, Any],
embodiment: str | None = None,
action_head: str | None = None,
data_override: dict[str, Any] | None = None,
) -> G05Config:
model = hydra.get("model", {})
data = data_override if data_override is not None else hydra.get("data", {})
if not isinstance(data, dict):
raise ValueError(
"The checkpoint Hydra config contains an unresolved data include. "
"Pass --author-source so the converter can package the pinned data contract."
)
arch = _resolve_model_arch(hydra, data_override)
action_codec_metadata = model.get("tokenizer")
if not isinstance(action_codec_metadata, dict):
action_codec_metadata = hydra.get("tokenizer")
if not isinstance(action_codec_metadata, dict):
action_codec_metadata = {}
fixed_embodiment = {
"g05-libero": "libero",
"g05-robotwin20": "robotwin",
"g05-so101": "so100",
}.get(profile)
if fixed_embodiment is not None:
if embodiment is not None and embodiment != fixed_embodiment:
raise ValueError(f"{profile} requires embodiment={fixed_embodiment!r}.")
embodiment = fixed_embodiment
elif profile == "g05-base":
available = sorted((data.get("processors") or {}).keys())
if embodiment is None:
raise ValueError(f"g05-base requires a concrete --embodiment; choose one of {available}.")
if embodiment not in available:
raise ValueError(f"g05-base embodiment must be one of {available}, got {embodiment!r}.")
else:
raise ValueError(f"Unsupported G0.5 profile {profile!r}.")
assert embodiment is not None
data_processor = (data.get("processors") or {}).get(embodiment, {})
# Match the author's build_processors merge order: embodiment data first,
# then task-level model.processor overrides.
processor_metadata = {**data_processor, **model.get("processor", {})}
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))
if action_head is None:
action_head = "flow" if continuous else "actioncodec"
if action_head not in {"flow", "actioncodec"}:
raise ValueError("action_head must be 'flow' or 'actioncodec'.")
if action_head == "flow" and not continuous:
raise ValueError("The selected checkpoint does not enable the continuous flow action head.")
if action_head == "actioncodec" and not discrete:
raise ValueError("The selected checkpoint does not enable the autoregressive ActionCodec head.")
if profile not in {"g05-base", "g05-so101"} and action_head != "flow":
raise ValueError(f"The released {profile} checkpoint supports only --action-head flow.")
norm_name = str(processor_metadata.get("norm_default_mode", "")).lower()
checkpoint_normalization = {
"q01/q99": "q01_q99",
"z-score": "z_score",
"z-score-tail": "z_score_tail_mixed",
"identity": "identity",
"dummy": "identity",
}.get(norm_name)
arch = dict(arch)
arch.pop("_target_", None)
n_obs_steps = int(processor_metadata.get("num_obs_steps", data.get("obs_size", 1)))
raw_num_input_images = arch.get("num_input_images", len(G05_CAMERA_PROFILES[embodiment]) * n_obs_steps)
num_input_images = (
int(raw_num_input_images)
if isinstance(raw_num_input_images, int | float)
else len(G05_CAMERA_PROFILES[embodiment]) * n_obs_steps
)
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,
n_action_steps=10,
normalization_mode="q01_q99",
normalization_clip=(-5.0, 5.0),
use_stepwise_action_norm=True,
libero_gripper_binarize=True,
camera_order=G05_CAMERA_PROFILES["libero"],
num_input_images=num_input_images,
author_model_config=arch,
processor_metadata=processor_metadata,
action_codec_metadata=action_codec_metadata,
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="robotwin",
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,
n_action_steps=min(8, horizon),
normalization_mode="q01_q99",
normalization_clip=(-5.0, 5.0),
use_stepwise_action_norm=True,
camera_order=G05_CAMERA_PROFILES["robotwin"],
num_input_images=num_input_images,
author_model_config=arch,
processor_metadata=processor_metadata,
action_codec_metadata=action_codec_metadata,
author_source_revision=G05_SOURCE_REVISION,
source_checkpoint_revision=G05_HUB_REVISION,
license="other",
tags=["g05", "robotics", "non-commercial"],
)
if profile == "g05-so101":
shape_meta = processor_metadata.get("shape_meta") or {}
raw_state_dim = sum(int(item["shape"]) for item in shape_meta.get("state", []))
raw_action_dim = sum(int(item["shape"]) for item in shape_meta.get("action", []))
action_feature_names = tuple(
f"{item['key']}.{index}"
for item in shape_meta.get("action", [])
for index in range(int(item["shape"]))
)
return G05Config(
checkpoint_profile=profile,
embodiment="so100",
action_head=action_head,
runtime_system="system2",
predict_cot=predict_cot,
discrete_action=discrete,
continuous_action=continuous,
return_continuous_action=action_head == "flow",
raw_state_dim=raw_state_dim,
raw_action_dim=raw_action_dim,
chunk_size=horizon,
n_action_steps=min(16, horizon),
normalization_mode="q01_q99",
normalization_clip=(-5.0, 5.0),
use_relative_actions=True,
action_feature_names=action_feature_names,
use_stepwise_action_norm=True,
camera_order=G05_CAMERA_PROFILES["so100"],
optional_camera_keys=("observation.images.wrist_left",),
num_input_images=num_input_images,
author_model_config=arch,
processor_metadata=processor_metadata,
action_codec_metadata=action_codec_metadata,
author_source_revision=G05_SOURCE_REVISION,
source_checkpoint_revision=G05_HUB_REVISION,
license="other",
tags=["g05", "robotics", "so101", "non-commercial"],
)
if checkpoint_normalization != "z_score_tail_mixed":
raise ValueError("The pinned g05-base R1 contract requires z-score-tail normalization.")
exceptions = processor_metadata.get("norm_exception_mode") or {}
exception_modes = {mode for category in exceptions.values() for mode in (category or {}).values()}
if exception_modes - {"q01/q99"}:
raise ValueError(f"Unsupported g05-base normalization exceptions: {exceptions}.")
shape_meta = processor_metadata.get("shape_meta") or {}
raw_state_dim = sum(int(item["shape"]) for item in shape_meta.get("state", []))
raw_action_dim = sum(int(item["shape"]) for item in shape_meta.get("action", []))
action_feature_names = tuple(
f"{item['key']}.{index}"
for item in shape_meta.get("action", [])
for index in range(int(item["shape"]))
)
return G05Config(
checkpoint_profile="g05-base",
embodiment=embodiment,
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",
policy_action_dim=int(arch.get("action_dim", 27)),
policy_state_dim=int(arch.get("proprio_dim", 27)),
raw_state_dim=raw_state_dim,
raw_action_dim=raw_action_dim,
chunk_size=horizon,
n_action_steps=min(16, horizon),
normalization_mode=checkpoint_normalization,
normalization_clip=(-5.0, 5.0),
use_relative_actions=True,
relative_exclude_joints=("gripper",),
action_feature_names=action_feature_names,
use_stepwise_action_norm=bool(processor_metadata.get("use_stepwise_action_norm", False)),
camera_order=G05_CAMERA_PROFILES[embodiment],
n_obs_steps=int(processor_metadata.get("num_obs_steps", 1)),
num_input_images=num_input_images,
author_model_config=arch,
processor_metadata=processor_metadata,
action_codec_metadata=action_codec_metadata,
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 []
camera_size_config = processor_metadata.get("camera_size_config") or {}
by_lerobot_key = {
item.get("lerobot_key"): tuple(camera_size_config.get(item.get("camera_type"), item["shape"][-2:]))
for item in images
if item.get("lerobot_key") and len(item.get("shape") or ()) >= 3
}
if by_lerobot_key and all(key in by_lerobot_key for key in camera_order):
return {key: by_lerobot_key[key] for key in camera_order}
if len(images) == len(camera_order):
return {
key: tuple(camera_size_config.get(item.get("camera_type"), item["shape"][-2:]))
for key, item in zip(camera_order, images, strict=True)
}
return {}
def _write_model_card(output_dir: Path, config: G05Config) -> None:
system = "System 2 (native CoT + actions)" if config.runtime_system == "system2" else "System 1"
output_dir.joinpath("README.md").write_text(
f"""---
license: other
license_name: g05-community-license
license_link: https://huggingface.co/OpenGalaxea/G05/blob/main/licenses/LICENSE-G0.5
library_name: lerobot
tags:
- robotics
- lerobot
- g05
- non-commercial
---
# LeRobot conversion of OpenGalaxea G0.5
This is a mechanically converted `{config.checkpoint_profile}` checkpoint for LeRobot.
It uses the `{config.action_head}` action head and exposes {system}.
The weights, configuration, ActionCodec tokenizer, and processor assets are derivative
materials licensed under the **G0.5 Community License Agreement
(Non-Commercial + Limited Patent License)**, not Apache-2.0. Use is limited to the
purposes allowed by `LICENSE-G0.5`.
## Modification notice
LeRobot changed the checkpoint container, tensor key prefixes, portable sidecar paths,
and serialized preprocessing/postprocessing metadata. Model tensor values were not
trained or numerically altered. See `conversion_report.json` for the complete key map
and source revisions.
This repository does not imply endorsement by Galaxea.
"""
)
def _resolve_data_include(hydra: dict[str, Any], author_source: Path | None) -> dict[str, Any] | None:
if isinstance(hydra.get("data"), dict):
return None
if author_source is None:
return None
candidate = author_source / "configs" / "data" / "robotwin.yaml"
if not candidate.is_file():
raise FileNotFoundError(f"Could not resolve checkpoint data include from {candidate}.")
payload = yaml.safe_load(candidate.read_text())
if not isinstance(payload, dict):
raise TypeError(f"Expected a data mapping in {candidate}.")
return payload
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 {}
if config.normalization_mode == "z_score_tail_mixed":
exceptions = config.processor_metadata.get("norm_exception_mode") or {}
default_mode = str(config.processor_metadata.get("norm_default_mode"))
result: dict[str, dict[str, torch.Tensor]] = {}
for source_category, feature_name in (("state", OBS_STATE), ("action", ACTION)):
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."
)
prefix = (
"stepwise" if source_category == "action" and config.use_stepwise_action_norm else "global"
)
collected: dict[str, list[torch.Tensor]] = {
name: [] for name in ("mean", "std", "tail_q01", "tail_q99", "tail_mean", "tail_mask")
}
for component in component_meta:
key = component["key"]
stats = category_stats.get(key)
if not isinstance(stats, dict):
raise ValueError(f"Missing checkpoint statistics for {source_category}.{key}.")
mode = (exceptions.get(source_category) or {}).get(key, default_mode)
q01 = torch.as_tensor(stats[f"{prefix}_q01"], dtype=torch.float32)
q99 = torch.as_tensor(stats[f"{prefix}_q99"], dtype=torch.float32)
mean = torch.as_tensor(stats[f"{prefix}_mean"], dtype=torch.float32)
std = torch.as_tensor(stats[f"{prefix}_std"], dtype=torch.float32)
if mode == "q01/q99":
collected["mean"].append((q01 + q99) / 2)
collected["std"].append((q99 - q01) / 2)
tail_mask = torch.zeros(q01.shape[-1], dtype=torch.bool)
elif mode == "z-score-tail":
collected["mean"].append(mean)
collected["std"].append(std)
tail_mask = torch.ones(q01.shape[-1], dtype=torch.bool)
else:
raise ValueError(
f"Unsupported checkpoint normalization mode {source_category}.{key}={mode!r}."
)
collected["tail_q01"].append(q01)
collected["tail_q99"].append(q99)
collected["tail_mean"].append(mean)
collected["tail_mask"].append(tail_mask)
result[feature_name] = {name: torch.cat(parts, dim=-1) for name, parts in collected.items()}
return result
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,
embodiment: str | None = None,
action_head: str | None = None,
author_source: Path | None = None,
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 = next(
(
path
for path in (
source_dir / "action_tokenizer.pt",
source_dir / "checkpoints" / "action_tokenizer.pt",
source_dir.parent / "action_tokenizer.pt",
)
if path.is_file()
),
source_dir / "action_tokenizer.pt",
)
processor_path = next(
(
path
for path in (
source_dir / "hf_processor",
source_dir / "qwen3_5_2b_base_processor",
source_dir.parent / "qwen3_5_2b_base_processor",
)
if path.is_dir()
),
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())
data_override = _resolve_data_include(hydra, author_source)
config = _profile_config(
profile,
hydra,
embodiment=embodiment,
action_head=action_head,
data_override=data_override,
)
camera_sizes = _camera_sizes(config.processor_metadata, config.camera_order)
if camera_sizes:
config.camera_sizes = camera_sizes
config.validate_features()
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)
report.shared_aliases = save_converted_state_dict(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)
packaged_hydra = dict(hydra)
if data_override is not None:
packaged_hydra["data"] = data_override
(output_dir / "author_config.yaml").write_text(yaml.safe_dump(packaged_hydra, sort_keys=False))
shutil.copy2(license_file, output_dir / "LICENSE-G0.5")
if author_source is not None:
for name in ("LICENSE_QWEN3_5.txt", "THIRD_PARTY_NOTICES.md"):
source = author_source / name
if source.is_file():
shutil.copy2(source, output_dir / name)
(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\n"
"Modification notice: LeRobot converted the checkpoint container, tensor key "
"prefixes, portable sidecar paths, and processor metadata. Tensor values were "
"not trained or numerically altered.\n"
)
_write_model_card(output_dir, config)
(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", "g05-so101"),
required=True,
)
parser.add_argument(
"--embodiment",
choices=("galaxea_r1lite", "galaxea_r1pro"),
help="Required concrete embodiment for g05-base.",
)
parser.add_argument(
"--action-head",
choices=("flow", "actioncodec"),
help="Select one enabled g05-base output head; benchmark checkpoints are flow-only.",
)
parser.add_argument("--license-file", type=Path, required=True)
parser.add_argument(
"--author-source",
type=Path,
help="Pinned GalaxeaVLA checkout used to resolve official config includes and notices.",
)
args = parser.parse_args()
report = convert_checkpoint(
args.source_dir,
args.output_dir,
args.profile,
license_file=args.license_file,
embodiment=args.embodiment,
action_head=args.action_head,
author_source=args.author_source,
)
print(json.dumps(asdict(report), indent=2, sort_keys=True))
if __name__ == "__main__":
main()
+1 -2
View File
@@ -31,7 +31,7 @@ from .configuration_g05 import G05_POLICY_PARTS, 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 "
"G0.5 author_model_config is empty. Load a packaged checkpoint, or "
"inject a backend explicitly for testing."
)
try:
@@ -137,7 +137,6 @@ class G05Policy(PreTrainedPolicy):
"THIRD_PARTY_NOTICES.md",
"NOTICE",
"README.md",
"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():
+1 -1
View File
@@ -666,7 +666,7 @@ def make_g05_pre_post_processors(
if config.normalization_mode == "checkpoint" and not config.processor_metadata:
raise ValueError(
"normalization_mode='checkpoint' requires processor_metadata from the converted checkpoint."
"normalization_mode='checkpoint' requires processor_metadata from the packaged checkpoint."
)
mode = _normalization_mode(config)
if mode is NormalizationMode.QUANTILES and dataset_stats:
+51 -412
View File
@@ -2,7 +2,6 @@
from __future__ import annotations
import json
import os
from pathlib import Path
@@ -14,14 +13,6 @@ 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 (
_camera_sizes,
_profile_config,
convert_checkpoint,
convert_dataset_stats,
convert_state_dict,
save_converted_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
@@ -109,175 +100,43 @@ def test_system2_fm_only_builder_uses_exact_cot_template_without_action_tokens()
assert "<action_action" not in config.prompt_template
def _base_r1lite_hydra():
state_action = [
{"key": "left_arm", "shape": 6},
{"key": "left_gripper", "shape": 1},
{"key": "right_arm", "shape": 6},
{"key": "right_gripper", "shape": 1},
]
images = [
{
"key": "head_rgb",
"camera_type": "exterior",
"lerobot_key": "observation.images.head_rgb",
"shape": [3, 224, 224],
def test_so101_runtime_pads_optional_left_wrist():
config = G05Config(
checkpoint_profile="g05-so101",
embodiment="so100",
action_head="flow",
runtime_system="system2",
predict_cot=True,
discrete_action=True,
continuous_action=True,
return_continuous_action=True,
policy_action_dim=20,
policy_state_dim=20,
raw_action_dim=6,
raw_state_dim=6,
chunk_size=32,
n_action_steps=16,
normalization_mode="identity",
camera_order=(
"observation.images.exterior",
"observation.images.wrist_left",
"observation.images.wrist_right",
),
camera_sizes={
"observation.images.exterior": (8, 8),
"observation.images.wrist_left": (8, 8),
"observation.images.wrist_right": (8, 8),
},
{
"key": "left_wrist_rgb",
"camera_type": "wrist_left",
"lerobot_key": "observation.images.left_wrist_rgb",
"shape": [3, 224, 224],
optional_camera_keys=("observation.images.wrist_left",),
input_features={
OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(6,)),
"observation.images.exterior": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 8, 8)),
"observation.images.wrist_right": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 8, 8)),
},
{
"key": "right_wrist_rgb",
"camera_type": "wrist_right",
"lerobot_key": "observation.images.right_wrist_rgb",
"shape": [3, 224, 224],
},
]
return {
"model": {
"model_arch": {
"action_dim": 27,
"proprio_dim": 27,
"num_input_images": 18,
"horizon_steps": 32,
"predict_cot": True,
"discrete_action": True,
"continuous_action": True,
},
"processor": {
"num_obs_steps": 6,
"use_stepwise_action_norm": True,
"camera_size_config": {
"exterior": [256, 256],
"wrist_left": [256, 256],
"wrist_right": [256, 256],
},
"samples_builder": {
"_target_": "g05.data_processor.processor.samples_builder.MixedSamplesBuilder"
},
},
},
"data": {
"action_size": 32,
"processors": {
"galaxea_r1lite": {
"shape_meta": {
"state": state_action,
"action": state_action,
"images": images,
},
"norm_default_mode": "z-score-tail",
"norm_exception_mode": {
"state": {"left_gripper": "q01/q99", "right_gripper": "q01/q99"},
"action": {"left_gripper": "q01/q99", "right_gripper": "q01/q99"},
},
"action_filter": {
"_target_": (
"g05.data_processor.processor.galaxea_action_processor.R1LiteJointActionFilter"
),
"joint_threshold": 0.002,
"gripper_threshold": 0.01,
},
}
},
},
}
def _base_r1lite_stats():
result = {"state": {}, "action": {}}
for category in result:
for key, width in (
("left_arm", 6),
("left_gripper", 1),
("right_arm", 6),
("right_gripper", 1),
):
if category == "action":
shape = (32, width)
prefix = "stepwise"
else:
shape = (width,)
prefix = "global"
result[category][key] = {
f"{prefix}_mean": torch.zeros(shape).tolist(),
f"{prefix}_std": torch.ones(shape).tolist(),
f"{prefix}_q01": torch.full(shape, -1.0).tolist(),
f"{prefix}_q99": torch.full(shape, 1.0).tolist(),
}
return {"galaxea_r1lite": result}
def _so101_hydra():
arm = [{"key": "right_arm", "shape": 6}]
images = [
{
"key": name,
"camera_type": camera_type,
"lerobot_key": f"__so100_{name}__",
"shape": [3, 224, 224],
}
for name, camera_type in (
("exterior", "exterior"),
("wrist_left", "wrist_left"),
("wrist_right", "wrist_right"),
)
]
return {
"model": {
"model_arch": {
"action_dim": 20,
"proprio_dim": 20,
"num_input_images": 3,
"horizon_steps": 32,
"predict_cot": True,
"discrete_action": True,
"continuous_action": True,
"return_continuous_action": True,
},
"processor": {
"num_obs_steps": 1,
"use_stepwise_action_norm": True,
"norm_default_mode": "q01/q99",
"camera_size_config": {
"exterior": [256, 256],
"wrist_left": [256, 256],
"wrist_right": [256, 256],
},
"samples_builder": {
"_target_": "g05.data_processor.processor.samples_builder.MixedSamplesBuilder"
},
},
},
"data": {
"action_size": 32,
"processors": {
"so100": {
"shape_meta": {"state": arm, "action": arm, "images": images},
"norm_default_mode": "q01/q99",
"use_stepwise_action_norm": True,
}
},
},
}
def test_so101_profile_exposes_both_system2_heads_and_pads_left_wrist():
flow = _profile_config("g05-so101", _so101_hydra())
actioncodec = _profile_config("g05-so101", _so101_hydra(), action_head="actioncodec")
flow.camera_sizes = _camera_sizes(flow.processor_metadata, flow.camera_order)
flow.input_features = {
OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(6,)),
"observation.images.exterior": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 8, 8)),
"observation.images.wrist_right": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 8, 8)),
}
flow.output_features = {ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(6,))}
flow.device = "cpu"
flow.camera_sizes = dict.fromkeys(flow.camera_order, (8, 8))
preprocessor, _ = make_pre_post_processors(flow)
output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(6,))},
device="cpu",
)
preprocessor, _ = make_pre_post_processors(config)
processed = preprocessor(
{
@@ -288,83 +147,30 @@ def test_so101_profile_exposes_both_system2_heads_and_pads_left_wrist():
}
)
assert flow.embodiment == "so100"
assert flow.runtime_system == "system2" and flow.action_head == "flow"
assert actioncodec.runtime_system == "system2" and actioncodec.action_head == "actioncodec"
assert flow.optional_camera_keys == ("observation.images.wrist_left",)
assert processed["observation.images.wrist_left"].shape == (1, 3, 8, 8)
assert torch.all(processed["observation.images.wrist_left"] == -1)
assert not processed["action_dim_is_pad"][0, 10:16].any()
def test_robotwin_profile_resolves_official_data_include_and_uses_env_camera_names():
hydra = _so101_hydra()
hydra["data"] = "${oc.load:configs/data/_mixtures/robotwin.yaml}"
hydra["model"]["model_arch"].update(
{
"predict_cot": False,
"discrete_action": False,
"continuous_action": True,
"horizon_steps": "${data.action_size}",
"num_input_images": "${eval:'${model.model_arch.cond_steps} * 3'}",
"cond_steps": "${obs_image_steps:${data.obs_size}}",
}
def test_libero_runtime_executes_ten_step_window_and_binarizes_gripper():
config = G05Config(
checkpoint_profile="custom",
embodiment="libero",
action_head="flow",
discrete_action=False,
continuous_action=True,
return_continuous_action=True,
chunk_size=32,
n_action_steps=10,
normalization_mode="identity",
libero_gripper_binarize=True,
)
hydra["model"]["processor"]["num_obs_steps"] = 1
hydra["model"]["tokenizer"] = "${tokenizer}"
hydra["tokenizer"] = {"_target_": "g05.tokenizer.interface.vq_base.VQActionTokenizer"}
robotwin_data = {
"action_size": 32,
"obs_size": 1,
"processors": {
"robotwin": {
**_so101_hydra()["data"]["processors"]["so100"],
"shape_meta": {
"state": [
{"key": "left_arm", "shape": 6},
{"key": "left_gripper", "shape": 1},
{"key": "right_arm", "shape": 6},
{"key": "right_gripper", "shape": 1},
],
"action": [
{"key": "left_arm", "shape": 6},
{"key": "left_gripper", "shape": 1},
{"key": "right_arm", "shape": 6},
{"key": "right_gripper", "shape": 1},
],
"images": _so101_hydra()["data"]["processors"]["so100"]["shape_meta"]["images"],
},
}
},
}
with pytest.raises(ValueError, match="unresolved data include"):
_profile_config("g05-robotwin20", hydra)
config = _profile_config("g05-robotwin20", hydra, data_override=robotwin_data)
assert config.chunk_size == 32
assert config.n_action_steps == 8
assert config.action_codec_metadata == hydra["tokenizer"]
assert config.num_input_images == 3
assert config.camera_order == (
"observation.images.head_camera",
"observation.images.left_camera",
"observation.images.right_camera",
)
def test_libero_profile_executes_official_ten_step_open_loop_window():
hydra = _so101_hydra()
hydra["model"]["model_arch"]["num_input_images"] = 2
config = _profile_config("g05-libero", hydra)
assert config.chunk_size == 32
assert config.n_action_steps == 10
assert config.libero_gripper_binarize
_, postprocessor = make_pre_post_processors(config)
policy_action = torch.zeros(5, 20)
policy_action[:, 19] = torch.tensor([0.0, 0.5, 1.0, -0.2, 1.2])
env_action = postprocessor(policy_action)
torch.testing.assert_close(env_action[:, -1], torch.tensor([1.0, 1.0, -1.0, 1.0, -1.0]))
@@ -387,54 +193,6 @@ def test_select_action_discards_tail_beyond_execution_window():
assert calls == 2
def test_base_system2_requires_named_embodiment_and_roundtrips_mixed_tail_stats(tmp_path):
hydra = _base_r1lite_hydra()
with pytest.raises(ValueError, match="concrete --embodiment"):
_profile_config("g05-base", hydra)
config = _profile_config("g05-base", hydra, embodiment="galaxea_r1lite")
actioncodec_config = _profile_config(
"g05-base", hydra, embodiment="galaxea_r1lite", action_head="actioncodec"
)
config.camera_sizes = _camera_sizes(config.processor_metadata, config.camera_order)
config.camera_sizes = dict.fromkeys(config.camera_order, (8, 8))
stats = convert_dataset_stats(_base_r1lite_stats(), config)
preprocessor, postprocessor = make_pre_post_processors(config, dataset_stats=stats)
raw_state = torch.linspace(-2, 2, 14).repeat(1, 6, 1)
raw_action = raw_state[:, -1] + torch.linspace(-0.2, 0.2, 14).repeat(1, 32, 1)
raw_batch = {
OBS_STATE: raw_state,
ACTION: raw_action,
**{camera: torch.zeros(1, 6, 3, 8, 8, dtype=torch.uint8) for camera in config.camera_order},
"task": ["native system 2"],
}
processed = preprocessor(raw_batch)
restored = postprocessor(processed[ACTION])
assert config.runtime_system == "system2"
assert config.predict_cot and config.discrete_action and config.continuous_action
assert config.action_head == "flow" and actioncodec_config.action_head == "actioncodec"
assert actioncodec_config.runtime_system == "system2"
assert not actioncodec_config.return_continuous_action
assert config.policy_action_dim == 27
assert config.num_input_images == 18
assert "<prompt_text_!>" in config.prompt_template
# G0.5-base's 32-step head includes n_obs_steps - 1 alignment steps.
assert processed["action_dim_is_pad"].shape == (1, 27)
assert not processed["action_op_mask"].any()
assert restored.shape == (1, 27, 14)
torch.testing.assert_close(restored, raw_action[:, 5:], atol=2e-5, rtol=2e-5)
preprocessor.save_pretrained(tmp_path)
postprocessor.save_pretrained(tmp_path)
loaded_preprocessor, loaded_postprocessor = make_pre_post_processors(config, pretrained_path=tmp_path)
reloaded = loaded_preprocessor(raw_batch)
reloaded_restored = loaded_postprocessor(reloaded[ACTION])
torch.testing.assert_close(reloaded[ACTION], processed[ACTION])
torch.testing.assert_close(reloaded_restored, restored)
def test_libero_and_atomic4_are_distinct_validated_mappings():
with pytest.raises(ValueError, match="27D"):
G05Config(
@@ -763,13 +521,13 @@ def test_forward_backward_update_and_save_reload(tmp_path: Path):
def test_save_pretrained_copies_required_gated_sidecars_portably(tmp_path: Path):
source = tmp_path / "converted"
source = tmp_path / "checkpoint"
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"):
for name in ("LICENSE-G0.5", "NOTICE"):
(source / name).write_text("{}")
config = _config(
author_model_config={
@@ -804,125 +562,6 @@ def test_tiny_fixed_batch_overfit_reduces_loss():
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_conversion_records_and_deduplicates_tied_weight_aliases(tmp_path: Path):
tied = torch.zeros(2, 3)
aliases = save_converted_state_dict(
{
"backend.model.vlm.input_proj.weight": tied,
"backend.model.vlm.output_proj.weight": tied,
},
tmp_path / "model.safetensors",
)
assert aliases == {"backend.model.vlm.output_proj.weight": "backend.model.vlm.input_proj.weight"}
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
norm_default_mode: q01/q99
camera_size_config:
exterior: [256, 256]
wrist_right: [256, 256]
data:
action_size: 32
processors:
libero:
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, camera_type: exterior, lerobot_key: observation.images.image, shape: [3, 224, 224]}
- {key: wrist_image, camera_type: wrist_right, lerobot_key: observation.images.wrist_image, shape: [3, 224, 224]}
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()
assert (output / "README.md").is_file()
assert "Modification notice" in (output / "NOTICE").read_text()
config = PreTrainedConfig.from_pretrained(output)
assert isinstance(config, G05Config)
assert config.source_checkpoint_revision
assert config.prompt_template.startswith("<chat_user_prefix><image0_image_!><image1_image_!>")
assert config.camera_sizes == {
"observation.images.image": (256, 256),
"observation.images.wrist_image": (256, 256),
}
@pytest.mark.skipif(
not os.environ.get("LEROBOT_G05_CHECKPOINT"),
reason="requires an accepted gated OpenGalaxea/G05 checkpoint and author CUDA environment",
Generated
+1 -1
View File
@@ -3508,7 +3508,7 @@ requires-dist = [
{ name = "transformers", marker = "extra == 'transformers-dep'", specifier = ">=5.4.0,<5.6.0" },
{ name = "wandb", marker = "extra == 'training'", specifier = ">=0.24.0,<0.28.0" },
]
provides-extras = ["dataset", "training", "hardware", "viz", "core-scripts", "evaluation", "dataset-viz", "av-dep", "pygame-dep", "placo-dep", "transformers-dep", "grpcio-dep", "accelerate-dep", "can-dep", "peft-dep", "scipy-dep", "diffusers-dep", "qwen-vl-utils-dep", "matplotlib-dep", "pyserial-dep", "deepdiff-dep", "pynput-dep", "pyzmq-dep", "motorbridge-dep", "motorbridge-smart-servo-dep", "timm-dep", "feetech", "dynamixel", "damiao", "robstride", "openarms", "gamepad", "hopejr", "lekiwi", "unitree-g1", "reachy2", "rebot", "kinematics", "intelrealsense", "phone", "diffusion", "wallx", "pi", "molmoact2", "smolvla", "multi-task-dit", "groot", "sarm", "robometer", "topreward", "xvla", "eo1", "fastwam", "evo1", "hilserl", "vla-jepa", "lingbot-va", "async", "peft", "annotations", "dev", "notebook", "test", "video-benchmark", "aloha", "pusht", "libero", "metaworld", "all"]
provides-extras = ["dataset", "training", "hardware", "viz", "core-scripts", "evaluation", "dataset-viz", "av-dep", "pygame-dep", "placo-dep", "transformers-dep", "grpcio-dep", "accelerate-dep", "can-dep", "peft-dep", "scipy-dep", "diffusers-dep", "qwen-vl-utils-dep", "matplotlib-dep", "pyserial-dep", "deepdiff-dep", "pynput-dep", "pyzmq-dep", "motorbridge-dep", "motorbridge-smart-servo-dep", "timm-dep", "feetech", "dynamixel", "damiao", "robstride", "openarms", "gamepad", "hopejr", "lekiwi", "unitree-g1", "reachy2", "rebot", "kinematics", "intelrealsense", "phone", "diffusion", "wallx", "pi", "molmoact2", "smolvla", "multi-task-dit", "groot", "sarm", "robometer", "topreward", "xvla", "eo1", "fastwam", "evo1", "hilserl", "vla-jepa", "lingbot-va", "g05", "async", "peft", "annotations", "dev", "notebook", "test", "video-benchmark", "aloha", "pusht", "libero", "metaworld", "all"]
[[package]]
name = "librt"