exercise2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e7758886",
   "metadata": {},
   "source": [
    "# Bayesian Optimization Tutorial (Exercise): Using Ax package to tune a linear controller to stabilize a cart_pole.    \n",
    "**By Wenjie Xu**\n",
    "\n",
    "Throughout this exercise, you will see blocks of code. In some of the places, they contain an indicator:\n",
    "```python\n",
    "## [TODO]\n",
    "```\n",
    "\n",
    "This indicates the part of code that you need to implement yourself.  \n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# In this exercise, we give the objective reward function to maximize, and you are required to use Ax package to tune the linear controller. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b991c8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gym\n",
    "import numpy as np\n",
    "import time\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import Video, display, clear_output\n",
    "from matplotlib.animation import FuncAnimation\n",
    "from IPython.display import HTML\n",
    "import moviepy.editor as mpy \n",
    "import os\n",
    "from render import show_video # show video implements the visualization of frames"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28affe9c",
   "metadata": {},
   "source": [
    "# The goal is to stabilize the following cart-pole, both in position and angle."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e170e63",
   "metadata": {},
   "outputs": [],
   "source": [
    "display(Video('./render_videos/cart_pole_play.mp4', embed=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22b1495b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# [TODO]: import relevant packages including Ax\n",
    "# import Ax related packages\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4990f33b",
   "metadata": {},
   "source": [
    "# We implement the function to evaluate reward for you. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8075f08f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_cart_pole_experiment(Kp, Kd, Kpx, Kdx, num_episodes=5, max_steps=500):\n",
    "    # reward function takes Kp and Kd as inputs\n",
    "    \n",
    "    # Initialize the CartPole environment\n",
    "    env = gym.make('CartPole-v1')\n",
    "\n",
    "    def pd_controller(state):\n",
    "        # Extract the state variables\n",
    "        x, x_dot, theta, theta_dot = state\n",
    "    \n",
    "        # PD control law\n",
    "        force = Kp * theta + Kd * theta_dot + Kpx * x + Kdx * x_dot\n",
    "    \n",
    "        # Convert force to discrete action (left or right)\n",
    "        action = 1 if force > 0 else 0\n",
    "    \n",
    "        return action\n",
    "    \n",
    "    total_reward_list = []\n",
    "    for episode in range(num_episodes):\n",
    "        state = env.reset()\n",
    "        total_reward = 0\n",
    "        frames = []\n",
    "        for t in range(max_steps):\n",
    "            if True:\n",
    "                frame = env.render(mode='rgb_array')  # Capture frame\n",
    "                frames.append(frame)\n",
    "                    \n",
    "            # Get action from PD controller\n",
    "            action = pd_controller(state)\n",
    "        \n",
    "            # Step the environment\n",
    "            state, reward, done, _ = env.step(action)  \n",
    "            total_reward += max((reward - 70 * np.abs(state[2]) - 5 * np.abs(state[1])), 0) # penalize the pole angle and moving\n",
    "        \n",
    "            if done:\n",
    "                break\n",
    "    \n",
    "        print(f\"Episode {episode + 1}: Total Reward = {total_reward}\")\n",
    "        total_reward_list.append(total_reward)\n",
    "    env.close()\n",
    "    return np.min(total_reward_list), frames\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cec3f35",
   "metadata": {},
   "outputs": [],
   "source": [
    "# [TODO]: manually tune Kp and Kd yourself before call Ax package\n",
    "manual_kp = 2.0\n",
    "manual_kd = 1.0\n",
    "manual_kpx = 2.0\n",
    "manual_kdx = 1.0\n",
    "mean_total_reward, frames = run_cart_pole_experiment(manual_kp, manual_kd, manual_kpx, manual_kdx, num_episodes=5)\n",
    "show_video(frames)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "392cd46d",
   "metadata": {},
   "source": [
    "# Now you should implement the tuning algorithm by Ax package yourself"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d3890c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# [TODO]: define the Ax client and the interface to evaluate objective\n",
    "\n",
    "\n",
    "def pd_to_frames(pd):\n",
    "    kp = pd['kp']\n",
    "    kd = pd['kd']\n",
    "    kpx = pd['kpx']\n",
    "    kdx = pd['kdx']\n",
    "    mean_total_reward, frames = run_cart_pole_experiment(kp, kd, kpx, kdx, num_episodes=1)\n",
    "    return frames\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e0b4561",
   "metadata": {},
   "outputs": [],
   "source": [
    "# [TODO]: run the optimization loop\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c25c8734",
   "metadata": {},
   "outputs": [],
   "source": [
    "# [TODO] plot the optimization trace\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61bd8ada",
   "metadata": {},
   "outputs": [],
   "source": [
    "# [TODO] compare the running result with the parameters tuned by you manually and the parameters obtained by Ax \n",
    "\n",
    "# Example code of doing visualization\n",
    "# init_frames = pd_to_frames(init_parameters)\n",
    "# show_video(init_frames)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3813f66e",
   "metadata": {},
   "outputs": [],
   "source": [
    "best_frames = pd_to_frames(best_parameters)\n",
    "show_video(best_frames)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77c8d16a",
   "metadata": {},
   "source": [
    "# Bonus question.\n",
    "Tune the controller only using pairwise comparison data. That is, in each step, the algorithm proposes two solutions, and then you express which one is preferred. The algorithm uses this comparison data to search your favorite solution.\n",
    "\n",
    "# Hint.\n",
    "Refer to https://botorch.org/tutorials/preference_bo .\n"
   ]
  }
 ],
 "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}