rlightning.env

Environment module for reinforcement learning.

This module provides various environment implementations and utilities for managing vectorized and distributed environments in RL training.

Available components:

• BaseEnv: Abstract base class for all environments.

• EnvGroup: Group manager for multiple environments.

• EnvMeta: Metadata container for environment properties.

class rlightning.env.BaseEnv(config: EnvConfig, worker_index: int | None = 0, preprocess_fn: Callable | None = None)

Bases: ABC

Abstract base class for reinforcement learning environments.

This class defines the common interface for all RL environments, including methods for reset, step, and environment metadata.

config

Environment configuration object.

env_id

Unique identifier for this environment instance.

env

The underlying gymnasium environment.

num_envs

Number of parallel environments (1 by default).

max_episode_steps

Maximum steps per episode.

timing_raw

Dictionary for tracking timing statistics.

apply_evaluate_cfg() → None

Apply evaluation-time config overrides for this environment.

Default implementation is a no-op. Specific environments can override this to support temporary evaluate-only behaviors.

close() → None

Close the environment.

Override this method if special cleanup is needed. Default implementation does nothing.

collect_async() → List[EnvRet]

Asynchronous collect interface (only for RemoteEnvServer).

This interface also reserves for future integration with env that natively support asynchronous steps.

Raises:

NotImplementedError – Always, as this is not native supported and only for RemoteEnvServer from now.

finish_rollout() → None

Finish the rollout.

Override this method in subclasses to implement custom rollout finishing behavior.

get_action_space()

Retrieve action space. User can override this method as needed.

get_env_id() → str

Get the unique environment identifier.

Returns:

The unique identifier string for this environment.

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

Get episode-level environment statistics.

Computes sum and count for each recorded metric key so that the caller (e.g. EnvGroup) can aggregate across multiple environments.

Parameters:

reset – If True, clear recorded info after computing stats.

Returns:

Dict mapping metric name to [sum, count].

get_max_episode_steps() → int | None

Retrieve max episode steps

get_metadata() → EnvMeta

Get the environment meta information.

Returns:

the environment meta information

Return type:

EnvMeta

get_observation_space()

Retrieve observation space. User can override this method as needed.

init() → EnvMeta

Initialize the environment and return metadata.

Returns:

EnvMeta containing environment properties.

is_finish() → bool

Check if the environment should finish running.

Override this method in RemoteEnvClient subclasses to determine when to stop the environment loop.

Returns:

True if the environment should stop, False otherwise.

offload()

Offload the environment to free GPU memory.

print_timing_summary(reset: bool = False) → None

Print timing summary for profiling.

Parameters:

reset – If True, reset timing statistics after printing.

reload()

Reload the environment to load GPU memory.

abstractmethod reset(*args, **kwargs) → EnvRet | List[EnvRet] | List[MockModule('ray.ObjectRef')]

Reset the environment to initial state.

Parameters:

• *args – Variable positional arguments.

• **kwargs – Variable keyword arguments.

Returns:

Environment return containing observation and info.

restore_evaluate_cfg() → None

Restore environment members changed by apply_evaluate_cfg.

Default implementation is a no-op.

abstractmethod step(policy_resp: PolicyResponse) → EnvRet

Step the environment with the given action.

Parameters:

policy_resp – Policy response containing the action.

Returns:

Environment return containing observation, reward, done, and info.

step_async(policy_resp_list: List[PolicyResponse]) → None

Asynchronous step interface (only for RemoteEnvServer).

This interface also reserves for future integration with env that natively support asynchronous steps.

Parameters:

policy_resp_list – List of Policy response

Raises:

NotImplementedError – Always, as this is not native supported and only for RemoteEnvServer from now.

class rlightning.env.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.

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]

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

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

class rlightning.env.EnvMeta(env_id: str = None, action_space: MockModule('gymnasium.spaces.Space') | None = None, observation_space: MockModule('gymnasium.spaces.Space') | None = None, num_envs: int | None = None)

Bases: object

Environment metadata container.

Stores metadata about an environment including its spaces and configuration.

env_id

Unique identifier for the environment.

Type:

str

action_space

Gymnasium action space.

Type:

MockModule(‘gymnasium.spaces.Space’) | None

observation_space

Gymnasium observation space.

Type:

MockModule(‘gymnasium.spaces.Space’) | None

num_envs

Number of parallel environments.

Type:

int | None

action_space: MockModule('gymnasium.spaces.Space') | None = None

Action space of the environment.

env_id: str = None

The environment ID.

num_envs: int | None = None

The vectorized number of the environment.

observation_space: MockModule('gymnasium.spaces.Space') | None = None

Observation space of the environment.

Subpackages

• rlightning.env.remote_env
  • rlightning.env.remote_env.env_client
• rlightning.env.utils namespace
  • rlightning.env.utils.libero_utils
  • rlightning.env.utils.utils
  • rlightning.env.utils.venv

Submodules

• rlightning.env.ale_env
• rlightning.env.base_env
• rlightning.env.env_group
• rlightning.env.env_server
• rlightning.env.isaac_env
• rlightning.env.isaac_marl_env
• rlightning.env.libero_env
• rlightning.env.maniskill_env
• rlightning.env.mujoco_env
• rlightning.env.piper_env