fixed env eval setup

2026-07-31 18:35:09 +02:00
parent a3e46c3abf
commit 5317ef1299
5 changed files with 105 additions and 72 deletions
@@ -143,7 +143,7 @@ class JackBotEnv(gym.Env):
        else:
            self.commands = np.zeros((self.num_robots, 4), dtype=np.float32)

-        for _ in range(15):
+        for _ in range(100):
            self.sim_manager.step()

        if self.use_gui:
@@ -199,15 +199,30 @@ class JackBotEnv(gym.Env):
    def _compute_reward(self) -> Tuple[float, list[float]]:
        rewards = []
        for idx, pb_id in enumerate(self.pb_robots):
-            linear_vel, angular_vel = p.getBaseVelocity(pb_id, physicsClientId=self.sim_manager.physics_client)
-            _, orientation = p.getBasePositionAndOrientation(pb_id, physicsClientId=self.sim_manager.physics_client)
+            # Stop rewarding robots that have already collapsed or flipped
+            if self.failed_robots_mask[idx]:
+                rewards.append(-0.5)  # Penalty per step while collapsed
+                continue
+
+            linear_vel, angular_vel = p.getBaseVelocity(
+                pb_id, physicsClientId=self.sim_manager.physics_client
+            )
+            _, orientation = p.getBasePositionAndOrientation(
+                pb_id, physicsClientId=self.sim_manager.physics_client
+            )
            roll, pitch, _ = p.getEulerFromQuaternion(orientation)
            command = self.commands[idx]

            forward_reward = command[0] * linear_vel[0] + command[1] * linear_vel[1]
            rotation_reward = command[3] * angular_vel[2]
            stability_penalty = abs(roll) + abs(pitch)
-            action_penalty = float(np.sum(np.square(self.last_action.reshape(self.num_robots, -1)[idx]))) * 0.01
+            
+            # Calculate per-robot action penalty
+            action_dim_per_robot = 18
+            start_idx = idx * action_dim_per_robot
+            end_idx = start_idx + action_dim_per_robot
+            robot_action = self.last_action[start_idx:end_idx]
+            action_penalty = float(np.sum(np.square(robot_action))) * 0.01

            r_step = 0.1 + forward_reward + rotation_reward - 0.2 * stability_penalty - action_penalty
            rewards.append(r_step)