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

synth_ai/demos/demo_task_apps/math/modal_task_app.py

@@ -1,14 +1,13 @@
-from __future__ import annotations
-
 """Modal task app for Hendrycks MATH single-step RL environment."""
 
+from __future__ import annotations
+
 import os
+from collections.abc import Iterable
+from functools import lru_cache
 from pathlib import Path
 
 from modal import App, Image, Secret, asgi_app
-from functools import lru_cache
-from typing import Iterable
-
 from starlette.requests import Request
 
 try:  # Backward compatibility with older installed SDKs
@@ -25,7 +24,9 @@ _SYNTH_HOSTED = None
 try:
     probe = _HERE
     for _ in range(8):
-        candidate = (probe / "backend/app/routes/clustered_training/dev/synth_envs_hosted").resolve()
+        candidate = (
+            probe / "backend/app/routes/clustered_training/dev/synth_envs_hosted"
+        ).resolve()
         if candidate.exists():
             _SYNTH_HOSTED = candidate
             break
@@ -97,16 +98,17 @@ app = App("hendrycks-math-task-app")
 @asgi_app()
 def fastapi_app():
     import httpx
-    from fastapi import Body, HTTPException, status
-    from fastapi import FastAPI
+    from fastapi import Body, FastAPI, HTTPException, status
     from fastapi.middleware.cors import CORSMiddleware
     from fastapi.responses import JSONResponse
+
     try:
         from synth_ai.task.auth import (
             is_api_key_header_authorized,
             normalize_environment_api_key,
         )
     except Exception:  # pragma: no cover - fallback for older synth-ai builds
+
         def _normalize_env_key_fallback() -> str | None:
             key = os.getenv("ENVIRONMENT_API_KEY")
             if key:
@@ -130,7 +132,7 @@ def fastapi_app():
             for value in values:
                 if not isinstance(value, str):
                     continue
-                for chunk in value.split(','):
+                for chunk in value.split(","):
                     chunk = chunk.strip()
                     if chunk:
                         parts.append(chunk)
@@ -172,19 +174,27 @@ def fastapi_app():
 
     def _normalize_answer_text(s: str) -> str:
         import re as _re
+
         return _re.sub(r"[^0-9A-Za-z.+\-/*=]", "", (s or "").strip()).lower()
 
     def _extract_boxed(s: str) -> str:
         import re as _re
+
         m = list(_re.finditer(r"\\boxed\{([^}]+)\}", s or ""))
         return m[-1].group(1) if m else ""
 
     def _load_hendrycks_problem(seed: int, subject: str | None = None) -> tuple[str, str]:
         subj = subject or os.getenv("HENDRYCKS_MATH_CONFIG", "default")
-        ds = _hf_split(subj, os.getenv("HENDRYCKS_MATH_SPLIT", "test"), os.getenv("HENDRYCKS_MATH_SLICE"))
+        ds = _hf_split(
+            subj, os.getenv("HENDRYCKS_MATH_SPLIT", "test"), os.getenv("HENDRYCKS_MATH_SLICE")
+        )
         n = len(ds) if hasattr(ds, "__len__") else 0
         if n == 0 and subject not in {"", "default"}:
-            ds = _hf_split("default", os.getenv("HENDRYCKS_MATH_SPLIT", "test"), os.getenv("HENDRYCKS_MATH_SLICE"))
+            ds = _hf_split(
+                "default",
+                os.getenv("HENDRYCKS_MATH_SPLIT", "test"),
+                os.getenv("HENDRYCKS_MATH_SLICE"),
+            )
             n = len(ds) if hasattr(ds, "__len__") else 0
         if n == 0:
             raise RuntimeError("Hendrycks MATH dataset loaded empty")
@@ -225,7 +235,11 @@ def fastapi_app():
 
         def _resolve_env_keys() -> set[str]:
             keys: set[str] = set()
-            for alias in ("ENVIRONMENT_API_KEY", "dev_environment_api_key", "DEV_ENVIRONMENT_API_KEY"):
+            for alias in (
+                "ENVIRONMENT_API_KEY",
+                "dev_environment_api_key",
+                "DEV_ENVIRONMENT_API_KEY",
+            ):
                 value = os.environ.get(alias)
                 if value:
                     os.environ.setdefault("ENVIRONMENT_API_KEY", value)
@@ -250,8 +264,12 @@ def fastapi_app():
                 candidates.append(primary.strip())
             secondary = x_api_keys or headers.get("x-api-keys")
             if secondary:
-                candidates.extend([value.strip() for value in secondary.split(",") if value.strip()])
-            auth_header = authorization or headers.get("authorization") or headers.get("Authorization")
+                candidates.extend(
+                    [value.strip() for value in secondary.split(",") if value.strip()]
+                )
+            auth_header = (
+                authorization or headers.get("authorization") or headers.get("Authorization")
+            )
             if auth_header and auth_header.lower().startswith("bearer "):
                 token = auth_header.split(" ", 1)[1].strip()
                 if token:
@@ -274,7 +292,10 @@ def fastapi_app():
         async def info():
             return {
                 "service": {"base_url": os.getenv("SERVICE_BASE_URL", "")},
-                "inference": {"base_url": "", "endpoints": {"chat_completions": "/v1/chat/completions"}},
+                "inference": {
+                    "base_url": "",
+                    "endpoints": {"chat_completions": "/v1/chat/completions"},
+                },
             }
 
         @app.get("/health")
@@ -282,7 +303,10 @@ def fastapi_app():
             env_keys = _resolve_env_keys()
             env_key = next(iter(env_keys), None)
             if not env_key:
-                return JSONResponse(status_code=503, content={"status": "unhealthy", "detail": "Missing ENVIRONMENT_API_KEY"})
+                return JSONResponse(
+                    status_code=503,
+                    content={"status": "unhealthy", "detail": "Missing ENVIRONMENT_API_KEY"},
+                )
             # Authorize using all header variants; avoid typed Header params to prevent 422s
             authorized = is_api_key_header_authorized(request)
             if not authorized:
@@ -302,7 +326,10 @@ def fastapi_app():
             env_keys = _resolve_env_keys()
             env_key = next(iter(env_keys), None)
             if not env_key:
-                return JSONResponse(status_code=503, content={"status": "unhealthy", "detail": "Missing ENVIRONMENT_API_KEY"})
+                return JSONResponse(
+                    status_code=503,
+                    content={"status": "unhealthy", "detail": "Missing ENVIRONMENT_API_KEY"},
+                )
             authorized = is_api_key_header_authorized(request)
             if not authorized:
                 prefix = _log_env_key_prefix("health/rollout", env_key)
@@ -321,17 +348,22 @@ def fastapi_app():
         async def task_info(seed: int = 0, subject: str = "default"):
             """Return Hendrycks MATH problem/answer and tool schema for a seed."""
             q, a = _load_hendrycks_problem(int(seed), subject=subject)
-            tools = [{
-                "name": "submit_answer",
-                "description": "Provide the final numerical or algebraic answer for the current math problem.",
-                "parameters": {
-                    "type": "object",
-                    "properties": {
-                        "answer": {"type": "string", "description": "The proposed final answer"},
+            tools = [
+                {
+                    "name": "submit_answer",
+                    "description": "Provide the final numerical or algebraic answer for the current math problem.",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "answer": {
+                                "type": "string",
+                                "description": "The proposed final answer",
+                            },
+                        },
+                        "required": ["answer"],
                     },
-                    "required": ["answer"],
-                },
-            }]
+                }
+            ]
             return {
                 "seed": int(seed),
                 "subject": subject,
@@ -354,7 +386,7 @@ def fastapi_app():
         try:
             hdr = request.headers
             snapshot = {
-                "path": str(getattr(request, "url").path),
+                "path": str(request.url.path),
                 "have_x_api_key": bool(hdr.get("x-api-key")),
                 "have_x_api_keys": bool(hdr.get("x-api-keys")),
                 "have_authorization": bool(hdr.get("authorization")),
@@ -363,7 +395,9 @@ def fastapi_app():
             print("[422] validation", snapshot, flush=True)
         except Exception:
             pass
-        return JSONResponse(status_code=422, content={"status": "invalid", "detail": exc.errors()[:5]})
+        return JSONResponse(
+            status_code=422, content={"status": "invalid", "detail": exc.errors()[:5]}
+        )
 
     @api.get("/")
     async def root_probe():
@@ -376,27 +410,32 @@ def fastapi_app():
     env_key = (
         os.environ.get("ENVIRONMENT_API_KEY")
         or os.environ.get("DEV_ENVIRONMENT_API_KEY")
-        or os.environ.get("dev_environment_api_key")
+        or os.environ.get("DEV_ENVIRONMENT_API_KEY")
     )
     if not env_key:
         raise RuntimeError("ENVIRONMENT_API_KEY missing in task app environment")
 
-    OPENAI_REMOVE_FIELDS = ("stop_after_tool_calls", "thinking_mode", "thinking_budget", "reasoning")
-    OPENAI_REMOVE_SAMPLING_FIELDS = ("temperature", "top_p")
-    TOOL_CHOICE_FORCE = {"type": "function", "function": {"name": "submit_answer"}}
+    openai_remove_fields = (
+        "stop_after_tool_calls",
+        "thinking_mode",
+        "thinking_budget",
+        "reasoning",
+    )
+    openai_remove_sampling_fields = ("temperature", "top_p")
+    tool_choice_force = {"type": "function", "function": {"name": "submit_answer"}}
 
     def _prepare_openai_payload(model: str | None, payload: dict[str, object]) -> dict[str, object]:
         sanitized = dict(payload)
-        for key in OPENAI_REMOVE_FIELDS:
+        for key in openai_remove_fields:
             sanitized.pop(key, None)
         if model and "gpt-5" in model:
             if "max_tokens" in sanitized and "max_completion_tokens" not in sanitized:
                 sanitized["max_completion_tokens"] = sanitized.pop("max_tokens")
             else:
                 sanitized.pop("max_tokens", None)
-            for field in OPENAI_REMOVE_SAMPLING_FIELDS:
+            for field in openai_remove_sampling_fields:
                 sanitized.pop(field, None)
-                sanitized["tool_choice"] = TOOL_CHOICE_FORCE
+                sanitized["tool_choice"] = tool_choice_force
                 sanitized["parallel_tool_calls"] = False
             return sanitized
 
@@ -404,12 +443,18 @@ def fastapi_app():
     def proxy_chat_completions(request: dict[str, object] = Body(...)):
         key = os.environ.get("OPENAI_API_KEY")
         if not key:
-            raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="OPENAI_API_KEY missing")
+            raise HTTPException(
+                status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="OPENAI_API_KEY missing"
+            )
         model = request.get("model") if isinstance(request, dict) else None
-        payload = _prepare_openai_payload(model if isinstance(model, str) else None, request if isinstance(request, dict) else {})
+        payload = _prepare_openai_payload(
+            model if isinstance(model, str) else None, request if isinstance(request, dict) else {}
+        )
         headers = {"Authorization": f"Bearer {key}"}
         with httpx.Client(timeout=httpx.Timeout(180.0), follow_redirects=True) as client:
-            resp = client.post("https://api.openai.com/v1/chat/completions", json=payload, headers=headers)
+            resp = client.post(
+                "https://api.openai.com/v1/chat/completions", json=payload, headers=headers
+            )
             try:
                 data = resp.json()
             except Exception:
@@ -423,8 +468,8 @@ def fastapi_app():
     # Minimal math rollout endpoint: alternates agent/env; calls inference_url chat/completions
     @api.post("/rollout")
     def rollout(request: dict[str, object] = Body(...)):
-        from typing import Any
         import json as _json
+        from typing import Any
 
         run_id = str(request.get("run_id"))
         data = request if isinstance(request, dict) else {}
@@ -442,15 +487,25 @@ def fastapi_app():
         env_cfg = (env or {}).get("config") or {}
         # Prefer env.seed; fall back to env.config.seed -> default 0
         try:
-            seed_val = int((env or {}).get("seed")) if isinstance(env, dict) and (env or {}).get("seed") is not None else 0
+            seed_val = (
+                int((env or {}).get("seed"))
+                if isinstance(env, dict) and (env or {}).get("seed") is not None
+                else 0
+            )
         except Exception:
             seed_val = 0
         if seed_val == 0:
             try:
-                seed_val = int(env_cfg.get("seed")) if isinstance(env_cfg, dict) and env_cfg.get("seed") is not None else 0
+                seed_val = (
+                    int(env_cfg.get("seed"))
+                    if isinstance(env_cfg, dict) and env_cfg.get("seed") is not None
+                    else 0
+                )
             except Exception:
                 seed_val = 0
-        subject = (env_cfg.get("subject") if isinstance(env_cfg, dict) else None) or os.getenv("HENDRYCKS_MATH_CONFIG", "default")
+        subject = (env_cfg.get("subject") if isinstance(env_cfg, dict) else None) or os.getenv(
+            "HENDRYCKS_MATH_CONFIG", "default"
+        )
         # Load real Hendrycks problem text/solution (download if necessary). Crash on failure.
         qh, ah = _load_hendrycks_problem(seed_val, subject=subject)
         question = qh
@@ -468,7 +523,10 @@ def fastapi_app():
                     sanitized.pop("max_tokens", None)
                 for field in ("temperature", "top_p"):
                     sanitized.pop(field, None)
-                sanitized["tool_choice"] = {"type": "function", "function": {"name": "submit_answer"}}
+                sanitized["tool_choice"] = {
+                    "type": "function",
+                    "function": {"name": "submit_answer"},
+                }
                 sanitized["parallel_tool_calls"] = False
             return sanitized
 
@@ -503,25 +561,27 @@ def fastapi_app():
 
         user_prompt = (
             str(question)
-            if isinstance(question, (str, int, float)) and str(question).strip()
+            if isinstance(question, str | int | float) and str(question).strip()
             else "Solve the problem. Provide answer steps succinctly."
         )
         payload = {
             "model": model,
             "messages": [{"role": "user", "content": user_prompt}],
-            "tools": [{
-                "type": "function",
-                "function": {
-                    "name": "submit_answer",
-                    "parameters": {
-                        "type": "object",
-                        "properties": {
-                            "answer": {"type": "string"},
+            "tools": [
+                {
+                    "type": "function",
+                    "function": {
+                        "name": "submit_answer",
+                        "parameters": {
+                            "type": "object",
+                            "properties": {
+                                "answer": {"type": "string"},
+                            },
+                            "required": ["answer"],
                         },
-                        "required": ["answer"],
                     },
-                },
-            }],
+                }
+            ],
             "max_tokens": 256,
             "temperature": 0.2,
         }
@@ -529,13 +589,13 @@ def fastapi_app():
 
         try:
             tool_names = []
-            for t in (payload.get("tools") or []):
+            for t in payload.get("tools") or []:
                 if isinstance(t, dict):
                     fn = (t.get("function") or {}) if isinstance(t.get("function"), dict) else {}
                     name = fn.get("name")
                     if isinstance(name, str):
                         tool_names.append(name)
-            print(f"[math] system: <none>", flush=True)
+            print("[math] system: <none>", flush=True)
             print(f"[math] user: {user_prompt}", flush=True)
             print(f"[math] tools: {tool_names}", flush=True)
         except Exception:
@@ -547,7 +607,9 @@ def fastapi_app():
             if sk:
                 headers["Authorization"] = f"Bearer {sk}"
         with httpx.Client(timeout=httpx.Timeout(180.0), follow_redirects=True) as client:
-            resp = client.post(f"{inference_url}/v1/chat/completions", json=to_send, headers=headers)
+            resp = client.post(
+                f"{inference_url}/v1/chat/completions", json=to_send, headers=headers
+            )
             try:
                 data = resp.json()
             except Exception:
@@ -580,14 +642,21 @@ def fastapi_app():
 
         tool_answer = _parse_tool_answer(data)
         history.append({"answer": tool_answer})
-        steps.append({
-            "obs": {},
-            "tool_calls": [{"tool_name": "submit_answer", "arguments": _json.dumps({"answer": tool_answer})}],
-            "reward": None,
-            "done": False,
-            "truncated": False,
-            "info": None,
-        })
+        steps.append(
+            {
+                "obs": {},
+                "tool_calls": [
+                    {
+                        "tool_name": "submit_answer",
+                        "arguments": _json.dumps({"answer": tool_answer}),
+                    }
+                ],
+                "reward": None,
+                "done": False,
+                "truncated": False,
+                "info": None,
+            }
+        )
 
         # Evaluate answer correctness using tool output (or fall back to assistant text)
         reward_val = 0.0
@@ -605,25 +674,57 @@ def fastapi_app():
         except Exception:
             reward_val = 0.0
 
+        # Immediate, concise rollout logging mirroring RL format
+        try:
+            preview = tool_answer[:120] + (
+                "…" if isinstance(tool_answer, str) and len(tool_answer) > 120 else ""
+            )
+            components = {
+                "env": float(reward_val),
+                "rubric_event": 1.0 if bool(tool_answer.strip()) else 0.0,
+                "rubric_outcome": 1.0 if float(reward_val) > 0.0 else 0.0,
+            }
+            print(
+                "[MATH_ROLLOUT] run=",
+                run_id,
+                " seed=",
+                seed_val,
+                " subject=",
+                subject,
+                " tool=submit_answer answer=",
+                preview,
+                " reward=",
+                float(reward_val),
+                " components=",
+                components,
+                flush=True,
+            )
+        except Exception:
+            pass
+
         total_reward += float(reward_val)
-        steps.append({
-            "obs": {},
-            "tool_calls": [],
-            "reward": reward_val,
-            "done": True,
-            "truncated": False,
-            "info": None,
-        })
+        steps.append(
+            {
+                "obs": {},
+                "tool_calls": [],
+                "reward": reward_val,
+                "done": True,
+                "truncated": False,
+                "info": None,
+            }
+        )
 
         return {
             "run_id": run_id,
-            "trajectories": [{
-                "env_id": env_name,
-                "policy_id": (policy or {}).get("policy_name") or "math-react",
-                "steps": steps,
-                "final": {"observation": {}},
-                "length": len(steps),
-            }],
+            "trajectories": [
+                {
+                    "env_id": env_name,
+                    "policy_id": (policy or {}).get("policy_name") or "math-react",
+                    "steps": steps,
+                    "final": {"observation": {}},
+                    "length": len(steps),
+                }
+            ],
             "branches": {},
             "metrics": {
                 "episode_returns": [total_reward],

synth_ai/demos/demo_task_apps/math/task_app_entry.py

@@ -5,7 +5,6 @@ from __future__ import annotations
 from synth_ai.task.apps import ModalDeploymentConfig, TaskAppEntry, register_task_app
 from synth_ai.task.apps.math_single_step import build_config as base_build_config
 
-
 DEMO_MODAL_CONFIG = ModalDeploymentConfig(
     app_name="hendrycks-math-task-app",
     pip_packages=(
@@ -36,4 +35,3 @@ register_task_app(
         modal=DEMO_MODAL_CONFIG,
     )
 )
-

synth_ai/environments/examples/bandit/engine.py

@@ -191,7 +191,9 @@ class BanditEngine(StatefulEngine, IReproducibleEngine):
             step_count=self.step_count,
             max_steps=self.max_steps,
             last_arm=self.last_arm,
-            last_reward=float(reward) if reward is not None else (self.last_reward if self.step_count else None),
+            last_reward=float(reward)
+            if reward is not None
+            else (self.last_reward if self.step_count else None),
             cumulative_reward=float(self.total_reward),
             reward_history=self.reward_history.copy(),
             arm_pull_counts=self.arm_pull_counts.copy(),
@@ -238,7 +240,9 @@ class BanditEngine(StatefulEngine, IReproducibleEngine):
         engine.arm_probabilities = data.get("arm_probabilities", engine.arm_probabilities)
         engine.arm_means = data.get("arm_means", engine.arm_means)
         engine.arm_stds = data.get("arm_stds", engine.arm_stds)
-        engine.true_expected_rewards = list(data.get("true_expected_rewards", engine.true_expected_rewards))
+        engine.true_expected_rewards = list(
+            data.get("true_expected_rewards", engine.true_expected_rewards)
+        )
         engine.arm_count = len(engine.true_expected_rewards)
 
         engine.step_count = int(data.get("step_count", 0))
@@ -247,7 +251,9 @@ class BanditEngine(StatefulEngine, IReproducibleEngine):
         engine.last_arm = data.get("last_arm")
         engine.reward_history = list(data.get("reward_history", []))
         engine.arm_history = list(data.get("arm_history", []))
-        engine.arm_pull_counts = list(data.get("arm_pull_counts", [0 for _ in range(engine.arm_count)]))
+        engine.arm_pull_counts = list(
+            data.get("arm_pull_counts", [0 for _ in range(engine.arm_count)])
+        )
         engine.terminated = bool(data.get("terminated", False))
         engine.status = data.get("status", "in_progress")
 
@@ -287,7 +293,9 @@ class SynthBanditCheckpointObservationCallable(GetObservationCallable):
             "arm_count": pub.arm_count,
             "total_reward": priv.total_reward,
             "steps_taken": pub.step_count,
-            "best_expected_reward": max(priv.true_expected_rewards) if priv.true_expected_rewards else None,
+            "best_expected_reward": max(priv.true_expected_rewards)
+            if priv.true_expected_rewards
+            else None,
             "terminated": pub.terminated,
             "status": pub.status,
         }

synth_ai/environments/examples/bandit/taskset.py

@@ -156,10 +156,10 @@ async def create_bandit_taskset(
         )
 
         expected = _expected_rewards(metadata)
-        arm_count = len(expected) if expected else (
-            len(metadata.arm_probabilities or [])
-            or len(metadata.arm_means or [])
-            or 0
+        arm_count = (
+            len(expected)
+            if expected
+            else (len(metadata.arm_probabilities or []) or len(metadata.arm_means or []) or 0)
         )
         if arm_count == 0:
             arm_count = 1

synth_ai/environments/examples/crafter_classic/environment.py

@@ -2,11 +2,16 @@
 
 from __future__ import annotations
 
+import base64
 import dataclasses
 import logging
 import time
+from io import BytesIO
 from typing import Any, Dict, List, Optional, Union
 
+import numpy as np
+from PIL import Image
+
 # Import tracing abstractions
 from synth_ai.tracing_v3.abstractions import (
     RuntimeEvent,
@@ -43,6 +48,51 @@ from synth_ai.environments.reproducibility.core import ReproducibleEnvironment
 from synth_ai.environments.stateful.core import StatefulEnvironment
 
 
+def _convert_numpy_to_python(obj: Any) -> Any:
+    if isinstance(obj, np.integer):
+        return int(obj)
+    if isinstance(obj, np.floating):
+        return float(obj)
+    if isinstance(obj, np.ndarray):
+        return obj.tolist()
+    if isinstance(obj, dict):
+        return {k: _convert_numpy_to_python(v) for k, v in obj.items()}
+    if isinstance(obj, (list, tuple)):
+        return [_convert_numpy_to_python(item) for item in obj]
+    return obj
+
+
+def _encode_image_to_base64(image_array: Any) -> dict[str, Any] | None:
+    if not isinstance(image_array, np.ndarray):
+        return None
+    if image_array.ndim != 3 or image_array.shape[-1] not in (1, 3, 4):
+        return None
+    try:
+        array_uint8 = (
+            image_array.astype("uint8")
+            if image_array.dtype != np.uint8
+            else image_array  # pragma: no cover - fast path
+        )
+        mode = "L" if array_uint8.shape[-1] == 1 else "RGB"
+        if array_uint8.shape[-1] == 4:
+            mode = "RGBA"
+        img = Image.fromarray(array_uint8, mode=mode)
+        buffer = BytesIO()
+        img.save(buffer, format="PNG")
+        encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
+        width = int(array_uint8.shape[1])
+        height = int(array_uint8.shape[0])
+        return {
+            "format": "png",
+            "width": width,
+            "height": height,
+            "data": encoded,
+            "data_url": f"data:image/png;base64,{encoded}",
+        }
+    except Exception:
+        return None
+
+
 # --- Tool Definition ---
 class CrafterActionInput(BaseModel):
     action: int = Field(..., description="Integer action for the Crafter environment.")
@@ -362,7 +412,8 @@ class CrafterClassicEnvironment(StatefulEnvironment, ReproducibleEnvironment[Cra
         state_before = {"private_state": priv, "public_state": pub}
 
         active_obs_cb = obs_cb or SynthCrafterObservationCallable()
-        observation = await active_obs_cb.get_observation(pub, priv)
+        raw_observation = await active_obs_cb.get_observation(pub, priv)
+        observation = self._prepare_observation(raw_observation)
         if extra_obs and isinstance(observation, dict):
             observation.update(extra_obs)
 
@@ -385,6 +436,30 @@ class CrafterClassicEnvironment(StatefulEnvironment, ReproducibleEnvironment[Cra
 
         return observation
 
+    def _prepare_observation(self, observation: Any) -> dict[str, Any]:
+        obs_dict: dict[str, Any]
+        image_payload: dict[str, Any] | None = None
+
+        if isinstance(observation, dict):
+            image_payload = _encode_image_to_base64(observation.get("observation_image"))
+            sanitized = dict(observation)
+            sanitized.pop("observation_image", None)
+            obs_dict = _convert_numpy_to_python(sanitized) or {}
+        else:
+            obs_dict = _convert_numpy_to_python(observation) or {}
+
+        if not isinstance(obs_dict, dict):
+            obs_dict = {"value": obs_dict}
+
+        if image_payload:
+            obs_dict["observation_image_base64"] = image_payload["data"]
+            obs_dict["observation_image_format"] = image_payload["format"]
+            obs_dict["observation_image_width"] = image_payload["width"]
+            obs_dict["observation_image_height"] = image_payload["height"]
+            obs_dict["observation_image_data_url"] = image_payload["data_url"]
+
+        return obs_dict
+
     # ────────────────────────────────────────────────────────────────────
     # ReproducibleEnvironment plumbing
     # ────────────────────────────────────────────────────────────────────