feat(g05): add LeRobot training support

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

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:

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()
+            params = get_params()
-        return {"params": [parameter for parameter in self.parameters() if parameter.requires_grad]}
+            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]

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",