V2GEvUser.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from matplotlib import pyplot as plt\n",
    "from IPython.display import clear_output\n",
    "import math\n",
    "from time import time\n",
    "from stable_baselines3.common.callbacks import BaseCallback\n",
    "from ev2gym.models.ev2gym_env import EV2Gym\n",
    "from ev2gym.rl_agent.state import V2G_profit_max_loads\n",
    "from stable_baselines3.sac.policies import SACPolicy\n",
    "from stable_baselines3.td3.policies import TD3Policy\n",
    "from sb3_contrib.tqc.policies import TQCPolicy\n",
    "from stable_baselines3.common.policies import ActorCriticPolicy\n",
    "from logger.evaluation_log import save_eval_log\n",
    "from logger.test_log import save_test_log"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "SEED = 42"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cpu\"\n",
    "if torch.cuda.is_available():\n",
    "    device = \"cuda\"\n",
    "\n",
    "print(\"Using device:\", device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "Total_timesteps = int(1 * 1e6)\n",
    "EVAL_FREQ = 2048\n",
    "TEST_EPISODES = 100\n",
    "TEST_ENV_COUNT = 5\n",
    "INIT_WEIGHTS = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# EV2Gym Setting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_file = \"./config_files/V2GProfitPlusLoads10.yaml\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def reward_function(env, total_costs, user_satisfaction_list, *args):\n",
    "    reward = total_costs\n",
    "\n",
    "    for ev in env.EVs:\n",
    "        score = ev.get_user_satisfaction()\n",
    "        d_cal, d_cyc = ev.get_battery_degradation()\n",
    "\n",
    "        reward -= 2000 * ( d_cal + 2 * d_cyc )\n",
    "        reward -= 15 * math.exp(-10 * score)\n",
    "\n",
    "    return reward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_reward_function(env, total_costs, user_satisfaction_list, *args):\n",
    "    reward = total_costs\n",
    "    \n",
    "    for tr in env.transformers:\n",
    "        reward -= 100 * tr.get_how_overloaded()                  \n",
    "    \n",
    "    for score in user_satisfaction_list:        \n",
    "        reward -= 100 * math.exp(-10*score)\n",
    "        \n",
    "    return reward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = EV2Gym(config_file, \n",
    "             render_mode=False,\n",
    "             seed=SEED,\n",
    "             save_plots=False,\n",
    "             state_function=V2G_profit_max_loads,\n",
    "             reward_function=reward_function,\n",
    "             save_replay=False)\n",
    "\n",
    "test_env = EV2Gym(config_file,\n",
    "                render_mode=False,\n",
    "                seed=SEED,\n",
    "                save_plots=False,\n",
    "                state_function=V2G_profit_max_loads,\n",
    "                reward_function=test_reward_function,\n",
    "                save_replay=False)\n",
    "\n",
    "N_CS = env.cs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env_only_render = EV2Gym(config_file, \n",
    "                        render_mode=True,\n",
    "                        seed=SEED,\n",
    "                        save_plots=False,\n",
    "                        state_function=V2G_profit_max_loads,\n",
    "                        reward_function=reward_function,\n",
    "                        save_replay=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PreLoad"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Policy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_weights(m):\n",
    "    if isinstance(m, torch.nn.Linear):\n",
    "        torch.nn.init.xavier_uniform_(m.weight)\n",
    "        if m.bias is not None:\n",
    "            torch.nn.init.constant_(m.bias, -1)\n",
    "\n",
    "class CustomActorCriticPolicy(ActorCriticPolicy):\n",
    "    def _build(self, lr_schedule):\n",
    "        super()._build(lr_schedule)\n",
    "        self.mlp_extractor.apply(init_weights)\n",
    "        self.action_net.apply(init_weights)\n",
    "        self.value_net.apply(init_weights)\n",
    "\n",
    "class CustomSACPolicy(SACPolicy):\n",
    "    def __init__(self, *args, **kwargs):\n",
    "        super(CustomSACPolicy, self).__init__(*args, **kwargs)\n",
    "        # 커스텀 초기화 적용\n",
    "        self.actor.latent_pi.apply(init_weights)\n",
    "        self.critic.apply(init_weights)\n",
    "        self.actor.mu.apply(init_weights)\n",
    "\n",
    "class CustomTD3Policy(TD3Policy):\n",
    "    def __init__(self, *args, **kwargs):\n",
    "        super(CustomTD3Policy, self).__init__(*args, **kwargs)\n",
    "        # 커스텀 초기화 적용\n",
    "        self.actor.apply(init_weights)\n",
    "        self.critic.apply(init_weights)\n",
    "\n",
    "class CustomTQCPolicy(TQCPolicy):\n",
    "    def __init__(self, *args, **kwargs):\n",
    "        super(CustomTQCPolicy, self).__init__(*args, **kwargs)\n",
    "        # 커스텀 초기화 적용\n",
    "        self.actor.latent_pi.apply(init_weights)\n",
    "        self.critic.apply(init_weights)\n",
    "        self.actor.mu.apply(init_weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def display_plot(rewards1=[], rewards2=[], title=\"Episode\", label1=\"Train\", label2=\"Test\", erase = True):\n",
    "    if erase : clear_output(wait=True)\n",
    "    plt.ylabel(f'{title} reward')\n",
    "    plt.xlabel(title)\n",
    "\n",
    "    if len(rewards1) > 0:\n",
    "        print(f\"Last {label1} reward: {rewards1[-1]:.2e}\", end=\" \")\n",
    "        plt.plot(rewards1, label=label1)\n",
    "\n",
    "    if len(rewards2) > 0:\n",
    "        print(f\"Last {label2} reward: {rewards2[-1]:.2e}\", end=\" \")\n",
    "        plt.plot(rewards2, label=label2)\n",
    "    \n",
    "    print()\n",
    "\n",
    "    plt.legend()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Callback"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "EVAL_DETERMINISTIC = False\n",
    "\n",
    "class SBCallback(BaseCallback):\n",
    "    def __init__(self,total_timesteps, eval_freq, verbose=0):\n",
    "        super(SBCallback, self).__init__(verbose)\n",
    "        self.eval_freq = eval_freq\n",
    "        self.eval_rewards = [0]\n",
    "        self.total_timesteps = total_timesteps\n",
    "\n",
    "    def _on_step(self):\n",
    "\n",
    "        self.eval_rewards[-1] += float(self.locals[\"rewards\"])\n",
    "\n",
    "        if self.num_timesteps % self.eval_freq == 1:\n",
    "            self.eval_rewards.append(0)\n",
    "            display_plot(self.get_data(), title=\"Evaluation\")\n",
    "\n",
    "        print(f\"Step: {self.num_timesteps}/{self.total_timesteps} {int(self.num_timesteps * 100 / self.total_timesteps)}%\", end=\"\\r\")\n",
    "        return True\n",
    "    \n",
    "    def get_data(self):\n",
    "        return self.eval_rewards[1:-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PlotCallback():\n",
    "    def __init__(self, eval_freq):\n",
    "        self.eval_freq = eval_freq\n",
    "        self.eval_rewards = [0]\n",
    "\n",
    "    def __call__(self, reward, num_timesteps):\n",
    "        self.eval_rewards[-1] += reward\n",
    "\n",
    "        if num_timesteps % self.eval_freq == 1:\n",
    "            self.eval_rewards.append(0)\n",
    "            display_plot(self.get_data(), title=\"Evaluation\")\n",
    "        \n",
    "    def get_data(self):\n",
    "        return self.eval_rewards[1:-1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(agent, test_env = test_env, name=\"\"):\n",
    "    test_episode_rewards = []\n",
    "    stats = []\n",
    "    execute_times = []\n",
    "    for i in range(TEST_EPISODES):\n",
    "        state, _ = test_env.reset(seed=SEED+1+i)\n",
    "        episode_reward = 0\n",
    "        done = False\n",
    "        start_time = time()\n",
    "        while not done:\n",
    "            action, _ = agent.predict(state, deterministic=True)\n",
    "            next_state, reward, done, _, info = test_env.step(action)\n",
    "            state = next_state\n",
    "            episode_reward += reward\n",
    "        execute_times.append(time()-start_time)\n",
    "        stats.append(info)\n",
    "        test_episode_rewards.append(episode_reward)\n",
    "    display_plot([], test_episode_rewards, title=\"Episode\", erase=False)\n",
    "\n",
    "    result = {\n",
    "        \"name\" : name,\n",
    "        \"num_episodes\" : len(stats),\n",
    "        \"profits\" : sum([stat['total_profits'] for stat in stats])/len(stats),\n",
    "        \"user_satisfaction\" : sum([stat['average_user_satisfaction'] for stat in stats])/len(stats) * 100,\n",
    "        \"energy_charged\" : sum([stat['total_energy_charged'] for stat in stats])/len(stats),\n",
    "        \"energy_discharged\" : sum([stat['total_energy_discharged'] for stat in stats])/len(stats),\n",
    "        \"transformer_overload\" : sum([stat['total_transformer_overload'] for stat in stats])/len(stats),\n",
    "        \"battery_degradation\" : sum([stat['battery_degradation'] for stat in stats]),\n",
    "        \"battery_degradation_calendar\" : sum([stat['battery_degradation_calendar'] for stat in stats]),\n",
    "        \"battery_degradation_cycling\" : sum([stat['battery_degradation_cycling'] for stat in stats]),\n",
    "        \"execution_time\" : sum(execute_times)/len(execute_times),\n",
    "        \"reward\" : sum([stat['total_reward'] for stat in stats])/len(stats)/1000,\n",
    "        \"test_episode_rewards\" : test_episode_rewards\n",
    "    }\n",
    "\n",
    "    return result\n",
    "\n",
    "def print_test_result(result):\n",
    "\n",
    "    print(\"=====================================================\")\n",
    "    print(f' Average stats for {result[\"name\"]} algorithm, {result[\"num_episodes\"]} episodes')\n",
    "\n",
    "    # Profits / Costs\n",
    "    print(\"Profits/ Costs(€): %.1f\" % result[\"profits\"])\n",
    "\n",
    "    # User satisfaction (%)\n",
    "    print(\"User satisfaction(%%): %d\" % int(result[\"user_satisfaction\"]))\n",
    "\n",
    "    # Energy Charged (kWh)\n",
    "    print(\"Energy Charged(kWh): %d\" % int(result[\"energy_charged\"]))\n",
    "\n",
    "    # Energy Discharged (kWh)\n",
    "    print(\"Energy Discharged(kWh): %d\" % int(result[\"energy_discharged\"]))\n",
    "\n",
    "    # Transformer Overload (kWh)\n",
    "    print(\"Transformer Overload(kWh): %d\" % int(result[\"transformer_overload\"]))\n",
    "\n",
    "    # Total Battery Capacity Loss (e-3)\n",
    "    print(\"Total Battery Capacity Loss(e-3): %.2f\" % result[\"battery_degradation\"])\n",
    "\n",
    "    # Total Battery Degradation Calender (e-3)\n",
    "    print(\"Total Battery Degradation Calender(e-3): %.2f\" % result[\"battery_degradation_calendar\"])\n",
    "\n",
    "    # Total Battery Degradation Cycle (e-3)\n",
    "    print(\"Total Battery Degradation Cycle(e-3): %.2f\" % result[\"battery_degradation_cycling\"])\n",
    "\n",
    "    # Execution Time of episodes (s)\n",
    "    print(\"Execution Time of episodes(s): %.2f\" % result[\"execution_time\"])\n",
    "\n",
    "    # Reward (e3)\n",
    "    print(\"Reward(e3): %.2f\" % result[\"reward\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_envs(agent, name=\"\"):\n",
    "    results = []\n",
    "    for i in range(TEST_ENV_COUNT):\n",
    "        test_env = EV2Gym(config_file, \n",
    "                           render_mode=False,\n",
    "                           seed=SEED+i,\n",
    "                           save_plots=False,\n",
    "                           state_function=V2G_profit_max_loads,\n",
    "                           reward_function=test_reward_function,\n",
    "                           save_replay=False)\n",
    "        print(f\"Testing environment {i+1}\")\n",
    "        result = test(agent, test_env, f\"{name} - {i+1}\")\n",
    "        results.append(result)\n",
    "        print_test_result(result)\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "algorithm_results = []"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Select Algorithms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "algorithms = [\n",
    "    \"PPO\",\n",
    "    \"PPO_CAPG\",\n",
    "    \"PPO_MAXA\",\n",
    "    \"PPO_TANH\",\n",
    "    \"SB3_TRPO\",\n",
    "    \"SB3_SAC\",\n",
    "    # \"SB3_TQC\"\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Algorithms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### PPO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "from algorithms.Origin_PPO import Origin_PPO as PPO\n",
    "\n",
    "ppo = PPO(\n",
    "    env=env,\n",
    "    actor_lr=0.00003,\n",
    "    critic_lr=0.00003,\n",
    "    gamma=0.99,\n",
    "    epochs=10,\n",
    "    eps_clip=0.2,\n",
    "    rollout_length=2048,\n",
    "    seed=SEED,\n",
    "    device=device,\n",
    "    init_weight=init_weights if INIT_WEIGHTS else None\n",
    ")\n",
    "\n",
    "run_ppo = \"PPO\" in algorithms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if run_ppo:\n",
    "    train_ppo_eval_callback = PlotCallback(EVAL_FREQ)\n",
    "\n",
    "    train_ppo = ppo.learn(\n",
    "        total_timestamp=Total_timesteps, \n",
    "        eval_callback=train_ppo_eval_callback)\n",
    "    \n",
    "    train_ppo_result = (train_ppo_eval_callback.get_data(), \"PPO_User\")\n",
    "    \n",
    "    algorithm_results.append(train_ppo_result)\n",
    "    save_eval_log(train_ppo_result, f\"{N_CS}/PPO_User\")\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if run_ppo: \n",
    "    test_ppo_result = test_envs(ppo, \"PPO_User\")\n",
    "    save_test_log(test_ppo_result, f\"{N_CS}/PPO_User\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### PPO-CAPG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "from algorithms.Origin_PPO_CAPG import Origin_PPO_CAPG as PPO_CAPG\n",
    "\n",
    "ppo_capg = PPO_CAPG(\n",
    "    env=env,\n",
    "    actor_lr=0.00003,\n",
    "    critic_lr=0.00003,\n",
    "    gamma=0.99,\n",
    "    epochs=10,\n",
    "    eps_clip=0.2,\n",
    "    rollout_length=2048,\n",
    "    seed=SEED,\n",
    "    device=device,\n",
    "    init_weight=init_weights if INIT_WEIGHTS else None\n",
    ")\n",
    "\n",
    "run_ppo_capg = \"PPO_CAPG\" in algorithms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if run_ppo_capg:\n",
    "\n",
    "    train_ppo_capg_eval_callback = PlotCallback(EVAL_FREQ)\n",
    "\n",
    "    train_ppo_capg = ppo_capg.learn(\n",
    "        total_timestamp=Total_timesteps, \n",
    "        eval_callback=train_ppo_capg_eval_callback\n",
    "    )\n",
    "\n",
    "    train_ppo_capg_result = (train_ppo_capg_eval_callback.get_data(), \"PPO_CAPG_User\")\n",
    "\n",
    "    algorithm_results.append(train_ppo_capg_result)\n",
    "    save_eval_log(train_ppo_capg_result, f\"{N_CS}/PPO_CAPG_User\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if run_ppo_capg: \n",
    "    test_ppo_capg_result = test_envs(ppo_capg, \"PPO_CAPG_User\")\n",
    "    save_test_log(test_ppo_capg_result, f\"{N_CS}/PPO_CAPG_User\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### PPO_MAXA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "from algorithms.Origin_PPO_MAXA import Origin_PPO_MAXA as PPO_MAXA\n",
    "\n",
    "ppo_maxa = PPO_MAXA(\n",
    "    env=env,\n",
    "    actor_lr=0.00003,\n",
    "    critic_lr=0.00003,\n",
    "    gamma=0.99,\n",
    "    epochs=10,\n",
    "    eps_clip=0.2,\n",
    "    rollout_length=2048,\n",
    "    seed=SEED,\n",
    "    device=device,\n",
    "    init_weight=init_weights if INIT_WEIGHTS else None\n",
    ")\n",
    "\n",
    "run_ppo_maxa = \"PPO_MAXA\" in algorithms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if run_ppo_maxa:\n",
    "\n",
    "    train_ppo_maxa_eval_callback = PlotCallback(EVAL_FREQ)\n",
    "\n",
    "    train_ppo_maxa = ppo_maxa.learn(\n",
    "        total_timestamp=Total_timesteps, \n",
    "        eval_callback=train_ppo_maxa_eval_callback\n",
    "    )\n",
    "\n",
    "    train_ppo_maxa_result = (train_ppo_maxa_eval_callback.get_data(), \"PPO_MAXA_User\")\n",
    "\n",
    "    algorithm_results.append(train_ppo_maxa_result)\n",
    "    save_eval_log(train_ppo_maxa_result, f\"{N_CS}/PPO_MAXA_User\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if run_ppo_maxa: \n",
    "    test_ppo_maxa_result = test_envs(ppo_maxa, \"PPO_MAXA_User\")\n",
    "    save_test_log(test_ppo_maxa_result, f\"{N_CS}/PPO_MAXA_User\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### PPO_TANH"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "from algorithms.Origin_PPO_TANH import Origin_PPO_TANH as PPO_TANH\n",
    "\n",
    "ppo_tanh = PPO_TANH(\n",
    "    env=env,\n",
    "    actor_lr=0.00003,\n",
    "    critic_lr=0.00003,\n",
    "    gamma=0.99,\n",
    "    epochs=10,\n",
    "    eps_clip=0.2,\n",
    "    rollout_length=2048,\n",
    "    seed=SEED,\n",
    "    device=device,\n",
    "    init_weight=init_weights if INIT_WEIGHTS else None\n",
    ")\n",
    "\n",
    "run_ppo_tanh = \"PPO_TANH\" in algorithms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if run_ppo_tanh:\n",
    "\n",
    "    train_ppo_tanh_eval_callback = PlotCallback(EVAL_FREQ)\n",
    "\n",
    "    train_ppo_tanh = ppo_tanh.learn(\n",
    "        total_timestamp=Total_timesteps, \n",
    "        eval_callback=train_ppo_tanh_eval_callback\n",
    "    )\n",
    "\n",
    "    train_ppo_tanh_result = (train_ppo_tanh_eval_callback.get_data(), \"PPO_TANH_User\")\n",
    "\n",
    "    algorithm_results.append(train_ppo_tanh_result)\n",
    "    save_eval_log(train_ppo_tanh_result, f\"{N_CS}/PPO_TANH_User\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if run_ppo_tanh:\n",
    "    test_ppo_tanh_result = test_envs(ppo_tanh, \"PPO_TANH_User\")\n",
    "    save_test_log(test_ppo_tanh_result, f\"{N_CS}/PPO_TANH_User\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### SB3 TRPO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sb3_contrib import TRPO as SB3TRPO\n",
    "\n",
    "sb3trpo = SB3TRPO(\n",
    "    CustomActorCriticPolicy if INIT_WEIGHTS else \"MlpPolicy\", \n",
    "    env, \n",
    "    verbose=0, \n",
    "    device=device, \n",
    "    seed=SEED)\n",
    "run_sb3trpo = \"SB3_TRPO\" in algorithms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if run_sb3trpo:\n",
    "    train_sb3trpo_callback = SBCallback(\n",
    "        total_timesteps=Total_timesteps,\n",
    "        eval_freq=EVAL_FREQ,\n",
    "    )\n",
    "\n",
    "    sb3trpo.learn(total_timesteps=Total_timesteps, callback=train_sb3trpo_callback)\n",
    "\n",
    "    train_sb3trpo_result = (train_sb3trpo_callback.get_data(), \"SB3_TRPO_User\")\n",
    "\n",
    "    algorithm_results.append(train_sb3trpo_result)\n",
    "    save_eval_log(train_sb3trpo_result, f\"{N_CS}/SB3_TRPO_User\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if run_sb3trpo:\n",
    "    test_sb3trpo_result = test_envs(sb3trpo, \"SB3_TRPO_User\")\n",
    "    save_test_log(test_sb3trpo_result, f\"{N_CS}/SB3_TRPO_User\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### SB3 SAC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "from stable_baselines3 import SAC as SB3SAC\n",
    "\n",
    "sb3sac = SB3SAC(\n",
    "                CustomSACPolicy if INIT_WEIGHTS else \"MlpPolicy\",\n",
    "                env, \n",
    "                verbose=0, \n",
    "                device=device, \n",
    "                seed=SEED, \n",
    "                train_freq=(2048, \"step\"),\n",
    "                gradient_steps=100,\n",
    "                target_update_interval=100,\n",
    "                )\n",
    "\n",
    "run_sb3sac = \"SB3_SAC\" in algorithms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if run_sb3sac:\n",
    "    train_sb3sac_callback = SBCallback(\n",
    "        total_timesteps=Total_timesteps,\n",
    "        eval_freq=EVAL_FREQ,\n",
    "    )\n",
    "\n",
    "    sb3sac.learn(total_timesteps=Total_timesteps, callback=train_sb3sac_callback)\n",
    "\n",
    "    train_sb3sac_result = (train_sb3sac_callback.get_data(), \"SB3_SAC_User\")\n",
    "\n",
    "    algorithm_results.append(train_sb3sac_result)\n",
    "    save_eval_log(train_sb3sac_result, f\"{N_CS}/SB3_SAC_User\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if run_sb3sac:\n",
    "    test_sb3sac_result = test_envs(sb3sac, \"SB3_SAC_User\")\n",
    "    save_test_log(test_sb3sac_result, f\"{N_CS}/SB3_SAC_User\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### SB3 TQC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sb3_contrib import TQC as SB3TQC\n",
    "\n",
    "sb3tqc = SB3TQC(\n",
    "                CustomTQCPolicy if INIT_WEIGHTS else \"MlpPolicy\",\n",
    "                env, \n",
    "                verbose=0, \n",
    "                device=device, \n",
    "                seed=SEED,\n",
    "                train_freq=(2048, \"step\"),\n",
    "                gradient_steps=25,\n",
    "                target_update_interval=100,\n",
    "                stats_window_size=1000\n",
    "                )\n",
    "\n",
    "run_sb3tqc = \"SB3_TQC\" in algorithms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "if run_sb3tqc:\n",
    "    train_sb3tqc_callback = SBCallback(\n",
    "        total_timesteps=Total_timesteps,\n",
    "        eval_freq=EVAL_FREQ,\n",
    "    )\n",
    "\n",
    "    sb3tqc.learn(total_timesteps=Total_timesteps, callback=train_sb3tqc_callback)\n",
    "\n",
    "    train_sb3tqc_result = (train_sb3tqc_callback.get_data(), \"SB3_TQC_User\")\n",
    "\n",
    "    algorithm_results.append(train_sb3tqc_result)\n",
    "    save_eval_log(train_sb3tqc_result, f\"{N_CS}/SB3_TQC_User\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "if run_sb3tqc:\n",
    "    test_sb3tqc_result = test_envs(sb3tqc, \"SB3_TQC_User\")\n",
    "    save_test_log(test_sb3tqc_result, f\"{N_CS}/SB3_TQC_User\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Draw All"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "ignore_list = []\n",
    "group_size = 10\n",
    "skip = 0\n",
    "\n",
    "def draw_rewards_graphs(rewards_list, title=\"Title\"):\n",
    "    plt.ylabel('reward')\n",
    "    plt.xlabel(title)\n",
    "\n",
    "    for rewards in rewards_list:\n",
    "        if rewards[1] in ignore_list: continue\n",
    "        values = rewards[0][skip:]\n",
    "        plt.plot([ sum(values[i*group_size : (i+1)*group_size])/group_size for i in range(int(len(values)/group_size))], label=rewards[1])\n",
    "\n",
    "    plt.legend()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "draw_rewards_graphs(algorithm_results, title=\"Train\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ev2gym",
   "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}