From a6b1b7a5e4f36ebbc1600212da2920ca0c448426 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleksandar=20Milovanovi=C4=87?= Date: Wed, 6 Feb 2019 17:05:15 +0100 Subject: [PATCH 01/33] impala changes --- python/ray/rllib/agents/impala/impala.py | 3 + python/ray/rllib/agents/impala/vtrace.py | 385 +++++++++--------- .../agents/impala/vtrace_policy_graph.py | 138 +++++-- python/ray/rllib/models/action_dist.py | 26 +- python/ray/rllib/models/catalog.py | 10 +- 5 files changed, 331 insertions(+), 231 deletions(-) diff --git a/python/ray/rllib/agents/impala/impala.py b/python/ray/rllib/agents/impala/impala.py index 9221e3764600..d2a380ab045e 100644 --- a/python/ray/rllib/agents/impala/impala.py +++ b/python/ray/rllib/agents/impala/impala.py @@ -70,6 +70,9 @@ # max number of workers to broadcast one set of weights to "broadcast_interval": 1, + # Actions are chosen based on this distribution, if provided + "dist_type": None, + # Learning params. "grad_clip": 40.0, # either "adam" or "rmsprop" diff --git a/python/ray/rllib/agents/impala/vtrace.py b/python/ray/rllib/agents/impala/vtrace.py index ac5abf0e6592..50dfc813f147 100644 --- a/python/ray/rllib/agents/impala/vtrace.py +++ b/python/ray/rllib/agents/impala/vtrace.py @@ -11,14 +11,12 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Functions to compute V-trace off-policy actor critic targets. +"""Functions to compute V-trace off-policy actor critic targets. For details and theory see: - "IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures" by Espeholt, Soyer, Munos et al. - See https://arxiv.org/abs/1802.01561 for the full paper. """ @@ -32,204 +30,195 @@ nest = tf.contrib.framework.nest -VTraceFromLogitsReturns = collections.namedtuple('VTraceFromLogitsReturns', [ - 'vs', 'pg_advantages', 'log_rhos', 'behaviour_action_log_probs', - 'target_action_log_probs' -]) +VTraceFromLogitsReturns = collections.namedtuple( + 'VTraceFromLogitsReturns', + ['vs', 'pg_advantages', 'rhos', + 'behaviour_action_policy', 'target_action_policy']) VTraceReturns = collections.namedtuple('VTraceReturns', 'vs pg_advantages') -def log_probs_from_logits_and_actions(policy_logits, actions): - """Computes action log-probs from policy logits and actions. - - In the notation used throughout documentation and comments, T refers to the - time dimension ranging from 0 to T-1. B refers to the batch size and - NUM_ACTIONS refers to the number of actions. - - Args: - policy_logits: A float32 tensor of shape [T, B, NUM_ACTIONS] with - un-normalized log-probabilities parameterizing a softmax policy. - actions: An int32 tensor of shape [T, B] with actions. - - Returns: - A float32 tensor of shape [T, B] corresponding to the sampling log - probability of the chosen action w.r.t. the policy. - """ - policy_logits = tf.convert_to_tensor(policy_logits, dtype=tf.float32) - actions = tf.convert_to_tensor(actions, dtype=tf.int32) - - policy_logits.shape.assert_has_rank(3) - actions.shape.assert_has_rank(2) - - return -tf.nn.sparse_softmax_cross_entropy_with_logits( - logits=policy_logits, labels=actions) - - -def from_logits(behaviour_policy_logits, - target_policy_logits, - actions, - discounts, - rewards, - values, - bootstrap_value, - clip_rho_threshold=1.0, - clip_pg_rho_threshold=1.0, +def from_logits(behaviour_policy, target_policy, actions, + discounts, rewards, values, bootstrap_value, + clip_rho_threshold=1.0, clip_pg_rho_threshold=1.0, name='vtrace_from_logits'): r"""V-trace for softmax policies. - - Calculates V-trace actor critic targets for softmax polices as described in - - "IMPALA: Scalable Distributed Deep-RL with - Importance Weighted Actor-Learner Architectures" - by Espeholt, Soyer, Munos et al. - - Target policy refers to the policy we are interested in improving and - behaviour policy refers to the policy that generated the given - rewards and actions. - - In the notation used throughout documentation and comments, T refers to the - time dimension ranging from 0 to T-1. B refers to the batch size and - NUM_ACTIONS refers to the number of actions. - - Args: - behaviour_policy_logits: A float32 tensor of shape [T, B, NUM_ACTIONS] with - un-normalized log-probabilities parametrizing the softmax behaviour - policy. - target_policy_logits: A float32 tensor of shape [T, B, NUM_ACTIONS] with - un-normalized log-probabilities parametrizing the softmax target policy. - actions: An int32 tensor of shape [T, B] of actions sampled from the - behaviour policy. - discounts: A float32 tensor of shape [T, B] with the discount encountered - when following the behaviour policy. - rewards: A float32 tensor of shape [T, B] with the rewards generated by - following the behaviour policy. - values: A float32 tensor of shape [T, B] with the value function estimates - wrt. the target policy. - bootstrap_value: A float32 of shape [B] with the value function estimate at - time T. - clip_rho_threshold: A scalar float32 tensor with the clipping threshold for - importance weights (rho) when calculating the baseline targets (vs). - rho^bar in the paper. - clip_pg_rho_threshold: A scalar float32 tensor with the clipping threshold - on rho_s in \rho_s \delta log \pi(a|x) (r + \gamma v_{s+1} - V(x_s)). - name: The name scope that all V-trace operations will be created in. - - Returns: - A `VTraceFromLogitsReturns` namedtuple with the following fields: - vs: A float32 tensor of shape [T, B]. Can be used as target to train a - baseline (V(x_t) - vs_t)^2. - pg_advantages: A float 32 tensor of shape [T, B]. Can be used as an - estimate of the advantage in the calculation of policy gradients. - log_rhos: A float32 tensor of shape [T, B] containing the log importance - sampling weights (log rhos). - behaviour_action_log_probs: A float32 tensor of shape [T, B] containing - behaviour policy action log probabilities (log \mu(a_t)). - target_action_log_probs: A float32 tensor of shape [T, B] containing - target policy action probabilities (log \pi(a_t)). - """ - behaviour_policy_logits = tf.convert_to_tensor( - behaviour_policy_logits, dtype=tf.float32) - target_policy_logits = tf.convert_to_tensor( - target_policy_logits, dtype=tf.float32) - actions = tf.convert_to_tensor(actions, dtype=tf.int32) - - # Make sure tensor ranks are as expected. - # The rest will be checked by from_action_log_probs. - behaviour_policy_logits.shape.assert_has_rank(3) - target_policy_logits.shape.assert_has_rank(3) - actions.shape.assert_has_rank(2) - - with tf.name_scope( - name, - values=[ - behaviour_policy_logits, target_policy_logits, actions, - discounts, rewards, values, bootstrap_value - ]): - target_action_log_probs = log_probs_from_logits_and_actions( - target_policy_logits, actions) - behaviour_action_log_probs = log_probs_from_logits_and_actions( - behaviour_policy_logits, actions) - log_rhos = target_action_log_probs - behaviour_action_log_probs - vtrace_returns = from_importance_weights( - log_rhos=log_rhos, + Calculates V-trace actor critic targets for softmax polices as described in + "IMPALA: Scalable Distributed Deep-RL with + Importance Weighted Actor-Learner Architectures" + by Espeholt, Soyer, Munos et al. + Target policy refers to the policy we are interested in improving and + behaviour policy refers to the policy that generated the given + rewards and actions. + In the notation used throughout documentation and comments, T refers to the + time dimension ranging from 0 to T-1. B refers to the batch size and + NUM_ACTIONS refers to the number of actions. + Args: + behaviour_policy: A float32 tensor of shape [T, B, NUM_ACTIONS] with + un-normalized log-probabilities parametrizing the softmax behaviour + policy. + target_policy: A float32 tensor of shape [T, B, NUM_ACTIONS] with + un-normalized log-probabilities parametrizing the softmax target + policy. + actions: An int32 tensor of shape [T, B] of actions sampled from the + behaviour policy. + discounts: A float32 tensor of shape [T, B] with the discount encountered + when following the behaviour policy. + rewards: A float32 tensor of shape [T, B] with the rewards generated by + following the behaviour policy. + values: A float32 tensor of shape [T, B] with the value function + estimates wrt. the target policy. + bootstrap_value: A float32 of shape [B] with the value function estimate + at time T. + clip_rho_threshold: A scalar float32 tensor with the clipping threshold + for importance weights (rho) when calculating the baseline targets (vs) + rho^bar in the paper. + clip_pg_rho_threshold: A scalar float32 tensor with the clipping + threshold on rho_s in + \rho_s \delta log \pi(a|x) (r + \gamma v_{s+1} - V(x_s)). + name: The name scope that all V-trace operations will be created in. + Returns: + A `VTraceFromLogitsReturns` namedtuple with the following fields: + vs: A float32 tensor of shape [T, B]. Can be used as target to train a + baseline (V(x_t) - vs_t)^2. + pg_advantages: A float 32 tensor of shape [T, B]. Can be used as an + estimate of the advantage in the calculation of policy gradients. + log_rhos: A float32 tensor of shape [T, B] containing the log + importance sampling weights (log rhos). + behaviour_action_log_probs: A float32 tensor of shape [T, B] containing + behaviour policy action log probabilities (log \mu(a_t)). + target_action_log_probs: A float32 tensor of shape [T, B] containing + target policy action probabilities (log \pi(a_t)). + """ + + for i in range(len(behaviour_policy)): + behaviour_policy[i] = tf.convert_to_tensor( + behaviour_policy[i], dtype=tf.float32) + target_policy[i] = tf.convert_to_tensor( + target_policy[i], dtype=tf.float32) + actions[i] = tf.convert_to_tensor(actions[i], dtype=tf.int32) + + # Make sure tensor ranks are as expected. + # The rest will be checked by from_action_log_probs. + behaviour_policy[i].shape.assert_has_rank(3) + target_policy[i].shape.assert_has_rank(3) + actions[i].shape.assert_has_rank(2) + + with tf.name_scope(name, values=[behaviour_policy, target_policy, actions, + discounts, rewards, values, + bootstrap_value]): + target_action_policy = select_policy_values_using_actions( + target_policy, actions) + behaviour_action_policy = select_policy_values_using_actions( + behaviour_policy, actions) + + rhos = get_rhos(target_action_policy, behaviour_action_policy) + + vtrace_returns = _from_importance_weights( + rhos=rhos, discounts=discounts, rewards=rewards, values=values, bootstrap_value=bootstrap_value, clip_rho_threshold=clip_rho_threshold, clip_pg_rho_threshold=clip_pg_rho_threshold) + return VTraceFromLogitsReturns( - log_rhos=log_rhos, - behaviour_action_log_probs=behaviour_action_log_probs, - target_action_log_probs=target_action_log_probs, + rhos=rhos, + behaviour_action_policy=behaviour_action_policy, + target_action_policy=target_action_policy, **vtrace_returns._asdict()) -def from_importance_weights(log_rhos, - discounts, - rewards, - values, - bootstrap_value, - clip_rho_threshold=1.0, - clip_pg_rho_threshold=1.0, - name='vtrace_from_importance_weights'): - r"""V-trace from log importance weights. +def select_policy_values_using_actions(policy_logits, actions): + """ + Computes action log-probs from policy logits and actions. + In the notation used throughout documentation and comments, T refers to the + time dimension ranging from 0 to T-1. B refers to the batch size and + NUM_ACTIONS refers to the number of actions. + Args: + policy_logits: A float32 tensor of shape [T, B, NUM_ACTIONS] with + un-normalized log-probabilities parameterizing a softmax policy. + actions: An int32 tensor of shape [T, B] with actions. + Returns: + A float32 tensor of shape [T, B] corresponding to the sampling log + probability of the chosen action w.r.t. the policy. + """ + + log_probs = [] + for i in range(len(policy_logits)): + log_probs.append(-tf.nn.sparse_softmax_cross_entropy_with_logits( + logits=policy_logits[i], labels=actions[i])) + + return log_probs + - Calculates V-trace actor critic targets as described in - - "IMPALA: Scalable Distributed Deep-RL with - Importance Weighted Actor-Learner Architectures" - by Espeholt, Soyer, Munos et al. - - In the notation used throughout documentation and comments, T refers to the - time dimension ranging from 0 to T-1. B refers to the batch size and - NUM_ACTIONS refers to the number of actions. This code also supports the - case where all tensors have the same number of additional dimensions, e.g., - `rewards` is [T, B, C], `values` is [T, B, C], `bootstrap_value` is [B, C]. - - Args: - log_rhos: A float32 tensor of shape [T, B, NUM_ACTIONS] representing the - log importance sampling weights, i.e. - log(target_policy(a) / behaviour_policy(a)). V-trace performs operations - on rhos in log-space for numerical stability. - discounts: A float32 tensor of shape [T, B] with discounts encountered when - following the behaviour policy. - rewards: A float32 tensor of shape [T, B] containing rewards generated by - following the behaviour policy. - values: A float32 tensor of shape [T, B] with the value function estimates - wrt. the target policy. - bootstrap_value: A float32 of shape [B] with the value function estimate at - time T. - clip_rho_threshold: A scalar float32 tensor with the clipping threshold for - importance weights (rho) when calculating the baseline targets (vs). - rho^bar in the paper. If None, no clipping is applied. - clip_pg_rho_threshold: A scalar float32 tensor with the clipping threshold - on rho_s in \rho_s \delta log \pi(a|x) (r + \gamma v_{s+1} - V(x_s)). If - None, no clipping is applied. - name: The name scope that all V-trace operations will be created in. - - Returns: - A VTraceReturns namedtuple (vs, pg_advantages) where: - vs: A float32 tensor of shape [T, B]. Can be used as target to - train a baseline (V(x_t) - vs_t)^2. - pg_advantages: A float32 tensor of shape [T, B]. Can be used as the - advantage in the calculation of policy gradients. - """ - log_rhos = tf.convert_to_tensor(log_rhos, dtype=tf.float32) +def get_rhos(behaviour_action_log_probs, target_action_log_probs): + """With the selected policy values (logits or probs) subclasses compute + the rhos for calculating the vtrace.""" + log_rhos = [t - b for t, + b in zip(target_action_log_probs, behaviour_action_log_probs)] + log_rhos = [tf.convert_to_tensor(l, dtype=tf.float32) for l in log_rhos] + log_rhos = tf.reduce_sum(tf.stack(log_rhos), axis=0) + + return tf.exp(log_rhos) + + +def _from_importance_weights(rhos, discounts, rewards, values, bootstrap_value, + clip_rho_threshold=1.0, clip_pg_rho_threshold=1.0, + name='vtrace_from_importance_weights'): + r"""V-trace from log importance weights. + Calculates V-trace actor critic targets as described in + "IMPALA: Scalable Distributed Deep-RL with + Importance Weighted Actor-Learner Architectures" + by Espeholt, Soyer, Munos et al. + In the notation used throughout documentation and comments, T refers to the + time dimension ranging from 0 to T-1. B refers to the batch size and + NUM_ACTIONS refers to the number of actions. This code also supports the + case where all tensors have the same number of additional dimensions, e.g., + `rewards` is [T, B, C], `values` is [T, B, C], `bootstrap_value` is [B, C]. + Args: + rhos: A float32 tensor of shape [T, B, NUM_ACTIONS] representing the + importance sampling weights, + i.e. target_policy(a) / behaviour_policy(a). + discounts: A float32 tensor of shape [T, B] with discounts encountered + when following the behaviour policy. + rewards: A float32 tensor of shape [T, B] containing rewards generated by + following the behaviour policy. + values: A float32 tensor of shape [T, B] with the value function + estimates wrt. the target policy. + bootstrap_value: A float32 of shape [B] with the value function estimate + at time T. + clip_rho_threshold: A scalar float32 tensor with the clipping threshold + for importance weights (rho) when calculating the baseline targets (vs) + rho^bar in the paper. If None, no clipping is applied. + clip_pg_rho_threshold: A scalar float32 tensor with the clipping + threshold on rho_s in \rho_s \delta log \pi(a|x) + (r + \gamma v_{s+1} - V(x_s)). If None, no clipping is applied. + name: The name scope that all V-trace operations will be created in. + Returns: + A VTraceReturns namedtuple (vs, pg_advantages) where: + vs: A float32 tensor of shape [T, B]. Can be used as target to + train a baseline (V(x_t) - vs_t)^2. + pg_advantages: A float32 tensor of shape [T, B]. Can be used as the + advantage in the calculation of policy gradients. + """ + rhos = tf.convert_to_tensor(rhos, dtype=tf.float32) discounts = tf.convert_to_tensor(discounts, dtype=tf.float32) rewards = tf.convert_to_tensor(rewards, dtype=tf.float32) + rewards = tf.cast(rewards, dtype=tf.float32) values = tf.convert_to_tensor(values, dtype=tf.float32) - bootstrap_value = tf.convert_to_tensor(bootstrap_value, dtype=tf.float32) + bootstrap_value = tf.convert_to_tensor( + bootstrap_value, dtype=tf.float32) if clip_rho_threshold is not None: - clip_rho_threshold = tf.convert_to_tensor( - clip_rho_threshold, dtype=tf.float32) + clip_rho_threshold = tf.convert_to_tensor(clip_rho_threshold, + dtype=tf.float32) if clip_pg_rho_threshold is not None: - clip_pg_rho_threshold = tf.convert_to_tensor( - clip_pg_rho_threshold, dtype=tf.float32) + clip_pg_rho_threshold = tf.convert_to_tensor(clip_pg_rho_threshold, + dtype=tf.float32) # Make sure tensor ranks are consistent. - rho_rank = log_rhos.shape.ndims # Usually 2. + rho_rank = rhos.shape.ndims values.shape.assert_has_rank(rho_rank) bootstrap_value.shape.assert_has_rank(rho_rank - 1) discounts.shape.assert_has_rank(rho_rank) @@ -239,31 +228,41 @@ def from_importance_weights(log_rhos, if clip_pg_rho_threshold is not None: clip_pg_rho_threshold.shape.assert_has_rank(0) - with tf.name_scope( - name, - values=[log_rhos, discounts, rewards, values, bootstrap_value]): - rhos = tf.exp(log_rhos) + with tf.name_scope(name, values=[rhos, discounts, rewards, values, + bootstrap_value]): if clip_rho_threshold is not None: clipped_rhos = tf.minimum( clip_rho_threshold, rhos, name='clipped_rhos') else: clipped_rhos = rhos + tf.summary.histogram('clipped_rhos_1000', tf.minimum(1000.0, rhos)) + tf.summary.scalar( + 'num_of_clipped_rhos', + tf.reduce_sum( + tf.cast( + tf.equal( + clipped_rhos, + clip_rho_threshold), + tf.int32))) + tf.summary.scalar('size_of_clipped_rhos', tf.size(clipped_rhos)) + cs = tf.minimum(1.0, rhos, name='cs') # Append bootstrapped value to get [v1, ..., v_t+1] values_t_plus_1 = tf.concat( [values[1:], tf.expand_dims(bootstrap_value, 0)], axis=0) - deltas = clipped_rhos * ( - rewards + discounts * values_t_plus_1 - values) + deltas = clipped_rhos * \ + (rewards + discounts * values_t_plus_1 - values) - # All sequences are reversed, computation starts from the back. + # Note that all sequences are reversed, computation starts from the + # back. sequences = ( tf.reverse(discounts, axis=[0]), tf.reverse(cs, axis=[0]), tf.reverse(deltas, axis=[0]), ) - # V-trace vs are calculated through a scan from the back to the + # V-trace vs are calculated through a scan from the back to # beginning of the given trajectory. def scanfunc(acc, sequence_item): discount_t, c_t, delta_t = sequence_item @@ -278,23 +277,23 @@ def scanfunc(acc, sequence_item): back_prop=False, name='scan') # Reverse the results back to original order. - vs_minus_v_xs = tf.reverse(vs_minus_v_xs, [0], name='vs_minus_v_xs') + vs_minus_v_xs = tf.reverse( + vs_minus_v_xs, [0], name='vs_minus_v_xs') # Add V(x_s) to get v_s. vs = tf.add(vs_minus_v_xs, values, name='vs') # Advantage for policy gradient. - vs_t_plus_1 = tf.concat( - [vs[1:], tf.expand_dims(bootstrap_value, 0)], axis=0) + vs_t_plus_1 = tf.concat([ + vs[1:], tf.expand_dims(bootstrap_value, 0)], axis=0) if clip_pg_rho_threshold is not None: - clipped_pg_rhos = tf.minimum( - clip_pg_rho_threshold, rhos, name='clipped_pg_rhos') + clipped_pg_rhos = tf.minimum(clip_pg_rho_threshold, rhos, + name='clipped_pg_rhos') else: clipped_pg_rhos = rhos pg_advantages = ( clipped_pg_rhos * (rewards + discounts * vs_t_plus_1 - values)) # Make sure no gradients backpropagated through the returned values. - return VTraceReturns( - vs=tf.stop_gradient(vs), - pg_advantages=tf.stop_gradient(pg_advantages)) + return VTraceReturns(vs=tf.stop_gradient(vs), + pg_advantages=tf.stop_gradient(pg_advantages)) diff --git a/python/ray/rllib/agents/impala/vtrace_policy_graph.py b/python/ray/rllib/agents/impala/vtrace_policy_graph.py index 127c3f9c5365..a176e8288a56 100644 --- a/python/ray/rllib/agents/impala/vtrace_policy_graph.py +++ b/python/ray/rllib/agents/impala/vtrace_policy_graph.py @@ -18,7 +18,7 @@ from ray.rllib.utils.annotations import override from ray.rllib.utils.error import UnsupportedSpaceException from ray.rllib.utils.explained_variance import explained_variance -from ray.rllib.models.action_dist import Categorical +from ray.rllib.models.action_dist import MultiCategorical class VTraceLoss(object): @@ -45,12 +45,20 @@ def __init__(self, handle episode cut boundaries. Args: - actions: An int32 tensor of shape [T, B, NUM_ACTIONS]. + actions: An int32 tensor of shape [T, B, ACTION_SPACE]. actions_logp: A float32 tensor of shape [T, B]. actions_entropy: A float32 tensor of shape [T, B]. dones: A bool tensor of shape [T, B]. - behaviour_logits: A float32 tensor of shape [T, B, NUM_ACTIONS]. - target_logits: A float32 tensor of shape [T, B, NUM_ACTIONS]. + behaviour_logits: A list with length of ACTION_SPACE of float32 + tensors of shapes + [T, B, ACTION_SPACE[0]], + ..., + [T, B, ACTION_SPACE[-1]] + target_logits: A list with length of ACTION_SPACE of float32 + tensors of shapes + [T, B, ACTION_SPACE[0]], + ..., + [T, B, ACTION_SPACE[-1]] discount: A float32 scalar. rewards: A float32 tensor of shape [T, B]. values: A float32 tensor of shape [T, B]. @@ -61,9 +69,9 @@ def __init__(self, # Compute vtrace on the CPU for better perf. with tf.device("/cpu:0"): self.vtrace_returns = vtrace.from_logits( - behaviour_policy_logits=behaviour_logits, - target_policy_logits=target_logits, - actions=tf.cast(actions, tf.int32), + behaviour_policy=behaviour_logits, + target_policy=target_logits, + actions=tf.unstack(tf.cast(actions, tf.int32), axis=2), discounts=tf.to_float(~dones) * discount, rewards=rewards, values=values, @@ -101,6 +109,11 @@ def __init__(self, "Must use `truncate_episodes` batch mode with V-trace." self.config = config self.sess = tf.get_default_session() + self._is_discrete = False + self.grads = None + + output_hidden_shape = None + actions_shape = [None] # Create input placeholders if existing_inputs: @@ -110,16 +123,25 @@ def __init__(self, existing_seq_lens = existing_inputs[-1] else: if isinstance(action_space, gym.spaces.Discrete): - ac_size = action_space.n - actions = tf.placeholder(tf.int64, [None], name="ac") + self._is_discrete = True + output_hidden_shape = [action_space.n] + elif isinstance(action_space, + gym.spaces.multi_discrete.MultiDiscrete): + actions_shape = [None, len(action_space.nvec)] + output_hidden_shape = action_space.nvec else: raise UnsupportedSpaceException( "Action space {} is not supported for IMPALA.".format( action_space)) + + actions = tf.placeholder(tf.int64, actions_shape, name="ac") dones = tf.placeholder(tf.bool, [None], name="dones") rewards = tf.placeholder(tf.float32, [None], name="rewards") - behaviour_logits = tf.placeholder( - tf.float32, [None, ac_size], name="behaviour_logits") + behaviour_logits = tf.placeholder(tf.float32, + [None, sum(output_hidden_shape)], + name="behaviour_logits") + unpacked_behaviour_logits = tf.split( + behaviour_logits, output_hidden_shape, axis=1) observations = tf.placeholder( tf.float32, [None] + list(observation_space.shape)) existing_state_in = None @@ -127,7 +149,8 @@ def __init__(self, # Setup the policy dist_class, logit_dim = ModelCatalog.get_action_dist( - action_space, self.config["model"]) + action_space, self.config["model"], + dist_type=self.config["dist_type"]) prev_actions = ModelCatalog.get_action_placeholder(action_space) prev_rewards = tf.placeholder(tf.float32, [None], name="prev_reward") self.model = ModelCatalog.get_model( @@ -142,12 +165,30 @@ def __init__(self, self.config["model"], state_in=existing_state_in, seq_lens=existing_seq_lens) - action_dist = dist_class(self.model.outputs) + unpacked_outputs = tf.split( + self.model.outputs, output_hidden_shape, axis=1) + + dist_inputs = self.model.outputs if self._is_discrete else \ + unpacked_outputs + action_dist = dist_class(dist_inputs) + values = self.model.value_function() self.var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, tf.get_variable_scope().name) - def to_batches(tensor): + def make_time_major(tensor, drop_last=False): + """Swaps batch and trajectory axis. + Args: + tensor: A tensor or list of tensors to reshape. + drop_last: A bool indicating whether to drop the last + trajectory item. + Returns: + res: A tensor with swapped axes or a list of tensors with + swapped axes. + """ + if isinstance(tensor, list): + return [make_time_major(t, drop_last) for t in tensor] + if self.config["model"]["use_lstm"]: B = tf.shape(self.model.seq_lens)[0] T = tf.shape(tensor)[0] // B @@ -158,11 +199,16 @@ def to_batches(tensor): B = tf.shape(tensor)[0] // T rs = tf.reshape(tensor, tf.concat([[B, T], tf.shape(tensor)[1:]], axis=0)) + # swap B and T axes - return tf.transpose( + res = tf.transpose( rs, [1, 0] + list(range(2, 1 + int(tf.shape(tensor).shape[0])))) + if drop_last: + return res[:-1] + return res + if self.model.state_in: max_seq_len = tf.reduce_max(self.model.seq_lens) - 1 mask = tf.sequence_mask(self.model.seq_lens, max_seq_len) @@ -170,31 +216,55 @@ def to_batches(tensor): else: mask = tf.ones_like(rewards, dtype=tf.bool) + # Prepare actions for loss + loss_actions = tf.expand_dims( + actions, axis=1) if self._is_discrete else actions + logp_action = actions if self._is_discrete else tf.unstack( + actions, axis=1) + # Inputs are reshaped from [B * T] => [T - 1, B] for V-trace calc. self.loss = VTraceLoss( - actions=to_batches(actions)[:-1], - actions_logp=to_batches(action_dist.logp(actions))[:-1], - actions_entropy=to_batches(action_dist.entropy())[:-1], - dones=to_batches(dones)[:-1], - behaviour_logits=to_batches(behaviour_logits)[:-1], - target_logits=to_batches(self.model.outputs)[:-1], + actions=make_time_major(loss_actions, drop_last=True), + actions_logp=make_time_major(action_dist.logp(logp_action), + drop_last=True), + actions_entropy=make_time_major(action_dist.entropy(), + drop_last=True), + dones=make_time_major(dones, drop_last=True), + behaviour_logits=make_time_major( + unpacked_behaviour_logits, drop_last=True), + target_logits=make_time_major(unpacked_outputs, drop_last=True), discount=config["gamma"], - rewards=to_batches(rewards)[:-1], - values=to_batches(values)[:-1], - bootstrap_value=to_batches(values)[-1], - valid_mask=to_batches(mask)[:-1], + rewards=make_time_major(rewards, drop_last=True), + values=make_time_major(values, drop_last=True), + bootstrap_value=make_time_major(values)[-1], + valid_mask=make_time_major(mask, drop_last=True), vf_loss_coeff=self.config["vf_loss_coeff"], entropy_coeff=self.config["entropy_coeff"], clip_rho_threshold=self.config["vtrace_clip_rho_threshold"], clip_pg_rho_threshold=self.config["vtrace_clip_pg_rho_threshold"]) # KL divergence between worker and learner logits for debugging - model_dist = Categorical(self.model.outputs) - behaviour_dist = Categorical(behaviour_logits) - self.KLs = model_dist.kl(behaviour_dist) - self.mean_KL = tf.reduce_mean(self.KLs) - self.max_KL = tf.reduce_max(self.KLs) - self.median_KL = tf.contrib.distributions.percentile(self.KLs, 50.0) + model_dist = MultiCategorical(unpacked_outputs) + behaviour_dist = MultiCategorical(unpacked_behaviour_logits) + + kls = model_dist.kl(behaviour_dist) + if len(kls) > 1: + self.KL_stats = {} + + for i, kl in enumerate(kls): + self.KL_stats.update({ + f"mean_KL_{i}": tf.reduce_mean(kl), + f"max_KL_{i}": tf.reduce_max(kl), + f"median_KL_{i}": tf.contrib.distributions.percentile( + kl, 50.0), + }) + else: + self.KL_stats = { + "mean_KL": tf.reduce_mean(kls[0]), + "max_KL": tf.reduce_max(kls[0]), + "median_KL": tf.contrib.distributions.percentile( + kls[0], 50.0), + } # Initialize TFPolicyGraph loss_in = [ @@ -237,10 +307,8 @@ def to_batches(tensor): "vf_loss": self.loss.vf_loss, "vf_explained_var": explained_variance( tf.reshape(self.loss.vtrace_returns.vs, [-1]), - tf.reshape(to_batches(values)[:-1], [-1])), - "mean_KL": self.mean_KL, - "max_KL": self.max_KL, - "median_KL": self.median_KL, + tf.reshape(make_time_major(values, drop_last=True), [-1])), + **self.KL_stats, }, } diff --git a/python/ray/rllib/models/action_dist.py b/python/ray/rllib/models/action_dist.py index cad979201fa4..0a43d5ab7732 100644 --- a/python/ray/rllib/models/action_dist.py +++ b/python/ray/rllib/models/action_dist.py @@ -99,6 +99,28 @@ def sample(self): return tf.squeeze(tf.multinomial(self.inputs, 1), axis=1) +class MultiCategorical(ActionDistribution): + """Categorical distribution for discrete action spaces.""" + + def __init__(self, inputs): + self.cats = [Categorical(input_) for input_ in inputs] + + def logp(self, actions): + logps = tf.stack([cat.logp(act) + for cat, act in zip(self.cats, actions)]) + return tf.reduce_sum(logps, axis=0) + + def entropy(self): + return tf.stack([cat.entropy() for cat in self.cats], axis=1) + + def kl(self, other): + return [cat.kl(oth_cat) + for cat, oth_cat in zip(self.cats, other.cats)] + + def sample(self): + return tf.stack([cat.sample() for cat in self.cats], axis=1) + + class DiagGaussian(ActionDistribution): """Action distribution where each vector element is a gaussian. @@ -117,8 +139,8 @@ def __init__(self, inputs): def logp(self, x): return (-0.5 * tf.reduce_sum( tf.square((x - self.mean) / self.std), reduction_indices=[1]) - - 0.5 * np.log(2.0 * np.pi) * tf.to_float(tf.shape(x)[1]) - - tf.reduce_sum(self.log_std, reduction_indices=[1])) + 0.5 * np.log(2.0 * np.pi) * tf.to_float(tf.shape(x)[1]) - + tf.reduce_sum(self.log_std, reduction_indices=[1])) @override(ActionDistribution) def kl(self, other): diff --git a/python/ray/rllib/models/catalog.py b/python/ray/rllib/models/catalog.py index 474a0e9056b3..9c2c45701a92 100644 --- a/python/ray/rllib/models/catalog.py +++ b/python/ray/rllib/models/catalog.py @@ -12,7 +12,8 @@ _global_registry from ray.rllib.models.action_dist import ( - Categorical, Deterministic, DiagGaussian, MultiActionDistribution) + Categorical, MultiCategorical, Deterministic, DiagGaussian, + MultiActionDistribution) from ray.rllib.models.preprocessors import get_preprocessor from ray.rllib.models.fcnet import FullyConnectedNetwork from ray.rllib.models.visionnet import VisionNetwork @@ -132,6 +133,8 @@ def get_action_dist(action_space, config, dist_type=None): child_distributions=child_dist, action_space=action_space, input_lens=input_lens), sum(input_lens) + elif isinstance(action_space, gym.spaces.multi_discrete.MultiDiscrete): + return MultiCategorical, sum(action_space.nvec) raise NotImplementedError("Unsupported args: {} {}".format( action_space, dist_type)) @@ -165,6 +168,11 @@ def get_action_placeholder(action_space): tf.int64 if all_discrete else tf.float32, shape=(None, size), name="action") + elif isinstance(action_space, gym.spaces.multi_discrete.MultiDiscrete): + return tf.placeholder( + tf.as_dtype(action_space.dtype), + shape=(None, len(action_space.nvec)), + name="action") else: raise NotImplementedError("action space {}" " not supported".format(action_space)) From 088985b0bbdbae7a05693540872e5fbd19f5ac69 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleksandar=20Milovanovi=C4=87?= Date: Wed, 6 Feb 2019 17:16:07 +0100 Subject: [PATCH 02/33] fixed newlines --- python/ray/rllib/agents/impala/vtrace.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/python/ray/rllib/agents/impala/vtrace.py b/python/ray/rllib/agents/impala/vtrace.py index 50dfc813f147..f740cf5089b2 100644 --- a/python/ray/rllib/agents/impala/vtrace.py +++ b/python/ray/rllib/agents/impala/vtrace.py @@ -13,10 +13,13 @@ # limitations under the License. """Functions to compute V-trace off-policy actor critic targets. + For details and theory see: + "IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures" by Espeholt, Soyer, Munos et al. + See https://arxiv.org/abs/1802.01561 for the full paper. """ @@ -43,16 +46,21 @@ def from_logits(behaviour_policy, target_policy, actions, clip_rho_threshold=1.0, clip_pg_rho_threshold=1.0, name='vtrace_from_logits'): r"""V-trace for softmax policies. + Calculates V-trace actor critic targets for softmax polices as described in + "IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures" by Espeholt, Soyer, Munos et al. + Target policy refers to the policy we are interested in improving and behaviour policy refers to the policy that generated the given rewards and actions. + In the notation used throughout documentation and comments, T refers to the time dimension ranging from 0 to T-1. B refers to the batch size and NUM_ACTIONS refers to the number of actions. + Args: behaviour_policy: A float32 tensor of shape [T, B, NUM_ACTIONS] with un-normalized log-probabilities parametrizing the softmax behaviour @@ -77,6 +85,7 @@ def from_logits(behaviour_policy, target_policy, actions, threshold on rho_s in \rho_s \delta log \pi(a|x) (r + \gamma v_{s+1} - V(x_s)). name: The name scope that all V-trace operations will be created in. + Returns: A `VTraceFromLogitsReturns` namedtuple with the following fields: vs: A float32 tensor of shape [T, B]. Can be used as target to train a @@ -133,13 +142,16 @@ def from_logits(behaviour_policy, target_policy, actions, def select_policy_values_using_actions(policy_logits, actions): """ Computes action log-probs from policy logits and actions. + In the notation used throughout documentation and comments, T refers to the time dimension ranging from 0 to T-1. B refers to the batch size and NUM_ACTIONS refers to the number of actions. + Args: policy_logits: A float32 tensor of shape [T, B, NUM_ACTIONS] with un-normalized log-probabilities parameterizing a softmax policy. actions: An int32 tensor of shape [T, B] with actions. + Returns: A float32 tensor of shape [T, B] corresponding to the sampling log probability of the chosen action w.r.t. the policy. @@ -168,15 +180,19 @@ def _from_importance_weights(rhos, discounts, rewards, values, bootstrap_value, clip_rho_threshold=1.0, clip_pg_rho_threshold=1.0, name='vtrace_from_importance_weights'): r"""V-trace from log importance weights. + Calculates V-trace actor critic targets as described in + "IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures" by Espeholt, Soyer, Munos et al. + In the notation used throughout documentation and comments, T refers to the time dimension ranging from 0 to T-1. B refers to the batch size and NUM_ACTIONS refers to the number of actions. This code also supports the case where all tensors have the same number of additional dimensions, e.g., `rewards` is [T, B, C], `values` is [T, B, C], `bootstrap_value` is [B, C]. + Args: rhos: A float32 tensor of shape [T, B, NUM_ACTIONS] representing the importance sampling weights, @@ -196,6 +212,7 @@ def _from_importance_weights(rhos, discounts, rewards, values, bootstrap_value, threshold on rho_s in \rho_s \delta log \pi(a|x) (r + \gamma v_{s+1} - V(x_s)). If None, no clipping is applied. name: The name scope that all V-trace operations will be created in. + Returns: A VTraceReturns namedtuple (vs, pg_advantages) where: vs: A float32 tensor of shape [T, B]. Can be used as target to From 3b38ebb21f2f2b8df65b9fc62a02169ab51105f6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleksandar=20Milovanovi=C4=87?= Date: Thu, 7 Feb 2019 12:13:23 +0100 Subject: [PATCH 03/33] reformatting impalla.py --- python/ray/rllib/agents/impala/vtrace.py | 271 +++++++++++------------ 1 file changed, 135 insertions(+), 136 deletions(-) diff --git a/python/ray/rllib/agents/impala/vtrace.py b/python/ray/rllib/agents/impala/vtrace.py index f740cf5089b2..8fb06c44e342 100644 --- a/python/ray/rllib/agents/impala/vtrace.py +++ b/python/ray/rllib/agents/impala/vtrace.py @@ -11,7 +11,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - """Functions to compute V-trace off-policy actor critic targets. For details and theory see: @@ -33,73 +32,99 @@ nest = tf.contrib.framework.nest -VTraceFromLogitsReturns = collections.namedtuple( - 'VTraceFromLogitsReturns', - ['vs', 'pg_advantages', 'rhos', - 'behaviour_action_policy', 'target_action_policy']) +VTraceFromLogitsReturns = collections.namedtuple('VTraceFromLogitsReturns', [ + 'vs', 'pg_advantages', 'rhos', 'behaviour_action_policy', + 'target_action_policy' +]) VTraceReturns = collections.namedtuple('VTraceReturns', 'vs pg_advantages') -def from_logits(behaviour_policy, target_policy, actions, - discounts, rewards, values, bootstrap_value, - clip_rho_threshold=1.0, clip_pg_rho_threshold=1.0, - name='vtrace_from_logits'): - r"""V-trace for softmax policies. - - Calculates V-trace actor critic targets for softmax polices as described in - - "IMPALA: Scalable Distributed Deep-RL with - Importance Weighted Actor-Learner Architectures" - by Espeholt, Soyer, Munos et al. - - Target policy refers to the policy we are interested in improving and - behaviour policy refers to the policy that generated the given - rewards and actions. +def select_policy_values_using_actions(policy_logits, actions): + """ + Computes action log-probs from policy logits and actions. In the notation used throughout documentation and comments, T refers to the time dimension ranging from 0 to T-1. B refers to the batch size and NUM_ACTIONS refers to the number of actions. Args: - behaviour_policy: A float32 tensor of shape [T, B, NUM_ACTIONS] with - un-normalized log-probabilities parametrizing the softmax behaviour - policy. - target_policy: A float32 tensor of shape [T, B, NUM_ACTIONS] with - un-normalized log-probabilities parametrizing the softmax target - policy. - actions: An int32 tensor of shape [T, B] of actions sampled from the - behaviour policy. - discounts: A float32 tensor of shape [T, B] with the discount encountered - when following the behaviour policy. - rewards: A float32 tensor of shape [T, B] with the rewards generated by - following the behaviour policy. - values: A float32 tensor of shape [T, B] with the value function - estimates wrt. the target policy. - bootstrap_value: A float32 of shape [B] with the value function estimate - at time T. - clip_rho_threshold: A scalar float32 tensor with the clipping threshold - for importance weights (rho) when calculating the baseline targets (vs) - rho^bar in the paper. - clip_pg_rho_threshold: A scalar float32 tensor with the clipping - threshold on rho_s in - \rho_s \delta log \pi(a|x) (r + \gamma v_{s+1} - V(x_s)). - name: The name scope that all V-trace operations will be created in. + policy_logits: A float32 tensor of shape [T, B, NUM_ACTIONS] with + un-normalized log-probabilities parameterizing a softmax policy. + actions: An int32 tensor of shape [T, B] with actions. Returns: - A `VTraceFromLogitsReturns` namedtuple with the following fields: - vs: A float32 tensor of shape [T, B]. Can be used as target to train a - baseline (V(x_t) - vs_t)^2. - pg_advantages: A float 32 tensor of shape [T, B]. Can be used as an - estimate of the advantage in the calculation of policy gradients. - log_rhos: A float32 tensor of shape [T, B] containing the log - importance sampling weights (log rhos). - behaviour_action_log_probs: A float32 tensor of shape [T, B] containing - behaviour policy action log probabilities (log \mu(a_t)). - target_action_log_probs: A float32 tensor of shape [T, B] containing - target policy action probabilities (log \pi(a_t)). + A float32 tensor of shape [T, B] corresponding to the sampling log + probability of the chosen action w.r.t. the policy. """ + log_probs = [] + for i in range(len(policy_logits)): + log_probs.append(-tf.nn.sparse_softmax_cross_entropy_with_logits( + logits=policy_logits[i], labels=actions[i])) + + return log_probs + + +def from_logits(behaviour_policy, target_policy, actions, + discounts, rewards, values, bootstrap_value, + clip_rho_threshold=1.0, clip_pg_rho_threshold=1.0, + name='vtrace_from_logits'): + r"""V-trace for softmax policies. + + Calculates V-trace actor critic targets for softmax polices as described in + + "IMPALA: Scalable Distributed Deep-RL with + Importance Weighted Actor-Learner Architectures" + by Espeholt, Soyer, Munos et al. + + Target policy refers to the policy we are interested in improving and + behaviour policy refers to the policy that generated the given + rewards and actions. + + In the notation used throughout documentation and comments, T refers to the + time dimension ranging from 0 to T-1. B refers to the batch size and + NUM_ACTIONS refers to the number of actions. + + Args: + behaviour_policy: A float32 tensor of shape [T, B, NUM_ACTIONS] with + un-normalized log-probabilities parametrizing the softmax behaviour + policy. + target_policy: A float32 tensor of shape [T, B, NUM_ACTIONS] with + un-normalized log-probabilities parametrizing the softmax target + policy. + actions: An int32 tensor of shape [T, B] of actions sampled from the + behaviour policy. + discounts: A float32 tensor of shape [T, B] with the discount encountered + when following the behaviour policy. + rewards: A float32 tensor of shape [T, B] with the rewards generated by + following the behaviour policy. + values: A float32 tensor of shape [T, B] with the value function + estimates wrt. the target policy. + bootstrap_value: A float32 of shape [B] with the value function estimate + at time T. + clip_rho_threshold: A scalar float32 tensor with the clipping threshold + for importance weights (rho) when calculating the baseline targets (vs) + rho^bar in the paper. + clip_pg_rho_threshold: A scalar float32 tensor with the clipping + threshold on rho_s in + \rho_s \delta log \pi(a|x) (r + \gamma v_{s+1} - V(x_s)). + name: The name scope that all V-trace operations will be created in. + + Returns: + A `VTraceFromLogitsReturns` namedtuple with the following fields: + vs: A float32 tensor of shape [T, B]. Can be used as target to train a + baseline (V(x_t) - vs_t)^2. + pg_advantages: A float 32 tensor of shape [T, B]. Can be used as an + estimate of the advantage in the calculation of policy gradients. + log_rhos: A float32 tensor of shape [T, B] containing the log + importance sampling weights (log rhos). + behaviour_action_log_probs: A float32 tensor of shape [T, B] containing + behaviour policy action log probabilities (log \mu(a_t)). + target_action_log_probs: A float32 tensor of shape [T, B] containing + target policy action probabilities (log \pi(a_t)). + """ + for i in range(len(behaviour_policy)): behaviour_policy[i] = tf.convert_to_tensor( behaviour_policy[i], dtype=tf.float32) @@ -123,7 +148,7 @@ def from_logits(behaviour_policy, target_policy, actions, rhos = get_rhos(target_action_policy, behaviour_action_policy) - vtrace_returns = _from_importance_weights( + vtrace_returns = from_importance_weights( rhos=rhos, discounts=discounts, rewards=rewards, @@ -139,87 +164,50 @@ def from_logits(behaviour_policy, target_policy, actions, **vtrace_returns._asdict()) -def select_policy_values_using_actions(policy_logits, actions): - """ - Computes action log-probs from policy logits and actions. - - In the notation used throughout documentation and comments, T refers to the - time dimension ranging from 0 to T-1. B refers to the batch size and - NUM_ACTIONS refers to the number of actions. - - Args: - policy_logits: A float32 tensor of shape [T, B, NUM_ACTIONS] with - un-normalized log-probabilities parameterizing a softmax policy. - actions: An int32 tensor of shape [T, B] with actions. - - Returns: - A float32 tensor of shape [T, B] corresponding to the sampling log - probability of the chosen action w.r.t. the policy. - """ - - log_probs = [] - for i in range(len(policy_logits)): - log_probs.append(-tf.nn.sparse_softmax_cross_entropy_with_logits( - logits=policy_logits[i], labels=actions[i])) - - return log_probs - - -def get_rhos(behaviour_action_log_probs, target_action_log_probs): - """With the selected policy values (logits or probs) subclasses compute - the rhos for calculating the vtrace.""" - log_rhos = [t - b for t, - b in zip(target_action_log_probs, behaviour_action_log_probs)] - log_rhos = [tf.convert_to_tensor(l, dtype=tf.float32) for l in log_rhos] - log_rhos = tf.reduce_sum(tf.stack(log_rhos), axis=0) - - return tf.exp(log_rhos) - - -def _from_importance_weights(rhos, discounts, rewards, values, bootstrap_value, - clip_rho_threshold=1.0, clip_pg_rho_threshold=1.0, - name='vtrace_from_importance_weights'): - r"""V-trace from log importance weights. - - Calculates V-trace actor critic targets as described in - - "IMPALA: Scalable Distributed Deep-RL with - Importance Weighted Actor-Learner Architectures" - by Espeholt, Soyer, Munos et al. - - In the notation used throughout documentation and comments, T refers to the - time dimension ranging from 0 to T-1. B refers to the batch size and - NUM_ACTIONS refers to the number of actions. This code also supports the - case where all tensors have the same number of additional dimensions, e.g., - `rewards` is [T, B, C], `values` is [T, B, C], `bootstrap_value` is [B, C]. - - Args: - rhos: A float32 tensor of shape [T, B, NUM_ACTIONS] representing the - importance sampling weights, - i.e. target_policy(a) / behaviour_policy(a). - discounts: A float32 tensor of shape [T, B] with discounts encountered - when following the behaviour policy. - rewards: A float32 tensor of shape [T, B] containing rewards generated by - following the behaviour policy. - values: A float32 tensor of shape [T, B] with the value function - estimates wrt. the target policy. - bootstrap_value: A float32 of shape [B] with the value function estimate - at time T. - clip_rho_threshold: A scalar float32 tensor with the clipping threshold - for importance weights (rho) when calculating the baseline targets (vs) - rho^bar in the paper. If None, no clipping is applied. - clip_pg_rho_threshold: A scalar float32 tensor with the clipping - threshold on rho_s in \rho_s \delta log \pi(a|x) - (r + \gamma v_{s+1} - V(x_s)). If None, no clipping is applied. - name: The name scope that all V-trace operations will be created in. - - Returns: - A VTraceReturns namedtuple (vs, pg_advantages) where: - vs: A float32 tensor of shape [T, B]. Can be used as target to - train a baseline (V(x_t) - vs_t)^2. - pg_advantages: A float32 tensor of shape [T, B]. Can be used as the - advantage in the calculation of policy gradients. - """ +def from_importance_weights(rhos, discounts, rewards, values, bootstrap_value, + clip_rho_threshold=1.0, clip_pg_rho_threshold=1.0, + name='vtrace_from_importance_weights'): + r"""V-trace from log importance weights. + + Calculates V-trace actor critic targets as described in + + "IMPALA: Scalable Distributed Deep-RL with + Importance Weighted Actor-Learner Architectures" + by Espeholt, Soyer, Munos et al. + + In the notation used throughout documentation and comments, T refers to the + time dimension ranging from 0 to T-1. B refers to the batch size and + NUM_ACTIONS refers to the number of actions. This code also supports the + case where all tensors have the same number of additional dimensions, e.g., + `rewards` is [T, B, C], `values` is [T, B, C], `bootstrap_value` is [B, C]. + + Args: + rhos: A float32 tensor of shape [T, B, NUM_ACTIONS] representing the + importance sampling weights, + i.e. target_policy(a) / behaviour_policy(a). + discounts: A float32 tensor of shape [T, B] with discounts encountered + when following the behaviour policy. + rewards: A float32 tensor of shape [T, B] containing rewards generated by + following the behaviour policy. + values: A float32 tensor of shape [T, B] with the value function + estimates wrt. the target policy. + bootstrap_value: A float32 of shape [B] with the value function estimate + at time T. + clip_rho_threshold: A scalar float32 tensor with the clipping threshold + for importance weights (rho) when calculating the baseline targets (vs) + rho^bar in the paper. If None, no clipping is applied. + clip_pg_rho_threshold: A scalar float32 tensor with the clipping + threshold on rho_s in \rho_s \delta log \pi(a|x) + (r + \gamma v_{s+1} - V(x_s)). If None, no clipping is applied. + name: The name scope that all V-trace operations will be created in. + + Returns: + A VTraceReturns namedtuple (vs, pg_advantages) where: + vs: A float32 tensor of shape [T, B]. Can be used as target to + train a baseline (V(x_t) - vs_t)^2. + pg_advantages: A float32 tensor of shape [T, B]. Can be used as the + advantage in the calculation of policy gradients. + """ rhos = tf.convert_to_tensor(rhos, dtype=tf.float32) discounts = tf.convert_to_tensor(discounts, dtype=tf.float32) rewards = tf.convert_to_tensor(rewards, dtype=tf.float32) @@ -314,3 +302,14 @@ def scanfunc(acc, sequence_item): # Make sure no gradients backpropagated through the returned values. return VTraceReturns(vs=tf.stop_gradient(vs), pg_advantages=tf.stop_gradient(pg_advantages)) + + +def get_rhos(behaviour_action_log_probs, target_action_log_probs): + """With the selected policy values (logits or probs) subclasses compute + the rhos for calculating the vtrace.""" + log_rhos = [t - b for t, + b in zip(target_action_log_probs, behaviour_action_log_probs)] + log_rhos = [tf.convert_to_tensor(l, dtype=tf.float32) for l in log_rhos] + log_rhos = tf.reduce_sum(tf.stack(log_rhos), axis=0) + + return tf.exp(log_rhos) From 1858404009ed229e167ed6fbe07501fd46d97080 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleksandar=20Milovanovi=C4=87?= Date: Thu, 7 Feb 2019 12:19:30 +0100 Subject: [PATCH 04/33] aligned vtrace.py formatting some more --- python/ray/rllib/agents/impala/vtrace.py | 84 ++++++++++++------------ 1 file changed, 43 insertions(+), 41 deletions(-) diff --git a/python/ray/rllib/agents/impala/vtrace.py b/python/ray/rllib/agents/impala/vtrace.py index 8fb06c44e342..ce9470d95669 100644 --- a/python/ray/rllib/agents/impala/vtrace.py +++ b/python/ray/rllib/agents/impala/vtrace.py @@ -42,21 +42,21 @@ def select_policy_values_using_actions(policy_logits, actions): """ - Computes action log-probs from policy logits and actions. + Computes action log-probs from policy logits and actions. - In the notation used throughout documentation and comments, T refers to the - time dimension ranging from 0 to T-1. B refers to the batch size and - NUM_ACTIONS refers to the number of actions. + In the notation used throughout documentation and comments, T refers to the + time dimension ranging from 0 to T-1. B refers to the batch size and + NUM_ACTIONS refers to the number of actions. - Args: - policy_logits: A float32 tensor of shape [T, B, NUM_ACTIONS] with - un-normalized log-probabilities parameterizing a softmax policy. - actions: An int32 tensor of shape [T, B] with actions. + Args: + policy_logits: A float32 tensor of shape [T, B, NUM_ACTIONS] with + un-normalized log-probabilities parameterizing a softmax policy. + actions: An int32 tensor of shape [T, B] with actions. - Returns: - A float32 tensor of shape [T, B] corresponding to the sampling log - probability of the chosen action w.r.t. the policy. - """ + Returns: + A float32 tensor of shape [T, B] corresponding to the sampling log + probability of the chosen action w.r.t. the policy. + """ log_probs = [] for i in range(len(policy_logits)): @@ -66,9 +66,14 @@ def select_policy_values_using_actions(policy_logits, actions): return log_probs -def from_logits(behaviour_policy, target_policy, actions, - discounts, rewards, values, bootstrap_value, - clip_rho_threshold=1.0, clip_pg_rho_threshold=1.0, +def from_logits(behaviour_policy, + target_policy, actions, + discounts, + rewards, + values, + bootstrap_value, + clip_rho_threshold=1.0, + clip_pg_rho_threshold=1.0, name='vtrace_from_logits'): r"""V-trace for softmax policies. @@ -91,24 +96,22 @@ def from_logits(behaviour_policy, target_policy, actions, un-normalized log-probabilities parametrizing the softmax behaviour policy. target_policy: A float32 tensor of shape [T, B, NUM_ACTIONS] with - un-normalized log-probabilities parametrizing the softmax target - policy. + un-normalized log-probabilities parametrizing the softmax target policy. actions: An int32 tensor of shape [T, B] of actions sampled from the behaviour policy. discounts: A float32 tensor of shape [T, B] with the discount encountered when following the behaviour policy. rewards: A float32 tensor of shape [T, B] with the rewards generated by following the behaviour policy. - values: A float32 tensor of shape [T, B] with the value function - estimates wrt. the target policy. - bootstrap_value: A float32 of shape [B] with the value function estimate - at time T. - clip_rho_threshold: A scalar float32 tensor with the clipping threshold - for importance weights (rho) when calculating the baseline targets (vs) + values: A float32 tensor of shape [T, B] with the value function estimates + wrt. the target policy. + bootstrap_value: A float32 of shape [B] with the value function estimate at + time T. + clip_rho_threshold: A scalar float32 tensor with the clipping threshold for + importance weights (rho) when calculating the baseline targets (vs). rho^bar in the paper. - clip_pg_rho_threshold: A scalar float32 tensor with the clipping - threshold on rho_s in - \rho_s \delta log \pi(a|x) (r + \gamma v_{s+1} - V(x_s)). + clip_pg_rho_threshold: A scalar float32 tensor with the clipping threshold + on rho_s in \rho_s \delta log \pi(a|x) (r + \gamma v_{s+1} - V(x_s)). name: The name scope that all V-trace operations will be created in. Returns: @@ -117,8 +120,8 @@ def from_logits(behaviour_policy, target_policy, actions, baseline (V(x_t) - vs_t)^2. pg_advantages: A float 32 tensor of shape [T, B]. Can be used as an estimate of the advantage in the calculation of policy gradients. - log_rhos: A float32 tensor of shape [T, B] containing the log - importance sampling weights (log rhos). + log_rhos: A float32 tensor of shape [T, B] containing the log importance + sampling weights (log rhos). behaviour_action_log_probs: A float32 tensor of shape [T, B] containing behaviour policy action log probabilities (log \mu(a_t)). target_action_log_probs: A float32 tensor of shape [T, B] containing @@ -183,22 +186,21 @@ def from_importance_weights(rhos, discounts, rewards, values, bootstrap_value, Args: rhos: A float32 tensor of shape [T, B, NUM_ACTIONS] representing the - importance sampling weights, - i.e. target_policy(a) / behaviour_policy(a). - discounts: A float32 tensor of shape [T, B] with discounts encountered - when following the behaviour policy. + importance sampling weights, i.e. target_policy(a) / behaviour_policy(a). + discounts: A float32 tensor of shape [T, B] with discounts encountered when + following the behaviour policy. rewards: A float32 tensor of shape [T, B] containing rewards generated by following the behaviour policy. - values: A float32 tensor of shape [T, B] with the value function - estimates wrt. the target policy. - bootstrap_value: A float32 of shape [B] with the value function estimate - at time T. - clip_rho_threshold: A scalar float32 tensor with the clipping threshold - for importance weights (rho) when calculating the baseline targets (vs) + values: A float32 tensor of shape [T, B] with the value function estimates + wrt. the target policy. + bootstrap_value: A float32 of shape [B] with the value function estimate at + time T. + clip_rho_threshold: A scalar float32 tensor with the clipping threshold for + importance weights (rho) when calculating the baseline targets (vs). rho^bar in the paper. If None, no clipping is applied. - clip_pg_rho_threshold: A scalar float32 tensor with the clipping - threshold on rho_s in \rho_s \delta log \pi(a|x) - (r + \gamma v_{s+1} - V(x_s)). If None, no clipping is applied. + clip_pg_rho_threshold: A scalar float32 tensor with the clipping threshold + on rho_s in \rho_s \delta log \pi(a|x) (r + \gamma v_{s+1} - V(x_s)). If + None, no clipping is applied. name: The name scope that all V-trace operations will be created in. Returns: From 9840eb610570e595d59d51f321029ed7ef25ef39 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleksandar=20Milovanovi=C4=87?= Date: Thu, 7 Feb 2019 12:24:44 +0100 Subject: [PATCH 05/33] aligned formatting some more --- python/ray/rllib/agents/impala/vtrace.py | 46 ++++++++++++------------ 1 file changed, 23 insertions(+), 23 deletions(-) diff --git a/python/ray/rllib/agents/impala/vtrace.py b/python/ray/rllib/agents/impala/vtrace.py index ce9470d95669..24cd690b1aef 100644 --- a/python/ray/rllib/agents/impala/vtrace.py +++ b/python/ray/rllib/agents/impala/vtrace.py @@ -41,8 +41,7 @@ def select_policy_values_using_actions(policy_logits, actions): - """ - Computes action log-probs from policy logits and actions. + """Computes action log-probs from policy logits and actions. In the notation used throughout documentation and comments, T refers to the time dimension ranging from 0 to T-1. B refers to the batch size and @@ -67,7 +66,8 @@ def select_policy_values_using_actions(policy_logits, actions): def from_logits(behaviour_policy, - target_policy, actions, + target_policy, + actions, discounts, rewards, values, @@ -170,7 +170,7 @@ def from_logits(behaviour_policy, def from_importance_weights(rhos, discounts, rewards, values, bootstrap_value, clip_rho_threshold=1.0, clip_pg_rho_threshold=1.0, name='vtrace_from_importance_weights'): - r"""V-trace from log importance weights. + r"""V-trace from log importance weights. Calculates V-trace actor critic targets as described in @@ -215,14 +215,13 @@ def from_importance_weights(rhos, discounts, rewards, values, bootstrap_value, rewards = tf.convert_to_tensor(rewards, dtype=tf.float32) rewards = tf.cast(rewards, dtype=tf.float32) values = tf.convert_to_tensor(values, dtype=tf.float32) - bootstrap_value = tf.convert_to_tensor( - bootstrap_value, dtype=tf.float32) + bootstrap_value = tf.convert_to_tensor(bootstrap_value, dtype=tf.float32) if clip_rho_threshold is not None: - clip_rho_threshold = tf.convert_to_tensor(clip_rho_threshold, - dtype=tf.float32) + clip_rho_threshold = tf.convert_to_tensor( + clip_rho_threshold, dtype=tf.float32) if clip_pg_rho_threshold is not None: - clip_pg_rho_threshold = tf.convert_to_tensor(clip_pg_rho_threshold, - dtype=tf.float32) + clip_pg_rho_threshold = tf.convert_to_tensor( + clip_pg_rho_threshold, dtype=tf.float32) # Make sure tensor ranks are consistent. rho_rank = rhos.shape.ndims @@ -235,8 +234,9 @@ def from_importance_weights(rhos, discounts, rewards, values, bootstrap_value, if clip_pg_rho_threshold is not None: clip_pg_rho_threshold.shape.assert_has_rank(0) - with tf.name_scope(name, values=[rhos, discounts, rewards, values, - bootstrap_value]): + with tf.name_scope( + name, + values=[rhos, discounts, rewards, values, bootstrap_value]): if clip_rho_threshold is not None: clipped_rhos = tf.minimum( clip_rho_threshold, rhos, name='clipped_rhos') @@ -258,18 +258,17 @@ def from_importance_weights(rhos, discounts, rewards, values, bootstrap_value, # Append bootstrapped value to get [v1, ..., v_t+1] values_t_plus_1 = tf.concat( [values[1:], tf.expand_dims(bootstrap_value, 0)], axis=0) - deltas = clipped_rhos * \ - (rewards + discounts * values_t_plus_1 - values) + deltas = clipped_rhos * ( + rewards + discounts * values_t_plus_1 - values) - # Note that all sequences are reversed, computation starts from the - # back. + # All sequences are reversed, computation starts from the back. sequences = ( tf.reverse(discounts, axis=[0]), tf.reverse(cs, axis=[0]), tf.reverse(deltas, axis=[0]), ) - # V-trace vs are calculated through a scan from the back to + # V-trace vs are calculated through a scan from the back to the # beginning of the given trajectory. def scanfunc(acc, sequence_item): discount_t, c_t, delta_t = sequence_item @@ -291,19 +290,20 @@ def scanfunc(acc, sequence_item): vs = tf.add(vs_minus_v_xs, values, name='vs') # Advantage for policy gradient. - vs_t_plus_1 = tf.concat([ - vs[1:], tf.expand_dims(bootstrap_value, 0)], axis=0) + vs_t_plus_1 = tf.concat( + [vs[1:], tf.expand_dims(bootstrap_value, 0)], axis=0) if clip_pg_rho_threshold is not None: - clipped_pg_rhos = tf.minimum(clip_pg_rho_threshold, rhos, - name='clipped_pg_rhos') + clipped_pg_rhos = tf.minimum( + clip_pg_rho_threshold, rhos, name='clipped_pg_rhos') else: clipped_pg_rhos = rhos pg_advantages = ( clipped_pg_rhos * (rewards + discounts * vs_t_plus_1 - values)) # Make sure no gradients backpropagated through the returned values. - return VTraceReturns(vs=tf.stop_gradient(vs), - pg_advantages=tf.stop_gradient(pg_advantages)) + return VTraceReturns( + vs=tf.stop_gradient(vs), + pg_advantages=tf.stop_gradient(pg_advantages)) def get_rhos(behaviour_action_log_probs, target_action_log_probs): From e48f9ae53b954d5ba5cd640c768e222250442cde Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleksandar=20Milovanovi=C4=87?= Date: Thu, 7 Feb 2019 12:27:38 +0100 Subject: [PATCH 06/33] aligned formatting some more --- python/ray/rllib/agents/impala/vtrace.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/python/ray/rllib/agents/impala/vtrace.py b/python/ray/rllib/agents/impala/vtrace.py index 24cd690b1aef..2e19d6f13d79 100644 --- a/python/ray/rllib/agents/impala/vtrace.py +++ b/python/ray/rllib/agents/impala/vtrace.py @@ -167,8 +167,13 @@ def from_logits(behaviour_policy, **vtrace_returns._asdict()) -def from_importance_weights(rhos, discounts, rewards, values, bootstrap_value, - clip_rho_threshold=1.0, clip_pg_rho_threshold=1.0, +def from_importance_weights(rhos, + discounts, + rewards, + values, + bootstrap_value, + clip_rho_threshold=1.0, + clip_pg_rho_threshold=1.0, name='vtrace_from_importance_weights'): r"""V-trace from log importance weights. @@ -224,7 +229,7 @@ def from_importance_weights(rhos, discounts, rewards, values, bootstrap_value, clip_pg_rho_threshold, dtype=tf.float32) # Make sure tensor ranks are consistent. - rho_rank = rhos.shape.ndims + rho_rank = rhos.shape.ndims # Usually 2. values.shape.assert_has_rank(rho_rank) bootstrap_value.shape.assert_has_rank(rho_rank - 1) discounts.shape.assert_has_rank(rho_rank) @@ -283,8 +288,7 @@ def scanfunc(acc, sequence_item): back_prop=False, name='scan') # Reverse the results back to original order. - vs_minus_v_xs = tf.reverse( - vs_minus_v_xs, [0], name='vs_minus_v_xs') + vs_minus_v_xs = tf.reverse(vs_minus_v_xs, [0], name='vs_minus_v_xs') # Add V(x_s) to get v_s. vs = tf.add(vs_minus_v_xs, values, name='vs') From 3171c8a74a52d69d402acc7e13c58aeb73d381a7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleksandar=20Milovanovi=C4=87?= Date: Fri, 8 Feb 2019 12:08:55 +0100 Subject: [PATCH 07/33] fixed impala stuff --- python/ray/rllib/agents/impala/vtrace.py | 9 ++------- .../rllib/agents/impala/vtrace_policy_graph.py | 16 ++++++++-------- 2 files changed, 10 insertions(+), 15 deletions(-) diff --git a/python/ray/rllib/agents/impala/vtrace.py b/python/ray/rllib/agents/impala/vtrace.py index 2e19d6f13d79..153751cd7d33 100644 --- a/python/ray/rllib/agents/impala/vtrace.py +++ b/python/ray/rllib/agents/impala/vtrace.py @@ -218,7 +218,6 @@ def from_importance_weights(rhos, rhos = tf.convert_to_tensor(rhos, dtype=tf.float32) discounts = tf.convert_to_tensor(discounts, dtype=tf.float32) rewards = tf.convert_to_tensor(rewards, dtype=tf.float32) - rewards = tf.cast(rewards, dtype=tf.float32) values = tf.convert_to_tensor(values, dtype=tf.float32) bootstrap_value = tf.convert_to_tensor(bootstrap_value, dtype=tf.float32) if clip_rho_threshold is not None: @@ -251,12 +250,8 @@ def from_importance_weights(rhos, tf.summary.histogram('clipped_rhos_1000', tf.minimum(1000.0, rhos)) tf.summary.scalar( 'num_of_clipped_rhos', - tf.reduce_sum( - tf.cast( - tf.equal( - clipped_rhos, - clip_rho_threshold), - tf.int32))) + tf.reduce_sum(tf.cast(tf.equal(clipped_rhos, clip_rho_threshold), tf.int32)) + ) tf.summary.scalar('size_of_clipped_rhos', tf.size(clipped_rhos)) cs = tf.minimum(1.0, rhos, name='cs') diff --git a/python/ray/rllib/agents/impala/vtrace_policy_graph.py b/python/ray/rllib/agents/impala/vtrace_policy_graph.py index a176e8288a56..3c8a4d6e2a37 100644 --- a/python/ray/rllib/agents/impala/vtrace_policy_graph.py +++ b/python/ray/rllib/agents/impala/vtrace_policy_graph.py @@ -109,7 +109,7 @@ def __init__(self, "Must use `truncate_episodes` batch mode with V-trace." self.config = config self.sess = tf.get_default_session() - self._is_discrete = False + self._is_multidiscrete = False self.grads = None output_hidden_shape = None @@ -123,10 +123,10 @@ def __init__(self, existing_seq_lens = existing_inputs[-1] else: if isinstance(action_space, gym.spaces.Discrete): - self._is_discrete = True output_hidden_shape = [action_space.n] elif isinstance(action_space, gym.spaces.multi_discrete.MultiDiscrete): + self._is_multidiscrete = True actions_shape = [None, len(action_space.nvec)] output_hidden_shape = action_space.nvec else: @@ -168,8 +168,8 @@ def __init__(self, unpacked_outputs = tf.split( self.model.outputs, output_hidden_shape, axis=1) - dist_inputs = self.model.outputs if self._is_discrete else \ - unpacked_outputs + dist_inputs = unpacked_outputs if self._is_multidiscrete else \ + self.model.outputs action_dist = dist_class(dist_inputs) values = self.model.value_function() @@ -217,10 +217,10 @@ def make_time_major(tensor, drop_last=False): mask = tf.ones_like(rewards, dtype=tf.bool) # Prepare actions for loss - loss_actions = tf.expand_dims( - actions, axis=1) if self._is_discrete else actions - logp_action = actions if self._is_discrete else tf.unstack( - actions, axis=1) + loss_actions = actions if self._is_multidiscrete else tf.expand_dims( + actions, axis=1) + logp_action = tf.unstack( + actions, axis=1) if self._is_multidiscrete else actions # Inputs are reshaped from [B * T] => [T - 1, B] for V-trace calc. self.loss = VTraceLoss( From 9d62dd1ba7ae18d4818221c0fe605ab8b876989f Mon Sep 17 00:00:00 2001 From: pimpke Date: Fri, 8 Feb 2019 16:39:22 +0100 Subject: [PATCH 08/33] Address vtrace comments (#6) * Address vtrace comments * Update get_log_rhos method's comment --- python/ray/rllib/agents/impala/vtrace.py | 135 +++++++++++------- .../agents/impala/vtrace_policy_graph.py | 4 +- 2 files changed, 86 insertions(+), 53 deletions(-) diff --git a/python/ray/rllib/agents/impala/vtrace.py b/python/ray/rllib/agents/impala/vtrace.py index 153751cd7d33..f71c87b8fe3a 100644 --- a/python/ray/rllib/agents/impala/vtrace.py +++ b/python/ray/rllib/agents/impala/vtrace.py @@ -20,6 +20,12 @@ by Espeholt, Soyer, Munos et al. See https://arxiv.org/abs/1802.01561 for the full paper. + +In addition to the original paper's code, changes have been made +to support MultiDiscrete action spaces. behaviour_policy_logits, +target_policy_logits and actions parameters in the entry point +from_logits method are now lists of tensors instead of just +being tensors. """ from __future__ import absolute_import @@ -33,8 +39,8 @@ nest = tf.contrib.framework.nest VTraceFromLogitsReturns = collections.namedtuple('VTraceFromLogitsReturns', [ - 'vs', 'pg_advantages', 'rhos', 'behaviour_action_policy', - 'target_action_policy' + 'vs', 'pg_advantages', 'log_rhos', 'behaviour_action_log_probs', + 'target_action_log_probs' ]) VTraceReturns = collections.namedtuple('VTraceReturns', 'vs pg_advantages') @@ -45,16 +51,30 @@ def select_policy_values_using_actions(policy_logits, actions): In the notation used throughout documentation and comments, T refers to the time dimension ranging from 0 to T-1. B refers to the batch size and - NUM_ACTIONS refers to the number of actions. + ACTION_SPACE refers to the list of numbers each representing a number of actions. Args: - policy_logits: A float32 tensor of shape [T, B, NUM_ACTIONS] with - un-normalized log-probabilities parameterizing a softmax policy. - actions: An int32 tensor of shape [T, B] with actions. + policy_logits: A list with length of ACTION_SPACE of float32 + tensors of shapes + [T, B, ACTION_SPACE[0]], + ..., + [T, B, ACTION_SPACE[-1]] + with un-normalized log-probabilities parameterizing a softmax policy. + actions: A list with length of ACTION_SPACE of int32 + tensors of shapes + [T, B], + ..., + [T, B] + with actions. Returns: - A float32 tensor of shape [T, B] corresponding to the sampling log - probability of the chosen action w.r.t. the policy. + A list with length of ACTION_SPACE of float32 + tensors of shapes + [T, B], + ..., + [T, B] + corresponding to the sampling log probability + of the chosen action w.r.t. the policy. """ log_probs = [] @@ -65,8 +85,8 @@ def select_policy_values_using_actions(policy_logits, actions): return log_probs -def from_logits(behaviour_policy, - target_policy, +def from_logits(behaviour_policy_logits, + target_policy_logits, actions, discounts, rewards, @@ -89,16 +109,27 @@ def from_logits(behaviour_policy, In the notation used throughout documentation and comments, T refers to the time dimension ranging from 0 to T-1. B refers to the batch size and - NUM_ACTIONS refers to the number of actions. + ACTION_SPACE refers to the list of numbers each representing a number of actions. Args: - behaviour_policy: A float32 tensor of shape [T, B, NUM_ACTIONS] with - un-normalized log-probabilities parametrizing the softmax behaviour - policy. - target_policy: A float32 tensor of shape [T, B, NUM_ACTIONS] with - un-normalized log-probabilities parametrizing the softmax target policy. - actions: An int32 tensor of shape [T, B] of actions sampled from the - behaviour policy. + behaviour_policy_logits: A list with length of ACTION_SPACE of float32 + tensors of shapes + [T, B, ACTION_SPACE[0]], + ..., + [T, B, ACTION_SPACE[-1]] + with un-normalized log-probabilities parameterizing the softmax behaviour policy. + target_policy_logits: A list with length of ACTION_SPACE of float32 + tensors of shapes + [T, B, ACTION_SPACE[0]], + ..., + [T, B, ACTION_SPACE[-1]] + with un-normalized log-probabilities parameterizing the softmax target policy. + actions: A list with length of ACTION_SPACE of int32 + tensors of shapes + [T, B], + ..., + [T, B] + with actions sampled from the behaviour policy. discounts: A float32 tensor of shape [T, B] with the discount encountered when following the behaviour policy. rewards: A float32 tensor of shape [T, B] with the rewards generated by @@ -128,31 +159,31 @@ def from_logits(behaviour_policy, target policy action probabilities (log \pi(a_t)). """ - for i in range(len(behaviour_policy)): - behaviour_policy[i] = tf.convert_to_tensor( - behaviour_policy[i], dtype=tf.float32) - target_policy[i] = tf.convert_to_tensor( - target_policy[i], dtype=tf.float32) + for i in range(len(behaviour_policy_logits)): + behaviour_policy_logits[i] = tf.convert_to_tensor( + behaviour_policy_logits[i], dtype=tf.float32) + target_policy_logits[i] = tf.convert_to_tensor( + target_policy_logits[i], dtype=tf.float32) actions[i] = tf.convert_to_tensor(actions[i], dtype=tf.int32) # Make sure tensor ranks are as expected. # The rest will be checked by from_action_log_probs. - behaviour_policy[i].shape.assert_has_rank(3) - target_policy[i].shape.assert_has_rank(3) + behaviour_policy_logits[i].shape.assert_has_rank(3) + target_policy_logits[i].shape.assert_has_rank(3) actions[i].shape.assert_has_rank(2) - with tf.name_scope(name, values=[behaviour_policy, target_policy, actions, + with tf.name_scope(name, values=[behaviour_policy_logits, target_policy_logits, actions, discounts, rewards, values, bootstrap_value]): - target_action_policy = select_policy_values_using_actions( - target_policy, actions) - behaviour_action_policy = select_policy_values_using_actions( - behaviour_policy, actions) + target_action_log_probs = select_policy_values_using_actions( + target_policy_logits, actions) + behaviour_action_log_probs = select_policy_values_using_actions( + behaviour_policy_logits, actions) - rhos = get_rhos(target_action_policy, behaviour_action_policy) + log_rhos = get_log_rhos(target_action_log_probs, behaviour_action_log_probs) vtrace_returns = from_importance_weights( - rhos=rhos, + log_rhos=log_rhos, discounts=discounts, rewards=rewards, values=values, @@ -161,13 +192,13 @@ def from_logits(behaviour_policy, clip_pg_rho_threshold=clip_pg_rho_threshold) return VTraceFromLogitsReturns( - rhos=rhos, - behaviour_action_policy=behaviour_action_policy, - target_action_policy=target_action_policy, + log_rhos=log_rhos, + behaviour_action_log_probs=behaviour_action_log_probs, + target_action_log_probs=target_action_log_probs, **vtrace_returns._asdict()) -def from_importance_weights(rhos, +def from_importance_weights(log_rhos, discounts, rewards, values, @@ -184,14 +215,16 @@ def from_importance_weights(rhos, by Espeholt, Soyer, Munos et al. In the notation used throughout documentation and comments, T refers to the - time dimension ranging from 0 to T-1. B refers to the batch size and - NUM_ACTIONS refers to the number of actions. This code also supports the - case where all tensors have the same number of additional dimensions, e.g., - `rewards` is [T, B, C], `values` is [T, B, C], `bootstrap_value` is [B, C]. + time dimension ranging from 0 to T-1. B refers to the batch size. This code + also supports the case where all tensors have the same number of additional + dimensions, e.g., `rewards` is [T, B, C], `values` is [T, B, C], + `bootstrap_value` is [B, C]. Args: - rhos: A float32 tensor of shape [T, B, NUM_ACTIONS] representing the - importance sampling weights, i.e. target_policy(a) / behaviour_policy(a). + log_rhos: A float32 tensor of shape [T, B] representing the + log importance sampling weights, i.e. + log(target_policy(a) / behaviour_policy(a)). V-trace performs operations + on rhos in log-space for numerical stability. discounts: A float32 tensor of shape [T, B] with discounts encountered when following the behaviour policy. rewards: A float32 tensor of shape [T, B] containing rewards generated by @@ -215,7 +248,7 @@ def from_importance_weights(rhos, pg_advantages: A float32 tensor of shape [T, B]. Can be used as the advantage in the calculation of policy gradients. """ - rhos = tf.convert_to_tensor(rhos, dtype=tf.float32) + log_rhos = tf.convert_to_tensor(log_rhos, dtype=tf.float32) discounts = tf.convert_to_tensor(discounts, dtype=tf.float32) rewards = tf.convert_to_tensor(rewards, dtype=tf.float32) values = tf.convert_to_tensor(values, dtype=tf.float32) @@ -228,7 +261,7 @@ def from_importance_weights(rhos, clip_pg_rho_threshold, dtype=tf.float32) # Make sure tensor ranks are consistent. - rho_rank = rhos.shape.ndims # Usually 2. + rho_rank = log_rhos.shape.ndims # Usually 2. values.shape.assert_has_rank(rho_rank) bootstrap_value.shape.assert_has_rank(rho_rank - 1) discounts.shape.assert_has_rank(rho_rank) @@ -240,7 +273,8 @@ def from_importance_weights(rhos, with tf.name_scope( name, - values=[rhos, discounts, rewards, values, bootstrap_value]): + values=[log_rhos, discounts, rewards, values, bootstrap_value]): + rhos = tf.exp(log_rhos) if clip_rho_threshold is not None: clipped_rhos = tf.minimum( clip_rho_threshold, rhos, name='clipped_rhos') @@ -305,12 +339,11 @@ def scanfunc(acc, sequence_item): pg_advantages=tf.stop_gradient(pg_advantages)) -def get_rhos(behaviour_action_log_probs, target_action_log_probs): - """With the selected policy values (logits or probs) subclasses compute - the rhos for calculating the vtrace.""" - log_rhos = [t - b for t, - b in zip(target_action_log_probs, behaviour_action_log_probs)] +def get_log_rhos(behaviour_action_log_probs, target_action_log_probs): + """With the selected log_probs for multi-discrete actions of behaviour + and target policies we compute the log_rhos for calculating the vtrace.""" + log_rhos = [t - b for t, b in zip(target_action_log_probs, behaviour_action_log_probs)] log_rhos = [tf.convert_to_tensor(l, dtype=tf.float32) for l in log_rhos] log_rhos = tf.reduce_sum(tf.stack(log_rhos), axis=0) - return tf.exp(log_rhos) + return log_rhos diff --git a/python/ray/rllib/agents/impala/vtrace_policy_graph.py b/python/ray/rllib/agents/impala/vtrace_policy_graph.py index 3c8a4d6e2a37..28deffbcb003 100644 --- a/python/ray/rllib/agents/impala/vtrace_policy_graph.py +++ b/python/ray/rllib/agents/impala/vtrace_policy_graph.py @@ -69,8 +69,8 @@ def __init__(self, # Compute vtrace on the CPU for better perf. with tf.device("/cpu:0"): self.vtrace_returns = vtrace.from_logits( - behaviour_policy=behaviour_logits, - target_policy=target_logits, + behaviour_policy_logits=behaviour_logits, + target_policy_logits=target_logits, actions=tf.unstack(tf.cast(actions, tf.int32), axis=2), discounts=tf.to_float(~dones) * discount, rewards=rewards, From 659729598bfca9307db0082be1b8b3e88d5b7ef2 Mon Sep 17 00:00:00 2001 From: Stefan Pantic Date: Mon, 11 Feb 2019 13:52:13 +0100 Subject: [PATCH 09/33] Made APPO work with VTrace --- .../ray/rllib/agents/ppo/appo_policy_graph.py | 904 ++++++++++-------- 1 file changed, 481 insertions(+), 423 deletions(-) diff --git a/python/ray/rllib/agents/ppo/appo_policy_graph.py b/python/ray/rllib/agents/ppo/appo_policy_graph.py index f5533f137c79..452bc99dacc6 100644 --- a/python/ray/rllib/agents/ppo/appo_policy_graph.py +++ b/python/ray/rllib/agents/ppo/appo_policy_graph.py @@ -1,423 +1,481 @@ -"""Adapted from VTracePolicyGraph to use the PPO surrogate loss. - -Keep in sync with changes to VTracePolicyGraph.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import tensorflow as tf -import logging -import gym - -import ray -from ray.rllib.agents.impala import vtrace -from ray.rllib.evaluation.tf_policy_graph import TFPolicyGraph, \ - LearningRateSchedule -from ray.rllib.models.catalog import ModelCatalog -from ray.rllib.utils.error import UnsupportedSpaceException -from ray.rllib.utils.explained_variance import explained_variance -from ray.rllib.models.action_dist import Categorical -from ray.rllib.evaluation.postprocessing import compute_advantages - -logger = logging.getLogger(__name__) - - -class PPOSurrogateLoss(object): - """Loss used when V-trace is disabled. - - Arguments: - prev_actions_logp: A float32 tensor of shape [T, B]. - actions_logp: A float32 tensor of shape [T, B]. - actions_kl: A float32 tensor of shape [T, B]. - actions_entropy: A float32 tensor of shape [T, B]. - values: A float32 tensor of shape [T, B]. - valid_mask: A bool tensor of valid RNN input elements (#2992). - advantages: A float32 tensor of shape [T, B]. - value_targets: A float32 tensor of shape [T, B]. - """ - - def __init__(self, - prev_actions_logp, - actions_logp, - action_kl, - actions_entropy, - values, - valid_mask, - advantages, - value_targets, - vf_loss_coeff=0.5, - entropy_coeff=-0.01, - clip_param=0.3): - - logp_ratio = tf.exp(actions_logp - prev_actions_logp) - - surrogate_loss = tf.minimum( - advantages * logp_ratio, - advantages * tf.clip_by_value(logp_ratio, 1 - clip_param, - 1 + clip_param)) - - self.mean_kl = tf.reduce_mean(action_kl) - self.pi_loss = -tf.reduce_sum(surrogate_loss) - - # The baseline loss - delta = tf.boolean_mask(values - value_targets, valid_mask) - self.value_targets = value_targets - self.vf_loss = 0.5 * tf.reduce_sum(tf.square(delta)) - - # The entropy loss - self.entropy = tf.reduce_sum( - tf.boolean_mask(actions_entropy, valid_mask)) - - # The summed weighted loss - self.total_loss = (self.pi_loss + self.vf_loss * vf_loss_coeff + - self.entropy * entropy_coeff) - - -class VTraceSurrogateLoss(object): - def __init__(self, - actions, - prev_actions_logp, - actions_logp, - action_kl, - actions_entropy, - dones, - behaviour_logits, - target_logits, - discount, - rewards, - values, - bootstrap_value, - valid_mask, - vf_loss_coeff=0.5, - entropy_coeff=-0.01, - clip_rho_threshold=1.0, - clip_pg_rho_threshold=1.0, - clip_param=0.3): - """PPO surrogate loss with vtrace importance weighting. - - VTraceLoss takes tensors of shape [T, B, ...], where `B` is the - batch_size. The reason we need to know `B` is for V-trace to properly - handle episode cut boundaries. - - Arguments: - actions: An int32 tensor of shape [T, B, NUM_ACTIONS]. - prev_actions_logp: A float32 tensor of shape [T, B]. - actions_logp: A float32 tensor of shape [T, B]. - actions_kl: A float32 tensor of shape [T, B]. - actions_entropy: A float32 tensor of shape [T, B]. - dones: A bool tensor of shape [T, B]. - behaviour_logits: A float32 tensor of shape [T, B, NUM_ACTIONS]. - target_logits: A float32 tensor of shape [T, B, NUM_ACTIONS]. - discount: A float32 scalar. - rewards: A float32 tensor of shape [T, B]. - values: A float32 tensor of shape [T, B]. - bootstrap_value: A float32 tensor of shape [B]. - valid_mask: A bool tensor of valid RNN input elements (#2992). - """ - - # Compute vtrace on the CPU for better perf. - with tf.device("/cpu:0"): - self.vtrace_returns = vtrace.from_logits( - behaviour_policy_logits=behaviour_logits, - target_policy_logits=target_logits, - actions=tf.cast(actions, tf.int32), - discounts=tf.to_float(~dones) * discount, - rewards=rewards, - values=values, - bootstrap_value=bootstrap_value, - clip_rho_threshold=tf.cast(clip_rho_threshold, tf.float32), - clip_pg_rho_threshold=tf.cast(clip_pg_rho_threshold, - tf.float32)) - - logp_ratio = tf.exp(actions_logp - prev_actions_logp) - - advantages = self.vtrace_returns.pg_advantages - surrogate_loss = tf.minimum( - advantages * logp_ratio, - advantages * tf.clip_by_value(logp_ratio, 1 - clip_param, - 1 + clip_param)) - - self.mean_kl = tf.reduce_mean(action_kl) - self.pi_loss = -tf.reduce_sum(surrogate_loss) - - # The baseline loss - delta = tf.boolean_mask(values - self.vtrace_returns.vs, valid_mask) - self.value_targets = self.vtrace_returns.vs - self.vf_loss = 0.5 * tf.reduce_sum(tf.square(delta)) - - # The entropy loss - self.entropy = tf.reduce_sum( - tf.boolean_mask(actions_entropy, valid_mask)) - - # The summed weighted loss - self.total_loss = (self.pi_loss + self.vf_loss * vf_loss_coeff + - self.entropy * entropy_coeff) - - -class AsyncPPOPolicyGraph(LearningRateSchedule, TFPolicyGraph): - def __init__(self, - observation_space, - action_space, - config, - existing_inputs=None): - config = dict(ray.rllib.agents.impala.impala.DEFAULT_CONFIG, **config) - assert config["batch_mode"] == "truncate_episodes", \ - "Must use `truncate_episodes` batch mode with V-trace." - self.config = config - self.sess = tf.get_default_session() - - # Policy network model - dist_class, logit_dim = ModelCatalog.get_action_dist( - action_space, self.config["model"]) - - # Create input placeholders - if existing_inputs: - if self.config["vtrace"]: - actions, dones, behaviour_logits, rewards, observations, \ - prev_actions, prev_rewards = existing_inputs[:7] - existing_state_in = existing_inputs[7:-1] - existing_seq_lens = existing_inputs[-1] - else: - actions, dones, behaviour_logits, rewards, observations, \ - prev_actions, prev_rewards, adv_ph, value_targets = \ - existing_inputs[:9] - existing_state_in = existing_inputs[9:-1] - existing_seq_lens = existing_inputs[-1] - else: - actions = ModelCatalog.get_action_placeholder(action_space) - if (not isinstance(action_space, gym.spaces.Discrete) - and self.config["vtrace"]): - raise UnsupportedSpaceException( - "Action space {} is not supported with vtrace.".format( - action_space)) - dones = tf.placeholder(tf.bool, [None], name="dones") - rewards = tf.placeholder(tf.float32, [None], name="rewards") - behaviour_logits = tf.placeholder( - tf.float32, [None, logit_dim], name="behaviour_logits") - observations = tf.placeholder( - tf.float32, [None] + list(observation_space.shape)) - existing_state_in = None - existing_seq_lens = None - if not self.config["vtrace"]: - adv_ph = tf.placeholder( - tf.float32, name="advantages", shape=(None, )) - value_targets = tf.placeholder( - tf.float32, name="value_targets", shape=(None, )) - self.observations = observations - - # Setup the policy - prev_actions = ModelCatalog.get_action_placeholder(action_space) - prev_rewards = tf.placeholder(tf.float32, [None], name="prev_reward") - self.model = ModelCatalog.get_model( - { - "obs": observations, - "prev_actions": prev_actions, - "prev_rewards": prev_rewards, - }, - observation_space, - logit_dim, - self.config["model"], - state_in=existing_state_in, - seq_lens=existing_seq_lens) - - action_dist = dist_class(self.model.outputs) - prev_action_dist = dist_class(behaviour_logits) - - values = self.model.value_function() - self.value_function = values - self.var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, - tf.get_variable_scope().name) - - def to_batches(tensor): - if self.config["model"]["use_lstm"]: - B = tf.shape(self.model.seq_lens)[0] - T = tf.shape(tensor)[0] // B - else: - # Important: chop the tensor into batches at known episode cut - # boundaries. TODO(ekl) this is kind of a hack - T = self.config["sample_batch_size"] - B = tf.shape(tensor)[0] // T - rs = tf.reshape(tensor, - tf.concat([[B, T], tf.shape(tensor)[1:]], axis=0)) - # swap B and T axes - return tf.transpose( - rs, - [1, 0] + list(range(2, 1 + int(tf.shape(tensor).shape[0])))) - - if self.model.state_in: - max_seq_len = tf.reduce_max(self.model.seq_lens) - 1 - mask = tf.sequence_mask(self.model.seq_lens, max_seq_len) - mask = tf.reshape(mask, [-1]) - else: - mask = tf.ones_like(rewards) - - # Inputs are reshaped from [B * T] => [T - 1, B] for V-trace calc. - if self.config["vtrace"]: - logger.info("Using V-Trace surrogate loss (vtrace=True)") - self.loss = VTraceSurrogateLoss( - actions=to_batches(actions)[:-1], - prev_actions_logp=to_batches( - prev_action_dist.logp(actions))[:-1], - actions_logp=to_batches(action_dist.logp(actions))[:-1], - action_kl=prev_action_dist.kl(action_dist), - actions_entropy=to_batches(action_dist.entropy())[:-1], - dones=to_batches(dones)[:-1], - behaviour_logits=to_batches(behaviour_logits)[:-1], - target_logits=to_batches(self.model.outputs)[:-1], - discount=config["gamma"], - rewards=to_batches(rewards)[:-1], - values=to_batches(values)[:-1], - bootstrap_value=to_batches(values)[-1], - valid_mask=to_batches(mask)[:-1], - vf_loss_coeff=self.config["vf_loss_coeff"], - entropy_coeff=self.config["entropy_coeff"], - clip_rho_threshold=self.config["vtrace_clip_rho_threshold"], - clip_pg_rho_threshold=self.config[ - "vtrace_clip_pg_rho_threshold"], - clip_param=self.config["clip_param"]) - else: - logger.info("Using PPO surrogate loss (vtrace=False)") - self.loss = PPOSurrogateLoss( - prev_actions_logp=to_batches(prev_action_dist.logp(actions)), - actions_logp=to_batches(action_dist.logp(actions)), - action_kl=prev_action_dist.kl(action_dist), - actions_entropy=to_batches(action_dist.entropy()), - values=to_batches(values), - valid_mask=to_batches(mask), - advantages=to_batches(adv_ph), - value_targets=to_batches(value_targets), - vf_loss_coeff=self.config["vf_loss_coeff"], - entropy_coeff=self.config["entropy_coeff"], - clip_param=self.config["clip_param"]) - - # KL divergence between worker and learner logits for debugging - model_dist = Categorical(self.model.outputs) - behaviour_dist = Categorical(behaviour_logits) - self.KLs = model_dist.kl(behaviour_dist) - self.mean_KL = tf.reduce_mean(self.KLs) - self.max_KL = tf.reduce_max(self.KLs) - self.median_KL = tf.contrib.distributions.percentile(self.KLs, 50.0) - # Initialize TFPolicyGraph - loss_in = [ - ("actions", actions), - ("dones", dones), - ("behaviour_logits", behaviour_logits), - ("rewards", rewards), - ("obs", observations), - ("prev_actions", prev_actions), - ("prev_rewards", prev_rewards), - ] - if not self.config["vtrace"]: - loss_in.append(("advantages", adv_ph)) - loss_in.append(("value_targets", value_targets)) - LearningRateSchedule.__init__(self, self.config["lr"], - self.config["lr_schedule"]) - TFPolicyGraph.__init__( - self, - observation_space, - action_space, - self.sess, - obs_input=observations, - action_sampler=action_dist.sample(), - loss=self.model.loss() + self.loss.total_loss, - loss_inputs=loss_in, - state_inputs=self.model.state_in, - state_outputs=self.model.state_out, - prev_action_input=prev_actions, - prev_reward_input=prev_rewards, - seq_lens=self.model.seq_lens, - max_seq_len=self.config["model"]["max_seq_len"], - batch_divisibility_req=self.config["sample_batch_size"]) - - self.sess.run(tf.global_variables_initializer()) - - if self.config["vtrace"]: - values_batched = to_batches(values)[:-1] - else: - values_batched = to_batches(values) - self.stats_fetches = { - "stats": { - "model_loss": self.model.loss(), - "cur_lr": tf.cast(self.cur_lr, tf.float64), - "policy_loss": self.loss.pi_loss, - "entropy": self.loss.entropy, - "grad_gnorm": tf.global_norm(self._grads), - "var_gnorm": tf.global_norm(self.var_list), - "vf_loss": self.loss.vf_loss, - "vf_explained_var": explained_variance( - tf.reshape(self.loss.value_targets, [-1]), - tf.reshape(values_batched, [-1])), - "mean_KL": self.mean_KL, - "max_KL": self.max_KL, - "median_KL": self.median_KL, - }, - } - self.stats_fetches["kl"] = self.loss.mean_kl - - def optimizer(self): - if self.config["opt_type"] == "adam": - return tf.train.AdamOptimizer(self.cur_lr) - else: - return tf.train.RMSPropOptimizer(self.cur_lr, self.config["decay"], - self.config["momentum"], - self.config["epsilon"]) - - def gradients(self, optimizer): - grads = tf.gradients(self.loss.total_loss, self.var_list) - self.grads, _ = tf.clip_by_global_norm(grads, self.config["grad_clip"]) - clipped_grads = list(zip(self.grads, self.var_list)) - return clipped_grads - - def extra_compute_action_fetches(self): - out = {"behaviour_logits": self.model.outputs} - if not self.config["vtrace"]: - out["vf_preds"] = self.value_function - return out - - def extra_compute_grad_fetches(self): - return self.stats_fetches - - def value(self, ob, *args): - feed_dict = {self.observations: [ob], self.model.seq_lens: [1]} - assert len(args) == len(self.model.state_in), \ - (args, self.model.state_in) - for k, v in zip(self.model.state_in, args): - feed_dict[k] = v - vf = self.sess.run(self.value_function, feed_dict) - return vf[0] - - def postprocess_trajectory(self, - sample_batch, - other_agent_batches=None, - episode=None): - if not self.config["vtrace"]: - completed = sample_batch["dones"][-1] - if completed: - last_r = 0.0 - else: - next_state = [] - for i in range(len(self.model.state_in)): - next_state.append( - [sample_batch["state_out_{}".format(i)][-1]]) - last_r = self.value(sample_batch["new_obs"][-1], *next_state) - batch = compute_advantages( - sample_batch, - last_r, - self.config["gamma"], - self.config["lambda"], - use_gae=self.config["use_gae"]) - else: - batch = sample_batch - del batch.data["new_obs"] # not used, so save some bandwidth - return batch - - def get_initial_state(self): - return self.model.state_init - - def copy(self, existing_inputs): - return AsyncPPOPolicyGraph( - self.observation_space, - self.action_space, - self.config, - existing_inputs=existing_inputs) +"""Adapted from VTracePolicyGraph to use the PPO surrogate loss. + +Keep in sync with changes to VTracePolicyGraph.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import tensorflow as tf +import logging +import gym + +import ray +from ray.rllib.agents.impala import vtrace +from ray.rllib.evaluation.tf_policy_graph import TFPolicyGraph, \ + LearningRateSchedule +from ray.rllib.models.catalog import ModelCatalog +from ray.rllib.utils.error import UnsupportedSpaceException +from ray.rllib.utils.explained_variance import explained_variance +from ray.rllib.models.action_dist import Categorical +from ray.rllib.evaluation.postprocessing import compute_advantages + +logger = logging.getLogger(__name__) + + +class PPOSurrogateLoss(object): + """Loss used when V-trace is disabled. + + Arguments: + prev_actions_logp: A float32 tensor of shape [T, B]. + actions_logp: A float32 tensor of shape [T, B]. + action_kl: A float32 tensor of shape [T, B]. + actions_entropy: A float32 tensor of shape [T, B]. + values: A float32 tensor of shape [T, B]. + valid_mask: A bool tensor of valid RNN input elements (#2992). + advantages: A float32 tensor of shape [T, B]. + value_targets: A float32 tensor of shape [T, B]. + """ + + def __init__(self, + prev_actions_logp, + actions_logp, + action_kl, + actions_entropy, + values, + valid_mask, + advantages, + value_targets, + vf_loss_coeff=0.5, + entropy_coeff=-0.01, + clip_param=0.3): + + logp_ratio = tf.exp(actions_logp - prev_actions_logp) + + surrogate_loss = tf.minimum( + advantages * logp_ratio, + advantages * tf.clip_by_value(logp_ratio, 1 - clip_param, + 1 + clip_param)) + + self.mean_kl = tf.reduce_mean(action_kl) + self.pi_loss = -tf.reduce_sum(surrogate_loss) + + # The baseline loss + delta = tf.boolean_mask(values - value_targets, valid_mask) + self.value_targets = value_targets + self.vf_loss = 0.5 * tf.reduce_sum(tf.square(delta)) + + # The entropy loss + self.entropy = tf.reduce_sum( + tf.boolean_mask(actions_entropy, valid_mask)) + + # The summed weighted loss + self.total_loss = (self.pi_loss + self.vf_loss * vf_loss_coeff + + self.entropy * entropy_coeff) + + +class VTraceSurrogateLoss(object): + def __init__(self, + actions, + prev_actions_logp, + actions_logp, + action_kl, + actions_entropy, + dones, + behaviour_logits, + target_logits, + discount, + rewards, + values, + bootstrap_value, + valid_mask, + vf_loss_coeff=0.5, + entropy_coeff=-0.01, + clip_rho_threshold=1.0, + clip_pg_rho_threshold=1.0, + clip_param=0.3): + """PPO surrogate loss with vtrace importance weighting. + + VTraceLoss takes tensors of shape [T, B, ...], where `B` is the + batch_size. The reason we need to know `B` is for V-trace to properly + handle episode cut boundaries. + + Arguments: + actions: An int32 tensor of shape [T, B, NUM_ACTIONS]. + prev_actions_logp: A float32 tensor of shape [T, B]. + actions_logp: A float32 tensor of shape [T, B]. + action_kl: A float32 tensor of shape [T, B]. + actions_entropy: A float32 tensor of shape [T, B]. + dones: A bool tensor of shape [T, B]. + behaviour_logits: A float32 tensor of shape [T, B, NUM_ACTIONS]. + target_logits: A float32 tensor of shape [T, B, NUM_ACTIONS]. + discount: A float32 scalar. + rewards: A float32 tensor of shape [T, B]. + values: A float32 tensor of shape [T, B]. + bootstrap_value: A float32 tensor of shape [B]. + valid_mask: A bool tensor of valid RNN input elements (#2992). + """ + + # Compute vtrace on the CPU for better perf. + with tf.device("/cpu:0"): + self.vtrace_returns = vtrace.from_logits( + behaviour_policy_logits=behaviour_logits, + target_policy_logits=target_logits, + actions=tf.unstack(tf.cast(actions, tf.int32), axis=2), + discounts=tf.to_float(~dones) * discount, + rewards=rewards, + values=values, + bootstrap_value=bootstrap_value, + clip_rho_threshold=tf.cast(clip_rho_threshold, tf.float32), + clip_pg_rho_threshold=tf.cast(clip_pg_rho_threshold, + tf.float32)) + + logp_ratio = tf.exp(actions_logp - prev_actions_logp) + + advantages = self.vtrace_returns.pg_advantages + surrogate_loss = tf.minimum( + advantages * logp_ratio, + advantages * tf.clip_by_value(logp_ratio, 1 - clip_param, + 1 + clip_param)) + + self.mean_kl = tf.reduce_mean(action_kl) + self.pi_loss = -tf.reduce_sum(surrogate_loss) + + # The baseline loss + delta = tf.boolean_mask(values - self.vtrace_returns.vs, valid_mask) + self.value_targets = self.vtrace_returns.vs + self.vf_loss = 0.5 * tf.reduce_sum(tf.square(delta)) + + # The entropy loss + self.entropy = tf.reduce_sum( + tf.boolean_mask(actions_entropy, valid_mask)) + + # The summed weighted loss + self.total_loss = (self.pi_loss + self.vf_loss * vf_loss_coeff + + self.entropy * entropy_coeff) + + +class AsyncPPOPolicyGraph(LearningRateSchedule, TFPolicyGraph): + def __init__(self, + observation_space, + action_space, + config, + existing_inputs=None): + config = dict(ray.rllib.agents.impala.impala.DEFAULT_CONFIG, **config) + assert config["batch_mode"] == "truncate_episodes", \ + "Must use `truncate_episodes` batch mode with V-trace." + self.config = config + self.sess = tf.get_default_session() + self.grads = None + + is_discrete = False + output_hidden_shape = None + actions_shape = [None] + + # Policy network model + dist_class, logit_dim = ModelCatalog.get_action_dist( + action_space, self.config["model"]) + + # Create input placeholders + if existing_inputs: + if self.config["vtrace"]: + actions, dones, behaviour_logits, rewards, observations, \ + prev_actions, prev_rewards = existing_inputs[:7] + existing_state_in = existing_inputs[7:-1] + existing_seq_lens = existing_inputs[-1] + else: + actions, dones, behaviour_logits, rewards, observations, \ + prev_actions, prev_rewards, adv_ph, value_targets = \ + existing_inputs[:9] + existing_state_in = existing_inputs[9:-1] + existing_seq_lens = existing_inputs[-1] + else: + if isinstance(action_space, gym.spaces.Discrete): + is_discrete = True + output_hidden_shape = [action_space.n] + elif isinstance(action_space, + gym.spaces.multi_discrete.MultiDiscrete): + actions_shape = [None, len(action_space.nvec)] + output_hidden_shape = action_space.nvec + elif self.config["vtrace"]: + raise UnsupportedSpaceException( + "Action space {} is not supported for IMPALA.".format( + action_space)) + + actions = tf.placeholder(tf.int64, actions_shape, name="ac") + dones = tf.placeholder(tf.bool, [None], name="dones") + rewards = tf.placeholder(tf.float32, [None], name="rewards") + behaviour_logits = tf.placeholder( + tf.float32, [None, logit_dim], name="behaviour_logits") + observations = tf.placeholder( + tf.float32, [None] + list(observation_space.shape)) + existing_state_in = None + existing_seq_lens = None + + if not self.config["vtrace"]: + adv_ph = tf.placeholder( + tf.float32, name="advantages", shape=(None, )) + value_targets = tf.placeholder( + tf.float32, name="value_targets", shape=(None, )) + self.observations = observations + + # Unpack behaviour logits + unpacked_behaviour_logits = tf.split( + behaviour_logits, output_hidden_shape, axis=1) + + # Setup the policy + dist_class, logit_dim = ModelCatalog.get_action_dist( + action_space, self.config["model"], + dist_type=self.config["dist_type"]) + prev_actions = ModelCatalog.get_action_placeholder(action_space) + prev_rewards = tf.placeholder(tf.float32, [None], name="prev_reward") + self.model = ModelCatalog.get_model( + { + "obs": observations, + "prev_actions": prev_actions, + "prev_rewards": prev_rewards, + "is_training": self._get_is_training_placeholder(), + }, + observation_space, + logit_dim, + self.config["model"], + state_in=existing_state_in, + seq_lens=existing_seq_lens) + unpacked_outputs = tf.split( + self.model.outputs, output_hidden_shape, axis=1) + + dist_inputs = self.model.outputs if is_discrete else \ + unpacked_outputs + prev_dist_inputs = behaviour_logits if is_discrete else \ + unpacked_behaviour_logits + + action_dist = dist_class(dist_inputs) + prev_action_dist = dist_class(prev_dist_inputs) + + values = self.model.value_function() + self.value_function = values + self.var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, + tf.get_variable_scope().name) + + def make_time_major(tensor, drop_last=False): + """Swaps batch and trajectory axis. + + Args: + tensor: A tensor or list of tensors to reshape. + drop_last: A bool indicating whether to drop the last + trajectory item. + + Returns: + res: A tensor with swapped axes or a list of tensors with + swapped axes. + """ + if isinstance(tensor, list): + return [make_time_major(t, drop_last) for t in tensor] + + if self.config["model"]["use_lstm"]: + B = tf.shape(self.model.seq_lens)[0] + T = tf.shape(tensor)[0] // B + else: + # Important: chop the tensor into batches at known episode cut + # boundaries. TODO(ekl) this is kind of a hack + T = self.config["sample_batch_size"] + B = tf.shape(tensor)[0] // T + rs = tf.reshape(tensor, + tf.concat([[B, T], tf.shape(tensor)[1:]], axis=0)) + + # swap B and T axes + res = tf.transpose( + rs, + [1, 0] + list(range(2, 1 + int(tf.shape(tensor).shape[0])))) + + if drop_last: + return res[:-1] + return res + + if self.model.state_in: + max_seq_len = tf.reduce_max(self.model.seq_lens) - 1 + mask = tf.sequence_mask(self.model.seq_lens, max_seq_len) + mask = tf.reshape(mask, [-1]) + else: + mask = tf.ones_like(rewards) + + # Inputs are reshaped from [B * T] => [T - 1, B] for V-trace calc. + if self.config["vtrace"]: + logger.info("Using V-Trace surrogate loss (vtrace=True)") + + # Prepare actions for loss + loss_actions = tf.expand_dims( + actions, axis=1) if is_discrete else actions + logp_actions = actions if is_discrete else tf.unstack( + actions, axis=1) + + self.loss = VTraceSurrogateLoss( + actions=make_time_major(loss_actions, drop_last=True), + prev_actions_logp=make_time_major(prev_action_dist.logp( + logp_actions), drop_last=True), + actions_logp=make_time_major(action_dist.logp(logp_actions), + drop_last=True), + action_kl=prev_action_dist.kl(action_dist), + actions_entropy=make_time_major(action_dist.entropy(), + drop_last=True), + dones=make_time_major(dones, drop_last=True), + behaviour_logits=make_time_major(unpacked_behaviour_logits, + drop_last=True), + target_logits=make_time_major(unpacked_outputs, + drop_last=True), + discount=config["gamma"], + rewards=make_time_major(rewards, drop_last=True), + values=make_time_major(values, drop_last=True), + bootstrap_value=make_time_major(values)[-1], + valid_mask=make_time_major(mask, drop_last=True), + vf_loss_coeff=self.config["vf_loss_coeff"], + entropy_coeff=self.config["entropy_coeff"], + clip_rho_threshold=self.config["vtrace_clip_rho_threshold"], + clip_pg_rho_threshold=self.config[ + "vtrace_clip_pg_rho_threshold"], + clip_param=self.config["clip_param"]) + else: + logger.info("Using PPO surrogate loss (vtrace=False)") + self.loss = PPOSurrogateLoss( + prev_actions_logp=make_time_major( + prev_action_dist.logp(actions)), + actions_logp=make_time_major( + action_dist.logp(actions)), + action_kl=prev_action_dist.kl(action_dist), + actions_entropy=make_time_major( + action_dist.entropy()), + values=make_time_major(values), + valid_mask=make_time_major(mask), + advantages=make_time_major(adv_ph), + value_targets=make_time_major(value_targets), + vf_loss_coeff=self.config["vf_loss_coeff"], + entropy_coeff=self.config["entropy_coeff"], + clip_param=self.config["clip_param"]) + + # KL divergence between worker and learner logits for debugging + model_dist = Categorical(self.model.outputs) + behaviour_dist = Categorical(behaviour_logits) + self.KLs = model_dist.kl(behaviour_dist) + self.mean_KL = tf.reduce_mean(self.KLs) + self.max_KL = tf.reduce_max(self.KLs) + self.median_KL = tf.contrib.distributions.percentile(self.KLs, 50.0) + # Initialize TFPolicyGraph + loss_in = [ + ("actions", actions), + ("dones", dones), + ("behaviour_logits", behaviour_logits), + ("rewards", rewards), + ("obs", observations), + ("prev_actions", prev_actions), + ("prev_rewards", prev_rewards), + ] + if not self.config["vtrace"]: + loss_in.append(("advantages", adv_ph)) + loss_in.append(("value_targets", value_targets)) + LearningRateSchedule.__init__(self, self.config["lr"], + self.config["lr_schedule"]) + TFPolicyGraph.__init__( + self, + observation_space, + action_space, + self.sess, + obs_input=observations, + action_sampler=action_dist.sample(), + loss=self.model.loss() + self.loss.total_loss, + loss_inputs=loss_in, + state_inputs=self.model.state_in, + state_outputs=self.model.state_out, + prev_action_input=prev_actions, + prev_reward_input=prev_rewards, + seq_lens=self.model.seq_lens, + max_seq_len=self.config["model"]["max_seq_len"], + batch_divisibility_req=self.config["sample_batch_size"]) + + self.sess.run(tf.global_variables_initializer()) + + if self.config["vtrace"]: + values_batched = make_time_major(values, drop_last=True) + else: + values_batched = make_time_major(values) + self.stats_fetches = {"stats": { + "model_loss": self.model.loss(), + "cur_lr": tf.cast(self.cur_lr, tf.float64), + "policy_loss": self.loss.pi_loss, + "entropy": self.loss.entropy, + "grad_gnorm": tf.global_norm(self._grads), + "var_gnorm": tf.global_norm(self.var_list), + "vf_loss": self.loss.vf_loss, + "vf_explained_var": explained_variance( + tf.reshape(self.loss.value_targets, [-1]), + tf.reshape(values_batched, [-1])), + "mean_KL": self.mean_KL, + "max_KL": self.max_KL, + "median_KL": self.median_KL, + }, "kl": self.loss.mean_kl} + + def optimizer(self): + if self.config["opt_type"] == "adam": + return tf.train.AdamOptimizer(self.cur_lr) + else: + return tf.train.RMSPropOptimizer(self.cur_lr, self.config["decay"], + self.config["momentum"], + self.config["epsilon"]) + + def gradients(self, optimizer): + grads = tf.gradients(self.loss.total_loss, self.var_list) + self.grads, _ = tf.clip_by_global_norm(grads, self.config["grad_clip"]) + clipped_grads = list(zip(self.grads, self.var_list)) + return clipped_grads + + def extra_compute_action_fetches(self): + out = {"behaviour_logits": self.model.outputs} + if not self.config["vtrace"]: + out["vf_preds"] = self.value_function + return out + + def extra_compute_grad_fetches(self): + return self.stats_fetches + + def value(self, ob, *args): + feed_dict = {self.observations: [ob], self.model.seq_lens: [1]} + assert len(args) == len(self.model.state_in), \ + (args, self.model.state_in) + for k, v in zip(self.model.state_in, args): + feed_dict[k] = v + vf = self.sess.run(self.value_function, feed_dict) + return vf[0] + + def postprocess_trajectory(self, + sample_batch, + other_agent_batches=None, + episode=None): + if not self.config["vtrace"]: + completed = sample_batch["dones"][-1] + if completed: + last_r = 0.0 + else: + next_state = [] + for i in range(len(self.model.state_in)): + next_state.append( + [sample_batch["state_out_{}".format(i)][-1]]) + last_r = self.value(sample_batch["new_obs"][-1], *next_state) + batch = compute_advantages( + sample_batch, + last_r, + self.config["gamma"], + self.config["lambda"], + use_gae=self.config["use_gae"]) + else: + batch = sample_batch + del batch.data["new_obs"] # not used, so save some bandwidth + return batch + + def get_initial_state(self): + return self.model.state_init + + def copy(self, existing_inputs): + return AsyncPPOPolicyGraph( + self.observation_space, + self.action_space, + self.config, + existing_inputs=existing_inputs) From 1d3199196e7b80452675ec13cd38fc0e9f43241b Mon Sep 17 00:00:00 2001 From: Stefan Pantic Date: Mon, 11 Feb 2019 13:52:28 +0100 Subject: [PATCH 10/33] Variable is no longer a member --- .../agents/impala/vtrace_policy_graph.py | 20 ++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/python/ray/rllib/agents/impala/vtrace_policy_graph.py b/python/ray/rllib/agents/impala/vtrace_policy_graph.py index 28deffbcb003..0b69ceeaf060 100644 --- a/python/ray/rllib/agents/impala/vtrace_policy_graph.py +++ b/python/ray/rllib/agents/impala/vtrace_policy_graph.py @@ -109,9 +109,9 @@ def __init__(self, "Must use `truncate_episodes` batch mode with V-trace." self.config = config self.sess = tf.get_default_session() - self._is_multidiscrete = False self.grads = None + is_multidiscrete = False output_hidden_shape = None actions_shape = [None] @@ -126,7 +126,7 @@ def __init__(self, output_hidden_shape = [action_space.n] elif isinstance(action_space, gym.spaces.multi_discrete.MultiDiscrete): - self._is_multidiscrete = True + is_multidiscrete = True actions_shape = [None, len(action_space.nvec)] output_hidden_shape = action_space.nvec else: @@ -140,13 +140,15 @@ def __init__(self, behaviour_logits = tf.placeholder(tf.float32, [None, sum(output_hidden_shape)], name="behaviour_logits") - unpacked_behaviour_logits = tf.split( - behaviour_logits, output_hidden_shape, axis=1) observations = tf.placeholder( tf.float32, [None] + list(observation_space.shape)) existing_state_in = None existing_seq_lens = None + # Unpack behaviour logits + unpacked_behaviour_logits = tf.split( + behaviour_logits, output_hidden_shape, axis=1) + # Setup the policy dist_class, logit_dim = ModelCatalog.get_action_dist( action_space, self.config["model"], @@ -168,7 +170,7 @@ def __init__(self, unpacked_outputs = tf.split( self.model.outputs, output_hidden_shape, axis=1) - dist_inputs = unpacked_outputs if self._is_multidiscrete else \ + dist_inputs = unpacked_outputs if is_multidiscrete else \ self.model.outputs action_dist = dist_class(dist_inputs) @@ -217,15 +219,15 @@ def make_time_major(tensor, drop_last=False): mask = tf.ones_like(rewards, dtype=tf.bool) # Prepare actions for loss - loss_actions = actions if self._is_multidiscrete else tf.expand_dims( + loss_actions = actions if is_multidiscrete else tf.expand_dims( actions, axis=1) - logp_action = tf.unstack( - actions, axis=1) if self._is_multidiscrete else actions + logp_actions = tf.unstack( + actions, axis=1) if is_multidiscrete else actions # Inputs are reshaped from [B * T] => [T - 1, B] for V-trace calc. self.loss = VTraceLoss( actions=make_time_major(loss_actions, drop_last=True), - actions_logp=make_time_major(action_dist.logp(logp_action), + actions_logp=make_time_major(action_dist.logp(logp_actions), drop_last=True), actions_entropy=make_time_major(action_dist.entropy(), drop_last=True), From 252f6b3ec6ea88e9f2e464b52bc2891f0e8b5d2e Mon Sep 17 00:00:00 2001 From: Stefan Pantic Date: Mon, 11 Feb 2019 13:57:20 +0100 Subject: [PATCH 11/33] Optimized imports --- python/ray/rllib/agents/impala/vtrace_policy_graph.py | 5 ++--- python/ray/rllib/agents/ppo/appo_policy_graph.py | 8 ++++---- 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/python/ray/rllib/agents/impala/vtrace_policy_graph.py b/python/ray/rllib/agents/impala/vtrace_policy_graph.py index 0b69ceeaf060..5dc53acf7437 100644 --- a/python/ray/rllib/agents/impala/vtrace_policy_graph.py +++ b/python/ray/rllib/agents/impala/vtrace_policy_graph.py @@ -6,19 +6,18 @@ from __future__ import division from __future__ import print_function -import tensorflow as tf import gym - import ray +import tensorflow as tf from ray.rllib.agents.impala import vtrace from ray.rllib.evaluation.policy_graph import PolicyGraph from ray.rllib.evaluation.tf_policy_graph import TFPolicyGraph, \ LearningRateSchedule +from ray.rllib.models.action_dist import MultiCategorical from ray.rllib.models.catalog import ModelCatalog from ray.rllib.utils.annotations import override from ray.rllib.utils.error import UnsupportedSpaceException from ray.rllib.utils.explained_variance import explained_variance -from ray.rllib.models.action_dist import MultiCategorical class VTraceLoss(object): diff --git a/python/ray/rllib/agents/ppo/appo_policy_graph.py b/python/ray/rllib/agents/ppo/appo_policy_graph.py index 452bc99dacc6..85124a0cea4b 100644 --- a/python/ray/rllib/agents/ppo/appo_policy_graph.py +++ b/python/ray/rllib/agents/ppo/appo_policy_graph.py @@ -6,19 +6,19 @@ from __future__ import division from __future__ import print_function -import tensorflow as tf import logging -import gym +import gym import ray +import tensorflow as tf from ray.rllib.agents.impala import vtrace +from ray.rllib.evaluation.postprocessing import compute_advantages from ray.rllib.evaluation.tf_policy_graph import TFPolicyGraph, \ LearningRateSchedule +from ray.rllib.models.action_dist import Categorical from ray.rllib.models.catalog import ModelCatalog from ray.rllib.utils.error import UnsupportedSpaceException from ray.rllib.utils.explained_variance import explained_variance -from ray.rllib.models.action_dist import Categorical -from ray.rllib.evaluation.postprocessing import compute_advantages logger = logging.getLogger(__name__) From 5ef2e30c0d24e6c8782fc7665da6351364d437d4 Mon Sep 17 00:00:00 2001 From: Stefan Pantic Date: Mon, 11 Feb 2019 14:08:26 +0100 Subject: [PATCH 12/33] Changed is_discrete to is_multidiscrete, fixed KL distribution --- .../ray/rllib/agents/ppo/appo_policy_graph.py | 27 +++++++++---------- 1 file changed, 13 insertions(+), 14 deletions(-) diff --git a/python/ray/rllib/agents/ppo/appo_policy_graph.py b/python/ray/rllib/agents/ppo/appo_policy_graph.py index 85124a0cea4b..311c25c172ba 100644 --- a/python/ray/rllib/agents/ppo/appo_policy_graph.py +++ b/python/ray/rllib/agents/ppo/appo_policy_graph.py @@ -15,7 +15,7 @@ from ray.rllib.evaluation.postprocessing import compute_advantages from ray.rllib.evaluation.tf_policy_graph import TFPolicyGraph, \ LearningRateSchedule -from ray.rllib.models.action_dist import Categorical +from ray.rllib.models.action_dist import MultiCategorical from ray.rllib.models.catalog import ModelCatalog from ray.rllib.utils.error import UnsupportedSpaceException from ray.rllib.utils.explained_variance import explained_variance @@ -168,7 +168,7 @@ def __init__(self, self.sess = tf.get_default_session() self.grads = None - is_discrete = False + is_multidiscrete = False output_hidden_shape = None actions_shape = [None] @@ -191,16 +191,15 @@ def __init__(self, existing_seq_lens = existing_inputs[-1] else: if isinstance(action_space, gym.spaces.Discrete): - is_discrete = True output_hidden_shape = [action_space.n] elif isinstance(action_space, gym.spaces.multi_discrete.MultiDiscrete): + is_multidiscrete = True actions_shape = [None, len(action_space.nvec)] output_hidden_shape = action_space.nvec elif self.config["vtrace"]: raise UnsupportedSpaceException( - "Action space {} is not supported for IMPALA.".format( - action_space)) + "Action space {} is not supported for APPO with VTrace.".format(action_space)) actions = tf.placeholder(tf.int64, actions_shape, name="ac") dones = tf.placeholder(tf.bool, [None], name="dones") @@ -244,10 +243,10 @@ def __init__(self, unpacked_outputs = tf.split( self.model.outputs, output_hidden_shape, axis=1) - dist_inputs = self.model.outputs if is_discrete else \ - unpacked_outputs - prev_dist_inputs = behaviour_logits if is_discrete else \ - unpacked_behaviour_logits + dist_inputs = unpacked_outputs if is_multidiscrete else \ + self.model.outputs + prev_dist_inputs = unpacked_behaviour_logits if is_multidiscrete else \ + behaviour_logits action_dist = dist_class(dist_inputs) prev_action_dist = dist_class(prev_dist_inputs) @@ -304,10 +303,10 @@ def make_time_major(tensor, drop_last=False): logger.info("Using V-Trace surrogate loss (vtrace=True)") # Prepare actions for loss - loss_actions = tf.expand_dims( - actions, axis=1) if is_discrete else actions - logp_actions = actions if is_discrete else tf.unstack( + loss_actions = actions if is_multidiscrete else tf.expand_dims( actions, axis=1) + logp_actions = tf.unstack( + actions, axis=1) if is_multidiscrete else actions self.loss = VTraceSurrogateLoss( actions=make_time_major(loss_actions, drop_last=True), @@ -353,8 +352,8 @@ def make_time_major(tensor, drop_last=False): clip_param=self.config["clip_param"]) # KL divergence between worker and learner logits for debugging - model_dist = Categorical(self.model.outputs) - behaviour_dist = Categorical(behaviour_logits) + model_dist = MultiCategorical(unpacked_outputs) + behaviour_dist = MultiCategorical(unpacked_behaviour_logits) self.KLs = model_dist.kl(behaviour_dist) self.mean_KL = tf.reduce_mean(self.KLs) self.max_KL = tf.reduce_max(self.KLs) From cf5c1c5b642a2f8262f9ee51d36616794984e32c Mon Sep 17 00:00:00 2001 From: Stefan Pantic Date: Mon, 11 Feb 2019 14:13:28 +0100 Subject: [PATCH 13/33] Fixed KL divergence --- .../ray/rllib/agents/ppo/appo_policy_graph.py | 28 ++++++++++++++----- 1 file changed, 21 insertions(+), 7 deletions(-) diff --git a/python/ray/rllib/agents/ppo/appo_policy_graph.py b/python/ray/rllib/agents/ppo/appo_policy_graph.py index 311c25c172ba..92b2afac45d9 100644 --- a/python/ray/rllib/agents/ppo/appo_policy_graph.py +++ b/python/ray/rllib/agents/ppo/appo_policy_graph.py @@ -354,10 +354,26 @@ def make_time_major(tensor, drop_last=False): # KL divergence between worker and learner logits for debugging model_dist = MultiCategorical(unpacked_outputs) behaviour_dist = MultiCategorical(unpacked_behaviour_logits) - self.KLs = model_dist.kl(behaviour_dist) - self.mean_KL = tf.reduce_mean(self.KLs) - self.max_KL = tf.reduce_max(self.KLs) - self.median_KL = tf.contrib.distributions.percentile(self.KLs, 50.0) + + kls = model_dist.kl(behaviour_dist) + if len(kls) > 1: + self.KL_stats = {} + + for i, kl in enumerate(kls): + self.KL_stats.update({ + f"mean_KL_{i}": tf.reduce_mean(kl), + f"max_KL_{i}": tf.reduce_max(kl), + f"median_KL_{i}": tf.contrib.distributions.percentile( + kl, 50.0), + }) + else: + self.KL_stats = { + "mean_KL": tf.reduce_mean(kls[0]), + "max_KL": tf.reduce_max(kls[0]), + "median_KL": tf.contrib.distributions.percentile( + kls[0], 50.0), + } + # Initialize TFPolicyGraph loss_in = [ ("actions", actions), @@ -407,9 +423,7 @@ def make_time_major(tensor, drop_last=False): "vf_explained_var": explained_variance( tf.reshape(self.loss.value_targets, [-1]), tf.reshape(values_batched, [-1])), - "mean_KL": self.mean_KL, - "max_KL": self.max_KL, - "median_KL": self.median_KL, + **self.KL_stats, }, "kl": self.loss.mean_kl} def optimizer(self): From 54f4f79e89468f65ea531af7a008762b238f5a0d Mon Sep 17 00:00:00 2001 From: Stefan Pantic Date: Mon, 11 Feb 2019 14:15:40 +0100 Subject: [PATCH 14/33] Removed if statement --- python/ray/rllib/agents/ppo/appo_policy_graph.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/python/ray/rllib/agents/ppo/appo_policy_graph.py b/python/ray/rllib/agents/ppo/appo_policy_graph.py index 92b2afac45d9..0d47946566a6 100644 --- a/python/ray/rllib/agents/ppo/appo_policy_graph.py +++ b/python/ray/rllib/agents/ppo/appo_policy_graph.py @@ -408,10 +408,7 @@ def make_time_major(tensor, drop_last=False): self.sess.run(tf.global_variables_initializer()) - if self.config["vtrace"]: - values_batched = make_time_major(values, drop_last=True) - else: - values_batched = make_time_major(values) + values_batched = make_time_major(values, drop_last=self.config["vtrace"]) self.stats_fetches = {"stats": { "model_loss": self.model.loss(), "cur_lr": tf.cast(self.cur_lr, tf.float64), From 7cb1f97f8d59bdbf4e4168ab5312bbd1badd6fc1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleksandar=20Milovanovi=C4=87?= Date: Thu, 14 Feb 2019 12:00:19 +0100 Subject: [PATCH 15/33] revert appo file --- .../ray/rllib/agents/ppo/appo_policy_graph.py | 194 ++++++------------ 1 file changed, 63 insertions(+), 131 deletions(-) diff --git a/python/ray/rllib/agents/ppo/appo_policy_graph.py b/python/ray/rllib/agents/ppo/appo_policy_graph.py index 0d47946566a6..b902e6a6ac5b 100644 --- a/python/ray/rllib/agents/ppo/appo_policy_graph.py +++ b/python/ray/rllib/agents/ppo/appo_policy_graph.py @@ -6,19 +6,19 @@ from __future__ import division from __future__ import print_function +import tensorflow as tf import logging - import gym + import ray -import tensorflow as tf from ray.rllib.agents.impala import vtrace -from ray.rllib.evaluation.postprocessing import compute_advantages from ray.rllib.evaluation.tf_policy_graph import TFPolicyGraph, \ LearningRateSchedule -from ray.rllib.models.action_dist import MultiCategorical from ray.rllib.models.catalog import ModelCatalog from ray.rllib.utils.error import UnsupportedSpaceException from ray.rllib.utils.explained_variance import explained_variance +from ray.rllib.models.action_dist import Categorical +from ray.rllib.evaluation.postprocessing import compute_advantages logger = logging.getLogger(__name__) @@ -29,7 +29,7 @@ class PPOSurrogateLoss(object): Arguments: prev_actions_logp: A float32 tensor of shape [T, B]. actions_logp: A float32 tensor of shape [T, B]. - action_kl: A float32 tensor of shape [T, B]. + actions_kl: A float32 tensor of shape [T, B]. actions_entropy: A float32 tensor of shape [T, B]. values: A float32 tensor of shape [T, B]. valid_mask: A bool tensor of valid RNN input elements (#2992). @@ -104,7 +104,7 @@ def __init__(self, actions: An int32 tensor of shape [T, B, NUM_ACTIONS]. prev_actions_logp: A float32 tensor of shape [T, B]. actions_logp: A float32 tensor of shape [T, B]. - action_kl: A float32 tensor of shape [T, B]. + actions_kl: A float32 tensor of shape [T, B]. actions_entropy: A float32 tensor of shape [T, B]. dones: A bool tensor of shape [T, B]. behaviour_logits: A float32 tensor of shape [T, B, NUM_ACTIONS]. @@ -121,7 +121,7 @@ def __init__(self, self.vtrace_returns = vtrace.from_logits( behaviour_policy_logits=behaviour_logits, target_policy_logits=target_logits, - actions=tf.unstack(tf.cast(actions, tf.int32), axis=2), + actions=tf.cast(actions, tf.int32), discounts=tf.to_float(~dones) * discount, rewards=rewards, values=values, @@ -166,11 +166,6 @@ def __init__(self, "Must use `truncate_episodes` batch mode with V-trace." self.config = config self.sess = tf.get_default_session() - self.grads = None - - is_multidiscrete = False - output_hidden_shape = None - actions_shape = [None] # Policy network model dist_class, logit_dim = ModelCatalog.get_action_dist( @@ -190,18 +185,12 @@ def __init__(self, existing_state_in = existing_inputs[9:-1] existing_seq_lens = existing_inputs[-1] else: - if isinstance(action_space, gym.spaces.Discrete): - output_hidden_shape = [action_space.n] - elif isinstance(action_space, - gym.spaces.multi_discrete.MultiDiscrete): - is_multidiscrete = True - actions_shape = [None, len(action_space.nvec)] - output_hidden_shape = action_space.nvec - elif self.config["vtrace"]: + actions = ModelCatalog.get_action_placeholder(action_space) + if (not isinstance(action_space, gym.spaces.Discrete) + and self.config["vtrace"]): raise UnsupportedSpaceException( - "Action space {} is not supported for APPO with VTrace.".format(action_space)) - - actions = tf.placeholder(tf.int64, actions_shape, name="ac") + "Action space {} is not supported with vtrace.".format( + action_space)) dones = tf.placeholder(tf.bool, [None], name="dones") rewards = tf.placeholder(tf.float32, [None], name="rewards") behaviour_logits = tf.placeholder( @@ -210,7 +199,6 @@ def __init__(self, tf.float32, [None] + list(observation_space.shape)) existing_state_in = None existing_seq_lens = None - if not self.config["vtrace"]: adv_ph = tf.placeholder( tf.float32, name="advantages", shape=(None, )) @@ -218,14 +206,7 @@ def __init__(self, tf.float32, name="value_targets", shape=(None, )) self.observations = observations - # Unpack behaviour logits - unpacked_behaviour_logits = tf.split( - behaviour_logits, output_hidden_shape, axis=1) - # Setup the policy - dist_class, logit_dim = ModelCatalog.get_action_dist( - action_space, self.config["model"], - dist_type=self.config["dist_type"]) prev_actions = ModelCatalog.get_action_placeholder(action_space) prev_rewards = tf.placeholder(tf.float32, [None], name="prev_reward") self.model = ModelCatalog.get_model( @@ -233,44 +214,22 @@ def __init__(self, "obs": observations, "prev_actions": prev_actions, "prev_rewards": prev_rewards, - "is_training": self._get_is_training_placeholder(), }, observation_space, logit_dim, self.config["model"], state_in=existing_state_in, seq_lens=existing_seq_lens) - unpacked_outputs = tf.split( - self.model.outputs, output_hidden_shape, axis=1) - - dist_inputs = unpacked_outputs if is_multidiscrete else \ - self.model.outputs - prev_dist_inputs = unpacked_behaviour_logits if is_multidiscrete else \ - behaviour_logits - action_dist = dist_class(dist_inputs) - prev_action_dist = dist_class(prev_dist_inputs) + action_dist = dist_class(self.model.outputs) + prev_action_dist = dist_class(behaviour_logits) values = self.model.value_function() self.value_function = values self.var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, tf.get_variable_scope().name) - def make_time_major(tensor, drop_last=False): - """Swaps batch and trajectory axis. - - Args: - tensor: A tensor or list of tensors to reshape. - drop_last: A bool indicating whether to drop the last - trajectory item. - - Returns: - res: A tensor with swapped axes or a list of tensors with - swapped axes. - """ - if isinstance(tensor, list): - return [make_time_major(t, drop_last) for t in tensor] - + def to_batches(tensor): if self.config["model"]["use_lstm"]: B = tf.shape(self.model.seq_lens)[0] T = tf.shape(tensor)[0] // B @@ -281,16 +240,11 @@ def make_time_major(tensor, drop_last=False): B = tf.shape(tensor)[0] // T rs = tf.reshape(tensor, tf.concat([[B, T], tf.shape(tensor)[1:]], axis=0)) - # swap B and T axes - res = tf.transpose( + return tf.transpose( rs, [1, 0] + list(range(2, 1 + int(tf.shape(tensor).shape[0])))) - if drop_last: - return res[:-1] - return res - if self.model.state_in: max_seq_len = tf.reduce_max(self.model.seq_lens) - 1 mask = tf.sequence_mask(self.model.seq_lens, max_seq_len) @@ -301,32 +255,21 @@ def make_time_major(tensor, drop_last=False): # Inputs are reshaped from [B * T] => [T - 1, B] for V-trace calc. if self.config["vtrace"]: logger.info("Using V-Trace surrogate loss (vtrace=True)") - - # Prepare actions for loss - loss_actions = actions if is_multidiscrete else tf.expand_dims( - actions, axis=1) - logp_actions = tf.unstack( - actions, axis=1) if is_multidiscrete else actions - self.loss = VTraceSurrogateLoss( - actions=make_time_major(loss_actions, drop_last=True), - prev_actions_logp=make_time_major(prev_action_dist.logp( - logp_actions), drop_last=True), - actions_logp=make_time_major(action_dist.logp(logp_actions), - drop_last=True), + actions=to_batches(actions)[:-1], + prev_actions_logp=to_batches( + prev_action_dist.logp(actions))[:-1], + actions_logp=to_batches(action_dist.logp(actions))[:-1], action_kl=prev_action_dist.kl(action_dist), - actions_entropy=make_time_major(action_dist.entropy(), - drop_last=True), - dones=make_time_major(dones, drop_last=True), - behaviour_logits=make_time_major(unpacked_behaviour_logits, - drop_last=True), - target_logits=make_time_major(unpacked_outputs, - drop_last=True), + actions_entropy=to_batches(action_dist.entropy())[:-1], + dones=to_batches(dones)[:-1], + behaviour_logits=to_batches(behaviour_logits)[:-1], + target_logits=to_batches(self.model.outputs)[:-1], discount=config["gamma"], - rewards=make_time_major(rewards, drop_last=True), - values=make_time_major(values, drop_last=True), - bootstrap_value=make_time_major(values)[-1], - valid_mask=make_time_major(mask, drop_last=True), + rewards=to_batches(rewards)[:-1], + values=to_batches(values)[:-1], + bootstrap_value=to_batches(values)[-1], + valid_mask=to_batches(mask)[:-1], vf_loss_coeff=self.config["vf_loss_coeff"], entropy_coeff=self.config["entropy_coeff"], clip_rho_threshold=self.config["vtrace_clip_rho_threshold"], @@ -336,44 +279,25 @@ def make_time_major(tensor, drop_last=False): else: logger.info("Using PPO surrogate loss (vtrace=False)") self.loss = PPOSurrogateLoss( - prev_actions_logp=make_time_major( - prev_action_dist.logp(actions)), - actions_logp=make_time_major( - action_dist.logp(actions)), + prev_actions_logp=to_batches(prev_action_dist.logp(actions)), + actions_logp=to_batches(action_dist.logp(actions)), action_kl=prev_action_dist.kl(action_dist), - actions_entropy=make_time_major( - action_dist.entropy()), - values=make_time_major(values), - valid_mask=make_time_major(mask), - advantages=make_time_major(adv_ph), - value_targets=make_time_major(value_targets), + actions_entropy=to_batches(action_dist.entropy()), + values=to_batches(values), + valid_mask=to_batches(mask), + advantages=to_batches(adv_ph), + value_targets=to_batches(value_targets), vf_loss_coeff=self.config["vf_loss_coeff"], entropy_coeff=self.config["entropy_coeff"], clip_param=self.config["clip_param"]) # KL divergence between worker and learner logits for debugging - model_dist = MultiCategorical(unpacked_outputs) - behaviour_dist = MultiCategorical(unpacked_behaviour_logits) - - kls = model_dist.kl(behaviour_dist) - if len(kls) > 1: - self.KL_stats = {} - - for i, kl in enumerate(kls): - self.KL_stats.update({ - f"mean_KL_{i}": tf.reduce_mean(kl), - f"max_KL_{i}": tf.reduce_max(kl), - f"median_KL_{i}": tf.contrib.distributions.percentile( - kl, 50.0), - }) - else: - self.KL_stats = { - "mean_KL": tf.reduce_mean(kls[0]), - "max_KL": tf.reduce_max(kls[0]), - "median_KL": tf.contrib.distributions.percentile( - kls[0], 50.0), - } - + model_dist = Categorical(self.model.outputs) + behaviour_dist = Categorical(behaviour_logits) + self.KLs = model_dist.kl(behaviour_dist) + self.mean_KL = tf.reduce_mean(self.KLs) + self.max_KL = tf.reduce_max(self.KLs) + self.median_KL = tf.contrib.distributions.percentile(self.KLs, 50.0) # Initialize TFPolicyGraph loss_in = [ ("actions", actions), @@ -408,20 +332,28 @@ def make_time_major(tensor, drop_last=False): self.sess.run(tf.global_variables_initializer()) - values_batched = make_time_major(values, drop_last=self.config["vtrace"]) - self.stats_fetches = {"stats": { - "model_loss": self.model.loss(), - "cur_lr": tf.cast(self.cur_lr, tf.float64), - "policy_loss": self.loss.pi_loss, - "entropy": self.loss.entropy, - "grad_gnorm": tf.global_norm(self._grads), - "var_gnorm": tf.global_norm(self.var_list), - "vf_loss": self.loss.vf_loss, - "vf_explained_var": explained_variance( - tf.reshape(self.loss.value_targets, [-1]), - tf.reshape(values_batched, [-1])), - **self.KL_stats, - }, "kl": self.loss.mean_kl} + if self.config["vtrace"]: + values_batched = to_batches(values)[:-1] + else: + values_batched = to_batches(values) + self.stats_fetches = { + "stats": { + "model_loss": self.model.loss(), + "cur_lr": tf.cast(self.cur_lr, tf.float64), + "policy_loss": self.loss.pi_loss, + "entropy": self.loss.entropy, + "grad_gnorm": tf.global_norm(self._grads), + "var_gnorm": tf.global_norm(self.var_list), + "vf_loss": self.loss.vf_loss, + "vf_explained_var": explained_variance( + tf.reshape(self.loss.value_targets, [-1]), + tf.reshape(values_batched, [-1])), + "mean_KL": self.mean_KL, + "max_KL": self.max_KL, + "median_KL": self.median_KL, + }, + } + self.stats_fetches["kl"] = self.loss.mean_kl def optimizer(self): if self.config["opt_type"] == "adam": From 65c82d42e1f1943021229b8f47421b67b350efce Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleksandar=20Milovanovi=C4=87?= Date: Thu, 14 Feb 2019 12:05:16 +0100 Subject: [PATCH 16/33] revered stefans appo changes --- .../ray/rllib/agents/ppo/appo_policy_graph.py | 158 ++++++++++++------ 1 file changed, 108 insertions(+), 50 deletions(-) diff --git a/python/ray/rllib/agents/ppo/appo_policy_graph.py b/python/ray/rllib/agents/ppo/appo_policy_graph.py index b902e6a6ac5b..452bc99dacc6 100644 --- a/python/ray/rllib/agents/ppo/appo_policy_graph.py +++ b/python/ray/rllib/agents/ppo/appo_policy_graph.py @@ -29,7 +29,7 @@ class PPOSurrogateLoss(object): Arguments: prev_actions_logp: A float32 tensor of shape [T, B]. actions_logp: A float32 tensor of shape [T, B]. - actions_kl: A float32 tensor of shape [T, B]. + action_kl: A float32 tensor of shape [T, B]. actions_entropy: A float32 tensor of shape [T, B]. values: A float32 tensor of shape [T, B]. valid_mask: A bool tensor of valid RNN input elements (#2992). @@ -104,7 +104,7 @@ def __init__(self, actions: An int32 tensor of shape [T, B, NUM_ACTIONS]. prev_actions_logp: A float32 tensor of shape [T, B]. actions_logp: A float32 tensor of shape [T, B]. - actions_kl: A float32 tensor of shape [T, B]. + action_kl: A float32 tensor of shape [T, B]. actions_entropy: A float32 tensor of shape [T, B]. dones: A bool tensor of shape [T, B]. behaviour_logits: A float32 tensor of shape [T, B, NUM_ACTIONS]. @@ -121,7 +121,7 @@ def __init__(self, self.vtrace_returns = vtrace.from_logits( behaviour_policy_logits=behaviour_logits, target_policy_logits=target_logits, - actions=tf.cast(actions, tf.int32), + actions=tf.unstack(tf.cast(actions, tf.int32), axis=2), discounts=tf.to_float(~dones) * discount, rewards=rewards, values=values, @@ -166,6 +166,11 @@ def __init__(self, "Must use `truncate_episodes` batch mode with V-trace." self.config = config self.sess = tf.get_default_session() + self.grads = None + + is_discrete = False + output_hidden_shape = None + actions_shape = [None] # Policy network model dist_class, logit_dim = ModelCatalog.get_action_dist( @@ -185,12 +190,19 @@ def __init__(self, existing_state_in = existing_inputs[9:-1] existing_seq_lens = existing_inputs[-1] else: - actions = ModelCatalog.get_action_placeholder(action_space) - if (not isinstance(action_space, gym.spaces.Discrete) - and self.config["vtrace"]): + if isinstance(action_space, gym.spaces.Discrete): + is_discrete = True + output_hidden_shape = [action_space.n] + elif isinstance(action_space, + gym.spaces.multi_discrete.MultiDiscrete): + actions_shape = [None, len(action_space.nvec)] + output_hidden_shape = action_space.nvec + elif self.config["vtrace"]: raise UnsupportedSpaceException( - "Action space {} is not supported with vtrace.".format( + "Action space {} is not supported for IMPALA.".format( action_space)) + + actions = tf.placeholder(tf.int64, actions_shape, name="ac") dones = tf.placeholder(tf.bool, [None], name="dones") rewards = tf.placeholder(tf.float32, [None], name="rewards") behaviour_logits = tf.placeholder( @@ -199,6 +211,7 @@ def __init__(self, tf.float32, [None] + list(observation_space.shape)) existing_state_in = None existing_seq_lens = None + if not self.config["vtrace"]: adv_ph = tf.placeholder( tf.float32, name="advantages", shape=(None, )) @@ -206,7 +219,14 @@ def __init__(self, tf.float32, name="value_targets", shape=(None, )) self.observations = observations + # Unpack behaviour logits + unpacked_behaviour_logits = tf.split( + behaviour_logits, output_hidden_shape, axis=1) + # Setup the policy + dist_class, logit_dim = ModelCatalog.get_action_dist( + action_space, self.config["model"], + dist_type=self.config["dist_type"]) prev_actions = ModelCatalog.get_action_placeholder(action_space) prev_rewards = tf.placeholder(tf.float32, [None], name="prev_reward") self.model = ModelCatalog.get_model( @@ -214,22 +234,44 @@ def __init__(self, "obs": observations, "prev_actions": prev_actions, "prev_rewards": prev_rewards, + "is_training": self._get_is_training_placeholder(), }, observation_space, logit_dim, self.config["model"], state_in=existing_state_in, seq_lens=existing_seq_lens) + unpacked_outputs = tf.split( + self.model.outputs, output_hidden_shape, axis=1) + + dist_inputs = self.model.outputs if is_discrete else \ + unpacked_outputs + prev_dist_inputs = behaviour_logits if is_discrete else \ + unpacked_behaviour_logits - action_dist = dist_class(self.model.outputs) - prev_action_dist = dist_class(behaviour_logits) + action_dist = dist_class(dist_inputs) + prev_action_dist = dist_class(prev_dist_inputs) values = self.model.value_function() self.value_function = values self.var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, tf.get_variable_scope().name) - def to_batches(tensor): + def make_time_major(tensor, drop_last=False): + """Swaps batch and trajectory axis. + + Args: + tensor: A tensor or list of tensors to reshape. + drop_last: A bool indicating whether to drop the last + trajectory item. + + Returns: + res: A tensor with swapped axes or a list of tensors with + swapped axes. + """ + if isinstance(tensor, list): + return [make_time_major(t, drop_last) for t in tensor] + if self.config["model"]["use_lstm"]: B = tf.shape(self.model.seq_lens)[0] T = tf.shape(tensor)[0] // B @@ -240,11 +282,16 @@ def to_batches(tensor): B = tf.shape(tensor)[0] // T rs = tf.reshape(tensor, tf.concat([[B, T], tf.shape(tensor)[1:]], axis=0)) + # swap B and T axes - return tf.transpose( + res = tf.transpose( rs, [1, 0] + list(range(2, 1 + int(tf.shape(tensor).shape[0])))) + if drop_last: + return res[:-1] + return res + if self.model.state_in: max_seq_len = tf.reduce_max(self.model.seq_lens) - 1 mask = tf.sequence_mask(self.model.seq_lens, max_seq_len) @@ -255,21 +302,32 @@ def to_batches(tensor): # Inputs are reshaped from [B * T] => [T - 1, B] for V-trace calc. if self.config["vtrace"]: logger.info("Using V-Trace surrogate loss (vtrace=True)") + + # Prepare actions for loss + loss_actions = tf.expand_dims( + actions, axis=1) if is_discrete else actions + logp_actions = actions if is_discrete else tf.unstack( + actions, axis=1) + self.loss = VTraceSurrogateLoss( - actions=to_batches(actions)[:-1], - prev_actions_logp=to_batches( - prev_action_dist.logp(actions))[:-1], - actions_logp=to_batches(action_dist.logp(actions))[:-1], + actions=make_time_major(loss_actions, drop_last=True), + prev_actions_logp=make_time_major(prev_action_dist.logp( + logp_actions), drop_last=True), + actions_logp=make_time_major(action_dist.logp(logp_actions), + drop_last=True), action_kl=prev_action_dist.kl(action_dist), - actions_entropy=to_batches(action_dist.entropy())[:-1], - dones=to_batches(dones)[:-1], - behaviour_logits=to_batches(behaviour_logits)[:-1], - target_logits=to_batches(self.model.outputs)[:-1], + actions_entropy=make_time_major(action_dist.entropy(), + drop_last=True), + dones=make_time_major(dones, drop_last=True), + behaviour_logits=make_time_major(unpacked_behaviour_logits, + drop_last=True), + target_logits=make_time_major(unpacked_outputs, + drop_last=True), discount=config["gamma"], - rewards=to_batches(rewards)[:-1], - values=to_batches(values)[:-1], - bootstrap_value=to_batches(values)[-1], - valid_mask=to_batches(mask)[:-1], + rewards=make_time_major(rewards, drop_last=True), + values=make_time_major(values, drop_last=True), + bootstrap_value=make_time_major(values)[-1], + valid_mask=make_time_major(mask, drop_last=True), vf_loss_coeff=self.config["vf_loss_coeff"], entropy_coeff=self.config["entropy_coeff"], clip_rho_threshold=self.config["vtrace_clip_rho_threshold"], @@ -279,14 +337,17 @@ def to_batches(tensor): else: logger.info("Using PPO surrogate loss (vtrace=False)") self.loss = PPOSurrogateLoss( - prev_actions_logp=to_batches(prev_action_dist.logp(actions)), - actions_logp=to_batches(action_dist.logp(actions)), + prev_actions_logp=make_time_major( + prev_action_dist.logp(actions)), + actions_logp=make_time_major( + action_dist.logp(actions)), action_kl=prev_action_dist.kl(action_dist), - actions_entropy=to_batches(action_dist.entropy()), - values=to_batches(values), - valid_mask=to_batches(mask), - advantages=to_batches(adv_ph), - value_targets=to_batches(value_targets), + actions_entropy=make_time_major( + action_dist.entropy()), + values=make_time_major(values), + valid_mask=make_time_major(mask), + advantages=make_time_major(adv_ph), + value_targets=make_time_major(value_targets), vf_loss_coeff=self.config["vf_loss_coeff"], entropy_coeff=self.config["entropy_coeff"], clip_param=self.config["clip_param"]) @@ -333,27 +394,24 @@ def to_batches(tensor): self.sess.run(tf.global_variables_initializer()) if self.config["vtrace"]: - values_batched = to_batches(values)[:-1] + values_batched = make_time_major(values, drop_last=True) else: - values_batched = to_batches(values) - self.stats_fetches = { - "stats": { - "model_loss": self.model.loss(), - "cur_lr": tf.cast(self.cur_lr, tf.float64), - "policy_loss": self.loss.pi_loss, - "entropy": self.loss.entropy, - "grad_gnorm": tf.global_norm(self._grads), - "var_gnorm": tf.global_norm(self.var_list), - "vf_loss": self.loss.vf_loss, - "vf_explained_var": explained_variance( - tf.reshape(self.loss.value_targets, [-1]), - tf.reshape(values_batched, [-1])), - "mean_KL": self.mean_KL, - "max_KL": self.max_KL, - "median_KL": self.median_KL, - }, - } - self.stats_fetches["kl"] = self.loss.mean_kl + values_batched = make_time_major(values) + self.stats_fetches = {"stats": { + "model_loss": self.model.loss(), + "cur_lr": tf.cast(self.cur_lr, tf.float64), + "policy_loss": self.loss.pi_loss, + "entropy": self.loss.entropy, + "grad_gnorm": tf.global_norm(self._grads), + "var_gnorm": tf.global_norm(self.var_list), + "vf_loss": self.loss.vf_loss, + "vf_explained_var": explained_variance( + tf.reshape(self.loss.value_targets, [-1]), + tf.reshape(values_batched, [-1])), + "mean_KL": self.mean_KL, + "max_KL": self.max_KL, + "median_KL": self.median_KL, + }, "kl": self.loss.mean_kl} def optimizer(self): if self.config["opt_type"] == "adam": From 946df0177a6f683b4d6cebed943685bd8ca657dc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleksandar=20Milovanovi=C4=87?= Date: Thu, 14 Feb 2019 12:12:09 +0100 Subject: [PATCH 17/33] old appo policy graph --- .../ray/rllib/agents/ppo/appo_policy_graph.py | 164 ++++++------------ 1 file changed, 51 insertions(+), 113 deletions(-) diff --git a/python/ray/rllib/agents/ppo/appo_policy_graph.py b/python/ray/rllib/agents/ppo/appo_policy_graph.py index 452bc99dacc6..dffc3a6ba38e 100644 --- a/python/ray/rllib/agents/ppo/appo_policy_graph.py +++ b/python/ray/rllib/agents/ppo/appo_policy_graph.py @@ -1,5 +1,4 @@ """Adapted from VTracePolicyGraph to use the PPO surrogate loss. - Keep in sync with changes to VTracePolicyGraph.""" from __future__ import absolute_import @@ -25,11 +24,10 @@ class PPOSurrogateLoss(object): """Loss used when V-trace is disabled. - Arguments: prev_actions_logp: A float32 tensor of shape [T, B]. actions_logp: A float32 tensor of shape [T, B]. - action_kl: A float32 tensor of shape [T, B]. + actions_kl: A float32 tensor of shape [T, B]. actions_entropy: A float32 tensor of shape [T, B]. values: A float32 tensor of shape [T, B]. valid_mask: A bool tensor of valid RNN input elements (#2992). @@ -95,16 +93,14 @@ def __init__(self, clip_pg_rho_threshold=1.0, clip_param=0.3): """PPO surrogate loss with vtrace importance weighting. - VTraceLoss takes tensors of shape [T, B, ...], where `B` is the batch_size. The reason we need to know `B` is for V-trace to properly handle episode cut boundaries. - Arguments: actions: An int32 tensor of shape [T, B, NUM_ACTIONS]. prev_actions_logp: A float32 tensor of shape [T, B]. actions_logp: A float32 tensor of shape [T, B]. - action_kl: A float32 tensor of shape [T, B]. + actions_kl: A float32 tensor of shape [T, B]. actions_entropy: A float32 tensor of shape [T, B]. dones: A bool tensor of shape [T, B]. behaviour_logits: A float32 tensor of shape [T, B, NUM_ACTIONS]. @@ -121,7 +117,7 @@ def __init__(self, self.vtrace_returns = vtrace.from_logits( behaviour_policy_logits=behaviour_logits, target_policy_logits=target_logits, - actions=tf.unstack(tf.cast(actions, tf.int32), axis=2), + actions=tf.cast(actions, tf.int32), discounts=tf.to_float(~dones) * discount, rewards=rewards, values=values, @@ -166,11 +162,6 @@ def __init__(self, "Must use `truncate_episodes` batch mode with V-trace." self.config = config self.sess = tf.get_default_session() - self.grads = None - - is_discrete = False - output_hidden_shape = None - actions_shape = [None] # Policy network model dist_class, logit_dim = ModelCatalog.get_action_dist( @@ -190,19 +181,12 @@ def __init__(self, existing_state_in = existing_inputs[9:-1] existing_seq_lens = existing_inputs[-1] else: - if isinstance(action_space, gym.spaces.Discrete): - is_discrete = True - output_hidden_shape = [action_space.n] - elif isinstance(action_space, - gym.spaces.multi_discrete.MultiDiscrete): - actions_shape = [None, len(action_space.nvec)] - output_hidden_shape = action_space.nvec - elif self.config["vtrace"]: + actions = ModelCatalog.get_action_placeholder(action_space) + if (not isinstance(action_space, gym.spaces.Discrete) + and self.config["vtrace"]): raise UnsupportedSpaceException( - "Action space {} is not supported for IMPALA.".format( + "Action space {} is not supported with vtrace.".format( action_space)) - - actions = tf.placeholder(tf.int64, actions_shape, name="ac") dones = tf.placeholder(tf.bool, [None], name="dones") rewards = tf.placeholder(tf.float32, [None], name="rewards") behaviour_logits = tf.placeholder( @@ -211,7 +195,6 @@ def __init__(self, tf.float32, [None] + list(observation_space.shape)) existing_state_in = None existing_seq_lens = None - if not self.config["vtrace"]: adv_ph = tf.placeholder( tf.float32, name="advantages", shape=(None, )) @@ -219,14 +202,7 @@ def __init__(self, tf.float32, name="value_targets", shape=(None, )) self.observations = observations - # Unpack behaviour logits - unpacked_behaviour_logits = tf.split( - behaviour_logits, output_hidden_shape, axis=1) - # Setup the policy - dist_class, logit_dim = ModelCatalog.get_action_dist( - action_space, self.config["model"], - dist_type=self.config["dist_type"]) prev_actions = ModelCatalog.get_action_placeholder(action_space) prev_rewards = tf.placeholder(tf.float32, [None], name="prev_reward") self.model = ModelCatalog.get_model( @@ -234,44 +210,22 @@ def __init__(self, "obs": observations, "prev_actions": prev_actions, "prev_rewards": prev_rewards, - "is_training": self._get_is_training_placeholder(), }, observation_space, logit_dim, self.config["model"], state_in=existing_state_in, seq_lens=existing_seq_lens) - unpacked_outputs = tf.split( - self.model.outputs, output_hidden_shape, axis=1) - dist_inputs = self.model.outputs if is_discrete else \ - unpacked_outputs - prev_dist_inputs = behaviour_logits if is_discrete else \ - unpacked_behaviour_logits - - action_dist = dist_class(dist_inputs) - prev_action_dist = dist_class(prev_dist_inputs) + action_dist = dist_class(self.model.outputs) + prev_action_dist = dist_class(behaviour_logits) values = self.model.value_function() self.value_function = values self.var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, tf.get_variable_scope().name) - def make_time_major(tensor, drop_last=False): - """Swaps batch and trajectory axis. - - Args: - tensor: A tensor or list of tensors to reshape. - drop_last: A bool indicating whether to drop the last - trajectory item. - - Returns: - res: A tensor with swapped axes or a list of tensors with - swapped axes. - """ - if isinstance(tensor, list): - return [make_time_major(t, drop_last) for t in tensor] - + def to_batches(tensor): if self.config["model"]["use_lstm"]: B = tf.shape(self.model.seq_lens)[0] T = tf.shape(tensor)[0] // B @@ -282,16 +236,11 @@ def make_time_major(tensor, drop_last=False): B = tf.shape(tensor)[0] // T rs = tf.reshape(tensor, tf.concat([[B, T], tf.shape(tensor)[1:]], axis=0)) - # swap B and T axes - res = tf.transpose( + return tf.transpose( rs, [1, 0] + list(range(2, 1 + int(tf.shape(tensor).shape[0])))) - if drop_last: - return res[:-1] - return res - if self.model.state_in: max_seq_len = tf.reduce_max(self.model.seq_lens) - 1 mask = tf.sequence_mask(self.model.seq_lens, max_seq_len) @@ -302,32 +251,21 @@ def make_time_major(tensor, drop_last=False): # Inputs are reshaped from [B * T] => [T - 1, B] for V-trace calc. if self.config["vtrace"]: logger.info("Using V-Trace surrogate loss (vtrace=True)") - - # Prepare actions for loss - loss_actions = tf.expand_dims( - actions, axis=1) if is_discrete else actions - logp_actions = actions if is_discrete else tf.unstack( - actions, axis=1) - self.loss = VTraceSurrogateLoss( - actions=make_time_major(loss_actions, drop_last=True), - prev_actions_logp=make_time_major(prev_action_dist.logp( - logp_actions), drop_last=True), - actions_logp=make_time_major(action_dist.logp(logp_actions), - drop_last=True), + actions=to_batches(actions)[:-1], + prev_actions_logp=to_batches( + prev_action_dist.logp(actions))[:-1], + actions_logp=to_batches(action_dist.logp(actions))[:-1], action_kl=prev_action_dist.kl(action_dist), - actions_entropy=make_time_major(action_dist.entropy(), - drop_last=True), - dones=make_time_major(dones, drop_last=True), - behaviour_logits=make_time_major(unpacked_behaviour_logits, - drop_last=True), - target_logits=make_time_major(unpacked_outputs, - drop_last=True), + actions_entropy=to_batches(action_dist.entropy())[:-1], + dones=to_batches(dones)[:-1], + behaviour_logits=to_batches(behaviour_logits)[:-1], + target_logits=to_batches(self.model.outputs)[:-1], discount=config["gamma"], - rewards=make_time_major(rewards, drop_last=True), - values=make_time_major(values, drop_last=True), - bootstrap_value=make_time_major(values)[-1], - valid_mask=make_time_major(mask, drop_last=True), + rewards=to_batches(rewards)[:-1], + values=to_batches(values)[:-1], + bootstrap_value=to_batches(values)[-1], + valid_mask=to_batches(mask)[:-1], vf_loss_coeff=self.config["vf_loss_coeff"], entropy_coeff=self.config["entropy_coeff"], clip_rho_threshold=self.config["vtrace_clip_rho_threshold"], @@ -337,17 +275,14 @@ def make_time_major(tensor, drop_last=False): else: logger.info("Using PPO surrogate loss (vtrace=False)") self.loss = PPOSurrogateLoss( - prev_actions_logp=make_time_major( - prev_action_dist.logp(actions)), - actions_logp=make_time_major( - action_dist.logp(actions)), + prev_actions_logp=to_batches(prev_action_dist.logp(actions)), + actions_logp=to_batches(action_dist.logp(actions)), action_kl=prev_action_dist.kl(action_dist), - actions_entropy=make_time_major( - action_dist.entropy()), - values=make_time_major(values), - valid_mask=make_time_major(mask), - advantages=make_time_major(adv_ph), - value_targets=make_time_major(value_targets), + actions_entropy=to_batches(action_dist.entropy()), + values=to_batches(values), + valid_mask=to_batches(mask), + advantages=to_batches(adv_ph), + value_targets=to_batches(value_targets), vf_loss_coeff=self.config["vf_loss_coeff"], entropy_coeff=self.config["entropy_coeff"], clip_param=self.config["clip_param"]) @@ -394,24 +329,27 @@ def make_time_major(tensor, drop_last=False): self.sess.run(tf.global_variables_initializer()) if self.config["vtrace"]: - values_batched = make_time_major(values, drop_last=True) + values_batched = to_batches(values)[:-1] else: - values_batched = make_time_major(values) - self.stats_fetches = {"stats": { - "model_loss": self.model.loss(), - "cur_lr": tf.cast(self.cur_lr, tf.float64), - "policy_loss": self.loss.pi_loss, - "entropy": self.loss.entropy, - "grad_gnorm": tf.global_norm(self._grads), - "var_gnorm": tf.global_norm(self.var_list), - "vf_loss": self.loss.vf_loss, - "vf_explained_var": explained_variance( - tf.reshape(self.loss.value_targets, [-1]), - tf.reshape(values_batched, [-1])), - "mean_KL": self.mean_KL, - "max_KL": self.max_KL, - "median_KL": self.median_KL, - }, "kl": self.loss.mean_kl} + values_batched = to_batches(values) + self.stats_fetches = { + "stats": { + "model_loss": self.model.loss(), + "cur_lr": tf.cast(self.cur_lr, tf.float64), + "policy_loss": self.loss.pi_loss, + "entropy": self.loss.entropy, + "grad_gnorm": tf.global_norm(self._grads), + "var_gnorm": tf.global_norm(self.var_list), + "vf_loss": self.loss.vf_loss, + "vf_explained_var": explained_variance( + tf.reshape(self.loss.value_targets, [-1]), + tf.reshape(values_batched, [-1])), + "mean_KL": self.mean_KL, + "max_KL": self.max_KL, + "median_KL": self.median_KL, + }, + } + self.stats_fetches["kl"] = self.loss.mean_kl def optimizer(self): if self.config["opt_type"] == "adam": @@ -478,4 +416,4 @@ def copy(self, existing_inputs): self.observation_space, self.action_space, self.config, - existing_inputs=existing_inputs) + existing_inputs=existing_inputs) \ No newline at end of file From b6b2c52a9bb24ee3f1d1ba0e1e104ec49cab4000 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleksandar=20Milovanovi=C4=87?= Date: Thu, 14 Feb 2019 12:19:49 +0100 Subject: [PATCH 18/33] returned stefan appo changes and returned newline --- .../ray/rllib/agents/ppo/appo_policy_graph.py | 892 ++++++++++-------- 1 file changed, 474 insertions(+), 418 deletions(-) diff --git a/python/ray/rllib/agents/ppo/appo_policy_graph.py b/python/ray/rllib/agents/ppo/appo_policy_graph.py index dffc3a6ba38e..398129ead93d 100644 --- a/python/ray/rllib/agents/ppo/appo_policy_graph.py +++ b/python/ray/rllib/agents/ppo/appo_policy_graph.py @@ -1,419 +1,475 @@ -"""Adapted from VTracePolicyGraph to use the PPO surrogate loss. -Keep in sync with changes to VTracePolicyGraph.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import tensorflow as tf -import logging -import gym - -import ray -from ray.rllib.agents.impala import vtrace -from ray.rllib.evaluation.tf_policy_graph import TFPolicyGraph, \ - LearningRateSchedule -from ray.rllib.models.catalog import ModelCatalog -from ray.rllib.utils.error import UnsupportedSpaceException -from ray.rllib.utils.explained_variance import explained_variance -from ray.rllib.models.action_dist import Categorical -from ray.rllib.evaluation.postprocessing import compute_advantages - -logger = logging.getLogger(__name__) - - -class PPOSurrogateLoss(object): - """Loss used when V-trace is disabled. - Arguments: - prev_actions_logp: A float32 tensor of shape [T, B]. - actions_logp: A float32 tensor of shape [T, B]. - actions_kl: A float32 tensor of shape [T, B]. - actions_entropy: A float32 tensor of shape [T, B]. - values: A float32 tensor of shape [T, B]. - valid_mask: A bool tensor of valid RNN input elements (#2992). - advantages: A float32 tensor of shape [T, B]. - value_targets: A float32 tensor of shape [T, B]. - """ - - def __init__(self, - prev_actions_logp, - actions_logp, - action_kl, - actions_entropy, - values, - valid_mask, - advantages, - value_targets, - vf_loss_coeff=0.5, - entropy_coeff=-0.01, - clip_param=0.3): - - logp_ratio = tf.exp(actions_logp - prev_actions_logp) - - surrogate_loss = tf.minimum( - advantages * logp_ratio, - advantages * tf.clip_by_value(logp_ratio, 1 - clip_param, - 1 + clip_param)) - - self.mean_kl = tf.reduce_mean(action_kl) - self.pi_loss = -tf.reduce_sum(surrogate_loss) - - # The baseline loss - delta = tf.boolean_mask(values - value_targets, valid_mask) - self.value_targets = value_targets - self.vf_loss = 0.5 * tf.reduce_sum(tf.square(delta)) - - # The entropy loss - self.entropy = tf.reduce_sum( - tf.boolean_mask(actions_entropy, valid_mask)) - - # The summed weighted loss - self.total_loss = (self.pi_loss + self.vf_loss * vf_loss_coeff + - self.entropy * entropy_coeff) - - -class VTraceSurrogateLoss(object): - def __init__(self, - actions, - prev_actions_logp, - actions_logp, - action_kl, - actions_entropy, - dones, - behaviour_logits, - target_logits, - discount, - rewards, - values, - bootstrap_value, - valid_mask, - vf_loss_coeff=0.5, - entropy_coeff=-0.01, - clip_rho_threshold=1.0, - clip_pg_rho_threshold=1.0, - clip_param=0.3): - """PPO surrogate loss with vtrace importance weighting. - VTraceLoss takes tensors of shape [T, B, ...], where `B` is the - batch_size. The reason we need to know `B` is for V-trace to properly - handle episode cut boundaries. - Arguments: - actions: An int32 tensor of shape [T, B, NUM_ACTIONS]. - prev_actions_logp: A float32 tensor of shape [T, B]. - actions_logp: A float32 tensor of shape [T, B]. - actions_kl: A float32 tensor of shape [T, B]. - actions_entropy: A float32 tensor of shape [T, B]. - dones: A bool tensor of shape [T, B]. - behaviour_logits: A float32 tensor of shape [T, B, NUM_ACTIONS]. - target_logits: A float32 tensor of shape [T, B, NUM_ACTIONS]. - discount: A float32 scalar. - rewards: A float32 tensor of shape [T, B]. - values: A float32 tensor of shape [T, B]. - bootstrap_value: A float32 tensor of shape [B]. - valid_mask: A bool tensor of valid RNN input elements (#2992). - """ - - # Compute vtrace on the CPU for better perf. - with tf.device("/cpu:0"): - self.vtrace_returns = vtrace.from_logits( - behaviour_policy_logits=behaviour_logits, - target_policy_logits=target_logits, - actions=tf.cast(actions, tf.int32), - discounts=tf.to_float(~dones) * discount, - rewards=rewards, - values=values, - bootstrap_value=bootstrap_value, - clip_rho_threshold=tf.cast(clip_rho_threshold, tf.float32), - clip_pg_rho_threshold=tf.cast(clip_pg_rho_threshold, - tf.float32)) - - logp_ratio = tf.exp(actions_logp - prev_actions_logp) - - advantages = self.vtrace_returns.pg_advantages - surrogate_loss = tf.minimum( - advantages * logp_ratio, - advantages * tf.clip_by_value(logp_ratio, 1 - clip_param, - 1 + clip_param)) - - self.mean_kl = tf.reduce_mean(action_kl) - self.pi_loss = -tf.reduce_sum(surrogate_loss) - - # The baseline loss - delta = tf.boolean_mask(values - self.vtrace_returns.vs, valid_mask) - self.value_targets = self.vtrace_returns.vs - self.vf_loss = 0.5 * tf.reduce_sum(tf.square(delta)) - - # The entropy loss - self.entropy = tf.reduce_sum( - tf.boolean_mask(actions_entropy, valid_mask)) - - # The summed weighted loss - self.total_loss = (self.pi_loss + self.vf_loss * vf_loss_coeff + - self.entropy * entropy_coeff) - - -class AsyncPPOPolicyGraph(LearningRateSchedule, TFPolicyGraph): - def __init__(self, - observation_space, - action_space, - config, - existing_inputs=None): - config = dict(ray.rllib.agents.impala.impala.DEFAULT_CONFIG, **config) - assert config["batch_mode"] == "truncate_episodes", \ - "Must use `truncate_episodes` batch mode with V-trace." - self.config = config - self.sess = tf.get_default_session() - - # Policy network model - dist_class, logit_dim = ModelCatalog.get_action_dist( - action_space, self.config["model"]) - - # Create input placeholders - if existing_inputs: - if self.config["vtrace"]: - actions, dones, behaviour_logits, rewards, observations, \ - prev_actions, prev_rewards = existing_inputs[:7] - existing_state_in = existing_inputs[7:-1] - existing_seq_lens = existing_inputs[-1] - else: - actions, dones, behaviour_logits, rewards, observations, \ - prev_actions, prev_rewards, adv_ph, value_targets = \ - existing_inputs[:9] - existing_state_in = existing_inputs[9:-1] - existing_seq_lens = existing_inputs[-1] - else: - actions = ModelCatalog.get_action_placeholder(action_space) - if (not isinstance(action_space, gym.spaces.Discrete) - and self.config["vtrace"]): - raise UnsupportedSpaceException( - "Action space {} is not supported with vtrace.".format( - action_space)) - dones = tf.placeholder(tf.bool, [None], name="dones") - rewards = tf.placeholder(tf.float32, [None], name="rewards") - behaviour_logits = tf.placeholder( - tf.float32, [None, logit_dim], name="behaviour_logits") - observations = tf.placeholder( - tf.float32, [None] + list(observation_space.shape)) - existing_state_in = None - existing_seq_lens = None - if not self.config["vtrace"]: - adv_ph = tf.placeholder( - tf.float32, name="advantages", shape=(None, )) - value_targets = tf.placeholder( - tf.float32, name="value_targets", shape=(None, )) - self.observations = observations - - # Setup the policy - prev_actions = ModelCatalog.get_action_placeholder(action_space) - prev_rewards = tf.placeholder(tf.float32, [None], name="prev_reward") - self.model = ModelCatalog.get_model( - { - "obs": observations, - "prev_actions": prev_actions, - "prev_rewards": prev_rewards, - }, - observation_space, - logit_dim, - self.config["model"], - state_in=existing_state_in, - seq_lens=existing_seq_lens) - - action_dist = dist_class(self.model.outputs) - prev_action_dist = dist_class(behaviour_logits) - - values = self.model.value_function() - self.value_function = values - self.var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, - tf.get_variable_scope().name) - - def to_batches(tensor): - if self.config["model"]["use_lstm"]: - B = tf.shape(self.model.seq_lens)[0] - T = tf.shape(tensor)[0] // B - else: - # Important: chop the tensor into batches at known episode cut - # boundaries. TODO(ekl) this is kind of a hack - T = self.config["sample_batch_size"] - B = tf.shape(tensor)[0] // T - rs = tf.reshape(tensor, - tf.concat([[B, T], tf.shape(tensor)[1:]], axis=0)) - # swap B and T axes - return tf.transpose( - rs, - [1, 0] + list(range(2, 1 + int(tf.shape(tensor).shape[0])))) - - if self.model.state_in: - max_seq_len = tf.reduce_max(self.model.seq_lens) - 1 - mask = tf.sequence_mask(self.model.seq_lens, max_seq_len) - mask = tf.reshape(mask, [-1]) - else: - mask = tf.ones_like(rewards) - - # Inputs are reshaped from [B * T] => [T - 1, B] for V-trace calc. - if self.config["vtrace"]: - logger.info("Using V-Trace surrogate loss (vtrace=True)") - self.loss = VTraceSurrogateLoss( - actions=to_batches(actions)[:-1], - prev_actions_logp=to_batches( - prev_action_dist.logp(actions))[:-1], - actions_logp=to_batches(action_dist.logp(actions))[:-1], - action_kl=prev_action_dist.kl(action_dist), - actions_entropy=to_batches(action_dist.entropy())[:-1], - dones=to_batches(dones)[:-1], - behaviour_logits=to_batches(behaviour_logits)[:-1], - target_logits=to_batches(self.model.outputs)[:-1], - discount=config["gamma"], - rewards=to_batches(rewards)[:-1], - values=to_batches(values)[:-1], - bootstrap_value=to_batches(values)[-1], - valid_mask=to_batches(mask)[:-1], - vf_loss_coeff=self.config["vf_loss_coeff"], - entropy_coeff=self.config["entropy_coeff"], - clip_rho_threshold=self.config["vtrace_clip_rho_threshold"], - clip_pg_rho_threshold=self.config[ - "vtrace_clip_pg_rho_threshold"], - clip_param=self.config["clip_param"]) - else: - logger.info("Using PPO surrogate loss (vtrace=False)") - self.loss = PPOSurrogateLoss( - prev_actions_logp=to_batches(prev_action_dist.logp(actions)), - actions_logp=to_batches(action_dist.logp(actions)), - action_kl=prev_action_dist.kl(action_dist), - actions_entropy=to_batches(action_dist.entropy()), - values=to_batches(values), - valid_mask=to_batches(mask), - advantages=to_batches(adv_ph), - value_targets=to_batches(value_targets), - vf_loss_coeff=self.config["vf_loss_coeff"], - entropy_coeff=self.config["entropy_coeff"], - clip_param=self.config["clip_param"]) - - # KL divergence between worker and learner logits for debugging - model_dist = Categorical(self.model.outputs) - behaviour_dist = Categorical(behaviour_logits) - self.KLs = model_dist.kl(behaviour_dist) - self.mean_KL = tf.reduce_mean(self.KLs) - self.max_KL = tf.reduce_max(self.KLs) - self.median_KL = tf.contrib.distributions.percentile(self.KLs, 50.0) - # Initialize TFPolicyGraph - loss_in = [ - ("actions", actions), - ("dones", dones), - ("behaviour_logits", behaviour_logits), - ("rewards", rewards), - ("obs", observations), - ("prev_actions", prev_actions), - ("prev_rewards", prev_rewards), - ] - if not self.config["vtrace"]: - loss_in.append(("advantages", adv_ph)) - loss_in.append(("value_targets", value_targets)) - LearningRateSchedule.__init__(self, self.config["lr"], - self.config["lr_schedule"]) - TFPolicyGraph.__init__( - self, - observation_space, - action_space, - self.sess, - obs_input=observations, - action_sampler=action_dist.sample(), - loss=self.model.loss() + self.loss.total_loss, - loss_inputs=loss_in, - state_inputs=self.model.state_in, - state_outputs=self.model.state_out, - prev_action_input=prev_actions, - prev_reward_input=prev_rewards, - seq_lens=self.model.seq_lens, - max_seq_len=self.config["model"]["max_seq_len"], - batch_divisibility_req=self.config["sample_batch_size"]) - - self.sess.run(tf.global_variables_initializer()) - - if self.config["vtrace"]: - values_batched = to_batches(values)[:-1] - else: - values_batched = to_batches(values) - self.stats_fetches = { - "stats": { - "model_loss": self.model.loss(), - "cur_lr": tf.cast(self.cur_lr, tf.float64), - "policy_loss": self.loss.pi_loss, - "entropy": self.loss.entropy, - "grad_gnorm": tf.global_norm(self._grads), - "var_gnorm": tf.global_norm(self.var_list), - "vf_loss": self.loss.vf_loss, - "vf_explained_var": explained_variance( - tf.reshape(self.loss.value_targets, [-1]), - tf.reshape(values_batched, [-1])), - "mean_KL": self.mean_KL, - "max_KL": self.max_KL, - "median_KL": self.median_KL, - }, - } - self.stats_fetches["kl"] = self.loss.mean_kl - - def optimizer(self): - if self.config["opt_type"] == "adam": - return tf.train.AdamOptimizer(self.cur_lr) - else: - return tf.train.RMSPropOptimizer(self.cur_lr, self.config["decay"], - self.config["momentum"], - self.config["epsilon"]) - - def gradients(self, optimizer): - grads = tf.gradients(self.loss.total_loss, self.var_list) - self.grads, _ = tf.clip_by_global_norm(grads, self.config["grad_clip"]) - clipped_grads = list(zip(self.grads, self.var_list)) - return clipped_grads - - def extra_compute_action_fetches(self): - out = {"behaviour_logits": self.model.outputs} - if not self.config["vtrace"]: - out["vf_preds"] = self.value_function - return out - - def extra_compute_grad_fetches(self): - return self.stats_fetches - - def value(self, ob, *args): - feed_dict = {self.observations: [ob], self.model.seq_lens: [1]} - assert len(args) == len(self.model.state_in), \ - (args, self.model.state_in) - for k, v in zip(self.model.state_in, args): - feed_dict[k] = v - vf = self.sess.run(self.value_function, feed_dict) - return vf[0] - - def postprocess_trajectory(self, - sample_batch, - other_agent_batches=None, - episode=None): - if not self.config["vtrace"]: - completed = sample_batch["dones"][-1] - if completed: - last_r = 0.0 - else: - next_state = [] - for i in range(len(self.model.state_in)): - next_state.append( - [sample_batch["state_out_{}".format(i)][-1]]) - last_r = self.value(sample_batch["new_obs"][-1], *next_state) - batch = compute_advantages( - sample_batch, - last_r, - self.config["gamma"], - self.config["lambda"], - use_gae=self.config["use_gae"]) - else: - batch = sample_batch - del batch.data["new_obs"] # not used, so save some bandwidth - return batch - - def get_initial_state(self): - return self.model.state_init - - def copy(self, existing_inputs): - return AsyncPPOPolicyGraph( - self.observation_space, - self.action_space, - self.config, +"""Adapted from VTracePolicyGraph to use the PPO surrogate loss. +Keep in sync with changes to VTracePolicyGraph.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import tensorflow as tf +import logging +import gym + +import ray +from ray.rllib.agents.impala import vtrace +from ray.rllib.evaluation.tf_policy_graph import TFPolicyGraph, \ + LearningRateSchedule +from ray.rllib.models.catalog import ModelCatalog +from ray.rllib.utils.error import UnsupportedSpaceException +from ray.rllib.utils.explained_variance import explained_variance +from ray.rllib.models.action_dist import Categorical +from ray.rllib.evaluation.postprocessing import compute_advantages + +logger = logging.getLogger(__name__) + + +class PPOSurrogateLoss(object): + """Loss used when V-trace is disabled. + Arguments: + prev_actions_logp: A float32 tensor of shape [T, B]. + actions_logp: A float32 tensor of shape [T, B]. + action_kl: A float32 tensor of shape [T, B]. + actions_entropy: A float32 tensor of shape [T, B]. + values: A float32 tensor of shape [T, B]. + valid_mask: A bool tensor of valid RNN input elements (#2992). + advantages: A float32 tensor of shape [T, B]. + value_targets: A float32 tensor of shape [T, B]. + """ + + def __init__(self, + prev_actions_logp, + actions_logp, + action_kl, + actions_entropy, + values, + valid_mask, + advantages, + value_targets, + vf_loss_coeff=0.5, + entropy_coeff=-0.01, + clip_param=0.3): + + logp_ratio = tf.exp(actions_logp - prev_actions_logp) + + surrogate_loss = tf.minimum( + advantages * logp_ratio, + advantages * tf.clip_by_value(logp_ratio, 1 - clip_param, + 1 + clip_param)) + + self.mean_kl = tf.reduce_mean(action_kl) + self.pi_loss = -tf.reduce_sum(surrogate_loss) + + # The baseline loss + delta = tf.boolean_mask(values - value_targets, valid_mask) + self.value_targets = value_targets + self.vf_loss = 0.5 * tf.reduce_sum(tf.square(delta)) + + # The entropy loss + self.entropy = tf.reduce_sum( + tf.boolean_mask(actions_entropy, valid_mask)) + + # The summed weighted loss + self.total_loss = (self.pi_loss + self.vf_loss * vf_loss_coeff + + self.entropy * entropy_coeff) + + +class VTraceSurrogateLoss(object): + def __init__(self, + actions, + prev_actions_logp, + actions_logp, + action_kl, + actions_entropy, + dones, + behaviour_logits, + target_logits, + discount, + rewards, + values, + bootstrap_value, + valid_mask, + vf_loss_coeff=0.5, + entropy_coeff=-0.01, + clip_rho_threshold=1.0, + clip_pg_rho_threshold=1.0, + clip_param=0.3): + """PPO surrogate loss with vtrace importance weighting. + VTraceLoss takes tensors of shape [T, B, ...], where `B` is the + batch_size. The reason we need to know `B` is for V-trace to properly + handle episode cut boundaries. + Arguments: + actions: An int32 tensor of shape [T, B, NUM_ACTIONS]. + prev_actions_logp: A float32 tensor of shape [T, B]. + actions_logp: A float32 tensor of shape [T, B]. + action_kl: A float32 tensor of shape [T, B]. + actions_entropy: A float32 tensor of shape [T, B]. + dones: A bool tensor of shape [T, B]. + behaviour_logits: A float32 tensor of shape [T, B, NUM_ACTIONS]. + target_logits: A float32 tensor of shape [T, B, NUM_ACTIONS]. + discount: A float32 scalar. + rewards: A float32 tensor of shape [T, B]. + values: A float32 tensor of shape [T, B]. + bootstrap_value: A float32 tensor of shape [B]. + valid_mask: A bool tensor of valid RNN input elements (#2992). + """ + + # Compute vtrace on the CPU for better perf. + with tf.device("/cpu:0"): + self.vtrace_returns = vtrace.from_logits( + behaviour_policy_logits=behaviour_logits, + target_policy_logits=target_logits, + actions=tf.unstack(tf.cast(actions, tf.int32), axis=2), + discounts=tf.to_float(~dones) * discount, + rewards=rewards, + values=values, + bootstrap_value=bootstrap_value, + clip_rho_threshold=tf.cast(clip_rho_threshold, tf.float32), + clip_pg_rho_threshold=tf.cast(clip_pg_rho_threshold, + tf.float32)) + + logp_ratio = tf.exp(actions_logp - prev_actions_logp) + + advantages = self.vtrace_returns.pg_advantages + surrogate_loss = tf.minimum( + advantages * logp_ratio, + advantages * tf.clip_by_value(logp_ratio, 1 - clip_param, + 1 + clip_param)) + + self.mean_kl = tf.reduce_mean(action_kl) + self.pi_loss = -tf.reduce_sum(surrogate_loss) + + # The baseline loss + delta = tf.boolean_mask(values - self.vtrace_returns.vs, valid_mask) + self.value_targets = self.vtrace_returns.vs + self.vf_loss = 0.5 * tf.reduce_sum(tf.square(delta)) + + # The entropy loss + self.entropy = tf.reduce_sum( + tf.boolean_mask(actions_entropy, valid_mask)) + + # The summed weighted loss + self.total_loss = (self.pi_loss + self.vf_loss * vf_loss_coeff + + self.entropy * entropy_coeff) + + +class AsyncPPOPolicyGraph(LearningRateSchedule, TFPolicyGraph): + def __init__(self, + observation_space, + action_space, + config, + existing_inputs=None): + config = dict(ray.rllib.agents.impala.impala.DEFAULT_CONFIG, **config) + assert config["batch_mode"] == "truncate_episodes", \ + "Must use `truncate_episodes` batch mode with V-trace." + self.config = config + self.sess = tf.get_default_session() + self.grads = None + + is_discrete = False + output_hidden_shape = None + actions_shape = [None] + + # Policy network model + dist_class, logit_dim = ModelCatalog.get_action_dist( + action_space, self.config["model"]) + + # Create input placeholders + if existing_inputs: + if self.config["vtrace"]: + actions, dones, behaviour_logits, rewards, observations, \ + prev_actions, prev_rewards = existing_inputs[:7] + existing_state_in = existing_inputs[7:-1] + existing_seq_lens = existing_inputs[-1] + else: + actions, dones, behaviour_logits, rewards, observations, \ + prev_actions, prev_rewards, adv_ph, value_targets = \ + existing_inputs[:9] + existing_state_in = existing_inputs[9:-1] + existing_seq_lens = existing_inputs[-1] + else: + if isinstance(action_space, gym.spaces.Discrete): + is_discrete = True + output_hidden_shape = [action_space.n] + elif isinstance(action_space, + gym.spaces.multi_discrete.MultiDiscrete): + actions_shape = [None, len(action_space.nvec)] + output_hidden_shape = action_space.nvec + elif self.config["vtrace"]: + raise UnsupportedSpaceException( + "Action space {} is not supported for IMPALA.".format( + action_space)) + + actions = tf.placeholder(tf.int64, actions_shape, name="ac") + dones = tf.placeholder(tf.bool, [None], name="dones") + rewards = tf.placeholder(tf.float32, [None], name="rewards") + behaviour_logits = tf.placeholder( + tf.float32, [None, logit_dim], name="behaviour_logits") + observations = tf.placeholder( + tf.float32, [None] + list(observation_space.shape)) + existing_state_in = None + existing_seq_lens = None + + if not self.config["vtrace"]: + adv_ph = tf.placeholder( + tf.float32, name="advantages", shape=(None, )) + value_targets = tf.placeholder( + tf.float32, name="value_targets", shape=(None, )) + self.observations = observations + + # Unpack behaviour logits + unpacked_behaviour_logits = tf.split( + behaviour_logits, output_hidden_shape, axis=1) + + # Setup the policy + dist_class, logit_dim = ModelCatalog.get_action_dist( + action_space, self.config["model"], + dist_type=self.config["dist_type"]) + prev_actions = ModelCatalog.get_action_placeholder(action_space) + prev_rewards = tf.placeholder(tf.float32, [None], name="prev_reward") + self.model = ModelCatalog.get_model( + { + "obs": observations, + "prev_actions": prev_actions, + "prev_rewards": prev_rewards, + "is_training": self._get_is_training_placeholder(), + }, + observation_space, + logit_dim, + self.config["model"], + state_in=existing_state_in, + seq_lens=existing_seq_lens) + unpacked_outputs = tf.split( + self.model.outputs, output_hidden_shape, axis=1) + + dist_inputs = self.model.outputs if is_discrete else \ + unpacked_outputs + prev_dist_inputs = behaviour_logits if is_discrete else \ + unpacked_behaviour_logits + + action_dist = dist_class(dist_inputs) + prev_action_dist = dist_class(prev_dist_inputs) + + values = self.model.value_function() + self.value_function = values + self.var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, + tf.get_variable_scope().name) + + def make_time_major(tensor, drop_last=False): + """Swaps batch and trajectory axis. + Args: + tensor: A tensor or list of tensors to reshape. + drop_last: A bool indicating whether to drop the last + trajectory item. + Returns: + res: A tensor with swapped axes or a list of tensors with + swapped axes. + """ + if isinstance(tensor, list): + return [make_time_major(t, drop_last) for t in tensor] + + if self.config["model"]["use_lstm"]: + B = tf.shape(self.model.seq_lens)[0] + T = tf.shape(tensor)[0] // B + else: + # Important: chop the tensor into batches at known episode cut + # boundaries. TODO(ekl) this is kind of a hack + T = self.config["sample_batch_size"] + B = tf.shape(tensor)[0] // T + rs = tf.reshape(tensor, + tf.concat([[B, T], tf.shape(tensor)[1:]], axis=0)) + + # swap B and T axes + res = tf.transpose( + rs, + [1, 0] + list(range(2, 1 + int(tf.shape(tensor).shape[0])))) + + if drop_last: + return res[:-1] + return res + + if self.model.state_in: + max_seq_len = tf.reduce_max(self.model.seq_lens) - 1 + mask = tf.sequence_mask(self.model.seq_lens, max_seq_len) + mask = tf.reshape(mask, [-1]) + else: + mask = tf.ones_like(rewards) + + # Inputs are reshaped from [B * T] => [T - 1, B] for V-trace calc. + if self.config["vtrace"]: + logger.info("Using V-Trace surrogate loss (vtrace=True)") + + # Prepare actions for loss + loss_actions = tf.expand_dims( + actions, axis=1) if is_discrete else actions + logp_actions = actions if is_discrete else tf.unstack( + actions, axis=1) + + self.loss = VTraceSurrogateLoss( + actions=make_time_major(loss_actions, drop_last=True), + prev_actions_logp=make_time_major(prev_action_dist.logp( + logp_actions), drop_last=True), + actions_logp=make_time_major(action_dist.logp(logp_actions), + drop_last=True), + action_kl=prev_action_dist.kl(action_dist), + actions_entropy=make_time_major(action_dist.entropy(), + drop_last=True), + dones=make_time_major(dones, drop_last=True), + behaviour_logits=make_time_major(unpacked_behaviour_logits, + drop_last=True), + target_logits=make_time_major(unpacked_outputs, + drop_last=True), + discount=config["gamma"], + rewards=make_time_major(rewards, drop_last=True), + values=make_time_major(values, drop_last=True), + bootstrap_value=make_time_major(values)[-1], + valid_mask=make_time_major(mask, drop_last=True), + vf_loss_coeff=self.config["vf_loss_coeff"], + entropy_coeff=self.config["entropy_coeff"], + clip_rho_threshold=self.config["vtrace_clip_rho_threshold"], + clip_pg_rho_threshold=self.config[ + "vtrace_clip_pg_rho_threshold"], + clip_param=self.config["clip_param"]) + else: + logger.info("Using PPO surrogate loss (vtrace=False)") + self.loss = PPOSurrogateLoss( + prev_actions_logp=make_time_major( + prev_action_dist.logp(actions)), + actions_logp=make_time_major( + action_dist.logp(actions)), + action_kl=prev_action_dist.kl(action_dist), + actions_entropy=make_time_major( + action_dist.entropy()), + values=make_time_major(values), + valid_mask=make_time_major(mask), + advantages=make_time_major(adv_ph), + value_targets=make_time_major(value_targets), + vf_loss_coeff=self.config["vf_loss_coeff"], + entropy_coeff=self.config["entropy_coeff"], + clip_param=self.config["clip_param"]) + + # KL divergence between worker and learner logits for debugging + model_dist = Categorical(self.model.outputs) + behaviour_dist = Categorical(behaviour_logits) + self.KLs = model_dist.kl(behaviour_dist) + self.mean_KL = tf.reduce_mean(self.KLs) + self.max_KL = tf.reduce_max(self.KLs) + self.median_KL = tf.contrib.distributions.percentile(self.KLs, 50.0) + # Initialize TFPolicyGraph + loss_in = [ + ("actions", actions), + ("dones", dones), + ("behaviour_logits", behaviour_logits), + ("rewards", rewards), + ("obs", observations), + ("prev_actions", prev_actions), + ("prev_rewards", prev_rewards), + ] + if not self.config["vtrace"]: + loss_in.append(("advantages", adv_ph)) + loss_in.append(("value_targets", value_targets)) + LearningRateSchedule.__init__(self, self.config["lr"], + self.config["lr_schedule"]) + TFPolicyGraph.__init__( + self, + observation_space, + action_space, + self.sess, + obs_input=observations, + action_sampler=action_dist.sample(), + loss=self.model.loss() + self.loss.total_loss, + loss_inputs=loss_in, + state_inputs=self.model.state_in, + state_outputs=self.model.state_out, + prev_action_input=prev_actions, + prev_reward_input=prev_rewards, + seq_lens=self.model.seq_lens, + max_seq_len=self.config["model"]["max_seq_len"], + batch_divisibility_req=self.config["sample_batch_size"]) + + self.sess.run(tf.global_variables_initializer()) + + if self.config["vtrace"]: + values_batched = make_time_major(values, drop_last=True) + else: + values_batched = make_time_major(values) + self.stats_fetches = {"stats": { + "model_loss": self.model.loss(), + "cur_lr": tf.cast(self.cur_lr, tf.float64), + "policy_loss": self.loss.pi_loss, + "entropy": self.loss.entropy, + "grad_gnorm": tf.global_norm(self._grads), + "var_gnorm": tf.global_norm(self.var_list), + "vf_loss": self.loss.vf_loss, + "vf_explained_var": explained_variance( + tf.reshape(self.loss.value_targets, [-1]), + tf.reshape(values_batched, [-1])), + "mean_KL": self.mean_KL, + "max_KL": self.max_KL, + "median_KL": self.median_KL, + }, "kl": self.loss.mean_kl} + + def optimizer(self): + if self.config["opt_type"] == "adam": + return tf.train.AdamOptimizer(self.cur_lr) + else: + return tf.train.RMSPropOptimizer(self.cur_lr, self.config["decay"], + self.config["momentum"], + self.config["epsilon"]) + + def gradients(self, optimizer): + grads = tf.gradients(self.loss.total_loss, self.var_list) + self.grads, _ = tf.clip_by_global_norm(grads, self.config["grad_clip"]) + clipped_grads = list(zip(self.grads, self.var_list)) + return clipped_grads + + def extra_compute_action_fetches(self): + out = {"behaviour_logits": self.model.outputs} + if not self.config["vtrace"]: + out["vf_preds"] = self.value_function + return out + + def extra_compute_grad_fetches(self): + return self.stats_fetches + + def value(self, ob, *args): + feed_dict = {self.observations: [ob], self.model.seq_lens: [1]} + assert len(args) == len(self.model.state_in), \ + (args, self.model.state_in) + for k, v in zip(self.model.state_in, args): + feed_dict[k] = v + vf = self.sess.run(self.value_function, feed_dict) + return vf[0] + + def postprocess_trajectory(self, + sample_batch, + other_agent_batches=None, + episode=None): + if not self.config["vtrace"]: + completed = sample_batch["dones"][-1] + if completed: + last_r = 0.0 + else: + next_state = [] + for i in range(len(self.model.state_in)): + next_state.append( + [sample_batch["state_out_{}".format(i)][-1]]) + last_r = self.value(sample_batch["new_obs"][-1], *next_state) + batch = compute_advantages( + sample_batch, + last_r, + self.config["gamma"], + self.config["lambda"], + use_gae=self.config["use_gae"]) + else: + batch = sample_batch + del batch.data["new_obs"] # not used, so save some bandwidth + return batch + + def get_initial_state(self): + return self.model.state_init + + def copy(self, existing_inputs): + return AsyncPPOPolicyGraph( + self.observation_space, + self.action_space, + self.config, existing_inputs=existing_inputs) \ No newline at end of file From 56fe32e89403e651fe023f7df457b99a4b77470d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleksandar=20Milovanovi=C4=87?= Date: Thu, 14 Feb 2019 12:21:31 +0100 Subject: [PATCH 19/33] fixed newlines in appo_policy_graph --- python/ray/rllib/agents/ppo/appo_policy_graph.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/python/ray/rllib/agents/ppo/appo_policy_graph.py b/python/ray/rllib/agents/ppo/appo_policy_graph.py index 398129ead93d..3144432e95b6 100644 --- a/python/ray/rllib/agents/ppo/appo_policy_graph.py +++ b/python/ray/rllib/agents/ppo/appo_policy_graph.py @@ -1,4 +1,5 @@ """Adapted from VTracePolicyGraph to use the PPO surrogate loss. + Keep in sync with changes to VTracePolicyGraph.""" from __future__ import absolute_import @@ -24,6 +25,7 @@ class PPOSurrogateLoss(object): """Loss used when V-trace is disabled. + Arguments: prev_actions_logp: A float32 tensor of shape [T, B]. actions_logp: A float32 tensor of shape [T, B]. @@ -93,9 +95,11 @@ def __init__(self, clip_pg_rho_threshold=1.0, clip_param=0.3): """PPO surrogate loss with vtrace importance weighting. + VTraceLoss takes tensors of shape [T, B, ...], where `B` is the batch_size. The reason we need to know `B` is for V-trace to properly handle episode cut boundaries. + Arguments: actions: An int32 tensor of shape [T, B, NUM_ACTIONS]. prev_actions_logp: A float32 tensor of shape [T, B]. From 76046f38904aa844806c979efa11b440fe8ba7c9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleksandar=20Milovanovi=C4=87?= Date: Thu, 14 Feb 2019 12:40:59 +0100 Subject: [PATCH 20/33] aligned with action_dist changes in ray master --- python/ray/rllib/agents/impala/vtrace_policy_graph.py | 6 ++---- python/ray/rllib/models/action_dist.py | 6 +++++- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/python/ray/rllib/agents/impala/vtrace_policy_graph.py b/python/ray/rllib/agents/impala/vtrace_policy_graph.py index 978aeb3b795b..7e932aab25f8 100644 --- a/python/ray/rllib/agents/impala/vtrace_policy_graph.py +++ b/python/ray/rllib/agents/impala/vtrace_policy_graph.py @@ -219,14 +219,12 @@ def make_time_major(tensor, drop_last=False): # Prepare actions for loss loss_actions = actions if is_multidiscrete else tf.expand_dims( - actions, axis=1) - logp_actions = tf.unstack( - actions, axis=1) if is_multidiscrete else actions + actions, axis=1) # Inputs are reshaped from [B * T] => [T - 1, B] for V-trace calc. self.loss = VTraceLoss( actions=make_time_major(loss_actions, drop_last=True), - actions_logp=make_time_major(action_dist.logp(logp_actions), + actions_logp=make_time_major(action_dist.logp(actions), drop_last=True), actions_entropy=make_time_major(action_dist.entropy(), drop_last=True), diff --git a/python/ray/rllib/models/action_dist.py b/python/ray/rllib/models/action_dist.py index 65d1876797f0..ab9eedb6e59b 100644 --- a/python/ray/rllib/models/action_dist.py +++ b/python/ray/rllib/models/action_dist.py @@ -119,8 +119,12 @@ class MultiCategorical(ActionDistribution): def __init__(self, inputs): self.cats = [Categorical(input_) for input_ in inputs] + self.sample_op = self._build_sample_op() def logp(self, actions): + # If tensor is provided, unstack it into list + if isinstance(actions, tf.Tensor): + actions = tf.unstack(actions, axis=1) logps = tf.stack([cat.logp(act) for cat, act in zip(self.cats, actions)]) return tf.reduce_sum(logps, axis=0) @@ -132,7 +136,7 @@ def kl(self, other): return [cat.kl(oth_cat) for cat, oth_cat in zip(self.cats, other.cats)] - def sample(self): + def _build_sample_op(self): return tf.stack([cat.sample() for cat in self.cats], axis=1) From d017c9f4eabb3f728270082372cf8887476dcb51 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleksandar=20Milovanovi=C4=87?= Date: Thu, 14 Feb 2019 12:52:35 +0100 Subject: [PATCH 21/33] small appo fixes --- .../ray/rllib/agents/ppo/appo_policy_graph.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/python/ray/rllib/agents/ppo/appo_policy_graph.py b/python/ray/rllib/agents/ppo/appo_policy_graph.py index 2daa6b7a35f1..af2e9359e4bf 100644 --- a/python/ray/rllib/agents/ppo/appo_policy_graph.py +++ b/python/ray/rllib/agents/ppo/appo_policy_graph.py @@ -168,7 +168,7 @@ def __init__(self, self.sess = tf.get_default_session() self.grads = None - is_discrete = False + is_multidiscrete = False output_hidden_shape = None actions_shape = [None] @@ -191,10 +191,10 @@ def __init__(self, existing_seq_lens = existing_inputs[-1] else: if isinstance(action_space, gym.spaces.Discrete): - is_discrete = True output_hidden_shape = [action_space.n] elif isinstance(action_space, gym.spaces.multi_discrete.MultiDiscrete): + is_multidiscrete = True actions_shape = [None, len(action_space.nvec)] output_hidden_shape = action_space.nvec elif self.config["vtrace"]: @@ -244,10 +244,10 @@ def __init__(self, unpacked_outputs = tf.split( self.model.outputs, output_hidden_shape, axis=1) - dist_inputs = self.model.outputs if is_discrete else \ - unpacked_outputs - prev_dist_inputs = behaviour_logits if is_discrete else \ - unpacked_behaviour_logits + dist_inputs = unpacked_outputs if is_multidiscrete else \ + self.model.outputs + prev_dist_inputs = unpacked_behaviour_logits if is_multidiscrete else \ + behaviour_logits action_dist = dist_class(dist_inputs) prev_action_dist = dist_class(prev_dist_inputs) @@ -302,16 +302,14 @@ def make_time_major(tensor, drop_last=False): logger.info("Using V-Trace surrogate loss (vtrace=True)") # Prepare actions for loss - loss_actions = tf.expand_dims( - actions, axis=1) if is_discrete else actions - logp_actions = actions if is_discrete else tf.unstack( + loss_actions = actions if is_multidiscrete else tf.expand_dims( actions, axis=1) self.loss = VTraceSurrogateLoss( actions=make_time_major(loss_actions, drop_last=True), prev_actions_logp=make_time_major(prev_action_dist.logp( logp_actions), drop_last=True), - actions_logp=make_time_major(action_dist.logp(logp_actions), + actions_logp=make_time_major(action_dist.logp(actions), drop_last=True), action_kl=prev_action_dist.kl(action_dist), actions_entropy=make_time_major(action_dist.entropy(), From c02b9f56fd9d13a99a4b5ccc522a24885e225984 Mon Sep 17 00:00:00 2001 From: Eric Liang Date: Fri, 15 Feb 2019 14:01:06 -0800 Subject: [PATCH 22/33] add vtrace test --- python/ray/rllib/agents/impala/vtrace.py | 35 ++- python/ray/rllib/agents/impala/vtrace_test.py | 276 ++++++++++++++++++ 2 files changed, 308 insertions(+), 3 deletions(-) create mode 100644 python/ray/rllib/agents/impala/vtrace_test.py diff --git a/python/ray/rllib/agents/impala/vtrace.py b/python/ray/rllib/agents/impala/vtrace.py index f71c87b8fe3a..8c2b81a569af 100644 --- a/python/ray/rllib/agents/impala/vtrace.py +++ b/python/ray/rllib/agents/impala/vtrace.py @@ -46,7 +46,12 @@ VTraceReturns = collections.namedtuple('VTraceReturns', 'vs pg_advantages') -def select_policy_values_using_actions(policy_logits, actions): +def log_probs_from_logits_and_actions(policy_logits, actions): + return multi_log_probs_from_logits_and_actions( + [policy_logits], [actions])[0] + + +def multi_log_probs_from_logits_and_actions(policy_logits, actions): """Computes action log-probs from policy logits and actions. In the notation used throughout documentation and comments, T refers to the @@ -95,6 +100,30 @@ def from_logits(behaviour_policy_logits, clip_rho_threshold=1.0, clip_pg_rho_threshold=1.0, name='vtrace_from_logits'): + return multi_from_logits( + [behaviour_policy_logits], + [target_policy_logits], + [actions], + discounts, + rewards, + values, + bootstrap_value, + clip_rho_threshold=clip_rho_threshold, + clip_pg_rho_threshold=clip_pg_rho_threshold, + name=name) + + +def multi_from_logits( + behaviour_policy_logits, + target_policy_logits, + actions, + discounts, + rewards, + values, + bootstrap_value, + clip_rho_threshold=1.0, + clip_pg_rho_threshold=1.0, + name='vtrace_from_logits'): r"""V-trace for softmax policies. Calculates V-trace actor critic targets for softmax polices as described in @@ -175,9 +204,9 @@ def from_logits(behaviour_policy_logits, with tf.name_scope(name, values=[behaviour_policy_logits, target_policy_logits, actions, discounts, rewards, values, bootstrap_value]): - target_action_log_probs = select_policy_values_using_actions( + target_action_log_probs = multi_log_probs_from_logits_and_actions( target_policy_logits, actions) - behaviour_action_log_probs = select_policy_values_using_actions( + behaviour_action_log_probs = multi_log_probs_from_logits_and_actions( behaviour_policy_logits, actions) log_rhos = get_log_rhos(target_action_log_probs, behaviour_action_log_probs) diff --git a/python/ray/rllib/agents/impala/vtrace_test.py b/python/ray/rllib/agents/impala/vtrace_test.py new file mode 100644 index 000000000000..2de05d121924 --- /dev/null +++ b/python/ray/rllib/agents/impala/vtrace_test.py @@ -0,0 +1,276 @@ +# Copyright 2018 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for V-trace. + +For details and theory see: + +"IMPALA: Scalable Distributed Deep-RL with +Importance Weighted Actor-Learner Architectures" +by Espeholt, Soyer, Munos et al. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from absl.testing import parameterized +import numpy as np +import tensorflow as tf +import vtrace + + +def _shaped_arange(*shape): + """Runs np.arange, converts to float and reshapes.""" + return np.arange(np.prod(shape), dtype=np.float32).reshape(*shape) + + +def _softmax(logits): + """Applies softmax non-linearity on inputs.""" + return np.exp(logits) / np.sum(np.exp(logits), axis=-1, keepdims=True) + + +def _ground_truth_calculation(discounts, log_rhos, rewards, values, + bootstrap_value, clip_rho_threshold, + clip_pg_rho_threshold): + """Calculates the ground truth for V-trace in Python/Numpy.""" + vs = [] + seq_len = len(discounts) + rhos = np.exp(log_rhos) + cs = np.minimum(rhos, 1.0) + clipped_rhos = rhos + if clip_rho_threshold: + clipped_rhos = np.minimum(rhos, clip_rho_threshold) + clipped_pg_rhos = rhos + if clip_pg_rho_threshold: + clipped_pg_rhos = np.minimum(rhos, clip_pg_rho_threshold) + + # This is a very inefficient way to calculate the V-trace ground truth. + # We calculate it this way because it is close to the mathematical notation of + # V-trace. + # v_s = V(x_s) + # + \sum^{T-1}_{t=s} \gamma^{t-s} + # * \prod_{i=s}^{t-1} c_i + # * \rho_t (r_t + \gamma V(x_{t+1}) - V(x_t)) + # Note that when we take the product over c_i, we write `s:t` as the notation + # of the paper is inclusive of the `t-1`, but Python is exclusive. + # Also note that np.prod([]) == 1. + values_t_plus_1 = np.concatenate([values, bootstrap_value[None, :]], axis=0) + for s in range(seq_len): + v_s = np.copy(values[s]) # Very important copy. + for t in range(s, seq_len): + v_s += ( + np.prod(discounts[s:t], axis=0) * np.prod(cs[s:t], + axis=0) * clipped_rhos[t] * + (rewards[t] + discounts[t] * values_t_plus_1[t + 1] - values[t])) + vs.append(v_s) + vs = np.stack(vs, axis=0) + pg_advantages = ( + clipped_pg_rhos * (rewards + discounts * np.concatenate( + [vs[1:], bootstrap_value[None, :]], axis=0) - values)) + + return vtrace.VTraceReturns(vs=vs, pg_advantages=pg_advantages) + + +class LogProbsFromLogitsAndActionsTest(tf.test.TestCase, + parameterized.TestCase): + + @parameterized.named_parameters(('Batch1', 1), ('Batch2', 2)) + def test_log_probs_from_logits_and_actions(self, batch_size): + """Tests log_probs_from_logits_and_actions.""" + seq_len = 7 + num_actions = 3 + + policy_logits = _shaped_arange(seq_len, batch_size, num_actions) + 10 + actions = np.random.randint( + 0, num_actions - 1, size=(seq_len, batch_size), dtype=np.int32) + + action_log_probs_tensor = vtrace.log_probs_from_logits_and_actions( + policy_logits, actions) + + # Ground Truth + # Using broadcasting to create a mask that indexes action logits + action_index_mask = actions[..., None] == np.arange(num_actions) + + def index_with_mask(array, mask): + return array[mask].reshape(*array.shape[:-1]) + + # Note: Normally log(softmax) is not a good idea because it's not + # numerically stable. However, in this test we have well-behaved values. + ground_truth_v = index_with_mask( + np.log(_softmax(policy_logits)), action_index_mask) + + with self.test_session() as session: + self.assertAllClose(ground_truth_v, session.run(action_log_probs_tensor)) + + +class VtraceTest(tf.test.TestCase, parameterized.TestCase): + + @parameterized.named_parameters(('Batch1', 1), ('Batch5', 5)) + def test_vtrace(self, batch_size): + """Tests V-trace against ground truth data calculated in python.""" + seq_len = 5 + + # Create log_rhos such that rho will span from near-zero to above the + # clipping thresholds. In particular, calculate log_rhos in [-2.5, 2.5), + # so that rho is in approx [0.08, 12.2). + log_rhos = _shaped_arange(seq_len, batch_size) / (batch_size * seq_len) + log_rhos = 5 * (log_rhos - 0.5) # [0.0, 1.0) -> [-2.5, 2.5). + values = { + 'log_rhos': log_rhos, + # T, B where B_i: [0.9 / (i+1)] * T + 'discounts': + np.array([[0.9 / (b + 1) + for b in range(batch_size)] + for _ in range(seq_len)]), + 'rewards': + _shaped_arange(seq_len, batch_size), + 'values': + _shaped_arange(seq_len, batch_size) / batch_size, + 'bootstrap_value': + _shaped_arange(batch_size) + 1.0, + 'clip_rho_threshold': + 3.7, + 'clip_pg_rho_threshold': + 2.2, + } + + output = vtrace.from_importance_weights(**values) + + with self.test_session() as session: + output_v = session.run(output) + + ground_truth_v = _ground_truth_calculation(**values) + for a, b in zip(ground_truth_v, output_v): + self.assertAllClose(a, b) + + @parameterized.named_parameters(('Batch1', 1), ('Batch2', 2)) + def test_vtrace_from_logits(self, batch_size): + """Tests V-trace calculated from logits.""" + seq_len = 5 + num_actions = 3 + clip_rho_threshold = None # No clipping. + clip_pg_rho_threshold = None # No clipping. + + # Intentionally leaving shapes unspecified to test if V-trace can + # deal with that. + placeholders = { + # T, B, NUM_ACTIONS + 'behaviour_policy_logits': + tf.placeholder(dtype=tf.float32, shape=[None, None, None]), + # T, B, NUM_ACTIONS + 'target_policy_logits': + tf.placeholder(dtype=tf.float32, shape=[None, None, None]), + 'actions': + tf.placeholder(dtype=tf.int32, shape=[None, None]), + 'discounts': + tf.placeholder(dtype=tf.float32, shape=[None, None]), + 'rewards': + tf.placeholder(dtype=tf.float32, shape=[None, None]), + 'values': + tf.placeholder(dtype=tf.float32, shape=[None, None]), + 'bootstrap_value': + tf.placeholder(dtype=tf.float32, shape=[None]), + } + + from_logits_output = vtrace.from_logits( + clip_rho_threshold=clip_rho_threshold, + clip_pg_rho_threshold=clip_pg_rho_threshold, + **placeholders) + + target_log_probs = vtrace.log_probs_from_logits_and_actions( + placeholders['target_policy_logits'], placeholders['actions']) + behaviour_log_probs = vtrace.log_probs_from_logits_and_actions( + placeholders['behaviour_policy_logits'], placeholders['actions']) + log_rhos = target_log_probs - behaviour_log_probs + ground_truth = (log_rhos, behaviour_log_probs, target_log_probs) + + values = { + 'behaviour_policy_logits': + _shaped_arange(seq_len, batch_size, num_actions), + 'target_policy_logits': + _shaped_arange(seq_len, batch_size, num_actions), + 'actions': + np.random.randint(0, num_actions - 1, size=(seq_len, batch_size)), + 'discounts': + np.array( # T, B where B_i: [0.9 / (i+1)] * T + [[0.9 / (b + 1) + for b in range(batch_size)] + for _ in range(seq_len)]), + 'rewards': + _shaped_arange(seq_len, batch_size), + 'values': + _shaped_arange(seq_len, batch_size) / batch_size, + 'bootstrap_value': + _shaped_arange(batch_size) + 1.0, # B + } + + feed_dict = {placeholders[k]: v for k, v in values.items()} + with self.test_session() as session: + from_logits_output_v = session.run( + from_logits_output, feed_dict=feed_dict) + (ground_truth_log_rhos, ground_truth_behaviour_action_log_probs, + ground_truth_target_action_log_probs) = session.run( + ground_truth, feed_dict=feed_dict) + + # Calculate V-trace using the ground truth logits. + from_iw = vtrace.from_importance_weights( + log_rhos=ground_truth_log_rhos, + discounts=values['discounts'], + rewards=values['rewards'], + values=values['values'], + bootstrap_value=values['bootstrap_value'], + clip_rho_threshold=clip_rho_threshold, + clip_pg_rho_threshold=clip_pg_rho_threshold) + + with self.test_session() as session: + from_iw_v = session.run(from_iw) + + self.assertAllClose(from_iw_v.vs, from_logits_output_v.vs) + self.assertAllClose(from_iw_v.pg_advantages, + from_logits_output_v.pg_advantages) + self.assertAllClose(ground_truth_behaviour_action_log_probs, + from_logits_output_v.behaviour_action_log_probs) + self.assertAllClose(ground_truth_target_action_log_probs, + from_logits_output_v.target_action_log_probs) + self.assertAllClose(ground_truth_log_rhos, from_logits_output_v.log_rhos) + + def test_higher_rank_inputs_for_importance_weights(self): + """Checks support for additional dimensions in inputs.""" + placeholders = { + 'log_rhos': tf.placeholder(dtype=tf.float32, shape=[None, None, 1]), + 'discounts': tf.placeholder(dtype=tf.float32, shape=[None, None, 1]), + 'rewards': tf.placeholder(dtype=tf.float32, shape=[None, None, 42]), + 'values': tf.placeholder(dtype=tf.float32, shape=[None, None, 42]), + 'bootstrap_value': tf.placeholder(dtype=tf.float32, shape=[None, 42]) + } + output = vtrace.from_importance_weights(**placeholders) + self.assertEqual(output.vs.shape.as_list()[-1], 42) + + def test_inconsistent_rank_inputs_for_importance_weights(self): + """Test one of many possible errors in shape of inputs.""" + placeholders = { + 'log_rhos': tf.placeholder(dtype=tf.float32, shape=[None, None, 1]), + 'discounts': tf.placeholder(dtype=tf.float32, shape=[None, None, 1]), + 'rewards': tf.placeholder(dtype=tf.float32, shape=[None, None, 42]), + 'values': tf.placeholder(dtype=tf.float32, shape=[None, None, 42]), + # Should be [None, 42]. + 'bootstrap_value': tf.placeholder(dtype=tf.float32, shape=[None]) + } + with self.assertRaisesRegexp(ValueError, 'must have rank 2'): + vtrace.from_importance_weights(**placeholders) + + +if __name__ == '__main__': + tf.test.main() From fbbed63c434ee794633edbb3312bf273724b7646 Mon Sep 17 00:00:00 2001 From: Eric Liang Date: Fri, 15 Feb 2019 15:16:50 -0800 Subject: [PATCH 23/33] fix appo impala integration --- python/ray/rllib/agents/impala/vtrace_policy_graph.py | 2 +- python/ray/rllib/agents/ppo/appo_policy_graph.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/python/ray/rllib/agents/impala/vtrace_policy_graph.py b/python/ray/rllib/agents/impala/vtrace_policy_graph.py index 7e932aab25f8..6cb0f25e757e 100644 --- a/python/ray/rllib/agents/impala/vtrace_policy_graph.py +++ b/python/ray/rllib/agents/impala/vtrace_policy_graph.py @@ -67,7 +67,7 @@ def __init__(self, # Compute vtrace on the CPU for better perf. with tf.device("/cpu:0"): - self.vtrace_returns = vtrace.from_logits( + self.vtrace_returns = vtrace.multi_from_logits( behaviour_policy_logits=behaviour_logits, target_policy_logits=target_logits, actions=tf.unstack(tf.cast(actions, tf.int32), axis=2), diff --git a/python/ray/rllib/agents/ppo/appo_policy_graph.py b/python/ray/rllib/agents/ppo/appo_policy_graph.py index af2e9359e4bf..2087908e541a 100644 --- a/python/ray/rllib/agents/ppo/appo_policy_graph.py +++ b/python/ray/rllib/agents/ppo/appo_policy_graph.py @@ -118,7 +118,7 @@ def __init__(self, # Compute vtrace on the CPU for better perf. with tf.device("/cpu:0"): - self.vtrace_returns = vtrace.from_logits( + self.vtrace_returns = vtrace.multi_from_logits( behaviour_policy_logits=behaviour_logits, target_policy_logits=target_logits, actions=tf.unstack(tf.cast(actions, tf.int32), axis=2), @@ -475,4 +475,4 @@ def copy(self, existing_inputs): self.observation_space, self.action_space, self.config, - existing_inputs=existing_inputs) \ No newline at end of file + existing_inputs=existing_inputs) From bcb2113763e0f21188374f5cb7b51d58e5796446 Mon Sep 17 00:00:00 2001 From: Eric Liang Date: Fri, 15 Feb 2019 15:48:06 -0800 Subject: [PATCH 24/33] add to jenkins --- test/jenkins_tests/run_multi_node_tests.sh | 3 +++ 1 file changed, 3 insertions(+) diff --git a/test/jenkins_tests/run_multi_node_tests.sh b/test/jenkins_tests/run_multi_node_tests.sh index b513a3148644..9fdde32aa404 100755 --- a/test/jenkins_tests/run_multi_node_tests.sh +++ b/test/jenkins_tests/run_multi_node_tests.sh @@ -202,6 +202,9 @@ docker run --rm --shm-size=${SHM_SIZE} --memory=${MEMORY_SIZE} $DOCKER_SHA \ --stop '{"training_iteration": 2}' \ --config '{"num_workers": 1}' +docker run --rm --shm-size=${SHM_SIZE} --memory=${MEMORY_SIZE} $DOCKER_SHA \ + python /ray/python/ray/rllib/agents/impala/vtrace_test.py + docker run --rm --shm-size=${SHM_SIZE} --memory=${MEMORY_SIZE} $DOCKER_SHA \ python /ray/python/ray/rllib/train.py \ --env CartPole-v0 \ From 6e06ba6718c40984b508f5811798294315b506bf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleksandar=20Milovanovi=C4=87?= Date: Mon, 18 Feb 2019 11:47:53 +0100 Subject: [PATCH 25/33] fixing appo policy graph changes --- .../ray/rllib/agents/ppo/appo_policy_graph.py | 42 ++++++++++++------- 1 file changed, 26 insertions(+), 16 deletions(-) diff --git a/python/ray/rllib/agents/ppo/appo_policy_graph.py b/python/ray/rllib/agents/ppo/appo_policy_graph.py index 2087908e541a..041dc191f97a 100644 --- a/python/ray/rllib/agents/ppo/appo_policy_graph.py +++ b/python/ray/rllib/agents/ppo/appo_policy_graph.py @@ -17,7 +17,7 @@ from ray.rllib.models.catalog import ModelCatalog from ray.rllib.utils.error import UnsupportedSpaceException from ray.rllib.utils.explained_variance import explained_variance -from ray.rllib.models.action_dist import Categorical +from ray.rllib.models.action_dist import MultiCategorical from ray.rllib.evaluation.postprocessing import compute_advantages logger = logging.getLogger(__name__) @@ -199,8 +199,7 @@ def __init__(self, output_hidden_shape = action_space.nvec elif self.config["vtrace"]: raise UnsupportedSpaceException( - "Action space {} is not supported for IMPALA.".format( - action_space)) + "Action space {} is not supported for APPO with VTrace.".format(action_space)) actions = tf.placeholder(tf.int64, actions_shape, name="ac") dones = tf.placeholder(tf.bool, [None], name="dones") @@ -349,12 +348,28 @@ def make_time_major(tensor, drop_last=False): clip_param=self.config["clip_param"]) # KL divergence between worker and learner logits for debugging - model_dist = Categorical(self.model.outputs) - behaviour_dist = Categorical(behaviour_logits) - self.KLs = model_dist.kl(behaviour_dist) - self.mean_KL = tf.reduce_mean(self.KLs) - self.max_KL = tf.reduce_max(self.KLs) - self.median_KL = tf.contrib.distributions.percentile(self.KLs, 50.0) + model_dist = MultiCategorical(unpacked_outputs) + behaviour_dist = MultiCategorical(unpacked_behaviour_logits) + + kls = model_dist.kl(behaviour_dist) + if len(kls) > 1: + self.KL_stats = {} + + for i, kl in enumerate(kls): + self.KL_stats.update({ + f"mean_KL_{i}": tf.reduce_mean(kl), + f"max_KL_{i}": tf.reduce_max(kl), + f"median_KL_{i}": tf.contrib.distributions.percentile( + kl, 50.0), + }) + else: + self.KL_stats = { + "mean_KL": tf.reduce_mean(kls[0]), + "max_KL": tf.reduce_max(kls[0]), + "median_KL": tf.contrib.distributions.percentile( + kls[0], 50.0), + } + # Initialize TFPolicyGraph loss_in = [ ("actions", actions), @@ -390,10 +405,7 @@ def make_time_major(tensor, drop_last=False): self.sess.run(tf.global_variables_initializer()) - if self.config["vtrace"]: - values_batched = make_time_major(values, drop_last=True) - else: - values_batched = make_time_major(values) + values_batched = make_time_major(values, drop_last=self.config["vtrace"]) self.stats_fetches = {"stats": { "model_loss": self.model.loss(), "cur_lr": tf.cast(self.cur_lr, tf.float64), @@ -405,9 +417,7 @@ def make_time_major(tensor, drop_last=False): "vf_explained_var": explained_variance( tf.reshape(self.loss.value_targets, [-1]), tf.reshape(values_batched, [-1])), - "mean_KL": self.mean_KL, - "max_KL": self.max_KL, - "median_KL": self.median_KL, + **self.KL_stats, }, "kl": self.loss.mean_kl} def optimizer(self): From f161edf89e996ce734ac3aa3fd4cf079f8dfbfe9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleksandar=20Milovanovi=C4=87?= Date: Mon, 18 Feb 2019 13:51:28 +0100 Subject: [PATCH 26/33] fixed vtrace tests --- python/ray/rllib/agents/impala/vtrace.py | 57 ++++++++++++------- .../ray/rllib/agents/ppo/appo_policy_graph.py | 2 +- 2 files changed, 36 insertions(+), 23 deletions(-) diff --git a/python/ray/rllib/agents/impala/vtrace.py b/python/ray/rllib/agents/impala/vtrace.py index 8c2b81a569af..1f8174d77f26 100644 --- a/python/ray/rllib/agents/impala/vtrace.py +++ b/python/ray/rllib/agents/impala/vtrace.py @@ -24,8 +24,8 @@ In addition to the original paper's code, changes have been made to support MultiDiscrete action spaces. behaviour_policy_logits, target_policy_logits and actions parameters in the entry point -from_logits method are now lists of tensors instead of just -being tensors. +multi_from_logits method accepts lists of tensors instead of just +tensors. """ from __future__ import absolute_import @@ -100,7 +100,9 @@ def from_logits(behaviour_policy_logits, clip_rho_threshold=1.0, clip_pg_rho_threshold=1.0, name='vtrace_from_logits'): - return multi_from_logits( + """multi_from_logits wrapper used only for tests""" + + res = multi_from_logits( [behaviour_policy_logits], [target_policy_logits], [actions], @@ -112,18 +114,26 @@ def from_logits(behaviour_policy_logits, clip_pg_rho_threshold=clip_pg_rho_threshold, name=name) + return VTraceFromLogitsReturns( + vs = res.vs, + pg_advantages = res.pg_advantages, + log_rhos=res.log_rhos, + behaviour_action_log_probs=tf.squeeze(res.behaviour_action_log_probs, axis=0), + target_action_log_probs=tf.squeeze(res.target_action_log_probs, axis=0), + ) + def multi_from_logits( - behaviour_policy_logits, - target_policy_logits, - actions, - discounts, - rewards, - values, - bootstrap_value, - clip_rho_threshold=1.0, - clip_pg_rho_threshold=1.0, - name='vtrace_from_logits'): + behaviour_policy_logits, + target_policy_logits, + actions, + discounts, + rewards, + values, + bootstrap_value, + clip_rho_threshold=1.0, + clip_pg_rho_threshold=1.0, + name='vtrace_from_logits'): r"""V-trace for softmax policies. Calculates V-trace actor critic targets for softmax polices as described in @@ -209,7 +219,8 @@ def multi_from_logits( behaviour_action_log_probs = multi_log_probs_from_logits_and_actions( behaviour_policy_logits, actions) - log_rhos = get_log_rhos(target_action_log_probs, behaviour_action_log_probs) + log_rhos = get_log_rhos(target_action_log_probs, + behaviour_action_log_probs) vtrace_returns = from_importance_weights( log_rhos=log_rhos, @@ -307,16 +318,17 @@ def from_importance_weights(log_rhos, if clip_rho_threshold is not None: clipped_rhos = tf.minimum( clip_rho_threshold, rhos, name='clipped_rhos') + + tf.summary.histogram('clipped_rhos_1000', tf.minimum(1000.0, rhos)) + tf.summary.scalar( + 'num_of_clipped_rhos', + tf.reduce_sum(tf.cast( + tf.equal(clipped_rhos, clip_rho_threshold), tf.int32)) + ) + tf.summary.scalar('size_of_clipped_rhos', tf.size(clipped_rhos)) else: clipped_rhos = rhos - tf.summary.histogram('clipped_rhos_1000', tf.minimum(1000.0, rhos)) - tf.summary.scalar( - 'num_of_clipped_rhos', - tf.reduce_sum(tf.cast(tf.equal(clipped_rhos, clip_rho_threshold), tf.int32)) - ) - tf.summary.scalar('size_of_clipped_rhos', tf.size(clipped_rhos)) - cs = tf.minimum(1.0, rhos, name='cs') # Append bootstrapped value to get [v1, ..., v_t+1] values_t_plus_1 = tf.concat( @@ -371,7 +383,8 @@ def scanfunc(acc, sequence_item): def get_log_rhos(behaviour_action_log_probs, target_action_log_probs): """With the selected log_probs for multi-discrete actions of behaviour and target policies we compute the log_rhos for calculating the vtrace.""" - log_rhos = [t - b for t, b in zip(target_action_log_probs, behaviour_action_log_probs)] + log_rhos = [t - b for t, + b in zip(target_action_log_probs, behaviour_action_log_probs)] log_rhos = [tf.convert_to_tensor(l, dtype=tf.float32) for l in log_rhos] log_rhos = tf.reduce_sum(tf.stack(log_rhos), axis=0) diff --git a/python/ray/rllib/agents/ppo/appo_policy_graph.py b/python/ray/rllib/agents/ppo/appo_policy_graph.py index 041dc191f97a..212e52bd34ed 100644 --- a/python/ray/rllib/agents/ppo/appo_policy_graph.py +++ b/python/ray/rllib/agents/ppo/appo_policy_graph.py @@ -307,7 +307,7 @@ def make_time_major(tensor, drop_last=False): self.loss = VTraceSurrogateLoss( actions=make_time_major(loss_actions, drop_last=True), prev_actions_logp=make_time_major(prev_action_dist.logp( - logp_actions), drop_last=True), + actions), drop_last=True), actions_logp=make_time_major(action_dist.logp(actions), drop_last=True), action_kl=prev_action_dist.kl(action_dist), From 1dbed0898a7499d6912e74ac98d5247cfdabbb04 Mon Sep 17 00:00:00 2001 From: Eric Liang Date: Mon, 18 Feb 2019 11:28:07 -0800 Subject: [PATCH 27/33] lint and py2 compat --- python/ray/rllib/agents/impala/vtrace.py | 74 +-- .../agents/impala/vtrace_policy_graph.py | 33 +- python/ray/rllib/agents/impala/vtrace_test.py | 428 +++++++++--------- .../ray/rllib/agents/ppo/appo_policy_graph.py | 74 +-- python/ray/rllib/models/action_dist.py | 11 +- 5 files changed, 310 insertions(+), 310 deletions(-) diff --git a/python/ray/rllib/agents/impala/vtrace.py b/python/ray/rllib/agents/impala/vtrace.py index 1f8174d77f26..4031ee4c8d45 100644 --- a/python/ray/rllib/agents/impala/vtrace.py +++ b/python/ray/rllib/agents/impala/vtrace.py @@ -47,8 +47,8 @@ def log_probs_from_logits_and_actions(policy_logits, actions): - return multi_log_probs_from_logits_and_actions( - [policy_logits], [actions])[0] + return multi_log_probs_from_logits_and_actions([policy_logits], + [actions])[0] def multi_log_probs_from_logits_and_actions(policy_logits, actions): @@ -56,7 +56,8 @@ def multi_log_probs_from_logits_and_actions(policy_logits, actions): In the notation used throughout documentation and comments, T refers to the time dimension ranging from 0 to T-1. B refers to the batch size and - ACTION_SPACE refers to the list of numbers each representing a number of actions. + ACTION_SPACE refers to the list of numbers each representing a number of + actions. Args: policy_logits: A list with length of ACTION_SPACE of float32 @@ -103,9 +104,7 @@ def from_logits(behaviour_policy_logits, """multi_from_logits wrapper used only for tests""" res = multi_from_logits( - [behaviour_policy_logits], - [target_policy_logits], - [actions], + [behaviour_policy_logits], [target_policy_logits], [actions], discounts, rewards, values, @@ -115,25 +114,26 @@ def from_logits(behaviour_policy_logits, name=name) return VTraceFromLogitsReturns( - vs = res.vs, - pg_advantages = res.pg_advantages, - log_rhos=res.log_rhos, - behaviour_action_log_probs=tf.squeeze(res.behaviour_action_log_probs, axis=0), - target_action_log_probs=tf.squeeze(res.target_action_log_probs, axis=0), + vs=res.vs, + pg_advantages=res.pg_advantages, + log_rhos=res.log_rhos, + behaviour_action_log_probs=tf.squeeze( + res.behaviour_action_log_probs, axis=0), + target_action_log_probs=tf.squeeze( + res.target_action_log_probs, axis=0), ) -def multi_from_logits( - behaviour_policy_logits, - target_policy_logits, - actions, - discounts, - rewards, - values, - bootstrap_value, - clip_rho_threshold=1.0, - clip_pg_rho_threshold=1.0, - name='vtrace_from_logits'): +def multi_from_logits(behaviour_policy_logits, + target_policy_logits, + actions, + discounts, + rewards, + values, + bootstrap_value, + clip_rho_threshold=1.0, + clip_pg_rho_threshold=1.0, + name='vtrace_from_logits'): r"""V-trace for softmax policies. Calculates V-trace actor critic targets for softmax polices as described in @@ -148,7 +148,8 @@ def multi_from_logits( In the notation used throughout documentation and comments, T refers to the time dimension ranging from 0 to T-1. B refers to the batch size and - ACTION_SPACE refers to the list of numbers each representing a number of actions. + ACTION_SPACE refers to the list of numbers each representing a number of + actions. Args: behaviour_policy_logits: A list with length of ACTION_SPACE of float32 @@ -156,13 +157,15 @@ def multi_from_logits( [T, B, ACTION_SPACE[0]], ..., [T, B, ACTION_SPACE[-1]] - with un-normalized log-probabilities parameterizing the softmax behaviour policy. + with un-normalized log-probabilities parameterizing the softmax behaviour + policy. target_policy_logits: A list with length of ACTION_SPACE of float32 tensors of shapes [T, B, ACTION_SPACE[0]], ..., [T, B, ACTION_SPACE[-1]] - with un-normalized log-probabilities parameterizing the softmax target policy. + with un-normalized log-probabilities parameterizing the softmax target + policy. actions: A list with length of ACTION_SPACE of int32 tensors of shapes [T, B], @@ -211,9 +214,12 @@ def multi_from_logits( target_policy_logits[i].shape.assert_has_rank(3) actions[i].shape.assert_has_rank(2) - with tf.name_scope(name, values=[behaviour_policy_logits, target_policy_logits, actions, - discounts, rewards, values, - bootstrap_value]): + with tf.name_scope( + name, + values=[ + behaviour_policy_logits, target_policy_logits, actions, + discounts, rewards, values, bootstrap_value + ]): target_action_log_probs = multi_log_probs_from_logits_and_actions( target_policy_logits, actions) behaviour_action_log_probs = multi_log_probs_from_logits_and_actions( @@ -322,9 +328,9 @@ def from_importance_weights(log_rhos, tf.summary.histogram('clipped_rhos_1000', tf.minimum(1000.0, rhos)) tf.summary.scalar( 'num_of_clipped_rhos', - tf.reduce_sum(tf.cast( - tf.equal(clipped_rhos, clip_rho_threshold), tf.int32)) - ) + tf.reduce_sum( + tf.cast( + tf.equal(clipped_rhos, clip_rho_threshold), tf.int32))) tf.summary.scalar('size_of_clipped_rhos', tf.size(clipped_rhos)) else: clipped_rhos = rhos @@ -383,8 +389,10 @@ def scanfunc(acc, sequence_item): def get_log_rhos(behaviour_action_log_probs, target_action_log_probs): """With the selected log_probs for multi-discrete actions of behaviour and target policies we compute the log_rhos for calculating the vtrace.""" - log_rhos = [t - b for t, - b in zip(target_action_log_probs, behaviour_action_log_probs)] + log_rhos = [ + t - b + for t, b in zip(target_action_log_probs, behaviour_action_log_probs) + ] log_rhos = [tf.convert_to_tensor(l, dtype=tf.float32) for l in log_rhos] log_rhos = tf.reduce_sum(tf.stack(log_rhos), axis=0) diff --git a/python/ray/rllib/agents/impala/vtrace_policy_graph.py b/python/ray/rllib/agents/impala/vtrace_policy_graph.py index 6cb0f25e757e..0ce178e71d1c 100644 --- a/python/ray/rllib/agents/impala/vtrace_policy_graph.py +++ b/python/ray/rllib/agents/impala/vtrace_policy_graph.py @@ -136,9 +136,9 @@ def __init__(self, actions = tf.placeholder(tf.int64, actions_shape, name="ac") dones = tf.placeholder(tf.bool, [None], name="dones") rewards = tf.placeholder(tf.float32, [None], name="rewards") - behaviour_logits = tf.placeholder(tf.float32, - [None, sum(output_hidden_shape)], - name="behaviour_logits") + behaviour_logits = tf.placeholder( + tf.float32, [None, sum(output_hidden_shape)], + name="behaviour_logits") observations = tf.placeholder( tf.float32, [None] + list(observation_space.shape)) existing_state_in = None @@ -150,7 +150,8 @@ def __init__(self, # Setup the policy dist_class, logit_dim = ModelCatalog.get_action_dist( - action_space, self.config["model"], + action_space, + self.config["model"], dist_type=self.config["dist_type"]) prev_actions = ModelCatalog.get_action_placeholder(action_space) prev_rewards = tf.placeholder(tf.float32, [None], name="prev_reward") @@ -224,10 +225,10 @@ def make_time_major(tensor, drop_last=False): # Inputs are reshaped from [B * T] => [T - 1, B] for V-trace calc. self.loss = VTraceLoss( actions=make_time_major(loss_actions, drop_last=True), - actions_logp=make_time_major(action_dist.logp(actions), - drop_last=True), - actions_entropy=make_time_major(action_dist.entropy(), - drop_last=True), + actions_logp=make_time_major( + action_dist.logp(actions), drop_last=True), + actions_entropy=make_time_major( + action_dist.entropy(), drop_last=True), dones=make_time_major(dones, drop_last=True), behaviour_logits=make_time_major( unpacked_behaviour_logits, drop_last=True), @@ -252,17 +253,16 @@ def make_time_major(tensor, drop_last=False): for i, kl in enumerate(kls): self.KL_stats.update({ - f"mean_KL_{i}": tf.reduce_mean(kl), - f"max_KL_{i}": tf.reduce_max(kl), - f"median_KL_{i}": tf.contrib.distributions.percentile( - kl, 50.0), + "mean_KL_{}".format(i): tf.reduce_mean(kl), + "max_KL_{}".format(i): tf.reduce_max(kl), + "median_KL_{}".format(i): tf.contrib.distributions. + percentile(kl, 50.0), }) else: self.KL_stats = { "mean_KL": tf.reduce_mean(kls[0]), "max_KL": tf.reduce_max(kls[0]), - "median_KL": tf.contrib.distributions.percentile( - kls[0], 50.0), + "median_KL": tf.contrib.distributions.percentile(kls[0], 50.0), } # Initialize TFPolicyGraph @@ -298,7 +298,7 @@ def make_time_major(tensor, drop_last=False): self.sess.run(tf.global_variables_initializer()) self.stats_fetches = { - "stats": { + "stats": dict({ "cur_lr": tf.cast(self.cur_lr, tf.float64), "policy_loss": self.loss.pi_loss, "entropy": self.loss.entropy, @@ -308,8 +308,7 @@ def make_time_major(tensor, drop_last=False): "vf_explained_var": explained_variance( tf.reshape(self.loss.vtrace_returns.vs, [-1]), tf.reshape(make_time_major(values, drop_last=True), [-1])), - **self.KL_stats, - }, + }, **self.KL_stats), } @override(TFPolicyGraph) diff --git a/python/ray/rllib/agents/impala/vtrace_test.py b/python/ray/rllib/agents/impala/vtrace_test.py index 2de05d121924..f74798fffdbb 100644 --- a/python/ray/rllib/agents/impala/vtrace_test.py +++ b/python/ray/rllib/agents/impala/vtrace_test.py @@ -11,7 +11,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - """Tests for V-trace. For details and theory see: @@ -32,245 +31,238 @@ def _shaped_arange(*shape): - """Runs np.arange, converts to float and reshapes.""" - return np.arange(np.prod(shape), dtype=np.float32).reshape(*shape) + """Runs np.arange, converts to float and reshapes.""" + return np.arange(np.prod(shape), dtype=np.float32).reshape(*shape) def _softmax(logits): - """Applies softmax non-linearity on inputs.""" - return np.exp(logits) / np.sum(np.exp(logits), axis=-1, keepdims=True) + """Applies softmax non-linearity on inputs.""" + return np.exp(logits) / np.sum(np.exp(logits), axis=-1, keepdims=True) def _ground_truth_calculation(discounts, log_rhos, rewards, values, bootstrap_value, clip_rho_threshold, clip_pg_rho_threshold): - """Calculates the ground truth for V-trace in Python/Numpy.""" - vs = [] - seq_len = len(discounts) - rhos = np.exp(log_rhos) - cs = np.minimum(rhos, 1.0) - clipped_rhos = rhos - if clip_rho_threshold: - clipped_rhos = np.minimum(rhos, clip_rho_threshold) - clipped_pg_rhos = rhos - if clip_pg_rho_threshold: - clipped_pg_rhos = np.minimum(rhos, clip_pg_rho_threshold) - - # This is a very inefficient way to calculate the V-trace ground truth. - # We calculate it this way because it is close to the mathematical notation of - # V-trace. - # v_s = V(x_s) - # + \sum^{T-1}_{t=s} \gamma^{t-s} - # * \prod_{i=s}^{t-1} c_i - # * \rho_t (r_t + \gamma V(x_{t+1}) - V(x_t)) - # Note that when we take the product over c_i, we write `s:t` as the notation - # of the paper is inclusive of the `t-1`, but Python is exclusive. - # Also note that np.prod([]) == 1. - values_t_plus_1 = np.concatenate([values, bootstrap_value[None, :]], axis=0) - for s in range(seq_len): - v_s = np.copy(values[s]) # Very important copy. - for t in range(s, seq_len): - v_s += ( - np.prod(discounts[s:t], axis=0) * np.prod(cs[s:t], - axis=0) * clipped_rhos[t] * - (rewards[t] + discounts[t] * values_t_plus_1[t + 1] - values[t])) - vs.append(v_s) - vs = np.stack(vs, axis=0) - pg_advantages = ( - clipped_pg_rhos * (rewards + discounts * np.concatenate( - [vs[1:], bootstrap_value[None, :]], axis=0) - values)) - - return vtrace.VTraceReturns(vs=vs, pg_advantages=pg_advantages) + """Calculates the ground truth for V-trace in Python/Numpy.""" + vs = [] + seq_len = len(discounts) + rhos = np.exp(log_rhos) + cs = np.minimum(rhos, 1.0) + clipped_rhos = rhos + if clip_rho_threshold: + clipped_rhos = np.minimum(rhos, clip_rho_threshold) + clipped_pg_rhos = rhos + if clip_pg_rho_threshold: + clipped_pg_rhos = np.minimum(rhos, clip_pg_rho_threshold) + + # This is a very inefficient way to calculate the V-trace ground truth. + # We calculate it this way because it is close to the mathematical notation + # of + # V-trace. + # v_s = V(x_s) + # + \sum^{T-1}_{t=s} \gamma^{t-s} + # * \prod_{i=s}^{t-1} c_i + # * \rho_t (r_t + \gamma V(x_{t+1}) - V(x_t)) + # Note that when we take the product over c_i, we write `s:t` as the + # notation + # of the paper is inclusive of the `t-1`, but Python is exclusive. + # Also note that np.prod([]) == 1. + values_t_plus_1 = np.concatenate( + [values, bootstrap_value[None, :]], axis=0) + for s in range(seq_len): + v_s = np.copy(values[s]) # Very important copy. + for t in range(s, seq_len): + v_s += (np.prod(discounts[s:t], axis=0) * np.prod(cs[s:t], axis=0) + * clipped_rhos[t] * (rewards[t] + discounts[t] * + values_t_plus_1[t + 1] - values[t])) + vs.append(v_s) + vs = np.stack(vs, axis=0) + pg_advantages = (clipped_pg_rhos * (rewards + discounts * np.concatenate( + [vs[1:], bootstrap_value[None, :]], axis=0) - values)) + + return vtrace.VTraceReturns(vs=vs, pg_advantages=pg_advantages) class LogProbsFromLogitsAndActionsTest(tf.test.TestCase, parameterized.TestCase): + @parameterized.named_parameters(('Batch1', 1), ('Batch2', 2)) + def test_log_probs_from_logits_and_actions(self, batch_size): + """Tests log_probs_from_logits_and_actions.""" + seq_len = 7 + num_actions = 3 - @parameterized.named_parameters(('Batch1', 1), ('Batch2', 2)) - def test_log_probs_from_logits_and_actions(self, batch_size): - """Tests log_probs_from_logits_and_actions.""" - seq_len = 7 - num_actions = 3 - - policy_logits = _shaped_arange(seq_len, batch_size, num_actions) + 10 - actions = np.random.randint( - 0, num_actions - 1, size=(seq_len, batch_size), dtype=np.int32) + policy_logits = _shaped_arange(seq_len, batch_size, num_actions) + 10 + actions = np.random.randint( + 0, num_actions - 1, size=(seq_len, batch_size), dtype=np.int32) - action_log_probs_tensor = vtrace.log_probs_from_logits_and_actions( - policy_logits, actions) + action_log_probs_tensor = vtrace.log_probs_from_logits_and_actions( + policy_logits, actions) - # Ground Truth - # Using broadcasting to create a mask that indexes action logits - action_index_mask = actions[..., None] == np.arange(num_actions) + # Ground Truth + # Using broadcasting to create a mask that indexes action logits + action_index_mask = actions[..., None] == np.arange(num_actions) - def index_with_mask(array, mask): - return array[mask].reshape(*array.shape[:-1]) + def index_with_mask(array, mask): + return array[mask].reshape(*array.shape[:-1]) - # Note: Normally log(softmax) is not a good idea because it's not - # numerically stable. However, in this test we have well-behaved values. - ground_truth_v = index_with_mask( - np.log(_softmax(policy_logits)), action_index_mask) + # Note: Normally log(softmax) is not a good idea because it's not + # numerically stable. However, in this test we have well-behaved + # values. + ground_truth_v = index_with_mask( + np.log(_softmax(policy_logits)), action_index_mask) - with self.test_session() as session: - self.assertAllClose(ground_truth_v, session.run(action_log_probs_tensor)) + with self.test_session() as session: + self.assertAllClose(ground_truth_v, + session.run(action_log_probs_tensor)) class VtraceTest(tf.test.TestCase, parameterized.TestCase): - - @parameterized.named_parameters(('Batch1', 1), ('Batch5', 5)) - def test_vtrace(self, batch_size): - """Tests V-trace against ground truth data calculated in python.""" - seq_len = 5 - - # Create log_rhos such that rho will span from near-zero to above the - # clipping thresholds. In particular, calculate log_rhos in [-2.5, 2.5), - # so that rho is in approx [0.08, 12.2). - log_rhos = _shaped_arange(seq_len, batch_size) / (batch_size * seq_len) - log_rhos = 5 * (log_rhos - 0.5) # [0.0, 1.0) -> [-2.5, 2.5). - values = { - 'log_rhos': log_rhos, - # T, B where B_i: [0.9 / (i+1)] * T - 'discounts': - np.array([[0.9 / (b + 1) - for b in range(batch_size)] - for _ in range(seq_len)]), - 'rewards': - _shaped_arange(seq_len, batch_size), - 'values': - _shaped_arange(seq_len, batch_size) / batch_size, - 'bootstrap_value': - _shaped_arange(batch_size) + 1.0, - 'clip_rho_threshold': - 3.7, - 'clip_pg_rho_threshold': - 2.2, - } - - output = vtrace.from_importance_weights(**values) - - with self.test_session() as session: - output_v = session.run(output) - - ground_truth_v = _ground_truth_calculation(**values) - for a, b in zip(ground_truth_v, output_v): - self.assertAllClose(a, b) - - @parameterized.named_parameters(('Batch1', 1), ('Batch2', 2)) - def test_vtrace_from_logits(self, batch_size): - """Tests V-trace calculated from logits.""" - seq_len = 5 - num_actions = 3 - clip_rho_threshold = None # No clipping. - clip_pg_rho_threshold = None # No clipping. - - # Intentionally leaving shapes unspecified to test if V-trace can - # deal with that. - placeholders = { - # T, B, NUM_ACTIONS - 'behaviour_policy_logits': - tf.placeholder(dtype=tf.float32, shape=[None, None, None]), - # T, B, NUM_ACTIONS - 'target_policy_logits': - tf.placeholder(dtype=tf.float32, shape=[None, None, None]), - 'actions': - tf.placeholder(dtype=tf.int32, shape=[None, None]), - 'discounts': - tf.placeholder(dtype=tf.float32, shape=[None, None]), - 'rewards': - tf.placeholder(dtype=tf.float32, shape=[None, None]), - 'values': - tf.placeholder(dtype=tf.float32, shape=[None, None]), - 'bootstrap_value': - tf.placeholder(dtype=tf.float32, shape=[None]), - } - - from_logits_output = vtrace.from_logits( - clip_rho_threshold=clip_rho_threshold, - clip_pg_rho_threshold=clip_pg_rho_threshold, - **placeholders) - - target_log_probs = vtrace.log_probs_from_logits_and_actions( - placeholders['target_policy_logits'], placeholders['actions']) - behaviour_log_probs = vtrace.log_probs_from_logits_and_actions( - placeholders['behaviour_policy_logits'], placeholders['actions']) - log_rhos = target_log_probs - behaviour_log_probs - ground_truth = (log_rhos, behaviour_log_probs, target_log_probs) - - values = { - 'behaviour_policy_logits': - _shaped_arange(seq_len, batch_size, num_actions), - 'target_policy_logits': - _shaped_arange(seq_len, batch_size, num_actions), - 'actions': - np.random.randint(0, num_actions - 1, size=(seq_len, batch_size)), - 'discounts': - np.array( # T, B where B_i: [0.9 / (i+1)] * T - [[0.9 / (b + 1) - for b in range(batch_size)] + @parameterized.named_parameters(('Batch1', 1), ('Batch5', 5)) + def test_vtrace(self, batch_size): + """Tests V-trace against ground truth data calculated in python.""" + seq_len = 5 + + # Create log_rhos such that rho will span from near-zero to above the + # clipping thresholds. In particular, calculate log_rhos in + # [-2.5, 2.5), + # so that rho is in approx [0.08, 12.2). + log_rhos = _shaped_arange(seq_len, batch_size) / (batch_size * seq_len) + log_rhos = 5 * (log_rhos - 0.5) # [0.0, 1.0) -> [-2.5, 2.5). + values = { + 'log_rhos': log_rhos, + # T, B where B_i: [0.9 / (i+1)] * T + 'discounts': np.array([[0.9 / (b + 1) for b in range(batch_size)] + for _ in range(seq_len)]), + 'rewards': _shaped_arange(seq_len, batch_size), + 'values': _shaped_arange(seq_len, batch_size) / batch_size, + 'bootstrap_value': _shaped_arange(batch_size) + 1.0, + 'clip_rho_threshold': 3.7, + 'clip_pg_rho_threshold': 2.2, + } + + output = vtrace.from_importance_weights(**values) + + with self.test_session() as session: + output_v = session.run(output) + + ground_truth_v = _ground_truth_calculation(**values) + for a, b in zip(ground_truth_v, output_v): + self.assertAllClose(a, b) + + @parameterized.named_parameters(('Batch1', 1), ('Batch2', 2)) + def test_vtrace_from_logits(self, batch_size): + """Tests V-trace calculated from logits.""" + seq_len = 5 + num_actions = 3 + clip_rho_threshold = None # No clipping. + clip_pg_rho_threshold = None # No clipping. + + # Intentionally leaving shapes unspecified to test if V-trace can + # deal with that. + placeholders = { + # T, B, NUM_ACTIONS + 'behaviour_policy_logits': tf.placeholder( + dtype=tf.float32, shape=[None, None, None]), + # T, B, NUM_ACTIONS + 'target_policy_logits': tf.placeholder( + dtype=tf.float32, shape=[None, None, None]), + 'actions': tf.placeholder(dtype=tf.int32, shape=[None, None]), + 'discounts': tf.placeholder(dtype=tf.float32, shape=[None, None]), + 'rewards': tf.placeholder(dtype=tf.float32, shape=[None, None]), + 'values': tf.placeholder(dtype=tf.float32, shape=[None, None]), + 'bootstrap_value': tf.placeholder(dtype=tf.float32, shape=[None]), + } + + from_logits_output = vtrace.from_logits( + clip_rho_threshold=clip_rho_threshold, + clip_pg_rho_threshold=clip_pg_rho_threshold, + **placeholders) + + target_log_probs = vtrace.log_probs_from_logits_and_actions( + placeholders['target_policy_logits'], placeholders['actions']) + behaviour_log_probs = vtrace.log_probs_from_logits_and_actions( + placeholders['behaviour_policy_logits'], placeholders['actions']) + log_rhos = target_log_probs - behaviour_log_probs + ground_truth = (log_rhos, behaviour_log_probs, target_log_probs) + + values = { + 'behaviour_policy_logits': _shaped_arange(seq_len, batch_size, + num_actions), + 'target_policy_logits': _shaped_arange(seq_len, batch_size, + num_actions), + 'actions': np.random.randint( + 0, num_actions - 1, size=(seq_len, batch_size)), + 'discounts': np.array( # T, B where B_i: [0.9 / (i+1)] * T + [[0.9 / (b + 1) for b in range(batch_size)] for _ in range(seq_len)]), - 'rewards': - _shaped_arange(seq_len, batch_size), - 'values': - _shaped_arange(seq_len, batch_size) / batch_size, - 'bootstrap_value': - _shaped_arange(batch_size) + 1.0, # B - } - - feed_dict = {placeholders[k]: v for k, v in values.items()} - with self.test_session() as session: - from_logits_output_v = session.run( - from_logits_output, feed_dict=feed_dict) - (ground_truth_log_rhos, ground_truth_behaviour_action_log_probs, - ground_truth_target_action_log_probs) = session.run( - ground_truth, feed_dict=feed_dict) - - # Calculate V-trace using the ground truth logits. - from_iw = vtrace.from_importance_weights( - log_rhos=ground_truth_log_rhos, - discounts=values['discounts'], - rewards=values['rewards'], - values=values['values'], - bootstrap_value=values['bootstrap_value'], - clip_rho_threshold=clip_rho_threshold, - clip_pg_rho_threshold=clip_pg_rho_threshold) - - with self.test_session() as session: - from_iw_v = session.run(from_iw) - - self.assertAllClose(from_iw_v.vs, from_logits_output_v.vs) - self.assertAllClose(from_iw_v.pg_advantages, - from_logits_output_v.pg_advantages) - self.assertAllClose(ground_truth_behaviour_action_log_probs, - from_logits_output_v.behaviour_action_log_probs) - self.assertAllClose(ground_truth_target_action_log_probs, - from_logits_output_v.target_action_log_probs) - self.assertAllClose(ground_truth_log_rhos, from_logits_output_v.log_rhos) - - def test_higher_rank_inputs_for_importance_weights(self): - """Checks support for additional dimensions in inputs.""" - placeholders = { - 'log_rhos': tf.placeholder(dtype=tf.float32, shape=[None, None, 1]), - 'discounts': tf.placeholder(dtype=tf.float32, shape=[None, None, 1]), - 'rewards': tf.placeholder(dtype=tf.float32, shape=[None, None, 42]), - 'values': tf.placeholder(dtype=tf.float32, shape=[None, None, 42]), - 'bootstrap_value': tf.placeholder(dtype=tf.float32, shape=[None, 42]) - } - output = vtrace.from_importance_weights(**placeholders) - self.assertEqual(output.vs.shape.as_list()[-1], 42) - - def test_inconsistent_rank_inputs_for_importance_weights(self): - """Test one of many possible errors in shape of inputs.""" - placeholders = { - 'log_rhos': tf.placeholder(dtype=tf.float32, shape=[None, None, 1]), - 'discounts': tf.placeholder(dtype=tf.float32, shape=[None, None, 1]), - 'rewards': tf.placeholder(dtype=tf.float32, shape=[None, None, 42]), - 'values': tf.placeholder(dtype=tf.float32, shape=[None, None, 42]), - # Should be [None, 42]. - 'bootstrap_value': tf.placeholder(dtype=tf.float32, shape=[None]) - } - with self.assertRaisesRegexp(ValueError, 'must have rank 2'): - vtrace.from_importance_weights(**placeholders) + 'rewards': _shaped_arange(seq_len, batch_size), + 'values': _shaped_arange(seq_len, batch_size) / batch_size, + 'bootstrap_value': _shaped_arange(batch_size) + 1.0, # B + } + + feed_dict = {placeholders[k]: v for k, v in values.items()} + with self.test_session() as session: + from_logits_output_v = session.run( + from_logits_output, feed_dict=feed_dict) + (ground_truth_log_rhos, ground_truth_behaviour_action_log_probs, + ground_truth_target_action_log_probs) = session.run( + ground_truth, feed_dict=feed_dict) + + # Calculate V-trace using the ground truth logits. + from_iw = vtrace.from_importance_weights( + log_rhos=ground_truth_log_rhos, + discounts=values['discounts'], + rewards=values['rewards'], + values=values['values'], + bootstrap_value=values['bootstrap_value'], + clip_rho_threshold=clip_rho_threshold, + clip_pg_rho_threshold=clip_pg_rho_threshold) + + with self.test_session() as session: + from_iw_v = session.run(from_iw) + + self.assertAllClose(from_iw_v.vs, from_logits_output_v.vs) + self.assertAllClose(from_iw_v.pg_advantages, + from_logits_output_v.pg_advantages) + self.assertAllClose(ground_truth_behaviour_action_log_probs, + from_logits_output_v.behaviour_action_log_probs) + self.assertAllClose(ground_truth_target_action_log_probs, + from_logits_output_v.target_action_log_probs) + self.assertAllClose(ground_truth_log_rhos, + from_logits_output_v.log_rhos) + + def test_higher_rank_inputs_for_importance_weights(self): + """Checks support for additional dimensions in inputs.""" + placeholders = { + 'log_rhos': tf.placeholder( + dtype=tf.float32, shape=[None, None, 1]), + 'discounts': tf.placeholder( + dtype=tf.float32, shape=[None, None, 1]), + 'rewards': tf.placeholder( + dtype=tf.float32, shape=[None, None, 42]), + 'values': tf.placeholder(dtype=tf.float32, shape=[None, None, 42]), + 'bootstrap_value': tf.placeholder( + dtype=tf.float32, shape=[None, 42]) + } + output = vtrace.from_importance_weights(**placeholders) + self.assertEqual(output.vs.shape.as_list()[-1], 42) + + def test_inconsistent_rank_inputs_for_importance_weights(self): + """Test one of many possible errors in shape of inputs.""" + placeholders = { + 'log_rhos': tf.placeholder( + dtype=tf.float32, shape=[None, None, 1]), + 'discounts': tf.placeholder( + dtype=tf.float32, shape=[None, None, 1]), + 'rewards': tf.placeholder( + dtype=tf.float32, shape=[None, None, 42]), + 'values': tf.placeholder(dtype=tf.float32, shape=[None, None, 42]), + # Should be [None, 42]. + 'bootstrap_value': tf.placeholder(dtype=tf.float32, shape=[None]) + } + with self.assertRaisesRegexp(ValueError, 'must have rank 2'): + vtrace.from_importance_weights(**placeholders) if __name__ == '__main__': - tf.test.main() + tf.test.main() diff --git a/python/ray/rllib/agents/ppo/appo_policy_graph.py b/python/ray/rllib/agents/ppo/appo_policy_graph.py index 212e52bd34ed..fbf30beddcec 100644 --- a/python/ray/rllib/agents/ppo/appo_policy_graph.py +++ b/python/ray/rllib/agents/ppo/appo_policy_graph.py @@ -199,7 +199,8 @@ def __init__(self, output_hidden_shape = action_space.nvec elif self.config["vtrace"]: raise UnsupportedSpaceException( - "Action space {} is not supported for APPO with VTrace.".format(action_space)) + "Action space {} is not supported for APPO with VTrace.". + format(action_space)) actions = tf.placeholder(tf.int64, actions_shape, name="ac") dones = tf.placeholder(tf.bool, [None], name="dones") @@ -224,7 +225,8 @@ def __init__(self, # Setup the policy dist_class, logit_dim = ModelCatalog.get_action_dist( - action_space, self.config["model"], + action_space, + self.config["model"], dist_type=self.config["dist_type"]) prev_actions = ModelCatalog.get_action_placeholder(action_space) prev_rewards = tf.placeholder(tf.float32, [None], name="prev_reward") @@ -306,18 +308,18 @@ def make_time_major(tensor, drop_last=False): self.loss = VTraceSurrogateLoss( actions=make_time_major(loss_actions, drop_last=True), - prev_actions_logp=make_time_major(prev_action_dist.logp( - actions), drop_last=True), - actions_logp=make_time_major(action_dist.logp(actions), - drop_last=True), + prev_actions_logp=make_time_major( + prev_action_dist.logp(actions), drop_last=True), + actions_logp=make_time_major( + action_dist.logp(actions), drop_last=True), action_kl=prev_action_dist.kl(action_dist), - actions_entropy=make_time_major(action_dist.entropy(), - drop_last=True), + actions_entropy=make_time_major( + action_dist.entropy(), drop_last=True), dones=make_time_major(dones, drop_last=True), - behaviour_logits=make_time_major(unpacked_behaviour_logits, - drop_last=True), - target_logits=make_time_major(unpacked_outputs, - drop_last=True), + behaviour_logits=make_time_major( + unpacked_behaviour_logits, drop_last=True), + target_logits=make_time_major( + unpacked_outputs, drop_last=True), discount=config["gamma"], rewards=make_time_major(rewards, drop_last=True), values=make_time_major(values, drop_last=True), @@ -334,11 +336,9 @@ def make_time_major(tensor, drop_last=False): self.loss = PPOSurrogateLoss( prev_actions_logp=make_time_major( prev_action_dist.logp(actions)), - actions_logp=make_time_major( - action_dist.logp(actions)), + actions_logp=make_time_major(action_dist.logp(actions)), action_kl=prev_action_dist.kl(action_dist), - actions_entropy=make_time_major( - action_dist.entropy()), + actions_entropy=make_time_major(action_dist.entropy()), values=make_time_major(values), valid_mask=make_time_major(mask), advantages=make_time_major(adv_ph), @@ -357,17 +357,16 @@ def make_time_major(tensor, drop_last=False): for i, kl in enumerate(kls): self.KL_stats.update({ - f"mean_KL_{i}": tf.reduce_mean(kl), - f"max_KL_{i}": tf.reduce_max(kl), - f"median_KL_{i}": tf.contrib.distributions.percentile( - kl, 50.0), + "mean_KL_{}".format(i): tf.reduce_mean(kl), + "max_KL_{}".format(i): tf.reduce_max(kl), + "median_KL_{}".format(i): tf.contrib.distributions. + percentile(kl, 50.0), }) else: self.KL_stats = { "mean_KL": tf.reduce_mean(kls[0]), "max_KL": tf.reduce_max(kls[0]), - "median_KL": tf.contrib.distributions.percentile( - kls[0], 50.0), + "median_KL": tf.contrib.distributions.percentile(kls[0], 50.0), } # Initialize TFPolicyGraph @@ -405,20 +404,23 @@ def make_time_major(tensor, drop_last=False): self.sess.run(tf.global_variables_initializer()) - values_batched = make_time_major(values, drop_last=self.config["vtrace"]) - self.stats_fetches = {"stats": { - "model_loss": self.model.loss(), - "cur_lr": tf.cast(self.cur_lr, tf.float64), - "policy_loss": self.loss.pi_loss, - "entropy": self.loss.entropy, - "grad_gnorm": tf.global_norm(self._grads), - "var_gnorm": tf.global_norm(self.var_list), - "vf_loss": self.loss.vf_loss, - "vf_explained_var": explained_variance( - tf.reshape(self.loss.value_targets, [-1]), - tf.reshape(values_batched, [-1])), - **self.KL_stats, - }, "kl": self.loss.mean_kl} + values_batched = make_time_major( + values, drop_last=self.config["vtrace"]) + self.stats_fetches = { + "stats": dict({ + "model_loss": self.model.loss(), + "cur_lr": tf.cast(self.cur_lr, tf.float64), + "policy_loss": self.loss.pi_loss, + "entropy": self.loss.entropy, + "grad_gnorm": tf.global_norm(self._grads), + "var_gnorm": tf.global_norm(self.var_list), + "vf_loss": self.loss.vf_loss, + "vf_explained_var": explained_variance( + tf.reshape(self.loss.value_targets, [-1]), + tf.reshape(values_batched, [-1])), + "kl": self.loss.mean_kl, + }, **self.KL_stats) + } def optimizer(self): if self.config["opt_type"] == "adam": diff --git a/python/ray/rllib/models/action_dist.py b/python/ray/rllib/models/action_dist.py index d43c34d7ef4a..138fd9f8a6a8 100644 --- a/python/ray/rllib/models/action_dist.py +++ b/python/ray/rllib/models/action_dist.py @@ -125,16 +125,15 @@ def logp(self, actions): # If tensor is provided, unstack it into list if isinstance(actions, tf.Tensor): actions = tf.unstack(actions, axis=1) - logps = tf.stack([cat.logp(act) - for cat, act in zip(self.cats, actions)]) + logps = tf.stack( + [cat.logp(act) for cat, act in zip(self.cats, actions)]) return tf.reduce_sum(logps, axis=0) def entropy(self): return tf.stack([cat.entropy() for cat in self.cats], axis=1) def kl(self, other): - return [cat.kl(oth_cat) - for cat, oth_cat in zip(self.cats, other.cats)] + return [cat.kl(oth_cat) for cat, oth_cat in zip(self.cats, other.cats)] def _build_sample_op(self): return tf.stack([cat.sample() for cat in self.cats], axis=1) @@ -158,8 +157,8 @@ def __init__(self, inputs): def logp(self, x): return (-0.5 * tf.reduce_sum( tf.square((x - self.mean) / self.std), reduction_indices=[1]) - - 0.5 * np.log(2.0 * np.pi) * tf.to_float(tf.shape(x)[1]) - - tf.reduce_sum(self.log_std, reduction_indices=[1])) + 0.5 * np.log(2.0 * np.pi) * tf.to_float(tf.shape(x)[1]) - + tf.reduce_sum(self.log_std, reduction_indices=[1])) @override(ActionDistribution) def kl(self, other): From aa50f98b8409cf9e3f0ee50fabc5a8fb356b6189 Mon Sep 17 00:00:00 2001 From: Eric Liang Date: Mon, 18 Feb 2019 11:29:34 -0800 Subject: [PATCH 28/33] kl --- python/ray/rllib/agents/ppo/appo_policy_graph.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/ray/rllib/agents/ppo/appo_policy_graph.py b/python/ray/rllib/agents/ppo/appo_policy_graph.py index fbf30beddcec..963f5165af87 100644 --- a/python/ray/rllib/agents/ppo/appo_policy_graph.py +++ b/python/ray/rllib/agents/ppo/appo_policy_graph.py @@ -407,6 +407,7 @@ def make_time_major(tensor, drop_last=False): values_batched = make_time_major( values, drop_last=self.config["vtrace"]) self.stats_fetches = { + "kl": self.loss.mean_kl, "stats": dict({ "model_loss": self.model.loss(), "cur_lr": tf.cast(self.cur_lr, tf.float64), @@ -418,8 +419,7 @@ def make_time_major(tensor, drop_last=False): "vf_explained_var": explained_variance( tf.reshape(self.loss.value_targets, [-1]), tf.reshape(values_batched, [-1])), - "kl": self.loss.mean_kl, - }, **self.KL_stats) + }, **self.KL_stats), } def optimizer(self): From d32d253c63ca73944ada0dc887a34826e4f37b6f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleksandar=20Milovanovi=C4=87?= Date: Tue, 19 Feb 2019 09:15:21 +0100 Subject: [PATCH 29/33] removed dist_type as it is actually not needed for IMPALA --- python/ray/rllib/agents/impala/impala.py | 3 --- python/ray/rllib/agents/impala/vtrace_policy_graph.py | 3 +-- python/ray/rllib/agents/ppo/appo_policy_graph.py | 3 +-- 3 files changed, 2 insertions(+), 7 deletions(-) diff --git a/python/ray/rllib/agents/impala/impala.py b/python/ray/rllib/agents/impala/impala.py index 44d1d1057675..94294a1fc76d 100644 --- a/python/ray/rllib/agents/impala/impala.py +++ b/python/ray/rllib/agents/impala/impala.py @@ -72,9 +72,6 @@ # max number of workers to broadcast one set of weights to "broadcast_interval": 1, - # Actions are chosen based on this distribution, if provided - "dist_type": None, - # Learning params. "grad_clip": 40.0, # either "adam" or "rmsprop" diff --git a/python/ray/rllib/agents/impala/vtrace_policy_graph.py b/python/ray/rllib/agents/impala/vtrace_policy_graph.py index 0ce178e71d1c..80141cb9d303 100644 --- a/python/ray/rllib/agents/impala/vtrace_policy_graph.py +++ b/python/ray/rllib/agents/impala/vtrace_policy_graph.py @@ -151,8 +151,7 @@ def __init__(self, # Setup the policy dist_class, logit_dim = ModelCatalog.get_action_dist( action_space, - self.config["model"], - dist_type=self.config["dist_type"]) + self.config["model"]) prev_actions = ModelCatalog.get_action_placeholder(action_space) prev_rewards = tf.placeholder(tf.float32, [None], name="prev_reward") self.model = ModelCatalog.get_model( diff --git a/python/ray/rllib/agents/ppo/appo_policy_graph.py b/python/ray/rllib/agents/ppo/appo_policy_graph.py index 963f5165af87..481533dc3b51 100644 --- a/python/ray/rllib/agents/ppo/appo_policy_graph.py +++ b/python/ray/rllib/agents/ppo/appo_policy_graph.py @@ -226,8 +226,7 @@ def __init__(self, # Setup the policy dist_class, logit_dim = ModelCatalog.get_action_dist( action_space, - self.config["model"], - dist_type=self.config["dist_type"]) + self.config["model"]) prev_actions = ModelCatalog.get_action_placeholder(action_space) prev_rewards = tf.placeholder(tf.float32, [None], name="prev_reward") self.model = ModelCatalog.get_model( From 967db5ca49b2c49fbc5c8a017ad2bd0aaeb43211 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleksandar=20Milovanovi=C4=87?= Date: Tue, 19 Feb 2019 10:14:38 +0100 Subject: [PATCH 30/33] fixing issue with new gym version --- python/ray/rllib/agents/impala/vtrace_policy_graph.py | 3 ++- python/ray/rllib/agents/ppo/appo_policy_graph.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/python/ray/rllib/agents/impala/vtrace_policy_graph.py b/python/ray/rllib/agents/impala/vtrace_policy_graph.py index 80141cb9d303..0814bbfeb76b 100644 --- a/python/ray/rllib/agents/impala/vtrace_policy_graph.py +++ b/python/ray/rllib/agents/impala/vtrace_policy_graph.py @@ -8,6 +8,7 @@ import gym import ray +import numpy as np import tensorflow as tf from ray.rllib.agents.impala import vtrace from ray.rllib.evaluation.policy_graph import PolicyGraph @@ -127,7 +128,7 @@ def __init__(self, gym.spaces.multi_discrete.MultiDiscrete): is_multidiscrete = True actions_shape = [None, len(action_space.nvec)] - output_hidden_shape = action_space.nvec + output_hidden_shape = action_space.nvec.astype(np.int32) else: raise UnsupportedSpaceException( "Action space {} is not supported for IMPALA.".format( diff --git a/python/ray/rllib/agents/ppo/appo_policy_graph.py b/python/ray/rllib/agents/ppo/appo_policy_graph.py index 481533dc3b51..c580c4616cfb 100644 --- a/python/ray/rllib/agents/ppo/appo_policy_graph.py +++ b/python/ray/rllib/agents/ppo/appo_policy_graph.py @@ -6,6 +6,7 @@ from __future__ import division from __future__ import print_function +import numpy as np import tensorflow as tf import logging import gym @@ -196,7 +197,7 @@ def __init__(self, gym.spaces.multi_discrete.MultiDiscrete): is_multidiscrete = True actions_shape = [None, len(action_space.nvec)] - output_hidden_shape = action_space.nvec + output_hidden_shape = action_space.nvec.astype(np.int32) elif self.config["vtrace"]: raise UnsupportedSpaceException( "Action space {} is not supported for APPO with VTrace.". From 9584a7c04623562df21cec73e6bed84a7968ef1d Mon Sep 17 00:00:00 2001 From: Eric Liang Date: Wed, 20 Feb 2019 09:53:54 -0800 Subject: [PATCH 31/33] lint --- python/ray/rllib/agents/impala/vtrace_policy_graph.py | 3 +-- python/ray/rllib/agents/ppo/appo_policy_graph.py | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/python/ray/rllib/agents/impala/vtrace_policy_graph.py b/python/ray/rllib/agents/impala/vtrace_policy_graph.py index 0814bbfeb76b..700f00fa5326 100644 --- a/python/ray/rllib/agents/impala/vtrace_policy_graph.py +++ b/python/ray/rllib/agents/impala/vtrace_policy_graph.py @@ -151,8 +151,7 @@ def __init__(self, # Setup the policy dist_class, logit_dim = ModelCatalog.get_action_dist( - action_space, - self.config["model"]) + action_space, self.config["model"]) prev_actions = ModelCatalog.get_action_placeholder(action_space) prev_rewards = tf.placeholder(tf.float32, [None], name="prev_reward") self.model = ModelCatalog.get_model( diff --git a/python/ray/rllib/agents/ppo/appo_policy_graph.py b/python/ray/rllib/agents/ppo/appo_policy_graph.py index c580c4616cfb..a52982d6fdb9 100644 --- a/python/ray/rllib/agents/ppo/appo_policy_graph.py +++ b/python/ray/rllib/agents/ppo/appo_policy_graph.py @@ -226,8 +226,7 @@ def __init__(self, # Setup the policy dist_class, logit_dim = ModelCatalog.get_action_dist( - action_space, - self.config["model"]) + action_space, self.config["model"]) prev_actions = ModelCatalog.get_action_placeholder(action_space) prev_rewards = tf.placeholder(tf.float32, [None], name="prev_reward") self.model = ModelCatalog.get_model( From 8999621d173362d0e413fc3edaacf3d5438b5aa9 Mon Sep 17 00:00:00 2001 From: Eric Liang Date: Wed, 20 Feb 2019 10:05:10 -0800 Subject: [PATCH 32/33] fix multigpu test --- .../agents/impala/vtrace_policy_graph.py | 27 +++++++++---------- 1 file changed, 12 insertions(+), 15 deletions(-) diff --git a/python/ray/rllib/agents/impala/vtrace_policy_graph.py b/python/ray/rllib/agents/impala/vtrace_policy_graph.py index 700f00fa5326..a1d280df1d11 100644 --- a/python/ray/rllib/agents/impala/vtrace_policy_graph.py +++ b/python/ray/rllib/agents/impala/vtrace_policy_graph.py @@ -111,9 +111,18 @@ def __init__(self, self.sess = tf.get_default_session() self.grads = None - is_multidiscrete = False - output_hidden_shape = None - actions_shape = [None] + if isinstance(action_space, gym.spaces.Discrete): + is_multidiscrete = False + actions_shape = [None] + output_hidden_shape = [action_space.n] + elif isinstance(action_space, gym.spaces.multi_discrete.MultiDiscrete): + is_multidiscrete = True + actions_shape = [None, len(action_space.nvec)] + output_hidden_shape = action_space.nvec.astype(np.int32) + else: + raise UnsupportedSpaceException( + "Action space {} is not supported for IMPALA.".format( + action_space)) # Create input placeholders if existing_inputs: @@ -122,18 +131,6 @@ def __init__(self, existing_state_in = existing_inputs[7:-1] existing_seq_lens = existing_inputs[-1] else: - if isinstance(action_space, gym.spaces.Discrete): - output_hidden_shape = [action_space.n] - elif isinstance(action_space, - gym.spaces.multi_discrete.MultiDiscrete): - is_multidiscrete = True - actions_shape = [None, len(action_space.nvec)] - output_hidden_shape = action_space.nvec.astype(np.int32) - else: - raise UnsupportedSpaceException( - "Action space {} is not supported for IMPALA.".format( - action_space)) - actions = tf.placeholder(tf.int64, actions_shape, name="ac") dones = tf.placeholder(tf.bool, [None], name="dones") rewards = tf.placeholder(tf.float32, [None], name="rewards") From eb18cff0b4e203bbdf8978e1af46ae8cc02a35b1 Mon Sep 17 00:00:00 2001 From: Eric Liang Date: Fri, 1 Mar 2019 14:27:08 -0800 Subject: [PATCH 33/33] fix tests --- .../ray/rllib/agents/ppo/appo_policy_graph.py | 28 ++++++++----------- .../ray/rllib/evaluation/policy_evaluator.py | 2 +- 2 files changed, 13 insertions(+), 17 deletions(-) diff --git a/python/ray/rllib/agents/ppo/appo_policy_graph.py b/python/ray/rllib/agents/ppo/appo_policy_graph.py index 4b09811bc3d9..362f93b0721a 100644 --- a/python/ray/rllib/agents/ppo/appo_policy_graph.py +++ b/python/ray/rllib/agents/ppo/appo_policy_graph.py @@ -169,9 +169,18 @@ def __init__(self, self.sess = tf.get_default_session() self.grads = None - is_multidiscrete = False - output_hidden_shape = None - actions_shape = [None] + if isinstance(action_space, gym.spaces.Discrete): + is_multidiscrete = False + actions_shape = [None] + output_hidden_shape = [action_space.n] + elif isinstance(action_space, gym.spaces.multi_discrete.MultiDiscrete): + is_multidiscrete = True + actions_shape = [None, len(action_space.nvec)] + output_hidden_shape = action_space.nvec.astype(np.int32) + else: + raise UnsupportedSpaceException( + "Action space {} is not supported for APPO.", + format(action_space)) # Policy network model dist_class, logit_dim = ModelCatalog.get_action_dist( @@ -191,18 +200,6 @@ def __init__(self, existing_state_in = existing_inputs[9:-1] existing_seq_lens = existing_inputs[-1] else: - if isinstance(action_space, gym.spaces.Discrete): - output_hidden_shape = [action_space.n] - elif isinstance(action_space, - gym.spaces.multi_discrete.MultiDiscrete): - is_multidiscrete = True - actions_shape = [None, len(action_space.nvec)] - output_hidden_shape = action_space.nvec.astype(np.int32) - elif self.config["vtrace"]: - raise UnsupportedSpaceException( - "Action space {} is not supported for APPO with VTrace.". - format(action_space)) - actions = tf.placeholder(tf.int64, actions_shape, name="ac") dones = tf.placeholder(tf.bool, [None], name="dones") rewards = tf.placeholder(tf.float32, [None], name="rewards") @@ -408,7 +405,6 @@ def make_time_major(tensor, drop_last=False): values, drop_last=self.config["vtrace"]) self.stats_fetches = { "stats": dict({ - "model_loss": self.model.loss(), "cur_lr": tf.cast(self.cur_lr, tf.float64), "policy_loss": self.loss.pi_loss, "entropy": self.loss.entropy, diff --git a/python/ray/rllib/evaluation/policy_evaluator.py b/python/ray/rllib/evaluation/policy_evaluator.py index 7a8d3234569d..71e2009a1882 100644 --- a/python/ray/rllib/evaluation/policy_evaluator.py +++ b/python/ray/rllib/evaluation/policy_evaluator.py @@ -658,7 +658,7 @@ def _build_policy_map(self, policy_dict, policy_config): return policy_map, preprocessors def __del__(self): - if isinstance(self.sampler, AsyncSampler): + if hasattr(self, "sampler") and isinstance(self.sampler, AsyncSampler): self.sampler.shutdown = True