Update app.py

app.py

@@ -123,23 +123,29 @@ def optimize_signal_rl(congestion_level):
         # Train the model
         model.learn(total_timesteps=1000)
 
-        # Reset environment
-        obs, info = env.reset()  # For Gymnasium 1.0.0, reset() returns (obs, info)
-        obs = obs[0]  # Extract first observation from batch
+        # Reset the environment
+        obs, info = env.reset()  # For Gymnasium 1.0.0
+        obs = obs[0]  # Extract single observation from the batched observation
 
         for _ in range(10):
             # Predict action
-            action, _ = model.predict(obs)
+            action, _ = model.predict(obs, deterministic=True)
 
             # Take a step
-            obs, reward, done, truncated, info = env.step(action)
-            obs = obs[0]  # Extract first observation from batch
-            reward = reward[0]  # Extract first reward from batch
-            done = done[0]  # Extract first done flag from batch
-            truncated = truncated[0]  # Extract first truncated flag from batch
-            info = info[0]  # Extract first info from batch
-
-            # End simulation if episode is done or truncated
+            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
 
@@ -151,6 +157,7 @@ def optimize_signal_rl(congestion_level):
         return "Error in RL Optimization"
 
 
+
 def process_traffic_image(image):
     """
     Orchestrates the traffic analysis workflow.