`nemo_rl.environments.nemo_gym`#

Module Contents#

Classes#

`NemoGymCompatibleConfig`	Configuration fields required to select the NeMo Gym rollout path.
`NemoGymConfig`
`NemoGym`	This environment class isn’t really used for training. It’s really meant as an integration wrapper around NeMo-Gym that hooks into the existing NeMo RL resource management via ray. So there is still one source of truth for resource management in NeMo RL.

Functions#

`should_use_nemo_gym`	Determine whether NeMo Gym should handle rollouts and validation.
`_has_nan_generation_logprobs`	Return whether a postprocessed rollout contains NaN policy logprobs.
`_typed_gym_failure`	Map a NeMo-Gym HTTP failure onto a typed, PICKLABLE failure, or None if not one.
`get_nemo_gym_uv_cache_dir`	Return the uv cache directory inside a container, or None outside one.
`get_nemo_gym_venv_dir`	Return the NeMo Gym venv directory from NEMO_GYM_VENV_DIR, or None.
`_detect_invalid_tool_call_and_malformed_thinking`	Flag a NeMo-Gym output item as an invalid tool call / malformed thinking.
`_looks_like_image_src`	True when `src` plausibly points at an image the loader can open.
`get_pad_dynamic_image_shapes`	Return nemo_gym’s pad_dynamic_image_shapes from an env config, or False.
`_extract_input_images_from_message`	Pull PIL images out of a non-assistant Responses-API item.
`_is_trainable_output_item`	Report whether an output item becomes a trainable assistant turn.
`_index_per_turn_images`	Bin server-returned images by the trainable turn that saw them.
`_image_sources_equal`
`_without_initial_image_sources`	Copy Responses messages and remove one ordered copy of initial images.
`_attach_multimodal_data_to_user_message`	Attach per-turn multimodal tensors to `user_message`.
`extract_reward_components`	Return per-component rewards from a NeMo Gym verify result, or None.
`build_reward_component_columns`	Build `reward/<name>` batch columns from per-sample reward-component dicts.
`validate_reward_components_match_scalar`	Assert each multi-reward result sets `reward == sum(reward_components)`.
`setup_nemo_gym_config`
`spinup_nemo_gym_actor`	Spin up the NeMo-Gym actor against the given generation server URLs.

Data#

`_ROUTED_EXPERTS_DTYPES`
`DEFAULT_INVALID_TOOL_CALL_PATTERNS`
`DEFAULT_THINKING_TAGS`
`_IMAGE_SRC_PREFIXES`

API#

nemo_rl.environments.nemo_gym._ROUTED_EXPERTS_DTYPES#: None

nemo_rl.environments.nemo_gym.DEFAULT_INVALID_TOOL_CALL_PATTERNS#: [‘<tool_call>’, ‘</tool_call>’, ‘<function_call>’, ‘</function_call>’]

nemo_rl.environments.nemo_gym.DEFAULT_THINKING_TAGS#: [’’, ‘’]

class nemo_rl.environments.nemo_gym.NemoGymCompatibleConfig#

Bases: typing.Protocol

Configuration fields required to select the NeMo Gym rollout path.

property env: dict[str, Any]#

property policy: nemo_rl.models.policy.PolicyConfig#

nemo_rl.environments.nemo_gym.should_use_nemo_gym( master_config: nemo_rl.environments.nemo_gym.NemoGymCompatibleConfig, ) → bool#: Determine whether NeMo Gym should handle rollouts and validation.

nemo_rl.environments.nemo_gym._has_nan_generation_logprobs(result: dict) → bool#: Return whether a postprocessed rollout contains NaN policy logprobs.

nemo_rl.environments.nemo_gym._typed_gym_failure(error: Exception) → Optional[Exception]#

Map a NeMo-Gym HTTP failure onto a typed, PICKLABLE failure, or None if not one.

Classification has to happen here, on the raising side, because run_rollouts runs inside the NemoGym Ray actor and the exception must survive the actor boundary to reach the retry policy on the driver.

It does not survive. aiohttp’s raise_for_status passes headers=self.headers, and those are a CIMultiDictProxy, which cloudpickle cannot serialize – so Ray drops the cause and the driver receives a bare RayTaskError with no type and no .status. Every gym HTTP failure then classified DATA, capping the gym path at max_data_attempts_per_prompt (2) and leaving max_attempts_per_prompt (5) unreachable on the very path whose dead-endpoint scenario motivates it. Two things made that the dominant case rather than a corner: Gym’s middleware turns inner-server failures into 500 – exactly the status the INFRA branch is for – and its transport layer retries disconnects in an uncapped loop, so those never arrive at all.

GymTransportError and RolloutDataFailure take a single str, so they pickle cleanly and classify_rollout_failure’s explicit-class fast path wins on the far side.

Returns None when the exception carries no HTTP status, leaving the caller to re-raise it untouched.

nemo_rl.environments.nemo_gym.get_nemo_gym_uv_cache_dir() → str | None#

Return the uv cache directory inside a container, or None outside one.

Inside a container (NRL_CONTAINER=1), returns the uv cache location so Gym stores its caches in the expected shared path. Returns None outside a container, meaning the caller should omit this arg and let Gym create the cache locally (the default when you may not be able to write to /opt).

nemo_rl.environments.nemo_gym.get_nemo_gym_venv_dir() → str | None#

Return the NeMo Gym venv directory from NEMO_GYM_VENV_DIR, or None.

Returns the value of NEMO_GYM_VENV_DIR if set, otherwise None. When None the caller should omit this arg and let Gym create venvs locally (the default when a container is not used since you may not be able to write to /opt).

class nemo_rl.environments.nemo_gym.NemoGymConfig#

Bases: typing.TypedDict

model_name: str#: None

base_urls: List[str]#: None

initial_global_config_dict: Dict[str, Any]#: None

port_range_low: NotRequired[int]#: None

port_range_high: NotRequired[int]#: None

invalid_tool_call_patterns: NotRequired[List[str] | None]#: None

thinking_tags: NotRequired[List[str] | None]#: None

require_routed_experts: NotRequired[bool]#: None

routed_experts_dtype: NotRequired[str]#: None

use_fastokens: NotRequired[bool]#: None

tokenizer_config: NotRequired[Optional[nemo_rl.models.policy.TokenizerConfig]]#: None

pad_dynamic_image_shapes: NotRequired[bool]#: None

nemo_rl.environments.nemo_gym._detect_invalid_tool_call_and_malformed_thinking( output_item_dict: dict[str, Any], invalid_tool_call_patterns: list[str] | None = None, thinking_tags: list[str] | None = None, ) → tuple[bool, bool]#

Flag a NeMo-Gym output item as an invalid tool call / malformed thinking.

Inspects the final output item of a model turn. For a final content message, any thinking tag is malformed (thinking should never leak into the answer); for a reasoning summary, only a repeated tag (count > 1) is malformed (a single pair is expected). A textual tool-call pattern in either indicates an invalid (unexecuted) tool call.

Returns:: (is_invalid_tool_call, has_malformed_thinking).

nemo_rl.environments.nemo_gym._IMAGE_SRC_PREFIXES#: (‘data:image/’, ‘http://’, ‘https://’, ‘file://’)

nemo_rl.environments.nemo_gym._looks_like_image_src(src: str) → bool#

True when src plausibly points at an image the loader can open.

Guards against tool responses (e.g. {"x": 0.65, "y": 0.83} from a click tool) that are strings but not image URLs. Without this, the indexer forwards the JSON payload to resolve_to_image → PIL.open, which treats it as a filesystem path and raises FileNotFoundError.

nemo_rl.environments.nemo_gym.get_pad_dynamic_image_shapes( env_config: collections.abc.Mapping[str, Any], ) → bool#

Return nemo_gym’s pad_dynamic_image_shapes from an env config, or False.

Takes master_config.env rather than the whole config: interpreting NeMo-Gym settings belongs with the environment, and callers outside it only need the resolved boolean.

The NemoGym actor reads the same key from its own config for the per-turn attach. The initial-payload attach runs in the driver instead, so it has to be read here and passed down, or multi-image prompts would be processed under different rules on the two paths.

nemo_rl.environments.nemo_gym._extract_input_images_from_message( item: dict, ) → list[PIL.Image.Image]#

Pull PIL images out of a non-assistant Responses-API item.

Handles both content-list items (user / tool messages carrying input_image/image/image_url parts) and function_call_output items whose output field is an image data URL. Tool outputs that are non-image strings (e.g. structured JSON returned by tools like click(x, y)) contribute zero images to the bucket.

nemo_rl.environments.nemo_gym._is_trainable_output_item(item: dict) → bool#

Report whether an output item becomes a trainable assistant turn.

The postprocess loop skips items whose generation_token_ids is missing or empty, so per-turn image binning has to use the same predicate or the two walks disagree and every later turn gets the wrong images.

nemo_rl.environments.nemo_gym._index_per_turn_images( output: list[dict], input_messages: list[dict] | None = None, ) → list[list[PIL.Image.Image]]#

Bin server-returned images by the trainable turn that saw them.

Walks the Responses-API items in order and flushes pending into a per-turn bucket each time it hits an item carrying truthy generation_token_ids — matching the exact gate that _postprocess_nemo_gym_to_nemo_rl_result uses to decide which items become trainable turns. Every other item (user turns, tool messages, function_call_output, non-trainable reasoning) contributes its images to pending for the next trainable turn. This ensures the returned list has one entry per trainable turn, aligned with the postprocess loop’s turn_idx even when the trainable item’s role is not assistant (e.g. a reasoning-only response, or a function_call).

input_messages is the initial responses_create_params.input list — images there (e.g. a single-shot user prompt for tool-based envs like circle-click) are consumed by the first trainable turn’s tokenized prompt and must land in the first bucket. Agents like gym_v_agent that keep input empty and inject observations as function_call_output items are unaffected — the seed is a no-op when input_messages is empty.

nemo_rl.environments.nemo_gym._image_sources_equal(left: Any, right: Any) → bool#

nemo_rl.environments.nemo_gym._without_initial_image_sources( messages: Any, initial_sources: list[Any], ) → tuple[Any, bool]#: Copy Responses messages and remove one ordered copy of initial images.

nemo_rl.environments.nemo_gym._attach_multimodal_data_to_user_message( user_message: dict, *, images: list[PIL.Image.Image], processor: Any, pad_dynamic_image_shapes: bool = False, ) → None#

Attach per-turn multimodal tensors to user_message.

The processor is only invoked to extract multimodal tensors (pixel_values, imgs_sizes, num_patches, etc.); its text output is discarded — vLLM’s tokens remain the trajectory. We therefore feed it the minimal placeholder text it needs to count image regions: one processor.image_token per image. Passing the vLLM-decoded text does not work because that text already contains expanded <img>...<image>*N...</img> regions, and the processor would try to re-expand every embedded <image>.

class nemo_rl.environments.nemo_gym.NemoGym(cfg: nemo_rl.environments.nemo_gym.NemoGymConfig)#

Bases: nemo_rl.environments.interfaces.EnvironmentInterface

This environment class isn’t really used for training. It’s really meant as an integration wrapper around NeMo-Gym that hooks into the existing NeMo RL resource management via ray. So there is still one source of truth for resource management in NeMo RL.

Initialization

_require_spinup() → None#: Raise a diagnosable error if this instance never ran :meth:_spinup.

health_check() → None#

Raise if the Gym head server or any subprocess server has died.

Thin wrapper over NeMo-Gym’s own RunHelper.poll, which is what gym env start calls every 60s from run_forever. NeMo-RL only calls rh.start, so without this the check Gym already implements never runs and a dead tool server surfaces as unexplained rollout timeouts instead of a named process.

_spinup() → None#

Start the NeMo-Gym head server and rollout collection helper.

Deferred from init so the actor can be created cheaply (and scheduled onto reserved nodes) and spun up explicitly once the vLLM server URLs are available, overlapping with vLLM model loading.

set_tokenizer(tokenizer: transformers.PreTrainedTokenizerBase) → None#

Install the tokenizer run_rollouts postprocesses with.

Called once per actor at spinup. It used to be a run_rollouts argument, which meant Ray deserialized a tokenizer per prompt on this actor’s task-execution thread. That thread holds the GIL, so it blocked the actor’s event loop and no rollout could issue its first HTTP request until its own copy finished loading.

The cost is not marginal. A tokenizer of this shape measured 7.45 MB on the wire, 286 ms to serialize and 1052 ms to deserialize, so a SingleController recipe admitting ~1000 prompts per step spends roughly 18 minutes deserializing a single admission burst, against a step that should take minutes. Runs on that shape stalled without completing a rollout.

Measured on one CPU node with 1024 concurrent calls against one actor: 153 of 1024 prompts finished in 300 s passing the tokenizer per call, versus all 1024 in 16 s holding it here. Passing an ObjectRef instead does not help – Ray caches the object buffer, not the deserialized value, so it still pays per task.

async run_rollouts( nemo_gym_examples: list[dict], timer_prefix: str, deduplicate_multimodal_data: bool = False, ) → collections.abc.AsyncGenerator[tuple[int, dict, dict | None], None]#: Stream postprocessed rollouts as NeMo-Gym tasks complete.

_postprocess_nemo_gym_to_nemo_rl_result( nemo_gym_row: dict, nemo_gym_result: dict, tokenizer: transformers.PreTrainedTokenizerBase, *, include_initial_multimodal_data: bool = True, ) → dict#

shutdown() → None#

abstractmethod step(message_log_batch, metadata)#

abstractmethod global_post_process_and_metrics(batch)#

nemo_rl.environments.nemo_gym.extract_reward_components( nemo_gym_result: dict, ) → Dict[str, float] | None#

Return per-component rewards from a NeMo Gym verify result, or None.

Single-reward NeMo Gym environments return only a scalar reward. Multi-reward environments additionally return reward_components: a mapping of component-name -> score. These are surfaced as reward/<name> batch keys and consumed by GDPO (see nemo_rl.algorithms.advantage_estimator.GDPOAdvantageEstimator).

Returns None when the environment is single-reward (no reward_components), so callers fall back to the scalar reward path unchanged.

nemo_rl.environments.nemo_gym.build_reward_component_columns( component_dicts: List[Dict[str, float] | None], ) → Dict[str, torch.Tensor]#

Build reward/<name> batch columns from per-sample reward-component dicts.

Takes the union of component names across the batch in sorted (deterministic) order and, for each, emits a reward/<name> tensor with one entry per sample. A component absent on a given sample is filled with 0.0 so every column covers all samples (the per-prompt baseline requires each component present for all responses).

Keys are prefixed reward/ so they are exactly what nemo_rl.algorithms.utils.get_gdpo_reward_component_keys selects (it matches startswith("reward/") and sorts by name); the name carries the component identity, so no positional index is needed. Returns an empty dict when no sample has components.

nemo_rl.environments.nemo_gym.validate_reward_components_match_scalar( nemo_gym_results: List[dict], ) → None#

Assert each multi-reward result sets reward == sum(reward_components).

A multi-reward verifier must set the scalar reward to the sum of its reward_components so single-reward (GRPO) consumers and GDPO read the same aggregate. We keep the verifier’s scalar reward as total_reward rather than silently overwriting it with the component sum, so a verifier that violates this contract must be surfaced here instead of masked.

Raises ValueError on the first violating result. A no-op for single-reward results (those without reward_components).

nemo_rl.environments.nemo_gym.setup_nemo_gym_config(config, tokenizer) → None#

nemo_rl.environments.nemo_gym.spinup_nemo_gym_actor( env_configs: dict[str, Any], base_urls: list[str], model_name: str, *, tokenizer: transformers.PreTrainedTokenizerBase, enable_router_replay: bool, routed_experts_dtype: str, use_fastokens: bool, ) → Any#

Spin up the NeMo-Gym actor against the given generation server URLs.

When env_configs[“nemo_gym”][“num_gpu_nodes”] > 0, the actor is scheduled with soft NodeAffinity to the current Ray node so its colocated GPU resources land where the caller expects.

Parameters:

env_configs – The master_config.env mapping; env_configs[“nemo_gym”] supplies the Gym global config plus NeMo-RL detection knobs (invalid_tool_call_patterns, thinking_tags, num_gpu_nodes).
base_urls – Per-DP-rank OpenAI-compatible server base URLs from the generation backend.
model_name – Served model name the Gym rollouts should target.
tokenizer – Installed on the actor once, here, rather than passed per rollout call. See NemoGym.set_tokenizer for why that distinction is the difference between a working run and a stalled one.
enable_router_replay – Sets require_routed_experts on the NemoGymConfig.
routed_experts_dtype – Dtype name for R3 routed_experts tensors (“int8”/”int16”/”int32”), resolved by the caller from the model’s expert count.
use_fastokens – Forwarded from policy.tokenizer.use_fastokens so the rollout actor patches its tokenizer consistently with the driver.

Returns:

The spun-up NemoGym Ray actor handle (_spinup already awaited).

nemo_rl.environments.nemo_gym#