fix(g05): honor autoregressive loss weights

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

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

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())