multiMixEnv.py

# Copyright (c) 2019-2020, RTE (https://www.rte-france.com)
# See AUTHORS.txt
# This Source Code Form is subject to the terms of the Mozilla Public License, version 2.0.
# If a copy of the Mozilla Public License, version 2.0 was not distributed with this file,
# you can obtain one at http://mozilla.org/MPL/2.0/.
# SPDX-License-Identifier: MPL-2.0
# This file is part of Grid2Op, Grid2Op a testbed platform to model sequential decision making in power systems.

import os
import warnings
import numpy as np
import copy
from typing import Any, Dict, Tuple, Union, List, Literal

from grid2op.dtypes import dt_int, dt_float
from grid2op.Space import GridObjects, RandomObject, DEFAULT_N_BUSBAR_PER_SUB
from grid2op.Exceptions import EnvError, Grid2OpException
from grid2op.Observation import BaseObservation


class MultiMixEnvironment(GridObjects, RandomObject):
    """
    This class represent a single powergrid configuration,
    backed by multiple environments parameters and chronics

    It implements most of the :class:`BaseEnv` public interface:
    so it can be used as a more classic environment.

    MultiMixEnvironment environments behave like a superset of the environment: they
    are made of sub environments (called mixes) that are grid2op regular :class:`Environment`.
    You might think the MultiMixEnvironment as a dictionary of :class:`Environment` that implements
    some of the :class:`BaseEnv` interface such as :func:`BaseEnv.step` or :func:`BaseEnv.reset`.

    By default, each time you call the "step" function a different mix is used. Mixes, by default
    are looped through always in the same order. You can see the Examples section for information
    about control of these


    Examples
    --------
    In this section we present some common use of the MultiMix environment.

    **Basic Usage**

    You can think of a MultiMixEnvironment as any :class:`Environment`. So this is a perfectly
    valid way to use a MultiMix:

    .. code-block:: python

        import grid2op
        from grid2op.Agent import RandomAgent

        # we use an example of a multimix dataset attached with grid2op pacakage
        multimix_env = grid2op.make("l2rpn_neurips_2020_track2", test=True)

        # define an agent like in any environment
        agent = RandomAgent(multimix_env.action_space)

        # and now you can do the open ai gym loop
        NB_EPISODE = 10
        for i in range(NB_EPISODE):
            obs = multimix_env.reset()
            # each time "reset" is called, another mix is used.
            reward = multimix_env.reward_range[0]
            done = False
            while not done:
                act = agent.act(obs, reward, done)
                obs, reward, done, info = multimix_env.step(act)

    **Use each mix one after the other**

    In case you want to study each mix independently, you can iterate through the MultiMix
    in a pythonic way. This makes it easy to perform, for example, 10 episode for a given mix
    before passing to the next one.

    .. code-block:: python

        import grid2op
        from grid2op.Agent import RandomAgent

        # we use an example of a multimix dataset attached with grid2op pacakage
        multimix_env = grid2op.make("l2rpn_neurips_2020_track2", test=True)

        NB_EPISODE = 10
        for mix in multimix_env:
            # mix is a regular environment, you can do whatever you want with it
            # for example
            for i in range(NB_EPISODE):
                obs = multimix_env.reset()
                # each time "reset" is called, another mix is used.
                reward = multimix_env.reward_range[0]
                done = False
                while not done:
                    act = agent.act(obs, reward, done)
                    obs, reward, done, info = multimix_env.step(act)


    **Selecting a given Mix**

    Sometimes it might be interesting to study only a given mix.
    For that you can use the `[]` operator to select only a given mix (which is a grid2op environment)
    and use it as you would.

    This can be done with:

    .. code-block:: python

        import grid2op
        from grid2op.Agent import RandomAgent

        # we use an example of a multimix dataset attached with grid2op pacakage
        multimix_env = grid2op.make("l2rpn_neurips_2020_track2", test=True)

        # define an agent like in any environment
        agent = RandomAgent(multimix_env.action_space)

        # list all available mixes:
        mixes_names = list(multimix_env.keys())

        # and now supposes we want to study only the first one
        mix = multimix_env[mixes_names[0]]

        # and now you can do the open ai gym loop, or anything you want with it
        NB_EPISODE = 10
        for i in range(NB_EPISODE):
            obs = mix.reset()
            # each time "reset" is called, another mix is used.
            reward = mix.reward_range[0]
            done = False
            while not done:
                act = agent.act(obs, reward, done)
                obs, reward, done, info = mix.step(act)

    **Using the Runner**

    For MultiMixEnvironment using the :class:`grid2op.Runner.Runner` cannot be done in a
    straightforward manner. Here we give an example on how to do it.

    .. code-block:: python

        import os
        import grid2op
        from grid2op.Agent import RandomAgent

        # we use an example of a multimix dataset attached with grid2op pacakage
        multimix_env = grid2op.make("l2rpn_neurips_2020_track2", test=True)

        # you can use the runner as following
        PATH = "PATH/WHERE/YOU/WANT/TO/SAVE/THE/RESULTS"
        for mix in multimix_env:
            runner = Runner(**mix.get_params_for_runner(), agentClass=RandomAgent)
            runner.run(nb_episode=1,
                       path_save=os.path.join(PATH,mix.name))

    """

    KEYS_RESET_OPTIONS = {"time serie id"}
    
    def __init__(
        self,
        envs_dir,
        logger=None,
        experimental_read_from_local_dir=False,
        n_busbar=DEFAULT_N_BUSBAR_PER_SUB,
        _add_to_name="",  # internal, for test only, do not use !
        _compat_glop_version=None,  # internal, for test only, do not use !
        _test=False,
        **kwargs,
    ):
        GridObjects.__init__(self)
        RandomObject.__init__(self)
        self.current_env = None
        self.env_index = None
        self.mix_envs = []
        self._env_dir = os.path.abspath(envs_dir)
        self.__closed = False
        # Special case handling for backend
        # TODO: with backend.copy() instead !
        backendClass = None
        backend_kwargs = {}
        if "backend" in kwargs:
            backendClass = type(kwargs["backend"])
            if hasattr(kwargs["backend"], "_my_kwargs"):
                # was introduced in grid2op 1.7.1
                backend_kwargs = kwargs["backend"]._my_kwargs
            del kwargs["backend"]

        # Inline import to prevent cyclical import
        from grid2op.MakeEnv.Make import make

        # TODO reuse same observation_space and action_space in all the envs maybe ?
        try:
            for env_dir in sorted(os.listdir(envs_dir)):
                env_path = os.path.join(envs_dir, env_dir)
                if not os.path.isdir(env_path):
                    continue
                this_logger = (
                    logger.getChild(f"MultiMixEnvironment_{env_dir}")
                    if logger is not None
                    else None
                )
                # Special case for backend
                if backendClass is not None:
                    try:
                        # should pass with grid2op >= 1.7.1
                        bk = backendClass(**backend_kwargs)
                    except TypeError as exc_:
                        # with grid2Op version prior to 1.7.1
                        # you might have trouble with 
                        # "TypeError: __init__() got an unexpected keyword argument 'can_be_copied'"
                        msg_ = ("Impossible to create a backend for each mix using the "
                                "backend key-word arguments. Falling back to creating "
                                "with no argument at all (default behaviour with grid2op <= 1.7.0).")
                        warnings.warn(msg_)
                        bk = backendClass()
                    env = make(
                        env_path,
                        backend=bk,
                        _add_to_name=_add_to_name,
                        _compat_glop_version=_compat_glop_version,
                        n_busbar=n_busbar,
                        test=_test,
                        logger=this_logger,
                        experimental_read_from_local_dir=experimental_read_from_local_dir,
                        **kwargs,
                    )
                else:
                    env = make(
                        env_path,
                        n_busbar=n_busbar,
                        _add_to_name=_add_to_name,
                        _compat_glop_version=_compat_glop_version,
                        test=_test,
                        logger=this_logger,
                        experimental_read_from_local_dir=experimental_read_from_local_dir,
                        **kwargs,
                    )
                self.mix_envs.append(env)
        except Exception as exc_:
            err_msg = "MultiMix environment creation failed: {}".format(exc_)
            raise EnvError(err_msg)

        if len(self.mix_envs) == 0:
            err_msg = "MultiMix envs_dir did not contain any valid env"
            raise EnvError(err_msg)

        self.env_index = 0
        self.current_env = self.mix_envs[self.env_index]
        # Make sure GridObject class attributes are set from first env
        # Should be fine since the grid is the same for all envs
        multi_env_name = os.path.basename(os.path.abspath(envs_dir)) + _add_to_name
        save_env_name = self.current_env.env_name
        self.current_env.env_name = multi_env_name
        self.__class__ = self.init_grid(self.current_env)
        self.current_env.env_name = save_env_name

    def get_path_env(self):
        """
        Get the path that allows to create this environment.

        It can be used for example in `grid2op.utils.underlying_statistics` to save the information directly inside
        the environment data.

        """
        if self.__closed:
            raise EnvError("This environment is closed, you cannot use it.")
        return self._env_dir

    @property
    def current_index(self):
        return self.env_index

    def __len__(self):
        return len(self.mix_envs)

    def __iter__(self):
        """
        Operator __iter__ overload to make a ``MultiMixEnvironment`` iterable

        .. code-block:: python

            import grid2op
            from grid2op.Environment import MultiMixEnvironment
            from grid2op.Runner import Runner

            mm_env = MultiMixEnvironment("/path/to/multi/dataset/folder")

            for env in mm_env:
                run_p = env.get_params_for_runner()
                runner = Runner(**run_p)
                runner.run(nb_episode=1, max_iter=-1)
        """
        self.env_index = 0
        return self

    def __next__(self):
        if self.env_index < len(self.mix_envs):
            r = self.mix_envs[self.env_index]
            self.env_index = self.env_index + 1
            return r
        else:
            self.env_index = 0
            raise StopIteration

    def __getattr__(self, name):
        # TODO what if name is an integer ? make it possible to loop with integer here
        return getattr(self.current_env, name)

    def keys(self):
        for mix in self.mix_envs:
            yield mix.name

    def values(self):
        for mix in self.mix_envs:
            yield mix

    def items(self):
        for mix in self.mix_envs:
            yield mix.name, mix

    def copy(self):
        if self.__closed:
            raise EnvError("This environment is closed, you cannot use it.")
        mix_envs = self.mix_envs
        self.mix_envs = None
        current_env = self.current_env
        self.current_env = None

        cls = self.__class__
        res = cls.__new__(cls)
        for k in self.__dict__:
            if k == "mix_envs" or k == "current_env":
                # this is handled elsewhere
                continue
            setattr(res, k, copy.deepcopy(getattr(self, k)))
        res.mix_envs = [mix.copy() for mix in mix_envs]
        res.current_env = res.mix_envs[res.env_index]

        self.mix_envs = mix_envs
        self.current_env = current_env
        return res

    def __getitem__(self, key):
        """
        Operator [] overload for accessing underlying mixes by name

        .. code-block:: python

            import grid2op
            from grid2op.Environment import MultiMixEnvironment

            mm_env = MultiMixEnvironment("/path/to/multi/dataset/folder")

            mix1_env.name = mm_env["mix_1"]
            assert mix1_env == "mix_1"
            mix2_env.name = mm_env["mix_2"]
            assert mix2_env == "mix_2"
        """
        if self.__closed:
            raise EnvError("This environment is closed, you cannot use it.")
        # Search for key
        for mix in self.mix_envs:
            if mix.name == key:
                return mix

        # Not found by name
        raise KeyError

    def reset(self, 
              *,
              seed: Union[int, None] = None,
              random=False,
              options: Union[Dict[Union[str, Literal["time serie id"]], Union[int, str]], None] = None) -> BaseObservation:
        
        if self.__closed:
            raise EnvError("This environment is closed, you cannot use it.")
        
        if options is not None:
            for el in options:
                if el not in type(self).KEYS_RESET_OPTIONS:
                    raise EnvError(f"You tried to customize the `reset` call with some "
                                   f"`options` using the key `{el}` which is invalid. "
                                   f"Only keys in {sorted(list(type(self).KEYS_RESET_OPTIONS))} "
                                   f"can be used.")
                    
        if random:
            self.env_index = self.space_prng.randint(len(self.mix_envs))
        else:
            self.env_index = (self.env_index + 1) % len(self.mix_envs)

        self.current_env = self.mix_envs[self.env_index]
        
        if options is not None and "time serie id" in options:
            self.set_id(options["time serie id"])
                    
        if seed is not None:
            self.seed(seed)
        return self.current_env.reset()

    def seed(self, seed=None):
        """
        Set the seed of this :class:`Environment` for a better control
        and to ease reproducible experiments.

        Parameters
        ----------
        seed: ``int``
           The seed to set.

        Returns
        ---------
        seeds: ``list``
            The seed used to set the prng (pseudo random number generator)
            for all environments, and each environment ``tuple`` seeds

        """
        if self.__closed:
            raise EnvError("This environment is closed, you cannot use it.")
        try:
            seed = np.array(seed).astype(dt_int)
        except Exception as e:
            raise Grid2OpException(
                "Cannot to seed with the seed provided."
                "Make sure it can be converted to a"
                "numpy 32 bits integer."
            )

        s = super().seed(seed)
        seeds = [s]
        max_dt_int = np.iinfo(dt_int).max
        for env in self.mix_envs:
            env_seed = self.space_prng.randint(max_dt_int)
            env_seeds = env.seed(env_seed)
            seeds.append(env_seeds)
        return seeds

    def set_chunk_size(self, new_chunk_size):
        if self.__closed:
            raise EnvError("This environment is closed, you cannot use it.")
        for mix in self.mix_envs:
            mix.set_chunk_size(new_chunk_size)

    def set_id(self, id_):
        if self.__closed:
            raise EnvError("This environment is closed, you cannot use it.")
        for mix in self.mix_envs:
            mix.set_id(id_)

    def deactivate_forecast(self):
        if self.__closed:
            raise EnvError("This environment is closed, you cannot use it.")
        for mix in self.mix_envs:
            mix.deactivate_forecast()

    def reactivate_forecast(self):
        if self.__closed:
            raise EnvError("This environment is closed, you cannot use it.")
        for mix in self.mix_envs:
            mix.reactivate_forecast()

    def set_thermal_limit(self, thermal_limit):
        """
        Set the thermal limit effectively.
        Will propagate to all underlying mixes
        """
        if self.__closed:
            raise EnvError("This environment is closed, you cannot use it.")
        for mix in self.mix_envs:
            mix.set_thermal_limit(thermal_limit)

    def __enter__(self):
        """
        Support *with-statement* for the environment.

        """
        return self

    def __exit__(self, *args):
        """
        Support *with-statement* for the environment.

        """
        self.close()
        # propagate exception
        return False

    def close(self):
        if self.__closed:
            return

        for mix in self.mix_envs:
            mix.close()
        self.__closed = True

    def attach_layout(self, grid_layout):
        if self.__closed:
            raise EnvError("This environment is closed, you cannot use it.")
        for mix in self.mix_envs:
            mix.attach_layout(grid_layout)

    def __del__(self):
        """when the environment is garbage collected, free all the memory, including cross reference to itself in the observation space."""
        if not self.__closed:
            self.close()
            
    def generate_classes(self):
        # TODO this is not really a good idea, as the multi-mix itself is not read from the
        # files !
        for mix in self.mix_envs:
            mix.generate_classes()