c8e7d02cba
Small typing improvement, related to https://github.com/thu-ml/tianshou/pull/915#discussion_r1329734222
25 lines
595 B
Python
25 lines
595 B
Python
from typing import Any
|
|
|
|
import cloudpickle
|
|
import gymnasium
|
|
import numpy as np
|
|
|
|
from tianshou.env.pettingzoo_env import PettingZooEnv
|
|
|
|
ENV_TYPE = gymnasium.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)
|