ENV 2D distance_based, goal_oriented.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pygame\n",
    "import random\n",
    "import math\n",
    "\n",
    "pygame.init()\n",
    "\n",
    "# Constants\n",
    "WINDOW_SIZE = (600, 600)\n",
    "WHITE = (255, 255, 255)\n",
    "BLACK = (0, 0, 0)\n",
    "RED = (255, 0, 0)\n",
    "GREEN = (0, 255, 0)\n",
    "BLUE = (0, 0, 255)\n",
    "YELLOW = (255, 255, 0)\n",
    "FONT_COLOR = (0, 0, 0)\n",
    "FONT_SIZE = 24\n",
    "step_size = 3\n",
    "tolerance = 10\n",
    "ARROW_LENGTH = 20  # Length of direction arrows\n",
    "\n",
    "# Set up display\n",
    "screen = pygame.display.set_mode(WINDOW_SIZE)\n",
    "pygame.display.set_caption(\"2D Environment with Moving Target\")\n",
    "\n",
    "# Load font for rendering text\n",
    "font = pygame.font.Font(None, FONT_SIZE)\n",
    "\n",
    "# Agent positions and goal flags\n",
    "agent1_pos = [300, 300]  # W (green)\n",
    "agent2_pos = [300, 300]  # H (blue, user-controlled)\n",
    "agent3_pos = [300, 300]  # X (red, weighted average)\n",
    "# Add direction vectors for each agent\n",
    "agent1_dir = [0, 0]\n",
    "agent2_dir = [0, 0]\n",
    "agent3_dir = [0, 0]\n",
    "reached_goal = False\n",
    "\n",
    "# Target (goal) is placed randomly in the 2D space\n",
    "target_pos = [random.randint(0, WINDOW_SIZE[0]), random.randint(0, WINDOW_SIZE[1])]\n",
    "\n",
    "def distance(pos1, pos2):\n",
    "    return math.sqrt((pos1[0] - pos2[0])**2 + (pos1[1] - pos2[1])**2)\n",
    "\n",
    "def calculate_gamma():\n",
    "    max_distance = math.sqrt(WINDOW_SIZE[0]**2 + WINDOW_SIZE[1]**2)\n",
    "    distance_w = distance(agent1_pos, target_pos)\n",
    "    distance_h = distance(agent2_pos, target_pos)\n",
    "    \n",
    "    # Normalize distances\n",
    "    normalized_distance_w = distance_w / max_distance\n",
    "    normalized_distance_h = distance_h / max_distance\n",
    "    \n",
    "    # Calculate gamma based on relative distances\n",
    "    if normalized_distance_w + normalized_distance_h == 0:\n",
    "        return 0.5\n",
    "    else:\n",
    "        gamma = normalized_distance_h / (normalized_distance_w + normalized_distance_h)\n",
    "        \n",
    "    gamma = gamma ** 2\n",
    "    return gamma\n",
    "\n",
    "def normalize_vector(dx, dy):\n",
    "    length = math.sqrt(dx**2 + dy**2)\n",
    "    if length == 0:\n",
    "        return 0, 0\n",
    "    return dx/length, dy/length\n",
    "\n",
    "def draw_arrow(surface, color, start_pos, direction, length=ARROW_LENGTH):\n",
    "    if direction[0] == 0 and direction[1] == 0:\n",
    "        return\n",
    "        \n",
    "    # Calculate end point\n",
    "    end_x = start_pos[0] + direction[0] * length\n",
    "    end_y = start_pos[1] + direction[1] * length\n",
    "    \n",
    "    # Draw main line\n",
    "    pygame.draw.line(surface, color, start_pos, (end_x, end_y), 2)\n",
    "    \n",
    "    # Calculate arrow head\n",
    "    arrow_size = 7\n",
    "    angle = math.atan2(direction[1], direction[0])\n",
    "    arrow1_x = end_x - arrow_size * math.cos(angle + math.pi/6)\n",
    "    arrow1_y = end_y - arrow_size * math.sin(angle + math.pi/6)\n",
    "    arrow2_x = end_x - arrow_size * math.cos(angle - math.pi/6)\n",
    "    arrow2_y = end_y - arrow_size * math.sin(angle - math.pi/6)\n",
    "    \n",
    "    # Draw arrow head\n",
    "    pygame.draw.line(surface, color, (end_x, end_y), (arrow1_x, arrow1_y), 2)\n",
    "    pygame.draw.line(surface, color, (end_x, end_y), (arrow2_x, arrow2_y), 2)\n",
    "\n",
    "def move_agent1():\n",
    "    global agent1_pos, reached_goal, agent1_dir\n",
    "    dx = target_pos[0] - agent1_pos[0]\n",
    "    dy = target_pos[1] - agent1_pos[1]\n",
    "    distance_to_target = math.sqrt(dx**2 + dy**2)\n",
    "    \n",
    "    # Update direction vector\n",
    "    agent1_dir = normalize_vector(dx, dy)\n",
    "    \n",
    "    if distance_to_target > step_size:\n",
    "        move_ratio = step_size / distance_to_target\n",
    "        agent1_pos[0] += dx * move_ratio\n",
    "        agent1_pos[1] += dy * move_ratio\n",
    "    else:\n",
    "        agent1_pos[0] = target_pos[0]\n",
    "        agent1_pos[1] = target_pos[1]\n",
    "        reached_goal = True\n",
    "\n",
    "def move_agent2(dx, dy):\n",
    "    global agent2_pos, agent2_dir\n",
    "    new_x = agent2_pos[0] + dx\n",
    "    new_y = agent2_pos[1] + dy\n",
    "    agent2_pos[0] = max(0, min(WINDOW_SIZE[0], new_x))\n",
    "    agent2_pos[1] = max(0, min(WINDOW_SIZE[1], new_y))\n",
    "    \n",
    "    # Update direction vector\n",
    "    if dx != 0 or dy != 0:\n",
    "        agent2_dir = normalize_vector(dx, dy)\n",
    "\n",
    "def update_agent3():\n",
    "    global agent3_pos, agent3_dir\n",
    "    gamma = calculate_gamma()\n",
    "    new_x = gamma * agent1_pos[0] + (1 - gamma) * agent2_pos[0]\n",
    "    new_y = gamma * agent1_pos[1] + (1 - gamma) * agent2_pos[1]\n",
    "    \n",
    "    # Update direction vector\n",
    "    dx = new_x - agent3_pos[0]\n",
    "    dy = new_y - agent3_pos[1]\n",
    "    agent3_dir = normalize_vector(dx, dy)\n",
    "    \n",
    "    agent3_pos[0] = new_x\n",
    "    agent3_pos[1] = new_y\n",
    "    return gamma\n",
    "\n",
    "def reset():\n",
    "    global agent1_pos, agent2_pos, agent3_pos, reached_goal, target_pos, agent1_dir, agent2_dir, agent3_dir\n",
    "    agent1_pos = [300, 300]\n",
    "    agent2_pos = [300, 300]\n",
    "    agent3_pos = [300, 300]\n",
    "    agent1_dir = [0, 0]\n",
    "    agent2_dir = [0, 0]\n",
    "    agent3_dir = [0, 0]\n",
    "    reached_goal = False\n",
    "    target_pos = [random.randint(0, WINDOW_SIZE[0]), random.randint(0, WINDOW_SIZE[1])]\n",
    "\n",
    "def render(gamma):\n",
    "    screen.fill(WHITE)\n",
    "    \n",
    "    # Draw agents and their direction arrows\n",
    "    pygame.draw.circle(screen, GREEN, (int(agent1_pos[0]), int(agent1_pos[1])), 10)\n",
    "    pygame.draw.circle(screen, BLUE, (int(agent2_pos[0]), int(agent2_pos[1])), 10)\n",
    "    pygame.draw.circle(screen, RED, (int(agent3_pos[0]), int(agent3_pos[1])), 10)\n",
    "    pygame.draw.circle(screen, YELLOW, (int(target_pos[0]), int(target_pos[1])), 10)\n",
    "    \n",
    "    # Draw direction arrows\n",
    "    draw_arrow(screen, GREEN, (int(agent1_pos[0]), int(agent1_pos[1])), agent1_dir)\n",
    "    draw_arrow(screen, BLUE, (int(agent2_pos[0]), int(agent2_pos[1])), agent2_dir)\n",
    "    draw_arrow(screen, RED, (int(agent3_pos[0]), int(agent3_pos[1])), agent3_dir)\n",
    "\n",
    "    # Draw labels\n",
    "    screen.blit(font.render(\"W\", True, BLACK), (agent1_pos[0] - 5, agent1_pos[1] - 12))\n",
    "    screen.blit(font.render(\"H\", True, BLACK), (agent2_pos[0] - 5, agent2_pos[1] - 12))\n",
    "    screen.blit(font.render(\"X\", True, BLACK), (agent3_pos[0] - 5, agent3_pos[1] - 12))\n",
    "\n",
    "    # Display gamma and formula\n",
    "    gamma_text = font.render(f\"Gamma: {gamma:.2f}\", True, FONT_COLOR)\n",
    "    screen.blit(gamma_text, (10, 10))\n",
    "    formula_text = font.render(f\"X = {gamma:.2f}W + {1-gamma:.2f}H\", True, FONT_COLOR)\n",
    "    screen.blit(formula_text, (10, 40))\n",
    "\n",
    "    if reached_goal:\n",
    "        reset_text = font.render(\"Goal Reached! Press R to reset\", True, FONT_COLOR)\n",
    "        screen.blit(reset_text, (150, 70))\n",
    "\n",
    "    pygame.display.update()\n",
    "\n",
    "# Main\n",
    "running = True\n",
    "clock = pygame.time.Clock()\n",
    "gamma = 0.5  # Initialize gamma\n",
    "\n",
    "while running:\n",
    "    for event in pygame.event.get():\n",
    "        if event.type == pygame.QUIT:\n",
    "            running = False\n",
    "        if event.type == pygame.KEYDOWN:\n",
    "            if event.key == pygame.K_r:\n",
    "                reset()\n",
    "\n",
    "    if not reached_goal:\n",
    "        keys = pygame.key.get_pressed()\n",
    "        dx, dy = 0, 0\n",
    "        if keys[pygame.K_LEFT]:\n",
    "            dx -= step_size\n",
    "        if keys[pygame.K_RIGHT]:\n",
    "            dx += step_size\n",
    "        if keys[pygame.K_UP]:\n",
    "            dy -= step_size\n",
    "        if keys[pygame.K_DOWN]:\n",
    "            dy += step_size\n",
    "\n",
    "        if dx != 0 or dy != 0:\n",
    "            move_agent2(dx, dy)\n",
    "            move_agent1()\n",
    "            gamma = update_agent3()\n",
    "\n",
    "    render(gamma)\n",
    "    clock.tick(30)\n",
    "\n",
    "pygame.quit()"
   ]
  }
 ],
 "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
}