rlightning.policy

class rlightning.policy.BasePolicy(config: PolicyConfig, role_type: PolicyRole)

Bases: torch.nn.Module, RayActorMixin, WeightBufferMixin, ABC

Abstract base class for reinforcement learning policies.

This class provides the common interface for all RL policies, including methods for initialization, rollout, training, and weight management.

check_idle() → None

Check if the policy is idle.

Called by policy_group to determine if the actor is available for new requests.

Raises:

AssertionError – If called with async rollout mode.

abstractmethod construct_network(env_meta: Any, *args: Any, **kwargs: Any) → None

Construct the neural network architecture.

Parameters:

• env_meta – Environment metadata for network configuration.

• *args – Variable positional arguments.

• **kwargs – Variable keyword arguments.

get_num_requests() → int

Get current number of in-flight requests.

Returns:

Number of in-flight requests for async rollout routing.

get_param_fingerprint() → Dict[str, float]

Return per-parameter L2 norms as a lightweight weight fingerprint.

Used by PolicyGroup.verify_eval_weight_consistency() to detect weight-sync failures without shipping full tensors across processes.

Returns:

Dict mapping "<module_name>.<param_name>" to the float32 L2 norm of each parameter tensor.

abstractmethod get_trainable_parameters() → Dict[str, Dict[str, MockModule('torch.Tensor')]]

Return a dict of module state dicts.

init_eval(eval_config: Any | None = None, env_meta: Any | None = None) → None

Initialize the policy for evaluation.

Parameters:

• eval_config – Evaluation configuration.

• env_meta – Environment metadata.

init_train(train_config: TrainConfig, env_meta: Any | None = None) → None

Initialize the policy for training.

Sets up the network, finds trainable models, wraps with DDP if needed, and initializes the optimizer.

Parameters:

• train_config – Training configuration.

• env_meta – Environment metadata.

is_initialized() → bool

Return True if the policy has been initialized.

abstractmethod load_state_dict(state_dict: Dict[str, MockModule('torch.Tensor')], *args: Any, **kwargs: Any) → None

Load trainable parameters from state_dict.

notify_update_weights() → None

Signal that new weights are available for update.

postprocess = MockModule('torch.inference_mode')

print_timing_summary(reset: bool = False) → None

Print timing summary for profiling.

Parameters:

reset – If True, reset timing statistics after printing.

reset_training_state(train_config: TrainConfig, env_meta: Any | None = None, seed: int | None = None) → None

Reset model + optimizer to initial state after warm_up.

rollout_step = MockModule('torch.inference_mode')

save_checkpoint(path: str) → None

Save checkpoint of the policy. User can override this method in subclass if needed.

Parameters:

path (str) – The path to save the checkpoint.

abstractmethod setup_optimizer(optim_cfg: Any) → None

Set up the optimizer for training.

Parameters:

optim_cfg – Optimizer configuration.

abstractmethod train(*args: Any, **kwargs: Any) → Any

Run a training step for the policy.

abstractmethod update_dataset(data: MockModule('tensordict.TensorDict')) → None

Update internal dataset from sampled buffer data.

update_weights() → None

Continuously update weights from weight_buffer.

Runs in a daemon thread, waiting for update signals and applying new weights when available.

wait_for_weight_update_done(timeout: float | None = None) → bool

Block until the background weight-update thread finishes its latest cycle.

Parameters:

timeout – Maximum seconds to wait. None means wait indefinitely.

Returns:

True if the update completed within timeout, False if it timed out. Always returns True when there is no background update thread (weight_buffer_strategy == "None").

class rlightning.policy.PolicyGroup(train_policy_list: List[MockModule('ray.actor.ActorHandle')], eval_policy_list: List[MockModule('ray.actor.ActorHandle')], config: PolicyConfig)

Bases: object

Manager for collections of train and eval policies.

Coordinates multiple policies for distributed training and evaluation, handling weight distribution, communication groups, and rollout batching.

convert_eval_to_train(num: int) → None

Convert evaluation policies to training policies.

Parameters:

num – Number of policies to convert.

convert_train_to_eval(num: int) → None

Convert training policies to evaluation policies.

Parameters:

num – Number of policies to convert.

property eval_list: List[MockModule('ray.actor.ActorHandle')]

Return the list of policies for evaluation

Returns:

List of evaluation policy instances.

init(train_config: TrainConfig, env_meta: EnvMeta | None = None, eval_config: Any | None = None) → None

Initialize both training and evaluation policies.

init_comm_group(backend: str = 'nccl', is_colocated: bool = False) → None

Initialize communication groups for distributed training.

Sets up distributed environment, intra-node groups, weight transfer groups, and DDP communication groups.

Parameters:

• backend – Communication backend (‘nccl’ or ‘gloo’).

• is_colocated – If True, initialize only for training policies.

init_eval(eval_config: Any | None = None, env_meta: EnvMeta | None = None) → List

Initialize all evaluation policies.

Parameters:

• eval_config – Evaluation configuration.

• env_meta – Environment metadata.

Returns:

List of initialization results.

init_placement_info(is_colocated: bool = False) → None

Initialize the placement info of the policy group with the consideration of given policy config.

placement_info is a dict that mapping from ray_node_id to a dict of training/evaluation policy (actor).

init_train(train_config: Any, env_meta: EnvMeta | None = None) → List

Initialize all training policies.

Parameters:

• train_config – Training configuration.

• env_meta – Environment metadata.

Returns:

List of initialization results.

init_weight_buffer() → None

Initialize the weight buffer based on the configured strategy.

The weights buffer is created by the buffer strategy: None: no weight buffer Double: each eval policy has its own weight buffer Shared: shared weight buffer between eval policies on the same node Sharded: not supported now

Raises:

ValueError – If the weight buffer strategy is unsupported.

notify_update_weights() → None

Signal all eval policies that new weights are available.

offload_eval_model()

Offload the eval model to cpu.

offload_model_optimizer()

Offload the model optimizer to cpu.

offload_model_param_and_grad(offload_grad: bool = True, offload_optimizer: bool = True)

Offload the model parameters and gradients to cpu.

pop(role: PolicyRole)

Pop the last policy worker from the list.

postprocess(batched_env_ret: BatchedData | None = None, batched_policy_resp: BatchedData | None = None) → BatchedData

Submit postprocess tasks to eval policies.

Parameters:

• batched_env_ret – Batched environment returns.

• batched_policy_resp – Batched policy responses.

Returns:

BatchedData containing processed results.

print_timing_summary(reset: bool = False) → None

Print the timing summary of the policy group.

push(policy_worker: MockModule('ray.actor.ActorHandle'), role: PolicyRole) → None

Add new policy worker to existing group.

Parameters:

• policy_worker (BasePolicy) – Worker instance

• role (PolicyRole) – Policy role, indicating training or evaluation.

reload_eval_model()

Reload the eval model to gpu.

reload_model_param_and_grad(load_grad: bool = True, load_optimizer: bool = True)

Reload the model parameters and gradients to gpu.

reset_training_state(train_config: TrainConfig, env_meta: Any | None = None, seed: int | None = None) → None

Reset training policies after warm_up.

rollout_batch(batched_env_ret: BatchedData) → BatchedData

Perform batched rollout across eval policies.

Parameters:

batched_env_ret – Batched environment returns.

Returns:

BatchedData containing policy responses.

save_checkpoint(path: str) → None

Save the checkpoint of train policy.

If multiple train policies exist, only save the checkpoint of the first one.

Parameters:

path (str) – The path to save the checkpoint.

send_weights() → None

Broadcast weights from train to eval policies.

Send weights from train to eval policies based on the configured weight buffer strategy.

shutdown() → None

Shutdown the policy group and cleanup resources.

sync_weights()

Broadcast the state dict to the eval policy and store into the weight_buffer for later loading.

train(*args: Any, **kwargs: Any) → Dict[str, float]

Train all training policies.

Parameters:

• *args – Variable positional arguments.

• **kwargs – Variable keyword arguments.

Returns:

Training info from the first policy.

property train_list: List[MockModule('ray.actor.ActorHandle')]

Return the list of policies for training

Returns:

List of training policy instances.

update_dataset(sample_data: List) → List[MockModule('ray.ObjectRef')]

Update the dataset in the policy group by getting sampled data from the buffer.

Parameters:

sample_data – List of sampled data objects, one per train policy.

verify_eval_weight_consistency(rtol: float = 0.001) → bool

Wait for eval weight updates and verify they match the first train policy.

For each eval policy, per-parameter L2 norms are collected and compared against those of the first train policy. A relative tolerance rtol is applied: a mismatch is reported when

|train_norm - eval_norm| / max(|train_norm|, 1e-8) > rtol

Parameters:

rtol – Relative tolerance for per-parameter norm comparison.

Returns:

True if every eval policy is consistent with the train policy, False if any mismatch is detected.

wait_for_eval_weight_update(timeout: float = 60.0) → None

Block until every eval policy finishes its background weight update.

No-op when weight_buffer_strategy == "None" because weights are transferred synchronously and no background thread is involved.

Parameters:

timeout – Per-policy wait timeout in seconds.

Subpackages

• rlightning.policy.rsl_rl_policy
  • rlightning.policy.rsl_rl_policy.ppo_vtrace
  • rlightning.policy.rsl_rl_policy.rsl_rl_policy
• rlightning.policy.simple_ppo_policy
  • rlightning.policy.simple_ppo_policy.ppo_policy
• rlightning.policy.utils namespace
  • rlightning.policy.utils.losses
  • rlightning.policy.utils.router
  • rlightning.policy.utils.utils
  • rlightning.policy.utils.vtrace
• rlightning.policy.vla_policy
  • rlightning.policy.vla_policy.ppo_policy

Submodules

• rlightning.policy.base_policy
• rlightning.policy.policy_group
• rlightning.policy.supervised_policy