Update app.py

app.py

@@ -90,35 +90,41 @@ class TrafficSimEnv(gym.Env):
         return obs, {}  # Returns observation and info (info can be empty)
 
 
-    def step(self, action):
-        """
-        Takes an action and updates the environment state.
-        Returns 5 values: observation, reward, done, truncated, and info.
-        """
+    def reset(self, seed=None, options=None):
+        self.steps = 0
+        self.current_signal = 30  # Reset the signal at the beginning
+        congestion_map = {"Low": 0, "Medium": 1, "High": 2}
+        self.congestion_numeric = congestion_map.get(self.congestion_level, 0)
+        obs = np.array([self.congestion_numeric, self.current_signal], dtype=np.float32)
+        return obs, {}  # Returns observation and info (info can remain empty)
 
+    def step(self, action):
         print(f"Action taken: {action}")
         print(f"Signal before change: {self.current_signal}")
         print(f"Congestion Level: {self.congestion_numeric}")
         
         # Signal changes based on the action
-        signal_change = {0: -5, 1: 0, 2: 5}[action]
-        self.current_signal = max(20, min(60, self.current_signal + signal_change))
-
+        if action == 0:
+            self.current_signal = max(20, self.current_signal - 5)  # Decrease signal
+        elif action == 2:
+            self.current_signal = min(60, self.current_signal + 5)  # Increase signal
+        # Action 1 is to maintain the current signal
+        
         # Reward Calculation based on congestion level and signal
         reward = self._calculate_reward()
 
         # Increment the step count
         self.steps += 1
-        done = self.steps >= 10  # End condition
+        done = self.steps >= 10  # End condition: after 10 steps
         truncated = False  # Default to False; can change based on custom conditions
 
-        # Ensure we always return 5 values: obs, reward, done, truncated, info
+        # Observation after the action
         obs = np.array([self.congestion_numeric, self.current_signal], dtype=np.float32)
         print(f"Observation: {obs}")
 
         info = {}  # Additional info (can remain empty or contain any useful data)
 
-        return obs, reward, done, truncated, info  # Must return 5 values here
+        return obs, reward, done, truncated, info
 
     def _calculate_reward(self):
         if self.congestion_level == "High":