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

examples/vlm/crafter_image_only_agent.py

@@ -0,0 +1,207 @@
+#!/usr/bin/env python3
+"""
+Run a minimal Crafter agent that emits image-only prompts and saves rendered frames.
+
+This script demonstrates the multimodal observation pipeline by:
+  1. Initialising a `CrafterClassicEnvironment` with a deterministic seed.
+  2. Capturing `observation_image_base64` at each step and writing PNG frames.
+  3. Building OpenAI-style user messages that contain only an image part.
+  4. Emitting a small JSONL preview of the messages so they can be inspected or fed
+     directly into the fine-tuning dataset builder.
+
+Usage:
+    uv run python examples/vlm/crafter_image_only_agent.py --seed 7 --steps 5
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import base64
+import json
+import random
+from collections.abc import Iterable
+from pathlib import Path
+from typing import Any
+from uuid import uuid4
+
+from synth_ai.environments.environment.tools import EnvToolCall
+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
+
+ACTION_NAME_TO_ID = {
+    "noop": 0,
+    "move_left": 1,
+    "move_right": 2,
+    "move_up": 3,
+    "move_down": 4,
+    "do": 5,
+    "sleep": 6,
+    "place_stone": 7,
+    "place_table": 8,
+    "place_furnace": 9,
+    "place_plant": 10,
+    "make_wood_pickaxe": 11,
+    "make_stone_pickaxe": 12,
+    "make_iron_pickaxe": 13,
+    "make_wood_sword": 14,
+    "make_stone_sword": 15,
+    "make_iron_sword": 16,
+}
+
+
+def _build_task_instance(seed: int) -> CrafterTaskInstance:
+    """Construct a minimal Crafter task instance with the requested seed."""
+
+    impetus = Impetus(instructions="Explore the world and survive.")
+    intent = Intent(
+        rubric={"goal": "Unlock achievements and stay alive."},
+        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,
+    )
+    # Attach environment config expected by the engine
+    instance.config = {"seed": seed, "length": 256, "area": [64, 64]}
+    return instance
+
+
+def _select_actions(action_names: Iterable[str], steps: int) -> list[int]:
+    resolved: list[int] = []
+    names = list(action_names)
+    if not names:
+        names = ["move_right", "move_down", "move_left", "move_up", "do"]
+    for idx in range(steps):
+        name = names[idx % len(names)]
+        action_id = ACTION_NAME_TO_ID.get(name)
+        if action_id is None:
+            raise ValueError(f"Unknown Crafter action: {name}")
+        resolved.append(action_id)
+    return resolved
+
+
+def _save_base64_png(data: str, path: Path) -> None:
+    """Decode a base64 string (with or without data URL prefix) and write to disk."""
+
+    if data.startswith("data:"):
+        _, _, encoded = data.partition(",")
+    else:
+        encoded = data
+    path.write_bytes(base64.b64decode(encoded))
+
+
+def _build_image_only_message(data_url: str) -> dict[str, Any]:
+    return {
+        "role": "user",
+        "content": [{"type": "image_url", "image_url": {"url": data_url}}],
+    }
+
+
+async def run(args: argparse.Namespace) -> None:
+    output_dir = Path(args.output_dir).resolve()
+    frames_dir = output_dir / "frames"
+    frames_dir.mkdir(parents=True, exist_ok=True)
+    messages_path = output_dir / "image_only_messages.jsonl"
+
+    task_instance = _build_task_instance(args.seed)
+    env = CrafterClassicEnvironment(task_instance)
+
+    # Initialise environment
+    raw_obs = await env.initialize()
+    observation = getattr(raw_obs, "observation", raw_obs)
+
+    action_ids = _select_actions(args.actions, args.steps)
+    records: list[dict[str, Any]] = []
+
+    for step_idx in range(args.steps):
+        obs_dict = observation if isinstance(observation, dict) else {}
+        image_b64 = obs_dict.get("observation_image_base64")
+        data_url = obs_dict.get("observation_image_data_url")
+
+        if image_b64:
+            frame_path = frames_dir / f"step_{step_idx:03d}.png"
+            _save_base64_png(image_b64, frame_path)
+
+        if data_url:
+            message = _build_image_only_message(data_url)
+        else:
+            message = {
+                "role": "user",
+                "content": [{"type": "text", "text": "Image missing from observation."}],
+            }
+
+        records.append(
+            {
+                "step": step_idx,
+                "action_id": action_ids[step_idx],
+                "message": message,
+                "observation_keys": sorted(obs_dict.keys()),
+            }
+        )
+
+        # For the very first step, show the message structure
+        if step_idx == 0:
+            print("=== Image-only message example ===")
+            print(json.dumps(message, indent=2))
+
+        tool_call = EnvToolCall(tool="interact", args={"action": int(action_ids[step_idx])})
+        env_step = await env.step(tool_call)
+        observation = getattr(env_step, "observation", env_step)
+
+    # Wrap up and dump the preview JSONL
+    await env.terminate()
+    with messages_path.open("w", encoding="utf-8") as fh:
+        for record in records:
+            fh.write(json.dumps(record, ensure_ascii=False) + "\n")
+
+    print(f"Saved {len(records)} frames -> {frames_dir}")
+    print(f"Saved image-only message preview -> {messages_path}")
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--seed", type=int, default=7, help="Crafter environment seed")
+    parser.add_argument("--steps", type=int, default=5, help="Number of env steps to capture")
+    parser.add_argument(
+        "--actions",
+        nargs="*",
+        default=["move_right", "move_down", "move_left", "move_up", "do"],
+        help="Sequence of Crafter action names to cycle through",
+    )
+    default_output = Path("examples/vlm/temp")
+    parser.add_argument(
+        "--output-dir",
+        type=Path,
+        default=default_output,
+        help=f"Directory for frames and message preview (default: {default_output})",
+    )
+    parser.add_argument(
+        "--randomise",
+        action="store_true",
+        help="Shuffle the provided action sequence before running",
+    )
+    args = parser.parse_args()
+    if args.randomise:
+        random.shuffle(args.actions)
+    return args
+
+
+if __name__ == "__main__":
+    asyncio.run(run(parse_args()))

examples/vlm/crafter_openai_vlm_agent.py

@@ -0,0 +1,277 @@
+#!/usr/bin/env python3
+"""
+Crafter agent that calls the OpenAI Chat Completions API with image + text prompts.
+
+The harness mirrors the text-based agent workflow from `examples/warming_up_to_rl`:
+  * Uses the Crafter policy to build prompts, maintain history, and parse tool calls.
+  * Executes actions against the Synth Crafter environment (no HTTP task app required).
+  * Persists every rendered frame to `examples/vlm/temp/` so you can inspect exactly
+    what the VLM saw.
+
+Requirements:
+  - `OPENAI_API_KEY` environment variable.
+  - `openai` Python package (installed via project dependencies).
+
+Usage:
+  uv run python examples/vlm/crafter_openai_vlm_agent.py \
+      --model gpt-4o-mini-2024-07-18 --seeds 10 --steps 10
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import base64
+import json
+import os
+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 OpenAI
+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
+
+DEFAULT_OUTPUT = Path("examples/vlm/temp")
+FRAME_SUBDIR = "openai_agent_frames"
+
+
+class EpisodeResult:
+    def __init__(self, seed: int) -> None:
+        self.seed = seed
+        self.steps_taken: int = 0
+        self.achievements: set[str] = set()
+        self.total_reward: float = 0.0
+        self.tool_calls: int = 0
+
+    def record_observation(self, observation: dict[str, Any]) -> None:
+        obs = observation.get("observation") if isinstance(observation, dict) else None
+        if not isinstance(obs, dict):
+            return
+        ach = obs.get("achievements_status")
+        if isinstance(ach, dict):
+            for name, unlocked in ach.items():
+                if unlocked:
+                    self.achievements.add(str(name))
+        reward = obs.get("reward_last_step")
+        if isinstance(reward, (int, float)):
+            self.total_reward += float(reward)
+
+
+def _ensure_client() -> OpenAI:
+    api_key = os.getenv("OPENAI_API_KEY")
+    if not api_key:
+        raise RuntimeError("OPENAI_API_KEY must be set for OpenAI calls")
+    return OpenAI(api_key=api_key)
+
+
+def _build_task_instance(seed: int) -> CrafterTaskInstance:
+    impetus = Impetus(instructions="Explore, survive, and unlock achievements.")
+    intent = Intent(
+        rubric={"goal": "Maximise Crafter 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,
+    )
+    instance.config = {"seed": seed, "length": 256, "area": [64, 64]}
+    return instance
+
+
+def _decode_and_save_image(observation: dict[str, Any], path: Path) -> None:
+    obs = observation.get("observation") if isinstance(observation, dict) else None
+    if not isinstance(obs, dict):
+        return
+    base64_data = obs.get("observation_image_base64")
+    if not isinstance(base64_data, str) or not base64_data:
+        return
+    path.parent.mkdir(parents=True, exist_ok=True)
+    try:
+        path.write_bytes(base64.b64decode(base64_data))
+    except Exception:
+        # Best-effort; corrupted frames should not halt rollout
+        pass
+
+
+def _normalise_openai_request(payload: dict[str, Any], model: str, temperature: float) -> dict[str, Any]:
+    request = dict(payload)
+    request["model"] = model
+    request.setdefault("temperature", temperature)
+    request.setdefault("max_tokens", 512)
+    # Remove vendor-specific knobs unsupported by OpenAI
+    request.pop("stop_after_tool_calls", None)
+    request.pop("thinking_mode", None)
+    request.pop("thinking_budget", None)
+    max_completion = request.pop("max_completion_tokens", None)
+    if max_completion is not None:
+        request.setdefault("max_tokens", max_completion)
+    return request
+
+
+async def _run_episode(
+    *,
+    seed: int,
+    client: OpenAI,
+    model: str,
+    max_steps: int,
+    output_dir: Path,
+    temperature: float,
+) -> EpisodeResult:
+    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})
+
+    episode_result = EpisodeResult(seed=seed)
+
+    observation_packet = await wrapper.initialize()
+    episode_result.record_observation(observation_packet)
+
+    frames_root = output_dir / FRAME_SUBDIR / f"seed_{seed:04d}"
+    _decode_and_save_image(observation_packet, frames_root / "step_000.png")
+
+    for step_idx in range(max_steps):
+        obs_dict = observation_packet.get("observation")
+        if not isinstance(obs_dict, dict):
+            break
+
+        obs_text = policy._format_observation_for_llm(observation_packet)  # noqa: SLF001
+        tool_calls, meta = await policy.step(
+            observation_text=obs_text,
+            metadata={"raw_observation": observation_packet},
+        )
+        if "inference_request" not in meta:
+            break
+
+        episode_result.steps_taken += 1
+        inference_request = _normalise_openai_request(
+            meta["inference_request"],
+            model=model,
+            temperature=temperature,
+        )
+
+        response = 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:
+            print(
+                f"Seed {seed}: no tool calls returned by model; ending episode early at step {step_idx}."
+            )
+            break
+
+        episode_result.tool_calls += 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}"
+            )
+        episode_result.record_observation(env_response)
+
+        policy._append_assistant_turn(  # noqa: SLF001
+            assistant_text,
+            assistant_tool_calls,
+            env_response,
+        )
+
+        frame_path = frames_root / f"step_{step_idx + 1:03d}.png"
+        _decode_and_save_image(env_response, frame_path)
+
+        if env_response.get("done"):
+            break
+        observation_packet = env_response
+
+    await wrapper.terminate()
+    return episode_result
+
+
+async def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--model", default="gpt-4o-mini-2024-07-18", help="OpenAI model id")
+    parser.add_argument("--seeds", type=int, default=10, help="Number of random seeds to evaluate")
+    parser.add_argument("--steps", type=int, default=10, help="Max steps per seed")
+    parser.add_argument("--temperature", type=float, default=0.6, help="Sampling temperature")
+    parser.add_argument(
+        "--output-dir",
+        type=Path,
+        default=DEFAULT_OUTPUT,
+        help=f"Directory for saved frames and summaries (default: {DEFAULT_OUTPUT})",
+    )
+    args = parser.parse_args()
+
+    client = _ensure_client()
+    results: list[EpisodeResult] = []
+
+    seeds = list(range(args.seeds))
+    print(f"Running {len(seeds)} Crafter episodes with model={args.model}")
+
+    for seed in seeds:
+        result = await _run_episode(
+            seed=seed,
+            client=client,
+            model=args.model,
+            max_steps=args.steps,
+            output_dir=args.output_dir,
+            temperature=args.temperature,
+        )
+        results.append(result)
+        print(
+            f"Seed {seed:02d}: steps={result.steps_taken}, "
+            f"achievements={len(result.achievements)}, "
+            f"tool_calls={result.tool_calls}, reward≈{result.total_reward:.3f}"
+        )
+
+    summary = {
+        "model": args.model,
+        "episodes": len(results),
+        "mean_steps": round(
+            sum(res.steps_taken for res in results) / max(len(results), 1), 2
+        ),
+        "mean_achievements": round(
+            sum(len(res.achievements) for res in results) / max(len(results), 1), 2
+        ),
+        "total_tool_calls": sum(res.tool_calls for res in results),
+        "output_dir": str(args.output_dir / FRAME_SUBDIR),
+    }
+
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+    summary_path = args.output_dir / "openai_agent_summary.json"
+    summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8")
+
+    print("\nSummary")
+    print("-------")
+    print(json.dumps(summary, indent=2))
+    print(f"\nFrames saved in: {summary['output_dir']}")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/vlm/filter_image_rows.py

@@ -0,0 +1,63 @@
+#!/usr/bin/env python3
+"""
+Filter SFT JSONL rows to those that contain image content.
+
+This is a convenience wrapper around `examples/warming_up_to_rl/export_trace_sft.py`
+output now that each record's metadata includes `has_image`, `user_has_image`, and
+`assistant_has_image`.
+
+Usage:
+    uv run python examples/vlm/filter_image_rows.py \
+        --input examples/sft/ft_data/crafter_traces.jsonl \
+        --output examples/vlm/output/crafter_vlm_dataset.jsonl
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--input", type=Path, required=True, help="Source JSONL dataset")
+    parser.add_argument("--output", type=Path, required=True, help="Filtered JSONL path")
+    parser.add_argument(
+        "--include-assistant",
+        action="store_true",
+        help="Require the assistant message to include an image as well",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    src = args.input
+    dst = args.output
+    dst.parent.mkdir(parents=True, exist_ok=True)
+
+    kept = 0
+    total = 0
+    with src.open("r", encoding="utf-8") as reader, dst.open("w", encoding="utf-8") as writer:
+        for line in reader:
+            total += 1
+            try:
+                record = json.loads(line)
+            except json.JSONDecodeError:
+                continue
+            metadata = record.get("metadata") or {}
+            has_user_image = bool(metadata.get("user_has_image"))
+            has_assistant_image = bool(metadata.get("assistant_has_image"))
+            if not has_user_image:
+                continue
+            if args.include_assistant and not has_assistant_image:
+                continue
+            writer.write(json.dumps(record, ensure_ascii=False) + "\n")
+            kept += 1
+
+    print(f"Filtered {kept} / {total} rows with user images -> {dst}")
+
+
+if __name__ == "__main__":
+    main()