bouncing_ball.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "%matplotlib inline\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def load_pickle(filename):\n",
    "    with open(filename, 'rb') as f:\n",
    "        arr = pickle.load(f)\n",
    "    return arr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "train = load_pickle('bouncing_balls_training_data.pkl')\n",
    "train.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "plt.imshow(train[0,0,:,:],cmap=matplotlib.cm.Greys_r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sigma0_np = np.array([0.2, 0.4, 0.6], dtype=np.float32)*7\n",
    "sigma0_np_mtx = np.diag(sigma0_np)\n",
    "mu0_np = np.array([2, 4, 6], dtype=np.float32)\n",
    "\n",
    "sigma1_np = np.array([0.1, 0.3, 0.5], dtype=np.float32)\n",
    "sigma1_np_mtx = np.diag(sigma1_np)\n",
    "mu1_np = np.array([1, 3, 5], dtype=np.float32)\n",
    "\n",
    "k = 1.0\n",
    "\n",
    "z1 = np.trace(np.matmul(np.linalg.inv(sigma1_np_mtx), sigma0_np_mtx))\n",
    "z2 = np.matmul(np.matmul(np.transpose(mu1_np-mu0_np), np.linalg.inv(sigma1_np_mtx)), (mu1_np-mu0_np))\n",
    "z3 = np.log(np.linalg.det(sigma1_np_mtx)/np.linalg.det(sigma0_np_mtx))\n",
    "Dkl = 0.5 * (z1 + z2 - k + z3)\n",
    "Dkl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def kl_divergence_gaussians(q_mu, q_sigma, p_mu, p_sigma):\n",
    "#https://hk.saowen.com/a/7404b78cd5b980e16e08423192e35ec18ecb3cb243d310a9d9a194747d9ee1ba    \n",
    "    r = q_mu - p_mu\n",
    "    return np.sum(np.log(p_sigma) - np.log(q_sigma) - .5 * (1. - (q_sigma**2 + r**2) / p_sigma**2), axis=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def kl_divergence_gaussians(q_mu, q_sigma, p_mu, p_sigma):\n",
    "    r = q_mu - p_mu\n",
    "    return np.sum(np.log(p_sigma) - np.log(q_sigma) - .5 * (1. - (q_sigma**2/p_sigma**2 + r**2/p_sigma**2)), axis=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def kl_divergence_gaussians_log(q_mu, q_sigma_log, p_mu, p_sigma_log):\n",
    "    r = q_mu - p_mu\n",
    "    return np.sum(p_sigma_log - q_sigma_log - .5 * (1. - (np.exp(q_sigma_log*2)/np.exp(p_sigma_log*2) + r**2/np.exp(p_sigma_log*2))), axis=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def kl_divergence_gaussians_tf(q_mu, q_sigma, p_mu, p_sigma):\n",
    "    with tf.Session() as sess:\n",
    "        zz = tf.reduce_sum(tf.distributions.kl_divergence(\n",
    "        tf.distributions.Normal(loc=q_mu, scale=q_sigma),\n",
    "        tf.distributions.Normal(loc=p_mu, scale=p_sigma)), axis=-1).eval()\n",
    "    return zz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def kl_divergence_gaussians_tf2(q_mu, q_sigma, p_mu, p_sigma):\n",
    "#https://hk.saowen.com/a/7404b78cd5b980e16e08423192e35ec18ecb3cb243d310a9d9a194747d9ee1ba    \n",
    "    with tf.Session() as sess:\n",
    "        r = q_mu - p_mu\n",
    "        return tf.reduce_sum(tf.log(p_sigma) - tf.log(q_sigma) - .5 * (1. - (q_sigma**2 + r**2) / p_sigma**2), axis=-1).eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def kl_divergence_gaussians_tf3(q_mu, q_sigma, p_mu, p_sigma):\n",
    "    with tf.Session() as sess:\n",
    "        num = tf.square(q_mu - p_mu) + tf.square(q_sigma) - tf.square(p_sigma)\n",
    "        den = 2 * tf.square(p_sigma) + 1e-8\n",
    "        kl = tf.reduce_sum(num/den + tf.log(p_sigma) - tf.log(q_sigma), reduction_indices=[2,3,4])\n",
    "        return kl.eval()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#print(kl_divergence_gaussians(mu0_np, sigma0_np, mu1_np, sigma1_np))\n",
    "print(kl_divergence_gaussians_tf(mu0_np, sigma0_np, mu1_np, sigma1_np))\n",
    "print(kl_divergence_gaussians_tf3(mu0_np, sigma0_np, mu1_np, sigma1_np))\n",
    "#print(kl_divergence_gaussians_tf2(mu0_np, sigma0_np, mu1_np, sigma1_np))\n",
    "#print(kl_divergence_gaussians_log(mu0_np, np.log(sigma0_np), mu1_np, np.log(sigma1_np)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def gauss_KL(mu1, logstd1, mu2, logstd2):\n",
    "    \"\"\" Returns KL divergence among two multivariate Gaussians, component-wise.\n",
    "\n",
    "    It assumes the covariance matrix is diagonal. All inputs have shape (n,a).\n",
    "    It is not necessary to know the number of actions because reduce_sum will\n",
    "    sum over this to get the `d` constant offset. The part consisting of the\n",
    "    trace in the formula is blended with the mean difference squared due to the\n",
    "    common \"denominator\" of var2_na.  This forumula generalizes for an arbitrary\n",
    "    number of actions.  I think mu2 and logstd2 should represent the policy\n",
    "    before the update.\n",
    "\n",
    "    Returns the KL divergence for each of the n components in the minibatch,\n",
    "    then we do a reduce_mean outside this.\n",
    "    \"\"\"\n",
    "    var1_na = tf.exp(2.*logstd1)\n",
    "    var2_na = tf.exp(2.*logstd2)\n",
    "    tmp_matrix = 2.*(logstd2 - logstd1) + (var1_na + tf.square(mu1-mu2))/var2_na - 1\n",
    "    #kl_n = tf.reduce_sum(0.5 * tmp_matrix, axis=[1]) # Don't forget the 1/2 !!\n",
    "    kl_n = 0.5 * tmp_matrix\n",
    "    assert_op = tf.Assert(tf.reduce_all(kl_n >= -0.0000001), [kl_n]) \n",
    "    with tf.control_dependencies([assert_op]):\n",
    "        kl_n = tf.identity(kl_n)\n",
    "    return kl_n "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def KL(mu_q, cov_q, mu_prior, cov_prior):\n",
    "        \"\"\"\n",
    "        https://github.com/clinicalml/dmm/blob/master/model_th/dmm.py\n",
    "        KL(q_t||p_t) = 0.5*(log|sigmasq_p| -log|sigmasq_q|  -D + Tr(sigmasq_p^-1 sigmasq_q)\n",
    "                        + (mu_p-mu_q)^T sigmasq_p^-1 (mu_p-mu_q))\n",
    "        \"\"\"\n",
    "        diff_mu = mu_prior-mu_q\n",
    "        KL      = tf.log(cov_prior)-tf.log(cov_q) - 1. + cov_q/cov_prior + diff_mu**2/cov_prior\n",
    "        KL_t    = 0.5*tf.reduce_sum(KL)\n",
    "        return KL_t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "\n",
    "mu0 = tf.placeholder(tf.float32, shape=(3,))\n",
    "sigma0 = tf.placeholder(tf.float32, shape=(3,))\n",
    "mu1 = tf.placeholder(tf.float32, shape=(3,))\n",
    "sigma1 = tf.placeholder(tf.float32, shape=(3,))\n",
    "\n",
    "\n",
    "# https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence\n",
    "k = 1.0\n",
    "tf1 = tf.reduce_sum(1.0/sigma1 * sigma0)\n",
    "tf2 = tf.reduce_sum((mu1-mu0)*1.0/sigma1*(mu1-mu0))\n",
    "tf3 = tf.cast(tf.log(tf.reduce_prod(sigma1/sigma0)), tf.float32)\n",
    "Dkl = 0.5 * (tf1 + tf2 - k + tf3)\n",
    "\n",
    "Dkl2 = tf.reduce_sum(gauss_KL(mu0, tf.diag(tf.log(sigma0)), mu1, tf.diag(tf.log(sigma1))))\n",
    "Dkl3 = KL(mu0, sigma0, mu1, sigma1)\n",
    "\n",
    "with tf.Session() as sess:\n",
    "    \n",
    "    sigma0_vec = np.array([0.2, 0.4, 0.6])\n",
    "    mu0_vec = np.array([2, 4, 6])\n",
    "    sigma1_vec = np.array([0.1, 0.3, 0.5])\n",
    "    mu1_vec = np.array([1, 3, 5])\n",
    "\n",
    "    dkl1, dkl2, dkl3 = sess.run([Dkl, Dkl2, Dkl3], feed_dict={mu0: mu0_vec, sigma0:sigma0_vec, mu1: mu1_vec, sigma1:sigma1_vec})\n",
    "    print(dkl1, dkl2, dkl3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "zz = tf.placeholder(tf.float32, shape=(15,))\n",
    "yy = tf.split(zz, num_or_size_splits=3, axis=0)\n",
    "with tf.Session() as sess:\n",
    "    qq=sess.run(yy, feed_dict={zz:np.arange(0,15)})\n",
    "print(qq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def initial_state_module(ex):\n",
    "    \"\"\"\n",
    "    \"\"\"\n",
    "    with tf.variable_scope('initial_state_module'):\n",
    "        shape = [-1,3]+ex.get_shape().as_list()[1:]\n",
    "        print(shape)\n",
    "        e = tf.reshape(ex, shape)\n",
    "        ez = tf.unstack(e, axis=1)\n",
    "        c = tf.concat(ez, axis=-1)\n",
    "    return c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "ww = np.ones((10,10,1))\n",
    "wwx = np.stack([ww*i for i in range(15)])\n",
    "wwx.shape\n",
    "ex = tf.placeholder(tf.float32, shape=(15,10,10,1))\n",
    "qq = initial_state_module(ex)\n",
    "with tf.Session() as sess:\n",
    "    qq_=sess.run(qq, feed_dict={ex:wwx})\n",
    "print(qq_.shape)\n",
    "print(qq_[0,0,0,0],qq_[0,0,0,1],qq_[0,0,0,2])\n",
    "print(qq_[1,0,0,0],qq_[1,0,0,1],qq_[1,0,0,2])\n",
    "print(qq_[2,0,0,0],qq_[2,0,0,1],qq_[2,0,0,2])\n",
    "print(qq_[3,0,0,0],qq_[3,0,0,1],qq_[3,0,0,2])\n",
    "print(qq_[4,0,0,0],qq_[4,0,0,1],qq_[4,0,0,2])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import train_env_model as mx\n",
    "import tensorflow as tf\n",
    "with tf.Session() as sess:\n",
    "    with tf.variable_scope('env_model'):\n",
    "        env_model = mx.EnvModel((13, 80, 80, 3), 1, 10)\n",
    "        #env_model = EnvModel((13, 80, 80, 3), 1, 10)\n",
    "\n",
    "    reg_loss = tf.reduce_sum(env_model.regularization_loss)\n",
    "    rec_loss = tf.reduce_sum(env_model.reconstruction_loss)\n",
    "    loss = reg_loss + rec_loss\n",
    "    env_model.train_op = tf.train.AdamOptimizer().minimize(loss)\n",
    "\n",
    "    sess.run(tf.global_variables_initializer())\n",
    "\n",
    "    train = load_pickle('bouncing_balls_testing_data.pkl')\n",
    "    train = np.expand_dims(train, 4)\n",
    "    train = np.repeat(train, 3, axis=4)\n",
    "    obs = train[0:2,:13,:,:,:]\n",
    "    actions = np.zeros((2,13), dtype=np.float32)\n",
    "    m = env_model\n",
    "\n",
    "    print('Training')\n",
    "    feed_dict = {env_model.obs:obs, env_model.actions:actions}\n",
    "    #_, obs_hat_, next_state_, rec_loss_, reg_loss_, loss_ = sess.run([env_model.train_op, env_model.obs_hat, env_model.next_state, rec_loss, reg_loss, loss], feed_dict=feed_dict)\n",
    "    obs_hat_, initial_state_, next_state_, reg_loss_, rec_loss_, eoi_, mu_, sigma_, mu_hat_, sigma_hat_, zz_ = sess.run(\n",
    "      [m.obs_hat, m.initial_state, m.next_state, m.regularization_loss, m.reconstruction_loss, m.encoded_obs_init, m.mu, m.sigma, m.mu_hat, m.sigma_hat, m.zz], feed_dict=feed_dict)\n",
    "\n",
    "    print('Reconstruction loss: ', rec_loss_)\n",
    "    print('Regularization loss: ', reg_loss_)\n",
    "    #print('Total loss: ', loss_)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "_ = plt.hist(zz_.flatten(),bins=np.arange(0,10,0.1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def kl_divergence_gaussians(q_mu, q_sigma, p_mu, p_sigma):\n",
    "    r = q_mu - p_mu\n",
    "    a = np.log(p_sigma) - np.log(q_sigma)\n",
    "    return  - .5 * (1. - (q_sigma**2 + r**2) / p_sigma**2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def kl_divergence_gaussians(q_mu, q_sigma, p_mu, p_sigma):\n",
    "    r = q_mu - p_mu\n",
    "    a = np.log(p_sigma) - np.log(q_sigma)\n",
    "    b = q_sigma**2/p_sigma**2\n",
    "    c = r**2/p_sigma**2\n",
    "    return a - .5 * (1. - (b + c)), a,b,c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def kl_divergence_gaussians_log(q_mu, q_sigma_log, p_mu, p_sigma_log):\n",
    "    r = q_mu - p_mu\n",
    "    return np.sum(p_sigma_log - q_sigma_log - .5 * (1. - (np.exp(q_sigma_log*2)/np.exp(p_sigma_log*2) + r**2/np.exp(p_sigma_log*2))), axis=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "kl, a,b,c = kl_divergence_gaussians(mu_hat_, sigma_hat_, mu_, sigma_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "kl_divergence_gaussians_log(mu_hat_, sigma_hat_, mu_, sigma_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "print(np.mean(np.isnan(a) + np.isinf(a)))\n",
    "print(np.mean(np.isnan(b) + np.isinf(b)))\n",
    "print(np.mean(np.isnan(c) + np.isinf(c)))\n",
    "print(np.mean(np.isnan(kl) + np.isinf(kl)))\n"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}