fix(g05): wire runtime system selection

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