rlightning.env.env_group

Environment group module for managing multiple environments.

This module provides the EnvGroup class for coordinating multiple environments in distributed RL training, with support for synchronous and asynchronous stepping, auto-reset, and progress tracking.

class rlightning.env.env_group.EnvGroup(env_cfg_list: ~typing.List[~rlightning.utils.config.config.EnvConfig], preprocess_fn: ~typing.Callable | ~typing.List[~typing.Callable] | None = <function default_env_preprocess_fn>)

Bases: object

Group manager for multiple reinforcement learning environments.

This class coordinates multiple environments for distributed RL training, supporting both synchronous and asynchronous stepping, auto-reset with step counting, and progress tracking.

env_list

List of local environment instances.

env_servers

List of remote environment server instances.

env_ids

List of unique environment identifiers.

id_to_env

Mapping from env_id to environment instance.

env_to_id

Mapping from environment instance to env_id.

step_counter

Counter for tracking steps in auto-reset context.

Type:

rlightning.env.env_group.EnvStepCounter

action_spaces: List[MockModule('gymnasium.Spaces')]

apply_evaluate_cfg() → None

Apply evaluation-time config overrides for all environments.

auto_reset(max_episode_steps: int | None = None) → Generator[None, None, None]

Context manager for enforcing a maximum number of steps per environment. When used, environments that reach max episode steps during step or step_async will be automatically reset. Progress tracking and step counting are handled internally.

Parameters:

max_episode_steps (Optional[int]) – Maximum number of steps per environment episode. If None, uses each environment’s configured max_episode_steps.

Examples

>>> with env_group.auto_reset(max_episode_steps=100):
>>>     for _ in range(1000):
>>>         batched_policy_resp = policy_group.rollout_batch(batched_env_ret)
>>>         batched_env_ret = env_group.step(batched_policy_resp)
>>>         # do other work

classmethod build_env(env_cfg: EnvConfig, preprocess_fn: Callable | None = <function default_env_preprocess_fn>, worker_index: int | None = None) → Tuple[MockModule('ray.actor.ActorHandle') | BaseEnv, bool]

Create an environment instance with given config and preprocess function.

The returns are the environment instance and a boolean indicating whether the environment is a remote environment server.

Parameters:

• cls – Environment group class.

• env_cfg (EnvConfig) – Environment configuration.

• preprocess_fn (Optional[Callable]) – Preprocess function for observations.

• worker_index (Optional[int]) – Env worker index for placement grouping.

Returns:

Environment instance. bool: Whether the environment is a remote environment server.

Return type:

Union[ray.actor.ActorHandle, BaseEnv]

close() → None

Close all environments.

collect_async(timeout: float | None = None, wait_all: bool = False) → Tuple[BatchedData, List[bool]]

Collects the results of previously submitted asynchronous step operations.

This method blocks until at least one environment’s computation is complete and its result (env_ret) is available. If timeout is assigned, it will wait and try to get more until the specified timeout is reached. Returns the batched results and truncation flags for the environments that have completed their steps.

Parameters:

• timeout (Optional[int]) – Maximum time to wait for more results in seconds.

• wait_all (bool) – If True, waits for all submitted operations to complete before returning. Default is False.

Returns:

Batched env_ret from the environments that have completed their step operations. List[bool]: List of truncation flags indicating whether each environment was truncated.

Return type:

BatchedData

Example

>>> env_group.step_async(batched_policy_resp)
>>> # do other work
>>> batched_env_ret, truncations = env_group.collect_async()

get_action_spaces() → List[MockModule('gymnasium.Space')]

Retrieve a list of action spaces

Returns:

List of action spaces for each environment.

get_env_stats(reset: bool = False) → Dict[str, float]

Retrieve aggregated episode-level environment statistics.

Collects per-env stats from all environments and computes the mean for each metric key.

Parameters:

reset – If True, clear recorded info in each env after retrieval.

Returns:

Dict mapping metric name to its mean value.

get_observation_spaces() → List[MockModule('gymnasium.Space')]

Retrieve a list of environment observation spaces

Returns:

List of observation spaces for each environment.

get_stats() → Dict[str, float]

Retrieve throughput statistics

init() → List[EnvMeta]

Actually init for all environments. It returns the metadata of all environments.

Returns:

List of metadata for each environment.

Return type:

List[EnvMeta]

observation_spaces: List[MockModule('gymnasium.Spaces')]

offload()

Offload the environment group to free GPU memory.

print_timing_summary(reset: bool = False) → None

Print the timing summary of the environment group.

reload()

Reload the environment group to load GPU memory.

reset(*args: Any, **kwargs: Any) → Tuple[BatchedData, List[bool]]

Reset all environments.

Parameters:

• *args – Arguments to pass to the reset function of each environment.

• **kwargs – Keyword arguments to pass to the reset function of each environment.

Returns:

The batched env_ret from the environments. List[bool]: List of truncation flags indicating whether each environment was truncated.

Return type:

BatchedData

restore_evaluate_cfg() → None

Restore environment members changed by apply_evaluate_cfg.

size() → int

Get the number of environments in the group.

Returns:

The number of environments.

Return type:

int

step(batched_policy_resp: BatchedData) → Tuple[BatchedData, List[bool]]

Call step function for environments provided in batched_policy_resp synchronously. If this calling is wrappered by max_rollout_context context manager, it will also handle auto reset and progress updates.

Parameters:

batched_policy_resp (BatchedData) – The batched policy_resp from the policies.

Returns:

The batched env_ret from the environments. List[bool]: List of truncation flags indicating whether each environment was truncated.

Return type:

BatchedData

step_async(batched_policy_resp: BatchedData) → None

Asynchronously submits step operations for environments specified in batched_policy_resp. It is non-blocking and doesn’t return the results immediately. The results are stored as futures and can be retrieved later using collect_async.

Parameters:

batched_policy_resp (BatchedData) – Batched policy responses for the environments.

Examples

>>> env_group.step_async(batched_policy_resp)
>>> # do other work
>>> batched_env_ret, truncations = env_group.collect_async()

step_counter: EnvStepCounter

throughput = <rlightning.env.env_group.ThroughputTracker object>

class rlightning.env.env_group.EnvStepCounter(env_ids: List[str], max_steps_list: List[int])

Bases: object

Tracks the number of steps taken in each environment.

get_max_steps(env_id: str) → int

Get the maximum number of steps for the environment.

Parameters:

env_id (str) – The ID of the environment.

Returns:

The maximum number of steps for the environment.

Return type:

int

get_steps(env_id: str) → int

Get the number of steps taken in the environment.

Parameters:

env_id (str) – The ID of the environment.

Returns:

The number of steps taken in the environment.

Return type:

int

get_truncations(env_ids: List[str]) → List[bool]

Get whether the environment has reached its maximum steps.

Parameters:

env_ids (List[str]) – List of environment IDs to check.

Returns:

List of booleans indicating whether each environment has reached its

maximum steps.

Return type:

List[bool]

is_reached_max_steps(env_id: str) → bool

Check if the environment has reached its maximum steps. This method is used

Parameters:

env_id (str) – The ID of the environment.

Returns:

True if the environment has reached its maximum steps, False otherwise.

Return type:

bool

reset(env_id: str) → None

Reset the step count for the specified environment or all environments.

Parameters:

env_id (int, optional) – The ID of the environment to reset. If None, all environments are reset.

Returns:

None

reset_all() → None

Reset the step count for all environments.

step(env_id: str) → None

Increment the step count for the provided environment.

Parameters:

env_id (str) – The ID of the environment to step.

step_all() → None

Increment the step count for all environments.

class rlightning.env.env_group.ThroughputTracker(log_interval_s: float = 5.0)

Bases: object

Track and log throughput metrics over a rolling time window. Default time window is 1 second.

get_stats()

record(name: str, count: int, step: int | None = None) → None

record_async(name: str, count: int, step: int | None = None) → None