base

base

Base contract for configurable action heads.

BaseActionHead

BaseActionHead(input_dimension, blocks=None)

Bases: ABC, Module

Abstract base class for action heads with block-based processing and output projection.

The output dimension is set lazily via set_output_dim() because action heads are instantiated from config with only the embedding dimension known. The output dimension depends on the ActionSpace (resolved by ActionDecoder during policy assembly) or the tokenizer vocabulary size (for tokenized decoders).

Subclasses must implement forward() with their specific return type.

Initialize base action head.

Parameters:

Name	Type	Description	Default
input_dimension	int	Input embedding dimension from decoder.	required
blocks	list[ActionHeadBlock] | None	Blocks to apply before output projection.	None

Source code in src/versatil/models/decoding/action_heads/base.py

def __init__(
    self,
    input_dimension: int,
    blocks: list[ActionHeadBlock] | None = None,
) -> None:
    """Initialize base action head.

    Args:
        input_dimension: Input embedding dimension from decoder.
        blocks: Blocks to apply before output projection.
    """
    super().__init__()
    self.input_dimension = input_dimension
    self._output_dim: int | None = None
    if blocks is None:
        blocks = []
    self.blocks = nn.ModuleList(blocks)
    self.output_proj: nn.Linear | None = None

output_dim property writable

output_dim

Get output dimension. Raises if not set.

set_output_dim

set_output_dim(dim)

Set output dimension and create output projection layer.

Parameters:

Name	Type	Description	Default
dim	int	Output action dimension.	required

Source code in src/versatil/models/decoding/action_heads/base.py

def set_output_dim(self, dim: int) -> None:
    """Set output dimension and create output projection layer.

    Args:
        dim: Output action dimension.
    """
    self._output_dim = dim
    hidden_dimension = self._get_hidden_dim()
    self.output_proj = nn.Linear(hidden_dimension, dim)

forward abstractmethod

forward(action_embedding)

Forward pass. Subclasses define return type.

Source code in src/versatil/models/decoding/action_heads/base.py

@abstractmethod
def forward(
    self,
    action_embedding: torch.Tensor,
) -> torch.Tensor | dict[str, torch.Tensor]:
    """Forward pass. Subclasses define return type."""
    raise NotImplementedError