`nemo_rl.environments.nemo_gym`#

Module Contents#

Classes#

`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#

`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.
`setup_nemo_gym_config`

Data#

`DEFAULT_INVALID_TOOL_CALL_PATTERNS`
`DEFAULT_THINKING_TAGS`

API#

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#: [’’, ‘’]

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

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

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

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

async run_rollouts( nemo_gym_examples: list[dict], tokenizer: transformers.PreTrainedTokenizerBase, timer_prefix: str, ) → list[dict]#

_postprocess_nemo_gym_to_nemo_rl_result( nemo_gym_result: dict, tokenizer: transformers.PreTrainedTokenizerBase, ) → dict#

shutdown() → None#

abstractmethod step(message_log_batch, metadata)#

abstractmethod global_post_process_and_metrics(batch)#

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

nemo_rl.environments.nemo_gym#