Spaces:
Runtime error
Runtime error
File size: 1,636 Bytes
d746b98 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
from typing import Optional, List, Union, Sequence, Type, Any
import gym
import numpy as np
from ray.rllib import BaseEnv
from stable_baselines3.common.vec_env import VecEnv
from stable_baselines3.common.vec_env.base_vec_env import VecEnvIndices, VecEnvStepReturn, VecEnvObs
from MyRemoteVectorEnv import MyRemoteVectorEnv
class WrapperRay(VecEnv):
def __init__(self, make_env, num_workers, per_worker_env, device):
self.one_env = make_env(0)
self.remote: BaseEnv = MyRemoteVectorEnv(make_env, num_workers, per_worker_env, self.one_env.observation_space, device)
super(WrapperRay, self).__init__(num_workers * per_worker_env, self.one_env.observation_space, self.one_env.action_space)
def reset(self) -> VecEnvObs:
return self.remote.poll()[0]
def step_async(self, actions: np.ndarray) -> None:
self.remote.send_actions(actions)
def step_wait(self) -> VecEnvStepReturn:
return self.remote.poll()
def close(self) -> None:
self.remote.stop()
def get_attr(self, attr_name: str, indices: VecEnvIndices = None) -> List[Any]:
pass
def set_attr(self, attr_name: str, value: Any, indices: VecEnvIndices = None) -> None:
pass
def env_method(self, method_name: str, *method_args, indices: VecEnvIndices = None, **method_kwargs) -> List[Any]:
pass
def env_is_wrapped(self, wrapper_class: Type[gym.Wrapper], indices: VecEnvIndices = None) -> List[bool]:
pass
def get_images(self) -> Sequence[np.ndarray]:
pass
def seed(self, seed: Optional[int] = None) -> List[Union[None, int]]:
pass
|