moe

moe

Mixture of Experts (MoE) decoder for action prediction.

MoEDecoder

MoEDecoder(base_expert, num_experts, gating_feature_key, inference_gating_key=None, gating_input_dim=None, gating_hidden_dims=None, gating_activation=value, routing_type=value, top_k=2, temperature=1.0, learnable_temperature=False, gating_dropout=0.1, gating_normalization=True)

Bases: BaseMixtureOfExperts, ActionDecoder

Mixture of action decoder experts.

Source code in src/versatil/models/decoding/decoders/moe.py

def __init__(
    self,
    base_expert: ActionDecoder,
    num_experts: int,
    gating_feature_key: str,
    inference_gating_key: str | None = None,
    gating_input_dim: int | None = None,
    gating_hidden_dims: list[int] | None = None,
    gating_activation: str = ActivationFunction.SILU.value,
    routing_type: str = MoERoutingType.SOFT.value,
    top_k: int = 2,
    temperature: float = 1.0,
    learnable_temperature: bool = False,
    gating_dropout: float = 0.1,
    gating_normalization: bool = True,
) -> None:
    expert_list = self._create_experts_from_config(
        base_expert=base_expert, num_experts=num_experts
    )
    self.action_head_layout = base_expert.action_head_layout
    gating_keys = [gating_feature_key]
    if (
        inference_gating_key is not None
        and inference_gating_key != gating_feature_key
    ):
        gating_keys.append(inference_gating_key)
    moe_input = replace(
        base_expert.decoder_input,
        keys=[
            *base_expert.decoder_input.keys,
            *[
                key
                for key in gating_keys
                if key not in base_expert.decoder_input.keys
            ],
        ],
    )
    ActionDecoder.__init__(
        self,
        decoder_input=moe_input,
        observation_space=base_expert.observation_space,
        action_space=base_expert.action_space,
        action_heads=dict(base_expert.action_heads),
        device=str(base_expert.device),
        observation_horizon=base_expert.observation_horizon,
        prediction_horizon=base_expert.prediction_horizon,
    )

    BaseMixtureOfExperts.__init__(
        self,
        num_experts=num_experts,
        device=base_expert.device,
        gating_input_dim=gating_input_dim,
        gating_activation_function=gating_activation,
        gating_hidden_dims=gating_hidden_dims,
        routing_type=routing_type,
        top_k=top_k,
        temperature=temperature,
        learnable_temperature=learnable_temperature,
        gating_dropout=gating_dropout,
        gating_normalization=gating_normalization,
    )
    self.expert_decoders = nn.ModuleList(expert_list)
    self._base_expert_input_keys = set(base_expert.decoder_input.keys)
    self.num_experts = num_experts
    self.gating_feature_key = gating_feature_key
    self.inference_gating_key = (
        inference_gating_key
        if inference_gating_key is not None
        else gating_feature_key
    )
    self.action_keys = self._get_routed_output_keys(base_expert=base_expert)
    self.action_heads = self.expert_decoders[0].action_heads

get_auxiliary_output_keys

get_auxiliary_output_keys()

MoE decoder produces routing weights and per-expert outputs.

Source code in src/versatil/models/decoding/decoders/moe.py

def get_auxiliary_output_keys(self) -> set[str]:
    """MoE decoder produces routing weights and per-expert outputs."""
    keys = super().get_auxiliary_output_keys()
    keys.add(DecoderOutputKey.ROUTING_WEIGHTS.value)
    keys.add(DecoderOutputKey.EXPERT_OUTPUTS.value)
    return keys

forward

forward(features, actions=None)

Forward pass through mixture of expert decoders.

Parameters:

Name	Type	Description	Default
features	dict[str, Tensor]	Dictionary of input features	required
actions	dict[str, Tensor] | None	Optional ground-truth actions (for training)	None

Returns:

Type	Description
dict[str, Tensor | list[dict[str, Tensor]]]	Dictionary containing: - Combined predictions from routed experts (action keys) - routing_weights: Computed routing weights - expert_outputs: Individual expert prediction dictionaries

Raises:

Type	Description
KeyError	If the gating key for the current mode (gating_feature_key in training, inference_gating_key in eval) is missing from features.

Source code in src/versatil/models/decoding/decoders/moe.py

def forward(
    self,
    features: dict[str, torch.Tensor],
    actions: dict[str, torch.Tensor] | None = None,
) -> dict[str, torch.Tensor | list[dict[str, torch.Tensor]]]:
    """Forward pass through mixture of expert decoders.

    Args:
        features: Dictionary of input features
        actions: Optional ground-truth actions (for training)

    Returns:
        Dictionary containing:
            - Combined predictions from routed experts (action keys)
            - routing_weights: Computed routing weights
            - expert_outputs: Individual expert prediction dictionaries

    Raises:
        KeyError: If the gating key for the current mode (gating_feature_key
            in training, inference_gating_key in eval) is missing from
            ``features``.
    """
    if self.training:
        gating_key = self.gating_feature_key
    else:
        gating_key = self.inference_gating_key
    if gating_key not in features:
        raise KeyError(
            f"MoE gating feature '{gating_key}' is missing from decoder "
            f"features {sorted(features)}. Algorithm-injected latents are "
            f"exposed under '{LatentKey.POSTERIOR_LATENT.value}' during "
            "both training and inference."
        )
    gating_feature = features[gating_key]  # (B, embedding dimension)
    mixing_probabilities = self.compute_routing_weights(
        gating_feature
    )  # (B, num_experts)
    expert_features = {
        key: value
        for key, value in features.items()
        if key != gating_key or key in self._base_expert_input_keys
    }
    expert_outputs = [
        expert(expert_features, actions) for expert in self.expert_decoders
    ]
    combined_outputs = self._combine_expert_outputs(
        expert_outputs=expert_outputs, weights=mixing_probabilities
    )
    combined_outputs[DecoderOutputKey.ROUTING_WEIGHTS.value] = mixing_probabilities
    combined_outputs[DecoderOutputKey.EXPERT_OUTPUTS.value] = expert_outputs
    return combined_outputs