# NeMo RL (https://www.agentenvframework.com/docs/training/nemo-rl) > Train using NVIDIA's NeMo RL ## How a rollout runs 1. NeMo RL hands a row of the split’s JSONL file to NeMo Gym, whose `simple_agent` seeds a session on the resources server with it. 2. The server starts the row’s task as an episode: it deploys the env, loads the task, and opens the agent’s turn, while your model writes its first reply. 3. Your model plays the turn: `simple_agent` posts each tool call to the tool’s route on the server, which runs it on the env through the episode. 4. Your model ends its turn with `end_turn`, or with a reply that calls no tool, which `simple_agent` posts to `/verify` to end it. The task’s verifier scores the episode. 5. `/verify` answers with the score as the reward, and masks a rollout the environment failed; the episode is torn down and the session closed. Parts of the scene: - **NeMo RL and NeMo Gym**: NeMo RL’s GRPO hands each row to NeMo Gym’s `simple_agent`, which plays it with your model, served from NeMo RL’s own vLLM. NeMo RL owns the model, the generation and the update; the environment owns the reward. - **A dataset row**: A line of `data/.jsonl`: the task’s prompt and tools as `responses_create_params`, and the `split`, `index` and `task_id` the server plays. The prompt is yours to change; the task stays the same. - **The resources server**: `agent_env.training.nemo_gym`’s server, which Gym runs from `resources_servers//` in a venv of its own. `/seed_session` starts a row’s task, a route per tool calls it, and `/verify` ends the turn and answers with the reward. - **The session**: One Gym session per rollout, by cookie, holding one session on the bundle’s server, so one episode. Gym runs a step’s rollouts all at once, so the bundle’s server needs room for every one. - **The OpenEnv server**: `agent_env.openenv.build_app(bundle)` serves every task of a bundle as an OpenEnv env, the split `all` plus one per eval. Serve it yourself and name its URL, or let the resources server serve the bundle. - **The env**: The task’s env, deployed fresh for this episode in a sandbox, here a container on your Docker. Its tools are what the model calls. - **The agent’s turn**: The task prompts one agent, and your model plays it from outside: the episode deploys a proxy in the agent’s place and hands the turn’s prompt and tools to the server. - **The verifier**: The task’s own verifier step scores the episode once the turn ends, the same score an evaluation run records. With several scores, `reward=` picks the one to train on. - **The reward**: `/verify` returns the score as `reward`. A rollout the environment failed gets `mask_sample` and a `failure_reason` instead, and NeMo RL leaves it out of the loss. GRPO compares the reward with the other rollouts of the same row. NeMo RL trains on environments through [NeMo Gym](https://docs.nvidia.com/nemo/gym/), which runs each environment as a resources server: Gym's agent seeds a session on it with a dataset row, calls the row's tools, and asks it to verify the rollout for a reward. Four pieces connect NeMo RL to a bundle: | Piece | What it does | | --------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------- | | `agent_env.openenv.build_app(bundle)` | Serves every task of a bundle as an [OpenEnv](https://www.agentenvframework.com/docs/openenv.md) env: the split `all`, plus one split per eval. | | `python -m agent_env.training.nemo_gym` | Writes the bundle's resources server into a Gym root: the server, its Gym config, and a JSONL file of rows per split. | | `AgentEnvResourcesServer` | The resources server. A row's seed starts its task as an episode, each tool has a route, and `/verify` answers with the task's score. | | NeMo Gym's `simple_agent` | Plays each row with your model, served from NeMo RL's vLLM, and hands the reward back to NeMo RL. | ## Install NeMo RL runs on Linux with NVIDIA GPUs. Clone its main branch with its submodules: NeMo Gym is one of them, in `3rdparty/Gym-workspace/Gym`. ```bash title="Terminal" git clone --recursive https://github.com/NVIDIA-NeMo/RL.git nemo-rl cd nemo-rl uv venv ``` The resources server needs NeMo Gym 0.7 or newer, which NeMo RL's main branch has; its v0.7.0 release container pairs with NeMo Gym 0.4. Set up CUDA 13 and cuDNN as [NeMo RL's README](https://github.com/NVIDIA-NeMo/RL#installation-prerequisites-and-quick-start) describes. The page's run used main at `dd2e251`. Then, in an environment of its own on the machine where the envs will deploy, with Docker running, install the framework to serve the bundle and write the server: ```bash title="Terminal" pip install 'agentenv-framework[nemo-gym]' ``` On Python 3.11 to 3.13.13 the extra brings OpenEnv only, which is all serving and writing take. When NeMo Gym starts the server, it builds the server's own venv, with the same framework version and the checkout's NeMo Gym in it. > [!WARNING] > **On L40S and other Ada GPUs** > > NeMo RL's Dockerfile builds for Hopper and Blackwell. The page's run built NeMo RL on the host for > four L40S, with these changes to its `pyproject.toml` and then `uv lock`: > > * Drop `deep_ep` and `deep_gemm` from the `vllm` and `automodel` extras: DeepEP doesn't build for Ada. > * Change `nemo-automodel[moe]` to `nemo-automodel`, since its `moe` extra brings DeepEP back. > * Raise flash-attn's `MAX_JOBS` from 4 to 16 to build faster, where memory allows: the run's machine had 384 GB. > > Build with `TORCH_CUDA_ARCH_LIST=8.9` and `NVTE_CUDA_ARCHS=89` set. The first `uv run` builds > NeMo RL's environment, which took about 40 minutes. ## Train on GridSim GridSim is the bundle the [TRL page](https://www.agentenvframework.com/docs/training/trl.md#train-on-gridsim) trains on. Each task runs one grid battery for one day in the ERCOT day-ahead market: the model reads the day's 24 hourly prices and submits a schedule, one integer an hour, to charge, sell or idle. The env grades the schedule against the day's exact optimum: 0 for nothing useful, 0.2 for idling all day, 1 for the optimum. The bundle has 2,000 train days and 200 eval days. ### Serve the bundle Download the bundle, then serve it on the machine NeMo RL runs on: ```bash title="Terminal" hf download earakely-scale/GridSimEnv-AgentEnv --repo-type dataset --include "bundles/gridsim-v1/*" --local-dir . ``` ```python title="serve.py" import uvicorn from agent_env.openenv import build_app app = build_app("bundles/gridsim-v1", max_concurrent_envs=64) uvicorn.run(app, host="127.0.0.1", port=8000) ``` Gym plays all of a step's rollouts at once, so `max_concurrent_envs` should cover a step: 32 here. Past the limit, rollouts wait for a free episode, as the 200 validation rollouts do here. ### Write the resources server From the `nemo-rl` folder, write the server into NeMo RL's Gym checkout: ```bash title="Terminal" python -m agent_env.training.nemo_gym http://127.0.0.1:8000 3rdparty/Gym-workspace/Gym ``` ```text title="Output" wrote 3rdparty/Gym-workspace/Gym/resources_servers/gridsim_v1 ``` The command plays one task of each toolset to list the tools its rows offer, so it deploys an env once. It writes: ```text title="resources_servers/gridsim_v1" app.py runs AgentEnvResourcesServer pyproject.toml the server's venv: the framework, and NeMo Gym from the checkout configs/gridsim_v1.yaml the server, simple_agent, and the environment server NeMo RL calls data/all.jsonl 2,200 rows data/train.jsonl 2,000 rows data/eval.jsonl 200 rows ``` ```yaml title="configs/gridsim_v1.yaml" gridsim_v1: resources_servers: gridsim_v1: entrypoint: app.py domain: agent description: gridsim-v1's tasks, served by AgentEnv url: http://127.0.0.1:8000 gridsim_v1_simple_agent: responses_api_agents: simple_agent: entrypoint: app.py resources_server: type: resources_servers name: gridsim_v1 model_server: type: responses_api_models name: policy_model gridsim_v1_environment_server: environment_servers: legacy_agent: entrypoint: app.py agent_server: type: responses_api_agents name: gridsim_v1_simple_agent ``` To skip `serve.py`, pass the bundle's folder instead of a URL. The resources server then serves the bundle itself, at `--max-concurrent-envs` episodes at once. ### Warm start with SFT Played as it comes, Qwen3-4B idles: it scores 0.18 on the 200 eval days. Warm-start it with the TRL page's [`sft.py`](https://www.agentenvframework.com/docs/training/trl.md#warm-start-with-sft), in an environment with TRL, run from the `nemo-rl` folder so it saves `qwen3-4b-gridsim-sft` where the config below looks. It took 25 minutes on one L40S. ### Train with GRPO NeMo RL takes a YAML config. This run's starts from NeMo RL's `examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml` and uses the TRL run's settings: the same 32 rollouts a step, learning rate, schedule and LoRA, on the DTensor backend instead of Megatron. The parts that matter: ```yaml title="grpo_gridsim.yaml (excerpt)" grpo: num_prompts_per_step: 4 # 32 rollouts a step: 4 days, 8 rollouts each num_generations_per_prompt: 8 max_num_steps: 40 val_period: 20 val_at_start: true val_num_generations_per_prompt: 1 loss_fn: reference_policy_kl_penalty: 0 ratio_clip_max: 0.28 policy: model_name: qwen3-4b-gridsim-sft tokenizer: chat_template_kwargs: {enable_thinking: false} max_total_sequence_length: 3072 automodel_cfg: enabled: true lora_cfg: {enabled: true, match_all_linear: true, dim: 32, alpha: 64} megatron_cfg: enabled: false optimizer: name: torch.optim.AdamW kwargs: {lr: 2.0e-05, weight_decay: 0.0} generation: max_new_tokens: 1024 temperature: 1.1 lora_refit_mode: merged vllm_cfg: gpu_memory_utilization: 0.5 http_server_serving_chat_kwargs: {enable_auto_tools: true, tool_parser: hermes} data: train: data_path: 3rdparty/Gym-workspace/Gym/resources_servers/gridsim_v1/data/train.jsonl validation: data_path: 3rdparty/Gym-workspace/Gym/resources_servers/gridsim_v1/data/eval.jsonl env: should_use_nemo_gym: true nemo_gym: config_paths: - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml - resources_servers/gridsim_v1/configs/gridsim_v1.yaml gridsim_v1_simple_agent: responses_api_agents: simple_agent: {max_steps: 4} cluster: gpus_per_node: 4 ``` The whole file, as it ran with its paths made relative to the checkout, is [`grpo_gridsim.yaml`](/training/grpo_gridsim.yaml). Save it in the `nemo-rl` folder and run: ```bash title="Terminal" uv run python examples/nemo_gym/run_grpo_nemo_gym.py --config grpo_gridsim.yaml ``` NeMo RL starts the Gym servers itself: the resources server, in the venv Gym builds for it on the first start, `simple_agent`, and the environment server NeMo RL sends rows to. It serves your model to them from its vLLM as `policy_model`. The `hermes` tool parser turns Qwen3's tool calls into the calls `simple_agent` makes. ```text title="Output (step 12 of the GridSim run, trimmed)" ========================= Step 12/40 ========================= 📊 Training Results: • Loss: 0.0075 • Generation KL Error: 0.0044 • Avg Reward: 0.7447 • Mean Generation Length: 783.1562 ⏱️ Timing: • Total step time: 60.86s • generation: 29.38s (48.3%) • policy_training: 14.51s (23.8%) ``` `Avg Reward` is the mean of the step's 32 rewards. NeMo RL reports the eval mean as `Accuracy`. ### What it learns | Qwen3-4B | eval reward (200 days) | |---|---| | after SFT (GRPO step 0) | 0.546 | | GRPO step 20 | 0.679 | | GRPO step 40 | 0.672 | Mean train reward of each GRPO step's 32 rollouts, by step: 1: 0.701, 2: 0.722, 3: 0.437, 4: 0.421, 5: 0.563, 6: 0.805, 7: 0.660, 8: 0.547, 9: 0.801, 10: 0.713, 11: 0.753, 12: 0.745, 13: 0.833, 14: 0.831, 15: 0.801, 16: 0.773, 17: 0.894, 18: 0.620, 19: 0.903, 20: 0.835, 21: 0.847, 22: 0.475, 23: 0.793, 24: 0.707, 25: 0.782, 26: 0.856, 27: 0.927, 28: 0.512, 29: 0.375, 30: 0.866, 31: 0.793, 32: 0.828, 33: 0.953, 34: 0.782, 35: 0.848, 36: 0.661, 37: 0.657, 38: 0.616, 39: 0.834, 40: 0.677. | Qwen3-4B, eval reward on the 200 eval days | NeMo RL | TRL | | ------------------------------------------ | ------- | ---- | | After SFT (GRPO step 0) | 0.55 | 0.57 | | After 20 GRPO steps | 0.68 | 0.56 | | After 40 GRPO steps | 0.67 | 0.61 | Through NeMo RL, the first 20 GRPO steps add 0.13, more than one rollout a day can tell apart (about ±0.06), and the next 20 keep it. GridSim's greedy rule scores 0.76 on these days. Both runs train the same SFT model on the same rows, with the same LoRA, learning rate, schedule and rollouts a step. One difference is known: TRL's `dr_grpo` loss divides each step's sum by `max_completion_length`, 1,024 tokens, while NeMo RL's token-level loss divides it by the tokens the rollouts generated, which are fewer, so NeMo RL's updates were larger. This is one run of each, so read the gap as a pointer, not a ranking of the trainers. The GRPO run took 53 minutes on four L40S, an AWS g6e.12xlarge that also served the envs: about 61 seconds a step, half of it generation, and two to four minutes for each of the three evals. ## Collect rollouts with NeMo Gym alone NeMo Gym also plays the rows without NeMo RL, against any OpenAI-compatible endpoint with tool calling, to evaluate a model or collect rollouts; for Qwen3 on vLLM, serve it with `--enable-auto-tool-choice --tool-call-parser hermes`. Write the server into a root of your own, which Gym reads through `NEMO_GYM_EXTRA_ROOTS`, here serving the bundle itself. NeMo Gym needs Python 3.13.14 or newer: ```bash title="Terminal 1" pip install 'agentenv-framework[nemo-gym]' python -m agent_env.training.nemo_gym bundles/gridsim-v1 gym-root --max-concurrent-envs 16 export NEMO_GYM_EXTRA_ROOTS=$PWD/gym-root gym env start --config gym-root/resources_servers/gridsim_v1/configs/gridsim_v1.yaml \ --model-type vllm_model --model Qwen/Qwen3-4B --model-url http://127.0.0.1:8001/v1 --model-api-key dummy ``` ```bash title="Terminal 2, with the same NEMO_GYM_EXTRA_ROOTS" gym eval run --no-serve --agent gridsim_v1_simple_agent \ --input gym-root/resources_servers/gridsim_v1/data/eval.jsonl --output rollouts.jsonl --concurrency 16 ``` Each line of `rollouts.jsonl` is a rollout with its `reward`, `mask_sample` and `failure_reason`. Played this way, Gym's rewards match what the Episode API scores for the same schedules. ## What Gym gets Each line of `data/.jsonl` is a row for one task of the split: ```json title="data/train.jsonl, line 1 (prompts and tools cut)" { "responses_create_params": { "input": [ {"role": "system", "content": "You operate a grid battery in the ERCOT day-ahead market (DAM) for one operating day…"}, {"role": "user", "content": "{\"task_id\": \"dam-v1-2021-05-06-HB_HUBAVG-1h-r\", \"tier\": \"foresight\", \"market\": \"ERCOT DAM\"…"} ], "tools": [ {"type": "function", "name": "submit_schedule", "description": "…", "parameters": {…}, "strict": false}, {"type": "function", "name": "end_turn", "description": "…", "parameters": {…}, "strict": false} ] }, "agent_ref": {"type": "responses_api_agents", "name": "gridsim_v1_simple_agent"}, "split": "train", "index": 0, "task_id": "@local/…/gridsim-v1/dam-v1-2021-05-06-HB_HUBAVG-1h-r" } ``` The prompt is yours to change; `split`, `index` and `task_id` name the task the server plays. A row's tools are every tool its task's turns offer, as the env describes them, plus `end_turn`. NeMo RL keeps the whole row with the rollout. ## What the server does | Route | What it does | | --------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `POST /seed_session` | Starts the row's task as an episode. It answers before the env is deployed, so the deploy overlaps the model's first reply; the first tool call waits for it. | | `POST /` | Calls the tool in the episode, and answers with its text or error as the model sees it. `end_turn` answers with the next turn's prompt. | | `POST /verify` | Ends the turns the model left open, its last reply ending the first, and answers with the episode's score as `reward`. | | `POST /close_session` | Ends the session's episode, unscored. | Gym calls a tool at the route of its name, so writing the server refuses a tool named after one of Gym's own routes, or whose name isn't one path segment, with the fix in the message. For a task with several turns, tell the model in the system prompt to end each turn with `end_turn`. ## Masked rollouts A rollout the environment failed gets `mask_sample` and a `failure_reason` from `/verify`, with a reward of 0, and NeMo RL leaves it out of the loss. That covers a deploy that failed, a call that missed its deadline, a server that went away, and a rollout whose row shows other tools than its task offers. Each one is logged: ```text title="Log" WARNING agent_env.training.nemo_gym.server: rollout masked: submit_schedule didn't answer within 60s, and the episode ended, since what the call does may still be under way ``` The page's run masked none of its 1,880 rollouts. ## Deadlines and sizing Set these on the resources server in its config, `configs/.yaml`: | Setting | Default | What it bounds | | --------------------- | ------- | ------------------------------------------------------------------------------------------------------------------------------------------- | | `tool_timeout` | 60 s | One tool call. The server ends a call that misses it, and the episode with it, and masks the rollout. | | `reset_timeout` | 600 s | The deploy and the opening of the first turn, each `end_turn`, and a rollout's wait for a free episode. Past it, the rollout is masked. | | `idle_timeout` | 600 s | A session with no request. Gym never verifies a rollout whose run failed or was cancelled, so the server ends its episode to free the slot. | | `max_concurrent_envs` | 8 | With `bundle:`, the episodes the server it starts runs at once. | The resources server runs with one worker, since a rollout's session lives in its process; a config that sets `num_workers` above 1 is refused. When rollouts wait for a free episode, the server warns once: ```text title="Log" WARNING agent_env.training._openenv: the server at http://127.0.0.1:8000 is at its limit of 64 sessions at once, so rollouts wait their turn: serve the bundle with a higher MAX_CONCURRENT_ENVS (or max_concurrent_envs=) to run more at once ``` ## What's supported | | | | ----------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | NeMo RL | GRPO through NeMo Gym's `legacy_agent` environment server, trained on GPUs from main at `dd2e251`. | | NeMo Gym | 0.7 and later. A CI job plays rollouts through the server on the locked version and the latest, every week. | | Multi-turn tasks | `end_turn` hands the model each next prompt. | | Bundles with several toolsets | Each row lists its own task's tools. | | Sandboxes | `build_app(..., sandbox=)`, or `--sandbox` when the server serves the bundle, deploys every episode where `agent-env run --sandbox` would, from your Docker to [any provider](https://www.agentenvframework.com/docs/sandboxes.md). | ## Limits * **NeMo RL's main branch,** until a release container carries NeMo Gym 0.7. * **One worker** for the resources server, which holds every rollout's session. * **One deploy per rollout.** Every rollout deploys its envs fresh, and the deploy counts in its time.