Fix CartPole for equivalence to vector implementation (#1121)

Author: pseudo-rnd-thoughts

gymnasium/envs/classic_control/cartpole.py

@@ -157,7 +157,7 @@ def __init__(
         self.screen = None
         self.clock = None
         self.isopen = True
-        self.state = None
+        self.state: np.ndarray | None = None
 
         self.steps_beyond_terminated = None
 
@@ -168,16 +168,17 @@ def step(self, action):
         assert self.state is not None, "Call reset before using step method."
         x, x_dot, theta, theta_dot = self.state
         force = self.force_mag if action == 1 else -self.force_mag
-        costheta = math.cos(theta)
-        sintheta = math.sin(theta)
+        costheta = np.cos(theta)
+        sintheta = np.sin(theta)
 
         # For the interested reader:
         # https://coneural.org/florian/papers/05_cart_pole.pdf
         temp = (
-            force + self.polemass_length * theta_dot**2 * sintheta
+            force + self.polemass_length * np.square(theta_dot) * sintheta
         ) / self.total_mass
         thetaacc = (self.gravity * sintheta - costheta * temp) / (
-            self.length * (4.0 / 3.0 - self.masspole * costheta**2 / self.total_mass)
+            self.length
+            * (4.0 / 3.0 - self.masspole * np.square(costheta) / self.total_mass)
         )
         xacc = temp - self.polemass_length * thetaacc * costheta / self.total_mass
 
@@ -192,7 +193,7 @@ def step(self, action):
             theta_dot = theta_dot + self.tau * thetaacc
             theta = theta + self.tau * theta_dot
 
-        self.state = (x, x_dot, theta, theta_dot)
+        self.state = np.array((x, x_dot, theta, theta_dot), dtype=np.float64)
 
         terminated = bool(
             x < -self.x_threshold
@@ -202,33 +203,25 @@ def step(self, action):
         )
 
         if not terminated:
-            if self._sutton_barto_reward:
-                reward = 0.0
-            elif not self._sutton_barto_reward:
-                reward = 1.0
+            reward = 0.0 if self._sutton_barto_reward else 1.0
         elif self.steps_beyond_terminated is None:
             # Pole just fell!
             self.steps_beyond_terminated = 0
-            if self._sutton_barto_reward:
-                reward = -1.0
-            else:
-                reward = 1.0
+
+            reward = -1.0 if self._sutton_barto_reward else 1.0
         else:
             if self.steps_beyond_terminated == 0:
                 logger.warn(
-                    "You are calling 'step()' even though this "
-                    "environment has already returned terminated = True. You "
-                    "should always call 'reset()' once you receive 'terminated = "
-                    "True' -- any further steps are undefined behavior."
+                    "You are calling 'step()' even though this environment has already returned terminated = True. "
+                    "You should always call 'reset()' once you receive 'terminated = True' -- any further steps are undefined behavior."
                 )
             self.steps_beyond_terminated += 1
-            if self._sutton_barto_reward:
-                reward = -1.0
-            else:
-                reward = 0.0
+
+            reward = -1.0 if self._sutton_barto_reward else 0.0
 
         if self.render_mode == "human":
             self.render()
+
         # truncation=False as the time limit is handled by the `TimeLimit` wrapper added during `make`
         return np.array(self.state, dtype=np.float32), reward, terminated, False, {}
 
@@ -439,10 +432,11 @@ def step(
         # For the interested reader:
         # https://coneural.org/florian/papers/05_cart_pole.pdf
         temp = (
-            force + self.polemass_length * theta_dot**2 * sintheta
+            force + self.polemass_length * np.square(theta_dot) * sintheta
         ) / self.total_mass
         thetaacc = (self.gravity * sintheta - costheta * temp) / (
-            self.length * (4.0 / 3.0 - self.masspole * costheta**2 / self.total_mass)
+            self.length
+            * (4.0 / 3.0 - self.masspole * np.square(costheta) / self.total_mass)
         )
         xacc = temp - self.polemass_length * thetaacc * costheta / self.total_mass
 
@@ -470,7 +464,7 @@ def step(
 
         truncated = self.steps >= self.max_episode_steps
 
-        if self._sutton_barto_reward is True:
+        if self._sutton_barto_reward:
             reward = -np.array(terminated, dtype=np.float32)
         else:
             reward = np.ones_like(terminated, dtype=np.float32)
@@ -484,7 +478,7 @@ def step(
         terminated[self.prev_done] = False
         truncated[self.prev_done] = False
 
-        self.prev_done = terminated | truncated
+        self.prev_done = np.logical_or(terminated, truncated)
 
         return self.state.T.astype(np.float32), reward, terminated, truncated, {}
 
@@ -497,9 +491,8 @@ def reset(
         super().reset(seed=seed)
         # Note that if you use custom reset bounds, it may lead to out-of-bound
         # state/observations.
-        self.low, self.high = utils.maybe_parse_reset_bounds(
-            options, -0.05, 0.05  # default low
-        )  # default high
+        # -0.05 and 0.05 is the default low and high bounds
+        self.low, self.high = utils.maybe_parse_reset_bounds(options, -0.05, 0.05)
         self.state = self.np_random.uniform(
             low=self.low, high=self.high, size=(4, self.num_envs)
         )

gymnasium/wrappers/vector/dict_info_to_list.py

@@ -37,7 +37,6 @@ class DictInfoToList(VectorWrapper):
     Example for vector environments:
         >>> import numpy as np
         >>> import gymnasium as gym
-        >>> from gymnasium.spaces import Dict, Box
         >>> envs = gym.make_vec("CartPole-v1", num_envs=3)
         >>> obs, info = envs.reset(seed=123)
         >>> info
@@ -61,7 +60,7 @@ class DictInfoToList(VectorWrapper):
         >>> _ = envs.action_space.seed(123)
         >>> _, _, _, _, infos = envs.step(envs.action_space.sample())
         >>> infos
-        [{'x_position': np.float64(0.03332210900362942), 'x_velocity': np.float64(-0.06296527291998533), 'reward_run': np.float64(-0.06296527291998533), 'reward_ctrl': np.float32(-0.24503504)}, {'x_position': np.float64(0.10172354684460168), 'x_velocity': np.float64(0.8934584807363618), 'reward_run': np.float64(0.8934584807363618), 'reward_ctrl': np.float32(-0.21944423)}]
+        [{'x_position': np.float64(0.0333221090036294), 'x_velocity': np.float64(-0.06296527291998574), 'reward_run': np.float64(-0.06296527291998574), 'reward_ctrl': np.float32(-0.24503504)}, {'x_position': np.float64(0.10172354684460168), 'x_velocity': np.float64(0.8934584807363618), 'reward_run': np.float64(0.8934584807363618), 'reward_ctrl': np.float32(-0.21944423)}]
 
     Change logs:
      * v0.24.0 - Initially added as ``VectorListInfo``

tests/envs/test_env_implementation.py

@@ -283,11 +283,15 @@ def test_cartpole_vector_equiv():
     assert env.action_space == envs.single_action_space
     assert env.observation_space == envs.single_observation_space
 
-    # reset
+    # for seed in range(0, 10_000):
     seed = np.random.randint(0, 1000)
+
+    # reset
     obs, info = env.reset(seed=seed)
     vec_obs, vec_info = envs.reset(seed=seed)
 
+    env.action_space.seed(seed=seed)
+
     assert obs in env.observation_space
     assert vec_obs in envs.observation_space
     assert np.all(obs == vec_obs[0])
@@ -315,24 +319,27 @@ def test_cartpole_vector_equiv():
 
         assert np.all(env.unwrapped.state == envs.unwrapped.state[:, 0])
 
-        if term:
+        if term or trunc:
             break
 
-    obs, info = env.reset()
-    # the vector action shouldn't matter as autoreset
-    vec_obs, vec_reward, vec_term, vec_trunc, vec_info = envs.step(
-        envs.action_space.sample()
-    )
+    # if the sub-environment episode ended
+    if term or trunc:
+        obs, info = env.reset()
+        # the vector action shouldn't matter as autoreset
+        assert envs.unwrapped.prev_done
+        vec_obs, vec_reward, vec_term, vec_trunc, vec_info = envs.step(
+            envs.action_space.sample()
+        )
 
-    assert obs in env.observation_space
-    assert vec_obs in envs.observation_space
-    assert np.all(obs == vec_obs[0])
-    assert vec_reward == np.array([0])
-    assert vec_term == np.array([False])
-    assert vec_trunc == np.array([False])
-    assert info == vec_info
+        assert obs in env.observation_space
+        assert vec_obs in envs.observation_space
+        assert np.all(obs == vec_obs[0])
+        assert vec_reward == np.array([0])
+        assert vec_term == np.array([False])
+        assert vec_trunc == np.array([False])
+        assert info == vec_info
 
-    assert np.all(env.unwrapped.state == envs.unwrapped.state[:, 0])
+        assert np.all(env.unwrapped.state == envs.unwrapped.state[:, 0])
 
     env.close()
     envs.close()