fix(g05): wire runtime system selection

docs/source/g05.mdx

@@ -94,6 +94,9 @@ lerobot-rollout \
   --mode=action
 ```
 
+`--direct_subtask` explicitly selects the action-only System 1 path, even when
+the checkpoint supports System 2 with `predict_cot=true`.
+
 System 2 is available only when the packaged checkpoint metadata has
 `predict_cot=true`. The adapter forwards the operator task byte-for-byte and
 returns CoT telemetry and the matching action chunk atomically. It never samples

src/lerobot/policies/g05/inference/g05_adapter.py

@@ -54,14 +54,15 @@ class G05PolicyAdapter(BaseLanguageAdapter):
 
     The preferred policy hook is::
 
-        predict_action_chunk_with_runtime(observation, *, task) ->
+        predict_action_chunk_with_runtime(observation, *, task, system_mode) ->
             (action_chunk, {"cot_text": ...})
 
     ``task`` is exactly :attr:`RuntimeState.task`, including whitespace and
     Unicode.  The policy must pass it to the author prompt builder unchanged
-    before applying checkpoint-specific formatting.  A structured return is
-    required for batch-safe CoT handling; ``predict_action_chunk`` remains a
-    compatibility fallback for System 1.
+    before applying checkpoint-specific formatting. ``system_mode`` selects
+    action-only System 1 or unified CoT-plus-action System 2 for that call. A
+    structured return is required for batch-safe CoT handling;
+    ``predict_action_chunk`` remains a compatibility fallback for System 1.
     """
 
     def __init__(
@@ -78,6 +79,10 @@ class G05PolicyAdapter(BaseLanguageAdapter):
     def _resolve_system_mode(self, requested: str | None) -> str:
         config = getattr(self.policy, "config", None)
         raw_mode = requested
+        if raw_mode is None and not self.gen.enable_subtask:
+            # The shared CLI maps --direct_subtask to enable_subtask=False.
+            # For G0.5 this explicitly selects the action-only System 1 path.
+            raw_mode = "system1"
         if raw_mode is None:
             raw_mode = _read_config(config, "runtime_system_mode")
         if raw_mode is None:
@@ -125,7 +130,7 @@ class G05PolicyAdapter(BaseLanguageAdapter):
 
         runtime_hook = getattr(self.policy, "predict_action_chunk_with_runtime", None)
         if callable(runtime_hook):
-            output = runtime_hook(batch, task=state.task)
+            output = runtime_hook(batch, task=state.task, system_mode=self.system_mode)
         else:
             if self.system_mode == "system2":
                 raise RuntimeError(

src/lerobot/policies/g05/modeling_g05.py

@@ -32,7 +32,7 @@ from .configuration_g05 import (
     make_g05_cot_prompt_template,
     make_g05_prompt_template,
 )
-from .native_g05 import G05NativeBackend
+from .native_g05 import G05_RUNTIME_PREDICT_COT, G05NativeBackend
 
 
 def _native_backend(config: G05Config) -> nn.Module:
@@ -395,10 +395,19 @@ class G05Policy(PreTrainedPolicy):
         }[fields]
         return True
 
-    def _prepare_author_batch(self, batch: Mapping[str, Any], task: str | None = None) -> dict[str, Any]:
+    def _prepare_author_batch(
+        self,
+        batch: Mapping[str, Any],
+        task: str | None = None,
+        *,
+        predict_cot: bool | None = None,
+    ) -> dict[str, Any]:
+        run_predict_cot = self.config.predict_cot if predict_cot is None else predict_cot
         prepare = getattr(self.backend, "prepare_lerobot_batch", None)
         if callable(prepare):
-            return prepare(batch, task=task, config=self.config)
+            prepared = prepare(batch, task=task, config=self.config)
+            prepared[G05_RUNTIME_PREDICT_COT] = run_predict_cot
+            return prepared
 
         state = batch.get(OBS_STATE)
         if not isinstance(state, Tensor):
@@ -431,12 +440,22 @@ class G05Policy(PreTrainedPolicy):
             )
 
         samples = []
+        flow_only = "<action_action" not in self.config.prompt_template
+        inference_template = (
+            self.config.prompt_template
+            if run_predict_cot
+            else make_g05_prompt_template(
+                self.config.num_prompt_images,
+                predict_cot=False,
+                flow_only=flow_only,
+            )
+        )
         for index, raw_task in enumerate(tasks):
             proprio = state[index]
             if proprio.ndim == 1:
                 proprio = proprio.unsqueeze(0)
             sample = {
-                "template": self.config.prompt_template,
+                "template": inference_template,
                 # This is the author InputPreprocessor command slot. Keep it byte-for-byte
                 # unchanged; checkpoint-specific chat formatting occurs downstream.
                 "command": raw_task,
@@ -449,7 +468,7 @@ class G05Policy(PreTrainedPolicy):
             frequency = self.config.processor_metadata.get("frequency")
             if frequency is not None:
                 sample["frequency"] = frequency
-            if self.config.predict_cot:
+            if run_predict_cot:
                 rendered_recipe = "messages" in batch
                 applied_recipe_cot = rendered_recipe and self._apply_recipe_cot(
                     sample, batch, index, batch_size
@@ -512,12 +531,27 @@ class G05Policy(PreTrainedPolicy):
         prepared = dict(batch)
         prepared["samples"] = samples
         prepared["pixel_values"] = pixel_values
+        prepared[G05_RUNTIME_PREDICT_COT] = run_predict_cot
         return prepared
 
     def _run_inference(
-        self, batch: Mapping[str, Any], *, task: str | None = None
+        self,
+        batch: Mapping[str, Any],
+        *,
+        task: str | None = None,
+        system_mode: str | None = None,
     ) -> tuple[Tensor, dict[str, Any]]:
-        prepared = self._prepare_author_batch(batch, task=task)
+        if system_mode is None:
+            system_mode = self.config.runtime_system
+        if system_mode not in {"system1", "system2"}:
+            raise ValueError("G0.5 system_mode must be 'system1' or 'system2'.")
+        if system_mode == "system2" and not self.config.predict_cot:
+            raise ValueError("G0.5 System 2 requires predict_cot=True in the packaged checkpoint.")
+        prepared = self._prepare_author_batch(
+            batch,
+            task=task,
+            predict_cot=system_mode == "system2",
+        )
         predict = getattr(self.backend, "predict_action", None)
         device = next(self.backend.parameters()).device
         with torch.autocast(
@@ -537,19 +571,22 @@ class G05Policy(PreTrainedPolicy):
             action = result.get(ACTION)
         if not isinstance(action, Tensor):
             raise ValueError(f"G0.5 {self.config.action_head} output is missing its action tensor.")
-        metadata = {
-            key: result[key]
-            for key in ("cot_text", "generated_ids", "decoded_action_tokens", "ar_absent_keys", "_timing")
-            if key in result
-        }
+        metadata_keys = ("decoded_action_tokens", "ar_absent_keys", "_timing")
+        if system_mode == "system2":
+            metadata_keys = ("cot_text", "generated_ids", *metadata_keys)
+        metadata = {key: result[key] for key in metadata_keys if key in result}
         return action, metadata
 
     def predict_action_chunk_with_runtime(
-        self, batch: dict[str, Any], *, task: str
+        self,
+        batch: dict[str, Any],
+        *,
+        task: str,
+        system_mode: str | None = None,
     ) -> tuple[Tensor, dict[str, Any]]:
-        """Return System 1 actions and same-pass System 2 telemetry atomically."""
+        """Run the selected system and return its action plus same-pass telemetry."""
 
-        return self._run_inference(batch, task=task)
+        return self._run_inference(batch, task=task, system_mode=system_mode)
 
     @torch.no_grad()
     def predict_action_chunk(self, batch: dict[str, Any], **kwargs) -> Tensor:

src/lerobot/policies/g05/native_g05.py

@@ -32,6 +32,8 @@ from lerobot.utils.constants import ACTION
 from .action_codec_g05 import G05NativeActionCodec
 from .processing_g05 import IGNORE_INDEX, G05SequenceBatch, G05Tokenizer, G05TokenType
 
+G05_RUNTIME_PREDICT_COT = "g05_runtime_predict_cot"
+
 
 def _qwen_text_config(values: Mapping[str, Any], *, vocab_size: int | None = None):
     """Translate the serialized G0.5 Qwen config into a Transformers config."""
@@ -1083,7 +1085,8 @@ class G05NativeBackend(nn.Module):
         result: dict[str, Any] = {}
         last_hidden = hidden_states[:, -1]
         token_types = sequence.token_types
-        if bool(self.model_config.get("predict_cot", False)):
+        predict_cot = bool(batch.get(G05_RUNTIME_PREDICT_COT, self.model_config.get("predict_cot", False)))
+        if predict_cot:
             generated, cache, last_hidden, token_types, positions = self._generate_text(
                 last_hidden,
                 token_types=sequence.token_types,

tests/policies/g05/test_g05.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import os
 from pathlib import Path
+from types import SimpleNamespace
 
 import pytest
 import torch
@@ -14,6 +15,7 @@ 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.modeling_g05 import G05Policy
+from lerobot.policies.g05.native_g05 import G05_RUNTIME_PREDICT_COT, G05NativeBackend
 from lerobot.processor import PolicyProcessorPipeline
 from lerobot.utils.constants import ACTION, OBS_STATE, POLICY_PREPROCESSOR_DEFAULT_NAME
 
@@ -23,9 +25,11 @@ class TinyG05Backend(nn.Module):
         super().__init__()
         self.proj = nn.Linear(20, 20)
         self.last_samples = None
+        self.last_runtime_predict_cot = None
 
     def predict_action(self, batch):
         self.last_samples = batch["samples"]
+        self.last_runtime_predict_cot = batch[G05_RUNTIME_PREDICT_COT]
         state = batch[OBS_STATE]
         if state.ndim == 2:
             state = state.unsqueeze(1)
@@ -389,8 +393,97 @@ def test_exact_raw_task_reaches_author_command_and_head_selection():
     action, metadata = policy.predict_action_chunk_with_runtime(_policy_batch(), task=raw_task)
 
     assert backend.last_samples[0]["command"] == raw_task
+    assert backend.last_runtime_predict_cot is False
     assert action.shape == (1, 4, 20)
-    assert metadata["cot_text"] == ["Subtask: move carefully"]
+    assert "cot_text" not in metadata
+
+
+def test_same_predict_cot_checkpoint_switches_prompt_and_backend_runtime_path():
+    backend = TinyG05Backend()
+    policy = G05Policy(_config(predict_cot=True, runtime_system="system2"), backend=backend)
+
+    _, system1_metadata = policy.predict_action_chunk_with_runtime(
+        _policy_batch(),
+        task="pick",
+        system_mode="system1",
+    )
+    system1_sample = backend.last_samples[0]
+    assert backend.last_runtime_predict_cot is False
+    assert "prompt" not in system1_sample
+    assert "<atomic_task_text>" not in system1_sample["template"]
+    assert "cot_text" not in system1_metadata
+
+    _, system2_metadata = policy.predict_action_chunk_with_runtime(
+        _policy_batch(),
+        task="pick",
+        system_mode="system2",
+    )
+    system2_sample = backend.last_samples[0]
+    assert backend.last_runtime_predict_cot is True
+    assert system2_sample["prompt"] == "predict subtask"
+    assert "<atomic_task_text>" in system2_sample["template"]
+    assert system2_metadata["cot_text"] == ["Subtask: move carefully"]
+
+
+def test_system1_config_disables_cot_on_predict_cot_checkpoint_without_override():
+    backend = TinyG05Backend()
+    policy = G05Policy(_config(predict_cot=True, runtime_system="system1"), backend=backend)
+
+    _, metadata = policy.predict_action_chunk_with_runtime(_policy_batch(), task="pick")
+
+    assert backend.last_runtime_predict_cot is False
+    assert "<atomic_task_text>" not in backend.last_samples[0]["template"]
+    assert "cot_text" not in metadata
+
+
+def test_native_backend_uses_per_call_cot_gate_instead_of_checkpoint_default():
+    class TinyNativeBackend(G05NativeBackend):
+        def __init__(self):
+            nn.Module.__init__(self)
+            self.model_config = {
+                "predict_cot": True,
+                "continuous_action": True,
+                "discrete_action": False,
+                "ar": {"max_new_tokens": 4},
+            }
+            self.processor = SimpleNamespace(
+                encode_inference=lambda samples, device: SimpleNamespace(
+                    token_types=torch.zeros(len(samples), 1)
+                ),
+                eov_token_id=2,
+                decode=lambda ids: "Subtask: pick",
+            )
+            self.generated = 0
+
+        def _prefill(self, sequence, pixel_values, proprio):
+            batch_size = len(proprio)
+            return (
+                torch.zeros(batch_size, 1, 4),
+                object(),
+                torch.zeros(3, batch_size, 1, dtype=torch.long),
+            )
+
+        def _generate_text(self, last_hidden, *, token_types, positions, cache, **kwargs):
+            self.generated += 1
+            generated = torch.tensor([[1, 2]] * last_hidden.shape[0])
+            return generated, cache, last_hidden, token_types, positions
+
+        def _infer_flow(self, *, token_types, **kwargs):
+            return torch.zeros(token_types.shape[0], 4, 20)
+
+    backend = TinyNativeBackend()
+    batch = {
+        "samples": [{"proprio": torch.zeros(1, 20)}],
+        "pixel_values": {"camera": torch.zeros(1, 1, 3, 8, 8)},
+    }
+
+    system1 = backend.predict_action({**batch, G05_RUNTIME_PREDICT_COT: False})
+    assert backend.generated == 0
+    assert "cot_text" not in system1
+
+    system2 = backend.predict_action({**batch, G05_RUNTIME_PREDICT_COT: True})
+    assert backend.generated == 1
+    assert system2["cot_text"] == ["Subtask: pick"]
 
 
 def test_author_action_payload_fills_required_tokenizer_metadata():

tests/runtime/test_g05_adapter.py

@@ -17,6 +17,7 @@ from types import SimpleNamespace
 import pytest
 
 from lerobot.runtime import LanguageConditionedRuntime, RuntimeState
+from lerobot.runtime.adapter import GenerationConfig
 
 
 class FakeG05Policy:
@@ -25,6 +26,7 @@ class FakeG05Policy:
             predict_cot=predict_cot,
             discrete_action=discrete_action,
             continuous_action=continuous_action,
+            runtime_system="system2" if predict_cot else "system1",
         )
         self.calls = []
 
@@ -62,8 +64,8 @@ def test_system2_surfaces_same_pass_cot_and_action():
         def __init__(self):
             super().__init__(predict_cot=True, continuous_action=True)
 
-        def predict_action_chunk_with_runtime(self, observation, *, task):
-            self.calls.append((observation, task))
+        def predict_action_chunk_with_runtime(self, observation, *, task, system_mode=None):
+            self.calls.append((observation, task, system_mode))
             return {
                 "action_chunk": ["fm0", "fm1"],
                 "cot_text": "BBox: cup [1,2,3,4]|\nSubtask: grasp the cup|Updated Memory: cup located",
@@ -78,6 +80,7 @@ def test_system2_surfaces_same_pass_cot_and_action():
     assert chunk == ["fm0", "fm1"]
     assert policy.calls[0][1] == "  clear the table  "
     assert policy.calls[0][0]["task"] == "  clear the table  "
+    assert policy.calls[0][2] == "system2"
     assert (
         state.language_context["cot_text"]
         == "BBox: cup [1,2,3,4]|\nSubtask: grasp the cup|Updated Memory: cup located"
@@ -93,7 +96,7 @@ def test_system2_accepts_batch_safe_tuple_metadata():
         def __init__(self):
             super().__init__(predict_cot=True)
 
-        def predict_action_chunk_with_runtime(self, observation, *, task):
+        def predict_action_chunk_with_runtime(self, observation, *, task, system_mode=None):
             return ("chunk", {"cot_text": ["Subtask: move left"], "plan": "first move left"})
 
     state = RuntimeState(task="move")
@@ -111,7 +114,7 @@ def test_system2_reasoning_does_not_invalidate_same_pass_action_chunk():
         def __init__(self):
             super().__init__(predict_cot=True)
 
-        def predict_action_chunk_with_runtime(self, observation, *, task):
+        def predict_action_chunk_with_runtime(self, observation, *, task, system_mode=None):
             return (["a0", "a1"], {"cot_text": "Subtask: pick cup"})
 
     executed = []
@@ -149,3 +152,26 @@ def test_system2_requires_structured_single_pass_hook():
     adapter = G05PolicyAdapter(FakeG05Policy(predict_cot=True))
     with pytest.raises(RuntimeError, match="predict_action_chunk_with_runtime"):
         adapter.select_action({}, RuntimeState(task="pick"))
+
+
+def test_direct_subtask_selects_system1_on_system2_checkpoint():
+    from lerobot.policies.g05.inference.g05_adapter import G05PolicyAdapter
+
+    class SwitchablePolicy(FakeG05Policy):
+        def __init__(self):
+            super().__init__(predict_cot=True, continuous_action=True)
+
+        def predict_action_chunk_with_runtime(self, observation, *, task, system_mode=None):
+            self.calls.append(system_mode)
+            return ("chunk", {"cot_text": "Subtask: should not be generated"})
+
+    policy = SwitchablePolicy()
+    adapter = G05PolicyAdapter(policy, GenerationConfig(enable_subtask=False))
+    state = RuntimeState(task="pick")
+
+    chunk = adapter.select_action({}, state)
+
+    assert adapter.system_mode == "system1"
+    assert policy.calls == ["system1"]
+    assert chunk == "chunk"
+    assert "cot_text" not in state.language_context