rlightning.engine

Engine module for reinforcement learning training loops.

This module provides various RL engine implementations including synchronous, asynchronous, and specialized engines for different training paradigms.

Available engines:

• SyncRLEngine: Synchronous reinforcement learning engine.

• AsyncRLEngine: Asynchronous reinforcement learning engine.

• EvaluationEngine: Engine for policy evaluation.

• RSLRLEngine: Engine for RSL-RL based training.

class rlightning.engine.AsyncRLEngine(config: MainConfig, env_group: EnvGroup, policy_group: PolicyGroup, buffer: DataBuffer)

Bases: BaseEngine

Asynchronous reinforcement learning engine.

This engine implements asynchronous training where: - Rollout thread: Collects experience from environments. - Training thread: Updates policy using collected experience. - Weight update thread: Broadcasts updated weights to eval policies.

evaluate() → None

Evaluate is not implemented for async engine.

rollout() → None

Perform rollout to collect experience from environments.

Runs the rollout loop, collecting experience by stepping through environments and storing transitions in the buffer. Runs until the done_flag is set.

run() → None

Launch asynchronous training threads.

Starts three worker threads (rollout, train, weight update) and waits for them to complete. Prints timing summary when finished.

train() → None

Perform training on collected experience.

Runs the training loop for the configured number of epochs, sampling from the buffer and updating the policy. Sets the new_weights_ready flag after each training step and done_flag when training completes.

update_dataset() → None

Update dataset from buffer to train policy.

warm_up()

init and dummy run for constructing RL dataflow Performs a dummy rollout and training iteration to ensure proper dataflow construction.

class rlightning.engine.AsyncRSLRLEngine(config: MainConfig, env_group: EnvGroup, policy_group: PolicyGroup, buffer: DataBuffer)

Bases: AsyncRLEngine

Async RSL-RL engine.

rollout(*args, **kwargs)

Perform rollout to collect experience from environments.

Runs the rollout loop, collecting experience by stepping through environments and storing transitions in the buffer. Runs until the done_flag is set.

sync_weights() → None

sync weights from train policy to eval policy.

Parameters:

policy_group – Policy group containing train and eval policies.

train() → None

Perform training on collected experience.

Runs the training loop for the configured number of epochs, sampling from the buffer and updating the policy. Sets the new_weights_ready flag after each training step and done_flag when training completes.

update_dataset(*args, **kwargs)

Update dataset from buffer to train policy.

warm_up()

init and dummy run for constructing RL dataflow Performs a dummy rollout and training iteration to ensure proper dataflow construction.

class rlightning.engine.BaseEngine(config: MainConfig, env_group: EnvGroup | None = None, policy_group: PolicyGroup | None = None, buffer: DataBuffer | None = None)

Bases: ABC

Abstract base class for reinforcement learning engines.

This class defines the common interface for all RL engines, including methods for warm-up, training, rollout, weight updates, and timing.

abstractmethod evaluate(*args: Any, **kwargs: Any) → None

User defined evaluation to collect experience.

Parameters:

• *args – Variable positional arguments.

• **kwargs – Variable keyword arguments.

iter_epochs(num_epochs: int) → iter

An iterator that yields epoch numbers up to num_epochs.

Parameters:

num_epochs (int) – The number of epochs to iterate through.

Returns:

An iterator over epoch numbers.

Return type:

iter

print_timing_summary(reset: bool = False) → None

Print timing summary for profiling.

Parameters:

reset – If True, reset timing statistics after printing.

abstractmethod rollout(*args: Any, **kwargs: Any) → None

User defined rollout to collect experience.

Parameters:

• *args – Variable positional arguments.

• **kwargs – Variable keyword arguments.

abstractmethod run() → None

Run the RL training flow.

This method should execute the complete RL training flow, coordinating rollout, training, and weight updates.

sync_weights() → None

sync weights from train policy to eval policy.

Parameters:

policy_group – Policy group containing train and eval policies.

abstractmethod train() → None

User defined training on collected experience.

Parameters:

• *args – Variable positional arguments.

• **kwargs – Variable keyword arguments.

abstractmethod update_dataset() → None

User defined dataset update from buffer to train policy.

abstractmethod warm_up() → None

Initialize and perform dummy run for constructing RL dataflow.

This method should initialize all components (environment, policy, buffer) and perform a dummy training iteration to ensure the dataflow is properly constructed.

class rlightning.engine.EvaluationEngine(*args: Any, **kwargs: Any)

Bases: BaseEngine

Evaluation engine for policy testing.

This engine runs policy evaluation without training, useful for testing trained policies on various environments.

get_metrics(policy_resp: PolicyResponse) → None

Extract and process metrics from policy response.

Parameters:

policy_resp – Policy response containing evaluation info.

rollout() → None

Rollout one complete episode and collect metrics.

run = MockModule('torch.inference_mode')

sync_weights() → None

No-op weight update method for evaluation engine.

train() → None

No-op training method for evaluation engine.

warm_up() → None

Initialize environments and evaluation policy.

class rlightning.engine.RSLRLEngine(config, env_group=None, policy_group=None, buffer=None)

Bases: SyncRLEngine

RSL-RL training engine.

This engine extends SyncRLEngine to support RSL-RL style policies, with modified rollout and training loops for on-policy learning.

rollout(obj_set: str, prefix: str = '', is_eval: bool = False) → None

Perform rollout to collect experience from environments.

Collects experience by stepping through environments, applying policy post-processing before storing transitions in the buffer.

Parameters:

• obj_set – Object set identifier for environment reset options.

• prefix – Prefix for logging metrics.

• is_eval – If True, skips buffer storage (evaluation only).

train() → None

Perform training on collected experience.

Samples data from the buffer, updates the dataset, and trains the policy. Designed for on-policy training with batch_size=-1.

Raises:

ValueError – If buffer is empty when training.

update_dataset() → None

Update the dataset in the policy group from the buffer.

warm_up()

Initialize runtime state without running training.

class rlightning.engine.SyncRLEngine(config: MainConfig, env_group: EnvGroup = None, policy_group: PolicyGroup | None = None, buffer: DataBuffer | None = None)

Bases: BaseEngine

Synchronous reinforcement learning engine.

This engine implements synchronous training where rollout and training happen sequentially. Supports both on-policy and off-policy algorithms.

evaluate(obj_set: str, prefix: str = '') → None

Perform evaluation to collect experience from environments.

Uses eval_env_list inside env_group if available, otherwise falls back to the training env_list.

Parameters:

• obj_set – Object set identifier for environment reset options.

• prefix – Prefix for logging metrics.

rollout(obj_set: str, prefix: str = '') → None

Perform rollout to collect experience from environments.

Parameters:

• obj_set – Object set identifier for environment reset options.

• prefix – Prefix for logging metrics.

run() → None

Run the main training loop.

Executes the training loop for the configured number of epochs, performing rollout, training, and periodic evaluation.

train() → None

Perform training on collected experience.

Samples data from the buffer, updates the dataset, and trains the policy. Logs training info metrics.

Raises:

ValueError – If buffer size is smaller than batch size when sampling without replacement.

update_dataset() → None

Update the dataset in the policy group from the buffer.

warm_up()

init and dummy run for constructing RL dataflow

Submodules

• rlightning.engine.async_rl_engine
• rlightning.engine.async_rsl_rl_engine
• rlightning.engine.base_engine
• rlightning.engine.evaluation_engine
• rlightning.engine.rsl_rl_engine
• rlightning.engine.sync_rl_engine