synth-ai 0.2.9.dev5__py3-none-any.whl → 0.2.10__py3-none-any.whl

examples/vlm/README.md

@@ -0,0 +1,68 @@
+# Crafter VLM Pipeline
+
+This folder captures the reference workflow for fine-tuning Crafter policies with
+multimodal (text + image) prompts. It stitches together the new image-aware tracing
+plumbing with lightweight utilities for dataset curation and training.
+
+## Quick Start
+
+1. **Verify image capture**  
+   ```
+   uv run python examples/vlm/crafter_image_only_agent.py --seed 7 --steps 5
+   ```
+   This writes PNG frames to `examples/vlm/output/frames/` and produces a JSONL preview
+   of OpenAI-style image-only user messages.
+
+2. **Collect traced rollouts**  
+   Use the Crafter task app (or your existing pipeline) with tracing enabled. The new
+   tracing schema automatically records `observation_image_base64` and stores image parts
+   in LM call records.
+
+3. **Export multimodal SFT rows**  
+   ```
+   uv run python examples/warming_up_to_rl/export_trace_sft.py \
+       --db traces/v3/synth_ai.db \
+       --output examples/vlm/output/crafter_traces_full.jsonl
+   ```
+   The exporter now emits `metadata.has_image`, `metadata.user_has_image`, and
+   `metadata.assistant_has_image` flags per turn.
+
+4. **Filter to image-rich turns**  
+   ```
+   uv run python examples/vlm/filter_image_rows.py \
+       --input examples/vlm/output/crafter_traces_full.jsonl \
+       --output examples/vlm/output/crafter_vlm_dataset.jsonl
+   ```
+
+5. **(Optional) Split validation or augment**, then upload using the standard CLI:
+   ```
+   uv run python examples/warming_up_to_rl/run_fft_and_save.py \
+       --toml examples/vlm/configs/crafter_vlm_gpt4o.toml \
+       --data examples/vlm/output/crafter_vlm_dataset.jsonl
+   ```
+
+## Config & Utilities
+
+| File | Purpose |
+| --- | --- |
+| `configs/crafter_vlm_gpt4o.toml` | Sample Synth job targeting an image-capable model (`openai/gpt-4o-mini`). Set `job.data` or pass `--data` explicitly. |
+| `crafter_image_only_agent.py` | Captures frames and builds image-only prompts for sanity checks. |
+| `filter_image_rows.py` | Extracts rows with image parts from exported JSONL datasets. |
+
+## Notes & Next Steps
+
+- The training config assumes full-finetuning (`mode = "sft_offline"`). Adjust the
+  model id, hardware, or hyperparameters to match available infrastructure.
+- Dataset rows emitted by `export_trace_sft.py` already contain OpenAI multimodal
+  content parts like:
+  ```json
+  {
+    "role": "user",
+    "content": [
+      {"type": "text", "text": "..."},
+      {"type": "image_url", "image_url": {"url": "data:image/png;base64,..." }}
+    ]
+  }
+  ```
+- See `PROPOSAL.md` for a deeper dive into outstanding work (longer rollouts,
+  richer multimodal augmentations, evaluation ideas).

examples/vlm/configs/crafter_vlm_gpt4o.toml

@@ -0,0 +1,44 @@
+[job]
+model = "openai/gpt-4o-mini-2024-07-18"
+modalities = ["text", "image"]
+# data = "examples/vlm/output/crafter_vlm_dataset.jsonl"
+description = "Crafter VLM SFT (text + image prompts)"
+
+[compute]
+gpu_type = "A100"
+gpu_count = 1
+nodes = 1
+
+[data]
+topology = {}
+# validation_path = "examples/vlm/output/crafter_vlm_dataset.val.jsonl"
+
+[training]
+mode = "sft_offline"
+use_qlora = false
+
+[training.validation]
+enabled = true
+evaluation_strategy = "steps"
+eval_steps = 50
+save_best_model_at_end = true
+metric_for_best_model = "val.loss"
+greater_is_better = false
+
+[hyperparameters]
+n_epochs = 1
+train_kind = "fft"
+per_device_batch = 1
+gradient_accumulation_steps = 32
+sequence_length = 4096
+learning_rate = 1e-5
+warmup_ratio = 0.03
+weight_decay = 0.01
+
+[hyperparameters.parallelism]
+use_deepspeed = true
+deepspeed_stage = 2
+fsdp = false
+bf16 = true
+fp16 = false
+activation_checkpointing = true

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()