openadapt-ml 0.2.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

openadapt_ml/benchmarks/trace_export.py

@@ -0,0 +1,631 @@
+"""Export WAA benchmark traces as training data.
+
+This module provides functionality to filter and export successful WAA benchmark
+traces in a format suitable for VLM fine-tuning. It converts benchmark execution
+traces to the openadapt-ml Episode format.
+
+Usage:
+    # Via CLI
+    uv run python -m openadapt_ml.benchmarks.cli export-traces --status passed --output training_data/
+
+    # Via Python
+    from openadapt_ml.benchmarks.trace_export import export_traces, TraceExporter
+
+    # Export all passing traces
+    exporter = TraceExporter(
+        benchmark_dir=Path("benchmark_results/waa_eval_20241214"),
+        output_dir=Path("training_data"),
+        status_filter="passed",
+    )
+    episodes = exporter.export()
+
+    # Or use convenience function
+    episodes = export_traces(
+        benchmark_dir="benchmark_results/waa_eval_20241214",
+        output_dir="training_data",
+        status_filter="passed",
+    )
+
+Directory structure created:
+    training_data/
+    |-- episodes/
+    |   |-- episode_001.json       # Episode schema format
+    |   |-- episode_002.json
+    |   |-- ...
+    |-- screenshots/
+    |   |-- episode_001/
+    |   |   |-- step_000.png
+    |   |   |-- step_001.png
+    |   |-- episode_002/
+    |-- manifest.json              # Index of all exported episodes
+    |-- training_samples.jsonl     # JSONL format for training
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import shutil
+from dataclasses import dataclass, field
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Literal
+
+from openadapt_ml.benchmarks.viewer import (
+    load_benchmark_metadata,
+    load_benchmark_summary,
+    load_task_results,
+)
+from openadapt_ml.schema import (
+    Action,
+    ActionType,
+    BenchmarkSource,
+    Coordinates,
+    Episode,
+    Observation,
+    Step,
+    save_episode,
+)
+
+logger = logging.getLogger(__name__)
+
+
+StatusFilter = Literal["passed", "failed", "all"]
+
+
+@dataclass
+class ExportStats:
+    """Statistics from a trace export operation."""
+
+    total_tasks: int = 0
+    exported_tasks: int = 0
+    skipped_tasks: int = 0
+    total_steps: int = 0
+    exported_screenshots: int = 0
+    errors: list[str] = field(default_factory=list)
+
+
+@dataclass
+class TraceExporter:
+    """Export WAA benchmark traces as training data.
+
+    Filters and converts benchmark execution traces to Episode format,
+    copies screenshots, and creates training-ready data files.
+
+    Args:
+        benchmark_dir: Path to benchmark results directory containing metadata.json,
+            summary.json, and tasks/ subdirectory.
+        output_dir: Output directory for exported training data.
+        status_filter: Filter by task status ("passed", "failed", "all").
+        copy_screenshots: Whether to copy screenshots to output directory.
+        create_jsonl: Whether to create training_samples.jsonl file.
+        viewport_size: Default viewport size (width, height) for normalizing coordinates.
+    """
+
+    benchmark_dir: Path
+    output_dir: Path
+    status_filter: StatusFilter = "passed"
+    copy_screenshots: bool = True
+    create_jsonl: bool = True
+    viewport_size: tuple[int, int] = (1920, 1200)
+
+    def __post_init__(self):
+        self.benchmark_dir = Path(self.benchmark_dir)
+        self.output_dir = Path(self.output_dir)
+
+    def export(self) -> list[Episode]:
+        """Export traces according to configuration.
+
+        Returns:
+            List of Episode objects created from the traces.
+        """
+        # Load benchmark data
+        metadata = load_benchmark_metadata(self.benchmark_dir)
+        load_benchmark_summary(self.benchmark_dir)
+        tasks = load_task_results(self.benchmark_dir)
+
+        logger.info(
+            f"Loaded {len(tasks)} tasks from {self.benchmark_dir.name} "
+            f"(model: {metadata.get('model_id', 'unknown')})"
+        )
+
+        # Filter tasks
+        filtered_tasks = self._filter_tasks(tasks)
+        logger.info(
+            f"Filtered to {len(filtered_tasks)} tasks with status={self.status_filter}"
+        )
+
+        if not filtered_tasks:
+            logger.warning("No tasks match the filter criteria")
+            return []
+
+        # Create output directories
+        self._setup_output_dirs()
+
+        # Convert and export
+        episodes = []
+        stats = ExportStats(total_tasks=len(tasks))
+
+        for i, task in enumerate(filtered_tasks):
+            try:
+                episode = self._convert_task_to_episode(task, i, metadata)
+                episodes.append(episode)
+
+                # Save episode JSON
+                episode_path = (
+                    self.output_dir / "episodes" / f"{episode.episode_id}.json"
+                )
+                save_episode(episode, episode_path)
+
+                # Copy screenshots if enabled
+                if self.copy_screenshots:
+                    self._copy_task_screenshots(task, episode.episode_id)
+                    stats.exported_screenshots += len(task.get("screenshots", []))
+
+                stats.exported_tasks += 1
+                stats.total_steps += len(episode.steps)
+
+                logger.debug(
+                    f"Exported episode {episode.episode_id}: "
+                    f"{len(episode.steps)} steps, success={episode.success}"
+                )
+
+            except Exception as e:
+                error_msg = (
+                    f"Failed to export task {task.get('task_id', 'unknown')}: {e}"
+                )
+                logger.error(error_msg)
+                stats.errors.append(error_msg)
+                stats.skipped_tasks += 1
+
+        # Create manifest
+        self._create_manifest(episodes, metadata, stats)
+
+        # Create JSONL training file
+        if self.create_jsonl:
+            self._create_training_jsonl(episodes)
+
+        # Log summary
+        logger.info(
+            f"Export complete: {stats.exported_tasks}/{stats.total_tasks} tasks, "
+            f"{stats.total_steps} steps, {stats.exported_screenshots} screenshots"
+        )
+        if stats.errors:
+            logger.warning(f"{len(stats.errors)} errors during export")
+
+        return episodes
+
+    def _filter_tasks(self, tasks: list[dict[str, Any]]) -> list[dict[str, Any]]:
+        """Filter tasks by status.
+
+        Args:
+            tasks: List of task dictionaries from load_task_results.
+
+        Returns:
+            Filtered list of tasks.
+        """
+        if self.status_filter == "all":
+            return tasks
+
+        filtered = []
+        for task in tasks:
+            execution = task.get("execution", {})
+            success = execution.get("success", False)
+
+            if self.status_filter == "passed" and success:
+                filtered.append(task)
+            elif self.status_filter == "failed" and not success:
+                filtered.append(task)
+
+        return filtered
+
+    def _setup_output_dirs(self) -> None:
+        """Create output directory structure."""
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        (self.output_dir / "episodes").mkdir(exist_ok=True)
+        if self.copy_screenshots:
+            (self.output_dir / "screenshots").mkdir(exist_ok=True)
+
+    def _convert_task_to_episode(
+        self,
+        task: dict[str, Any],
+        index: int,
+        metadata: dict[str, Any],
+    ) -> Episode:
+        """Convert a benchmark task to Episode format.
+
+        Args:
+            task: Task dictionary from load_task_results.
+            index: Task index for episode ID generation.
+            metadata: Benchmark metadata.
+
+        Returns:
+            Episode instance.
+        """
+        definition = task.get("definition", {})
+        execution = task.get("execution", {})
+        screenshots = task.get("screenshots", [])
+        execution_steps = execution.get("steps", [])
+
+        task_id = task.get("task_id", f"task_{index:03d}")
+        episode_id = f"waa_{task_id}"
+
+        # Convert execution steps to Episode steps
+        steps = []
+        for step_idx, step_data in enumerate(execution_steps):
+            step = self._convert_step(step_data, step_idx, screenshots)
+            steps.append(step)
+
+        return Episode(
+            episode_id=episode_id,
+            task_id=task_id,
+            instruction=definition.get("instruction", ""),
+            goal=definition.get("instruction", ""),
+            steps=steps,
+            success=execution.get("success", False),
+            final_reward=execution.get("score", 0.0),
+            source=BenchmarkSource.WAA,
+            source_file=str(self.benchmark_dir / "tasks" / task_id),
+            agent_model=metadata.get("model_id", "unknown"),
+            environment="Windows 11",
+            tags=[
+                definition.get("domain", "unknown"),
+                "waa",
+                "benchmark",
+            ],
+            metadata={
+                "benchmark_name": metadata.get("benchmark_name", "waa"),
+                "run_name": metadata.get("run_name"),
+                "domain": definition.get("domain"),
+                "num_steps": execution.get("num_steps", len(steps)),
+                "total_time_seconds": execution.get("total_time_seconds"),
+                "error": execution.get("error"),
+                "reason": execution.get("reason"),
+                "evaluation_spec": definition.get("evaluation_spec"),
+            },
+        )
+
+    def _convert_step(
+        self,
+        step_data: dict[str, Any],
+        step_idx: int,
+        screenshots: list[str],
+    ) -> Step:
+        """Convert a benchmark execution step to Episode Step format.
+
+        Args:
+            step_data: Step data from execution.json.
+            step_idx: Step index.
+            screenshots: List of screenshot paths.
+
+        Returns:
+            Step instance.
+        """
+        action_data = step_data.get("action", {})
+
+        # Build observation
+        screenshot_path = None
+        if step_idx < len(screenshots):
+            screenshot_path = screenshots[step_idx]
+        elif step_data.get("screenshot_path"):
+            screenshot_path = step_data["screenshot_path"]
+
+        observation = Observation(
+            screenshot_path=screenshot_path,
+            screen_size=self.viewport_size,
+        )
+
+        # Convert action type
+        action_type = self._map_action_type(action_data.get("type", "click"))
+
+        # Build action with coordinates
+        action_kwargs: dict[str, Any] = {
+            "type": action_type,
+            "raw": action_data,
+        }
+
+        # Handle coordinates - convert to normalized if pixel values
+        x = action_data.get("x")
+        y = action_data.get("y")
+        if x is not None and y is not None:
+            # Check if already normalized (0-1 range)
+            if 0 <= x <= 1 and 0 <= y <= 1:
+                action_kwargs["normalized_coordinates"] = (x, y)
+            else:
+                # Assume pixel coordinates, normalize
+                norm_x = x / self.viewport_size[0]
+                norm_y = y / self.viewport_size[1]
+                action_kwargs["normalized_coordinates"] = (norm_x, norm_y)
+                # Also store pixel coordinates
+                action_kwargs["coordinates"] = Coordinates(x=int(x), y=int(y))
+
+        # Handle text for type action
+        if action_data.get("text"):
+            action_kwargs["text"] = action_data["text"]
+
+        # Handle key for key action
+        if action_data.get("key"):
+            action_kwargs["key"] = action_data["key"]
+
+        # Handle modifiers
+        if action_data.get("modifiers"):
+            action_kwargs["modifiers"] = action_data["modifiers"]
+
+        # Handle scroll
+        if action_data.get("scroll_direction"):
+            action_kwargs["scroll_direction"] = action_data["scroll_direction"]
+        if action_data.get("scroll_amount"):
+            action_kwargs["scroll_amount"] = int(action_data["scroll_amount"])
+
+        # Handle drag end coordinates
+        end_x = action_data.get("end_x")
+        end_y = action_data.get("end_y")
+        if end_x is not None and end_y is not None:
+            if 0 <= end_x <= 1 and 0 <= end_y <= 1:
+                action_kwargs["normalized_end"] = (end_x, end_y)
+            else:
+                norm_end_x = end_x / self.viewport_size[0]
+                norm_end_y = end_y / self.viewport_size[1]
+                action_kwargs["normalized_end"] = (norm_end_x, norm_end_y)
+                action_kwargs["end_coordinates"] = Coordinates(
+                    x=int(end_x), y=int(end_y)
+                )
+
+        # Handle element targeting
+        if action_data.get("target_node_id"):
+            from openadapt_ml.schema import UIElement
+
+            action_kwargs["element"] = UIElement(
+                element_id=action_data.get("target_node_id"),
+                role=action_data.get("target_role"),
+                name=action_data.get("target_name"),
+            )
+
+        action = Action(**action_kwargs)
+
+        return Step(
+            step_index=step_idx,
+            observation=observation,
+            action=action,
+            reasoning=step_data.get("reasoning"),
+            timestamp=step_data.get("timestamp"),
+        )
+
+    def _map_action_type(self, action_type_str: str) -> ActionType:
+        """Map benchmark action type string to ActionType enum.
+
+        Args:
+            action_type_str: Action type string from benchmark.
+
+        Returns:
+            ActionType enum value.
+        """
+        mapping = {
+            "click": ActionType.CLICK,
+            "double_click": ActionType.DOUBLE_CLICK,
+            "right_click": ActionType.RIGHT_CLICK,
+            "type": ActionType.TYPE,
+            "key": ActionType.KEY,
+            "scroll": ActionType.SCROLL,
+            "drag": ActionType.DRAG,
+            "hover": ActionType.HOVER,
+            "wait": ActionType.WAIT,
+            "done": ActionType.DONE,
+            "answer": ActionType.DONE,
+            "failed": ActionType.FAIL,
+            "fail": ActionType.FAIL,
+        }
+        return mapping.get(action_type_str.lower(), ActionType.CLICK)
+
+    def _copy_task_screenshots(self, task: dict[str, Any], episode_id: str) -> None:
+        """Copy task screenshots to output directory.
+
+        Args:
+            task: Task dictionary.
+            episode_id: Episode ID for output subdirectory.
+        """
+        screenshots = task.get("screenshots", [])
+        if not screenshots:
+            return
+
+        # Create episode screenshot directory
+        episode_screenshots_dir = self.output_dir / "screenshots" / episode_id
+        episode_screenshots_dir.mkdir(parents=True, exist_ok=True)
+
+        for i, rel_path in enumerate(screenshots):
+            src_path = self.benchmark_dir / rel_path
+            if src_path.exists():
+                dest_path = episode_screenshots_dir / f"step_{i:03d}.png"
+                shutil.copy2(src_path, dest_path)
+
+    def _create_manifest(
+        self,
+        episodes: list[Episode],
+        metadata: dict[str, Any],
+        stats: ExportStats,
+    ) -> None:
+        """Create manifest.json with export metadata.
+
+        Args:
+            episodes: List of exported episodes.
+            metadata: Benchmark metadata.
+            stats: Export statistics.
+        """
+        manifest = {
+            "export_timestamp": datetime.utcnow().isoformat(),
+            "source_benchmark": metadata.get("benchmark_name", "waa"),
+            "source_run": metadata.get("run_name"),
+            "source_model": metadata.get("model_id"),
+            "status_filter": self.status_filter,
+            "statistics": {
+                "total_tasks": stats.total_tasks,
+                "exported_tasks": stats.exported_tasks,
+                "skipped_tasks": stats.skipped_tasks,
+                "total_steps": stats.total_steps,
+                "exported_screenshots": stats.exported_screenshots,
+                "errors": len(stats.errors),
+            },
+            "episodes": [
+                {
+                    "episode_id": ep.episode_id,
+                    "task_id": ep.task_id,
+                    "instruction": ep.instruction,
+                    "num_steps": len(ep.steps),
+                    "success": ep.success,
+                    "domain": ep.metadata.get("domain") if ep.metadata else None,
+                }
+                for ep in episodes
+            ],
+        }
+
+        manifest_path = self.output_dir / "manifest.json"
+        with open(manifest_path, "w") as f:
+            json.dump(manifest, f, indent=2)
+
+        logger.info(f"Created manifest: {manifest_path}")
+
+    def _create_training_jsonl(self, episodes: list[Episode]) -> None:
+        """Create JSONL file for training.
+
+        Each line contains a training sample with:
+        - instruction: Task instruction
+        - screenshot_path: Path to screenshot
+        - action: Action taken
+        - reasoning: Optional reasoning
+
+        Args:
+            episodes: List of exported episodes.
+        """
+        jsonl_path = self.output_dir / "training_samples.jsonl"
+
+        with open(jsonl_path, "w") as f:
+            for episode in episodes:
+                for step in episode.steps:
+                    sample = {
+                        "episode_id": episode.episode_id,
+                        "task_id": episode.task_id,
+                        "instruction": episode.instruction,
+                        "step_index": step.step_index,
+                        "screenshot_path": step.observation.screenshot_path,
+                        "action_type": step.action.type.value,
+                        "action": {
+                            "type": step.action.type.value,
+                            "coordinates": (
+                                {
+                                    "x": step.action.coordinates.x,
+                                    "y": step.action.coordinates.y,
+                                }
+                                if step.action.coordinates
+                                else None
+                            ),
+                            "normalized_coordinates": step.action.normalized_coordinates,
+                            "text": step.action.text,
+                            "key": step.action.key,
+                            "modifiers": step.action.modifiers,
+                            "scroll_direction": step.action.scroll_direction,
+                            "scroll_amount": step.action.scroll_amount,
+                        },
+                        "reasoning": step.reasoning,
+                        "domain": episode.metadata.get("domain")
+                        if episode.metadata
+                        else None,
+                        "success": episode.success,
+                    }
+                    f.write(json.dumps(sample) + "\n")
+
+        logger.info(f"Created training JSONL: {jsonl_path}")
+
+
+def export_traces(
+    benchmark_dir: str | Path,
+    output_dir: str | Path,
+    status_filter: StatusFilter = "passed",
+    copy_screenshots: bool = True,
+    create_jsonl: bool = True,
+    viewport_size: tuple[int, int] = (1920, 1200),
+) -> list[Episode]:
+    """Convenience function to export benchmark traces.
+
+    Args:
+        benchmark_dir: Path to benchmark results directory.
+        output_dir: Output directory for exported training data.
+        status_filter: Filter by task status ("passed", "failed", "all").
+        copy_screenshots: Whether to copy screenshots to output directory.
+        create_jsonl: Whether to create training_samples.jsonl file.
+        viewport_size: Default viewport size for normalizing coordinates.
+
+    Returns:
+        List of Episode objects created from the traces.
+
+    Example:
+        episodes = export_traces(
+            benchmark_dir="benchmark_results/waa_eval_20241214",
+            output_dir="training_data",
+            status_filter="passed",
+        )
+        print(f"Exported {len(episodes)} episodes")
+    """
+    exporter = TraceExporter(
+        benchmark_dir=Path(benchmark_dir),
+        output_dir=Path(output_dir),
+        status_filter=status_filter,
+        copy_screenshots=copy_screenshots,
+        create_jsonl=create_jsonl,
+        viewport_size=viewport_size,
+    )
+    return exporter.export()
+
+
+def list_available_runs(
+    benchmark_results_dir: str | Path = "benchmark_results",
+) -> list[dict[str, Any]]:
+    """List available benchmark runs for export.
+
+    Args:
+        benchmark_results_dir: Base directory containing benchmark results.
+
+    Returns:
+        List of dictionaries with run information.
+    """
+    results_dir = Path(benchmark_results_dir)
+    if not results_dir.exists():
+        return []
+
+    runs = []
+    for run_dir in sorted(results_dir.iterdir()):
+        if not run_dir.is_dir():
+            continue
+
+        metadata_path = run_dir / "metadata.json"
+        summary_path = run_dir / "summary.json"
+
+        run_info = {
+            "run_name": run_dir.name,
+            "path": str(run_dir),
+        }
+
+        if metadata_path.exists():
+            with open(metadata_path) as f:
+                metadata = json.load(f)
+            run_info.update(
+                {
+                    "benchmark_name": metadata.get("benchmark_name"),
+                    "model_id": metadata.get("model_id"),
+                    "created_at": metadata.get("created_at"),
+                }
+            )
+
+        if summary_path.exists():
+            with open(summary_path) as f:
+                summary = json.load(f)
+            run_info.update(
+                {
+                    "num_tasks": summary.get("num_tasks", 0),
+                    "num_success": summary.get("num_success", 0),
+                    "success_rate": summary.get("success_rate", 0.0),
+                }
+            )
+
+        runs.append(run_info)
+
+    return runs