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

synth_ai/environments/examples/crafter_custom/environment.py

@@ -0,0 +1,312 @@
+"""CrafterCustomEnvironment — Custom Crafter with configurable world generation."""
+
+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 synth_ai.environments.examples.crafter_classic.config_logging import safe_compare
+
+# Import tracing abstractions
+from synth_ai.tracing_v3.abstractions import (
+    RuntimeEvent,
+    SessionEventMarkovBlanketMessage,
+    TimeRecord,
+)
+
+logger = logging.getLogger(__name__)
+
+# Import the base Crafter components
+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
+
+
+# Use the same tool and observation classes as CrafterClassic
+from synth_ai.environments.examples.crafter_classic.environment import (
+    CrafterActionInput,
+    CrafterInteractTool,
+    SynthCrafterObservationCallable,
+)
+
+
+class CrafterCustomEnvironment(StatefulEnvironment, ReproducibleEnvironment[CrafterEngine]):
+    """Custom Crafter environment with configurable world generation."""
+
+    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 = "CrafterCustom"
+        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()
+        )
+
+        # Ensure task instance has world configuration
+        if hasattr(task_instance, "metadata"):
+            logger.info(
+                f"Creating CrafterCustom with world_config: {getattr(task_instance.metadata, 'world_config', 'default')}"
+            )
+
+        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)
+
+        # Register tool with a unique name for this environment
+        tool_name = f"{self.name.lower()}_interact"
+        if tool_name not in TOOL_REGISTRY:
+            # Create a copy of the tool with the custom name
+            self._interact_tool.name = tool_name
+            register_tool(self._interact_tool)
+
+    # ────────────────────────────────────────────────────────────────────
+    # Lifecycle helpers
+    # ────────────────────────────────────────────────────────────────────
+
+    async def initialize(self) -> InternalObservation:  # type: ignore[override]
+        priv, pub = await self.engine._reset_engine()
+        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
+        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)}")
+
+        # Accept both "interact" and "craftercustom_interact"
+        if agent_call.tool not in ["interact", f"{self.name.lower()}_interact"]:
+            raise ValueError(
+                f"Unknown tool: {agent_call.tool}. Expected 'interact' or '{self.name.lower()}_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_custom_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]
+        agent_call = self.validate_tool_calls(tool_calls)
+        tool_result: ToolResult = await self._interact_tool(agent_call)
+
+        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
+
+        return await self._to_observation(
+            priv_state, pub_state, self.custom_step_observation_callable
+        )
+
+    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
+
+    async def get_metadata(self) -> Dict[str, Any]:
+        """Get metadata about the current environment configuration."""
+        metadata = {
+            "environment_type": "CrafterCustom",
+            "engine_seed": getattr(self.engine.env, "_seed", None),
+            "world_area": self.engine.env._area,
+            "max_steps": self.engine.env._length,
+            "current_step": self.engine.env._step,
+        }
+
+        # Add task instance metadata
+        if hasattr(self.task_instance, "metadata"):
+            task_metadata = self.task_instance.metadata
+            metadata.update(
+                {
+                    "difficulty": getattr(task_metadata, "difficulty", None),
+                    "world_config": getattr(task_metadata, "world_config", None),
+                    "world_config_path": getattr(task_metadata, "world_config_path", None),
+                    "num_trees_radius": getattr(task_metadata, "num_trees_radius", None),
+                    "num_cows_radius": getattr(task_metadata, "num_cows_radius", None),
+                    "num_hostiles_radius": getattr(task_metadata, "num_hostiles_radius", None),
+                }
+            )
+
+        # Add current world statistics
+        if hasattr(self.engine, "env") and hasattr(self.engine.env, "_world"):
+            world = self.engine.env._world
+            object_counts = {}
+
+            for obj in world._objects:
+                if obj is None:
+                    continue
+                obj_type = type(obj).__name__
+                object_counts[obj_type] = object_counts.get(obj_type, 0) + 1
+
+            metadata["world_object_counts"] = object_counts
+
+        return metadata
+
+    # ────────────────────────────────────────────────────────────────────
+    # 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"
+    ) -> "CrafterCustomEnvironment":
+        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_custom/run_dataset.py

@@ -0,0 +1,305 @@
+#!/usr/bin/env python3
+"""
+Run script for Crafter dataset instances
+"""
+
+import json
+import argparse
+import random
+from pathlib import Path
+from typing import List, Optional, Dict, Any
+import uuid
+import os
+import sys
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))))
+
+from crafter import Env
+
+
+class CrafterDatasetRunner:
+    """Run Crafter instances from a dataset"""
+
+    def __init__(self, dataset_path: Path = Path("dataset")):
+        self.dataset_path = dataset_path
+
+    def load_dataset(self, dataset_name: str) -> Dict[str, Any]:
+        """Load a dataset from disk"""
+        dataset_dir = self.dataset_path / dataset_name
+
+        # Load metadata
+        with open(dataset_dir / "metadata.json", "r") as f:
+            metadata = json.load(f)
+
+        # Load instances
+        with open(dataset_dir / "instances.json", "r") as f:
+            instances = json.load(f)
+
+        return {"metadata": metadata, "instances": instances}
+
+    def filter_instances(
+        self,
+        instances: List[Dict[str, Any]],
+        difficulties: Optional[List[str]] = None,
+        impetus_types: Optional[List[str]] = None,
+        split: Optional[str] = None,
+        split_info: Optional[Dict[str, Any]] = None,
+    ) -> List[Dict[str, Any]]:
+        """Filter instances based on criteria"""
+        filtered = instances
+
+        # Filter by difficulty
+        if difficulties:
+            filtered = [inst for inst in filtered if inst["metadata"]["difficulty"] in difficulties]
+
+        # Filter by impetus type
+        if impetus_types:
+            filtered = [inst for inst in filtered if self._get_impetus_type(inst) in impetus_types]
+
+        # Filter by split
+        if split and split_info:
+            if split == "train":
+                val_ids = set(split_info["val_instance_ids"])
+                test_ids = set(split_info["test_instance_ids"])
+                filtered = [
+                    inst
+                    for inst in filtered
+                    if inst["id"] not in val_ids and inst["id"] not in test_ids
+                ]
+            elif split == "val":
+                val_ids = set(split_info["val_instance_ids"])
+                filtered = [inst for inst in filtered if inst["id"] in val_ids]
+            elif split == "test":
+                test_ids = set(split_info["test_instance_ids"])
+                filtered = [inst for inst in filtered if inst["id"] in test_ids]
+
+        return filtered
+
+    def _get_impetus_type(self, instance: Dict[str, Any]) -> str:
+        """Determine impetus type from instructions"""
+        instructions = instance["impetus"]["instructions"].lower()
+        if "speedrun" in instructions:
+            return "speedrun"
+        elif "focus on" in instructions:
+            return "focused"
+        else:
+            return "general"
+
+    def run_instance(
+        self, instance: Dict[str, Any], render: bool = False, max_steps: int = 1000, agent_fn=None
+    ):
+        """Run a single instance"""
+
+        # Extract parameters
+        difficulty = instance["metadata"]["difficulty"]
+        seed = instance["metadata"]["world_seed"]
+
+        print(f"\n{'=' * 60}")
+        print(f"Running instance: {instance['id']}")
+        print(f"Difficulty: {difficulty}")
+        print(f"Seed: {seed}")
+        print(f"Instructions: {instance['impetus']['instructions']}")
+        if instance["impetus"].get("achievement_focus"):
+            print(f"Focus: {', '.join(instance['impetus']['achievement_focus'])}")
+        print(f"{'=' * 60}")
+
+        # Create environment
+        env = Env(seed=seed, world_config=difficulty)
+
+        obs = env.reset()
+
+        # Run agent or random policy
+        total_reward = 0
+        achievements = set()
+
+        for step in range(max_steps):
+            if agent_fn:
+                action = agent_fn(obs, instance)
+            else:
+                action = env.action_space.sample()
+
+            obs, reward, done, info = env.step(action)
+            total_reward += reward
+
+            # Track achievements
+            if "achievements" in info:
+                for ach, unlocked in info["achievements"].items():
+                    if unlocked:
+                        achievements.add(ach)
+
+            if done:
+                break
+
+        # Evaluate based on intent
+        success = self._evaluate_instance(instance, achievements, total_reward, step)
+
+        print(f"\nResults:")
+        print(f"Steps: {step}")
+        print(f"Total reward: {total_reward}")
+        print(f"Achievements: {len(achievements)} - {list(achievements)}")
+        print(f"Success: {success}")
+
+        return {
+            "instance_id": instance["id"],
+            "difficulty": difficulty,
+            "seed": seed,
+            "steps": step,
+            "total_reward": total_reward,
+            "achievements": list(achievements),
+            "success": success,
+        }
+
+    def _evaluate_instance(
+        self, instance: Dict[str, Any], achievements: set, total_reward: float, steps: int
+    ) -> bool:
+        """Evaluate if instance was successful based on intent"""
+        intent = instance["intent"]
+
+        # Check minimum score
+        if intent.get("minimum_score"):
+            if len(achievements) < intent["minimum_score"]:
+                return False
+
+        # Check target achievements
+        if intent.get("target_achievements"):
+            targets = set(intent["target_achievements"])
+            if not achievements.intersection(targets):
+                return False
+
+        return True
+
+    def run_batch(
+        self,
+        dataset_name: str,
+        num_instances: int = 10,
+        difficulties: Optional[List[str]] = None,
+        impetus_types: Optional[List[str]] = None,
+        split: Optional[str] = None,
+        render: bool = False,
+        max_steps: int = 1000,
+        agent_fn=None,
+    ):
+        """Run a batch of instances"""
+
+        # Load dataset
+        dataset = self.load_dataset(dataset_name)
+        instances = dataset["instances"]
+
+        # Filter instances
+        filtered = self.filter_instances(
+            instances,
+            difficulties=difficulties,
+            impetus_types=impetus_types,
+            split=split,
+            split_info=dataset["metadata"].get("split_info"),
+        )
+
+        if not filtered:
+            print("No instances match the filter criteria!")
+            return []
+
+        # Sample instances
+        if num_instances > len(filtered):
+            print(f"Only {len(filtered)} instances available, running all")
+            selected = filtered
+        else:
+            selected = random.sample(filtered, num_instances)
+
+        print(f"\nRunning {len(selected)} instances from {dataset_name}")
+        print(f"Difficulties: {difficulties or 'all'}")
+        print(f"Impetus types: {impetus_types or 'all'}")
+        print(f"Split: {split or 'all'}")
+
+        # Run instances
+        results = []
+        for instance in selected:
+            result = self.run_instance(
+                instance, render=render, max_steps=max_steps, agent_fn=agent_fn
+            )
+            results.append(result)
+
+        # Summary statistics
+        self._print_summary(results)
+
+        return results
+
+    def _print_summary(self, results: List[Dict[str, Any]]):
+        """Print summary statistics"""
+        print(f"\n{'=' * 60}")
+        print("SUMMARY")
+        print(f"{'=' * 60}")
+
+        # Group by difficulty
+        by_difficulty = {}
+        for result in results:
+            diff = result["difficulty"]
+            if diff not in by_difficulty:
+                by_difficulty[diff] = []
+            by_difficulty[diff].append(result)
+
+        print(f"\nResults by difficulty:")
+        print(
+            f"{'Difficulty':<15} {'Count':<8} {'Success':<10} {'Avg Steps':<12} {'Avg Achievements'}"
+        )
+        print("-" * 60)
+
+        for diff in sorted(by_difficulty.keys()):
+            diff_results = by_difficulty[diff]
+            count = len(diff_results)
+            success_rate = sum(1 for r in diff_results if r["success"]) / count
+            avg_steps = sum(r["steps"] for r in diff_results) / count
+            avg_achievements = sum(len(r["achievements"]) for r in diff_results) / count
+
+            print(
+                f"{diff:<15} {count:<8} {success_rate:<10.1%} {avg_steps:<12.1f} {avg_achievements:.1f}"
+            )
+
+        # Overall stats
+        total_success = sum(1 for r in results if r["success"])
+        print(
+            f"\nOverall success rate: {total_success}/{len(results)} ({total_success / len(results):.1%})"
+        )
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Run Crafter dataset instances")
+    parser.add_argument("dataset", help="Dataset name")
+    parser.add_argument(
+        "-n", "--num-instances", type=int, default=10, help="Number of instances to run"
+    )
+    parser.add_argument(
+        "-d",
+        "--difficulties",
+        nargs="+",
+        choices=["easy", "normal", "hard", "peaceful", "resource_rich"],
+        help="Filter by difficulties",
+    )
+    parser.add_argument(
+        "-t",
+        "--impetus-types",
+        nargs="+",
+        choices=["general", "focused", "speedrun"],
+        help="Filter by impetus types",
+    )
+    parser.add_argument(
+        "-s", "--split", choices=["train", "val", "test"], help="Filter by dataset split"
+    )
+    parser.add_argument("--render", action="store_true", help="Render the environment")
+    parser.add_argument("--max-steps", type=int, default=1000, help="Maximum steps per episode")
+
+    args = parser.parse_args()
+
+    runner = CrafterDatasetRunner()
+    runner.run_batch(
+        dataset_name=args.dataset,
+        num_instances=args.num_instances,
+        difficulties=args.difficulties,
+        impetus_types=args.impetus_types,
+        split=args.split,
+        render=args.render,
+        max_steps=args.max_steps,
+    )
+
+
+if __name__ == "__main__":
+    main()