Skip to content

phase_act

phase_act

Phase-conditioned ACT decoder with MoE routing.

Extends the base ACT architecture to support phase-based expert routing. The phase classifier head produces routing logits that are used to route position and gripper predictions through phase-specific expert networks.

PhaseACT 

PhaseACT(input_keys, action_space, action_heads, observation_space, observation_horizon, prediction_horizon, device, embedding_dimension=256, number_of_heads=8, feedforward_dimension=512, number_of_encoder_layers=6, number_of_decoder_layers=6, activation='relu', dropout_rate=0.1, normalize_before=False, phase_routing_key=value)

Bases: ACT

Phase-conditioned Action Chunking Transformer.

This decoder extends ACT with phase-based expert routing. It expects: - A phase classifier head (output key: 'phase_label') - MoE action heads that use phase predictions for routing

During forward pass
  1. Phase classifier predicts phase logits/probabilities
  2. MoE action heads use phase logits as routing weights
  3. Each expert specializes in one surgical phase

Initialize PhaseACT decoder.

Parameters:

Name Type Description Default
input_keys list[str]

List of feature keys expected from encoder pipeline.

 required
action_space ActionSpace

Action space configuration.

 required
action_heads dict[str, ActionHead]

Dictionary of action head modules.

 required
observation_space ObservationSpace

Observation space configuration.

 required
observation_horizon int

Number of observation timesteps.

 required
prediction_horizon int

Number of actions to predict.

 required
device str

Device to run the model on.

 required
embedding_dimension int

Transformer hidden dimension.

 256
number_of_heads int

Number of attention heads.

 8
feedforward_dimension int

Feedforward network dimension.

 512
number_of_encoder_layers int

Number of transformer encoder layers.

 6
number_of_decoder_layers int

Number of transformer decoder layers.

 6
activation str

Activation function name.

 'relu'
dropout_rate float

Dropout probability.

 0.1
normalize_before bool

Use pre-normalization.

 False
phase_routing_key str

Key for the phase classifier head that provides routing weights.

 value
Source code in src/versatil/models/decoding/decoders/factory/phase_act.py 
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
def __init__(
    self,
    input_keys: list[str],
    action_space: ActionSpace,
    action_heads: dict[str, ActionHead],
    observation_space: ObservationSpace,
    observation_horizon: int,
    prediction_horizon: int,
    device: str,
    embedding_dimension: int = 256,
    number_of_heads: int = 8,
    feedforward_dimension: int = 512,
    number_of_encoder_layers: int = 6,
    number_of_decoder_layers: int = 6,
    activation: str = "relu",
    dropout_rate: float = 0.1,
    normalize_before: bool = False,
    phase_routing_key: str = TSOObsKey.PHASE_LABEL.value,
) -> None:
    """Initialize PhaseACT decoder.

    Args:
        input_keys: List of feature keys expected from encoder pipeline.
        action_space: Action space configuration.
        action_heads: Dictionary of action head modules.
        observation_space: Observation space configuration.
        observation_horizon: Number of observation timesteps.
        prediction_horizon: Number of actions to predict.
        device: Device to run the model on.
        embedding_dimension: Transformer hidden dimension.
        number_of_heads: Number of attention heads.
        feedforward_dimension: Feedforward network dimension.
        number_of_encoder_layers: Number of transformer encoder layers.
        number_of_decoder_layers: Number of transformer decoder layers.
        activation: Activation function name.
        dropout_rate: Dropout probability.
        normalize_before: Use pre-normalization.
        phase_routing_key: Key for the phase classifier head that provides routing weights.
    """
    resolved_action_heads = resolve_dict_keys(action_heads)
    if phase_routing_key not in resolved_action_heads:
        raise ValueError(
            f"PhaseACT requires '{phase_routing_key}' head for routing, "
            f"but only found: {list(resolved_action_heads.keys())}"
        )
    if not any(
        isinstance(resolved_action_heads[key], MoEHead)
        for key in resolved_action_heads
        if key != phase_routing_key
    ):
        raise ValueError(
            "PhaseACT requires at least one MoE action head for phase-based routing."
        )

    super().__init__(
        input_keys=input_keys,
        action_space=action_space,
        action_heads=resolved_action_heads,
        observation_space=observation_space,
        observation_horizon=observation_horizon,
        prediction_horizon=prediction_horizon,
        device=device,
        embedding_dimension=embedding_dimension,
        number_of_heads=number_of_heads,
        feedforward_dimension=feedforward_dimension,
        number_of_encoder_layers=number_of_encoder_layers,
        number_of_decoder_layers=number_of_decoder_layers,
        activation=activation,
        dropout_rate=dropout_rate,
        normalize_before=normalize_before,
    )
    self.phase_routing_key = phase_routing_key
    self._initialize_moe_experts()

get_callbacks 

get_callbacks(experiment_config)

Provide confusion matrix callback for phase classification monitoring.

Source code in src/versatil/models/decoding/decoders/factory/phase_act.py 
109
110
111
112
113
114
115
def get_callbacks(self, experiment_config: ExperimentConfig) -> list:
    """Provide confusion matrix callback for phase classification monitoring."""
    return [
        ConfusionMatrixCallback(
            log_every_n_epochs=experiment_config.val_every,
        )
    ]