[WIP] feat: MCP adapter for Gymnasium

.github/workflows/pre-commit.yml

@@ -15,7 +15,7 @@ jobs:
     steps:
       - uses: actions/checkout@v7
       - uses: actions/setup-python@v7
-      - run: pip install pre-commit
+      - run: pip install pre-commit pillow "fastmcp>=2,<5"
       - run: pre-commit --version
       - run: pre-commit install
       - run: pre-commit run --all-files

.github/workflows/run-tests.yml

@@ -108,5 +108,22 @@ jobs:
           docker build -f bin/android_env.Dockerfile \
             --build-arg PYTHON_VERSION='${{ matrix.python-version }}' \
             --tag shimmy-android-env-docker .
-      - name: Run android_env tests
+      - name: Run AndroidEnv tests
         run: docker run shimmy-android-env-docker pytest tests/test_android_env.py
+
+  optional-test-gymnasium-mcp:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ['3.10', '3.12']
+    steps:
+      - uses: actions/checkout@v6
+
+      # AndroidEnv
+      - run: |
+          docker build -f bin/gymnasium-mcp.Dockerfile \
+            --build-arg PYTHON_VERSION='${{ matrix.python-version }}' \
+            --tag shimmy-gymnasium-mcp-docker .
+      - name: Run Gymnasium MCP tests
+        run: docker run shimmy-gymnasium-mcp-docker pytest tests/test_gymnasium_mcp.py

bin/gymnasium-mcp.Dockerfile

@@ -0,0 +1,16 @@
+# A Dockerfile that sets up an openspiel
+
+# if PYTHON_VERSION is not specified as a build argument, set it to 3.10.
+ARG PYTHON_VERSION
+ARG PYTHON_VERSION=${PYTHON_VERSION:-3.10}
+FROM python:$PYTHON_VERSION
+
+SHELL ["/bin/bash", "-o", "pipefail", "-c"]
+
+RUN pip install --upgrade pip
+
+COPY . /usr/local/shimmy/
+WORKDIR /usr/local/shimmy/
+
+# Install Shimmy
+RUN pip install ".[mcp, testing]" --no-cache-dir

setup.py

@@ -34,6 +34,7 @@ def get_version():
 header_count, long_description = get_description()
 
 extras = {
+    "mcp": ["fastmcp>=2,<5", "pillow"],
     "gym-v21": ["gym==0.21.0", "pyglet==1.5.11", "numpy<2.0"],
     "gym-v26": ["gym==0.26.2", "numpy<2.0"],
     "dm-control": ["dm-control>=1.0.10", "imageio", "h5py>=3.7.0"],
@@ -68,6 +69,9 @@ def get_version():
     install_requires=["numpy>=1.18.0", "gymnasium>=1.0.0"],
     tests_require=extras["testing"],
     extras_require=extras,
+    entry_points={
+        "console_scripts": ["gym-mcp=shimmy.mcp_adapters.gymnasium_mcp:main"],
+    },
     classifiers=[
         "Programming Language :: Python :: 3.10",
         "Programming Language :: Python :: 3.11",

shimmy/mcp_adapters/__init__.py

@@ -0,0 +1,5 @@
+"""Optional MCP adapters."""
+
+from shimmy.mcp_adapters.gymnasium_interface import GymnasiumMCPAdapter
+
+__all__ = ["GymnasiumMCPAdapter"]

shimmy/mcp_adapters/gymnasium_interface.py

@@ -0,0 +1,215 @@
+"""Expose a single Gymnasium environment through FastMCP."""
+
+import json
+from functools import cache
+from threading import RLock
+from typing import Any, cast
+
+import gymnasium
+import numpy as np
+from numpy.typing import NDArray
+
+from shimmy.mcp_adapters.spaces import (
+    decode,
+    describe,
+    encode,
+    image_content,
+    image_layout,
+    json_value,
+)
+
+try:
+    from importlib import import_module
+
+    from fastmcp import FastMCP
+    from mcp.types import TextContent
+
+    import_module("PIL.Image")
+except ImportError as e:
+    raise ImportError("MCP support requires: pip install 'shimmy[mcp]'") from e
+
+
+class GymnasiumMCPAdapter:
+    """Wrap an environment with MCP tools, resources, and episode guidance."""
+
+    def __init__(
+        self,
+        env: gymnasium.Env,
+        name: str | None = None,
+        *,
+        max_image_bytes: int = 5_000_000,
+        **kwargs,
+    ):
+        """Create the server.
+
+        Args:
+            env: Environment owned by the caller.
+            name: Optional server name override.
+            max_image_bytes: Maximum base64 bytes per image.
+            **kwargs: Forwarded to FastMCP.
+        """
+        if max_image_bytes <= 0:
+            raise ValueError("max_image_bytes must be positive")
+        self.env = env
+        self.max_image_bytes = max_image_bytes
+        self._lock = RLock()
+        self._closed = False
+        self.total_reward = 0.0
+        env_id = getattr(env.spec, "id", None) or "unknown"
+        # TODO: improve this default instruction
+        kwargs.setdefault(
+            "instructions",
+            "Read gymnasium://observation_space and gymnasium://action_space, reset, "
+            "then step with a JSON action until terminated or truncated. sample_action "
+            "provides valid actions. Image observations need to be check via the "
+            "`render` tool'. Close when finished.",
+        )
+        self.mcp = FastMCP(
+            name if name is not None else f"Gymnasium[{env_id}]", **kwargs
+        )
+        for method in (
+            self.reset,
+            self.step,
+            self.render,
+            self.close,
+            self.sample_action,
+        ):
+            self.mcp.tool()(method)
+        self._resource("spec", lambda: env.spec.to_json() if env.spec else "{}")
+        self._resource(
+            "metadata",
+            lambda: json.dumps(
+                json_value({**env.metadata, "render_mode": env.render_mode}),
+                allow_nan=False,
+            ),
+        )
+        self._resource(
+            "observation_space",
+            lambda: json.dumps(
+                describe(env.observation_space, images=True), allow_nan=False
+            ),
+        )
+        self._resource(
+            "action_space",
+            lambda: json.dumps(describe(env.action_space), allow_nan=False),
+        )
+        self.mcp.prompt()(self.play_episode)
+
+    def _resource(self, name, factory):
+        @cache
+        def read() -> str:
+            with self._lock:
+                return factory()
+
+        self.mcp.resource(
+            f"gymnasium://{name}", name=name, mime_type="application/json"
+        )(read)
+
+    def __getattr__(self, name):
+        """Delegate server methods to FastMCP."""
+        return getattr(self.mcp, name)
+
+    def _observation(self, observation, **fields):
+        result = {
+            "observation": encode(self.env.observation_space, observation, images=True),
+            **json_value(fields),
+        }
+        return [
+            TextContent(type="text", text=json.dumps(result, allow_nan=False)),
+        ]
+
+    def reset(self, seed: int | None = None, options: dict | None = None) -> Any:
+        """Start an episode and return its observation and info."""
+        with self._lock:
+            observation, info = self.env.reset(seed=seed, options=options)
+            self._closed = False
+            self.total_reward = 0.0
+            return self._observation(
+                observation, info=info, total_reward=self.total_reward
+            )
+
+    def step(self, action: Any) -> Any:
+        """Apply one JSON action and return the transition with render hints for images."""
+        with self._lock:
+            try:
+                action = decode(self.env.action_space, action)
+            except ValueError as exc:
+                return {"error": True, "message": str(exc)}
+            observation, reward, terminated, truncated, info = self.env.step(action)
+            self.total_reward += float(reward)
+            return self._observation(
+                observation,
+                reward=reward,
+                total_reward=self.total_reward,
+                terminated=terminated,
+                truncated=truncated,
+                info=info,
+            )
+
+    def render(self) -> Any:
+        """Return the current render as text, images, or status metadata."""
+        with self._lock:
+            value = self.env.render()
+            if isinstance(value, str):
+                return TextContent(type="text", text=value)
+            is_list_expected = self.env.render_mode and self.env.render_mode.endswith(
+                "_list"
+            )
+            frames = value if is_list_expected else [value]
+            assert (
+                frames is not None
+            ), f"render_mode of '{self.env.render_mode}' should not produce None"
+            if all(isinstance(elem, str) for elem in frames):
+                return TextContent(type="text", text=str(frames))
+            content = [
+                TextContent(
+                    type="text", text=json.dumps({"render_mode": self.env.render_mode})
+                )
+            ]
+            for frame in frames:
+                layout = (
+                    image_layout(gymnasium.spaces.Box(0, 255, frame.shape, np.uint8))
+                    if isinstance(frame, np.ndarray) and frame.dtype == np.uint8
+                    else None
+                )
+                if layout:
+                    content.append(
+                        image_content(
+                            cast(NDArray[np.uint8], frame),
+                            layout,
+                            self.max_image_bytes,
+                        )
+                    )
+                else:
+                    content.append(
+                        TextContent(
+                            type="text",
+                            text=json.dumps(
+                                {
+                                    "render_mode": self.env.render_mode,
+                                    "data": json_value(frame),
+                                },
+                                allow_nan=False,
+                            ),
+                        )
+                    )
+            return content
+
+    def close(self) -> dict:
+        """Close the environment."""
+        with self._lock:
+            if not self._closed:
+                self.env.close()
+                self._closed = True
+            return {"ok": True}
+
+    def sample_action(self) -> Any:
+        """Sample a valid JSON action."""
+        with self._lock:
+            return encode(self.env.action_space, self.env.action_space.sample())
+
+    def play_episode(self, goal: str | None = None) -> str:
+        """Explain how to play one episode."""
+        return f"{self.mcp.instructions}\n" + (
+            f"Goal: {goal}" if goal else "Play one episode."
+        )

shimmy/mcp_adapters/gymnasium_mcp.py

@@ -0,0 +1,95 @@
+"""Run a Gymnasium environment over MCP."""
+
+import argparse
+import importlib
+import json
+
+import gymnasium
+
+from shimmy.mcp_adapters.gymnasium_interface import GymnasiumMCPAdapter
+
+
+def try_load_json(value: str):
+    """Attempt to load value as json."""
+    try:
+        return json.loads(value)
+    except json.JSONDecodeError:
+        return value
+
+
+class GymEnvironmentArgument(argparse.Action):
+    """Process keyword arguments that should passthrough to ``gym.make``."""
+
+    def __call__(self, parser, namespace, values, option_string=None):
+        """Add arbitrary named arguments into ``namespace.environment_kwargs``."""
+        kwargs = getattr(namespace, "environment_kwargs", None)
+        if kwargs is None:
+            kwargs = namespace.environment_kwargs = {}
+        if self.dest in kwargs:
+            raise KeyError(f"Duplicate environment argument: {self.dest}")
+        kwargs[self.dest] = values
+
+
+def main(argv: list[str] | None = None) -> None:
+    """Start MCP, forwarding extra flags and JSON kwargs to gymnasium.make."""
+    parser = argparse.ArgumentParser(
+        description=__doc__,
+        allow_abbrev=False,
+        epilog="Extra --key value or --key=value flags become environment kwargs. "
+        "Values use JSON types when valid, otherwise strings; bare flags mean true. "
+        "Hyphens become underscores. Duplicate keys raise KeyError.",
+    )
+    parser.add_argument("-i", "--import", dest="imports", default="")
+    parser.add_argument(
+        "-t", "--transport", default="stdio", help="MCP transport (default: stdio)"
+    )
+    parser.add_argument(
+        "-r",
+        "--render-mode",
+        action=GymEnvironmentArgument,
+        help="Gymnasium render mode",
+    )
+    parser.add_argument(
+        "--kwargs",
+        type=json.loads,
+        default={},
+        help="JSON object of gymnasium.make arguments",
+    )
+    parser.add_argument("env_id")
+    _, unknown = parser.parse_known_args(argv)
+    for flag in dict.fromkeys(token.split("=", 1)[0] for token in unknown):
+        if (
+            flag.startswith("-")
+            and flag != "--"
+            and isinstance(try_load_json(flag), str)
+        ):
+            parser.add_argument(
+                flag,
+                dest=flag.lstrip("-").replace("-", "_"),
+                type=try_load_json,
+                nargs="?",
+                const=True,
+                action=GymEnvironmentArgument,
+            )
+    args = parser.parse_args(argv)
+    if not isinstance(args.kwargs, dict):
+        parser.error("--kwargs must be a JSON object")
+    kwargs = getattr(args, "environment_kwargs", {})
+    duplicates = kwargs.keys() & args.kwargs.keys()
+    if duplicates:
+        raise KeyError(
+            f"Duplicate environment arguments: {', '.join(sorted(duplicates))}"
+        )
+    kwargs.update(args.kwargs)
+    for module in args.imports.split(","):
+        if module.strip():
+            importlib.import_module(module.strip())
+    env = gymnasium.make(args.env_id, **kwargs)
+    try:
+        GymnasiumMCPAdapter(env).run(transport=args.transport)
+    finally:
+        env.close()
+
+
+if __name__ == "__main__":
+    main()