def make_vec_envs()

in gala/envs.py [0:0]

• 24 lines of code
• 7 McCabe index (conditional complexity)

def make_vec_envs(env_name, seed, num_processes, gamma, log_dir, device,
                  allow_early_resets, num_frame_stack=None, rank=0,
                  signature='', max_steps=None):
    print('log-dir', log_dir)
    envs = [
        make_env(env_name, seed, (rank * num_processes) + i, log_dir,
                 allow_early_resets, signature, max_steps)
        for i in range(num_processes)
    ]

    if len(envs) > 1:
        envs = ShmemVecEnv(envs)
    else:
        envs = DummyVecEnv(envs)

    if len(envs.observation_space.shape) == 1:
        if gamma is None:
            envs = VecNormalize(envs, ret=False)
        else:
            envs = VecNormalize(envs, gamma=gamma)

    envs = VecPyTorch(envs, device)

    if num_frame_stack is not None:
        envs = VecPyTorchFrameStack(envs, num_frame_stack, device)
    elif len(envs.observation_space.shape) == 3:
        envs = VecPyTorchFrameStack(envs, 4, device)

    return envs