Cartpole.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gym\n",
    "import numpy as np\n",
    "import random\n",
    "from matplotlib import pyplot as plt\n",
    "from collections import deque\n",
    "import xgboost as xgb\n",
    "from sklearn.model_selection import train_test_split, KFold, cross_val_score\n",
    "import itertools"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize class\n",
    "Class is initialized with 4 XGBRegressor models, one for each state feature, and one XGBClassifier model.\n",
    "Training data is collected in state-action pairs. \n",
    "Each XGBRegressor model is fit using the training data with one timestep offset, so for a given model current state feature is predicted by previous state-action.\n",
    "A custom reward function is used where selected state feature closest to zero is given highest reward - in this case the pole angle is used as selected state feature for the custom reward function.\n",
    "Action at current state is selected based on predicted state at timestep t+n with the highest reward.\n",
    "\n",
    "On a successfully completed episode data is passed to the XGBClassifier to be trained. A \"switch\" variable is used to gradually, after each successful epsidode, hand over the action selection to this second layer. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class XGBAgent(object):\n",
    "\n",
    "    def __init__(self, ma_training_data_size, mb_training_data_size, ma_mb_switch):\n",
    "\n",
    "        self.env = gym.make(\"CartPole-v0\")\n",
    "\n",
    "        self.action_space = self.env.action_space.n\n",
    "        self.observation_space_n = self.env.observation_space.shape[0]\n",
    "\n",
    "        self.model_a_params = {\n",
    "            'learning_rate':0.05\n",
    "            , 'colsample_bytree':1\n",
    "            , 'objective':'reg:linear'\n",
    "            , 'n_estimators':1000\n",
    "            , 'max_depth':10\n",
    "        }\n",
    "\n",
    "        self.model_a = [xgb.XGBRegressor(**self.model_a_params) \n",
    "                        for model in range(self.observation_space_n)]\n",
    "\n",
    "        self.model_b_params = {\n",
    "            'learning_rate':0.05\n",
    "            , 'colsample_bytree':0.7\n",
    "            , 'objective':'binary:logistic'\n",
    "            , 'n_estimators':1000\n",
    "            , 'max_depth':10\n",
    "        }\n",
    "\n",
    "        self.model_b = xgb.XGBClassifier(**self.model_b_params)\n",
    "\n",
    "        self.model_b_validation_score = 0\n",
    "\n",
    "        self.model_a_features = [\n",
    "            [0,1,2,3,4]\n",
    "            , [1,4]\n",
    "            , [2,3,4]\n",
    "            , [3,4]\n",
    "        ]\n",
    "        \n",
    "        self.model_a_data_columns = ['Level', 'Sid', 'Pid', 'State', 'Action']\n",
    "\n",
    "        self.rewarded_features = [[2, 1]]\n",
    "        self.action_index = 4\n",
    "        self.episode_index = 5\n",
    "\n",
    "        self.training_data = deque(maxlen=ma_training_data_size)\n",
    "        self.training_data_b = deque(maxlen=mb_training_data_size)\n",
    "        \n",
    "        self.steps = 0\n",
    "        self.episodes = 0\n",
    "\n",
    "        self.ma_mb_switch = ma_mb_switch\n",
    "\n",
    "    def fit_model_a(self):\n",
    "\n",
    "        for idx, m in enumerate(self.model_a):\n",
    "            \n",
    "            X, Y = [], []\n",
    "            for e in range(self.episodes):\n",
    "                \n",
    "                X_append = [x[self.model_a_features[idx]] for x in self.training_data if x[self.episode_index] == e][:-1]\n",
    "                for x in X_append: X.append(x)\n",
    "                \n",
    "                y_append = [y[idx] for y in self.training_data if y[self.episode_index] == e][1:]\n",
    "                for y in y_append: Y.append(y)\n",
    "            \n",
    "            X, Y = np.array(X), np.array(Y)\n",
    "\n",
    "            X_train, X_test, y_train, y_test = train_test_split(X, Y, test_size=0.25)\n",
    "\n",
    "            self.model_a[idx].fit(\n",
    "                X_train, y_train\n",
    "                , early_stopping_rounds=50\n",
    "                , eval_set=[(X_test, y_test)]\n",
    "                , eval_metric='rmse'\n",
    "                , verbose=False)\n",
    "\n",
    "    def fit_model_b(self):\n",
    "        \n",
    "        if self.model_b_validation_score >= 0.96: return\n",
    "\n",
    "        X, Y = [], []\n",
    "\n",
    "        for row in self.training_data_b:\n",
    "            X.append(row[:self.action_index])\n",
    "            Y.append(row[self.action_index])\n",
    "\n",
    "        X, Y = np.array(X), np.array(Y)\n",
    "\n",
    "        if len(self.training_data_b) < self.training_data_b.maxlen:\n",
    "            self.model_b.fit(X, Y)\n",
    "        else:\n",
    "            kfold = KFold(n_splits=5)\n",
    "            results = cross_val_score(self.model_b, X, Y, cv=kfold)\n",
    "            if results.mean() > self.model_b_validation_score:\n",
    "                self.model_b.fit(X, Y)\n",
    "                self.model_b_validation_score = results.mean()\n",
    "\n",
    "    def adjust_ma_mb_switch(self, ma_mb_switch_factor, ma_mb_switch_minimum, ma_mb_switch):\n",
    "        self.ma_mb_switch = max(self.ma_mb_switch * ma_mb_switch_factor, ma_mb_switch_minimum)\n",
    "\n",
    "    def transfer_data_to_model_b(self, mb_threshold_for_model_fit):\n",
    "        batch = list(itertools.islice(self.training_data, len(self.training_data)-mb_threshold_for_model_fit, len(self.training_data)))\n",
    "        for row in batch: self.training_data_b.append(row)\n",
    " \n",
    "    def receive_custom_reward(self, state):\n",
    "        reward = 0\n",
    "        for feature in self.rewarded_features:\n",
    "            reward += 1 - min(((state[feature[0]])**2)**0.5, 1)\n",
    "            reward *= feature[1]\n",
    "        return reward\n",
    "\n",
    "    def run_episode(self, test=False, ma_prediction_timesteps=2):\n",
    "\n",
    "        terminal_state = None\n",
    "        episode_steps = 0\n",
    "        state = self.env.reset()\n",
    "        \n",
    "        while True:\n",
    "            \n",
    "            action = random.randrange(self.action_space) \n",
    "\n",
    "            if test:\n",
    "\n",
    "                if self.ma_mb_switch > np.random.rand():\n",
    "                \n",
    "                    level = 0\n",
    "                    sid = 0\n",
    "                    pid = 0\n",
    "                    best_reward = 0\n",
    "\n",
    "                    data = [dict(zip(\n",
    "                                self.model_a_data_columns\n",
    "                                , [level, sid, None, state, None]))]\n",
    "\n",
    "                    while level < ma_prediction_timesteps:\n",
    "\n",
    "                        states = [row['State'] for row in data if row['Level'] == level]\n",
    "                        level += 1\n",
    "\n",
    "                        for s in states:\n",
    "                            for a in range(self.action_space):\n",
    "\n",
    "                                predicted_state = []\n",
    "\n",
    "                                for idx, m in enumerate(self.model_a):\n",
    "\n",
    "                                    x_state = np.append(np.array(s).copy(), a)\n",
    "                                    x_state = x_state[self.model_a_features[idx]]\n",
    "\n",
    "                                    predicted_state_feature = self.model_a[idx].predict([x_state])[0]\n",
    "                                    predicted_state.append(predicted_state_feature)\n",
    "\n",
    "                                sid += 1\n",
    "                                \n",
    "                                data.append(\n",
    "                                    dict(zip(\n",
    "                                        self.model_a_data_columns\n",
    "                                        , [level, sid, pid, np.array(predicted_state), a])))\n",
    "\n",
    "                            pid += 1\n",
    "                    \n",
    "                    leaf_states = [{'Pid': row['Pid'], 'State': row['State']}\n",
    "                                    for row in data if row['Level'] == level]\n",
    "                    \n",
    "                    rewards = np.array([self.receive_custom_reward(row['State']) \n",
    "                                for row in leaf_states])\n",
    "                    \n",
    "                    best_reward_idx = np.random.choice(np.flatnonzero(rewards == rewards.max()))\n",
    "\n",
    "                    pid_search = leaf_states[best_reward_idx]\n",
    "                    \n",
    "                    while level > 1:\n",
    "                        level -= 1\n",
    "                        pid_state = [row for row in data if row['Sid'] == pid_search['Pid']][0]\n",
    "                        pid_search = pid_state.copy()\n",
    "                    action = pid_state['Action']\n",
    "\n",
    "                else:\n",
    "\n",
    "                    action = np.int(self.model_b.predict([state])[0])\n",
    "            \n",
    "            next_state, _, done, _ = self.env.step(action)\n",
    "\n",
    "            if done:\n",
    "                terminal_state = state\n",
    "                break\n",
    "\n",
    "            episode_steps += 1\n",
    "\n",
    "            row = np.append(np.append(state, action), np.array([self.episodes]))\n",
    "\n",
    "            self.training_data.append(row)\n",
    "\n",
    "            state = next_state\n",
    "        \n",
    "        self.steps += episode_steps\n",
    "        self.episodes += 1\n",
    "\n",
    "        return episode_steps, terminal_state "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "MA_TRAINING_DATA_SIZE = 800\n",
    "MA_PREDICTION_TIMESTEPS = 2\n",
    "\n",
    "MB_TRAINING_DATA_SIZE = 2200\n",
    "MB_THRESHOLD_FOR_MODEL_FIT = 195\n",
    "\n",
    "MA_MB_SWITCH = 1\n",
    "MA_MB_SWITCH_MINIMUM = 0\n",
    "MA_MB_SWITCH_FACTOR = 0.8\n",
    "\n",
    "MAXIMUM_EPISODES = 500\n",
    "EPISODES_TO_SOLVE = 100"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run main"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXwAAAD8CAYAAAB0IB+mAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDMuMC4zLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvnQurowAAE0lJREFUeJzt3X2wZVV95vHvQ9MMglomciUiXltQMV2WKGmRl5RMozUxQTGxxEhFJjJJ9VSNCohlCjPOTOJMIqlIymgcTY/El8QYKV4qREScEdDIKEqDL0ADGmwUJSJOAshMAt3+5o+zb7gh3efuvr3XuX3Z30/VrXP2Pmef/Zzu5ncXa6+9VqoKSdKj3z4rHUCSNBsWfEkaCQu+JI2EBV+SRsKCL0kjYcGXpJGw4EvSSFjwJWkkLPiSNBL7rnSAxQ466KBat27dSseQpFVjy5Yt91TVXJ/37lUFf926dVx33XUrHUOSVo0kd/R9r106kjQSFnxJGgkLviSNhAVfkkbCgi9JI9F0lE6SbcD9wA5ge1VtaHk+SdKuzWJY5saqumcG55EkTWGXjiSNROuCX8Cnk2xJsqnxuSRJU7Tu0jm+qr6X5EnA/0xyS1V9bvEbul8EmwDm5+cbx5FWt3XnXNb8HNvOPan5OVaTR9OfedMWflV9r3u8G7gEOHon79lcVRuqasPcXK/pICRJy9Cs4Cc5MMnjFp4D/wa4sdX5JEnTtezSORi4JMnCef68qj7V8HySpCmaFfyquh04stXnS5J2j8MyJWkkLPiSNBIWfEkaCQu+JI2EBV+SRsKCL0kjYcGXpJGw4EvSSFjwJWkkLPiSNBIWfEkaCQu+JI2EBV+SRmK3Cn6SfZI8vlUYSVI7Sxb8JH+e5PHdIiY3A7cmeUv7aJKkIfVp4a+vqvuAXwQ+CcwDpzVNJUkaXJ+CvzbJWiYF/y+r6iGg2saSJA2tT8H/Y2AbcCDwuSRPA+5rGUqSNLwllzisqncD7160644kG9tFkiS10Oei7cFJzk9yebe9HvjV5skkSYPq06XzIeAK4JBu+zbgrFaBJElt9Cn4B1XVBcCPAapqO7CjaSpJ0uD6FPwHkjyRbmROkmOAe5umkiQNbsmLtsDZwKXA4UmuAeaAVzVNJUkaXJ9ROtcnOQE4AghwazcWX5K0iuyy4Cd55S5eelYSquriRpkkSQ1Ma+G/vHt8EnAccGW3vRG4GrDgS9IqssuCX1WnAyT5BJP5dO7qtp8MvHc28SRJQ+kzSmfdQrHvfB94VqM8kqRG+ozSuTrJFcDHmAzNfA1wVdNUkqTB9Rml84YkvwS8qNu1uaouaRtLkjS0Pi18gP8NbGfSwv9SuziSpFb6TJ72aiZF/lXAq4Frk/S+8SrJmiQ3dBd/JUkrpE8L/z8CL6iquwGSzAH/C7iw5znOBLYCroUrSSuozyidfRaKfeeHPY8jyaHAScAHlpFNkjSgPi38Ty0apQPwy0zWtu3jXcBvAI9bRjZJ0oD6jNJ5SzfNws8ymUun1yidJC8D7q6qLUn+9ZT3bQI2AczPz/fNLWnG1p1zWfNzbDv3pObnGLM+F20PZLJ4+dnA+4Ed3aLmSzkeODnJNuAvgBOT/Nkj31RVm6tqQ1VtmJub2730kqTe+vTFfw74V0mewuRi7elMVsGaqqreWlWHVtU6JjdrXVlVr92DrJKkPdCn4Keq/i/wSuA9VfVLwPq2sSRJQ+tV8JMcC/wKsNCJ1/eGLQCq6uqqetnuhpMkDadPwT8LeCtwSVXdlOQwnEtHkladPqN0Pgt8dtH27cAZLUNJkoY3bcWrd1XVWUn+im4B88Wq6uSmySRJg5rWwv/T7vGdswgiSWpr2opXW7rHzybZD3g2k5b+rVX14IzySZIGsmQffpKTmNxw9TdM7rR9epJ/X1WXtw4nSRpOn+GV5wEbq+qbAEkOZzI804IvSatIn2GZdy8U+87twN27erMkae/Up4V/U5JPAhcw6cM/BfhyN6EaVXVxw3ySpIH0Kfj7A98HTui2fwD8JPByJr8ALPiStAr0ufHq9FkEkSS11Wd65Gcl+UySG7vt5yZ5W/tokqQh9blo+z+YzKXzEEBVfY3JdMeSpFWkT8E/oKq+9Ih921uEkSS106fg39ONvS+AJK8C7mqaSpI0uD6jdF4PbAaeneS7wLeYzI0vSVpF+ozSuR14Sbe27T5VdX/7WJKkofVeuaqqHmgZRJLUVp8+fEnSo8AuC36SU7rHp88ujiSplWkt/Ld2jxfNIogkqa1pffg/THIVk/nvL33kiy5xKEmry7SCfxJwFJOlDs+bTRxJUivTljh8EPhikuOq6gdJHjfZXT+aXTxJ0lD6jNI5OMkNwI3AzUm2JHlO41ySpIH1KfibgbOr6mlVNQ+8udsnSVpF+hT8A6vqqoWNqroaOLBZIklSE33utL09yX9icvEW4LVM5tORJK0ifVr4/w6YY7KU4cXAQYCrYEnSKtNn8rS/A86YQRZJUkPOpSNJI2HBl6SRmFrwk6xJ8qZZhZEktTO14FfVDuAVy/ngJPsn+VKSrya5KclvLyuhJGkQfYZlXpPkj4CPA/+0CEpVXb/Ecf8InFhVP0qyFvh8ksur6ovLjytJWq4+Bf+47vHti/YVcOK0g6qqgIV5d9Z2P7W7ASVJw+gzLHPjcj88yRpgC/AM4L1Vde1yP0uStGeWLPhJDgZ+Fzikqn4+yXrg2Ko6f6lju2sAz0vyBOCSJM+pqhsf8fmbgE0A8/Pzy/kOK27dOZc1P8e2c09qfo7VZCX/zP371mrVZ1jmh4ArgEO67duAs3bnJFX198DVwEt38trmqtpQVRvm5uZ252MlSbuhT8E/qKouAH4MUFXbgR1LHZRkrmvZk+QxwEuAW/YgqyRpD/S5aPtAkifSXXBNcgxwb4/jngx8uOvH3we4oKo+seykkqQ90qfgnw1cChye5BomE6m9aqmDquprwPP3LJ4kaSh9Rulcn+QE4AggwK1V9VDzZJKkQfUZpbM/8B+An2XSrfPXSd5fVf/QOpwkaTh9unQ+AtwPvKfbPpXJYiintAolSRpen4J/RFUduWj7qiRfbRVIktRGn2GZN3QjcwBI8kLgmnaRJEkt7LKFn+TrTPrs1wL/Nsm3u5fmgZtnkE2SNKBpXTovm1kKSVJzuyz4VXXHwvMkPwE89RHvv+NfHCRJ2mv1GZb5X4HXAX/Dw9MbLzk9siRp79JnlM6rgcOr6sHWYSRJ7fQZpXMj8ITWQSRJbfVp4b+DydDMG5ksWwhAVZ3cLJUkaXB9Cv6Hgd8Dvk43RbIkafXpU/Dvqap3N08iSWqqT8HfkuQdTKZIXtylc32zVJKkwfUp+Atz2h+zaJ/DMiVplekzH/7GWQSRJLXV58ar/7yz/VX19uHjSJJa6bWm7aLn+zOZY2drmziSpFb6dOmct3g7yTuZXMCVJK0ife60faQDgMOGDiJJaqtPH/7CvPgAa4A5wP57SVpl+vThL54Xfzvw/ara3iiPJKmRJbt0unnx7wQeYtLCPyTJfOtgkqRh9enSeSPwX4Dv8/BcOgU8t2EuSdLA+nTpnAkcUVU/bB1GktROn1E63wHubR1EktRWnxb+7cDVSS7jn0+e9gfNUkmSBten4H+7+9mv+5EkrUJ97rT97VkEkSS1tZw7bSVJq5AFX5JGYmrBT7ImyZuW88FJnprkqiRbk9yU5MzlRZQkDWFqwa+qHcArlvnZ24E3V9VPM1kt6/VJ1i/zsyRJe6jPKJ1rkvwR8HEWzY2/1Jq2VXUXcFf3/P4kW4GnADcvP64kabn6FPzjusfFM2Tu1pq2SdYxWRv32r7HSJKG1XxN2ySPBS4Czqqq+3by+iZgE8D8/PLnZFt3zmXLPravbeee1Pwcu2slv/dY/8w1e/5bG8aSo3SSHJzk/CSXd9vrk/xanw9PspZJsf9oVV28s/dU1eaq2lBVG+bm5nYnuyRpN/QZlvkh4ArgkG77NuCspQ5KEuB8YKvTMEjSyutT8A+qqgvopkbuFj/Z0eO444HTgBOTfKX7+YXlR5Uk7Yk+F20fSPJEumUOkxxDj9kzq+rzQPYsniRpKH0K/tnApcDhSa5hsqbtKU1TSZIG16fg3wScABzBpMV+K07JIEmrTp/C/YWq2l5VN1XVjVX1EPCF1sEkScPaZQs/yU8xuTP2MUmez8P98Y8HDphBNknSgKZ16fwc8DrgUOA8Hi749wG/2TaWJGlo0wr+E6pqY5K3VdV/m1kiSVIT0/rwT+8eXzmLIJKktqa18Lcm2QbMJfnaov0Bqqqe2zSZJGlQuyz4VXVqd+H2CuDk2UWSJLUwdRx+Vf0tcOSMskiSGlryxqskzwTeAawH9l/YX1WHNcwlSRpYnxuvPgi8j8mShRuBjwB/2jKUJGl4fQr+Y6rqM0Cq6o6q+i12Y7UrSdLeoc9cOv+QZB/gG0neAHwXeFLbWJKkofVp4Z/FZCqFM4CfAV4L/GrLUJKk4fVZ0/bL3dMf8fDNWJKkVcZpjiVpJCz4kjQSFnxJGolp8+G/h24d252pqjOaJJIkNTGthX8dsIXJ3bVHAd/ofp4H7GgfTZI0pGmTp30YIMnrgI3d0oYkeT/w6ZmkkyQNpk8f/iHA4xZtP7bbJ0laRfrcaXsucEOSq7rtE4DfapZIktREnxuvPpjkcuCF3a5zummTJUmryJJdOkkCvAQ4sqr+EtgvydHNk0mSBtWnD/+/A8cCp3bb9wPvbZZIktREnz78F1bVUUluAKiqv0uyX+NckqSB9WnhP5RkDd1NWEnmgB83TSVJGlyfgv9u4BLgSUl+B/g88LtNU0mSBtdnlM5Hk2wBXgwE+MWq2to8mSRpUH0WMf9D4ONV5YVaSVrF+nTpXA+8Lck3k/x+kg2tQ0mShrdkwa+qD1fVLwBHA7cBv5fkG0sdl+RPktyd5MYBckqS9tDuzIf/DODZwDrglh7v/xDw0t2PJElqoc+dtgst+rcDNwE/U1UvX+q4qvoc8H/2PKIkaQh9brz6FnBsVd3TIkCSTcAmgPn5+RankCQxfcWro7qnXwLmk/yzalxV1w8RoKo2A5sBNmzYsMsVtiRJe2ZaC/+8Ka8VcOLAWSRJDU1b8WrjLINIktrqc9F2bZIzklzY/bwhydoex30M+AJwRJI7k/zaEIElScvT56Lt+4C1TKZJBjit2/fr0w6qqlOnvS5Jmq0+Bf8FVXXkou0rk3y1VSBJUht9brzakeTwhY0khwE72kWSJLXQp4X/FuCqJLczmS3zacDpTVNJkgbXZ3rkzyR5JnAEk4J/S1X9Y/NkkqRB7bJLJ8kLkvwUQFfgn8dkeoXfT/KTM8onSRrItD78PwYeBEjyIuBc4CPAvXR3xkqSVo9pXTprqmph8rNfBjZX1UXARUm+0j6aJGlI01r4a5Is/EJ4MXDlotf6XOyVJO1FphXujwGfTXIP8P+AvwZI8gwm3TqSpFVk2lw6v5PkM8CTgU9X1cJMlvsAb5xFOEnScKZ2zVTVF3ey77Z2cSRJrezOEoeSpFXMgi9JI2HBl6SRsOBL0khY8CVpJCz4kjQSFnxJGgkLviSNhAVfkkbCgi9JI2HBl6SRsOBL0khY8CVpJCz4kjQSFnxJGgkLviSNhAVfkkbCgi9JI2HBl6SRsOBL0khY8CVpJJoW/CQvTXJrkm8mOafluSRJ0zUr+EnWAO8Ffh5YD5yaZH2r80mSpmvZwj8a+GZV3V5VDwJ/Abyi4fkkSVO0LPhPAb6zaPvObp8kaQWkqtp8cHIK8HNV9evd9mnA0VX1xke8bxOwqds8Ari1SaB/6SDgnhmda2/i9x6fsX73sXzvp1XVXJ837tswxJ3AUxdtHwp875FvqqrNwOaGOXYqyXVVtWHW511pfu/xGet3H+v3nqZll86XgWcmeXqS/YDXAJc2PJ8kaYpmLfyq2p7kDcAVwBrgT6rqplbnkyRN17JLh6r6JPDJlufYAzPvRtpL+L3HZ6zffazfe5eaXbSVJO1dnFpBkkZilAV/jFM+JHlqkquSbE1yU5IzVzrTLCVZk+SGJJ9Y6SyzkuQJSS5Mckv3937sSmeahSRv6v6N35jkY0n2X+lMe4vRFfwRT/mwHXhzVf00cAzw+pF87wVnAltXOsSM/SHwqap6NnAkI/j+SZ4CnAFsqKrnMBkw8pqVTbX3GF3BZ6RTPlTVXVV1fff8fib/8Y/izuckhwInAR9Y6SyzkuTxwIuA8wGq6sGq+vuVTTUz+wKPSbIvcAA7uf9nrMZY8Ec/5UOSdcDzgWtXNsnMvAv4DeDHKx1khg4DfgB8sOvK+kCSA1c6VGtV9V3gncC3gbuAe6vq0yubau8xxoKfnewbzVClJI8FLgLOqqr7VjpPa0leBtxdVVtWOsuM7QscBbyvqp4PPAA86q9XJfkJJv/H/nTgEODAJK9d2VR7jzEW/F5TPjwaJVnLpNh/tKouXuk8M3I8cHKSbUy6705M8mcrG2km7gTurKqF/4u7kMkvgEe7lwDfqqofVNVDwMXAcSucaa8xxoI/yikfkoRJf+7WqvqDlc4zK1X11qo6tKrWMfm7vrKqHvUtvqr6W+A7SY7odr0YuHkFI83Kt4FjkhzQ/Zt/MSO4WN1X0ztt90YjnvLheOA04OtJvtLt+83ubmg9Or0R+GjXsLkdOH2F8zRXVdcmuRC4nsnItBvwjtt/4p22kjQSY+zSkaRRsuBL0khY8CVpJCz4kjQSFnxJGgkLviSNhAVfkkbCgi9JI/H/AaibBiTXtq02AAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Solved on average after 3.1 episodes\n"
     ]
    }
   ],
   "source": [
    "if __name__ == \"__main__\":\n",
    "\n",
    "    solved_envs = []\n",
    "    for _ in range(10):\n",
    "        \n",
    "        episode_list = []\n",
    "\n",
    "        xgba = XGBAgent(MA_TRAINING_DATA_SIZE, MB_TRAINING_DATA_SIZE, MA_MB_SWITCH)\n",
    "        \n",
    "        episode_steps = xgba.run_episode(test=False, ma_prediction_timesteps=0)\n",
    "\n",
    "        while True:\n",
    "\n",
    "            xgba.fit_model_a()\n",
    "            episode_steps, terminal_state = xgba.run_episode(test=True, ma_prediction_timesteps=MA_PREDICTION_TIMESTEPS)\n",
    "            episode_list.append(episode_steps)\n",
    "\n",
    "            if episode_steps >= MB_THRESHOLD_FOR_MODEL_FIT:\n",
    "                xgba.transfer_data_to_model_b(MB_THRESHOLD_FOR_MODEL_FIT)\n",
    "                xgba.adjust_ma_mb_switch(MA_MB_SWITCH_FACTOR, MA_MB_SWITCH_MINIMUM, MA_MB_SWITCH)\n",
    "                xgba.fit_model_b()\n",
    "            else:\n",
    "                pass\n",
    "                #print(xgba.episodes , episode_steps, xgba.ma_mb_switch, 'ts:', terminal_state)\n",
    "  \n",
    "            if len(episode_list) > EPISODES_TO_SOLVE:\n",
    "\n",
    "                if np.mean(episode_list[-EPISODES_TO_SOLVE:]) >= 195.0:\n",
    "                    \n",
    "                    #print ('Solved after', xgba.episodes-EPISODES_TO_SOLVE, 'episodes')\n",
    "                    solved_envs.append(xgba.episodes-EPISODES_TO_SOLVE)\n",
    "                    break\n",
    "\n",
    "            if xgba.episodes > MAXIMUM_EPISODES:\n",
    "                print('Not solved!')\n",
    "                solved_envs.append(MAXIMUM_EPISODES)\n",
    "                break\n",
    "    \n",
    "    plt.bar([idx for idx, x in enumerate(solved_envs)],solved_envs)\n",
    "    plt.ylabel('Solved after number of episodes')\n",
    "    plt.show()\n",
    "    \n",
    "    print('Solved on average after',np.mean(solved_envs),'episodes')\n",
    "    #xgba.env.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "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.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}