conditional

conditional

Conditional action heads.

ConditionalActionHead

ConditionalActionHead(input_dimension, conditioning_dimension, blocks=None)

Bases: BaseActionHead

Action head whose blocks receive a conditioning vector.

Initialize the conditional action head.

Parameters:

Name	Type	Description	Default
input_dimension	int	Input action-token embedding dimension.	required
conditioning_dimension	int	Conditioning vector dimension.	required
blocks	list[ConditionalActionHeadBlock] | None	Conditional blocks applied before the output projection.	None

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

def __init__(
    self,
    input_dimension: int,
    conditioning_dimension: int,
    blocks: list[ConditionalActionHeadBlock] | None = None,
) -> None:
    """Initialize the conditional action head.

    Args:
        input_dimension: Input action-token embedding dimension.
        conditioning_dimension: Conditioning vector dimension.
        blocks: Conditional blocks applied before the output projection.
    """
    super().__init__(input_dimension=input_dimension, blocks=blocks)
    self.conditioning_dimension = conditioning_dimension

forward

forward(action_embedding, condition)

Project conditioned action embeddings to action predictions.

Parameters:

Name	Type	Description	Default
action_embedding	Tensor	Action-token embeddings with shape (B, prediction_horizon, input_dimension).	required
condition	Tensor	Conditioning tensor with shape (B, conditioning_dimension).	required

Returns:

Type	Description
Tensor	Action predictions with shape
Tensor	(B, prediction_horizon, output_dim).

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

def forward(
    self,
    action_embedding: torch.Tensor,
    condition: torch.Tensor,
) -> torch.Tensor:
    """Project conditioned action embeddings to action predictions.

    Args:
        action_embedding: Action-token embeddings with shape
            ``(B, prediction_horizon, input_dimension)``.
        condition: Conditioning tensor with shape ``(B, conditioning_dimension)``.

    Returns:
        Action predictions with shape
        ``(B, prediction_horizon, output_dim)``.
    """
    if self.output_proj is None:
        raise RuntimeError("output_dim not set. Call set_output_dim() first.")
    for block in self.blocks:
        action_embedding = block(action_embedding, condition)
    return self.output_proj(action_embedding)