Creating Environment Wrappers for Existing RL Benchmarks like Atari or Chess with OpenEnv

OpenEnv provides a Gymnasium-style API that transforms any RL benchmark into a fully-featured, containerized execution environment using a three-layer architecture of typed data models, FastAPI servers, and thin client wrappers.

OpenEnv by Hugging Face simplifies the process of creating environment wrappers for existing RL benchmarks like Atari or chess by providing a standardized, Docker-native framework. This library abstracts away infrastructure complexity, allowing researchers to wrap legacy environments—whether they use the Arcade Learning Environment (ALE) or python-chess—into isolated, WebSocket-connected services that expose a clean, type-safe Python API.

Architecture Overview

The wrapper architecture consists of three distinct layers that together create an isolated execution environment:

1. Data-model layer – Action, Observation, and State dataclasses that define the strict interaction contract between agent and environment.
2. Server layer – A FastAPI service running inside a Docker container that implements the environment logic (e.g., ALE for Atari, python-chess for Chess) and exposes /reset, /step, and /state endpoints over WebSocket.
3. Client layer – A thin subclass of openenv.core.EnvClient that serializes actions, parses results, and manages the WebSocket connection or automatically spins up the Docker container.

The EnvClient base class lives in src/openenv/core/env_client.py, while the generic StepResult type is defined in src/openenv/core/client_types.py.

Building an Atari Wrapper

To wrap the Arcade Learning Environment (ALE), you implement a six-step pipeline that bridges the legacy API with OpenEnv's modern interface.

1. Scaffold the Project

Run the CLI initializer to create the folder structure and minimal configuration:

openenv init atari_env

This generates the openenv.yaml manifest and directory layout according to the openenv init implementation.

2. Define Typed Models

Create envs/atari_env/models.py to declare AtariAction, AtariObservation, and AtariState dataclasses. These Pydantic models validate every request and response, ensuring type safety across the client-server boundary.

3. Implement the Server Logic

In envs/atari_env/server/atari_environment.py, wrap the ALE interface by mapping incoming AtariAction objects to ALE method calls and constructing observation payloads. This file contains the core environment state machine.

4. Expose FastAPI Endpoints

The envs/atari_env/server/app.py file creates HTTP and WebSocket entry points for /reset, /step, and /state. This leverages the create_web_interface_app function from src/openenv/core/env_server/web_interface.py to handle connection management.

5. Containerize with Docker

The envs/atari_env/server/Dockerfile installs ale-py and launches the FastAPI application. This ensures the Atari environment runs in complete isolation with all native dependencies properly configured.

6. Write the Client Subclass

Finally, envs/atari_env/client.py subclasses EnvClient and implements three critical methods:

• _step_payload – Serializes actions for transmission
• _parse_result – Deserializes server responses into StepResult objects
• _parse_state – Handles state-specific parsing logic

Building a Chess Wrapper

Creating a Chess wrapper follows the identical pattern but substitutes the server implementation for python-chess and the moonfish engine.

The key files mirror the Atari structure:

• envs/chess_env/models.py – Defines ChessAction, ChessObservation, and ChessState
• envs/chess_env/server/chess_environment.py – Glue code that loads chess boards, applies UCI moves, and returns legal-move lists
• envs/chess_env/server/app.py – FastAPI entry point with identical endpoint structure
• envs/chess_env/client.py – ChessEnv subclass of EnvClient
• envs/chess_env/Dockerfile – Build image with python-chess and moonfish dependencies

All implementations reuse the generic core components from src/openenv/core/, ensuring API uniformity across different RL benchmarks.

Key Design Invariants

OpenEnv enforces four critical design constraints across all wrapper implementations:

• Stateless client – The client never stores episode data; the server maintains the canonical state in envs/<env_name>/server/<environment>.py.
• Typed contracts – Every request/response validates against strict dataclass schemas defined in the respective models.py files.
• Container-agnostic – You can swap Docker providers (local Docker, Docker Swarm, Kubernetes) without modifying client code, thanks to the abstraction in src/openenv/core/containers/runtime/uv_provider.py.
• WebSocket-first – Low-latency persistent connections are the default transport; HTTP fallback remains optional.

Code Examples

Basic Async Usage (Local Server)

When running the server locally on port 8000, use the async context manager for non-blocking interaction:

import asyncio
from atari_env import AtariEnv, AtariAction

async def run():
    # Assumes local server: python -m atari_env.server.app

    async with AtariEnv(base_url="http://localhost:8000") as env:
        # Initialise episode

        result = await env.reset()
        print("Screen shape:", result.observation.screen_shape)

        # Randomly choose a legal action

        action_id = result.observation.legal_actions[0]
        result = await env.step(AtariAction(action_id=action_id))
        print("Reward:", result.reward, "Done:", result.done)

asyncio.run(run())

The AtariEnv implementation resides in envs/atari_env/client.py.

Synchronous Usage via .sync()

For synchronous training loops, use the .sync() wrapper method defined in src/openenv/core/env_client.py:

from atari_env import AtariEnv, AtariAction

with AtariEnv(base_url="http://localhost:8000").sync() as env:
    result = env.reset()
    result = env.step(AtariAction(action_id=2))   # e.g., UP in Pong

    print("Reward:", result.reward)

Docker-Managed Client

Eliminate manual docker run commands by using from_docker_image, which automatically spins up containers:

import asyncio
from atari_env import AtariEnv, AtariAction

async def train():
    # Automatically spins up the container

    client = await AtariEnv.from_docker_image("atari-env:latest")
    async with client:
        for ep in range(5):
            result = await client.reset()
            while not result.done:
                action = AtariAction(action_id=result.observation.legal_actions[0])
                result = await client.step(action)
            print(f"Episode {ep} finished with reward {result.reward}")

asyncio.run(train())

This functionality is implemented in src/openenv/core/env_client.py with container orchestration handled by src/openenv/core/containers/runtime/uv_provider.py.

Chess Environment Example

The Chess client follows identical patterns but uses UCI notation for actions:

import asyncio
from chess_env import ChessEnv, ChessAction

async def play():
    async with ChessEnv(base_url="http://localhost:8000") as env:
        obs = await env.reset()
        print("Starting FEN:", obs.observation.fen)

        while not obs.done:
            move = obs.observation.legal_moves[0]          # naive policy

            obs = await env.step(ChessAction(move=move))

        print("Game result:", obs.observation.result)

asyncio.run(play())

Reference the implementation in envs/chess_env/client.py.

Summary

• OpenEnv standardizes RL benchmark wrapping through a three-layer architecture: typed models, FastAPI servers, and thin clients.
• Atari and Chess wrappers demonstrate the pattern: define models in models.py, implement logic in server/<environment>.py, expose via FastAPI in server/app.py, and interface via client.py.
• Containerization is first-class; use AtariEnv.from_docker_image() or ChessEnv.from_docker_image() for zero-config deployment.
• Client code remains agnostic to container runtime thanks to abstractions in src/openenv/core/containers/.
• WebSocket transport provides low-latency communication, while the StepResult type in src/openenv/core/client_types.py ensures consistent return values across all environments.

Frequently Asked Questions

How does OpenEnv differ from standard Gymnasium wrappers?

Standard Gymnasium wrappers run in-process, requiring all dependencies (like ALE or chess engines) to be installed in your training environment. OpenEnv wraps these benchmarks in Docker containers, exposing them via WebSocket APIs. This isolation prevents dependency conflicts and allows environments written in different languages or requiring specific system libraries to run uniformly.

Can I use OpenEnv without Docker?

Yes. While Docker is the recommended deployment method for isolation, you can point the client to any running FastAPI server using the base_url parameter. Simply run the server locally (e.g., python -m atari_env.server.app) and instantiate the client with AtariEnv(base_url="http://localhost:8000") to bypass containerization during development.

What protocols does OpenEnv use for client-server communication?

OpenEnv uses WebSocket connections for low-latency, persistent communication between client and server, with HTTP endpoints available for /reset, /step, and /state operations. The client implementation in src/openenv/core/env_client.py handles connection lifecycle management, automatic reconnection, and serialization of dataclass payloads.

How do I handle custom action spaces in my wrapper?

Define your action space as a Pydantic dataclass in envs/<your_env>/models.py (e.g., ChessAction, AtariAction). Your EnvClient subclass must implement _step_payload to serialize this dataclass into JSON for transmission, and the server must deserialize it in the corresponding endpoint. The type system ensures only valid actions reach the environment logic.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}