fix test

automl · Aug 10, 2023 · 8edacbe · 8edacbe
1 parent 6cffcfc
commit 8edacbe
Show file tree

Hide file tree

Showing 7 changed files with 28 additions and 11 deletions.
diff --git a/carl/__init__.py b/carl/__init__.py
@@ -6,7 +6,7 @@
 import datetime
 
 name = "CARL"
-package_name = "carl-gym"
+package_name = "carl-bench"
 author = __author__
 
 author_email = "benjamins@tnt.uni-hannover.de"

diff --git a/carl/envs/gymnasium/classic_control/carl_acrobot.py b/carl/envs/gymnasium/classic_control/carl_acrobot.py
@@ -85,8 +85,7 @@ def reset(
             size=(2,),
         )
         self.env.unwrapped.state = np.concatenate([angles, velocities])
-        return (
-            np.array(
+        state = np.array(
                 [
                     np.cos(self.env.unwrapped.state[0]),
                     np.sin(self.env.unwrapped.state[0]),
@@ -96,6 +95,8 @@ def reset(
                     self.env.unwrapped.state[3],
                 ],
                 dtype=np.float32,
-            ),
-            {},
-        )
+            )
+        info = {}
+        state = self._add_context_to_state(state)
+        info["context_id"] = self.context_id
+        return state, info
diff --git a/carl/envs/gymnasium/classic_control/carl_cartpole.py b/carl/envs/gymnasium/classic_control/carl_cartpole.py
@@ -52,4 +52,8 @@ def reset(
             high=self.context["initial_state_upper"],
             size=(4,),
         )
-        return np.array(self.env.unwrapped.state, dtype=np.float32), {}
+        state = np.array(self.env.unwrapped.state, dtype=np.float32)
+        info = {}
+        state = self._add_context_to_state(state)
+        info["context_id"] = self.context_id
+        return state, info
diff --git a/carl/envs/gymnasium/classic_control/carl_mountaincar.py b/carl/envs/gymnasium/classic_control/carl_mountaincar.py
@@ -65,4 +65,8 @@ def reset(
             high=self.context["max_velocity_start"],
         )
         self.env.unwrapped.state = np.array([position, velocity])
-        return np.array(self.env.unwrapped.state, dtype=np.float32), {}
+        state = np.array(self.env.unwrapped.state, dtype=np.float32)
+        info = {}
+        state = self._add_context_to_state(state)
+        info["context_id"] = self.context_id
+        return state, info
diff --git a/carl/envs/gymnasium/classic_control/carl_mountaincarcontinuous.py b/carl/envs/gymnasium/classic_control/carl_mountaincarcontinuous.py
@@ -62,4 +62,8 @@ def reset(
             high=self.context["max_velocity_start"],
         )
         self.env.unwrapped.state = np.array([position, velocity])
-        return np.array(self.env.unwrapped.state, dtype=np.float32), {}
+        state = np.array(self.env.unwrapped.state, dtype=np.float32)
+        info = {}
+        state = self._add_context_to_state(state)
+        info["context_id"] = self.context_id
+        return state, info
diff --git a/carl/envs/gymnasium/classic_control/carl_pendulum.py b/carl/envs/gymnasium/classic_control/carl_pendulum.py
@@ -47,4 +47,8 @@ def reset(
         theta = self.env.np_random.uniform(high=self.context["initial_angle_max"])
         thetadot = self.env.np_random.uniform(high=self.context["initial_velocity_max"])
         self.env.unwrapped.state = np.array([theta, thetadot], dtype=np.float32)
-        return np.array([np.cos(theta), np.sin(theta), thetadot], dtype=np.float32), {}
+        state = np.array([np.cos(theta), np.sin(theta), thetadot], dtype=np.float32)
+        info = {}
+        state = self._add_context_to_state(state)
+        info["context_id"] = self.context_id
+        return state, info
diff --git a/examples/vary_initial_state_distributions_with_classic_control.ipynb b/examples/vary_initial_state_distributions_with_classic_control.ipynb
@@ -173,7 +173,7 @@
     "        _states = []\n",
     "        for i in tqdm(range(n_initial_states)):\n",
     "            s, _ = env.reset()\n",
-    "            _states.append(s)\n",
+    "            _states.append(s[\"obs\"])\n",
     "            _renders.append(env.render())\n",
     "        renders[env_cls.__name__] = _renders\n",
     "        states[env_cls.__name__] = np.array(_states)\n",