DPG for CartPole.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from collections import OrderedDict, namedtuple\n",
    "import time\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "import numpy as np\n",
    "import theano as th\n",
    "import theano.tensor as T\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Warning: IOStream has no fileno.\n"
     ]
    }
   ],
   "source": [
    "from tdlearn.examples import PendulumSwingUpCartPole\n",
    "from utils.cartpole import CartPole\n",
    "from utils import VariableStore, Linear, SGD, momentum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1446156040\n"
     ]
    }
   ],
   "source": [
    "seed = int(time.time())\n",
    "print seed\n",
    "rng = T.shared_randomstreams.RandomStreams(seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Set up continuous MDP instance.\n",
    "mdp = PendulumSwingUpCartPole()\n",
    "STATE_DIM = mdp.dim_S\n",
    "ACTION_DIM = mdp.dim_A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "EXPLORE_RANGE = 0.5\n",
    "\n",
    "Critic = namedtuple(\"Critic\", [\"pred\", \"targets\", \"cost\", \"updates\"])\n",
    "\n",
    "\n",
    "class DPGModel(object):\n",
    "    \n",
    "    def __init__(self, state_dim, action_dim, explore_range=0.5, track=True,\n",
    "                 _parent=None, name=\"dpg\"):\n",
    "        self.state_dim = state_dim\n",
    "        self.action_dim = action_dim\n",
    "        self.explore_range = explore_range\n",
    "        self.name = name\n",
    "        \n",
    "        self.parent = _parent\n",
    "        if _parent is None:\n",
    "            self._vs_actor = VariableStore(\"%s/vs_a\" % name)\n",
    "            self._vs_critic = VariableStore(\"%s/vs_c\" % name)\n",
    "            self._vs_prefix = self.name\n",
    "            self._make_vars()\n",
    "        else:\n",
    "            self._vs_actor = VariableStore.snapshot(_parent._vs_actor)\n",
    "            self._vs_critic = VariableStore.snapshot(_parent._vs_critic)\n",
    "            self._vs_prefix = _parent.name\n",
    "            self._pull_vars(_parent)\n",
    "        \n",
    "        self._make_graph()\n",
    "        self._make_updates()\n",
    "        self._make_functions()\n",
    "        \n",
    "        if track:\n",
    "            self.track = DPGModel(state_dim, action_dim, explore_range,\n",
    "                                  track=False, _parent=self,\n",
    "                                  name=\"%s_track\" % name)\n",
    "        \n",
    "    def _make_vars(self):\n",
    "        self.X = T.matrix(\"X\")\n",
    "\n",
    "        # Optionally directly provide actions predicted\n",
    "        self.actions = T.matrix(\"actions\")\n",
    "        # Q target values\n",
    "        self.q_targets = T.vector(\"q_targets\")\n",
    "        # Learning rates\n",
    "        self.lr_actor = T.scalar(\"lr_actor\")\n",
    "        self.lr_critic = T.scalar(\"lr_critic\")\n",
    "        \n",
    "    def _pull_vars(self, parent):\n",
    "        self.X = parent.X\n",
    "        self.actions = parent.actions\n",
    "        self.q_targets = parent.q_targets\n",
    "        self.lr_actor = parent.lr_actor\n",
    "        self.lr_critic = parent.lr_critic\n",
    "        \n",
    "        # Target network: tracking coefficient\n",
    "        self.tau = T.scalar(\"tau\")\n",
    "        \n",
    "    def _make_graph(self):\n",
    "        # Deterministic policy: linear map\n",
    "        self.a_pred = Linear(self.X, self.state_dim, self.action_dim,\n",
    "                             self._vs_actor, name=\"%s/a\" % self._vs_prefix)\n",
    "\n",
    "        # Exploration policy: add noise\n",
    "        self.a_explore = self.a_pred + rng.normal(self.a_pred.shape,\n",
    "                                                  0, self.explore_range, ndim=2)\n",
    "\n",
    "        # Create a few different Critic instances (Q-functions). These all\n",
    "        # share parameters; they only differ in the sources of their inputs.\n",
    "        #\n",
    "        # Critic 1: actions given\n",
    "        self.critic_given = self._make_critic(self.actions, self.q_targets)\n",
    "        # Critic 2: with deterministic policy\n",
    "        self.critic_det = self._make_critic(self.a_pred, self.q_targets)\n",
    "        # Critic 3: with noised / exploration policy\n",
    "        self.critic_exp = self._make_critic(self.a_explore, self.q_targets)\n",
    "        \n",
    "    def _make_critic(self, actions, targets):\n",
    "        # Q-function is a linear map on state+action pair.\n",
    "        hidden_dim = 200 # DEV\n",
    "        q_hid = Linear(T.concatenate([self.X, actions], axis=1),\n",
    "                       self.state_dim + self.action_dim, hidden_dim,\n",
    "                       self._vs_critic, \"%s/q/hid\" % self._vs_prefix)\n",
    "        q_pred = Linear(q_hid, hidden_dim, 1, self._vs_critic,\n",
    "                       \"%s/q/pred\" % self._vs_prefix)\n",
    "        q_pred = q_pred.reshape((-1,))\n",
    "        \n",
    "        # MSE loss on TD backup targets.\n",
    "        q_cost = ((targets - q_pred) ** 2).mean()\n",
    "        q_updates = SGD(q_cost, self._vs_critic.vars.values(),\n",
    "                             self.lr_critic)\n",
    "        \n",
    "        return Critic(q_pred, targets, q_cost, q_updates)\n",
    "    \n",
    "    def _make_updates(self):\n",
    "        # Actor-critic learning w/ critic 3\n",
    "        # NB, need to flatten all timesteps into a single batch\n",
    "        self.updates = OrderedDict(self.critic_exp.updates)\n",
    "        # Add policy gradient updates\n",
    "        self.updates.update(SGD(-self.critic_exp.pred.mean(),\n",
    "                                     self._vs_actor.vars.values(),\n",
    "                                     self.lr_actor))\n",
    "        \n",
    "        # Target network: update w.r.t. parent\n",
    "        if self.parent is not None:\n",
    "            self.target_updates = OrderedDict()\n",
    "            for vs, parent_vs in [(self._vs_actor, self.parent._vs_actor),\n",
    "                                  (self._vs_critic, self.parent._vs_critic)]:\n",
    "                for param_name, param_var in vs.vars.iteritems():\n",
    "                    self.target_updates[param_var] = (\n",
    "                        self.tau * vs.vars[param_name]\n",
    "                        + (1 - self.tau) * parent_vs.vars[param_name])\n",
    "        \n",
    "    def _make_functions(self):\n",
    "        # On-policy action prediction function\n",
    "        self.f_action_on = th.function([self.X], self.a_pred)\n",
    "        # Off-policy action prediction function\n",
    "        self.f_action_off = th.function([self.X], self.a_explore)\n",
    "\n",
    "        # Q-function\n",
    "        self.f_q = th.function([self.X, self.actions], self.critic_given.pred)\n",
    "\n",
    "        # Actor-critic update\n",
    "        self.f_update = th.function([self.X, self.q_targets,\n",
    "                                     self.lr_actor, self.lr_critic],\n",
    "                                    (self.critic_exp.cost, self.critic_exp.pred),\n",
    "                                    updates=self.updates)\n",
    "        \n",
    "        # Target networks only: update w.r.t. parent\n",
    "        if self.parent is not None:\n",
    "            self.f_track_update = th.function([self.tau],\n",
    "                                              updates=self.target_updates)\n",
    "        \n",
    "        \n",
    "dpg = DPGModel(STATE_DIM, ACTION_DIM, EXPLORE_RANGE, track=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def preprocess_state(state, clip_range=20):\n",
    "    state = np.clip(state, -clip_range, clip_range)\n",
    "    return state\n",
    "\n",
    "def preprocess_action(action, clip_range=20):\n",
    "    action = np.clip(action, -clip_range, clip_range)\n",
    "    return action\n",
    "\n",
    "def update_buffers(buffers, states, actions, rewards, states_next):\n",
    "    R_states, R_actions, R_rewards, R_states_next = buffers\n",
    "    \n",
    "    R_states = np.append(R_states, states, axis=0)\n",
    "    R_states_next = np.append(R_states_next, states_next, axis=0)\n",
    "    R_actions = np.append(R_actions, actions)\n",
    "    R_rewards = np.append(R_rewards, rewards)\n",
    "    \n",
    "    buffers = (R_states, R_actions, R_rewards, R_states_next)\n",
    "    return buffers\n",
    "\n",
    "def run_episode(mdp, buffers, max_len=100, on_policy=True):\n",
    "    policy = dpg.f_action_on if on_policy else dpg.f_action_off\n",
    "    def policy_fn(state):\n",
    "        state = preprocess_state(state).reshape((-1, mdp.dim_S))\n",
    "        return policy(state)\n",
    "    \n",
    "    states, actions, rewards, states_next = [], [], [], []\n",
    "    trajectory_gen = mdp.sample_transition(max_len, policy_fn)\n",
    "\n",
    "    for s, a, s_n, r in trajectory_gen:\n",
    "        states.append(preprocess_state(s))\n",
    "        actions.append(preprocess_action(a))\n",
    "        states_next.append(preprocess_state(s_n))\n",
    "        rewards.append(r)\n",
    "    \n",
    "    if buffers is not None:\n",
    "        buffers = update_buffers(buffers, states, actions, rewards, states_next)\n",
    "    return buffers, (states, actions, rewards, states_next)\n",
    "\n",
    "def train_batch(dpg, buffers, batch_size, gamma=0.9, tau=0.5, lr_actor=0.01,\n",
    "                lr_critic=0.001):\n",
    "    R_states, R_actions, R_rewards, R_states_next = buffers\n",
    "    \n",
    "    if len(R_states) - 1 < batch_size:\n",
    "        # Not enough data. Keep collecting trajectories.\n",
    "        return 0.0\n",
    "\n",
    "    # Sample a training minibatch.\n",
    "    idxs = np.random.choice(len(R_states) - 1, size=batch_size, replace=False)\n",
    "    b_states, b_actions, b_rewards, b_states_next = \\\n",
    "        R_states[idxs], R_actions[idxs], R_rewards[idxs], R_states_next[idxs]\n",
    "    \n",
    "    # Compute targets (TD error backups) given current Q function.\n",
    "    next_actions = dpg.track.f_action_on(b_states_next)\n",
    "    b_targets = b_rewards + gamma * dpg.track.f_q(b_states_next, next_actions).reshape((-1,))\n",
    "    \n",
    "    # SGD update.\n",
    "    cost_t, _ = dpg.f_update(b_states, b_targets, lr_actor, lr_critic)\n",
    "    # Update tracking model.\n",
    "    dpg.track.f_track_update(tau)\n",
    "    \n",
    "    return cost_t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\t 0.599997\t  0.328305\t\t12.595794\t0.259071\n",
      "1\t 0.599939\t  0.319896\t\t12.260872\t0.370451\n",
      "2\t 0.599743\t  0.290961\t\t12.172077\t0.368476\n",
      "3\t 0.599905\t  0.178679\t\t9.661598\t0.358782\n",
      "4\t 0.599901\t  0.226255\t\t10.432234\t0.363795\n",
      "5\t 0.599917\t  0.233444\t\t9.726625\t0.360371\n",
      "6\t 0.391674\t  0.224148\t\t10.877763\t0.323790\n",
      "7\t-0.464671\t  0.184228\t\t4.264116\t0.129724\n",
      "8\t 0.413142\t  0.239857\t\t10.812671\t0.341001\n",
      "9\t 0.552881\t  0.196289\t\t11.697633\t0.360778\n",
      "10\t 0.599951\t  0.231479\t\t12.753005\t0.254156\n",
      "11\t 0.599552\t  0.192708\t\t12.728482\t0.348176\n",
      "12\t-0.590935\t  0.200312\t\t1.069553\t0.033105\n",
      "13\t 0.599234\t  0.220308\t\t12.386711\t0.347487\n",
      "14\t 0.599659\t  0.228133\t\t12.297399\t0.360820\n",
      "15\t-0.592502\t  0.166926\t\t1.010888\t0.029385\n",
      "16\t-0.546987\t  0.188946\t\t2.516226\t0.076800\n",
      "17\t-0.068684\t  0.221445\t\t7.667111\t0.239110\n",
      "18\t 0.305567\t  0.217385\t\t10.869443\t0.319755\n",
      "19\t 0.599979\t  0.206876\t\t12.307619\t0.351475\n",
      "20\t-0.415514\t  0.217404\t\t4.706374\t0.145215\n",
      "21\t 0.599977\t  0.199785\t\t10.118516\t0.344086\n",
      "22\t-0.053827\t  0.168195\t\t7.860006\t0.249461\n",
      "23\t 0.599962\t  0.172367\t\t8.993717\t0.341591\n",
      "24\t 0.599993\t  0.197140\t\t12.053496\t0.360582\n",
      "25\t 0.594469\t  0.171876\t\t11.934060\t0.344840\n",
      "26\t 0.484486\t  0.201900\t\t11.809356\t0.333738\n",
      "27\t-0.579385\t  0.223968\t\t1.581147\t0.048990\n",
      "28\t 0.599986\t  0.236223\t\t12.771008\t0.250919\n",
      "29\t 0.599969\t  0.215365\t\t12.050170\t0.354735\n",
      "30\t-0.545031\t  0.189776\t\t2.701067\t0.080824\n",
      "31\t 0.257644\t  0.204540\t\t10.382238\t0.308167\n",
      "32\t 0.599968\t  0.226874\t\t12.263999\t0.354975\n",
      "33\t-0.378367\t  0.213382\t\t5.312224\t0.155248\n",
      "34\t-0.224109\t  0.234970\t\t6.878834\t0.213085\n",
      "35\t-0.522290\t  0.198808\t\t3.228816\t0.095991\n",
      "36\t-0.005388\t  0.212074\t\t8.888014\t0.259851\n",
      "37\t 0.599953\t  0.221142\t\t12.321065\t0.361324\n",
      "38\t 0.239483\t  0.211208\t\t9.995638\t0.297675\n",
      "39\t 0.178166\t  0.224408\t\t9.351666\t0.285294\n",
      "40\t 0.599986\t  0.208106\t\t11.777468\t0.358521\n",
      "41\t-0.395037\t  0.173567\t\t4.915930\t0.152498\n",
      "42\t 0.599958\t  0.193893\t\t12.487279\t0.341941\n",
      "43\t 0.434580\t  0.183571\t\t11.443061\t0.324526\n",
      "44\t 0.599736\t  0.202570\t\t12.657559\t0.345552\n",
      "45\t 0.599964\t  0.212744\t\t9.901651\t0.347712\n",
      "46\t 0.194931\t  0.190180\t\t10.358542\t0.290489\n",
      "47\t 0.599765\t  0.193065\t\t12.433393\t0.349090\n",
      "48\t 0.390245\t  0.207635\t\t11.206905\t0.325591\n",
      "49\t 0.530799\t  0.207705\t\t11.970534\t0.338878\n",
      "50\t 0.031556\t  0.180234\t\t9.242545\t0.264372\n",
      "51\t 0.430501\t  0.209084\t\t10.685352\t0.336153\n",
      "52\t 0.599774\t  0.207419\t\t12.400205\t0.352229\n",
      "53\t 0.599965\t  0.192851\t\t10.005239\t0.352397\n",
      "54\t-0.110194\t  0.182773\t\t7.890953\t0.221177\n",
      "55\t-0.030835\t  0.190278\t\t7.629322\t0.231553\n",
      "56\t-0.258319\t  0.220775\t\t6.641678\t0.193752\n",
      "57\t 0.599897\t  0.210115\t\t10.492384\t0.357925\n",
      "58\t 0.564275\t  0.223229\t\t11.948150\t0.347917\n",
      "59\t 0.599970\t  0.191429\t\t8.211848\t0.340445\n",
      "60\t-0.029194\t  0.219477\t\t7.858218\t0.243139\n",
      "61\t 0.447278\t  0.187959\t\t10.686070\t0.331949\n",
      "62\t-0.162717\t  0.203513\t\t7.606630\t0.214275\n",
      "63\t-0.572294\t  0.206786\t\t1.856771\t0.056153\n",
      "64\t 0.460498\t  0.221258\t\t11.526676\t0.326665\n",
      "65\t 0.599967\t  0.184114\t\t12.199491\t0.347589\n",
      "66\t-0.423766\t  0.205725\t\t4.792469\t0.139193\n",
      "67\t-0.112655\t  0.222777\t\t8.081460\t0.234598\n",
      "68\t 0.525189\t  0.216404\t\t11.120501\t0.342564\n",
      "69\t 0.575246\t  0.177410\t\t12.133763\t0.350101\n",
      "70\t 0.452486\t  0.236699\t\t11.731262\t0.324720\n",
      "71\t 0.075712\t  0.201050\t\t8.748885\t0.268031\n",
      "72\t-0.525895\t  0.197111\t\t3.015140\t0.089980\n",
      "73\t 0.599992\t  0.197953\t\t12.478413\t0.352846\n",
      "74\t 0.089294\t  0.150958\t\t8.697000\t0.266715\n",
      "75\t-0.285404\t  0.190868\t\t6.451623\t0.174080\n",
      "76\t-0.535668\t  0.190366\t\t2.950416\t0.088686\n",
      "77\t 0.270473\t  0.177360\t\t10.181094\t0.300357\n",
      "78\t 0.599991\t  0.200271\t\t12.693633\t0.255357\n",
      "79\t-0.450455\t  0.197548\t\t4.433230\t0.129914\n",
      "80\t 0.599990\t  0.177988\t\t12.277361\t0.352777\n",
      "81\t 0.444205\t  0.170776\t\t10.523335\t0.325912\n",
      "82\t-0.467756\t  0.207613\t\t4.266544\t0.124807\n",
      "83\t-0.513095\t  0.212054\t\t3.382622\t0.102075\n",
      "84\t-0.505847\t  0.186177\t\t3.092634\t0.099086\n",
      "85\t 0.578868\t  0.197959\t\t12.056735\t0.348287\n",
      "86\t 0.225896\t  0.217237\t\t10.566827\t0.296482\n",
      "87\t-0.588226\t  0.184301\t\t1.265606\t0.037684\n",
      "88\t-0.054949\t  0.194118\t\t8.230231\t0.248555\n",
      "89\t 0.599977\t  0.190260\t\t12.155600\t0.342326\n",
      "90\t-0.082439\t  0.194997\t\t7.647666\t0.223748\n",
      "91\t-0.593704\t  0.199114\t\t0.928379\t0.026288\n",
      "92\t 0.599915\t  0.165987\t\t12.350608\t0.353523\n",
      "93\t 0.599984\t  0.174077\t\t12.401706\t0.339681\n",
      "94\t-0.358423\t  0.190803\t\t5.589679\t0.156411\n",
      "95\t 0.553715\t  0.203964\t\t11.952614\t0.342914\n",
      "96\t-0.482185\t  0.181780\t\t3.898021\t0.113374\n",
      "97\t 0.120603\t  0.190653\t\t9.767459\t0.282142\n",
      "98\t-0.258965\t  0.202950\t\t6.601491\t0.185983\n",
      "99\t 0.549612\t  0.175675\t\t12.019994\t0.340424\n",
      "100\t 0.148701\t  0.187003\t\t9.474109\t0.273563\n",
      "101\t 0.575049\t  0.208930\t\t12.089935\t0.348013\n",
      "102\t 0.599849\t  0.234728\t\t12.485543\t0.337578\n",
      "103\t 0.591109\t  0.179766\t\t12.384061\t0.356153\n",
      "104\t 0.599996\t  0.182232\t\t12.687205\t0.341120\n",
      "105\t-0.599176\t  0.182609\t\t0.309091\t0.009056\n",
      "106\t 0.285557\t  0.177189\t\t10.400118\t0.304705\n",
      "107\t-0.551850\t  0.210957\t\t2.351034\t0.071485\n",
      "108\t-0.544747\t  0.211702\t\t2.413693\t0.074182\n",
      "109\t 0.599994\t  0.204748\t\t12.140569\t0.351464\n",
      "110\t 0.599919\t  0.204718\t\t8.642788\t0.340540\n",
      "111\t 0.573061\t  0.188073\t\t11.952599\t0.343768\n",
      "112\t-0.417469\t  0.195758\t\t4.726230\t0.141620\n",
      "113\t 0.555762\t  0.219444\t\t11.546141\t0.350250\n",
      "114\t 0.174768\t  0.220308\t\t9.754240\t0.277838\n",
      "115\t 0.283150\t  0.217663\t\t9.834035\t0.292033\n",
      "116\t 0.599994\t  0.186608\t\t12.552665\t0.345447\n",
      "117\t 0.471949\t  0.191394\t\t10.906878\t0.336683\n",
      "118\t 0.599403\t  0.209004\t\t12.136038\t0.339143\n",
      "119\t 0.039847\t  0.180438\t\t9.215306\t0.260387\n",
      "120\t 0.402735\t  0.201063\t\t11.497406\t0.323506\n",
      "121\t 0.599909\t  0.172044\t\t8.846065\t0.342967\n",
      "122\t 0.384186\t  0.188657\t\t10.759556\t0.326084\n",
      "123\t-0.084081\t  0.191201\t\t7.639628\t0.228152\n",
      "124\t 0.102808\t  0.170718\t\t9.304964\t0.272990\n",
      "125\t 0.033322\t  0.211344\t\t9.242628\t0.261994\n",
      "126\t-0.522875\t  0.200342\t\t3.280126\t0.095904\n",
      "127\t 0.599984\t  0.229332\t\t12.721751\t0.253976\n",
      "128\t 0.466764\t  0.176863\t\t11.857283\t0.326650\n",
      "129\t 0.542826\t  0.183902\t\t12.066418\t0.346708\n",
      "130\t 0.201185\t  0.171604\t\t9.233882\t0.290334\n",
      "131\t-0.265347\t  0.208368\t\t6.458826\t0.199882\n",
      "132\t 0.051730\t  0.195370\t\t8.707200\t0.265809\n",
      "133\t-0.336065\t  0.222151\t\t5.776697\t0.162849\n",
      "134\t 0.051880\t  0.216825\t\t8.700050\t0.266232\n",
      "135\t 0.599311\t  0.170481\t\t12.138924\t0.344997\n",
      "136\t 0.599992\t  0.190223\t\t12.323598\t0.356359\n",
      "137\t 0.483964\t  0.174137\t\t11.590123\t0.329185\n",
      "138\t-0.397411\t  0.198092\t\t5.021367\t0.153584\n",
      "139\t 0.366755\t  0.225634\t\t10.393059\t0.318555\n",
      "140\t-0.476956\t  0.220443\t\t3.757795\t0.118338\n",
      "141\t 0.378106\t  0.208717\t\t11.063595\t0.319298\n",
      "142\t 0.284848\t  0.198643\t\t10.266949\t0.307043\n",
      "143\t 0.599764\t  0.184633\t\t12.560180\t0.271199\n",
      "144\t 0.449939\t  0.209208\t\t11.412005\t0.332787"
     ]
    }
   ],
   "source": [
    "# Keep a replay buffer of states, actions, rewards, targets\n",
    "R_states = np.empty((0, STATE_DIM), dtype=th.config.floatX)\n",
    "R_actions = np.empty((0,), dtype=np.int32)\n",
    "R_rewards = np.empty((0,), dtype=np.int32)\n",
    "R_states_next = np.empty_like(R_states)\n",
    "buffers = (R_states, R_actions, R_rewards, R_states_next)\n",
    "\n",
    "avg_rewards = []\n",
    "max_rewards = []\n",
    "q_costs = []\n",
    "\n",
    "BATCH_SIZE = 64\n",
    "LR_ACTOR = 0.01\n",
    "LR_CRITIC = 0.001\n",
    "GAMMA = 0.99\n",
    "TAU = 0.001\n",
    "\n",
    "for t in xrange(1000):\n",
    "    # Sample a trajectory off-policy, then update the actor and critic.\n",
    "    buffers, _ = run_episode(mdp, buffers, on_policy=False)\n",
    "    cost_t = train_batch(dpg, buffers, BATCH_SIZE, GAMMA, TAU,\n",
    "                         LR_ACTOR, LR_CRITIC)\n",
    "    \n",
    "    # Evaluate actor by sampling a trajectory on-policy.\n",
    "    _, (states, actions, rewards, _) = run_episode(mdp, None, on_policy=True)\n",
    "    \n",
    "    rewards = np.array(rewards)\n",
    "    avg_reward, max_reward = rewards.mean(), rewards.max()\n",
    "    avg_rewards.append(avg_reward)\n",
    "    max_rewards.append(max_reward)\n",
    "    q_costs.append(cost_t)\n",
    "    print \"%i\\t% 4f\\t%10f\\t\\t%f\\t%f\" % (t, max_reward, cost_t, np.max(states), np.max(actions))\n",
    "    \n",
    "    if not np.isfinite(cost_t):\n",
    "        print \"STOP: Non-finite cost\"\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "plt.figure(1)\n",
    "plt.plot(avg_rewards, \"g\")\n",
    "plt.xlabel(\"Iteration\")\n",
    "plt.ylabel(\"Average reward achieved\")\n",
    "\n",
    "plt.figure(2)\n",
    "plt.plot(max_rewards, \"b\")\n",
    "plt.xlabel(\"Iteration\")\n",
    "plt.ylabel(\"Max reward achieved\")\n",
    "\n",
    "plt.figure(3)\n",
    "plt.plot(q_costs, \"r\")\n",
    "plt.xlabel(\"Iteration\")\n",
    "plt.ylabel(\"Q cost\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "states = [state for state, _, _, _\n",
    "          in mdp.sample_transition(100, lambda s: dpg.f_action_on(s.reshape((-1, STATE_DIM))))]\n",
    "states = np.array(states)\n",
    "mdp.animate_trace(states).save(\"test.mp4\")\n",
    "plt.clf()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}