MAPF_playground.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "import matplotlib.animation as animation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Agent:\n",
    "    def __init__(self, start, goal):\n",
    "        self.start = start\n",
    "        self.goal = goal\n",
    "        self.path = []  # Stores the path of the agent\n",
    "\n",
    "\n",
    "def a_star(graph, start, goal):\n",
    "    open_set = {start}\n",
    "    closed_set = set()\n",
    "    g_values = {start: 0}\n",
    "    f_values = {start: heuristic(start, goal)}\n",
    "    came_from = {}\n",
    "\n",
    "    while open_set:\n",
    "        current = min(open_set, key=lambda node: f_values[node])\n",
    "\n",
    "        if current == goal:\n",
    "            return reconstruct_path(start, goal, came_from)\n",
    "\n",
    "        open_set.remove(current)\n",
    "        closed_set.add(current)\n",
    "\n",
    "        for neighbor in get_neighbors(graph, current):\n",
    "            if neighbor in closed_set:\n",
    "                continue\n",
    "\n",
    "            tentative_g = g_values[current] + 1\n",
    "\n",
    "            if neighbor not in open_set or tentative_g < g_values[neighbor]:\n",
    "                came_from[neighbor] = current\n",
    "                g_values[neighbor] = tentative_g\n",
    "                f_values[neighbor] = tentative_g + heuristic(neighbor, goal)\n",
    "\n",
    "                if neighbor not in open_set:\n",
    "                    open_set.add(neighbor)\n",
    "\n",
    "    return None\n",
    "\n",
    "\n",
    "def heuristic(node, goal):\n",
    "    return abs(node[0] - goal[0]) + abs(node[1] - goal[1])\n",
    "\n",
    "\n",
    "def get_neighbors(graph, node):\n",
    "    neighbors = []\n",
    "    for i, j in [(0, 1), (1, 0), (0, -1), (-1, 0)]:\n",
    "        new_node = (node[0] + i, node[1] + j)\n",
    "        if 0 <= new_node[0] < len(graph) and 0 <= new_node[1] < len(graph[0]) and graph[new_node[0]][new_node[1]] == 0:\n",
    "            neighbors.append(new_node)\n",
    "    return neighbors\n",
    "\n",
    "\n",
    "def reconstruct_path(start, goal, came_from):\n",
    "    path = [goal]\n",
    "    current = goal\n",
    "    while current != start:\n",
    "        current = came_from[current]\n",
    "        path.append(current)\n",
    "    path.reverse()\n",
    "    return path\n",
    "\n",
    "\n",
    "# Write a new function that finds conflicts between agents\n",
    "def find_conflicts(agents):\n",
    "    conflicts = []\n",
    "    for i in range(len(agents)):\n",
    "        for j in range(i + 1, len(agents)):\n",
    "            for pos_i, pos_j in zip(agents[i].path, agents[j].path):\n",
    "                if pos_i == pos_j:\n",
    "                    conflicts.append((i, j))\n",
    "                    break\n",
    "    return conflicts\n",
    "\n",
    "\n",
    "def resolve_conflict(agent1, agent2):\n",
    "    # Implement conflict resolution logic here\n",
    "    agent1.start, agent2.start = agent2.start, agent1.start\n",
    "\n",
    "\n",
    "def cbs(map, agents):\n",
    "    for agent in agents:\n",
    "            agent.path = a_star(map, agent.start, agent.goal)\n",
    "\n",
    "    while True:\n",
    "        # Step 2: Check for conflicts\n",
    "        conflicts = find_conflicts(agents)\n",
    "\n",
    "        # If there are no conflicts, the solution is found\n",
    "        if not conflicts:\n",
    "            return [agent.path for agent in agents]\n",
    "\n",
    "        # Step 3: Resolve conflicts\n",
    "        for i, j in conflicts:\n",
    "            resolve_conflict(agents[i], agents[j])\n",
    "\n",
    "        # Re-plan paths\n",
    "        for agent in agents:\n",
    "            agent.path = a_star(map, agent.start, agent.goal)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Write a function to visualise the solution as a gif with the agents moving one step in the path at each timestep\n",
    "def visualize_solution(map, agents, path):\n",
    "    # Initialize the plot\n",
    "    map = np.array(map)\n",
    "    path = np.array(path)\n",
    "\n",
    "    fig, ax = plt.subplots(figsize=(map.shape[0], map.shape[1]))\n",
    "    ax.set_title('MAPF Solution')\n",
    "    ax.set_xticks(np.arange(0, map.shape[0], 1))\n",
    "    ax.set_yticks(np.arange(0, map.shape[1], 1))\n",
    "    ax.grid()\n",
    "    ax.imshow(map, cmap='Greys', vmin=0, vmax=1)\n",
    "\n",
    "    # Initialize the agents\n",
    "    agent_markers = []\n",
    "    for agent in agents:\n",
    "        agent_markers.append(ax.scatter(agent.start[0], agent.start[1], marker='o', color='blue'))\n",
    "\n",
    "    # Initialize the paths\n",
    "    path_markers = []\n",
    "    for i in range(len(agents)):\n",
    "        path_markers.append(ax.plot([], [], color='blue', linewidth=1)[0])\n",
    "\n",
    "    # Animate the solution\n",
    "    def animate(i):\n",
    "        for j, agent in enumerate(agents):\n",
    "            if i < len(path[j]):\n",
    "                agent_markers[j].set_offsets([path[j][i][1], path[j][i][0]])\n",
    "                path_markers[j].set_data([p[1] for p in path[j][:i+1]], [p[0] for p in path[j][:i+1]])\n",
    "\n",
    "    anim = animation.FuncAnimation(fig, animate, frames=100, interval=100, blit=False)\n",
    "    anim.save('solution.gif', writer='imagemagick', fps=1)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Agent 1 path: [(4, 0), (4, 1), (4, 2), (4, 3), (4, 4)]\n",
      "Agent 2 path: [(0, 0), (0, 1), (0, 2), (0, 3), (0, 4)]\n"
     ]
    }
   ],
   "source": [
    "\n",
    "# Define the map and agents\n",
    "map =  [[0, 0, 0, 0, 0],\n",
    "        [0, 0, 1, 0, 0],\n",
    "        [0, 0, 0, 0, 0],\n",
    "        [0, 0, 1, 0, 0],\n",
    "        [0, 0, 0, 0, 0]]\n",
    "\n",
    "agents = [Agent((0, 0), (4, 4)), Agent((4, 0), (0, 4))]\n",
    "\n",
    "# Run CBS algorithm\n",
    "paths = cbs(map, agents)\n",
    "\n",
    "# Print the paths\n",
    "for i, path in enumerate(paths):\n",
    "    print(f\"Agent {i + 1} path:\", path)\n",
    "\n",
    "# for agent in agents:\n",
    "#     agent.path = a_star(map, agent.start, agent.goal)\n",
    "#     print(agent.path)\n",
    "\n",
    "# conflicts = find_conflicts(agents)\n",
    "# print(conflicts)\n",
    "\n",
    "# visualize_solution(map, agents, [agent.path for agent in agents])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}