[RLlib] ONNX Export support when using RLModule API

ray-project / ray

Ray is a unified framework for scaling AI and Python applications. Ray consists of a core distributed runtime and a set of AI Libraries for accelerating ML workloads.

https://ray.io

Apache License 2.0

33.29k stars 5.63k forks source link

[RLlib] ONNX Export support when using RLModule API #45526

Open H-Park opened 4 months ago

H-Park commented 4 months ago

Description

When training using the RLModule API , policy.export is unable to export to ONNX, citing

ValueError: ONNX export not supported for RLModule API.

Use case

As a user of Ray RL Lib, I want to be able to train using the RLModule API and then export to onnx.

simonsays1980 commented 4 months ago

@H-Park This has no high priority for us right now. As all RLModules are nn.Module s can you try to export directly via torch?