Update app.py

app.py

@@ -66,47 +66,61 @@ class TrafficSimEnv(gym.Env):
     def __init__(self, congestion_level):
         super(TrafficSimEnv, self).__init__()
         self.congestion_level = congestion_level
+
+        # Define observation space: [congestion_level, signal_duration]
+        self.observation_space = spaces.Box(low=np.array([0, 0]), high=np.array([10, 60]), dtype=np.float32)
+
+        # Define action space: [increase, decrease, maintain]
         self.action_space = spaces.Discrete(3)
-        self.observation_space = spaces.Box(low=np.array([0, 20]), high=np.array([2, 60]), dtype=np.float32)
-        self.current_signal = 30
-        self.steps = 0
 
-    def reset(self, seed=None, options=None):
-        self.steps = 0
-        self.current_signal = 30
-        congestion_map = {"Low": 0, "Medium": 1, "High": 2}
-        self.congestion_numeric = congestion_map.get(self.congestion_level, 0)
+        # Initial state
+        self.state = np.array([congestion_level, 30], dtype=np.float32)
+        self.done = False
+
+    def reset(self, *, seed=None, options=None):
+        super().reset(seed=seed)
 
-        obs = np.array([self.congestion_numeric, self.current_signal], dtype=np.float32)
-        info = {}  # Empty dictionary as info
-        print(f"RESET: obs={obs}, info={info}")  # Debugging print
-        return obs, info
+        self.state = np.array([self.congestion_level, 30], dtype=np.float32)
+        self.done = False
+        return self.state, {}
 
     def step(self, action):
-        signal_change = {0: -5, 1: 0, 2: 5}[action]
-        self.current_signal = max(20, min(60, self.current_signal + signal_change))
+        if self.done:
+            raise RuntimeError("Cannot call step() on a terminated environment. Please reset the environment.")
+
+        # Extract state components
+        congestion, signal_duration = self.state
+
+        # Apply action
+        if action == 0:  # Decrease signal duration
+            signal_duration = max(10, signal_duration - 5)
+        elif action == 1:  # Maintain signal duration
+            signal_duration = signal_duration
+        elif action == 2:  # Increase signal duration
+            signal_duration = min(60, signal_duration + 5)
+
+        # Update congestion (simple simulation logic for this example)
+        if signal_duration > 30:
+            congestion += 1
+        else:
+            congestion -= 1
 
-        reward = self._calculate_reward()
+        # Set rewards (example logic)
+        if 20 <= signal_duration <= 40:
+            reward = 0
+        else:
+            reward = -abs(signal_duration - 30)
 
-        self.steps += 1
-        done = self.steps >= 10
-        truncated = False
+        # Check if done
+        self.done = congestion <= 0 or congestion >= 10
 
-        obs = np.array([self.congestion_numeric, self.current_signal], dtype=np.float32)
-        info = {}  # Info dictionary (can be populated with useful debugging data)
-        print(f"STEP: action={action}, obs={obs}, reward={reward}, done={done}, truncated={truncated}, info={info}")
-        return obs, reward, done, truncated, info
+        # Update state
+        self.state = np.array([congestion, signal_duration], dtype=np.float32)
 
-    def _calculate_reward(self):
-        if self.congestion_level == "High":
-            return -abs(40 - self.current_signal)
-        elif self.congestion_level == "Medium":
-            return -abs(30 - self.current_signal)
-        else:
-            return -abs(20 - self.current_signal)
+        return self.state, reward, self.done, False, {}
 
-    def render(self, mode="human"):
-        print(f"Current Signal: {self.current_signal}s")
+    def render(self):
+        print(f"State: {self.state}")
 
     def close(self):
         pass
@@ -114,7 +128,7 @@ class TrafficSimEnv(gym.Env):
 # Prior to commit had a lot of errors regarding expected output errors
 def optimize_signal_rl(congestion_level):
     try:
-        # Create the environment wrapped in DummyVecEnv
+        # Create environment
         env = DummyVecEnv([lambda: TrafficSimEnv(congestion_level)])
 
         # Initialize PPO model
@@ -124,40 +138,34 @@ def optimize_signal_rl(congestion_level):
         model.learn(total_timesteps=1000)
 
         # Reset the environment
-        obs, info = env.reset()  # For Gymnasium 1.0.0
-        obs = obs[0]  # Extract single observation from the batched observation
+        obs, info = env.reset()
+        obs = obs[0]  # Extract first observation from batch
 
         for _ in range(10):
-            # Predict action
             action, _ = model.predict(obs, deterministic=True)
+            obs, reward, done, truncated, info = env.step(action)
+
+            # Ensure single observation extraction
+            obs = obs[0]
+            reward = reward[0]
+            done = done[0]
+            truncated = truncated[0]
+            info = info[0]
+
+            # Log each step for debugging
+            print(f"STEP: action={action}, obs={obs}, reward={reward}, done={done}, truncated={truncated}, info={info}")
 
-            # Take a step
-            step_output = env.step(action)
-
-            # Handle batched output of step()
-            if len(step_output) == 5:
-                obs, reward, done, truncated, info = step_output
-                obs = obs[0]  # Extract single observation
-                reward = reward[0]  # Extract single reward
-                done = done[0]  # Extract single done flag
-                truncated = truncated[0]  # Extract single truncated flag
-                info = info[0]  # Extract single info
-            else:
-                raise ValueError(f"Unexpected step output: {step_output}")
-
-            # End simulation if done or truncated
             if done or truncated:
                 break
 
-        # Extract optimal signal duration
         optimal_duration = int(obs[1]) if len(obs) > 1 else 30
         return f"Green for {optimal_duration}s, Red for {60 - optimal_duration}s"
+
     except Exception as e:
         logging.error(f"Error optimizing signal with RL: {e}")
         return "Error in RL Optimization"
 
 
-
 def process_traffic_image(image):
     """
     Orchestrates the traffic analysis workflow.