From 65a8b404841cfa483164512ded68d7c955fbe503 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Thu, 30 Jul 2026 12:28:40 +0200 Subject: [PATCH] fix(g05): honor autoregressive loss weights --- docs/source/g05.mdx | 3 + src/lerobot/policies/g05/native_g05.py | 50 +++++++++++++-- tests/policies/g05/test_g05.py | 84 ++++++++++++++++++++++++++ 3 files changed, 131 insertions(+), 6 deletions(-) diff --git a/docs/source/g05.mdx b/docs/source/g05.mdx index 1ae23de41..ec03dd9ff 100644 --- a/docs/source/g05.mdx +++ b/docs/source/g05.mdx @@ -151,6 +151,9 @@ G0.5 implements LeRobot's training surface natively: `forward` computes assistant-token cross entropy and flow-matching loss, the policy exposes VLM/vision/action optimizer groups, and the checkpoint can be saved, resumed, and loaded by the normal LeRobot scripts. +The objective preserves the packaged `ar.ce_weight`, optional +`ar.ce_z_loss_scale`, and `fm.fm_weight`; the autoregressive loss jointly covers +CoT text and ActionCodec targets. For example, fine-tune the private SO-101 checkpoint on a LeRobot dataset: ### Training Command Example diff --git a/src/lerobot/policies/g05/native_g05.py b/src/lerobot/policies/g05/native_g05.py index 7a956ff5e..dfed2f23c 100644 --- a/src/lerobot/policies/g05/native_g05.py +++ b/src/lerobot/policies/g05/native_g05.py @@ -50,6 +50,34 @@ from .processing_g05 import IGNORE_INDEX, G05SequenceBatch, G05Tokenizer, G05Tok G05_RUNTIME_PREDICT_COT = "g05_runtime_predict_cot" +def _autoregressive_ce_loss( + logits: Tensor, + labels: Tensor, + *, + ce_weight: float, + z_loss_scale: float, +) -> Tensor: + """Apply G0.5's checkpoint-configured autoregressive objective.""" + + if ce_weight < 0: + raise ValueError("G0.5 ar.ce_weight must be non-negative.") + if z_loss_scale < 0: + raise ValueError("G0.5 ar.ce_z_loss_scale must be non-negative.") + if logits.ndim != 2 or labels.ndim != 1 or logits.shape[0] != labels.shape[0]: + raise ValueError( + "G0.5 autoregressive CE expects logits [N,V] and labels [N], " + f"got {tuple(logits.shape)} and {tuple(labels.shape)}." + ) + if logits.shape[0] == 0 or ce_weight == 0: + return logits.sum() * 0 + + token_loss = functional.cross_entropy(logits, labels, reduction="none") + if z_loss_scale: + log_z = torch.logsumexp(logits.float(), dim=-1) + token_loss = token_loss.float() + z_loss_scale * log_z.square() + return token_loss.mean() * ce_weight + + def _qwen_text_config(values: Mapping[str, Any], *, vocab_size: int | None = None): """Translate the serialized G0.5 Qwen config into a Transformers config.""" @@ -1171,6 +1199,13 @@ class G05NativeBackend(nn.Module): hidden_states, cache, positions = self._prefill(sequence, pixel_values, proprio) loss_dict: dict[str, Tensor] = {} + ar_config = self.model_config.get("ar") or {} + ce_weight = float(ar_config.get("ce_weight", 1.0)) + z_loss_scale = float(ar_config.get("ce_z_loss_scale", 0.0)) + if ce_weight < 0: + raise ValueError("G0.5 ar.ce_weight must be non-negative.") + if z_loss_scale < 0: + raise ValueError("G0.5 ar.ce_z_loss_scale must be non-negative.") skip_ce = ( bool(self.model_config.get("continuous_action", False)) and not bool(self.model_config.get("discrete_action", False)) @@ -1179,12 +1214,15 @@ class G05NativeBackend(nn.Module): if not skip_ce: shift_labels = sequence.labels[:, 1:] valid = shift_labels != IGNORE_INDEX - if valid.any(): - logits = self.model.vlm.logits(hidden_states[:, :-1]) - loss_dict["ce_loss"] = functional.cross_entropy( - logits.reshape(-1, logits.shape[-1]), - shift_labels.reshape(-1), - ignore_index=IGNORE_INDEX, + if valid.any() and ce_weight: + shift_hidden = hidden_states[:, :-1].reshape(-1, hidden_states.shape[-1]) + valid_hidden = shift_hidden[valid.reshape(-1)] + valid_labels = shift_labels.reshape(-1)[valid.reshape(-1)] + loss_dict["ce_loss"] = _autoregressive_ce_loss( + self.model.vlm.logits(valid_hidden), + valid_labels, + ce_weight=ce_weight, + z_loss_scale=z_loss_scale, ) else: loss_dict["ce_loss"] = hidden_states.sum() * 0 diff --git a/tests/policies/g05/test_g05.py b/tests/policies/g05/test_g05.py index fc323f369..08a4a2b07 100644 --- a/tests/policies/g05/test_g05.py +++ b/tests/policies/g05/test_g05.py @@ -105,6 +105,55 @@ class GroupedTinyG05Backend(TinyG05Backend): ] +class TinyLanguageTrainingBackend(G05NativeBackend): + def __init__(self, *, ce_weight: float, z_loss_scale: float = 0.0): + nn.Module.__init__(self) + self.model_config = { + "ar": {"ce_weight": ce_weight, "ce_z_loss_scale": z_loss_scale}, + "continuous_action": False, + "discrete_action": True, + "predict_cot": True, + } + self.head = nn.Linear(2, 3, bias=False) + self.head.weight.data.copy_( + torch.tensor( + [ + [1.0, -0.5], + [-0.25, 0.75], + [0.5, 0.25], + ] + ) + ) + self.model = SimpleNamespace(vlm=SimpleNamespace(logits=self.head)) + self.hidden = nn.Parameter( + torch.tensor( + [ + [0.5, -0.5], + [1.0, 0.25], + [-0.25, 0.75], + ] + ) + ) + self.processor = SimpleNamespace( + encode_train=lambda samples, device, action_codec: SimpleNamespace( + labels=torch.tensor([[-100, 0, 2]], device=device), + token_types=torch.zeros(1, 3, device=device), + split_index=3, + ) + ) + self.action_tokenizer = None + + def _proprio(self, samples, device): + return torch.zeros(len(samples), 1, 2, device=device) + + def _prefill(self, sequence, pixel_values, proprio): + return ( + self.hidden.unsqueeze(0), + object(), + torch.zeros(3, 1, 3, dtype=torch.long), + ) + + def _features(): return { OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(7,)), @@ -500,6 +549,41 @@ def test_native_backend_uses_per_call_cot_gate_instead_of_checkpoint_default(): assert system2["cot_text"] == ["Subtask: pick"] +def test_native_training_applies_ar_loss_config_and_reaches_language_head(): + ce_weight = 0.25 + z_loss_scale = 0.2 + backend = TinyLanguageTrainingBackend(ce_weight=ce_weight, z_loss_scale=z_loss_scale) + batch = { + "samples": [{}], + "pixel_values": {"camera": torch.zeros(1, 1, 3, 2, 2)}, + } + + loss, metrics = backend(batch) + + logits = backend.head(backend.hidden[:2]) + labels = torch.tensor([0, 2]) + expected = ( + ce_weight + * ( + torch.nn.functional.cross_entropy(logits, labels, reduction="none") + + z_loss_scale * torch.logsumexp(logits, dim=-1).square() + ).mean() + ) + torch.testing.assert_close(loss, expected) + torch.testing.assert_close(metrics["ce_loss"], expected) + loss.backward() + language_head_grad = backend.head.weight.grad + assert language_head_grad is not None + assert torch.isfinite(language_head_grad).all() + assert language_head_grad.abs().sum() > 0 + + disabled = TinyLanguageTrainingBackend(ce_weight=0.0) + disabled_loss, disabled_metrics = disabled(batch) + assert disabled_loss.requires_grad + torch.testing.assert_close(disabled_loss, torch.zeros_like(disabled_loss)) + torch.testing.assert_close(disabled_metrics["ce_loss"], torch.zeros_like(disabled_loss)) + + def test_author_action_payload_fills_required_tokenizer_metadata(): policy = G05Policy(_config(), backend=TinyG05Backend())