minigame10a.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Minigame 10a: Choose One Element To Refine\n",
    "\n",
    "This is essentially our \"global environment\" for refinement where the action is choosing one element at a time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "from math import cos,sin\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/rwa/pyvenv/lib/python3.6/site-packages/tensorflow/python/compat/v2_compat.py:96: disable_resource_variables (from tensorflow.python.ops.variable_scope) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "non-resource variables are not supported in the long term\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "import gym\n",
    "from gym import spaces, utils\n",
    "import numpy as np\n",
    "import ray\n",
    "import ray.rllib.agents.ppo as ppo\n",
    "import ray.rllib.agents.dqn as dqn\n",
    "from os.path import expanduser, join\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from glvis import glvis, to_stream"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from mfem import path\n",
    "import mfem.ser as mfem"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Start up rllib"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define some synthetic test functions: steps and bumps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rotate(x,theta):\n",
    "    x0 = x[0]\n",
    "    y0 = x[1]\n",
    "    x1 = x0*cos(theta)-y0*sin(theta)\n",
    "    y1 = x0*sin(theta)+y0*cos(theta)\n",
    "    return [x1,y1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def step(x):\n",
    "    x0 = x[0]\n",
    "    if (x0 < 0.0):\n",
    "        return 1.0\n",
    "    else:\n",
    "        return 0.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rotated_step(x, theta):\n",
    "    xr = rotate(x,theta)\n",
    "    return step(xr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def bump(x):\n",
    "    rsq = x[0]**2 +x[1]**2\n",
    "    v = math.exp(-rsq)\n",
    "    assert v <= 1.0\n",
    "    assert v >= 0.0\n",
    "    return v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def smooth_step(x):\n",
    "    return 0.5*(1.0 +math.tanh(x[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rotated_smooth_step(x,theta):\n",
    "    xr = rotate(x,theta)\n",
    "    return smooth_step(xr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create classes where we can set the parameters and then eval a bunch of points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Step(mfem.PyCoefficient):\n",
    "    \n",
    "    def SetParams(self):\n",
    "        self.theta = random.uniform(0.0, 2.0*math.pi)\n",
    "        self.dx = [random.uniform(-1.0, 1.0),random.uniform(-1.0, 1.0)]\n",
    "        \n",
    "    def EvalValue(self, x):\n",
    "        return rotated_step(x+self.dx, self.theta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Bump(mfem.PyCoefficient):\n",
    "    \n",
    "    def SetParams(self):\n",
    "        self.width = random.uniform(0.1,0.5)\n",
    "        self.xc = [0.5,0.5]\n",
    "        self.dx = [random.uniform(-0.5, 0.5),random.uniform(-0.5, 0.5)]\n",
    "        y1 = random.uniform(0.1,0.9)\n",
    "        y2 = random.uniform(0.1,0.9)\n",
    "        self.floor = min(y1,y2)\n",
    "        self.ceiling = max(y1,y2)\n",
    "        self.height = self.ceiling -self.floor\n",
    "        \n",
    "    def Print(self):\n",
    "        print(\"width = %f\" % self.width)\n",
    "        print(\"xc = %f,%f\" % (self.xc[0],self.xc[1]))\n",
    "        print(\"dx = %f,%f\" % (self.dx[0],self.dx[1]))\n",
    "        print(\"floor = %f\" % self.floor)\n",
    "        print(\"ceil = %f\"  % self.ceiling)\n",
    "        print(\"height = %f\"% self.height)\n",
    "        \n",
    "    def EvalValue(self, x):\n",
    "        v = self.floor +self.height*bump((x-self.xc+self.dx)/self.width)\n",
    "        #assert v > 0.0\n",
    "        #assert v < 1.0\n",
    "        return v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TwoBump(mfem.PyCoefficient):\n",
    "    \n",
    "    def SetParams(self):\n",
    "        self.width1 = random.uniform(0.1,0.5)\n",
    "        self.width2 = random.uniform(0.1,0.5)\n",
    "        self.xc1 = [0.5,0.5]\n",
    "        self.xc2 = [0.5,0.5]\n",
    "        self.dx1 = [random.uniform(-0.5, 0.5),random.uniform(-0.5, 0.5)]\n",
    "        self.dx2 = [random.uniform(-0.5, 0.5),random.uniform(-0.5, 0.5)]\n",
    "\n",
    "    def EvalValue(self, x):\n",
    "        return 0.5*(bump((x-self.xc1+self.dx1)/self.width1)+bump((x-self.xc2+self.dx2)/self.width2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SmoothStep(mfem.PyCoefficient):\n",
    "    \n",
    "    def SetParams(self):\n",
    "        self.width = random.uniform(5.0, 10.0)\n",
    "        self.xc = [0.5,0.5]\n",
    "        self.dx = random.uniform(-0.5,0.5)\n",
    "        self.theta = random.uniform(0.0, 2.0*math.pi)\n",
    "        self.height = random.uniform(0.0, 1.0)\n",
    "        y1 = random.uniform(0.0,1.0)\n",
    "        y2 = random.uniform(0.0,1.0)\n",
    "        self.floor = min(y1,y2)\n",
    "        self.ceiling = max(y1,y2)\n",
    "        self.height = self.ceiling -self.floor\n",
    "\n",
    "    def EvalValue(self, x):\n",
    "        x -= self.xc\n",
    "        x += self.dx\n",
    "        return self.floor +self.height*rotated_smooth_step(x*self.width, self.theta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BumpAndSmoothStep(mfem.PyCoefficient):\n",
    "    \n",
    "    def SetParams(self):\n",
    "        self.bump = Bump()\n",
    "        self.bump.SetParams()\n",
    "        self.smooth_step = SmoothStep()\n",
    "        self.smooth_step.SetParams()\n",
    "        self.alpha = random.uniform(0.0, 1.0)\n",
    "\n",
    "    def EvalValue(self, x):\n",
    "        return self.alpha*self.bump.EvalValue(x)+ (1-self.alpha)*self.smooth_step.EvalValue(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BumpNarrowWide(mfem.PyCoefficient):\n",
    "    \n",
    "    def SetParams(self):\n",
    "        a = random.uniform(0.0,1.0)\n",
    "        if (a < 0.5):\n",
    "            self.width = 0.2\n",
    "            self.height = 1.0\n",
    "        else:\n",
    "            self.width = 0.4\n",
    "            self.height = 0.1\n",
    "        self.xc = [0.5,0.5]\n",
    "        self.dx = [random.uniform(-0.5, 0.5),random.uniform(-0.5, 0.5)]\n",
    "\n",
    "    def EvalValue(self, x):\n",
    "        return self.height*bump((x-self.xc+self.dx)/self.width)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DiscreteBump(mfem.PyCoefficient):\n",
    "    \n",
    "    def SetParams(self):\n",
    "        nlocs = 100\n",
    "        nmesh = 5\n",
    "        i = random.randrange(nlocs)\n",
    "        j = random.randrange(nlocs)\n",
    "        dx = 1.0/nlocs\n",
    "        dx_mesh = 1.0/nmesh\n",
    "        self.xc = [0.,0.]\n",
    "        self.xc[0] = i*dx+0.5*dx\n",
    "        self.xc[1] = j*dx+0.5*dx\n",
    "        self.width = dx_mesh/2\n",
    "        #print(\"(%d,%d)\" % (i,j))\n",
    "        #print(\"(%f,%f)\" % (self.xc[0],self.xc[1]))\n",
    "        \n",
    "    def Print(self):\n",
    "        pass\n",
    "        \n",
    "    def EvalValue(self, x):\n",
    "        v = bump((x-self.xc)/self.width)\n",
    "        assert v >= 0.0\n",
    "        assert v <= 1.0\n",
    "        return v"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualize an instance of the test function. Note that each instance has randomly chosen parameters.  For the steps, it's a rotation angle and a displacement.  For the bumps, it's a width and a displacement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh1 = mfem.Mesh('inline-quad-5.mesh')\n",
    "mesh1.UniformRefinement()\n",
    "fec1 = mfem.L2_FECollection(p=1, dim=2)\n",
    "fes1 = mfem.FiniteElementSpace(mesh1, fec1)\n",
    "u1 = mfem.GridFunction(fes1)\n",
    "c1 = Bump()\n",
    "c1.SetParams()\n",
    "u1.ProjectCoefficient(c1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "39e234bf0be3402fa32625cddd084f82",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "glvis()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "glvis(to_stream(mesh1,u1) + 'keys Rjlmc',600,600)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create the gym environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AMRGame(gym.Env):\n",
    "    \n",
    "    class u0_coeff(mfem.PyCoefficient):\n",
    "        \n",
    "        def SetParams(self):\n",
    "            #self.fn = BumpAndSmoothStep()\n",
    "            self.fn = Bump()\n",
    "\n",
    "            self.fn.SetParams()\n",
    "            \n",
    "        def Print(self):\n",
    "            self.fn.Print()\n",
    "            \n",
    "        def EvalValue(self, x):\n",
    "            v = self.fn.EvalValue(x)\n",
    "            #assert v >= 0.0\n",
    "            #assert v <= 1.0\n",
    "            return self.fn.EvalValue(x)\n",
    "        \n",
    "    # In RLlib, you need the config arg\n",
    "    def __init__(self,config):\n",
    "        self.meshfile = 'inline-quad-5.mesh'\n",
    "        \n",
    "        # keep a copy of the unrefined mesh so we can restore it from memory\n",
    "        self.mesh0 = mfem.Mesh(self.meshfile)\n",
    "        self.mesh = mfem.Mesh(self.meshfile)\n",
    "        \n",
    "        # The only reason we need to create an fespace and gf here\n",
    "        # is to find the sizes needed for the action and observation spaces\n",
    "        dim = self.mesh.Dimension()\n",
    "        self.order = 1\n",
    "        self.fec = mfem.L2_FECollection(self.order, dim)\n",
    "        self.fes = mfem.FiniteElementSpace(self.mesh, self.fec)\n",
    "        self.u = mfem.GridFunction(self.fes);\n",
    "\n",
    "        # actions are: refine each element, or do nothing\n",
    "        self.action_space = spaces.Discrete(self.mesh.GetNE())\n",
    "        \n",
    "        # observation is a set of sample points forming a 2D image\n",
    "        n = int(math.sqrt(self.mesh.GetNE()))\n",
    "        self.obsx = 2*n\n",
    "        self.obsy = 2*n\n",
    "        self.obsx = 42\n",
    "        self.obsy = 42\n",
    "        \n",
    "        # add a little extra range on the space to account for interpolation errors\n",
    "        self.observation_space = spaces.Box(-1.0, 2.0, shape=(self.obsx,self.obsy,1))\n",
    "        self.get_obs_points()\n",
    "        \n",
    "        self.n = 0\n",
    "        \n",
    "        # call reset to create the first synthetic function\n",
    "        self.reset()\n",
    "        \n",
    "        #self.gl = GlvisWidget(get_solnstream(self.mesh,self.u))\n",
    "        \n",
    "    # precompute the observation points and the elements and integration points we need\n",
    "    def get_obs_points(self):\n",
    "        n = math.sqrt(self.mesh.GetNE())\n",
    "        dx = 1.0/self.obsx\n",
    "        dy = 1.0/self.obsy\n",
    "        self.sample_pts = []\n",
    "        self.sample_els = []\n",
    "        self.sample_ips = []\n",
    "        for j in range(self.obsy):\n",
    "            for i in range(self.obsx):\n",
    "                pt = [i*dx+0.5*dx,j*dy+0.5*dy]\n",
    "                self.sample_pts.append(pt)\n",
    "                n, el, ip = self.mesh.FindPoints([pt])\n",
    "                #assert n == 1\n",
    "                #assert ip[0].x > 0.0\n",
    "                #assert ip[0].x < 1.0\n",
    "                #assert ip[0].y > 0.0\n",
    "                #assert ip[0].y < 1.0\n",
    "                #assert el[0] >= 0\n",
    "                #assert el[0] < self.mesh.GetNE()\n",
    "                # copy these so they won't be destroyed when mesh goes away?\n",
    "                ip0 = mfem.IntegrationPoint()\n",
    "                ip0.x = ip[0].x\n",
    "                ip0.y = ip[0].y\n",
    "                self.sample_els.append(el[0])\n",
    "                self.sample_ips.append(ip0)\n",
    "                \n",
    "    def get_obs(self):\n",
    "        state = np.empty((self.obsx,self.obsy,1))\n",
    "        k = 0\n",
    "        for j in range(self.obsy):\n",
    "            for i in range(self.obsx):\n",
    "                #assert k < len(self.sample_els)\n",
    "                #assert k < len(self.sample_ips)\n",
    "                el = self.sample_els[k]\n",
    "                #assert el >= 0\n",
    "                #assert el < self.mesh.GetNE()\n",
    "                ip = self.sample_ips[k]\n",
    "                #assert ip.x >= 0.0, \"k={}, i={}, j={}, x is {}\".format(k,i,j,ip.x)\n",
    "                #assert ip.x <= 1.0, \"k={}, i={}, j={}, x is {}\".format(k,i,j,ip.x)\n",
    "                #assert ip.y >= 0.0, \"k={}, i={}, j={}, y is {}\".format(k,i,j,ip.y)\n",
    "                #assert ip.y <= 1.0, \"k={}, i={}, j={}, y is {}\".format(k,i,j,ip.y)\n",
    "                v = self.u.GetValue(self.sample_els[k],self.sample_ips[k])\n",
    "                state[i][j] = v\n",
    "                if (v > 2.0 or v < -1.0):\n",
    "                    print(\"element %d\" % self.sample_els[k])\n",
    "                    print(\"ip.x = %f\" % self.sample_ips[k].x)\n",
    "                    print(\"ip.y = %f\" % self.sample_ips[k].y)\n",
    "                    print(\"%d,%d -> %f\" % (i,j,v))\n",
    "                    print(\"%d,%d -> %f\" % (i,j,state[i][j]))\n",
    "                    self.u0.Print()\n",
    "                k += 1\n",
    "        self.state = state\n",
    "        return state\n",
    "                    \n",
    "    # Compute L2 error wrt to the analytic fn definition\n",
    "    def get_error(self):\n",
    "        err = self.u.ComputeL2Error(self.u0)\n",
    "        return err\n",
    "    \n",
    "    # Manually refine the elements in the array elems\n",
    "    def refine_elems(self, elems):\n",
    "        self.mesh.GeneralRefinement(mfem.intArray(elems))\n",
    "        self.fes.Update()\n",
    "        self.u.Update()\n",
    "        self.u.ProjectCoefficient(self.u0)\n",
    "            \n",
    "    # action is the number of the element to refine\n",
    "    def step(self, action):\n",
    "        self.n += 1\n",
    "        #print(\"begin step %d: action is %d\" % (self.n,action))\n",
    "        err1 = self.get_error()\n",
    "        #print(\"err1 = %f\" % err1)\n",
    "        self.refine_elems([action])\n",
    "        err2 = self.get_error()\n",
    "        #print(\"err2 = %f\" % err2)\n",
    "        reward = err1-err2\n",
    "        done = True\n",
    "        #print(\"step reward = %f\" % reward)\n",
    "        #obs = self.get_obs()\n",
    "        \n",
    "        # use old state\n",
    "        return np.array(self.state), reward, done, {}\n",
    "    \n",
    "    # similar to reset, but do not choose a new function\n",
    "    def reinit(self):\n",
    "        #print(\"reinit\")\n",
    "        del self.mesh\n",
    "        self.mesh = mfem.Mesh(self.mesh0)\n",
    "\n",
    "        del self.fes\n",
    "        self.fes = mfem.FiniteElementSpace(self.mesh, self.fec)\n",
    "\n",
    "        del self.u\n",
    "        self.u = mfem.GridFunction(self.fes)\n",
    "        self.u.ProjectCoefficient(self.u0)\n",
    "        \n",
    "        #self.get_obs_points()\n",
    "        obs = self.get_obs()\n",
    "        \n",
    "        return np.array(obs)\n",
    "    \n",
    "    # every reset of the env chooses a new synthetic function\n",
    "    def reset(self):\n",
    "        #print(\"reset\")\n",
    "        self.u0 = self.u0_coeff()\n",
    "        self.u0.SetParams()\n",
    "        return self.reinit()\n",
    "    \n",
    "    def render(self):\n",
    "        return glvis(to_stream(self.mesh,self.u) + 'keys Rjlmc',600,600)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Instantiate the environment and sanity check it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "545c98f5ab154c3b89f0826629b399b6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "glvis()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "env = AMRGame(None)\n",
    "env.render()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "02c819e4172b4a1a9b7051cccb566fe8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "glvis()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "env.reinit()\n",
    "obs, reward, done, _ = env.step(2)\n",
    "env.render()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.324055291670144e-07"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "reward"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ok, try training a policy:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ray.shutdown()\n",
    "ray.init(ignore_reinit_error=True)\n",
    "config = ppo.DEFAULT_CONFIG.copy()\n",
    "config['train_batch_size'] = int(1e4)\n",
    "config['num_workers'] = 3\n",
    "#config['kl_coeff']\n",
    "config"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "config['train_batch_size'] = int(2e2)\n",
    "agent = ppo.PPOTrainer(config, env=AMRGame)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "%%time\n",
    "for n in range(1):\n",
    "    result = agent.train()\n",
    "    print(\"episode reward mean: %f \" % result[\"episode_reward_mean\"])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "policy = agent.get_policy()\n",
    "model = policy.model\n",
    "print(model.base_model.summary())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a convenience function for applying a policy to a given observation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "def apply_policy(model, obs):\n",
    "    action = agent.compute_action(obs, explore=False) # use deterministic mode\n",
    "    state, reward, done, info = env.step(action)\n",
    "    #print(\"policy chooses action %d with reward %f\" % (action, reward))\n",
    "    return action, reward"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "obs = env.reset()\n",
    "action, reward = apply_policy(model, obs)\n",
    "action, reward"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "obs = env.reinit()\n",
    "action, reward = apply_policy(model, obs)\n",
    "action, reward"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Brute force search for the best choice by trying each one, remembering to reset the environment after each action and after we're done."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def find_optimal(obs):\n",
    "    u0 = mfem.Vector(obs)\n",
    "    maxr = 0.0;\n",
    "    maxel = -1;\n",
    "    env.reinit()\n",
    "    ne = env.mesh.GetNE()\n",
    "    for n in range(ne):\n",
    "        env.reinit()\n",
    "        state, reward, done, info = env.step(n)\n",
    "        if reward > maxr:\n",
    "            maxr = reward\n",
    "            maxel = n\n",
    "    #print(\"max reward is %f by refining element %d\" % (maxr, maxel))\n",
    "    env.reinit()\n",
    "    return maxel, maxr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f021239c4caa4f85b7d04a7fb09e6da8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "glvis()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "obs = env.reset()\n",
    "maxel, maxr = find_optimal(obs)\n",
    "env.refine_elems([maxel])\n",
    "env.render()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare with what the policy does:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "env.reinit()\n",
    "apply_policy(model,obs)\n",
    "env.render()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define an error estimator based on the difference between the discontinuous and continuous representations. This is only valid for L2 FE spaces."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_dgjumps(env):\n",
    "    \n",
    "    mesh = env.mesh\n",
    "    u = env.u\n",
    "    \n",
    "    # put the L2 gridfunction into a coefficient so we can project it\n",
    "    u_disc_coeff = mfem.GridFunctionCoefficient(u)\n",
    "    h1_fec = mfem.H1_FECollection(p=1, dim=2)\n",
    "    h1_fes = mfem.FiniteElementSpace(mesh, h1_fec)\n",
    "    u_h1 = mfem.GridFunction(h1_fes)\n",
    "    u_h1.ProjectDiscCoefficient(u_disc_coeff, mfem.GridFunction.ARITHMETIC)\n",
    "    \n",
    "    # put the H1 smoothed function into a coefficient\n",
    "    u_h1_coeff = mfem.GridFunctionCoefficient(u_h1)\n",
    "    \n",
    "    # create a 0-order L2 field to hold errors\n",
    "    l2_0_fec = mfem.L2_FECollection(p=0,dim=2)\n",
    "    l2_0_fes = mfem.FiniteElementSpace(mesh,l2_0_fec)\n",
    "\n",
    "    # Compute elementwise \"errors\" between continuous and discontinuous fields\n",
    "    err_gf = mfem.GridFunction(l2_0_fes);\n",
    "    u.ComputeElementL2Errors(u_h1_coeff, err_gf);\n",
    "    \n",
    "    best_action = np.argmax(err_gf.GetDataArray())\n",
    "    \n",
    "    state, reward, done, info = env.step(best_action)\n",
    "    env.reinit()\n",
    "\n",
    "    return best_action, reward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "60e6b8e61e05422888d20ca41e9da7b4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "glvis()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "env.reset()\n",
    "action, reward = find_dgjumps(env)\n",
    "env.step(action)\n",
    "env.render()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run a more systematic evaluation using an ensemble of samples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval_ensemble(model, ntrials):\n",
    "    ncorrect = 0.0\n",
    "    sumsq = 0.0\n",
    "    maxerrsq = 0.0\n",
    "    dg_ncorrect = 0.0\n",
    "    dg_sumsq = 0.0\n",
    "    dg_maxerrsq = 0.0\n",
    "    for n in range(ntrials):\n",
    "        obs = env.reset()\n",
    "        bestaction, bestreward = find_optimal(obs)\n",
    "        dgaction, dgreward = find_dgjumps(env)\n",
    "        action, reward = apply_policy(model,obs)\n",
    "        err = bestreward-reward\n",
    "        maxerrsq = max(err*err,maxerrsq)\n",
    "        sumsq += err*err\n",
    "        dg_err = bestreward-dgreward\n",
    "        dg_maxerrsq = max(dg_err*dg_err,dg_maxerrsq)\n",
    "        dg_sumsq += dg_err*dg_err\n",
    "        if (bestaction == action):\n",
    "            ncorrect += 1\n",
    "        if (bestaction == dgaction):\n",
    "            dg_ncorrect += 1\n",
    "    rms = math.sqrt(sumsq/ntrials)\n",
    "    corr = 100.*ncorrect/ntrials\n",
    "    print(\"policy rms error: \",rms,flush=True)\n",
    "    print(\"policy max sq error: \",maxerrsq,flush=True)\n",
    "    print(\"policy % correct: \",corr,flush=True)\n",
    "    dg_rms = math.sqrt(dg_sumsq/ntrials)\n",
    "    dg_corr = 100.*dg_ncorrect/ntrials\n",
    "    print(\"dg rms error: \",dg_rms,flush=True)\n",
    "    print(\"dg max sq error: \",dg_maxerrsq,flush=True)\n",
    "    print(\"dg % correct: \",dg_corr,flush=True)\n",
    "    return rms, math.sqrt(maxerrsq), corr, dg_rms, math.sqrt(dg_maxerrsq), dg_corr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "eval_ensemble(model, 100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run a few eval sample sizes to get a sense of how many are needed to estimate the metrics of the policy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "eval_ensemble(model, 200)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "eval_ensemble(model, 400)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see if the training process is making progress:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-02-26 16:03:27,215\tINFO services.py:1173 -- View the Ray dashboard at \u001b[1m\u001b[32mhttp://127.0.0.1:8265\u001b[39m\u001b[22m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'num_workers': 3,\n",
       " 'num_envs_per_worker': 1,\n",
       " 'create_env_on_driver': False,\n",
       " 'rollout_fragment_length': 200,\n",
       " 'batch_mode': 'truncate_episodes',\n",
       " 'num_gpus': 0,\n",
       " 'train_batch_size': 1000,\n",
       " 'model': {'fcnet_hiddens': [256, 256],\n",
       "  'fcnet_activation': 'tanh',\n",
       "  'conv_filters': None,\n",
       "  'conv_activation': 'relu',\n",
       "  'free_log_std': False,\n",
       "  'no_final_linear': False,\n",
       "  'vf_share_layers': True,\n",
       "  'use_lstm': False,\n",
       "  'max_seq_len': 20,\n",
       "  'lstm_cell_size': 256,\n",
       "  'lstm_use_prev_action': False,\n",
       "  'lstm_use_prev_reward': False,\n",
       "  '_time_major': False,\n",
       "  'framestack': True,\n",
       "  'dim': 84,\n",
       "  'grayscale': False,\n",
       "  'zero_mean': True,\n",
       "  'custom_model': None,\n",
       "  'custom_model_config': {},\n",
       "  'custom_action_dist': None,\n",
       "  'custom_preprocessor': None,\n",
       "  'lstm_use_prev_action_reward': -1},\n",
       " 'optimizer': {},\n",
       " 'gamma': 0.99,\n",
       " 'horizon': None,\n",
       " 'soft_horizon': False,\n",
       " 'no_done_at_end': False,\n",
       " 'env_config': {},\n",
       " 'env': None,\n",
       " 'normalize_actions': False,\n",
       " 'clip_rewards': None,\n",
       " 'clip_actions': True,\n",
       " 'preprocessor_pref': 'deepmind',\n",
       " 'lr': 5e-05,\n",
       " 'monitor': False,\n",
       " 'log_level': 'WARN',\n",
       " 'callbacks': ray.rllib.agents.callbacks.DefaultCallbacks,\n",
       " 'ignore_worker_failures': False,\n",
       " 'log_sys_usage': True,\n",
       " 'fake_sampler': False,\n",
       " 'framework': 'tf',\n",
       " 'eager_tracing': False,\n",
       " 'explore': True,\n",
       " 'exploration_config': {'type': 'StochasticSampling'},\n",
       " 'evaluation_interval': None,\n",
       " 'evaluation_num_episodes': 10,\n",
       " 'in_evaluation': False,\n",
       " 'evaluation_config': {},\n",
       " 'evaluation_num_workers': 0,\n",
       " 'custom_eval_function': None,\n",
       " 'sample_async': False,\n",
       " '_use_trajectory_view_api': True,\n",
       " 'observation_filter': 'NoFilter',\n",
       " 'synchronize_filters': True,\n",
       " 'tf_session_args': {'intra_op_parallelism_threads': 2,\n",
       "  'inter_op_parallelism_threads': 2,\n",
       "  'gpu_options': {'allow_growth': True},\n",
       "  'log_device_placement': False,\n",
       "  'device_count': {'CPU': 1},\n",
       "  'allow_soft_placement': True},\n",
       " 'local_tf_session_args': {'intra_op_parallelism_threads': 8,\n",
       "  'inter_op_parallelism_threads': 8},\n",
       " 'compress_observations': False,\n",
       " 'collect_metrics_timeout': 180,\n",
       " 'metrics_smoothing_episodes': 100,\n",
       " 'remote_worker_envs': False,\n",
       " 'remote_env_batch_wait_ms': 0,\n",
       " 'min_iter_time_s': 0,\n",
       " 'timesteps_per_iteration': 0,\n",
       " 'seed': None,\n",
       " 'extra_python_environs_for_driver': {},\n",
       " 'extra_python_environs_for_worker': {},\n",
       " 'num_cpus_per_worker': 1,\n",
       " 'num_gpus_per_worker': 0,\n",
       " 'custom_resources_per_worker': {},\n",
       " 'num_cpus_for_driver': 1,\n",
       " 'memory': 0,\n",
       " 'object_store_memory': 0,\n",
       " 'memory_per_worker': 0,\n",
       " 'object_store_memory_per_worker': 0,\n",
       " 'input': 'sampler',\n",
       " 'input_evaluation': ['is', 'wis'],\n",
       " 'postprocess_inputs': False,\n",
       " 'shuffle_buffer_size': 0,\n",
       " 'output': None,\n",
       " 'output_compress_columns': ['obs', 'new_obs'],\n",
       " 'output_max_file_size': 67108864,\n",
       " 'multiagent': {'policies': {},\n",
       "  'policy_mapping_fn': None,\n",
       "  'policies_to_train': None,\n",
       "  'observation_fn': None,\n",
       "  'replay_mode': 'independent'},\n",
       " 'logger_config': None,\n",
       " 'replay_sequence_length': 1,\n",
       " 'use_critic': True,\n",
       " 'use_gae': True,\n",
       " 'lambda': 1.0,\n",
       " 'kl_coeff': 0.2,\n",
       " 'sgd_minibatch_size': 128,\n",
       " 'shuffle_sequences': True,\n",
       " 'num_sgd_iter': 30,\n",
       " 'lr_schedule': None,\n",
       " 'vf_share_layers': False,\n",
       " 'vf_loss_coeff': 1.0,\n",
       " 'entropy_coeff': 0.0,\n",
       " 'entropy_coeff_schedule': None,\n",
       " 'clip_param': 0.3,\n",
       " 'vf_clip_param': 10.0,\n",
       " 'grad_clip': None,\n",
       " 'kl_target': 0.01,\n",
       " 'simple_optimizer': False,\n",
       " '_fake_gpus': False}"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ray.shutdown()\n",
    "ray.init(ignore_reinit_error=True)\n",
    "\n",
    "total_episodes = 2.e6\n",
    "batch_size = 1000\n",
    "nbatches = int(total_episodes/batch_size)\n",
    "\n",
    "config = ppo.DEFAULT_CONFIG.copy()\n",
    "config['train_batch_size'] = batch_size\n",
    "config['num_workers'] = 3\n",
    "config['num_gpus'] = 0\n",
    "#config['lr'] = 1.e-5\n",
    "config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-02-26 16:03:29,670\tINFO trainer.py:591 -- Tip: set framework=tfe or the --eager flag to enable TensorFlow eager execution\n",
      "2021-02-26 16:03:29,671\tINFO trainer.py:618 -- Current log_level is WARN. For more information, set 'log_level': 'INFO' / 'DEBUG' or use the -v and -vv flags.\n",
      "\u001b[2m\u001b[36m(pid=5580)\u001b[0m WARNING:tensorflow:From /home/rwa/pyvenv/lib/python3.6/site-packages/tensorflow/python/compat/v2_compat.py:96: disable_resource_variables (from tensorflow.python.ops.variable_scope) is deprecated and will be removed in a future version.\n",
      "\u001b[2m\u001b[36m(pid=5580)\u001b[0m Instructions for updating:\n",
      "\u001b[2m\u001b[36m(pid=5580)\u001b[0m non-resource variables are not supported in the long term\n",
      "\u001b[2m\u001b[36m(pid=5578)\u001b[0m WARNING:tensorflow:From /home/rwa/pyvenv/lib/python3.6/site-packages/tensorflow/python/compat/v2_compat.py:96: disable_resource_variables (from tensorflow.python.ops.variable_scope) is deprecated and will be removed in a future version.\n",
      "\u001b[2m\u001b[36m(pid=5578)\u001b[0m Instructions for updating:\n",
      "\u001b[2m\u001b[36m(pid=5578)\u001b[0m non-resource variables are not supported in the long term\n",
      "\u001b[2m\u001b[36m(pid=5573)\u001b[0m WARNING:tensorflow:From /home/rwa/pyvenv/lib/python3.6/site-packages/tensorflow/python/compat/v2_compat.py:96: disable_resource_variables (from tensorflow.python.ops.variable_scope) is deprecated and will be removed in a future version.\n",
      "\u001b[2m\u001b[36m(pid=5573)\u001b[0m Instructions for updating:\n",
      "\u001b[2m\u001b[36m(pid=5573)\u001b[0m non-resource variables are not supported in the long term\n",
      "2021-02-26 16:03:35,696\tWARNING util.py:43 -- Install gputil for GPU system monitoring.\n",
      "\u001b[2m\u001b[36m(pid=5580)\u001b[0m 2021-02-26 16:03:35,705\tWARNING deprecation.py:30 -- DeprecationWarning: `env_index` has been deprecated. Use `episode.env_id` instead. This will raise an error in the future!\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "training batch 0 of size 1000\n",
      "WARNING:tensorflow:From /home/rwa/pyvenv/lib/python3.6/site-packages/ray/rllib/policy/tf_policy.py:852: Variable.load (from tensorflow.python.ops.variables) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Prefer Variable.assign which has equivalent behavior in 2.X.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(pid=5573)\u001b[0m WARNING:tensorflow:From /home/rwa/pyvenv/lib/python3.6/site-packages/ray/rllib/policy/tf_policy.py:852: Variable.load (from tensorflow.python.ops.variables) is deprecated and will be removed in a future version.\n",
      "\u001b[2m\u001b[36m(pid=5573)\u001b[0m Instructions for updating:\n",
      "\u001b[2m\u001b[36m(pid=5573)\u001b[0m Prefer Variable.assign which has equivalent behavior in 2.X.\n",
      "\u001b[2m\u001b[36m(pid=5580)\u001b[0m WARNING:tensorflow:From /home/rwa/pyvenv/lib/python3.6/site-packages/ray/rllib/policy/tf_policy.py:852: Variable.load (from tensorflow.python.ops.variables) is deprecated and will be removed in a future version.\n",
      "\u001b[2m\u001b[36m(pid=5580)\u001b[0m Instructions for updating:\n",
      "\u001b[2m\u001b[36m(pid=5580)\u001b[0m Prefer Variable.assign which has equivalent behavior in 2.X.\n",
      "\u001b[2m\u001b[36m(pid=5578)\u001b[0m WARNING:tensorflow:From /home/rwa/pyvenv/lib/python3.6/site-packages/ray/rllib/policy/tf_policy.py:852: Variable.load (from tensorflow.python.ops.variables) is deprecated and will be removed in a future version.\n",
      "\u001b[2m\u001b[36m(pid=5578)\u001b[0m Instructions for updating:\n",
      "\u001b[2m\u001b[36m(pid=5578)\u001b[0m Prefer Variable.assign which has equivalent behavior in 2.X.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "training batch 1 of size 1000\n",
      "training batch 2 of size 1000\n",
      "training batch 3 of size 1000\n",
      "training batch 4 of size 1000\n",
      "/home/rwa/ray_results/PPO_AMRGame_2021-02-26_16-03-29rmbq4q_r/checkpoint_5/checkpoint-5\n",
      "policy rms error:  0.005723741959471854\n",
      "policy max sq error:  0.0009568858256583815\n",
      "policy % correct:  72.0\n",
      "dg rms error:  0.0036886531445395033\n",
      "dg max sq error:  0.00018448834829059654\n",
      "dg % correct:  64.0\n",
      "training batch 5 of size 1000\n",
      "training batch 6 of size 1000\n",
      "training batch 7 of size 1000\n",
      "training batch 8 of size 1000\n",
      "training batch 9 of size 1000\n",
      "/home/rwa/ray_results/PPO_AMRGame_2021-02-26_16-03-29rmbq4q_r/checkpoint_10/checkpoint-10\n",
      "policy rms error:  0.0003924991999552904\n",
      "policy max sq error:  2.641714267611632e-06\n",
      "policy % correct:  76.0\n",
      "dg rms error:  0.0029788947472461527\n",
      "dg max sq error:  0.0001883538823050839\n",
      "dg % correct:  74.0\n",
      "training batch 10 of size 1000\n",
      "training batch 11 of size 1000\n",
      "training batch 12 of size 1000\n",
      "training batch 13 of size 1000\n",
      "training batch 14 of size 1000\n",
      "/home/rwa/ray_results/PPO_AMRGame_2021-02-26_16-03-29rmbq4q_r/checkpoint_15/checkpoint-15\n",
      "policy rms error:  0.0002717430700145064\n",
      "policy max sq error:  1.244473326475883e-06\n",
      "policy % correct:  86.0\n",
      "dg rms error:  0.003172426072475546\n",
      "dg max sq error:  0.00036262287883386543\n",
      "dg % correct:  80.0\n",
      "training batch 15 of size 1000\n",
      "training batch 16 of size 1000\n",
      "training batch 17 of size 1000\n",
      "training batch 18 of size 1000\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-31-a7cc7a3970a0>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     22\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mn\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnbatches\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     23\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"training batch %d of size %d\"\u001b[0m \u001b[0;34m%\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mn\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mconfig\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'train_batch_size'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 24\u001b[0;31m     \u001b[0magent\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     25\u001b[0m     \u001b[0mepisode\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0mconfig\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'train_batch_size'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     26\u001b[0m     \u001b[0mcheckpoint_episode\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0mconfig\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'train_batch_size'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/pyvenv/lib/python3.6/site-packages/ray/rllib/agents/trainer.py\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    492\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0m_\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mMAX_WORKER_FAILURE_RETRIES\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    493\u001b[0m             \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 494\u001b[0;31m                 \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mTrainable\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    495\u001b[0m             \u001b[0;32mexcept\u001b[0m \u001b[0mRayError\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    496\u001b[0m                 \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconfig\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"ignore_worker_failures\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/pyvenv/lib/python3.6/site-packages/ray/tune/trainable.py\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    181\u001b[0m         \"\"\"\n\u001b[1;32m    182\u001b[0m         \u001b[0mstart\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtime\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtime\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 183\u001b[0;31m         \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    184\u001b[0m         \u001b[0;32massert\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mresult\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdict\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"step() needs to return a dict.\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    185\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/pyvenv/lib/python3.6/site-packages/ray/rllib/agents/trainer_template.py\u001b[0m in \u001b[0;36mstep\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    145\u001b[0m         \u001b[0;34m@\u001b[0m\u001b[0moverride\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mTrainer\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    146\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 147\u001b[0;31m             \u001b[0mres\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnext\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain_exec_impl\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    148\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mres\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    149\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/pyvenv/lib/python3.6/site-packages/ray/util/iter.py\u001b[0m in \u001b[0;36m__next__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    754\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__next__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    755\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_build_once\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 756\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mnext\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbuilt_iterator\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    757\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    758\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__str__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/pyvenv/lib/python3.6/site-packages/ray/util/iter.py\u001b[0m in \u001b[0;36mapply_foreach\u001b[0;34m(it)\u001b[0m\n\u001b[1;32m    781\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    782\u001b[0m             \u001b[0;32mdef\u001b[0m \u001b[0mapply_foreach\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mit\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 783\u001b[0;31m                 \u001b[0;32mfor\u001b[0m \u001b[0mitem\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mit\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    784\u001b[0m                     \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mitem\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_NextValueNotReady\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    785\u001b[0m                         \u001b[0;32myield\u001b[0m \u001b[0mitem\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/pyvenv/lib/python3.6/site-packages/ray/util/iter.py\u001b[0m in \u001b[0;36mapply_foreach\u001b[0;34m(it)\u001b[0m\n\u001b[1;32m    781\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    782\u001b[0m             \u001b[0;32mdef\u001b[0m \u001b[0mapply_foreach\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mit\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 783\u001b[0;31m                 \u001b[0;32mfor\u001b[0m \u001b[0mitem\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mit\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    784\u001b[0m                     \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mitem\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_NextValueNotReady\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    785\u001b[0m                         \u001b[0;32myield\u001b[0m \u001b[0mitem\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/pyvenv/lib/python3.6/site-packages/ray/util/iter.py\u001b[0m in \u001b[0;36mapply_filter\u001b[0;34m(it)\u001b[0m\n\u001b[1;32m    841\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mfilter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mCallable\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mT\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbool\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;34m\"LocalIterator[T]\"\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    842\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0mapply_filter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mit\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 843\u001b[0;31m             \u001b[0;32mfor\u001b[0m \u001b[0mitem\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mit\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    844\u001b[0m                 \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_metrics_context\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    845\u001b[0m                     \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mitem\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_NextValueNotReady\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mitem\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/pyvenv/lib/python3.6/site-packages/ray/util/iter.py\u001b[0m in \u001b[0;36mapply_filter\u001b[0;34m(it)\u001b[0m\n\u001b[1;32m    841\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mfilter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mCallable\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mT\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbool\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;34m\"LocalIterator[T]\"\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    842\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0mapply_filter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mit\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 843\u001b[0;31m             \u001b[0;32mfor\u001b[0m \u001b[0mitem\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mit\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    844\u001b[0m                 \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_metrics_context\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    845\u001b[0m                     \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mitem\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_NextValueNotReady\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mitem\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/pyvenv/lib/python3.6/site-packages/ray/util/iter.py\u001b[0m in \u001b[0;36mapply_foreach\u001b[0;34m(it)\u001b[0m\n\u001b[1;32m    781\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    782\u001b[0m             \u001b[0;32mdef\u001b[0m \u001b[0mapply_foreach\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mit\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 783\u001b[0;31m                 \u001b[0;32mfor\u001b[0m \u001b[0mitem\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mit\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    784\u001b[0m                     \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mitem\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_NextValueNotReady\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    785\u001b[0m                         \u001b[0;32myield\u001b[0m \u001b[0mitem\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/pyvenv/lib/python3.6/site-packages/ray/util/iter.py\u001b[0m in \u001b[0;36mapply_foreach\u001b[0;34m(it)\u001b[0m\n\u001b[1;32m    781\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    782\u001b[0m             \u001b[0;32mdef\u001b[0m \u001b[0mapply_foreach\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mit\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 783\u001b[0;31m                 \u001b[0;32mfor\u001b[0m \u001b[0mitem\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mit\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    784\u001b[0m                     \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mitem\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_NextValueNotReady\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    785\u001b[0m                         \u001b[0;32myield\u001b[0m \u001b[0mitem\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/pyvenv/lib/python3.6/site-packages/ray/util/iter.py\u001b[0m in \u001b[0;36mapply_foreach\u001b[0;34m(it)\u001b[0m\n\u001b[1;32m    789\u001b[0m                         \u001b[0;32mwhile\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    790\u001b[0m                             \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_metrics_context\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 791\u001b[0;31m                                 \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mitem\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    792\u001b[0m                             \u001b[0;32myield\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    793\u001b[0m                             \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mresult\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_NextValueNotReady\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/pyvenv/lib/python3.6/site-packages/ray/rllib/execution/train_ops.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, samples)\u001b[0m\n\u001b[1;32m    206\u001b[0m                         batch_fetches = optimizer.optimize(\n\u001b[1;32m    207\u001b[0m                             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msess\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpermutation\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mbatch_index\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 208\u001b[0;31m                             self.per_device_batch_size)\n\u001b[0m\u001b[1;32m    209\u001b[0m                         \u001b[0;32mfor\u001b[0m \u001b[0mk\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mv\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mbatch_fetches\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mLEARNER_STATS_KEY\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitems\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    210\u001b[0m                             \u001b[0miter_extra_fetches\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mk\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mv\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/pyvenv/lib/python3.6/site-packages/ray/rllib/execution/multi_gpu_impl.py\u001b[0m in \u001b[0;36moptimize\u001b[0;34m(self, sess, batch_index)\u001b[0m\n\u001b[1;32m    255\u001b[0m             \u001b[0mfetches\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mupdate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtower\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mloss_graph\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_get_grad_and_stats_fetches\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    256\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 257\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0msess\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfetches\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfeed_dict\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mfeed_dict\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    258\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    259\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mget_common_loss\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/pyvenv/lib/python3.6/site-packages/tensorflow/python/client/session.py\u001b[0m in \u001b[0;36mrun\u001b[0;34m(self, fetches, feed_dict, options, run_metadata)\u001b[0m\n\u001b[1;32m    966\u001b[0m     \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    967\u001b[0m       result = self._run(None, fetches, feed_dict, options_ptr,\n\u001b[0;32m--> 968\u001b[0;31m                          run_metadata_ptr)\n\u001b[0m\u001b[1;32m    969\u001b[0m       \u001b[0;32mif\u001b[0m \u001b[0mrun_metadata\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    970\u001b[0m         \u001b[0mproto_data\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtf_session\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTF_GetBuffer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrun_metadata_ptr\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/pyvenv/lib/python3.6/site-packages/tensorflow/python/client/session.py\u001b[0m in \u001b[0;36m_run\u001b[0;34m(self, handle, fetches, feed_dict, options, run_metadata)\u001b[0m\n\u001b[1;32m   1189\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mfinal_fetches\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mfinal_targets\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mhandle\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mfeed_dict_tensor\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1190\u001b[0m       results = self._do_run(handle, final_targets, final_fetches,\n\u001b[0;32m-> 1191\u001b[0;31m                              feed_dict_tensor, options, run_metadata)\n\u001b[0m\u001b[1;32m   1192\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1193\u001b[0m       \u001b[0mresults\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/pyvenv/lib/python3.6/site-packages/tensorflow/python/client/session.py\u001b[0m in \u001b[0;36m_do_run\u001b[0;34m(self, handle, target_list, fetch_list, feed_dict, options, run_metadata)\u001b[0m\n\u001b[1;32m   1367\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mhandle\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1368\u001b[0m       return self._do_call(_run_fn, feeds, fetches, targets, options,\n\u001b[0;32m-> 1369\u001b[0;31m                            run_metadata)\n\u001b[0m\u001b[1;32m   1370\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1371\u001b[0m       \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_do_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0m_prun_fn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhandle\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfeeds\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfetches\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/pyvenv/lib/python3.6/site-packages/tensorflow/python/client/session.py\u001b[0m in \u001b[0;36m_do_call\u001b[0;34m(self, fn, *args)\u001b[0m\n\u001b[1;32m   1373\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0m_do_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1374\u001b[0m     \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1375\u001b[0;31m       \u001b[0;32mreturn\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1376\u001b[0m     \u001b[0;32mexcept\u001b[0m \u001b[0merrors\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mOpError\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1377\u001b[0m       \u001b[0mmessage\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcompat\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mas_text\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmessage\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/pyvenv/lib/python3.6/site-packages/tensorflow/python/client/session.py\u001b[0m in \u001b[0;36m_run_fn\u001b[0;34m(feed_dict, fetch_list, target_list, options, run_metadata)\u001b[0m\n\u001b[1;32m   1358\u001b[0m       \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_extend_graph\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1359\u001b[0m       return self._call_tf_sessionrun(options, feed_dict, fetch_list,\n\u001b[0;32m-> 1360\u001b[0;31m                                       target_list, run_metadata)\n\u001b[0m\u001b[1;32m   1361\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1362\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_prun_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mhandle\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfeed_dict\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfetch_list\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/pyvenv/lib/python3.6/site-packages/tensorflow/python/client/session.py\u001b[0m in \u001b[0;36m_call_tf_sessionrun\u001b[0;34m(self, options, feed_dict, fetch_list, target_list, run_metadata)\u001b[0m\n\u001b[1;32m   1451\u001b[0m     return tf_session.TF_SessionRun_wrapper(self._session, options, feed_dict,\n\u001b[1;32m   1452\u001b[0m                                             \u001b[0mfetch_list\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtarget_list\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1453\u001b[0;31m                                             run_metadata)\n\u001b[0m\u001b[1;32m   1454\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1455\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0m_call_tf_sessionprun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhandle\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfeed_dict\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfetch_list\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "agent = ppo.PPOTrainer(config, env=AMRGame)\n",
    "#agent.restore(\"/home/rwa/ray_results/PPO_AMRGame_2021-02-26_14-31-49nivcscnb/checkpoint_10/checkpoint-10\")\n",
    "policy = agent.get_policy()\n",
    "model = policy.model\n",
    "\n",
    "rms = []\n",
    "cor = []\n",
    "maxerr = []\n",
    "\n",
    "dg_rms = []\n",
    "dg_cor = []\n",
    "dg_maxerr = []\n",
    "\n",
    "checkpoint_period = 5000\n",
    "\n",
    "neval = 50\n",
    "eval_period = 5000\n",
    "\n",
    "episode = 0\n",
    "checkpoint_episode = 0\n",
    "eval_episode = 0\n",
    "for n in range(nbatches):\n",
    "    print(\"training batch %d of size %d\" % (n,config['train_batch_size']))\n",
    "    agent.train()\n",
    "    episode += config['train_batch_size']\n",
    "    checkpoint_episode += config['train_batch_size']\n",
    "    if (checkpoint_episode >= checkpoint_period):\n",
    "        checkpoint_episode = 0\n",
    "        checkpoint_path = agent.save()\n",
    "        print(checkpoint_path)\n",
    "        \n",
    "    eval_episode += config['train_batch_size']\n",
    "    if (eval_episode >= eval_period):\n",
    "        eval_episode = 0\n",
    "        rms1, maxerr1, cor1, dg_rms1, dg_maxerr1, dg_cor1 = eval_ensemble(model, neval)\n",
    "        rms.append(rms1)\n",
    "        maxerr.append(maxerr1)\n",
    "        cor.append(cor1)\n",
    "        dg_rms.append(dg_rms1)\n",
    "        dg_maxerr.append(dg_maxerr1)\n",
    "        dg_cor.append(dg_cor1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "isteps = list(range(nbatches))\n",
    "asteps = [i*config['train_batch_size'] for i in isteps]\n",
    "import matplotlib.pyplot as plt\n",
    "ax = plt.subplot(211)\n",
    "ax.set_ylim(0.00001,0.01)\n",
    "ax.set_ylabel('Error')\n",
    "line1, = plt.semilogy(asteps,rms[:nbatches], marker='o')\n",
    "line2, = plt.semilogy(asteps,dg_rms[:nbatches], marker='x')\n",
    "line3, = plt.semilogy(asteps,maxerr[:nbatches], marker='.')\n",
    "line4, = plt.semilogy(asteps,dg_maxerr[:nbatches], marker='+')\n",
    "\n",
    "line1.set_label('RL rms')\n",
    "line2.set_label('DG rms')\n",
    "line3.set_label('RL max')\n",
    "line4.set_label('DG max')\n",
    "ax.legend()\n",
    "plt.ticklabel_format(style='sci', axis='x', scilimits=(0,0))\n",
    "\n",
    "ax = plt.subplot(212)\n",
    "ax.set_ylim(0,100)\n",
    "ax.set_ylabel('% correct')\n",
    "ax.set_xlabel('training episodes')\n",
    "line1, = plt.plot(asteps,cor[:nbatches], marker='o')\n",
    "line2, = plt.plot(asteps,dg_cor[:nbatches], marker='x')\n",
    "line1.set_label('RL policy')\n",
    "line2.set_label('DG')\n",
    "ax.legend()\n",
    "plt.ticklabel_format(style='sci', axis='x', scilimits=(0,0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "rms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's look for cases where the policy gets it right and the DG method gets it wrong."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for n in range(500):\n",
    "    obs = env.reset()\n",
    "    opt_action, opt_reward = find_optimal(obs)\n",
    "    dg_action, dg_reward = find_dgjumps(env)\n",
    "    pol_action, pol_reward = apply_policy(model, obs)\n",
    "    if ((pol_action == opt_action) and (dg_action != opt_action)):\n",
    "        break\n",
    "env.reinit()\n",
    "env.step(pol_action)\n",
    "env.render()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.reinit()\n",
    "env.step(dg_action)\n",
    "env.render()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "minigame",
   "language": "python",
   "name": "minigame"
  },
  "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}