fix: parse chat completions in rewards

Normalize TRL chat completion outputs before applying JSON reward parsing so GRPO can score conversational generations.

Made-with: Cursor

training/train_grpo.py

@@ -112,8 +112,26 @@ def build_dataset(n_episodes: int, seed: int = SEED) -> "datasets.Dataset":
 
 # ── Reward functions ───────────────────────────────────────────────────────────
 
+def _completion_to_text(completion) -> str:
+    """Normalize TRL string or chat-message completions into assistant text."""
+    if isinstance(completion, str):
+        return completion
+    if isinstance(completion, dict):
+        return str(completion.get("content", completion))
+    if isinstance(completion, list):
+        parts = []
+        for item in completion:
+            if isinstance(item, dict):
+                parts.append(str(item.get("content", "")))
+            else:
+                parts.append(str(item))
+        return "\n".join(part for part in parts if part)
+    return str(completion)
+
+
 def _parse_completion(text: str) -> dict | None:
     """Extract first JSON object from completion text."""
+    text = _completion_to_text(text)
     # Strip any <think>...</think> blocks (chain-of-thought models)
     text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()
     # Try direct parse first