DQNExportModel

DQNExportModel(model: DQNModel, export_format: Literal['q_values', 'action'], action_masker: ActionMasker | None = None) -> None

Wrapper module for exporting a DQNModel to ONNX with a simplified output format.

Transforms the raw DQNModel output into either Q-values or a single action index, handling both categorical and non-categorical DQN variants via get_q_values().

When an action mask is provided, untrained actions are set to -1e9 before computing the argmax or returning Q-values. This ensures the exported model's behaviour matches training, where only masked actions are considered.

Initialize the DQNExportModel.

Parameters:

• model (DQNModel) – The DQNModel to wrap for export.
• export_format (Literal['q_values', 'action']) – The output format - "q_values" for Q-values per action, or "action" for the greedy action index.
• action_masker (ActionMasker | None) – Optional ActionMasker to apply before action selection. When provided, the mask is baked into the ONNX graph.

Source