NeMo RL

How a rollout runsCopy Anchor Link

NeMo RL trains on environments through 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:

InstallCopy Anchor Link

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.

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 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:

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.

Train on GridSimCopy Anchor Link

GridSim is the bundle the TRL page 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 bundleCopy Anchor Link

Download the bundle, then serve it on the machine NeMo RL runs on:

hf download earakely-scale/GridSimEnv-AgentEnv --repo-type dataset --include "bundles/gridsim-v1/*" --local-dir .

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 serverCopy Anchor Link

From the nemo-rl folder, write the server into NeMo RL's Gym checkout:

python -m agent_env.training.nemo_gym http://127.0.0.1:8000 3rdparty/Gym-workspace/Gym

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:

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

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 SFTCopy Anchor Link

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, 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 GRPOCopy Anchor Link

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:

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. Save it in the nemo-rl folder and run:

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.

========================= 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 learnsCopy Anchor Link

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 aloneCopy Anchor Link

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:

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

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 getsCopy Anchor Link

Each line of data/<split>.jsonl is a row for one task of the split:

{
  "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 doesCopy Anchor Link

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 rolloutsCopy Anchor Link

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:

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 sizingCopy Anchor Link

Set these on the resources server in its config, configs/<name>.yaml:

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:

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 supportedCopy Anchor Link

LimitsCopy Anchor Link

• 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.