def gibson_env(hparams):
basic_config = cfg_env(config_file=hparams.task_config,
config_dir=os.path.join(dirname, 'configs'))
scenes = PointNavDatasetV1.get_scenes_to_load(basic_config.DATASET)
config_env = cfg_env(config_file=hparams.task_config,
config_dir=os.path.join(dirname, 'configs'))
config_env.defrost()
if len(scenes) > 0:
random.shuffle(scenes)
config_env.DATASET.POINTNAVV1.CONTENT_SCENES = scenes
for sensor in hparams.sensors:
assert sensor in ["RGB_SENSOR", "DEPTH_SENSOR"]
config_env.SIMULATOR.AGENT_0.SENSORS = hparams.sensors
config_env.freeze()
config_baseline = cfg_baseline()
dataset = PointNavDatasetV1(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)
return env
def construct_envs(args):
env_configs = []
baseline_configs = []
args_list = []
# TODO check params consistency here
basic_config = cfg_env(config_paths=[args.task_config])
print("loading scenes ...")
scenes = PointNavDatasetV1.get_scenes_to_load(basic_config.DATASET)
if len(scenes) > 0:
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))
print("using ", args.num_processes, " processes and ", scene_split_size,
" scenes per process")
for i in range(args.num_processes):
config_env = cfg_env(config_paths=[args.task_config])
config_env.defrost()
if len(scenes) > 0:
config_env.DATASET.CONTENT_SCENES = scenes[i *
scene_split_size:(i +
1) *
scene_split_size]
if i < args.num_processes_on_first_gpu:
gpu_id = 0
else:
gpu_id = int((i - args.num_processes_on_first_gpu) //
args.num_processes_per_gpu) + args.sim_gpu_id
# gpu_id = min(torch.cuda.device_count() - 1, gpu_id)
gpu_id = 0
config_env.SIMULATOR.HABITAT_SIM_V0.GPU_DEVICE_ID = gpu_id
config_env.freeze()
env_configs.append(config_env)
config_baseline = cfg_baseline()
baseline_configs.append(config_baseline)
args_list.append(args)
envs = VectorEnv(
make_env_fn=make_env_fn,
env_fn_args=tuple(
tuple(
zip(args_list, env_configs, baseline_configs,
range(args.num_processes)))),
)
# envs = make_env_fn(args_list[0], env_configs[0], config_baseline=baseline_configs[0], rank=42)
print("returning with environment")
return envs
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
def construct_envs(args):
env_configs = []
baseline_configs = []
basic_config = cfg_env(config_file=args.task_config)
scenes = PointNavDatasetV1.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))
for i in range(args.num_processes):
config_env = cfg_env(config_file=args.task_config)
config_env.defrost()
if len(scenes) > 0:
config_env.DATASET.POINTNAVV1.CONTENT_SCENES = scenes[
i * scene_split_size:(i + 1) * scene_split_size]
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
def main():
basic_config = cfg_env(config_paths=
["env/habitat/habitat_lab/configs/" + args.task_config])
basic_config.defrost()
basic_config.DATASET.SPLIT = args.split
basic_config.freeze()
scenes = PointNavDatasetV1.get_scenes_to_load(basic_config.DATASET)
config_env = cfg_env(config_paths=
["env/habitat/habitat_lab/configs/" + args.task_config])
config_env.defrost()
config_env.DATASET.CONTENT_SCENES = scenes
gpu_id = 0
config_env.SIMULATOR.HABITAT_SIM_V0.GPU_DEVICE_ID = gpu_id
agent_sensors = []
agent_sensors.append("RGB_SENSOR")
agent_sensors.append("DEPTH_SENSOR")
config_env.SIMULATOR.AGENT_0.SENSORS = agent_sensors
config_env.ENVIRONMENT.MAX_EPISODE_STEPS = args.max_episode_length
config_env.ENVIRONMENT.ITERATOR_OPTIONS.SHUFFLE = False
config_env.SIMULATOR.RGB_SENSOR.WIDTH = args.env_frame_width
config_env.SIMULATOR.RGB_SENSOR.HEIGHT = args.env_frame_height
config_env.SIMULATOR.RGB_SENSOR.HFOV = args.hfov
config_env.SIMULATOR.RGB_SENSOR.POSITION = [0, args.camera_height, 0]
config_env.SIMULATOR.DEPTH_SENSOR.WIDTH = args.env_frame_width
config_env.SIMULATOR.DEPTH_SENSOR.HEIGHT = args.env_frame_height
config_env.SIMULATOR.DEPTH_SENSOR.HFOV = args.hfov
config_env.SIMULATOR.DEPTH_SENSOR.POSITION = [0, args.camera_height, 0]
config_env.SIMULATOR.TURN_ANGLE = 10
config_env.DATASET.SPLIT = args.split
config_env.freeze()
env = make_env_fn(args, config_env, 0)
obs, inf = env.reset()
print("done")
def make_habitat_vector_env(num_processes=2,
target_dim=7,
preprocessing_fn=None,
log_dir=None,
visdom_name='main',
visdom_log_file=None,
visdom_server='localhost',
visdom_port='8097',
vis_interval=200,
scenes=None,
val_scenes=['Greigsville', 'Pablo', 'Mosquito'],
num_val_processes=0,
swap_building_k_episodes=10,
gpu_devices=[0],
collate_obs_before_transform=False,
map_kwargs={},
reward_kwargs={},
seed=42):
assert map_kwargs[
'map_building_size'] > 0, 'Map building size must be positive!'
default_reward_kwargs = {
'slack_reward': -0.01,
'success_reward': 10,
'use_visit_penalty': False,
'visit_penalty_coef': 0,
'penalty_eps': 999,
}
for k, v in default_reward_kwargs.items():
if k not in reward_kwargs:
reward_kwargs[k] = v
habitat_path = os.path.dirname(os.path.dirname(habitat.__file__))
task_config = os.path.join(habitat_path,
'configs/tasks/pointnav_gibson_train.yaml')
basic_config = cfg_env(config_file=task_config)
basic_config.defrost()
basic_config.DATASET.POINTNAVV1.DATA_PATH = os.path.join(
habitat_path, basic_config.DATASET.POINTNAVV1.DATA_PATH)
basic_config.freeze()
if scenes is None:
scenes = PointNavDatasetV1.get_scenes_to_load(basic_config.DATASET)
random.shuffle(scenes)
val_task_config = os.path.join(
habitat_path, 'configs/tasks/pointnav_gibson_val_mini.yaml')
val_cfg = cfg_env(config_file=val_task_config)
val_cfg.defrost()
val_cfg.DATASET.SPLIT = "val"
val_cfg.freeze()
scenes = [s for s in scenes if s not in val_scenes]
if num_val_processes > 0 and len(val_scenes) % num_val_processes != 0:
warnings.warn(
"Please make num_val_processes ({}) evenly divide len(val_scenes) ({}) or some buildings may be overrepresented"
.format(num_val_processes, len(val_scenes)))
env_configs = []
baseline_configs = []
encoders = []
target_dims = []
is_val = []
# Assign specific buildings to each process
train_process_scenes = [[]
for _ in range(num_processes - num_val_processes)]
for i, scene in enumerate(scenes):
train_process_scenes[i % len(train_process_scenes)].append(scene)
if num_val_processes > 0:
val_process_scenes = [[] for _ in range(num_val_processes)]
for i, scene in enumerate(val_scenes):
val_process_scenes[i % len(val_process_scenes)].append(scene)
for i in range(num_processes):
config_env = cfg_env(task_config)
config_env.defrost()
config_env.DATASET.POINTNAVV1.DATA_PATH = os.path.join(
habitat_path, basic_config.DATASET.POINTNAVV1.DATA_PATH)
if i < num_processes - num_val_processes:
config_env.DATASET.SPLIT = 'train'
config_env.DATASET.POINTNAVV1.CONTENT_SCENES = train_process_scenes[
i]
else:
config_env.DATASET.SPLIT = 'val'
val_i = i - (num_processes - num_val_processes)
config_env.DATASET.POINTNAVV1.CONTENT_SCENES = val_process_scenes[
val_i]
print("Env {}:".format(i),
config_env.DATASET.POINTNAVV1.CONTENT_SCENES)
config_env.SIMULATOR.HABITAT_SIM_V0.GPU_DEVICE_ID = gpu_devices[
i % len(gpu_devices)]
agent_sensors = ["RGB_SENSOR"]
config_env.SIMULATOR.AGENT_0.SENSORS = agent_sensors
config_env.SIMULATOR.SCENE = os.path.join(habitat_path,
config_env.SIMULATOR.SCENE)
config_env.freeze()
env_configs.append(config_env)
config_baseline = cfg_baseline()
baseline_configs.append(config_baseline)
encoders.append(preprocessing_fn)
target_dims.append(target_dim)
should_record = [(i == 0 or i == (num_processes - num_val_processes))
for i in range(num_processes)]
envs = HabitatPreprocessVectorEnv(
make_env_fn=make_env_fn,
env_fn_args=tuple(
tuple(
zip(
env_configs,
baseline_configs,
range(num_processes),
target_dims,
[log_dir for _ in range(num_processes)],
[visdom_name for _ in range(num_processes)],
[visdom_log_file for _ in range(num_processes)],
[vis_interval for _ in range(num_processes)],
[visdom_server for _ in range(num_processes)],
[visdom_port for _ in range(num_processes)],
[swap_building_k_episodes for _ in range(num_processes)],
[map_kwargs for _ in range(num_processes)],
[reward_kwargs for _ in range(num_processes)],
should_record,
[seed + i for i in range(num_processes)],
))),
preprocessing_fn=preprocessing_fn,
collate_obs_before_transform=collate_obs_before_transform)
envs.observation_space = envs.observation_spaces[0]
envs.action_space = spaces.Discrete(3)
envs.reward_range = None
envs.metadata = None
envs.is_embodied = True
return envs
def run_training():
parser = ppo_args()
args = parser.parse_args()
random.seed(args.seed)
device = torch.device("cuda:{}".format(args.pth_gpu_id))
logger.add_filehandler(args.log_file)
if not os.path.isdir(args.checkpoint_folder):
os.makedirs(args.checkpoint_folder)
for p in sorted(list(vars(args))):
logger.info("{}: {}".format(p, getattr(args, p)))
envs = construct_envs(args)
task_cfg = cfg_env(config_paths=args.task_config)
actor_critic = Policy(
observation_space=envs.observation_spaces[0],
action_space=envs.action_spaces[0],
hidden_size=args.hidden_size,
goal_sensor_uuid=task_cfg.TASK.GOAL_SENSOR_UUID,
)
actor_critic.to(device)
agent = PPO(
actor_critic,
args.clip_param,
args.ppo_epoch,
args.num_mini_batch,
args.value_loss_coef,
args.entropy_coef,
lr=args.lr,
eps=args.eps,
max_grad_norm=args.max_grad_norm,
)
logger.info("agent number of parameters: {}".format(
sum(param.numel() for param in agent.parameters())))
observations = envs.reset()
batch = batch_obs(observations)
rollouts = RolloutStorage(
args.num_steps,
envs.num_envs,
envs.observation_spaces[0],
envs.action_spaces[0],
args.hidden_size,
)
for sensor in rollouts.observations:
rollouts.observations[sensor][0].copy_(batch[sensor])
rollouts.to(device)
episode_rewards = torch.zeros(envs.num_envs, 1)
episode_counts = torch.zeros(envs.num_envs, 1)
current_episode_reward = torch.zeros(envs.num_envs, 1)
window_episode_reward = deque()
window_episode_counts = deque()
t_start = time()
env_time = 0
pth_time = 0
count_steps = 0
count_checkpoints = 0
for update in range(args.num_updates):
if args.use_linear_lr_decay:
update_linear_schedule(agent.optimizer, update, args.num_updates,
args.lr)
agent.clip_param = args.clip_param * (1 - update / args.num_updates)
for step in range(args.num_steps):
t_sample_action = time()
# sample actions
with torch.no_grad():
step_observation = {
k: v[step]
for k, v in rollouts.observations.items()
}
(
values,
actions,
actions_log_probs,
recurrent_hidden_states,
) = actor_critic.act(
step_observation,
rollouts.recurrent_hidden_states[step],
rollouts.masks[step],
)
pth_time += time() - t_sample_action
t_step_env = time()
outputs = envs.step([a[0].item() for a in actions])
observations, rewards, dones, infos = [
list(x) for x in zip(*outputs)
]
env_time += time() - t_step_env
t_update_stats = time()
batch = batch_obs(observations)
rewards = torch.tensor(rewards, dtype=torch.float)
rewards = rewards.unsqueeze(1)
masks = torch.tensor([[0.0] if done else [1.0] for done in dones],
dtype=torch.float)
current_episode_reward += rewards
episode_rewards += (1 - masks) * current_episode_reward
episode_counts += 1 - masks
current_episode_reward *= masks
rollouts.insert(
batch,
recurrent_hidden_states,
actions,
actions_log_probs,
values,
rewards,
masks,
)
count_steps += envs.num_envs
pth_time += time() - t_update_stats
if len(window_episode_reward) == args.reward_window_size:
window_episode_reward.popleft()
window_episode_counts.popleft()
window_episode_reward.append(episode_rewards.clone())
window_episode_counts.append(episode_counts.clone())
t_update_model = time()
with torch.no_grad():
last_observation = {
k: v[-1]
for k, v in rollouts.observations.items()
}
next_value = actor_critic.get_value(
last_observation,
rollouts.recurrent_hidden_states[-1],
rollouts.masks[-1],
).detach()
rollouts.compute_returns(next_value, args.use_gae, args.gamma,
args.tau)
value_loss, action_loss, dist_entropy = agent.update(rollouts)
rollouts.after_update()
pth_time += time() - t_update_model
# log stats
if update > 0 and update % args.log_interval == 0:
logger.info("update: {}\tfps: {:.3f}\t".format(
update, count_steps / (time() - t_start)))
logger.info("update: {}\tenv-time: {:.3f}s\tpth-time: {:.3f}s\t"
"frames: {}".format(update, env_time, pth_time,
count_steps))
window_rewards = (window_episode_reward[-1] -
window_episode_reward[0]).sum()
window_counts = (window_episode_counts[-1] -
window_episode_counts[0]).sum()
if window_counts > 0:
logger.info("Average window size {} reward: {:3f}".format(
len(window_episode_reward),
(window_rewards / window_counts).item(),
))
else:
logger.info("No episodes finish in current window")
# checkpoint model
if update % args.checkpoint_interval == 0:
checkpoint = {"state_dict": agent.state_dict()}
torch.save(
checkpoint,
os.path.join(
args.checkpoint_folder,
"ckpt.{}.pth".format(count_checkpoints),
),
)
count_checkpoints += 1
def construct_envs(args):
env_configs = []
# baseline_configs = []
args_list = []
basic_config = cfg_env(
config_paths=["env/habitat/habitat_lab/configs/" + args.task_config])
basic_config.defrost()
basic_config.DATASET.SPLIT = args.split
basic_config.freeze()
scenes = PointNavDatasetV1.get_scenes_to_load(basic_config.DATASET)
if len(scenes) > 0:
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))
for i in range(args.num_processes):
config_env = cfg_env(config_paths=[
"env/habitat/habitat_lab/configs/" + args.task_config
])
config_env.defrost()
if len(scenes) > 0:
config_env.DATASET.CONTENT_SCENES = scenes[i *
scene_split_size:(i +
1) *
scene_split_size]
if i < args.num_processes_on_first_gpu:
gpu_id = 0
else:
gpu_id = int((i - args.num_processes_on_first_gpu) //
args.num_processes_per_gpu) + args.sim_gpu_id
gpu_id = min(torch.cuda.device_count() - 1, gpu_id)
config_env.SIMULATOR.HABITAT_SIM_V0.GPU_DEVICE_ID = gpu_id
agent_sensors = []
agent_sensors.append("RGB_SENSOR")
agent_sensors.append("DEPTH_SENSOR")
config_env.SIMULATOR.AGENT_0.SENSORS = agent_sensors
config_env.ENVIRONMENT.MAX_EPISODE_STEPS = args.max_episode_length
config_env.ENVIRONMENT.ITERATOR_OPTIONS.SHUFFLE = False
config_env.SIMULATOR.RGB_SENSOR.WIDTH = args.env_frame_width
config_env.SIMULATOR.RGB_SENSOR.HEIGHT = args.env_frame_height
config_env.SIMULATOR.RGB_SENSOR.HFOV = args.hfov
config_env.SIMULATOR.RGB_SENSOR.POSITION = [0, args.camera_height, 0]
config_env.SIMULATOR.DEPTH_SENSOR.WIDTH = args.env_frame_width
config_env.SIMULATOR.DEPTH_SENSOR.HEIGHT = args.env_frame_height
config_env.SIMULATOR.DEPTH_SENSOR.HFOV = args.hfov
config_env.SIMULATOR.DEPTH_SENSOR.POSITION = [0, args.camera_height, 0]
config_env.SIMULATOR.TURN_ANGLE = 10
config_env.DATASET.SPLIT = args.split
config_env.freeze()
env_configs.append(config_env)
# Baseline configs are not used
# config_baseline = cfg_baseline()
# baseline_configs.append(config_baseline)
args_list.append(args)
envs = VectorEnv(
make_env_fn=make_env_fn,
env_fn_args=tuple(
tuple(zip(args_list, env_configs, range(args.num_processes)))),
)
return envs
def construct_envs(args):
env_configs = []
args_list = []
basic_config = cfg_env(config_paths=["envs/habitat/configs/"
+ args.task_config])
basic_config.defrost()
basic_config.DATASET.SPLIT = args.split
basic_config.DATASET.DATA_PATH = \
basic_config.DATASET.DATA_PATH.replace("v1", args.version)
basic_config.DATASET.EPISODES_DIR = \
basic_config.DATASET.EPISODES_DIR.replace("v1", args.version)
basic_config.freeze()
scenes = basic_config.DATASET.CONTENT_SCENES
if "*" in basic_config.DATASET.CONTENT_SCENES:
content_dir = os.path.join(basic_config.DATASET.EPISODES_DIR.format(
split=args.split), "content")
scenes = _get_scenes_from_folder(content_dir)
if len(scenes) > 0:
assert len(scenes) >= args.num_processes, (
"reduce the number of processes as there "
"aren't enough number of scenes"
)
scene_split_sizes = [int(np.floor(len(scenes) / args.num_processes))
for _ in range(args.num_processes)]
for i in range(len(scenes) % args.num_processes):
scene_split_sizes[i] += 1
print("Scenes per thread:")
for i in range(args.num_processes):
config_env = cfg_env(config_paths=["envs/habitat/configs/"
+ args.task_config])
config_env.defrost()
if len(scenes) > 0:
config_env.DATASET.CONTENT_SCENES = scenes[
sum(scene_split_sizes[:i]):
sum(scene_split_sizes[:i + 1])
]
print("Thread {}: {}".format(i, config_env.DATASET.CONTENT_SCENES))
if i < args.num_processes_on_first_gpu:
gpu_id = 0
else:
gpu_id = int((i - args.num_processes_on_first_gpu)
// args.num_processes_per_gpu) + args.sim_gpu_id
gpu_id = min(torch.cuda.device_count() - 1, gpu_id)
config_env.SIMULATOR.HABITAT_SIM_V0.GPU_DEVICE_ID = gpu_id
agent_sensors = []
agent_sensors.append("RGB_SENSOR")
agent_sensors.append("DEPTH_SENSOR")
# agent_sensors.append("SEMANTIC_SENSOR")
config_env.SIMULATOR.AGENT_0.SENSORS = agent_sensors
# Reseting episodes manually, setting high max episode length in sim
config_env.ENVIRONMENT.MAX_EPISODE_STEPS = 10000000
config_env.ENVIRONMENT.ITERATOR_OPTIONS.SHUFFLE = False
config_env.SIMULATOR.RGB_SENSOR.WIDTH = args.env_frame_width
config_env.SIMULATOR.RGB_SENSOR.HEIGHT = args.env_frame_height
config_env.SIMULATOR.RGB_SENSOR.HFOV = args.hfov
config_env.SIMULATOR.RGB_SENSOR.POSITION = [0, args.camera_height, 0]
config_env.SIMULATOR.DEPTH_SENSOR.WIDTH = args.env_frame_width
config_env.SIMULATOR.DEPTH_SENSOR.HEIGHT = args.env_frame_height
config_env.SIMULATOR.DEPTH_SENSOR.HFOV = args.hfov
config_env.SIMULATOR.DEPTH_SENSOR.MIN_DEPTH = args.min_depth
config_env.SIMULATOR.DEPTH_SENSOR.MAX_DEPTH = args.max_depth
config_env.SIMULATOR.DEPTH_SENSOR.POSITION = [0, args.camera_height, 0]
# config_env.SIMULATOR.SEMANTIC_SENSOR.WIDTH = args.env_frame_width
# config_env.SIMULATOR.SEMANTIC_SENSOR.HEIGHT = args.env_frame_height
# config_env.SIMULATOR.SEMANTIC_SENSOR.HFOV = args.hfov
# config_env.SIMULATOR.SEMANTIC_SENSOR.POSITION = \
# [0, args.camera_height, 0]
config_env.SIMULATOR.TURN_ANGLE = args.turn_angle
config_env.DATASET.SPLIT = args.split
config_env.DATASET.DATA_PATH = \
config_env.DATASET.DATA_PATH.replace("v1", args.version)
config_env.DATASET.EPISODES_DIR = \
config_env.DATASET.EPISODES_DIR.replace("v1", args.version)
config_env.freeze()
env_configs.append(config_env)
args_list.append(args)
envs = VectorEnv(
make_env_fn=make_env_fn,
env_fn_args=tuple(
tuple(
zip(args_list, env_configs, range(args.num_processes))
)
),
)
return envs
def make_habitat_vector_env(
scenario='PointNav',
num_processes=2,
target_dim=7,
preprocessing_fn=None,
log_dir=None,
visdom_name='main',
visdom_log_file=None,
visdom_server='localhost',
visdom_port='8097',
vis_interval=200,
train_scenes=None,
val_scenes=None,
num_val_processes=0,
swap_building_k_episodes=10,
gpu_devices=[0],
map_kwargs={},
reward_kwargs={},
seed=42,
test_mode=False,
debug_mode=False,
scenario_kwargs={},
):
assert map_kwargs[
'map_building_size'] > 0, 'Map building size must be positive!'
default_reward_kwargs = {
'slack_reward': -0.01,
'success_reward': 10,
'use_visit_penalty': False,
'visit_penalty_coef': 0,
'penalty_eps': 999,
'sparse': False,
'dist_coef': 1.0,
}
for k, v in default_reward_kwargs.items():
if k not in reward_kwargs:
reward_kwargs[k] = v
habitat_path = os.path.dirname(os.path.dirname(habitat.__file__))
if scenario == 'PointNav' or scenario == 'Exploration':
task_config = os.path.join(habitat_path,
'configs/tasks/pointnav_gibson_train.yaml')
# only difference is that Exploration needs DEPTH_SENSOR but that is added in the Env
# task_config = os.path.join(habitat_path, 'configs/tasks/exploration_gibson.yaml')
env_configs = []
baseline_configs = []
encoders = []
target_dims = []
is_val = []
# Assign specific episodes to each process
config_env = cfg_env(task_config)
# Load dataset
print('Loading val dataset (partition by episode)...')
datasetfile_path = config_env.DATASET.POINTNAVV1.DATA_PATH.format(
split='val')
dataset = PointNavDatasetV1()
with gzip.open(datasetfile_path, "rt") as f:
dataset.from_json(f.read())
val_datasets = get_splits(dataset, max(num_val_processes, 1))
# for d in val_datasets:
# d.episodes = [d.episodes[0]]
print('Loaded.')
print('Loading train dataset (partition by building)...')
train_datasets = []
if num_processes - num_val_processes > 0:
# dataset = PointNavDatasetV1(config_env.DATASET)
train_datasets = [
None for _ in range(num_processes - num_val_processes)
]
print('Loaded.')
# Assign specific buildings to each process
if num_processes > num_val_processes:
train_process_scenes = [[] for _ in range(num_processes -
num_val_processes)]
if train_scenes is None:
train_scenes = PointNavDatasetV1.get_scenes_to_load(
config_env.DATASET)
random.shuffle(train_scenes)
for i, scene in enumerate(train_scenes):
train_process_scenes[i % len(train_process_scenes)].append(scene)
# If n processes > n envs, some processes can use all envs
for j, process in enumerate(train_process_scenes):
if len(process) == 0:
train_process_scenes[j] = list(train_scenes)
get_scenes = lambda d: list(
Counter([e.scene_id.split('/')[-1].split(".")[0]
for e in d.episodes]).items())
for i in range(num_processes):
config_env = cfg_env(task_config)
config_env.defrost()
if i < num_processes - num_val_processes:
config_env.DATASET.SPLIT = 'train'
# config_env.DATASET.POINTNAVV1.CONTENT_SCENES = get_scenes(train_datasets[i])
config_env.DATASET.POINTNAVV1.CONTENT_SCENES = train_process_scenes[
i]
else:
val_i = i - (num_processes - num_val_processes)
config_env.DATASET.SPLIT = 'val'
if val_scenes is not None:
config_env.DATASET.POINTNAVV1.CONTENT_SCENES = val_scenes
else:
config_env.DATASET.POINTNAVV1.CONTENT_SCENES = get_scenes(
val_datasets[val_i])
print("Env {}:".format(i),
config_env.DATASET.POINTNAVV1.CONTENT_SCENES)
config_env.SIMULATOR.HABITAT_SIM_V0.GPU_DEVICE_ID = gpu_devices[
i % len(gpu_devices)]
config_env.SIMULATOR.SCENE = os.path.join(habitat_path,
config_env.SIMULATOR.SCENE)
config_env.SIMULATOR.AGENT_0.SENSORS = ["RGB_SENSOR"]
# Now define the config for the sensor
# config.TASK.AGENT_POSITION_SENSOR = habitat.Config()
# config.TASK.AGENT_POSITION_SENSOR.TYPE = "agent_position_sensor"
# config.TASK.SENSORS.append("AGENT_POSITION_SENSOR")
config_env.TASK.MEASUREMENTS.append('COLLISIONS')
config_env.freeze()
env_configs.append(config_env)
config_baseline = cfg_baseline()
baseline_configs.append(config_baseline)
encoders.append(preprocessing_fn)
target_dims.append(target_dim)
should_record = [(i == 0 or i == (num_processes - num_val_processes))
for i in range(num_processes)]
if debug_mode:
env = make_env_fn(
scenario,
env_configs[0],
baseline_configs[0],
0,
0,
1,
target_dim,
log_dir,
visdom_name,
visdom_log_file,
vis_interval,
visdom_server,
visdom_port,
swap_building_k_episodes,
map_kwargs,
reward_kwargs,
False,
seed, # TODO set should_record to True
test_mode,
(train_datasets + val_datasets)[0],
scenario_kwargs)
envs = PreprocessEnv(env, preprocessing_fn=preprocessing_fn)
else:
envs = HabitatPreprocessVectorEnv(
make_env_fn=make_env_fn,
env_fn_args=tuple(
tuple(
zip(
[scenario for _ in range(num_processes)],
env_configs,
baseline_configs,
range(num_processes),
[num_val_processes for _ in range(num_processes)],
[num_processes for _ in range(num_processes)],
target_dims,
[log_dir for _ in range(num_processes)],
[visdom_name for _ in range(num_processes)],
[visdom_log_file for _ in range(num_processes)],
[vis_interval for _ in range(num_processes)],
[visdom_server for _ in range(num_processes)],
[visdom_port for _ in range(num_processes)],
[
swap_building_k_episodes
for _ in range(num_processes)
],
[map_kwargs for _ in range(num_processes)],
[reward_kwargs for _ in range(num_processes)],
should_record,
[seed + i for i in range(num_processes)],
[test_mode for _ in range(num_processes)],
train_datasets + val_datasets,
[scenario_kwargs for _ in range(num_processes)],
))),
preprocessing_fn=preprocessing_fn,
)
envs.observation_space = envs.observation_spaces[0]
envs.action_space = spaces.Discrete(3)
envs.reward_range = None
envs.metadata = None
envs.is_embodied = True
return envs