ma_policy/ma_policy.py

import tensorflow as tf
import numpy as np
import gym
import logging
import sys
from copy import deepcopy
from functools import partial
from collections import OrderedDict
from baselines.common.distributions import make_pdtype

from ma_policy.util import listdict2dictnp, normc_initializer, shape_list, l2_loss
from ma_policy.variable_schema import VariableSchema, BATCH, TIMESTEPS
from ma_policy.normalizers import EMAMeanStd
from ma_policy.graph_construct import construct_tf_graph, construct_schemas_zero_state


class MAPolicy(object):
    '''
        Args:
            ob_space: gym observation space of a SINGLE agent. Expects a dict space.
            ac_space: gym action space. Expects a dict space where each item is a tuple of action
                spaces
            network_spec: list of layers. See construct_tf_graph for details.
            v_network_spec: optional. If specified it is the network spec of the value function.
            trainable_vars: optional. List of variable name segments that should be trained.
            not_trainable_vars: optional. List of variable name segements that should not be
                trained. trainable_vars supercedes this if both are specified.
    '''
    def __init__(self, scope, *, ob_space, ac_space, network_spec, v_network_spec=None,
                 stochastic=True, reuse=False, build_act=True,
                 trainable_vars=None, not_trainable_vars=None,
                 gaussian_fixed_var=True, weight_decay=0.0, ema_beta=0.99999,
                 **kwargs):
        self.reuse = reuse
        self.scope = scope
        self.ob_space = ob_space
        self.ac_space = deepcopy(ac_space)
        self.network_spec = network_spec
        self.v_network_spec = v_network_spec or self.network_spec
        self.stochastic = stochastic
        self.trainable_vars = trainable_vars
        self.not_trainable_vars = not_trainable_vars
        self.gaussian_fixed_var = gaussian_fixed_var
        self.weight_decay = weight_decay
        self.kwargs = kwargs
        self.build_act = build_act
        self._reset_ops = []
        self._auxiliary_losses = []
        self._running_mean_stds = {}
        self._ema_beta = ema_beta
        self.training_stats = []

        assert isinstance(self.ac_space, gym.spaces.Dict)
        assert isinstance(self.ob_space, gym.spaces.Dict)
        assert 'observation_self' in self.ob_space.spaces

        # Action space will come in as a MA action space. Convert to a SA action space.
        self.ac_space.spaces = {k: v.spaces[0] for k, v in self.ac_space.spaces.items()}

        self.pdtypes = {k: make_pdtype(s) for k, s in self.ac_space.spaces.items()}

        # Create input schemas for each action type
        self.input_schemas = {
            k: VariableSchema(shape=[BATCH, TIMESTEPS] + pdtype.sample_shape(),
                              dtype=pdtype.sample_dtype())
            for k, pdtype in self.pdtypes.items()
        }

        # Creat input schemas for each observation
        for k, v in self.ob_space.spaces.items():
            self.input_schemas[k] = VariableSchema(shape=[BATCH, TIMESTEPS] + list(v.shape),
                                                   dtype=tf.float32)

        # Setup schemas and zero state for layers with state
        v_state_schemas, v_zero_states = construct_schemas_zero_state(
            self.v_network_spec, self.ob_space, 'vpred_net')
        pi_state_schemas, pi_zero_states = construct_schemas_zero_state(
            self.network_spec, self.ob_space, 'policy_net')

        self.state_keys = list(v_state_schemas.keys()) + list(pi_state_schemas.keys())
        self.input_schemas.update(v_state_schemas)
        self.input_schemas.update(pi_state_schemas)
        self.zero_state = {}
        self.zero_state.update(v_zero_states)
        self.zero_state.update(pi_zero_states)

        if build_act:
            with tf.variable_scope(self.scope, reuse=self.reuse):
                self.phs = {name: schema.placeholder(name=name)
                            for name, schema in self.get_input_schemas().items()}
            self.build(self.phs)

    def build(self, inputs):
        with tf.variable_scope(self.scope, reuse=self.reuse):
            self.full_scope_name = tf.get_variable_scope().name
            self._init(inputs, **self.kwargs)

    def _init(self, inputs, gaussian_fixed_var=True, **kwargs):
        '''
            Args:
                inputs (dict): input dictionary containing tf tensors
                gaussian_fixed_var (bool): If True the policies variance won't be conditioned on state
        '''
        taken_actions = {k: inputs[k] for k in self.pdtypes.keys()}

        #  Copy inputs to not overwrite. Don't need to pass actions to policy, so exlcude these
        processed_inp = {k: v for k, v in inputs.items() if k not in self.pdtypes.keys()}

        self._normalize_inputs(processed_inp)

        self.state_out = OrderedDict()

        # Value network
        (vpred,
         vpred_state_out,
         vpred_reset_ops) = construct_tf_graph(
            processed_inp, self.v_network_spec, scope='vpred_net', act=self.build_act)

        self._init_vpred_head(vpred, processed_inp, 'vpred_out0', "value0")

        # Policy network
        (pi,
         pi_state_out,
         pi_reset_ops) = construct_tf_graph(
            processed_inp, self.network_spec, scope='policy_net', act=self.build_act)

        self.state_out.update(vpred_state_out)
        self.state_out.update(pi_state_out)
        self._reset_ops += vpred_reset_ops + pi_reset_ops
        self._init_policy_out(pi, taken_actions)
        if self.weight_decay != 0.0:
            kernels = [var for var in self.get_trainable_variables() if 'kernel' in var.name]
            w_norm_sum = tf.reduce_sum([tf.nn.l2_loss(var) for var in kernels])
            w_norm_loss = w_norm_sum * self.weight_decay
            self.add_auxiliary_loss('weight_decay', w_norm_loss)

        # set state to zero state
        self.reset()

    def _init_policy_out(self, pi, taken_actions):
        with tf.variable_scope('policy_out'):
            self.pdparams = {}
            for k in self.pdtypes.keys():
                with tf.variable_scope(k):
                    if self.gaussian_fixed_var and isinstance(self.ac_space.spaces[k], gym.spaces.Box):
                        mean = tf.layers.dense(pi["main"],
                                               self.pdtypes[k].param_shape()[0] // 2,
                                               kernel_initializer=normc_initializer(0.01),
                                               activation=None)
                        logstd = tf.get_variable(name="logstd",
                                                 shape=[1, self.pdtypes[k].param_shape()[0] // 2],
                                                 initializer=tf.zeros_initializer())
                        self.pdparams[k] = tf.concat([mean, mean * 0.0 + logstd], axis=2)
                    elif k in pi:
                        # This is just for the case of entity specific actions
                        if isinstance(self.ac_space.spaces[k], (gym.spaces.Discrete)):
                            assert pi[k].get_shape()[-1] == 1
                            self.pdparams[k] = pi[k][..., 0]
                        elif isinstance(self.ac_space.spaces[k], (gym.spaces.MultiDiscrete)):
                            assert np.prod(pi[k].get_shape()[-2:]) == self.pdtypes[k].param_shape()[0],\
                                f"policy had shape {pi[k].get_shape()} for action {k}, but required {self.pdtypes[k].param_shape()}"
                            new_shape = shape_list(pi[k])[:-2] + [np.prod(pi[k].get_shape()[-2:]).value]
                            self.pdparams[k] = tf.reshape(pi[k], shape=new_shape)
                        else:
                            assert False
                    else:
                        self.pdparams[k] = tf.layers.dense(pi["main"],
                                                           self.pdtypes[k].param_shape()[0],
                                                           kernel_initializer=normc_initializer(0.01),
                                                           activation=None)

            with tf.variable_scope('pds'):
                self.pds = {k: pdtype.pdfromflat(self.pdparams[k])
                            for k, pdtype in self.pdtypes.items()}

            with tf.variable_scope('sampled_action'):
                self.sampled_action = {k: pd.sample() if self.stochastic else pd.mode()
                                       for k, pd in self.pds.items()}
            with tf.variable_scope('sampled_action_logp'):
                self.sampled_action_logp = sum([self.pds[k].logp(self.sampled_action[k])
                                                for k in self.pdtypes.keys()])
            with tf.variable_scope('entropy'):
                self.entropy = sum([pd.entropy() for pd in self.pds.values()])
            with tf.variable_scope('taken_action_logp'):
                self.taken_action_logp = sum([self.pds[k].logp(taken_actions[k])
                                              for k in self.pdtypes.keys()])

    def _init_vpred_head(self, vpred, processed_inp, vpred_scope, feedback_name):
        with tf.variable_scope(vpred_scope):
            _vpred = tf.layers.dense(vpred['main'], 1, activation=None,
                                     kernel_initializer=tf.contrib.layers.xavier_initializer())
            _vpred = tf.squeeze(_vpred, -1)
            normalize_axes = (0, 1)
            loss_fn = partial(l2_loss, mask=processed_inp.get(feedback_name + "_mask", None))
            rms_class = partial(EMAMeanStd, beta=self._ema_beta)

            rms_shape = [dim for i, dim in enumerate(_vpred.get_shape()) if i not in normalize_axes]
            self.value_rms = rms_class(shape=rms_shape, scope='value0filter')
            self.scaled_value_tensor = self.value_rms.mean + _vpred * self.value_rms.std
            self.add_running_mean_std(rms=self.value_rms, name='feedback.value0', axes=normalize_axes)

    def _normalize_inputs(self, processed_inp):
        with tf.variable_scope('normalize_self_obs'):
            ob_rms_self = EMAMeanStd(shape=self.ob_space.spaces['observation_self'].shape,
                                     scope="obsfilter", beta=self._ema_beta, per_element_update=False)
            self.add_running_mean_std("observation_self", ob_rms_self, axes=(0, 1))
            normalized = (processed_inp['observation_self'] - ob_rms_self.mean) / ob_rms_self.std
            clipped = tf.clip_by_value(normalized, -5.0, 5.0)
            processed_inp['observation_self'] = clipped

        for key in self.ob_space.spaces.keys():
            if key == 'observation_self':
                continue
            elif 'mask' in key:  # Don't normalize observation masks
                pass
            else:
                with tf.variable_scope(f'normalize_{key}'):
                    ob_rms = EMAMeanStd(shape=self.ob_space.spaces[key].shape[1:],
                                        scope=f"obsfilter/{key}", beta=self._ema_beta, per_element_update=False)
                    normalized = (processed_inp[key] - ob_rms.mean) / ob_rms.std
                    processed_inp[key] = tf.clip_by_value(normalized, -5.0, 5.0)
                    self.add_running_mean_std(key, ob_rms, axes=(0, 1, 2))

    def get_input_schemas(self):
        return self.input_schemas.copy()

    def process_state_batch(self, states):
        '''
            Batch states together.
            args:
                states -- list (batch) of dicts of states with shape (n_agent, dim state).
        '''
        new_states = listdict2dictnp(states, keepdims=True)
        return new_states

    def process_observation_batch(self, obs):
        '''
            Batch obs together.
            Args:
                obs -- list of lists (batch, time), where elements are dictionary observations
        '''

        new_obs = deepcopy(obs)
        # List tranpose -- now in (time, batch)
        new_obs = list(map(list, zip(*new_obs)))
        # Convert list of list of dicts to dict of numpy arrays
        new_obs = listdict2dictnp([listdict2dictnp(batch, keepdims=True) for batch in new_obs])
        # Flatten out the agent dimension, so batches look like normal SA batches
        new_obs = {k: self.reshape_ma_observations(v) for k, v in new_obs.items()}

        return new_obs

    def reshape_ma_observations(self, obs):
        # Observations with shape (time, batch)
        if len(obs.shape) == 2:
            batch_first_ordering = (1, 0)
        # Observations with shape (time, batch, dim obs)
        elif len(obs.shape) == 3:
            batch_first_ordering = (1, 0, 2)
        # Observations with shape (time, batch, n_entity, dim obs)
        elif len(obs.shape) == 4:
            batch_first_ordering = (1, 0, 2, 3)
        else:
            raise ValueError(f"Obs dim {obs.shape}. Only supports dim 3 or 4")
        new_obs = obs.copy().transpose(batch_first_ordering)  # (n_agent, batch, time, dim obs)

        return new_obs

    def prepare_input(self, observation, state_in, taken_action=None):
        ''' Add in time dimension to observations, assumes that first dimension of observation is
            already the batch dimension and does not need to be added.'''
        obs = deepcopy(observation)
        obs.update(state_in)
        if taken_action is not None:
            obs.update(taken_action)
        return obs

    def act(self, observation, extra_feed_dict={}):
        outputs = {
            'ac': self.sampled_action,
            'ac_logp': self.sampled_action_logp,
            'vpred': self.scaled_value_tensor,
            'state': self.state_out}
        # Add timestep dimension to observations
        obs = deepcopy(observation)
        n_agents = observation['observation_self'].shape[0]

        # Make sure that there are as many states as there are agents.
        # This should only happen with the zero state.
        for k, v in self.state.items():
            assert v.shape[0] == 1 or v.shape[0] == n_agents
            if v.shape[0] == 1 and v.shape[0] != n_agents:
                self.state[k] = np.repeat(v, n_agents, 0)

        # Add time dimension to obs
        for k, v in obs.items():
            obs[k] = np.expand_dims(v, 1)
        inputs = self.prepare_input(observation=obs, state_in=self.state)
        feed_dict = {self.phs[k]: v for k, v in inputs.items()}
        feed_dict.update(extra_feed_dict)

        outputs = tf.get_default_session().run(outputs, feed_dict)
        self.state = outputs['state']

        # Remove time dimension from outputs
        def preprocess_act_output(act_output):
            if isinstance(act_output, dict):
                return {k: np.squeeze(v, 1) for k, v in act_output.items()}
            else:
                return np.squeeze(act_output, 1)

        info = {'vpred': preprocess_act_output(outputs['vpred']),
                'ac_logp': preprocess_act_output(outputs['ac_logp']),
                'state': outputs['state']}

        return preprocess_act_output(outputs['ac']), info

    def get_variables(self):
        variables = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, self.full_scope_name + '/')
        return variables

    def get_trainable_variables(self):
        variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, self.full_scope_name + '/')
        if self.trainable_vars is not None:
            variables = [v for v in variables
                         if any([tr_v in v.name for tr_v in self.trainable_vars])]
        elif self.not_trainable_vars is not None:
            variables = [v for v in variables
                         if not any([tr_v in v.name for tr_v in self.not_trainable_vars])]
        variables = [v for v in variables if 'not_trainable' not in v.name]
        return variables

    def reset(self):
        self.state = deepcopy(self.zero_state)
        if tf.get_default_session() is not None:
            tf.get_default_session().run(self._reset_ops)

    def set_state(self, state):
        self.state = deepcopy(state)

    def auxiliary_losses(self):
        """ Any extra losses internal to the policy, automatically added to the total loss."""
        return self._auxiliary_losses

    def add_auxiliary_loss(self, name, loss):
        self.training_stats.append((name, 'scalar', loss, lambda x: x))
        self._auxiliary_losses.append(loss)

    def add_running_mean_std(self, name, rms, axes=(0, 1)):
        """
        Add a RunningMeanStd/EMAMeanStd object to the policy's list. It will then get updated during optimization.
        :param name: name of the input field to update from.
        :param rms: RMS object to update.
        :param axes: axes of the input to average over.
            RMS's shape should be equal to input's shape after axes are removed.
            e.g. if inputs is [5, 6, 7, 8] and axes is [0, 2], then RMS's shape should be [6, 8].
        :return:
        """
        self._running_mean_stds[name] = {'rms': rms, 'axes': axes}