def record_episodic_info(self, worker_name=None, **stats): video = stats.pop('video', None) if 'epslen' in stats: self.env_step += np.sum(stats['epslen']) if worker_name is not None: stats = {f'{k}_{worker_name}': v for k, v in stats.items()} with self._locker: self.store(**stats) if video is not None: video_summary(f'{self.name}/sim', video, step=self.env_step)
def train(agent, env, eval_env, replay): collect_fn = pkg.import_module('agent', algo=agent.name).collect collect = functools.partial(collect_fn, replay) _, step = replay.count_episodes() step = max(agent.env_step, step) runner = Runner(env, agent, step=step) def random_actor(*args, **kwargs): prev_action = random_actor.prev_action random_actor.prev_action = action = env.random_action() return action, {'prev_action': prev_action} random_actor.prev_action = np.zeros_like(env.random_action()) \ if isinstance(env.random_action(), np.ndarray) else 0 while not replay.good_to_learn(): step = runner.run(action_selector=random_actor, step_fn=collect) to_log = Every(agent.LOG_PERIOD) to_eval = Every(agent.EVAL_PERIOD) print('Training starts...') while step < int(agent.MAX_STEPS): start_step = step start_t = time.time() agent.learn_log(step) step = runner.run(step_fn=collect, nsteps=agent.TRAIN_PERIOD) duration = time.time() - start_t agent.store(fps=(step - start_step) / duration, tps=agent.N_UPDATES / duration) if to_eval(step): with TempStore(agent.get_states, agent.reset_states): score, epslen, video = evaluate(eval_env, agent, record=agent.RECORD, size=(64, 64)) if agent.RECORD: video_summary(f'{agent.name}/sim', video, step=step) agent.store(eval_score=score, eval_epslen=epslen) if to_log(step): agent.log(step) agent.save()
def train(agent, env, eval_env, buffer): collect_fn = pkg.import_module('agent', algo=agent.name).collect collect = functools.partial(collect_fn, buffer) step = agent.env_step runner = Runner(env, agent, step=step, nsteps=agent.N_STEPS) exp_buffer = get_expert_data(f'{buffer.DATA_PATH}-{env.name}') if step == 0 and agent.is_obs_normalized: print('Start to initialize running stats...') for _ in range(10): runner.run(action_selector=env.random_action, step_fn=collect) agent.update_obs_rms(np.concatenate(buffer['obs'])) agent.update_reward_rms(buffer['reward'], buffer['discount']) buffer.reset() buffer.clear() agent.save(print_terminal_info=True) runner.step = step # print("Initial running stats:", *[f'{k:.4g}' for k in agent.get_running_stats() if k]) to_log = Every(agent.LOG_PERIOD, agent.LOG_PERIOD) to_eval = Every(agent.EVAL_PERIOD) rt = Timer('run') tt = Timer('train') et = Timer('eval') lt = Timer('log') print('Training starts...') while step < agent.MAX_STEPS: start_env_step = agent.env_step agent.before_run(env) with rt: step = runner.run(step_fn=collect) agent.store(fps=(step - start_env_step) / rt.last()) buffer.reshape_to_sample() agent.disc_learn_log(exp_buffer) buffer.compute_reward_with_func(agent.compute_reward) buffer.reshape_to_store() # NOTE: normalizing rewards here may introduce some inconsistency # if normalized rewards is fed as an input to the network. # One can reconcile this by moving normalization to collect # or feeding the network with unnormalized rewards. # The latter is adopted in our implementation. # However, the following line currently doesn't store # a copy of unnormalized rewards agent.update_reward_rms(buffer['reward'], buffer['discount']) buffer.update('reward', agent.normalize_reward(buffer['reward']), field='all') agent.record_last_env_output(runner.env_output) value = agent.compute_value() buffer.finish(value) start_train_step = agent.train_step with tt: agent.learn_log(step) agent.store(tps=(agent.train_step - start_train_step) / tt.last()) buffer.reset() if to_eval(agent.train_step) or step > agent.MAX_STEPS: with TempStore(agent.get_states, agent.reset_states): with et: eval_score, eval_epslen, video = evaluate( eval_env, agent, n=agent.N_EVAL_EPISODES, record=agent.RECORD, size=(64, 64)) if agent.RECORD: video_summary(f'{agent.name}/sim', video, step=step) agent.store(eval_score=eval_score, eval_epslen=eval_epslen) if to_log(agent.train_step) and agent.contains_stats('score'): with lt: agent.store( **{ 'train_step': agent.train_step, 'time/run': rt.total(), 'time/train': tt.total(), 'time/eval': et.total(), 'time/log': lt.total(), 'time/run_mean': rt.average(), 'time/train_mean': tt.average(), 'time/eval_mean': et.average(), 'time/log_mean': lt.average(), }) agent.log(step) agent.save()
def train(agent, env, eval_env, buffer): def collect(env, step, reset, next_obs, **kwargs): buffer.add(**kwargs) step = agent.env_step runner = Runner(env, agent, step=step, nsteps=agent.N_STEPS) actsel = lambda *args, **kwargs: np.random.randint( 0, env.action_dim, size=env.n_envs) if not agent.rnd_rms_restored(): print('Start to initialize observation running stats...') for _ in range(50): runner.run(action_selector=actsel, step_fn=collect) agent.update_obs_rms(buffer['obs']) buffer.reset() buffer.clear() agent.save() runner.step = step to_log = Every(agent.LOG_PERIOD, agent.LOG_PERIOD) to_eval = Every(agent.EVAL_PERIOD) print('Training starts...') while step < agent.MAX_STEPS: start_env_step = agent.env_step with Timer('env') as rt: step = runner.run(step_fn=collect) agent.store(fps=(step - start_env_step) / rt.last()) agent.record_last_env_output(runner.env_output) value_int, value_ext = agent.compute_value() obs = buffer.get_obs(runner.env_output.obs) assert obs.shape[:2] == (env.n_envs, agent.N_STEPS + 1) assert obs.dtype == np.uint8 agent.update_obs_rms(obs[:, :-1]) norm_obs = agent.normalize_obs(obs) # compute intrinsic reward from the next normalized obs reward_int = agent.compute_int_reward(norm_obs[:, 1:]) agent.update_int_reward_rms(reward_int) reward_int = agent.normalize_int_reward(reward_int) buffer.finish(reward_int, norm_obs[:, :-1], value_int, value_ext) agent.store( reward_int_max=np.max(reward_int), reward_int_min=np.min(reward_int), reward_int=np.mean(reward_int), reward_int_std=np.std(reward_int), ) start_train_step = agent.train_step with Timer('train') as tt: agent.learn_log(step) agent.store(tps=(agent.train_step - start_train_step) / tt.last()) buffer.reset() if to_eval(agent.train_step): with TempStore(agent.get_states, agent.reset_states): scores, epslens, video = evaluate(eval_env, agent, record=True, video_len=4500) video_summary(f'{agent.name}/sim', video, step=step) if eval_env.n_envs == 1: rews_int, rews_ext = agent.retrieve_eval_rewards() assert len(rews_ext) == len(rews_int) == video.shape[1], ( len(rews_ext), len(rews_int), video.shape[1]) n = 10 idxes_int = rews_int.argsort()[::-1][:n] idxes_ext = rews_ext.argsort()[::-1][:n] assert idxes_int.shape == idxes_ext.shape, ( idxes_int.shape, idxes_ext.shape) imgs_int = video[0, idxes_int] imgs_ext = video[0, idxes_ext] rews_int = rews_int[idxes_int] rews_ext = rews_ext[idxes_ext] terms = { **{ f'eval/reward_int_{i}': rews_int[i] for i in range(0, n) }, **{ f'eval/reward_ext_{i}': rews_ext[i] for i in range(0, n) }, } agent.store(**terms) imgs = np.concatenate([imgs_int[:n], imgs_ext[:n]], 0) image_summary(f'{agent.name}/img', imgs, step=step) # info = eval_env.info()[0] # episode = info.get('episode', {'visited_rooms': 1}) # agent.store(visited_rooms_max=len(episode['visited_rooms'])) agent.histogram_summary( {'eval/action': agent.retrieve_eval_actions()}, step=step) agent.store(eval_score=scores, eval_epslen=epslens) if to_log(agent.train_step) and agent.contains_stats('score'): agent.store( **{ 'episodes': runner.episodes, 'train_step': agent.train_step, 'time/run': rt.total(), 'time/train': tt.total() }) agent.log(step) agent.save()
def train(agent, env, eval_env, replay): collect_fn = pkg.import_module('agent', algo=agent.name).collect collect = functools.partial(collect_fn, replay) env_step = agent.env_step runner = Runner(env, agent, step=env_step, nsteps=agent.TRAIN_PERIOD) while not replay.good_to_learn(): env_step = runner.run( # NOTE: random action below makes a huge difference for Mujoco tasks # by default, we don't use it as it's not a conventional practice. # action_selector=env.random_action, step_fn=collect) to_eval = Every(agent.EVAL_PERIOD) to_log = Every(agent.LOG_PERIOD, agent.LOG_PERIOD) to_eval = Every(agent.EVAL_PERIOD) to_record = Every(agent.EVAL_PERIOD * 10) rt = Timer('run') tt = Timer('train') et = Timer('eval') lt = Timer('log') print('Training starts...') while env_step <= int(agent.MAX_STEPS): with rt: env_step = runner.run(step_fn=collect) with tt: agent.learn_log(env_step) if to_eval(env_step): with TempStore(agent.get_states, agent.reset_states): with et: record = agent.RECORD and to_record(env_step) eval_score, eval_epslen, video = evaluate( eval_env, agent, n=agent.N_EVAL_EPISODES, record=agent.RECORD, size=(64, 64)) if record: video_summary(f'{agent.name}/sim', video, step=env_step) agent.store(eval_score=eval_score, eval_epslen=eval_epslen) if to_log(env_step): with lt: fps = rt.average() * agent.TRAIN_PERIOD tps = tt.average() * agent.N_UPDATES agent.store( env_step=agent.env_step, train_step=agent.train_step, fps=fps, tps=tps, ) agent.store( **{ 'train_step': agent.train_step, 'time/run': rt.total(), 'time/train': tt.total(), 'time/eval': et.total(), 'time/log': lt.total(), 'time/run_mean': rt.average(), 'time/train_mean': tt.average(), 'time/eval_mean': et.average(), 'time/log_mean': lt.average(), }) agent.log(env_step) agent.save()