fix: increase max_new_tokens to 2048 and make configurable via GRPOConfig (#62)

* fix: align GRPO prompt format with SFT training format

The GRPO rollout prompt was missing the "Thought:" line and action
history that the SFT training uses. Models fine-tuned via SFT output
"Thought: ...\nAction: CLICK(...)" but the GRPO prompt didn't
prompt for this format, causing verbose free-form output that
couldn't be parsed → reward 0.0 → zero gradients.

Changes:
- Add "Thought:" and "Action:" prompt lines matching SFT format
- Add action_history parameter for step context
- Parser extracts action from "Action: ..." line before regex matching
- Parser handles JSON format {"action_type": "click", "coordinate": [x,y]}
- Debug logging of raw VLM output for zero-reward diagnosis

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

* fix: increase max_new_tokens to 2048 and make configurable

The default of 100 tokens truncated reasoning models mid-thought,
producing unparseable output → DONE → reward 0.0 → zero gradients.
Caused 4 failed training runs (~20 GPU-hours wasted).

- Add max_new_tokens to GRPOConfig (default 2048)
- Use config value instead of hardcoded 100
- Add truncation warning when generation hits the limit

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

---------

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

Author: abrichr

openadapt_ml/training/grpo/config.py

@@ -86,5 +86,12 @@ class GRPOConfig:
     save_every_steps: int = 50
     output_dir: str = "checkpoints/grpo"
 
+    # Generation
+    max_new_tokens: int = 2048  # Token budget per step. Reasoning models need
+    # 1000+ tokens (thought + action). 100 truncates mid-reasoning → unparseable.
+
+    # Task configs
+    task_dir: str | None = None  # Directory of TaskConfig YAMLs for milestone rewards
+
     # Stuck detection
     stuck_window: int = 3

openadapt_ml/training/grpo/trainer.py

@@ -500,7 +500,7 @@ def agent_fn(obs: Any) -> BenchmarkAction:
             with torch.no_grad():
                 outputs = model.generate(
                     **inputs,
-                    max_new_tokens=100,
+                    max_new_tokens=self._config.max_new_tokens,
                     temperature=temperature,
                     do_sample=True,
                 )
@@ -517,6 +517,15 @@ def agent_fn(obs: Any) -> BenchmarkAction:
                 except Exception:
                     pass
 
+            # Warn if output was likely truncated (hit max_new_tokens)
+            gen_len = outputs[0].shape[0] - inputs["input_ids"].shape[1]
+            if gen_len >= self._config.max_new_tokens - 1:
+                logger.warning(
+                    "Generation hit max_new_tokens=%d — output may be truncated. "
+                    "Increase config.max_new_tokens if actions aren't parsed.",
+                    self._config.max_new_tokens,
+                )
+
             action = _parse_vlm_output_to_action(decoded, screen_size=screen_size)
 
             # Store raw VLM output for accurate loss computation