def _learn_loop(self, multi_env, step_callback=None): """ Main training loop. :param step_callback: a hacky callback that takes a dictionary with all local variables as an argument. Allows you too look inside the training process. """ step = initial_step = tf.train.global_step(self.session, tf.train.get_global_step()) env_steps = self.total_env_steps.eval(session=self.session) batch_size = self.params.rollout * self.params.num_envs img_obs, timer_obs = extract_keys(multi_env.initial_obs(), 'obs', 'timer') adv_running_mean_std = RunningMeanStd(max_past_samples=10000) def end_of_training(s, es): return s >= self.params.train_for_steps or es > self.params.train_for_env_steps while not end_of_training(step, env_steps): timing = AttrDict({'experience': time.time(), 'batch': time.time()}) experience_start = time.time() env_steps_before_batch = env_steps batch_obs, batch_timer = [img_obs], [timer_obs] env_steps += len(img_obs) batch_actions, batch_values, batch_rewards, batch_dones, batch_next_obs = [], [], [], [], [] for rollout_step in range(self.params.rollout): actions, values = self._policy_step_timer(img_obs, timer_obs) batch_actions.append(actions) batch_values.append(values) # wait for all the workers to complete an environment step next_obs, rewards, dones, infos = multi_env.step(actions) next_img_obs, next_timer = extract_keys(next_obs, 'obs', 'timer') # calculate curiosity bonus bonuses = self._prediction_curiosity_bonus(img_obs, actions, next_img_obs) rewards += bonuses batch_rewards.append(rewards) batch_dones.append(dones) batch_next_obs.append(next_img_obs) img_obs = next_img_obs timer_obs = next_timer if infos is not None and 'num_frames' in infos[0]: env_steps += sum((info['num_frames'] for info in infos)) else: env_steps += multi_env.num_envs if rollout_step != self.params.rollout - 1: # we don't need the newest observation in the training batch, already have enough batch_obs.append(img_obs) batch_timer.append(timer_obs) assert len(batch_obs) == len(batch_rewards) assert len(batch_obs) == len(batch_next_obs) batch_rewards = np.asarray(batch_rewards, np.float32).swapaxes(0, 1) batch_dones = np.asarray(batch_dones, np.bool).swapaxes(0, 1) batch_values = np.asarray(batch_values, np.float32).swapaxes(0, 1) # Last value won't be valid for envs with done=True (because env automatically resets and shows 1st # observation of the next episode. But that's okay, because we should never use last_value in this case. last_values = self._estimate_values_timer(img_obs, timer_obs) gamma = self.params.gamma disc_rewards = [] for i in range(len(batch_rewards)): env_rewards = self._calc_discounted_rewards(gamma, batch_rewards[i], batch_dones[i], last_values[i]) disc_rewards.extend(env_rewards) disc_rewards = np.asarray(disc_rewards, np.float32) # convert observations and estimations to meaningful n-step batches batch_obs_shape = (self.params.rollout * multi_env.num_envs,) + img_obs[0].shape batch_obs = np.asarray(batch_obs, np.float32).swapaxes(0, 1).reshape(batch_obs_shape) batch_next_obs = np.asarray(batch_next_obs, np.float32).swapaxes(0, 1).reshape(batch_obs_shape) batch_actions = np.asarray(batch_actions, np.int32).swapaxes(0, 1).flatten() batch_timer = np.asarray(batch_timer, np.float32).swapaxes(0, 1).flatten() batch_values = batch_values.flatten() advantages = disc_rewards - batch_values if self.params.normalize_adv: adv_running_mean_std.update(advantages) advantages = (advantages - adv_running_mean_std.mean) / (np.sqrt(adv_running_mean_std.var) + EPS) advantages = np.clip(advantages, -self.params.clip_advantage, self.params.clip_advantage) timing.experience = time.time() - timing.experience timing.train = time.time() step = self._curious_train_step( step, env_steps, batch_obs, batch_timer, batch_actions, batch_values, disc_rewards, advantages, batch_next_obs, ) self._maybe_save(step, env_steps) timing.train = time.time() - timing.train avg_reward = multi_env.calc_avg_rewards(n=self.params.stats_episodes) avg_length = multi_env.calc_avg_episode_lengths(n=self.params.stats_episodes) fps = (env_steps - env_steps_before_batch) / (time.time() - timing.batch) self._maybe_print(step, avg_reward, avg_length, fps, timing) self._maybe_aux_summaries(step, env_steps, avg_reward, avg_length) self._maybe_update_avg_reward(avg_reward, multi_env.stats_num_episodes()) if step_callback is not None: step_callback(locals(), globals())
def learn(self, step_callback=None): """ Main training loop. :param step_callback: a hacky callback that takes a dictionary with all local variables as an argument. Allows you too look inside the training process. """ step = initial_step = tf.train.global_step(self.session, tf.train.get_global_step()) env_steps = env_steps_initial = self.total_env_steps.eval( session=self.session) batch_size = self.params.rollout * self.params.num_envs multi_env = MultiEnv( self.params.num_envs, self.params.num_workers, make_env_func=self.make_env_func, stats_episodes=self.params.stats_episodes, ) observations = extract_key(multi_env.initial_obs(), 'obs') def end_of_training(s): return s >= self.params.train_for_steps while not end_of_training(step): timing = AttrDict({ 'experience': time.time(), 'batch': time.time() }) experience_start = time.time() env_steps_before_batch = env_steps batch_obs = [observations] env_steps += len(observations) batch_actions, batch_values, batch_rewards, batch_dones = [], [], [], [] for rollout_step in range(self.params.rollout): actions, values = self._policy_step(observations) batch_actions.append(actions) batch_values.append(values) # wait for all the workers to complete an environment step observations, rewards, dones, infos = multi_env.step(actions) observations = extract_key(observations, 'obs') batch_rewards.append(rewards) batch_dones.append(dones) if infos is not None and 'num_frames' in infos[0]: env_steps += sum((info['num_frames'] for info in infos)) else: env_steps += multi_env.num_envs if rollout_step != self.params.rollout - 1: # we don't need the newest observation in the training batch, already have enough batch_obs.append(observations) assert len(batch_obs) == len(batch_rewards) batch_rewards = np.asarray(batch_rewards, np.float32).swapaxes(0, 1) batch_dones = np.asarray(batch_dones, np.bool).swapaxes(0, 1) last_values = self._estimate_values(observations) gamma = self.params.gamma discounted_rewards = [] for env_rewards, env_dones, last_value in zip( batch_rewards, batch_dones, last_values): discounted_rewards.extend( self._calc_discounted_rewards(gamma, env_rewards, env_dones, last_value)) # convert observations and estimations to meaningful n-step batches batch_obs_shape = (self.params.rollout * multi_env.num_envs, ) + observations[0].shape batch_obs = np.asarray(batch_obs, np.float32).swapaxes( 0, 1).reshape(batch_obs_shape) batch_actions = np.asarray(batch_actions, np.int32).swapaxes(0, 1).flatten() batch_values = np.asarray(batch_values, np.float32).swapaxes(0, 1).flatten() timing.experience = time.time() - timing.experience timing.train = time.time() step = self._train_step(step, env_steps, batch_obs, batch_actions, batch_values, discounted_rewards) self._maybe_save(step, env_steps) timing.train = time.time() - timing.train avg_reward = multi_env.calc_avg_rewards( n=self.params.stats_episodes) avg_length = multi_env.calc_avg_episode_lengths( n=self.params.stats_episodes) fps = (env_steps - env_steps_before_batch) / (time.time() - timing.batch) self._maybe_print(step, avg_reward, avg_length, fps, timing) self._maybe_aux_summaries(step, env_steps, avg_reward, avg_length) self._maybe_update_avg_reward(avg_reward, env_steps - env_steps_initial) if step_callback is not None: step_callback(locals(), globals()) log.info('Done!') multi_env.close()