Core Exporter Functionality

This section contains detailed API documentation for the core exporter functionality of exploy.

Actor

Abstract interface and utilities for exportable actors.

This module defines the ExportableActor base class for policy networks that can be exported to ONNX, along with helpers like make_exportable_actor and add_actor_memory.

class exploy.exporter.core.actor.ExportableActor

Bases: Module, ABC

Abstract interface for an actor that can be exported to ONNX.

__init__()

Initialize internal Module state, shared by both nn.Module and ScriptModule.

abstractmethod forward(obs)

Given a batch of observations, compute the corresponding actions.

Parameters:

obs (Tensor) – A tensor of shape (batch_size, obs_dim) containing the observations.

Return type:

Tensor

reset(dones)

Reset the actor’s internal state (e.g., RNN hidden states) based on the done flags.

Parameters:

dones (Tensor) – A tensor of shape (batch_size,) containing boolean flags indicating which environments have been reset.

get_state()

Get the actor’s internal state as a tuple of tensors, or None if there is no state.

Return type:

tuple[Tensor, ...] | None

exploy.exporter.core.actor.make_exportable_actor(actor)

Convert a torch.nn.Module actor to an ExportableActor.

Parameters:

actor (Module) – The actor to convert.

Return type:

ExportableActor

exploy.exporter.core.actor.add_actor_memory(context_manager, get_hidden_states_func)

Add inputs for actor hidden states.

Parameters:

• context_manager (ContextManager) – The context manager to add the inputs to.

• get_hidden_states_func (Callable[[], tuple[Tensor, ...]]) – A function that returns a tuple of hidden state tensors, used to get the hidden states to add as inputs.

Exportable Environment

Abstract base class for exportable environments.

This module defines the ExportableEnvironment base class that provides the standardized interface required by the exporter to trace observation computation, action processing, and simulation stepping.

class exploy.exporter.core.exportable_environment.ExportableEnvironment

Bases: ABC

__init__()

context_manager()

Return type:

ContextManager

abstractmethod compute_observations()

Compute and return the observations of the environment.

Return type:

Tensor

abstractmethod process_actions(actions)

Process actions.

abstractmethod apply_actions()

Apply processed actions (e.g., joint targets) to the environment

abstractmethod prepare_export()

Prepare the environment for export. Called before each export.

abstractmethod empty_actor_observations()

Return type:

Tensor

abstractmethod empty_actions()

Return type:

Tensor

abstractmethod metadata()

Return metadata about the environment required for export.

Return type:

dict[str, str]

abstract property decimation: int

Return metadata about the environment required for export.

abstractmethod register_evaluation_hooks(update, evaluate_substep)

Register evaluation hooks for this environment.

abstractmethod step(actions)

Step the environment forward by one step. Returns the next observations and a boolean indicating if the environment was reset.

Return type:

tuple[Tensor, bool]

abstractmethod get_observation_names()

Get the names of the observations in the environment.

Return type:

list[str]

abstractmethod observations_reset()

Get the observations after an environment reset.

Return type:

Tensor

register_command_update(command_update)

Register callable to update the commands in the environment before observations are computed.

Parameters:

command_update (Callable[[], None]) – A callable that updates the commands.

property command_updates: list[Callable[[], None]]

Get the registered command updates.

Components

Core building blocks for environment export.

This module defines the component abstractions (Input, Output, Memory, Group, Connection) used to structure and manage data flow during policy export.

class exploy.exporter.core.components.Input(name, get_from_env_cb, metadata=None)

Bases: object

Abstraction for controller inputs.

The input encapsulates a callback that retrieves data from the environment and provides it as a tensor for use in the generation of a computational graph.

__init__(name, get_from_env_cb, metadata=None)

Constuct an Input.

Parameters:

• name (str) – Identifier for this input.

• get_from_env_cb (Callable[[], torch.Tensor]) – Callback function that retrieves the input from the environment as a torch.Tensor.

• metadata (Any) – Optional metadata associated with this input (e.g., shape, data type, semantic information).

property data: Tensor

Get internal data as a torch tensor.

property data_numpy: ndarray

Get internal data as a numpy array.

property metadata: Any

Return metadata about this handler’s data

property name: str

Return the name of this input.

property id: int

Return the unique identifier of this input.

read()

Get the latest data from the environment by calling the callback.

Return type:

None

property get_from_env_cb: Callable[[], Tensor]

Get the callback function that retrieves the input from the environment.

class exploy.exporter.core.components.Output(name, get_from_env_cb, metadata=None)

Bases: object

Abstract interface for outputs.

__init__(name, get_from_env_cb, metadata=None)

property metadata: Any

Get metadata about this output.

property name: str

Return the name of this output.

property value: Tensor

Get the latest value from the environment by calling the callback.

property value_numpy: ndarray

Get the latest value from the environment as a numpy array by calling the callback.

property get_from_env_cb: Callable[[], Tensor]

Get the callback function that retrieves the output from the environment.

class exploy.exporter.core.components.Memory(name, get_from_env_cb, metadata=None)

Bases: object

Handle memory inputs and outputs.

This class abstracts how to get and set values used in an environment that has memory, for example actions and previous actions. Values are retrieved by passing callables.

A memory element pairs an Input and an Output that share the same environment callback: the output of one inference step is fed back as the input of the next one.

__init__(name, get_from_env_cb, metadata=None)

Construct a Memory element.

Parameters:

• name (str) – Identifier for this memory element.

• get_from_env_cb (Callable[[], torch.Tensor]) – Callback function that retrieves the memory value from the environment as a torch.Tensor.

• metadata (Any) – Optional metadata associated with this memory element. It is stored on the memory element itself, under the name shared by its input and output.

property name: str

Return the name shared by this memory element’s input and output.

property metadata: Any

Get metadata about this memory element.

property input: Input

Return the input side of this memory element.

property output: Output

Return the output side of this memory element.

class exploy.exporter.core.components.Group(name, items, metadata=None)

Bases: object

Abstraction for grouping related inputs and outputs together.

__init__(name, items, metadata=None)

property name: str

Return the name of this group.

property items: list[Input | Output]

Get the items in this group.

property metadata: Any

Get metadata about this group.

class exploy.exporter.core.components.Connection(name, getter, setter)

Bases: object

Abstraction for connecting existing inputs to data sources.

__init__(name, getter, setter)

property name: str

The name of this connection.

write()

Write data to the environment by calling the setter with the value from the getter.

Return type:

None

Context Manager

Manages components for environment export.

The ContextManager class organizes and manages all inputs, outputs, memory components, and connections needed during the ONNX export process.

class exploy.exporter.core.context_manager.ContextManager

Bases: object

Manages all components (inputs, outputs, memory) for an environment export.

__init__()

add_component(component)

Add a component (Input, Output, Memory, or Connection) to the context manager.

A Memory is registered as its input and output sides, so it is picked up by both get_input_components() and get_output_components().

Parameters:

component (Input | Output | Memory | Connection) – The component to add. Can be an Input, Output, Memory, or Connection.

Raises:

AssertionError – If a component with the same name already exists.

Return type:

None

add_components(components)

Add multiple components to the context manager.

Parameters:

components (list[Input | Output | Memory | Connection]) – A list of components to add. Each can be an Input, Output, Memory, or Connection.

Raises:

AssertionError – If any component has a name that already exists.

Return type:

None

add_group(group)

Add a group of components to the context manager.

Recursively adds all items in the group, including nested groups.

Parameters:

group (Group) – The group to add containing inputs, outputs, or nested groups.

Raises:

AssertionError – If a group or component with the same name already exists.

Return type:

None

get_input_components()

Get all input components including the input side of memory components.

Return type:

list[Input]

Returns:

A list of all Input components.

get_connection_components()

Get all connection components.

Return type:

list[Connection]

Returns:

A list of all Connection components.

get_output_components()

Get all output components including the output side of memory components.

Return type:

list[Output]

Returns:

A list of all Output components.

get_memory_components()

Get all memory components.

Return type:

list[Memory]

Returns:

A list of all Memory components.

get_component_by_name(name)

Get a component by its name.

Parameters:

name (str) – The name of the component to find.

Return type:

Input | Output | None

Returns:

The component with the given name, or None if not found.

read_inputs()

Read and update all input components from the environment.

Calls the read() method on each input component to refresh their data.

Return type:

None

write_connections()

Write all connection components to transfer data.

Calls the write() method on each connection component to transfer data from getters to setters.

Return type:

None

get_inputs(to_numpy=False)

Get all input data as a dictionary.

Parameters:

to_numpy (bool) – If True, return data as numpy arrays. If False, return as torch tensors.

Return type:

dict[str, Tensor] | dict[str, ndarray]

Returns:

Dictionary mapping input names to their data (as torch.Tensor or np.ndarray).

get_input_names()

Get the names of all input components.

Return type:

list[str]

Returns:

A list of input component names.

get_outputs(to_numpy=False)

Get all output data as a dictionary.

Parameters:

to_numpy (bool) – If True, return data as numpy arrays. If False, return as torch tensors.

Return type:

dict[str, Tensor] | dict[str, ndarray]

Returns:

Dictionary mapping output names to their values (as torch.Tensor or np.ndarray).

get_output_names()

Get the names of all output components.

Return type:

list[str]

Returns:

A list of output component names.

property metadata: dict[str, Any]

assert_unique_name(name)

Assert that the given name is unique across all components and groups.

Parameters:

name (str) – The name to check for uniqueness.

Raises:

KeyError – If the name already exists in groups, inputs, or outputs.

assert_unique_id(id, name=None)

Assert that the given id is unique across all components.

Parameters:

id (int) – The id to check for uniqueness.

Raises:

KeyError – If the id already exists in any component.

add_module(module)

Register a PyTorch module with this context manager.

This helper allows adding pretrained models as components of the context manager so they can be treated as submodules of the exporter and properly included when exporting to ONNX.

Parameters:

module (Module) – The PyTorch module to register.

Return type:

None

property modules: tuple[Module, ...]

Get all registered modules.

Returns:

A tuple of all registered torch modules.

Evaluator

Validation and testing utilities for exported ONNX models.

This module provides the evaluate function to compare ONNX model outputs against the original environment and PyTorch policy for correctness verification.

exploy.exporter.core.evaluator.evaluate(env, context_manager, session_wrapper, num_episodes, max_episode_steps, verbose=True, reset_from_onnx_counter_steps=50, atol=1e-05, rtol=1e-05, pause_on_failure=True)

Evaluate an ONNX exported model against an ExportableEnvironment stepped through a SessionWrapper.

This function runs the simulation for a specified number of episodes, each with a maximum number of steps, and compares the outputs of the ONNX model with the environment’s state and actor’s actions at each step. This is useful for verifying the correctness of the ONNX export.

Parameters:

• env (ExportableEnvironment) – The environment to run the evaluation in.

• context_manager (ContextManager) – The context manager handling inputs and outputs.

• session_wrapper (SessionWrapper) – An ONNX session wrapper.

• num_episodes (int) – The number of episodes to run the evaluation for.

• max_episode_steps (int) – The maximum number of steps per episode.

• verbose (bool) – Whether to print verbose output during evaluation. Defaults to True.

• reset_from_onnx_counter_steps (int) –

  Set after how many steps we should set memory inputs from ONNX instead of using the environment’s state.

  Note: we do this to avoid numerical error accumulation that would occur if we only every use the ONNX inference outputs as memory fed back as ONNX inference inputs, while all other inputs are set directly from the environment’s state.

  Note: this value is chosen arbitrarily.

• atol (float) – Absolute tolerance used to compare tensors.

• rtol (float) – Relative tolerance used to compare tensors.

• pause_on_failure (bool) – Whether to pause on each failed step and wait for user input before continuing. Defaults to True.

Return type:

tuple[bool, Tensor]

Returns:

A tuple containing a boolean indicating if the evaluation was successful and the final observations tensor.

exploy.exporter.core.evaluator.evaluate_episode(env, context_manager, session_wrapper, max_num_steps, verbose=True, reset_from_onnx_counter_steps=50, atol=1e-05, rtol=1e-05, pause_on_failure=True)

Run evaluation for a single episode, comparing the ONNX model outputs against the environment.

Steps the environment and the ONNX session in lockstep, comparing observations, actions, and outputs at each step. Useful for fine-grained inspection of a single episode, e.g. in a debugging loop, without the outer episode iteration of evaluate().

Parameters:

• env (ExportableEnvironment) – The environment to run the evaluation in.

• context_manager (ContextManager) – The context manager handling inputs and outputs.

• session_wrapper (SessionWrapper) – An ONNX session wrapper.

• max_num_steps (int) – The maximum number of steps to run.

• verbose (bool) – Whether to print verbose output during evaluation. Defaults to True.

• reset_from_onnx_counter_steps (int) –

  Set after how many steps we should set memory inputs from ONNX instead of using the environment’s state.

  Note: we do this to avoid numerical error accumulation that would occur if we only ever use the ONNX inference outputs as memory fed back as ONNX inference inputs, while all other inputs are set directly from the environment’s state.

  Note: this value is chosen arbitrarily.

• atol (float) – Absolute tolerance used to compare tensors.

• rtol (float) – Relative tolerance used to compare tensors.

• pause_on_failure (bool) – Whether to pause on each failed step and wait for user input before continuing. Defaults to True.

Returns:

A tuple containing a boolean indicating if the episode evaluation was successful and the final observations tensor.

Exporter Module

Core ONNX export functionality for RL policies.

This module provides the main entry point for exporting trained policies to ONNX format, including the export_environment_as_onnx function and the OnnxEnvironmentExporter class.

exploy.exporter.core.exporter.export_environment_as_onnx(env, actor, path, filename='policy.onnx', model_source=None, verbose=False, opset_version=20, ir_version=11)

Export policy into a Torch ONNX file.

Parameters:

• env (ExportableEnvironment) – The environment to be exported.

• actor (Module) – The actor torch module.

• path (str) – The path to the saving directory.

• filename (str) – The name of exported ONNX file. Defaults to “policy.onnx”.

• model_source (dict | None) – Information about the policy’s origin (e.g., wandb, local file, etc.), added to the ONNX metadata.

• verbose (bool) – Whether to print the model summary. Defaults to False.

• opset_version (int) – Version of the operator specification referenced by the ONNX graph. Needs to be compatible with ONNX Runtime in deployment environment, check https://onnxruntime.ai/docs/reference/compatibility.html

• ir_version (int) – Version of the intermediate representation specifications. Needs to be compatible with ONNX Runtime in deployment environment, check https://onnxruntime.ai/docs/reference/compatibility.html

class exploy.exporter.core.exporter.ExportMode(*values)

Bases: Enum

Default = 0

ProcessActions = 1

class exploy.exporter.core.exporter.OnnxEnvironmentExporter(env, actor, opset_version, ir_version, verbose=False)

Bases: Module

Exporter of actor-critic into ONNX file using the environment’s managers.

__init__(env, actor, opset_version, ir_version, verbose=False)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(input_data)

Use the robot’s state to compute policy actions, joint position targets, and policy

observations, and outputs that support history.

Parameters:

input_data (dict[str, Tensor]) – A dictionary containing all input tensors required for the forward pass. The expected keys and shapes of the tensors depend on the environment’s context manager and the policy’s computational graph.

Notes

• Dictionary inputs are flattened by the torch ONNX exporter implementation.

• Only inputs that are part of the computational graph will be required when using the resulting ONNX file for inference. For example, if pos_base_in_w is not used by any of the observation functions, it will not be a required input. This can be verified by querying the ONNX input names when using the ONNX runtime framework.

Returns:

Environment outputs. A tuple of desired joint positions (i.e., processed actions), actions (i.e., unprocessed actions), memory (containing the previous actions for example).

Return type:

tuple[torch.Tensor, …]

register_modules()

Register all modules from the environment’s context manager.

This method iterates over all modules in the environment’s context manager and registers them using sequential names. This ensures that all relevant modules are included in the ONNX export, allowing the exported model to function correctly when loaded in an ONNX runtime environment.

Calling this method multiple times will not re-register already registered modules. The modules’ names are based on insertion order.

export(onnx_path, onnx_file_name, model_source)

Export to ONNX.

Parameters:

• onnx_path (str) – The path to the folder that will contain the ONNX file.

• onnx_file_name (str) – The name (including the ONNX extension) of the exported file.

• model_source (dict) – Information about the policy’s origin (e.g., wandb, local file, etc.), added to the ONNX metadata.

Session Wrapper

ONNX Runtime inference session management.

The SessionWrapper class provides a convenient interface for loading and running ONNX models with ONNX Runtime, managing input/output handling and session configuration.

class exploy.exporter.core.session_wrapper.SessionWrapper(onnx_folder, onnx_file_name, actor=None, optimize=True)

Bases: object

Manage a torch Module and its associated ONNX inference session.

__init__(onnx_folder, onnx_file_name, actor=None, optimize=True)

Construct a SessionWrapper to use it for policy inference.

Parameters:

• onnx_folder (Path) – The folder containing an ONNX file to load.

• onnx_file_name (str) – The name of the ONNX file contained in ONNX_folder.

• actor (ExportableActor | None) – An ExportableActor representing the actor.

• optimize (bool) – If true, optimize the ONNX graph, save it to file, and use it for inference.

property onnx_file_path: Path

__call__(**kwargs)

Run ONNX inference with the given inputs.

Parameters:

**kwargs – Keyword arguments where keys are input names and values are input data.

Returns:

List of output arrays from the ONNX model inference.

get_actor()

Get the original ExportableActor object used by this session wrapper.

Return type:

ExportableActor | None

Returns:

The ExportableActor representing the actor, or None if not provided.

get_output_value(output_name)

Get a specific output value from the last inference run.

Parameters:

output_name (str) – The name of the output to retrieve.

Returns:

The numpy array corresponding to the requested output.

Raises:

KeyError – If the output_name is not in the model’s outputs.

reset()

Reset the internal results to zeros to avoid stale data at environment reset.

Tensor Proxy

Tensor list abstraction for improved ONNX export.

The TensorProxy class manages lists of tensors and exposes them as a single stacked tensor, improving the structure of exported computational graphs.

class exploy.exporter.core.tensor_proxy.TensorProxy(tensor, split_dim)

Bases: object

Manage a list of tensors and expose them to the user as a stacked tensor.

This class takes a tensor, splits it along a specified dimension, and exposes it as if it were the original tensor. For example, this class allows the user to implement the following:

# Make a tensor of all body positions. body_pos = torch.rand((batch_dim, num_bodies, 3))

# Wrap the tensor in a TensorProxy object, splitting along dimension 1 (num_bodies). body_pos_proxy = TensorProxy(body_pos, split_dim=1)

Now, the user can index into body_pos_proxy as if it was the original body_pos. Indexing will index into one of the split tensors created from unbinding along the split dimension.

One use case for this implementation is to split the body_state_w tensor from ArticulationData into separate body tensors to improve exporting.

__init__(tensor, split_dim)

__getitem__(idx)

Index into a TensorProxy as if the user was indexing into the un-split list of tensors.

__setitem__(idx, value)

Set into a TensorProxy as if the user was indexing into the un-split list of tensors.

to_tensor()

Convert the data stored in a TensorProxy into a torch.Tensor.

Return type:

Tensor

classmethod __torch_function__(func, types, args=(), kwargs=None)

Allow using this class with the torch API.

Torch provides a mechanism to treat any object as a torch.Tensor. This is enabled by implementing the __torch_function__ method for any python class.

For more details on how to implement __torch_function__, see:

https://docs.pytorch.org/docs/stable/notes/extending.html#extending-torch-python-api

For a discussion on a concrete implementation of __torch_function__, see:

https://github.com/docarray/notes/blob/main/blog/02-this-weeks-in-docarray-01.md#__torch_function__–or-how-to-give-pytorch-a-little-bit-more-confidence

property tensors: list[Tensor]

Get the list of tensors stored in this TensorProxy.

exploy.exporter.core.tensor_proxy.to_tensor(value)

Convert a TensorProxy to a torch.Tensor; pass plain tensors through unchanged.

Useful before calling functions that use tensor methods internally (e.g. framework math utilities), which TensorProxy.__torch_function__ cannot intercept.

Return type:

Tensor

exploy.exporter.core.tensor_proxy.args_to_tensor(args)

Convert each element in a sequence to torch.Tensor, preserving the sequence structure.

Derived Tensors

Materialization of derived quantities as ONNX-traceable leaf tensors.

The DerivedTensor class bundles a derived quantity’s compute function and split dimension. Materializing it against a data-source proxy stores the computed value as a leaf tensor so it can become an ONNX graph input; calling it returns the stored leaf during export and computes the quantity live otherwise.

class exploy.exporter.core.derived_tensors.DerivedTensor(compute_fn, split_dim=None)

Bases: object

A derived quantity that can be materialized as an ONNX-traceable leaf tensor.

The ONNX exporter wires graph inputs by object identity: an Input becomes a graph input only if the exact tensor object it holds is read during tracing. A quantity that is computed on access produces a fresh tensor per call and can never become a graph input. Materializing a derived quantity computes it once and stores the clone, so subsequent reads return the stored object and inputs bound to it (or to its per-body slices when stored as a TensorProxy) are traced as ONNX graph inputs, cutting the graph at the derived quantity.

Materialized values are stored inside this object, keyed by the identity of all arguments they were computed from (typically a single data-source proxy). This keeps the data-source proxies pure framework adaptors (no dynamic attributes) and gives the stored values the same lifecycle as the export data-source swap: materialize with the data-source proxy during export, and once the proxy is discarded (after cleanup() restores the framework’s native data) calls fall through to computing the quantity live.

Bundling the compute function and split dimension in one object guarantees that materialization (during export) and live computation (during training or evaluation) can never use mismatched functions.

Typical usage: define the quantity once at module level, materialize it on the data-source proxy when registering inputs, and call it wherever the value is needed:

body_link_lin_vel_b = DerivedTensor(compute_fn=..., split_dim=1)

# Input registration (export script):
body_link_lin_vel_b.materialize(entity.data)
Input(..., get_from_env_cb=lambda: body_link_lin_vel_b(entity.data)[:, idx])

# Observation term (training and export):
def obs_term(env, asset_cfg):
    return body_link_lin_vel_b(env.scene[asset_cfg.name].data)[:, asset_cfg.body_ids]

The compute function may take any signature — including that of a class-based observation term — since materialize() and __call__() forward *args and **kwargs to it unchanged.

__init__(compute_fn, split_dim=None)

Construct a DerivedTensor.

Parameters:

• compute_fn (Callable[..., Tensor]) – Callback that computes the quantity. Any callable works — a free function, a bound method, or a callable class instance; materialize and __call__ forward their arguments to it unchanged.

• split_dim (int | None) – If given, materialized values are wrapped in a TensorProxy split along this dimension so that each slice is an independent leaf (e.g. one ONNX input per body).

materialize(*args, **kwargs)

Compute this quantity once and store it as a leaf tensor.

Subsequent calls of this DerivedTensor with the same arguments return the stored value, preserving the object identity required for ONNX tracing. The stored value lives as long as the arguments do. Materializing the same arguments again returns the already stored value without recomputing.

Only materialize export data-source proxies: their state is a frozen snapshot, so the stored value stays consistent with them. Materializing live framework data would freeze the quantity while the underlying state keeps changing.

Parameters:

• *args – Positional arguments forwarded to the compute function, typically an exploy data-source proxy.

• **kwargs – Keyword arguments forwarded to the compute function.

Return type:

Tensor | TensorProxy

Returns:

The stored tensor, or TensorProxy if split_dim was given.

__call__(*args, **kwargs)

Return the value materialized for these arguments if present, else compute it live.

This lets the same accessor work both during export, where the quantity has been materialized on the data-source proxy (returning the stored leaf preserves the object identity required for ONNX tracing), and outside of export (e.g. evaluation or training), where the quantity is computed from live state instead.

Parameters:

• *args – Positional arguments forwarded to the compute function, typically an exploy data-source proxy or native framework data.

• **kwargs – Keyword arguments forwarded to the compute function.

Return type:

Tensor | TensorProxy

Returns:

The stored tensor/TensorProxy if materialized for these arguments, else compute_fn(*args, **kwargs).