DLR-RM / DLR-RM/stable-baselines3

[Feature Request] Multi-Cluster Support for SubprocVecEnv

Open
#1,345 10 comments 1 reaction 1 assignee Claimed by @araffin View on GitHub
enhancement
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.