pipeline

pipeline

Encoding pipeline for multi-modal observation encoding and fusion.

EncodingPipeline

EncodingPipeline(encoders, observation_space, fusion_stages=None)

Bases: Module

Pipeline that encodes inputs and fuses them hierarchically.

All encoder outputs and fusion outputs are available in the final feature dictionary — nothing is consumed or removed.

Initialize the encoding pipeline.

Parameters:

Name	Type	Description	Default
encoders	dict[str, EncodingMixin]	Dictionary of instantiated encoders keyed by name.	required
observation_space	ObservationSpace	Observation space with task metadata.	required
fusion_stages	list[FusionModule] | None	List of instantiated fusion modules.	None

Source code in src/versatil/models/encoding/pipeline.py

def __init__(
    self,
    encoders: dict[str, EncodingMixin],
    observation_space: ObservationSpace,
    fusion_stages: list[FusionModule] | None = None,
):
    """Initialize the encoding pipeline.

    Args:
        encoders: Dictionary of instantiated encoders keyed by name.
        observation_space: Observation space with task metadata.
        fusion_stages: List of instantiated fusion modules.
    """
    super().__init__()
    self.observation_space = observation_space
    self.encoders = nn.ModuleDict()
    self.conditional_encoders = nn.ModuleDict()
    #: all feature_name (both encoders and fusion layers) -> FeatureMetadata
    self._feature_registry: dict[str, FeatureMetadata] = {}
    self._encoder_feature_keys: dict[str, list[str]] = {}
    #: Set by the workspace from experiment.precision via set_output_dtype.
    self.output_dtype: torch.dtype | None = None
    self._setup_encoders(encoders=encoders)
    self._setup_fusion_modules(fusion_stages=fusion_stages)
    self._validate_pipeline()

set_output_dtype

set_output_dtype(output_dtype)

Pin the dtype of features emitted by forward.

Called once by the workspace after Hydra instantiation, with the pipeline-wide precision (experiment.precision).

Source code in src/versatil/models/encoding/pipeline.py

def set_output_dtype(self, output_dtype: torch.dtype | None) -> None:
    """Pin the dtype of features emitted by ``forward``.

    Called once by the workspace after Hydra instantiation, with the
    pipeline-wide precision (``experiment.precision``).
    """
    self.output_dtype = output_dtype

forward

forward(observation)

Encode observations through all encoders and apply fusion stages.

Runs non-conditional encoders first, then conditional encoders (which may depend on earlier outputs as conditioning), then fusion modules. All outputs are kept in the returned dictionary. Encoders whose input keys are missing from the observation are skipped with a warning.

Parameters:

Name	Type	Description	Default
observation	dict[str, Tensor]	Dictionary mapping observation keys to tensors.	required

Returns:

Type	Description
dict[str, Tensor]	Dictionary mapping prefixed feature names to tensors. Includes all
dict[str, Tensor]	encoder outputs and all fusion outputs.

Source code in src/versatil/models/encoding/pipeline.py

def forward(
    self,
    observation: dict[str, torch.Tensor],
) -> dict[str, torch.Tensor]:
    """Encode observations through all encoders and apply fusion stages.

    Runs non-conditional encoders first, then conditional encoders (which may
    depend on earlier outputs as conditioning), then fusion modules. All outputs
    are kept in the returned dictionary. Encoders whose input keys are missing
    from the observation are skipped with a warning.

    Args:
        observation: Dictionary mapping observation keys to tensors.

    Returns:
        Dictionary mapping prefixed feature names to tensors. Includes all
        encoder outputs and all fusion outputs.
    """
    flat_obs = self._flatten_observation_dict(observation)
    features = {}
    for encoder_name, encoder in self.encoders.items():
        input_keys = encoder.input_specification.keys
        missing_keys, encoder_input = self._get_encoder_runtime_inputs(
            input_keys=input_keys,
            flat_observation=flat_obs,
        )
        if missing_keys:
            logging.warning(
                f"Encoder '{encoder_name}' skipped: missing {missing_keys}"
            )
            continue
        encoded = encoder(encoder_input)
        for feature_key in self._encoder_feature_keys[encoder_name]:
            if feature_key in encoded:
                features[f"{encoder_name}_{feature_key}"] = encoded[feature_key]

    for encoder_name, encoder in self.conditional_encoders.items():
        input_keys = encoder.input_specification.keys
        missing_keys, encoder_input = self._get_encoder_runtime_inputs(
            input_keys=input_keys,
            flat_observation=flat_obs,
        )
        if missing_keys:
            logging.warning(
                f"Conditional encoder '{encoder_name}' skipped: missing {missing_keys}"
            )
            continue
        condition_key = encoder.condition_key
        if condition_key not in features:
            logging.warning(
                f"Conditional encoder '{encoder_name}' skipped: conditioning "
                f"feature '{condition_key}' was not produced (its encoder "
                "may have been skipped)"
            )
            continue
        encoded = encoder(encoder_input, features[condition_key])
        for feature_key in self._encoder_feature_keys[encoder_name]:
            if feature_key in encoded:
                features[f"{encoder_name}_{feature_key}"] = encoded[feature_key]

    for fusion_module in self.fusion_stages:
        missing_features = [
            feature_name
            for feature_name in fusion_module.input_features
            if feature_name not in features
        ]
        if missing_features:
            raise ValueError(
                f"Fusion stage '{fusion_module.output_name}' requires "
                f"features {missing_features} that were not produced; a "
                "contributing encoder was likely skipped because its "
                "observation keys are missing."
            )
        input_features = [
            features[feat_name] for feat_name in fusion_module.input_features
        ]
        features[fusion_module.output_name] = fusion_module(input_features)

    if self.output_dtype is not None:
        features = dict_apply(
            features,
            lambda x: x.to(self.output_dtype) if torch.is_floating_point(x) else x,
        )
    return features

get_feature_names

get_feature_names()

Get all feature names produced by the encoding pipeline.

Source code in src/versatil/models/encoding/pipeline.py

def get_feature_names(self) -> list[str]:
    """Get all feature names produced by the encoding pipeline."""
    return list(self._feature_registry.keys())

get_features

get_features()

Get all feature metadata produced by the encoding pipeline.

Source code in src/versatil/models/encoding/pipeline.py

def get_features(self) -> dict[str, FeatureMetadata]:
    """Get all feature metadata produced by the encoding pipeline."""
    return dict(self._feature_registry)

get_features_to_dimensions

get_features_to_dimensions()

Get a dictionary of feature names to dimensions.

Source code in src/versatil/models/encoding/pipeline.py

def get_features_to_dimensions(self) -> dict[str, tuple[int, ...]]:
    """Get a dictionary of feature names to dimensions."""
    return {key: meta.dimension for key, meta in self._feature_registry.items()}

set_tokenizer

set_tokenizer(tokenizer=None)

Set tokenizer and validate vocab sizes.

Parameters:

Name	Type	Description	Default
tokenizer	Tokenizer | None	Tokenizer instance from data pipeline (can be None).	None

Raises:

Type	Description
ValueError	If observation tokenizer vocab size doesn't match encoder vocab sizes.

Source code in src/versatil/models/encoding/pipeline.py

def set_tokenizer(self, tokenizer: Tokenizer | None = None):
    """Set tokenizer and validate vocab sizes.

    Args:
        tokenizer: Tokenizer instance from data pipeline (can be None).

    Raises:
        ValueError: If observation tokenizer vocab size doesn't match encoder vocab sizes.
    """
    all_encoders = {**self.encoders, **self.conditional_encoders}
    for encoder_name, encoder in all_encoders.items():
        if not encoder.input_specification.requires_tokenized:
            continue
        if tokenizer is None or tokenizer.observation_tokenizer is None:
            raise ValueError(
                f"Encoder '{encoder_name}' requires tokenized input, "
                f"but no observation tokenizer is available."
            )
        data_vocab_size = tokenizer.observation_tokenizer.vocab_size
        base_vocab_size = (
            tokenizer.observation_tokenizer.language_tokenizer.vocab_size
        )
        encoder_vocab_size = encoder.get_vocab_size()
        # Accept if the encoder covers either the total tokenizer vocab
        # (base + added) or at least the base vocabulary. Some tokenizers
        # (e.g. EmbeddingGemma) can add extra tokens that never appear in
        # text prompts, so tolerating the gap is safe.
        if (
            encoder_vocab_size < data_vocab_size
            and encoder_vocab_size < base_vocab_size
        ):
            raise ValueError(
                f"Vocab size mismatch: Observation tokenizer has vocab_size={data_vocab_size} "
                f"(base={base_vocab_size}), but encoder '{encoder_name}' only supports "
                f"vocab_size={encoder_vocab_size}. The encoder's embedding matrix must be "
                f"at least as large as the tokenizer's base vocabulary. "
                f"Observation tokenizer model: {tokenizer.observation_tokenizer.tokenizer_model}"
            )

__repr__

__repr__()

Pretty print the pipeline structure.

Source code in src/versatil/models/encoding/pipeline.py

def __repr__(self) -> str:
    """Pretty print the pipeline structure."""
    lines = ["EncodingPipeline(", "  Encoders:"]
    for enc_name, encoder in {**self.encoders, **self.conditional_encoders}.items():
        input_keys = encoder.input_specification.keys
        output_keys = [
            f"{enc_name}_{k}" for k in self._encoder_feature_keys[enc_name]
        ]
        lines.append(
            f"    {enc_name}: {input_keys} -> {output_keys} "
            f"({encoder.__class__.__name__})"
        )
    if self.fusion_stages:
        lines.append("  Fusion stages:")
        for i, fusion in enumerate(self.fusion_stages):
            inputs = ", ".join(fusion.input_features)
            lines.append(
                f"    {i}: [{inputs}] -> {fusion.output_name} "
                f"({fusion.__class__.__name__})"
            )
    lines.append(")")
    return "\n".join(lines)