{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "import numpy as np\n",
    "import gym # == 0.26.2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ReplayBuffer(object):\n",
    "    def __init__(self, max_size, input_shape, n_actions, discrete = False):\n",
    "        self.mem_size = max_size\n",
    "        self.mem_cntr = 0\n",
    "        self.discrete = discrete\n",
    "        self.state_memory = np.zeros((self.mem_size, input_shape))\n",
    "        self.new_state_memory = np.zeros((self.mem_size, input_shape))\n",
    "        dtype = np.int8 if discrete else np.float32\n",
    "        self.action_memory = np.zeros((self.mem_size, n_actions), dtype=dtype)\n",
    "        self.reward_memory = np.zeros(self.mem_size)\n",
    "        self.terminal_memory = np.zeros(self.mem_size, dtype=np.float32)\n",
    "\n",
    "    def store_transition(self, state, action, reward, state_, done):\n",
    "        index = self.mem_cntr % self.mem_size\n",
    "        self.state_memory[index] = state\n",
    "        self.new_state_memory[index] = state_\n",
    "        self.reward_memory[index] = reward\n",
    "        self.terminal_memory[index] = 1 - int(done)\n",
    "        if self.discrete:\n",
    "            actions = np.zeros(self.action_memory.shape[1])\n",
    "            actions[action] = 1.0\n",
    "            self.action_memory[index] = actions\n",
    "        else:\n",
    "            self.action_memory[index] = action\n",
    "        self.mem_cntr += 1\n",
    "\n",
    "    def sample_buffer(self, batch_size):\n",
    "        max_mem = min(self.mem_cntr, self.mem_size)\n",
    "        batch = np.random.choice(max_mem, batch_size)\n",
    "        states = self.state_memory[batch]\n",
    "        states_ = self.new_state_memory[batch]\n",
    "        rewards = self.reward_memory[batch]\n",
    "        actions = self.action_memory[batch]\n",
    "        terminal = self.terminal_memory[batch]\n",
    "\n",
    "        return states, actions, rewards, states_, terminal\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_dqn(lr, n_actions, input_dims, fc1_dims, fc2_dims):\n",
    "    model = tf.keras.Sequential([\n",
    "        tf.keras.layers.Dense(fc1_dims, activation=tf.keras.activations.relu, input_shape=(input_dims,)),\n",
    "        tf.keras.layers.Dense(fc2_dims, activation=tf.keras.activations.relu),\n",
    "        tf.keras.layers.Dense(n_actions)\n",
    "    ])\n",
    "    model.compile(optimizer=tf.keras.optimizers.legacy.Adam(learning_rate=lr), loss=\"mse\")\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Agent(object):\n",
    "    def __init__(self, alpha, gamma, n_actions, epsilon, batch_size, input_dims, epsilon_dec=0.98, epsilon_end=0.01, mem_size=1000000, fname=\"dqn_model.keras\"):\n",
    "        self.action_space = [i for i in range(n_actions)]\n",
    "        self.n_actions = n_actions\n",
    "        self.gamma = gamma\n",
    "        self.epsilon = epsilon\n",
    "        self.epsilon_dec = epsilon_dec\n",
    "        self.epsilon_min = epsilon_end\n",
    "        self.batch_size = batch_size\n",
    "        self.model_file = fname\n",
    "\n",
    "        self.memory = ReplayBuffer(mem_size, input_dims, n_actions, discrete=True)\n",
    "        self.q_eval = build_dqn(alpha, n_actions, input_dims, 256, 256)\n",
    "\n",
    "    def remember(self, state, action, reward, new_state, done):\n",
    "        self.memory.store_transition(state, action, reward, new_state, done)\n",
    "\n",
    "    def choose_action(self, state):\n",
    "        state = state[np.newaxis, :]\n",
    "        rand = np.random.random()\n",
    "        if rand < self.epsilon:\n",
    "            action = np.random.choice(self.action_space)\n",
    "        else:\n",
    "            actions = self.q_eval.predict(state, verbose=0)\n",
    "            action = np.argmax(actions)\n",
    "\n",
    "        return action\n",
    "    \n",
    "    def learn(self):\n",
    "        if self.memory.mem_cntr < self.batch_size:\n",
    "            return\n",
    "        state, action, reward, new_state, done = self.memory.sample_buffer(self.batch_size)\n",
    "\n",
    "        action_values = np.array(self.action_space, dtype=np.int8)\n",
    "        action_indices = np.dot(action, action_values)\n",
    "\n",
    "\n",
    "        q_eval = self.q_eval.predict(state, verbose=0)\n",
    "        q_next = self.q_eval.predict(new_state, verbose=0)\n",
    "\n",
    "        q_target = q_eval.copy()\n",
    "\n",
    "        batch_index = np.arange(self.batch_size, dtype=np.int32)\n",
    "\n",
    "        q_target[batch_index, action_indices] = reward + self.gamma*np.max(q_next, axis=1)*done\n",
    "\n",
    "        _ = self.q_eval.fit(state, q_target, verbose=0)\n",
    "\n",
    "        self.epsilon = self.epsilon * self.epsilon_dec if self.epsilon > self.epsilon_min else self.epsilon_min\n",
    "\n",
    "    def save_model(self):\n",
    "        self.q_eval.save(self.model_file)\n",
    "    \n",
    "    def load_model(self):\n",
    "        self.q_eval = tf.keras.models.load_model(self.model_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.make('CartPole-v1')\n",
    "n_games = 50\n",
    "agent = Agent(gamma=0.99, epsilon=1.0, alpha=0.0005, input_dims=4, n_actions=2, mem_size=1000000, batch_size=64, epsilon_end=0.01)\n",
    "\n",
    "scores = []\n",
    "eps_history = []\n",
    "\n",
    "for i in range(n_games):\n",
    "    done = False\n",
    "    score = 0\n",
    "    observation, info = env.reset()\n",
    "    while not done:\n",
    "        action = agent.choose_action(observation)\n",
    "        observation_, reward, terminated, truncated, info = env.step(action)\n",
    "        done = terminated or truncated\n",
    "        score += reward\n",
    "        agent.remember(observation, action, reward, observation_, done)\n",
    "        observation = observation_\n",
    "        agent.learn()\n",
    "\n",
    "    eps_history.append(agent.epsilon)\n",
    "    scores.append(score)\n",
    "\n",
    "    avg_score = np.mean(scores[max(0, i-100):(i+1)])\n",
    "    print(f\"episode={i}, {score=:0.2f}, {avg_score=:0.2f}, {terminated=}, {truncated=}\")\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Aufgabe 3\n",
    "\n",
    "Ob _Exploration_ oder _Exploitation_ stärker gewichtet ist, bestimmt bei uns der Faktor _epsilon_.\n",
    "Dieser beginnt bei 1, und wird bei jedem Training mit `epsilon_dec` multipliziert (und somit reduziert, bis\n",
    "zu einem Minimalwert). Dadurch wird unser Agent anfangs Aktionen öfters Aktionen wagen, die nicht\n",
    "optimal sind, aber dafür neue Zustände offenbaren. Durch Ausprobieren haben wir uns für einen Wert von `epsilon_dec = 0.98`\n",
    "entschieden, da dieser gute Resultate lieferte."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.plot(scores)"
   ]
  }
 ],
 "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.10.1"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "bd385fe162c5ca0c84973b7dd5c518456272446b2b64e67c2a69f949ca7a1754"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
