ExportModel

ExportModel(export_format: str, action_masker: ActionMasker | None = None) -> None

ABC for wrapping a model for ONNX export.

Handles action masking and output format (per-action scores or greedy action index). Subclasses implement _get_scores to extract the relevant per-action output from the underlying model.

Initialize the ExportModel.

Parameters:

• export_format (str) – The output format string. If "action", the forward pass returns the greedy action index. Otherwise, returns per-action scores.
• action_masker (ActionMasker | None) – Optional ActionMasker to apply before action selection. When provided, the mask is baked into the ONNX graph.

Source

Methods

forward

forward(x: torch.Tensor) -> torch.Tensor

Perform a forward pass and return the output in the configured export format.

Parameters:

• x (Tensor) – The input observation tensor.

Returns:

• Tensor – Per-action scores of shape (batch, actions) if export_format is not "action", or action indices of shape (batch,) if export_format is "action".

Source