def init(self):
"""
Actually instantiate the env instances.
Also creates ActorState objects that hold the state of individual actors in (potentially) multi-agent envs.
"""
for env_i in range(self.num_envs):
vector_idx = self.split_idx * self.num_envs + env_i
# global env id within the entire system
env_id = self.worker_idx * self.cfg.num_envs_per_worker + vector_idx
env_config = AttrDict(
worker_index=self.worker_idx,
vector_index=vector_idx,
env_id=env_id,
)
# log.info('Creating env %r... %d-%d-%d', env_config, self.worker_idx, self.split_idx, env_i)
env = make_env_func(self.cfg, env_config=env_config)
env.seed(env_id)
self.envs.append(env)
actor_states_env, episode_rewards_env = [], []
for agent_idx in range(self.num_agents):
agent_traj_tensors = self.traj_tensors.index(
(env_i, agent_idx))
actor_state = ActorState(
self.cfg,
env,
self.worker_idx,
self.split_idx,
env_i,
agent_idx,
agent_traj_tensors,
self.num_traj_buffers,
self.policy_outputs,
self.policy_output_tensors[env_i, agent_idx],
self.pbt_reward_shaping,
self.policy_mgr,
)
actor_states_env.append(actor_state)
episode_rewards_env.append(0.0)
self.actor_states.append(actor_states_env)
self.episode_rewards.append(episode_rewards_env)
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)
if self.cfg.sampler_worker_gpus:
set_gpus_for_process(
proc_idx,
num_gpus_per_process=1,
process_type='sampler_proc',
gpu_mask=self.cfg.sampler_worker_gpus,
)
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 = make_env_func(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]
# sample a lot of random actions once, otherwise it is pretty slow in Python
total_random_actions = 500
actions = [[
env.action_space.sample() for _ in range(env.num_agents)
] for _ in range(total_random_actions)]
action_i = 0
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):
with timing.add_time(f'{env_key[env_idx]}.step'):
obs, rewards, dones, infos = env.step(
actions[action_i])
action_i = (action_i +
1) % total_random_actions
num_frames = sum(
[info.get('num_frames', 1) for info in infos])
total_env_frames += num_frames
episode_length[env_idx] += num_frames
if all(dones):
episode_lengths[env_idx].append(
episode_length[env_idx] / env.num_agents)
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))
if proc_idx == 0:
log.debug('Memory usage: %.4f Mb',
memory_consumption_mb())
# 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 __init__(self, cfg):
super().__init__(cfg)
# we should not use CUDA in the main thread, only on the workers
set_global_cuda_envvars(cfg)
tmp_env = make_env_func(self.cfg, env_config=None)
self.obs_space = tmp_env.observation_space
self.action_space = tmp_env.action_space
self.num_agents = tmp_env.num_agents
self.reward_shaping_scheme = None
if self.cfg.with_pbt:
self.reward_shaping_scheme = get_default_reward_shaping(tmp_env)
tmp_env.close()
# shared memory allocation
self.traj_buffers = SharedBuffers(self.cfg, self.num_agents,
self.obs_space, self.action_space)
self.actor_workers = None
self.report_queue = MpQueue(40 * 1000 * 1000)
self.policy_workers = dict()
self.policy_queues = dict()
self.learner_workers = dict()
self.workers_by_handle = None
self.policy_inputs = [[] for _ in range(self.cfg.num_policies)]
self.policy_outputs = dict()
for worker_idx in range(self.cfg.num_workers):
for split_idx in range(self.cfg.worker_num_splits):
self.policy_outputs[(worker_idx, split_idx)] = dict()
self.policy_avg_stats = dict()
self.policy_lag = [dict() for _ in range(self.cfg.num_policies)]
self.last_timing = dict()
self.env_steps = dict()
self.samples_collected = [0 for _ in range(self.cfg.num_policies)]
self.total_env_steps_since_resume = 0
# currently this applies only to the current run, not experiment as a whole
# to change this behavior we'd need to save the state of the main loop to a filesystem
self.total_train_seconds = 0
self.last_report = time.time()
self.last_experiment_summaries = 0
self.report_interval = 5.0 # sec
self.experiment_summaries_interval = self.cfg.experiment_summaries_interval # sec
self.avg_stats_intervals = (2, 12, 60
) # 10 seconds, 1 minute, 5 minutes
self.fps_stats = deque([], maxlen=max(self.avg_stats_intervals))
self.throughput_stats = [
deque([], maxlen=5) for _ in range(self.cfg.num_policies)
]
self.avg_stats = dict()
self.stats = dict() # regular (non-averaged) stats
init_wandb(self.cfg)
self.writers = dict()
writer_keys = list(range(self.cfg.num_policies))
for key in writer_keys:
summary_dir = join(summaries_dir(experiment_dir(cfg=self.cfg)),
str(key))
summary_dir = ensure_dir_exists(summary_dir)
self.writers[key] = SummaryWriter(summary_dir, flush_secs=20)
self.pbt = PopulationBasedTraining(self.cfg,
self.reward_shaping_scheme,
self.writers)