[Feature Request] Multiple environments per process

[neighthan]

SubprocVecEnv allows us to run multiple instances of an environment in parallel using subprocesses. It can be nice to also run multiple environments in each subprocess. The idea here is that if the time for each environment to step isn't always the same, we'll have to wait for the slowest subprocess to finish before we can move on (the straggler effect). By having each subprocess run multiple environments (sequentially), we can reduce this effect (the more environments, the likelier that the total time for a process is near the average stepping time). The paper Accelerated Methods for Deep Reinforcement Learning proposes this (among other things).

I have this working already in my fork (though I'm only 50% of the way done running the tests). I've only done a couple simple and relatively short performance tests on my machine, but if you set n_envs_per_process = 1 then the performance is essentially the same as now (the only difference is unpacking a singleton list in a few places). With 4 processes and 3 environments in each, I saw an 18% increase in FPS over 4 processes with 1 environment each. I wouldn't claim that increase is standard as I haven't tested much, but I don't think performance should be hurt.

Let me know you're interested in this feature, and I'll submit a PR. Right now, I've changed SubprocVecEnv so that it's capable of using any number of environments per subprocess, but if you'd rather, I could create a new class instead (SequentialSubprocVecEnv? Not sure what to name it) to leave the current one untouched.

[araffin]

With 4 processes and 3 environments in each, I saw an 18% increase in FPS over 4 processes with 1 environment each.

Well, you should compare at least with the same number of environment. Also, try to compare it with DummyVecEnv, it is in fact usually competitive.

Let me know you're interested in this feature, and I'll submit a PR.

I think we would be interested if there is a real improvement, and I would prefer a new type of SubprocVecEnv.

[neighthan]

After some further testing, I'm not sure there's enough improvement to be worth it, and I've switched to a PyTorch RL library, so I'll close this. If anybody is interested in this in the future, feel free to ping me, and I'll point you at my code for it as a starting point.

[jbulow]

I found this issue when trying to propose the same feature. I have created a SubprocDummyVecEnv that, as the name suggests, encapsulates a DummyVecEnv within a SubprocVecEnv. With SubprocVecEnv and DummyVecEnv I could increase performance upto around 60 environments. With my (rough) SubprocDummyVecEnv I can 600 environments. fps increase is about a factor 5-6. If of interest I can do more formal performance tests and clean up the implementation. The design is very simple. The results from DummyVecEnvs are concatenated at the SubprocVecEnv level and the actions are partitioned on the number of subprocs. At subproc-level the number of subprocs * envs_per_dummyvecenv is registered as the total number of environments.

Sample code for using SubprocDummyVecEnv:

def make_env(env_id, rank, subrank, seed=0):
    def _init():
        env = gym.make(env_id)
        env = Monitor(env, f'{log_dir}/{rank}_{subrank}_monitor.csv', allow_early_resets=True)
        env.seed(seed + rank)
        return env
    return _init

def make_dummyvecenv(env_id, rank, seed=0):
    def _init():
        env = DummyVecEnv([make_env(env_id, rank, i) for i in range(num_envs_per_cpu)])
        return env
    set_global_seeds(seed)
    return _init

num_cpu = 25
num_envs_per_cpu = 25
env = SubprocDummyVecEnv([make_dummyvecenv(env_id, i) for i in range(num_cpu)], start_method='fork')

[Miffyli]

I can confirm that, outside stable-baselines, this kind of setups can speed up gathering samples considerably (multiple times faster), depending on the environment. I am not 100% sure how well it would scale up with stable-baselines, but the preliminary results sound promising :).

I think the name of the new VecEnv should be something different though. Dummy was/is used for "something simple that implements the necessary API", and SubprocDummyVecEnv sounds like simple version of SubprocVecEnv. Perhaps using word "grouped" in there (e.g. GroupedSubprocVecEnv).

When you are happy with the quality of your code feel free to open a PR for this :). Like @araffin mentioned above (months ago), it would be good to have some performance comparisons as well with e.g. standard control problems.

[Robbert96]

@jbulow could you post the code for your SubprocDummyVecEnv class?

[jbulow]

Sorry, lost the code.

[araffin]

you can take a look at https://github.com/Stable-Baselines-Team/stable-baselines3-contrib/blob/master/sb3_contrib/common/vec_env/async_eval.py (here it is async evaluation for the algorithm ARS but the main idea is the same, at the end, it allows to run n envs per process).