add goal and collision code

Author: eugenevinitsky

index.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 59,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -15,7 +15,30 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "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": [
@@ -61,19 +84,28 @@
     "        self.tasks = np.array(self.tasks)\n",
     "\n",
     "class BatchedEnvironment:\n",
-    "    def __init__(self, environments):\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"
+    "        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": 55,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [
-    "p = [\"/home/aarav/aaravpandya/example_problems/city.domain/paris_200.json\", \"/home/aarav/aaravpandya/example_problems/game.domain/brc202d_200.json\"]\n",
+    "p = [\"./example_problems/city.domain/paris_200.json\"]\n",
     "envs = []\n",
     "for e in p:\n",
     "    envs.append(Env(e))"
@@ -107,7 +139,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "amazon",
+   "display_name": "Python 3.10.12 64-bit",
    "language": "python",
    "name": "python3"
   },
@@ -121,7 +153,12 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.5"
+   "version": "3.10.12"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "e7370f93d1d0cde622a1f8e1c04877d8463912d04d973331ad4851f04de6915a"
+   }
   }
  },
  "nbformat": 4,