synth-ai 0.2.4.dev3__py3-none-any.whl → 0.2.4.dev5__py3-none-any.whl

synth_ai/environments/examples/crafter_classic/environment.py

@@ -0,0 +1,364 @@
+"""CrafterClassicEnvironment — thin wrapper exposing CrafterEngine via StatefulEnvironment API."""
+
+from __future__ import annotations
+
+from typing import List, Optional, Any, Dict, Union
+import dataclasses
+import logging
+import time
+
+# Import logging configuration to suppress JAX debug messages
+from .config_logging import safe_compare
+
+# Import tracing abstractions
+from synth_ai.tracing_v3.abstractions import (
+    RuntimeEvent,
+    SessionEventMarkovBlanketMessage,
+    TimeRecord,
+)
+
+logger = logging.getLogger(__name__)
+
+from synth_ai.environments.examples.crafter_classic.engine import (
+    CrafterEngine,
+    CrafterPrivateState,
+    CrafterPublicState,
+    CrafterEngineSnapshot,
+)
+from synth_ai.environments.examples.crafter_classic.taskset import CrafterTaskInstance
+from synth_ai.environments.environment.shared_engine import (
+    GetObservationCallable,
+    InternalObservation,
+)
+from synth_ai.environments.reproducibility.core import ReproducibleEnvironment
+from synth_ai.environments.stateful.core import StatefulEnvironment
+from synth_ai.environments.environment.tools import (
+    AbstractTool,
+    EnvToolCall,
+    ToolResult,
+    TOOL_REGISTRY,
+    register_tool,
+)
+from pydantic import BaseModel, Field
+
+
+# --- Tool Definition ---
+class CrafterActionInput(BaseModel):
+    action: int = Field(..., description="Integer action for the Crafter environment.")
+
+
+class CrafterInteractTool(AbstractTool):
+    name = "interact"
+    description = "Performs an action in the Crafter environment."
+    call_schema = CrafterActionInput
+    result_schema = ToolResult
+
+    def __init__(self, engine: CrafterEngine, session_tracer: Optional[Any] = None):
+        self.engine = engine
+        self.session_tracer = session_tracer
+
+    async def __call__(self, call: EnvToolCall) -> ToolResult:
+        try:
+            # Store state before execution
+            state_before = {"action_args": call.args}
+
+            validated_args = self.call_schema(**call.args)
+            action_to_pass = self.engine._validate_action_engine(validated_args.action)
+
+            # Execute the engine step
+            priv_state, pub_state = await self.engine._step_engine(action_to_pass)
+
+            # Store state after execution
+            state_after = {
+                "engine_result": {"private_state": priv_state, "public_state": pub_state}
+            }
+
+            # Record runtime event for tool execution
+            if (
+                self.session_tracer
+                and hasattr(self.session_tracer, "current_session")
+                and self.session_tracer.current_session
+            ):
+                runtime_execution_event = RuntimeEvent()
+                runtime_execution_event.time_record = TimeRecord()
+                runtime_execution_event.time_record.event_time = time.time()
+                runtime_execution_event.time_record.message_time = None
+                runtime_execution_event.system_instance_id = "crafter_interact_tool"
+                runtime_execution_event.system_state_before = state_before
+                runtime_execution_event.system_state_after = state_after
+                runtime_execution_event.actions = [action_to_pass]
+                runtime_execution_event.metadata = {"execution_step": "engine_action"}
+                # Add directly to event history, bypassing timestep requirement
+                self.session_tracer.current_session.add_event(runtime_execution_event)
+
+            return ToolResult(
+                ok=True,
+                payload={
+                    "public_state": pub_state,
+                    "private_state": priv_state,
+                },
+            )
+        except Exception as e:
+            pub_state_on_error = self.engine._get_public_state_from_env()  # Use engine helper
+            # Get a safe private state for error cases
+            health_dead = safe_compare(0, self.engine.env._player.health, ">=")
+            step_exceeded = safe_compare(self.engine.env._length, self.engine.env._step, "<=")
+            priv_state_on_error = self.engine._get_private_state_from_env(
+                0, health_dead, step_exceeded
+            )
+            return ToolResult(
+                ok=False,
+                error=str(e),
+                payload={
+                    "public_state": pub_state_on_error,
+                    "private_state": priv_state_on_error,
+                },
+            )
+
+
+# Default observation callable (can be customized via __init__)
+class SynthCrafterObservationCallable(GetObservationCallable):
+    async def get_observation(
+        self, pub: CrafterPublicState, priv: CrafterPrivateState
+    ) -> InternalObservation:
+        # Example: return a dictionary combining public and selected private info
+        # Actual observation structure depends on agent's needs.
+        obs_dict: Dict[str, Any] = dataclasses.asdict(pub)  # type: ignore
+        obs_dict["reward_last_step"] = priv.reward_last_step
+        obs_dict["total_reward_episode"] = priv.total_reward_episode
+        obs_dict["terminated"] = priv.terminated
+        obs_dict["truncated"] = priv.truncated
+        if pub.error_info:
+            obs_dict["tool_error"] = pub.error_info
+        return obs_dict
+
+
+class CrafterClassicEnvironment(StatefulEnvironment, ReproducibleEnvironment[CrafterEngine]):
+    """Environment wrapper bridging agent tool‑calls to `crafter.Env` dynamics."""
+
+    def __init__(
+        self,
+        task_instance: "CrafterTaskInstance",
+        custom_step_obs: Optional[GetObservationCallable] = None,
+        custom_ckpt_obs: Optional[GetObservationCallable] = None,
+        session_tracer: Optional[Any] = None,  # SessionTracer from higher level
+    ) -> None:
+        self.name = "CrafterClassic"
+        self.task_instance = task_instance
+        self.custom_step_observation_callable = custom_step_obs or SynthCrafterObservationCallable()
+        self.custom_checkpoint_observation_callable = (
+            custom_ckpt_obs or SynthCrafterObservationCallable()
+        )
+        self.engine = CrafterEngine(task_instance)
+        self.session_tracer = session_tracer  # Store tracer for runtime events
+
+        self._interact_tool = CrafterInteractTool(self.engine, session_tracer=session_tracer)
+        if self._interact_tool.name not in TOOL_REGISTRY:
+            register_tool(self._interact_tool)
+
+    # ────────────────────────────────────────────────────────────────────
+    # Lifecycle helpers
+    # ────────────────────────────────────────────────────────────────────
+
+    async def initialize(self, seed: Optional[int] = None) -> InternalObservation:  # type: ignore[override]
+        # Check if seed was provided in task instance metadata
+        if (
+            seed is None
+            and hasattr(self.task_instance, "metadata")
+            and hasattr(self.task_instance.metadata, "seed")
+        ):
+            seed = self.task_instance.metadata.seed
+        # Check if seed was provided in initial_engine_snapshot
+        elif (
+            seed is None
+            and hasattr(self.task_instance, "initial_engine_snapshot")
+            and isinstance(self.task_instance.initial_engine_snapshot, dict)
+        ):
+            seed = self.task_instance.initial_engine_snapshot.get("seed")
+
+        # Initialize with seed from various sources
+
+        priv, pub = await self.engine._reset_engine(seed=seed)
+        return await self._to_observation(priv, pub, self.custom_step_observation_callable)
+
+    async def terminate(self) -> InternalObservation:  # type: ignore[override]
+        pub = self.engine._get_public_state_from_env()
+        priv = self.engine._get_private_state_from_env(0, True, False)  # Terminated state
+        priv.terminated = True
+        obs_dict = {"status": "Environment terminated."}
+        return await self._to_observation(
+            priv, pub, self.custom_step_observation_callable, extra_obs=obs_dict
+        )
+
+    # ────────────────────────────────────────────────────────────────────
+    # Step + checkpoint
+    # ────────────────────────────────────────────────────────────────────
+
+    def validate_tool_calls(
+        self, tool_calls: Union[EnvToolCall, List[EnvToolCall], List[List[EnvToolCall]]]
+    ) -> EnvToolCall:
+        # Store the original tool calls for tracing
+        state_before = {"tool_calls": tool_calls}
+
+        # Normalize and validate to a single EnvToolCall (same as Sokoban)
+        if isinstance(tool_calls, list):
+            if not tool_calls:
+                raise ValueError("Received empty list of tool calls.")
+            if isinstance(tool_calls[0], list):
+                if not tool_calls[0]:
+                    raise ValueError("Received empty inner list of tool calls.")
+                agent_call = tool_calls[0][0]
+            else:
+                agent_call = tool_calls[0]
+        elif isinstance(tool_calls, EnvToolCall):
+            agent_call = tool_calls
+        else:
+            raise TypeError(f"Unexpected type for tool_calls: {type(tool_calls)}")
+
+        if not isinstance(agent_call, EnvToolCall):
+            raise TypeError(f"Processed call is not EnvToolCall: {type(agent_call)}")
+        if agent_call.tool != "interact":
+            raise ValueError(f"Unknown tool: {agent_call.tool}. Expected 'interact'.")
+
+        # Record runtime event for tool call validation
+        if (
+            self.session_tracer
+            and hasattr(self.session_tracer, "current_session")
+            and self.session_tracer.current_session
+        ):
+            runtime_validation_event = RuntimeEvent()
+            runtime_validation_event.time_record = TimeRecord()
+            runtime_validation_event.time_record.event_time = time.time()
+            runtime_validation_event.time_record.message_time = None
+            runtime_validation_event.system_instance_id = "crafter_environment"
+            runtime_validation_event.system_state_before = state_before
+            runtime_validation_event.system_state_after = {"validated_call": agent_call}
+            runtime_validation_event.metadata = {"validation_step": "tool_call_validation"}
+            # Add directly to event history, bypassing timestep requirement
+            self.session_tracer.current_session.add_event(runtime_validation_event)
+
+        return agent_call
+
+    async def step(
+        self, tool_calls: Union[EnvToolCall, List[EnvToolCall], List[List[EnvToolCall]]]
+    ) -> InternalObservation:  # type: ignore[override]
+        step_start_time = time.time()
+        agent_call = self.validate_tool_calls(tool_calls)
+        interact_start = time.time()
+        tool_result: ToolResult = await self._interact_tool(agent_call)
+        interact_time = time.time() - interact_start
+
+        payload_dict = tool_result.payload
+        pub_state: CrafterPublicState
+        priv_state: CrafterPrivateState
+
+        if tool_result.ok:
+            # payload contains the actual state objects from the interact tool
+            priv_state = payload_dict.get("private_state")
+            pub_state = payload_dict.get("public_state")
+
+            # Validate we got the expected state objects
+            if not isinstance(priv_state, CrafterPrivateState) or not isinstance(
+                pub_state, CrafterPublicState
+            ):
+                logger.error(
+                    f"Invalid state types in payload: priv={type(priv_state)}, pub={type(pub_state)}"
+                )
+                # Fall back to getting current state
+                pub_state = self.engine._get_public_state_from_env()
+                health_dead = safe_compare(0, self.engine.env._player.health, ">=")
+                step_exceeded = safe_compare(self.engine.env._length, self.engine.env._step, "<=")
+                priv_state = self.engine._get_private_state_from_env(0, health_dead, step_exceeded)
+                pub_state.error_info = "Invalid state types in tool result"
+        else:
+            # Tool call failed, use states from payload if available, otherwise get current state
+            priv_state = payload_dict.get("private_state")
+            pub_state = payload_dict.get("public_state")
+
+            if not isinstance(priv_state, CrafterPrivateState) or not isinstance(
+                pub_state, CrafterPublicState
+            ):
+                # Fall back to getting current state
+                pub_state = self.engine._get_public_state_from_env()
+                health_dead = safe_compare(0, self.engine.env._player.health, ">=")
+                step_exceeded = safe_compare(self.engine.env._length, self.engine.env._step, "<=")
+                priv_state = self.engine._get_private_state_from_env(0, health_dead, step_exceeded)
+
+            if tool_result.error:
+                pub_state.error_info = tool_result.error
+
+        obs = await self._to_observation(
+            priv_state, pub_state, self.custom_step_observation_callable
+        )
+        total_step_time = time.time() - step_start_time
+        logger.info(
+            f"CrafterClassic step completed in {total_step_time:.3f}s (interact: {interact_time:.3f}s)"
+        )
+        return obs
+
+    async def checkpoint(self) -> InternalObservation:  # type: ignore[override]
+        engine_snapshot: CrafterEngineSnapshot = await self.engine._serialize_engine()
+        priv = self.engine._get_private_state_from_env(0, False, False)  # Get current state for obs
+        pub = self.engine._get_public_state_from_env()
+        obs_data = await self._to_observation(
+            priv, pub, self.custom_checkpoint_observation_callable
+        )
+        if isinstance(obs_data, dict):
+            obs_data["engine_snapshot_data"] = engine_snapshot.model_dump()
+        return obs_data
+
+    # ────────────────────────────────────────────────────────────────────
+    # Helpers
+    # ────────────────────────────────────────────────────────────────────
+
+    async def _to_observation(
+        self,
+        priv: CrafterPrivateState,
+        pub: CrafterPublicState,
+        obs_cb: Optional[GetObservationCallable],
+        extra_obs: Optional[Dict[str, Any]] = None,
+    ) -> InternalObservation:
+        # Store state before observation generation
+        state_before = {"private_state": priv, "public_state": pub}
+
+        active_obs_cb = obs_cb or SynthCrafterObservationCallable()
+        observation = await active_obs_cb.get_observation(pub, priv)
+        if extra_obs and isinstance(observation, dict):
+            observation.update(extra_obs)
+
+        # Record runtime event for observation generation
+        if (
+            self.session_tracer
+            and hasattr(self.session_tracer, "current_session")
+            and self.session_tracer.current_session
+        ):
+            runtime_obs_event = RuntimeEvent()
+            runtime_obs_event.time_record = TimeRecord()
+            runtime_obs_event.time_record.event_time = time.time()
+            runtime_obs_event.time_record.message_time = None
+            runtime_obs_event.system_instance_id = "observation_generator"
+            runtime_obs_event.system_state_before = state_before
+            runtime_obs_event.system_state_after = {"observation": observation}
+            runtime_obs_event.metadata = {"observation_step": "state_to_obs_conversion"}
+            # Add directly to event history, bypassing timestep requirement
+            self.session_tracer.current_session.add_event(runtime_obs_event)
+
+        return observation
+
+    # ────────────────────────────────────────────────────────────────────
+    # ReproducibleEnvironment plumbing
+    # ────────────────────────────────────────────────────────────────────
+
+    async def _serialize_engine(self) -> CrafterEngineSnapshot:
+        return await self.engine._serialize_engine()
+
+    @classmethod
+    async def _deserialize_engine(
+        cls, snapshot: CrafterEngineSnapshot, task_instance: "CrafterTaskInstance"
+    ) -> "CrafterClassicEnvironment":
+        eng = await CrafterEngine._deserialize_engine(snapshot, task_instance)
+        env = cls(task_instance)
+        env.engine = eng
+        # CRITICAL: Update the interact tool to use the new engine!
+        env._interact_tool.engine = eng
+        return env

synth_ai/environments/examples/crafter_classic/taskset.py

@@ -0,0 +1,233 @@
+"""Procedural Crafter taskset generation with seed filtering by world traits.
+Run this to build a TaskInstanceSet with reproducible initial snapshots.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import random
+from dataclasses import dataclass, asdict, fields
+from typing import Dict, List
+from uuid import UUID, uuid4
+
+import numpy as np
+import crafter
+from crafter import objects
+
+from synth_ai.environments.tasks.core import (
+    Impetus,
+    Intent,
+    SplitInfo,
+    Task,
+    TaskInstance,
+    TaskInstanceMetadata,
+    TaskInstanceSet,
+)
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Config
+# ──────────────────────────────────────────────────────────────────────────────
+TASK = Task(
+    global_premises="Procedural Crafter seed generation",
+    global_constraints="",
+    global_objectives="Survive and unlock achievements.",
+    shared_env_params={},
+)
+
+AREA = (64, 64)
+LEN = 10000
+RADIUS = 10  # Manhattan distance for local trait count
+SEED_START = 0
+NUM_INSTANCES = 50
+
+# Desired trait ranges per difficulty tier
+TRAIT_BOUNDS = {
+    "easy": {
+        "min_trees": 4,
+        "max_hostiles": 0,
+    },
+    "medium": {
+        "min_trees": 2,
+        "max_hostiles": 2,
+    },
+    "hard": {
+        "min_trees": 0,
+        "max_hostiles": 5,
+    },
+}
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Metadata + instance helpers
+# ──────────────────────────────────────────────────────────────────────────────
+
+
+from typing import Optional
+
+
+@dataclass
+class CrafterTaskInstanceMetadata(TaskInstanceMetadata):
+    difficulty: str
+    seed: int
+    num_trees_radius: int
+    num_cows_radius: int
+    num_hostiles_radius: int
+    world_config: Optional[str] = "normal"  # 'easy', 'normal', 'hard', 'peaceful'
+    world_config_path: Optional[str] = None  # Path to custom JSON config
+
+
+@dataclass
+class CrafterTaskInstance(TaskInstance):
+    async def serialize(self) -> dict:  # identical to Sokoban pattern
+        data = asdict(self)
+        if isinstance(data.get("id"), UUID):
+            data["id"] = str(data["id"])
+        if "intent" in data and data["intent"] is not None:
+            data["intent"]["deterministic_eval_functions"] = []
+        return data
+
+    @classmethod
+    async def deserialize(cls, data: dict) -> "CrafterTaskInstance":
+        if "id" in data:
+            try:
+                data["id"] = UUID(str(data["id"]))
+            except Exception:
+                pass
+        if "impetus" in data and isinstance(data["impetus"], dict):
+            impetus_data = data["impetus"]
+            # Ensure instructions field exists with default if missing
+            if "instructions" not in impetus_data:
+                impetus_data["instructions"] = "Survive and unlock achievements"
+            data["impetus"] = Impetus(**impetus_data)
+        if "intent" in data and isinstance(data["intent"], dict):
+            intent_data = data["intent"]
+            # Ensure required fields exist with defaults if missing
+            if "rubric" not in intent_data:
+                intent_data["rubric"] = {"goal": "Unlock achievements"}
+            if "gold_trajectories" not in intent_data:
+                intent_data["gold_trajectories"] = None
+            if "gold_state_diff" not in intent_data:
+                intent_data["gold_state_diff"] = {}
+            intent_data["deterministic_eval_functions"] = []
+            data["intent"] = Intent(**intent_data)
+        if "metadata" in data and isinstance(data["metadata"], dict):
+            metadata_data = data["metadata"]
+            # Ensure required fields exist with defaults if missing
+            if "difficulty" not in metadata_data:
+                metadata_data["difficulty"] = "medium"
+            if "seed" not in metadata_data:
+                metadata_data["seed"] = 0
+            if "num_trees_radius" not in metadata_data:
+                metadata_data["num_trees_radius"] = 0
+            if "num_cows_radius" not in metadata_data:
+                metadata_data["num_cows_radius"] = 0
+            if "num_hostiles_radius" not in metadata_data:
+                metadata_data["num_hostiles_radius"] = 0
+            data["metadata"] = CrafterTaskInstanceMetadata(**metadata_data)
+        keep = {f.name for f in fields(cls)}
+        return cls(**{k: v for k, v in data.items() if k in keep})
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Trait extraction util
+# ──────────────────────────────────────────────────────────────────────────────
+
+
+def world_traits(env: crafter.Env, radius: int = RADIUS) -> Dict[str, int]:
+    player = env._player  # type: ignore[attr-defined]
+    pos = np.array(player.pos)
+    counts = {"trees": 0, "cows": 0, "hostiles": 0}
+    for obj in env._world._objects:  # type: ignore[attr-defined]
+        if obj is None or obj is player:
+            continue
+        if np.abs(obj.pos - pos).sum() > radius:
+            continue
+        if isinstance(obj, objects.Plant) and getattr(obj, "kind", "") == "tree":
+            counts["trees"] += 1
+        elif isinstance(obj, objects.Cow):
+            counts["cows"] += 1
+        elif isinstance(obj, (objects.Zombie, objects.Skeleton)):
+            counts["hostiles"] += 1
+    return counts
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Main generator
+# ──────────────────────────────────────────────────────────────────────────────
+
+
+async def create_crafter_taskset(num_instances: int = NUM_INSTANCES) -> TaskInstanceSet:
+    instances: List[CrafterTaskInstance] = []
+    seed = SEED_START
+    while len(instances) < num_instances:
+        env = crafter.Env(area=AREA, length=LEN, seed=seed)
+        _ = env.reset()
+        traits = world_traits(env)
+        # assign difficulty tier first match
+        difficulty: str | None = None
+        for diff, bounds in TRAIT_BOUNDS.items():
+            if (
+                traits["trees"] >= bounds["min_trees"]
+                and traits["hostiles"] <= bounds["max_hostiles"]
+            ):
+                difficulty = diff
+                break
+        if difficulty is None:
+            seed += 1
+            continue
+        # build instance
+        impetus = Impetus(instructions=f"Survive and unlock achievements. Difficulty={difficulty}.")
+        intent = Intent(
+            rubric={"goal": "Unlock as many achievements as possible."},
+            gold_trajectories=None,
+            gold_state_diff={},
+        )
+        metadata = CrafterTaskInstanceMetadata(
+            difficulty=difficulty,
+            seed=seed,
+            num_trees_radius=traits["trees"],
+            num_cows_radius=traits["cows"],
+            num_hostiles_radius=traits["hostiles"],
+        )
+        instance = CrafterTaskInstance(
+            id=uuid4(),
+            impetus=impetus,
+            intent=intent,
+            metadata=metadata,
+            is_reproducible=True,
+            initial_engine_snapshot=None,  # will be filled lazily when env starts
+        )
+        instances.append(instance)
+        seed += 1
+
+    # simple random split 80/10/10
+    random.shuffle(instances)
+    n = len(instances)
+    val_ids = {inst.id for inst in instances[int(0.8 * n) : int(0.9 * n)]}
+    test_ids = {inst.id for inst in instances[int(0.9 * n) :]}
+    split = SplitInfo(val_instance_ids=val_ids, test_instance_ids=test_ids, _is_split_defined=True)
+
+    return TaskInstanceSet(
+        name="Crafter Procedural TaskSet",
+        description="Crafter seeds filtered by local world traits around spawn.",
+        instances=instances,
+        split_info=split,
+    )
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# CLI example
+# ──────────────────────────────────────────────────────────────────────────────
+
+if __name__ == "__main__":
+    import json
+    import pathlib
+
+    async def _main():
+        ts = await create_crafter_taskset(30)
+        serial = await asyncio.gather(*(inst.serialize() for inst in ts.instances))
+        out = pathlib.Path("dataset/crafter_instances.json")
+        out.parent.mkdir(parents=True, exist_ok=True)
+        out.write_text(json.dumps(serial, indent=2))
+        print(f"Saved {len(serial)} instances → {out}")
+
+    asyncio.run(_main())