OakEmu-RL v0.1.0

Appearance

DQNModel ​

python
DQNModel(observation_shape: tuple[int, ...], action_shape: int, model_config: DQNModelConfiguration | None = None) -> None

DQN model with convolutional and fully connected layers.

Initialize the DQNModel.

Parameters:

  • observation_shape (tuple[int, ...]) – The shape of the environment observations.
  • action_shape (int) – The number of possible actions.
  • model_config (DQNModelConfiguration | None) – Optional configuration for the model architecture.

Source

Methods ​

forward ​

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

Perform a forward pass through the model.

Parameters:

  • x (Tensor) – The input tensor in the source observation dtype.

Returns:

  • Tensor – The output Q-values (non-categorical) or probabilities (categorical). For categorical, returns shape (batch, actions, atoms) with probabilities. For non-categorical, returns shape (batch, actions) with Q-values.

Source

get_q_values ​

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

Compute Q-values for the given input states.

For categorical DQN, this computes the expected Q-value from the probability distribution: Q(s, a) = sum_i(p_i * z_i) where z_i is the support atom value and p_i is the probability of that atom.

For non-categorical DQN, this is equivalent to forward().

Parameters:

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

Returns:

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

Source

reset_noise ​

python
reset_noise() -> None

Reset noise for all noisy layers in the model. Uses batched noise generation to reduce kernel launch overhead.

Source

set_advantage_mask ​

python
set_advantage_mask(action_masker: ActionMasker) -> None

Set a mask for the duelling advantage mean computation.

When action masking is used, only a subset of actions receive meaningful gradient updates. The advantage mean should be computed over only those actions so the centring property holds for the trained action subspace.

Parameters:

  • action_masker (ActionMasker) – The ActionMasker whose mask will be used to compute the centred advantage.

Source

Values ​

action_shape ​

python
action_shape = action_shape

Source

advantage_head ​

python
advantage_head = nn.Sequential(*adv_layers)

Source

categorical ​

python
categorical = model_config.categorical is not None

Source

conv ​

python
conv = build_conv_stack(input_channels=channels, conv_filters=model_config.conv_filters, activation=conv_activation)

Source

dueling ​

python
dueling = model_config.dueling

Source

fc ​

python
fc = nn.Sequential(*fc_layers)

Source

noisy_layers ​

python
noisy_layers: list[NoisyLinear] = [m for m in self.modules() if isinstance(m, NoisyLinear)]

Source

num_atoms ​

python
num_atoms = model_config.categorical.num_atoms

Source

v_max ​

python
v_max = model_config.categorical.v_max

Source

v_min ​

python
v_min = model_config.categorical.v_min

Source

value_head ​

python
value_head = nn.Sequential(*val_layers)

Source