From 24e3564168d7afa86bdfc9741fe2273950c6f768 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Tue, 28 Jul 2026 21:34:29 +0200 Subject: [PATCH] feat(g05): add LeRobot training support --- docs/source/g05.mdx | 52 ++++++++++- src/lerobot/policies/g05/configuration_g05.py | 3 + src/lerobot/policies/g05/modeling_g05.py | 31 ++++++- tests/policies/g05/test_g05.py | 88 ++++++++++++++++++- 4 files changed, 168 insertions(+), 6 deletions(-) diff --git a/docs/source/g05.mdx b/docs/source/g05.mdx index e002c9250..4153525bd 100644 --- a/docs/source/g05.mdx +++ b/docs/source/g05.mdx @@ -112,11 +112,61 @@ The checkpoint is non-commercial and may be private; authenticate with `hf auth login` before loading it. Do not add `--direct_subtask` when inspecting the checkpoint's native System 2 CoT telemetry. +## Fine-tune with `lerobot-train` + +G0.5 implements LeRobot's training surface: `forward` runs the author training +backend, the policy exposes the author VLM/vision/action optimizer groups, and +the checkpoint can be saved, resumed, and loaded by the normal LeRobot scripts. +For example, fine-tune the private SO-101 checkpoint on a LeRobot dataset: + +```bash +export HF_USER=your_hf_username + +lerobot-train \ + --dataset.repo_id=${HF_USER}/my_so101_dataset \ + --policy.path=lerobot/g05_so101 \ + --policy.device=cuda \ + --policy.repo_id=${HF_USER}/g05_so101_finetuned \ + --policy.private=true \ + --output_dir=outputs/train/g05_so101 \ + --job_name=g05_so101 \ + --batch_size=16 \ + --steps=10000 \ + --save_freq=1000 +``` + +The SO-101 recipe uses AdamW at `8e-5` with 1,000 warmup steps. The packaged +LIBERO and RoboTwin configurations use their released `1e-5` recipe, with +1,000 and 500 warmup steps respectively. All profiles preserve G0.5's six +decay/no-decay parameter groups and the configured VLM and vision learning-rate +multipliers. Override these only when deliberately changing the author recipe: + +```bash +--policy.optimizer_lr=2e-5 \ +--policy.optimizer_backbone_lr_multiplier=0.5 \ +--policy.optimizer_vision_lr_multiplier=0.1 +``` + +The dataset must expose the state, action, camera, and task features matching the +selected checkpoint contract in the table above. For SO-101, use camera names +`exterior` and `wrist_right`; the optional `wrist_left` input is zero-filled. +Training System 2 language targets additionally requires the checkpoint's +annotated CoT fields; a normal LeRobot recording supplies action supervision but +does not synthesize CoT labels. + +Resume a saved run with the standard LeRobot checkpoint: + +```bash +lerobot-train \ + --config_path=outputs/train/g05_so101/checkpoints/last/pretrained_model/train_config.json \ + --resume=true +``` + ## Validation status CPU unit tests cover factory loading, config incompatibilities, prompt pass-through, LIBERO and `atomic_4` mappings, padding masks, inverse action projection, a finite -forward/backward/update, and save/reload parity: +forward/backward/update, author optimizer-group wiring, and save/reload parity: ```bash uv run pytest tests/policies/g05 tests/runtime/test_g05_adapter.py -q diff --git a/src/lerobot/policies/g05/configuration_g05.py b/src/lerobot/policies/g05/configuration_g05.py index e34d6da8c..385c2b2a3 100644 --- a/src/lerobot/policies/g05/configuration_g05.py +++ b/src/lerobot/policies/g05/configuration_g05.py @@ -210,6 +210,9 @@ class G05Config(PreTrainedConfig): optimizer_betas: tuple[float, float] = (0.9, 0.95) optimizer_weight_decay: float = 0.01 optimizer_grad_clip_norm: float = 1.0 + optimizer_backbone_lr_multiplier: float = 1.0 + optimizer_vision_lr_multiplier: float = 1.0 + optimizer_apply_decay_on_norm_and_bias: bool = False scheduler_warmup_steps: int = 500 def __post_init__(self) -> None: diff --git a/src/lerobot/policies/g05/modeling_g05.py b/src/lerobot/policies/g05/modeling_g05.py index 2bdf673d7..4e0b9ced5 100644 --- a/src/lerobot/policies/g05/modeling_g05.py +++ b/src/lerobot/policies/g05/modeling_g05.py @@ -22,6 +22,7 @@ from huggingface_hub import snapshot_download from torch import Tensor, nn from lerobot.configs.policies import PreTrainedConfig +from lerobot.optim.optimizers import OptimizerParams from lerobot.policies.pretrained import PreTrainedPolicy from lerobot.utils.constants import ACTION, OBS_STATE @@ -190,11 +191,21 @@ class G05Policy(PreTrainedPolicy): move_tokenizer(device) return result - def get_optim_params(self) -> dict: + def get_optim_params(self) -> OptimizerParams: + get_param_groups = getattr(self.backend, "get_optim_param_groups", None) + if callable(get_param_groups): + return get_param_groups( + lr=self.config.optimizer_lr, + weight_decay=self.config.optimizer_weight_decay, + apply_decay_on_norm_and_bias=self.config.optimizer_apply_decay_on_norm_and_bias, + backbone_lr_multiplier=self.config.optimizer_backbone_lr_multiplier, + vision_lr_multiplier=self.config.optimizer_vision_lr_multiplier, + ) get_params = getattr(self.backend, "get_optim_params", None) if callable(get_params): - return get_params() - return {"params": [parameter for parameter in self.parameters() if parameter.requires_grad]} + params = get_params() + return [params] if isinstance(params, dict) and "params" in params else params + return [parameter for parameter in self.parameters() if parameter.requires_grad] @staticmethod def _task_values(batch: Mapping[str, Any], task: str | None, batch_size: int) -> list[str]: @@ -210,6 +221,14 @@ class G05Policy(PreTrainedPolicy): "or model-local sampling is performed." ) + @staticmethod + def _batch_item(value: Any, index: int, batch_size: int) -> Any: + if isinstance(value, Tensor) and value.ndim > 0 and value.shape[0] == batch_size: + return value[index] + if isinstance(value, list | tuple) and len(value) == batch_size: + return value[index] + return value + def _prepare_author_batch(self, batch: Mapping[str, Any], task: str | None = None) -> dict[str, Any]: prepare = getattr(self.backend, "prepare_lerobot_batch", None) if callable(prepare): @@ -266,6 +285,12 @@ class G05Policy(PreTrainedPolicy): sample["frequency"] = frequency if self.config.predict_cot: sample["prompt"] = "predict subtask" + atomic_task = batch.get("atomic_task") + if atomic_task is not None: + atomic_task = str(self._batch_item(atomic_task, index, batch_size)) + sample["atomic_task"] = ( + atomic_task if atomic_task.startswith("Subtask:") else f"Subtask: {atomic_task}" + ) for image_index in range(self.config.num_prompt_images): camera = self.config.camera_order[image_index % len(self.config.camera_order)] sample[f"image{image_index}"] = self.config.camera_sizes[camera] diff --git a/tests/policies/g05/test_g05.py b/tests/policies/g05/test_g05.py index 03a030873..b0856ba37 100644 --- a/tests/policies/g05/test_g05.py +++ b/tests/policies/g05/test_g05.py @@ -43,6 +43,50 @@ class TinyG05Backend(nn.Module): return loss, {"fm_loss": loss.detach()} +class GroupedTinyG05Backend(TinyG05Backend): + def __init__(self): + super().__init__() + self.action_scale = nn.Parameter(torch.ones(())) + self.vision_scale = nn.Parameter(torch.ones(())) + self.optim_kwargs = None + + def get_optim_param_groups( + self, + lr, + weight_decay, + apply_decay_on_norm_and_bias=False, + backbone_lr_multiplier=1.0, + vision_lr_multiplier=1.0, + ): + self.optim_kwargs = { + "lr": lr, + "weight_decay": weight_decay, + "apply_decay_on_norm_and_bias": apply_decay_on_norm_and_bias, + "backbone_lr_multiplier": backbone_lr_multiplier, + "vision_lr_multiplier": vision_lr_multiplier, + } + return [ + { + "params": [self.proj.weight, self.proj.bias], + "lr": lr * backbone_lr_multiplier, + "weight_decay": weight_decay, + "name": "backbone_decay", + }, + { + "params": [self.action_scale], + "lr": lr, + "weight_decay": 0.0, + "name": "action_no_decay", + }, + { + "params": [self.vision_scale], + "lr": lr * backbone_lr_multiplier * vision_lr_multiplier, + "weight_decay": 0.0, + "name": "vision_no_decay", + }, + ] + + def _features(): return { OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(7,)), @@ -436,6 +480,18 @@ def test_author_action_payload_fills_required_tokenizer_metadata(): } +def test_system2_training_target_is_forwarded_without_replacing_operator_task(): + config = _config(predict_cot=True, runtime_system="system2") + policy = G05Policy(config, backend=TinyG05Backend()) + batch = _policy_batch(" operator task\n") + batch["atomic_task"] = ["grasp the cup"] + + prepared = policy._prepare_author_batch(batch) + + assert prepared["samples"][0]["command"] == " operator task\n" + assert prepared["samples"][0]["atomic_task"] == "Subtask: grasp the cup" + + def test_author_inference_payload_synthesizes_required_dummy_action(): policy = G05Policy(_config(), backend=TinyG05Backend()) batch = _policy_batch() @@ -500,7 +556,7 @@ def test_batch_two_preserves_each_raw_task_and_every_camera_slot(): def test_forward_backward_update_and_save_reload(tmp_path: Path): policy = G05Policy(_config(), backend=TinyG05Backend()) - optimizer = torch.optim.AdamW(policy.get_optim_params()["params"], lr=1e-3) + optimizer = torch.optim.AdamW(policy.get_optim_params(), lr=1e-3) loss, metrics = policy(_policy_batch("train")) loss.backward() grad_norm = torch.stack( @@ -550,7 +606,7 @@ def test_save_pretrained_copies_required_gated_sidecars_portably(tmp_path: Path) def test_tiny_fixed_batch_overfit_reduces_loss(): policy = G05Policy(_config(), backend=TinyG05Backend()) - optimizer = torch.optim.AdamW(policy.get_optim_params()["params"], lr=5e-2) + optimizer = torch.optim.AdamW(policy.get_optim_params(), lr=5e-2) batch = _policy_batch("overfit") initial = policy(batch)[0].item() for _ in range(20): @@ -562,6 +618,34 @@ def test_tiny_fixed_batch_overfit_reduces_loss(): assert final < initial * 0.25 +def test_training_preset_uses_author_optimizer_parameter_groups(): + config = _config( + optimizer_lr=2e-4, + optimizer_weight_decay=0.03, + optimizer_backbone_lr_multiplier=0.5, + optimizer_vision_lr_multiplier=0.2, + optimizer_apply_decay_on_norm_and_bias=True, + ) + backend = GroupedTinyG05Backend() + policy = G05Policy(config, backend=backend) + + optimizer = config.get_optimizer_preset().build(policy.get_optim_params()) + + assert backend.optim_kwargs == { + "lr": 2e-4, + "weight_decay": 0.03, + "apply_decay_on_norm_and_bias": True, + "backbone_lr_multiplier": 0.5, + "vision_lr_multiplier": 0.2, + } + assert [group["name"] for group in optimizer.param_groups] == [ + "backbone_decay", + "action_no_decay", + "vision_no_decay", + ] + assert [group["lr"] for group in optimizer.param_groups] == pytest.approx([1e-4, 2e-4, 2e-5]) + + @pytest.mark.skipif( not os.environ.get("LEROBOT_G05_CHECKPOINT"), reason="requires an accepted gated OpenGalaxea/G05 checkpoint and author CUDA environment",