examples/reinforcement_learning_pendulum.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reinforcement Learning for the Inverted Pendulum\n",
    "\n",
    "Perform approximate policy iteration in an actor-critic framework for the inverted pendulum."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import division, print_function\n",
    "\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import gpflow\n",
    "import safe_learning\n",
    "import matplotlib.pyplot as plt\n",
    "import time\n",
    "import os\n",
    "\n",
    "from matplotlib.colors import ListedColormap\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "from scipy.linalg import block_diag\n",
    "from utilities import InvertedPendulum, compute_closedloop_response, get_parameter_change, binary_cmap, compute_roa, reward_rollout\n",
    "\n",
    "# Nice progress bars\n",
    "try:\n",
    "    from tqdm import tqdm\n",
    "except ImportError:\n",
    "    tqdm = lambda x: x\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## User Options"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Options(object):\n",
    "    def __init__(self, **kwargs):\n",
    "        super(Options, self).__init__()\n",
    "        self.__dict__.update(kwargs)\n",
    "\n",
    "OPTIONS = Options(np_dtype              = safe_learning.config.np_dtype,\n",
    "                  tf_dtype              = safe_learning.config.dtype,\n",
    "                  saturate              = True,                            # apply saturation constraints to the control input\n",
    "                  eps                   = 1e-8,                            # numerical tolerance\n",
    "                  use_linear_dynamics   = False,                           # use the linearized form of the dynamics as the true dynamics (for testing)\n",
    "                  dpi                   = 200,\n",
    "                  num_cores             = 4,\n",
    "                  num_sockets           = 1)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TensorFlow Session\n",
    "\n",
    "Customize the TensorFlow session for the current device."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"KMP_BLOCKTIME\"]    = str(0)\n",
    "os.environ[\"KMP_SETTINGS\"]     = str(1)\n",
    "os.environ[\"KMP_AFFINITY\"]     = 'granularity=fine,noverbose,compact,1,0'\n",
    "os.environ[\"OMP_NUM_THREADS\"]  = str(OPTIONS.num_cores)\n",
    "\n",
    "config = tf.ConfigProto(intra_op_parallelism_threads  = OPTIONS.num_cores,\n",
    "                        inter_op_parallelism_threads  = OPTIONS.num_sockets,\n",
    "                        allow_soft_placement          = False,\n",
    "                        device_count                  = {'CPU': OPTIONS.num_cores})\n",
    "\n",
    "try:\n",
    "    session.close()\n",
    "except NameError:\n",
    "    pass\n",
    "session = tf.InteractiveSession(config=config)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dynamics\n",
    "\n",
    "Define the nonlinear and linearized forms of the inverted pendulum dynamics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Constants\n",
    "dt = 0.01   # sampling time\n",
    "g = 9.81    # gravity\n",
    "\n",
    "# System parameters\n",
    "m = 0.15    # pendulum mass\n",
    "L = 0.5     # pole length\n",
    "b = 0.1     # rotational friction\n",
    "\n",
    "# State and action normalizers\n",
    "theta_max = np.deg2rad(30)                 # angular position [rad]\n",
    "omega_max = np.sqrt(g / L)                 # angular velocity [rad/s]\n",
    "u_max     = g * m * L * np.sin(theta_max)  # torque [N.m], control action\n",
    "\n",
    "state_norm = (theta_max, omega_max)\n",
    "action_norm = (u_max, )\n",
    "\n",
    "# Dimensions and domains\n",
    "state_dim     = 2\n",
    "action_dim    = 1\n",
    "state_limits  = np.array([[-1., 1.]] * state_dim)\n",
    "action_limits = np.array([[-1., 1.]] * action_dim)\n",
    "\n",
    "# Initialize system class and its linearization\n",
    "pendulum = InvertedPendulum(m, L, b, dt, [state_norm, action_norm])\n",
    "A, B = pendulum.linearize()\n",
    "\n",
    "if OPTIONS.use_linear_dynamics:\n",
    "    dynamics = safe_learning.functions.LinearSystem((A, B), name='dynamics')\n",
    "else:\n",
    "    dynamics = pendulum.__call__\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reward Function\n",
    "\n",
    "Define a positive-definite reward function over the state-action space $\\mathcal{X} \\times \\mathcal{U}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Q = 0.1 * np.identity(state_dim).astype(OPTIONS.np_dtype)     # state cost matrix\n",
    "R = 0.1 * np.identity(action_dim).astype(OPTIONS.np_dtype)    # action cost matrix\n",
    "\n",
    "# Quadratic reward (- cost) function\n",
    "reward_function = safe_learning.QuadraticFunction(block_diag(- Q, - R), name='reward_function')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parametric Policy and Value Function\n",
    "\n",
    "Define a parametric value function $V_{\\bf \\theta} : \\mathcal{X} \\to \\mathbb{R}$ and policy $\\pi_{\\bf \\delta} : \\mathcal{X} \\to \\mathcal{U}$ as neural networks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Policy\n",
    "layer_dims = [64, 64, action_dim]\n",
    "activations = [tf.nn.relu, tf.nn.relu, None]\n",
    "if OPTIONS.saturate:\n",
    "    activations[-1] = tf.nn.tanh\n",
    "# Remove bias terms to ensure the policy maps the zero state to the zero input\n",
    "policy = safe_learning.functions.NeuralNetwork(layer_dims, activations, name='policy', use_bias=False)\n",
    "\n",
    "# Value function\n",
    "layer_dims = [64, 64, 1]\n",
    "activations = [tf.nn.relu, tf.nn.relu, None]\n",
    "value_function = safe_learning.functions.NeuralNetwork(layer_dims, activations, name='value_function', use_bias=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Approximate Policy Evaluation\n",
    "\n",
    "Fix the policy, and learn the corresponding value function."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fixed Policy\n",
    "\n",
    "Fix the policy $\\pi$ to the LQR solution for the linearized system, possibly with saturation constraints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "K, P = safe_learning.utilities.dlqr(A, B, Q, R)\n",
    "policy_lqr = safe_learning.functions.LinearSystem((-K, ), name='policy_lqr')\n",
    "if OPTIONS.saturate:\n",
    "    policy_lqr = safe_learning.Saturation(policy_lqr, -1, 1)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TensorFlow Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define new graph\n",
    "states = tf.placeholder(OPTIONS.tf_dtype, shape=[None, state_dim], name='states')\n",
    "actions = policy_lqr(states)\n",
    "rewards = reward_function(states, actions)\n",
    "future_states = dynamics(states, actions)\n",
    "\n",
    "# Use the parametric value function\n",
    "values = value_function(states)\n",
    "future_values = value_function(future_states)\n",
    "\n",
    "with tf.name_scope('approximate_policy_evaluation'):\n",
    "    # Discount factor and scaling\n",
    "    gamma = tf.placeholder(OPTIONS.tf_dtype, shape=[], name='discount_factor')\n",
    "    max_state = np.ones((1, state_dim))\n",
    "    max_action = np.ones((1, action_dim))\n",
    "    r_max = np.linalg.multi_dot((max_state, Q, max_state.T)) + np.linalg.multi_dot((max_action, R, max_action.T))\n",
    "#     scaling = (1 - gamma) / r_max.ravel()\n",
    "    scaling = 1 / r_max.ravel()\n",
    "    \n",
    "    # Objective function\n",
    "    target = tf.stop_gradient(rewards + gamma * future_values, name='target')\n",
    "    objective = scaling * tf.reduce_mean(tf.abs(values - target), name='objective')\n",
    "    \n",
    "    # Optimizer settings\n",
    "    learning_rate = tf.placeholder(OPTIONS.tf_dtype, shape=[], name='learning_rate')\n",
    "    optimizer = tf.train.GradientDescentOptimizer(learning_rate)\n",
    "    training_update = optimizer.minimize(objective, var_list=value_function.parameters)\n",
    "    \n",
    "with tf.name_scope('state_sampler'):\n",
    "    batch_size = tf.placeholder(tf.int32, shape=[], name='batch_size')\n",
    "    batch = tf.random_uniform([batch_size, state_dim], -1, 1, dtype=OPTIONS.tf_dtype, name='batch')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Value Function Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.run(tf.variables_initializer(value_function.parameters))\n",
    "\n",
    "# Uniformly sampled test set\n",
    "test_size = 1e3\n",
    "test_set = batch.eval({batch_size: test_size})\n",
    "\n",
    "# Keep track of the test set loss and parameter changes during training\n",
    "test_loss = []\n",
    "param_changes = []\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training\n",
    "\n",
    "Changing the discount factor affects the results immensely, due to the effect of diverging trajectories on the value function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training hyperparameters\n",
    "feed_dict = {\n",
    "    states:         np.zeros((1, state_dim)), # placeholder\n",
    "    gamma:          0.95,\n",
    "    learning_rate:  0.005,\n",
    "    batch_size:     100,   \n",
    "}\n",
    "max_iters = 500\n",
    "\n",
    "old_params = session.run(value_function.parameters)\n",
    "for _ in tqdm(range(max_iters)):\n",
    "    feed_dict[states] = batch.eval(feed_dict)\n",
    "    session.run(training_update, feed_dict)\n",
    "    \n",
    "    new_params = session.run(value_function.parameters)\n",
    "    param_changes.append(get_parameter_change(old_params, new_params, 'inf'))\n",
    "    old_params = list(new_params)\n",
    "    \n",
    "    feed_dict[states] = test_set\n",
    "    test_loss.append(objective.eval(feed_dict))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(10, 2), dpi=OPTIONS.dpi)\n",
    "fig.subplots_adjust(wspace=0.3)\n",
    "\n",
    "axes[0].plot(test_loss, '.-r')\n",
    "axes[0].set_xlabel(r'SGD iteration $k$')\n",
    "axes[0].set_ylabel(r'test loss')\n",
    "\n",
    "axes[1].plot(param_changes, '.-r')\n",
    "axes[1].set_xlabel(r'SGD iteration $k$')\n",
    "axes[1].set_ylabel(r'$||{\\bf \\theta}_k - {\\bf \\theta}_{k-1}||_\\infty$')\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### True Dynamics: Trajectories, ROA, and Value Function\n",
    "\n",
    "Compute the true largest region of attraction (ROA) by forward-simulating the closed-loop dynamics on a state space grid. Also, compute the true value function $V_\\pi$ on this grid by summing discounted rewards along these forward-simulated trajectories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Number of states along each dimension\n",
    "num_states = 101\n",
    "\n",
    "# State grid\n",
    "grid_limits = np.array([[-1., 1.], ] * state_dim)\n",
    "grid = safe_learning.GridWorld(grid_limits, num_states)\n",
    "\n",
    "# ROA and closed-loop dynamics\n",
    "closed_loop_dynamics = lambda x: future_states.eval({states: x})\n",
    "horizon = 500\n",
    "tol = 1e-2\n",
    "roa, trajectories = compute_roa(grid, closed_loop_dynamics, horizon, tol, no_traj=False)\n",
    "\n",
    "# Estimate true value function with a reward rollout\n",
    "reward_eval = lambda x: rewards.eval({states: x})\n",
    "discount = feed_dict[gamma]\n",
    "horizon = 1000\n",
    "tol = 1e-2\n",
    "true_values = reward_rollout(grid, closed_loop_dynamics, reward_eval, discount, horizon, tol)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting\n",
    "\n",
    "Visually compare the true value function $V_\\pi$ and the neural network $V_{\\bf \\theta}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(6, 5), dpi=OPTIONS.dpi)\n",
    "ax = fig.add_subplot(111, projection='3d')\n",
    "ax.set_xlabel(r'$\\theta$ [deg]')\n",
    "ax.set_ylabel(r'$\\omega$ [deg/s]')\n",
    "\n",
    "scaled_discrete_points = [norm * points for norm, points in zip(np.rad2deg(state_norm), grid.discrete_points)]\n",
    "xx, yy = np.meshgrid(*scaled_discrete_points)\n",
    "\n",
    "# ROA\n",
    "z = roa.reshape(grid.num_points)\n",
    "ax.contourf(xx, yy, z, cmap=binary_cmap('green', 0.7), zdir='z', offset=0)\n",
    "\n",
    "# Cost (- value) functions\n",
    "z = - true_values.reshape(grid.num_points)\n",
    "surf_true = ax.plot_surface(xx, yy, z, color='b', alpha=0.65)\n",
    "\n",
    "z = - values.eval({states: grid.all_points}).reshape(grid.num_points)\n",
    "surf_approx = ax.plot_surface(xx, yy, z, color='r', alpha=0.65)\n",
    "\n",
    "for surf in (surf_true, surf_approx):\n",
    "    surf._facecolors2d = surf._facecolors3d\n",
    "    surf._edgecolors2d = surf._edgecolors3d\n",
    "\n",
    "# Legend\n",
    "proxy = [plt.Rectangle((0,0), 1, 1, fc=c) for c in [(0., 0., 1., 0.65), (1., 0., 0., 0.65), (0., 1., 0., 0.65)]]    \n",
    "ax.legend(proxy, [r'$-V_{\\pi}({\\bf x})$', r'$-V_{\\bf \\theta}({\\bf x})$', 'ROA'])\n",
    "\n",
    "ax.view_init(None, -45)\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Approximate Policy Improvement\n",
    "\n",
    "Fix the value function, and learn the corresponding policy. Changing the discount factor affects the results immensely, due to the effect of diverging trajectories on the value function. However, a high discount factor encourages the policy to yield a large ROA for the closed-loop system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "discount = 0.965\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fixed Value Function\n",
    "\n",
    "Use supervised learning to closely approximate a fixed value function with a neural network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "horizon = 1000\n",
    "tol = 1e-2\n",
    "true_values = reward_rollout(grid, closed_loop_dynamics, reward_eval, discount, horizon, tol)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with tf.name_scope('supervised_value_function_learning'):\n",
    "    tf_true_values = tf.placeholder(OPTIONS.tf_dtype, shape=[None, 1], name='true_values')\n",
    "    loss = tf.abs(values - tf_true_values)\n",
    "    objective = tf.reduce_mean(loss, name='objective')\n",
    "    \n",
    "    learning_rate = tf.placeholder(OPTIONS.tf_dtype, shape=[], name='learning_rate')\n",
    "    optimizer = tf.train.GradientDescentOptimizer(learning_rate)\n",
    "    training_update = optimizer.minimize(objective, var_list=value_function.parameters)\n",
    "\n",
    "test_set = grid.all_points\n",
    "test_values = true_values.reshape((-1, 1))\n",
    "test_loss = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Feed dict with hyperparameters\n",
    "feed_dict = {\n",
    "    states:         np.zeros((1, state_dim)), # placeholder\n",
    "    tf_true_values: np.zeros((1, 1)),         # placeholder\n",
    "    learning_rate:  0.005\n",
    "}\n",
    "max_iters  = 1000\n",
    "batch_size = 1000\n",
    "\n",
    "for i in tqdm(range(max_iters)):\n",
    "    batch_idx = np.random.randint(grid.nindex, size=batch_size)\n",
    "    feed_dict[states] = grid.all_points[batch_idx]\n",
    "    feed_dict[tf_true_values] = true_values[batch_idx].reshape((-1, 1))\n",
    "    session.run(training_update, feed_dict)\n",
    "    \n",
    "    feed_dict[states] = test_set\n",
    "    feed_dict[tf_true_values] = test_values\n",
    "    test_loss.append(objective.eval(feed_dict))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TensorFlow Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use the parametric policy\n",
    "states = tf.placeholder(OPTIONS.tf_dtype, shape=[None, state_dim], name='states')\n",
    "actions = policy(states)\n",
    "rewards = reward_function(states, actions)\n",
    "future_states = dynamics(states, actions)\n",
    "\n",
    "# Future values according to parametric policy and true value function\n",
    "future_values = value_function(future_states)\n",
    "\n",
    "with tf.name_scope('approximate_policy_evaluation'):\n",
    "    # Discount factor and scaling\n",
    "    gamma = tf.placeholder(OPTIONS.tf_dtype, shape=[], name='discount_factor')\n",
    "    max_state = np.ones((1, state_dim))\n",
    "    max_action = np.ones((1, action_dim))\n",
    "    r_max = np.linalg.multi_dot((max_state, Q, max_state.T)) + np.linalg.multi_dot((max_action, R, max_action.T))\n",
    "    scaling = (1 - gamma) / r_max.ravel()\n",
    "    \n",
    "    # Objective function    \n",
    "    objective = - scaling * tf.reduce_mean(rewards + gamma * future_values, name='objective')\n",
    "    \n",
    "    # Optimizer settings\n",
    "    learning_rate = tf.placeholder(OPTIONS.tf_dtype, shape=[], name='learning_rate')\n",
    "    optimizer = tf.train.GradientDescentOptimizer(learning_rate)\n",
    "    training_update = optimizer.minimize(objective, var_list=policy.parameters)\n",
    "    \n",
    "with tf.name_scope('state_sampler'):\n",
    "    batch_size = tf.placeholder(tf.int32, shape=[], name='batch_size')\n",
    "    batch = tf.random_uniform([batch_size, state_dim], -1, 1, dtype=OPTIONS.tf_dtype, name='batch')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Policy Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.run(tf.variables_initializer(policy.parameters))\n",
    "\n",
    "# Uniformly sampled test set\n",
    "test_size = 1e3\n",
    "test_set = batch.eval({batch_size: test_size})\n",
    "\n",
    "# Keep track of the test set loss and parameter changes during training\n",
    "test_loss = []\n",
    "param_changes = []\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Feed dict with hyperparameters\n",
    "feed_dict = {\n",
    "    states:         np.zeros((1, state_dim)), # placeholder\n",
    "    gamma:          discount,                 # use the same discount factor from the value function!\n",
    "    learning_rate:  0.6,\n",
    "    batch_size:     100,   \n",
    "}\n",
    "max_iters = 1000\n",
    "\n",
    "old_params = session.run(policy.parameters)\n",
    "for i in tqdm(range(max_iters)):\n",
    "    feed_dict[states] = batch.eval(feed_dict)\n",
    "    session.run(training_update, feed_dict)\n",
    "    \n",
    "    new_params = session.run(policy.parameters)\n",
    "    param_changes.append(get_parameter_change(old_params, new_params, 'inf'))\n",
    "    old_params = list(new_params)\n",
    "    \n",
    "    feed_dict[states] = test_set\n",
    "    test_loss.append(objective.eval(feed_dict))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(10, 2), dpi=OPTIONS.dpi)\n",
    "fig.subplots_adjust(wspace=0.3)\n",
    "\n",
    "axes[0].plot(test_loss, '.-r')\n",
    "axes[0].set_xlabel(r'SGD iteration $k$')\n",
    "axes[0].set_ylabel(r'test loss')\n",
    "\n",
    "axes[1].plot(param_changes, '.-r')\n",
    "axes[1].set_xlabel(r'SGD iteration $k$')\n",
    "axes[1].set_ylabel(r'$||{\\bf \\delta}_k - {\\bf \\delta}_{k-1}||_\\infty$')\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting\n",
    "\n",
    "Visually compare the true policy $\\pi$ and the neural network $\\pi_{\\bf \\delta}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(6, 5), dpi=OPTIONS.dpi)\n",
    "ax = fig.add_subplot(111, projection='3d')\n",
    "ax.set_xlabel(r'$\\theta$ [deg]')\n",
    "ax.set_ylabel(r'$\\omega$ [deg/s]')\n",
    "    \n",
    "# Policies\n",
    "approx = u_max * actions.eval({states: grid.all_points}).reshape(grid.num_points)\n",
    "truth = u_max * policy_lqr(grid.all_points).eval().reshape(grid.num_points)\n",
    "for zz, c in zip([approx, truth], ['r', 'b']):\n",
    "    surf = ax.plot_surface(xx, yy, zz, color=c, alpha=0.65)\n",
    "    surf._facecolors2d = surf._facecolors3d\n",
    "    surf._edgecolors2d = surf._edgecolors3d\n",
    "\n",
    "# Legend\n",
    "proxy = [plt.Rectangle((0,0), 1, 1, fc=c) for c in [(0., 0., 1., 0.65), (1., 0., 0., 0.65)]]    \n",
    "ax.legend(proxy, [r'$\\pi({\\bf x})$ [N]', r'$\\pi_{\\bf \\delta}({\\bf x})$ [N]'])\n",
    "ax.view_init(None, -45)\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### New ROA\n",
    "\n",
    "Visually compare the new ROA yielded by the parametric policy to the one yielded by the LQR policy, possibly with saturation constraints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "closed_loop_dynamics = lambda x: future_states.eval({states: x})\n",
    "roa_horizon = 600\n",
    "tol = 0.1\n",
    "new_roa, new_trajectories = compute_roa(grid, closed_loop_dynamics, roa_horizon, tol, equilibrium=None, no_traj=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(4, 4), dpi=OPTIONS.dpi)\n",
    "\n",
    "norms = np.rad2deg(state_norm)\n",
    "plot_limits = np.column_stack((- norms, norms))\n",
    "ax.set_aspect(norms[0] / norms[1])\n",
    "ax.set_xlim(plot_limits[0])\n",
    "ax.set_ylim(plot_limits[1])\n",
    "ax.set_xlabel(r'$\\theta$ [deg]')\n",
    "ax.set_ylabel(r'$\\omega$ [deg/s]')\n",
    "    \n",
    "# Compare ROAs\n",
    "cmap = ListedColormap([(1., 1., 1., 0.), (0., 1., 0., 0.65), (1., 0., 0., 0.65)])\n",
    "z = (roa.astype(int) + new_roa.astype(int)).reshape(grid.num_points)\n",
    "im = ax.imshow(z.T, origin='lower', extent=plot_limits.ravel(), aspect=norms[0] / norms[1], cmap=cmap)\n",
    "\n",
    "# Sub-sample discretization for faster and clearer plotting\n",
    "N_traj = 13\n",
    "skip = int(grid.num_points[0] / N_traj)\n",
    "sub_idx = np.arange(grid.nindex).reshape(grid.num_points)\n",
    "sub_idx = sub_idx[::skip, ::skip].ravel()\n",
    "sub_trajectories = new_trajectories[sub_idx, :, :]\n",
    "\n",
    "# Trajectories\n",
    "for n in range(sub_trajectories.shape[0]):\n",
    "    theta = sub_trajectories[n, 0, :] * norms[0]\n",
    "    omega = sub_trajectories[n, 1, :] * norms[1]\n",
    "    ax.plot(theta, omega, 'k--', linewidth=0.6)\n",
    "sub_states = grid.all_points[sub_idx]\n",
    "dx_dt = (future_states.eval({states: sub_states}) - sub_states) / dt\n",
    "dx_dt = dx_dt / np.linalg.norm(dx_dt, ord=2, axis=1, keepdims=True)\n",
    "ax.quiver(sub_states[:, 0] * norms[0], sub_states[:, 1] * norms[1], dx_dt[:, 0], dx_dt[:, 1], scale=None, pivot='mid', headwidth=4, headlength=8, color='k')\n",
    "\n",
    "# Legend\n",
    "proxy = [plt.Rectangle((0,0), 1, 1, fc=c) for c in [(0., 1., 0., 0.65), (1., 0., 0., 0.65)]]    \n",
    "legend = ax.legend(proxy, [r'ROA for $\\pi$', r'ROA for $\\pi_{\\bf \\delta}$'], loc='upper right')\n",
    "legend.get_frame().set_alpha(1.)\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Approximate Policy Iteration\n",
    "\n",
    "Train the policy $\\pi_{\\bf \\delta}$ and value function $V_{\\bf \\theta}$ in tandem with approximate policy iteration. Once again, changing the discount factor strongly affects the results. A low discount factor encourages a well-behaved value function, while a high discount factor encourages the policy to yield a larger ROA."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TensorFlow Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use parametric policy and value function\n",
    "states = tf.placeholder(OPTIONS.tf_dtype, shape=[None, state_dim], name='states')\n",
    "actions = policy(states)\n",
    "rewards = reward_function(states, actions)\n",
    "values = value_function(states)\n",
    "future_states = dynamics(states, actions)\n",
    "future_values = value_function(future_states)\n",
    "\n",
    "# Compare with LQR solution, possibly with saturation constraints\n",
    "actions_lqr = policy_lqr(states)\n",
    "rewards_lqr = reward_function(states, actions_lqr)\n",
    "future_states_lqr = dynamics(states, actions_lqr)\n",
    "\n",
    "# Discount factor and scaling\n",
    "max_state = np.ones((1, state_dim))\n",
    "max_action = np.ones((1, action_dim))\n",
    "r_max = np.linalg.multi_dot((max_state, Q, max_state.T)) + np.linalg.multi_dot((max_action, R, max_action.T))\n",
    "gamma = tf.placeholder(OPTIONS.tf_dtype, shape=[], name='discount_factor')\n",
    "eval_scaling = 1 / r_max.ravel()\n",
    "impv_scaling = (1 - gamma) / r_max.ravel()\n",
    "\n",
    "# Policy evaluation\n",
    "with tf.name_scope('value_optimization'):\n",
    "    value_learning_rate = tf.placeholder(OPTIONS.tf_dtype, shape=[], name='learning_rate')\n",
    "    target = tf.stop_gradient(rewards + gamma * future_values, name='target')\n",
    "    value_objective = eval_scaling * tf.reduce_mean(tf.abs(values - target), name='objective')\n",
    "    optimizer = tf.train.GradientDescentOptimizer(value_learning_rate)\n",
    "    value_update = optimizer.minimize(value_objective, var_list=value_function.parameters)\n",
    "\n",
    "# Policy improvement\n",
    "with tf.name_scope('policy_optimization'):\n",
    "    policy_learning_rate = tf.placeholder(OPTIONS.tf_dtype, shape=[], name='learning_rate')\n",
    "    policy_objective = - impv_scaling * tf.reduce_mean(rewards + gamma * future_values, name='objective')\n",
    "    optimizer = tf.train.GradientDescentOptimizer(policy_learning_rate)\n",
    "    policy_update = optimizer.minimize(policy_objective, var_list=policy.parameters)\n",
    "\n",
    "# Sampling    \n",
    "with tf.name_scope('state_sampler'):\n",
    "    batch_size = tf.placeholder(tf.int32, shape=[], name='batch_size')\n",
    "    batch = tf.random_uniform([batch_size, state_dim], -1, 1, dtype=OPTIONS.tf_dtype, name='batch')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Policy and Value Function Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize parameters\n",
    "session.run(tf.variables_initializer(value_function.parameters))\n",
    "session.run(tf.variables_initializer(policy.parameters))\n",
    "\n",
    "# Uniformly sampled test set\n",
    "test_size = 1e3\n",
    "test_set = batch.eval({batch_size: test_size})\n",
    "\n",
    "# Keep track of the test set losses and parameter changes during training\n",
    "value_test_loss = []\n",
    "value_param_changes = []\n",
    "\n",
    "policy_test_loss = []\n",
    "policy_param_changes = []\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Feed dict with hyperparameters\n",
    "feed_dict = {\n",
    "    states:                np.zeros((1, state_dim)), # placeholder\n",
    "    gamma:                 0.965,\n",
    "    value_learning_rate:   0.005,\n",
    "    policy_learning_rate:  0.6,\n",
    "    batch_size:            int(1e2),\n",
    "}\n",
    "max_iters    = 200\n",
    "value_iters  = 100\n",
    "policy_iters = 10\n",
    "\n",
    "old_value_params = session.run(value_function.parameters)\n",
    "old_policy_params = session.run(policy.parameters)\n",
    "\n",
    "for i in tqdm(range(max_iters)):\n",
    "    \n",
    "    # Policy evaluation (value function update)\n",
    "    for _ in range(value_iters):\n",
    "        feed_dict[states] = batch.eval(feed_dict)\n",
    "        session.run(value_update, feed_dict)\n",
    "        #\n",
    "    feed_dict[states] = test_set\n",
    "    value_test_loss.append(value_objective.eval(feed_dict))\n",
    "    new_value_params = session.run(value_function.parameters)\n",
    "    value_param_changes.append(get_parameter_change(old_value_params, new_value_params, 'inf'))\n",
    "    old_value_params = list(new_value_params)\n",
    "\n",
    "    # Policy improvement (policy update)\n",
    "    for _ in range(policy_iters):\n",
    "        feed_dict[states] = batch.eval(feed_dict)\n",
    "        session.run(policy_update, feed_dict)\n",
    "        #\n",
    "    feed_dict[states] = test_set\n",
    "    policy_test_loss.append(policy_objective.eval(feed_dict))\n",
    "    new_policy_params = session.run(policy.parameters)\n",
    "    policy_param_changes.append(get_parameter_change(old_policy_params, new_policy_params, 'inf'))\n",
    "    old_policy_params = list(new_policy_params)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(10, 2), dpi=OPTIONS.dpi)\n",
    "fig.subplots_adjust(wspace=0.3)\n",
    "\n",
    "axes[0].plot(value_test_loss, '.-r')\n",
    "axes[0].set_xlabel(r'Policy iteration $k$')\n",
    "axes[0].set_ylabel(r'test loss (policy evaluation)')\n",
    "\n",
    "axes[1].plot(value_param_changes, '.-r')\n",
    "axes[1].set_xlabel(r'Policy iteration $k$')\n",
    "axes[1].set_ylabel(r'$||{\\bf \\theta}_k - {\\bf \\theta}_{k-1}||_\\infty$')\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(10, 2), dpi=OPTIONS.dpi)\n",
    "fig.subplots_adjust(wspace=0.3)\n",
    "\n",
    "axes[0].plot(policy_test_loss, '.-r')\n",
    "axes[0].set_xlabel(r'Policy iteration $k$')\n",
    "axes[0].set_ylabel(r'test loss (policy improvement)')\n",
    "\n",
    "axes[1].plot(policy_param_changes, '.-r')\n",
    "axes[1].set_xlabel(r'Policy iteration $k$')\n",
    "axes[1].set_ylabel(r'$||{\\bf \\delta}_k - {\\bf \\delta}_{k-1}||_\\infty$')\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(10, 4), dpi=OPTIONS.dpi)\n",
    "\n",
    "# Value functions ----------------------------------------------- #\n",
    "ax = fig.add_subplot(121, projection='3d')\n",
    "ax.set_xlabel(r'$\\theta$ [deg]')\n",
    "ax.set_ylabel(r'$\\omega$ [deg/s]')\n",
    "\n",
    "approx = - values.eval({states: grid.all_points}).reshape(grid.num_points)\n",
    "truth = - true_values.reshape(grid.num_points)\n",
    "for zz, c in zip([approx, truth], ['r', 'b']):\n",
    "    surf = ax.plot_surface(xx, yy, zz, color=c, alpha=0.65)\n",
    "    surf._facecolors2d = surf._facecolors3d\n",
    "    surf._edgecolors2d = surf._edgecolors3d\n",
    "\n",
    "# ROA\n",
    "z = roa.reshape(grid.num_points)\n",
    "ax.contourf(xx, yy, z, cmap=binary_cmap('green', 0.65), zdir='z', offset=0)\n",
    "\n",
    "# Legend\n",
    "proxy = [plt.Rectangle((0,0), 1, 1, fc=c) for c in [(0., 0., 1., 0.65), (1., 0., 0., 0.65), (0., 1., 0., 0.65)]]    \n",
    "ax.legend(proxy, [r'$-V_{\\pi}({\\bf x})$', r'$-V_{\\bf \\theta}({\\bf x})$', r'ROA for $\\pi$'])\n",
    "ax.view_init(None, -45)\n",
    "\n",
    "\n",
    "# Policies ------------------------------------------------------ #\n",
    "ax = fig.add_subplot(122, projection='3d')\n",
    "ax.set_xlabel(r'$\\theta$ [deg]')\n",
    "ax.set_ylabel(r'$\\omega$ [deg/s]')\n",
    "\n",
    "approx = u_max * actions.eval({states: grid.all_points}).reshape(grid.num_points)\n",
    "truth = u_max * policy_lqr(grid.all_points).eval().reshape(grid.num_points)\n",
    "for zz, c in zip([approx, truth], ['r', 'b']):\n",
    "    surf = ax.plot_surface(xx, yy, zz, color=c, alpha=0.65)\n",
    "    surf._facecolors2d = surf._facecolors3d\n",
    "    surf._edgecolors2d = surf._edgecolors3d\n",
    "\n",
    "# Legend\n",
    "proxy = [plt.Rectangle((0,0), 1, 1, fc=c) for c in [(0., 0., 1., 0.65), (1., 0., 0., 0.65)]]    \n",
    "ax.legend(proxy, [r'$\\pi({\\bf x})$ [N]', r'$\\pi_{\\bf \\delta}({\\bf x})$ [N]'])\n",
    "ax.view_init(None, -45)\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### New ROA\n",
    "\n",
    "Visually compare the new ROA yielded by the parametric policy to the one yielded by the LQR policy, possibly with saturation constraints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "closed_loop_dynamics = lambda x: future_states.eval({states: x})\n",
    "roa_horizon = 800\n",
    "tol = 0.01\n",
    "new_roa, new_trajectories = compute_roa(grid, closed_loop_dynamics, roa_horizon, tol, equilibrium=None, no_traj=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(4, 4), dpi=OPTIONS.dpi)\n",
    "\n",
    "norms = np.rad2deg(state_norm)\n",
    "plot_limits = np.column_stack((- norms, norms))\n",
    "ax.set_aspect(norms[0] / norms[1])\n",
    "ax.set_xlim(plot_limits[0])\n",
    "ax.set_ylim(plot_limits[1])\n",
    "ax.set_xlabel(r'$\\theta$ [deg]')\n",
    "ax.set_ylabel(r'$\\omega$ [deg/s]')\n",
    "    \n",
    "# Compare ROAs\n",
    "cmap = ListedColormap([(1., 1., 1., 0.), (0., 1., 0., 0.65), (1., 0., 0., 0.65)])\n",
    "z = (roa.astype(int) + new_roa.astype(int)).reshape(grid.num_points)\n",
    "im = ax.imshow(z.T, origin='lower', extent=plot_limits.ravel(), aspect=norms[0] / norms[1], cmap=cmap)\n",
    "\n",
    "# Sub-sample discretization for faster and clearer plotting\n",
    "N_traj = 13\n",
    "skip = int(grid.num_points[0] / N_traj)\n",
    "sub_idx = np.arange(grid.nindex).reshape(grid.num_points)\n",
    "sub_idx = sub_idx[::skip, ::skip].ravel()\n",
    "sub_trajectories = new_trajectories[sub_idx, :, :]\n",
    "\n",
    "# Trajectories\n",
    "for n in range(sub_trajectories.shape[0]):\n",
    "    theta = sub_trajectories[n, 0, :] * norms[0]\n",
    "    omega = sub_trajectories[n, 1, :] * norms[1]\n",
    "    ax.plot(theta, omega, 'k--', linewidth=0.6)\n",
    "sub_states = grid.all_points[sub_idx]\n",
    "dx_dt = (future_states.eval({states: sub_states}) - sub_states) / dt\n",
    "dx_dt = dx_dt / np.linalg.norm(dx_dt, ord=2, axis=1, keepdims=True)\n",
    "ax.quiver(sub_states[:, 0] * norms[0], sub_states[:, 1] * norms[1], dx_dt[:, 0], dx_dt[:, 1], scale=None, pivot='mid', headwidth=4, headlength=8, color='k')\n",
    "\n",
    "# Legend\n",
    "proxy = [plt.Rectangle((0,0), 1, 1, fc=c) for c in [(0., 1., 0., 0.65), (1., 0., 0., 0.65)]]    \n",
    "legend = ax.legend(proxy, [r'ROA for $\\pi$', r'ROA for $\\pi_{\\bf \\delta}$'], loc='upper right')\n",
    "legend.get_frame().set_alpha(1.)\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}