add frozen lake server to package

xzrderek · xzrderek · commit 91121cc20c69 · 2025-10-03T11:15:55.000-07:00
diff --git a/eval_protocol/mcp_servers/frozen_lake/frozen_lake_adapter.py b/eval_protocol/mcp_servers/frozen_lake/frozen_lake_adapter.py
@@ -0,0 +1,160 @@
+"""
+FrozenLake Environment Adapter
+
+This adapter implements the EnvironmentAdapter interface for FrozenLake environments,
+enabling integration with the MCP-Gym framework.
+"""
+
+from typing import Any, Dict, Optional, Tuple
+
+from gymnasium.envs.toy_text.frozen_lake import FrozenLakeEnv, generate_random_map
+
+from eval_protocol.mcp.adapter import EnvironmentAdapter
+
+
+class FrozenLakeAdapter(EnvironmentAdapter):
+    """FrozenLake adapter for MCP-Gym framework."""
+
+    ACTION_NAMES = ["LEFT", "DOWN", "RIGHT", "UP"]
+
+    def create_environment(self, config: Optional[Dict[str, Any]] = None) -> FrozenLakeEnv:
+        """
+        Create FrozenLake environment.
+
+        Args:
+            config: Configuration dictionary with optional 'map_name' and 'seed'
+
+        Returns:
+            FrozenLake environment instance
+        """
+        print(f"🔍 FrozenLakeAdapter.create_environment: config: {config}")
+        config = config or {}
+
+        # Determine grid size from config
+        grid_size = 4
+        if "map_name" in config:
+            if "8x8" in config["map_name"]:
+                grid_size = 8
+
+        # Generate random map if seed is provided
+        seed = config.get("seed")
+        print(f"🔍 FrozenLakeAdapter.create_environment: extracted seed: {seed} (type: {type(seed)})")
+        print(f"🔍 FrozenLakeAdapter.create_environment: grid_size: {grid_size}")
+
+        if seed is not None:
+            print(f"🔍 FrozenLakeAdapter.create_environment: Generating map with seed {seed}")
+            desc = generate_random_map(size=grid_size, p=0.8, seed=seed)
+            print(f"🔍 FrozenLakeAdapter.create_environment: Generated map desc: {desc}")
+        else:
+            print("🔍 FrozenLakeAdapter.create_environment: Generating map without seed")
+            desc = generate_random_map(size=grid_size, p=0.8)
+            print(f"🔍 FrozenLakeAdapter.create_environment: Generated map desc: {desc}")
+
+        env = FrozenLakeEnv(desc=desc, is_slippery=False, render_mode="ansi")
+        print("🔍 FrozenLakeAdapter.create_environment: Created FrozenLakeEnv")
+        return env
+
+    def create_environment_with_seed(
+        self, config: Optional[Dict[str, Any]] = None, seed: Optional[int] = None
+    ) -> Tuple[FrozenLakeEnv, int, Dict[str, Any]]:
+        """
+        Create FrozenLake environment with seed and return initial state.
+
+        Args:
+            config: Configuration dictionary
+            seed: Seed for reproducible environments
+
+        Returns:
+            Tuple of (environment, initial_observation, initial_info)
+        """
+        print(f"🔍 FrozenLakeAdapter.create_environment_with_seed: config: {config}, seed: {seed}")
+        config = config or {}
+
+        # Add seed to config for environment creation
+        env_config = {**config, "seed": seed}
+        print(f"🔍 FrozenLakeAdapter.create_environment_with_seed: env_config: {env_config}")
+
+        env = self.create_environment(env_config)
+        print(f"🔍 FrozenLakeAdapter.create_environment_with_seed: created env, calling reset with seed: {seed}")
+        obs, info = env.reset(seed=seed)
+        print(f"🔍 FrozenLakeAdapter.create_environment_with_seed: reset returned obs: {obs}, info: {info}")
+
+        return env, obs, info
+
+    def reset_environment(self, env: FrozenLakeEnv, seed: Optional[int] = None) -> Tuple[int, Dict[str, Any]]:
+        """
+        Reset environment.
+
+        Args:
+            env: Environment instance
+            seed: Optional seed for reset
+
+        Returns:
+            Tuple of (observation, info)
+        """
+        return env.reset(seed=seed)
+
+    def step_environment(self, env: FrozenLakeEnv, action: int) -> Tuple[int, float, bool, bool, Dict[str, Any]]:
+        """
+        Execute environment step.
+
+        Args:
+            env: Environment instance
+            action: Action index
+
+        Returns:
+            Tuple of (observation, reward, terminated, truncated, info)
+        """
+        return env.step(action)
+
+    def close_environment(self, env: FrozenLakeEnv) -> None:
+        """
+        Close environment.
+
+        Args:
+            env: Environment instance
+        """
+        # FrozenLake doesn't need explicit cleanup
+        pass
+
+    def parse_action(self, action_str: str) -> int:
+        """
+        Parse action string to integer.
+
+        Args:
+            action_str: Action string (LEFT, DOWN, RIGHT, UP)
+
+        Returns:
+            Action index
+
+        Raises:
+            ValueError: If action is invalid
+        """
+        action_str = action_str.strip().upper()
+        if action_str not in self.ACTION_NAMES:
+            raise ValueError(f"Invalid action '{action_str}'. Valid actions: {self.ACTION_NAMES}")
+        return self.ACTION_NAMES.index(action_str)
+
+    def format_observation(self, observation: int) -> int:
+        """
+        Format observation for JSON serialization.
+
+        Args:
+            observation: Raw observation from environment
+
+        Returns:
+            Formatted observation
+        """
+        return int(observation)
+
+    def get_default_config(self) -> Dict[str, Any]:
+        """
+        Get default configuration.
+
+        Returns:
+            Default configuration dictionary
+        """
+        return {
+            "map_name": "4x4",
+            "is_slippery": False,
+        }
diff --git a/eval_protocol/mcp_servers/frozen_lake/frozen_lake_mcp.py b/eval_protocol/mcp_servers/frozen_lake/frozen_lake_mcp.py
@@ -0,0 +1,102 @@
+"""
+FrozenLake MCP-Gym Implementation
+
+This module implements the north star vision for MCP-Gym environments,
+providing a clean, simple implementation of FrozenLake using the McpGym base class.
+
+Key Features:
+- Multi-session support with session-based control plane state
+- Data plane: Tool responses contain only observations
+- Control plane: Server-side state management keyed by session ID
+- Rollout system can query control plane state for termination logic
+
+Example usage:
+    from frozen_lake_mcp import FrozenLakeMcp
+
+    server = FrozenLakeMcp(seed=42)
+    server.run()
+"""
+
+from typing import Any, Dict, Optional
+
+from frozen_lake_adapter import FrozenLakeAdapter
+from mcp.server.fastmcp import Context
+
+from eval_protocol.mcp import McpGym
+
+
+class FrozenLakeMcp(McpGym):
+    """
+    FrozenLake MCP-Gym environment implementing the north star vision.
+
+    This demonstrates the clean, simple API for MCP-Gym environments:
+    - Inherit from McpGym (which inherits from GymProductionServer)
+    - Use proper EnvironmentAdapter pattern
+    - Register tools with @self.mcp.tool() decorator
+    - Compatible with CondaServerProcessManager
+    - Multi-session support with session-based control plane state
+    """
+
+    def __init__(self, seed: Optional[int] = None, **kwargs):
+        """Initialize FrozenLake MCP-Gym environment."""
+        adapter = FrozenLakeAdapter()
+        super().__init__("FrozenLake-v1", adapter, seed, **kwargs)
+
+        # Multi-session support is now handled by the base class
+
+    # Session management methods are now handled by the base class
+
+    def _register_tools(self):
+        """Register domain-specific MCP tools."""
+
+        @self.mcp.tool(
+            name="lake_move",
+            description="Move on the frozen lake. Actions: LEFT, DOWN, RIGHT, UP. "
+            "Returns only observation data; control plane state managed server-side.",
+        )
+        def lake_move(action: str, ctx: Context) -> Dict[str, Any]:
+            """
+            Move in the FrozenLake environment.
+
+            Args:
+                action: Direction to move (LEFT, DOWN, RIGHT, UP)
+                ctx: MCP context (proper FastMCP context)
+
+            Returns:
+                Dictionary with observation data ONLY (data plane).
+                Control plane state managed server-side per session.
+            """
+            # Validate action
+            if not action or not isinstance(action, str):
+                raise ValueError(
+                    f"Invalid action parameter: '{action}'. "
+                    f"Must be a non-empty string. Valid actions: LEFT, DOWN, RIGHT, UP"
+                )
+
+            action = action.strip().upper()
+
+            # Parse action
+            try:
+                action_int = self.adapter.parse_action(action)
+            except ValueError as e:
+                raise ValueError(str(e))
+
+            # Get session ID and session data
+            session_id = self._get_session_id(ctx)
+            session_data = self._get_or_create_session(ctx)
+
+            # Execute environment step using base class method
+            observation_data = self._execute_session_environment_step(session_id, action_int)
+            observation_data["action"] = action
+
+            # Log move (no control plane data in logs)
+            print(f"🎮 Session {session_id[:16]}...: {action} → position {session_data['obs']}")
+
+            return observation_data
+
+    def format_observation(self, obs: int, env: Any) -> Dict[str, Any]:
+        """Format observation for MCP response (data plane only)."""
+        return {
+            "position": int(obs),
+            "grid": env.render(),
+        }
diff --git a/eval_protocol/mcp_servers/frozen_lake/server.py b/eval_protocol/mcp_servers/frozen_lake/server.py
@@ -0,0 +1,52 @@
+#!/usr/bin/env python3
+"""
+FrozenLake MCP-Gym Server
+
+This script launches the FrozenLake MCP-Gym server using the proper MCP-Gym framework.
+Compatible with CondaServerProcessManager for isolated execution.
+
+Usage:
+    python server.py --port 9004 --seed 42
+"""
+
+import argparse
+import os
+import sys
+from pathlib import Path
+
+# Add root directory to path so we can import eval_protocol
+sys.path.insert(0, str(Path(__file__).parent.parent.parent))
+
+from frozen_lake_mcp import FrozenLakeMcp
+
+
+def main():
+    """Run the FrozenLake MCP server."""
+    parser = argparse.ArgumentParser(description="FrozenLake MCP Server")
+    parser.add_argument(
+        "--transport",
+        choices=["streamable-http", "stdio"],
+        default="streamable-http",
+        help="Transport protocol to use",
+    )
+    parser.add_argument("--port", type=int, default=8000, help="Port for HTTP transport")
+    parser.add_argument("--seed", type=int, default=None, help="Seed for the environment")
+
+    args = parser.parse_args()
+
+    # Set environment variable for HTTP port (required by FastMCP)
+    if args.transport == "streamable-http":
+        os.environ["PORT"] = str(args.port)
+
+    # Create and run server
+    server = FrozenLakeMcp(seed=args.seed)
+
+    print(f"🚀 Starting FrozenLake MCP server on port {args.port}")
+    print(f"🌱 Seed: {args.seed}")
+    print(f"📡 Transport: {args.transport}")
+
+    server.run(transport=args.transport)
+
+
+if __name__ == "__main__":
+    main()
