Skip to content

Conversation

@Jadeiin
Copy link

@Jadeiin Jadeiin commented Jan 24, 2026

Hi there! Thanks for keeping skrl so useful—really appreciate the work here. This PR introduces a lightweight PyTorch exporter utility so policies can be packaged for deployment as TorchScript or ONNX. The helper _TorchPolicyExporter wraps a policy together with optional observation and state preprocessors, exposing a simple forward(observations, states) -> actions interface for inference. Two user-facing entry points live in exporter.py: export_policy_as_jit to export TorchScript and export_policy_as_onnx to export ONNX, both handling tracing/export inputs, dummy fallbacks, device handling, and optional optimization toggles.

Some limitations:

  • PyTorch backend only for these exporters
  • Recurrent models (RNN/LSTM/GRU) are not supported yet
  • More documentation may be needed to introduce the exporter module.

Glad to have more discussions on this PR.

@Toni-SM
Copy link
Owner

Toni-SM commented Feb 7, 2026

Hi @Jadeiin

Such feature is absolutely welcome. I will take a look at the implementation and get back to you with feedback.

@renanmb
Copy link

renanmb commented Feb 7, 2026

I made it work for me with limitations. The major problem was figuring out those strings. In my case I wanted to use Gaussian Mixin and the exported policy did not had as u can see below an example of what happened.

I am mentioning this because my code was fairly similar to yours so did u check if the exported policy matches ? and yes, the RNN/LSTM/GRU are quite challenging

Figure_1 Figure_2

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants