model_double_gravity_ddcnn.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# From Pixels to Torques: Policy Learning using\n",
    "# Deep Dynamical Convolutional Neural Networks (DDCNN)\n",
    "\n",
    "### by John-Alexander M. Assael, Marc P. Deisenroth"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "if itorch then\n",
    "    arg = {}\n",
    "end\n",
    "\n",
    "cmd = torch.CmdLine()\n",
    "cmd:text()\n",
    "cmd:text('From Pixels to Torques:')\n",
    "cmd:text('Policy Learning using Deep Dynamical Convolutional Neural Networks (DDCNN)')\n",
    "cmd:text('by John-Alexander M. Assael, Marc P. Deisenroth')\n",
    "cmd:text()\n",
    "cmd:text('Options')\n",
    "\n",
    "-- general options:\n",
    "cmd:option('-seed', 1, 'initial random seed')\n",
    "cmd:option('-threads', 4, 'number of threads')\n",
    "\n",
    "-- gpu\n",
    "cmd:option('-cuda', false, 'cuda')\n",
    "\n",
    "-- model\n",
    "cmd:option('-lambda', 1, 'lambda')\n",
    "cmd:option('-action_size', 2, 'action size')\n",
    "\n",
    "-- training\n",
    "cmd:option('-batch_size', 20, 'batch size')\n",
    "cmd:option('-hist_len', 2, 'history length')\n",
    "cmd:option('-learningRate', 3e-4, 'learning rate')\n",
    "\n",
    "-- get current path\n",
    "require 'sys'\n",
    "dname, fname = sys.fpath()\n",
    "cmd:option('-save', dname, 'save path')\n",
    "cmd:option('-load', false, 'load pretrained model')\n",
    "\n",
    "cmd:option('-v', false, 'be verbose')\n",
    "cmd:text()\n",
    "\n",
    "opt = cmd:parse(arg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import Packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "require 'base'\n",
    "require 'hdf5'\n",
    "require 'image'\n",
    "require 'nngraph'\n",
    "require 'optim'\n",
    "require 'nn'\n",
    "require 'KLDistCriterion'\n",
    "require 'KLDCriterion'\n",
    "require 'LinearO'\n",
    "require 'AddCons'\n",
    "require 'Reparametrize'\n",
    "require 'unsup'\n",
    "Plot = require 'itorch.Plot'\n",
    "\n",
    "-- Cuda initialisation\n",
    "if opt.cuda then\n",
    "    require 'cutorch'\n",
    "    require 'cunn'\n",
    "    cutorch.setDevice(1)\n",
    "    print(cutorch.getDeviceProperties(1))\n",
    "end\n",
    "\n",
    "torch.manualSeed(opt.seed)\n",
    "torch.setnumthreads(opt.threads)\n",
    "-- Set float as default type\n",
    "torch.setdefaulttensortype('torch.FloatTensor') "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "# Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "function disp_img(img)\n",
    "    if itorch then\n",
    "        if opt.y_mean ~= nil then\n",
    "            img = g_destandarize(img:float(), opt.y_mean, opt.y_std)\n",
    "        end\n",
    "        itorch.image(image.scale(img:float():reshape(opt.img_w, opt.img_h), 256))\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "local myFile = hdf5.open('data/pendulum_double_gravity_48.h5', 'r')\n",
    "\n",
    "local y_all = myFile:read('train_y'):all():float()\n",
    "local u_all = myFile:read('train_u'):all():float():reshape(y_all:size(1), opt.action_size)\n",
    "\n",
    "myFile:close()\n",
    "\n",
    "-- Scale images\n",
    "-- local new_size = 10\n",
    "-- local prev_size = torch.sqrt(y_all:size(2))\n",
    "-- y_all = image.scale(y_all:reshape(y_all:size(1), prev_size,prev_size), new_size, new_size):reshape(y_all:size(1), new_size^2)\n",
    "\n",
    "-- Train Test\n",
    "local y = y_all[{{1,20001}}]\n",
    "local u = u_all[{{1,20001}}]\n",
    "\n",
    "local ys = y_all[{{20001,21001}}]\n",
    "local us = u_all[{{20001,21001}}]\n",
    "\n",
    "\n",
    "-- Update parameters\n",
    "opt.img_w = torch.sqrt(y:size(2))\n",
    "opt.img_h = torch.sqrt(y:size(2))\n",
    "opt.max_seq_length = y:size(1) - 1\n",
    "\n",
    "-- Store data\n",
    "state_train = {\n",
    "  x = transfer_data(y),\n",
    "  u = transfer_data(u)\n",
    "}\n",
    "\n",
    "state_test = {\n",
    "  x = transfer_data(ys),\n",
    "  u = transfer_data(us)\n",
    "}\n",
    "\n",
    "print('Train=' .. state_train.x:size(1) .. ' Test=' .. state_test.x:size(1) .. ' (' .. opt.img_w .. 'x' .. opt.img_h .. ')')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "idx=1\n",
    "disp_img(state_train.x[idx])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define Model Architecture"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "function create_network()\n",
    "    \n",
    "    opt.latent_dims = 4\n",
    "    local enc_dims = 400\n",
    "    local trans_dims = 200\n",
    "    \n",
    "    -- Model Specific parameters\n",
    "    local f_maps_1 = 32\n",
    "    local f_size_1 = 5\n",
    "    local f_maps_2 = 32\n",
    "    local f_size_2 = 5\n",
    "    local f_maps_3 = 32\n",
    "    local f_size_3 = 5\n",
    "    \n",
    "    -- Encoder\n",
    "    encoder = nn.Sequential()\n",
    "    encoder:add(nn.Reshape(1, opt.img_w, opt.img_h))\n",
    "    encoder:add(nn.SpatialConvolutionMM(1, f_maps_1, f_size_1, f_size_1))\n",
    "    encoder:add(nn.ReLU())\n",
    "    encoder:add(nn.SpatialMaxPooling(2,2,2,2))\n",
    "    \n",
    "    --layer 2\n",
    "    encoder:add(nn.SpatialConvolutionMM(f_maps_1, f_maps_2, f_size_2, f_size_2))\n",
    "    encoder:add(nn.ReLU())\n",
    "    encoder:add(nn.SpatialMaxPooling(2,2,2,2))\n",
    "    \n",
    "    --layer 3\n",
    "    encoder:add(nn.SpatialConvolutionMM(f_maps_2, f_maps_3, f_size_3, f_size_3))\n",
    "    encoder:add(nn.ReLU())\n",
    "    -- encoder:add(nn.SpatialMaxPooling(2,2,2,2))\n",
    "    \n",
    "    encoder:add(nn.Reshape(f_maps_3*5*5))\n",
    "    encoder:add(nn.LinearO(f_maps_3*5*5, enc_dims))\n",
    "    encoder:add(nn.ReLU())\n",
    "        \n",
    "    encoder:add(nn.LinearO(enc_dims, enc_dims))\n",
    "    encoder:add(nn.ReLU())\n",
    "    \n",
    "    encoder:add(nn.LinearO(enc_dims, opt.latent_dims))\n",
    "       \n",
    "    -- Decoder\n",
    "    decoder = nn.Sequential()\n",
    "    decoder:add(nn.LinearO(opt.latent_dims, enc_dims))\n",
    "    decoder:add(nn.ReLU())\n",
    "\n",
    "    decoder:add(nn.LinearO(enc_dims, enc_dims))\n",
    "    decoder:add(nn.ReLU())    \n",
    "    \n",
    "    decoder:add(nn.LinearO(enc_dims, f_maps_3*6*6))\n",
    "    decoder:add(nn.ReLU())\n",
    "    \n",
    "    decoder:add(nn.Reshape(f_maps_3, 6, 6))\n",
    "    \n",
    "    -- layer 3\n",
    "    decoder:add(nn.SpatialUpSamplingNearest(2))\n",
    "    decoder:add(nn.SpatialConvolutionMM(f_maps_3, f_maps_3, f_size_3-1, f_size_3-1))\n",
    "    decoder:add(nn.ReLU())\n",
    "        \n",
    "    -- layer 2\n",
    "    decoder:add(nn.SpatialUpSamplingNearest(2))\n",
    "    decoder:add(nn.SpatialConvolutionMM(f_maps_3, f_maps_2, f_size_2-1, f_size_2-1))\n",
    "    decoder:add(nn.ReLU())\n",
    "    \n",
    "    -- layer 1\n",
    "    decoder:add(nn.SpatialUpSamplingNearest(2))\n",
    "    decoder:add(nn.SpatialConvolutionMM(f_maps_2, f_maps_1, f_size_2, f_size_2))\n",
    "    decoder:add(nn.ReLU())\n",
    "    \n",
    "    decoder:add(nn.SpatialUpSamplingNearest(2))\n",
    "    decoder:add(nn.SpatialConvolutionMM(f_maps_1, 1, f_size_2, f_size_2))\n",
    "    \n",
    "    decoder:add(nn.Sigmoid())\n",
    "    --decoder:add(nn.View(opt.img_w^2))\n",
    "    \n",
    "    \n",
    "    -- Clone enc-dec\n",
    "    local encoder2 = encoder:clone(\"weight\", \"bias\", \"gradWeight\", \"gradBias\")\n",
    "    local decoder2 = decoder:clone(\"weight\", \"bias\", \"gradWeight\", \"gradBias\")\n",
    "    \n",
    "    -- Define model\n",
    "    local x_t_prev = nn.Identity()():annotate{name = 'x_t_prev'}\n",
    "    local x_t = nn.Identity()():annotate{name = 'x_t'}\n",
    "    local u_t = nn.Identity()():annotate{name = 'u_t'}\n",
    "    \n",
    "    -- Define Encoder Module\n",
    "    local z_t_prev = encoder2(x_t_prev):annotate{name = 'z_t_prev'}\n",
    "    local z_t = encoder(x_t):annotate{name = 'z_t'}\n",
    "    \n",
    "        \n",
    "    -- transition\n",
    "    trans = nn.Sequential()\n",
    "    trans:add(nn.LinearO(opt.action_size+opt.latent_dims*2, trans_dims))\n",
    "    trans:add(nn.ReLU())\n",
    "    trans:add(nn.LinearO(trans_dims, trans_dims))\n",
    "    trans:add(nn.ReLU())\n",
    "    trans:add(nn.LinearO(trans_dims, opt.latent_dims))\n",
    "    \n",
    "    \n",
    "    local dynamics_all = trans(nn.JoinTable(2)({z_t_prev, z_t, nn.Reshape(opt.action_size)(u_t)})):annotate{name = 'dynamics'}\n",
    "\n",
    "    -- Define Output\n",
    "    local decoder_x_t_next = decoder(dynamics_all):annotate{name = 'decoder_x_t_next'}\n",
    "    local decoder_x_t_cur = decoder2(z_t):annotate{name = 'decoder_x_t_cur'}\n",
    "    \n",
    "    -- Create model\n",
    "    \n",
    "    model = nn.gModule({x_t_prev, x_t, u_t}, {z_t_prev, z_t, dynamics_all, decoder_x_t_cur, decoder_x_t_next})\n",
    "    \n",
    "    -- create_links(model)\n",
    "    \n",
    "    return model\n",
    "end\n",
    "\n",
    "\n",
    "function create_links(model)\n",
    "    encoder = model.forwardnodes[5].data.module\n",
    "    trans = model.forwardnodes[13].data.module\n",
    "    decoder = model.forwardnodes[14].data.module\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup Network function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "function setup()\n",
    "    print(\"Creating Neural Net.\")\n",
    "    model = create_network()\n",
    "    params, gradParams = model:getParameters()\n",
    "    \n",
    "    criterion = nn.BCECriterion()\n",
    "    criterion.sizeAverage = false\n",
    "    \n",
    "    criterion_mse = nn.MSECriterion()\n",
    "    criterion_mse.sizeAverage = false\n",
    "\n",
    "end\n",
    "\n",
    "function setup_load()\n",
    "    \n",
    "    print(\"Loading Neural Net.\")\n",
    "    \n",
    "    load_model()\n",
    "    \n",
    "    create_links(model)\n",
    "    \n",
    "    params, gradParams = model:getParameters()\n",
    "    \n",
    "    opt.load = true\n",
    "    \n",
    "    dname, fname = sys.fpath()\n",
    "    opt.save = dname\n",
    "    \n",
    "    criterion = nn.BCECriterion()\n",
    "    criterion.sizeAverage = false\n",
    "    \n",
    "    criterion_mse = nn.MSECriterion()\n",
    "    criterion_mse.sizeAverage = false\n",
    "\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "function save_model()\n",
    "    -- save/log current net\n",
    "    local filename = paths.concat(opt.save, 'model/relu_double_gravity_ddcnn.t7')\n",
    "    os.execute('mkdir -p ' .. paths.dirname(filename))\n",
    "    if paths.filep(filename) then\n",
    "        os.execute('mv ' .. filename .. ' ' .. filename .. '.old')\n",
    "    end\n",
    "    -- print('<trainer> saving network to '..filename)\n",
    "    torch.save(filename, {model, opt, optim_config, train_err, test_err})\n",
    "end\n",
    "\n",
    "function load_model()\n",
    "    model, opt, optim_config, train_err, test_err = unpack(torch.load('model/relu_double_gravity_ddcnn.t7'))\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Initialize Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(\"Network parameters:\")\n",
    "print(opt)\n",
    "\n",
    "if opt.load then\n",
    "    setup_load()\n",
    "else\n",
    "    setup()\n",
    "    optim_config = { learningRate = opt.learningRate,\n",
    "                     beta2 = 0.9\n",
    "                    }\n",
    "    train_err = {}\n",
    "    test_err = {}\n",
    "end\n",
    "\n",
    "epoch = #train_err"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# Train Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "function g_create_batch(dataset)\n",
    "    local batches = {}\n",
    "    \n",
    "    -- shuffle at each epoch\n",
    "    local shuffle = torch.randperm(dataset.x:size(1)):long()\n",
    "    \n",
    "    for t = 1,dataset.x:size(1),opt.batch_size do\n",
    "\n",
    "        -- Count size\n",
    "        local mini_batch_size = 0\n",
    "        for i = t,math.min(t+opt.batch_size-1,dataset.x:size(1)) do\n",
    "            local idx = shuffle[i]\n",
    "            if idx - 1 >= 1 and idx+1 <= dataset.x:size(1) then\n",
    "                mini_batch_size = mini_batch_size + 1\n",
    "            end\n",
    "        end\n",
    "\n",
    "        -- create mini batch\n",
    "        local batch_x_prev = torch.Tensor(mini_batch_size, opt.img_w^2)\n",
    "        local batch_x_cur = torch.Tensor(mini_batch_size, opt.img_w^2)\n",
    "        local batch_u = torch.Tensor(mini_batch_size, opt.action_size)\n",
    "        local batch_y = torch.Tensor(mini_batch_size, opt.img_w^2)\n",
    "\n",
    "        local cur_idx = 1\n",
    "\n",
    "        for i = t,math.min(t+opt.batch_size-1,dataset.x:size(1)) do\n",
    "            \n",
    "            local idx = shuffle[i]\n",
    "\n",
    "            -- Filter batches\n",
    "            if idx - 1 >= 1 and idx+1 <= dataset.x:size(1) then\n",
    "                \n",
    "                -- load new sample\n",
    "                batch_x_prev[cur_idx] = dataset.x[idx-1]\n",
    "                batch_x_cur[cur_idx] = dataset.x[idx]\n",
    "                batch_y[cur_idx] = dataset.x[idx+1]\n",
    "                batch_u[cur_idx] = dataset.u[idx]\n",
    "\n",
    "                cur_idx = cur_idx + 1\n",
    "            end\n",
    "        end\n",
    "\n",
    "        table.insert(batches, {batch_x_prev, batch_x_cur, batch_u, batch_y})\n",
    "    end\n",
    "    \n",
    "    dataset.batch = batches\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "function train(dataset)\n",
    "\n",
    "    g_create_batch(state_train)\n",
    "\n",
    "    -- epoch tracker\n",
    "    epoch = epoch or 0\n",
    "\n",
    "    -- local vars\n",
    "    local err = {all=0, bce=0, bce_1=0, mse=0}\n",
    "\n",
    "    -- shuffle at each epoch\n",
    "    local shuffle = torch.randperm(#dataset.batch):long()\n",
    "\n",
    "    for t = 1,#dataset.batch do\n",
    "        \n",
    "        -- xlua.progress(t, #dataset.batch)\n",
    "\n",
    "        -- create mini batch\n",
    "        local batch_x_prev = dataset.batch[shuffle[t]][1]\n",
    "        local batch_x_cur = dataset.batch[shuffle[t]][2]\n",
    "        local batch_u = dataset.batch[shuffle[t]][3]\n",
    "        local batch_y = dataset.batch[shuffle[t]][4]\n",
    "\n",
    "        local batch_size = batch_y:size(1)\n",
    "\n",
    "        -- create closure to evaluate f(X) and df/dX\n",
    "        local feval = function(x)\n",
    "            \n",
    "            -- get new parameters\n",
    "            if x ~= params then\n",
    "                params:copy(x)\n",
    "            end\n",
    "\n",
    "            -- reset gradients\n",
    "            gradParams:zero()\n",
    "            \n",
    "            -- reset errors\n",
    "            local mse_err, bce_err, bce_1_err = 0, 0, 0\n",
    "            \n",
    "            local z_t_next_true = encoder:forward(batch_y)\n",
    "\n",
    "            -- evaluate function for complete mini batch                                                \n",
    "            local z_t_prev, z_t_cur, z_t_next, x_t, x_t_next = unpack(model:forward({batch_x_prev, batch_x_cur, batch_u}))  \n",
    "\n",
    "            -- BCE x_t\n",
    "            bce_err = bce_err + criterion:forward(x_t, batch_x_cur)\n",
    "            local d_x_t = criterion:backward(x_t, batch_x_cur):clone()\n",
    "            \n",
    "            -- BCE x_t+1\n",
    "            bce_1_err = bce_1_err + criterion:forward(x_t_next, batch_y)\n",
    "            local d_x_t1 = criterion:backward(x_t_next, batch_y):clone()  \n",
    "            \n",
    "            -- MSE z_t+1\n",
    "            mse_err = mse_err + criterion_mse:forward(z_t_next, z_t_next_true) * opt.lambda\n",
    "            local d_z_t_next = criterion_mse:backward(z_t_next, z_t_next_true):clone():mul(opt.lambda)\n",
    "            \n",
    "            -- Backpropagate\n",
    "            model:backward({batch_x_prev, batch_x_cur, batch_u}, {\n",
    "                    torch.zeros(batch_size, opt.latent_dims),\n",
    "                    torch.zeros(batch_size, opt.latent_dims),\n",
    "                    torch.zeros(batch_size, opt.latent_dims),\n",
    "                    d_x_t,\n",
    "                    d_x_t1\n",
    "                })\n",
    "            \n",
    "            local trans_in = torch.cat(torch.cat(z_t_prev, z_t_cur), batch_u)\n",
    "            trans:forward(trans_in)\n",
    "            trans:backward(trans_in, d_z_t_next)\n",
    "            \n",
    "            -- Accumulate errors\n",
    "            err.mse = err.mse + mse_err\n",
    "            err.bce = err.bce + bce_err\n",
    "            err.bce_1 = err.bce_1 + bce_1_err\n",
    "            err.all = err.all + bce_err + bce_1_err + mse_err\n",
    "                        \n",
    "            -- normalize gradients and f(X)\n",
    "            local batcherr = (bce_err + bce_1_err + mse_err) / batch_size\n",
    "            gradParams:div(batch_size)\n",
    "                \n",
    "            -- print(bce_err/batch_size, bce_1_err/batch_size, mse_err/batch_size)\n",
    "                \n",
    "            -- return f and df/dX\n",
    "            return batcherr, gradParams\n",
    "        end\n",
    "        \n",
    "        if batch_size > 0 then\n",
    "            optim.adam(feval, params, optim_config)\n",
    "            -- optim.adagrad(feval, params, optim_config)\n",
    "            -- optim.rmsprop(feval, params, optim_config)\n",
    "        end\n",
    "        \n",
    "    end\n",
    "    \n",
    "    -- Normalise errors\n",
    "    err.all = err.all / (dataset.x:size(1) - 2)\n",
    "    err.mse = err.mse / (dataset.x:size(1) - 2)\n",
    "    err.bce = err.bce / (dataset.x:size(1) - 2)\n",
    "    err.bce_1 = err.bce_1 / (dataset.x:size(1) - 2)\n",
    "    \n",
    "    epoch = epoch + 1\n",
    "\n",
    "    return err\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "-- epochs to run\n",
    "opt.max_epoch = 300\n",
    "\n",
    "-- start time\n",
    "local beginning_time = torch.tic()\n",
    "\n",
    "-- iterate through epochs\n",
    "for e = 1, opt.max_epoch do\n",
    "    \n",
    "    -- local vars\n",
    "    local time = sys.clock()\n",
    "    \n",
    "    -- train for 1 epoch\n",
    "    local err = train(state_train)\n",
    "        \n",
    "    train_err[#train_err+1] = err\n",
    "    \n",
    "    -- time taken\n",
    "    time = sys.clock() - time\n",
    "    \n",
    "    -- display stats\n",
    "    if (epoch) % 1 == 0 then\n",
    "        \n",
    "        local since_beginning = g_d(torch.toc(beginning_time) / 60)\n",
    "        print('epoch=' .. (epoch) ..\n",
    "          ', Train err=' .. g_f3(train_err[#train_err].all) ..\n",
    "          ', mse=' .. g_f3(train_err[#train_err].mse) ..\n",
    "          ', bce=' .. g_f3(train_err[#train_err].bce) ..\n",
    "          ', bce_1=' .. g_f3(train_err[#train_err].bce_1) ..\n",
    "          ', t/epoch = ' .. g_f3(time) .. ' sec' ..\n",
    "          ', t = ' .. since_beginning .. ' mins.')\n",
    "\n",
    "\n",
    "        if (epoch) % 5 == 0 then\n",
    "            save_model()\n",
    "        end\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot Performance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "function get_error(err, criterion)\n",
    "    local criterion = criterion or 'all'\n",
    "    local arr = torch.zeros(#err-1)\n",
    "    for i=2,#err do arr[i-1] = err[i][criterion] end    \n",
    "    return arr\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "colors = {'blue', 'green', 'red', 'purple', 'orange', 'magenta', 'cyan'}\n",
    "plot = Plot()\n",
    "plot:title(string.format('Double pendulum with gravity - %d epochs', #train_err))\n",
    "plot = plot:line(torch.range(2,#train_err), get_error(train_err,'all'), colors[1], 'L(D)')\n",
    "plot = plot:line(torch.range(2,#train_err), get_error(train_err,'mse'), colors[2], '|| z-z_goal ||^2')\n",
    "plot = plot:line(torch.range(2,#train_err), get_error(train_err,'bce'), colors[3], 'log p(x_t|z_t)')\n",
    "plot = plot:line(torch.range(2,#train_err), get_error(train_err,'bce_1'), colors[4], 'log p(x_t+1|z_t+1)')\n",
    "plot:legend(true):redraw()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate Predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "local idx = 600\n",
    "local dataset = state_test\n",
    "\n",
    "-- create mini batch\n",
    "local dataset = state_train\n",
    "local batch_x_prev = dataset.x:narrow(1, idx-1, 1)\n",
    "local batch_x_cur = dataset.x:narrow(1, idx, 1)\n",
    "local batch_x_next = dataset.x:narrow(1, idx+1, 1)\n",
    "local batch_u = dataset.u:narrow(1, idx, 1)\n",
    "\n",
    "local z_t_prev, z_t_cur, z_t_next, x_t, x_t_next = unpack(model:forward({batch_x_prev, batch_x_cur, batch_u}))  \n",
    "\n",
    "disp_img(x_t)\n",
    "disp_img(dataset.x[idx])\n",
    "disp_img(x_t_next)\n",
    "disp_img(dataset.x[idx+1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "local dataset = state_test\n",
    "local steps = 5\n",
    "local idx = 610\n",
    "\n",
    "-- create mini batch\n",
    "local dataset = state_train\n",
    "local x_prev = dataset.x:narrow(1, idx-1, 1)\n",
    "local x_cur = dataset.x:narrow(1, idx, 1)\n",
    "\n",
    "image.savePNG(string.format(\"preds/true-%05d.png\", 0), dataset.x[idx]:view(opt.img_w,opt.img_w))\n",
    "\n",
    "for i=0,steps-1 do\n",
    "    local batch_u = dataset.u:narrow(1, idx+i, 1)\n",
    "\n",
    "    local z_t_prev, z_t_cur, z_t_next, x_t, x_t_next = unpack(model:forward({x_prev, x_cur, batch_u}))  \n",
    "    \n",
    "    x_prev = x_cur:clone()\n",
    "    x_cur = x_t_next:clone():view(1, opt.img_w^2)\n",
    "    \n",
    "    image.savePNG(string.format(\"preds/true-%05d.png\", (i+1)), dataset.x[idx+1+i]:view(opt.img_w,opt.img_w))\n",
    "    image.savePNG(string.format(\"preds/pred-%05d.png\", (i+1)), x_t_next:view(opt.img_w,opt.img_w))\n",
    "    \n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# License"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright (C) 2015 John-Alexander M. Assael, Marc P. Deisenroth\n",
    "\n",
    "The MIT License (MIT)\n",
    "\n",
    "Permission is hereby granted, free of charge, to any person obtaining a copy of\n",
    "this software and associated documentation files (the \"Software\"), to deal in\n",
    "the Software without restriction, including without limitation the rights to\n",
    "use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies\n",
    "of the Software, and to permit persons to whom the Software is furnished to do\n",
    "so, subject to the following conditions:\n",
    "\n",
    "The above copyright notice and this permission notice shall be included in all\n",
    "copies or substantial portions of the Software.\n",
    "\n",
    "THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n",
    "IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n",
    "FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n",
    "AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n",
    "LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n",
    "OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\n",
    "SOFTWARE."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "iTorch",
   "language": "lua",
   "name": "itorch"
  },
  "language_info": {
   "name": "lua",
   "version": "20100"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}