def make_env_fn(config_env, config_baseline, rank):
dataset = make_dataset(config_env.DATASET.TYPE, config=config_env.DATASET)
config_env.defrost()
config_env.SIMULATOR.SCENE = dataset.episodes[0].scene_id
config_env.freeze()
env = NavRLEnv(config_env=config_env,
config_baseline=config_baseline,
dataset=dataset)
env.seed(rank)
return env
def construct_envs(args):
env_configs = []
baseline_configs = []
basic_config = cfg_env(config_paths=args.task_config, opts=args.opts)
dataset = make_dataset(basic_config.DATASET.TYPE)
scenes = dataset.get_scenes_to_load(basic_config.DATASET)
if len(scenes) > 0:
random.shuffle(scenes)
assert len(scenes) >= args.num_processes, (
"reduce the number of processes as there "
"aren't enough number of scenes"
)
scene_split_size = int(np.floor(len(scenes) / args.num_processes))
scene_splits = [[] for _ in range(args.num_processes)]
for j, s in enumerate(scenes):
scene_splits[j % len(scene_splits)].append(s)
assert sum(map(len, scene_splits)) == len(scenes)
for i in range(args.num_processes):
config_env = cfg_env(config_paths=args.task_config, opts=args.opts)
config_env.defrost()
if len(scenes) > 0:
config_env.DATASET.CONTENT_SCENES = scene_splits[i]
config_env.SIMULATOR.HABITAT_SIM_V0.GPU_DEVICE_ID = args.sim_gpu_id
agent_sensors = args.sensors.strip().split(",")
for sensor in agent_sensors:
assert sensor in ["RGB_SENSOR", "DEPTH_SENSOR"]
config_env.SIMULATOR.AGENT_0.SENSORS = agent_sensors
config_env.freeze()
env_configs.append(config_env)
config_baseline = cfg_baseline()
baseline_configs.append(config_baseline)
logger.info("config_env: {}".format(config_env))
envs = habitat.VectorEnv(
make_env_fn=make_env_fn,
env_fn_args=tuple(
tuple(
zip(env_configs, baseline_configs, range(args.num_processes))
)
),
)
return envs