254 lines
8.7 KiB
Python
254 lines
8.7 KiB
Python
from abc import ABC, abstractmethod
|
|
from collections.abc import Callable, Sequence
|
|
from enum import Enum
|
|
from typing import Any, TypeAlias, cast
|
|
|
|
import gymnasium as gym
|
|
|
|
from tianshou.env import (
|
|
BaseVectorEnv,
|
|
DummyVectorEnv,
|
|
RayVectorEnv,
|
|
ShmemVectorEnv,
|
|
SubprocVectorEnv,
|
|
)
|
|
from tianshou.highlevel.persistence import Persistence
|
|
from tianshou.utils.net.common import TActionShape
|
|
from tianshou.utils.string import ToStringMixin
|
|
|
|
TObservationShape: TypeAlias = int | Sequence[int]
|
|
|
|
|
|
class EnvType(Enum):
|
|
"""Enumeration of environment types."""
|
|
|
|
CONTINUOUS = "continuous"
|
|
DISCRETE = "discrete"
|
|
|
|
def is_discrete(self) -> bool:
|
|
return self == EnvType.DISCRETE
|
|
|
|
def is_continuous(self) -> bool:
|
|
return self == EnvType.CONTINUOUS
|
|
|
|
def assert_continuous(self, requiring_entity: Any) -> None:
|
|
if not self.is_continuous():
|
|
raise AssertionError(f"{requiring_entity} requires continuous environments")
|
|
|
|
def assert_discrete(self, requiring_entity: Any) -> None:
|
|
if not self.is_discrete():
|
|
raise AssertionError(f"{requiring_entity} requires discrete environments")
|
|
|
|
|
|
class VectorEnvType(Enum):
|
|
DUMMY = "dummy"
|
|
"""Vectorized environment without parallelization; environments are processed sequentially"""
|
|
SUBPROC = "subproc"
|
|
"""Parallelization based on `subprocess`"""
|
|
SUBPROC_SHARED_MEM = "shmem"
|
|
"""Parallelization based on `subprocess` with shared memory"""
|
|
RAY = "ray"
|
|
"""Parallelization based on the `ray` library"""
|
|
|
|
def create_venv(self, factories: list[Callable[[], gym.Env]]) -> BaseVectorEnv:
|
|
match self:
|
|
case VectorEnvType.DUMMY:
|
|
return DummyVectorEnv(factories)
|
|
case VectorEnvType.SUBPROC:
|
|
return SubprocVectorEnv(factories)
|
|
case VectorEnvType.SUBPROC_SHARED_MEM:
|
|
return ShmemVectorEnv(factories)
|
|
case VectorEnvType.RAY:
|
|
return RayVectorEnv(factories)
|
|
case _:
|
|
raise NotImplementedError(self)
|
|
|
|
|
|
class Environments(ToStringMixin, ABC):
|
|
"""Represents (vectorized) environments."""
|
|
|
|
def __init__(self, env: gym.Env, train_envs: BaseVectorEnv, test_envs: BaseVectorEnv):
|
|
self.env = env
|
|
self.train_envs = train_envs
|
|
self.test_envs = test_envs
|
|
self.persistence: Sequence[Persistence] = []
|
|
|
|
@staticmethod
|
|
def from_factory_and_type(
|
|
factory_fn: Callable[[], gym.Env],
|
|
env_type: EnvType,
|
|
venv_type: VectorEnvType,
|
|
num_training_envs: int,
|
|
num_test_envs: int,
|
|
) -> "Environments":
|
|
"""Creates a suitable subtype instance from a factory function that creates a single instance and the type of environment (continuous/discrete).
|
|
|
|
:param factory_fn: the factory for a single environment instance
|
|
:param env_type: the type of environments created by `factory_fn`
|
|
:param venv_type: the vector environment type to use for parallelization
|
|
:param num_training_envs: the number of training environments to create
|
|
:param num_test_envs: the number of test environments to create
|
|
:return: the instance
|
|
"""
|
|
train_envs = venv_type.create_venv([factory_fn] * num_training_envs)
|
|
test_envs = venv_type.create_venv([factory_fn] * num_test_envs)
|
|
env = factory_fn()
|
|
match env_type:
|
|
case EnvType.CONTINUOUS:
|
|
return ContinuousEnvironments(env, train_envs, test_envs)
|
|
case EnvType.DISCRETE:
|
|
return DiscreteEnvironments(env, train_envs, test_envs)
|
|
case _:
|
|
raise ValueError(f"Environment type {env_type} not handled")
|
|
|
|
def _tostring_includes(self) -> list[str]:
|
|
return []
|
|
|
|
def _tostring_additional_entries(self) -> dict[str, Any]:
|
|
return self.info()
|
|
|
|
def info(self) -> dict[str, Any]:
|
|
return {
|
|
"action_shape": self.get_action_shape(),
|
|
"state_shape": self.get_observation_shape(),
|
|
}
|
|
|
|
def set_persistence(self, *p: Persistence) -> None:
|
|
"""Associates the given persistence handlers which may persist and restore environment-specific information.
|
|
|
|
:param p: persistence handlers
|
|
"""
|
|
self.persistence = p
|
|
|
|
@abstractmethod
|
|
def get_action_shape(self) -> TActionShape:
|
|
pass
|
|
|
|
@abstractmethod
|
|
def get_observation_shape(self) -> TObservationShape:
|
|
pass
|
|
|
|
def get_action_space(self) -> gym.Space:
|
|
return self.env.action_space
|
|
|
|
def get_observation_space(self) -> gym.Space:
|
|
return self.env.observation_space
|
|
|
|
@abstractmethod
|
|
def get_type(self) -> EnvType:
|
|
pass
|
|
|
|
|
|
class ContinuousEnvironments(Environments):
|
|
"""Represents (vectorized) continuous environments."""
|
|
|
|
def __init__(self, env: gym.Env, train_envs: BaseVectorEnv, test_envs: BaseVectorEnv):
|
|
super().__init__(env, train_envs, test_envs)
|
|
self.state_shape, self.action_shape, self.max_action = self._get_continuous_env_info(env)
|
|
|
|
@staticmethod
|
|
def from_factory(
|
|
factory_fn: Callable[[], gym.Env],
|
|
venv_type: VectorEnvType,
|
|
num_training_envs: int,
|
|
num_test_envs: int,
|
|
) -> "ContinuousEnvironments":
|
|
"""Creates an instance from a factory function that creates a single instance.
|
|
|
|
:param factory_fn: the factory for a single environment instance
|
|
:param venv_type: the vector environment type to use for parallelization
|
|
:param num_training_envs: the number of training environments to create
|
|
:param num_test_envs: the number of test environments to create
|
|
:return: the instance
|
|
"""
|
|
return cast(
|
|
ContinuousEnvironments,
|
|
Environments.from_factory_and_type(
|
|
factory_fn,
|
|
EnvType.CONTINUOUS,
|
|
venv_type,
|
|
num_training_envs,
|
|
num_test_envs,
|
|
),
|
|
)
|
|
|
|
def info(self) -> dict[str, Any]:
|
|
d = super().info()
|
|
d["max_action"] = self.max_action
|
|
return d
|
|
|
|
@staticmethod
|
|
def _get_continuous_env_info(
|
|
env: gym.Env,
|
|
) -> tuple[tuple[int, ...], tuple[int, ...], float]:
|
|
if not isinstance(env.action_space, gym.spaces.Box):
|
|
raise ValueError(
|
|
"Only environments with continuous action space are supported here. "
|
|
f"But got env with action space: {env.action_space.__class__}.",
|
|
)
|
|
state_shape = env.observation_space.shape or env.observation_space.n # type: ignore
|
|
if not state_shape:
|
|
raise ValueError("Observation space shape is not defined")
|
|
action_shape = env.action_space.shape
|
|
max_action = env.action_space.high[0]
|
|
return state_shape, action_shape, max_action
|
|
|
|
def get_action_shape(self) -> TActionShape:
|
|
return self.action_shape
|
|
|
|
def get_observation_shape(self) -> TObservationShape:
|
|
return self.state_shape
|
|
|
|
def get_type(self) -> EnvType:
|
|
return EnvType.CONTINUOUS
|
|
|
|
|
|
class DiscreteEnvironments(Environments):
|
|
"""Represents (vectorized) discrete environments."""
|
|
|
|
def __init__(self, env: gym.Env, train_envs: BaseVectorEnv, test_envs: BaseVectorEnv):
|
|
super().__init__(env, train_envs, test_envs)
|
|
self.observation_shape = env.observation_space.shape or env.observation_space.n # type: ignore
|
|
self.action_shape = env.action_space.shape or env.action_space.n # type: ignore
|
|
|
|
@staticmethod
|
|
def from_factory(
|
|
factory_fn: Callable[[], gym.Env],
|
|
venv_type: VectorEnvType,
|
|
num_training_envs: int,
|
|
num_test_envs: int,
|
|
) -> "DiscreteEnvironments":
|
|
"""Creates an instance from a factory function that creates a single instance.
|
|
|
|
:param factory_fn: the factory for a single environment instance
|
|
:param venv_type: the vector environment type to use for parallelization
|
|
:param num_training_envs: the number of training environments to create
|
|
:param num_test_envs: the number of test environments to create
|
|
:return: the instance
|
|
"""
|
|
return cast(
|
|
DiscreteEnvironments,
|
|
Environments.from_factory_and_type(
|
|
factory_fn,
|
|
EnvType.CONTINUOUS,
|
|
venv_type,
|
|
num_training_envs,
|
|
num_test_envs,
|
|
),
|
|
)
|
|
|
|
def get_action_shape(self) -> TActionShape:
|
|
return self.action_shape
|
|
|
|
def get_observation_shape(self) -> TObservationShape:
|
|
return self.observation_shape
|
|
|
|
def get_type(self) -> EnvType:
|
|
return EnvType.DISCRETE
|
|
|
|
|
|
class EnvFactory(ToStringMixin, ABC):
|
|
@abstractmethod
|
|
def create_envs(self, num_training_envs: int, num_test_envs: int) -> Environments:
|
|
pass
|