GEM Framework: Agentic RL and MCP Environments

graph TD
    Agent[Agentic LLM / Policy] -->|1. Generate tools/call| GEMCore[GEM Environment Runner]
    GEMCore -->|2. Dispatch JSON-RPC| MCPStdio[MCP Server: Docker / Postgres / Git]
    MCPStdio -->|3. Return Tool Result| GEMCore
    GEMCore -->|4. Run Test Assertions| Verifier[Reward Model / Verifier]
    Verifier -->|5. Compute Reward r_t| Optimizer[PPO / GRPO Loss Function]
    GEMCore -->|6. Append Observation o_t+1| Agent

# gem_mcp_env.py
import asyncio
import json
from typing import Dict, Any, Tuple
from dataclasses import dataclass

@dataclass
class StepResult:
    observation: Dict[str, Any]
    reward: float
    done: bool
    info: Dict[str, Any]

class MCPGEMEnvironment:
    """
    Gymnasium-compatible wrapper connecting GEM RL training loops
    to Model Context Protocol (MCP) servers over stdio.
    """
    def __init__(self, server_command: str, server_args: list[str], max_steps: int = 15):
        self.server_command = server_command
        self.server_args = server_args
        self.max_steps = max_steps
        self.current_step = 0
        self.process = None

    async def reset(self) -> Dict[str, Any]:
        """Start a fresh MCP server subprocess for a clean episode sandbox."""
        self.current_step = 0
        if self.process:
            self.process.terminate()
            await self.process.wait()

        self.process = await asyncio.create_subprocess_exec(
            self.server_command,
            *self.server_args,
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )

        # Initialize protocol handshake
        init_payload = {
            "jsonrpc": "2.0",
            "id": 1,
            "method": "initialize",
            "params": {
                "protocolVersion": "2024-11-05",
                "capabilities": {"tools": {}},
                "clientInfo": {"name": "GEM-Gymnasium", "version": "1.0.0"}
            }
        }
        await self._send(init_payload)
        init_response = await self._recv()

        # Discover initial tools state
        tools_req = {"jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": {}}
        await self._send(tools_req)
        tools_response = await self._recv()

        return {
            "prompt": "Restore database state and fix failed SQL indexes.",
            "available_tools": tools_response.get("result", {}).get("tools", [])
        }

    async def step(self, tool_call_action: Dict[str, Any]) -> StepResult:
        """Execute agent's tool call against MCP server and compute reward."""
        self.current_step += 1
        
        rpc_req = {
            "jsonrpc": "2.0",
            "id": self.current_step + 10,
            "method": "tools/call",
            "params": {
                "name": tool_call_action.get("name"),
                "arguments": tool_call_action.get("arguments", {})
            }
        }

        await self._send(rpc_req)
        rpc_res = await self._recv()

        # Verify whether the action advanced the task
        observation = rpc_res.get("result", rpc_res.get("error", {}))
        reward, done, info = self._evaluate_state(tool_call_action, observation)

        if self.current_step >= self.max_steps:
            done = True

        return StepResult(observation=observation, reward=reward, done=done, info=info)

    def _evaluate_state(self, action: Dict[str, Any], observation: Dict[str, Any]) -> Tuple[float, bool, Dict[str, Any]]:
        # Verifiable reward calculation based on execution success
        if "error" in observation:
            return -0.2, False, {"error": observation["error"]}
        
        # Check task completion criteria
        output_text = json.dumps(observation)
        if "QUERY PLAN" in output_text and "Index Scan" in output_text:
            return 1.0, True, {"status": "success_optimized"}

        return 0.1, False, {"status": "progressing"}

    async def _send(self, payload: Dict[str, Any]):
        data = json.dumps(payload) + "\n"
        self.process.stdin.write(data.encode())
        await self.process.stdin.drain()

    async def _recv(self) -> Dict[str, Any]:
        line = await self.process.stdout.readline()
        if not line:
            return {"error": "Server connection terminated"}
        return json.loads(line.decode())

{
  "gem_runner_config": {
    "parallel_workers": 32,
    "mcp_server": {
      "command": "uvx",
      "args": ["mcp-server-sqlite", "--db-path", ":memory:"]
    },
    "timeout_per_step_ms": 3000,
    "max_episode_steps": 12,
    "reward_discount_factor": 0.99
  }
}