{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "nwaAZRu1NTiI" }, "source": [ "# Policy Gradient\n", "\n", "#### This version implements Policy Gradient with Keras to solve cartpole\n" ] }, { "cell_type": "code", "execution_count": 13, "metadata": { "id": "Nm5rvpUZNxDp" }, "outputs": [], "source": [ "# %%capture\n", "# !pip install gym==0.22\n", "# !pip install pygame\n", "# !apt install python-opengl\n", "# !apt install ffmpeg\n", "# !apt install xvfb\n", "# !pip install pyvirtualdisplay\n", "# !pip install pyglet==1.5.1" ] }, { "cell_type": "code", "execution_count": 14, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "LNXxxKojNTiL", "outputId": "c48489ab-d67f-448e-9362-746a4e6bcba2" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "2.9.2\n" ] }, { "data": { "text/plain": [ "" ] }, "execution_count": 14, "metadata": {}, "output_type": "execute_result" } ], "source": [ "import tensorflow as tf\n", "from tensorflow.keras import layers, Model, Input\n", "from tensorflow.keras.utils import to_categorical\n", "import tensorflow.keras.backend as K\n", "\n", "import gym\n", "from gym import spaces\n", "from gym.utils import seeding\n", "from gym import wrappers\n", "\n", "from tqdm.notebook import tqdm\n", "from collections import deque\n", "import numpy as np\n", "import random\n", "from matplotlib import pyplot as plt\n", "from sklearn.preprocessing import MinMaxScaler\n", "\n", "import io\n", "import base64\n", "from IPython.display import HTML, Video\n", "print(tf.__version__)\n", "\n", "# Virtual display\n", "from pyvirtualdisplay import Display\n", "\n", "virtual_display = Display(visible=0, size=(1400, 900))\n", "virtual_display.start()" ] }, { "cell_type": "code", "execution_count": 17, "metadata": { "id": "c84LoGsXNnJo" }, "outputs": [], "source": [ "# custom model to be able to run a custom loss with parameters\n", "class CustomModel(tf.keras.Model):\n", " def custom_loss(self,y, y_pred, d_returns):\n", " # print(\"y\", y.shape)\n", " # K.print_tensor(y)\n", " # print(\"y Pred\", y_pred.shape) \n", " # K.print_tensor(y_pred)\n", " # print(\"d_retur\", d_returns.shape) \n", " # K.print_tensor(d_returns)\n", " # crossentropy \n", " log_like = y * K.log(y_pred)\n", " # print(\"-log_like\", log_like.shape) \n", " # K.print_tensor(log_like)\n", " # print(\"-Log_lik * d_returns\")\n", " # K.print_tensor(-log_like * d_returns)\n", " # print(\"k_sum\")\n", " # K.print_tensor(K.sum(-log_like * d_returns ))\n", " return K.sum(-log_like * d_returns )\n", " \n", " def train_step(self, data):\n", " # Unpack the data. Its structure depends on your model and\n", " # on what you pass to `fit()`.\n", " if len(data) == 3:\n", " x, y, sample_weight = data\n", " else:\n", " sample_weight = None\n", " x, y = data\n", "\n", " # check if we passed the d_return\n", " if isinstance(x, tuple):\n", " x, d_return = x\n", "\n", " with tf.GradientTape() as tape:\n", " y_pred = self(x, training=True) # Forward pass\n", " # Compute the loss value.\n", " y = tf.cast(y, tf.float32)\n", " loss = self.custom_loss(y, y_pred, d_return)\n", "\n", " # Compute gradients\n", " trainable_vars = self.trainable_variables\n", " gradients = tape.gradient(loss, trainable_vars)\n", "\n", " # Update weights\n", " self.optimizer.apply_gradients(zip(gradients, trainable_vars))\n", "\n", " # Update the metrics.\n", " # Metrics are configured in `compile()`.\n", " self.compiled_metrics.update_state(y, y_pred, sample_weight=sample_weight)\n", "\n", " # Return a dict mapping metric names to current value.\n", " # Note that it will include the loss (tracked in self.metrics).\n", " return {m.name: m.result() for m in self.metrics}" ] }, { "cell_type": "code", "execution_count": 18, "metadata": { "id": "sF8L5d-GNnJp" }, "outputs": [], "source": [ "class Policy:\n", " def __init__(self, env=None, action_size=2):\n", "\n", " self.action_size = action_size\n", "\n", " # Hyperparameters\n", " self.gamma = 0.95 # Discount rate\n", "\n", " self.learning_rate = 1e-2\n", " \n", " # Construct DQN models\n", " self.env = env\n", " self.action_size = action_size\n", " self.action_space = [i for i in range(action_size)]\n", " print(\"action space\",self.action_space)\n", " # self.saved_log_probs = None\n", " self.model= self._build_model()\n", " self.model.summary()\n", "\n", " def _build_model(self):\n", " \n", " x = Input(shape=(4,), name='x_input')\n", " # y_true = Input( shape=(2,), name='y_true' )\n", " d_returns = Input(shape=[1], name='d_returns')\n", "\n", " l = layers.Dense(16, activation = 'relu')(x)\n", " l = layers.Dense(16, activation = 'relu')(l)\n", " y_pred = layers.Dense(self.action_size, activation = 'softmax', name='y_pred')(l)\n", " \n", " optimizer = tf.keras.optimizers.Adam(learning_rate=self.learning_rate)\n", "\n", " # model_train = Model( inputs=[x], outputs=[y_pred], name='train_only' )\n", " model_train = CustomModel( inputs=x, outputs=y_pred, name='train_only' )\n", " # model_predict = Model( inputs=x, outputs=y_pred, name='predict_only' )\n", " model_train.compile(loss=None, optimizer=optimizer, metrics = ['accuracy'])\n", " # model_train.compile(loss=None, optimizer=optimizer, metrics = ['accuracy'], run_eagerly = True)\n", "\n", " return model_train\n", "\n", "\n", " def act(self, state):\n", " # print(\"Act state\",state)\n", " probs = self.model.predict(np.array([state]), verbose=0)[0]\n", " # print(\"probs\",probs)\n", " action = np.random.choice(self.action_space, p=probs)\n", " # print(\"Action\",action)\n", " # return the action and the log of the probability \n", " # return action, np.log(probs[action])\n", " return action\n", "\n", "\n", " # this implements the reinforce \n", " def learn(self, n_training_episodes=None, max_t=None, print_every=100):\n", " # Help us to calculate the score during the training\n", " scores_deque = deque(maxlen=100)\n", " scores = []\n", " # Line 3 of pseudocode\n", " for i_episode in range(1, n_training_episodes+1):\n", " # saved_log_probs = []\n", " saved_actions = []\n", " saved_state = []\n", " rewards = []\n", " state = self.env.reset()\n", " # Line 4 of pseudocode\n", " for t in range(max_t):\n", " saved_state.append(state)\n", " action = self.act(state)\n", " # action, log_prob = self.act(state)\n", " # saved_log_probs.append(log_prob)\n", " saved_actions.append(action)\n", " state, reward, done, _ = self.env.step(action)\n", " rewards.append(reward)\n", " if done:\n", " break \n", " scores_deque.append(sum(rewards))\n", " scores.append(sum(rewards))\n", " \n", " # Line 6 of pseudocode: calculate the return\n", " returns = deque(maxlen=max_t) \n", " n_steps = len(rewards) \n", " # Compute the discounted returns at each timestep,\n", " # as \n", " # the sum of the gamma-discounted return at time t (G_t) + the reward at time t\n", " #\n", " # In O(N) time, where N is the number of time steps\n", " # (this definition of the discounted return G_t follows the definition of this quantity \n", " # shown at page 44 of Sutton&Barto 2017 2nd draft)\n", " # G_t = r_(t+1) + r_(t+2) + ...\n", " \n", " # Given this formulation, the returns at each timestep t can be computed \n", " # by re-using the computed future returns G_(t+1) to compute the current return G_t\n", " # G_t = r_(t+1) + gamma*G_(t+1)\n", " # G_(t-1) = r_t + gamma* G_t\n", " # (this follows a dynamic programming approach, with which we memorize solutions in order \n", " # to avoid computing them multiple times)\n", " \n", " # This is correct since the above is equivalent to (see also page 46 of Sutton&Barto 2017 2nd draft)\n", " # G_(t-1) = r_t + gamma*r_(t+1) + gamma*gamma*r_(t+2) + ...\n", " \n", " \n", " ## Given the above, we calculate the returns at timestep t as: \n", " # gamma[t] * return[t] + reward[t]\n", " #\n", " ## We compute this starting from the last timestep to the first, in order\n", " ## to employ the formula presented above and avoid redundant computations that would be needed \n", " ## if we were to do it from first to last.\n", " \n", " ## Hence, the queue \"returns\" will hold the returns in chronological order, from t=0 to t=n_steps\n", " ## thanks to the appendleft() function which allows to append to the position 0 in constant time O(1)\n", " ## a normal python list would instead require O(N) to do this.\n", " for t in range(n_steps)[::-1]:\n", " disc_return_t = (returns[0] if len(returns)>0 else 0)\n", " returns.appendleft( self.gamma*disc_return_t + rewards[t] ) \n", " \n", " ## standardization of the returns is employed to make training more stable\n", " eps = np.finfo(np.float32).eps.item()\n", " ## eps is the smallest representable float, which is \n", " # added to the standard deviation of the returns to avoid numerical instabilities \n", " returns = np.array(returns)\n", " returns = (returns - returns.mean()) / (returns.std() + eps)\n", " # self.saved_log_probs = saved_log_probs\n", " \n", " # Line 7:\n", " saved_state = np.array(saved_state)\n", " # print(\"Saved state\", saved_state, saved_state.shape)\n", " saved_actions = np.array(to_categorical(saved_actions, num_classes=self.action_size))\n", " # print(\"Saved actions\", saved_actions, saved_actions.shape)\n", " returns = returns.reshape(-1,1)\n", " # print(\"Returns\", returns, returns.shape)\n", " # this is the trick part, we send a tuple so the CustomModel is able to split the x and use \n", " # the returns inside to calculate the custom loss\n", " self.model.train_on_batch(x=(saved_state,returns), y=saved_actions)\n", "\n", " # policy_loss = []\n", " # for action, log_prob, disc_return in zip(saved_actions, saved_log_probs, returns):\n", " # policy_loss.append(-log_prob * disc_return)\n", " # policy_loss = torch.cat(policy_loss).sum()\n", " \n", " # # Line 8: PyTorch prefers gradient descent \n", " # optimizer.zero_grad()\n", " # policy_loss.backward()\n", " # optimizer.step()\n", " \n", " if i_episode % print_every == 0:\n", " print('Episode {}\\tAverage Score: {:.2f}'.format(i_episode, np.mean(scores_deque)))\n", " \n", " return scores\n", "\n", " #\n", " # Loads a saved model\n", " #https://medium.com/@Bloomore/how-to-write-a-custom-loss-function-with-additional-arguments-in-keras-5f193929f7a0\n", " #\n", " def load(self, name):\n", " self.model.load_weights(name)\n", "\n", " #\n", " # Saves parameters of a trained model\n", " #\n", " def save(self, name):\n", " self.model.save_weights(name)\n", "\n", " def play(self, state):\n", " return np.argmax(self.model.predict(np.array([state]), verbose=0)[0])" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "z64C2rO7NnJq", "outputId": "fd4c1942-c7b6-49af-cead-e6a48ee987d0" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "action space [0, 1]\n", "Model: \"train_only\"\n", "_________________________________________________________________\n", " Layer (type) Output Shape Param # \n", "=================================================================\n", " x_input (InputLayer) [(None, 4)] 0 \n", " \n", " dense_6 (Dense) (None, 16) 80 \n", " \n", " dense_7 (Dense) (None, 16) 272 \n", " \n", " y_pred (Dense) (None, 2) 34 \n", " \n", "=================================================================\n", "Total params: 386\n", "Trainable params: 386\n", "Non-trainable params: 0\n", "_________________________________________________________________\n", "Episode 100\tAverage Score: 66.31\n", "Episode 200\tAverage Score: 161.58\n", "Episode 300\tAverage Score: 282.58\n" ] } ], "source": [ "env = gym.make('CartPole-v1')\n", "\n", "model = Policy(env=env, action_size=2)\n", "# model.learn(total_steps=6_000)\n", "\n", "model.learn(n_training_episodes=1000, max_t=1000, print_every=100)\n", "env.close()\n" ] }, { "cell_type": "code", "execution_count": 11, "metadata": { "id": "7zS2PuLSNnJr" }, "outputs": [], "source": [ "model.save(\"./alt/policy_grad_cartpole.h5\")" ] }, { "cell_type": "code", "execution_count": 15, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "vgDMDokeNnJr", "outputId": "13a744e9-5119-4c5f-d681-39b0933fd661" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "action space [0, 1]\n", "Model: \"train_only\"\n", "_________________________________________________________________\n", " Layer (type) Output Shape Param # \n", "=================================================================\n", " x_input (InputLayer) [(None, 4)] 0 \n", " \n", " dense_4 (Dense) (None, 16) 80 \n", " \n", " dense_5 (Dense) (None, 16) 272 \n", " \n", " y_pred (Dense) (None, 2) 34 \n", " \n", "=================================================================\n", "Total params: 386\n", "Trainable params: 386\n", "Non-trainable params: 0\n", "_________________________________________________________________\n", "Total reward 189.0\n" ] } ], "source": [ "eval_env = gym.make('CartPole-v1')\n", "model = Policy(env=eval_env, action_size=2)\n", "model.load(\"./alt/policy_grad_cartpole.h5\")\n", "eval_env = wrappers.Monitor(eval_env, \"./alt/gym-results\", force=True)\n", "state = eval_env.reset()\n", "total_reward = 0\n", "for _ in range(1000):\n", " action = model.play(state)\n", " observation, reward, done, info = eval_env.step(action)\n", " total_reward +=reward\n", " state = observation\n", " if done: \n", " print(f\"Total reward {total_reward}\")\n", " break\n", "eval_env.close()" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "2HxBAnLQUoZW" }, "outputs": [], "source": [] } ], "metadata": { "accelerator": "GPU", "colab": { "provenance": [] }, "gpuClass": "standard", "kernelspec": { "display_name": "Python 3.9.16 ('rl3')", "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.9.16" }, "orig_nbformat": 4, "vscode": { "interpreter": { "hash": "9070e15ca35f8308b0c5d51e893fc04d77e428fe4d803a6d9ae4f68a65d8ce17" } } }, "nbformat": 4, "nbformat_minor": 0 }