getting_started/11_stable_baselines3_integration.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "90b9341f",
   "metadata": {},
   "source": [
    "# Grid2Op integration with stable baselines3 framework\n",
    "\n",
    "Try me out interactively with: [![Binder](./img/badge_logo.svg)](https://mybinder.org/v2/gh/Grid2Op/grid2op/master)\n",
    "\n",
    "\n",
    "**objectives** This notebooks briefly explains how to use grid2op with stable baselines 3 RL framework. Make sure to read the previous notebook [11_IntegrationWithExistingRLFrameworks](./11_IntegrationWithExistingRLFrameworks.ipynb) for a deeper dive into what happens. We only show the working solution here.\n",
    "\n",
    "<font color='red'> This explains the ideas and shows a \"self contained\" somewhat minimal example of use of stable baselines 3 framework with grid2op. It is not meant to be fully generic, code might need to be adjusted.</font> \n",
    "\n",
    "This notebook is more an \"example of what works\" rather than a deep dive tutorial.\n",
    "\n",
    "See stable-baselines3.readthedocs.io/ for a more detailed information.\n",
    "\n",
    "This notebook is tested with grid2op 1.10 and stable baselines 2.3.2 on an ubuntu 20.04 machine.\n",
    "\n",
    "It is organised as followed:\n",
    "\n",
    "- [0 Some tips to get started](#0-some-tips-to-get-started) : is a reminder on what you can do to make things work. Indeed, this notebook explains \"how to use grid2op with stable baselines\" but not \"how to create a working agent able to operate a real powergrid in real time with stable baselines\". We wish we could explain the later...\n",
    "- [1 Create the \"Grid2opEnvWrapper\" class](#1-create-the-grid2openvwraper-class) : explain how to create the main grid2op env class that you can use a \"gymnasium\" environment. \n",
    "- [2 Create an environment, and train a first policy](#2-create-an-environment-and-train-a-first-policy): show how to create an environment from the class above (is pretty easy)\n",
    "- [3 Evaluate the trained agent ](#3-evaluate-the-trained-agent): show how to evaluate the trained \"agent\"\n",
    "- [4 Some customizations](#4-some-customizations): explain how to perform some customization of your agent / environment / policy\n",
    "\n",
    "## 0 Some tips to get started\n",
    "\n",
    "<font color='red'> It is unlikely that \"simply\" using a RL algorithm on a grid2op environment will lead to good results for the vast majority of environments.</font>\n",
    "\n",
    "To make RL algorithms work with more or less sucess you might want to:\n",
    "\n",
    "  1) ajust the observation space: in particular selecting the right information for your agent. Too much information\n",
    "     and the size of the observation space will blow up and your agent will not learn anything. Not enough\n",
    "     information and your agent will not be able to capture anything.\n",
    "     \n",
    "  2) customize the action space: dealing with both discrete and continuous values is often a challenge. So maybe you      want to focus on only one type of action. And in all cases, try to still reduce the amount of actions your\n",
    "     agent \n",
    "     can perform. Indeed, for \"larger\" grids (118 substations, as a reference the french grid counts more than 6.000\n",
    "     such substations...) and by limiting 2 busbars per substation (as a reference, for some subsations, you have more\n",
    "     than 12 such \"busbars\") your agent will have the opportunity to choose between more than 60.000 different discrete\n",
    "     actions each steps. This is way too large for current RL algorithm as far as we know (and proposed environment are\n",
    "     small in comparison to real one)\n",
    "     \n",
    "  3) customize the reward: the default reward might not work great for you. Ultimately, what TSO's or ISO's want is\n",
    "     to operate the grid safely, as long as possible with a cost as low as possible. This is of course really hard to\n",
    "     catch everything in one single reward signal. Customizing the reward is also really important because the \"do\n",
    "     nothing\" policy often leads to really good results (much better than random actions) which makes exploration \n",
    "     different actions...). So you kind of want to incentivize your agent to perform some actions at some point.\n",
    "  \n",
    "  4) use fast simulator: even if you target an industrial application with industry grade simulators, we still would\n",
    "     advise you to use (at early stage of training at least) fast simulator for the vast majority of the training\n",
    "     process and then maybe to fine tune on better one.\n",
    "  \n",
    "  5) combine RL with some heuristics: it's super easy to implement things like \"if there is no issue, then do\n",
    "     nothing\". This can be quite time consuming to learn though. Don't hesitate to check out the \"l2rpn-baselines\"\n",
    "     repository for already \"kind of working\" heuristics\n",
    "     \n",
    "And finally don't hesitate to check solution proposed by winners of past l2rpn competitions in l2rpn-baselines.\n",
    "\n",
    "You can also ask question on our discord or on our github.\n",
    "\n",
    "\n",
    "## 1 Create the \"Grid2opEnvWrapper\" class\n",
    "\n",
    "### 1.1 Easy but not easily customizable"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae59e1f5",
   "metadata": {},
   "source": [
    "### 1.2 Similar to ray / rllib with same type of configuration\n",
    "\n",
    "In the next cell, we define a custom environment (that will internally use the `GymEnv` grid2op class) that is needed for ray / rllib.\n",
    "\n",
    "Indeed, in order to work with ray / rllib you need to define a custom wrapper on top of the GymEnv wrapper. You then have:\n",
    "\n",
    "- self._g2op_env which is the default grid2op environment, receiving grid2op Action and producing grid2op Observation.\n",
    "- self._gym_env which is a the grid2op defined `gymnasium Environment` that cannot be directly used with ray / rllib\n",
    "- `Grid2opEnv` which is a the wrapper on top of `self._gym_env` to make it usable with ray / rllib.\n",
    "\n",
    "Ray / rllib expects the gymnasium environment to inherit from `gymnasium.Env` and to be initialized with a given configuration. This is why you need to create the `Grid2opEnv` wrapper on top of `GymEnv`.\n",
    "\n",
    "In the initialization of `Grid2opEnv`, the `env_config` variable is a dictionary that can take as key-word arguments:\n",
    "\n",
    "- `backend_cls` : what is the class of the backend. If not provided, it will use `LightSimBackend` from the `lightsim2grid` package\n",
    "- `backend_options`: what options will be used to create the backend for your environment. Your backend will be created by calling\n",
    "   `backend_cls(**backend_options)`, for example if you want to build `LightSimBackend(detailed_info_for_cascading_failure=False)` you can pass `{\"backend_cls\": LightSimBackend, \"backend_options\": {\"detailed_info_for_cascading_failure\": False}}`\n",
    "- `env_name` : name of the grid2op environment you want to use, by default it uses `\"l2rpn_case14_sandbox\"`\n",
    "- `env_is_test` : whether to add `test=True` when creating the grid2op environment (if `env_is_test` is True it will add `test=True` when calling `grid2op.make(..., test=True)`) otherwise it uses `test=False`\n",
    "- `obs_attr_to_keep` : in this wrapper we only allow your agent to see a Box as an observation. This parameter allows you to control which attributes of the grid2op observation will be present in the agent observation space. By default it's `[\"rho\", \"p_or\", \"gen_p\", \"load_p\"]` which is \"kind of random\" and is probably not suited for every agent.\n",
    "- `act_type` : controls the type of actions your agent will be able to perform. Already coded in this notebook are:\n",
    "   - `\"discrete\"` to use a `Discrete` action space\n",
    "   - `\"box\"` to use a `Box` action space\n",
    "   - `\"multi_discrete\"` to use a `MultiDiscrete` action space\n",
    "- `act_attr_to_keep` :  that allows you to customize the action space. If not provided, it defaults to:\n",
    "  - `[\"set_line_status_simple\", \"set_bus\"]` if `act_type` is `\"discrete\"` \n",
    "  - `[\"redispatch\", \"set_storage\", \"curtail\"]` if `act_type` is `\"box\"` \n",
    "  - `[\"one_line_set\", \"one_sub_set\"]` if `act_type` is `\"multi_discrete\"`\n",
    "\n",
    "If you want to add more customization, for example the reward function, the parameters of the environment etc. etc. feel free to get inspired by this code and extend it. Any PR on this regard is more than welcome."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55e043a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "from typing import Dict, Literal, Any\n",
    "import json\n",
    "\n",
    "from gymnasium import Env\n",
    "from gymnasium.spaces import Discrete, MultiDiscrete, Box\n",
    "\n",
    "import grid2op\n",
    "from grid2op.gym_compat import GymEnv, BoxGymObsSpace, DiscreteActSpace, BoxGymActSpace, MultiDiscreteActSpace\n",
    "from lightsim2grid import LightSimBackend\n",
    "\n",
    "\n",
    "class Grid2opEnvWrapper(Env):\n",
    "    def __init__(self,\n",
    "                 env_config: Dict[Literal[\"backend_cls\",\n",
    "                                          \"backend_options\",\n",
    "                                          \"env_name\",\n",
    "                                          \"env_is_test\",\n",
    "                                          \"obs_attr_to_keep\",\n",
    "                                          \"act_type\",\n",
    "                                          \"act_attr_to_keep\"],\n",
    "                                  Any] = None):\n",
    "        super().__init__()\n",
    "        if env_config is None:\n",
    "            env_config = {}\n",
    "\n",
    "        # handle the backend\n",
    "        backend_cls = LightSimBackend\n",
    "        if \"backend_cls\" in env_config:\n",
    "            backend_cls = env_config[\"backend_cls\"]\n",
    "        backend_options = {}\n",
    "        if \"backend_options\" in env_config:\n",
    "            backend_options = env_config[\"backend_options\"]\n",
    "        backend = backend_cls(**backend_options)\n",
    "\n",
    "        # create the grid2op environment\n",
    "        env_name = \"l2rpn_case14_sandbox\"\n",
    "        if \"env_name\" in env_config:\n",
    "            env_name = env_config[\"env_name\"]\n",
    "        if \"env_is_test\" in env_config:\n",
    "            is_test = bool(env_config[\"env_is_test\"])\n",
    "        else:\n",
    "            is_test = False\n",
    "        self._g2op_env = grid2op.make(env_name, backend=backend, test=is_test)\n",
    "        # NB by default this might be really slow (when the environment is reset)\n",
    "        # see https://grid2op.readthedocs.io/en/latest/data_pipeline.html for maybe 10x speed ups !\n",
    "        # TODO customize reward or action_class for example !\n",
    "\n",
    "        # create the gym env (from grid2op)\n",
    "        self._gym_env = GymEnv(self._g2op_env)\n",
    "\n",
    "        # customize observation space\n",
    "        obs_attr_to_keep = [\"rho\", \"p_or\", \"gen_p\", \"load_p\"]\n",
    "        if \"obs_attr_to_keep\" in env_config:\n",
    "            obs_attr_to_keep = copy.deepcopy(env_config[\"obs_attr_to_keep\"])\n",
    "        self._gym_env.observation_space.close()\n",
    "        self._gym_env.observation_space = BoxGymObsSpace(self._g2op_env.observation_space,\n",
    "                                                         attr_to_keep=obs_attr_to_keep\n",
    "                                                         )\n",
    "        # export observation space for the Grid2opEnv\n",
    "        self.observation_space = Box(shape=self._gym_env.observation_space.shape,\n",
    "                                     low=self._gym_env.observation_space.low,\n",
    "                                     high=self._gym_env.observation_space.high)\n",
    "\n",
    "        # customize the action space\n",
    "        act_type = \"discrete\"\n",
    "        if \"act_type\" in env_config:\n",
    "            act_type = env_config[\"act_type\"]\n",
    "\n",
    "        self._gym_env.action_space.close()\n",
    "        if act_type == \"discrete\":\n",
    "            # user wants a discrete action space\n",
    "            act_attr_to_keep =  [\"set_line_status_simple\", \"set_bus\"]\n",
    "            if \"act_attr_to_keep\" in env_config:\n",
    "                act_attr_to_keep = copy.deepcopy(env_config[\"act_attr_to_keep\"])\n",
    "            self._gym_env.action_space = DiscreteActSpace(self._g2op_env.action_space,\n",
    "                                                          attr_to_keep=act_attr_to_keep)\n",
    "            self.action_space = Discrete(self._gym_env.action_space.n)\n",
    "        elif act_type == \"box\":\n",
    "            # user wants continuous action space\n",
    "            act_attr_to_keep =  [\"redispatch\", \"set_storage\", \"curtail\"]\n",
    "            if \"act_attr_to_keep\" in env_config:\n",
    "                act_attr_to_keep = copy.deepcopy(env_config[\"act_attr_to_keep\"])\n",
    "            self._gym_env.action_space = BoxGymActSpace(self._g2op_env.action_space,\n",
    "                                                        attr_to_keep=act_attr_to_keep)\n",
    "            self.action_space = Box(shape=self._gym_env.action_space.shape,\n",
    "                                    low=self._gym_env.action_space.low,\n",
    "                                    high=self._gym_env.action_space.high)\n",
    "        elif act_type == \"multi_discrete\":\n",
    "            # user wants a multi-discrete action space\n",
    "            act_attr_to_keep = [\"one_line_set\", \"one_sub_set\"]\n",
    "            if \"act_attr_to_keep\" in env_config:\n",
    "                act_attr_to_keep = copy.deepcopy(env_config[\"act_attr_to_keep\"])\n",
    "            self._gym_env.action_space = MultiDiscreteActSpace(self._g2op_env.action_space,\n",
    "                                                               attr_to_keep=act_attr_to_keep)\n",
    "            self.action_space = MultiDiscrete(self._gym_env.action_space.nvec)\n",
    "        else:\n",
    "            raise NotImplementedError(f\"action type '{act_type}' is not currently supported.\")\n",
    "            \n",
    "            \n",
    "    def reset(self, seed=None, options=None):\n",
    "        # use default _gym_env (from grid2op.gym_compat module)\n",
    "        # NB: here you can also specify \"default options\" when you reset, for example:\n",
    "        # - limiting the duration of the episode \"max step\"\n",
    "        # - starting at different steps  \"init ts\"\n",
    "        # - study difficult scenario   \"time serie id\"\n",
    "        # - specify an initial state of your grid \"init state\"\n",
    "        return self._gym_env.reset(seed=seed, options=options)\n",
    "        \n",
    "    def step(self, action):\n",
    "        # use default _gym_env (from grid2op.gym_compat module)\n",
    "        return self._gym_env.step(action)\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a93964d8",
   "metadata": {},
   "source": [
    "## 2 Create an environment, and train a first policy\n",
    "\n",
    "In this section we quickly show :\n",
    "\n",
    "- how to create the gym environment, which is an instance from `Grid2opEnvWrapper` defined above\n",
    "- how to train a PPO policy using stable baselines3\n",
    "\n",
    "This part, for stable baselines is really small."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38629107",
   "metadata": {},
   "outputs": [],
   "source": [
    "from stable_baselines3 import PPO\n",
    "\n",
    "gym_env = Grid2opEnvWrapper()\n",
    "sb3_algo1 = PPO(\"MlpPolicy\", gym_env, verbose=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89be6372",
   "metadata": {},
   "outputs": [],
   "source": [
    "sb3_algo1.learn(total_timesteps=1024)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a8f9717",
   "metadata": {},
   "source": [
    "## 3 Evaluate the trained agent\n",
    "\n",
    "This notebook is a simple quick introduction for stable baselines only. So we don't really recall everything that has been said previously.\n",
    "\n",
    "Please consult the section `0) Recommended initial steps` of the notebook [11_IntegrationWithExistingRLFrameworks](./11_IntegrationWithExistingRLFrameworks.ipynb) for more information.\n",
    "\n",
    "**TLD;DR** grid2op offers the possibility to test your agent on scenarios / episodes different from the one it has been trained. We greatly encourage you to use this functionality.\n",
    "\n",
    "There are two main ways to evaluate your agent:\n",
    "\n",
    "- you stay in the \"gymnasium\" world (see [here](#31-staying-in-the-gymnasium-ecosystem) ) and you evaluate your policy directly just like you would any other gymnasium compatible environment. Simple, easy but without support for some grid2op features\n",
    "- you \"get back\" to the \"grid2op\" world (detailed [here](#32-using-the-grid2op-ecosystem)) by \"converting\" your NN policy into something that is able to output grid2op like action. This introduces yet again a \"wrapper\" but you can benefit from all grid2op features, such as the `Runner` to save an inspect what your policy has done.\n",
    "\n",
    "<font color='red'> We show here just a simple examples to \"get easily started\". For much better working agents, you can have a look at l2rpn-baselines code. There you have classes that maps the environment, the agents etc. to grid2op directly (you don't have to copy paste any wrapper).</font> \n",
    "\n",
    "\n",
    "\n",
    "### 3.1 staying in the gymnasium ecosystem\n",
    "\n",
    "You can do pretty much what you want, but you have to do it yourself, or use any of the \"Wrappers\" available in gymnasium https://gymnasium.farama.org/main/api/wrappers/ (*eg* https://gymnasium.farama.org/main/api/wrappers/misc_wrappers/#gymnasium.wrappers.RecordEpisodeStatistics) or in your RL framework.\n",
    "\n",
    "For the sake of simplicity, we show how to do things \"manually\" even though we do not recommend to do it like that."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05f5e188",
   "metadata": {},
   "outputs": [],
   "source": [
    "nb_episode_test = 2\n",
    "seeds_test_env = (0, 1)    # same size as nb_episode_test\n",
    "seeds_test_agent = (3, 4)  # same size as nb_episode_test\n",
    "ts_ep_test =  (0, 1)       # same size as nb_episode_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da0e7990",
   "metadata": {},
   "outputs": [],
   "source": [
    "ep_infos = {}  # information that will be saved\n",
    "\n",
    "\n",
    "for ep_test_num in range(nb_episode_test):\n",
    "    init_obs, init_infos = gym_env.reset(seed=seeds_test_env[ep_test_num],\n",
    "                                         options={\"time serie id\": ts_ep_test[ep_test_num]})\n",
    "    sb3_algo1.set_random_seed(seeds_test_agent[ep_test_num])\n",
    "    done = False\n",
    "    cum_reward = 0\n",
    "    step_survived = 0\n",
    "    obs = init_obs\n",
    "    while not done:\n",
    "        act, _states = sb3_algo1.predict(obs, deterministic=True)\n",
    "        obs, reward, terminated, truncated, info = gym_env.step(act)\n",
    "        step_survived += 1\n",
    "        cum_reward += float(reward)\n",
    "        done = terminated or truncated\n",
    "    ep_infos[ep_test_num] = {\"time serie id\": ts_ep_test[ep_test_num],\n",
    "                             \"time serie folder\": gym_env._gym_env.init_env.chronics_handler.get_id(),\n",
    "                             \"env seed\": seeds_test_env[ep_test_num],\n",
    "                             \"agent seed\": seeds_test_agent[ep_test_num],\n",
    "                             \"steps survived\": step_survived,\n",
    "                             \"total steps\": int(gym_env._gym_env.init_env.max_episode_duration()),\n",
    "                             \"cum reward\": cum_reward}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f248fdc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# \"prettyprint\" the dictionnary above\n",
    "\n",
    "print(json.dumps(ep_infos, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a32899b",
   "metadata": {},
   "source": [
    "As you might have seen, it's not easy this way to retrieve some useful information about the grid2op environment if these informations are not passed to the policy.\n",
    "\n",
    "For example, we need to call `gym_env._gym_env.init_env` to access the underlying grid2op environment... You have to convert some things from int32 or float32 to float or int otherwise json complains, you have to control yourself the seeds to have reproducible results etc.\n",
    "\n",
    "It's a quick way to have something working but it might be perfected."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fde71911",
   "metadata": {},
   "source": [
    "### 3.2 using the grid2op ecosystem\n",
    "\n",
    "This second method brings it closer to grid2op ecosystem, you will be able to use it with the grid2op `Runner`, save the results and read it back with other tools such as grid2viz and do the evaluation in parrallel without too much trouble (and with high reproducibility).\n",
    "\n",
    "With this method, you build a grid2op agent and this agent can then be used like every other grid2op agent. For example you can compare it with heuristic agents, agent based on optimization etc.\n",
    "\n",
    "This way of doing things also allows you to customize when the neural network policy is used. For example, you might chose to use it only when the grid is \"unsafe\" (and if the grid is safe you use an \"expert\" rules).\n",
    "\n",
    "This is more flexible than the previous one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50625005",
   "metadata": {},
   "outputs": [],
   "source": [
    "from grid2op.Agent import BaseAgent\n",
    "from grid2op.Runner import Runner\n",
    "\n",
    "class Grid2opAgentWrapper(BaseAgent):\n",
    "    def __init__(self,\n",
    "                 gym_env: Grid2opEnvWrapper,\n",
    "                 trained_agent: PPO):\n",
    "        self.gym_env = gym_env\n",
    "        BaseAgent.__init__(self, gym_env._gym_env.init_env.action_space)\n",
    "        self.trained_agent = trained_agent\n",
    "        \n",
    "    def act(self, obs, reward, done):\n",
    "        # you can customize it here to call the NN policy `trained_agent`\n",
    "        # only in some cases, depending on the observation for example\n",
    "        gym_obs = self.gym_env._gym_env.observation_space.to_gym(obs)\n",
    "        gym_act, _states = self.trained_agent.predict(gym_obs, deterministic=True)\n",
    "        grid2op_act = self.gym_env._gym_env.action_space.from_gym(gym_act)\n",
    "        return grid2op_act\n",
    "    \n",
    "    def seed(self, seed):\n",
    "        # implement the seed function\n",
    "        if seed is None:\n",
    "            return\n",
    "        seed_int = int(seed)\n",
    "        if seed_int != seed:\n",
    "            raise RuntimeError(\"Seed must be convertible to an integer\")\n",
    "        self.trained_agent.set_random_seed(seed_int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99e84f4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "my_agent = Grid2opAgentWrapper(gym_env, sb3_algo1)\n",
    "runner = Runner(**gym_env._g2op_env.get_params_for_runner(),\n",
    "                agentClass=None,\n",
    "                agentInstance=my_agent)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18b461cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "res = runner.run(nb_episode=nb_episode_test,\n",
    "                 env_seeds=seeds_test_env,\n",
    "                 agent_seeds=seeds_test_agent,\n",
    "                 episode_id=ts_ep_test,\n",
    "                 add_detailed_output=True\n",
    "                 )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe880aac",
   "metadata": {},
   "outputs": [],
   "source": [
    "res"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fce9ed9",
   "metadata": {},
   "source": [
    "See the documentation or the notebook [05 StudyYourAgent](./05_StudyYourAgent.ipynb) on how to use grid2op tools to study your agent, its decisions etc."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49bf6095",
   "metadata": {},
   "source": [
    "## 4 Some customizations\n",
    "\n",
    "### 4.1 Train a PPO agent using 4 \"runners\" to make the rollouts\n",
    "\n",
    "This, for now, only works on linux based computers. Hopefully this will work on windows and macos as soon as possible.\n",
    "\n",
    "This allows to use some \"parralellism\" during the training: your agent will interact \"at the same time\" with 4 environments allowing it to gather experience faster. But in this case, its training is always done in the \"main\" process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2036ac9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from stable_baselines3.common.env_util import make_vec_env\n",
    "vec_env = make_vec_env(lambda : Grid2opEnvWrapper(), n_envs=4)\n",
    "sb3_algo2 = PPO(\"MlpPolicy\", vec_env, verbose=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d8ac595",
   "metadata": {},
   "outputs": [],
   "source": [
    "sb3_algo2.learn(total_timesteps=1024)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8fc163cd",
   "metadata": {},
   "source": [
    "### 4.2 Use non default parameters to make the grid2op environment\n",
    "\n",
    "In this third example, we will train a policy using the \"box\" action space, and on another environment (`l2rpn_idf_2023` instead of `l2rpn_case14_sandbox`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13740e53",
   "metadata": {},
   "outputs": [],
   "source": [
    "# see https://docs.ray.io/en/latest/rllib/package_ref/doc/ray.rllib.algorithms.algorithm_config.AlgorithmConfig.html\n",
    "\n",
    "# Use a \"Box\" action space (mainly to use redispatching, curtailment and storage units)\n",
    "env_config3 = {\"env_name\": \"l2rpn_idf_2023\",\n",
    "               \"env_is_test\": True,\n",
    "               \"act_type\": \"box\",\n",
    "              }\n",
    "gym_env3 = Grid2opEnvWrapper(env_config3)\n",
    "sb3_algo3 = PPO(\"MlpPolicy\", gym_env3, verbose=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93ac61ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "sb3_algo3.learn(total_timesteps=1024)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00790379",
   "metadata": {},
   "source": [
    "And now a policy using the \"multi discrete\" action space: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cd44edb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# see https://docs.ray.io/en/latest/rllib/package_ref/doc/ray.rllib.algorithms.algorithm_config.AlgorithmConfig.html\n",
    "\n",
    "# Use a \"Box\" action space (mainly to use redispatching, curtailment and storage units)\n",
    "env_config4 = {\"env_name\": \"l2rpn_idf_2023\",\n",
    "               \"env_is_test\": True,\n",
    "               \"act_type\": \"multi_discrete\",\n",
    "              }\n",
    "gym_env4 = Grid2opEnvWrapper(env_config4)\n",
    "sb3_algo4 = PPO(\"MlpPolicy\", gym_env4, verbose=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d18be5ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "sb3_algo4.learn(total_timesteps=1024)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cf2dd58",
   "metadata": {},
   "source": [
    "### 4.3 Customize the policy (number of layers, size of layers etc.)\n",
    "\n",
    "This notebook does not aim at covering all possibilities offered by ray / rllib. For that you need to refer to the ray / rllib documentation.\n",
    "\n",
    "We will simply show how to change the size of the neural network used as a policy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa7cc345",
   "metadata": {},
   "outputs": [],
   "source": [
    "# see https://docs.ray.io/en/latest/rllib/package_ref/doc/ray.rllib.algorithms.algorithm_config.AlgorithmConfig.html\n",
    "\n",
    "gym_env5 = Grid2opEnvWrapper()\n",
    "sb3_algo5 = PPO(\"MlpPolicy\",\n",
    "                gym_env5,\n",
    "                verbose=0,\n",
    "                policy_kwargs={\"net_arch\": [32, 32, 32]}\n",
    "               )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51d435e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "sb3_algo5.learn(total_timesteps=1024)"
   ]
  }
 ],
 "metadata": {
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}