`nemo_rl.environments.utils`#

Module Contents#

Classes#

EnvRegistryEntry

Functions#

`chunk_list_to_workers`	Chunk a list into a list of lists, where each sublist is assigned to a worker. Keeps ordering of elements.
`create_env`
`register_env`

Data#

ENV_REGISTRY

API#

class nemo_rl.environments.utils.EnvRegistryEntry#

Bases: typing.TypedDict

actor_class_fqn: str#: None

default_processor: NotRequired[str]#: None

nemo_rl.environments.utils.ENV_REGISTRY: Dict[str, nemo_rl.environments.utils.EnvRegistryEntry]#: None

nemo_rl.environments.utils.chunk_list_to_workers( to_chunk: list[Any], num_workers: int, ) → list[list[Any]]#

Chunk a list into a list of lists, where each sublist is assigned to a worker. Keeps ordering of elements.

If the list is not divisible by the number of workers, the last worker may have fewer elements. If there are more workers than elements, the first len(list) workers will have a single element each, and the remaining workers will have empty lists.

Parameters:

list – The list to be chunked.
num_workers – The number of workers to distribute the list to.

Returns:

A list of lists, where each sublist contains elements assigned to a worker.

Examples:

>>> from nemo_rl.environments.utils import chunk_list_to_workers
>>> chunk_list_to_workers([1, 2, 3, 4, 5], 3)
[[1, 2], [3, 4], [5]]

nemo_rl.environments.utils.create_env( env_name: str, env_configs: dict, ) → nemo_rl.environments.interfaces.EnvironmentInterface#

nemo_rl.environments.utils.register_env(env_name: str, actor_class_fqn: str) → None#

nemo_rl.environments.utils#