Made CartPoleVectorEnv compliant with Farama-Foundation#785 by delayi…

…ng the reset of done environments until the next step. Fixes Farama-Foundation#914.

gymnasium/envs/classic_control/cartpole.py

@@ -364,6 +364,7 @@ def __init__(
         self.kinematics_integrator = "euler"
 
         self.steps = np.zeros(num_envs, dtype=np.int32)
+        self.prev_done = np.zeros(num_envs, dtype=np.bool_)
 
         # Angle at which to fail the episode
         self.theta_threshold_radians = 12 * 2 * math.pi / 360
@@ -445,16 +446,18 @@ def step(
 
         truncated = self.steps >= self.max_episode_steps
 
-        done = terminated | truncated
+        reward = np.ones_like(terminated, dtype=np.float32)
 
-        if any(done):
-            # This code was generated by copilot, need to check if it works
-            self.state[:, done] = self.np_random.uniform(
-                low=self.low, high=self.high, size=(4, done.sum())
+        # Reset all environments which terminated or were truncated in the last step
+        self.state[:, self.prev_done] = self.np_random.uniform(
+                low=self.low, high=self.high, size=(4, self.prev_done.sum())
             ).astype(np.float32)
-            self.steps[done] = 0
+        self.steps[self.prev_done] = 0
+        reward[self.prev_done] = 0.0
+        terminated[self.prev_done] = False
+        truncated[self.prev_done] = False
 
-        reward = np.ones_like(terminated, dtype=np.float32)
+        self.prev_done = terminated | truncated
 
         if self.render_mode == "human":
             self.render()
@@ -478,6 +481,7 @@ def reset(
         ).astype(np.float32)
         self.steps_beyond_terminated = None
         self.steps = np.zeros(self.num_envs, dtype=np.int32)
+        self.prev_done = np.zeros(self.num_envs, dtype=np.bool_)
 
         if self.render_mode == "human":
             self.render()