ml-agents/ml-agents/mlagents/trainers/ppo/policy.py

import logging

import numpy as np
from mlagents.trainers.ppo.models import PPOModel
from mlagents.trainers.policy import Policy

logger = logging.getLogger("mlagents.trainers")


class PPOPolicy(Policy):
    def __init__(self, seed, brain, trainer_params, sess, is_training):
        """
        Policy for Proximal Policy Optimization Networks.
        :param seed: Random seed.
        :param brain: Assigned Brain object.
        :param trainer_params: Defined training parameters.
        :param sess: TensorFlow session.
        :param is_training: Whether the model should be trained.
        """
        super().__init__(seed, brain, trainer_params, sess)
        self.has_updated = False
        self.use_curiosity = bool(trainer_params['use_curiosity'])
        self.model = PPOModel(brain,
                              lr=float(trainer_params['learning_rate']),
                              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,
                              use_curiosity=bool(trainer_params['use_curiosity']),
                              curiosity_strength=float(trainer_params['curiosity_strength']),
                              curiosity_enc_size=float(trainer_params['curiosity_enc_size']),
                              scope=self.variable_scope, seed=seed)

        self.inference_dict = {'action': self.model.output, 'log_probs': self.model.all_log_probs,
                               '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']:
            self.inference_dict['update_mean'] = self.model.update_mean
            self.inference_dict['update_variance'] = self.model.update_variance

        self.update_dict = {'value_loss': self.model.value_loss,
                            'policy_loss': self.model.policy_loss,
                            'update_batch': self.model.update_batch}
        if self.use_curiosity:
            self.update_dict['forward_loss'] = self.model.forward_loss
            self.update_dict['inverse_loss'] = self.model.inverse_loss

    def evaluate(self, brain_info):
        """
        Evaluates policy for the agent experiences provided.
        :param brain_info: BrainInfo object containing inputs.
        :return: Outputs from network as defined by self.inference_dict.
        """
        feed_dict = {self.model.batch_size: len(brain_info.vector_observations),
                     self.model.sequence_length: 1}
        if self.use_recurrent:
            if not self.use_continuous_act:
                feed_dict[self.model.prev_action] = brain_info.previous_vector_actions.reshape(
                    [-1, len(self.model.act_size)])
            if brain_info.memories.shape[1] == 0:
                brain_info.memories = self.make_empty_memory(len(brain_info.agents))
            feed_dict[self.model.memory_in] = brain_info.memories
        feed_dict = self._fill_eval_dict(feed_dict, brain_info)
        run_out = self._execute_model(feed_dict, self.inference_dict)
        return run_out

    def update(self, mini_batch, num_sequences):
        """
        Updates model using buffer.
        :param num_sequences: Number of trajectories in batch.
        :param mini_batch: Experience batch.
        :return: Output from update process.
        """
        feed_dict = {self.model.batch_size: num_sequences,
                     self.model.sequence_length: self.sequence_length,
                     self.model.mask_input: mini_batch['masks'].flatten(),
                     self.model.returns_holder: mini_batch['discounted_returns'].flatten(),
                     self.model.old_value: mini_batch['value_estimates'].flatten(),
                     self.model.advantage: mini_batch['advantages'].reshape([-1, 1]),
                     self.model.all_old_log_probs: mini_batch['action_probs'].reshape(
                         [-1, sum(self.model.act_size)])}
        if self.use_continuous_act:
            feed_dict[self.model.output_pre] = mini_batch['actions_pre'].reshape(
                [-1, self.model.act_size[0]])
        else:
            feed_dict[self.model.action_holder] = mini_batch['actions'].reshape(
                [-1, len(self.model.act_size)])
            if self.use_recurrent:
                feed_dict[self.model.prev_action] = mini_batch['prev_action'].reshape(
                    [-1, len(self.model.act_size)])
            feed_dict[self.model.action_masks] = mini_batch['action_mask'].reshape(
                [-1, sum(self.brain.vector_action_space_size)])
        if self.use_vec_obs:
            feed_dict[self.model.vector_in] = mini_batch['vector_obs'].reshape(
                [-1, self.vec_obs_size])
            if self.use_curiosity:
                feed_dict[self.model.next_vector_in] = mini_batch['next_vector_in'].reshape(
                    [-1, self.vec_obs_size])
        if self.model.vis_obs_size > 0:
            for i, _ in enumerate(self.model.visual_in):
                _obs = mini_batch['visual_obs%d' % i]
                if self.sequence_length > 1 and self.use_recurrent:
                    (_batch, _seq, _w, _h, _c) = _obs.shape
                    feed_dict[self.model.visual_in[i]] = _obs.reshape([-1, _w, _h, _c])
                else:
                    feed_dict[self.model.visual_in[i]] = _obs
            if self.use_curiosity:
                for i, _ in enumerate(self.model.visual_in):
                    _obs = mini_batch['next_visual_obs%d' % i]
                    if self.sequence_length > 1 and self.use_recurrent:
                        (_batch, _seq, _w, _h, _c) = _obs.shape
                        feed_dict[self.model.next_visual_in[i]] = _obs.reshape([-1, _w, _h, _c])
                    else:
                        feed_dict[self.model.next_visual_in[i]] = _obs
        if self.use_recurrent:
            mem_in = mini_batch['memory'][:, 0, :]
            feed_dict[self.model.memory_in] = mem_in
        self.has_updated = True
        run_out = self._execute_model(feed_dict, self.update_dict)
        return run_out

    def get_intrinsic_rewards(self, curr_info, next_info):
        """
        Generates intrinsic reward used for Curiosity-based training.
        :BrainInfo curr_info: Current BrainInfo.
        :BrainInfo next_info: Next BrainInfo.
        :return: Intrinsic rewards for all agents.
        """
        if self.use_curiosity:
            if len(curr_info.agents) == 0:
                return []

            feed_dict = {self.model.batch_size: len(next_info.vector_observations),
                         self.model.sequence_length: 1}
            if self.use_continuous_act:
                feed_dict[self.model.output] = next_info.previous_vector_actions
            else:
                feed_dict[self.model.action_holder] = next_info.previous_vector_actions
            for i in range(self.model.vis_obs_size):
                feed_dict[self.model.visual_in[i]] = curr_info.visual_observations[i]
                feed_dict[self.model.next_visual_in[i]] = next_info.visual_observations[i]
            if self.use_vec_obs:
                feed_dict[self.model.vector_in] = curr_info.vector_observations
                feed_dict[self.model.next_vector_in] = next_info.vector_observations
            if self.use_recurrent:
                if curr_info.memories.shape[1] == 0:
                    curr_info.memories = self.make_empty_memory(len(curr_info.agents))
                feed_dict[self.model.memory_in] = curr_info.memories
            intrinsic_rewards = self.sess.run(self.model.intrinsic_reward,
                                              feed_dict=feed_dict) * float(self.has_updated)
            return intrinsic_rewards
        else:
            return None

    def get_value_estimate(self, brain_info, idx):
        """
        Generates value estimates for bootstrapping.
        :param brain_info: BrainInfo to be used for bootstrapping.
        :param idx: Index in BrainInfo of agent.
        :return: Value estimate.
        """
        feed_dict = {self.model.batch_size: 1, self.model.sequence_length: 1}
        for i in range(len(brain_info.visual_observations)):
            feed_dict[self.model.visual_in[i]] = [brain_info.visual_observations[i][idx]]
        if self.use_vec_obs:
            feed_dict[self.model.vector_in] = [brain_info.vector_observations[idx]]
        if self.use_recurrent:
            if brain_info.memories.shape[1] == 0:
                brain_info.memories = self.make_empty_memory(len(brain_info.agents))
            feed_dict[self.model.memory_in] = [brain_info.memories[idx]]
        if not self.use_continuous_act and self.use_recurrent:
            feed_dict[self.model.prev_action] = brain_info.previous_vector_actions[idx].reshape(
                [-1, len(self.model.act_size)])
        value_estimate = self.sess.run(self.model.value, feed_dict)
        return value_estimate

    def get_last_reward(self):
        """
        Returns the last reward the trainer has had
        :return: the new last reward
        """
        return self.sess.run(self.model.last_reward)

    def update_reward(self, new_reward):
        """
        Updates reward value for policy.
        :param new_reward: New reward to save.
        """
        self.sess.run(self.model.update_reward,
                      feed_dict={self.model.new_reward: new_reward})