def construct_envs()

in rl/common/env_utils.py [0:0]

• 38 lines of code
• 18 McCabe index (conditional complexity)

def construct_envs(config, env_class):
    r"""Create VectorEnv object with specified config and env class type.
    To allow better performance, dataset are split into small ones for
    each individual env, grouped by scenes.
    Args:
        config: configs that contain num_processes as well as information
        necessary to create individual environments.
        env_class: class type of the envs to be created.
        specified by the trainer
        randseed: use for debugging
    Returns:
        VectorEnv object created according to specification.
    """

    if config.MODE == 'eval':
        test_episodes = [[] for _ in range(config.NUM_PROCESSES)]
        for idx, episode in enumerate(config.ENV.TEST_EPISODES):
            test_episodes[idx%config.NUM_PROCESSES].append(episode)

        # if there are fewer than num_process episodes, reduce num_processes
        test_episodes = [eps for eps in test_episodes if len(eps)>0]
        # duplicate last episode to pause thread
        test_episodes = [eps + [eps[-1]] for eps in test_episodes]

        config.defrost()
        config.NUM_PROCESSES = len(test_episodes)
        config.ENV.TEST_EPISODES = test_episodes
        config.freeze()

    num_processes = config.NUM_PROCESSES
    env_classes = [env_class for _ in range(num_processes)]

    displays = [None]
    if config.X_DISPLAY is not None:
        displays = config.X_DISPLAY.strip(':').split(',')
    displays = displays*(num_processes//len(displays) + 1)

    # create {num_processes} configs, one for each environment
    configs = []
    for i in range(num_processes):
        proc_config = config.clone()
        proc_config.defrost()

        proc_config.X_DISPLAY = displays[i]
        
        # if testing, force the environment to iterate through a fixed set of episodes
        if config.MODE=='eval':
            proc_config.ENV.TEST_EPISODES = config.ENV.TEST_EPISODES[i]

        proc_config.freeze()
        configs.append(proc_config)

    # initialize the ranks for each processes. Ranks are used to seed the env
    if config.MODE=='train':
        ranks = [np.random.randint(1000) for _ in range(num_processes)]
    elif config.MODE=='eval':
        ranks = list(range(num_processes))

    envs = VectorEnv(
        make_env_fn=make_env_fn,
        env_fn_args=tuple(
            tuple(zip(configs, env_classes, ranks))
        ),
    )
    return envs