import logging
from typing import TYPE_CHECKING
import gymnasium as gym
from gymnasium import Env, Wrapper
from gymnasium.vector import SyncVectorEnv, AsyncVectorEnv
import numpy as np
from numpy.typing import NDArray
from tqdm.auto import tqdm
from collections.abc import Callable, Sequence
from typing import Any
from .labeling import StateLabeler, StateLabel
if TYPE_CHECKING:
from ..policy.policy import PolicyClass
logger = logging.getLogger(__name__)
[docs]
class VeriGymEnv(gym.Env):
"""
Abstract wrapper around `gymnasium.Env` for a unified VeriGym interface.
Parent to all types of VeriGym environments.
"""
def __init__(self):
self.state_labeler = StateLabeler(set())
[docs]
def simulate(
self, policy: "PolicyClass", n_steps: int = 1, verbose=True, seed=None
) -> list[list[tuple[NDArray, NDArray, NDArray, NDArray, bool]]]:
"""
Simulate the environment for `n_steps` using the provided `policy`.
Args:
policy: A function that takes an observation and returns an action.
n_steps (int): Number of steps to simulate.
Returns:
A list of trajectories containing a list of tuples (state, action, reward, next_state, terminated)
for each step. `terminated` is True only if `next_state` is terminal (e.g. the pole fell in CartPole);
a trajectory that was cut off by truncation (e.g. a time limit) ends with `terminated=False`.
"""
dataset, trajectory = [], []
self.seed_spaces(state_seed=seed, action_seed=seed)
state, info = self.reset(seed=seed)
for _ in tqdm(range(n_steps), desc="Simulating", disable=not verbose):
action = policy.get_action(state)
next_state, reward, terminated, truncated, info = self.step(action)
next_state, action, reward = (
np.array(next_state),
np.array(action),
np.array(reward),
)
trajectory.append((state, action, reward, next_state, bool(terminated)))
state = next_state
if terminated or truncated:
dataset.append(trajectory)
state, info = self.reset() # no need for resetting with the seed here, keep the running random generators.
trajectory = []
if len(trajectory) > 0:
# Append remaining trajectories
dataset.append(trajectory)
return dataset
def has_state_labels(self):
return len(self.state_labeler.labels) > 0
def add_state_label(self, label: StateLabel):
self.state_labeler.add_state_label(label)
def add_state_labels(self, labels: list[StateLabel]):
for label in labels:
self.state_labeler.add_state_label(label)
def get_labels_of_state(self, state):
return self.state_labeler.get_labels_of_state(state)
def seed_spaces(self, state_seed=None, action_seed=None):
if state_seed is not None:
self.observation_space.seed(state_seed)
if action_seed is not None:
self.action_space.seed(action_seed)
[docs]
@classmethod
def make_vec(cls,
num_envs: int = 1,
vectorization_mode: str | None = "sync",
vector_kwargs: dict[str, Any] | None = None,
wrappers: Sequence[Callable[[Env], Wrapper]] | None = None,
wrapper_kwargs: list | None = None,
make_callback = None, **kwargs):
"""
Builds vectorized VeriGymEnvs.
Parameters
----------
num_envs : int
How many envs to contain in the vectorized env.
render_mode : str
Environment render mode. Currently can only be None.
vectorization_mode : str
The vectorization mode. Can be "sync" or "async"
vector_kwargs : dict
Further arguments to apply to vectorization.
wrappers : Sequence
List of wrappers to be applied to the environments.
**kwargs :
Allows to include environment-specific parameters.
Returns
-------
env : gym.SyncVectorEnv(VeriGymEnv) if vectorization_mode=="sync" or gym.AsyncVectorEnv(VeriGymEnv) if vectorization_mode=="async"
The vectorized FrameworkExplicitEnvs.
"""
if kwargs is None:
kwargs = {}
if vector_kwargs is None:
vector_kwargs = {}
if wrappers is None:
wrappers = []
if wrapper_kwargs is None:
wrapper_kwargs = []
else:
assert len(wrappers) == len(wrapper_kwargs), "Wrappers and wrapper_kwargs do not match."
if make_callback is None:
make_callback = cls
def _make_env(func, wrappers, wrapper_kwargs, **env_kwargs):
env = func(**env_kwargs)
for wrapper, w_kwarg in zip(wrappers, wrapper_kwargs):
env = wrapper(env, **w_kwarg)
return env
env_fns = [lambda: _make_env(make_callback, wrappers, wrapper_kwargs, **kwargs) for _ in range(num_envs)]
if vectorization_mode == "sync":
return SyncVectorEnv(env_fns, **vector_kwargs)
elif vectorization_mode == "async":
return AsyncVectorEnv(env_fns, **vector_kwargs)
else:
raise ValueError()