[Feature Request] Multiple environments per process

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.

[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')

[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.