Heads

Heads#

This module provides various head layers that convert network outputs into probability distributions over actions. These heads are commonly used in reinforcement learning to model stochastic or deterministic policies.

Descriptions#

GaussianHead#

  • Splits the input into mean and log-variance to form a Normal distribution.

  • Returns a sample (action), its log-probability, and the underlying distribution.

SquashedGaussianHead#

  • Outputs a Tanh-squashed Gaussian distribution for bounded action spaces.

  • Supports both training (stochastic sampling) and evaluation (mean of samples) modes.

  • Returns the action, log-probability (if training), and distribution.

CategoricalHead#

  • Converts logits into a Categorical distribution for discrete actions.

  • Uses softmax to derive probabilities.

  • Returns a sampled action, its log-probability, and the distribution.

DeterministicHead#

  • Pass-through head for deterministic policies.

  • Directly returns the input as the action.

Classes#

class objectrl.nets.layers.heads.GaussianHead(n: int)[source]#

Bases: Module

Outputs a Gaussian distribution from a given input tensor. The input is split into mean and log-variance components which define a Normal distribution. The class returns the sampled action, its log-probability, and the distribution.

Parameters:

n (int) – The dimensionality of the output action space.

_n#

Dimensionality of the action space.

Type:

int

__init__(n: int) → None[source]#

Initializes the Gaussian head with the specified action dimensionality.

Parameters:

n (int) – Dimensionality of the action space.

Returns:

None

forward(x: Tensor) → dict[str, Any][source]#

Forward pass to compute the Gaussian distribution and sample an action.

Parameters:

x (Tensor) – Input tensor containing concatenated mean and log-variance.

Returns:

Dictionary containing the sampled action, its log-probability, mean, and distribution

Return type:

dict

class objectrl.nets.layers.heads.SquashedGaussianHead(n: int, upper_clamp: float = -2.0, n_samples: int = 100)[source]#

Bases: Module

Outputs a Tanh-squashed Gaussian distribution, commonly used for bounded actions in reinforcement learning.

Parameters:

  • n (int) – Dimensionality of the action space.

  • upper_clamp (float) – Upper clamp on log-variance values.

  • n_samples (int) – Number of samples used in evaluation mode.

  • Attributes – _n (int): Dimensionality of the action space. _upper_clamp (float): Upper clamp for log-variance. _n_samples (int): Number of samples for evaluation.

__init__(n: int, upper_clamp: float = -2.0, n_samples: int = 100) → None[source]#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor, is_training: float = True) → dict[str, Any][source]#

Forward pass producing a squashed Gaussian distribution.

Parameters:

  • x (Tensor) – Input tensor with mean and log-variance concatenated.

  • is_training (bool) – Whether in training mode (stochastic sampling) or not.

Returns:

Dictionary containing the sampled action, its log-probability, and the distribution.

Return type:

dict

class objectrl.nets.layers.heads.CategoricalHead(n: int)[source]#

Bases: Module

Outputs a Categorical distribution for discrete action spaces.

Parameters:

n (int) – Number of categories (discrete actions).

Returns:

Dictionary containing the sampled action, its log-probability, and the distribution.

Return type:

dict

__init__(n: int) → None[source]#

Initializes the Categorical head with the specified number of categories.

Parameters:

n (int) – Number of categories (discrete actions).

Returns:

None

forward(x: Tensor) → dict[str, Any][source]#

Forward pass to compute categorical distribution and sample an action.

Parameters:

x (Tensor) – Input logits for categorical distribution.

Returns:

Dictionary containing the sampled action, its log-probability, and the distribution.

Return type:

dict

class objectrl.nets.layers.heads.DeterministicHead(n: int)[source]#

Bases: Module

Pass-through head for deterministic outputs.

Parameters:

n (int) – Dimensionality of the output.

Returns:

Dictionary containing the action.

Return type:

dict

__init__(n: int) → None[source]#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor) → dict[str, Any][source]#

Returns the input as the action.

Parameters:

x (Tensor) – Input tensor.

Returns:

Dictionary containing the action.

Return type:

dict

Usage Example#

from nets.layers.heads import GaussianHead

head = GaussianHead(n=4)
output = head(torch.randn(1, 8))  # 4 for mean, 4 for logvar
print(output["action"].shape)  # torch.Size([1, 4])

Notes#

  • These heads abstract the output distribution from policy networks.

  • Useful in actor-critic methods for exploration and entropy regularization.

  • SquashedGaussianHead ensures outputs stay in [-1, 1] using TanhTransform.

Contents