examples/reinforcement_learning_cartpole.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reinforcement Learning for the Cart-Pole\n",
    "\n",
    "Perform approximate policy iteration in an actor-critic framework for the cart-pole (i.e., inverted pendulum on a cart)."
   ]
  },
  {
   "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 CartPole, compute_closedloop_response, get_parameter_change, find_nearest, reward_rollout, compute_roa, binary_cmap\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.175    # pendulum mass\n",
    "M = 1.732    # cart mass\n",
    "L = 0.28     # pole length\n",
    "b = 0.01     # rotational friction\n",
    "\n",
    "# State and action normalizers\n",
    "x_max         = 0.5                                 # linear position [m]\n",
    "theta_max     = np.deg2rad(30)                      # angular position [rad]\n",
    "x_dot_max     = 2                                   # linear velocity [m/s]\n",
    "theta_dot_max = np.deg2rad(30)                      # angular velocity [rad/s]\n",
    "u_max         = (m + M) * (x_dot_max ** 2) / x_max  # linear force [N], control action\n",
    "\n",
    "state_norm = (x_max, theta_max, x_dot_max, theta_dot_max)\n",
    "action_norm = (u_max,)\n",
    "\n",
    "# Dimensions and domains\n",
    "state_dim     = 4\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",
    "cartpole = CartPole(m, M, L, b, dt, [state_norm, action_norm])\n",
    "A, B = cartpole.linearize()\n",
    "\n",
    "if OPTIONS.use_linear_dynamics:\n",
    "    dynamics = safe_learning.functions.LinearSystem((A, B), name='dynamics')\n",
    "else:\n",
    "    dynamics = cartpole.__call__\n",
    "\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",
    "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=False)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LQR Policy\n",
    "\n",
    "We compare our results to the LQR solution for the linearized system later."
   ]
  },
  {
   "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": [
    "# 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",
    "\n",
    "val_scaling = 1 / r_max.ravel()\n",
    "pol_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 = pol_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 = - pol_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": [
    "## Approximate Policy Iteration\n",
    "\n",
    "Train the policy $\\pi_{\\bf \\delta}$ and value function $V_{\\bf \\theta}$ in tandem with approximate policy iteration. 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. We compare $\\pi_{\\bf \\delta}$ to the LQR policy $\\pi$ with saturation constraints, and $V_{\\bf \\theta}$ to the LQR value function $V_\\pi$ and the value function $V_{\\pi_{\\bf \\delta}}$ induced by the parametric policy $\\pi_{\\bf \\delta}$. We compute $V_{\\pi_{\\bf \\delta}}$ as a rollout sum of discounted rewards at states in a state space discretization."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.run(tf.global_variables_initializer())\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",
    "value_test_loss = []\n",
    "value_param_changes = []\n",
    "policy_test_loss = []\n",
    "policy_param_changes = []\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training hyperparameters\n",
    "max_iters    = 200\n",
    "value_iters  = 100\n",
    "policy_iters = 10\n",
    "feed_dict = {\n",
    "    states:                test_set,\n",
    "    gamma:                 0.99,\n",
    "    value_learning_rate:   0.2,\n",
    "    policy_learning_rate:  0.5,\n",
    "    batch_size:            1e2,\n",
    "}\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",
    "    # Policy evaluation (value update)\n",
    "    for _ in range(value_iters):\n",
    "        feed_dict[states] = batch.eval(feed_dict)\n",
    "        session.run(value_update, 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))\n",
    "    old_value_params = 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",
    "    new_policy_params = session.run(policy.parameters)\n",
    "    policy_param_changes.append(get_parameter_change(old_policy_params, new_policy_params))\n",
    "    old_policy_params = new_policy_params\n",
    "    \n",
    "    # Record objectives\n",
    "    feed_dict[states] = test_set\n",
    "    value_test_loss.append(value_objective.eval(feed_dict))\n",
    "    policy_test_loss.append(policy_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(2, 2, figsize=(12, 5), dpi=OPTIONS.dpi)\n",
    "fig.subplots_adjust(wspace=0.3, hspace=0.4)\n",
    "\n",
    "ax = axes[0,0]\n",
    "ax.plot(value_test_loss, '.-r')\n",
    "ax.set_xlabel(r'Policy iteration $k$')\n",
    "ax.set_ylabel(r'test loss (policy evaluation)')\n",
    "\n",
    "ax = axes[0,1]\n",
    "ax.plot(value_param_changes, '.-r')\n",
    "ax.set_xlabel(r'Policy iteration $k$')\n",
    "ax.set_ylabel(r'$||{\\bf \\theta}_k - {\\bf \\theta}_{k-1}||_\\infty$')\n",
    "\n",
    "ax = axes[1,0]\n",
    "ax.plot(policy_test_loss, '.-b')\n",
    "ax.set_xlabel(r'Policy iteration $k$')\n",
    "ax.set_ylabel(r'test loss (policy improvement)')\n",
    "\n",
    "ax = axes[1,1]\n",
    "ax.plot(policy_param_changes, '.-b')\n",
    "ax.set_xlabel(r'Policy iteration $k$')\n",
    "ax.set_ylabel(r'$||{\\bf \\delta}_k - {\\bf \\delta}_{k-1}||_\\infty$')\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Estimated Value Functions and ROAs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Number of states along each dimension\n",
    "num_states = 51\n",
    "\n",
    "# State grid\n",
    "grid_limits = np.array([[-1., 1.], ] * state_dim)\n",
    "grid = safe_learning.GridWorld(grid_limits, num_states)\n",
    "\n",
    "# Estimate value functions and ROAs with rollout\n",
    "roa_horizon     = 2000\n",
    "rollout_horizon = 500\n",
    "roa_tol         = 0.1\n",
    "rollout_tol     = 0.01\n",
    "discount        = feed_dict[gamma]  # use the same discount factor from training!\n",
    "pivot_state     = np.asarray([0., 0., 0., 0.], dtype=OPTIONS.np_dtype)\n",
    "\n",
    "# Snap pivot_state to the closest grid point\n",
    "pivot_index = np.zeros_like(pivot_state, dtype=int)\n",
    "for d in range(grid.ndim):\n",
    "    pivot_index[d], pivot_state[d] = find_nearest(grid.discrete_points[d], pivot_state[d])\n",
    "\n",
    "# Get 2d-planes of the discretization (x vs. v, theta vs. omega) according to pivot_state\n",
    "planes = [[1, 3], [0, 2]]\n",
    "grid_slices = []\n",
    "for p in planes:\n",
    "    grid_slices.append(np.logical_and(grid.all_points[:, p[0]] == pivot_state[p[0]], \n",
    "                                      grid.all_points[:, p[1]] == pivot_state[p[1]]).ravel())\n",
    "\n",
    "# LQR solution (\\pi and V_\\pi)\n",
    "closed_loop_dynamics = lambda x: future_states_lqr.eval({states: x})\n",
    "reward_eval          = lambda x: rewards_lqr.eval({states: x})\n",
    "true_values          = [reward_rollout(grid.all_points[mask], closed_loop_dynamics, reward_eval, discount, rollout_horizon, rollout_tol) for mask in grid_slices]\n",
    "true_roas            = [compute_roa(grid.all_points[mask], closed_loop_dynamics, roa_horizon, roa_tol) for mask in grid_slices]\n",
    "\n",
    "# Parametric policy's value function V_{\\pi_\\delta}\n",
    "closed_loop_dynamics = lambda x: future_states.eval({states: x})\n",
    "reward_eval          = lambda x: rewards.eval({states: x})\n",
    "est_values           = [reward_rollout(grid.all_points[mask], closed_loop_dynamics, reward_eval, discount, rollout_horizon, rollout_tol) for mask in grid_slices]\n",
    "est_roas             = [compute_roa(grid.all_points[mask], closed_loop_dynamics, roa_horizon, roa_tol) for mask in grid_slices]\n",
    "\n",
    "# Parametric value function V_\\theta\n",
    "par_values = [values.eval({states: grid.all_points[mask]}) for mask in grid_slices]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "planes = [[0, 2], [1, 3]]\n",
    "norms = np.asarray([x_max, np.rad2deg(theta_max), x_dot_max, np.rad2deg(theta_dot_max)])\n",
    "scaled_discrete_points = [norm * points for norm, points in zip(norms, grid.discrete_points)]\n",
    "\n",
    "fig = plt.figure(figsize=(12, 12), dpi=OPTIONS.dpi)\n",
    "\n",
    "for i, p in enumerate(planes):\n",
    "    ax = fig.add_subplot(221 + i, projection='3d')\n",
    "    if i == 0:\n",
    "        ax.set_title(r'$\\theta = {:g}$'.format(pivot_state[1]) + r', $\\dot\\theta = {:g}$'.format(pivot_state[3]) + '\\n')\n",
    "        ax.set_xlabel(r'$x$ [m]')\n",
    "        ax.set_ylabel(r'$\\dot{x}$ [m/s]')\n",
    "    else:\n",
    "        ax.set_title(r'$x= {:g}$'.format(pivot_state[0]) + r', $\\dot x = {:g}$'.format(pivot_state[2]) + '\\n')\n",
    "        ax.set_xlabel(r'$\\theta$ [deg]')\n",
    "        ax.set_ylabel(r'$\\dot{\\theta}$ [deg/s]')\n",
    "    ax.view_init(None, -45)\n",
    "\n",
    "    xx, yy = np.meshgrid(*[scaled_discrete_points[p[0]], scaled_discrete_points[p[1]]])\n",
    "\n",
    "    for j, (values, color) in enumerate(zip([true_values, est_values, par_values], [(0, 0, 1, 0.6), (0, 1, 0, 0.8), (1, 0, 0, 0.65)])):\n",
    "        z = - values[i].reshape(grid.num_points[p])\n",
    "        surf = ax.plot_surface(xx, yy, z, color=color)\n",
    "        surf._facecolors2d = surf._facecolors3d\n",
    "        surf._edgecolors2d = surf._edgecolors3d\n",
    "    proxy = [plt.Rectangle((0,0), 1, 1, fc=c) for c in [(0, 0, 1, 0.6), (0, 1, 0, 0.8), (1, 0, 0, 0.65)]]    \n",
    "    ax.legend(proxy, [r'$-V_{\\pi}({\\bf x})$', r'$-V_{\\pi_{\\bf \\delta}}({\\bf x})$', r'$-V_{\\bf \\theta}({\\bf x})$'])\n",
    "\n",
    "\n",
    "for i, (p, mask) in enumerate(zip(planes, grid_slices)):\n",
    "    ax = fig.add_subplot(223 + i, projection='3d')\n",
    "    if i == 0:\n",
    "        ax.set_title(r'$\\theta = {:g}$'.format(pivot_state[1]) + r', $\\dot\\theta = {:g}$'.format(pivot_state[3]) + '\\n')\n",
    "        ax.set_xlabel(r'$x$ [m]')\n",
    "        ax.set_ylabel(r'$\\dot{x}$ [m/s]') \n",
    "    else:\n",
    "        ax.set_title(r'$x= {:g}$'.format(pivot_state[0]) + r', $\\dot x = {:g}$'.format(pivot_state[2]) + '\\n')\n",
    "        ax.set_xlabel(r'$\\theta$ [deg]')\n",
    "        ax.set_ylabel(r'$\\dot{\\theta}$ [deg/s]')\n",
    "    ax.view_init(None, -45)\n",
    "    \n",
    "    xx, yy = np.meshgrid(*[scaled_discrete_points[p[0]], scaled_discrete_points[p[1]]])\n",
    "    acts = u_max * actions.eval({states: grid.all_points[mask]})\n",
    "    true_acts = u_max * actions_lqr.eval({states: grid.all_points[mask]})\n",
    "\n",
    "    ax.plot_surface(xx, yy, true_acts.reshape(grid.num_points[p]), color='blue', alpha=0.55)\n",
    "    ax.plot_surface(xx, yy, acts.reshape(grid.num_points[p]), color='red', alpha=0.75)\n",
    "\n",
    "    z = est_roas[i].reshape(grid.num_points[p])\n",
    "    ax.contourf(xx, yy, z, cmap=binary_cmap('green', 0.65), zdir='z', offset=-u_max)\n",
    "\n",
    "    proxy = [plt.Rectangle((0,0), 1, 1, fc=c) for c in [(0, 0, 1, 0.6), (1, 0, 0, 0.65), (0., 1., 0., 0.65)]]\n",
    "    ax.legend(proxy, [r'$\\pi({\\bf x})$ [N]', r'$\\pi_{\\bf \\delta}({\\bf x})$ [N]', r'ROA for $\\pi_{\\bf \\delta}$'])\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
}