def set_gpus_for_process(process_idx, num_gpus_per_process, process_type, gpu_mask=None): available_gpus = get_available_gpus() if gpu_mask is not None: assert len(available_gpus) >= len(available_gpus) available_gpus = [available_gpus[g] for g in gpu_mask] num_gpus = len(available_gpus) gpus_to_use = [] if num_gpus == 0: os.environ[CUDA_ENVVAR] = '' log.debug('Not using GPUs for %s process %d', process_type, process_idx) else: first_gpu_idx = process_idx * num_gpus_per_process for i in range(num_gpus_per_process): index_mod_num_gpus = (first_gpu_idx + i) % num_gpus gpus_to_use.append(available_gpus[index_mod_num_gpus]) os.environ[CUDA_ENVVAR] = ','.join([str(g) for g in gpus_to_use]) log.info( 'Set environment var %s to %r for %s process %d', CUDA_ENVVAR, os.environ[CUDA_ENVVAR], process_type, process_idx, ) log.debug('Visible devices: %r', torch.cuda.device_count()) return gpus_to_use
def cat(self, dict_of_tensor_arrays, macro_batch_size, use_pinned_memory, timing): """ Here 'macro_batch' is the overall size of experience per iteration. Macro-batch = mini-batch * num_batches_per_iteration """ tensor_batch = self.batch_pool.get() if tensor_batch is not None: old_batch_size = tensor_batch_size(tensor_batch) if old_batch_size != macro_batch_size: # this can happen due to PBT changing batch size during the experiment log.warning('Tensor macro-batch size changed from %d to %d!', old_batch_size, macro_batch_size) log.warning('Discarding the cached tensor batch!') del tensor_batch tensor_batch = None if tensor_batch is None: tensor_batch = copy_dict_structure(dict_of_tensor_arrays) log.info('Allocating new CPU tensor batch (could not get from the pool)') for d1, cache_d, key, tensor_arr, _ in iter_dicts_recursively(dict_of_tensor_arrays, tensor_batch): cache_d[key] = torch.cat(tensor_arr, dim=0) if use_pinned_memory: cache_d[key] = cache_d[key].pin_memory() else: with timing.add_time('batcher_mem'): for d1, cache_d, key, tensor_arr, cache_t in iter_dicts_recursively(dict_of_tensor_arrays, tensor_batch): offset = 0 for t in tensor_arr: first_dim = t.shape[0] cache_t[offset:offset + first_dim].copy_(t) offset += first_dim return tensor_batch
def _parse_info(self, info, done): if self.reward_shaping_scheme is None: # skip reward calculation return 0.0 # by default these are negative values if no weapon is selected selected_weapon = info.get('SELECTED_WEAPON', 0.0) selected_weapon = int(max(0, selected_weapon)) selected_weapon_ammo = float(max(0.0, info.get('SELECTED_WEAPON_AMMO', 0.0))) self.selected_weapon.append(selected_weapon) was_dead = self.prev_dead is_alive = not info.get('DEAD', 0.0) just_respawned = was_dead and is_alive shaping_reward = 0.0 if not done and not just_respawned: shaping_reward, deltas = self._delta_rewards(info) shaping_reward += self._selected_weapon_rewards( selected_weapon, selected_weapon_ammo, deltas, ) if abs(shaping_reward) > 2.5 and not self.print_once: log.info('Large shaping reward %.3f for %r', shaping_reward, deltas) self.print_once = True if done and 'FRAGCOUNT' in self.reward_structure: sorted_rew = sorted(self.reward_structure.items(), key=operator.itemgetter(1)) sum_rew = sum(r for key, r in sorted_rew) sorted_rew = {key: f'{r:.3f}' for key, r in sorted_rew} log.info('Sum rewards: %.3f, reward structure: %r', sum_rew, sorted_rew) return shaping_reward
def reset(self): if self._episode_recording_dir is not None and self._record_id > 0: # save actions to text file with open(join(self._episode_recording_dir, 'actions.json'), 'w') as actions_file: json.dump(self._recorded_actions, actions_file) # rename previous episode dir reward = self._recorded_episode_reward + self._recorded_episode_shaping_reward new_dir_name = self._episode_recording_dir + f'_r{reward:.2f}' os.rename(self._episode_recording_dir, new_dir_name) log.info( 'Finished recording %s (rew %.3f, shaping %.3f)', new_dir_name, reward, self._recorded_episode_shaping_reward, ) dir_name = f'ep_{self._record_id:03d}_p{self._player_id}' self._episode_recording_dir = join(self._record_to, dir_name) ensure_dir_exists(self._episode_recording_dir) self._record_id += 1 self._frame_id = 0 self._recorded_episode_reward = 0 self._recorded_episode_shaping_reward = 0 self._recorded_actions = [] return self.env.reset()
def _init(self): """ Initialize env runners, that actually do all the work. Also we're doing some utility stuff here, e.g. setting process affinity (this is a performance optimization). """ log.info('Initializing envs for env runner %d...', self.worker_idx) if self.cfg.force_envs_single_thread: from threadpoolctl import threadpool_limits threadpool_limits(limits=1, user_api=None) if self.cfg.set_workers_cpu_affinity: set_process_cpu_affinity(self.worker_idx, self.cfg.num_workers) psutil.Process().nice(min(self.cfg.default_niceness + 10, 20)) self.env_runners = [] for split_idx in range(self.num_splits): env_runner = VectorEnvRunner( self.cfg, self.vector_size // self.num_splits, self.worker_idx, split_idx, self.num_agents, self.shared_buffers, self.reward_shaping, ) env_runner.init() self.env_runners.append(env_runner)
def calc_num_trajectory_buffers(self): """ This calculates the number of shared trajectory (rollout) buffers required by the system to operate without interruptions. This consists of: 1) at least one trajectory buffer per agent, such that we always have a location to save new experience 2) a few trajectory buffers to hold data currently processed by the learner (including potential backlog) 3) (potentially) some extra trajectory buffers to keep the system operational. These might be required i.e. in multi-agent envs when some agents are deactivated during the rollout. Such agents together with the learner can hold on to many buffers, such that active agents might not have enough free buffers to continue collecting the experience. """ # Add a traj buffer for each agent num_traj_buffers = (self.cfg.num_workers + 1) * self.cfg.num_envs_per_worker * self.num_agents max_minibatches_to_accumulate = self.cfg.num_minibatches_to_accumulate if max_minibatches_to_accumulate == -1: # default value max_minibatches_to_accumulate = 2 * self.cfg.num_batches_per_iteration # Let each learner accumulate enough full sets of experience to pause learning max_experience_on_learners = max_minibatches_to_accumulate * self.cfg.batch_size * self.cfg.num_policies num_traj_buffers += max_experience_on_learners / self.cfg.rollout # Configurable excess ratio to be safe assert self.cfg.traj_buffers_excess_ratio >= 1.0 num_traj_buffers = self.cfg.traj_buffers_excess_ratio * num_traj_buffers num_traj_buffers = int(math.ceil(num_traj_buffers)) log.info('Using a total of %d trajectory buffers', num_traj_buffers) return num_traj_buffers
def _init(self, init_info): log.info('Initializing env for player %d, init_info: %r...', self.player_id, init_info) env = init_multiplayer_env(self.make_env_func, self.player_id, self.env_config, init_info) if self.reset_on_init: env.reset() return env
def _set_game_mode(self, mode): if mode == 'replay': self.game.set_mode(Mode.PLAYER) else: if self.async_mode: log.info('Starting in async mode! Use this only for testing, otherwise PLAYER mode is much faster') self.game.set_mode(Mode.ASYNC_PLAYER) else: self.game.set_mode(Mode.PLAYER)
def _vizdoom_variables_bug_workaround(self, info, done): """Some variables don't get reset to zero on game.new_episode(). This fixes it (also check overflow?).""" if done and 'DAMAGECOUNT' in info: log.info('DAMAGECOUNT value on done: %r', info.get('DAMAGECOUNT')) if self._last_episode_info is not None: bugged_vars = ['DEATHCOUNT', 'HITCOUNT', 'DAMAGECOUNT'] for v in bugged_vars: if v in info: info[v] -= self._last_episode_info.get(v, 0)
def run(cfg): cfg = maybe_load_from_checkpoint(cfg) algo = DmlabLevelGenerator(cfg) algo.initialize() status = algo.run() algo.finalize() log.info('Exit...') return status
def main(): env_name = 'doom_battle' env = create_env(env_name, cfg=default_cfg(env=env_name)) env.reset() done = False while not done: env.render() obs, rew, done, info = env.step(env.action_space.sample()) log.info('Done!')
def test_minigrid_env(self): env_name = 'MiniGrid-Empty-Random-5x5-v0' env = create_env(env_name, cfg=default_cfg(env=env_name)) log.info('Env action space: %r', env.action_space) log.info('Env obs space: %r', env.observation_space) env.reset() total_rew = 0 for i in range(1000): obs, rew, done, info = env.step(env.action_space.sample()) total_rew += rew if done: env.reset()
def _handle_reset(self): """ Reset all envs, one split at a time (double-buffering), and send requests to policy workers to get actions for the very first env step. """ for split_idx, env_runner in enumerate(self.env_runners): policy_inputs = env_runner.reset(self.report_queue) self._enqueue_policy_request(split_idx, policy_inputs) log.info('Finished reset for worker %d', self.worker_idx) safe_put(self.report_queue, dict(finished_reset=self.worker_idx), queue_name='report')
def test_doom_multiagent_parallel(self): num_workers = 16 workers = [] for i in range(num_workers): log.info('Starting worker #%d', i) worker = Process(target=self.doom_multiagent, args=(self.make_standard_dm, i, 200)) worker.start() workers.append(worker) time.sleep(0.01) for i in range(num_workers): workers[i].join()
def __init__(self, env, initial_difficulty=None): super().__init__(env) self._min_difficulty = 0 self._max_difficulty = 150 self._difficulty_step = 10 self._curr_difficulty = 20 if initial_difficulty is None else initial_difficulty self._difficulty_std = 10 log.info('Starting with bot difficulty %d', self._curr_difficulty) self._adaptive_curriculum = True if initial_difficulty == self._max_difficulty: log.debug('Starting at max difficulty, disable adaptive skill curriculum') self._adaptive_curriculum = False
def calc_num_trajectory_buffers(self): # calculate how many buffers are required per env runner to collect one "macro batch" for training # once macro batch is collected, all buffers will be released # we could have just copied the tensors on the learner to avoid this complicated logic, but it's better for # performance to keep data in shared buffers until they're needed samples_per_iteration = self.cfg.num_batches_per_iteration * self.cfg.batch_size * self.cfg.num_policies num_traj_buffers = samples_per_iteration / (self.cfg.num_workers * self.cfg.num_envs_per_worker * self.num_agents * self.cfg.rollout) # make sure we definitely have enough buffers to actually never wait # usually it'll be just two buffers and we swap back and forth num_traj_buffers *= 3 # make sure we have at least two to swap between so we never actually have to wait num_traj_buffers = math.ceil(max(num_traj_buffers, self.cfg.min_traj_buffers_per_worker)) log.info('Using %d sets of trajectory buffers', num_traj_buffers) return num_traj_buffers
def _update_weights(self, timing): learner_policy_version = self.policy_versions[self.policy_id].item() if self.latest_policy_version < learner_policy_version and self.shared_model_weights is not None: with timing.timeit('weight_update'): with self.policy_lock: self.actor_critic.load_state_dict(self.shared_model_weights) self.latest_policy_version = learner_policy_version if self.num_policy_updates % 10 == 0: log.info( 'Updated weights on worker %d-%d, policy_version %d (%.5f)', self.policy_id, self.worker_idx, self.latest_policy_version, timing.weight_update, ) self.num_policy_updates += 1
def dmlab_ensure_global_cache_initialized(experiment_dir, all_levels_for_experiment, num_policies, level_cache_dir): global DMLAB_GLOBAL_LEVEL_CACHE assert multiprocessing.current_process().name == 'MainProcess', \ 'make sure you initialize DMLab cache before child processes are forked' DMLAB_GLOBAL_LEVEL_CACHE = [] for policy_id in range(num_policies): # level cache is of course shared between independently training policies # it's easiest to achieve log.info('Initializing level cache for policy %d...', policy_id) cache = DmlabLevelCacheGlobal(level_cache_dir, experiment_dir, all_levels_for_experiment, policy_id) DMLAB_GLOBAL_LEVEL_CACHE.append(cache)
def replay(env, rec_path): doom = env.unwrapped doom.mode = 'replay' doom._ensure_initialized() doom.game.replay_episode(rec_path) episode_reward = 0 start = time.time() while not doom.game.is_episode_finished(): doom.game.advance_action() r = doom.game.get_last_reward() episode_reward += r log.info('Episode reward: %.3f, time so far: %.1f s', episode_reward, time.time() - start) log.info('Finishing replay') doom.close()
def __init__( self, worker_idx, policy_id, cfg, obs_space, action_space, shared_buffers, policy_queue, actor_queues, report_queue, task_queue, policy_lock, resume_experience_collection_cv ): log.info('Initializing policy worker %d for policy %d', worker_idx, policy_id) self.worker_idx = worker_idx self.policy_id = policy_id self.cfg = cfg self.obs_space = obs_space self.action_space = action_space self.device = None self.actor_critic = None self.shared_model_weights = None self.policy_lock = policy_lock self.resume_experience_collection_cv = resume_experience_collection_cv self.policy_queue = policy_queue self.actor_queues = actor_queues self.report_queue = report_queue # queue other components use to talk to this particular worker self.task_queue = task_queue self.initialized = False self.terminate = False self.initialized_event = multiprocessing.Event() self.initialized_event.clear() self.shared_buffers = shared_buffers self.tensors_individual_transitions = self.shared_buffers.tensors_individual_transitions self.policy_versions = shared_buffers.policy_versions self.stop_experience_collection = shared_buffers.stop_experience_collection self.latest_policy_version = -1 self.num_policy_updates = 0 self.requests = [] self.total_num_samples = 0 self.process = TorchProcess(target=self._run, daemon=True)
def reset(self, report_queue): """ Do the very first reset for all environments in a vector. Populate shared memory with initial obs. Note that this is called only once, at the very beginning of training. After this the envs should auto-reset. :param report_queue: we use report queue to monitor reset progress (see appo.py). This can be a lengthy process. :return: first requests for policy workers (to generate actions for the very first env step) """ for env_i, e in enumerate(self.envs): observations = e.reset() if self.cfg.decorrelate_envs_on_one_worker: env_i_split = self.num_envs * self.split_idx + env_i decorrelate_steps = self.cfg.rollout * env_i_split + self.cfg.rollout * random.randint( 0, 4) log.info('Decorrelating experience for %d frames...', decorrelate_steps) for decorrelate_step in range(decorrelate_steps): actions = [ e.action_space.sample() for _ in range(self.num_agents) ] observations, rew, dones, info = e.step(actions) for agent_i, obs in enumerate(observations): actor_state = self.actor_states[env_i][agent_i] actor_state.set_trajectory_data(dict(obs=obs), self.traj_buffer_idx, self.rollout_step) # rnn state is already initialized at zero # log.debug( # 'Reset progress w:%d-%d finished %d/%d, still initializing envs...', # self.worker_idx, self.split_idx, env_i + 1, len(self.envs), # ) safe_put(report_queue, dict(initialized_env=(self.worker_idx, self.split_idx, env_i)), queue_name='report') policy_request = self._format_policy_request() return policy_request
def __init__(self, num_agents, make_env_func, env_config, skip_frames): gym.Env.__init__(self) RewardShapingInterface.__init__(self) self.num_agents = num_agents log.debug('Multi agent env, num agents: %d', self.num_agents) self.skip_frames = skip_frames # number of frames to skip (1 = no skip) env = make_env_func( player_id=-1 ) # temporary env just to query observation_space and stuff self.action_space = env.action_space self.observation_space = env.observation_space self.default_reward_shaping = get_default_reward_shaping(env) env.close() self.current_reward_shaping = [ self.default_reward_shaping for _ in range(self.num_agents) ] self.make_env_func = make_env_func self.safe_init = env_config is not None and env_config.get( 'safe_init', False) if self.safe_init: sleep_seconds = env_config.worker_index * 1.0 log.info( 'Sleeping %.3f seconds to avoid creating all envs at once', sleep_seconds) time.sleep(sleep_seconds) log.info('Done sleeping at %d', env_config.worker_index) self.env_config = env_config self.workers = None # only needed when rendering self.enable_rendering = False self.last_obs = None self.reset_on_init = True self.initialized = False
def _advance_rollouts(self, data, timing): """ Process incoming request from policy worker. Use the data (policy outputs, actions) to advance the simulation by one step on the corresponding VectorEnvRunner. If we successfully managed to advance the simulation, send requests to policy workers to get actions for the next step. If we completed the entire rollout, also send request to the learner! :param data: request from the policy worker, containing actions and other policy outputs :param timing: profiling stuff """ split_idx = data['split_idx'] runner = self.env_runners[split_idx] policy_request, complete_rollouts, episodic_stats = runner.advance_rollouts( data, timing) with timing.add_time('complete_rollouts'): if complete_rollouts: self._enqueue_complete_rollouts(split_idx, complete_rollouts) if self.num_complete_rollouts == 0 and not self.cfg.benchmark: # we just finished our first complete rollouts, perfect time to wait for experience derorrelation # this guarantees that there won't be any "old" trajectories when we awaken delay = (float(self.worker_idx) / self.cfg.num_workers ) * self.cfg.decorrelate_experience_max_seconds log.info( 'Worker %d, sleep for %.3f sec to decorrelate experience collection', self.worker_idx, delay, ) time.sleep(delay) log.info('Worker %d awakens!', self.worker_idx) self.num_complete_rollouts += len( complete_rollouts['rollouts']) with timing.add_time('enqueue_policy_requests'): if policy_request is not None: self._enqueue_policy_request(split_idx, policy_request) if episodic_stats: self._report_stats(episodic_stats)
def test_env_performance(make_env, env_type, verbose=False): t = Timing() with t.timeit('init'): env = make_env(AttrDict({'worker_index': 0, 'vector_index': 0})) total_num_frames, frames = 10000, 0 with t.timeit('first_reset'): env.reset() t.reset = t.step = 1e-9 num_resets = 0 with t.timeit('experience'): while frames < total_num_frames: done = False start_reset = time.time() env.reset() t.reset += time.time() - start_reset num_resets += 1 while not done and frames < total_num_frames: start_step = time.time() if verbose: env.render() time.sleep(1.0 / 40) obs, rew, done, info = env.step(env.action_space.sample()) if verbose: log.info('Received reward %.3f', rew) t.step += time.time() - start_step frames += num_env_steps([info]) fps = total_num_frames / t.experience log.debug('%s performance:', env_type) log.debug('Took %.3f sec to collect %d frames on one CPU, %.1f FPS', t.experience, total_num_frames, fps) log.debug('Avg. reset time %.3f s', t.reset / num_resets) log.debug('Timing: %s', t) env.close()
def report(self, env_frames): now = time.time() self.last_report = now self.fps_stats.append((now, env_frames)) if len(self.fps_stats) <= 1: return fps = [] for avg_interval in self.avg_stats_intervals: past_moment, past_frames = self.fps_stats[max( 0, len(self.fps_stats) - 1 - avg_interval)] fps.append((env_frames - past_frames) / (now - past_moment)) fps_str = [] for interval, fps_value in zip(self.avg_stats_intervals, fps): fps_str.append( f'{int(interval * self.report_every_sec)} sec: {fps_value:.1f}' ) fps_str = f'({", ".join(fps_str)})' log.info('Sampling FPS: %s. Total frames collected: %d', fps_str, env_frames)
def step(self, action): obs, rew, done, info = self.env.step(action) if obs is None: return obs, rew, done, info self.orig_env_reward += rew shaping_rew = self._parse_info(info, done) rew += shaping_rew self.total_shaping_reward += shaping_rew if self.verbose: log.info('Original env reward before shaping: %.3f', self.orig_env_reward) player_id = 1 if hasattr(self.env.unwrapped, 'player_id'): player_id = self.env.unwrapped.player_id log.info( 'Total shaping reward is %.3f for %d (done %d)', self.total_shaping_reward, player_id, done, ) # remember new variable values for var_name in self.reward_shaping_scheme['delta'].keys(): self.prev_vars[var_name] = info.get(var_name, 0.0) self.prev_dead = not not info.get('DEAD', 0.0) # float -> bool if done: if self.true_reward_func is None: true_reward = self.orig_env_reward else: true_reward = self.true_reward_func(info) info['true_reward'] = true_reward return obs, rew, done, info
def doom_multiagent(make_multi_env, worker_index, num_steps=1000): env_config = AttrDict({ 'worker_index': worker_index, 'vector_index': 0, 'safe_init': False }) multi_env = make_multi_env(env_config) obs = multi_env.reset() visualize = False start = time.time() for i in range(num_steps): actions = [multi_env.action_space.sample()] * len(obs) obs, rew, dones, infos = multi_env.step(actions) if visualize: multi_env.render() if i % 100 == 0 or any(dones): log.info('Rew %r done %r info %r', rew, dones, infos) if all(dones): multi_env.reset() took = time.time() - start log.info('Took %.3f seconds for %d steps', took, num_steps) log.info('Server steps per second: %.1f', num_steps / took) log.info('Observations fps: %.1f', num_steps * multi_env.num_agents / took) log.info( 'Environment fps: %.1f', num_steps * multi_env.num_agents * multi_env.skip_frames / took) multi_env.close()
def test_voxel_env(self): env_name = 'voxel_env_Sokoban' env = create_env(env_name, cfg=default_cfg(env=env_name)) log.info('Env action space: %r', env.action_space) log.info('Env obs space: %r', env.observation_space) env.reset() total_rew = 0 for i in range(1000): obs, rew, done, info = env.step( [env.action_space.sample() for _ in range(env.num_agents)]) total_rew += sum(rew) log.info('Total rew: %.3f', total_rew)
def run(run_description, args): experiments = run_description.experiments max_parallel = args.max_parallel log.info('Starting processes with base cmds: %r', [e.cmd for e in experiments]) log.info('Max parallel processes is %d', max_parallel) log.info( 'Monitor log files using\n\n\ttail -f train_dir/%s/**/**/sf_log.txt\n\n', run_description.run_name) processes = [] processes_per_gpu = {g: [] for g in range(args.num_gpus)} experiments = run_description.generate_experiments(args.train_dir) next_experiment = next(experiments, None) def find_least_busy_gpu(): least_busy_gpu = None gpu_available_processes = 0 for gpu_id in range(args.num_gpus): available_processes = args.experiments_per_gpu - len( processes_per_gpu[gpu_id]) if available_processes > gpu_available_processes: gpu_available_processes = available_processes least_busy_gpu = gpu_id return least_busy_gpu, gpu_available_processes def can_squeeze_another_process(): if len(processes) >= max_parallel: return False if args.experiments_per_gpu > 0: least_busy_gpu, gpu_available_processes = find_least_busy_gpu() if gpu_available_processes <= 0: return False return True failed_processes = [] last_log_time = 0 log_interval = 3 # seconds while len(processes) > 0 or next_experiment is not None: while can_squeeze_another_process() and next_experiment is not None: cmd, name, root_dir, exp_env_vars = next_experiment cmd_tokens = cmd.split(' ') # workaround to make sure we're running the correct python executable from our virtual env if cmd_tokens[0].startswith('python'): cmd_tokens[0] = sys.executable log.debug('Using Python executable %s', cmd_tokens[0]) ensure_dir_exists(join(args.train_dir, root_dir)) envvars = os.environ.copy() best_gpu = None if args.experiments_per_gpu > 0: best_gpu, best_gpu_available_processes = find_least_busy_gpu() log.info( 'The least busy gpu is %d where we can run %d more processes', best_gpu, best_gpu_available_processes, ) envvars['CUDA_VISIBLE_DEVICES'] = f'{best_gpu}' log.info('Starting process %r', cmd_tokens) if exp_env_vars is not None: for key, value in exp_env_vars.items(): log.info('Adding env variable %r %r', key, value) envvars[str(key)] = str(value) process = subprocess.Popen(cmd_tokens, stdout=None, stderr=None, env=envvars) process.gpu_id = best_gpu process.proc_cmd = cmd processes.append(process) if process.gpu_id is not None: processes_per_gpu[process.gpu_id].append(process.proc_cmd) log.info('Started process %s on GPU %r', process.proc_cmd, process.gpu_id) log.info('Waiting for %d seconds before starting next process', args.pause_between) time.sleep(args.pause_between) next_experiment = next(experiments, None) remaining_processes = [] for process in processes: if process.poll() is None: remaining_processes.append(process) continue else: if process.gpu_id is not None: processes_per_gpu[process.gpu_id].remove(process.proc_cmd) log.info('Process %r finished with code %r', process.proc_cmd, process.returncode) if process.returncode != 0: failed_processes.append( (process.proc_cmd, process.pid, process.returncode)) log.error('WARNING: RETURN CODE IS %r', process.returncode) processes = remaining_processes if time.time() - last_log_time > log_interval: if failed_processes: log.error( 'Failed processes: %s', ', '.join([ f'PID: {p[1]} code: {p[2]}' for p in failed_processes ])) last_log_time = time.time() time.sleep(0.1) log.info('Done!') return 0
def enjoy(cfg, max_num_frames=1e9): cfg = load_from_checkpoint(cfg) render_action_repeat = cfg.render_action_repeat if cfg.render_action_repeat is not None else cfg.env_frameskip if render_action_repeat is None: log.warning('Not using action repeat!') render_action_repeat = 1 log.debug('Using action repeat %d during evaluation', render_action_repeat) cfg.env_frameskip = 1 # for evaluation cfg.num_envs = 1 def make_env_func(env_config): return create_env(cfg.env, cfg=cfg, env_config=env_config) env = make_env_func(AttrDict({'worker_index': 0, 'vector_index': 0})) # env.seed(0) is_multiagent = is_multiagent_env(env) if not is_multiagent: env = MultiAgentWrapper(env) if hasattr(env.unwrapped, 'reset_on_init'): # reset call ruins the demo recording for VizDoom env.unwrapped.reset_on_init = False actor_critic = create_actor_critic(cfg, env.observation_space, env.action_space) device = torch.device('cpu' if cfg.device == 'cpu' else 'cuda') actor_critic.model_to_device(device) policy_id = cfg.policy_index checkpoints = LearnerWorker.get_checkpoints( LearnerWorker.checkpoint_dir(cfg, policy_id)) checkpoint_dict = LearnerWorker.load_checkpoint(checkpoints, device) actor_critic.load_state_dict(checkpoint_dict['model']) episode_rewards = [deque([], maxlen=100) for _ in range(env.num_agents)] true_rewards = [deque([], maxlen=100) for _ in range(env.num_agents)] num_frames = 0 last_render_start = time.time() def max_frames_reached(frames): return max_num_frames is not None and frames > max_num_frames obs = env.reset() rnn_states = torch.zeros( [env.num_agents, get_hidden_size(cfg)], dtype=torch.float32, device=device) episode_reward = np.zeros(env.num_agents) finished_episode = [False] * env.num_agents with torch.no_grad(): while not max_frames_reached(num_frames): obs_torch = AttrDict(transform_dict_observations(obs)) for key, x in obs_torch.items(): obs_torch[key] = torch.from_numpy(x).to(device).float() policy_outputs = actor_critic(obs_torch, rnn_states, with_action_distribution=True) # sample actions from the distribution by default actions = policy_outputs.actions action_distribution = policy_outputs.action_distribution if isinstance(action_distribution, ContinuousActionDistribution): if not cfg.continuous_actions_sample: # TODO: add similar option for discrete actions actions = action_distribution.means actions = actions.cpu().numpy() rnn_states = policy_outputs.rnn_states for _ in range(render_action_repeat): if not cfg.no_render: target_delay = 1.0 / cfg.fps if cfg.fps > 0 else 0 current_delay = time.time() - last_render_start time_wait = target_delay - current_delay if time_wait > 0: # log.info('Wait time %.3f', time_wait) time.sleep(time_wait) last_render_start = time.time() env.render() obs, rew, done, infos = env.step(actions) episode_reward += rew num_frames += 1 for agent_i, done_flag in enumerate(done): if done_flag: finished_episode[agent_i] = True episode_rewards[agent_i].append( episode_reward[agent_i]) true_rewards[agent_i].append(infos[agent_i].get( 'true_reward', episode_reward[agent_i])) log.info( 'Episode finished for agent %d at %d frames. Reward: %.3f, true_reward: %.3f', agent_i, num_frames, episode_reward[agent_i], true_rewards[agent_i][-1]) rnn_states[agent_i] = torch.zeros( [get_hidden_size(cfg)], dtype=torch.float32, device=device) episode_reward[agent_i] = 0 # if episode terminated synchronously for all agents, pause a bit before starting a new one if all(done): if not cfg.no_render: env.render() time.sleep(0.05) if all(finished_episode): finished_episode = [False] * env.num_agents avg_episode_rewards_str, avg_true_reward_str = '', '' for agent_i in range(env.num_agents): avg_rew = np.mean(episode_rewards[agent_i]) avg_true_rew = np.mean(true_rewards[agent_i]) if not np.isnan(avg_rew): if avg_episode_rewards_str: avg_episode_rewards_str += ', ' avg_episode_rewards_str += f'#{agent_i}: {avg_rew:.3f}' if not np.isnan(avg_true_rew): if avg_true_reward_str: avg_true_reward_str += ', ' avg_true_reward_str += f'#{agent_i}: {avg_true_rew:.3f}' log.info('Avg episode rewards: %s, true rewards: %s', avg_episode_rewards_str, avg_true_reward_str) log.info( 'Avg episode reward: %.3f, avg true_reward: %.3f', np.mean([ np.mean(episode_rewards[i]) for i in range(env.num_agents) ]), np.mean([ np.mean(true_rewards[i]) for i in range(env.num_agents) ])) # VizDoom multiplayer stuff # for player in [1, 2, 3, 4, 5, 6, 7, 8]: # key = f'PLAYER{player}_FRAGCOUNT' # if key in infos[0]: # log.debug('Score for player %d: %r', player, infos[0][key]) env.close() return ExperimentStatus.SUCCESS, np.mean(episode_rewards)