diff --git a/loco_mujoco/environments/gymnasium.py b/loco_mujoco/environments/gymnasium.py
index a93b311..0472e00 100644
--- a/loco_mujoco/environments/gymnasium.py
+++ b/loco_mujoco/environments/gymnasium.py
@@ -41,10 +41,9 @@ def __init__(self, env_name, render_mode=None, **kwargs):
 
         self.metadata["render_fps"] = 1.0 / self._env.dt
 
-        self.observation_space = self._convert_space(self._env.info.observation_space)
+        self._set_observation_space()
         self._set_action_space()
 
-
     def step(self, action):
         """
         Run one timestep of the environment's dynamics.
@@ -151,6 +150,13 @@ def unwrapped(self):
         """
         return self._env
 
+    def _set_observation_space(self):
+        """
+        Setter for the observation space.
+        """
+        self.observation_space = self._convert_space(self._env.info.observation_space)
+        return self.observation_space
+
     def _set_action_space(self):
         """ 
         Setter for the action space.