Add Enum observation space unflattening (from one hot encoding)

iksnagreb · Oct 4, 2021 · 8a4c607 · 8a4c607
1 parent 784937e
commit 8a4c607
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/gym_csgo/spaces/enum.py b/gym_csgo/spaces/enum.py
@@ -41,3 +41,9 @@ def flatten_enum(space, x):
     onehot[space.values.tolist().index(x)] = 1
     # Return one hot encoded flat enum space
     return onehot
+
+# Reverse the flatten_enum operation
+@gym.spaces.unflatten.register(Enum)
+def unflatten_enum(space, x):
+    # Lookup enum value at encoded position
+    return space.values[np.nonzero(x)[0][0]]