From 60a8c81745a3fa80abb6c4cf62f198dc15ff4ce3 Mon Sep 17 00:00:00 2001
From: Benjamin Piwowarski <benjamin@piwowarski.fr>
Date: Mon, 20 Nov 2023 15:38:05 +0100
Subject: [PATCH] Fix so that flatten works

---
 src/pystk2_gymnasium/envs.py     | 2 +-
 src/pystk2_gymnasium/wrappers.py | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/pystk2_gymnasium/envs.py b/src/pystk2_gymnasium/envs.py
index aa710b2..7ef9628 100644
--- a/src/pystk2_gymnasium/envs.py
+++ b/src/pystk2_gymnasium/envs.py
@@ -188,7 +188,7 @@ def reset_race(
         self.current_track = self.default_track
         if self.current_track is None:
             self.current_track = self.TRACKS[random.randint(0, len(self.TRACKS))]
-            logging.info("Selected %s", self.current_track)
+            logging.debug("Selected %s", self.current_track)
         self.config = pystk2.RaceConfig(
             num_kart=self.num_kart,
             seed=random.randint(2**16),
diff --git a/src/pystk2_gymnasium/wrappers.py b/src/pystk2_gymnasium/wrappers.py
index 5243ee6..efac909 100644
--- a/src/pystk2_gymnasium/wrappers.py
+++ b/src/pystk2_gymnasium/wrappers.py
@@ -31,7 +31,8 @@ def __init__(self, space: gym.Space):
         highs = []
         counts = []
 
-        for key, value in space.items():
+        # Combine keys (sort them before hand)
+        for key, value in sorted(space.items(), key=lambda x: x[0]):
             # Ignore the AI action
             if key == "action":
                 continue