rlightning.policy.base_policy

Base policy module.

This module provides the abstract base class for all policies, defining the interface for training, evaluation, and rollout operations.

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

dataset: Any | None

device: str

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_init: bool

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.

model_list: List[Tuple[str, MockModule('torch.nn.Module')]]

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_mode: str

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.

timing_raw: Dict[str, Dict[str, Any]]

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").

world_size: int

class rlightning.policy.base_policy.PolicyRole(value)

Bases: StrEnum

Policy role enumeration.

EVAL = 'eval'

TRAIN = 'train'

rlightning.policy.base_policy.clone_checkpoint_value(value: Any) → Any

Recursively clone checkpoint payloads into CPU-owned tensors.

rlightning.policy.base_policy.infer_train_dp_world_size() → int

Infer TRAIN_DATA_PARALLEL world size, falling back safely to 1.

rlightning.policy.base_policy.set_nested_attr(obj: Any, path: str, value: Any) → None

Set a dotted-path attribute on an object.