diff --git a/benchmarks/g05_checkpoint_parity.py b/benchmarks/g05_checkpoint_parity.py new file mode 100644 index 000000000..1dff5b3cd --- /dev/null +++ b/benchmarks/g05_checkpoint_parity.py @@ -0,0 +1,307 @@ +#!/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) -> dict[str, Any]: + images_by_key = {meta["key"]: meta for meta in processor.shape_meta["images"]} + 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), + **{meta["lerobot_key"]: raw["images"][meta["key"]][-1] for meta in images_by_key.values()}, + "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() + 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)) 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(): + 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) + port_env_action = postprocessor(port_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() diff --git a/docs/source/g05.mdx b/docs/source/g05.mdx index 2dea16ee1..acb97036e 100644 --- a/docs/source/g05.mdx +++ b/docs/source/g05.mdx @@ -100,11 +100,29 @@ lerobot-rollout \ CPU unit tests cover factory loading, config incompatibilities, prompt pass-through, LIBERO and `atomic_4` mappings, padding masks, inverse action projection, strict conversion diagnostics, a finite forward/backward/update, and save/reload parity. -Numerical author-oracle parity and 50-episode LIBERO/RoboTwin evaluation require -licensed checkpoint access and suitable CUDA hardware; no benchmark number is -claimed by this integration until those gates run. +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: ```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. diff --git a/src/lerobot/policies/g05/configuration_g05.py b/src/lerobot/policies/g05/configuration_g05.py index c346a6f59..d7d567d1e 100644 --- a/src/lerobot/policies/g05/configuration_g05.py +++ b/src/lerobot/policies/g05/configuration_g05.py @@ -96,6 +96,22 @@ G05_EMBODIMENT_MAPPINGS: dict[str, dict[str, tuple[int, ...]]] = { }, } +G05_POLICY_PARTS: dict[int, dict[str, int]] = { + 20: { + "left_control": 9, + "left_gripper": 1, + "right_control": 9, + "right_gripper": 1, + }, + 27: { + "left_control": 9, + "left_gripper": 1, + "right_control": 9, + "right_gripper": 1, + "lower_body": 7, + }, +} + _PROFILE_DEFAULTS = { "g05-base": ("checkpoint", 20, 16), "g05-libero": ("q01_q99", 20, 32), @@ -128,6 +144,7 @@ class G05Config(PreTrainedConfig): raw_state_dim: int = 7 chunk_size: int = 16 normalization_mode: str = "checkpoint" + normalization_clip: tuple[float, float] | None = None use_stepwise_action_norm: bool = False gripper_indices: tuple[int, ...] = (6,) camera_order: tuple[str, ...] = field(default_factory=lambda: G05_CAMERA_PROFILES["libero"]) @@ -160,15 +177,24 @@ class G05Config(PreTrainedConfig): super().__post_init__() self.camera_order = tuple(self.camera_order) self.camera_sizes = {key: tuple(size) for key, size in (self.camera_sizes or {}).items()} + if self.normalization_clip is not None: + self.normalization_clip = tuple(self.normalization_clip) + if len(self.normalization_clip) != 2 or self.normalization_clip[0] >= self.normalization_clip[1]: + raise ValueError("normalization_clip must be an increasing (minimum, maximum) pair.") if not self.camera_sizes and self.embodiment in G05_CAMERA_SIZE_PROFILES: self.camera_sizes = G05_CAMERA_SIZE_PROFILES[self.embodiment].copy() if self.num_input_images == 0: self.num_input_images = len(self.camera_order) * self.n_obs_steps if not self.prompt_template: + samples_builder = self.processor_metadata.get("samples_builder") or {} + if isinstance(samples_builder, dict): + samples_builder_target = str(samples_builder.get("_target_", "")) + else: + samples_builder_target = str(samples_builder) self.prompt_template = make_g05_prompt_template( self.num_input_images, predict_cot=self.predict_cot, - flow_only=self.continuous_action and not self.discrete_action, + flow_only=samples_builder_target.endswith("FMOnly"), ) if self.checkpoint_profile not in _PROFILE_DEFAULTS and self.checkpoint_profile != "custom": raise ValueError( @@ -187,6 +213,10 @@ class G05Config(PreTrainedConfig): raise ValueError("The flow runtime requires continuous_action=True.") if self.action_head == "flow" and not self.return_continuous_action: raise ValueError("The flow runtime requires return_continuous_action=True.") + if self.policy_action_dim not in G05_POLICY_PARTS: + raise ValueError( + f"No named G0.5 shared action layout for policy_action_dim={self.policy_action_dim}." + ) if not (self.discrete_action or self.continuous_action): raise ValueError("At least one G0.5 action path must be enabled.") if self.embodiment not in G05_EMBODIMENT_MAPPINGS: diff --git a/src/lerobot/policies/g05/convert_g05_checkpoint.py b/src/lerobot/policies/g05/convert_g05_checkpoint.py index fe6eb4a56..f3fa48751 100644 --- a/src/lerobot/policies/g05/convert_g05_checkpoint.py +++ b/src/lerobot/policies/g05/convert_g05_checkpoint.py @@ -54,6 +54,7 @@ _REQUIRED_PREFIXES = ( @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) @@ -113,6 +114,29 @@ def convert_state_dict( 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): @@ -126,11 +150,19 @@ def _profile_config(profile: str, hydra: dict[str, Any]) -> G05Config: model = hydra.get("model", {}) arch = model.get("model_arch", {}) data = hydra.get("data", {}) + embodiment = { + "g05-libero": "libero", + "g05-robotwin20": "robotwin20", + }.get(profile, "libero") + 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)) - norm_name = str(model.get("processor", {}).get("norm_default_mode", "")).lower() + norm_name = str(processor_metadata.get("norm_default_mode", "")).lower() checkpoint_normalization = { "q01/q99": "q01_q99", "z-score": "z_score", @@ -152,11 +184,12 @@ def _profile_config(profile: str, hydra: dict[str, Any]) -> G05Config: return_continuous_action=True, chunk_size=horizon, normalization_mode="q01_q99", + normalization_clip=(-5.0, 5.0), use_stepwise_action_norm=True, camera_order=G05_CAMERA_PROFILES["libero"], num_input_images=num_input_images, author_model_config=arch, - processor_metadata=model.get("processor", {}), + processor_metadata=processor_metadata, action_codec_metadata=model.get("tokenizer", hydra.get("tokenizer", {})), author_source_revision=G05_SOURCE_REVISION, source_checkpoint_revision=G05_HUB_REVISION, @@ -177,11 +210,12 @@ def _profile_config(profile: str, hydra: dict[str, Any]) -> G05Config: raw_action_dim=14, chunk_size=horizon, normalization_mode="q01_q99", + normalization_clip=(-5.0, 5.0), use_stepwise_action_norm=True, camera_order=G05_CAMERA_PROFILES["robotwin20"], num_input_images=num_input_images, author_model_config=arch, - processor_metadata=model.get("processor", {}), + processor_metadata=processor_metadata, action_codec_metadata=model.get("tokenizer", hydra.get("tokenizer", {})), author_source_revision=G05_SOURCE_REVISION, source_checkpoint_revision=G05_HUB_REVISION, @@ -191,7 +225,7 @@ def _profile_config(profile: str, hydra: dict[str, Any]) -> G05Config: action_head = "actioncodec" if discrete else "flow" if checkpoint_normalization is None: raise ValueError("g05-base conversion requires the resolved checkpoint processor.norm_default_mode.") - exceptions = model.get("processor", {}).get("norm_exception_mode") + exceptions = processor_metadata.get("norm_exception_mode") if exceptions: raise ValueError( "g05-base has per-part normalization exceptions; select a concrete benchmark " @@ -208,11 +242,12 @@ def _profile_config(profile: str, hydra: dict[str, Any]) -> G05Config: return_continuous_action=action_head == "flow", chunk_size=horizon, normalization_mode=checkpoint_normalization, - use_stepwise_action_norm=bool(model.get("processor", {}).get("use_stepwise_action_norm", False)), + normalization_clip=(-5.0, 5.0), + use_stepwise_action_norm=bool(processor_metadata.get("use_stepwise_action_norm", False)), camera_order=G05_CAMERA_PROFILES["libero"], num_input_images=num_input_images, author_model_config=arch, - processor_metadata=model.get("processor", {}), + processor_metadata=processor_metadata, action_codec_metadata=model.get("tokenizer", hydra.get("tokenizer", {})), author_source_revision=G05_SOURCE_REVISION, source_checkpoint_revision=G05_HUB_REVISION, @@ -225,8 +260,9 @@ 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(item["shape"][-2:]) + 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 } @@ -320,7 +356,7 @@ def convert_checkpoint( report.fail_if_invalid() output_dir.mkdir(parents=True, exist_ok=True) - save_file(converted, output_dir / "model.safetensors") + 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) diff --git a/src/lerobot/policies/g05/modeling_g05.py b/src/lerobot/policies/g05/modeling_g05.py index f5b74ed6a..eae524df5 100644 --- a/src/lerobot/policies/g05/modeling_g05.py +++ b/src/lerobot/policies/g05/modeling_g05.py @@ -25,7 +25,7 @@ from lerobot.configs.policies import PreTrainedConfig from lerobot.policies.pretrained import PreTrainedPolicy from lerobot.utils.constants import ACTION, OBS_STATE -from .configuration_g05 import G05Config +from .configuration_g05 import G05_POLICY_PARTS, G05Config def _author_backend(config: G05Config) -> nn.Module: @@ -162,6 +162,17 @@ class G05Policy(PreTrainedPolicy): if callable(reset): reset() + def to(self, *args, **kwargs) -> G05Policy: + """Move the author ActionCodec sidecar along with the policy module.""" + + result = super().to(*args, **kwargs) + action_tokenizer = getattr(self.backend, "action_tokenizer", None) + move_tokenizer = getattr(action_tokenizer, "to", None) + if callable(move_tokenizer): + device = next(self.backend.parameters()).device + move_tokenizer(device) + return result + def get_optim_params(self) -> dict: get_params = getattr(self.backend, "get_optim_params", None) if callable(get_params): @@ -242,7 +253,11 @@ class G05Policy(PreTrainedPolicy): camera = self.config.camera_order[image_index % len(self.config.camera_order)] sample[f"image{image_index}"] = self.config.camera_sizes[camera] action = batch.get(ACTION) - if " 1 else action_op_mask + ) + else: + action_payload["action_op_mask"] = ~action_dim_is_pad[index] + action_payload["parts_meta"] = batch.get( + "action_parts_meta", G05_POLICY_PARTS[self.config.policy_action_dim] + ) + sample["action"] = action_payload samples.append(sample) prepared = dict(batch) prepared["samples"] = samples diff --git a/src/lerobot/policies/g05/processor_g05.py b/src/lerobot/policies/g05/processor_g05.py index 1a7701d6d..f6d1158d1 100644 --- a/src/lerobot/policies/g05/processor_g05.py +++ b/src/lerobot/policies/g05/processor_g05.py @@ -14,7 +14,7 @@ from dataclasses import dataclass from typing import Any import torch -import torch.nn.functional as functional +import torchvision.transforms.functional as vision_functional from lerobot.configs.types import FeatureType, NormalizationMode, PipelineFeatureType, PolicyFeature from lerobot.processor import ( @@ -41,7 +41,7 @@ from lerobot.utils.constants import ( POLICY_PREPROCESSOR_DEFAULT_NAME, ) -from .configuration_g05 import G05_EMBODIMENT_MAPPINGS, G05Config +from .configuration_g05 import G05_EMBODIMENT_MAPPINGS, G05_POLICY_PARTS, G05Config def _copy_feature_tree( @@ -80,13 +80,15 @@ class G05ImageTransformStep(ProcessorStep): if image.ndim < 3 or image.shape[-3] != 3: raise ValueError(f"G0.5 camera {key!r} must end in [3,H,W], got {image.shape}.") was_floating_point = torch.is_floating_point(image) - image = image.float() - if not was_floating_point: - image = image / 255.0 flat = image.reshape(-1, *image.shape[-3:]) target_size = self.camera_sizes[key] if tuple(flat.shape[-2:]) != target_size: - flat = functional.interpolate(flat, size=target_size, mode="bilinear", align_corners=False) + # The author prepends torchvision Resize before its uint8-to-float + # transform. Preserve its antialiasing and uint8 quantization. + flat = vision_functional.resize(flat, list(target_size)) + flat = flat.float() + if not was_floating_point: + flat = flat / 255.0 mean = flat.new_tensor(self.mean).view(1, 3, 1, 1) std = flat.new_tensor(self.std).view(1, 3, 1, 1) observation[key] = ((flat - mean) / std).reshape(*image.shape[:-3], 3, *target_size) @@ -161,6 +163,16 @@ class G05EmbodimentProjectionStep(ProcessorStep): state_mask[..., list(self.mapping["state"])] = False complementary = dict(transition.get(TransitionKey.COMPLEMENTARY_DATA) or {}) complementary["proprio_dim_is_pad"] = state_mask + action_mask = torch.ones( + *raw_state.shape[:-1], + self.policy_action_dim, + dtype=torch.bool, + device=observation[OBS_STATE].device, + ) + action_mask[..., list(self.mapping["action"])] = False + complementary["action_dim_is_pad"] = action_mask + complementary["action_op_mask"] = ~action_mask + complementary["action_parts_meta"] = G05_POLICY_PARTS[self.policy_action_dim].copy() complementary["g05_camera_order"] = self.camera_order transition[TransitionKey.COMPLEMENTARY_DATA] = complementary action = transition.get(TransitionKey.ACTION) @@ -174,7 +186,7 @@ class G05EmbodimentProjectionStep(ProcessorStep): ) action_mask[..., list(self.mapping["action"])] = False complementary = dict(transition.get(TransitionKey.COMPLEMENTARY_DATA) or {}) - complementary["action_dim_is_pad"] = action_mask + complementary.setdefault("action_dim_is_pad", action_mask) if "action_is_pad" not in complementary: complementary["action_is_pad"] = torch.zeros( *action.shape[:-1], dtype=torch.bool, device=action.device @@ -204,6 +216,37 @@ class G05EmbodimentProjectionStep(ProcessorStep): } +@dataclass +@ProcessorStepRegistry.register(name="g05_normalization_clamp") +class G05NormalizationClampStep(ProcessorStep): + """Match the author's finite clamp after normalization.""" + + minimum: float = -5.0 + maximum: float = 5.0 + + def _clamp(self, value: torch.Tensor) -> torch.Tensor: + return value.clamp(self.minimum, self.maximum).nan_to_num(nan=0.0, posinf=0.0, neginf=0.0) + + def __call__(self, transition: EnvTransition) -> EnvTransition: + transition = transition.copy() + observation = dict(transition.get(TransitionKey.OBSERVATION) or {}) + if OBS_STATE in observation: + observation[OBS_STATE] = self._clamp(observation[OBS_STATE]) + transition[TransitionKey.OBSERVATION] = observation + action = transition.get(TransitionKey.ACTION) + if isinstance(action, torch.Tensor): + transition[TransitionKey.ACTION] = self._clamp(action) + return transition + + def transform_features( + self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]] + ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]: + return features + + def get_config(self) -> dict[str, Any]: + return {"minimum": self.minimum, "maximum": self.maximum} + + @dataclass @ProcessorStepRegistry.register(name="g05_inverse_action_projection") class G05InverseActionProjectionStep(ProcessorStep): @@ -313,28 +356,36 @@ def make_g05_pre_post_processors( FeatureType.VISUAL: NormalizationMode.IDENTITY, } + steps: list[ProcessorStep] = [ + AddBatchDimensionProcessorStep(), + G05ImageTransformStep( + camera_order=config.camera_order, + camera_sizes=config.camera_sizes, + mean=config.image_mean, + std=config.image_std, + ), + G05EmbodimentProjectionStep( + embodiment=config.embodiment, + policy_state_dim=config.policy_state_dim, + policy_action_dim=config.policy_action_dim, + camera_order=config.camera_order, + ), + NormalizerProcessorStep( + features=policy_features, + norm_map=norm_map, + stats=projected_stats, + ), + ] + if config.normalization_clip is not None: + steps.append( + G05NormalizationClampStep( + minimum=config.normalization_clip[0], + maximum=config.normalization_clip[1], + ) + ) + steps.append(DeviceProcessorStep(device=config.device)) preprocessor = PolicyProcessorPipeline[dict[str, Any], dict[str, Any]]( - steps=[ - AddBatchDimensionProcessorStep(), - G05ImageTransformStep( - camera_order=config.camera_order, - camera_sizes=config.camera_sizes, - mean=config.image_mean, - std=config.image_std, - ), - G05EmbodimentProjectionStep( - embodiment=config.embodiment, - policy_state_dim=config.policy_state_dim, - policy_action_dim=config.policy_action_dim, - camera_order=config.camera_order, - ), - NormalizerProcessorStep( - features=policy_features, - norm_map=norm_map, - stats=projected_stats, - ), - DeviceProcessorStep(device=config.device), - ], + steps=steps, name=POLICY_PREPROCESSOR_DEFAULT_NAME, to_transition=batch_to_transition, to_output=transition_to_batch, diff --git a/tests/policies/g05/test_g05.py b/tests/policies/g05/test_g05.py index 12f6930b3..113bfe126 100644 --- a/tests/policies/g05/test_g05.py +++ b/tests/policies/g05/test_g05.py @@ -14,7 +14,11 @@ from lerobot.configs.policies import PreTrainedConfig from lerobot.configs.types import FeatureType, PolicyFeature from lerobot.policies.factory import get_policy_class, make_policy_config, make_pre_post_processors from lerobot.policies.g05.configuration_g05 import G05_EMBODIMENT_MAPPINGS, G05Config -from lerobot.policies.g05.convert_g05_checkpoint import convert_checkpoint, convert_state_dict +from lerobot.policies.g05.convert_g05_checkpoint import ( + convert_checkpoint, + 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 @@ -82,6 +86,25 @@ def test_factory_wiring_is_lazy(): assert get_policy_class("g05") is G05Policy +def test_system2_fm_only_builder_uses_exact_cot_template_without_action_tokens(): + config = _config( + action_head="flow", + runtime_system="system2", + predict_cot=True, + discrete_action=False, + continuous_action=True, + return_continuous_action=True, + processor_metadata={ + "samples_builder": { + "_target_": ("g05.data_processor.processor.samples_builder.SubtaskCoTBuilderFMOnly") + } + }, + ) + + assert "\n|Action: " in config.prompt_template + assert "") + assert config.prompt_template.startswith("") + assert config.camera_sizes == { + "observation.images.image": (256, 256), + "observation.images.wrist_image": (256, 256), + } @pytest.mark.skipif(