Update stable_baselines3_example.py

Sets to use the new onnx export method with support for SAC and PPO.
edbeeching · Aug 22, 2024 · 7806f26 · 7806f26
1 parent 1f8b277
commit 7806f26
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/stable_baselines3_example.py b/examples/stable_baselines3_example.py
@@ -8,7 +8,7 @@
 from stable_baselines3.common.vec_env.vec_monitor import VecMonitor
 
 from godot_rl.core.utils import can_import
-from godot_rl.wrappers.onnx.stable_baselines_export import export_ppo_model_as_onnx
+from godot_rl.wrappers.onnx.stable_baselines_export import export_model_as_onnx
 from godot_rl.wrappers.stable_baselines_wrapper import StableBaselinesGodotEnv
 
 # To download the env source and binary:
@@ -115,7 +115,7 @@ def handle_onnx_export():
     if args.onnx_export_path is not None:
         path_onnx = pathlib.Path(args.onnx_export_path).with_suffix(".onnx")
         print("Exporting onnx to: " + os.path.abspath(path_onnx))
-        export_ppo_model_as_onnx(model, str(path_onnx))
+        export_model_as_onnx(model, str(path_onnx), True)
 
 
 def handle_model_save():