updated evaluation code for previous env changes

outdated code from previous changes on env

ml/env.py

@@ -115,19 +115,26 @@ class JackBotEnv(gym.Env):
         self.robot_rewards = [0.0 for _ in range(self.num_robots)]
         self.failed_robots_mask = [False for _ in range(self.num_robots)]
 
-        if not self._first_reset:
-            p.resetSimulation(physicsClientId=self.sim_manager.physics_client)
-            p.setGravity(0, 0, -9.81, physicsClientId=self.sim_manager.physics_client)
-            p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0, physicsClientId=self.sim_manager.physics_client)
-            self.hud.reset()
-            self.leader_crown.reset()
-            self.plane, self.pb_robots, self.robot_joint_indices = self.sim_manager.load_scene(
-                self.urdf_path, self.num_robots, self.robot_spacing, self._robot_base_position
+        for idx, (pb_id, robot_obj) in enumerate(zip(self.pb_robots, self.robots)):
+            # Get default spawn position
+            spawn_pos = self._robot_base_position(idx, self.num_robots, self.robot_spacing)
+            spawn_orn = [0, 0, 0, 1]
+
+            # Teleport base back to start
+            p.resetBasePositionAndOrientation(
+                pb_id, spawn_pos, spawn_orn, physicsClientId=self.sim_manager.physics_client
             )
-            for robot_obj, pb_id in zip(self.robots, self.pb_robots):
-                robot_obj.backend = PyBulletBackend(self.sim_manager, body_id=pb_id)
-        else:
-            self._first_reset = False
+            p.resetBaseVelocity(
+                pb_id, linearVelocity=[0, 0, 0], angularVelocity=[0, 0, 0],
+                physicsClientId=self.sim_manager.physics_client
+            )
+
+            # Reset joint angles directly without reloading URDF
+            robot_obj.reset_to_init()
+
+            # Restore original default visual color (clears failure dark gray)
+            if self.use_gui:
+                self._set_robot_color(pb_id, [1.0, 1.0, 1.0, 1.0])
 
         self.last_action = np.zeros(self.action_space.shape[0], dtype=np.float32)
 
@@ -136,13 +143,14 @@ class JackBotEnv(gym.Env):
         else:
             self.commands = np.zeros((self.num_robots, 4), dtype=np.float32)
 
-        for robot in self.robots:
-            robot.reset_to_init()
-
-        for _ in range(100):
+        for _ in range(15):
             self.sim_manager.step()
 
-        self._update_hud()
+        if self.use_gui:
+            self.hud.reset()
+            self.leader_crown.reset()
+            self._update_hud()
+
         return self._get_obs(), {}
 
     def _get_obs(self) -> np.ndarray: