export

export

Model export utilities for post-training compression.

export_policy

export_policy(exportable, example_inputs)

Export an ExportablePolicy with dynamic batch dimension.

Runs one eager forward pass before tracing to materialize any lazily-initialized modules (e.g. FeatureProjection layers). torch.export's FX tracer silently drops nn.ModuleDict mutations, so all projection layers must exist before tracing begins.

Parameters:

Name	Type	Description	Default
exportable	ExportablePolicy	The ExportablePolicy wrapping the policy.	required
example_inputs	tuple[Tensor, ...]	Example input tensors for tracing.	required

Returns:

Type	Description
Module	Exported FX GraphModule.

Source code in src/versatil/post_training_compression/export.py

def export_policy(
    exportable: ExportablePolicy,
    example_inputs: tuple[torch.Tensor, ...],
) -> nn.Module:
    """Export an ExportablePolicy with dynamic batch dimension.

    Runs one eager forward pass before tracing to materialize any
    lazily-initialized modules (e.g. FeatureProjection layers).
    torch.export's FX tracer silently drops nn.ModuleDict mutations,
    so all projection layers must exist before tracing begins.

    Args:
        exportable: The ExportablePolicy wrapping the policy.
        example_inputs: Example input tensors for tracing.

    Returns:
        Exported FX GraphModule.
    """
    logging.info("Materializing lazy modules with eager forward pass...")
    with torch.no_grad():
        exportable(*example_inputs)

    return _export_with_dynamic_batch(
        model=exportable,
        example_inputs=example_inputs,
        dynamic_shapes_key="observation_tensors",
    ).module()

build_example_inputs

build_example_inputs(exportable, observation_space, observation_horizon, tokenizer=None)

Build example inputs from observation space metadata.

Parameters:

Name	Type	Description	Default
exportable	ExportablePolicy	ExportablePolicy defining required observation keys.	required
observation_space	ObservationSpace	Observation space with camera/proprio metadata.	required
observation_horizon	int	Number of temporal observation frames.	required
tokenizer	Tokenizer | None	Tokenizer for language token sequence length.	None

Returns:

Type	Description
tuple[Tensor, ...]	Tuple of example input tensors matching exportable.observation_keys.

Source code in src/versatil/post_training_compression/export.py

def build_example_inputs(
    exportable: ExportablePolicy,
    observation_space: ObservationSpace,
    observation_horizon: int,
    tokenizer: Tokenizer | None = None,
) -> tuple[torch.Tensor, ...]:
    """Build example inputs from observation space metadata.

    Args:
        exportable: ExportablePolicy defining required observation keys.
        observation_space: Observation space with camera/proprio metadata.
        observation_horizon: Number of temporal observation frames.
        tokenizer: Tokenizer for language token sequence length.

    Returns:
        Tuple of example input tensors matching exportable.observation_keys.
    """
    observation_shapes: dict[str, tuple[int, ...]] = {}
    observation_dtypes: dict[str, torch.dtype] = {}

    for key, camera_meta in observation_space.cameras.items():
        observation_shapes[key] = (
            observation_horizon,
            camera_meta.channels,
            camera_meta.image_height,
            camera_meta.image_width,
        )

    for key, state_meta in observation_space.numerical_observations.items():
        observation_shapes[key] = (observation_horizon, state_meta.dimension)
        observation_dtypes[key] = torch.from_numpy(
            np.empty(0, dtype=np.dtype(state_meta.dtype))
        ).dtype

    if tokenizer is not None and tokenizer.observation_tokenizer is not None:
        token_length = tokenizer.observation_tokenizer.max_token_len
        observation_shapes[SampleKey.TOKENIZED_OBSERVATIONS.value] = (
            observation_horizon,
            token_length,
        )
        observation_dtypes[SampleKey.TOKENIZED_OBSERVATIONS.value] = torch.long
        observation_shapes[SampleKey.IS_PAD_OBSERVATION.value] = (
            observation_horizon,
            token_length,
        )
        observation_dtypes[SampleKey.IS_PAD_OBSERVATION.value] = torch.bool

    return exportable.get_example_inputs(
        observation_shapes=observation_shapes,
        observation_dtypes=observation_dtypes,
        batch_size=2,
    )