action_head

action_head

Configuration classes for modular action heads.

ActionHeadBlockConfig dataclass

ActionHeadBlockConfig(_target_=MISSING)

Base configuration for action head blocks.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.

LayerNormBlockConfig dataclass

LayerNormBlockConfig(_target_='versatil.models.decoding.action_heads.LayerNormBlock', input_dimension=MISSING)

Bases: ActionHeadBlockConfig

Configuration for layer-normalization action-head block.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
input_dimension	int	Input and output feature dimension.

MLPBlockConfig dataclass

MLPBlockConfig(_target_='versatil.models.decoding.action_heads.MLPBlock', input_dimension=MISSING, hidden_dimensions=None, output_dim=None, activation='silu', dropout=0.1, normalization=True)

Bases: ActionHeadBlockConfig

Configuration for MLP block in action head.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
input_dimension	int	Input dimension.
hidden_dimensions	list[int] | None	List of hidden dimensions.
output_dim	int | None	Output dimension (None to keep same as last hidden).
activation	str	Activation function name.
dropout	float	Dropout rate.
normalization	bool	Whether to apply layer normalization before MLP.

AttentionBlockConfig dataclass

AttentionBlockConfig(_target_='versatil.models.decoding.action_heads.AttentionBlock', embedding_dimension=MISSING, number_of_heads=8, dropout=0.1, normalization=True)

Bases: ActionHeadBlockConfig

Configuration for attention block in action head.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
embedding_dimension	int	Embedding dimension.
number_of_heads	int	Number of attention heads.
dropout	float	Dropout rate.
normalization	bool	Whether to apply layer normalization.

ResidualBlockConfig dataclass

ResidualBlockConfig(_target_='versatil.models.decoding.action_heads.ResidualBlock', block=MISSING, dropout=0.1)

Bases: ActionHeadBlockConfig

Configuration for residual wrapper block.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
block	dict[str, Any]	Block to wrap with residual connection.
dropout	float	Dropout rate after block.

AdaNormBlockConfig dataclass

AdaNormBlockConfig(_target_='versatil.models.decoding.action_heads.AdaNormBlock', input_dimension=MISSING, conditioning_dimension=MISSING, activation='silu')

Bases: ActionHeadBlockConfig

Configuration for adaptive normalization action-head block.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
input_dimension	int	Action embedding feature dimension.
conditioning_dimension	int	Conditioning vector dimension.
activation	str	Activation used inside the modulation projection.

ActionHeadConfig dataclass

ActionHeadConfig(_target_='versatil.models.decoding.action_heads.ActionHead', input_dimension=MISSING, blocks=None)

Configuration for a single action head.

Note

output dimension is set by the decoder based on the action key.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
input_dimension	int	Set from decoder embedding_dimension.
blocks	list[dict[str, Any]] | None	Head blocks applied in order.

ConditionalActionHeadConfig dataclass

ConditionalActionHeadConfig(_target_='versatil.models.decoding.action_heads.ConditionalActionHead', input_dimension=MISSING, conditioning_dimension=MISSING, blocks=None)

Configuration for a conditioned action head.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
input_dimension	int	Input action-token embedding dimension.
conditioning_dimension	int	Conditioning vector dimension.
blocks	list[dict[str, Any]] | None	Conditional blocks applied before the output projection.

GaussianHeadConfig dataclass

GaussianHeadConfig(_target_='versatil.models.decoding.action_heads.GaussianHead', input_dimension=MISSING, blocks=None, min_logvar=-10.0, max_logvar=4.0)

Configuration for GaussianHead that outputs mean and logvar.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
input_dimension	int	Input embedding dimension from decoder.
blocks	list[dict[str, Any]] | None	Blocks to apply before output projection.
min_logvar	float	Minimum value for logvar clamping.
max_logvar	float	Maximum value for logvar clamping.

MixtureOfExpertsHeadConfig dataclass

MixtureOfExpertsHeadConfig(_target_='versatil.models.decoding.action_heads.MoEHead', device='${policy.device}', experts=None, base_expert=None, num_experts=MISSING, gating_input_dim=None, gating_hidden_dims=None, routing_type=value, top_k=2, temperature=1.0, learnable_temperature=False, gating_dropout=0.1, gating_normalization=True)

Configuration for Mixture of Experts action head.

Supports two modes: 1. Explicit experts: Pass list of ActionHeadConfig 2. Base expert cloning: Pass base_expert and num_experts (recommended)

Note

base_expert is instantiated by Hydra, then cloned num_experts times by MoEHead to create separate expert networks with independent weights. output_dim is set by the decoder based on the action key.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
device	str	Device to place the module on.
experts	list[ActionHeadConfig] | None	Optional pre-instantiated expert action heads.
base_expert	ActionHeadConfig | None	Single expert instance to clone num_experts times.
num_experts	int	Number of experts to create from base_expert (optional for lazy init).
gating_input_dim	int | None	Input dimension for gating network (None for external routing).
gating_hidden_dims	list[int] | None	Hidden layer dimensions for gating network.
routing_type	str	Routing strategy ("soft" or "top_k").
top_k	int	Number of experts to use for top-k routing.
temperature	float	Temperature for softmax scaling of routing weights.
learnable_temperature	bool	Whether temperature should be a learnable parameter.
gating_dropout	float	Dropout rate in gating network.
gating_normalization	bool	Whether to normalize inputs to gating network.