OakEmu-RL v0.1.0

Appearance

CategoricalProjection ​

python
CategoricalProjection(v_min: float, v_max: float, num_atoms: int, gamma: float, n_step: int, device: torch.device) -> None

Handles categorical distribution projection (L2-projection) for distributional RL (C51). Extracted from Learner for better separation of concerns.

Initialize the CategoricalProjection.

Parameters:

  • v_min (float) – Minimum value of the support.
  • v_max (float) – Maximum value of the support.
  • num_atoms (int) – Number of atoms in the distribution.
  • gamma (float) – Discount factor.
  • n_step (int) – Number of steps for n-step returns.
  • device (device) – The device to perform computations on.

Source

Methods ​

project ​

python
project(next_dist: torch.Tensor, rewards: torch.Tensor, dones: torch.Tensor, support: torch.Tensor, n_steps: torch.Tensor | int | None = None) -> torch.Tensor

Perform the categorical projection (L2-projection) of the target distribution.

Parameters:

  • next_dist (Tensor) – The distribution of the next state's chosen action with shape (batch_size, num_atoms).
  • rewards (Tensor) – The received rewards with shape (batch_size, 1).
  • dones (Tensor) – Whether the episode finished with shape (batch_size, 1).
  • support (Tensor) – The support tensor from the model.
  • n_steps (Tensor | int | None) – Optional number of steps for each transition. If None, uses self.n_step.

Returns:

  • Tensor – The projected target distribution with shape (batch_size, num_atoms).

Source

Values ​

delta_z ​

python
delta_z = (v_max - v_min) / (num_atoms - 1)

Source

device ​

python
device = device

Source

gamma ​

python
gamma = gamma

Source

n_step ​

python
n_step = n_step

Source

num_atoms ​

python
num_atoms = num_atoms

Source

v_max ​

python
v_max = v_max

Source

v_min ​

python
v_min = v_min

Source