synth-ai 0.2.9.dev4__py3-none-any.whl → 0.2.9.dev6__py3-none-any.whl

examples/vlm/run_crafter_vlm_benchmark.py

@@ -0,0 +1,316 @@
+#!/usr/bin/env python3
+"""
+Benchmark Crafter performance across prompt modalities (text-only, image-only, both).
+
+For each mode we:
+  * Run 20 seeded episodes (configurable) with GPT-4o mini via OpenAI Chat Completions.
+  * Execute the returned tool calls in the local Crafter environment.
+  * Record achievements/steps and save every rendered frame under `examples/vlm/temp/`.
+
+Concurrency is capped by an asyncio semaphore (default parallelism = 10).
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import base64
+import json
+import os
+from collections import Counter, defaultdict
+from dataclasses import dataclass
+from enum import Enum
+from pathlib import Path
+from typing import Any
+from uuid import uuid4
+
+from examples.warming_up_to_rl.task_app.synth_envs_hosted.envs.crafter.environment import (
+    CrafterEnvironmentWrapper,
+)
+from examples.warming_up_to_rl.task_app.synth_envs_hosted.envs.crafter.policy import CrafterPolicy
+from openai import AsyncOpenAI
+from synth_ai.environments.examples.crafter_classic.environment import CrafterClassicEnvironment
+from synth_ai.environments.examples.crafter_classic.taskset import (
+    CrafterTaskInstance,
+    CrafterTaskInstanceMetadata,
+)
+from synth_ai.environments.tasks.core import Impetus, Intent
+
+OUTPUT_ROOT = Path("examples/vlm/temp")
+
+
+class Mode(str, Enum):
+    TEXT = "text"
+    IMAGE = "image"
+    BOTH = "both"
+
+
+@dataclass
+class EpisodeResult:
+    mode: Mode
+    seed: int
+    steps_taken: int
+    achievements: set[str]
+    total_reward: float
+    tool_calls: int
+
+
+def _ensure_openai_client(api_key: str | None) -> AsyncOpenAI:
+    if not api_key:
+        raise RuntimeError(
+            "OPENAI_API_KEY must be set to run the VLM benchmark (export the key or add to your .env)."
+        )
+    return AsyncOpenAI(api_key=api_key)
+
+
+def _build_task_instance(seed: int) -> CrafterTaskInstance:
+    impetus = Impetus(instructions="Explore, survive, and unlock achievements.")
+    intent = Intent(rubric={"goal": "Unlock achievements"}, gold_trajectories=None, gold_state_diff={})
+    metadata = CrafterTaskInstanceMetadata(
+        difficulty="custom",
+        seed=seed,
+        num_trees_radius=0,
+        num_cows_radius=0,
+        num_hostiles_radius=0,
+    )
+    instance = CrafterTaskInstance(
+        id=uuid4(),
+        impetus=impetus,
+        intent=intent,
+        metadata=metadata,
+        is_reproducible=True,
+        initial_engine_snapshot=None,
+    )
+    # Engine expects these config keys
+    instance.config = {"seed": seed, "length": 256, "area": [64, 64]}
+    return instance
+
+
+def _save_observation_frame(observation_packet: dict[str, Any], dest_path: Path) -> None:
+    obs = observation_packet.get("observation")
+    if not isinstance(obs, dict):
+        return
+    image_b64 = obs.get("observation_image_base64")
+    if not isinstance(image_b64, str) or not image_b64:
+        return
+    try:
+        dest_path.parent.mkdir(parents=True, exist_ok=True)
+        dest_path.write_bytes(base64.b64decode(image_b64))
+    except Exception:
+        pass  # best effort
+
+
+def _strip_image_fields(observation_packet: dict[str, Any]) -> dict[str, Any]:
+    stripped = json.loads(json.dumps(observation_packet))
+    obs = stripped.get("observation")
+    if isinstance(obs, dict):
+        for key in list(obs.keys()):
+            if key.startswith("observation_image"):
+                obs.pop(key, None)
+    return stripped
+
+
+def _make_image_only_request(request: dict[str, Any]) -> dict[str, Any]:
+    cloned = json.loads(json.dumps(request))
+    for message in cloned.get("messages", []):
+        if message.get("role") != "user":
+            continue
+        content = message.get("content")
+        if isinstance(content, list):
+            image_parts = [
+                item
+                for item in content
+                if isinstance(item, dict) and item.get("type") in {"image_url", "image"}
+            ]
+            message["content"] = image_parts or content
+        elif isinstance(content, str):
+            # No structured parts available; leave as empty string
+            message["content"] = ""
+    return cloned
+
+
+async def _run_episode(
+    *,
+    mode: Mode,
+    seed: int,
+    client: AsyncOpenAI,
+    model: str,
+    max_steps: int,
+    temperature: float,
+    semaphore: asyncio.Semaphore,
+) -> EpisodeResult:
+    async with semaphore:
+        task_instance = _build_task_instance(seed)
+        env = CrafterClassicEnvironment(task_instance)
+        wrapper = CrafterEnvironmentWrapper(env, seed=seed)
+
+        policy = CrafterPolicy(inference_url="openai://chat-completions", model=model)
+        await policy.initialize({"use_tools": True, "model": model})
+
+        observation_packet = await wrapper.initialize()
+        achievements: set[str] = set()
+        total_reward = 0.0
+        steps_taken = 0
+        tool_calls_total = 0
+
+        frames_dir = OUTPUT_ROOT / f"{mode.value}_frames" / f"seed_{seed:04d}"
+        _save_observation_frame(observation_packet, frames_dir / "step_000.png")
+
+        for step_idx in range(max_steps):
+            obs_dict = observation_packet.get("observation")
+            if not isinstance(obs_dict, dict):
+                break
+
+            observation_for_policy: dict[str, Any]
+            metadata_payload: dict[str, Any] = {}
+
+            if mode == Mode.TEXT:
+                observation_for_policy = _strip_image_fields(observation_packet)
+            else:
+                observation_for_policy = json.loads(json.dumps(observation_packet))
+                metadata_payload["raw_observation"] = observation_packet
+
+            obs_text = policy._format_observation_for_llm(observation_for_policy)  # noqa: SLF001
+            _, meta = await policy.step(
+                observation_text=obs_text,
+                metadata=metadata_payload,
+            )
+            inference_request = json.loads(json.dumps(meta["inference_request"]))
+
+            if mode == Mode.IMAGE:
+                inference_request = _make_image_only_request(inference_request)
+
+            inference_request.update(
+                {
+                    "model": model,
+                    "temperature": temperature,
+                    "max_tokens": inference_request.get("max_tokens", 512),
+                }
+            )
+            inference_request.pop("stop_after_tool_calls", None)
+            inference_request.pop("thinking_mode", None)
+            inference_request.pop("thinking_budget", None)
+
+            response = await client.chat.completions.create(**inference_request)
+            response_dict = response.model_dump()
+
+            assistant_tool_calls = CrafterPolicy.parse_response_to_tool_calls(
+                response_dict,
+                use_tools=policy.use_tools,
+            )
+            if not assistant_tool_calls:
+                break
+
+            tool_calls_total += len(assistant_tool_calls)
+            assistant_message = response_dict["choices"][0].get("message") or {}
+            assistant_text = assistant_message.get("content")
+
+            env_response = await wrapper.step(assistant_tool_calls)
+            if not isinstance(env_response, dict):
+                raise RuntimeError(f"Unexpected environment response type: {type(env_response)!r}")
+
+            policy._append_assistant_turn(  # noqa: SLF001
+                assistant_text,
+                assistant_tool_calls,
+                env_response,
+            )
+
+            steps_taken += 1
+            obs = env_response.get("observation")
+            if isinstance(obs, dict):
+                ach = obs.get("achievements_status")
+                if isinstance(ach, dict):
+                    for name, unlocked in ach.items():
+                        if unlocked:
+                            achievements.add(str(name))
+                reward = obs.get("reward_last_step")
+                if isinstance(reward, (int, float)):
+                    total_reward += float(reward)
+
+            _save_observation_frame(env_response, frames_dir / f"step_{step_idx + 1:03d}.png")
+
+            if env_response.get("done"):
+                break
+            observation_packet = env_response
+
+        await wrapper.terminate()
+        return EpisodeResult(
+            mode=mode,
+            seed=seed,
+            steps_taken=steps_taken,
+            achievements=achievements,
+            total_reward=total_reward,
+            tool_calls=tool_calls_total,
+        )
+
+
+def _summarise(results: list[EpisodeResult]) -> dict[str, Any]:
+    grouped: dict[Mode, list[EpisodeResult]] = defaultdict(list)
+    for result in results:
+        grouped[result.mode].append(result)
+
+    summary: dict[str, Any] = {}
+    for mode, mode_results in grouped.items():
+        if not mode_results:
+            continue
+        mean_steps = sum(r.steps_taken for r in mode_results) / len(mode_results)
+        mean_achievements = sum(len(r.achievements) for r in mode_results) / len(mode_results)
+        achievement_counts = Counter()
+        for res in mode_results:
+            achievement_counts.update(res.achievements)
+        summary[mode.value] = {
+            "episodes": len(mode_results),
+            "mean_steps": round(mean_steps, 2),
+            "mean_achievements": round(mean_achievements, 2),
+            "total_tool_calls": sum(r.tool_calls for r in mode_results),
+            "achievements": {name: count for name, count in sorted(achievement_counts.items())},
+        }
+    return summary
+
+
+async def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--model", default="gpt-4o-mini-2024-07-18", help="OpenAI model id to benchmark")
+    parser.add_argument("--seeds", type=int, default=20, help="Number of seeds per mode")
+    parser.add_argument("--steps", type=int, default=10, help="Max steps per episode")
+    parser.add_argument("--temperature", type=float, default=0.6, help="Sampling temperature")
+    parser.add_argument("--concurrency", type=int, default=10, help="Max concurrent OpenAI calls")
+    args = parser.parse_args()
+
+    api_key = os.getenv("OPENAI_API_KEY")
+    client = _ensure_openai_client(api_key)
+    semaphore = asyncio.Semaphore(max(1, args.concurrency))
+
+    OUTPUT_ROOT.mkdir(parents=True, exist_ok=True)
+
+    tasks: list[asyncio.Task[EpisodeResult]] = []
+    for mode in (Mode.TEXT, Mode.IMAGE, Mode.BOTH):
+        for seed in range(args.seeds):
+            task = asyncio.create_task(
+                _run_episode(
+                    mode=mode,
+                    seed=seed,
+                    client=client,
+                    model=args.model,
+                    max_steps=args.steps,
+                    temperature=args.temperature,
+                    semaphore=semaphore,
+                )
+            )
+            tasks.append(task)
+
+    results = await asyncio.gather(*tasks)
+    summary = _summarise(results)
+
+    summary_path = OUTPUT_ROOT / "vlm_benchmark_summary.json"
+    summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8")
+
+    print("\nBenchmark Summary")
+    print("-----------------")
+    print(json.dumps(summary, indent=2))
+    print(f"\nFrames stored under: {OUTPUT_ROOT}/<mode>_frames/seed_xxxx/")
+    print(f"Summary saved to: {summary_path}")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/warming_up_to_rl/analyze_trace_db.py

@@ -9,7 +9,7 @@ import sqlite3
 import sys
 from collections import Counter, defaultdict
 from pathlib import Path
-from typing import Any, Dict, List, Set, Tuple
+from typing import Any
 
 Row = sqlite3.Row
 
@@ -56,7 +56,7 @@ def fetch_model_usage(conn: sqlite3.Connection) -> list[dict[str, Any]]:
 def _parse_json(value: Any) -> Any:
     if value is None:
         return None
-    if isinstance(value, (dict, list)):
+    if isinstance(value, dict | list):
         return value
     try:
         return json.loads(value)
@@ -64,7 +64,7 @@ def _parse_json(value: Any) -> Any:
         return None
 
 
-AchievementMap = dict[Tuple[str, int], dict[str, list[str]]]
+AchievementMap = dict[tuple[str, int], dict[str, list[str]]]
 
 
 def fetch_achievement_data(
@@ -162,7 +162,7 @@ def fetch_achievement_data(
         achievement_name_counts.update(achievement_set)
 
     achievement_size_counts: Counter = Counter()
-    for session_id, count in unique_counts_per_session.items():
+    for _session_id, count in unique_counts_per_session.items():
         achievement_size_counts[count] += 1
 
     return (
@@ -229,7 +229,9 @@ def format_model_stats(stats: list[dict[str, Any]]) -> str:
     if not stats:
         return "No model usage recorded."
     lines = ["Model usage (by LLM calls):"]
-    header = f"{'Model':30} {'Provider':10} {'Calls':>7} {'Tokens (in/out)':>20} {'Avg latency ms':>15}"
+    header = (
+        f"{'Model':30} {'Provider':10} {'Calls':>7} {'Tokens (in/out)':>20} {'Avg latency ms':>15}"
+    )
     lines.append(header)
     lines.append("-" * len(header))
     for item in stats:
@@ -243,9 +245,7 @@ def format_model_stats(stats: list[dict[str, Any]]) -> str:
     return "\n".join(lines)
 
 
-def format_achievement_summary(
-    name_counts: Counter, size_counts: Counter
-) -> str:
+def format_achievement_summary(name_counts: Counter, size_counts: Counter) -> str:
     lines = ["Unique achievements unlocked:"]
     if name_counts:
         top = name_counts.most_common()
@@ -295,7 +295,7 @@ def format_reward_summary(outcome: dict[str, Any], breakdown: list[dict[str, Any
 
 
 def compute_model_achievement_stats(
-    conn: sqlite3.Connection, session_unique_sets: dict[str, Set[str]]
+    conn: sqlite3.Connection, session_unique_sets: dict[str, set[str]]
 ) -> dict[str, dict[str, Any]]:
     """Aggregate unique-achievement stats per model."""
 
@@ -349,7 +349,9 @@ def format_model_achievement_stats(model_stats: dict[str, dict[str, Any]]) -> st
         return "Achievement stats by model:\n  (no model sessions recorded)"
 
     lines = ["Achievement stats by model:"]
-    for model_name in sorted(model_stats.keys(), key=lambda m: model_stats[m]["sessions"], reverse=True):
+    for model_name in sorted(
+        model_stats.keys(), key=lambda m: model_stats[m]["sessions"], reverse=True
+    ):
         stats = model_stats[model_name]
         providers = ", ".join(sorted(stats["providers"])) if stats["providers"] else "-"
         sessions = stats["sessions"]

examples/warming_up_to_rl/configs/rl_from_base_qwen4b.toml

@@ -42,9 +42,13 @@ base = "Qwen/Qwen3-4B"
 label = "crafter-rl-from-base"
 
 [rollout]
+env_name = "crafter"
 max_turns = 10
 episodes_per_batch = 64
-policy_name = "crafter"
+policy_name = "crafter-react"
+max_concurrent_rollouts = 8
+batches_per_step = 2
+ops = ["agent", "env"]
 
 [evaluation]
 # Run baseline evaluation over the first 100 seeds every 20 training iterations
@@ -55,6 +59,12 @@ seeds = [
 ]
 
 [training]
+num_epochs = 1
+iterations_per_epoch = 10
+batch_size = 16
+group_size = 4
+gradient_accumulation_steps = 1
+learning_rate = 5e-5
 log_interval = 1
 weight_sync_interval = 1
 # Additional RL hyperparameters can go here