diff --git a/magbotsim/rl_envs/object_manipulation/pushing/state_based_global_pushing_env.py b/magbotsim/rl_envs/object_manipulation/pushing/state_based_global_pushing_env.py
index e7beee1..039e65c 100644
--- a/magbotsim/rl_envs/object_manipulation/pushing/state_based_global_pushing_env.py
+++ b/magbotsim/rl_envs/object_manipulation/pushing/state_based_global_pushing_env.py
@@ -1025,7 +1025,7 @@ def _geoms(self, name: str) -> np.ndarray:
                     break
 
                 geom = self.model.geom(geom_id)
-                geoms[geom_idx, 0] = geom.type
+                geoms[geom_idx, 0] = geom.type.item()
                 geoms[geom_idx, 1:3] = geom.size[:2]
                 geoms[geom_idx, 3:5] = geom.pos[:2]
                 geom_idx += 1