synth-ai 0.2.16__py3-none-any.whl → 0.2.19__py3-none-any.whl

examples/blog_posts/pokemon_vl/pokemon_vl_baseline.py

@@ -0,0 +1,326 @@
+"""Pokemon Red baseline file for Game Boy emulation evaluation."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+from synth_ai.baseline import BaselineConfig, BaselineTaskRunner, DataSplit, TaskResult
+from synth_ai.inference import InferenceClient
+import os
+import httpx
+
+try:
+    from synth_ai.environments.examples.red.environment import PokemonRedEnvironment
+    from synth_ai.environments.examples.red.taskset import (
+        PokemonRedTaskInstance,
+        PokemonRedTaskInstanceMetadata,
+    )
+    POKEMON_RED_AVAILABLE = True
+except ImportError:
+    POKEMON_RED_AVAILABLE = False
+
+
+class PokemonRedTaskRunner(BaselineTaskRunner):
+    """Task runner for Pokemon Red Game Boy emulation."""
+    
+    def __init__(self, policy_config: Dict[str, Any], env_config: Dict[str, Any]):
+        super().__init__(policy_config, env_config)
+        
+        if not POKEMON_RED_AVAILABLE:
+            raise ImportError(
+                "Pokemon Red environment not available. "
+                "Install synth-ai with Pokemon Red support."
+            )
+        
+        # Store config for inference
+        self.model = policy_config["model"]
+        self.temperature = policy_config.get("temperature", 0.0)
+        self.max_tokens = policy_config.get("max_tokens", 512)
+        self.inference_url = policy_config.get("inference_url")
+        
+        # Tool definition
+        self.tools = [{
+            "type": "function",
+            "function": {
+                "name": "execute_sequence",
+                "description": "Execute multiple button presses in sequence",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "actions": {
+                            "type": "array",
+                            "items": {
+                                "type": "object",
+                                "properties": {
+                                    "button": {
+                                        "type": "string",
+                                        "enum": ["UP", "DOWN", "LEFT", "RIGHT", "A", "B", "START", "SELECT"],
+                                    },
+                                    "frames": {
+                                        "type": "integer",
+                                        "minimum": 1,
+                                        "maximum": 120,
+                                        "description": "Frames to hold button (60fps)",
+                                    },
+                                },
+                                "required": ["button", "frames"],
+                            },
+                            "minItems": 1,
+                            "maxItems": 20,
+                        },
+                    },
+                    "required": ["actions"],
+                },
+            },
+        }]
+    
+    def _format_observation(self, obs: Dict[str, Any], step: int, max_steps: int) -> str:
+        """Format observation for LLM."""
+        lines = [
+            f"Pokemon Red - Step {step}/{max_steps}",
+            "",
+        ]
+        
+        # Position
+        if "map_id" in obs:
+            lines.append(f"Location: Map {obs['map_id']}")
+        if "player_x" in obs and "player_y" in obs:
+            lines.append(f"Position: ({obs['player_x']}, {obs['player_y']})")
+        
+        # Party
+        if "party_count" in obs:
+            lines.append(f"Party Size: {obs['party_count']}")
+        if "party_pokemon" in obs and obs["party_pokemon"]:
+            pokemon = obs["party_pokemon"][0]
+            lines.append(
+                f"First Pokemon: Level {pokemon.get('level', '?')}, "
+                f"HP {pokemon.get('hp_current', '?')}/{pokemon.get('hp_max', '?')}"
+            )
+        
+        # Battle
+        if obs.get("in_battle"):
+            lines.append("=== IN BATTLE ===")
+            if "enemy_hp_current" in obs:
+                lines.append(
+                    f"Enemy HP: {obs['enemy_hp_current']}/{obs.get('enemy_hp_max', '?')}"
+                )
+            if "battle_turn" in obs:
+                lines.append(f"Battle Turn: {obs['battle_turn']}")
+        
+        # Progress
+        if "badges" in obs:
+            lines.append(f"Badges: {obs['badges']}")
+        if "money" in obs:
+            lines.append(f"Money: ${obs['money']}")
+        
+        # Dialogue
+        if obs.get("text_box_active"):
+            lines.append("Text box is active - press A to advance dialogue")
+        
+        lines.append("")
+        lines.append("What actions should we take?")
+        
+        return "\n".join(lines)
+    
+    async def run_task(self, seed: int) -> TaskResult:
+        """Run a single Pokemon Red episode."""
+        
+        # Create task instance
+        rom_path = self.env_config.get("rom_path")
+        if not rom_path:
+            raise ValueError("rom_path required in env_config for Pokemon Red")
+        
+        init_state_path = self.env_config.get("init_state_path")
+        max_steps = self.env_config.get("max_steps", 500)
+        
+        metadata = PokemonRedTaskInstanceMetadata(
+            seed=seed,
+            rom_path=rom_path,
+            init_state_path=init_state_path,
+            reward_type=self.env_config.get("reward_type", "pallet_town_progression"),
+        )
+        
+        task_instance = PokemonRedTaskInstance(
+            id=f"pokemon-red-{seed}",
+            metadata=metadata,
+        )
+        
+        # Create environment
+        env = PokemonRedEnvironment(task_instance=task_instance)
+        
+        # Initialize environment
+        raw_obs = await env.initialize()
+        observation = getattr(raw_obs, "observation", raw_obs) if hasattr(raw_obs, "observation") else raw_obs
+        obs_dict = observation if isinstance(observation, dict) else {}
+        
+        # Episode loop
+        total_reward = 0.0
+        total_steps = 0
+        event_rewards: List[Dict[str, Any]] = []
+        battle_won = False
+        game_over = False
+        
+        for step in range(max_steps):
+            # Format observation
+            prompt = self._format_observation(obs_dict, step, max_steps)
+            
+            # Add image if available
+            messages = [{"role": "user", "content": prompt}]
+            if obs_dict.get("observation_image_base64"):
+                messages[0]["content"] = [
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": f"data:image/png;base64,{obs_dict['observation_image_base64']}"
+                        },
+                    },
+                    {"type": "text", "text": prompt},
+                ]
+            
+            # Get action from LLM
+            if self.inference_url and self.inference_url.startswith("http"):
+                api_key = os.getenv("SYNTH_API_KEY") or os.getenv("OPENAI_API_KEY") or ""
+                base_url = self.inference_url.rstrip("/")
+                if not base_url.endswith("/api"):
+                    base_url = f"{base_url}/api" if "/api" not in base_url else base_url
+                client = InferenceClient(base_url=base_url, api_key=api_key)
+                response = await client.create_chat_completion(
+                    model=self.model,
+                    messages=messages,
+                    tools=self.tools,
+                    tool_choice={"type": "function", "function": {"name": "execute_sequence"}},
+                    temperature=self.temperature,
+                    max_tokens=self.max_tokens,
+                )
+            else:
+                api_key = os.getenv("OPENAI_API_KEY") or os.getenv("GROQ_API_KEY") or ""
+                base_url = "https://api.openai.com/v1" if "openai" in self.model.lower() else "https://api.groq.com/openai/v1"
+                async with httpx.AsyncClient() as http_client:
+                    resp = await http_client.post(
+                        f"{base_url}/chat/completions",
+                        json={
+                            "model": self.model,
+                            "messages": messages,
+                            "tools": self.tools,
+                            "tool_choice": {"type": "function", "function": {"name": "execute_sequence"}},
+                            "temperature": self.temperature,
+                            "max_tokens": self.max_tokens,
+                        },
+                        headers={"Authorization": f"Bearer {api_key}"} if api_key else {},
+                    )
+                    response = resp.json()
+            
+            # Extract actions
+            actions = []
+            tool_calls = []
+            if "choices" in response and len(response["choices"]) > 0:
+                message = response["choices"][0].get("message", {})
+                tool_calls = message.get("tool_calls", [])
+            elif "tool_calls" in response:
+                tool_calls = response["tool_calls"]
+            
+            if tool_calls:
+                tool_call = tool_calls[0]
+                actions = tool_call["function"]["arguments"].get("actions", [])
+            
+            if not actions:
+                break
+            
+            # Execute actions
+            for action_spec in actions:
+                if total_steps >= max_steps:
+                    break
+                
+                # Convert to tool call format
+                from synth_ai.environments.environment.tools import EnvToolCall
+                
+                tool_call = EnvToolCall(
+                    name="execute_sequence",
+                    arguments={"actions": [action_spec]},
+                )
+                
+                # Step environment
+                step_result = await env.step([tool_call])
+                total_steps += 1
+                
+                # Get observation
+                step_obs = (
+                    getattr(step_result, "observation", step_result)
+                    if hasattr(step_result, "observation")
+                    else step_result
+                )
+                obs_dict = step_obs if isinstance(step_obs, dict) else {}
+                
+                # Extract reward
+                reward = getattr(step_result, "reward", 0.0)
+                total_reward += reward
+                
+                if reward > 0:
+                    event_rewards.append({
+                        "step": total_steps,
+                        "reward": reward,
+                    })
+                
+                # Check termination
+                if getattr(step_result, "terminated", False) or getattr(step_result, "truncated", False):
+                    game_over = True
+                    break
+                
+                # Check battle outcome
+                if obs_dict.get("battle_outcome") == 1:
+                    battle_won = True
+                elif obs_dict.get("battle_outcome") == 2:
+                    game_over = True
+            
+            if game_over:
+                break
+        
+        # Cleanup
+        await env.terminate()
+        
+        return TaskResult(
+            seed=seed,
+            success=True,
+            outcome_reward=total_reward,
+            event_rewards=event_rewards,
+            total_steps=total_steps,
+            metadata={
+                "battle_won": battle_won,
+                "game_over": game_over,
+                "final_map": obs_dict.get("map_id"),
+                "badges": obs_dict.get("badges", 0),
+                "party_size": obs_dict.get("party_count", 0),
+            },
+        )
+
+
+# Define baseline config (only if Pokemon Red is available)
+if POKEMON_RED_AVAILABLE:
+    pokemon_vl_baseline = BaselineConfig(
+        baseline_id="pokemon_vl",
+        name="Pokemon VL - Pokemon Red",
+        description="Pokemon Red Game Boy emulation baseline for vision-language agents",
+        task_runner=PokemonRedTaskRunner,
+        splits={
+            "train": DataSplit(name="train", seeds=list(range(20))),
+            "val": DataSplit(name="val", seeds=list(range(20, 25))),
+            "test": DataSplit(name="test", seeds=list(range(25, 30))),
+        },
+        default_policy_config={
+            "model": "groq:llama-3.1-70b-versatile",
+            "temperature": 0.0,
+            "max_tokens": 512,
+        },
+        default_env_config={
+            "rom_path": None,  # Must be provided
+            "init_state_path": None,  # Optional
+            "reward_type": "pallet_town_progression",
+            "max_steps": 500,
+        },
+        metadata={
+            "environment": "pokemon_red",
+            "task_type": "emulation",
+            "requires_rom": True,
+        },
+    )
+

examples/blog_posts/pokemon_vl/run_eval_extract_images.py

@@ -0,0 +1,209 @@
+#!/usr/bin/env python3
+"""Run pokemon_vl eval with gpt-5-nano and extract images from trajectory response.
+
+This script bypasses the trace validation issue by extracting images directly from
+the trajectory steps in the rollout response.
+"""
+
+import argparse
+import asyncio
+import base64
+import json
+import os
+from pathlib import Path
+
+import httpx
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+async def run_eval_and_extract_images(
+    task_app_url: str,
+    output_dir: Path,
+    seed: int = 0,
+    max_turns: int = 10,
+    model: str = "gpt-5-nano",
+):
+    """Run eval and extract images from trajectory."""
+    output_dir.mkdir(parents=True, exist_ok=True)
+    
+    async with httpx.AsyncClient(timeout=300.0) as client:
+        # Build rollout request
+        rollout_request = {
+            "run_id": f"gpt5nano_eval_seed_{seed}",
+            "env": {
+                "env_name": "pokemon_red",
+                "seed": seed,
+                "config": {
+                    "split": "train",
+                    "index": seed,
+                    "env_params": {"max_steps_per_episode": 100},
+                },
+            },
+            "policy": {
+                "policy_name": "pokemon_vl_qwen3_vl",
+                "config": {
+                    "model": model,
+                    "provider": "openai",
+                    "inference_url": "https://api.openai.com/v1",
+                    "temperature": 0.7,
+                    "top_p": 0.95,
+                    "max_tokens": 512,
+                    "use_vision": True,
+                    "image_only_mode": False,
+                    "max_llm_calls": max_turns,
+                },
+            },
+            "ops": ["policy"] * max_turns,
+            "mode": "eval",
+            "record": {
+                "return_trace": True,
+                "trace_format": "full",
+            },
+        }
+        
+        print(f"Running eval with gpt-5-nano (seed={seed})...")
+        response = await client.post(f"{task_app_url}/rollout", json=rollout_request)
+        response.raise_for_status()
+        result = response.json()
+        
+        # Extract trajectory
+        trajectories = result.get("trajectories", [])
+        if not trajectories:
+            print("Error: No trajectories in response")
+            return
+        
+        trajectory = trajectories[0]
+        steps = trajectory.get("steps", [])
+        
+        print(f"✓ Received {len(steps)} steps")
+        print(f"Extracting images (filtering intermediate text box frames)...")
+        
+        # First pass: collect all images with their state
+        image_data = []
+        for idx, step in enumerate(steps):
+            obs = step.get("obs", {})
+            img_b64 = obs.get("observation_image_base64")
+            
+            if not img_b64:
+                continue
+            
+            try:
+                img_data = base64.b64decode(img_b64)
+                map_id = obs.get("map_id", "?")
+                player_x = obs.get("player_x", "?")
+                player_y = obs.get("player_y", "?")
+                text_box_active = obs.get("text_box_active", False)
+                
+                image_data.append({
+                    "idx": idx,
+                    "img_data": img_data,
+                    "map_id": map_id,
+                    "player_x": player_x,
+                    "player_y": player_y,
+                    "text_box_active": text_box_active,
+                })
+            except Exception as e:
+                print(f"  Error decoding step {idx}: {e}")
+                continue
+        
+        # Second pass: filter out intermediate text box frames
+        # Keep: text_box_active=False OR the last frame of a text box sequence
+        filtered_images = []
+        for i, img_info in enumerate(image_data):
+            text_box_active = img_info["text_box_active"]
+            prev_text_box_active = image_data[i - 1]["text_box_active"] if i > 0 else False
+            next_text_box_active = image_data[i + 1]["text_box_active"] if i + 1 < len(image_data) else False
+            
+            # Keep if:
+            # 1. Not in a text box (text_box_active=False)
+            # 2. Last frame of text box sequence (text_box_active=True and next is False)
+            # 3. Last frame overall and in text box (no next frame)
+            if not text_box_active:
+                # Always keep non-text-box frames
+                filtered_images.append(img_info)
+            elif text_box_active and (not next_text_box_active or i + 1 >= len(image_data)):
+                # Keep final frame of text box sequence (transition out or end of trajectory)
+                filtered_images.append(img_info)
+            # Otherwise skip intermediate text box loading frames
+        
+        # Save filtered images
+        image_count = 0
+        for img_info in filtered_images:
+            try:
+                map_id = img_info["map_id"]
+                player_x = img_info["player_x"]
+                player_y = img_info["player_y"]
+                text_box_active = img_info["text_box_active"]
+                idx = img_info["idx"]
+                
+                pos_str = f"Map{map_id}_{player_x},{player_y}"
+                textbox_str = "True" if text_box_active else "False"
+                filename = f"step_{idx:03d}_pos_{pos_str}_textbox_{textbox_str}.png"
+                
+                filepath = output_dir / filename
+                filepath.write_bytes(img_info["img_data"])
+                
+                print(f"  Saved: {filename}")
+                image_count += 1
+            except Exception as e:
+                print(f"  Error saving step {img_info['idx']}: {e}")
+                continue
+        
+        print(f"\n  Filtered: {len(image_data)} -> {len(filtered_images)} images (removed {len(image_data) - len(filtered_images)} intermediate text box frames)")
+        
+        print(f"\n✓ Extracted {image_count} images to {output_dir}/")
+        
+        # Also save metrics
+        metrics = result.get("metrics", {})
+        if metrics:
+            metrics_file = output_dir / "metrics.json"
+            with open(metrics_file, "w") as f:
+                json.dump(metrics, f, indent=2)
+            print(f"✓ Saved metrics to {metrics_file}")
+
+
+async def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--task-app-url",
+        default="http://127.0.0.1:8914",
+        help="Task app URL",
+    )
+    parser.add_argument(
+        "--output-dir",
+        default="examples/blog_posts/pokemon_vl/images_gpt5",
+        help="Output directory for images",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=0,
+        help="Random seed",
+    )
+    parser.add_argument(
+        "--max-turns",
+        type=int,
+        default=10,
+        help="Maximum turns",
+    )
+    parser.add_argument(
+        "--model",
+        default="gpt-5-nano",
+        help="Model name",
+    )
+    args = parser.parse_args()
+    
+    await run_eval_and_extract_images(
+        args.task_app_url,
+        Path(args.output_dir),
+        args.seed,
+        args.max_turns,
+        args.model,
+    )
+
+
+if __name__ == "__main__":
+    asyncio.run(main())
+