synth-ai 0.2.8.dev4__py3-none-any.whl → 0.2.23.dev3__py3-none-any.whl

examples/blog_posts/pokemon_vl/extract_images.py

@@ -0,0 +1,239 @@
+#!/usr/bin/env python3
+"""Extract images from pokemon_vl trace database or trace JSON file and save to images_gpt5 directory.
+
+Usage:
+    # From trace database:
+    python extract_images.py --trace-db traces/v3/pokemon_vl_gpt5nano.db
+
+    # From trace JSON file:
+    python extract_images.py --trace-json trace.json
+"""
+
+import argparse
+import base64
+import json
+import sqlite3
+from pathlib import Path
+from typing import Any
+
+from synth_ai.tracing_v3.trace_utils import load_session_trace
+
+
+def extract_image_urls_from_content(content: Any) -> list[str]:
+    """Extract image URLs from message content."""
+    urls = []
+    
+    if isinstance(content, list):
+        for part in content:
+            if isinstance(part, dict):
+                if part.get("type") == "image_url" and "image_url" in part:
+                    url = part["image_url"].get("url")
+                    if isinstance(url, str) and url.startswith("data:image"):
+                        urls.append(url)
+                elif part.get("type") == "image":
+                    img = part.get("image")
+                    if isinstance(img, str) and img.startswith("data:image"):
+                        urls.append(img)
+    elif isinstance(content, str):
+        # Check if it's a JSON string
+        try:
+            parsed = json.loads(content)
+            return extract_image_urls_from_content(parsed)
+        except:
+            pass
+    
+    return urls
+
+
+def extract_state_info_from_message(message: dict[str, Any]) -> dict[str, Any]:
+    """Extract state info from message metadata or content."""
+    metadata = message.get("metadata", {})
+    state = {}
+    
+    # Try to get state from metadata
+    if "system_state_before" in metadata:
+        state_before = metadata["system_state_before"]
+        if isinstance(state_before, dict):
+            obs = state_before.get("obs", {})
+            state.update({
+                "position": obs.get("position", "?"),
+                "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),
+            })
+    
+    # Try to extract from content text
+    content = message.get("content", "")
+    if isinstance(content, str) and "position" in content:
+        try:
+            # Look for state summary in content
+            if "State summary:" in content:
+                parts = content.split("State summary:")
+                if len(parts) > 1:
+                    import ast
+                    state_str = parts[1].split("'")[0] if "'" not in parts[1] else parts[1]
+                    try:
+                        state_dict = ast.literal_eval(state_str.split("'")[0] if "'" in state_str else state_str)
+                        if isinstance(state_dict, dict):
+                            state.update({
+                                "position": state_dict.get("position", "?"),
+                                "map_id": state_dict.get("map_id", "?"),
+                                "player_x": state_dict.get("player_x", "?"),
+                                "player_y": state_dict.get("player_y", "?"),
+                                "text_box_active": state_dict.get("text_box_active", False),
+                            })
+                    except:
+                        pass
+        except:
+            pass
+    
+    return state
+
+
+def extract_images_from_trace_dict(trace: dict[str, Any], output_dir: Path):
+    """Extract images from a trace dictionary."""
+    output_dir.mkdir(parents=True, exist_ok=True)
+    
+    # Get messages from trace
+    messages = trace.get("markov_blanket_message_history", []) or trace.get("messages", [])
+    
+    if not messages:
+        print(f"  No messages found in trace")
+        return 0
+    
+    print(f"  Found {len(messages)} messages")
+    
+    image_count = 0
+    step_idx = 0
+    for msg_idx, msg in enumerate(messages):
+        # Extract images from message content
+        content = msg.get("content", "")
+        image_urls = extract_image_urls_from_content(content)
+        
+        if not image_urls:
+            continue
+        
+        # Extract state info for filename
+        state = extract_state_info_from_message(msg)
+        
+        for img_idx, img_url in enumerate(image_urls):
+            # Extract base64 data
+            if img_url.startswith("data:image"):
+                # Format: data:image/png;base64,<data>
+                parts = img_url.split(",", 1)
+                if len(parts) != 2:
+                    continue
+                
+                b64_data = parts[1]
+                try:
+                    img_data = base64.b64decode(b64_data)
+                    
+                    # Create filename
+                    pos_str = f"{state.get('map_id', '?')}_{state.get('player_x', '?')},{state.get('player_y', '?')}"
+                    textbox_str = "True" if state.get("text_box_active") else "False"
+                    filename = f"step_{step_idx:03d}_pos_{pos_str}_textbox_{textbox_str}.png"
+                    
+                    filepath = output_dir / filename
+                    filepath.write_bytes(img_data)
+                    
+                    print(f"  Saved: {filename}")
+                    image_count += 1
+                    step_idx += 1
+                except Exception as e:
+                    print(f"  Error decoding image: {e}")
+                    continue
+    
+    return image_count
+
+
+def extract_images_from_trace_db(trace_db: str, output_dir: Path, model_filter: str | None = None):
+    """Extract images from trace database and save to output directory."""
+    conn = sqlite3.connect(trace_db)
+    conn.row_factory = sqlite3.Row
+    
+    # Get all session IDs
+    query = "SELECT session_id, metadata FROM session_traces"
+    if model_filter:
+        query += " WHERE metadata LIKE ?"
+        params = (f'%{model_filter}%',)
+    else:
+        params = ()
+    
+    rows = conn.execute(query, params).fetchall()
+    
+    if not rows:
+        print(f"No traces found in {trace_db}")
+        return
+    
+    print(f"Found {len(rows)} trace(s)")
+    
+    total_images = 0
+    for row in rows:
+        session_id = row["session_id"]
+        print(f"\nProcessing session: {session_id}")
+        
+        try:
+            trace = load_session_trace(conn, session_id)
+        except Exception as e:
+            print(f"  Error loading trace: {e}")
+            continue
+        
+        count = extract_images_from_trace_dict(trace, output_dir)
+        total_images += count
+    
+    conn.close()
+    print(f"\n✓ Extracted {total_images} images to {output_dir}/")
+
+
+def extract_images_from_trace_json(trace_json: Path, output_dir: Path):
+    """Extract images from trace JSON file."""
+    print(f"Loading trace from {trace_json}")
+    
+    with open(trace_json) as f:
+        trace = json.load(f)
+    
+    # Handle trace wrapped in "session_trace" key
+    if "session_trace" in trace:
+        trace = trace["session_trace"]
+    
+    count = extract_images_from_trace_dict(trace, output_dir)
+    print(f"\n✓ Extracted {count} images to {output_dir}/")
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--trace-db",
+        help="Path to trace database",
+    )
+    parser.add_argument(
+        "--trace-json",
+        type=Path,
+        help="Path to trace JSON file",
+    )
+    parser.add_argument(
+        "--output-dir",
+        default="examples/blog_posts/pokemon_vl/images_gpt5",
+        help="Output directory for images",
+    )
+    parser.add_argument(
+        "--model-filter",
+        help="Filter traces by model name (optional)",
+    )
+    args = parser.parse_args()
+    
+    output_dir = Path(args.output_dir)
+    
+    if args.trace_json:
+        extract_images_from_trace_json(args.trace_json, output_dir)
+    elif args.trace_db:
+        extract_images_from_trace_db(args.trace_db, output_dir, args.model_filter)
+    else:
+        parser.error("Must provide either --trace-db or --trace-json")
+
+
+if __name__ == "__main__":
+    main()
+
+

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,
+        },
+    )
+