OakEmu-RL v0.1.0

Appearance

ThreadSafeModel ​

python
ThreadSafeModel(learner_model: torch.nn.Module, device: torch.device)

A thread-safe container around any nn.Module that manages locking and synchronization. Uses explicit method forwarding instead of getattr for better clarity and IDE support.

Initialize the ThreadSafeModel container.

Parameters:

  • learner_model (Module) – The underlying model.
  • device (device) – The device the model is on.

Source

Properties ​

model ​

python
model: torch.nn.Module

Get the underlying model.

Returns:

  • Module – The wrapped nn.Module.

Source

Methods ​

get_lock ​

python
get_lock() -> threading.Lock

Get the lock for manual synchronization. Use with caution.

Returns:

  • Lock – The threading.Lock instance.

Source

get_q_values ​

python
get_q_values(x: torch.Tensor) -> torch.Tensor

Compute Q-values for the given input states.

Delegates to the underlying model's get_q_values method. Only works if the underlying model has a get_q_values() method.

Parameters:

  • x (Tensor) – The input tensor of shape (batch, *observation_shape).

Returns:

  • Tensor – Q-values of shape (batch, actions).

Raises:

  • AttributeError – If the underlying model does not have a get_q_values() method.

Source

get_state_dict ​

python
get_state_dict(reset_event: bool = False, to_cpu: bool = True) -> dict[str, torch.Tensor]

Retrieve the model's state dictionary safely.

Parameters:

  • reset_event (bool) – Whether to reset the updated event after retrieving the state.
  • to_cpu (bool) – Whether to move the state dictionary tensors to the CPU.

Returns:

  • dict[str, Tensor] – The model's state dictionary.

Source

is_updated_event_set ​

python
is_updated_event_set() -> bool

Check if the model has been updated since the last event reset.

Returns:

  • bool – True if the updated event is set, False otherwise.

Source

parameters ​

python
parameters() -> Iterator[torch.nn.Parameter]

Get the parameters of the underlying model.

Returns:

  • Iterator[Parameter] – Iterator over model parameters.

Source

perform_update ​

python
perform_update(optimizer: torch.optim.Optimizer, loss: torch.Tensor) -> None

Perform an optimization step safely under the lock and set the updated event.

Parameters:

  • optimizer (Optimizer) – The optimizer to use.
  • loss (Tensor) – The loss to backpropagate.

Source

reset_noise ​

python
reset_noise() -> None

Reset noise for all noisy layers in the model.

Thread-safe: acquires the lock to prevent concurrent reads during noise reset. Only works if the underlying model has a reset_noise() method.

Source

set_updated_event ​

python
set_updated_event() -> None

Set the updated event to signal that the model has been updated.

Source

sync_to ​

python
sync_to(target_model: torch.nn.Module) -> None

Synchronise the target model with the learner model's current weights safely.

Parameters:

  • target_model (Module) – The model to synchronise.

Source