def train_distance(params, env_id):
def make_env_func():
e = create_env(env_id)
return e
agent = AgentTMAX(make_env_func, params)
agent.initialize()
multi_env = None
try:
multi_env = MultiEnv(
params.num_envs,
params.num_workers,
make_env_func=agent.make_env_func,
stats_episodes=params.stats_episodes,
)
train_loop(agent, multi_env)
except (Exception, KeyboardInterrupt, SystemExit):
log.exception('Interrupt...')
finally:
log.info('Closing env...')
if multi_env is not None:
multi_env.close()
agent.finalize()
return 0
def safe_get(q, timeout=1e6, msg='Queue timeout'):
"""Using queue.get() with timeout is necessary, otherwise KeyboardInterrupt is not handled."""
while True:
try:
return q.get(timeout=timeout)
except Empty:
log.exception(msg)
def evaluate_experiment(env_id, experiment_name, num_envs=96):
# fixed seeds
random.seed(0)
np.random.seed(0)
tf.random.set_random_seed(0)
params = AgentTMAX.Params(experiment_name)
params = params.load()
params.seed = 0
# for faster evaluation
params.num_envs = num_envs
params.num_workers = 32 if num_envs >= 32 else num_envs
def make_env_func():
e = create_env(env_id, skip_frames=True)
e.seed(0)
return e
agent = AgentTMAX(make_env_func, params)
agent.initialize()
rate, speed = 0, -1
multi_env = None
try:
multi_env = MultiEnv(
params.num_envs,
params.num_workers,
make_env_func=make_env_func,
stats_episodes=params.stats_episodes,
)
success, avg_speed = evaluate_locomotion_agent(agent, multi_env)
log.info('Finished evaluating experiment %s', experiment_name)
rate = np.mean(success)
speed = -1
avg_speed = [s for s in avg_speed if s > 0]
if len(avg_speed) > 0:
speed = np.mean(avg_speed)
log.info('Success rate %.1f%%, avg. speed %.2f edges/frame', rate * 100, speed)
except (Exception, KeyboardInterrupt, SystemExit):
log.exception('Interrupt...')
finally:
log.info('Closing env...')
if multi_env is not None:
multi_env.close()
agent.finalize()
return rate, speed
def learn(self, step_callback=None):
try:
multi_env = MultiEnv(
self.params.num_envs,
self.params.num_workers,
make_env_func=self.make_env_func,
stats_episodes=self.params.stats_episodes,
)
self._learn_loop(multi_env, step_callback)
except Exception as exc:
log.exception(exc)
finally:
log.info('Closing env...')
multi_env.close()
def load_checkpoint(checkpoints, device):
if len(checkpoints) <= 0:
log.warning('No checkpoints found')
return None
else:
latest_checkpoint = checkpoints[-1]
# extra safety mechanism to recover from spurious filesystem errors
num_attempts = 3
for attempt in range(num_attempts):
try:
log.warning('Loading state from checkpoint %s...', latest_checkpoint)
checkpoint_dict = torch.load(latest_checkpoint, map_location=device)
return checkpoint_dict
except Exception:
log.exception(f'Could not load from checkpoint, attempt {attempt}')
def read_seeds_file(filename, has_keys):
seeds = []
with open(filename, 'r') as seed_file:
lines = seed_file.readlines()
for line in lines:
try:
if has_keys:
seed, cache_key = line.split(' ')
else:
seed = line
seed = int(seed)
seeds.append(seed)
except Exception:
log.error(
'Could not read seed value from the file! File potentially corrupted'
)
log.exception('Exception when reading seeds file')
return seeds
def learn(self):
status = TrainStatus.SUCCESS
multi_env = None
try:
multi_env = MultiEnv(
self.params.num_envs,
self.params.num_workers,
make_env_func=self.make_env_func,
stats_episodes=self.params.stats_episodes,
)
self._learn_loop(multi_env)
except (Exception, KeyboardInterrupt, SystemExit):
log.exception('Interrupt...')
status = TrainStatus.FAILURE
finally:
log.info('Closing env...')
if multi_env is not None:
multi_env.close()
return status
def main():
def make_env():
env = DmlabGymEnv()
return env
num_envs = 64
num_workers = 16
multi_env = MultiEnv(num_envs, num_workers, make_env, stats_episodes=100)
num_resets = 0
try:
while True:
multi_env.reset()
num_resets += 1
num_envs_generated = num_resets * num_envs
log.info('Generated %d environments...', num_envs_generated)
except (Exception, KeyboardInterrupt, SystemExit):
log.exception('Interrupt...')
finally:
log.info('Closing env...')
multi_env.close()
return 0
def _run(self):
"""
Main loop of the actor worker (rollout worker).
Process tasks (mainly ROLLOUT_STEP) until we get the termination signal, which usually means end of training.
Currently there is no mechanism to restart dead workers if something bad happens during training. We can only
retry on the initial reset(). This is definitely something to work on.
"""
log.info('Initializing vector env runner %d...', self.worker_idx)
# workers should ignore Ctrl+C because the termination is handled in the event loop by a special msg
signal.signal(signal.SIGINT, signal.SIG_IGN)
torch.multiprocessing.set_sharing_strategy('file_system')
timing = Timing()
last_report = time.time()
with torch.no_grad():
while not self.terminate:
try:
try:
with timing.add_time('waiting'), timing.timeit('wait_actor'):
tasks = self.task_queue.get_many(timeout=0.1)
except Empty:
tasks = []
for task in tasks:
task_type, data = task
if task_type == TaskType.INIT:
self._init()
continue
if task_type == TaskType.TERMINATE:
self._terminate()
break
# handling actual workload
if task_type == TaskType.ROLLOUT_STEP:
if 'work' not in timing:
timing.waiting = 0 # measure waiting only after real work has started
with timing.add_time('work'), timing.timeit('one_step'):
self._advance_rollouts(data, timing)
elif task_type == TaskType.RESET:
with timing.add_time('reset'):
self._handle_reset()
elif task_type == TaskType.PBT:
self._process_pbt_task(data)
if time.time() - last_report > 5.0 and 'one_step' in timing:
timing_stats = dict(wait_actor=timing.wait_actor, step_actor=timing.one_step)
memory_mb = memory_consumption_mb()
stats = dict(memory_actor=memory_mb)
self.report_queue.put(dict(timing=timing_stats, stats=stats))
last_report = time.time()
except RuntimeError as exc:
log.warning('Error while processing data w: %d, exception: %s', self.worker_idx, exc)
log.warning('Terminate process...')
self.terminate = True
self.report_queue.put(dict(critical_error=self.worker_idx))
except KeyboardInterrupt:
self.terminate = True
except:
log.exception('Unknown exception in rollout worker')
self.terminate = True
if self.worker_idx <= 1:
time.sleep(0.1)
log.info(
'Env runner %d, CPU aff. %r, rollouts %d: timing %s',
self.worker_idx, psutil.Process().cpu_affinity(), self.num_complete_rollouts, timing,
)
def run(self):
"""
This function contains the main loop of the algorithm, as well as initialization/cleanup code.
:return: ExperimentStatus (SUCCESS, FAILURE, INTERRUPTED). Useful in testing.
"""
status = ExperimentStatus.SUCCESS
if os.path.isfile(done_filename(self.cfg)):
log.warning(
'Training already finished! Remove "done" file to continue training'
)
return status
self.init_workers()
self.init_pbt()
self.finish_initialization()
log.info('Collecting experience...')
timing = Timing()
with timing.timeit('experience'):
# noinspection PyBroadException
try:
while not self._should_end_training():
try:
reports = self.report_queue.get_many(timeout=0.1)
for report in reports:
self.process_report(report)
except Empty:
pass
if time.time() - self.last_report > self.report_interval:
self.report()
now = time.time()
self.total_train_seconds += now - self.last_report
self.last_report = now
self.pbt.update(self.env_steps, self.policy_avg_stats)
except Exception:
log.exception('Exception in driver loop')
status = ExperimentStatus.FAILURE
except KeyboardInterrupt:
log.warning(
'Keyboard interrupt detected in driver loop, exiting...')
status = ExperimentStatus.INTERRUPTED
for learner in self.learner_workers.values():
# timeout is needed here because some environments may crash on KeyboardInterrupt (e.g. VizDoom)
# Therefore the learner train loop will never do another iteration and will never save the model.
# This is not an issue with normal exit, e.g. due to desired number of frames reached.
learner.save_model(timeout=5.0)
all_workers = self.actor_workers
for workers in self.policy_workers.values():
all_workers.extend(workers)
all_workers.extend(self.learner_workers.values())
child_processes = list_child_processes()
time.sleep(0.1)
log.debug('Closing workers...')
for i, w in enumerate(all_workers):
w.close()
time.sleep(0.01)
for i, w in enumerate(all_workers):
w.join()
log.debug('Workers joined!')
# VizDoom processes often refuse to die for an unidentified reason, so we're force killing them with a hack
kill_processes(child_processes)
fps = self.total_env_steps_since_resume / timing.experience
log.info('Collected %r, FPS: %.1f', self.env_steps, fps)
log.info('Timing: %s', timing)
if self._should_end_training():
with open(done_filename(self.cfg), 'w') as fobj:
fobj.write(f'{self.env_steps}')
time.sleep(0.5)
log.info('Done!')
return status
def sample(self, proc_idx):
# workers should ignore Ctrl+C because the termination is handled in the event loop by a special msg
signal.signal(signal.SIGINT, signal.SIG_IGN)
timing = Timing()
from threadpoolctl import threadpool_limits
with threadpool_limits(limits=1, user_api=None):
if self.cfg.set_workers_cpu_affinity:
set_process_cpu_affinity(proc_idx, self.cfg.num_workers)
initial_cpu_affinity = psutil.Process().cpu_affinity(
) if platform != 'darwin' else None
psutil.Process().nice(10)
with timing.timeit('env_init'):
envs = []
env_key = ['env' for _ in range(self.cfg.num_envs_per_worker)]
for env_idx in range(self.cfg.num_envs_per_worker):
global_env_id = proc_idx * self.cfg.num_envs_per_worker + env_idx
env_config = AttrDict(worker_index=proc_idx,
vector_index=env_idx,
env_id=global_env_id)
env = create_env(self.cfg.env,
cfg=self.cfg,
env_config=env_config)
log.debug(
'CPU affinity after create_env: %r',
psutil.Process().cpu_affinity()
if platform != 'darwin' else 'MacOS - None')
env.seed(global_env_id)
envs.append(env)
# this is to track the performance for individual DMLab levels
if hasattr(env.unwrapped, 'level_name'):
env_key[env_idx] = env.unwrapped.level_name
episode_length = [0 for _ in envs]
episode_lengths = [deque([], maxlen=20) for _ in envs]
try:
with timing.timeit('first_reset'):
for env_idx, env in enumerate(envs):
env.reset()
log.info('Process %d finished resetting %d/%d envs',
proc_idx, env_idx + 1, len(envs))
self.report_queue.put(
dict(proc_idx=proc_idx, finished_reset=True))
self.start_event.wait()
with timing.timeit('work'):
last_report = last_report_frames = total_env_frames = 0
while not self.terminate.value and total_env_frames < self.cfg.sample_env_frames_per_worker:
for env_idx, env in enumerate(envs):
action = env.action_space.sample()
with timing.add_time(f'{env_key[env_idx]}.step'):
obs, reward, done, info = env.step(action)
num_frames = info.get('num_frames', 1)
total_env_frames += num_frames
episode_length[env_idx] += num_frames
if done:
with timing.add_time(
f'{env_key[env_idx]}.reset'):
env.reset()
episode_lengths[env_idx].append(
episode_length[env_idx])
episode_length[env_idx] = 0
with timing.add_time('report'):
now = time.time()
if now - last_report > self.report_every_sec:
last_report = now
frames_since_last_report = total_env_frames - last_report_frames
last_report_frames = total_env_frames
self.report_queue.put(
dict(proc_idx=proc_idx,
env_frames=frames_since_last_report))
# Extra check to make sure cpu affinity is preserved throughout the execution.
# I observed weird effect when some environments tried to alter affinity of the current process, leading
# to decreased performance.
# This can be caused by some interactions between deep learning libs, OpenCV, MKL, OpenMP, etc.
# At least user should know about it if this is happening.
cpu_affinity = psutil.Process().cpu_affinity(
) if platform != 'darwin' else None
assert initial_cpu_affinity == cpu_affinity, \
f'Worker CPU affinity was changed from {initial_cpu_affinity} to {cpu_affinity}!' \
f'This can significantly affect performance!'
except:
log.exception('Unknown exception')
log.error('Unknown exception in worker %d, terminating...',
proc_idx)
self.report_queue.put(dict(proc_idx=proc_idx, crash=True))
time.sleep(proc_idx * 0.01 + 0.01)
log.info('Process %d finished sampling. Timing: %s', proc_idx,
timing)
for env_idx, env in enumerate(envs):
if len(episode_lengths[env_idx]) > 0:
log.warning('Level %s avg episode len %d',
env_key[env_idx],
np.mean(episode_lengths[env_idx]))
for env in envs:
env.close()
def _run(self):
# workers should ignore Ctrl+C because the termination is handled in the event loop by a special msg
signal.signal(signal.SIGINT, signal.SIG_IGN)
psutil.Process().nice(min(self.cfg.default_niceness + 2, 20))
cuda_envvars(self.policy_id)
torch.multiprocessing.set_sharing_strategy('file_system')
timing = Timing()
with timing.timeit('init'):
# initialize the Torch modules
log.info('Initializing model on the policy worker %d-%d...',
self.policy_id, self.worker_idx)
torch.set_num_threads(1)
if self.cfg.device == 'gpu':
# we should already see only one CUDA device, because of env vars
assert torch.cuda.device_count() == 1
self.device = torch.device('cuda', index=0)
else:
self.device = torch.device('cpu')
self.actor_critic = create_actor_critic(self.cfg, self.obs_space,
self.action_space, timing)
self.actor_critic.model_to_device(self.device)
for p in self.actor_critic.parameters():
p.requires_grad = False # we don't train anything here
log.info('Initialized model on the policy worker %d-%d!',
self.policy_id, self.worker_idx)
last_report = last_cache_cleanup = time.time()
last_report_samples = 0
request_count = deque(maxlen=50)
# very conservative limit on the minimum number of requests to wait for
# this will almost guarantee that the system will continue collecting experience
# at max rate even when 2/3 of workers are stuck for some reason (e.g. doing a long env reset)
# Although if your workflow involves very lengthy operations that often freeze workers, it can be beneficial
# to set min_num_requests to 1 (at a cost of potential inefficiency, i.e. policy worker will use very small
# batches)
min_num_requests = self.cfg.num_workers // (
self.cfg.num_policies * self.cfg.policy_workers_per_policy)
min_num_requests //= 3
min_num_requests = max(1, min_num_requests)
# Again, very conservative timer. Only wait a little bit, then continue operation.
wait_for_min_requests = 0.025
while not self.terminate:
try:
while self.stop_experience_collection[self.policy_id]:
with self.resume_experience_collection_cv:
self.resume_experience_collection_cv.wait(timeout=0.05)
waiting_started = time.time()
while len(self.requests) < min_num_requests and time.time(
) - waiting_started < wait_for_min_requests:
try:
with timing.timeit('wait_policy'), timing.add_time(
'wait_policy_total'):
policy_requests = self.policy_queue.get_many(
timeout=0.005)
self.requests.extend(policy_requests)
except Empty:
pass
self._update_weights(timing)
with timing.timeit('one_step'), timing.add_time(
'handle_policy_step'):
if self.initialized:
if len(self.requests) > 0:
request_count.append(len(self.requests))
self._handle_policy_steps(timing)
try:
task_type, data = self.task_queue.get_nowait()
# task from the task_queue
if task_type == TaskType.INIT:
self._init()
elif task_type == TaskType.TERMINATE:
self.terminate = True
break
elif task_type == TaskType.INIT_MODEL:
self._init_model(data)
self.task_queue.task_done()
except Empty:
pass
if time.time() - last_report > 3.0 and 'one_step' in timing:
timing_stats = dict(wait_policy=timing.wait_policy,
step_policy=timing.one_step)
samples_since_last_report = self.total_num_samples - last_report_samples
stats = memory_stats('policy_worker', self.device)
if len(request_count) > 0:
stats['avg_request_count'] = np.mean(request_count)
self.report_queue.put(
dict(
timing=timing_stats,
samples=samples_since_last_report,
policy_id=self.policy_id,
stats=stats,
))
last_report = time.time()
last_report_samples = self.total_num_samples
if time.time() - last_cache_cleanup > 300.0 or (
not self.cfg.benchmark
and self.total_num_samples < 1000):
if self.cfg.device == 'gpu':
torch.cuda.empty_cache()
last_cache_cleanup = time.time()
except KeyboardInterrupt:
log.warning('Keyboard interrupt detected on worker %d-%d',
self.policy_id, self.worker_idx)
self.terminate = True
except:
log.exception('Unknown exception on policy worker')
self.terminate = True
time.sleep(0.2)
log.info('Policy worker avg. requests %.2f, timing: %s',
np.mean(request_count), timing)
def sample(self, proc_idx):
# workers should ignore Ctrl+C because the termination is handled in the event loop by a special msg
signal.signal(signal.SIGINT, signal.SIG_IGN)
timing = Timing()
psutil.Process().nice(10)
num_envs = len(DMLAB30_LEVELS_THAT_USE_LEVEL_CACHE)
assert self.cfg.num_workers % num_envs == 0, f'should have an integer number of workers per env, e.g. {1 * num_envs}, {2 * num_envs}, etc...'
assert self.cfg.num_envs_per_worker == 1, 'use populate_cache with 1 env per worker'
with timing.timeit('env_init'):
env_key = 'env'
env_desired_num_levels = 0
global_env_id = proc_idx * self.cfg.num_envs_per_worker
env_config = AttrDict(worker_index=proc_idx, vector_index=0, env_id=global_env_id)
env = create_env(self.cfg.env, cfg=self.cfg, env_config=env_config)
env.seed(global_env_id)
# this is to track the performance for individual DMLab levels
if hasattr(env.unwrapped, 'level_name'):
env_key = env.unwrapped.level_name
env_level = env.unwrapped.level
approx_num_episodes_per_1b_frames = DMLAB30_APPROX_NUM_EPISODES_PER_BILLION_FRAMES[
env_key]
num_billions = DESIRED_TRAINING_LENGTH / int(1e9)
num_workers_for_env = self.cfg.num_workers // num_envs
env_desired_num_levels = int(
(approx_num_episodes_per_1b_frames * num_billions) / num_workers_for_env)
env_num_levels_generated = len(dmlab_level_cache.DMLAB_GLOBAL_LEVEL_CACHE[0].
all_seeds[env_level]) // num_workers_for_env
log.warning('Worker %d (env %s) generated %d/%d levels!',
proc_idx,
env_key,
env_num_levels_generated,
env_desired_num_levels)
time.sleep(4)
env.reset()
env_uses_level_cache = env.unwrapped.env_uses_level_cache
self.report_queue.put(dict(proc_idx=proc_idx, finished_reset=True))
self.start_event.wait()
try:
with timing.timeit('work'):
last_report = last_report_frames = total_env_frames = 0
while not self.terminate.value and total_env_frames < self.cfg.sample_env_frames_per_worker:
action = env.action_space.sample()
with timing.add_time(f'{env_key}.step'):
env.step(action)
total_env_frames += 1
with timing.add_time(f'{env_key}.reset'):
env.reset()
env_num_levels_generated += 1
log.debug('Env %s done %d/%d resets',
env_key,
env_num_levels_generated,
env_desired_num_levels)
if env_num_levels_generated >= env_desired_num_levels:
log.debug('%s finished %d/%d resets, sleeping...',
env_key,
env_num_levels_generated,
env_desired_num_levels)
time.sleep(30) # free up CPU time for other envs
# if env does not use level cache, there is no need to run it
# let other workers proceed
if not env_uses_level_cache:
log.debug('Env %s does not require cache, sleeping...', env_key)
time.sleep(200)
with timing.add_time('report'):
now = time.time()
if now - last_report > self.report_every_sec:
last_report = now
frames_since_last_report = total_env_frames - last_report_frames
last_report_frames = total_env_frames
self.report_queue.put(
dict(proc_idx=proc_idx, env_frames=frames_since_last_report))
if get_free_disk_space_mb(self.cfg) < 3 * 1024:
log.error('Not enough disk space! %d',
get_free_disk_space_mb(self.cfg))
time.sleep(200)
except:
log.exception('Unknown exception')
log.error('Unknown exception in worker %d, terminating...', proc_idx)
self.report_queue.put(dict(proc_idx=proc_idx, crash=True))
time.sleep(proc_idx * 0.1 + 0.1)
log.info('Process %d finished sampling. Timing: %s', proc_idx, timing)
env.close()
