Implementation:OpenRLHF OpenRLHF NemoGym AgentExecutor

Knowledge Sources	OpenRLHF
Domains	Reinforcement_Learning, Agent_Execution, Inference
Last Updated	2026-02-07 10:40 GMT

Overview

Concrete tool for executing NeMo Gym agent rollouts using a vLLM inference backend with OpenAI-compatible API endpoints.

Description

The AgentExecutor class integrates vLLM inference with NVIDIA NeMo Gym's rollout collection system for reinforcement learning training. It extends AgentExecutorBase and sets up a local FastAPI server wrapping vLLM with OpenAI-compatible endpoints (/v1/chat/completions, /tokenize, /health), then initializes NeMo Gym services to collect rollouts from a math-solving agent. The companion prepare_data function formats math problems into the NeMo Gym data item structure.

Usage

Use this executor when training language models with NeMo Gym's reinforcement learning pipeline where the policy model is served via vLLM. It is designed for math problem-solving tasks that require agent-environment interaction through NeMo Gym's rollout collection infrastructure.

Code Reference

Source Location

Repository: OpenRLHF
File: examples/python/agent_func_nemogym_executor.py
Lines: 1-427

Signature

def prepare_data(
    question: str,
    expected_answer: str,
    sampling_params: SamplingParams,
) -> Dict[str, Any]:
    """Prepare data item for NeMo Gym rollout collection."""

class AgentExecutor(AgentExecutorBase):
    """Agent executor integrating vLLM inference with NeMo Gym rollout collection."""

    def _wait_for_vllm_server(self) -> bool: ...
    def _start_vllm_server(self, llm_engine, hf_tokenizer) -> None: ...
    def _start_nemogym_services(self, hf_tokenizer) -> None: ...

    async def execute(
        self,
        prompt: str,
        label: str,
        sampling_params: SamplingParams,
        max_length: int,
        llm_engine,
        hf_tokenizer,
    ) -> dict: ...

    def shutdown(self) -> None: ...

Import

from examples.python.agent_func_nemogym_executor import AgentExecutor, prepare_data

I/O Contract

Inputs

Name	Type	Required	Description
prompt	str	Yes	The math problem question text
label	str	Yes	Expected correct answer for reward evaluation
sampling_params	SamplingParams	Yes	vLLM sampling parameters for generation
max_length	int	Yes	Maximum sequence length for generation
llm_engine	LLMEngine	Yes	vLLM engine instance (used for first-time server setup)
hf_tokenizer	PreTrainedTokenizer	Yes	HuggingFace tokenizer for the policy model

Outputs

Name	Type	Description
execute() returns	dict	Contains prompt, label, observation_tokens (List[int]), reward (float), scores (float), extra_logs (dict), action_ranges (List[Tuple]), rollout_log_probs (List[float])

Usage Examples

Basic Rollout Execution

from examples.python.agent_func_nemogym_executor import AgentExecutor
from vllm import SamplingParams

# Create executor
executor = AgentExecutor()

# Configure sampling parameters
sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.9,
    max_tokens=1024,
    logprobs=5,
)

# Execute a rollout (called inside the PPO experience maker)
result = await executor.execute(
    prompt="What is 2 + 3?",
    label="5",
    sampling_params=sampling_params,
    max_length=2048,
    llm_engine=vllm_engine,
    hf_tokenizer=tokenizer,
)

# Result contains tokens, rewards, and log probabilities for RL training
print(result["reward"])            # Reward score
print(result["observation_tokens"])  # Generated token IDs
print(result["rollout_log_probs"])   # Log probabilities per token

# Cleanup
executor.shutdown()

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment