Exports

Simulusion provides a unified export system that converts trained models and agents into deployable formats — torch.export programs, AOTInductor packages, ONNX, TensorRT and (legacy) TorchScript — without requiring each model to implement its own export logic. For making inference fast, choosing what to export for a recurrent world model, and packaging it with its config, see Deployment and export.

Overview

The export system is built around three levels of API:

Supported formats

torch.export options pass straight through, for example dynamic_shapes={"x": {0: torch.export.Dim("batch")}} to keep the batch size symbolic.

Loading and verifying artifacts

from simulusion import export_model, load_exported, verify_export

path = export_model(module, "step.pt2", format="exported_program", example_inputs=x)
max_err = verify_export(module, path, x)  # raises AssertionError on mismatch
runner = load_exported(path)  # format inferred from the suffix
runner = load_exported("step_aoti.pt2", format="aoti")

load_exported returns a plain callable for every format. ONNX artifacts run through ONNX Runtime when it is installed.

Quick start

Exporting any nn.Module

simulusion.export_model(module, path, ...) exports any torch.nn.Module, including your own. Simulusion does not add methods to torch.nn.Module; its public model classes get .export() from ExportableAgentMixin, which you can also inherit in your own classes (see Custom agents).

import torch
from simulusion import export_model


class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(64, 10)

    def forward(self, x):
        return self.linear(x)


model = MyModel()
model.eval()

# torch.export program (recommended)
export_model(
    model, "model.pt2", format="exported_program", example_inputs=torch.zeros(1, 64)
)

# ONNX (requires example_inputs)
export_model(model, "model.onnx", format="onnx", example_inputs=torch.zeros(1, 64))

Simulusion's own models (for example simulusion.create_model("genie-small")) can call the same thing as a method: genie.export("genie.pt2", format="exported_program", example_inputs=video).

Exporting a trained agent

High-level agents support the same .export() method directly:

import simulusion

agent = simulusion.create_model("dreamer", env="walker-walk", total_steps=1000)
# ... train the agent ...

# Export the policy to ONNX
agent.export("policy.onnx", format="onnx")

# Export the RSSM world model to TorchScript
agent.export("rssm.pt", format="torchscript", target="rssm")

# Export the observation encoder to ONNX
agent.export("encoder.onnx", format="onnx", target="obs_encoder")

The system automatically resolves which submodule to export and infers the correct example inputs for each agent type.

Target resolution

When you call .export(), the system needs to decide which nn.Module to serialize. It uses a priority-based resolution strategy:

1. Agent-specific defaults

Each agent type has a preferred default target:

2. Explicit target= parameter

Override the default by naming a specific submodule:

# Export by attribute path
agent.export("reward.pt", format="torchscript", target="dreamer.reward_model")

# If the attribute name is unique, the short name works
agent.export("value.pt", format="torchscript", target="value_model")

If the short name matches multiple modules, the system raises an error and lists the available fully qualified paths.

3. Single-module fallback

If the object is itself an nn.Module, it is exported directly. If it contains exactly one nn.Module attribute, that attribute is exported. If it contains multiple modules, the system picks the first match from this priority list: actor, policy, actor_critic, rssm, model, world_model, encoder.

Format-specific details

ONNX

agent.export(
    "policy.onnx",
    format="onnx",
    input_names=["latent"],
    output_names=["action"],
    dynamic_axes={"latent": {0: "batch"}, "action": {0: "batch"}},
    opset_version=17,  # default
)

See the PyTorch ONNX export docs for all supported keyword arguments.

TorchScript

Two modes, controlled by whether example_inputs is provided:

# Trace (requires example_inputs)
agent.export("traced.pt", format="torchscript", example_inputs=torch.zeros(1, 230))

# Script (no example_inputs needed)
agent.export("scripted.pt", format="torchscript")

TensorRT

Requires the optional torch_tensorrt package:

pip install torch_tensorrt

agent.export(
    "policy.trt",
    format="tensorrt",
    example_inputs=torch.zeros(1, 230, device="cuda"),
    enabled_precisions={torch.float16},  # FP16 inference
)

Custom agents

Inherit ExportableAgentMixin to give a custom agent (or your own nn.Module) the .export() method:

from simulusion import ExportableAgentMixin


class MyAgent(ExportableAgentMixin):
    def __init__(self):
        self.policy = torch.nn.Linear(64, 5)
        self.encoder = torch.nn.Linear(1024, 64)

The mixin will automatically discover self.policy and prefer it as the default target. Pass target="encoder" to export a different submodule.

Custom example input inference

If the auto-inferred example inputs are wrong for your agent, pass them explicitly:

agent.export(
    "policy.onnx",
    format="onnx",
    example_inputs=torch.zeros(1, 128),  # your custom shape
)

Or add inference support by implementing a matching pattern in _infer_example_inputs in simulusion/export.py.

Low-level API

For scripting or batch export, use the standalone functions directly:

from simulusion import export_any, export_model

# export_any resolves the target module from any object
export_any(agent, "policy.onnx", format="onnx")

# export_model exports a raw nn.Module
export_model(agent.policy, "policy.pt", format="torchscript")

Common pitfalls

1. Missing .eval(): Export always sets the module to eval mode before tracing and restores the original mode afterwards. Call .eval() manually if you are inspecting the exported graph afterwards.

2. Dynamic control flow with ONNX: ONNX requires tracing. If your module has if statements or loops that depend on tensor values, use TorchScript instead.

3. CUDA tensors for TensorRT: TensorRT export requires example inputs on the same device as the module. Pass CUDA tensors as example inputs.

4. Multiple matches for short names: If you see "matched multiple modules", use the fully qualified path, e.g. target="dreamer.actor" instead of target="actor".

See Also

• Inference — running exported models in production
• Public API — export_any and export_model factory functions