Fix unwrapped of observation and action spaces

LucasAlegre · Dec 4, 2024 · 57d8fee · 57d8fee
1 parent b39f316
commit 57d8fee
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 12 deletions.
diff --git a/morl_baselines/common/morl_algorithm.py b/morl_baselines/common/morl_algorithm.py
@@ -211,18 +211,19 @@ def extract_env_info(self, env: Optional[gym.Env]) -> None:
             self.env = env
             if isinstance(self.env.observation_space, spaces.Discrete):
                 self.observation_shape = (1,)
-                self.observation_dim = self.env.unwrapped.observation_space.n
+                self.observation_dim = self.env.observation_space.n
             else:
-                self.observation_shape = self.env.unwrapped.observation_space.shape
-                self.observation_dim = self.env.unwrapped.observation_space.shape[0]
+                self.observation_shape = self.env.observation_space.shape
+                self.observation_dim = self.env.observation_space.shape[0]
 
-            self.action_space = env.unwrapped.action_space
-            if isinstance(self.env.unwrapped.action_space, (spaces.Discrete, spaces.MultiBinary)):
+            self.action_space = env.action_space
+            if isinstance(self.env.action_space, (spaces.Discrete, spaces.MultiBinary)):
                 self.action_shape = (1,)
-                self.action_dim = self.env.unwrapped.action_space.n
+                self.action_dim = self.env.action_space.n
             else:
-                self.action_shape = self.env.unwrapped.action_space.shape
-                self.action_dim = self.env.unwrapped.action_space.shape[0]
+                self.action_shape = self.env.action_space.shape
+                self.action_dim = self.env.action_space.shape[0]
+
             self.reward_dim = self.env.unwrapped.reward_space.shape[0]
 
     @abstractmethod

diff --git a/tests/test_algos.py b/tests/test_algos.py
@@ -196,8 +196,8 @@ def test_gpi_pd():
 
 
 def test_gpi_pd_continuous_action():
-    env = mo_gym.make("mo-hopper-v4", cost_objective=False, max_episode_steps=500)
-    eval_env = mo_gym.make("mo-hopper-v4", cost_objective=False, max_episode_steps=500)
+    env = mo_gym.make("mo-hopper-v5", cost_objective=False, max_episode_steps=500)
+    eval_env = mo_gym.make("mo-hopper-v5", cost_objective=False, max_episode_steps=500)
 
     agent = GPIPDContinuousAction(
         env,
@@ -278,8 +278,8 @@ def test_pcn():
 
 
 def test_capql():
-    env = mo_gym.make("mo-hopper-v4", cost_objective=False, max_episode_steps=500)
-    eval_env = mo_gym.make("mo-hopper-v4", cost_objective=False, max_episode_steps=500)
+    env = mo_gym.make("mo-hopper-v5", cost_objective=False, max_episode_steps=500)
+    eval_env = mo_gym.make("mo-hopper-v5", cost_objective=False, max_episode_steps=500)
 
     agent = CAPQL(
         env,