Merge pull request #9 from reward-protocol/derekx/refactor-environment

Environment Creation moved from GymProductionServer to McpGym

Author: xzrderek

examples/lunar_lander_mcp/mcp_server/lunar_lander_mcp_server.py

@@ -34,6 +34,7 @@
 
 from reward_kit.mcp import GymProductionServer
 
+# TODO: FAST FOLLOW. refactor this entire file to use McpGym, leaving logic below incorrect for now.
 
 class LunarLanderProdServer(GymProductionServer):
     """LunarLander production server with visual rendering support."""
@@ -188,4 +189,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()

examples/taxi_mcp_complete/mcp_server/taxi_mcp_server.py

@@ -25,8 +25,9 @@
 from mcp.server.fastmcp import Context
 from taxi_adapter import TaxiAdapter
 
-from reward_protocol.mcp import GymProductionServer
+from reward_kit.mcp import GymProductionServer
 
+# TODO: FAST FOLLOW. refactor this entire file to use McpGym, leaving logic below incorrect for now.
 
 class TaxiProdServer(GymProductionServer):
     """Taxi production server using unified framework."""
@@ -53,6 +54,7 @@ def taxi_move(action: str, ctx: Context) -> Dict[str, Any]:
             """
             # Extract seed from client info and reinitialize if needed
             self.extract_seed_from_context(ctx)
+            # note for later, should be: session_data = self._get_or_create_session(ctx)
 
             # Validate action
             if not action or not isinstance(action, str):
@@ -169,4 +171,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()

reward_kit/mcp/gym_production_server.py

@@ -54,9 +54,6 @@ def __init__(self, name: str, adapter: EnvironmentAdapter):
         """
         self.adapter = adapter
 
-        # For backward compatibility, keep single-session support
-        self.env, self.obs, _info = self._new_env()
-
         # Multi-session support
         self.sessions = (
             {}
@@ -74,20 +71,6 @@ def __init__(self, name: str, adapter: EnvironmentAdapter):
         self._register_resources()
         self._register_tools()
 
-    def _new_env(self, seed: Optional[int] = None) -> Tuple[Any, Any, Dict]:
-        """Create new environment and return initial state."""
-        if hasattr(self.adapter, "create_environment_with_seed"):
-            env, obs, info = self.adapter.create_environment_with_seed(
-                self.adapter.get_default_config(), seed=seed
-            )
-        else:
-            env = self.adapter.create_environment(self.adapter.get_default_config())
-            obs, info = self.adapter.reset_environment(env, seed=seed)
-        return env, obs, info
-
-    def _render(self, obs) -> Dict[str, Any]:
-        """Format observation using subclass implementation."""
-        return self.format_observation(obs, self.env)
 
     def _register_resources(self):
         """Register standard MCP resources."""
@@ -250,28 +233,6 @@ def _get_or_create_session(self, ctx: Context) -> Dict[str, Any]:
 
             return self.sessions[session_id]
 
-    def extract_seed_from_context(self, ctx: Context) -> Optional[int]:
-        """
-        Extract seed from MCP client info if available.
-
-        NOTE: This method is kept for backward compatibility. New code should use
-        _get_or_create_session() which handles seed extraction automatically.
-        """
-        if hasattr(ctx, "session") and hasattr(ctx.session, "client_params"):
-            client_params = ctx.session.client_params
-            if hasattr(client_params, "clientInfo"):
-                client_info = client_params.clientInfo
-                if client_info and hasattr(client_info, "_extra"):
-                    extra_data = client_info._extra
-                    if extra_data and isinstance(extra_data, dict):
-                        seed = extra_data.get("seed")
-                        if seed is not None:
-                            print(f"🌱 Reinitializing with seed from client: {seed}")
-                            self.env, self.obs, _info = self._new_env(seed=seed)
-                            return seed
-
-        return None
-
     # Abstract methods that subclasses must implement
 
     @abstractmethod
@@ -293,4 +254,4 @@ def run(self, transport: str = "streamable-http", **kwargs):
         print("🔗 Initial state resource: game://initial_state")
 
         # Run the server
-        self.mcp.run(transport=transport, **kwargs)
+        self.mcp.run(transport=transport, **kwargs)

reward_kit/mcp/mcpgym.py

@@ -98,8 +98,7 @@ def __init__(
         }
 
         # Reset with seed if provided
-        if seed is not None:
-            self.env, self.obs, _info = self._new_env(seed=seed)
+        self.env, self.obs, _info = self._new_env(seed=seed)
 
         # Discover and register control plane endpoints
         self._discover_and_register_control_plane_endpoints()
@@ -434,3 +433,19 @@ def format_observation(obs: Any, env: Any) -> Dict[str, Any]:
             Formatted observation dictionary (DATA PLANE ONLY)
         """
         pass
+
+    def _new_env(self, seed: Optional[int] = None) -> Tuple[Any, Any, Dict]:
+        """Create new environment and return initial state."""
+        config = self.adapter.get_default_config()
+        
+        try:
+            env, obs, info = self.adapter.create_environment_with_seed(config, seed=seed)
+        except AttributeError:
+            env = self.adapter.create_environment(config)
+            obs, info = self.adapter.reset_environment(env, seed=seed)
+        
+        return env, obs, info
+
+    def _render(self, obs) -> Dict[str, Any]:
+        """Format observation using subclass implementation."""
+        return self.format_observation(obs, self.env)