feat: add --task-dir support for milestone-based rewards in standalone GRPO trainer (#60)

abrichr · claude · web-flow · commit 7d095dabbcf3 · 2026-03-22T12:53:11.000-04:00
* fix: include image placeholder in chat template for VLM GRPO

Qwen2.5-VL requires &lt;|image_pad|&gt; tokens in the input. These are
inserted by apply_chat_template only when messages include
{"type": "image"} content blocks.

Fixed both agent_fn and _compute_rollout_loss.

Co-Authored-By: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;

* feat: add --task-dir support for milestone-based rewards in standalone GRPO trainer

- GRPOConfig: add task_dir field
- reward.py: evaluate_milestones_screenshot() for client-side reward
- trainer.py: load TaskConfigs, auto-populate task_ids, override rewards
- rollout_collector.py: pass task_configs to env
- No WAA evaluate endpoint needed — rewards computed via VLM judge

Co-Authored-By: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;

---------

Co-authored-by: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/openadapt_ml/training/grpo/__init__.py b/openadapt_ml/training/grpo/__init__.py
@@ -47,6 +47,7 @@
 from openadapt_ml.training.grpo.reward import (
     binary_task_success,
     compute_group_advantages,
+    evaluate_milestones_screenshot,
 )
 from openadapt_ml.training.grpo.rollout_collector import (
     GRPORolloutCollector,
@@ -86,6 +87,7 @@ def __getattr__(name: str):
     "Rollout",
     "binary_task_success",
     "compute_group_advantages",
+    "evaluate_milestones_screenshot",
     "policy_gradient_loss",
     "grpo_loss",
     "parse_vlm_output_to_action",
diff --git a/openadapt_ml/training/grpo/config.py b/openadapt_ml/training/grpo/config.py
@@ -37,6 +37,11 @@ class GRPOConfig:
         server_url: URL of the WAA server for live environment interaction.
         evaluate_url: URL of the evaluate server. If None, defaults to server_url.
         task_ids: List of WAA task IDs to train on.
+        task_dir: Path to a directory of YAML task config files. When set,
+            the trainer loads TaskConfig objects and uses milestone-based
+            reward evaluation locally (no /evaluate endpoint needed).
+            If task_ids is empty, task IDs are auto-populated from the
+            loaded configs.
         learning_rate: Optimizer learning rate for LoRA parameter updates.
         num_training_steps: Total number of GRPO training steps (outer loop).
         save_every_steps: Checkpoint frequency.
@@ -69,6 +74,12 @@ class GRPOConfig:
     task_ids: list[str] = field(default_factory=list)
     screen_size: tuple[int, int] = (1920, 1080)  # (width, height)
 
+    # Task configuration directory (YAML files with milestones for dense rewards).
+    # When set, the trainer loads TaskConfig objects from this directory and
+    # uses milestone-based reward evaluation locally, without needing the
+    # WAA /evaluate endpoint.  Requires openadapt-evals to be installed.
+    task_dir: str | None = None
+
     # Training
     learning_rate: float = 5e-6
     num_training_steps: int = 1000
diff --git a/openadapt_ml/training/grpo/reward.py b/openadapt_ml/training/grpo/reward.py
@@ -6,10 +6,19 @@
 GRPO computes advantages relative to the group mean rather than using
 a learned value function, which is simpler and works well for sparse
 binary rewards (task success/failure).
+
+Also provides ``evaluate_milestones_screenshot``, a standalone utility
+that evaluates milestone-based rewards from a screenshot without needing
+the WAA /evaluate endpoint.  This is the local-evaluation path used by
+the standalone GRPO trainer when ``--task-dir`` is set.
 """
 
 from __future__ import annotations
 
+import logging
+
+logger = logging.getLogger(__name__)
+
 
 def binary_task_success(score: float, threshold: float = 0.5) -> float:
     """Convert evaluator score to binary reward.
@@ -54,3 +63,90 @@ def compute_group_advantages(rewards: list[float]) -> list[float]:
         return [0.0] * n
 
     return [(r - mean) / (std + eps) for r in rewards]
+
+
+def evaluate_milestones_screenshot(
+    task_config: object,
+    screenshot_bytes: bytes,
+    vlm_model: str = "gpt-4.1-mini",
+    vlm_provider: str = "openai",
+) -> float:
+    """Evaluate milestone-based rewards from a screenshot (no server needed).
+
+    Iterates over the milestones in a TaskConfig and evaluates each
+    ``screenshot``-type milestone using a VLM judge.  Non-screenshot
+    milestones are skipped (they require a live server).
+
+    This is a standalone utility that can be called independently of the
+    trainer, e.g.::
+
+        from openadapt_ml.training.grpo.reward import evaluate_milestones_screenshot
+        reward = evaluate_milestones_screenshot(task_config, screenshot_bytes)
+
+    Args:
+        task_config: A ``TaskConfig`` instance (from ``openadapt_evals.task_config``).
+            Must have a ``milestones`` attribute (list of ``Milestone`` objects).
+        screenshot_bytes: PNG screenshot bytes to evaluate against.
+        vlm_model: VLM model name for the judge.
+        vlm_provider: VLM provider (``"openai"`` or ``"anthropic"``).
+
+    Returns:
+        Fraction of screenshot milestones that passed (0.0 to 1.0).
+        Returns 0.0 if there are no milestones or no screenshot milestones.
+    """
+    milestones = getattr(task_config, "milestones", None)
+    if not milestones:
+        return 0.0
+
+    # Only evaluate screenshot-type milestones locally
+    screenshot_milestones = [
+        ms for ms in milestones
+        if getattr(ms.check, "check", None) == "screenshot"
+    ]
+    if not screenshot_milestones:
+        return 0.0
+
+    try:
+        from openadapt_evals.vlm_evaluator import vlm_judge
+    except ImportError:
+        logger.warning(
+            "openadapt-evals is not installed; cannot evaluate screenshot "
+            "milestones. Install with: pip install openadapt-evals"
+        )
+        return 0.0
+
+    passed = 0
+    for ms in screenshot_milestones:
+        description = getattr(ms.check, "description", None) or ""
+        if not description:
+            continue
+        try:
+            success, _confidence = vlm_judge(
+                screenshot_bytes,
+                description,
+                model=vlm_model,
+                provider=vlm_provider,
+            )
+            if success:
+                passed += 1
+            logger.debug(
+                "Milestone '%s': %s",
+                getattr(ms, "name", "?"),
+                "PASS" if success else "FAIL",
+            )
+        except Exception as exc:
+            logger.warning(
+                "Milestone '%s' evaluation failed: %s",
+                getattr(ms, "name", "?"),
+                exc,
+            )
+
+    total = len(screenshot_milestones)
+    score = passed / total if total > 0 else 0.0
+    logger.info(
+        "Milestone evaluation: %d/%d screenshot milestones passed (%.2f)",
+        passed,
+        total,
+        score,
+    )
+    return score
diff --git a/openadapt_ml/training/grpo/rollout_collector.py b/openadapt_ml/training/grpo/rollout_collector.py
@@ -68,19 +68,27 @@ class GRPORolloutCollector:
 
     Args:
         config: GRPO training configuration.
+        task_configs: Optional dict mapping task_id -> TaskConfig. When
+            provided, task configs are loaded into the RLEnvironment for
+            milestone-based dense reward evaluation.
 
     Raises:
         ImportError: If openadapt-evals is not installed.
     """
 
-    def __init__(self, config: GRPOConfig) -> None:
+    def __init__(
+        self,
+        config: GRPOConfig,
+        task_configs: dict[str, Any] | None = None,
+    ) -> None:
         if RLEnvironment is None:
             raise ImportError(
                 "openadapt-evals is required for rollout collection. "
                 "Install it with: uv add openadapt-evals"
             )
 
         self._config = config
+        self._task_configs = task_configs or {}
         self._adapter = WAALiveAdapter(
             WAALiveConfig(
                 server_url=config.server_url,
@@ -123,6 +131,11 @@ def collect_group(
 
         rollouts: list[Rollout] = []
 
+        # Load task config into the environment for dense milestone rewards
+        if task_id in self._task_configs:
+            tc = self._task_configs[task_id]
+            self._env.load_task_config(tc)
+
         for i in range(self._config.num_rollouts_per_step):
             logger.info(
                 "Collecting rollout %d/%d for task %s",
diff --git a/openadapt_ml/training/grpo/trainer.py b/openadapt_ml/training/grpo/trainer.py
@@ -45,12 +45,24 @@
 
 from openadapt_ml.datasets.next_action import SYSTEM_PROMPT
 from openadapt_ml.training.grpo.config import GRPOConfig
-from openadapt_ml.training.grpo.reward import compute_group_advantages
+from openadapt_ml.training.grpo.reward import (
+    compute_group_advantages,
+    evaluate_milestones_screenshot,
+)
 from openadapt_ml.training.grpo.rollout_collector import (
     GRPORolloutCollector,
     Rollout,
 )
 
+# Optional import for TaskConfig (openadapt-evals may not be installed)
+try:
+    from openadapt_evals.task_config import TaskConfig
+
+    _HAS_TASK_CONFIG = True
+except ImportError:
+    TaskConfig = None  # type: ignore[assignment, misc]
+    _HAS_TASK_CONFIG = False
+
 logger = logging.getLogger(__name__)
 
 DEFAULT_SCREEN_SIZE: tuple[int, int] = (1920, 1080)
@@ -301,6 +313,106 @@ def __init__(self, config: GRPOConfig) -> None:
         self._optimizer: Any = None
         self._collector: GRPORolloutCollector | None = None
         self._step: int = 0
+        self._task_configs: dict[str, Any] = {}
+
+        # Load task configs from --task-dir if specified
+        if config.task_dir:
+            self._load_task_configs(config.task_dir)
+
+    def _load_task_configs(self, task_dir: str) -> None:
+        """Load TaskConfig YAMLs from a directory.
+
+        Populates ``self._task_configs`` (keyed by task ID) and auto-fills
+        ``config.task_ids`` if it was left empty.
+
+        Args:
+            task_dir: Path to directory containing YAML/JSON task configs.
+
+        Raises:
+            ImportError: If openadapt-evals is not installed.
+            FileNotFoundError: If the directory does not exist.
+        """
+        if not _HAS_TASK_CONFIG:
+            raise ImportError(
+                "openadapt-evals is required for --task-dir support. "
+                "Install with: pip install openadapt-evals"
+            )
+
+        task_dir_path = Path(task_dir)
+        if not task_dir_path.is_dir():
+            raise FileNotFoundError(f"Task directory not found: {task_dir}")
+
+        configs = TaskConfig.from_dir(str(task_dir_path))
+        if not configs:
+            raise ValueError(f"No task configs found in {task_dir}")
+
+        for tc in configs:
+            self._task_configs[tc.id] = tc
+            logger.info(
+                "Loaded task config: %s (%s) — %d milestones",
+                tc.id,
+                tc.name[:50],
+                len(tc.milestones),
+            )
+
+        # Auto-populate task_ids if empty
+        if not self._config.task_ids:
+            self._config.task_ids = list(self._task_configs.keys())
+            logger.info(
+                "Auto-populated task_ids from task_dir: %s",
+                self._config.task_ids,
+            )
+
+    def _compute_milestone_reward(
+        self,
+        task_id: str,
+        screenshot_bytes: bytes,
+    ) -> float:
+        """Compute milestone-based reward for a task using VLM judge.
+
+        Evaluates screenshot-type milestones locally without needing the
+        WAA /evaluate endpoint.  Falls back to 0.0 if the task has no
+        milestones or the task_id is not found in loaded configs.
+
+        Args:
+            task_id: The task ID to look up in loaded configs.
+            screenshot_bytes: PNG screenshot bytes to evaluate.
+
+        Returns:
+            Fraction of screenshot milestones passed (0.0 to 1.0).
+        """
+        task_config = self._task_configs.get(task_id)
+        if task_config is None:
+            return 0.0
+        return evaluate_milestones_screenshot(task_config, screenshot_bytes)
+
+    def _compute_milestone_reward_from_rollout(
+        self,
+        rollout: Rollout,
+    ) -> float | None:
+        """Extract the last screenshot from a rollout and compute milestone reward.
+
+        Returns None if no task config or no screenshot is available,
+        signalling the caller to keep the existing reward.
+        """
+        task_config = self._task_configs.get(rollout.task_id)
+        if task_config is None or not getattr(task_config, "milestones", None):
+            return None
+
+        # Find the last step with a screenshot
+        screenshot_bytes: bytes | None = None
+        for step in reversed(rollout.steps):
+            obs = getattr(step, "observation", None)
+            if obs is not None:
+                ss = getattr(obs, "screenshot", None)
+                if ss:
+                    screenshot_bytes = ss
+                    break
+
+        if not screenshot_bytes:
+            return None
+
+        return evaluate_milestones_screenshot(task_config, screenshot_bytes)
 
     def _make_agent_fn(self) -> Callable:
         """Create agent closure: observation -> BenchmarkAction.
@@ -381,20 +493,26 @@ def train(self) -> str:
         if not self._config.task_ids:
             raise ValueError(
                 "config.task_ids must be non-empty. Provide at least one "
-                "WAA task ID to train on."
+                "WAA task ID to train on, or use --task-dir to load from "
+                "YAML files."
             )
 
         logger.info("Starting GRPO training")
         logger.info("  Model: %s", self._config.model_name)
         logger.info("  Tasks: %s", self._config.task_ids)
+        logger.info("  Task dir: %s", self._config.task_dir or "(none)")
+        logger.info("  Task configs loaded: %d", len(self._task_configs))
         logger.info("  Rollouts/step: %d", self._config.num_rollouts_per_step)
         logger.info("  Training steps: %d", self._config.num_training_steps)
 
         # Setup
         self._model, self._processor = _load_model_and_processor(self._config)
         trainable = [p for p in self._model.parameters() if p.requires_grad]
         self._optimizer = torch.optim.AdamW(trainable, lr=self._config.learning_rate)
-        self._collector = GRPORolloutCollector(self._config)
+        self._collector = GRPORolloutCollector(
+            self._config,
+            task_configs=self._task_configs if self._task_configs else None,
+        )
 
         Path(self._config.output_dir).mkdir(parents=True, exist_ok=True)
         agent_fn = self._make_agent_fn()
@@ -409,6 +527,16 @@ def train(self) -> str:
             self._model.eval()
             rollouts = self._collector.collect_group(agent_fn=agent_fn, task_id=task_id)
 
+            # If task configs with milestones are loaded, override the
+            # binary rewards with milestone-based dense rewards.
+            if self._task_configs:
+                for rollout in rollouts:
+                    milestone_reward = self._compute_milestone_reward_from_rollout(
+                        rollout
+                    )
+                    if milestone_reward is not None:
+                        rollout.reward = max(rollout.reward, milestone_reward)
+
             # Train (gradient update)
             self._model.train()
             metrics = self._training_step(rollouts)
