success_matching_for_cartpole.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6060d4d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gym\n",
    "import numpy as np\n",
    "from tqdm.notebook import tqdm\n",
    "from scipy.stats import ttest_ind \n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from optimization import Success_Matching, Parallel_Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c59816f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def cart_pole_v1_trial(env, params: np.array, tolerance=50, viz=True):\n",
    "    state = env.reset()\n",
    "    sum_reward = 0\n",
    "    for _ in range(env._max_episode_steps):\n",
    "        action = int(np.heaviside(np.array(state) @ params, 1))\n",
    "        state, reward, done, info = env.step(action)\n",
    "        sum_reward += reward\n",
    "        if done:\n",
    "            if tolerance > 0 and viz:\n",
    "                tolerance -= 1\n",
    "            else:\n",
    "                break\n",
    "        if viz: \n",
    "            env.render()\n",
    "    return sum_reward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f54f6293",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.9/site-packages/gym/envs/classic_control/cartpole.py:163: UserWarning: \u001b[33mWARN: You are calling 'step()' even though this environment has already returned done = True. You should always call 'reset()' once you receive 'done = True' -- any further steps are undefined behavior.\u001b[0m\n",
      "  logger.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "reward before training 10.0\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a6399d3c81a145fe8619a4b5a08dd859",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/10 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "final reward 500.0\n"
     ]
    }
   ],
   "source": [
    "# rollout functions\n",
    "reward_function = lambda env, params: cart_pole_v1_trial(env, params, viz=False)\n",
    "viz_function = lambda env, params: cart_pole_v1_trial(env, params, viz=True)\n",
    "\n",
    "# parallel env\n",
    "p_env = Parallel_Environment(\"CartPole-v1\", reward_function, viz_function)\n",
    "\n",
    "# optimizer\n",
    "n_states = p_env.observation_space_shape[0]\n",
    "sm = Success_Matching(n_states)\n",
    "\n",
    "# performance before training\n",
    "print(\"reward before training\", p_env.vizualize(sm.mean))\n",
    "\n",
    "\n",
    "# training\n",
    "best_params = sm.train(\n",
    "    p_env,\n",
    "    population_size=40, \n",
    "    epochs=10,\n",
    "    # viz=True\n",
    ")\n",
    "\n",
    "# performance after training\n",
    "print(\"final reward\", p_env.vizualize(sm.mean))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.1 64-bit",
   "language": "python",
   "name": "python39164bitaf7b9d806b6842969a72794a16bc5513"
  },
  "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}