1"""
2This module contains a base class for EPANET control environments --
3i.e. controlling hydraulic actuators such as pumps and valves or single chemical (no EPANET-MSX support!).
4"""
5from typing import Optional, Any
6import warnings
7import numpy as np
8from epyt_flow.simulation import ScenarioConfig
9from gymnasium.spaces import Dict
10from gymnasium.spaces.utils import flatten_space
11
12from .rl_env import RlEnv
13from .actions.pump_speed_actions import PumpSpeedAction
14from .actions.quality_actions import ChemicalInjectionAction
15from .actions.actuator_state_actions import PumpStateAction, ValveStateAction
16
17
[docs]
18class EpanetControlEnv(RlEnv):
19 """
20 Base class for hydraulic control environments
21 (incl. basic quality that can be simulated with EPANET only).
22
23 Parameters
24 ----------
25 scenario_config : `epyt_flow.simulation.ScenarioConfig <https://epyt-flow.readthedocs.io/en/stable/epyt_flow.simulation.html#epyt_flow.simulation.scenario_config.ScenarioConfig>`_
26 Configuration of the scenario.
27 pumps_speed_actions : list[:class:`~epyt_control.actions.pump_speed_actions.PumpSpeedAction`], optional
28 List of pumps where the speed has to be controlled.
29
30 The default is None.
31 pumps_state_actions : list[:class:`~epyt_control.actions.actuator_state_actions.PumpStateAction`], optional
32 Lisst of pumps where the state has to be controlled.
33
34 The default is None.
35 valves_state_actions : list[:class:`~epyt_control.actions.actuator_state_actions.ValveStateAction`], optional
36 List of valves that have to be controlled.
37
38 The default is None.
39 chemical_injection_actions : list[:class:`~epyt_control.actions.quality_actions.ChemicalInjectionAction`], optional
40 List chemical injection actions -- i.e. places in the network where the
41 injection of the chemical has to be controlled.
42
43 The default is None.
44 """
45 def __init__(self, scenario_config: ScenarioConfig,
46 pumps_speed_actions: Optional[list[PumpSpeedAction]] = None,
47 pumps_state_actions: Optional[list[PumpStateAction]] = None,
48 valves_state_actions: Optional[list[ValveStateAction]] = None,
49 chemical_injection_actions: Optional[list[ChemicalInjectionAction]] = None,
50 **kwds):
51 if pumps_speed_actions is not None:
52 if not isinstance(pumps_speed_actions, list):
53 raise TypeError("'pumps_speed_actions' must be an instance of " +
54 "'list[PumpSpeedAction]' but not of " +
55 f"'{type(pumps_speed_actions)}'")
56 if any(not isinstance(pump_speed_action, PumpSpeedAction)
57 for pump_speed_action in pumps_speed_actions):
58 raise TypeError("All items in 'pumps_speed_actions' must be an instance of " +
59 "'PumpSpeedAction'")
60 if pumps_state_actions is not None:
61 if not isinstance(pumps_state_actions, list):
62 raise TypeError("'pumps_state_actions' must be an instance of " +
63 "'list[PumpStateAction]' but not of " +
64 f"'{type(pumps_state_actions)}'")
65 if any(not isinstance(pump_state_action, PumpStateAction)
66 for pump_state_action in pumps_state_actions):
67 raise TypeError("All items in 'pumps_state_actions' must be an instance of " +
68 "'PumpStateAction'")
69 if valves_state_actions is not None:
70 if not isinstance(valves_state_actions, list):
71 raise TypeError("'valves_state_actions' must be an instance of " +
72 "'list[ValveAction]' but not of " +
73 f"'{type(valves_state_actions)}'")
74 if any(not isinstance(valve_state_action, ValveStateAction)
75 for valve_state_action in valves_state_actions):
76 raise TypeError("All items in 'valves_state_actions' must " +
77 "be an instance of 'ValveStateAction'")
78 if chemical_injection_actions is not None:
79 if not isinstance(chemical_injection_actions, list):
80 raise TypeError("'chemical_injection_actions' must be an instance of " +
81 "'list[ChemicalInjectionAction]' but not of " +
82 f"'{type(chemical_injection_actions)}'")
83 if any(not isinstance(chemical_injection_action, ChemicalInjectionAction)
84 for chemical_injection_action in chemical_injection_actions):
85 raise TypeError("All items in 'chemical_injection_actions' " +
86 "must be an instance of 'ChemicalInjectionAction'")
87
88 self._pumps_speed_actions = pumps_speed_actions
89 self._pumps_state_actions = pumps_state_actions
90 self._valves_state_actions = valves_state_actions
91 self._chemical_injection_actions = chemical_injection_actions
92
93 action_space = {}
94 my_actions = []
95 if self._pumps_speed_actions is not None:
96 my_actions += self._pumps_speed_actions
97 action_space |= {f"{action_space.pump_id}-speed": action_space.to_gym_action_space()
98 for action_space in self._pumps_speed_actions}
99 if self._pumps_state_actions is not None:
100 my_actions += self._pumps_state_actions
101 action_space |= {f"{action_space.pump_id}-state": action_space.to_gym_action_space()
102 for action_space in self._pumps_state_actions}
103 if self._valves_state_actions is not None:
104 my_actions += self._valves_state_actions
105 action_space |= {f"{action_space.valve_id}-state": action_space.to_gym_action_space()
106 for action_space in self._valves_state_actions}
107 if self._valves_state_actions is not None:
108 my_actions += self._valves_state_actions
109 action_space |= {f"{action_space.valve_id}-state": action_space.to_gym_action_space()
110 for action_space in self._valves_state_actions}
111 if self._chemical_injection_actions is not None:
112 my_actions += self._chemical_injection_actions
113 action_space |= {f"{action_space.node_id}-chem": action_space.to_gym_action_space()
114 for action_space in self._chemical_injection_actions}
115
116 gym_action_space = flatten_space(Dict(action_space))
117
118 super().__init__(scenario_config=scenario_config, gym_action_space=gym_action_space,
119 action_space=my_actions, **kwds)
120
121
122HydraulicControlEnv = EpanetControlEnv
123
124
[docs]
125class MultiConfigEpanetControlEnv(EpanetControlEnv):
126 """
127 Base class for hydraulic control environments (incl. basic quality that can be simulated
128 with EPANET only) that can handle multiple scenario configurations -- those scenarios are
129 utilized in a round-robin scheduling scheme (i.e. autorest=True).
130
131 Parameters
132 ----------
133 scenario_configs : list[`epyt_flow.simulation.ScenarioConfig <https://epyt-flow.readthedocs.io/en/stable/epyt_flow.simulation.html#epyt_flow.simulation.scenario_config.ScenarioConfig>`_]
134 List of all scenario configurations that are used in this environment.
135 pumps_speed_actions : list[:class:`~epyt_control.actions.pump_speed_actions.PumpSpeedAction`], optional
136 List of pumps where the speed has to be controlled.
137
138 The default is None.
139 pumps_state_actions : list[:class:`~epyt_control.actions.actuator_state_actions.PumpStateAction`], optional
140 Lisst of pumps where the state has to be controlled.
141
142 The default is None.
143 valves_state_actions : list[:class:`~epyt_control.actions.actuator_state_actions.ValveStateAction`], optional
144 List of valves that have to be controlled.
145
146 The default is None.
147 chemical_injection_actions : list[:class:`~epyt_control.actions.quality_actions.ChemicalInjectionAction`], optional
148 List chemical injection actions -- i.e. places in the network where the
149 injection of the chemical has to be controlled.
150
151 The default is None.
152 autoreset : `bool`, optional
153 If True, the environment is automatically reset when the episode ends. In this case, `terminated` will always be `False`, so the environment can't be wrapped in a vectorized environment from Stable Baselines 3.
154 If False, the environment's `step` method returns `terminated=True` at the end of an episode, and `reset` has to be called to start the next scenario.
155 The default is True.
156 reload_scenario_when_reset : `bool`, optional
157 If True, the scenario (incl. the .inp and .msx file) is reloaded from the hard disk.
158 If False, only the simulation is reset.
159
160 The default is True.
161 """
162 def __init__(self, scenario_configs: list[ScenarioConfig],
163 pumps_speed_actions: Optional[list[PumpSpeedAction]] = None,
164 pumps_state_actions: Optional[list[PumpStateAction]] = None,
165 valves_state_actions: Optional[list[ValveStateAction]] = None,
166 chemical_injection_actions: Optional[list[ChemicalInjectionAction]] = None,
167 autoreset: bool = True,
168 reload_scenario_when_reset: bool = True):
169 if not isinstance(scenario_configs, list):
170 raise TypeError("'scenario_configs' must be an instance of " +
171 "epyt_flow.simulation.ScenarioConfig but " +
172 f"not of '{type(scenario_configs)}'")
173 if any(not isinstance(scenario_config, ScenarioConfig)
174 for scenario_config in scenario_configs):
175 raise TypeError("All items in 'scenario_config' must be instances of " +
176 "epyt_flow.simulation.ScenarioConfig")
177
178 if len(scenario_configs) > 10:
179 warnings.warn("You are using many scenarios. You might face issues w.r.t. " +
180 "memory consumption as well as with the maximum number of open files " +
181 "allowed by the operating system.", UserWarning)
182
183 self._scenario_configs = scenario_configs
184 self._scenario_sims = [None] * len(scenario_configs)
185 self._current_scenario_idx = 0
186
187 super().__init__(self._scenario_configs[self._current_scenario_idx], pumps_speed_actions,
188 pumps_state_actions, valves_state_actions, chemical_injection_actions,
189 autoreset=autoreset,
190 reload_scenario_when_reset=reload_scenario_when_reset)
191
[docs]
192 def reset(self, seed: Optional[int] = None, options: Optional[dict[str, Any]] = None
193 ) -> tuple[np.ndarray, dict]:
194 # Back up current simulation
195 self._scenario_sims[self._current_scenario_idx] = self._scenario_sim
196
197 # Move on to next scenario
198 self._current_scenario_idx = (self._current_scenario_idx + 1) % len(self._scenario_configs)
199 self._scenario_config = self._scenario_configs[self._current_scenario_idx]
200 self._scenario_sim = self._scenario_sims[self._current_scenario_idx]
201
202 return super().reset(seed, options)
203
204
205MultiConfigHydraulicControlEnv = MultiConfigEpanetControlEnv