Initial Diamond CSGO AI deployment

Dockerfile

@@ -0,0 +1,38 @@
+# Use Python 3.9 slim image
+FROM python:3.9-slim
+
+# Set working directory
+WORKDIR /app
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    build-essential \
+    curl \
+    software-properties-common \
+    git \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy requirements first for better caching
+COPY requirements.txt .
+
+# Install Python dependencies
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy source code
+COPY . .
+
+# Create necessary directories
+RUN mkdir -p csgo/spawn config checkpoints cache
+
+# Set environment variables
+ENV PYTHONPATH=/app/src:/app
+ENV CUDA_VISIBLE_DEVICES=""
+
+# Expose port
+EXPOSE 7860
+
+# Health check
+HEALTHCHECK CMD curl --fail http://localhost:7860/ || exit 1
+
+# Run the application
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -1,11 +1,76 @@
----
-title: Diamond Ai Player
-emoji: 🏆
-colorFrom: yellow
-colorTo: indigo
-sdk: docker
-pinned: false
-license: apache-2.0
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Diamond CSGO AI Player 🎮
+
+A web-based demo of the Diamond AI agent playing Counter-Strike: Global Offensive using diffusion models and reinforcement learning.
+
+## Features
+
+- **Real-time Keyboard Input**: Use standard WASD controls and other keys to interact
+- **AI Agent**: Pre-trained agent using diffusion-based world models
+- **Web Interface**: No installation required, play directly in your browser
+- **Live Visualization**: See the AI's perspective and actions in real-time
+
+## Controls
+
+### Movement
+- **W** - Move Forward
+- **A** - Move Left  
+- **S** - Move Back
+- **D** - Move Right
+- **Space** - Jump
+- **Ctrl** - Crouch
+- **Shift** - Walk
+
+### Actions
+- **1, 2, 3** - Switch Weapons
+- **R** - Reload
+- **Arrow Keys** - Camera Movement
+- **Left/Right Click** - Primary/Secondary Fire
+
+### Game Controls
+- **M** - Switch between Human/AI control
+- **Enter** - Reset Environment
+
+## How to Play
+
+1. Click on the game canvas to focus it
+2. Use keyboard controls to play
+3. The AI agent will respond to your inputs in real-time
+4. Switch to AI mode to watch the agent play autonomously
+
+## Technical Details
+
+This demo uses:
+- **FastAPI + WebSocket** for real-time communication
+- **PyTorch** for AI model inference
+- **Diffusion Models** for next-frame prediction
+- **World Model Environment** for simulation
+
+The agent was trained using the Diamond framework, which combines:
+- Diffusion-based world models
+- Actor-critic reinforcement learning  
+- Multi-step planning and imagination
+
+## Model Information
+
+The AI agent uses several neural networks:
+- **Denoiser**: Diffusion model for generating next observations
+- **Upsampler**: High-resolution image generation
+- **Reward/End Model**: Predicting game outcomes
+- **Actor-Critic**: Action selection and value estimation
+
+## Citation
+
+This work is based on the Diamond framework. If you use this code, please cite:
+
+```bibtex
+@article{diamond2024,
+  title={Diamond: Diffusion for World Modeling},
+  author={[Authors]},
+  journal={[Journal]},
+  year={2024}
+}
+```
+
+## License
+
+See LICENSE file for details.

app.py

@@ -0,0 +1,969 @@
+"""
+Web-based Diamond CSGO AI Player for Hugging Face Spaces
+Uses FastAPI + WebSocket for real-time keyboard input and game streaming
+"""
+
+import asyncio
+import base64
+import io
+import json
+import logging
+import os
+from pathlib import Path
+from typing import Dict, List, Optional, Set
+
+import cv2
+import numpy as np
+import torch
+import uvicorn
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from fastapi.responses import HTMLResponse
+from fastapi.staticfiles import StaticFiles
+from hydra import compose, initialize
+from hydra.utils import instantiate
+from omegaconf import DictConfig, OmegaConf
+from PIL import Image
+
+# Import your modules
+from src.agent import Agent
+from src.csgo.web_action_processing import WebCSGOAction, web_keys_to_csgo_action_names
+from src.envs import WorldModelEnv
+from src.game.web_play_env import WebPlayEnv
+from config_web import web_config
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Global variables
+app = FastAPI(title="Diamond CSGO AI Player")
+connected_clients: Set[WebSocket] = set()
+
+class WebKeyMap:
+    """Map web key codes to pygame-like keys for CSGO actions"""
+    WEB_TO_CSGO = {
+        'KeyW': 'w',
+        'KeyA': 'a', 
+        'KeyS': 's',
+        'KeyD': 'd',
+        'Space': 'space',
+        'ControlLeft': 'left ctrl',
+        'ShiftLeft': 'left shift',
+        'Digit1': '1',
+        'Digit2': '2', 
+        'Digit3': '3',
+        'KeyR': 'r',
+        'ArrowUp': 'camera_up',
+        'ArrowDown': 'camera_down',
+        'ArrowLeft': 'camera_left',
+        'ArrowRight': 'camera_right'
+    }
+
+class WebGameEngine:
+    """Web-compatible game engine that replaces pygame functionality"""
+    
+    def __init__(self):
+        self.play_env: Optional[WebPlayEnv] = None
+        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)
+        self.frame_count = 0
+        self.ai_frame_count = 0
+        self.last_ai_time = 0
+        self.start_time = 0  # Track when AI started for proper FPS calculation
+        self.pressed_keys: Set[str] = set()
+        self.mouse_x = 0
+        self.mouse_y = 0
+        self.l_click = False
+        self.r_click = False
+        self.should_reset = False
+        self.cached_obs = None  # Cache last observation for frame skipping
+        self.first_inference_done = False  # Track if first inference completed
+        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
+        import time
+        self.time_module = time
+    
+    async def _download_model_async(self, url, filepath):
+        """Download model asynchronously with progress tracking"""
+        import asyncio
+        import concurrent.futures
+        import urllib.request
+        import os
+        
+        def download_with_progress():
+            """Download function that runs in thread pool"""
+            def progress_hook(block_num, block_size, total_size):
+                if total_size > 0:
+                    progress = min(100, (block_num * block_size * 100) / total_size)
+                    self.download_progress = int(progress)
+                    if progress % 10 == 0:  # Log every 10%
+                        logger.info(f"Download progress: {self.download_progress}%")
+            
+            urllib.request.urlretrieve(url, filepath, reporthook=progress_hook)
+            self.download_progress = 100
+            
+        # Run download in thread pool to avoid blocking
+        loop = asyncio.get_event_loop()
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            await loop.run_in_executor(executor, download_with_progress)
+        
+        logger.info("Model download completed!")
+        
+    async def initialize_models(self):
+        """Initialize the AI models and environment"""
+        try:
+            import torch
+            logger.info("Initializing models...")
+            
+            # Setup environment and paths
+            web_config.setup_environment_variables()
+            web_config.create_default_configs()
+            
+            config_path = web_config.get_config_path()
+            logger.info(f"Using config path: {config_path}")
+            
+            # Convert to relative path for Hydra
+            import os
+            relative_config_path = os.path.relpath(config_path)
+            logger.info(f"Relative config path: {relative_config_path}")
+                
+            with initialize(version_base="1.3", config_path=relative_config_path):
+                cfg = compose(config_name="trainer")
+            
+            # Override config for deployment
+            cfg.agent = OmegaConf.load(config_path / "agent" / "csgo.yaml")
+            cfg.env = OmegaConf.load(config_path / "env" / "csgo.yaml") 
+            
+            # Use CPU if no GPU available (for free HF spaces)
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            logger.info(f"Using device: {device}")
+            
+            # Load model checkpoint
+            checkpoint_path = web_config.get_checkpoint_path()
+            if not checkpoint_path.exists():
+                logger.warning(f"No checkpoint found at {checkpoint_path} - using dummy mode")
+                self._init_dummy_mode()
+                return True
+            
+            # Get spawn directory
+            spawn_dir = web_config.get_spawn_dir()
+            
+            # Initialize agent
+            num_actions = cfg.env.num_actions
+            agent = Agent(instantiate(cfg.agent, num_actions=num_actions)).to(device).eval()
+            
+            # Try to load checkpoint (remote or local)
+            try:
+                # First try to download from Hugging Face Hub using direct URL
+                try:
+                    import torch.hub
+                    import os
+                    logger.info("Downloading model from Hugging Face Hub...")
+                    
+                    # Direct download URL (change 'blob' to 'resolve' for direct download)
+                    model_url = "https://huggingface.co/Etadingrui/diamond-1B/resolve/main/agent_epoch_00003.pt"
+                    
+                    # Download to cache directory
+                    cache_dir = "./cache"
+                    os.makedirs(cache_dir, exist_ok=True)
+                    model_cache_path = os.path.join(cache_dir, "agent_epoch_00003.pt")
+                    
+                    # Download if not cached
+                    if not os.path.exists(model_cache_path):
+                        logger.info(f"Downloading 1.53GB model to {model_cache_path}...")
+                        self.loading_status = "Downloading AI model from Hugging Face Hub..."
+                        
+                        # Download with progress tracking in a separate thread
+                        await self._download_model_async(model_url, model_cache_path)
+                    else:
+                        logger.info(f"Using cached model from {model_cache_path}")
+                        self.loading_status = "Loading cached model..."
+                        
+                    # Use the agent's load method which expects a file path
+                    self.loading_status = "Loading model weights..."
+                    agent.load(model_cache_path)
+                    logger.info(f"Successfully loaded checkpoint from HF Hub")
+                    
+                except Exception as hub_error:
+                    logger.warning(f"Failed to download from HF Hub: {hub_error}")
+                    
+                    # Fallback to local checkpoint if available
+                    if checkpoint_path.exists():
+                        logger.info(f"Falling back to local checkpoint: {checkpoint_path}")
+                        agent.load(checkpoint_path)
+                        logger.info(f"Successfully loaded local checkpoint: {checkpoint_path}")
+                    else:
+                        raise FileNotFoundError("No model checkpoint available (local or remote)")
+                        
+            except Exception as e:
+                logger.error(f"Failed to load any checkpoint: {e}")
+                self._init_dummy_mode()
+                return True
+            
+            # Initialize world model environment
+            try:
+                sl = cfg.agent.denoiser.inner_model.num_steps_conditioning
+                if agent.upsampler is not None:
+                    sl = max(sl, cfg.agent.upsampler.inner_model.num_steps_conditioning)
+                wm_env_cfg = instantiate(cfg.world_model_env, num_batches_to_preload=1)
+                wm_env = WorldModelEnv(agent.denoiser, agent.upsampler, agent.rew_end_model, 
+                                     spawn_dir, 1, sl, wm_env_cfg, return_denoising_trajectory=True)
+                
+                # Create play environment
+                self.play_env = WebPlayEnv(agent, wm_env, False, False, False)
+                
+                # 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
+                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)...")
+                    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:
+                            wm_env.upsample_next_obs = torch.compile(wm_env.upsample_next_obs, mode="reduce-overhead")
+                        logger.info("Model compilation enabled successfully!")
+                    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.")
+                
+                # Reset environment
+                self.obs, _ = self.play_env.reset()
+                self.cached_obs = self.obs  # Initialize cache
+                
+                logger.info("Models initialized successfully!")
+                logger.info(f"Initial observation shape: {self.obs.shape if self.obs is not None else 'None'}")
+                self.models_ready = True
+                self.loading_status = "Ready!"
+                return True
+                
+            except Exception as e:
+                logger.error(f"Failed to initialize world model environment: {e}")
+                self._init_dummy_mode()
+                self.models_ready = True
+                self.loading_status = "Using dummy mode"
+                return True
+            
+        except Exception as e:
+            logger.error(f"Failed to initialize models: {e}")
+            import traceback
+            traceback.print_exc()
+            self._init_dummy_mode()
+            self.models_ready = True
+            self.loading_status = "Error - using dummy mode"
+            return True
+    
+    def _init_dummy_mode(self):
+        """Initialize dummy mode for testing without models"""
+        logger.info("Initializing dummy mode...")
+        
+        # Create a test observation
+        height, width = 150, 600
+        img_array = np.zeros((height, width, 3), dtype=np.uint8)
+        
+        # Add test pattern
+        for y in range(height):
+            for x in range(width):
+                img_array[y, x, 0] = (x % 256)  # Red gradient
+                img_array[y, x, 1] = (y % 256)  # Green gradient  
+                img_array[y, x, 2] = ((x + y) % 256)  # Blue pattern
+        
+        # Convert to torch tensor in expected format [-1, 1]
+        tensor = torch.from_numpy(img_array).float().permute(2, 0, 1)  # CHW format
+        tensor = tensor.div(255).mul(2).sub(1)  # Convert to [-1, 1] range
+        tensor = tensor.unsqueeze(0)  # Add batch dimension
+        
+        self.obs = tensor
+        self.play_env = None  # No real environment in dummy mode
+        logger.info("Dummy mode initialized with test pattern")
+    
+    
+    def step_environment(self):
+        """Step the environment with current input state (with intelligent frame skipping)"""
+        if self.play_env is None:
+            # Dummy mode - just return current observation
+            return self.obs, 0.0, False, False, {"mode": "dummy"}
+            
+        try:
+            # Check if reset is requested
+            if self.should_reset:
+                self.reset_environment()
+                self.should_reset = False
+                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()
+            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,
+                    mouse_x=self.mouse_x,
+                    mouse_y=self.mouse_y,
+                    l_click=self.l_click,
+                    r_click=self.r_click
+                )
+                inference_time = self.time_module.time() - inference_start
+                
+                # Log first inference completion
+                if not self.first_inference_done:
+                    self.first_inference_done = True
+                    logger.info(f"First AI inference completed in {inference_time:.2f}s - subsequent inferences will be faster!")
+                
+                # Cache the new observation and update timing
+                self.cached_obs = next_obs
+                self.last_ai_time = current_time
+                self.ai_frame_count += 1
+                
+                # Add AI performance info
+                info = info or {}
+                info["ai_inference"] = True
+                
+                # Calculate proper AI FPS: frames / elapsed time since start
+                elapsed_time = current_time - self.start_time
+                if elapsed_time > 0 and self.ai_frame_count > 0:
+                    ai_fps = self.ai_frame_count / elapsed_time
+                    # Cap at reasonable maximum (shouldn't exceed 100 FPS for AI inference)
+                    info["ai_fps"] = min(ai_fps, 100.0)
+                else:
+                    info["ai_fps"] = 0
+                
+                info["inference_time"] = inference_time
+                
+                return next_obs, reward, done, truncated, info
+            else:
+                # Use cached observation for smoother display without AI overhead
+                obs_to_return = self.cached_obs if self.cached_obs is not None else self.obs
+                
+                # Calculate AI FPS for cached frames too
+                elapsed_time = current_time - self.start_time
+                if elapsed_time > 0 and self.ai_frame_count > 0:
+                    ai_fps = min(self.ai_frame_count / elapsed_time, 100.0)  # Cap at 100 FPS
+                else:
+                    ai_fps = 0
+                
+                return obs_to_return, 0.0, False, False, {"cached": True, "ai_fps": ai_fps}
+            
+        except Exception as e:
+            logger.error(f"Error stepping environment: {e}")
+            obs_to_return = self.cached_obs if self.cached_obs is not None else self.obs
+            return obs_to_return, 0.0, False, False, {"error": str(e)}
+    
+    def reset_environment(self):
+        """Reset the environment"""
+        try:
+            if self.play_env is not None:
+                self.obs, _ = self.play_env.reset()
+                self.cached_obs = self.obs  # Update cache
+                logger.info("Environment reset successfully")
+            else:
+                # Dummy mode - recreate test pattern
+                self._init_dummy_mode()
+                self.cached_obs = self.obs  # Update cache
+                logger.info("Dummy environment reset")
+        except Exception as e:
+            logger.error(f"Error resetting environment: {e}")
+    
+    def request_reset(self):
+        """Request environment reset on next step"""
+        self.should_reset = True
+        logger.info("Environment reset requested")
+    
+    def start_game(self):
+        """Start the game"""
+        self.game_started = True
+        self.start_time = self.time_module.time()  # Reset start time for FPS calculation
+        self.ai_frame_count = 0  # Reset AI frame count
+        logger.info("Game started")
+    
+    def pause_game(self):
+        """Pause/stop the game"""
+        self.game_started = False
+        logger.info("Game paused")
+    
+    def obs_to_base64(self, obs: torch.Tensor) -> str:
+        """Convert observation tensor to base64 image for web display"""
+        if obs is None:
+            return ""
+            
+        try:
+            # Convert tensor to PIL Image
+            if obs.ndim == 4 and obs.size(0) == 1:
+                img_array = obs[0].add(1).div(2).mul(255).byte().permute(1, 2, 0).cpu().numpy()
+            else:
+                img_array = obs.add(1).div(2).mul(255).byte().permute(1, 2, 0).cpu().numpy()
+            
+            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}"
+            
+        except Exception as e:
+            logger.error(f"Error converting observation to base64: {e}")
+            return ""
+    
+    async def game_loop(self):
+        """Main game loop that runs continuously"""
+        self.running = True
+        
+        while self.running:
+            try:
+                # Check if models are ready
+                if not self.models_ready:
+                    # Send loading status to clients
+                    if connected_clients:
+                        loading_data = {
+                            'type': 'loading',
+                            'status': self.loading_status,
+                            'progress': self.download_progress,
+                            'ready': False
+                        }
+                        disconnected = set()
+                        for client in connected_clients.copy():
+                            try:
+                                await client.send_text(json.dumps(loading_data))
+                            except:
+                                disconnected.add(client)
+                        connected_clients.difference_update(disconnected)
+                    
+                    await asyncio.sleep(0.5)  # Check every 500ms during loading
+                    continue
+                
+                # Always send frames, but only step environment if game is started
+                should_send_frame = True
+                
+                if not self.game_started:
+                    # Game not started - just send current observation without stepping
+                    if self.obs is not None and connected_clients:
+                        should_send_frame = True
+                    else:
+                        should_send_frame = False
+                    await asyncio.sleep(0.1)
+                else:
+                    # Game is started - step environment
+                    if self.play_env is None:
+                        await asyncio.sleep(0.1)
+                        continue
+                    
+                    # Step environment with current input state
+                    next_obs, reward, done, truncated, info = self.step_environment()
+                    
+                    if done or truncated:
+                        # Auto-reset when episode ends
+                        self.reset_environment()
+                    else:
+                        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:
+                    # 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)
+                
+                self.frame_count += 1
+                await asyncio.sleep(1.0 / self.fps)  # Control FPS
+                
+            except Exception as e:
+                logger.error(f"Error in game loop: {e}")
+                await asyncio.sleep(0.1)
+
+# Global game engine instance
+game_engine = WebGameEngine()
+
+@app.on_event("startup")
+async def startup_event():
+    """Initialize models when the app starts"""
+    # Start the game loop immediately (it will handle loading state)
+    asyncio.create_task(game_engine.game_loop())
+    
+    # Initialize models in background (non-blocking)
+    asyncio.create_task(game_engine.initialize_models())
+
+@app.get("/", response_class=HTMLResponse)
+async def get_homepage():
+    """Serve the main game interface"""
+    html_content = """
+    <!DOCTYPE html>
+    <html>
+    <head>
+        <title>Diamond CSGO AI Player</title>
+        <style>
+            body {
+                margin: 0;
+                padding: 20px;
+                background: #1a1a1a;
+                color: white;
+                font-family: 'Courier New', monospace;
+                text-align: center;
+            }
+            #gameCanvas {
+                border: 2px solid #00ff00;
+                background: #000;
+                margin: 20px auto;
+                display: block;
+            }
+            #controls {
+                margin: 20px;
+                display: grid;
+                grid-template-columns: 1fr 1fr;
+                gap: 20px;
+                max-width: 800px;
+                margin: 20px auto;
+            }
+            .control-section {
+                background: #2a2a2a;
+                padding: 15px;
+                border-radius: 8px;
+                border: 1px solid #444;
+            }
+            .key-display {
+                background: #333;
+                border: 1px solid #555;
+                padding: 5px 10px;
+                margin: 2px;
+                border-radius: 4px;
+                display: inline-block;
+                min-width: 30px;
+            }
+            .key-pressed {
+                background: #00ff00;
+                color: #000;
+            }
+            #status {
+                margin: 10px;
+                padding: 10px;
+                background: #2a2a2a;
+                border-radius: 4px;
+            }
+            .info {
+                color: #00ff00;
+                margin: 5px 0;
+            }
+        </style>
+    </head>
+    <body>
+        <h1>🎮 Diamond CSGO AI Player</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>
+        
+        <!-- Model Download Progress -->
+        <div id="downloadSection" style="display: none; margin: 20px;">
+            <p id="downloadStatus" style="color: #ffaa00; margin: 10px 0;">📥 Downloading AI model...</p>
+            <div style="background: #333; border-radius: 10px; padding: 3px; width: 100%; max-width: 600px; margin: 0 auto;">
+                <div id="progressBar" style="background: linear-gradient(90deg, #00ff00, #88ff00); height: 20px; border-radius: 7px; width: 0%; transition: width 0.3s;"></div>
+            </div>
+            <p id="progressText" style="color: #aaa; font-size: 14px; margin: 5px 0;">0% - Initializing...</p>
+        </div>
+        
+        <canvas id="gameCanvas" width="600" height="150" tabindex="0"></canvas>
+        
+        <div id="status">
+            <div class="info">Status: <span id="connectionStatus">Connecting...</span></div>
+            <div class="info">Game: <span id="gameStatus">Click to Start</span></div>
+            <div class="info">Frame: <span id="frameCount">0</span> | AI FPS: <span id="aiFps">0</span></div>
+            <div class="info">Reward: <span id="reward">0</span></div>
+        </div>
+        
+        <div id="controls">
+            <div class="control-section">
+                <h3>Movement</h3>
+                <div>
+                    <span class="key-display" id="key-w">W</span> Forward<br>
+                    <span class="key-display" id="key-a">A</span> Left
+                    <span class="key-display" id="key-s">S</span> Back
+                    <span class="key-display" id="key-d">D</span> Right<br>
+                    <span class="key-display" id="key-space">Space</span> Jump
+                    <span class="key-display" id="key-ctrl">Ctrl</span> Crouch
+                    <span class="key-display" id="key-shift">Shift</span> Walk
+                </div>
+            </div>
+            
+            <div class="control-section">
+                <h3>Actions</h3>
+                <div>
+                    <span class="key-display" id="key-1">1</span> Weapon 1<br>
+                    <span class="key-display" id="key-2">2</span> Weapon 2
+                    <span class="key-display" id="key-3">3</span> Weapon 3<br>
+                    <span class="key-display" id="key-r">R</span> Reload<br>
+                    <span class="key-display" id="key-arrows">↑↓←→</span> Camera<br>
+                    <span class="key-display" id="key-enter">Enter</span> Reset Game<br>
+                    <span class="key-display" id="key-esc">Esc</span> Pause/Quit
+                </div>
+            </div>
+        </div>
+        
+        <script>
+            const canvas = document.getElementById('gameCanvas');
+            const ctx = canvas.getContext('2d');
+            const statusEl = document.getElementById('connectionStatus');
+            const gameStatusEl = document.getElementById('gameStatus');
+            const frameEl = document.getElementById('frameCount');
+            const aiFpsEl = document.getElementById('aiFps');
+            const rewardEl = document.getElementById('reward');
+            const loadingEl = document.getElementById('loadingIndicator');
+            const downloadSectionEl = document.getElementById('downloadSection');
+            const downloadStatusEl = document.getElementById('downloadStatus');
+            const progressBarEl = document.getElementById('progressBar');
+            const progressTextEl = document.getElementById('progressText');
+            
+            let ws = null;
+            let pressedKeys = new Set();
+            let gameStarted = false;
+            
+            // Key mapping
+            const keyDisplayMap = {
+                'KeyW': 'key-w',
+                'KeyA': 'key-a', 
+                'KeyS': 'key-s',
+                'KeyD': 'key-d',
+                'Space': 'key-space',
+                'ControlLeft': 'key-ctrl',
+                'ShiftLeft': 'key-shift',
+                'Digit1': 'key-1',
+                'Digit2': 'key-2',
+                'Digit3': 'key-3',
+                'KeyR': 'key-r',
+                'ArrowUp': 'key-arrows',
+                'ArrowDown': 'key-arrows',
+                'ArrowLeft': 'key-arrows',
+                'ArrowRight': 'key-arrows',
+                'Enter': 'key-enter',
+                'Escape': 'key-esc'
+            };
+            
+            function connectWebSocket() {
+                const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
+                const wsUrl = `${protocol}//${window.location.host}/ws`;
+                
+                ws = new WebSocket(wsUrl);
+                
+                ws.onopen = function(event) {
+                    statusEl.textContent = 'Connected';
+                    statusEl.style.color = '#00ff00';
+                };
+                
+                ws.onmessage = function(event) {
+                    const data = JSON.parse(event.data);
+                    
+                    if (data.type === 'loading') {
+                        // Handle loading status
+                        downloadSectionEl.style.display = 'block';
+                        downloadStatusEl.textContent = data.status;
+                        
+                        if (data.progress !== undefined) {
+                            progressBarEl.style.width = data.progress + '%';
+                            progressTextEl.textContent = data.progress + '% - ' + data.status;
+                        } else {
+                            progressTextEl.textContent = data.status;
+                        }
+                        
+                        gameStatusEl.textContent = 'Loading Models...';
+                        gameStatusEl.style.color = '#ffaa00';
+                        
+                    } else if (data.type === 'frame') {
+                        // Hide loading indicators once we get frames
+                        downloadSectionEl.style.display = 'none';
+                        // Update frame display
+                        if (data.image) {
+                            const img = new Image();
+                            img.onload = function() {
+                                ctx.clearRect(0, 0, canvas.width, canvas.height);
+                                ctx.drawImage(img, 0, 0, canvas.width, canvas.height);
+                            };
+                            img.src = data.image;
+                        }
+                        
+                        frameEl.textContent = data.frame_count;
+                        rewardEl.textContent = data.reward.toFixed(2);
+                        
+                        // Update AI FPS display and hide loading indicator once AI starts
+                        if (data.ai_fps !== undefined && data.ai_fps !== null) {
+                            // Ensure FPS value is reasonable
+                            const aiFps = Math.min(Math.max(data.ai_fps, 0), 100);
+                            aiFpsEl.textContent = aiFps.toFixed(1);
+                            
+                            // Color code AI FPS for performance indication
+                            if (aiFps >= 8) {
+                                aiFpsEl.style.color = '#00ff00';  // Green for good performance
+                            } else if (aiFps >= 5) {
+                                aiFpsEl.style.color = '#ffff00';  // Yellow for moderate performance
+                            } else if (aiFps > 0) {
+                                aiFpsEl.style.color = '#ff0000';  // Red for poor performance
+                            } else {
+                                aiFpsEl.style.color = '#888888';  // Gray for inactive
+                            }
+                            
+                            // Hide loading indicator once AI inference starts working
+                            if (aiFps > 0 && gameStarted) {
+                                loadingEl.style.display = 'none';
+                                gameStatusEl.textContent = 'Playing';
+                                gameStatusEl.style.color = '#00ff00';
+                            }
+                        }
+                    }
+                };
+                
+                ws.onclose = function(event) {
+                    statusEl.textContent = 'Disconnected';
+                    statusEl.style.color = '#ff0000';
+                    setTimeout(connectWebSocket, 1000); // Reconnect after 1 second
+                };
+                
+                ws.onerror = function(event) {
+                    statusEl.textContent = 'Error';
+                    statusEl.style.color = '#ff0000';
+                };
+            }
+            
+            function sendKeyState() {
+                if (ws && ws.readyState === WebSocket.OPEN) {
+                    ws.send(JSON.stringify({
+                        type: 'keys',
+                        keys: Array.from(pressedKeys)
+                    }));
+                }
+            }
+            
+            function startGame() {
+                if (ws && ws.readyState === WebSocket.OPEN) {
+                    ws.send(JSON.stringify({
+                        type: 'start'
+                    }));
+                    gameStarted = true;
+                    gameStatusEl.textContent = 'Starting AI...';
+                    gameStatusEl.style.color = '#ffff00';
+                    loadingEl.style.display = 'block';
+                    console.log('Game started');
+                }
+            }
+            
+            function pauseGame() {
+                if (ws && ws.readyState === WebSocket.OPEN) {
+                    ws.send(JSON.stringify({
+                        type: 'pause'
+                    }));
+                    gameStarted = false;
+                    gameStatusEl.textContent = 'Paused - Click to Resume';
+                    gameStatusEl.style.color = '#ffff00';
+                    console.log('Game paused');
+                }
+            }
+            
+            function updateKeyDisplay() {
+                // Reset all key displays
+                Object.values(keyDisplayMap).forEach(id => {
+                    const el = document.getElementById(id);
+                    if (el) el.classList.remove('key-pressed');
+                });
+                
+                // Highlight pressed keys
+                pressedKeys.forEach(key => {
+                    const displayId = keyDisplayMap[key];
+                    if (displayId) {
+                        const el = document.getElementById(displayId);
+                        if (el) el.classList.add('key-pressed');
+                    }
+                });
+            }
+            
+            // Focus canvas and handle keyboard events
+            canvas.addEventListener('click', () => {
+                canvas.focus();
+                if (!gameStarted) {
+                    startGame();
+                }
+            });
+            
+            canvas.addEventListener('keydown', (event) => {
+                event.preventDefault();
+                
+                // Handle special keys
+                if (event.code === 'Enter') {
+                    if (ws && ws.readyState === WebSocket.OPEN) {
+                        ws.send(JSON.stringify({
+                            type: 'reset'
+                        }));
+                        console.log('Environment reset requested');
+                    }
+                    // Add to pressedKeys for visual feedback
+                    pressedKeys.add(event.code);
+                    updateKeyDisplay();
+                    
+                    // Remove Enter from pressedKeys after a short delay for visual feedback
+                    setTimeout(() => {
+                        pressedKeys.delete(event.code);
+                        updateKeyDisplay();
+                    }, 200);
+                } else if (event.code === 'Escape') {
+                    pauseGame();
+                    // Add to pressedKeys for visual feedback
+                    pressedKeys.add(event.code);
+                    updateKeyDisplay();
+                    
+                    // Remove ESC from pressedKeys after a short delay for visual feedback
+                    setTimeout(() => {
+                        pressedKeys.delete(event.code);
+                        updateKeyDisplay();
+                    }, 200);
+                } else {
+                    // Only send game keys if game is started
+                    if (gameStarted) {
+                        pressedKeys.add(event.code);
+                        updateKeyDisplay();
+                        sendKeyState();
+                    }
+                }
+            });
+            
+            canvas.addEventListener('keyup', (event) => {
+                event.preventDefault();
+                
+                // Don't handle special keys release (handled in keydown with timeout)
+                if (event.code !== 'Enter' && event.code !== 'Escape') {
+                    if (gameStarted) {
+                        pressedKeys.delete(event.code);
+                        updateKeyDisplay();
+                        sendKeyState();
+                    }
+                }
+            });
+            
+            // Handle mouse events for clicks
+            canvas.addEventListener('mousedown', (event) => {
+                if (ws && ws.readyState === WebSocket.OPEN) {
+                    ws.send(JSON.stringify({
+                        type: 'mouse',
+                        button: event.button,
+                        action: 'down',
+                        x: event.offsetX,
+                        y: event.offsetY
+                    }));
+                }
+            });
+            
+            canvas.addEventListener('mouseup', (event) => {
+                if (ws && ws.readyState === WebSocket.OPEN) {
+                    ws.send(JSON.stringify({
+                        type: 'mouse',
+                        button: event.button,
+                        action: 'up',
+                        x: event.offsetX,
+                        y: event.offsetY
+                    }));
+                }
+            });
+            
+            // Initialize
+            connectWebSocket();
+            canvas.focus();
+        </script>
+    </body>
+    </html>
+    """
+    return html_content
+
+@app.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    """Handle WebSocket connections for real-time game communication"""
+    await websocket.accept()
+    connected_clients.add(websocket)
+    
+    try:
+        while True:
+            # Receive messages from client
+            data = await websocket.receive_text()
+            message = json.loads(data)
+            
+            if message['type'] == 'keys':
+                # Update pressed keys
+                game_engine.pressed_keys = set(message['keys'])
+                
+            elif message['type'] == 'reset':
+                # Handle environment reset request
+                game_engine.request_reset()
+                
+            elif message['type'] == 'start':
+                # Handle game start request
+                game_engine.start_game()
+                
+            elif message['type'] == 'pause':
+                # Handle game pause request
+                game_engine.pause_game()
+                
+            elif message['type'] == 'mouse':
+                # Handle mouse events
+                if message['action'] == 'down':
+                    if message['button'] == 0:  # Left click
+                        game_engine.l_click = True
+                    elif message['button'] == 2:  # Right click
+                        game_engine.r_click = True
+                elif message['action'] == 'up':
+                    if message['button'] == 0:  # Left click
+                        game_engine.l_click = False
+                    elif message['button'] == 2:  # Right click
+                        game_engine.r_click = False
+                        
+                # Update mouse position (relative to canvas)
+                game_engine.mouse_x = message.get('x', 0) - 300  # Center at 300px
+                game_engine.mouse_y = message.get('y', 0) - 150  # Center at 150px
+                
+    except WebSocketDisconnect:
+        connected_clients.discard(websocket)
+    except Exception as e:
+        logger.error(f"WebSocket error: {e}")
+        connected_clients.discard(websocket)
+
+if __name__ == "__main__":
+    # For local development
+    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)
+

config/agent/csgo.yaml

@@ -0,0 +1,34 @@
+_target_: agent.AgentConfig
+
+denoiser:
+  _target_: models.diffusion.DenoiserConfig
+  sigma_data: 0.5
+  sigma_offset_noise: 0.1
+  noise_previous_obs: true
+  upsampling_factor: null
+  inner_model:
+    _target_: models.diffusion.InnerModelConfig
+    img_channels: 3
+    num_steps_conditioning: 4
+    cond_channels: 2048
+    depths: [2, 2, 2, 2]
+    channels: [128, 256, 512, 1024]
+    attn_depths: [0, 0, 1, 1]
+
+upsampler:
+  _target_: models.diffusion.DenoiserConfig
+  sigma_data: 0.5
+  sigma_offset_noise: 0.1
+  noise_previous_obs: false
+  upsampling_factor: 5
+  inner_model:
+    _target_: models.diffusion.InnerModelConfig
+    img_channels: 3
+    num_steps_conditioning: 1
+    cond_channels: 2048
+    depths: [2, 2, 2, 2]
+    channels: [64, 64, 128, 256]
+    attn_depths: [0, 0, 0, 1]
+
+rew_end_model: null
+actor_critic: null

config/env/csgo.yaml

@@ -0,0 +1,7 @@
+train:
+  id: csgo
+  size: [150, 600]
+num_actions: 51
+path_data_low_res: /tmp/dummy_data_low_res
+path_data_full_res: /tmp/dummy_data_full_res
+keymap: csgo

config/trainer.yaml

@@ -0,0 +1,9 @@
+defaults:
+  - _self_
+  - env: csgo
+  - agent: csgo
+  - world_model_env: fast
+
+static_dataset:
+  path: /tmp/dummy_data_low_res
+  ignore_sample_weights: True

config/world_model_env/fast.yaml

@@ -0,0 +1,17 @@
+_target_: envs.WorldModelEnvConfig
+horizon: 1000
+num_batches_to_preload: 1
+diffusion_sampler_next_obs:
+  _target_: models.diffusion.DiffusionSamplerConfig
+  num_steps_denoising: 6  # Balanced: better quality than 3, faster than 10
+  sigma_min: 0.002
+  sigma_max: 5.0
+  rho: 7
+  order: 1
+diffusion_sampler_upsampling:
+  _target_: models.diffusion.DiffusionSamplerConfig
+  num_steps_denoising: 4  # Balanced: better quality than 2, faster than 5
+  sigma_min: 0.002
+  sigma_max: 5.0
+  rho: 7
+  order: 1

config_web.py

@@ -0,0 +1,208 @@
+"""
+Configuration helper for web deployment
+Handles path resolution and model loading for deployment
+"""
+
+import os
+from pathlib import Path
+from typing import Optional
+import logging
+
+logger = logging.getLogger(__name__)
+
+class WebConfig:
+    """Configuration manager for web deployment"""
+    
+    def __init__(self, base_path: Optional[Path] = None):
+        if base_path is None:
+            base_path = Path.cwd()
+        self.base_path = Path(base_path)
+        
+    def get_config_path(self) -> Path:
+        """Get configuration directory path"""
+        # Try multiple possible locations
+        possible_paths = [
+            self.base_path / "config",
+            self.base_path / "src" / ".." / "config",
+            Path(__file__).parent / "config"
+        ]
+        
+        for path in possible_paths:
+            if path.exists():
+                return path.resolve()
+                
+        # Create default config directory
+        config_path = self.base_path / "config"
+        config_path.mkdir(exist_ok=True)
+        return config_path
+    
+    def get_checkpoint_path(self) -> Path:
+        """Find and return the best available checkpoint"""
+        # Try different possible locations and names
+        possible_checkpoints = [
+            self.base_path / "agent_epoch_00003.pt",
+            self.base_path / "agent_epoch_00003.pt",
+            self.base_path / "checkpoints" / "agent_epoch_00003.pt", 
+            self.base_path / "checkpoints" / "agent_epoch_00003.pt",
+            self.base_path / "checkpoints" / "latest.pt",
+        ]
+        
+        for ckpt_path in possible_checkpoints:
+            if ckpt_path.exists():
+                logger.info(f"Found checkpoint: {ckpt_path}")
+                return ckpt_path
+        
+        # If no checkpoint found, create a dummy message
+        logger.warning("No checkpoint found - you may need to download models")
+        return self.base_path / "checkpoints" / "model_not_found.pt"
+    
+    def get_spawn_dir(self) -> Path:
+        """Get spawn data directory"""
+        spawn_dir = self.base_path / "csgo" / "spawn"
+        spawn_dir.mkdir(parents=True, exist_ok=True)
+        
+        # Create dummy spawn data if it doesn't exist
+        spawn_subdir = spawn_dir / "0"
+        spawn_subdir.mkdir(exist_ok=True)
+        
+        # Create dummy files if they don't exist
+        dummy_files = ["act.npy", "full_res.npy", "info.json", "low_res.npy", "next_act.npy"]
+        for filename in dummy_files:
+            file_path = spawn_subdir / filename
+            if not file_path.exists():
+                if filename.endswith('.npy'):
+                    import numpy as np
+                    np.save(file_path, np.zeros((1, 10)))  # Dummy array
+                elif filename.endswith('.json'):
+                    import json
+                    with open(file_path, 'w') as f:
+                        json.dump({"dummy": True}, f)
+        
+        return spawn_dir
+    
+    def setup_environment_variables(self):
+        """Set up environment variables for deployment"""
+        # Disable CUDA if not available (for CPU-only deployment)
+        if not self.has_cuda():
+            os.environ["CUDA_VISIBLE_DEVICES"] = ""
+            
+        # Set Python path
+        python_path = str(self.base_path / "src")
+        current_path = os.environ.get("PYTHONPATH", "")
+        if python_path not in current_path:
+            os.environ["PYTHONPATH"] = f"{python_path}:{current_path}" if current_path else python_path
+    
+    def has_cuda(self) -> bool:
+        """Check if CUDA is available"""
+        try:
+            import torch
+            return torch.cuda.is_available()
+        except ImportError:
+            return False
+    
+    def create_default_configs(self):
+        """Create default configuration files if they don't exist"""
+        config_dir = self.get_config_path()
+        
+        # Create agent config
+        agent_dir = config_dir / "agent"
+        agent_dir.mkdir(exist_ok=True)
+        
+        agent_config_path = agent_dir / "csgo.yaml"
+        if not agent_config_path.exists():
+            with open(agent_config_path, 'w') as f:
+                f.write("""_target_: agent.AgentConfig
+
+denoiser:
+  _target_: models.diffusion.DenoiserConfig
+  sigma_data: 0.5
+  sigma_offset_noise: 0.1
+  noise_previous_obs: true
+  upsampling_factor: null
+  inner_model:
+    _target_: models.diffusion.InnerModelConfig
+    img_channels: 3
+    num_steps_conditioning: 4
+    cond_channels: 2048
+    depths: [2, 2, 2, 2]
+    channels: [128, 256, 512, 1024]
+    attn_depths: [0, 0, 1, 1]
+
+upsampler:
+  _target_: models.diffusion.DenoiserConfig
+  sigma_data: 0.5
+  sigma_offset_noise: 0.1
+  noise_previous_obs: false
+  upsampling_factor: 5
+  inner_model:
+    _target_: models.diffusion.InnerModelConfig
+    img_channels: 3
+    num_steps_conditioning: 1
+    cond_channels: 2048
+    depths: [2, 2, 2, 2]
+    channels: [64, 64, 128, 256]
+    attn_depths: [0, 0, 0, 1]
+
+rew_end_model: null
+actor_critic: null
+""")
+        
+        # Create env config
+        env_dir = config_dir / "env"
+        env_dir.mkdir(exist_ok=True)
+        
+        env_config_path = env_dir / "csgo.yaml"
+        if not env_config_path.exists():
+            with open(env_config_path, 'w') as f:
+                f.write("""train:
+  id: csgo
+  size: [150, 600]
+num_actions: 51
+path_data_low_res: /tmp/dummy_data_low_res
+path_data_full_res: /tmp/dummy_data_full_res
+keymap: csgo
+""")
+        
+        # Create world model env config
+        wm_env_dir = config_dir / "world_model_env"
+        wm_env_dir.mkdir(exist_ok=True)
+        
+        wm_config_path = wm_env_dir / "fast.yaml"
+        if not wm_config_path.exists():
+            with open(wm_config_path, 'w') as f:
+                f.write("""_target_: envs.WorldModelEnvConfig
+horizon: 1000
+num_batches_to_preload: 1
+diffusion_sampler_next_obs:
+  _target_: models.diffusion.DiffusionSamplerConfig
+  num_steps_denoising: 10
+  sigma_min: 0.002
+  sigma_max: 5.0
+  rho: 7
+  order: 1
+diffusion_sampler_upsampling:
+  _target_: models.diffusion.DiffusionSamplerConfig
+  num_steps_denoising: 5
+  sigma_min: 0.002
+  sigma_max: 5.0
+  rho: 7
+  order: 1
+""")
+        
+        # Create trainer config
+        trainer_config_path = config_dir / "trainer.yaml"
+        if not trainer_config_path.exists():
+            with open(trainer_config_path, 'w') as f:
+                f.write("""defaults:
+  - _self_
+  - env: csgo
+  - agent: csgo
+  - world_model_env: fast
+
+static_dataset:
+  path: /tmp/dummy_data_low_res
+  ignore_sample_weights: True
+""")
+
+# Global config instance
+web_config = WebConfig()

csgo/spawn/0/act.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:11830620c54f47d0ee6a9f904e68516980f8cd5af488572bd6e9e4815e8be52d
+size 332

csgo/spawn/0/full_res.npy

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

csgo/spawn/0/info.json

@@ -0,0 +1 @@
+{"original_file_id": "4001-4200/hdf5_dm_july2021_4143.hdf5", "timestep_start": 540}

csgo/spawn/0/low_res.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4d775f579f104caf9e195fa12bdf302e54c3c0f8938483ace0fa2cb75b694be1
+size 43328

csgo/spawn/0/next_act.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:762d9db84444e12912a8e535d10b29783db59f6a7f97579c933496354ffb4bb6
+size 10328

packages.txt

@@ -0,0 +1,3 @@
+build-essential
+curl
+git

requirements.txt

@@ -0,0 +1,32 @@
+# Core ML dependencies
+torch>=1.13.0
+torchvision>=0.14.0
+torchaudio>=0.13.0
+numpy>=1.21.0
+
+# Configuration management
+hydra-core>=1.2.0
+omegaconf>=2.2.0
+
+# Web framework for deployment
+fastapi>=0.68.0
+uvicorn>=0.15.0
+websockets>=10.0
+
+# Image processing
+opencv-python-headless>=4.5.0
+Pillow>=8.0.0
+
+# Hugging Face integration
+huggingface-hub>=0.10.0
+
+# Data handling
+h5py>=3.7.0
+
+# Optional: for better performance
+# torch-audio  # if needed for audio processing
+
+# Development dependencies (uncomment for local development)
+# pytest>=6.0.0
+# black>=21.0.0
+# isort>=5.0.0

src/agent.py

@@ -0,0 +1,74 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional, Union
+
+import torch
+import torch.nn as nn
+
+from envs import TorchEnv, WorldModelEnv
+from models.actor_critic import ActorCritic, ActorCriticConfig, ActorCriticLossConfig
+from models.diffusion import Denoiser, DenoiserConfig, SigmaDistributionConfig
+from models.rew_end_model import RewEndModel, RewEndModelConfig
+from utils import extract_state_dict
+
+
+@dataclass
+class AgentConfig:
+    denoiser: DenoiserConfig
+    upsampler: Optional[DenoiserConfig]
+    rew_end_model: Optional[RewEndModelConfig]
+    actor_critic: Optional[ActorCriticConfig]
+    num_actions: int
+
+    def __post_init__(self) -> None:
+        self.denoiser.inner_model.num_actions = self.num_actions
+        if self.upsampler is not None:
+            self.upsampler.inner_model.num_actions = self.num_actions
+        if self.rew_end_model is not None:
+            self.rew_end_model.num_actions = self.num_actions
+        if self.actor_critic is not None:
+            self.actor_critic.num_actions = self.num_actions
+
+
+class Agent(nn.Module):
+    def __init__(self, cfg: AgentConfig) -> None:
+        super().__init__()
+        self.denoiser = Denoiser(cfg.denoiser)
+        self.upsampler = Denoiser(cfg.upsampler) if cfg.upsampler is not None else None
+        self.rew_end_model = RewEndModel(cfg.rew_end_model) if cfg.rew_end_model is not None else None
+        self.actor_critic = ActorCritic(cfg.actor_critic) if cfg.actor_critic is not None else None
+
+    @property
+    def device(self):
+        return self.denoiser.device
+
+    def setup_training(
+        self,
+        sigma_distribution_cfg: SigmaDistributionConfig,
+        sigma_distribution_cfg_upsampler: Optional[SigmaDistributionConfig],
+        actor_critic_loss_cfg: Optional[ActorCriticLossConfig],
+        rl_env: Optional[Union[TorchEnv, WorldModelEnv]],
+    ) -> None:
+        self.denoiser.setup_training(sigma_distribution_cfg)
+        if self.upsampler is not None:
+            self.upsampler.setup_training(sigma_distribution_cfg_upsampler)
+        if self.actor_critic is not None:
+            self.actor_critic.setup_training(rl_env, actor_critic_loss_cfg)
+
+    def load(
+        self,
+        path_to_ckpt: Path,
+        load_denoiser: bool = True,
+        load_upsampler: bool = True,
+        load_rew_end_model: bool = True,
+        load_actor_critic: bool = True,
+    ) -> None:
+        sd = torch.load(Path(path_to_ckpt), map_location=self.device)
+        if load_denoiser:
+            self.denoiser.load_state_dict(extract_state_dict(sd, "denoiser"))
+        if load_upsampler:
+            self.upsampler.load_state_dict(extract_state_dict(sd, "upsampler"))
+        if load_rew_end_model and self.rew_end_model is not None:
+            self.rew_end_model.load_state_dict(extract_state_dict(sd, "rew_end_model"))
+        if load_actor_critic and self.actor_critic is not None:
+            self.actor_critic.load_state_dict(extract_state_dict(sd, "actor_critic"))

src/coroutines/__init__.py

@@ -0,0 +1,11 @@
+from functools import wraps
+
+
+def coroutine(func):
+    @wraps(func)
+    def primer(*args, **kwargs):
+        gen = func(*args, **kwargs)
+        next(gen)
+        return gen
+
+    return primer

src/coroutines/collector.py

@@ -0,0 +1,126 @@
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Generator, Optional
+
+import torch
+import torch.nn as nn
+from tqdm import tqdm
+
+from . import coroutine
+from data import Episode, Dataset
+from envs import TorchEnv
+from .env_loop import make_env_loop
+from utils import Logs
+
+
+@coroutine
+def make_collector(
+    env: TorchEnv,
+    model: nn.Module,
+    dataset: Dataset,
+    epsilon: float = 0.0,
+    reset_every_collect: bool = False,
+    verbose: bool = True,
+) -> Generator[Logs, int, None]:
+    num_envs = env.num_envs
+
+    env_loop, buffer, episode_ids, dead = (None,) * 4
+    num_steps, num_episodes, to_log, pbar = (None,) * 4
+
+    def setup_new_collect():
+        nonlocal num_steps, num_episodes, buffer, to_log, pbar
+        num_steps = 0
+        num_episodes = 0
+        buffer = defaultdict(list)
+        to_log = []
+        pbar = tqdm(
+            total=num_to_collect.total,
+            unit=num_to_collect.unit,
+            desc=f"Collect {dataset.name}",
+            disable=not verbose,
+        )
+
+    def reset():
+        nonlocal env_loop, episode_ids, dead
+        env_loop = make_env_loop(env, model, epsilon)
+        episode_ids = defaultdict(lambda: None)
+        dead = [None] * num_envs
+
+    num_to_collect = yield
+    setup_new_collect()
+    reset()
+
+    while True:
+        with torch.no_grad():
+            all_obs, act, rew, end, trunc, *_, [infos] = env_loop.send(1)
+
+        num_steps += num_envs
+        pbar.update(num_envs if num_to_collect.steps is not None else 0)
+
+        for i, (o, a, r, e, t) in enumerate(zip(all_obs, act, rew, end, trunc)):
+            buffer[i].append((o, a, r, e, t))
+            dead[i] = (e + t).clip(max=1).item()
+
+        num_episodes += sum(dead)
+
+        can_stop = num_to_collect.can_stop(num_steps, num_episodes)
+
+        count_dead = 0
+        for i in range(num_envs):
+            # Store incomplete episodes only when reset_every_collect is set to False (train)
+            add_to_dataset = dead[i] or (can_stop and not reset_every_collect)
+            if add_to_dataset:
+                info = {"final_observation": infos["final_observation"][count_dead]} if dead[i] else {}
+                ep = Episode(*(torch.cat(x, dim=0) for x in zip(*buffer[i])), info).to("cpu")
+                if episode_ids[i] is not None:
+                    ep = dataset.load_episode(episode_ids[i]) + ep
+                episode_ids[i] = dataset.add_episode(ep, episode_id=episode_ids[i])
+
+            if dead[i]:
+                to_log.append(
+                    {
+                        f"{dataset.name}/episode_id": episode_ids[i],
+                        **ep.compute_metrics(),
+                    }
+                )
+                buffer[i] = []
+                episode_ids[i] = None
+                pbar.update(1 if num_to_collect.episodes is not None else 0)
+
+            count_dead += dead[i]
+
+        if can_stop:
+            pbar.close()
+            metrics = {
+                "num_steps": dataset.num_steps,
+                "counts/rew_-1": dataset.counts_rew[0],
+                "counts/rew__0": dataset.counts_rew[1],
+                "counts/rew_+1": dataset.counts_rew[2],
+                "counts/end_0": dataset.counts_end[0],
+                "counts/end_1": dataset.counts_end[1],
+            }
+            to_log.append({f"{dataset.name}/{k}": v for k, v in metrics.items()})
+            num_to_collect = yield to_log
+            setup_new_collect()
+            if reset_every_collect:
+                reset()
+
+
+@dataclass
+class NumToCollect:
+    steps: Optional[int] = None
+    episodes: Optional[int] = None
+
+    def __post_init__(self) -> None:
+        assert (self.steps is None) != (self.episodes is None)
+
+    def can_stop(self, num_steps: int, num_episodes: int) -> bool:
+        return num_steps >= self.steps if self.steps is not None else num_episodes >= self.episodes
+
+    @property
+    def unit(self) -> str:
+        return "steps" if self.steps is not None else "eps"
+
+    @property
+    def total(self) -> int:
+        return self.steps if self.steps is not None else self.episodes

src/coroutines/env_loop.py

@@ -0,0 +1,74 @@
+import random
+from typing import Generator, Tuple, Union
+
+import torch
+import torch.nn as nn
+from torch.distributions.categorical import Categorical
+
+from . import coroutine
+from envs import TorchEnv, WorldModelEnv
+
+
+@coroutine
+def make_env_loop(
+    env: Union[TorchEnv, WorldModelEnv], model: nn.Module, epsilon: float = 0.0
+) -> Generator[Tuple[torch.Tensor, ...], int, None]:
+    num_steps = yield
+
+    hx = torch.zeros(env.num_envs, model.lstm_dim, device=model.device)
+    cx = torch.zeros(env.num_envs, model.lstm_dim, device=model.device)
+
+    seed = random.randint(0, 2**31 - 1)
+    obs, _ = env.reset(seed=[seed + i for i in range(env.num_envs)])
+
+    while True:
+        hx, cx = hx.detach(), cx.detach()
+        all_ = []
+        infos = []
+        n = 0
+
+        while n < num_steps:
+            logits_act, val, (hx, cx) = model.predict_act_value(obs, (hx, cx))
+            act = Categorical(logits=logits_act).sample()
+
+            if random.random() < epsilon:
+                act = torch.randint(low=0, high=env.num_actions, size=(obs.size(0),), device=obs.device)
+
+            next_obs, rew, end, trunc, info = env.step(act)
+
+            if n > 0:
+                val_bootstrap = val.detach().clone()
+                if dead.any():
+                    val_bootstrap[dead] = val_final_obs
+                all_[-1][-1] = val_bootstrap
+
+            dead = torch.logical_or(end, trunc)
+
+            if dead.any():
+                with torch.no_grad():
+                    _, val_final_obs, _ = model.predict_act_value(info["final_observation"], (hx[dead], cx[dead]))
+                reset_gate = 1 - dead.float().unsqueeze(1)
+                hx = hx * reset_gate
+                cx = cx * reset_gate
+                if "burnin_obs" in info:
+                    burnin_obs = info["burnin_obs"]
+                    for i in range(burnin_obs.size(1)):
+                        _, _, (hx[dead], cx[dead]) = model.predict_act_value(burnin_obs[:, i], (hx[dead], cx[dead]))
+
+            all_.append([obs, act, rew, end, trunc, logits_act, val, None])
+            infos.append(info)
+
+            obs = next_obs
+            n += 1
+
+        with torch.no_grad():
+            _, val_bootstrap, _ = model.predict_act_value(next_obs, (hx, cx))  # do not update hx/cx
+
+        if dead.any():
+            val_bootstrap[dead] = val_final_obs
+
+        all_[-1][-1] = val_bootstrap
+
+        all_obs, act, rew, end, trunc, logits_act, val, val_bootstrap = (torch.stack(x, dim=1) for x in zip(*all_))
+
+        num_steps = yield all_obs, act, rew, end, trunc, logits_act, val, val_bootstrap, infos

src/csgo/action_processing.py

@@ -0,0 +1,230 @@
+"""
+Credits: some parts are taken and modified from the file `config.py` from https://github.com/TeaPearce/Counter-Strike_Behavioural_Cloning/
+"""
+
+from dataclasses import dataclass 
+from typing import Dict, List, Set, Tuple
+
+import numpy as np
+import pygame
+import torch
+
+from .keymap import CSGO_FORBIDDEN_COMBINATIONS, CSGO_KEYMAP
+
+
+@dataclass
+class CSGOAction:
+    keys: List[int]
+    mouse_x: float
+    mouse_y: float
+    l_click: bool
+    r_click: bool
+
+    def __post_init__(self) -> None:
+        self.keys = filter_keys_pressed_forbidden(self.keys)
+        self.process_mouse()
+
+    @property
+    def key_names(self) -> List[str]:
+        return [pygame.key.name(key) for key in self.keys] 
+
+    def process_mouse(self) -> None:
+        # Clip and match mouse to closest in list of possibles
+        x = np.clip(self.mouse_x, MOUSE_X_LIM[0], MOUSE_X_LIM[1])
+        y = np.clip(self.mouse_y, MOUSE_Y_LIM[0], MOUSE_Y_LIM[1])
+        self.mouse_x = min(MOUSE_X_POSSIBLES, key=lambda x_: abs(x_ - x))
+        self.mouse_y = min(MOUSE_Y_POSSIBLES, key=lambda x_: abs(x_ - y))
+        
+        # Use arrows to override mouse movements
+        for key in self.key_names:        
+            if key == "left":
+                self.mouse_x = -60
+            elif key == "right":
+                self.mouse_x = +60
+            elif key == "up":
+                self.mouse_y = -50
+            elif key == "down":
+                self.mouse_y = +50
+
+
+def print_csgo_action(action: CSGOAction) -> Tuple[str]:
+    action_names = [CSGO_KEYMAP[k] for k in action.keys] if len(action.keys) > 0 else []
+    action_names = [x for x in action_names if not x.startswith("camera_")]
+    keys   = " + ".join(action_names)     
+    mouse  = str((action.mouse_x, action.mouse_y)) * (action.mouse_x != 0 or action.mouse_y != 0)
+    clicks = "L" * action.l_click + " + " * (action.l_click and action.r_click) + "R" * action.r_click
+    return keys, mouse, clicks     
+
+
+MOUSE_X_POSSIBLES = [
+    -1000,
+    -500,
+    -300,
+    -200,
+    -100,
+    -60,
+    -30,
+    -20,
+    -10,
+    -4,
+    -2,
+    0,
+    2,
+    4,
+    10,
+    20,
+    30,
+    60,
+    100,
+    200,
+    300,
+    500,
+    1000,
+]
+
+MOUSE_Y_POSSIBLES = [
+    -200,
+    -100,
+    -50,
+    -20,
+    -10,
+    -4,
+    -2,
+    0,
+    2,
+    4,
+    10,
+    20,
+    50,
+    100,
+    200,
+]
+
+MOUSE_X_LIM = (MOUSE_X_POSSIBLES[0], MOUSE_X_POSSIBLES[-1])
+MOUSE_Y_LIM = (MOUSE_Y_POSSIBLES[0], MOUSE_Y_POSSIBLES[-1])
+N_KEYS = 11  # number of keyboard outputs, w,s,a,d,space,ctrl,shift,1,2,3,r
+N_CLICKS = 2  # number of mouse buttons, left, right
+N_MOUSE_X = len(MOUSE_X_POSSIBLES)  # number of outputs on mouse x axis
+N_MOUSE_Y = len(MOUSE_Y_POSSIBLES)  # number of outputs on mouse y axis
+
+
+def encode_csgo_action(csgo_action: CSGOAction, device: torch.device) -> torch.Tensor:
+
+    # mouse_x = csgo_action.mouse_x
+    # mouse_y = csgo_action.mouse_y
+
+    keys_pressed_onehot = np.zeros(N_KEYS)
+    mouse_x_onehot = np.zeros(N_MOUSE_X)
+    mouse_y_onehot = np.zeros(N_MOUSE_Y)
+    l_click_onehot = np.zeros(1)
+    r_click_onehot = np.zeros(1)
+
+    for key in csgo_action.key_names:
+        if key == "w":
+            keys_pressed_onehot[0] = 1
+        elif key == "a":
+            keys_pressed_onehot[1] = 1
+        elif key == "s":
+            keys_pressed_onehot[2] = 1
+        elif key == "d":
+            keys_pressed_onehot[3] = 1
+        elif key == "space":
+            keys_pressed_onehot[4] = 1
+        elif key == "left ctrl":
+            keys_pressed_onehot[5] = 1
+        elif key == "left shift":
+            keys_pressed_onehot[6] = 1
+        elif key == "1":
+            keys_pressed_onehot[7] = 1
+        elif key == "2":
+            keys_pressed_onehot[8] = 1
+        elif key == "3":
+            keys_pressed_onehot[9] = 1
+        elif key == "r":
+            keys_pressed_onehot[10] = 1
+
+    l_click_onehot[0] = int(csgo_action.l_click)
+    r_click_onehot[0] = int(csgo_action.r_click)
+
+    mouse_x_onehot[MOUSE_X_POSSIBLES.index(csgo_action.mouse_x)] = 1
+    mouse_y_onehot[MOUSE_Y_POSSIBLES.index(csgo_action.mouse_y)] = 1
+
+    assert mouse_x_onehot.sum() == 1
+    assert mouse_y_onehot.sum() == 1
+
+    return torch.tensor(
+        np.concatenate((
+            keys_pressed_onehot,
+            l_click_onehot,
+            r_click_onehot,
+            mouse_x_onehot,
+            mouse_y_onehot,
+        )),
+        device=device,
+        dtype=torch.float32,
+    )
+    
+
+def decode_csgo_action(y_preds: torch.Tensor) -> CSGOAction:
+    y_preds = y_preds.squeeze()
+    keys_pred = y_preds[0:N_KEYS]
+    l_click_pred = y_preds[N_KEYS : N_KEYS + 1]
+    r_click_pred = y_preds[N_KEYS + 1 : N_KEYS + N_CLICKS]
+    mouse_x_pred = y_preds[N_KEYS + N_CLICKS : N_KEYS + N_CLICKS + N_MOUSE_X]
+    mouse_y_pred = y_preds[
+        N_KEYS + N_CLICKS + N_MOUSE_X : N_KEYS + N_CLICKS + N_MOUSE_X + N_MOUSE_Y
+    ]
+
+    keys_pressed = []
+    keys_pressed_onehot = np.round(keys_pred)
+    if keys_pressed_onehot[0] == 1:
+        keys_pressed.append("w")
+    if keys_pressed_onehot[1] == 1:
+        keys_pressed.append("a")
+    if keys_pressed_onehot[2] == 1:
+        keys_pressed.append("s")
+    if keys_pressed_onehot[3] == 1:
+        keys_pressed.append("d")
+    if keys_pressed_onehot[4] == 1:
+        keys_pressed.append("space")
+    if keys_pressed_onehot[5] == 1:
+        keys_pressed.append("left ctrl")
+    if keys_pressed_onehot[6] == 1:
+        keys_pressed.append("left shift")
+    if keys_pressed_onehot[7] == 1:
+        keys_pressed.append("1")
+    if keys_pressed_onehot[8] == 1:
+        keys_pressed.append("2")
+    if keys_pressed_onehot[9] == 1:
+        keys_pressed.append("3")
+    if keys_pressed_onehot[10] == 1:
+        keys_pressed.append("r")
+
+    l_click = int(np.round(l_click_pred))
+    r_click = int(np.round(r_click_pred))
+
+    id = np.argmax(mouse_x_pred)
+    mouse_x = MOUSE_X_POSSIBLES[id]
+    id = np.argmax(mouse_y_pred)
+    mouse_y = MOUSE_Y_POSSIBLES[id]
+
+    keys_pressed = [pygame.key.key_code(x) for x in keys_pressed]
+
+    return CSGOAction(keys_pressed, mouse_x, mouse_y, bool(l_click), bool(r_click))
+
+
+def filter_keys_pressed_forbidden(keys_pressed: List[int], keymap: Dict[int, str] = CSGO_KEYMAP, forbidden_combinations: List[Set[str]] = CSGO_FORBIDDEN_COMBINATIONS) -> List[int]:
+    keys = set()
+    names = set()
+    for key in keys_pressed:
+        if key not in keymap:
+            continue
+        name = keymap[key]
+        keys.add(key)
+        names.add(name)
+        for forbidden in forbidden_combinations:
+            if forbidden.issubset(names):
+                keys.remove(key)
+                names.remove(name)
+                break
+    return list(filter(lambda key: key in keys, keys_pressed))

src/csgo/keymap.py

@@ -0,0 +1,33 @@
+import pygame
+
+
+CSGO_KEYMAP = {
+    pygame.K_w: "up",
+    pygame.K_d: "right",
+    pygame.K_a: "left",
+    pygame.K_s: "down",
+    pygame.K_SPACE: "jump",
+    pygame.K_LCTRL: "crouch",
+    pygame.K_LSHIFT: "walk",
+    pygame.K_1: "weapon1",
+    pygame.K_2: "weapon2",
+    pygame.K_3: "weapon3",
+    pygame.K_r: "reload",
+
+    # Override mouse movement with arrows
+    pygame.K_UP: "camera_up",
+    pygame.K_RIGHT: "camera_right",
+    pygame.K_LEFT: "camera_left",
+    pygame.K_DOWN: "camera_down",
+}
+
+
+CSGO_FORBIDDEN_COMBINATIONS = [
+    {"up", "down"},
+    {"left", "right"},
+    {"weapon1", "weapon2"},
+    {"weapon1", "weapon3"},
+    {"weapon2", "weapon3"},
+    {"camera_up", "camera_down"},
+    {"camera_left", "camera_right"},
+]

src/csgo/web_action_processing.py

@@ -0,0 +1,167 @@
+"""
+Web-compatible action processing for CSGO actions
+Converts web keyboard inputs to CSGO actions without pygame dependency
+"""
+
+from dataclasses import dataclass
+from typing import Dict, List, Set, Tuple
+
+import numpy as np
+import torch
+
+# Web key code to CSGO action mapping
+WEB_KEYMAP = {
+    'KeyW': "up",
+    'KeyD': "right", 
+    'KeyA': "left",
+    'KeyS': "down",
+    'Space': "jump",
+    'ControlLeft': "crouch",
+    'ShiftLeft': "walk",
+    'Digit1': "weapon1",
+    'Digit2': "weapon2",
+    'Digit3': "weapon3",
+    'KeyR': "reload",
+    'ArrowUp': "camera_up",
+    'ArrowRight': "camera_right",
+    'ArrowLeft': "camera_left",
+    'ArrowDown': "camera_down",
+}
+
+# Forbidden key combinations (same logic as original)
+WEB_FORBIDDEN_COMBINATIONS = [
+    {"up", "down"},
+    {"left", "right"},
+    {"weapon1", "weapon2"},
+    {"weapon1", "weapon3"},
+    {"weapon2", "weapon3"},
+    {"camera_up", "camera_down"},
+    {"camera_left", "camera_right"},
+]
+
+@dataclass
+class WebCSGOAction:
+    """Web-compatible CSGO action without pygame dependencies"""
+    key_names: List[str]  # Use string names instead of pygame key codes
+    mouse_x: float
+    mouse_y: float
+    l_click: bool
+    r_click: bool
+
+    def __post_init__(self) -> None:
+        self.key_names = filter_web_keys_forbidden(self.key_names)
+        self.process_mouse()
+
+    def process_mouse(self) -> None:
+        """Process mouse movement with discretization"""
+        # Import mouse constants
+        from .action_processing import MOUSE_X_POSSIBLES, MOUSE_Y_POSSIBLES, MOUSE_X_LIM, MOUSE_Y_LIM
+        
+        # Clip and match mouse to closest in list of possibles
+        x = np.clip(self.mouse_x, MOUSE_X_LIM[0], MOUSE_X_LIM[1])
+        y = np.clip(self.mouse_y, MOUSE_Y_LIM[0], MOUSE_Y_LIM[1])
+        self.mouse_x = min(MOUSE_X_POSSIBLES, key=lambda x_: abs(x_ - x))
+        self.mouse_y = min(MOUSE_Y_POSSIBLES, key=lambda x_: abs(x_ - y))
+        
+        # Use arrow keys to override mouse movements
+        for key_name in self.key_names:        
+            if key_name == "camera_left":
+                self.mouse_x = -60
+            elif key_name == "camera_right":
+                self.mouse_x = +60
+            elif key_name == "camera_up":
+                self.mouse_y = -50
+            elif key_name == "camera_down":
+                self.mouse_y = +50
+
+def filter_web_keys_forbidden(key_names: List[str]) -> List[str]:
+    """Filter out forbidden key combinations"""
+    names = set(key_names)
+    filtered_names = []
+    
+    for key_name in key_names:
+        # Check if adding this key would create a forbidden combination
+        test_names = set(filtered_names + [key_name])
+        is_forbidden = False
+        
+        for forbidden in WEB_FORBIDDEN_COMBINATIONS:
+            if forbidden.issubset(test_names):
+                is_forbidden = True
+                break
+        
+        if not is_forbidden:
+            filtered_names.append(key_name)
+    
+    return filtered_names
+
+def web_keys_to_csgo_action_names(pressed_web_keys: Set[str]) -> List[str]:
+    """Convert set of pressed web keys to CSGO action names"""
+    action_names = []
+    for web_key in pressed_web_keys:
+        if web_key in WEB_KEYMAP:
+            action_names.append(WEB_KEYMAP[web_key])
+    return action_names
+
+def encode_web_csgo_action(web_action: WebCSGOAction, device: torch.device) -> torch.Tensor:
+    """Encode web CSGO action to tensor format (compatible with original encoding)"""
+    from .action_processing import MOUSE_X_POSSIBLES, MOUSE_Y_POSSIBLES, N_KEYS, N_CLICKS, N_MOUSE_X, N_MOUSE_Y
+    
+    keys_pressed_onehot = np.zeros(N_KEYS)
+    mouse_x_onehot = np.zeros(N_MOUSE_X)
+    mouse_y_onehot = np.zeros(N_MOUSE_Y)
+    l_click_onehot = np.zeros(1)
+    r_click_onehot = np.zeros(1)
+
+    # Map action names to one-hot encoding
+    for action_name in web_action.key_names:
+        if action_name == "up":  # w key
+            keys_pressed_onehot[0] = 1
+        elif action_name == "left":  # a key
+            keys_pressed_onehot[1] = 1
+        elif action_name == "down":  # s key
+            keys_pressed_onehot[2] = 1
+        elif action_name == "right":  # d key
+            keys_pressed_onehot[3] = 1
+        elif action_name == "jump":  # space
+            keys_pressed_onehot[4] = 1
+        elif action_name == "crouch":  # ctrl
+            keys_pressed_onehot[5] = 1
+        elif action_name == "walk":  # shift
+            keys_pressed_onehot[6] = 1
+        elif action_name == "weapon1":  # 1
+            keys_pressed_onehot[7] = 1
+        elif action_name == "weapon2":  # 2
+            keys_pressed_onehot[8] = 1
+        elif action_name == "weapon3":  # 3
+            keys_pressed_onehot[9] = 1
+        elif action_name == "reload":  # r
+            keys_pressed_onehot[10] = 1
+
+    l_click_onehot[0] = int(web_action.l_click)
+    r_click_onehot[0] = int(web_action.r_click)
+
+    mouse_x_onehot[MOUSE_X_POSSIBLES.index(web_action.mouse_x)] = 1
+    mouse_y_onehot[MOUSE_Y_POSSIBLES.index(web_action.mouse_y)] = 1
+
+    assert mouse_x_onehot.sum() == 1
+    assert mouse_y_onehot.sum() == 1
+
+    return torch.tensor(
+        np.concatenate((
+            keys_pressed_onehot,
+            l_click_onehot,
+            r_click_onehot,
+            mouse_x_onehot,
+            mouse_y_onehot,
+        )),
+        device=device,
+        dtype=torch.float32,
+    )
+
+def print_web_csgo_action(action: WebCSGOAction) -> Tuple[str, str, str]:
+    """Print web CSGO action in readable format"""
+    action_names = [name for name in action.key_names if not name.startswith("camera_")]
+    keys = " + ".join(action_names)     
+    mouse = str((action.mouse_x, action.mouse_y)) * (action.mouse_x != 0 or action.mouse_y != 0)
+    clicks = "L" * action.l_click + " + " * (action.l_click and action.r_click) + "R" * action.r_click
+    return keys, mouse, clicks

src/data/__init__.py

@@ -0,0 +1,6 @@
+from .batch import Batch
+from .batch_sampler import BatchSampler
+from .dataset import Dataset, CSGOHdf5Dataset
+from .episode import Episode
+from .segment import Segment, SegmentId
+from .utils import collate_segments_to_batch, DatasetTraverser, make_segment

src/data/batch.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import Any, Dict, List
+
+import torch
+
+from .segment import SegmentId
+
+
+@dataclass
+class Batch:
+    obs: torch.ByteTensor
+    act: torch.LongTensor
+    rew: torch.FloatTensor
+    end: torch.LongTensor
+    trunc: torch.LongTensor
+    mask_padding: torch.BoolTensor
+    info: List[Dict[str, Any]]
+    segment_ids: List[SegmentId]
+
+    def pin_memory(self) -> Batch:
+        return Batch(**{k: v if k in ("segment_ids", "info") else v.pin_memory() for k, v in self.__dict__.items()})
+
+    def to(self, device: torch.device) -> Batch:
+        return Batch(**{k: v if k in ("segment_ids", "info") else v.to(device) for k, v in self.__dict__.items()})

src/data/batch_sampler.py

@@ -0,0 +1,72 @@
+from typing import Generator, List, Optional
+
+import numpy as np
+import torch
+
+from .dataset import CSGOHdf5Dataset, Dataset
+from .segment import SegmentId
+
+
+class BatchSampler(torch.utils.data.Sampler):
+    def __init__(
+        self,
+        dataset: Dataset,
+        rank: int,
+        world_size: int,
+        batch_size: int,
+        seq_length: int,
+        sample_weights: Optional[List[float]] = None,
+        can_sample_beyond_end: bool = False,
+    ) -> None:
+        super().__init__(dataset)
+        assert isinstance(dataset, (Dataset, CSGOHdf5Dataset))
+        self.dataset = dataset
+        self.rank = rank
+        self.world_size = world_size
+        self.sample_weights = sample_weights
+        self.batch_size = batch_size
+        self.seq_length = seq_length
+        self.can_sample_beyond_end = can_sample_beyond_end
+
+    def __len__(self):
+        raise NotImplementedError
+
+    def __iter__(self) -> Generator[List[SegmentId], None, None]:
+        while True:
+            yield self.sample()
+
+    def sample(self) -> List[SegmentId]:
+        num_episodes = self.dataset.num_episodes
+
+        if (self.sample_weights is None) or num_episodes < len(self.sample_weights):
+            weights = self.dataset.lengths / self.dataset.num_steps
+        else:
+            weights = self.sample_weights
+            num_weights = len(self.sample_weights)
+            assert all([0 <= x <= 1 for x in weights]) and sum(weights) == 1
+            sizes = [
+                num_episodes // num_weights + (num_episodes % num_weights) * (i == num_weights - 1)
+                for i in range(num_weights)
+            ]
+            weights = [w / s for (w, s) in zip(weights, sizes) for _ in range(s)]
+
+        episodes_partition = np.arange(self.rank, num_episodes, self.world_size)
+        weights = np.array(weights[self.rank::self.world_size])
+        max_eps = self.batch_size
+        episode_ids = np.random.choice(episodes_partition, size=max_eps, replace=True, p=weights / weights.sum())
+        episode_ids = episode_ids.repeat(self.batch_size // max_eps)
+        timesteps = np.random.randint(low=0, high=self.dataset.lengths[episode_ids])
+
+        # padding allowed, both before start and after end
+        if self.can_sample_beyond_end:
+            starts = timesteps - np.random.randint(0, self.seq_length, len(timesteps))
+            stops = starts + self.seq_length
+
+        # padding allowed only before start
+        else:
+            stops = np.minimum(
+                self.dataset.lengths[episode_ids], timesteps + 1 + np.random.randint(0, self.seq_length, len(timesteps))
+            )
+            starts = stops - self.seq_length
+
+        return [SegmentId(*x) for x in zip(episode_ids, starts, stops)]

src/data/dataset.py

@@ -0,0 +1,202 @@
+from collections import Counter
+import multiprocessing as mp
+from pathlib import Path
+import shutil
+from typing import Any, Dict, List, Optional
+
+import h5py
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch.utils.data import Dataset as TorchDataset
+
+from .episode import Episode
+from .segment import Segment, SegmentId
+from .utils import make_segment
+from utils import StateDictMixin
+
+
+class Dataset(StateDictMixin, TorchDataset):
+    def __init__(
+        self,
+        directory: Path,
+        dataset_full_res: Optional[TorchDataset],
+        name: Optional[str] = None,
+        cache_in_ram: bool = False,
+        use_manager: bool = False,
+        save_on_disk: bool = True,
+    ) -> None:
+        super().__init__()
+
+        # State
+        self.is_static = False
+        self.num_episodes = None
+        self.num_steps = None
+        self.start_idx = None
+        self.lengths = None
+        self.counter_rew = None
+        self.counter_end = None
+
+        self._directory = Path(directory).expanduser()
+        self._name = name if name is not None else self._directory.stem
+        self._cache_in_ram = cache_in_ram
+        self._save_on_disk = save_on_disk
+        self._default_path = self._directory / "info.pt"
+        self._cache = mp.Manager().dict() if use_manager else {}
+        self._reset()
+
+        self._dataset_full_res = dataset_full_res
+
+    def __len__(self) -> int:
+        return self.num_steps
+
+    def __getitem__(self, segment_id: SegmentId) -> Segment:
+        episode = self.load_episode(segment_id.episode_id)
+        segment = make_segment(episode, segment_id, should_pad=True)
+        if self._dataset_full_res is not None:
+            segment_id_full_res = SegmentId(episode.info["original_file_id"], segment_id.start, segment_id.stop)
+            segment.info["full_res"] = self._dataset_full_res[segment_id_full_res].obs
+        elif "full_res" in segment.info:
+            segment.info["full_res"] = segment.info["full_res"][segment_id.start:segment_id.stop]
+        return segment
+        
+    def __str__(self) -> str:
+        return f"{self.name}: {self.num_episodes} episodes, {self.num_steps} steps."
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def counts_rew(self) -> List[int]:
+        return [self.counter_rew[r] for r in [-1, 0, 1]]
+
+    @property
+    def counts_end(self) -> List[int]:
+        return [self.counter_end[e] for e in [0, 1]]
+
+    def _reset(self) -> None:
+        self.num_episodes = 0
+        self.num_steps = 0
+        self.start_idx = np.array([], dtype=np.int64)
+        self.lengths = np.array([], dtype=np.int64)
+        self.counter_rew = Counter()
+        self.counter_end = Counter()
+        self._cache.clear()
+
+    def clear(self) -> None:
+        self.assert_not_static()
+        if self._directory.is_dir():
+            shutil.rmtree(self._directory)
+        self._reset()
+
+    def load_episode(self, episode_id: int) -> Episode:
+        if self._cache_in_ram and episode_id in self._cache:
+            episode = self._cache[episode_id]
+        else:
+            episode = Episode.load(self._get_episode_path(episode_id))
+            if self._cache_in_ram:
+                self._cache[episode_id] = episode
+        return episode
+
+    def add_episode(self, episode: Episode, *, episode_id: Optional[int] = None) -> int:
+        self.assert_not_static()
+        episode = episode.to("cpu")
+
+        if episode_id is None:
+            episode_id = self.num_episodes
+            self.start_idx = np.concatenate((self.start_idx, np.array([self.num_steps])))
+            self.lengths = np.concatenate((self.lengths, np.array([len(episode)])))
+            self.num_steps += len(episode)
+            self.num_episodes += 1
+
+        else:
+            assert episode_id < self.num_episodes
+            old_episode = self.load_episode(episode_id)
+            incr_num_steps = len(episode) - len(old_episode)
+            self.lengths[episode_id] = len(episode)
+            self.start_idx[episode_id + 1 :] += incr_num_steps
+            self.num_steps += incr_num_steps
+            self.counter_rew.subtract(old_episode.rew.sign().tolist())
+            self.counter_end.subtract(old_episode.end.tolist())
+
+        self.counter_rew.update(episode.rew.sign().tolist())
+        self.counter_end.update(episode.end.tolist())
+
+        if self._save_on_disk:
+            episode.save(self._get_episode_path(episode_id))
+
+        if self._cache_in_ram:
+            self._cache[episode_id] = episode
+
+        return episode_id
+
+    def _get_episode_path(self, episode_id: int) -> Path:
+        n = 3  # number of hierarchies
+        powers = np.arange(n)
+        subfolders = np.floor((episode_id % 10 ** (1 + powers)) / 10**powers) * 10**powers
+        subfolders = [int(x) for x in subfolders[::-1]]
+        subfolders = "/".join([f"{x:0{n - i}d}" for i, x in enumerate(subfolders)])
+        return self._directory / subfolders / f"{episode_id}.pt"
+
+    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+        super().load_state_dict(state_dict)
+        self._cache.clear()
+
+    def assert_not_static(self) -> None:
+        assert not self.is_static, "Trying to modify a static dataset."
+
+    def save_to_default_path(self) -> None:
+        self._default_path.parent.mkdir(exist_ok=True, parents=True)
+        torch.save(self.state_dict(), self._default_path)
+
+    def load_from_default_path(self) -> None:
+        if self._default_path.is_file():
+            self.load_state_dict(torch.load(self._default_path))
+
+
+class CSGOHdf5Dataset(StateDictMixin, TorchDataset):
+    def __init__(self, directory: Path) -> None:
+        super().__init__()
+        filenames = sorted(Path(directory).rglob("*.hdf5"), key=lambda x: int(x.stem.split("_")[-1]))
+        self._filenames = {f"{x.parent.name}/{x.name}": x for x in filenames}
+        self._length_one_episode = 1000
+        self.num_episodes = len(self._filenames)
+        self.num_steps = self._length_one_episode * self.num_episodes
+        self.lengths = np.array([self._length_one_episode] * self.num_episodes, dtype=np.int64)
+
+    def __len__(self) -> int:
+        return self.num_steps
+  
+    def save_to_default_path(self) -> None:
+        pass
+
+    def __getitem__(self, segment_id: SegmentId) -> Segment:
+        assert segment_id.start < self._length_one_episode and segment_id.stop > 0 and segment_id.start < segment_id.stop
+        pad_len_right = max(0, segment_id.stop - self._length_one_episode)
+        pad_len_left = max(0, -segment_id.start)
+
+        start = max(0, segment_id.start)
+        stop = min(self._length_one_episode, segment_id.stop)
+        mask_padding = torch.cat((torch.zeros(pad_len_left), torch.ones(stop - start), torch.zeros(pad_len_right))).bool()
+
+        with h5py.File(self._filenames[segment_id.episode_id], "r") as f:
+            obs = torch.stack([torch.tensor(f[f"frame_{i}_x"][:]).flip(2).permute(2, 0, 1).div(255).mul(2).sub(1) for i in range(start, stop)])
+            act = torch.tensor(np.array([f[f"frame_{i}_y"][:] for i in range(start, stop)]))
+            states = torch.stack([torch.tensor(f[f"frame_{i}_observation"][:]) for i in range(start, stop)])
+            ego_state = torch.stack([torch.tensor(f[f"frame_{i}_ego_state"][:]) for i in range(start, stop)])
+
+        def pad(x):
+            right = F.pad(x, [0 for _ in range(2 * x.ndim - 1)] + [pad_len_right]) if pad_len_right > 0 else x
+            return F.pad(right, [0 for _ in range(2 * x.ndim - 2)] + [pad_len_left, 0]) if pad_len_left > 0 else right
+
+        obs = pad(obs)
+        act = pad(act)
+        rew = torch.zeros(obs.size(0))
+        end = torch.zeros(obs.size(0), dtype=torch.uint8)
+        trunc = torch.zeros(obs.size(0), dtype=torch.uint8)
+        return Segment(obs, act, rew, end, trunc, mask_padding, states=states, ego_state=ego_state, info={}, id=SegmentId(segment_id.episode_id, start, stop))
+    
+    def load_episode(self, episode_id: int) -> Episode:  # used by DatasetTraverser
+        s = self[SegmentId(episode_id, 0, self._length_one_episode)]
+        return Episode(s.obs, s.act, s.rew, s.end, s.trunc, s.info)

src/data/episode.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+import torch
+
+
+@dataclass
+class Episode:
+    obs: torch.FloatTensor
+    act: torch.LongTensor
+    rew: torch.FloatTensor
+    end: torch.ByteTensor
+    trunc: torch.ByteTensor
+    info: Dict[str, Any]
+    states: torch.FloatTensor
+    ego_state: torch.FloatTensor
+
+    def __len__(self) -> int:
+        return self.obs.size(0)
+
+    def __add__(self, other: Episode) -> Episode:
+        assert self.dead.sum() == 0
+        d = {k: torch.cat((v, other.__dict__[k]), dim=0) for k, v in self.__dict__.items() if k != "info"}
+        return Episode(**d, info=merge_info(self.info, other.info))
+
+    def to(self, device) -> Episode:
+        return Episode(**{k: v.to(device) if k != "info" else v for k, v in self.__dict__.items()})
+
+    @property
+    def dead(self) -> torch.ByteTensor:
+        return (self.end + self.trunc).clip(max=1)
+
+    def compute_metrics(self) -> Dict[str, Any]:
+        return {"length": len(self), "return": self.rew.sum().item()}
+
+    @classmethod
+    def load(cls, path: Path, map_location: Optional[torch.device] = None) -> Episode:
+        return cls(
+            **{
+                k: v.div(255).mul(2).sub(1) if k == "obs" else v
+                for k, v in torch.load(Path(path), map_location=map_location).items()
+            }
+        )
+
+    def save(self, path: Path) -> None:
+        path = Path(path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        d = {k: v.add(1).div(2).mul(255).byte() if k == "obs" else v for k, v in self.__dict__.items()}
+        torch.save(d, path.with_suffix(".tmp"))
+        path.with_suffix(".tmp").rename(path)
+
+
+def merge_info(info_a, info_b):
+    keys_a = set(info_a)
+    keys_b = set(info_b)
+    intersection = keys_a & keys_b
+    info = {
+        **{k: info_a[k] for k in keys_a if k not in intersection},
+        **{k: info_b[k] for k in keys_b if k not in intersection},
+        **{k: torch.cat((info_a[k], info_b[k]), dim=0) for k in intersection},
+    }
+    return info

src/data/segment.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import Any, Dict, Union
+
+import torch
+
+
+@dataclass
+class SegmentId:
+    episode_id: Union[int, str]
+    start: int
+    stop: int
+
+
+@dataclass
+class Segment:
+    obs: torch.FloatTensor
+    act: torch.LongTensor
+    rew: torch.FloatTensor
+    end: torch.ByteTensor
+    trunc: torch.ByteTensor
+    mask_padding: torch.BoolTensor
+    states: torch.FloatTensor
+    ego_state: torch.FloatTensor
+    info: Dict[str, Any]
+    id: SegmentId
+
+    @property
+    def effective_size(self):
+        return self.mask_padding.sum().item()

src/data/utils.py

@@ -0,0 +1,89 @@
+import math
+from typing import Generator, List
+
+import torch
+import torch.nn.functional as F
+
+from .batch import Batch
+from .episode import Episode
+from .segment import Segment, SegmentId
+
+
+def collate_segments_to_batch(segments: List[Segment]) -> Batch:
+    attrs = ("obs", "act", "rew", "end", "trunc", "mask_padding")
+    stack = (torch.stack([getattr(s, x) for s in segments]) for x in attrs)
+    return Batch(*stack, [s.info for s in segments], [s.id for s in segments])
+
+
+def make_segment(episode: Episode, segment_id: SegmentId, should_pad: bool = True) -> Segment:
+    assert segment_id.start < len(episode) and segment_id.stop > 0 and segment_id.start < segment_id.stop
+    pad_len_right = max(0, segment_id.stop - len(episode))
+    pad_len_left = max(0, -segment_id.start)
+    assert pad_len_right == pad_len_left == 0 or should_pad
+
+    def pad(x):
+        right = F.pad(x, [0 for _ in range(2 * x.ndim - 1)] + [pad_len_right]) if pad_len_right > 0 else x
+        return F.pad(right, [0 for _ in range(2 * x.ndim - 2)] + [pad_len_left, 0]) if pad_len_left > 0 else right
+
+    start = max(0, segment_id.start)
+    stop = min(len(episode), segment_id.stop)
+    mask_padding = torch.cat((torch.zeros(pad_len_left), torch.ones(stop - start), torch.zeros(pad_len_right))).bool()
+
+    return Segment(
+        pad(episode.obs[start:stop]),
+        pad(episode.act[start:stop]),
+        pad(episode.rew[start:stop]),
+        pad(episode.end[start:stop]),
+        pad(episode.trunc[start:stop]),
+        mask_padding,
+        pad(episode.states[start:stop]),
+        pad(episode.ego_state[start:stop]),
+        info=episode.info,
+        id=SegmentId(segment_id.episode_id, start, stop),
+    )
+
+
+class DatasetTraverser:
+    def __init__(self, dataset, batch_num_samples: int, chunk_size: int) -> None:
+        self.dataset = dataset
+        self.batch_num_samples = batch_num_samples
+        self.chunk_size = chunk_size
+
+    def __len__(self):
+        return math.ceil(
+            sum(
+                [
+                    math.ceil(self.dataset.lengths[episode_id] / self.chunk_size)
+                    - int(self.dataset.lengths[episode_id] % self.chunk_size == 1)
+                    for episode_id in range(self.dataset.num_episodes)
+                ]
+            )
+            / self.batch_num_samples
+        )
+
+    def __iter__(self) -> Generator[Batch, None, None]:
+        chunks = []
+        for episode_id in range(self.dataset.num_episodes):
+            episode = self.dataset.load_episode(episode_id)
+            segments = []
+            for i in range(math.ceil(len(episode) / self.chunk_size)):
+                start = i * self.chunk_size
+                stop = (i + 1) * self.chunk_size
+                segment = make_segment(
+                    episode,
+                    SegmentId(episode_id, start, stop),
+                    should_pad=True,
+                )
+                segment_id_full_res = SegmentId(episode.info["original_file_id"], start, stop)
+                segment.info["full_res"] = self.dataset._dataset_full_res[segment_id_full_res].obs
+                chunks.append(segment)
+            if chunks[-1].effective_size < 2:
+                chunks.pop()
+
+            while len(chunks) >= self.batch_num_samples:
+                yield collate_segments_to_batch(chunks[: self.batch_num_samples])
+                chunks = chunks[self.batch_num_samples :]
+
+        if len(chunks) > 0:
+            yield collate_segments_to_batch(chunks)
+

src/envs/__init__.py

@@ -0,0 +1,2 @@
+from .env import make_atari_env, TorchEnv
+from .world_model_env import WorldModelEnv, WorldModelEnvConfig