synth-ai 0.1.9__py3-none-any.whl → 0.2.1.dev0__py3-none-any.whl

synth_ai/environments/examples/red/environment.py

@@ -0,0 +1,235 @@
+from __future__ import annotations
+from typing import List, Optional, Any, Dict, Union
+from pydantic import BaseModel, Field
+
+# Import logging configuration to suppress JAX debug messages
+
+from .engine import (
+    PokemonRedEngine,
+    PokemonRedPrivateState,
+    PokemonRedPublicState,
+    PokemonRedEngineSnapshot,
+)
+from .taskset import PokemonRedTaskInstance, INSTANCE as DEFAULT_TASK_INSTANCE
+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,
+)
+
+
+# Tool input schemas
+class PressButtonInput(BaseModel):
+    button: str = Field(
+        ..., description="Game Boy button: A, B, UP, DOWN, LEFT, RIGHT, START, SELECT"
+    )
+    frames: int = Field(1, description="Number of frames to hold the button")
+
+
+# Tool definitions
+class PressButtonTool(AbstractTool):
+    name = "press_button"
+    description = "Press a Game Boy button for the specified number of frames"
+    call_schema = PressButtonInput
+    result_schema = ToolResult
+
+    def __init__(self, engine: PokemonRedEngine):
+        self.engine = engine
+
+    async def __call__(self, call: EnvToolCall) -> ToolResult:
+        try:
+            validated_args = self.call_schema(**call.args)
+            action = {"button": validated_args.button, "frames": validated_args.frames}
+            priv_state, pub_state = await self.engine._step_engine(action)
+            return ToolResult(
+                ok=True,
+                payload={
+                    "public": pub_state,
+                    "private": priv_state,
+                },
+            )
+        except Exception as e:
+            # Get current state for error context
+            priv_state, pub_state = self.engine._create_states(reward=0.0)
+            return ToolResult(
+                ok=False,
+                error=str(e),
+                payload={"public": pub_state},
+            )
+
+
+# Observation callable for Pokemon Red
+class PokemonRedObservationCallable(GetObservationCallable):
+    async def get_observation(
+        self, pub: PokemonRedPublicState, priv: PokemonRedPrivateState
+    ) -> InternalObservation:
+        """Convert Pokemon Red states to agent observation"""
+        from .engine_helpers.state_extraction import (
+            get_badge_count,
+            format_position,
+            format_hp_status,
+        )
+
+        badge_count = get_badge_count(pub.badges)
+        position = format_position(pub.player_x, pub.player_y, pub.map_id)
+        hp_status = format_hp_status(pub.party_hp_current, pub.party_hp_max)
+
+        obs = {
+            "position": position,
+            "badges_earned": badge_count,
+            "badges_bitfield": pub.badges,
+            "hp_status": hp_status,
+            "party_level": pub.party_level,
+            "party_xp": pub.party_xp,
+            "in_battle": pub.in_battle,
+            "step_count": pub.step_count,
+            "reward_last_step": priv.reward_last_step,
+            "total_reward": priv.total_reward,
+            "terminated": priv.terminated,
+        }
+
+        if pub.error_info:
+            obs["error"] = pub.error_info
+
+        return obs
+
+
+class PokemonRedEnvironment(StatefulEnvironment, ReproducibleEnvironment[PokemonRedEngine]):
+    """Pokemon Red stateful game environment for AI agents"""
+
+    def __init__(
+        self,
+        task_instance: Optional[PokemonRedTaskInstance] = None,
+        custom_step_obs: Optional[GetObservationCallable] = None,
+        custom_ckpt_obs: Optional[GetObservationCallable] = None,
+    ):
+        self.name = "PokemonRed"
+        self.task_instance = task_instance or DEFAULT_TASK_INSTANCE
+        self.custom_step_observation_callable = custom_step_obs or PokemonRedObservationCallable()
+        self.custom_checkpoint_observation_callable = (
+            custom_ckpt_obs or PokemonRedObservationCallable()
+        )
+        self.engine = PokemonRedEngine(self.task_instance)
+
+        # Register tools
+        self._press_button_tool = PressButtonTool(self.engine)
+        if self._press_button_tool.name not in TOOL_REGISTRY:
+            register_tool(self._press_button_tool)
+
+    async def initialize(self) -> InternalObservation:
+        """Initialize the Pokemon Red environment"""
+        priv, pub = await self.engine._reset_engine()
+        return await self._to_observation(priv, pub, self.custom_step_observation_callable)
+
+    async def terminate(self) -> InternalObservation:
+        """Terminate the environment"""
+        priv, pub = self.engine._create_states(reward=0.0, terminated=True)
+        obs_dict = {
+            "terminated": True,
+            "message": "Pokemon Red environment terminated.",
+        }
+        return await self._to_observation(
+            priv, pub, self.custom_step_observation_callable, extra_obs=obs_dict
+        )
+
+    def validate_tool_calls(
+        self, tool_calls: Union[EnvToolCall, List[EnvToolCall], List[List[EnvToolCall]]]
+    ) -> EnvToolCall:
+        """Validate and normalize tool calls to 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)}")
+        if agent_call.tool != "press_button":
+            raise ValueError(f"Unknown tool: {agent_call.tool}. Expected 'press_button'.")
+
+        return agent_call
+
+    async def step(
+        self, tool_calls: Union[EnvToolCall, List[EnvToolCall], List[List[EnvToolCall]]]
+    ) -> InternalObservation:
+        """Execute one step in the Pokemon Red environment"""
+        agent_call = self.validate_tool_calls(tool_calls)
+        tool_result: ToolResult = await self._press_button_tool(agent_call)
+
+        payload_dict = tool_result.payload
+        if not tool_result.ok or not isinstance(payload_dict, dict):
+            # Fallback if tool execution failed
+            priv_state, pub_state = self.engine._create_states(reward=0.0)
+            if tool_result.error and hasattr(pub_state, "error_info"):
+                pub_state.error_info = tool_result.error
+        else:
+            # Extract states from successful tool execution - now they're dataclass objects
+            priv_state = payload_dict.get("private")
+            pub_state = payload_dict.get("public")
+
+            if priv_state is None or pub_state is None:
+                priv_state, pub_state = self.engine._create_states(reward=0.0)
+                if tool_result.error and hasattr(pub_state, "error_info"):
+                    pub_state.error_info = tool_result.error
+            else:
+                # States are already dataclass objects, no need to reconstruct
+                if tool_result.error and hasattr(pub_state, "error_info"):
+                    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:
+        """Create a checkpoint of the current environment state"""
+        engine_snapshot: PokemonRedEngineSnapshot = await self.engine._serialize_engine()
+        priv, pub = self.engine._create_states(reward=0.0)
+        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 _to_observation(
+        self,
+        priv: PokemonRedPrivateState,
+        pub: PokemonRedPublicState,
+        obs_cb: Optional[GetObservationCallable],
+        extra_obs: Optional[Dict[str, Any]] = None,
+    ) -> InternalObservation:
+        """Convert states to observation using the specified callback"""
+        active_obs_cb = obs_cb or PokemonRedObservationCallable()
+        observation = await active_obs_cb.get_observation(pub, priv)
+        if extra_obs and isinstance(observation, dict):
+            observation.update(extra_obs)
+        return observation
+
+    # ReproducibleEnvironment methods
+    async def _serialize_engine(self) -> PokemonRedEngineSnapshot:
+        return await self.engine._serialize_engine()
+
+    @classmethod
+    async def _deserialize_engine(
+        cls, snapshot: PokemonRedEngineSnapshot, task_instance: PokemonRedTaskInstance
+    ) -> "PokemonRedEnvironment":
+        eng = await PokemonRedEngine._deserialize_engine(snapshot, task_instance)
+        env = cls(task_instance)
+        env.engine = eng
+        return env

synth_ai/environments/examples/red/taskset.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from pathlib import Path
+import uuid
+from synth_ai.environments.tasks.core import (
+    Task,
+    TaskInstance,
+    Impetus,
+    Intent,
+    TaskInstanceMetadata,
+)
+
+# Define the main task for Pokemon Red
+TASK = Task(
+    global_premises="You are playing Pokemon Red. Start in Pewter City with a level-10 Pikachu.",
+    global_constraints="No glitches or exploits. Play within normal game mechanics.",
+    global_objectives="Defeat Brock at the Pewter Gym to earn the Boulder Badge.",
+    shared_env_params={},
+)
+
+# Path to initial save state (would contain a save near Pewter Gym)
+INITIAL_SNAPSHOT = Path(__file__).parent / "snapshots" / "pewter_start.state"
+
+
+@dataclass
+class PokemonRedTaskInstance(TaskInstance):
+    """Task instance for Pokemon Red challenges"""
+
+    async def serialize(self) -> dict:
+        """Serialize the task instance to a dictionary"""
+        return {
+            "id": str(self.id),
+            "impetus": {"instructions": self.impetus.instructions},
+            "intent": {
+                "rubric": self.intent.rubric,
+                "gold_trajectories": None,
+                "gold_state_diff": self.intent.gold_state_diff,
+            },
+            "metadata": {},
+            "is_reproducible": self.is_reproducible,
+            "initial_engine_snapshot": str(self.initial_engine_snapshot)
+            if self.initial_engine_snapshot
+            else None,
+        }
+
+    @classmethod
+    async def deserialize(cls, data: dict) -> "PokemonRedTaskInstance":
+        """Deserialize a task instance from a dictionary"""
+        return cls(
+            id=uuid.UUID(data["id"]),
+            impetus=Impetus(instructions=data["impetus"]["instructions"]),
+            intent=Intent(
+                rubric=data["intent"]["rubric"],
+                gold_trajectories=None,
+                gold_state_diff=data["intent"]["gold_state_diff"],
+            ),
+            metadata=TaskInstanceMetadata(),
+            is_reproducible=data["is_reproducible"],
+            initial_engine_snapshot=None,
+        )
+
+
+# Main task instance - beat Brock for Boulder Badge
+INSTANCE = PokemonRedTaskInstance(
+    id=uuid.UUID("12345678-1234-5678-9abc-123456789abc"),
+    impetus=Impetus(
+        instructions="Navigate to Pewter Gym and defeat Brock to earn the Boulder Badge. Use strategic Pokemon battles and item management."
+    ),
+    intent=Intent(
+        rubric="Successfully obtain the Boulder Badge by defeating Brock at Pewter Gym. Efficiency measured by minimal steps and strategic Pokemon usage.",
+        gold_trajectories=None,
+        gold_state_diff={"badges": 1},
+    ),
+    metadata=TaskInstanceMetadata(),
+    is_reproducible=True,
+    initial_engine_snapshot=INITIAL_SNAPSHOT if INITIAL_SNAPSHOT.exists() else None,
+)

synth_ai/environments/examples/red/test_fixes.py

@@ -0,0 +1,125 @@
+#!/usr/bin/env python3
+"""
+Test script to verify red environment fixes.
+Tests JAX logging suppression and error handling.
+"""
+
+import asyncio
+import logging
+import sys
+
+from synth_ai.environments.examples.red.environment import PokemonRedEnvironment
+from synth_ai.environments.examples.red.taskset import INSTANCE as POKEMON_TASK
+from synth_ai.environments.environment.tools import EnvToolCall
+
+
+class PressButtonCall(EnvToolCall):
+    """Helper class for creating button press calls"""
+
+    def __init__(self, button: str, frames: int = 1):
+        super().__init__(tool="press_button", args={"button": button, "frames": frames})
+
+
+async def test_environment_setup():
+    """Test that the environment can be set up without errors."""
+    print("Testing Pokemon Red environment setup...")
+
+    try:
+        # Create environment instance
+        env = PokemonRedEnvironment(POKEMON_TASK)
+        print("✅ Environment created successfully")
+
+        # Try to initialize
+        obs = await env.initialize()
+        print("✅ Environment initialized successfully")
+        print(f"Initial observation keys: {list(obs.keys())}")
+
+        # Try a simple step
+        obs = await env.step(PressButtonCall("A"))
+        print("✅ Environment step executed successfully")
+        print(
+            f"Step observation: step_count={obs.get('step_count')}, terminated={obs.get('terminated')}"
+        )
+
+        # Terminate
+        final_obs = await env.terminate()
+        print("✅ Environment terminated successfully")
+
+        return True
+
+    except Exception as e:
+        print(f"❌ Failed to setup environment: {e}")
+        logging.exception("Failed to setup environment, aborting test")
+        return False
+
+
+def test_logging_configuration():
+    """Test that logging is properly configured."""
+    print("Testing logging configuration...")
+
+    # Check that JAX loggers are set to WARNING level
+    jax_loggers = [
+        "jax._src.cache_key",
+        "jax._src.compilation_cache",
+        "jax._src.compiler",
+        "jax._src.dispatch",
+    ]
+
+    for logger_name in jax_loggers:
+        logger = logging.getLogger(logger_name)
+        if logger.level >= logging.WARNING:
+            print(f"✅ {logger_name} logger level: {logging.getLevelName(logger.level)}")
+        else:
+            print(
+                f"❌ {logger_name} logger level: {logging.getLevelName(logger.level)} (should be WARNING or higher)"
+            )
+
+    # Test that debug messages are suppressed
+    jax_logger = logging.getLogger("jax._src.cache_key")
+    jax_logger.debug("This debug message should not appear")
+    print("✅ JAX debug logging appears to be suppressed")
+
+
+def test_safe_compare():
+    """Test the safe comparison function."""
+    print("Testing safe comparison function...")
+
+    from synth_ai.environments.examples.red.config_logging import safe_compare
+
+    # Test cases
+    test_cases = [
+        ("5", 3, ">", True),  # String vs int
+        (5, "3", ">", True),  # Int vs string
+        ("abc", 5, ">", False),  # Invalid string vs int
+        ("5", "3", ">", True),  # String vs string (numeric)
+        ("abc", "def", ">", False),  # String vs string (alphabetic)
+        (5, 3, ">", True),  # Normal int comparison
+    ]
+
+    for left, right, op, expected in test_cases:
+        result = safe_compare(left, right, op)
+        status = "✅" if result == expected else "❌"
+        print(f"{status} safe_compare({left}, {right}, '{op}') = {result} (expected {expected})")
+
+
+async def main():
+    """Main test function."""
+    print("Running Pokemon Red environment fixes test...\n")
+
+    # Test logging configuration
+    test_logging_configuration()
+    print()
+
+    # Test safe comparison
+    test_safe_compare()
+    print()
+
+    # Test environment setup
+    success = await test_environment_setup()
+
+    print(f"\nOverall test result: {'✅ PASSED' if success else '❌ FAILED'}")
+    return 0 if success else 1
+
+
+if __name__ == "__main__":
+    sys.exit(asyncio.run(main()))

synth_ai/environments/examples/red/test_fixes_mock.py

@@ -0,0 +1,148 @@
+#!/usr/bin/env python3
+"""
+Mock test script to verify red environment fixes without ROM file.
+Tests JAX logging suppression and error handling.
+"""
+
+import logging
+import sys
+from unittest.mock import Mock, patch
+
+
+def test_logging_configuration():
+    """Test that logging is properly configured."""
+    print("Testing logging configuration...")
+
+    # Import configuration to trigger setup
+    from synth_ai.environments.examples.red.config_logging import configure_logging
+
+    configure_logging()
+
+    # Check that JAX loggers are set to WARNING level
+    jax_loggers = [
+        "jax._src.cache_key",
+        "jax._src.compilation_cache",
+        "jax._src.compiler",
+        "jax._src.dispatch",
+    ]
+
+    success = True
+    for logger_name in jax_loggers:
+        logger = logging.getLogger(logger_name)
+        if logger.level >= logging.WARNING:
+            print(f"✅ {logger_name} logger level: {logging.getLevelName(logger.level)}")
+        else:
+            print(
+                f"❌ {logger_name} logger level: {logging.getLevelName(logger.level)} (should be WARNING or higher)"
+            )
+            success = False
+
+    # Test that debug messages are suppressed
+    jax_logger = logging.getLogger("jax._src.cache_key")
+    jax_logger.debug("This debug message should not appear")
+    print("✅ JAX debug logging appears to be suppressed")
+
+    return success
+
+
+def test_safe_compare():
+    """Test the safe comparison function."""
+    print("Testing safe comparison function...")
+
+    from synth_ai.environments.examples.red.config_logging import safe_compare
+
+    # Test cases that previously would cause the string vs int error
+    test_cases = [
+        ("5", 3, ">", True),  # String vs int
+        (5, "3", ">", True),  # Int vs string
+        ("abc", 5, ">", False),  # Invalid string vs int
+        ("5", "3", ">", True),  # String vs string (numeric)
+        ("abc", "def", ">", False),  # String vs string (alphabetic)
+        (5, 3, ">", True),  # Normal int comparison
+        ("10", 5, ">=", True),  # String number >= int
+        (3, "10", "<=", True),  # Int <= string number
+    ]
+
+    success = True
+    for left, right, op, expected in test_cases:
+        result = safe_compare(left, right, op)
+        status = "✅" if result == expected else "❌"
+        print(f"{status} safe_compare({left}, {right}, '{op}') = {result} (expected {expected})")
+        if result != expected:
+            success = False
+
+    return success
+
+
+def test_state_creation_error_handling():
+    """Test that state creation handles type errors gracefully."""
+    print("Testing state creation error handling...")
+
+    from synth_ai.environments.examples.red.engine import PokemonRedEngine
+    from synth_ai.environments.examples.red.taskset import INSTANCE as POKEMON_TASK
+
+    try:
+        # Mock the PyBoy emulator to avoid ROM requirement
+        with patch("examples.red.engine.PyBoy") as mock_pyboy:
+            mock_emulator = Mock()
+            mock_pyboy.return_value = mock_emulator
+
+            # Create engine instance
+            engine = PokemonRedEngine(POKEMON_TASK)
+
+            # Mock extract_game_state to return problematic data that could cause comparison errors
+            with patch.object(engine, "_extract_current_state") as mock_extract:
+                # Test with string badges that could cause comparison error
+                mock_extract.return_value = {
+                    "map_id": "1",  # String instead of int
+                    "player_x": "10",
+                    "player_y": "20",
+                    "badges": "abc",  # Non-numeric string
+                    "in_battle": "false",  # String instead of bool
+                    "party_level": "5",
+                    "party_hp_current": "50",
+                    "party_hp_max": "50",
+                    "party_xp": "100",
+                }
+
+                # This should not crash due to our error handling
+                priv_state, pub_state = engine._create_states(0.0, False)
+
+                print("✅ State creation handles problematic data gracefully")
+                print(f"✅ Created states: badges={pub_state.badges}, map_id={pub_state.map_id}")
+
+                # Test with completely invalid data
+                mock_extract.side_effect = Exception("Memory read error")
+                priv_state, pub_state = engine._create_states(0.0, False)
+                print("✅ State creation handles extraction errors gracefully")
+
+                return True
+
+    except Exception as e:
+        print(f"❌ State creation error handling failed: {e}")
+        return False
+
+
+def main():
+    """Main test function."""
+    print("Running Pokemon Red environment fixes test (mock version)...\n")
+
+    # Test logging configuration
+    logging_ok = test_logging_configuration()
+    print()
+
+    # Test safe comparison
+    compare_ok = test_safe_compare()
+    print()
+
+    # Test error handling
+    error_handling_ok = test_state_creation_error_handling()
+    print()
+
+    success = logging_ok and compare_ok and error_handling_ok
+    print(f"Overall test result: {'✅ PASSED' if success else '❌ FAILED'}")
+    return 0 if success else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

synth_ai/environments/examples/red/units/__init__.py

@@ -0,0 +1 @@
+# Unit tests for Pokemon Red environment

synth_ai/environments/examples/red/units/test_basic_functionality.py

@@ -0,0 +1,97 @@
+import pytest
+from synth_ai.environments.examples.red.environment import PokemonRedEnvironment
+from synth_ai.environments.examples.red.taskset import INSTANCE as POKEMON_TASK
+from synth_ai.environments.environment.tools import EnvToolCall
+
+
+class PressButtonCall(EnvToolCall):
+    """Helper class for creating button press calls"""
+
+    def __init__(self, button: str, frames: int = 1):
+        super().__init__(tool="press_button", args={"button": button, "frames": frames})
+
+
+@pytest.mark.asyncio
+async def test_pokemon_red_basic():
+    """Test basic Pokemon Red environment functionality"""
+    env = PokemonRedEnvironment(POKEMON_TASK)
+
+    # Initialize environment
+    obs = await env.initialize()
+    assert "position" in obs
+    assert "badges_earned" in obs
+    assert obs["badges_earned"] == 0  # Should start with no badges
+
+    # Test a few button presses
+    obs = await env.step(PressButtonCall("A"))
+    assert "step_count" in obs
+    assert obs["step_count"] == 1
+
+    obs = await env.step(PressButtonCall("RIGHT", 2))
+    assert obs["step_count"] == 2
+
+    # Test termination
+    final_obs = await env.terminate()
+    assert final_obs["terminated"] is True
+
+
+@pytest.mark.asyncio
+async def test_pokemon_red_multiple_actions():
+    """Test sequence of actions in Pokemon Red"""
+    env = PokemonRedEnvironment(POKEMON_TASK)
+
+    obs = await env.initialize()
+    initial_reward = obs["total_reward"]
+
+    # Sequence of movements and actions
+    actions = [
+        PressButtonCall("RIGHT"),
+        PressButtonCall("UP"),
+        PressButtonCall("A"),
+        PressButtonCall("DOWN"),
+        PressButtonCall("B"),
+    ]
+
+    for action in actions:
+        obs = await env.step(action)
+        assert "position" in obs
+        assert "hp_status" in obs
+        assert "party_level" in obs
+
+    # Should have accumulated some reward (mostly negative from step penalty)
+    assert obs["total_reward"] <= initial_reward  # Step penalties
+    assert obs["step_count"] == len(actions)
+
+
+@pytest.mark.asyncio
+async def test_pokemon_red_checkpointing():
+    """Test environment checkpointing functionality"""
+    env = PokemonRedEnvironment(POKEMON_TASK)
+
+    # Initialize and take some steps
+    await env.initialize()
+    await env.step(PressButtonCall("RIGHT"))
+    await env.step(PressButtonCall("A"))
+
+    # Create checkpoint
+    checkpoint_obs = await env.checkpoint()
+    assert "engine_snapshot_data" in checkpoint_obs
+    assert checkpoint_obs["step_count"] == 2
+
+    # Verify checkpoint contains expected data
+    snapshot_data = checkpoint_obs["engine_snapshot_data"]
+    assert "state_data" in snapshot_data
+    assert "total_reward" in snapshot_data
+    assert "step_count" in snapshot_data
+
+
+@pytest.mark.asyncio
+async def test_pokemon_red_invalid_button():
+    """Test handling of invalid button inputs"""
+    env = PokemonRedEnvironment(POKEMON_TASK)
+    await env.initialize()
+
+    # Test with invalid button - should handle gracefully
+    obs = await env.step(PressButtonCall("INVALID_BUTTON"))
+    # Should still return valid observation even if action failed
+    assert "position" in obs