single_output

single_output

ActionHead classes for composing blocks into action prediction heads.

ActionHead

ActionHead(input_dimension, blocks=None)

Bases: BaseActionHead

Single-output action head returning one tensor.

Converts decoder embeddings (B, horizon, embedding_dimension) to action predictions (B, horizon, action_dim) via blocks and a final linear projection.

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

forward

forward(action_embedding)

Convert embeddings to actions.

Parameters:

Name	Type	Description	Default
action_embedding	Tensor	(B, prediction_horizon, embedding_dimension) or (B, embedding_dimension)	required

Returns:

Type	Description
Tensor	Action predictions with same batch/horizon dims, last dim is action_dim.

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

def forward(self, action_embedding: torch.Tensor) -> torch.Tensor:
    """Convert embeddings to actions.

    Args:
        action_embedding: (B, prediction_horizon, embedding_dimension) or (B, embedding_dimension)

    Returns:
        Action predictions with same batch/horizon dims, last dim is action_dim.
    """
    if self.output_proj is None:
        raise RuntimeError("output_dim not set. Call set_output_dim() first.")
    action_embedding = self._apply_blocks(action_embedding)
    return self.output_proj(action_embedding)