simple_training_example.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from env.balancebot_env import BalancebotEnv\n",
    "from stable_baselines import PPO2\n",
    "from stable_baselines.common.policies import ActorCriticPolicy\n",
    "from stable_baselines.common.vec_env import DummyVecEnv\n",
    "from stable_baselines.common.vec_env import SubprocVecEnv\n",
    "from stable_baselines.bench import Monitor\n",
    "from stable_baselines.a2c.utils import linear\n",
    "\n",
    "import os\n",
    "import time\n",
    "import torch\n",
    "import tensorflow as tf\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log_dir = \"/tmp/gym/{}\".format(int(time.time()))\n",
    "os.makedirs(log_dir, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the environment\n",
    "def make_env(rank):\n",
    "    def _init():\n",
    "        env = BalancebotEnv(render=False)\n",
    "        env = Monitor(env, os.path.join(log_dir, str(rank)))\n",
    "        return env\n",
    "    return _init\n",
    "\n",
    "num_cpu = 16\n",
    "env = SubprocVecEnv([make_env(rank=i) for i in range(num_cpu)])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def const_linear(input_tensor, scope, np_w, np_b):\n",
    "    in_shape = input_tensor.get_shape().as_list()\n",
    "    n_input = in_shape[1]\n",
    "    assert n_input == np_w.shape[0], \\\n",
    "        f\"incompatible weight shape: {in_shape} and {np_w.shape}\"\n",
    "    n_output = np_w.shape[1]\n",
    "    assert n_output == np_b.shape[0], \\\n",
    "        f'incompatible bias shape: {np_w.shape} and {np_b.shape}'\n",
    "    with tf.variable_scope(scope):\n",
    "        tf_w = tf.constant(np_w, name='w', dtype=input_tensor.dtype)\n",
    "        tf_b = tf.constant(np_b, name='b', dtype=input_tensor.dtype)\n",
    "        return tf.matmul(input_tensor, tf_w) + tf_b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the RL Agwnt\n",
    "class CustomPolicy(ActorCriticPolicy):\n",
    "    np_pi_w = None\n",
    "    np_pi_b = None\n",
    "    np_vf_w = None\n",
    "    np_vf_b = None\n",
    "    var_layers = [16]\n",
    "    \n",
    "    def __init__(self, sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse=False):\n",
    "        super(CustomPolicy, self).__init__(sess, ob_space, \n",
    "                                           ac_space, n_env, n_steps, \n",
    "                                           n_batch, n_lstm=256,\n",
    "                                           reuse=reuse, scale=False)\n",
    "        if any([self.np_pi_w is None,\n",
    "                self.np_pi_b is None,\n",
    "                self.np_vf_w is None,\n",
    "                self.np_vf_b is None]):\n",
    "            raise ValueError('Run CustomPolicy.set_weights first before using this CustomPolicy')\n",
    "        with tf.variable_scope('model', reuse=reuse):\n",
    "            activ = tf.tanh\n",
    "            processed_x = tf.layers.flatten(self.processed_x)\n",
    "            pi_h = activ(const_linear(processed_x, 'pi_fc0', self.np_pi_w, self.np_pi_b))\n",
    "            vf_h = activ(const_linear(processed_x, 'vf_fc0', self.np_vf_w, self.np_vf_b))\n",
    "            for i, layer_size in enumerate(self.var_layers, 1):\n",
    "                pi_h = activ(linear(pi_h, 'pi_fc' + str(i), n_hidden=layer_size, init_scale=np.sqrt(2)))\n",
    "                vf_h = activ(linear(vf_h, 'vf_fc' + str(i), n_hidden=layer_size, init_scale=np.sqrt(2)))\n",
    "            value_fn = linear(vf_h, 'vf', 1)\n",
    "            pi_latent = pi_h\n",
    "            vf_latent = vf_h\n",
    "\n",
    "            self.proba_distribution, self.policy, self.q_value = \\\n",
    "                self.pdtype.proba_distribution_from_latent(pi_latent, vf_latent, init_scale=0.01)\n",
    "\n",
    "        self.value_fn = value_fn\n",
    "        self.initial_state = None\n",
    "        self._setup_init()\n",
    "    \n",
    "    def step(self, obs, state=None, mask=None, deterministic=False):\n",
    "        if deterministic:\n",
    "            action, value, neglogp = self.sess.run([self.deterministic_action, self._value, self.neglogp],\n",
    "                                                   {self.obs_ph: obs})\n",
    "        else:\n",
    "            action, value, neglogp = self.sess.run([self.action, self._value, self.neglogp],\n",
    "                                                   {self.obs_ph: obs})\n",
    "        return action, value, self.initial_state, neglogp\n",
    "\n",
    "    def proba_step(self, obs, state=None, mask=None):\n",
    "        return self.sess.run(self.policy_proba, {self.obs_ph: obs})\n",
    "\n",
    "    def value(self, obs, state=None, mask=None):\n",
    "        return self.sess.run(self._value, {self.obs_ph: obs})\n",
    "    \n",
    "    @classmethod\n",
    "    def custom_set_weights(cls, np_pi_w, np_pi_b, np_vf_w, np_vf_b, var_layers=None):\n",
    "        cls.np_pi_w = np_pi_w\n",
    "        cls.np_pi_b = np_pi_b\n",
    "        cls.np_vf_w = np_vf_w\n",
    "        cls.np_vf_b = np_vf_b\n",
    "        if var_layers:\n",
    "            cls.var_layers = var_layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# put your pytorch numpy weights here\n",
    "CustomPolicy.custom_set_weights(np.random.rand(4, 32), np.random.rand(32),\n",
    "                                np.random.rand(4, 32), np.random.rand(32))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = PPO2(CustomPolicy, env, verbose=1, tensorboard_log=log_dir+\"/tensorboard\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "observation_dim = env.observation_space.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weight_encoder = torch.rand(observation_dim, 32,requires_grad=False).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weight_encoder.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How do I put the weight of encoder into model ?\n",
    "## I want to put the weight into 'pi_fc0' and 'vf_fc0' and set them to untrainable\n",
    " "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Network](assets/network.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train and Save the agent\n",
    "model.learn(total_timesteps=1e3, tb_log_name=\"PPO2\")\n",
    "# model.save(\"ppo_save\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log_dir+\"/tensorboard/\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## You can open tensorboard at terminal\n",
    "## For example:\n",
    "### tensorboard --logdir log_dir+\"/tensorboard\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# delete trained model to demonstrate loading\n",
    "del model "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the evaluation env\n",
    "env = DummyVecEnv([lambda: BalancebotEnv(render=False)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the trained agent\n",
    "model = PPO2.load(\"ppo_save\", env=env, policy=CustomPolicy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Enjoy trained agent\n",
    "for ep in range(10):\n",
    "    obs = env.reset()\n",
    "    dones = False\n",
    "    while not dones:\n",
    "        action, _states = model.predict(obs)\n",
    "        obs, rewards, dones, info = env.step(action)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.5"
  },
  "toc": {
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": "block",
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}