Make3dAndTaxi.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/silviu/anaconda3/lib/python3.6/site-packages/statsmodels/tools/_testing.py:19: FutureWarning: pandas.util.testing is deprecated. Use the functions in the public API at pandas.testing instead.\n",
      "  import pandas.util.testing as tm\n"
     ]
    }
   ],
   "source": [
    "import random, numpy as np, matplotlib as mpl, matplotlib.pyplot as plt, seaborn as sns, multiprocessing as mp, time\n",
    "from collections import defaultdict\n",
    "mpl.rcParams['lines.linewidth'] = 1\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3d Gridworld"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "125000"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "D = 50\n",
    "N_STATES = D*D*D\n",
    "N_STATES"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def idx_to_coord(idx):\n",
    "  return idx // (D*D), (idx % (D*D)) // D, idx % D\n",
    "\n",
    "def coord_to_idx(coord):\n",
    "  x, y, z = coord\n",
    "  return x*D*D + y*D + z\n",
    "\n",
    "def neighbors_c(coord):\n",
    "  x, y, z = coord\n",
    "  return [\n",
    "    ((x + 1) % D, y, z),\n",
    "    ((x - 1) % D, y, z),\n",
    "    (x, (y + 1) % D, z),\n",
    "    (x, (y - 1) % D, z),\n",
    "    (x, y, (z + 1) % D),\n",
    "    (x, y, (z - 1) % D)\n",
    "  ]\n",
    "\n",
    "def neighbors_i(idx):\n",
    "  return list(map(coord_to_idx, neighbors_c(idx_to_coord(idx))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "edge_dict = dict()\n",
    "for node in range(N_STATES):\n",
    "  for neighbor in neighbors_i(node):\n",
    "    edge_dict[\"{} {}\".format(*sorted((node, neighbor)))] = 1 #np.random.randint(1, 101) / 100."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('3d1.edgelist', 'w') as f:\n",
    "  for k, v in edge_dict.items():\n",
    "    f.write('{} {}\\n'.format(k, v))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2d Taxi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "391250"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "D = 25\n",
    "N_STATES = D*D * D*D + D*D\n",
    "N_STATES"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def state_to_coord(state):\n",
    "  return hash(state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 5 Dim states: first two are gripper pos, next 2 are obj pos, if they are \n",
    "# the same, last coord indicates whether object is held\n",
    "state_list = []\n",
    "sh_dict = {}\n",
    "hs_dict = {}\n",
    "for a in range(D):\n",
    "  for b in range(D):\n",
    "    for c in range(D):\n",
    "      for d in range(D):\n",
    "        s = (a, b, c, d, 0)\n",
    "        sh_dict[s] = hash(s)\n",
    "        hs_dict[hash(s)] = s\n",
    "        state_list.append(s)\n",
    "        if a == c and b == d:\n",
    "          s = (a, b, c, d, 1)\n",
    "          sh_dict[s] = hash(s)\n",
    "          hs_dict[hash(s)] = s\n",
    "          state_list.append(s)\n",
    "assert(len(sh_dict) == len(hs_dict) == len(state_list) == N_STATES)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "hi_dict = {h:i for i, h in enumerate(hs_dict.keys())}\n",
    "ih_dict = {i:h for h, i in hi_dict.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "D = 25"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def twoDneighbors(x, y):\n",
    "  return list(set([(max(0, x-1), y), (x, max(0, y-1)), (min(D-1, x+1), y), (x, min(D-1, y+1))]) - set([(x, y)]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(23, 24), (24, 23)]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "twoDneighbors(24, 24)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def idx_to_coord(idx):\n",
    "  h = ih_dict[idx]\n",
    "  s = hs_dict[h]\n",
    "  return s\n",
    "\n",
    "def coord_to_idx(coord):\n",
    "  h = sh_dict[coord]\n",
    "  i = hi_dict[h]\n",
    "  return i\n",
    "\n",
    "def neighbors_c(coord):\n",
    "  ns = []\n",
    "  c = coord\n",
    "  # Gripper holding box\n",
    "  if c[4] == 1:\n",
    "    # Gripper drops box\n",
    "    ns.append((c[0],c[1],c[2],c[3],0))\n",
    "    # Gripper and box move together\n",
    "    for n in twoDneighbors(c[0], c[1]):\n",
    "      ns.append((n[0], n[1], n[0], n[1], 1))\n",
    "  # Gripper not holding box\n",
    "  elif c[4] == 0:\n",
    "    # Gripper picks up box if at same position\n",
    "    if c[0] == c[2] and c[1] == c[3]:\n",
    "      ns.append((c[0],c[1],c[2],c[3],1))\n",
    "    # Gripper moves, leaving the box where it is\n",
    "    for n in twoDneighbors(c[0], c[1]):\n",
    "      ns.append((n[0],n[1],c[2],c[3], 0))\n",
    "  return ns\n",
    "\n",
    "def neighbors_i(idx):\n",
    "  return list(map(coord_to_idx, neighbors_c(idx_to_coord(idx))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "edge_dict = dict()\n",
    "for node in range(N_STATES):\n",
    "  for neighbor in neighbors_i(node):\n",
    "    edge_dict[\"{} {}\".format(*sorted((node, neighbor)))] = 1 #np.random.randint(1, 101) / 100."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('taxi1.edgelist', 'w') as f:\n",
    "  for k, v in edge_dict.items():\n",
    "    f.write('{} {}\\n'.format(k, v))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3d1.edgelist               \u001b[0m\u001b[01;32mMake3dAndTaxi.ipynb\u001b[0m*   \u001b[01;34m__pycache__\u001b[0m/\r\n",
      "custom_metric_loss_ops.py  metrics_pytorch.py     readme.md\r\n",
      "data.py                    metrics_tf1.py         taxi1.edgelist\r\n",
      "experiment.py              \u001b[01;32mPrepXYDPickles.ipynb\u001b[0m*\r\n"
     ]
    }
   ],
   "source": [
    "ls"
   ]
  },
  {
   "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.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}