rlightning.weights.weight_buffer_mixin

class rlightning.weights.weight_buffer_mixin.WeightBufferMixin

Bases: object

Mixin class for weight buffer management functionality. This mixin provides methods for initializing, updating, and managing weight buffers for policy weight transfer between train and eval policies.

clear_memory(sync=False)

Clear the memory.

get_weight_buffer()

Get the weight buffer instance.

init_weight_buffer(shared_weight_buffer=None)

Initialize the weight buffer.

offload_model()

offload model to cpu

offload_model_param_and_grad(offload_grad=False)

offload model param and grad to cpu

offload_optimizer()

offload optimizer to cpu

recv_weights(metadata_by_dtype)

Receive the weights from the train policy. Only call this method from train policy, in weight_transfer_manager.send_weights()

recv_weights_ipc(data)

Receive the weights from the eval policy using IPC.

recv_weights_local(state_dict)

This method is used to receive weights locally without using ray in running on a single process.

reload_model()

load model to device

reload_model_param_and_grad(load_grad=False)

load model param and grad to device

reload_optimizer()

load optimizer to device

send_weights(eval_policys, shared_weight_buffer=None)

Send the weights to the eval policy.

send_weights_ipc()

Send the weights to the eval policy using IPC.

sync_weights_layer_by_layers(module_name, name, dtype, shape, is_last: bool = False)

Sync the weight from the train policy layer by layer.

update_weights()

Update the weights from weight_buffer. This method runs in a daemon thread and continuously waits for update signals.

update_weights_from_buffer()

Update weights from the weight buffer.