9c100e0705
- change the internal API name of worker: send_action -> send, get_result -> recv (align with envpool) - add a timing test for venvs.reset() to make sure the concurrent execution - change venvs.reset() logic Co-authored-by: Jiayi Weng <trinkle23897@gmail.com>
88 lines
2.6 KiB
Python
88 lines
2.6 KiB
Python
from abc import ABC, abstractmethod
|
|
from typing import Any, Callable, List, Optional, Tuple, Union
|
|
|
|
import gym
|
|
import numpy as np
|
|
|
|
|
|
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 = self.get_env_attr("action_space") # noqa: B009
|
|
|
|
@abstractmethod
|
|
def get_env_attr(self, key: str) -> Any:
|
|
pass
|
|
|
|
@abstractmethod
|
|
def set_env_attr(self, key: str, value: Any) -> None:
|
|
pass
|
|
|
|
@abstractmethod
|
|
def send(self, action: Optional[np.ndarray]) -> None:
|
|
"""Send action signal to low-level worker.
|
|
|
|
When action is None, it indicates sending "reset" signal; otherwise
|
|
it indicates "step" signal. The paired return value from "recv"
|
|
function is determined by such kind of different signal.
|
|
"""
|
|
pass
|
|
|
|
def recv(
|
|
self
|
|
) -> Union[Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray], np.ndarray]:
|
|
"""Receive result from low-level worker.
|
|
|
|
If the last "send" function sends a NULL action, it only returns a
|
|
single observation; otherwise it returns a tuple of (obs, rew, done,
|
|
info).
|
|
"""
|
|
return self.result
|
|
|
|
def reset(self) -> np.ndarray:
|
|
self.send(None)
|
|
return self.recv() # type: ignore
|
|
|
|
def step(
|
|
self, action: np.ndarray
|
|
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
|
|
"""Perform one timestep of the environment's dynamic.
|
|
|
|
"send" and "recv" are coupled in sync simulation, so users only call
|
|
"step" function. But they can be called separately in async
|
|
simulation, i.e. someone calls "send" first, and calls "recv" later.
|
|
"""
|
|
self.send(action)
|
|
return self.recv() # type: ignore
|
|
|
|
@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()
|