dit_block_action_transformer

dit_block_action_transformer

DiT Block action transformer decoder with pooled conditioning.

Uses DiTBlock Policy architecture, a diffusion transformer with an encoder that pools encoder output to a single conditioning vector. Supports encoder caching for inference optimization.

DiTBlockActionTransformer

DiTBlockActionTransformer(input_keys, action_space, action_heads, observation_space, observation_horizon, prediction_horizon, device, max_sequence_length=1024, embedding_dimension=512, timestep_embedding_dimension=256, number_of_heads=8, number_of_key_value_heads=None, number_of_encoder_layers=6, number_of_decoder_layers=6, feedforward_dimension=2048, activation=value, normalization_type=value, attention_type=value, dropout_rate=0.1, attention_dropout=0.0, positional_encoding_type=value, use_gating=True)

Bases: BaseParallelTransformerDecoder

Diffusion action transformer decoder using DiTBlock with pooled conditioning.

This architecture: - Processes observation tokens through encoder with mean pooling - Conditions decoder via the sum of pooled vector + timestep embedding (AdaLN) - Caches pooled encoder output during inference

Initialize DiTBlock action 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 (for history).	required
prediction_horizon	int	Number of actions to predict (horizon).	required
device	str	Device to run the model on.	required
max_sequence_length	int	Maximum sequence length for input tokens.	1024
embedding_dimension	int	Transformer hidden dimension.	512
timestep_embedding_dimension	int	Diffusion timestep embedding dimension.	256
number_of_heads	int	Number of attention heads.	8
number_of_key_value_heads	int | None	Number of K/V heads for GQA.	None
number_of_encoder_layers	int	Number of transformer encoder layers.	6
number_of_decoder_layers	int	Number of transformer decoder layers.	6
feedforward_dimension	int	Feedforward network dimension.	2048
activation	str	Activation function name.	value
normalization_type	str	Normalization type name.	value
attention_type	str	Attention type name (gqa, mha).	value
dropout_rate	float	Dropout probability for residual connections.	0.1
attention_dropout	float	Dropout probability for attention weights.	0.0
positional_encoding_type	str | None	Type of positional encoding.	value
use_gating	bool	Whether to use gating in AdaLN-Zero layers.	True

Source code in src/versatil/models/decoding/decoders/factory/dit_block_action_transformer.py

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,
    max_sequence_length: int = 1024,
    embedding_dimension: int = 512,
    timestep_embedding_dimension: int = 256,
    number_of_heads: int = 8,
    number_of_key_value_heads: int | None = None,
    number_of_encoder_layers: int = 6,
    number_of_decoder_layers: int = 6,
    feedforward_dimension: int = 2048,
    activation: str = ActivationFunction.SWIGLU.value,
    normalization_type: str = NormalizationType.RMS_NORM.value,
    attention_type: str = AttentionType.MULTI_HEAD.value,
    dropout_rate: float = 0.1,
    attention_dropout: float = 0.0,
    positional_encoding_type: str | None = PositionalEncodingType.ROPE.value,
    use_gating: bool = True,
) -> None:
    """Initialize DiTBlock action 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 (for history).
        prediction_horizon: Number of actions to predict (horizon).
        device: Device to run the model on.
        max_sequence_length: Maximum sequence length for input tokens.
        embedding_dimension: Transformer hidden dimension.
        timestep_embedding_dimension: Diffusion timestep embedding dimension.
        number_of_heads: Number of attention heads.
        number_of_key_value_heads: Number of K/V heads for GQA.
        number_of_encoder_layers: Number of transformer encoder layers.
        number_of_decoder_layers: Number of transformer decoder layers.
        feedforward_dimension: Feedforward network dimension.
        activation: Activation function name.
        normalization_type: Normalization type name.
        attention_type: Attention type name (gqa, mha).
        dropout_rate: Dropout probability for residual connections.
        attention_dropout: Dropout probability for attention weights.
        positional_encoding_type: Type of positional encoding.
        use_gating: Whether to use gating in AdaLN-Zero layers.
    """
    self.action_space = action_space
    self.observation_space = observation_space
    self.observation_horizon = observation_horizon
    self.prediction_horizon = prediction_horizon
    self.max_sequence_length = max_sequence_length
    self.timestep_embedding_dimension = timestep_embedding_dimension
    self.number_of_heads = number_of_heads
    self.number_of_key_value_heads = number_of_key_value_heads or number_of_heads
    self.number_of_encoder_layers = number_of_encoder_layers
    self.number_of_decoder_layers = number_of_decoder_layers
    self.feedforward_dimension = feedforward_dimension or (4 * embedding_dimension)
    self.activation = activation
    self.normalization_type = normalization_type
    self.attention_type = attention_type
    self.dropout_rate = dropout_rate
    self.attention_dropout = attention_dropout
    self.positional_encoding_type = positional_encoding_type
    self.use_gating = use_gating

    decoder_input = DecoderInput(
        keys=input_keys,
        raises_for_types=[FeatureType.SPATIAL.value],
        requires_actions=True,
    )
    super().__init__(
        decoder_input=decoder_input,
        action_space=action_space,
        action_heads=action_heads,
        observation_space=observation_space,
        prediction_horizon=prediction_horizon,
        observation_horizon=observation_horizon,
        device=device,
        embedding_dimension=embedding_dimension,
    )
    self._validate_conditional_action_head()
    self._build_transformer_components()

enable_encoder_cache

enable_encoder_cache()

Enable encoder caching for multi-step inference loops.

Called by algorithms at the start of predict() to enable caching across denoising steps within a single sample.

Source code in src/versatil/models/decoding/decoders/factory/dit_block_action_transformer.py

def enable_encoder_cache(self) -> None:
    """Enable encoder caching for multi-step inference loops.

    Called by algorithms at the start of predict() to enable caching
    across denoising steps within a single sample.
    """
    self._encoder_cache = None
    self._caching_enabled = True

disable_encoder_cache

disable_encoder_cache()

Disable encoder caching and clear cached values.

Source code in src/versatil/models/decoding/decoders/factory/dit_block_action_transformer.py

def disable_encoder_cache(self) -> None:
    """Disable encoder caching and clear cached values."""
    self._encoder_cache = None
    self._caching_enabled = False

forward

forward(features, actions=None)

Forward pass through DiTBlock transformer.

Parameters:

Name	Type	Description	Default
features	dict[str, Tensor]	Dictionary of encoded features plus timestep.	required
actions	dict[str, Tensor] | None	Dictionary of noise-injected actions.	None

Returns:

Type	Description
dict[str, Tensor]	Dictionary containing denoised predictions for each action head.

Raises:

Type	Description
ValueError	If timesteps or actions are missing.

Source code in src/versatil/models/decoding/decoders/factory/dit_block_action_transformer.py

def forward(
    self,
    features: dict[str, torch.Tensor],
    actions: dict[str, torch.Tensor] | None = None,
) -> dict[str, torch.Tensor]:
    """Forward pass through DiTBlock transformer.

    Args:
        features: Dictionary of encoded features plus timestep.
        actions: Dictionary of noise-injected actions.

    Returns:
        Dictionary containing denoised predictions for each action head.

    Raises:
        ValueError: If timesteps or actions are missing.
    """
    if actions is None:
        raise ValueError(
            "DiTBlockActionTransformer requires 'actions' parameter. "
            "The algorithm should provide noisy actions during forward pass."
        )
    noisy_actions = self.action_space.concatenate_action_tensors(
        actions=actions,
        prediction_horizon=self.prediction_horizon,
        owner_name=self.__class__.__name__,
    )
    timesteps = extract_timestep_conditioning(
        features=features,
        batch_size=noisy_actions.shape[0],
        action_device=noisy_actions.device,
    )
    observation_features = filter_timestep_feature(features=features)
    (
        observation_tokens,
        observation_positional_encodings,
        observation_padding_mask,
    ) = self._prepare_observation_tokens(observation_features)
    if observation_positional_encodings is not None:
        observation_tokens = observation_tokens + observation_positional_encodings
    noisy_embedding = self.noisy_input_projection(noisy_actions)  # (B, T, D)
    encoder_cache, action_hidden, action_conditioning = (
        self.transformer.forward_features(
            decoder_hidden_states=noisy_embedding,
            timesteps=timesteps,
            encoder_hidden_states=observation_tokens,
            encoder_padding_mask=observation_padding_mask,
            decoder_padding_mask=None,
            encoder_cache=self._encoder_cache if self._caching_enabled else None,
        )
    )  # (B, D), (B, T, D), (B, D)
    if self._caching_enabled:
        self._encoder_cache = encoder_cache
    noise_predictions = self._conditional_action_head()(
        action_hidden,
        action_conditioning,
    )
    return self.action_space.split_action_tensor(
        action_tensor=noise_predictions,
        owner_name=self.__class__.__name__,
    )