rlightning.env.base_env

Base environment module for reinforcement learning.

This module defines the abstract base class for all RL environments, providing the common interface for environment interactions.

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