From 95cb1ad6b7f9a0634bd72a85acd3f7cee8633046 Mon Sep 17 00:00:00 2001
From: Ivan-267 <61947090+Ivan-267@users.noreply.github.com>
Date: Thu, 22 Aug 2024 21:21:25 +0200
Subject: [PATCH] Update stable_baselines_export.py

Adds error handling for case of exporting a SAC policy without using single obs (currently not supported).
---
 godot_rl/wrappers/onnx/stable_baselines_export.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/godot_rl/wrappers/onnx/stable_baselines_export.py b/godot_rl/wrappers/onnx/stable_baselines_export.py
index 398e4f74..72c36863 100644
--- 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"],