{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training RL to do Cartpole Balancing\n",
    "\n",
    "This notebooks is part of [AI for Beginners Curriculum](http://aka.ms/ai-beginners). It has been inspired by [this blog post](https://medium.com/swlh/policy-gradient-reinforcement-learning-with-keras-57ca6ed32555), [official TensorFlow documentation](https://www.tensorflow.org/tutorials/reinforcement_learning/actor_critic) and [this Keras RL example](https://keras.io/examples/rl/actor_critic_cartpole/).\n",
    "\n",
    "In this example, we will use RL to train a model to balance a pole on a cart that can move left and right on horizontal scale. We will use [OpenAI Gym](https://www.gymlibrary.ml/) environment to simulate the pole.\n",
    "\n",
    "> **Note**: You can run this lesson's code locally (eg. from Visual Studio Code), in which case the simulation will open in a new window. When running the code online, you may need to make some tweaks to the code, as described [here](https://towardsdatascience.com/rendering-openai-gym-envs-on-binder-and-google-colab-536f99391cc7).\n",
    "\n",
    "We will start by making sure Gym is installed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Defaulting to user installation because normal site-packages is not writeable\n",
      "Requirement already satisfied: gym in /home/leo/.local/lib/python3.10/site-packages (0.25.0)\n",
      "Requirement already satisfied: pygame in /home/leo/.local/lib/python3.10/site-packages (2.1.2)\n",
      "Requirement already satisfied: gym-notices>=0.0.4 in /home/leo/.local/lib/python3.10/site-packages (from gym) (0.0.7)\n",
      "Requirement already satisfied: cloudpickle>=1.2.0 in /home/leo/.local/lib/python3.10/site-packages (from gym) (2.1.0)\n",
      "Requirement already satisfied: numpy>=1.18.0 in /usr/lib/python3/dist-packages (from gym) (1.21.5)\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "!{sys.executable} -m pip install gym pygame"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's create the CartPole environment and see how to operate on it. An environment has the following properties:\n",
    "\n",
    "* **Action space** is the set of possible actions that we can perform at each step of the simulation\n",
    "* **Observation space** is the space of observations that we can make"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Action space: Discrete(2)\n",
      "Observation space: Box([-4.8000002e+00 -3.4028235e+38 -4.1887903e-01 -3.4028235e+38], [4.8000002e+00 3.4028235e+38 4.1887903e-01 3.4028235e+38], (4,), float32)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/leo/.local/lib/python3.10/site-packages/gym/core.py:329: DeprecationWarning: \u001b[33mWARN: Initializing wrapper in old step API which returns one bool instead of two. It is recommended to set `new_step_api=True` to use new step API. This will be the default behaviour in future.\u001b[0m\n",
      "  deprecation(\n",
      "/home/leo/.local/lib/python3.10/site-packages/gym/wrappers/step_api_compatibility.py:39: DeprecationWarning: \u001b[33mWARN: Initializing environment in old step API which returns one bool instead of two. It is recommended to set `new_step_api=True` to use new step API. This will be the default behaviour in future.\u001b[0m\n",
      "  deprecation(\n"
     ]
    }
   ],
   "source": [
    "import gym\n",
    "import pygame\n",
    "import tqdm\n",
    "\n",
    "env = gym.make(\"CartPole-v1\")\n",
    "\n",
    "print(f\"Action space: {env.action_space}\")\n",
    "print(f\"Observation space: {env.observation_space}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see how the simulation works. The following loop runs the simulation, until `env.step` does not return the termination flag `done`. We will randomly chose actions using `env.action_space.sample()`, which means the experiment will probably fail very fast (CartPole environment terminates when the speed of CartPole, its position or angle are outside certain limits).\n",
    "\n",
    "> Simulation will open in the new window. You can run the code several times and see how it behaves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/leo/.local/lib/python3.10/site-packages/gym/core.py:57: DeprecationWarning: \u001b[33mWARN: You are calling render method, but you didn't specified the argument render_mode at environment initialization. To maintain backward compatibility, the environment will render in human mode.\n",
      "If you want to render in human mode, initialize the environment in this way: gym.make('EnvName', render_mode='human') and don't call the render method.\n",
      "See here for more information: https://www.gymlibrary.ml/content/api/\u001b[0m\n",
      "  deprecation(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 0.00425272 -0.19994313  0.00917169  0.34113726] -> 1.0\n",
      "[ 0.00025386 -0.00495286  0.01599443  0.05136059] -> 1.0\n",
      "[ 1.5480528e-04  1.8993615e-01  1.7021643e-02 -2.3623335e-01] -> 1.0\n",
      "[ 0.00395353  0.38481084  0.01229698 -0.5234989 ] -> 1.0\n",
      "[ 0.01164974  0.18951797  0.001827   -0.22696657] -> 1.0\n",
      "[ 0.0154401   0.38461378 -0.00271233 -0.51907265] -> 1.0\n",
      "[ 0.02313238  0.5797738  -0.01309379 -0.812609  ] -> 1.0\n",
      "[ 0.03472786  0.38483363 -0.02934597 -0.5240733 ] -> 1.0\n",
      "[ 0.04242453  0.580356   -0.03982743 -0.8258571 ] -> 1.0\n",
      "[ 0.05403165  0.38580072 -0.05634458 -0.54596174] -> 1.0\n",
      "[ 0.06174766  0.19151384 -0.06726381 -0.27155042] -> 1.0\n",
      "[ 0.06557794 -0.00258703 -0.07269482 -0.00081817] -> 1.0\n",
      "[ 0.0655262  -0.19659522 -0.07271118  0.26807207] -> 1.0\n",
      "[ 0.0615943  -0.00051497 -0.06734974 -0.04662942] -> 1.0\n",
      "[ 0.061584    0.19550486 -0.06828233 -0.3597784 ] -> 1.0\n",
      "[ 0.06549409  0.00141663 -0.0754779  -0.08938391] -> 1.0\n",
      "[ 0.06552242 -0.19254686 -0.07726558  0.17856352] -> 1.0\n",
      "[ 0.06167149  0.00359088 -0.0736943  -0.1374588 ] -> 1.0\n",
      "[ 0.0617433   0.19968675 -0.07644348 -0.45245075] -> 1.0\n",
      "[ 0.06573704  0.3958018  -0.0854925  -0.7682167 ] -> 1.0\n",
      "[ 0.07365308  0.20195423 -0.10085683 -0.50361156] -> 1.0\n",
      "[ 0.07769216  0.0083876  -0.11092906 -0.24433874] -> 1.0\n",
      "[ 0.07785992 -0.18498953 -0.11581583  0.01139782] -> 1.0\n",
      "[ 0.07416012  0.01158649 -0.11558788 -0.31546465] -> 1.0\n",
      "[ 0.07439185  0.20814891 -0.12189718 -0.64224803] -> 1.0\n",
      "[ 0.07855483  0.01491799 -0.13474214 -0.3903015 ] -> 1.0\n",
      "[ 0.07885319 -0.17806001 -0.14254816 -0.14295265] -> 1.0\n",
      "[ 0.07529199  0.01878517 -0.14540721 -0.47699296] -> 1.0\n",
      "[ 0.07566769 -0.17401667 -0.15494707 -0.23344138] -> 1.0\n",
      "[ 0.07218736  0.0229406  -0.1596159  -0.57071024] -> 1.0\n",
      "[ 0.07264617  0.21989843 -0.1710301  -0.9091196 ] -> 1.0\n",
      "[ 0.07704414  0.02745241 -0.1892125  -0.6747003 ] -> 1.0\n",
      "[ 0.07759319 -0.16460665 -0.20270652 -0.4470505 ] -> 1.0\n",
      "[ 0.07430106 -0.35637102 -0.21164753 -0.22448184] -> 1.0\n",
      "Total reward: 34.0\n"
     ]
    }
   ],
   "source": [
    "env.reset()\n",
    "\n",
    "done = False\n",
    "total_reward = 0\n",
    "while not done:\n",
    "   env.render()\n",
    "   obs, rew, done, info = env.step(env.action_space.sample())\n",
    "   total_reward += rew\n",
    "   print(f\"{obs} -> {rew}\")\n",
    "print(f\"Total reward: {total_reward}\")\n",
    "\n",
    "env.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Youn can notice that observations contain 4 numbers. They are:\n",
    "- Position of cart\n",
    "- Velocity of cart\n",
    "- Angle of pole\n",
    "- Rotation rate of pole\n",
    "\n",
    "`rew` is the reward we receive at each step. You can see that in CartPole environment you are rewarded 1 point for each simulation step, and the goal is to maximize total reward, i.e. the time CartPole is able to balance without falling.\n",
    "\n",
    "During reinforcement learning, our goal is to train a **policy** $\\pi$, that for each state $s$ will tell us which action $a$ to take, so essentially $a = \\pi(s)$.\n",
    "\n",
    "If you want probabilistic solution, you can think of policy as returning a set of probabilities for each action, i.e. $\\pi(a|s)$ would mean a probability that we should take action $a$ at state $s$.\n",
    "\n",
    "## Policy Gradient Method\n",
    "\n",
    "In simplest RL algorithm, called **Policy Gradient**, we will train a neural network to predict the next action."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.10/dist-packages/tensorflow/__init__.py:29: DeprecationWarning: The distutils package is deprecated and slated for removal in Python 3.12. Use setuptools or check PEP 632 for potential alternatives\n",
      "  import distutils as _distutils\n",
      "2022-07-24 16:50:47.597258: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory\n",
      "2022-07-24 16:50:47.597280: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.\n",
      "/usr/local/lib/python3.10/dist-packages/flatbuffers/compat.py:19: DeprecationWarning: the imp module is deprecated in favour of importlib and slated for removal in Python 3.12; see the module's documentation for alternative uses\n",
      "  import imp\n",
      "2022-07-24 16:50:49.838826: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:975] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2022-07-24 16:50:49.839078: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory\n",
      "2022-07-24 16:50:49.839143: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcublas.so.11'; dlerror: libcublas.so.11: cannot open shared object file: No such file or directory\n",
      "2022-07-24 16:50:49.839194: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcublasLt.so.11'; dlerror: libcublasLt.so.11: cannot open shared object file: No such file or directory\n",
      "2022-07-24 16:50:49.839245: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcufft.so.10'; dlerror: libcufft.so.10: cannot open shared object file: No such file or directory\n",
      "2022-07-24 16:50:49.839295: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcurand.so.10'; dlerror: libcurand.so.10: cannot open shared object file: No such file or directory\n",
      "2022-07-24 16:50:49.839345: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcusolver.so.11'; dlerror: libcusolver.so.11: cannot open shared object file: No such file or directory\n",
      "2022-07-24 16:50:49.839392: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcusparse.so.11'; dlerror: libcusparse.so.11: cannot open shared object file: No such file or directory\n",
      "2022-07-24 16:50:49.839441: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudnn.so.8'; dlerror: libcudnn.so.8: cannot open shared object file: No such file or directory\n",
      "2022-07-24 16:50:49.839449: W tensorflow/core/common_runtime/gpu/gpu_device.cc:1850] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.\n",
      "Skipping registering GPU devices...\n",
      "2022-07-24 16:50:49.839649: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "num_inputs = 4\n",
    "num_actions = 2\n",
    "\n",
    "model = keras.Sequential([\n",
    "    keras.layers.Dense(128, activation=\"relu\",input_shape=(num_inputs,)),\n",
    "    keras.layers.Dense(num_actions, activation=\"softmax\")\n",
    "])\n",
    "\n",
    "model.compile(loss='categorical_crossentropy', optimizer=keras.optimizers.Adam(learning_rate=0.01))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will train the network by running many experiments, and updating our network after each run. Let's define a function that will run the experiment and return the results (so-called **trace**) - all states, actions (and their recommended probabilities), and rewards:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_episode(max_steps_per_episode = 10000,render=False):    \n",
    "    states, actions, probs, rewards = [],[],[],[]\n",
    "    state = env.reset()\n",
    "    for _ in range(max_steps_per_episode):\n",
    "        if render:\n",
    "            env.render()\n",
    "        action_probs = model(np.expand_dims(state,0))[0]\n",
    "        action = np.random.choice(num_actions, p=np.squeeze(action_probs))\n",
    "        nstate, reward, done, info = env.step(action)\n",
    "        if done:\n",
    "            break\n",
    "        states.append(state)\n",
    "        actions.append(action)\n",
    "        probs.append(action_probs)\n",
    "        rewards.append(reward)\n",
    "        state = nstate\n",
    "    return np.vstack(states), np.vstack(actions), np.vstack(probs), np.vstack(rewards)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can run one episode with untrained network and observe that total reward (AKA length of episode) is very low:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total reward: 27.0\n"
     ]
    }
   ],
   "source": [
    "s,a,p,r = run_episode()\n",
    "print(f\"Total reward: {np.sum(r)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One of the tricky aspects of policy gradient algorithm is to use **discounted rewards**. The idea is that we compute the vector of total rewards at each step of the game, and during this process we discount the early rewards using some coefficient $gamma$. We also normalize the resulting vector, because we will use it as weight to affect our training: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "eps = 0.0001\n",
    "\n",
    "def discounted_rewards(rewards,gamma=0.99,normalize=True):\n",
    "    ret = []\n",
    "    s = 0\n",
    "    for r in rewards[::-1]:\n",
    "        s = r + gamma * s\n",
    "        ret.insert(0, s)\n",
    "    if normalize:\n",
    "        ret = (ret-np.mean(ret))/(np.std(ret)+eps)\n",
    "    return ret"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's do the actual training! We will run 300 episodes, and at each episode we will do the following:\n",
    "\n",
    "1. Run the experiment and collect the trace\n",
    "1. Calculate the difference (`gradients`) between the actions taken, and by predicted probabilities. The less the difference is, the more we are sure that we have taken the right action.\n",
    "1. Calculate discounted rewards and multiply gradients by discounted rewards - that will make sure that steps with higher rewards will make more effect on the final result than lower-rewarded ones\n",
    "1. Expected target actions for our neural network would be partly taken from the predicted probabilities during the run, and partly from calculated gradients. We will use `alpha` parameter to determine to which extent gradients and rewards are taken into account - this is called *learning rate* of reinforcement algorithm.\n",
    "1. Finally, we train our network on states and expected actions, and repeat the process "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 -> 29.0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-07-24 16:50:51.475024: W tensorflow/core/data/root_dataset.cc:247] Optimization loop failed: CANCELLED: Operation was cancelled\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100 -> 135.0\n",
      "200 -> 484.0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-07-24 16:51:35.910774: W tensorflow/core/data/root_dataset.cc:247] Optimization loop failed: CANCELLED: Operation was cancelled\n",
      "2022-07-24 16:51:37.151017: W tensorflow/core/data/root_dataset.cc:247] Optimization loop failed: CANCELLED: Operation was cancelled\n",
      "2022-07-24 16:51:39.284311: W tensorflow/core/data/root_dataset.cc:247] Optimization loop failed: CANCELLED: Operation was cancelled\n",
      "2022-07-24 16:51:42.235074: W tensorflow/core/data/root_dataset.cc:247] Optimization loop failed: CANCELLED: Operation was cancelled\n",
      "2022-07-24 16:51:44.691458: W tensorflow/core/data/root_dataset.cc:247] Optimization loop failed: CANCELLED: Operation was cancelled\n",
      "2022-07-24 16:51:48.381946: W tensorflow/core/data/root_dataset.cc:247] Optimization loop failed: CANCELLED: Operation was cancelled\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f40201e7b20>]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 28976 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "alpha = 1e-4\n",
    "\n",
    "history = []\n",
    "for epoch in range(300):\n",
    "    states, actions, probs, rewards = run_episode()\n",
    "    one_hot_actions = np.eye(2)[actions.T][0]\n",
    "    gradients = one_hot_actions-probs\n",
    "    dr = discounted_rewards(rewards)\n",
    "    gradients *= dr\n",
    "    target = alpha*np.vstack([gradients])+probs\n",
    "    model.train_on_batch(states,target)\n",
    "    history.append(np.sum(rewards))\n",
    "    if epoch%100==0:\n",
    "        print(f\"{epoch} -> {np.sum(rewards)}\")\n",
    "\n",
    "plt.plot(history)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's run the episode with rendering to see the result:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "ename": "error",
     "evalue": "display Surface quit",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31merror\u001b[0m                                     Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_44248/1459719159.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0m_\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mrun_episode\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrender\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m/tmp/ipykernel_44248/3855001447.py\u001b[0m in \u001b[0;36mrun_episode\u001b[0;34m(max_steps_per_episode, render)\u001b[0m\n\u001b[1;32m      4\u001b[0m     \u001b[0;32mfor\u001b[0m \u001b[0m_\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmax_steps_per_episode\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mrender\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 6\u001b[0;31m             \u001b[0menv\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrender\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      7\u001b[0m         \u001b[0maction_probs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexpand_dims\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstate\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      8\u001b[0m         \u001b[0maction\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrandom\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mchoice\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnum_actions\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mp\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msqueeze\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0maction_probs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.10/site-packages/gym/core.py\u001b[0m in \u001b[0;36mrender\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m     64\u001b[0m                 )\n\u001b[1;32m     65\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 66\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mrender_func\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     67\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     68\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mrender\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.10/site-packages/gym/core.py\u001b[0m in \u001b[0;36mrender\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    429\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mrender\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    430\u001b[0m         \u001b[0;34m\"\"\"Renders the environment.\"\"\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 431\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0menv\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrender\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    432\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    433\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mclose\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.10/site-packages/gym/core.py\u001b[0m in \u001b[0;36mrender\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m     64\u001b[0m                 )\n\u001b[1;32m     65\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 66\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mrender_func\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     67\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     68\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mrender\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.10/site-packages/gym/wrappers/order_enforcing.py\u001b[0m in \u001b[0;36mrender\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m     49\u001b[0m                 \u001b[0;34m\"set `disable_render_order_enforcing=True` on the OrderEnforcer wrapper.\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     50\u001b[0m             )\n\u001b[0;32m---> 51\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0menv\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrender\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     52\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     53\u001b[0m     \u001b[0;34m@\u001b[0m\u001b[0mproperty\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.10/site-packages/gym/core.py\u001b[0m in \u001b[0;36mrender\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m     64\u001b[0m                 )\n\u001b[1;32m     65\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 66\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mrender_func\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     67\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     68\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mrender\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.10/site-packages/gym/core.py\u001b[0m in \u001b[0;36mrender\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    429\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mrender\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    430\u001b[0m         \u001b[0;34m\"\"\"Renders the environment.\"\"\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 431\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0menv\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrender\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    432\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    433\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mclose\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.10/site-packages/gym/core.py\u001b[0m in \u001b[0;36mrender\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m     64\u001b[0m                 )\n\u001b[1;32m     65\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 66\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mrender_func\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     67\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     68\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mrender\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.10/site-packages/gym/wrappers/env_checker.py\u001b[0m in \u001b[0;36mrender\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m     53\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0menv_render_passive_checker\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0menv\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     54\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 55\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0menv\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrender\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m~/.local/lib/python3.10/site-packages/gym/core.py\u001b[0m in \u001b[0;36mrender\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m     64\u001b[0m                 )\n\u001b[1;32m     65\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 66\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mrender_func\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     67\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     68\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mrender\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.10/site-packages/gym/envs/classic_control/cartpole.py\u001b[0m in \u001b[0;36mrender\u001b[0;34m(self, mode)\u001b[0m\n\u001b[1;32m    215\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrenderer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_renders\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    216\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 217\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_render\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmode\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    218\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    219\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_render\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmode\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"human\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.10/site-packages/gym/envs/classic_control/cartpole.py\u001b[0m in \u001b[0;36m_render\u001b[0;34m(self, mode)\u001b[0m\n\u001b[1;32m    296\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    297\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msurf\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpygame\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtransform\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mflip\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msurf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 298\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mscreen\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mblit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msurf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    299\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mmode\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m\"human\"\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    300\u001b[0m             \u001b[0mpygame\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mevent\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpump\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31merror\u001b[0m: display Surface quit"
     ]
    }
   ],
   "source": [
    "_ = run_episode(render=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Hopefully, you can see that pole can now balance pretty well!\n",
    "\n",
    "## Actor-Critic Model\n",
    "\n",
    "Actor-Critic model is the further development of policy gradients, in which we build a neural network to learn both the policy and estimated rewards. The network will have two outputs (or you can view it as two separate networks):\n",
    "* **Actor** will recommend the action to take by giving us the state probability distribution, as in policy gradient model\n",
    "* **Critic** would estimate what the reward would be from those actions. It returns total estimated rewards in the future at the given state.\n",
    "\n",
    "Let's define such a model: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_inputs = 4\n",
    "num_actions = 2\n",
    "num_hidden = 128\n",
    "\n",
    "inputs = keras.layers.Input(shape=(num_inputs,))\n",
    "common = keras.layers.Dense(num_hidden, activation=\"relu\")(inputs)\n",
    "action = keras.layers.Dense(num_actions, activation=\"softmax\")(common)\n",
    "critic = keras.layers.Dense(1)(common)\n",
    "\n",
    "model = keras.Model(inputs=inputs, outputs=[action, critic])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We would need to slightly modify our `run_episode` function to return also critic results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_episode(max_steps_per_episode = 10000,render=False):    \n",
    "    states, actions, probs, rewards, critic = [],[],[],[],[]\n",
    "    state = env.reset()\n",
    "    for _ in range(max_steps_per_episode):\n",
    "        if render:\n",
    "            env.render()\n",
    "        action_probs, est_rew = model(np.expand_dims(state,0))\n",
    "        action = np.random.choice(num_actions, p=np.squeeze(action_probs[0]))\n",
    "        nstate, reward, done, info = env.step(action)\n",
    "        if done:\n",
    "            break\n",
    "        states.append(state)\n",
    "        actions.append(action)\n",
    "        probs.append(tf.math.log(action_probs[0,action]))\n",
    "        rewards.append(reward)\n",
    "        critic.append(est_rew[0,0])\n",
    "        state = nstate\n",
    "    return states, actions, probs, rewards, critic"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we will run the main training loop. We will use manual network training process by computing proper loss functions and updating network parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "running reward: 5.82 at episode 10\n",
      "running reward: 9.43 at episode 20\n",
      "running reward: 10.30 at episode 30\n",
      "running reward: 10.28 at episode 40\n",
      "running reward: 11.00 at episode 50\n",
      "running reward: 13.01 at episode 60\n",
      "running reward: 21.78 at episode 70\n",
      "running reward: 40.54 at episode 80\n",
      "running reward: 73.70 at episode 90\n",
      "running reward: 100.19 at episode 100\n",
      "running reward: 159.20 at episode 110\n",
      "Solved at episode 114!\n"
     ]
    }
   ],
   "source": [
    "optimizer = keras.optimizers.Adam(learning_rate=0.01)\n",
    "huber_loss = keras.losses.Huber()\n",
    "episode_count = 0\n",
    "running_reward = 0\n",
    "\n",
    "while True:  # Run until solved\n",
    "    state = env.reset()\n",
    "    episode_reward = 0\n",
    "    with tf.GradientTape() as tape:\n",
    "        _,_,action_probs, rewards, critic_values = run_episode()\n",
    "        episode_reward = np.sum(rewards)\n",
    "        \n",
    "        # Update running reward to check condition for solving\n",
    "        running_reward = 0.05 * episode_reward + (1 - 0.05) * running_reward\n",
    "\n",
    "        # Calculate discounted rewards that will be labels for our critic\n",
    "        dr = discounted_rewards(rewards)\n",
    "\n",
    "        # Calculating loss values to update our network\n",
    "        actor_losses = []\n",
    "        critic_losses = []\n",
    "        for log_prob, value, rew in zip(action_probs, critic_values, dr):\n",
    "            # When we took the action with probability `log_prob`, we received discounted reward of `rew`,\n",
    "            # while critic predicted it to be `value` \n",
    "            # First we calculate actor loss, to make actor predict actions that lead to higher rewards\n",
    "            diff = rew - value\n",
    "            actor_losses.append(-log_prob * diff)\n",
    "\n",
    "            # The critic loss is to minimize the difference between predicted reward `value` and actual\n",
    "            # discounted reward `rew`\n",
    "            critic_losses.append(\n",
    "                huber_loss(tf.expand_dims(value, 0), tf.expand_dims(rew, 0))\n",
    "            )\n",
    "\n",
    "        # Backpropagation\n",
    "        loss_value = sum(actor_losses) + sum(critic_losses)\n",
    "        grads = tape.gradient(loss_value, model.trainable_variables)\n",
    "        optimizer.apply_gradients(zip(grads, model.trainable_variables))\n",
    "\n",
    "    # Log details\n",
    "    episode_count += 1\n",
    "    if episode_count % 10 == 0:\n",
    "        template = \"running reward: {:.2f} at episode {}\"\n",
    "        print(template.format(running_reward, episode_count))\n",
    "\n",
    "    if running_reward > 195:  # Condition to consider the task solved\n",
    "        print(\"Solved at episode {}!\".format(episode_count))\n",
    "        break\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's run the episode and see how good our model is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = run_episode(render=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, let's close the environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Takeaway\n",
    "\n",
    "We have seen two RL algorithms in this demo: simple policy gradient, and more sophisticated actor-critic. You can see that those algorithms operate with abstract notions of state, action and reward - thus they can be applied to very different environments.\n",
    "\n",
    "Reinforcement learning allows us to learn the best strategy to solve the problem just by looking at the final reward. The fact that we do not need labelled datasets allows us to repeat simulations many times to optimize our models. However, there are still many challenges in RL, which you may learn if you decide to focus more on this interesting area of AI.   "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.4 64-bit",
   "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.10.4"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
