Create app.py

app.py

@@ -0,0 +1,181 @@
+from flask import Flask, render_template
+from flask_socketio import SocketIO, emit
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import numpy as np
+from PIL import Image, ImageDraw, io
+import time
+import threading
+import random
+
+app = Flask(__name__)
+socketio = SocketIO(app)
+
+# Initialize model with lower precision
+MODEL_NAME = "Qwen/Qwen-1_5B-Chat"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=torch.float16,
+    device_map="auto",
+    trust_remote_code=True
+)
+
+# Game Constants
+GRID_SIZE = 12  # Smaller grid for performance
+CELL_SIZE = 40
+COLORS = {
+    'background': 'white',
+    'grid': 'lightgray',
+    'snake': 'red',
+    'agent': 'blue',
+    'obstacle': 'gray'
+}
+
+class GameState:
+    def __init__(self):
+        self.snake = [6, 6]  # Center
+        self.agents = [[2, 2], [9, 9], [2, 9]]
+        self.obstacles = [[4, 4], [7, 7], [4, 7]]
+        self.scores = {'snake': 0, 'agents': 0}
+        self.history = []
+
+    def get_agent_state(self, agent_idx):
+        return {
+            'position': self.agents[agent_idx],
+            'snake_pos': self.snake,
+            'other_agents': [pos for i, pos in enumerate(self.agents) if i != agent_idx],
+            'obstacles': self.obstacles
+        }
+
+game = GameState()
+
+def get_model_decision(role, state):
+    """Get next move from Qwen model."""
+    if role == "snake":
+        prompt = f"You are a predator trying to catch prey. Your position is {state['position']}, prey positions are {state['other_agents']}. Choose one move from: UP, DOWN, LEFT, RIGHT, STAY. Just output the move word."
+    else:
+        prompt = f"You are prey avoiding a predator. Your position is {state['position']}, predator position is {state['snake_pos']}. Choose one move from: UP, DOWN, LEFT, RIGHT, STAY. Just output the move word."
+
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=10,
+        temperature=0.7,
+        do_sample=True
+    )
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    # Extract move from response
+    moves = ["UP", "DOWN", "LEFT", "RIGHT", "STAY"]
+    for move in moves:
+        if move in response.upper():
+            return move
+    return "STAY"
+
+def apply_move(position, move):
+    """Apply move while respecting grid boundaries."""
+    x, y = position.copy()
+    if move == "UP" and y > 0:
+        y -= 1
+    elif move == "DOWN" and y < GRID_SIZE - 1:
+        y += 1
+    elif move == "LEFT" and x > 0:
+        x -= 1
+    elif move == "RIGHT" and x < GRID_SIZE - 1:
+        x += 1
+    return [x, y]
+
+def create_game_image():
+    """Create game visualization."""
+    img = Image.new("RGB", (GRID_SIZE * CELL_SIZE, GRID_SIZE * CELL_SIZE), COLORS['background'])
+    draw = ImageDraw.Draw(img)
+
+    # Draw grid
+    for i in range(GRID_SIZE + 1):
+        draw.line([(i * CELL_SIZE, 0), (i * CELL_SIZE, GRID_SIZE * CELL_SIZE)], fill=COLORS['grid'])
+        draw.line([(0, i * CELL_SIZE), (GRID_SIZE * CELL_SIZE, i * CELL_SIZE)], fill=COLORS['grid'])
+
+    # Draw obstacles
+    for pos in game.obstacles:
+        draw.rectangle([
+            pos[0] * CELL_SIZE, pos[1] * CELL_SIZE,
+            (pos[0] + 1) * CELL_SIZE, (pos[1] + 1) * CELL_SIZE
+        ], fill=COLORS['obstacle'])
+
+    # Draw agents
+    for pos in game.agents:
+        center = ((pos[0] + 0.5) * CELL_SIZE, (pos[1] + 0.5) * CELL_SIZE)
+        radius = CELL_SIZE // 3
+        draw.ellipse([
+            center[0] - radius, center[1] - radius,
+            center[0] + radius, center[1] + radius
+        ], fill=COLORS['agent'])
+
+    # Draw snake
+    center = ((game.snake[0] + 0.5) * CELL_SIZE, (game.snake[1] + 0.5) * CELL_SIZE)
+    radius = CELL_SIZE // 3
+    draw.ellipse([
+        center[0] - radius, center[1] - radius,
+        center[0] + radius, center[1] + radius
+    ], fill=COLORS['snake'])
+
+    # Add scores
+    draw.text((10, 10), f"Snake: {game.scores['snake']} | Agents: {game.scores['agents']}", fill="black")
+
+    # Convert to bytes
+    img_byte_arr = io.BytesIO()
+    img.save(img_byte_arr, format='PNG')
+    img_byte_arr.seek(0)
+    return img_byte_arr
+
+def update_game():
+    """Update game state for one turn."""
+    # Snake's turn
+    snake_state = {'position': game.snake, 'other_agents': game.agents}
+    snake_move = get_model_decision('snake', snake_state)
+    new_pos = apply_move(game.snake, snake_move)
+    if new_pos not in game.obstacles:
+        game.snake = new_pos
+
+    # Agents' turns
+    for i in range(len(game.agents)):
+        agent_state = game.get_agent_state(i)
+        agent_move = get_model_decision('agent', agent_state)
+        new_pos = apply_move(game.agents[i], agent_move)
+        if new_pos not in game.obstacles:
+            game.agents[i] = new_pos
+
+    # Check captures
+    for i, agent_pos in enumerate(game.agents):
+        if agent_pos == game.snake:
+            game.scores['snake'] += 1
+            # Respawn agent
+            while True:
+                new_pos = [random.randint(0, GRID_SIZE - 1), random.randint(0, GRID_SIZE - 1)]
+                if new_pos not in game.obstacles and new_pos != game.snake:
+                    game.agents[i] = new_pos
+                    break
+
+def game_loop():
+    """Main game loop."""
+    while True:
+        update_game()
+        img_bytes = create_game_image()
+        socketio.emit('game_update', {
+            'image': img_bytes.getvalue().hex(),
+            'scores': game.scores
+        })
+        time.sleep(1.0)  # Slower updates to reduce resource usage
+
+@app.route('/')
+def index():
+    return render_template('index.html')
+
+@socketio.on('connect')
+def handle_connect():
+    print('Client connected')
+
+if __name__ == '__main__':
+    threading.Thread(target=game_loop, daemon=True).start()
+    socketio.run(app, host='0.0.0.0', port=7860)