DLR-RM / DLR-RM/stable-baselines3
[Feature Request] Multi-Cluster Support for SubprocVecEnv
- Dominant language
- Python
- Stars
- 13.8k
- Forks
- 2.2k
- Avg merge
- 1h 35m
- Merged PRs (30d)
- 2
Description
### 🚀 Feature
I propose enhancing the `SubprocVecEnv` to support multiple clusters. To achieve this, I have created a new class called `DistVecEnv` that is fully compatible with `SubprocVecEnv` as the following:
```python
class DistVecEnv(VecEnv):
"""
Creates a multiprocess vectorized wrapper for multiple environments, distributing each environment to its own
process, allowing significant speed up when the environment is computationally complex.
For performance reasons, if your environment is not IO bound, the number of environments should not exceed the
number of logical cores on your CPU.
:param env_fns: Environments to run in subprocesses
Notes
------
`Actor` is a support class for `DistVecEnv` that controls remote environments, similar to `_worker` for
`SubprocEnv`.
"""
def __init__(self, env_fns: List[Callable[[], gym.Env]]):
self.waiting = False
self.closed = False
self.ref_steps = []
self.actors: List[Actor] = [Actor.remote(env=env_fn) for env_fn in env_fns]
observation_space, action_space = ray.get(self.actors[0].get_spaces.remote())
VecEnv.__init__(self, len(env_fns), observation_space, action_space)
def step_async(self, actions: np.ndarray) -> None:
self.ref_results = [actor.step.remote(action) for actor, action in zip(self.actors, actions)]
self.waiting = True
def step_wait(self) -> VecEnvStepReturn:
results = ray.get(self.ref_results)
self.waiting = False
obs, rews, dones, infos = zip(*results)
return _flatten_obs(obs, self.observation_space), np.stack(rews), np.stack(dones), infos
def seed(self, seed: Optional[int] = None) -> List[Union[None, int]]:
if seed is None:
seed = np.random.randint(0, 2**32 - 1)
return ray.get([actor.seed.remote(seed + idx) for idx, actor in enumerate(self.actors)])
def reset(self) -> VecEnvObs:
obs = ray.get([actor.reset.remote() for actor in self.actors])
return _flatten_obs(obs, self.observation_space)
def close(self) -> None:
if self.closed:
return
ray.get([actor.close.remote() for actor in self.actors])
self.closed = True
def get_images(self) -> Sequence[np.ndarray]:
return ray.get([actor.render.remote("rgb_array") for actor in self.actors])
def get_attr(self, attr_name: str, indices: VecEnvIndices = None) -> List[Any]:
"""Return attribute from vectorized environment (see base class)."""
indices = self._get_indices(indices)
return ray.get([self.actors[index].get_attr.remote(attr_name) for index in indices])
def set_attr(self, attr_name: str, value: Any, indices: VecEnvIndices = None) -> None:
"""Set attribute inside vectorized environments (see base class)."""
indices = self._get_indices(indices)
ray.get([self.actors[index].set_attr.remote(attr_name, value) for index in indices])
def env_method(self, method_name: str, *method_args, indices: VecEnvIndices = None, **method_kwargs) -> List[Any]:
"""Call instance methods of vectorized environments."""
indices = self._get_indices(indices)
refs = [self.actors[index].env_method.remote(method_name, (method_args, method_kwargs)) for index in indices]
return ray.get(refs)
def env_is_wrapped(self, wrapper_class: Type[gym.Wrapper], indices: VecEnvIndices = None) -> List[bool]:
"""Check if worker environments are wrapped with a given wrapper"""
indices = self._get_indices(indices)
return ray.get([self.actors[index].is_wrapped.remote(wrapper_class) for index in indices])
```
Train a PPO agent on `CartPole-v1` using 8 environments.
```python
from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env
from stable_baselines3.common.vec_env import DistVecEnv
if __name__=="__main__":
env = make_vec_env("CartPole-v1", n_envs=8, vec_env_cls=DistVecEnv)
model = PPO("MlpPolicy", env, device="cpu")
model.learn(total_timesteps=25_000)
```
To see the implementation of `DistVecEnv`, please refer to my [gist](https://gist.github.com/snowyday/c5c1979a6f6bf1f7c9287b4e6d8390ef).
### Motivation
The current `SubporcVecEnv` implementation using `multiprocessing` can be limited by hardware and may not be suitable for scaling to larger cluster environments. Utilizing [Ray](https://github.com/ray-project/ray) for distributed computing can provide a more scalable solution.
### Pitch
This feature request aims to enhance the `SubprocVecEnv` class by creating a new class called `DistVecEnv` that uses the Ray library to distribute multiple environments to their own processes.
### Alternatives
Supporting MPI would be a complex and challenging task (please see [MPIVecEnv](https://github.com/Stable-Baselines-Team/stable-baselines3-contrib/issues/45)), and may not provide the scalability and ease of use that Ray offers.
### Additional context
Issues related to this feature request have been discussed for several years:
- [Question Using PPO2 on multiple cluster nodes (MPI)](https://github.com/hill-a/stable-baselines/issues/1054)
- [MPIVecEnv](https://github.com/Stable-Baselines-Team/stable-baselines3-contrib/issues/45)
### Checklist
- [X] I have checked that there is no similar [issue](https://github.com/DLR-RM/stable-baselines3/issues) in the repo
Contributor guide
Assessment
This issue has not been assessed yet.