exportable_policy

exportable_policy

Wrapper making VersatIL Policy compatible with torch.export.export().

ExportablePolicy

ExportablePolicy(encoding_pipeline, algorithm, decoder, observation_keys, action_keys)

Bases: Module

Wraps Policy components for torch.export with positional tensor I/O.

torch.export requires positional tensor args and tuple returns. This wrapper converts between Policy's dict-based interface and the positional interface needed for quantization export.

Initialize with policy components and key orderings.

Parameters:

Name	Type	Description	Default
encoding_pipeline	EncodingPipeline	The policy's encoding pipeline.	required
algorithm	DecodingAlgorithm	The policy's decoding algorithm.	required
decoder	ActionDecoder	The policy's action decoder.	required
observation_keys	list[str]	Sorted list of observation dict keys.	required
action_keys	list[str]	Action output keys in action-space metadata order.	required

Source code in src/versatil/models/exportable_policy.py

def __init__(
    self,
    encoding_pipeline: EncodingPipeline,
    algorithm: DecodingAlgorithm,
    decoder: ActionDecoder,
    observation_keys: list[str],
    action_keys: list[str],
) -> None:
    """Initialize with policy components and key orderings.

    Args:
        encoding_pipeline: The policy's encoding pipeline.
        algorithm: The policy's decoding algorithm.
        decoder: The policy's action decoder.
        observation_keys: Sorted list of observation dict keys.
        action_keys: Action output keys in action-space metadata order.
    """
    super().__init__()
    self.encoding_pipeline = encoding_pipeline
    self.algorithm = algorithm
    self.decoder = decoder
    self._observation_keys = observation_keys
    self._action_keys = action_keys

observation_keys property

observation_keys

Get observation key ordering.

action_keys property

action_keys

Get action key ordering.

forward

forward(*observation_tensors)

Forward pass with positional tensor I/O.

Parameters:

Name	Type	Description	Default
*observation_tensors	Tensor	Tensors in the same order as observation_keys.	()

Returns:

Type	Description
tuple[Tensor, ...]	Tuple of action tensors in the same order as action_keys.

Source code in src/versatil/models/exportable_policy.py

def forward(self, *observation_tensors: torch.Tensor) -> tuple[torch.Tensor, ...]:
    """Forward pass with positional tensor I/O.

    Args:
        *observation_tensors: Tensors in the same order as observation_keys.

    Returns:
        Tuple of action tensors in the same order as action_keys.
    """
    if len(observation_tensors) != len(self._observation_keys):
        raise ValueError(
            f"Expected {len(self._observation_keys)} observation tensors "
            f"matching keys {self._observation_keys}, "
            f"got {len(observation_tensors)}"
        )
    observation_dict = dict(
        zip(self._observation_keys, observation_tensors, strict=True)
    )
    features = build_algorithm_features(
        observation=observation_dict,
        encoding_pipeline=self.encoding_pipeline,
        decoder=self.decoder,
        algorithm_injected_keys=self.algorithm.injected_feature_keys(),
    )
    predictions = self.algorithm.predict(features=features, network=self.decoder)
    return tuple(predictions[key] for key in self._action_keys)

from_policy classmethod

from_policy(policy)

Create ExportablePolicy from a loaded Policy.

Delegates key derivation to Policy.input_keys and Policy.output_keys, keeping the single source of truth.

Parameters:

Name	Type	Description	Default
policy	Policy	A loaded, initialized Policy instance.	required

Returns:

Type	Description
ExportablePolicy	ExportablePolicy wrapping the policy's components.

Raises:

Type	Description
ValueError	If the policy predicts action tokens; exported forward passes return raw tensors and never detokenize, so the exported outputs would not match the action space.

Source code in src/versatil/models/exportable_policy.py

@classmethod
def from_policy(cls, policy: Policy) -> "ExportablePolicy":
    """Create ExportablePolicy from a loaded Policy.

    Delegates key derivation to Policy.input_keys and
    Policy.output_keys, keeping the single source of truth.

    Args:
        policy: A loaded, initialized Policy instance.

    Returns:
        ExportablePolicy wrapping the policy's components.

    Raises:
        ValueError: If the policy predicts action tokens; exported
            forward passes return raw tensors and never detokenize, so
            the exported outputs would not match the action space.
    """
    if policy.decoder.requires_tokenized_actions:
        raise ValueError(
            "Policies with tokenized-action decoders cannot be exported: "
            "the exported forward returns raw action tokens without "
            "detokenization."
        )
    return cls(
        encoding_pipeline=policy.encoding_pipeline,
        algorithm=policy.algorithm,
        decoder=policy.decoder,
        observation_keys=policy.input_keys,
        action_keys=policy.output_keys,
    )

get_example_inputs

get_example_inputs(observation_shapes, batch_size=1, observation_dtypes=None)

Generate example inputs with correct shapes for torch.export.

Parameters:

Name	Type	Description	Default
observation_shapes	dict[str, tuple[int, ...]]	Mapping from observation key to shape tuple (excluding batch dimension). Must cover all observation_keys.	required
batch_size	int	Batch dimension size.	1
observation_dtypes	dict[str, dtype] | None	Optional mapping from observation key to torch dtype. Defaults to torch.float32 for all keys.	None

Returns:

Type	Description
tuple[Tensor, ...]	Tuple of tensors matching observation_keys order.

Source code in src/versatil/models/exportable_policy.py

def get_example_inputs(
    self,
    observation_shapes: dict[str, tuple[int, ...]],
    batch_size: int = 1,
    observation_dtypes: dict[str, torch.dtype] | None = None,
) -> tuple[torch.Tensor, ...]:
    """Generate example inputs with correct shapes for torch.export.

    Args:
        observation_shapes: Mapping from observation key to shape tuple
            (excluding batch dimension). Must cover all observation_keys.
        batch_size: Batch dimension size.
        observation_dtypes: Optional mapping from observation key to
            torch dtype. Defaults to torch.float32 for all keys.

    Returns:
        Tuple of tensors matching observation_keys order.
    """
    if observation_dtypes is None:
        observation_dtypes = {}

    example_tensors = []
    for key in self._observation_keys:
        if key not in observation_shapes:
            raise ValueError(
                f"No shape provided for observation key {key!r}. "
                f"observation_shapes must cover all observation_keys. "
                f"Missing keys: {set(self._observation_keys) - set(observation_shapes.keys())}"
            )
        shape = (batch_size, *observation_shapes[key])
        dtype = observation_dtypes.get(key, torch.float32)
        example_tensors.append(torch.zeros(shape, dtype=dtype))

    return tuple(example_tensors)