Fix initial bugs

app.py

@@ -99,6 +99,7 @@ class WebGameEngine:
         self.models_ready = False  # Track if models are loaded
         self.download_progress = 0  # Track download progress (0-100)
         self.loading_status = "Initializing..."  # Loading status message
+        self.actor_critic_loaded = False  # Track if actor_critic was loaded with trained weights
         import time
         self.time_module = time
     
@@ -129,6 +130,9 @@ class WebGameEngine:
                 logger.info(f"Model has actor_critic weights: {has_actor_critic}")
                 agent.load_state_dict(state_dict, load_actor_critic=has_actor_critic)
                 
+                # Track if actor_critic was actually loaded with trained weights
+                self.actor_critic_loaded = has_actor_critic
+                
                 self.download_progress = 100
                 self.loading_status = "Model loaded successfully!"
                 logger.info("All model weights loaded successfully!")
@@ -220,6 +224,8 @@ class WebGameEngine:
                         self.loading_status = "Loading local checkpoint..."
                         agent.load(checkpoint_path)
                         logger.info(f"Successfully loaded local checkpoint: {checkpoint_path}")
+                        # Assume local checkpoint has actor_critic weights (may need verification)
+                        self.actor_critic_loaded = True
                     else:
                         logger.error(f"No local checkpoint found at: {checkpoint_path}")
                         raise FileNotFoundError("No model checkpoint available (local or remote)")
@@ -227,6 +233,7 @@ class WebGameEngine:
             except Exception as e:
                 logger.error(f"Failed to load any checkpoint: {e}")
                 self._init_dummy_mode()
+                self.actor_critic_loaded = False  # No actor_critic in dummy mode
                 return True
             
             # Initialize world model environment
@@ -242,12 +249,16 @@ class WebGameEngine:
                 self.play_env = WebPlayEnv(agent, wm_env, False, False, False)
                 
                 # Verify actor-critic is loaded and ready for inference
-                if agent.actor_critic is not None:
+                if agent.actor_critic is not None and self.actor_critic_loaded:
                     logger.info(f"Actor-critic model loaded with {agent.actor_critic.lstm_dim} LSTM dimensions")
                     logger.info(f"Actor-critic device: {agent.actor_critic.device}")
                     # Force AI control for web demo
                     self.play_env.is_human_player = False
                     logger.info("WebPlayEnv set to AI control mode")
+                elif agent.actor_critic is not None and not self.actor_critic_loaded:
+                    logger.warning("Actor-critic model exists but has no trained weights - using dummy mode!")
+                    self.play_env.is_human_player = True
+                    logger.info("WebPlayEnv set to human control mode (no trained weights)")
                 else:
                     logger.warning("No actor-critic model found - AI inference will not work!")
                     self.play_env.is_human_player = True
@@ -281,6 +292,7 @@ class WebGameEngine:
             except Exception as e:
                 logger.error(f"Failed to initialize world model environment: {e}")
                 self._init_dummy_mode()
+                self.actor_critic_loaded = False  # No actor_critic in dummy mode
                 self.models_ready = True
                 self.loading_status = "Using dummy mode"
                 return True
@@ -290,6 +302,7 @@ class WebGameEngine:
             import traceback
             traceback.print_exc()
             self._init_dummy_mode()
+            self.actor_critic_loaded = False  # No actor_critic in dummy mode
             self.models_ready = True
             self.loading_status = "Error - using dummy mode"
             return True
@@ -573,7 +586,7 @@ async def get_homepage():
     <!DOCTYPE html>
     <html>
     <head>
-        <title>Diamond CSGO AI Player</title>
+        <title>Physics-informed BEV World Model</title>
         <style>
             body {
                 margin: 0;
@@ -629,7 +642,7 @@ async def get_homepage():
         </style>
     </head>
     <body>
-        <h1>🎮 Diamond CSGO AI Player</h1>
+        <h1>🎮 Physics-informed BEV World Model</h1>
         <p><strong>Click the game canvas to start playing!</strong> Use ESC to pause, Enter to reset environment.</p>
         <p id="loadingIndicator" style="color: #ffff00; display: none;">🚀 Starting AI inference... This may take 5-15 seconds on first run.</p>
         

src/game/web_play_env.py

@@ -32,9 +32,13 @@ class WebPlayEnv(PlayEnv):
         self.is_human_player = False  # AI controls the actions
         self.human_input_override = False  # Can be set to True to allow human input
         
-        # Initialize LSTM hidden states for actor-critic
-        self.hx = torch.zeros(1, agent.actor_critic.lstm_dim, device=agent.device)
-        self.cx = torch.zeros(1, agent.actor_critic.lstm_dim, device=agent.device)
+        # Initialize LSTM hidden states for actor-critic (only if actor_critic exists)
+        if agent.actor_critic is not None:
+            self.hx = torch.zeros(1, agent.actor_critic.lstm_dim, device=agent.device)
+            self.cx = torch.zeros(1, agent.actor_critic.lstm_dim, device=agent.device)
+        else:
+            self.hx = None
+            self.cx = None
     
     def switch_controller(self) -> None:
         """Switch between AI and human control"""
@@ -97,9 +101,11 @@ class WebPlayEnv(PlayEnv):
                 if obs.device != self.agent.device:
                     obs = obs.to(self.agent.device, non_blocking=True)
                 
-                # Detach hidden states to prevent gradient tracking
-                self.hx = self.hx.detach()
-                self.cx = self.cx.detach()
+                # Detach hidden states to prevent gradient tracking (only if they exist)
+                if self.hx is not None:
+                    self.hx = self.hx.detach()
+                if self.cx is not None:
+                    self.cx = self.cx.detach()
                 
                 # Resize observation to match actor-critic expected encoder/LSTM input
                 # Count how many MaxPool2d layers are in the encoder to infer downsampling factor
@@ -145,10 +151,12 @@ class WebPlayEnv(PlayEnv):
         self.obs = next_obs
         self.t += 1
         
-        # Reset hidden states on episode end
+        # Reset hidden states on episode end (only if they exist)
         if end.any() or trunc.any():
-            self.hx.zero_()
-            self.cx.zero_()
+            if self.hx is not None:
+                self.hx.zero_()
+            if self.cx is not None:
+                self.cx.zero_()
         
         # Return the step results
         return next_obs, rew, end, trunc, env_info