600f4bbd55
Preparation for #914 and #920 Changes formatting to ruff and black. Remove python 3.8 ## Additional Changes - Removed flake8 dependencies - Adjusted pre-commit. Now CI and Make use pre-commit, reducing the duplication of linting calls - Removed check-docstyle option (ruff is doing that) - Merged format and lint. In CI the format-lint step fails if any changes are done, so it fulfills the lint functionality. --------- Co-authored-by: Jiayi Weng <jiayi@openai.com>
30 lines
780 B
Python
30 lines
780 B
Python
from typing import TYPE_CHECKING, Any, Union
|
|
|
|
import cloudpickle
|
|
import gymnasium
|
|
import numpy as np
|
|
|
|
from tianshou.env.pettingzoo_env import PettingZooEnv
|
|
|
|
if TYPE_CHECKING:
|
|
import gym
|
|
|
|
# TODO: remove gym entirely? Currently mypy complains in several places
|
|
# if gym.Env is removed from the Union
|
|
ENV_TYPE = Union[gymnasium.Env, "gym.Env", PettingZooEnv]
|
|
|
|
gym_new_venv_step_type = tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]
|
|
|
|
|
|
class CloudpickleWrapper:
|
|
"""A cloudpickle wrapper used in SubprocVectorEnv."""
|
|
|
|
def __init__(self, data: Any) -> None:
|
|
self.data = data
|
|
|
|
def __getstate__(self) -> str:
|
|
return cloudpickle.dumps(self.data)
|
|
|
|
def __setstate__(self, data: str) -> None:
|
|
self.data = cloudpickle.loads(data)
|