fix: constrained decoding preserves chain-of-thought reasoning (#200)

The Thought/Action format from SYSTEM_PROMPT is now enforced by the
constrained decoding regex:

  Thought: <up to 500 chars of reasoning>
  Action: CLICK(x=0.50, y=0.30)

This gives the model a reasoning budget while guaranteeing parseable
output. Prior regex had no prefix (model couldn't reason) or used
(.|\n)* (Outlines couldn't compile the DFA).

Also exposes _ACTION_RE (action-only regex) for use by the parser.

Tests updated: 30 pass (was 21).

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: abrichr

openadapt_evals/training/standalone/trainer.py

@@ -90,17 +90,21 @@ def __init__(
 
     # --- Constrained decoding -------------------------------------------
 
-    # Regex matching valid action formats.  No free-text prefix — the
-    # model MUST output an action as its very first token.  This is
-    # intentional: constrained decoding forces structured output.
-    # If the model needs chain-of-thought, disable constrained_decoding
-    # and rely on prompt instructions instead.
-    _ACTION_REGEX = (
+    # Regex matching the required Thought/Action format from SYSTEM_PROMPT.
+    # The model gets up to 500 chars of chain-of-thought reasoning, then
+    # MUST output exactly one valid action.  This preserves the model's
+    # ability to reason while guaranteeing parseable output.
+    #
+    # Format:  Thought: <reasoning>\nAction: <action>
+    _ACTION_RE = (
         r"CLICK\(x=0\.\d{1,3},\s*y=0\.\d{1,3}\)"
         r'|TYPE\(text="[^"]{0,200}"\)'
         r"|WAIT\(\)"
         r"|DONE\(\)"
     )
+    _ACTION_REGEX = (
+        r"Thought: [^\n]{1,500}\nAction: (" + _ACTION_RE + r")"
+    )
     # Sentinel: None = not yet attempted, list = success, False = failed
     _constrained_processor_cache: Any = None
 

tests/test_standalone_trainer.py

@@ -20,39 +20,65 @@
 
 
 class TestActionRegex:
-    """Verify the action format regex matches valid actions and rejects junk."""
+    """Verify the Thought/Action format regex matches valid output and rejects junk."""
 
-    regex = GRPOTrainer._ACTION_REGEX
+    full_regex = GRPOTrainer._ACTION_REGEX
+    action_regex = GRPOTrainer._ACTION_RE
+
+    # -- Full Thought/Action format tests --
 
     @pytest.mark.parametrize(
-        "action",
+        "output",
         [
-            "CLICK(x=0.50, y=0.30)",
-            "CLICK(x=0.0, y=0.0)",
-            "CLICK(x=0.999, y=0.123)",
-            'TYPE(text="hello world")',
-            'TYPE(text="")',
-            'TYPE(text="notepad")',
-            "WAIT()",
-            "DONE()",
+            "Thought: I need to click the start menu.\nAction: CLICK(x=0.50, y=0.30)",
+            "Thought: Type notepad in the search box.\nAction: TYPE(text=\"notepad\")",
+            "Thought: Wait for the UI to load.\nAction: WAIT()",
+            "Thought: The task is complete.\nAction: DONE()",
+            "Thought: Click the Chrome icon on the desktop to open Chrome.\nAction: CLICK(x=0.05, y=0.20)",
+            "Thought: x\nAction: CLICK(x=0.0, y=0.0)",
         ],
     )
-    def test_valid_actions_match(self, action: str) -> None:
-        assert re.match(self.regex, action), f"Expected match: {action!r}"
+    def test_valid_thought_action_matches(self, output: str) -> None:
+        assert re.match(self.full_regex, output), f"Expected match: {output!r}"
 
     @pytest.mark.parametrize(
         "text",
         [
+            # No Thought prefix
+            "CLICK(x=0.50, y=0.30)",
+            "Action: CLICK(x=0.50, y=0.30)",
+            # Free-text reasoning without structure
             "** Let me think about this...",
             "1. Analyze the user's goal",
             "The user wants to open Task Manager",
             "",
-            "CLICK",
-            "click(0.5, 0.3)",
+            # Missing Action line
+            "Thought: I should click here.",
+            # Wrong action format
+            "Thought: Click\nAction: click(0.5, 0.3)",
+            "Thought: Click\nAction: CLICK",
         ],
     )
     def test_invalid_text_rejected(self, text: str) -> None:
-        assert not re.match(self.regex, text), f"Should NOT match: {text!r}"
+        assert not re.match(self.full_regex, text), f"Should NOT match: {text!r}"
+
+    # -- Action-only regex tests (used by parser) --
+
+    @pytest.mark.parametrize(
+        "action",
+        [
+            "CLICK(x=0.50, y=0.30)",
+            "CLICK(x=0.0, y=0.0)",
+            "CLICK(x=0.999, y=0.123)",
+            'TYPE(text="hello world")',
+            'TYPE(text="")',
+            'TYPE(text="notepad")',
+            "WAIT()",
+            "DONE()",
+        ],
+    )
+    def test_action_only_regex_matches(self, action: str) -> None:
+        assert re.match(self.action_regex, action), f"Expected match: {action!r}"
 
 
 # ---------------------------------------------------------------------------