t3_sac_main.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z8saNlGCvwVJ",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "<h1><div style=\"text-align: center;\"> Reinforcement Learning </div></h1>\n",
    "<h2><div style=\"text-align: center;\"> Assigment T3: Soft Actor Critic - Value and Policy Combined Method\n",
    "</div></h2>\n",
    "<h2><div style=\"text-align: center;\"> Continuous Environments: Lunar_lander-v2  and\n",
    "Bipedal-walker-v3\n",
    "</div></h2>\n",
    "<h3><div style=\"text-align: center;\"> A Ascencio-Cabral\n",
    "</div></h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 24651,
     "status": "ok",
     "timestamp": 1617950006612,
     "user": {
      "displayName": "Aze Gln",
      "photoUrl": "",
      "userId": "12590864696663313384"
     },
     "user_tz": -60
    },
    "id": "NUYUeJzhwDpB",
    "outputId": "5af985bc-ca35-4326-dc33-3332e054bb67",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# comment out if running in colab\n",
    "# from google.colab import drive\n",
    "# drive.mount('/content/drive/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 27126,
     "status": "ok",
     "timestamp": 1617950009093,
     "user": {
      "displayName": "Aze Gln",
      "photoUrl": "",
      "userId": "12590864696663313384"
     },
     "user_tz": -60
    },
    "id": "2ss-PZT8wJ9I",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# change to your drive directory\n",
    "# %cd /content/drive/MyDrive/INM707/task_3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 30329,
     "status": "ok",
     "timestamp": 1617950012300,
     "user": {
      "displayName": "Aze Gln",
      "photoUrl": "",
      "userId": "12590864696663313384"
     },
     "user_tz": -60
    },
    "id": "zAWJA5u8zLMu",
    "outputId": "4a3f6018-7533-45d8-ee1c-530921388124",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# !pip install gym[all]\n",
    "# !pip install gym[box2d]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "<h2><div>Libraries</div></h2>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 33367,
     "status": "ok",
     "timestamp": 1617950015343,
     "user": {
      "displayName": "Aze Gln",
      "photoUrl": "",
      "userId": "12590864696663313384"
     },
     "user_tz": -60
    },
    "id": "F7wCp4HPvwVP",
    "outputId": "472a1e8f-dfa7-431b-fbdf-cfef311a3d3c",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import gym\n",
    "import numpy as np\n",
    "import time\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "from utils.sac import SoftActorCritic\n",
    "from main import warmup\n",
    "from utils.plotting import  Stats, plot_episode_stats\n",
    "from gym.wrappers.monitoring.video_recorder import VideoRecorder\n",
    "plt.style.use('ggplot')\n",
    "device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n",
    " "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EjUndByDvwVR",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "For this task I will implement a soft actor critic agent to run in the gym continous environments:\n",
    "\n",
    "-  LunarLanderContinuous-v2\n",
    "-  BipedalWalker-v3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gODt7BdXvwVR",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "<h2><div>1. Explore the environments</div></h2>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "<h3><div> 1.1 Lunar Lander observation and action space</div></h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 33359,
     "status": "ok",
     "timestamp": 1617950015345,
     "user": {
      "displayName": "Aze Gln",
      "photoUrl": "",
      "userId": "12590864696663313384"
     },
     "user_tz": -60
    },
    "id": "xgxvJRSovwVR",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "lunar_lander = gym.make('LunarLander-v2', continuous=True)\n",
    "bipedal = gym.make('BipedalWalker-v3')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 33352,
     "status": "ok",
     "timestamp": 1617950015345,
     "user": {
      "displayName": "Aze Gln",
      "photoUrl": "",
      "userId": "12590864696663313384"
     },
     "user_tz": -60
    },
    "id": "Cj2qxW_BvwVS",
    "outputId": "433a1efa-460b-49c2-ad2c-02143419529a",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "print(\"Lunar lander\")\n",
    "print('Action space: {}'.format(lunar_lander.action_space.shape[0]))\n",
    "print('Low threshold: {}, High threshold: {}'.format(lunar_lander.action_space.low,\n",
    "                                                     lunar_lander.action_space.high))\n",
    "print('Observation space: {}'.format(lunar_lander.observation_space.shape))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GP6OMMZBvwVU",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "**Lunar lander**\n",
    "\n",
    "The observation space contains 8 values:\n",
    "\n",
    "- x coodinate\n",
    "- y coordinate\n",
    "- velocity x \n",
    "- velocity y \n",
    "- angle orientation on the space\n",
    "- angular velocity, \n",
    "- Left leg touching the ground: 0 else 1 \n",
    "- Right leg touching the ground: 0 else 1\n",
    "\n",
    "The action is two floats:\n",
    "- main engine \n",
    "- left-right engines"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "<h3><div> 1.2 BipedalWalker-v3 observation and action space </div></h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "print(\"Bipedal\")\n",
    "print('Action space: {}'.format(bipedal.action_space.shape[0]))\n",
    "print('Observation space: {}'.format(bipedal.observation_space.shape))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sT2_0KGuvwVU",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "**Bipedal**\n",
    "\n",
    "The observation space contains 24 values:\n",
    "\n",
    "0. hull_angle\t0\t2*pi\t0.5\n",
    "1. hull_angularVelocity\t-inf\t+inf\n",
    "2. vel_x\t-1\t+1\n",
    "3. vel_y\t-1\t+1\n",
    "4. hip_joint_1_angle\t-inf\t+inf\n",
    "5. hip_joint_1_speed\t-inf\t+inf\n",
    "6. knee_joint_1_angle\t-inf\t+inf\n",
    "7. knee_joint_1_speed\t-inf\t+inf\n",
    "8. leg_1_ground_contact_flag\t0\t1\n",
    "9. hip_joint_2_angle\t-inf\t+inf\n",
    "10. hip_joint_2_speed\t-inf\t+inf\n",
    "11. knee_joint_2_angle\t-inf\t+inf\n",
    "12. knee_joint_2_speed\t-inf\t+inf\n",
    "13. leg_2_ground_contact_flag\t0\t1\n",
    "\n",
    "\n",
    "**14-23**  10 lidar readings\t-inf\t+inf\n",
    "\n",
    "\n",
    "The action space contains 4 values within the range -1 to +1:\n",
    "\n",
    "0.\t Hip_1 (Torque / Velocity)\t-1\t+1\n",
    "1.\t Knee_1 (Torque / Velocity)\t-1\t+1\n",
    "2.\t Hip_2 (Torque / Velocity)\t-1\t+1\n",
    "3.\t Knee_2 (Torque / Velocity)\t-1\t+1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "<h2><div> Section 2. Building the training and evaluation functions</div></h2>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8u3Ina2ovwVV",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "<h3><div>2.1 Training SAC function </div></h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 33344,
     "status": "ok",
     "timestamp": 1617950015346,
     "user": {
      "displayName": "Aze Gln",
      "photoUrl": "",
      "userId": "12590864696663313384"
     },
     "user_tz": -60
    },
    "id": "qsXauTs0vwVW",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def run_sac(agent_args, gym_env, episodes=50, exploration=10000, max_steps=None):\n",
    "    \"\"\"\n",
    "    :param agent_args: arguments to build the sac agent\n",
    "    :param gym_env: a gym environment\n",
    "    :param episodes: an integer indicating the number of training episodes\n",
    "    :param exploration: step tht the agent will take before training to feed the memory\n",
    "    :param max_steps: an integer with maximum number of steps to train per episode\n",
    "     or not stored weigths\n",
    "    :return: namedtuple containing the training stats and policy\n",
    "    \"\"\"\n",
    "\n",
    "    agent = SoftActorCritic(agent_args)\n",
    "\n",
    "    assert exploration > agent.batch_size, print('minimum required warmup: '\n",
    "                                              '{}'.format(agent.batch_size))\n",
    "\n",
    "    # Start warmup\n",
    "    warmup(exploration, sac_agent=agent, environment=gym_env)\n",
    "\n",
    "    # Logging stats\n",
    "    rewards = np.zeros(episodes)\n",
    "    loss_log = np.zeros(episodes)\n",
    "    steps_log = np.zeros(episodes)\n",
    "    returns_list = list()\n",
    "\n",
    "    best_return = gym_env.reward_range[0]\n",
    "\n",
    "    print('........Start training.........')\n",
    "\n",
    "    init_time = time.time()\n",
    "    for episode in range(episodes):\n",
    "        observation, _ = gym_env.reset()\n",
    "        returns = 0\n",
    "        total_loss = 0\n",
    "        terminated = False\n",
    "        steps = 0\n",
    "\n",
    "        while not terminated:\n",
    "\n",
    "            action = agent.act(observation)\n",
    "\n",
    "            next_observation, reward, terminated, truncated, _ = gym_env.step(action)\n",
    "\n",
    "            agent.store_memory(observation, action, next_observation, reward, terminated)\n",
    "\n",
    "            returns += reward\n",
    "            steps += 1\n",
    "\n",
    "            loss = agent.update_parameters()\n",
    "\n",
    "            total_loss += loss\n",
    "\n",
    "            returns_list.append(reward)\n",
    "\n",
    "            if max_steps is not None and steps == max_steps:\n",
    "                break\n",
    "\n",
    "            observation = next_observation\n",
    "\n",
    "        loss_log[episode] += total_loss\n",
    "        steps_log[episode] += steps\n",
    "        rewards[episode] += returns\n",
    "\n",
    "        avg_return = np.mean(rewards[-10:])\n",
    "        if avg_return > best_return:\n",
    "            best_return = avg_return\n",
    "            agent.save_check_point(episode)\n",
    "\n",
    "        print('Episode: {}/{} | Steps: {} | Reward: {:.2f} | '\n",
    "              'Loss: {:.3f} | Cumulative Steps: {}'.format(episode + 1,\n",
    "                episodes, steps, rewards[episode], loss_log[episode],  int(np.sum(steps_log))))\n",
    "\n",
    "    print('Training completed, total steps: {}, total training time: {}'\n",
    "          .format(int(np.sum(steps_log)), time.time() - init_time))\n",
    "    stats = Stats(length_episodes = steps_log, reward_episodes= rewards,\n",
    "                             episode_loss= loss_log)\n",
    "\n",
    "    torch.save({'episode': episodes, 'model': agent.policy.state_dict(),\n",
    "                    'optim': agent.policy_optim.state_dict()},\n",
    "               os.path.join(agent.path, 'policy_end_{}.pth'.format(episodes)))\n",
    "\n",
    "    return stats, agent.policy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "<h3><div> 2.2 Evaluate function - Optimization </div></h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def evaluate_sac(agent_policy, gym_env, episodes=50, video_path=None):\n",
    "    \"\"\"\n",
    "    :param agent_policy: trained agent algorithm\n",
    "    :param gym_env: a gym environment\n",
    "    :param episodes: integer with the number of episodes to train for\n",
    "    :param video_path: path to save videos\n",
    "    \"\"\"\n",
    "    agent_policy.eval()\n",
    "\n",
    "    for episode in range(episodes):\n",
    "\n",
    "        observation, _ = gym_env.reset()\n",
    "        terminated = False\n",
    "        ep_reward = 0\n",
    "        steps = 0\n",
    "\n",
    "        video = VideoRecorder(gym_env, enabled=True, metadata={'step_id': episode +1},\n",
    "                              path=f'{video_path}_ep_{episode + 1}.mp4')\n",
    "\n",
    "        while not terminated:\n",
    "\n",
    "            video.capture_frame()\n",
    "            observation = torch.from_numpy(observation).unsqueeze(dim=0).to(device)\n",
    "            _, action = agent_policy.sample_action(observation, reparam=False)\n",
    "            action = action.cpu().squeeze().numpy()\n",
    "            next_observation, reward, terminated, _, _ = gym_env.step(action)\n",
    "            ep_reward += reward\n",
    "            steps += 1\n",
    "            observation = next_observation\n",
    "        video.close()\n",
    "        video.enabled = False\n",
    "        gym_env.close()\n",
    "\n",
    "        print('Episode: {}/{} | Steps: {} | Reward: {:.2f} '.format(episode + 1, episodes,\n",
    "                                                                    steps, ep_reward))\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "\n",
    "<h2><div> 3. Start training </div></h2>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "<h3><div> 3.1 Build  environments and  directories </div></h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "envs = [lunar_lander, bipedal]\n",
    "env = envs[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Create directory for weights \n",
    "params_dir = list()\n",
    "for s in ['lunar', 'bipedal']:\n",
    "    param_dir = os.path.join(os.getcwd(), 'params', s)\n",
    "    os.makedirs(param_dir, exist_ok=True)\n",
    "    params_dir.append(param_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "<h3><div> 3.2 Agent paramenters </div></h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "agent_params = {'gamma': 0.99, 'state_dim': env.observation_space.shape[0],\n",
    "                'actions_dim': env.action_space.shape[0], 'min': -20, 'max': 2,\n",
    "                'q_lr': 0.0003, 'lr': 0.0003, 'hidden_size': [256, 256, 256],\n",
    "                'batch_size': 256, 'memory_size': 1000000, 'r_scale': 1,\n",
    "                'tau': 0.02, 'env': env, 'path': params_dir[1]}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "<h3><div> 3.3 Start training </div></h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "stats_log, policy = run_sac(agent_params, env, episodes=1000, exploration=10000, max_steps=1600)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "<h2><div> 4. Visualise results </div></h2>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "plot_episode_stats(stats_log, lr=agent_params['lr'], episodes=1000,\n",
    "                   smoothing_window=10, hideplot=False, name='bipedal')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "<h2><div> 5. Evaluation </div></h2>\n",
    "<h3><div> 5.1 Render environment to record evaluation </div></h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "lunar_ = gym.make('LunarLander-v2', continuous=True, render_mode='rgb_array')\n",
    "bipedal_ = gym.make('BipedalWalker-v3', render_mode='rgb_array')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "paths = list()\n",
    "for v in ['lunar', 'bipedal']:\n",
    "    path = os.path.join(os.getcwd(), 'videos', v)\n",
    "    os.makedirs(path, exist_ok=True)\n",
    "    paths.append(path)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "<h3><div> 5.2 Start Evalaution </div></h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Create path to videos\n",
    "path_ = paths[1]\n",
    "evaluate_sac(policy, bipedal_, episodes=2, video_path=path_)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "machine_shape": "hm",
   "name": "sac_main.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}