Fix initial bugs

app.py

@@ -28,7 +28,7 @@ import io
 import json
 import logging
 from pathlib import Path
-from typing import Dict, List, Optional, Set
+from typing import Dict, List, Optional, Set, Tuple
 
 import cv2
 import numpy as np
@@ -101,8 +101,13 @@ class WebGameEngine:
         self.obs = None
         self.running = False
         self.game_started = False
-        self.fps = 30  # Display FPS
-        self.ai_fps = 10  # AI inference FPS (slower than display for efficiency)
+        # Allow runtime tuning via environment variables
+        import os
+        self.fps = int(os.getenv("DISPLAY_FPS", "30"))  # Display FPS
+        # Increase default AI inference FPS; can be overridden with AI_FPS env var
+        self.ai_fps = int(os.getenv("AI_FPS", "15"))
+        # Send every Nth frame to the browser (1 = send all frames)
+        self.send_every = int(os.getenv("DISPLAY_SKIP", "1"))
         self.frame_count = 0
         self.ai_frame_count = 0
         self.last_ai_time = 0
@@ -121,6 +126,13 @@ class WebGameEngine:
         self.actor_critic_loaded = False  # Track if actor_critic was loaded with trained weights
         import time
         self.time_module = time
+
+        # Async inference queues to decouple GPU work from websocket I/O
+        import asyncio
+        self._in_queue: asyncio.Queue = asyncio.Queue(maxsize=1)
+        self._out_queue: asyncio.Queue = asyncio.Queue(maxsize=1)
+        # Flag to start worker once models are ready
+        self._worker_started = False
     
     async def _load_model_from_url_async(self, agent, device):
         """Load model from URL using torch.hub (HF Spaces compatible)"""
@@ -295,11 +307,10 @@ class WebGameEngine:
                     self.play_env.is_human_player = True
                     logger.info("WebPlayEnv set to human control mode (fallback)")
                 
-                # Model compilation causes 10-30s delay on first inference, so make it optional
-                # You can enable it by setting ENABLE_TORCH_COMPILE=1 environment variable
+                # Enable torch.compile by default like play.py does (can disable with DISABLE_TORCH_COMPILE=1)
                 import os
-                if device.type == "cuda" and os.getenv("ENABLE_TORCH_COMPILE", "0") == "1":
-                    logger.info("Compiling models for faster inference (will cause delay on first inference)...")
+                if device.type == "cuda" and os.getenv("DISABLE_TORCH_COMPILE", "0") != "1":
+                    logger.info("Compiling models for faster inference (like play.py --compile)...")
                     try:
                         wm_env.predict_next_obs = torch.compile(wm_env.predict_next_obs, mode="reduce-overhead")
                         if wm_env.upsample_next_obs is not None:
@@ -308,7 +319,7 @@ class WebGameEngine:
                     except Exception as e:
                         logger.warning(f"Model compilation failed: {e}")
                 else:
-                    logger.info("Model compilation disabled (faster startup). Set ENABLE_TORCH_COMPILE=1 to enable.")
+                    logger.info("Model compilation disabled. Set DISABLE_TORCH_COMPILE=0 to enable.")
                 
                 # Reset environment
                 self.obs, _ = self.play_env.reset()
@@ -377,27 +388,27 @@ class WebGameEngine:
                 self.last_ai_time = self.time_module.time()  # Reset AI timer
                 return self.obs, 0.0, False, False, {"reset": True}
             
-            # Intelligent frame skipping: only run AI inference at target FPS
             current_time = self.time_module.time()
+
+            # Push task to inference queue if needed
             time_since_last_ai = current_time - self.last_ai_time
             should_run_ai = time_since_last_ai >= (1.0 / self.ai_fps)
-            
-            if should_run_ai:
-                # Show loading indicator for first inference (can be slow)
-                if not self.first_inference_done:
-                    logger.info("Running first AI inference (may take 5-15 seconds)...")
-                
-                # Run AI inference
-                inference_start = self.time_module.time()
-                next_obs, reward, done, truncated, info = self.play_env.step_from_web_input(
-                    pressed_keys=self.pressed_keys,
+
+            if should_run_ai and self._in_queue.empty():
+                # Snapshot web input state
+                web_state = dict(
+                    pressed_keys=set(self.pressed_keys),
                     mouse_x=self.mouse_x,
                     mouse_y=self.mouse_y,
                     l_click=self.l_click,
-                    r_click=self.r_click
+                    r_click=self.r_click,
                 )
-                inference_time = self.time_module.time() - inference_start
-                
+                asyncio.create_task(self._in_queue.put((self.obs, web_state)))
+
+            # Check for completed inference
+            if not self._out_queue.empty():
+                (next_obs, reward, done, truncated, info, inference_time) = self._out_queue.get_nowait()
+ 
                 # Log first inference completion
                 if not self.first_inference_done:
                     self.first_inference_done = True
@@ -489,23 +500,80 @@ class WebGameEngine:
             img = Image.fromarray(img_array)
             
             # Resize for web display to match canvas size (optimized)
-            img = img.resize((600, 150), Image.NEAREST)  # NEAREST is faster than BICUBIC
-            
-            # Optimized base64 conversion with JPEG for better compression/speed
-            buffer = io.BytesIO()
-            img.save(buffer, format='JPEG', quality=85, optimize=True)  # JPEG is faster than PNG
-            img_str = base64.b64encode(buffer.getvalue()).decode()
-            return f"data:image/jpeg;base64,{img_str}"
+            img = img.resize((600, 150), Image.NEAREST)
+
+            # Choose codec via env var for flexibility (jpeg|png)
+            codec = os.getenv("IMG_CODEC", "jpeg").lower()
+            img_np = np.array(img)[:, :, ::-1]  # RGB -> BGR
+            if codec == "png":
+                success, encoded_img = cv2.imencode('.png', img_np, [cv2.IMWRITE_PNG_COMPRESSION, 1])
+                mime = 'png'
+            else:
+                # JPEG with quality 70 for speed/size balance
+                success, encoded_img = cv2.imencode('.jpg', img_np, [cv2.IMWRITE_JPEG_QUALITY, 70])
+                mime = 'jpeg'
+            if not success:
+                return ""
+            img_str = base64.b64encode(encoded_img).decode()
+            return f"data:image/{mime};base64,{img_str}"
             
         except Exception as e:
             logger.error(f"Error converting observation to base64: {e}")
             return ""
+
+    # ------------------------------------------------------------------
+    #   Faster binary encoder (JPEG/PNG) with OpenCV – no Pillow involved
+    # ------------------------------------------------------------------
+    def obs_to_bytes(self, obs: torch.Tensor) -> Tuple[bytes, str]:
+        """Return encoded image bytes and MIME (image/jpeg or image/png)."""
+        if obs is None:
+            return b"", "image/jpeg"
+
+        try:
+            # Keep operations on GPU as long as possible (like play.py)
+            if obs.ndim == 4 and obs.size(0) == 1:
+                img_tensor = obs[0]
+            else:
+                img_tensor = obs
+
+            # Resize on GPU first (faster than CPU resize)
+            img_tensor = torch.nn.functional.interpolate(
+                img_tensor.unsqueeze(0), size=(75, 300), mode='nearest'
+            ).squeeze(0)
+
+            # Convert to uint8 on GPU, then transfer to CPU once
+            img_np = (img_tensor.add(1).mul(127.5).clamp(0, 255).byte()
+                       .permute(1, 2, 0).contiguous().cpu().numpy())  # HWC uint8
+
+            # Encode with OpenCV
+            import os
+            codec = os.getenv("IMG_CODEC", "jpeg").lower()
+            if codec == "png":
+                ok, enc = cv2.imencode('.png', img_np, [cv2.IMWRITE_PNG_COMPRESSION, 1])
+                mime = "image/png"
+            else:
+                ok, enc = cv2.imencode('.jpg', img_np, [cv2.IMWRITE_JPEG_QUALITY, 75])
+                mime = "image/jpeg"
+            if not ok:
+                return b"", mime
+            return enc.tobytes(), mime
+        except Exception as e:
+            logger.error(f"obs_to_bytes error: {e}")
+            return b"", "image/jpeg"
     
     async def game_loop(self):
         """Main game loop that runs continuously"""
         self.running = True
-        
+        # Start inference worker once, when models are ready
+
         while self.running:
+            loop_start_time = self.time_module.time()
+            
+            # Spawn worker lazily after models initialized
+            if self.models_ready and not self._worker_started:
+                asyncio.create_task(self._inference_worker())
+                self._worker_started = True
+
             try:
                 # Check if models are ready
                 if not self.models_ready:
@@ -554,50 +622,104 @@ class WebGameEngine:
                         self.obs = next_obs
                 
                 # Send frame to all connected clients (regardless of game state)
-                if should_send_frame and connected_clients and self.obs is not None:
+                if should_send_frame and connected_clients and self.obs is not None and (self.frame_count % self.send_every == 0):
                     # Set default values for when game isn't running
                     if not self.game_started:
                         reward = 0.0
                         info = {"waiting": True}
                     # If game is started, reward and info should be set above
                     
-                    # Convert observation to base64
-                    image_data = self.obs_to_base64(self.obs)
-                    
-                    # Debug logging for first few frames
-                    if self.frame_count < 5:
-                        logger.info(f"Frame {self.frame_count}: obs shape={self.obs.shape if self.obs is not None else 'None'}, "
-                                  f"image_data_length={len(image_data) if image_data else 0}, "
-                                  f"game_started={self.game_started}")
-                    
-                    frame_data = {
-                        'type': 'frame',
-                        'image': image_data,
-                        'frame_count': self.frame_count,
-                        'reward': float(reward.item()) if hasattr(reward, 'item') else float(reward) if reward is not None else 0.0,
-                        'info': str(info) if info else "",
-                        'ai_fps': info.get('ai_fps', 0) if isinstance(info, dict) else 0,
-                        'is_ai_frame': info.get('ai_inference', False) if isinstance(info, dict) else False
-                    }
-                    
-                    # Send to all connected clients
-                    disconnected = set()
-                    for client in connected_clients.copy():
-                        try:
-                            await client.send_text(json.dumps(frame_data))
-                        except:
-                            disconnected.add(client)
-                    
-                    # Remove disconnected clients
-                    connected_clients.difference_update(disconnected)
+                    # Prefer binary frames if client agrees (feature flag)
+                    use_binary = os.getenv("BINARY_WS", "0") == "1"
+
+                    if use_binary:
+                        img_bytes, mime = self.obs_to_bytes(self.obs)
+                        meta = {
+                            'type': 'frame_meta',
+                            'mime': mime,
+                            'frame_count': self.frame_count,
+                            'reward': float(reward.item()) if hasattr(reward, 'item') else float(reward) if reward is not None else 0.0,
+                            'info': str(info) if info else "",
+                            'ai_fps': info.get('ai_fps', 0) if isinstance(info, dict) else 0,
+                            'is_ai_frame': info.get('ai_inference', False) if isinstance(info, dict) else False
+                        }
+                        disconnected = set()
+                        for client in connected_clients.copy():
+                            try:
+                                await client.send_text(json.dumps(meta))
+                                await client.send_bytes(img_bytes)
+                            except:
+                                disconnected.add(client)
+                        connected_clients.difference_update(disconnected)
+                    else:
+                        # Fallback to base64 JSON
+                        image_data = self.obs_to_base64(self.obs)
+
+                        if self.frame_count < 5:
+                            logger.info(
+                                f"Frame {self.frame_count}: base64_len={len(image_data)} ai={info.get('ai_fps',0):.1f}")
+
+                        frame_data = {
+                            'type': 'frame',
+                            'image': image_data,
+                            'frame_count': self.frame_count,
+                            'reward': float(reward.item()) if hasattr(reward, 'item') else float(reward) if reward is not None else 0.0,
+                            'info': str(info) if info else "",
+                            'ai_fps': info.get('ai_fps', 0) if isinstance(info, dict) else 0,
+                            'is_ai_frame': info.get('ai_inference', False) if isinstance(info, dict) else False
+                        }
+
+                        disconnected = set()
+                        for client in connected_clients.copy():
+                            try:
+                                await client.send_text(json.dumps(frame_data))
+                            except:
+                                disconnected.add(client)
+                        connected_clients.difference_update(disconnected)
                 
                 self.frame_count += 1
-                await asyncio.sleep(1.0 / self.fps)  # Control FPS
+
+                # Adaptive sleep so we don't waste idle time when GPU faster than display FPS
+                loop_elapsed = self.time_module.time() - loop_start_time
+                sleep_for = max((1.0 / self.fps) - loop_elapsed, 0)
+                if sleep_for:
+                    await asyncio.sleep(sleep_for)
                 
             except Exception as e:
                 logger.error(f"Error in game loop: {e}")
                 await asyncio.sleep(0.1)
 
+    async def _inference_worker(self):
+        """Runs AI inference in background to avoid blocking I/O."""
+        logger.info("Inference worker started")
+        next_inference_time = self.time_module.time()
+        
+        while True:
+            obs, web_state = await self._in_queue.get()
+
+            # Timing control: maintain steady AI_FPS like play.py's clock.tick()
+            now = self.time_module.time()
+            if now < next_inference_time:
+                await asyncio.sleep(next_inference_time - now)
+            next_inference_time += 1.0 / self.ai_fps
+
+            # Run inference directly in asyncio (not thread pool) with autocast for speed
+            try:
+                start = self.time_module.time()
+                
+                # Use FP16 autocast for faster inference (like play.py can do with modern GPUs)
+                from torch.cuda.amp import autocast
+                with autocast(dtype=torch.float16, enabled=torch.cuda.is_available()):
+                    res = self.play_env.step_from_web_input(**web_state)
+                
+                infer_t = self.time_module.time() - start
+                await self._out_queue.put((*res, infer_t))
+            except Exception as e:
+                logger.error(f"Inference worker error: {e}")
+                # Put a dummy result to avoid hanging
+                dummy_obs = self.obs if self.obs is not None else torch.zeros(3, 150, 600)
+                await self._out_queue.put((dummy_obs, 0.0, False, False, {"error": str(e)}, 0.0))
+
 # Global game engine instance
 game_engine = WebGameEngine()
 

src/csgo/spawn/0/act.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4170e93000886d1dfd379ed48f1c360897c80700db20996ea4cd5ba1464423eb
+size 208

src/csgo/spawn/0/full_res.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4170e93000886d1dfd379ed48f1c360897c80700db20996ea4cd5ba1464423eb
+size 208

src/csgo/spawn/0/info.json

@@ -0,0 +1 @@
+{"dummy": true}

src/csgo/spawn/0/low_res.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4170e93000886d1dfd379ed48f1c360897c80700db20996ea4cd5ba1464423eb
+size 208

src/csgo/spawn/0/next_act.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4170e93000886d1dfd379ed48f1c360897c80700db20996ea4cd5ba1464423eb
+size 208