toy_example.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import pickle, os\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "sns.set_style(\"whitegrid\",{'axes.grid' : False})\n",
    "\n",
    "from BayesianNeuralNetwork import *\n",
    "from HiPMDP import HiPMDP\n",
    "from __future__ import print_function\n",
    "from ExperienceReplay import ExperienceReplay\n",
    "from multiprocessing import Pool\n",
    "\n",
    "if not os.path.isdir('./results'):\n",
    "    os.mkdir('results')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate batch of transition data\n",
    "* Toy domain\n",
    "* 2 different task instances\n",
    "* Agent learning model-free using e-greedy policy based on DQN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "domain = 'grid'\n",
    "run_type = 'modelfree'\n",
    "num_batch_instances = 2\n",
    "preset_hidden_params = [{'latent_code':1},{'latent_code':2}]\n",
    "ddqn_learning_rate = 0.0005\n",
    "episode_count = 500\n",
    "bnn_hidden_layer_size = 25\n",
    "bnn_num_hidden_layers = 3\n",
    "bnn_network_weights = None\n",
    "eps_min = 0.15\n",
    "test_inst = None\n",
    "create_exp_batch = True\n",
    "state_diffs = True\n",
    "grid_beta = 0.1\n",
    "batch_generator_hipmdp = HiPMDP(domain,preset_hidden_params,\n",
    "                                ddqn_learning_rate=ddqn_learning_rate,\n",
    "                                episode_count=episode_count,\n",
    "                                run_type=run_type, eps_min=eps_min,\n",
    "                                create_exp_batch=create_exp_batch,\n",
    "                                num_batch_instances=num_batch_instances,\n",
    "                                grid_beta=grid_beta,\n",
    "                                print_output=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### *The authors firmly suggest that those walking through this example for the first time test and evaluate the HiP-MDP with the \"grid\" domain as its runs quicker. The remainder of this notebook is therefore built to support investigation into this simple toy domain*\n",
    "\n",
    "However, we provide example code of how to intialize a batch for use in either the Acrobot or HIV domains in the cell below. \n",
    "\n",
    "This example is meant to replace the code in lines 1-4 in the cell above. Testing the HiP-MDP entails selecting any additional parameter setting beyond the batch used here to train the initial BNN and latent weights.\n",
    "\n",
    "Pay particular attention to other parameters used for training the BNN below. They are hard coded here for demonstration purposes. You may want to reference HiPMDP.py (included in this repo) for domain specific parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "## Example code of how to initialize a batch generator for HIV or Acrobot\n",
    "# domain = 'hiv' # 'acrobot'\n",
    "# with open('preset_parameters/'+domain+'_preset_hidden_params','r') as f:\n",
    "#     preset_parameters = pickle.load(f)\n",
    "    \n",
    "# run_type = 'modelfree'\n",
    "# num_batch_instances = 5\n",
    "# preset_hidden_params = preset_parameters[:num_batch_instances]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This next cell will take a fair amount of time to run, it is generating several examples from the two different task instances to train the BNN and Q-network on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "(exp_buffer, networkweights, rewards, avg_rwd_per_ep, full_task_weights,\n",
    "     sys_param_set, mean_episode_errors, std_episode_errors) = batch_generator_hipmdp.run_experiment()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "with open('results/{}_exp_buffer'.format(domain),'w') as f:\n",
    "    pickle.dump(exp_buffer,f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# with open('results/{}_exp_buffer'.format(domain),'r') as f:\n",
    "#      exp_buffer = pickle.load(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Reformat Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Create numpy array \n",
    "exp_buffer_np = np.vstack(exp_buffer)\n",
    "# Collect the instances that each transition came from\n",
    "inst_indices = exp_buffer_np[:,4]\n",
    "inst_indices = inst_indices.astype(int)\n",
    "# Group experiences by instance\n",
    "# Create dictionary where keys are instance indexes and values are np.arrays experiences\n",
    "exp_dict = {}\n",
    "for idx in xrange(batch_generator_hipmdp.instance_count):\n",
    "    exp_dict[idx] = exp_buffer_np[inst_indices == idx]\n",
    "X = np.array([np.hstack([exp_buffer_np[tt,0],exp_buffer_np[tt,1]]) for tt in range(exp_buffer_np.shape[0])])\n",
    "y = np.array([exp_buffer_np[tt,3] for tt in range(exp_buffer_np.shape[0])])\n",
    "num_dims = 2\n",
    "num_actions = 4\n",
    "num_wb = 5\n",
    "if state_diffs:\n",
    "    # subtract previous state\n",
    "    y -= X[:,:num_dims]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train BNN and learn latent weights using batch data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Set up BNN and latent weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "relu = lambda x: np.maximum(x, 0.)\n",
    "param_set = {\n",
    "    'bnn_layer_sizes': [num_dims+num_actions+num_wb]+[bnn_hidden_layer_size]*bnn_num_hidden_layers+[num_dims],\n",
    "    'weight_count': num_wb,\n",
    "    'num_state_dims': num_dims,\n",
    "    'bnn_num_samples': 50,\n",
    "    'bnn_batch_size': 32,\n",
    "    'num_strata_samples': 5,\n",
    "    'bnn_training_epochs': 1,\n",
    "    'bnn_v_prior': 1.0,\n",
    "    'bnn_learning_rate': 0.00005,\n",
    "    'bnn_alpha':0.5,\n",
    "    'wb_num_epochs':1,\n",
    "    'wb_learning_rate':0.0005\n",
    "}\n",
    "# Initialize latent weights for each instance\n",
    "full_task_weights = np.random.normal(0.,0.1,(batch_generator_hipmdp.instance_count,num_wb))\n",
    "# Initialize BNN\n",
    "network = BayesianNeuralNetwork(param_set, nonlinearity=relu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Compute error before training\n",
    "l2_errors = network.get_td_error(np.hstack((X,full_task_weights[inst_indices])), y, location=0.0, scale=1.0, by_dim=False)\n",
    "print (\"Before training: Mean Error: {}, Std Error: {}\".format(np.mean(l2_errors),np.std(l2_errors)))\n",
    "np.mean(l2_errors),np.std(l2_errors)\n",
    "print (\"L2 Difference in latent weights between instances: {}\".format(np.sum((full_task_weights[0]-full_task_weights[1])**2)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Oscillate between training BNN and latent weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_random_sample(start,stop,size):\n",
    "    indices_set = set()\n",
    "    while len(indices_set) < size:\n",
    "        indices_set.add(np.random.randint(start,stop))\n",
    "    return np.array(list(indices_set))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We slowly train the BNN and latent weights $w_b$ so as to avoid prematurely fitting to the separate modes of input. We found that without taking this approach, the BNN associated the latent parameters as noise and ignored their contribution to the performance of the predictive model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# size of sample to compute error on\n",
    "sample_size = 10000\n",
    "for i in xrange(40):    \n",
    "    # Update BNN network weights\n",
    "    network.fit_network(exp_buffer_np, full_task_weights, 0, state_diffs=state_diffs,\n",
    "                        use_all_exp=True)\n",
    "    print('finished BNN update '+str(i))\n",
    "    if i % 4 == 0:\n",
    "        #get random sample of indices\n",
    "        sample_indices = get_random_sample(0,X.shape[0],sample_size)\n",
    "        l2_errors = network.get_td_error(np.hstack((X[sample_indices],full_task_weights[inst_indices[sample_indices]])), y[sample_indices], location=0.0, scale=1.0, by_dim=False)\n",
    "        print (\"After BNN update: iter: {}, Mean Error: {}, Std Error: {}\".format(i,np.mean(l2_errors),np.std(l2_errors)))\n",
    "    # Update latent weights\n",
    "    for inst in np.random.permutation(batch_generator_hipmdp.instance_count):\n",
    "        full_task_weights[inst,:] = network.optimize_latent_weighting_stochastic(\n",
    "            exp_dict[inst],np.atleast_2d(full_task_weights[inst,:]),0,state_diffs=state_diffs,use_all_exp=True)\n",
    "    print ('finished wb update '+str(i))\n",
    "    # Compute error on sample of transitions\n",
    "    if i % 4 == 0:\n",
    "        #get random sample of indices\n",
    "        sample_indices = get_random_sample(0,X.shape[0],sample_size)\n",
    "        l2_errors = network.get_td_error(np.hstack((X[sample_indices],full_task_weights[inst_indices[sample_indices]])), y[sample_indices], location=0.0, scale=1.0, by_dim=False)\n",
    "        print (\"After Latent update: iter: {}, Mean Error: {}, Std Error: {}\".format(i,np.mean(l2_errors),np.std(l2_errors)))\n",
    "        # We check to see if the latent updates are sufficiently different so as to avoid fitting [erroneously] to the same dynamics\n",
    "        print (\"L2 Difference in latent weights between instances: {}\".format(np.sum((full_task_weights[0]-full_task_weights[1])**2)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "network_weights = network.weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "with open('results/{}_network_weights'.format(domain), 'w') as f:\n",
    "    pickle.dump(network.weights, f)\n",
    "# with open('results/{}_network_weights'.format(domain), 'r') as f:\n",
    "#     network_weights = pickle.load(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learn dynamics and policy for new instance using HiP-MDP with embedded latent weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "results = {}\n",
    "run_type = 'full'\n",
    "create_exp_batch = False\n",
    "episode_count = 20 # reduce episode count for demonstration since HiPMDP learns policy quickly\n",
    "for run in xrange(5):\n",
    "    for test_inst in [0,1]:\n",
    "        test_hipmdp = HiPMDP(domain,preset_hidden_params, \n",
    "                         ddqn_learning_rate=ddqn_learning_rate, \n",
    "                         episode_count=episode_count,\n",
    "                         run_type=run_type,\n",
    "                         bnn_hidden_layer_size=bnn_hidden_layer_size,\n",
    "                         bnn_num_hidden_layers=bnn_num_hidden_layers,\n",
    "                         bnn_network_weights=network_weights, test_inst=test_inst,\n",
    "                         eps_min=eps_min, create_exp_batch=create_exp_batch,grid_beta=grid_beta,print_output=True)\n",
    "        results[(test_inst,run)] = test_hipmdp.run_experiment()\n",
    "        with open('results/{}_results'.format(domain),'w') as f:\n",
    "            pickle.dump(results,f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# with open('results/{}_results'.format(domain),'r') as f:\n",
    "#     results = pickle.load(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Group rewards, errors by instance\n",
    "reward_key = 'Reward'\n",
    "error_key = 'BNN Error'\n",
    "clean_results = {(0,reward_key):[],(0,error_key):[],(1,reward_key):[],(1,error_key):[]}\n",
    "for test_inst in [0,1]:\n",
    "    for run in xrange(5):\n",
    "        clean_results[(test_inst,reward_key)].append(results[(test_inst,run)][2])\n",
    "        clean_results[(test_inst,error_key)].append(results[(test_inst,run)][6])\n",
    "    clean_results[(test_inst,reward_key)] = np.vstack(clean_results[(test_inst,reward_key)])\n",
    "    clean_results[(test_inst,error_key)] = np.vstack(clean_results[(test_inst,error_key)])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def plot_results(clean_results, test_inst):\n",
    "    f, ax_array = plt.subplots(2,figsize=(7,7))\n",
    "    result_names = [reward_key,error_key]\n",
    "    for result_idx in xrange(2):\n",
    "        result = result_names[result_idx]\n",
    "        mean_result = np.mean(clean_results[(test_inst,result)], axis=0)\n",
    "        std_result = np.std(clean_results[(test_inst,result)], axis=0)\n",
    "        ax_array[result_idx].errorbar(x=np.arange(len(mean_result)),y=mean_result,yerr=std_result)\n",
    "        _ = ax_array[result_idx].set_ylim((np.min(mean_result)-0.01,np.max(mean_result)+0.01))\n",
    "        ax_array[result_idx].set_ylabel(result)\n",
    "    ax_array[1].set_xlabel(\"Episode\")\n",
    "    f.suptitle(\"Full HiP-MDP Training Results Instance {}\".format(test_inst),fontsize=12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "plot_results(clean_results, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "plot_results(clean_results, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [default]",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}