Update server/sentinel_env_environment.py

server/sentinel_env_environment.py

@@ -51,7 +51,7 @@ class SentinelEnv:
 
     def step(self, action: SentinelAction) -> Tuple[SentinelObservation, float, bool, Dict[str, Any]]:
         self._state.step_count += 1
-        reward = 0.0
+        reward = 0.05
         done = False
         msg = "No change detected."
         task = self._state.current_task_id
@@ -60,30 +60,28 @@ class SentinelEnv:
         if task == "easy-lockdown":
             if action.command == "lockdown" and action.target_id == "s3-vault":
                 self.resources[0].status = "private"
-                reward = 1.0
+                reward = 0.95
                 done = True
                 msg = "SUCCESS: S3 Bucket is now private."
             elif action.command == "scan" and action.target_id == "s3-vault":
-                # incremental reward for scanning before acting
                 reward = 0.1
                 msg = "Scan complete: S3 bucket is public. Apply lockdown."
             else:
-                reward = -0.1
+                reward = 0.05
                 msg = "Wrong action. Try: command='lockdown', target_id='s3-vault'."
 
         # ── TASK 2: MEDIUM — Least Privilege ──
         elif task == "medium-access":
             if action.command == "revoke_admin" and action.target_id == "user-dev-01":
                 self.resources[0].status = "read_only"
-                reward = 1.0
+                reward = 0.95
                 done = True
                 msg = "SUCCESS: Admin rights revoked. User set to Read-Only."
             elif action.command == "audit" and action.target_id == "user-dev-01":
-                # incremental reward for auditing before revoking
                 reward = 0.1
                 msg = "Audit complete: user-dev-01 has admin_access. Revoke it."
             else:
-                reward = -0.1
+                reward = 0.05
                 msg = "Wrong action. Try: command='revoke_admin', target_id='user-dev-01'."
 
         # ── TASK 3: HARD — Incident Response ──
@@ -93,29 +91,28 @@ class SentinelEnv:
                 for r in self.resources
             )
             if action.command == "investigate" and action.target_id == "attacker-ip":
-                # incremental reward for investigating before blocking
                 reward = 0.1
-                msg = "Investigation complete: brute force detected from attacker-ip. Block it."
+                msg = "Investigation complete: brute force from attacker-ip. Block it."
             elif action.command == "block_ip" and action.target_id == "attacker-ip":
                 for r in self.resources:
                     if r.id == "attacker-ip":
                         r.status = "blocked"
-                reward = 0.5
+                reward = 0.45
                 msg = "IP Blocked. Now close the open port on web-server."
             elif action.command == "close_port" and action.target_id == "web-server":
                 if ip_blocked:
                     for r in self.resources:
                         if r.id == "web-server":
                             r.status = "port_22_closed"
-                    reward = 0.5
+                    reward = 0.50
                     done = True
                     msg = "SUCCESS: Attack stopped and port secured."
                 else:
-                    reward = -0.1
+                    reward = 0.05
                     msg = "Port close failed. Block the attacker IP first."
             else:
-                reward = -0.1
-                msg = "Unknown action. Try block_ip on attacker-ip, then close_port on web-server."
+                reward = 0.05
+                msg = "Unknown action."
 
         self.logs.append(f"Step {self._state.step_count}: {msg} (reward={reward})")
         return self._generate_observation(msg), reward, done, {}