def create_reward_signals(self, reward_signal_configs): """ Create reward signals :param reward_signal_configs: Reward signal config. """ with self.graph.as_default(): with tf.variable_scope(TOWER_SCOPE_NAME, reuse=tf.AUTO_REUSE): for device_id, device in enumerate(self.devices): with tf.device(device): reward_tower = {} for reward_signal, config in reward_signal_configs.items( ): reward_tower[reward_signal] = create_reward_signal( self, self.towers[device_id], reward_signal, config) for k, v in reward_tower[ reward_signal].update_dict.items(): self.update_dict[k + "_" + str(device_id)] = v self.reward_signal_towers.append(reward_tower) for _, reward_tower in self.reward_signal_towers[0].items(): for _, update_key in reward_tower.stats_name_to_update_name.items( ): all_reward_signal_stats = tf.stack([ self.update_dict[update_key + "_" + str(i)] for i in range(len(self.towers)) ]) mean_reward_signal_stats = tf.reduce_mean( all_reward_signal_stats, 0) self.update_dict.update( {update_key: mean_reward_signal_stats}) self.reward_signals = self.reward_signal_towers[0]
def create_model(self, brain, trainer_params, reward_signal_configs, is_training, load, seed): """ Create PPO models, one on each device :param brain: Assigned Brain object. :param trainer_params: Defined training parameters. :param reward_signal_configs: Reward signal config :param seed: Random seed. """ self.devices = get_devices() with self.graph.as_default(): with tf.variable_scope("", reuse=tf.AUTO_REUSE): for device in self.devices: with tf.device(device): self.towers.append( PPOModel( brain=brain, lr=float(trainer_params["learning_rate"]), lr_schedule=LearningRateSchedule( trainer_params.get( "learning_rate_schedule", "linear")), h_size=int(trainer_params["hidden_units"]), epsilon=float(trainer_params["epsilon"]), beta=float(trainer_params["beta"]), max_step=float(trainer_params["max_steps"]), normalize=trainer_params["normalize"], use_recurrent=trainer_params["use_recurrent"], num_layers=int(trainer_params["num_layers"]), m_size=self.m_size, seed=seed, stream_names=list( reward_signal_configs.keys()), vis_encode_type=EncoderType( trainer_params.get("vis_encode_type", "simple")), )) self.towers[-1].create_ppo_optimizer() self.model = self.towers[0] avg_grads = self.average_gradients([t.grads for t in self.towers]) update_batch = self.model.optimizer.apply_gradients(avg_grads) avg_value_loss = tf.reduce_mean( tf.stack([model.value_loss for model in self.towers]), 0) avg_policy_loss = tf.reduce_mean( tf.stack([model.policy_loss for model in self.towers]), 0) self.inference_dict.update({ "action": self.model.output, "log_probs": self.model.all_log_probs, "value_heads": self.model.value_heads, "value": self.model.value, "entropy": self.model.entropy, "learning_rate": self.model.learning_rate, }) if self.use_continuous_act: self.inference_dict["pre_action"] = self.model.output_pre if self.use_recurrent: self.inference_dict["memory_out"] = self.model.memory_out if (is_training and self.use_vec_obs and trainer_params["normalize"] and not load): self.inference_dict[ "update_mean"] = self.model.update_normalization self.total_policy_loss = self.model.abs_policy_loss self.update_dict.update({ "value_loss": avg_value_loss, "policy_loss": avg_policy_loss, "update_batch": update_batch, })