Update stable_baselines_export.py

Adds error handling for case of exporting a SAC policy without using single obs (currently not supported).
edbeeching · Aug 22, 2024 · 95cb1ad · 95cb1ad
1 parent 1df043d
commit 95cb1ad
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/godot_rl/wrappers/onnx/stable_baselines_export.py b/godot_rl/wrappers/onnx/stable_baselines_export.py
@@ -52,6 +52,9 @@ def export_model_as_onnx(model, onnx_model_path: str, use_obs_array: bool = Fals
     dummy_input = None
     onnxable_model = None
 
+    if isinstance(model, SAC):
+        assert use_obs_array, "SAC ONNX export works with use_obs_array=True, MLPPolicy and SBGSingleObsEnv only."
+
     if isinstance(model, PPO):
         onnxable_model = OnnxablePolicy(
             ["obs"],