Generate Nemo Gym Env

generate-nemo-gym-env

Build the NeMo Gym variant of an env. NeMo Gym is NVIDIA's RL gym layer, optimized for Ray-based orchestration and post-episode grading. The Python package is nemo_gym (installed via pip install git+https://github.com/NVIDIA-NeMo/Gym).

Concept

NeMo Gym is NVIDIA's RL gym layer for LLM agents. It's built on Ray and ships a FastAPI-based SimpleResourcesServer that exposes one POST /<tool> endpoint per tool, plus the standard /seed_session (cookie-based session bootstrap) and /verify (post-episode grader). Targets docs.nvidia.com/nemo/gym/latest.

When the user has a shared domain module (<domain>.py) and wants a NeMo Gym variant, wrap it. Don't duplicate logic.

Archetypes

Recommended file layout

The user picks the actual paths. The canonical shape:

<env_dir>/nemo_gym/
├── pyproject.toml         # nemo_gym (git+) + e2b-* + fastapi + uvicorn + requests
├── __init__.py
├── Dockerfile             # Ray-aware multi-stage
├── configs/<env>.yaml     # NeMo Gym config (entrypoint, domain, description)
├── server.py              # SimpleResourcesServer subclass with tool endpoints
├── rollout.py             # raw requests + cookie session
└── README.md

Note: NeMo Gym requires Python 3.12+.

Implementation order

1. Server class — server.py

from nemo_gym.base_resources_server import (
    BaseResourcesServerConfig,
    BaseSeedSessionRequest, BaseSeedSessionResponse,
    BaseVerifyRequest, BaseVerifyResponse,
    SimpleResourcesServer,
)
from nemo_gym.server_utils import SESSION_ID_KEY
from fastapi import FastAPI, Request
from pydantic import BaseModel, Field
from typing import Any, Dict

class MyConfig(BaseResourcesServerConfig):
    pass

class GuessReq(BaseModel):
    word: str

class ToolResponse(BaseModel):
    output: str

class MyVerifyRequest(BaseVerifyRequest):
    ground_truth: list = []

class MyResourcesServer(SimpleResourcesServer):
    config: MyConfig
    sessions: Dict[str, Dict[str, Any]] = Field(default_factory=dict)

    def setup_webserver(self) -> FastAPI:
        app = super().setup_webserver()
        app.post("/guess")(self.guess)
        return app

    async def seed_session(self, body: BaseSeedSessionRequest) -> BaseSeedSessionResponse:
        return BaseSeedSessionResponse()

    def _sess(self, request: Request) -> Dict[str, Any]:
        sid = request.session[SESSION_ID_KEY]
        if sid not in self.sessions:
            self.sessions[sid] = {"game": WordleGame(), "step": 0}
        return self.sessions[sid]

    async def guess(self, body: GuessReq, request: Request) -> ToolResponse:
        sess = self._sess(request)
        feedback = sess["game"].guess(body.word)
        sess["step"] += 1
        return ToolResponse(output=feedback)

    async def verify(self, body: MyVerifyRequest) -> BaseVerifyResponse:
        # Compute reward from the response trajectory + ground truth
        expected = ""
        if body.ground_truth and isinstance(body.ground_truth, list):
            expected = body.ground_truth[0].get("expected_output", "")
        reward = 0.0
        for item in body.response.output:
            if hasattr(item, "type") and item.type == "function_call_output":
                if expected and expected in getattr(item, "output", ""):
                    reward = 1.0; break
        return BaseVerifyResponse(**body.model_dump(), reward=reward)

if __name__ == "__main__":
    MyResourcesServer.run_webserver()

Key contracts:

• One endpoint per tool. Register them in setup_webserver(). Pydantic models on the request body become the JSON shape.
• Sessions live in `self.sessions` keyed by request.session[SESSION_ID_KEY]. Lazy-init on first call. NeMo Gym sets the session cookie on POST /seed_session.
• `verify()` is the grader. Read body.ground_truth (passed by the trainer) and body.response.output (the trajectory). Return BaseVerifyResponse(**body.model_dump(), reward=...).

2. NeMo Gym config — configs/<name>.yaml

my_env_resources_server:
  resources_servers:
    my_env:
      entrypoint: server.py
      domain: agent
      description: "What this env does"

This is the file the NeMo Gym CLI looks for when launching via ng_run "+config_paths=[configs/my_env.yaml]".

3. Rollout — rollout.py

NeMo Gym has no Python client SDK. The rollout speaks raw HTTP via requests with a Session for cookie persistence:

import requests
session = requests.Session()
session.post(f"{ENV_URL}/seed_session", json={}).raise_for_status()
r = session.post(f"{ENV_URL}/guess", json={"word": "crane"})
result = r.json()["output"]

Tool definitions for the LLM are hardcoded in rollout.py (no introspection endpoint). Mirror the request schemas from server.py exactly.

4. Dockerfile

Multi-stage build. NeMo Gym pulls Ray and a fairly heavy stack — the Docker image is ~1.5GB. The container exposes port 11000 by default. For HF Spaces deployment, override to port 7860 (one-port limit on Spaces).

Validation gates

1. Import — uv run python -c "import os; os.environ.setdefault('E2B_API_KEY','x'); from server import MyResourcesServer" succeeds. 2. Local server — try uv run python server.py. Note: NeMo Gym's run_webserver() initializes a Ray cluster, which fails on shared SLURM / HF cluster nodes (gcs_server can't bind). On those machines, only Docker / HF Space deploy works. 3. Endpoint smoke — when running, curl http://localhost:11000/seed_session -X POST returns 200 and sets a session cookie. 4. Rollout — MAX_TURNS=3 uv run python rollout.py drives end-to-end against the deployed Space.

Common gotchas

• `No module named 'anyio'` — nemo_gym doesn't pin its full transitive set on every install. Add anyio>=4.0, attrs>=23.0, fastapi>=0.115, uvicorn, requests to your dependencies explicitly.
• `Address already in use` or `gcs_server` crash — Ray init failed. Almost always a shared cluster issue. Document this and tell the user to deploy via Space.
• Cookie not set on the rollout — make sure to use requests.Session(), not raw requests.post(). The session cookie is the SID handle.
• `/verify` returns reward 0 unexpectedly — ground_truth is wrapped in a list. Check body.ground_truth[0].get("expected_output") not body.ground_truth.get(...).
• Hardcoded tool schemas drift — when you change a server endpoint's Pydantic body, manually update the matching tool definition in rollout.py. There's no list_tools().

Reference

• references/architecture.md — Ray orchestration, dataset format with responses_create_params, deployment notes

Official documentation

• NVIDIA-NeMo/Gym — source repo
• NeMo Gym latest docs
• Creating a Resource Server tutorial
• Creating a Training Environment tutorial

NeMo Gym architecture (deep)

What NeMo Gym is

NVIDIA's RL gym layer for LLM agents. Built on Ray for orchestration. The Python package is nemo_gym (install via pip install git+https://github.com/NVIDIA-NeMo/Gym). It targets NVIDIA's NeMo training stack but works with TRL/GRPO via raw HTTP.

Wire protocol

The session cookie (set on /seed_session, named via SESSION_ID_KEY) is the only way to associate subsequent tool calls with state. There is no SDK client — rollouts speak raw HTTP.

SimpleResourcesServer lifecycle

1. `MyResourcesServer.run_webserver()` at __main__ — boots Ray, starts FastAPI, registers tools. 2. `setup_webserver(self)` — must call super().setup_webserver() first to get the FastAPI instance with /seed_session and /verify already registered, then app.post("/tool")(self.tool) for each tool. 3. `seed_session(self, body)` — called once per session by POST /seed_session. Return BaseSeedSessionResponse(). Lazy-init resources here or in the first tool call (recommended). 4. Tool methods — async, take (self, body, request), return a Pydantic response. Read session id via request.session[SESSION_ID_KEY]. 5. `verify(self, body)` — async, called once per episode by the trainer. Return BaseVerifyResponse(**body.model_dump(), reward=...). Always spread the body back — drops fields silently if you don't.

Session state pattern

sessions: Dict[str, Dict[str, Any]] = Field(default_factory=dict)

def _sess(self, request: Request) -> Dict[str, Any]:
    sid = request.session[SESSION_ID_KEY]
    if sid not in self.sessions:
        self.sessions[sid] = {"game": MyGame(), "step": 0}
    return self.sessions[sid]

Session entries persist for the server's lifetime by default. For long-running deployments, either prune on verify or add a TTL.

Dataset format

NeMo Gym expects datasets with two key fields:

• `responses_create_params` — a JSON-stringified OpenAI-Responses-API config (model, tools, system prompt). The trainer feeds this to the model.
• `ground_truth` — a list of dicts (typically one) carrying expected outputs / answer keys. verify() reads from this.

Example row:

{
    "responses_create_params": json.dumps({
        "model": "gpt-4o-mini",
        "input": [{"role": "user", "content": "Solve 2+2"}],
        "tools": [{"type": "function", "function": {"name": "guess", ...}}],
    }),
    "ground_truth": [{"expected_output": "4"}],
    "metadata": json.dumps({"task_id": "math-001"}),
}

Reward computation in /verify

body.response.output is a list of items emitted by the model:

Typical patterns:

Substring match — pass if expected appears anywhere:

expected = body.ground_truth[0].get("expected_output", "")
reward = 0.0
for item in body.response.output:
    if hasattr(item, "type") and item.type == "function_call_output":
        if expected.strip() in str(getattr(item, "output", "")).strip():
            reward = 1.0; break

Function-call match — pass if a specific tool was called with success:

for item in body.response.output:
    if getattr(item, "type", "") == "function_call" and item.name == "terminate":
        args = item.arguments or ""
        if "success" in args:
            reward = 1.0; break

Production deployment

Dockerfile is multi-stage; the runtime image is ~1.5GB because of Ray. Healthcheck via /seed_session. For HF Spaces:

• Port 7860 (one-port limit on Spaces)
• Set app_port: 7860 in README frontmatter
• HF Spaces handle Ray init reliably (unlike shared SLURM nodes)

Why run_webserver() fails on shared cluster nodes

NeMo Gym's run_webserver() calls ray.init(), which spawns a gcs_server process bound to specific ports. On shared SLURM / HF cluster nodes those ports are already taken, and the bind fails. The error looks like:

[gcs_server] Failed to bind on address ...

There's no fix from the env author's side — deploy via Docker / Space and connect over network.

What can go wrong

• Cookie isn't set on the client — use requests.Session(), not naked requests.post().
• `gcs_server` crash — Ray init failure on shared nodes; redirect to deployed Space.
• `No module named 'anyio'` / `attrs` — NeMo Gym's transitive deps drift. Pin them explicitly.
• Reward always 0 in verify — ground_truth is a list, not a dict. Use body.ground_truth[0].get(...).
• `request.session[SESSION_ID_KEY]` raises KeyError — /seed_session wasn't called first; fail fast with a clear error.