rlightning.policy.policy_group

Policy group module for managing multiple policy workers.

This module provides the PolicyGroup class for coordinating multiple train and eval policies, including weight distribution, communication group initialization, and rollout management.

class rlightning.policy.policy_group.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.

rlightning.policy.policy_group.build_transfer_info_for_colocated(train_list: List[BasePolicy], eval_list: List[BasePolicy])

Build transfer info for colocation.

Parameters:

• train_list – List of train policies.

• eval_list – List of eval policies.

Returns:

Transfer info for colocation.

rlightning.policy.policy_group.build_transfer_info_for_disaggregated(placement_info: Dict[str, Dict[str, List]], weight_buffer_strategy: str) → tuple[Dict[str, List[str]], Dict[str, List[Dict]]]

Build weight transfer info for disaggregated..

Determines how weights should be transferred from train to eval policies based on their node placement and the weight buffer strategy.

Parameters:

• placement_info – Dictionary mapping node IDs to policy lists.

• weight_buffer_strategy – Strategy for weight buffer management.

Returns:

• router_map: Mapping from train nodes to receiver nodes.

• transfer_info: Detailed transfer information per node.

Return type:

A tuple containing

Raises:

ValueError – If no train policy is found or unsupported strategy.