(self, **unused_kwargs)
return tl.Serial(
        tl.Fn('LossInput', LossInput, n_out=4),
        # Policy loss is expected to consume
        # (log_probs, advantages, old_log_probs, mask).
        SamplingAWRLoss(
            beta=self._beta, w_max=self._w_max, thresholds=self._thresholds,
            reweight=self._reweight,
            sampled_all_discrete=self._sample_all_discrete_actions)
    )
Policy loss.
Policy loss.
Policy loss .
def policy_loss(self, **unused_kwargs):
    """Policy loss."""
    def LossInput(dist_inputs, actions, q_values, act_log_probs, mask):  # pylint: disable=invalid-name
      """Calculates action log probabilities and normalizes advantages."""
      # (batch_size, n_samples, ...) -> (n_samples, batch_size, ...)
      q_values = jnp.swapaxes(q_values, 0, 1)
      mask = jnp.swapaxes(mask, 0, 1)
      actions = jnp.swapaxes(actions, 0, 1)
      act_log_probs = jnp.swapaxes(act_log_probs, 0, 1)

      # TODO(pkozakowski,lukaszkaiser): Try max here, or reweighting?
      if self._sample_all_discrete_actions:
        values = jnp.sum(q_values * jnp.exp(act_log_probs), axis=0)
      else:
        values = jnp.mean(q_values, axis=0)
      advantages = q_values - values  # Broadcasting values over n_samples
      advantages = self._preprocess_advantages(advantages)

      # Broadcast inputs and calculate log-probs
      dist_inputs = jnp.broadcast_to(
          dist_inputs, (self._q_value_n_samples,) + dist_inputs.shape)
      log_probs = self._policy_dist.log_prob(dist_inputs, actions)
      return (log_probs, advantages, act_log_probs, mask)

    return tl.Serial(
        tl.Fn('LossInput', LossInput, n_out=4),
        # Policy loss is expected to consume
        # (log_probs, advantages, old_log_probs, mask).
        SamplingAWRLoss(
            beta=self._beta, w_max=self._w_max, thresholds=self._thresholds,
            reweight=self._reweight,
            sampled_all_discrete=self._sample_all_discrete_actions)
    )
def policy_loss ( self , * * unused_kwargs ) : def LossInput ( dist_inputs , actions , q_values , act_log_probs , mask ) : # pylint: disable=invalid-name """Calculates action log probabilities and normalizes advantages.""" # (batch_size, n_samples, ...) -> (n_samples, batch_size, ...) q_values = jnp . swapaxes ( q_values , 0 , 1 ) mask = jnp . swapaxes ( mask , 0 , 1 ) actions = jnp . swapaxes ( actions , 0 , 1 ) act_log_probs = jnp . swapaxes ( act_log_probs , 0 , 1 ) # TODO(pkozakowski,lukaszkaiser): Try max here, or reweighting? if self . _sample_all_discrete_actions : values = jnp . sum ( q_values * jnp . exp ( act_log_probs ) , axis = 0 ) else : values = jnp . mean ( q_values , axis = 0 ) advantages = q_values - values # Broadcasting values over n_samples advantages = self . _preprocess_advantages ( advantages ) # Broadcast inputs and calculate log-probs dist_inputs = jnp . broadcast_to ( dist_inputs , ( self . _q_value_n_samples , ) + dist_inputs . shape ) log_probs = self . _policy_dist . log_prob ( dist_inputs , actions ) return ( log_probs , advantages , act_log_probs , mask ) return tl . Serial ( tl . Fn ( 'LossInput' , LossInput , n_out = 4 ) , # Policy loss is expected to consume # (log_probs, advantages, old_log_probs, mask). SamplingAWRLoss ( beta = self . _beta , w_max = self . _w_max , thresholds = self . _thresholds , reweight = self . _reweight , sampled_all_discrete = self . _sample_all_discrete_actions ) )