mirror of
https://github.com/huggingface/lerobot.git
synced 2026-08-08 17:39:44 +00:00
refactor(g05): keep checkpoint tooling out of runtime PR
This commit is contained in:
@@ -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
@@ -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()
|
||||
@@ -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():
|
||||
|
||||
@@ -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
@@ -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",
|
||||
|
||||
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user