Exporting Models for Deployment #

TorchWM provides a unified export system that converts trained models and agents into deployable formats — ONNX, TorchScript, and TensorRT — without requiring each model to implement its own export logic.

Overview #

The export system is built around three levels of API:

Level	Function / Method	When to use
Module method	`module.export(path, format, ...)`	Any `torch.nn.Module` — works automatically after importing `world_models.export`
Agent method	`agent.export(path, format, ...)`	High-level agents (`DreamerAgent`, `JEPAAgent`, `IRISAgent`) that inherit `ExportableAgentMixin`
Standalone	`export_any(obj, path, ...)` / `export_model(module, path, ...)`	When you need explicit control over which submodule is exported or want to bypass automatic target resolution

Supported formats #

Format	Extension	Use case
`"onnx"`	`.onnx`	Cross-platform inference, mobile, edge devices, TensorRT conversion
`"torchscript"` (aliases: `"jit"`, `"ts"`, `"pt"`, `"script"`)	`.pt`	Serving via LibTorch (C++), no Python dependency at inference time
`"tensorrt"` (aliases: `"trt"`)	`.pt`	NVIDIA GPU-optimized inference (requires `torch_tensorrt` package)

Quick start #

Exporting any `nn.Module`#

Importing world_models.export installs the .export() method on every torch.nn.Module instance once:

import torch
import world_models.export  # installs nn.Module.export

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()

# TorchScript (scripted — works without example_inputs)
model.export("model.pt", format="torchscript")

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

Exporting a trained agent #

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

import torchwm

agent = torchwm.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:

Agent	Default target	Example inputs
`DreamerAgent`	Policy head (actor)	`[batch, stoch_size + deter_size]`
`IRISAgent`	Actor-critic head	`[batch, 1, channels, h, w]`
`JEPAAgent`	Vision Transformer encoder	`[batch, 3, crop_size, crop_size]`

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:

Mode	When to use	Limitation
Tracing (with `example_inputs`)	Fast export of a fixed forward graph	May not handle dynamic control flow
Scripting (without `example_inputs`)	Full module graph with control flow	May fail on unsupported Python constructs

# 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 #

If you build a custom agent that is not an nn.Module, inherit ExportableAgentMixin to get the .export() method:

from torchwm 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 world_models/export.py.

Low-level API #

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

from torchwm 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 #

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.
Dynamic control flow with ONNX: ONNX requires tracing. If your module has if statements or loops that depend on tensor values, use TorchScript instead.
CUDA tensors for TensorRT: TensorRT export requires example inputs on the same device as the module. Pass CUDA tensors as example inputs.
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".