09692c84fe
Change the behavior of to_numpy and to_torch: from now on, dict is automatically converted to Batch and list is automatically converted to np.ndarray (if an error occurs, raise the exception instead of converting each element in the list).
68 lines
2.0 KiB
Python
68 lines
2.0 KiB
Python
import gym
|
|
import numpy as np
|
|
from abc import ABC, abstractmethod
|
|
from typing import Any, List, Tuple, Optional, Callable
|
|
|
|
|
|
class EnvWorker(ABC):
|
|
"""An abstract worker for an environment."""
|
|
|
|
def __init__(self, env_fn: Callable[[], gym.Env]) -> None:
|
|
self._env_fn = env_fn
|
|
self.is_closed = False
|
|
self.result: Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]
|
|
self.action_space = getattr(self, "action_space")
|
|
|
|
@abstractmethod
|
|
def __getattr__(self, key: str) -> Any:
|
|
pass
|
|
|
|
@abstractmethod
|
|
def reset(self) -> Any:
|
|
pass
|
|
|
|
@abstractmethod
|
|
def send_action(self, action: np.ndarray) -> None:
|
|
pass
|
|
|
|
def get_result(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
|
|
return self.result
|
|
|
|
def step(
|
|
self, action: np.ndarray
|
|
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
|
|
"""Perform one timestep of the environment's dynamic.
|
|
|
|
"send_action" and "get_result" are coupled in sync simulation, so
|
|
typically users only call "step" function. But they can be called
|
|
separately in async simulation, i.e. someone calls "send_action" first,
|
|
and calls "get_result" later.
|
|
"""
|
|
self.send_action(action)
|
|
return self.get_result()
|
|
|
|
@staticmethod
|
|
def wait(
|
|
workers: List["EnvWorker"], wait_num: int, timeout: Optional[float] = None
|
|
) -> List["EnvWorker"]:
|
|
"""Given a list of workers, return those ready ones."""
|
|
raise NotImplementedError
|
|
|
|
def seed(self, seed: Optional[int] = None) -> Optional[List[int]]:
|
|
return self.action_space.seed(seed) # issue 299
|
|
|
|
@abstractmethod
|
|
def render(self, **kwargs: Any) -> Any:
|
|
"""Render the environment."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def close_env(self) -> None:
|
|
pass
|
|
|
|
def close(self) -> None:
|
|
if self.is_closed:
|
|
return None
|
|
self.is_closed = True
|
|
self.close_env()
|