wehub-resource-sync
184 lines
6.8 KiB
Python
184 lines
6.8 KiB
Python
import logging
|
|
import time
|
|
|
|
import numpy as np
|
|
from gymnasium.envs.classic_control import CartPoleEnv
|
|
|
|
from ray.rllib.examples.envs.classes.multi_agent import make_multi_agent
|
|
from ray.rllib.utils.annotations import override
|
|
from ray.rllib.utils.error import EnvError
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class CartPoleCrashing(CartPoleEnv):
|
|
"""A CartPole env that crashes (or stalls) from time to time.
|
|
|
|
Useful for testing faulty sub-env (within a vectorized env) handling by
|
|
EnvRunners.
|
|
|
|
After crashing, the env expects a `reset()` call next (calling `step()` will
|
|
result in yet another error), which may or may not take a very long time to
|
|
complete. This simulates the env having to reinitialize some sub-processes, e.g.
|
|
an external connection.
|
|
|
|
The env can also be configured to stall (and do nothing during a call to `step()`)
|
|
from time to time for a configurable amount of time.
|
|
"""
|
|
|
|
def __init__(self, config=None):
|
|
super().__init__()
|
|
|
|
self.config = config if config is not None else {}
|
|
|
|
# Crash probability (in each `step()`).
|
|
self.p_crash = config.get("p_crash", 0.005)
|
|
# Crash probability when `reset()` is called.
|
|
self.p_crash_reset = config.get("p_crash_reset", 0.0)
|
|
# Crash exactly after every n steps. If a 2-tuple, will uniformly sample
|
|
# crash timesteps from in between the two given values.
|
|
self.crash_after_n_steps = config.get("crash_after_n_steps")
|
|
self._crash_after_n_steps = None
|
|
assert (
|
|
self.crash_after_n_steps is None
|
|
or isinstance(self.crash_after_n_steps, int)
|
|
or (
|
|
isinstance(self.crash_after_n_steps, tuple)
|
|
and len(self.crash_after_n_steps) == 2
|
|
)
|
|
)
|
|
# Only ever crash, if on certain worker indices.
|
|
faulty_indices = config.get("crash_on_worker_indices", None)
|
|
if faulty_indices and config.worker_index not in faulty_indices:
|
|
self.p_crash = 0.0
|
|
self.p_crash_reset = 0.0
|
|
self.crash_after_n_steps = None
|
|
|
|
# Stall probability (in each `step()`).
|
|
self.p_stall = config.get("p_stall", 0.0)
|
|
# Stall probability when `reset()` is called.
|
|
self.p_stall_reset = config.get("p_stall_reset", 0.0)
|
|
# Stall exactly after every n steps.
|
|
self.stall_after_n_steps = config.get("stall_after_n_steps")
|
|
self._stall_after_n_steps = None
|
|
# Amount of time to stall. If a 2-tuple, will uniformly sample from in between
|
|
# the two given values.
|
|
self.stall_time_sec = config.get("stall_time_sec")
|
|
assert (
|
|
self.stall_time_sec is None
|
|
or isinstance(self.stall_time_sec, (int, float))
|
|
or (
|
|
isinstance(self.stall_time_sec, tuple) and len(self.stall_time_sec) == 2
|
|
)
|
|
)
|
|
|
|
# Only ever stall, if on certain worker indices.
|
|
faulty_indices = config.get("stall_on_worker_indices", None)
|
|
if faulty_indices and config.worker_index not in faulty_indices:
|
|
self.p_stall = 0.0
|
|
self.p_stall_reset = 0.0
|
|
self.stall_after_n_steps = None
|
|
|
|
# Timestep counter for the ongoing episode.
|
|
self.timesteps = 0
|
|
|
|
# Time in seconds to initialize (in this c'tor).
|
|
sample = 0.0
|
|
if "init_time_s" in config:
|
|
sample = (
|
|
config["init_time_s"]
|
|
if not isinstance(config["init_time_s"], tuple)
|
|
else np.random.uniform(
|
|
config["init_time_s"][0], config["init_time_s"][1]
|
|
)
|
|
)
|
|
|
|
print(f"Initializing crashing env (with init-delay of {sample}sec) ...")
|
|
time.sleep(sample)
|
|
|
|
# Make sure envs don't crash at the same time.
|
|
self._rng = np.random.RandomState()
|
|
|
|
@override(CartPoleEnv)
|
|
def reset(self, *, seed=None, options=None):
|
|
# Reset timestep counter for the new episode.
|
|
self.timesteps = 0
|
|
self._crash_after_n_steps = None
|
|
|
|
# Should we crash?
|
|
if self._should_crash(p=self.p_crash_reset):
|
|
raise EnvError(
|
|
f"Simulated env crash on worker={self.config.worker_index} "
|
|
f"env-idx={self.config.vector_index} during `reset()`! "
|
|
"Feel free to use any other exception type here instead."
|
|
)
|
|
# Should we stall for a while?
|
|
self._stall_if_necessary(p=self.p_stall_reset)
|
|
|
|
return super().reset()
|
|
|
|
@override(CartPoleEnv)
|
|
def step(self, action):
|
|
# Increase timestep counter for the ongoing episode.
|
|
self.timesteps += 1
|
|
|
|
# Should we crash?
|
|
if self._should_crash(p=self.p_crash):
|
|
raise EnvError(
|
|
f"Simulated env crash on worker={self.config.worker_index} "
|
|
f"env-idx={self.config.vector_index} during `step()`! "
|
|
"Feel free to use any other exception type here instead."
|
|
)
|
|
# Should we stall for a while?
|
|
self._stall_if_necessary(p=self.p_stall)
|
|
|
|
return super().step(action)
|
|
|
|
def _should_crash(self, p):
|
|
rnd = self._rng.rand()
|
|
if rnd < p:
|
|
print("Crashing due to p(crash)!")
|
|
return True
|
|
elif self.crash_after_n_steps is not None:
|
|
if self._crash_after_n_steps is None:
|
|
self._crash_after_n_steps = (
|
|
self.crash_after_n_steps
|
|
if not isinstance(self.crash_after_n_steps, tuple)
|
|
else np.random.randint(
|
|
self.crash_after_n_steps[0], self.crash_after_n_steps[1]
|
|
)
|
|
)
|
|
if self._crash_after_n_steps == self.timesteps:
|
|
print("Crashing due to n timesteps reached!")
|
|
return True
|
|
|
|
return False
|
|
|
|
def _stall_if_necessary(self, p):
|
|
stall = False
|
|
if self._rng.rand() < p:
|
|
stall = True
|
|
elif self.stall_after_n_steps is not None:
|
|
if self._stall_after_n_steps is None:
|
|
self._stall_after_n_steps = (
|
|
self.stall_after_n_steps
|
|
if not isinstance(self.stall_after_n_steps, tuple)
|
|
else np.random.randint(
|
|
self.stall_after_n_steps[0], self.stall_after_n_steps[1]
|
|
)
|
|
)
|
|
if self._stall_after_n_steps == self.timesteps:
|
|
stall = True
|
|
|
|
if stall:
|
|
sec = (
|
|
self.stall_time_sec
|
|
if not isinstance(self.stall_time_sec, tuple)
|
|
else np.random.uniform(self.stall_time_sec[0], self.stall_time_sec[1])
|
|
)
|
|
print(f" -> will stall for {sec}sec ...")
|
|
time.sleep(sec)
|
|
|
|
|
|
MultiAgentCartPoleCrashing = make_multi_agent(lambda config: CartPoleCrashing(config))
|