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

synth_ai/environments/examples/minigrid/environment.py

@@ -0,0 +1,274 @@
+"""MiniGrid Environment implementation.
+
+This module provides a high-level interface for MiniGrid environments
+with tool-based interaction and flexible observation generation.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import numpy as np
+from pydantic import BaseModel, Field
+
+from synth_ai.environments.environment.tools import AbstractTool, EnvToolCall, ToolResult
+from synth_ai.environments.stateful.core import StatefulEnvironment
+from synth_ai.environments.reproducibility.core import ReproducibleEnvironment
+from synth_ai.environments.environment.shared_engine import (
+    GetObservationCallable,
+    InternalObservation,
+)
+from synth_ai.environments.examples.minigrid.engine import (
+    MiniGridEngine,
+    MiniGridPublicState,
+    MiniGridPrivateState,
+    MiniGridObservationCallable,
+    MiniGridCheckpointObservationCallable,
+)
+from synth_ai.environments.tasks.core import TaskInstance
+
+
+class MiniGridActionInput(BaseModel):
+    """Input model for MiniGrid actions."""
+
+    action: str = Field(
+        ...,
+        description="The action to take. Must be one of: 'left', 'right', 'forward', 'pickup', 'drop', 'toggle', 'done'",
+    )
+
+
+class MiniGridInteractTool(AbstractTool):
+    """Tool for interacting with the MiniGrid environment."""
+
+    name = "minigrid_act"
+    description = "Perform an action in the MiniGrid environment"
+    call_schema = MiniGridActionInput
+    result_schema = ToolResult
+
+    def __init__(self, engine: MiniGridEngine):
+        """Initialize the tool with a MiniGrid engine."""
+        self.engine = engine
+        self.action_map = {
+            "left": 0,  # Action 0 is counter-clockwise (left)
+            "right": 1,  # Action 1 is clockwise (right)
+            "forward": 2,
+            "pickup": 3,
+            "drop": 4,
+            "toggle": 5,
+            "done": 6,
+        }
+
+    async def __call__(self, call: EnvToolCall) -> ToolResult:
+        """Execute the action and return the result."""
+        try:
+            action_name = call.args.get("action", "").lower()
+
+            if action_name not in self.action_map:
+                return ToolResult(
+                    ok=False,
+                    error=f"Invalid action '{action_name}'. Valid actions are: {', '.join(self.action_map.keys())}",
+                    payload={},
+                )
+
+            action = self.action_map[action_name]
+
+            # Execute the action
+            private_state, public_state = await self.engine._step_engine(action)
+
+            # Build response
+            response_parts = [f"Action '{action_name}' executed."]
+
+            if private_state.reward_last != 0:
+                response_parts.append(f"Reward: {private_state.reward_last:.2f}")
+
+            if private_state.terminated:
+                response_parts.append("Episode terminated!")
+                if private_state.info.get("success", False):
+                    response_parts.append("Mission completed successfully!")
+            elif private_state.truncated:
+                response_parts.append("Episode truncated (max steps reached).")
+
+            return ToolResult(
+                ok=True,
+                payload={
+                    "message": " ".join(response_parts),
+                    "public_state": public_state,
+                    "private_state": private_state,
+                },
+            )
+        except Exception as e:
+            return ToolResult(ok=False, error=str(e), payload={})
+
+
+class MiniGridEnvironment(StatefulEnvironment, ReproducibleEnvironment[MiniGridEngine]):
+    """High-level MiniGrid environment with tool-based interaction."""
+
+    def __init__(
+        self,
+        task_instance: TaskInstance,
+        custom_step_obs: Optional[GetObservationCallable] = None,
+        custom_ckpt_obs: Optional[GetObservationCallable] = None,
+    ):
+        """Initialize the MiniGrid environment.
+
+        Args:
+            task_instance: Task instance containing configuration
+            custom_step_obs: Custom observation generator for steps
+            custom_ckpt_obs: Custom observation generator for checkpoints
+        """
+        self.name = "MiniGridEnvironment"
+        self.task_instance = task_instance
+        self.custom_step_observation_callable = custom_step_obs or MiniGridObservationCallable()
+        self.custom_checkpoint_observation_callable = (
+            custom_ckpt_obs or MiniGridCheckpointObservationCallable()
+        )
+
+        # Create engine
+        self.engine = MiniGridEngine(task_instance)
+
+        # Initialize tool
+        self._interact_tool = MiniGridInteractTool(self.engine)
+
+    async def initialize(self) -> InternalObservation:
+        """Initialize the environment and return initial observation."""
+        priv, pub = await self.engine._reset_engine()
+        return await self._to_observation(priv, pub, self.custom_step_observation_callable)
+
+    async def step(
+        self,
+        tool_calls: Union[List[Dict[str, Any]], List[EnvToolCall], Dict[str, Any], EnvToolCall],
+    ) -> InternalObservation:
+        """Process a tool call and return observation."""
+        validated_call = self.validate_tool_calls(tool_calls)
+        result = await self._interact_tool(validated_call)
+
+        if result.ok:
+            priv = result.payload["private_state"]
+            pub = result.payload["public_state"]
+            return await self._to_observation(priv, pub, self.custom_step_observation_callable)
+        else:
+            priv, pub = self.engine.get_current_states_for_observation()
+            return await self._to_observation(
+                priv,
+                pub,
+                self.custom_step_observation_callable,
+                extra_obs={"error": result.error},
+            )
+
+    async def terminate(self) -> InternalObservation:
+        """Terminate the environment and return final observation."""
+        priv, pub = self.engine.get_current_states_for_observation()
+        return await self._to_observation(priv, pub, self.custom_checkpoint_observation_callable)
+
+    async def checkpoint(self) -> InternalObservation:
+        """Create a checkpoint of the current state."""
+        priv, pub = self.engine.get_current_states_for_observation()
+        return await self._to_observation(priv, pub, self.custom_checkpoint_observation_callable)
+
+    def validate_tool_calls(
+        self,
+        tool_calls: Union[List[Dict[str, Any]], List[EnvToolCall], Dict[str, Any], EnvToolCall],
+    ) -> EnvToolCall:
+        """Validate and normalize tool calls."""
+        # If already an EnvToolCall, validate and return
+        if isinstance(tool_calls, EnvToolCall):
+            if tool_calls.tool != "minigrid_act":
+                raise ValueError(f"Unknown tool: {tool_calls.tool}. Expected 'minigrid_act'.")
+            return tool_calls
+
+        # Handle different input formats
+        if isinstance(tool_calls, dict):
+            # Single tool call
+            tool_call = tool_calls
+        elif isinstance(tool_calls, list) and len(tool_calls) > 0:
+            # List of tool calls - take the first one
+            first_item = tool_calls[0]
+            if isinstance(first_item, list) and len(first_item) > 0:
+                # Nested list
+                tool_call = first_item[0]
+            elif isinstance(first_item, EnvToolCall):
+                # Handle case where service sends list of EnvToolCall objects
+                if first_item.tool != "minigrid_act":
+                    raise ValueError(f"Unknown tool: {first_item.tool}. Expected 'minigrid_act'.")
+                return first_item
+            else:
+                tool_call = first_item
+        else:
+            raise ValueError("Invalid tool_calls format")
+
+        # At this point tool_call should be a dict
+        if isinstance(tool_call, EnvToolCall):
+            # Handle case where we somehow still have an EnvToolCall
+            if tool_call.tool != "minigrid_act":
+                raise ValueError(f"Unknown tool: {tool_call.tool}. Expected 'minigrid_act'.")
+            return tool_call
+
+        # Extract tool name and args - fail fast
+        if "tool" in tool_call:
+            tool_name = tool_call["tool"]
+        elif "name" in tool_call:
+            tool_name = tool_call["name"]
+        else:
+            raise ValueError("Tool call missing 'tool' or 'name' field")
+
+        # Handle different argument formats - fail fast
+        if "args" in tool_call:
+            args = tool_call["args"]
+        elif "parameters" in tool_call:
+            args = tool_call["parameters"]
+        elif "input" in tool_call:
+            if isinstance(tool_call["input"], str):
+                args = json.loads(tool_call["input"])
+            else:
+                args = tool_call["input"]
+        else:
+            raise ValueError("Tool call missing 'args', 'parameters', or 'input' field")
+
+        if tool_name != "minigrid_act":
+            raise ValueError(f"Unknown tool: {tool_name}. Expected 'minigrid_act'.")
+
+        # Create EnvToolCall
+        return EnvToolCall(
+            tool=tool_name,
+            args=args,
+        )
+
+    async def _to_observation(
+        self,
+        priv: MiniGridPrivateState,
+        pub: MiniGridPublicState,
+        observation_callable: GetObservationCallable,
+        extra_obs: Optional[Dict[str, Any]] = None,
+    ) -> InternalObservation:
+        """Convert states to observation using callable."""
+        obs = await observation_callable.get_observation(pub, priv)
+
+        # Attach full state objects for downstream analysis (fail fast)
+        obs["public"] = pub
+        obs["private"] = priv
+
+        if extra_obs:
+            obs.update(extra_obs)
+
+        return obs
+
+    async def _serialize_engine(self) -> Dict[str, Any]:
+        """Serialize the engine state."""
+        snapshot = await self.engine._serialize_engine()
+        return {
+            "task_instance_dict": snapshot.task_instance_dict,
+            "engine_snapshot": snapshot.engine_snapshot,
+        }
+
+    @classmethod
+    async def _deserialize_engine(cls, data: Dict[str, Any]) -> MiniGridEngine:
+        """Deserialize the engine state."""
+        from synth_ai.environments.examples.minigrid.engine import MiniGridEngineSnapshot
+
+        snapshot = MiniGridEngineSnapshot(
+            task_instance_dict=data["task_instance_dict"],
+            engine_snapshot=data["engine_snapshot"],
+        )
+
+        return await MiniGridEngine._deserialize_engine(snapshot)

synth_ai/environments/examples/minigrid/environment_mapping.py

@@ -0,0 +1,242 @@
+"""
+MiniGrid Environment Mapping Module
+
+This module provides functionality to map any integer seed to one of 60 MiniGrid
+environments using modulo arithmetic for deterministic and reproducible
+environment selection.
+"""
+
+from typing import Tuple
+
+# Environment mapping table (60 total environments)
+ENVIRONMENT_MAPPING = {
+    # Ultra-Easy (0-4)
+    0: "MiniGrid-Empty-5x5-v0",
+    1: "MiniGrid-Empty-6x6-v0",
+    2: "MiniGrid-Empty-Random-5x5-v0",
+    3: "MiniGrid-Empty-Random-6x6-v0",
+    4: "MiniGrid-GoToDoor-5x5-v0",
+    # Easy (5-14)
+    5: "MiniGrid-Empty-8x8-v0",
+    6: "MiniGrid-FourRooms-v0",
+    7: "MiniGrid-DoorKey-5x5-v0",
+    8: "MiniGrid-GoToDoor-6x6-v0",
+    9: "MiniGrid-GoToDoor-8x8-v0",
+    10: "MiniGrid-Fetch-5x5-N2-v0",
+    11: "MiniGrid-Fetch-6x6-N2-v0",
+    12: "MiniGrid-PutNear-6x6-N2-v0",
+    13: "MiniGrid-Unlock-v0",
+    14: "MiniGrid-UnlockPickup-v0",
+    # Medium (15-29)
+    15: "MiniGrid-DoorKey-6x6-v0",
+    16: "MiniGrid-DoorKey-8x8-v0",
+    17: "MiniGrid-MultiRoom-N2-S4-v0",
+    18: "MiniGrid-LavaGapS5-v0",
+    19: "MiniGrid-LavaGapS6-v0",
+    20: "MiniGrid-LavaGapS7-v0",
+    21: "MiniGrid-SimpleCrossingS9N1-v0",
+    22: "MiniGrid-SimpleCrossingS9N2-v0",
+    23: "MiniGrid-SimpleCrossingS9N3-v0",
+    24: "MiniGrid-Fetch-8x8-N3-v0",
+    25: "MiniGrid-PutNear-8x8-N3-v0",
+    26: "MiniGrid-RedBlueDoors-6x6-v0",
+    27: "MiniGrid-RedBlueDoors-8x8-v0",
+    28: "MiniGrid-BlockedUnlockPickup-v0",
+    29: "MiniGrid-KeyCorridorS3R1-v0",
+    # Hard (30-44)
+    30: "MiniGrid-DoorKey-16x16-v0",
+    31: "MiniGrid-MultiRoom-N4-S5-v0",
+    32: "MiniGrid-MultiRoom-N6-v0",
+    33: "MiniGrid-LavaCrossingS9N1-v0",
+    34: "MiniGrid-LavaCrossingS9N2-v0",
+    35: "MiniGrid-LavaCrossingS9N3-v0",
+    36: "MiniGrid-LavaCrossingS11N5-v0",
+    37: "MiniGrid-SimpleCrossingS11N5-v0",
+    38: "MiniGrid-KeyCorridorS3R2-v0",
+    39: "MiniGrid-KeyCorridorS3R3-v0",
+    40: "MiniGrid-KeyCorridorS4R3-v0",
+    41: "MiniGrid-KeyCorridorS5R3-v0",
+    42: "MiniGrid-KeyCorridorS6R3-v0",
+    43: "MiniGrid-MemoryS7-v0",
+    44: "MiniGrid-MemoryS9-v0",
+    # Ultra-Hard (45-54)
+    45: "MiniGrid-MemoryS11-v0",
+    46: "MiniGrid-MemoryS13-v0",
+    47: "MiniGrid-MemoryS13Random-v0",
+    48: "MiniGrid-MemoryS17Random-v0",
+    49: "MiniGrid-LockedRoom-v0",
+    50: "MiniGrid-ObstructedMaze-1Dlh-v0",
+    51: "MiniGrid-ObstructedMaze-1Dlhb-v0",
+    52: "MiniGrid-ObstructedMaze-2Dlhb-v0",
+    53: "MiniGrid-ObstructedMaze-Full-v0",
+    54: "MiniGrid-DistShift1-v0",
+    # Specialized (55-59)
+    55: "MiniGrid-DistShift2-v0",
+    56: "MiniGrid-Dynamic-Obstacles-8x8-v0",
+    57: "MiniGrid-Dynamic-Obstacles-16x16-v0",
+    58: "MiniGrid-Playground-v0",
+    59: "MiniGrid-Empty-16x16-v0",
+}
+
+# Difficulty mapping
+DIFFICULTY_MAPPING = {
+    "ultra-easy": (0, 4),
+    "easy": (5, 14),
+    "medium": (15, 29),
+    "hard": (30, 44),
+    "ultra-hard": (45, 54),
+    "specialized": (55, 59),
+}
+
+
+def get_environment_from_seed(seed: int, hash_seed: bool = True) -> str:
+    """
+    Map any integer seed to a MiniGrid environment name.
+
+    Args:
+        seed: Integer seed
+        hash_seed: If True, hash the seed for better distribution
+
+    Returns:
+        Environment name string
+    """
+    if hash_seed:
+        # Use hash for better distribution of sequential seeds
+        env_index = hash(seed) % 60
+    else:
+        # Simple modulo
+        env_index = seed % 60
+
+    return ENVIRONMENT_MAPPING[env_index]
+
+
+def get_difficulty_from_seed(seed: int, hash_seed: bool = True) -> str:
+    """
+    Get difficulty level for a given seed.
+
+    Args:
+        seed: Integer seed
+        hash_seed: If True, hash the seed for better distribution
+
+    Returns:
+        Difficulty level string
+    """
+    if hash_seed:
+        env_index = hash(seed) % 60
+    else:
+        env_index = seed % 60
+
+    if env_index <= 4:
+        return "ultra-easy"
+    elif env_index <= 14:
+        return "easy"
+    elif env_index <= 29:
+        return "medium"
+    elif env_index <= 44:
+        return "hard"
+    elif env_index <= 54:
+        return "ultra-hard"
+    else:
+        return "specialized"
+
+
+def get_minigrid_environment(seed: int, hash_seed: bool = True) -> Tuple[str, str]:
+    """
+    Get MiniGrid environment name and difficulty from seed.
+
+    Args:
+        seed: Integer seed
+        hash_seed: If True, hash the seed for better distribution
+
+    Returns:
+        (environment_name, difficulty_level)
+    """
+    env_name = get_environment_from_seed(seed, hash_seed)
+    difficulty = get_difficulty_from_seed(seed, hash_seed)
+
+    return env_name, difficulty
+
+
+def get_environment_by_difficulty(difficulty: str, seed: int = 0) -> str:
+    """
+    Get a random environment from a specific difficulty level.
+
+    Args:
+        difficulty: Difficulty level ("ultra-easy", "easy", "medium", "hard", "ultra-hard", "specialized")
+        seed: Seed for selecting within the difficulty range
+
+    Returns:
+        Environment name string
+    """
+    if difficulty not in DIFFICULTY_MAPPING:
+        raise ValueError(f"Unknown difficulty: {difficulty}")
+
+    start, end = DIFFICULTY_MAPPING[difficulty]
+    range_size = end - start + 1
+    env_index = start + (seed % range_size)
+
+    return ENVIRONMENT_MAPPING[env_index]
+
+
+def get_curriculum_environment(progress: float, seed: int = 0) -> Tuple[str, str]:
+    """
+    Select environment based on curriculum progress.
+
+    Args:
+        progress: Progress value from 0.0 to 1.0
+        seed: Seed for environment selection within difficulty
+
+    Returns:
+        (environment_name, difficulty_level)
+    """
+    if progress < 0.2:  # Early stage - ultra-easy
+        difficulty = "ultra-easy"
+    elif progress < 0.4:  # Beginning - easy
+        difficulty = "easy"
+    elif progress < 0.6:  # Intermediate - medium
+        difficulty = "medium"
+    elif progress < 0.8:  # Advanced - hard
+        difficulty = "hard"
+    else:  # Expert - ultra-hard/specialized
+        if progress < 0.9:
+            difficulty = "ultra-hard"
+        else:
+            difficulty = "specialized"
+
+    env_name = get_environment_by_difficulty(difficulty, seed)
+    return env_name, difficulty
+
+
+def validate_environment_name(env_name: str) -> bool:
+    """
+    Check if an environment name is supported.
+
+    Args:
+        env_name: Environment name to validate
+
+    Returns:
+        True if supported, False otherwise
+    """
+    return env_name in ENVIRONMENT_MAPPING.values()
+
+
+def get_all_environments() -> list[str]:
+    """Get list of all supported environment names."""
+    return list(ENVIRONMENT_MAPPING.values())
+
+
+def get_environments_by_difficulty(difficulty: str) -> list[str]:
+    """
+    Get all environments for a specific difficulty level.
+
+    Args:
+        difficulty: Difficulty level
+
+    Returns:
+        List of environment names
+    """
+    if difficulty not in DIFFICULTY_MAPPING:
+        raise ValueError(f"Unknown difficulty: {difficulty}")
+
+    start, end = DIFFICULTY_MAPPING[difficulty]
+    return [ENVIRONMENT_MAPPING[i] for i in range(start, end + 1)]