index.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import json\n",
    "import os\n",
    "from pathlib import Path\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pairwise(x: torch.Tensor, dim=1):\n",
    "    \"\"\"Split the N x A x 2 state vector to return x_ego (N x A x 2) and x_other (N x A x (A - 1) x 4)\n",
    "    where each column index in x_other corresponds to -i agents for agent i\"\"\"\n",
    "    x_ = x.unsqueeze(dim=dim)\n",
    "    x = x.unsqueeze(dim=dim + 1)\n",
    "    _, x_b = torch.broadcast_tensors(x, x_)\n",
    "    # Create a mask for the diagonal\n",
    "    mask = 1 - torch.eye(x.shape[dim], device=x.device).unsqueeze(0).unsqueeze(-1)\n",
    "\n",
    "    # Expand the mask to match the input size\n",
    "    mask = mask.expand(*x_b.shape)\n",
    "\n",
    "    # Apply the mask to the input tensor\n",
    "    result = x_b[mask.bool()].view(x.shape[:2] + (x.shape[1] - 1, x.shape[-1]))\n",
    "    return x.squeeze(dim=dim + 1), result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Env:\n",
    "    def __init__(self, envFile):\n",
    "        self.file_path = Path(envFile)\n",
    "        with open(envFile) as f:\n",
    "            self.data = json.load(f)\n",
    "        self.parseMap()\n",
    "        self.parseAgents()\n",
    "        self.parseTasks()\n",
    "        self.numTasksReveal = int(self.data[\"numTasksReveal\"])\n",
    "        self.taskAssignmentStrategy = self.data[\"taskAssignmentStrategy\"]\n",
    "\n",
    "    def parseMap(self):\n",
    "        with open(self.file_path.parent/self.data[\"mapFile\"]) as f:\n",
    "            lines = f.readlines()\n",
    "        type_map = lines[0].strip().split(\" \")\n",
    "        height = int(lines[1].strip().split(\" \")[1])\n",
    "        width = int(lines[2].strip().split(\" \")[1])\n",
    "        map_str = lines[3]\n",
    "        self.map_1darr = []\n",
    "        for line in lines[4:]:\n",
    "            for char in line.strip():\n",
    "                self.map_1darr.append(char)\n",
    "        \n",
    "    def parseAgents(self):\n",
    "        with open(self.file_path.parent/self.data[\"agentFile\"]) as f:\n",
    "            lines = f.readlines()\n",
    "        self.n_agents = int(lines[0].strip().split(\" \")[0])\n",
    "        self.agents = []\n",
    "        for line in lines[1:]:\n",
    "            self.agents.append(int(line.strip().split(\" \")[0]))\n",
    "        self.agents = np.array(self.agents)\n",
    "\n",
    "    def parseTasks(self):\n",
    "        with open(self.file_path.parent/self.data[\"taskFile\"]) as f:\n",
    "            lines = f.readlines()\n",
    "        self.n_tasks = int(lines[0].strip().split(\" \")[0])\n",
    "        self.tasks = []\n",
    "        for line in lines[1:]:\n",
    "            self.tasks.append(int(line.strip().split(\" \")[0]))\n",
    "        self.tasks = np.array(self.tasks)\n",
    "\n",
    "class BatchedEnvironment:\n",
    "    def __init__(self, environments, num_envs):\n",
    "        self.maps = torch.stack([torch.tensor(env.map, dtype=torch.float32) for env in environments])\n",
    "        self.curr_states = torch.stack([torch.tensor(env.agents) for env in environments])\n",
    "        self.goal_locations = torch.stack([torch.tensor(env.tasks) for env in environments])\n",
    "\n",
    "    def compute_goal_achieved(self):\n",
    "        return torch.lingalg.norm(self.curr_state - self.goal_locations, dim=-1) == 0\n",
    "    \n",
    "    def compute_collision(self):\n",
    "        \"\"\"Construct the agent by agent distance matrix and check for two agents occupying the same cell\"\"\"\n",
    "        x_ego, x_other = pairwise(self.curr_states)\n",
    "        collisions = torch.norm(x_ego.unsqueeze(dim=-2) - x_other, dim=-1) < 0.01\n",
    "        return collisions.any(dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = [\"./example_problems/city.domain/paris_200.json\"]\n",
    "envs = []\n",
    "for e in p:\n",
    "    envs.append(Env(e))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(200,)"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.12 64-bit",
   "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.10.12"
  },
  "vscode": {
   "interpreter": {
    "hash": "e7370f93d1d0cde622a1f8e1c04877d8463912d04d973331ad4851f04de6915a"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}