Controller_MK5.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'check_collision' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[4], line 239\u001b[0m\n\u001b[0;32m    236\u001b[0m     pygame\u001b[38;5;241m.\u001b[39mquit()\n\u001b[0;32m    238\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;18m__name__\u001b[39m \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m__main__\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[1;32m--> 239\u001b[0m     main()\n",
      "Cell \u001b[1;32mIn[4], line 213\u001b[0m, in \u001b[0;36mmain\u001b[1;34m()\u001b[0m\n\u001b[0;32m    211\u001b[0m \u001b[38;5;66;03m# Update environment\u001b[39;00m\n\u001b[0;32m    212\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m reached_goal:\n\u001b[1;32m--> 213\u001b[0m     h_dir, w_dir, x_dir \u001b[38;5;241m=\u001b[39m move_dot(human_input \u001b[38;5;241m*\u001b[39m MAX_SPEED)\n\u001b[0;32m    214\u001b[0m     recent_positions\u001b[38;5;241m.\u001b[39mappend((\u001b[38;5;241m*\u001b[39mdot_pos, time\u001b[38;5;241m.\u001b[39mtime()))\n\u001b[0;32m    216\u001b[0m \u001b[38;5;66;03m# Rendering\u001b[39;00m\n",
      "Cell \u001b[1;32mIn[4], line 148\u001b[0m, in \u001b[0;36mmove_dot\u001b[1;34m(human_input)\u001b[0m\n\u001b[0;32m    145\u001b[0m new_pos \u001b[38;5;241m=\u001b[39m dot_np \u001b[38;5;241m+\u001b[39m combined_dir \u001b[38;5;241m*\u001b[39m MAX_SPEED\n\u001b[0;32m    146\u001b[0m new_pos \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mclip(new_pos, [\u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m0\u001b[39m], FULL_VIEW_SIZE)\n\u001b[1;32m--> 148\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m check_collision(dot_np, new_pos):\n\u001b[0;32m    149\u001b[0m     dot_pos[:] \u001b[38;5;241m=\u001b[39m new_pos\u001b[38;5;241m.\u001b[39mtolist()  \u001b[38;5;66;03m# Maintain list type for pygame\u001b[39;00m\n\u001b[0;32m    151\u001b[0m \u001b[38;5;66;03m# Goal check\u001b[39;00m\n",
      "\u001b[1;31mNameError\u001b[0m: name 'check_collision' is not defined"
     ]
    }
   ],
   "source": [
    "import pygame\n",
    "import random\n",
    "import math\n",
    "import time\n",
    "import numpy as np\n",
    "from stable_baselines3 import PPO\n",
    "\n",
    "# Pygame 2.x with SDL2\n",
    "from pygame._sdl2.video import Window, Renderer, Texture\n",
    "\n",
    "pygame.init()\n",
    "pygame.joystick.init()\n",
    "\n",
    "###############################################################################\n",
    "# Config / Constants (Match Training Exactly)\n",
    "###############################################################################\n",
    "FULL_VIEW_SIZE = (1200, 800)   # Must match training\n",
    "RED_ONLY_SIZE  = (1200, 800)   \n",
    "MAX_SPEED = 3                  # Remove scaling factor\n",
    "NOISE_MAGNITUDE = 0.5          # Match training\n",
    "DOT_RADIUS = 30                # Original size from training\n",
    "TARGET_RADIUS = 10             \n",
    "GOAL_DETECTION_RADIUS = DOT_RADIUS + TARGET_RADIUS\n",
    "\n",
    "# Colors and visual settings\n",
    "WHITE = (255, 255, 255)\n",
    "BLACK = (0, 0, 0)\n",
    "RED = (255, 0, 0)\n",
    "GREEN = (0, 200, 0)\n",
    "BLUE = (0, 0, 255)\n",
    "YELLOW = (255, 255, 0)\n",
    "GRAY = (128, 128, 128)\n",
    "\n",
    "# Environment parameters (match training)\n",
    "NUM_GOALS = 3\n",
    "GHOST_TRAIL_DURATION = 3.0  \n",
    "RECENT_DIR_LOOKBACK = 1.0\n",
    "GOAL_SWITCH_THRESHOLD = 0.05\n",
    "START_POS = [FULL_VIEW_SIZE[0]//2, FULL_VIEW_SIZE[1]//2]\n",
    "\n",
    "# Initialize state\n",
    "dot_pos = START_POS.copy()\n",
    "gamma = 0.2\n",
    "reached_goal = False\n",
    "targets = []\n",
    "current_target_idx = 0\n",
    "obstacles = []\n",
    "recent_positions = []\n",
    "last_reset_time = time.time()\n",
    "\n",
    "###############################################################################\n",
    "# Neural Network Integration (Fixed)\n",
    "###############################################################################\n",
    "class GammaPredictor:\n",
    "    def __init__(self, model_path=\"ppo_dynamic_arbitration_simple\"):\n",
    "        policy_kwargs = dict(\n",
    "            net_arch=dict(pi=[256, 128], vf=[256, 128])\n",
    "        )\n",
    "        \n",
    "        self.model = PPO.load(\n",
    "            model_path,\n",
    "            custom_objects={\n",
    "                \"policy_kwargs\": policy_kwargs,\n",
    "                \"clip_range\": 0.2\n",
    "            }\n",
    "        )\n",
    "        self.max_dist = math.sqrt(FULL_VIEW_SIZE[0]**2 + FULL_VIEW_SIZE[1]**2)\n",
    "\n",
    "    def prepare_observation(self, dot_pos, target_pos, human_input):\n",
    "        dot_pos = np.array(dot_pos, dtype=np.float32)\n",
    "        target_pos = np.array(target_pos, dtype=np.float32)\n",
    "        human_input = np.array(human_input, dtype=np.float32)\n",
    "        \n",
    "        to_target = target_pos - dot_pos\n",
    "        dist = np.linalg.norm(to_target)\n",
    "        perfect_dir = to_target / dist if dist > 0 else np.zeros(2, dtype=np.float32)\n",
    "        \n",
    "        h_mag = np.linalg.norm(human_input)\n",
    "        human_dir = human_input / h_mag if h_mag > 0 else np.zeros(2, dtype=np.float32)\n",
    "        \n",
    "        return np.concatenate([\n",
    "            dot_pos,\n",
    "            human_dir,\n",
    "            target_pos,\n",
    "            perfect_dir,\n",
    "            [dist / self.max_dist]\n",
    "        ])\n",
    "\n",
    "    def predict_gamma(self, dot_pos, target_pos, human_input):\n",
    "        obs = self.prepare_observation(dot_pos, target_pos, human_input)\n",
    "        action, _ = self.model.predict(obs.reshape(1, -1), deterministic=True)\n",
    "        return np.clip(action.item(), 0.0, 0.4)  # Use .item() instead of float()\n",
    "\n",
    "gamma_predictor = GammaPredictor()\n",
    "\n",
    "###############################################################################\n",
    "# Environment Core Logic (Fixed)\n",
    "###############################################################################\n",
    "def generate_targets():\n",
    "    targets.clear()\n",
    "    for _ in range(NUM_GOALS):\n",
    "        while True:\n",
    "            # Match training's target generation range (100-1100x, 100-700y)\n",
    "            pos = [\n",
    "                random.randint(100, FULL_VIEW_SIZE[0]-100),\n",
    "                random.randint(100, FULL_VIEW_SIZE[1]-100)\n",
    "            ]\n",
    "            if all(distance(pos, o) > 100 for o in obstacles):  # Simple collision check\n",
    "                targets.append(pos)\n",
    "                break\n",
    "\n",
    "def move_dot(human_input):\n",
    "    global dot_pos, gamma, reached_goal, current_target_idx\n",
    "\n",
    "    # Convert positions to numpy arrays\n",
    "    dot_np = np.array(dot_pos, dtype=np.float32)\n",
    "    target_np = np.array(targets[current_target_idx], dtype=np.float32)\n",
    "    \n",
    "    # Predict gamma with proper typing\n",
    "    gamma = gamma_predictor.predict_gamma(\n",
    "        dot_pos=dot_np,\n",
    "        target_pos=target_np,\n",
    "        human_input=np.array(human_input, dtype=np.float32)\n",
    "    )\n",
    "\n",
    "    # Calculate directions with numpy arrays\n",
    "    h_vec = np.array(human_input, dtype=np.float32)\n",
    "    h_mag = np.linalg.norm(h_vec)\n",
    "    \n",
    "    if h_mag > 0:\n",
    "        noise = np.random.normal(0, NOISE_MAGNITUDE, 2).astype(np.float32)\n",
    "        h_dir = (h_vec + noise) / (h_mag + np.linalg.norm(noise))\n",
    "    else:\n",
    "        h_dir = np.zeros(2, dtype=np.float32)\n",
    "\n",
    "    w_dir = target_np - dot_np\n",
    "    w_mag = np.linalg.norm(w_dir)\n",
    "    w_dir = w_dir / w_mag if w_mag > 0 else np.zeros(2, dtype=np.float32)\n",
    "\n",
    "    # Combine directions\n",
    "    combined_dir = gamma * w_dir + (1 - gamma) * h_dir\n",
    "    combined_dir /= np.linalg.norm(combined_dir) if np.linalg.norm(combined_dir) > 0 else 1.0\n",
    "\n",
    "    # Update position\n",
    "    new_pos = dot_np + combined_dir * MAX_SPEED\n",
    "    new_pos = np.clip(new_pos, [0, 0], FULL_VIEW_SIZE)\n",
    "    \n",
    "    if not check_collision(dot_np, new_pos):\n",
    "        dot_pos[:] = new_pos.tolist()  # Maintain list type for pygame\n",
    "\n",
    "    # Goal check\n",
    "    if np.linalg.norm(new_pos - target_np) < GOAL_DETECTION_RADIUS:\n",
    "        reached_goal = True\n",
    "        pygame.time.set_timer(pygame.USEREVENT, 1000)\n",
    "\n",
    "    return h_dir, w_dir, combined_dir\n",
    "\n",
    "###############################################################################\n",
    "# Visualization & UI (No Changes Needed)\n",
    "###############################################################################\n",
    "# [Keep all rendering and window management code unchanged from original]\n",
    "# [Maintain existing reset(), draw_arrow(), render_full_view(), etc]\n",
    "\n",
    "###############################################################################\n",
    "# Main Loop with Fixed Input Handling\n",
    "###############################################################################\n",
    "def main():\n",
    "    window1 = Window(\"Full View\", size=FULL_VIEW_SIZE)\n",
    "    renderer1 = Renderer(window1, vsync=True)\n",
    "    window2 = Window(\"Red Arrow View\", size=RED_ONLY_SIZE)\n",
    "    renderer2 = Renderer(window2, vsync=True)\n",
    "\n",
    "    # Surfaces and initialization\n",
    "    surface_full = pygame.Surface(FULL_VIEW_SIZE, pygame.SRCALPHA)\n",
    "    surface_red = pygame.Surface(RED_ONLY_SIZE, pygame.SRCALPHA)\n",
    "    font = pygame.font.Font(None, 24)\n",
    "    \n",
    "    #generate_obstacles()\n",
    "    generate_targets()\n",
    "    \n",
    "    clock = pygame.time.Clock()\n",
    "    running = True\n",
    "    \n",
    "    while running:\n",
    "        # Event handling (remove manual gamma controls)\n",
    "        for event in pygame.event.get():\n",
    "            if event.type == pygame.QUIT:\n",
    "                running = False\n",
    "            elif event.type == pygame.KEYDOWN:\n",
    "                if event.key == pygame.K_r:\n",
    "                    reset()\n",
    "            elif event.type == pygame.USEREVENT and reached_goal:\n",
    "                reset()\n",
    "\n",
    "        # Movement handling\n",
    "        human_input = np.zeros(2)\n",
    "        keys = pygame.key.get_pressed()\n",
    "        if keys[pygame.K_LEFT]: human_input[0] -= 1\n",
    "        if keys[pygame.K_RIGHT]: human_input[0] += 1\n",
    "        if keys[pygame.K_UP]: human_input[1] -= 1\n",
    "        if keys[pygame.K_DOWN]: human_input[1] += 1\n",
    "        \n",
    "        # Joystick handling (remove gamma triggers)\n",
    "        if pygame.joystick.get_count() > 0:\n",
    "            joystick = pygame.joystick.Joystick(0)\n",
    "            axis_0 = joystick.get_axis(0)\n",
    "            axis_1 = joystick.get_axis(1)\n",
    "            if abs(axis_0) > 0.1 or abs(axis_1) > 0.1:\n",
    "                human_input = np.array([axis_0, axis_1])\n",
    "\n",
    "        # Update environment\n",
    "        if not reached_goal:\n",
    "            h_dir, w_dir, x_dir = move_dot(human_input * MAX_SPEED)\n",
    "            recent_positions.append((*dot_pos, time.time()))\n",
    "        \n",
    "        # Rendering\n",
    "        surface_full.fill(WHITE)\n",
    "        surface_red.fill(WHITE)\n",
    "        render_full_view(surface_full, h_dir, w_dir, x_dir)\n",
    "        render_red_only(surface_red, x_dir)\n",
    "        \n",
    "        # Update windows\n",
    "        tex1 = Texture.from_surface(renderer1, surface_full.convert_alpha())\n",
    "        tex2 = Texture.from_surface(renderer2, surface_red.convert_alpha())\n",
    "        \n",
    "        renderer1.clear()\n",
    "        tex1.draw()\n",
    "        renderer1.present()\n",
    "        \n",
    "        renderer2.clear()\n",
    "        tex2.draw()\n",
    "        renderer2.present()\n",
    "        \n",
    "        clock.tick(60)\n",
    "\n",
    "    pygame.quit()\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}