protomotions.utils.export_utils module

Tracing modules for exporting the deployable pipeline to ONNX.

These torch.nn.Module wrappers turn the training-time compute into a graph torch.onnx.export can trace. deployment/export_to_onnx.py drives them and writes the typed DeploymentContract sidecar.

Modules:

• ObservationExportModule: raw context tensors -> observation tensors, by running the configured observation kernels.

• ActionExportModule: raw policy action -> processed PD targets, by running the action config’s function with its tensors baked in as buffers.

• UnifiedPipelineModule: context -> observations -> policy -> action, the three stages chained into one traceable graph.

class protomotions.utils.export_utils.ActionExportModule(action_config, device)

Bases: <Mock object at 0x7a9277331090>[]

Module that wraps an ActionConfig’s function for ONNX export.

Takes raw actions from the policy and produces processed actions with stiffness/damping targets by running the config’s action function. The config’s tensor parameters become buffers and its constants are baked in, so the traced graph carries them.

__init__(action_config, device)

get_output_keys()

forward(action)

class protomotions.utils.export_utils.UnifiedPipelineModule(

observation_module,

policy_module,

action_module,

policy_in_keys,

passthrough_keys=None,

)

Bases: <Mock object at 0x7a927758bdd0>[]

Unified module that combines observations + policy + action processing.

Pipeline: Context -> Observations -> Policy -> Action Processing

Outputs:

• actions: Raw actions from the policy

• processed_action: PD targets (clamped and transformed)

• stiffness_targets: Per-DOF stiffness values

• damping_targets: Per-DOF damping values

__init__(

observation_module,

policy_module,

action_module,

policy_in_keys,

passthrough_keys=None,

)

get_all_input_keys()

forward(*all_tensors)

class protomotions.utils.export_utils.ObservationExportModule(observation_configs, sample_context, device)

Bases: <Mock object at 0x7a927c25b610>[]

Module that wraps observation functions for ONNX export.

This module takes raw context tensors as inputs and computes observations by calling the configured observation functions. It’s designed to be exported to ONNX for deployment.

Works with MdpComponent-based observation config format:

observation_components = {
    "max_coords_obs": MdpComponent(
        compute_func=compute_humanoid_max_coords_observations,
        dynamic_vars={"body_pos": EnvContext.current.rigid_body_pos, ...},
        static_params={"local_obs": True},
    ),
}

Parameters:

• observation_configs (Dict[str, Any]) – Dict of MdpComponent instances, one per observation.

• sample_context (Dict[str, Any]) – Sample context dict to determine input shapes and resolve mappings.

• device (torch.device) – Device for tensor operations.

Example

>>> from protomotions.envs.mdp_component import MdpComponent
>>> from protomotions.envs.context_views import EnvContext
>>> configs = {
...     "max_coords_obs": MdpComponent(
...         compute_func=compute_max_coords,
...         dynamic_vars={"body_pos": EnvContext.current.rigid_body_pos},
...     )
... }
>>> context = env.context
>>> module = ObservationExportModule(configs, context, device)

__init__(

observation_configs,

sample_context,

device,

)

get_input_keys()

Get ordered list of input context keys needed.

get_output_keys()

Get ordered list of output observation names.

forward(*args)

Compute all observations from input tensors.

Parameters:

*args – Input tensors in the order of get_input_keys().

Returns:

Tuple of observation tensors in the order of get_output_keys().

Return type:

tuple