(self)
return tl.Fn('PPOJointLoss', f)
Joint policy and value loss.
Joint policy and value loss.
Joint policy and value loss .
def joint_loss(self):
    """Joint policy and value loss."""
    def f(dist_inputs, values, returns, dones, rewards,
          actions, old_log_probs, mask):
      """Definition of the Proximal Policy Optimization loss."""
      del mask  # TODO(lukaszkaiser): make PPO work with Transformer
      # We have dist_inputs of the shape float32[128,1,18]
      assert len(dist_inputs.shape) == 3, (
          f'dist_inputs.shape was {dist_inputs.shape}'
          f'but expected length of the tensor shape is 3')
      # values of the shape float32[128,1,1]
      # returns of the shape float32[128,1,1]
      # dones of the shape int32[128,1,1]
      # rewards of the shape float32[128,1,1]
      # and old_log_probs of the shape float32[128,1]
      assert values.shape == returns.shape, (
          f'values.shape was {values.shape}'
          f'returns.shape was {returns.shape}')
      assert values.shape == dones.shape, (
          f'values.shape was {values.shape}'
          f'returns.shape was {dones.shape}')
      assert rewards.shape == dones.shape, (
          f'values.shape was {values.shape}'
          f'returns.shape was {dones.shape}')
      assert returns.shape[0:2] == old_log_probs.shape, (
          f'returns.shape was {returns.shape}'
          f'old_log_probs.shape was {old_log_probs.shape}')

      # actions is a tensor of the shape int32[128,1] in the case
      # of discrete actions and float32[128,1,6] in the case of
      # half-cheetah and other continuous actions
      # actions agree with returns/values on the first two coordinates
      # meaning batch and time
      assert actions.shape[0:2] == returns.shape[0:2], (
          f'actions.shape was {actions.shape} and '
          f'returns.shape was {returns.shape}')

      ppo_objective = rl_layers.PPOObjective(
          dist_inputs, stop_gradient(values), returns, dones, rewards,
          actions, old_log_probs,
          log_prob_fun=self._policy_dist.log_prob,
          epsilon=self._epsilon,
          normalize_advantages=self._normalize_advantages)

      # we insist that ppo_objective is a vector of shape [128,1]
      assert len(ppo_objective.shape) == 2, (
          f'ppo_objective was {ppo_objective}')
      # which agrees with returns/values/actions on the first two coordinates
      assert ppo_objective.shape[0:2] == values.shape[0:2], (
          f'ppo_objective.shape was {ppo_objective.shape} and '
          f'values.shape was {values.shape}')

      entropy_loss = rl_layers.EntropyLoss(
          dist_inputs,
          distribution=self._policy_dist,
          coeff=self._entropy_coeff,
      )

      assert jnp.ndim(entropy_loss) == 0, f'entropy_loss was {entropy_loss}'

      l2_value_loss = rl_layers.ValueLoss(
          values, returns, value_loss_coeff=self._value_loss_coeff)

      assert jnp.ndim(l2_value_loss) == 0, f'l2_value_loss was {l2_value_loss}'

      return -ppo_objective.mean() + l2_value_loss - entropy_loss

    return tl.Fn('PPOJointLoss', f)
def joint_loss ( self ) : def f ( dist_inputs , values , returns , dones , rewards , actions , old_log_probs , mask ) : """Definition of the Proximal Policy Optimization loss.""" del mask # TODO(lukaszkaiser): make PPO work with Transformer # We have dist_inputs of the shape float32[128,1,18] assert len ( dist_inputs . shape ) == 3 , ( f'dist_inputs.shape was {dist_inputs.shape}' f'but expected length of the tensor shape is 3' ) # values of the shape float32[128,1,1] # returns of the shape float32[128,1,1] # dones of the shape int32[128,1,1] # rewards of the shape float32[128,1,1] # and old_log_probs of the shape float32[128,1] assert values . shape == returns . shape , ( f'values.shape was {values.shape}' f'returns.shape was {returns.shape}' ) assert values . shape == dones . shape , ( f'values.shape was {values.shape}' f'returns.shape was {dones.shape}' ) assert rewards . shape == dones . shape , ( f'values.shape was {values.shape}' f'returns.shape was {dones.shape}' ) assert returns . shape [ 0 : 2 ] == old_log_probs . shape , ( f'returns.shape was {returns.shape}' f'old_log_probs.shape was {old_log_probs.shape}' ) # actions is a tensor of the shape int32[128,1] in the case # of discrete actions and float32[128,1,6] in the case of # half-cheetah and other continuous actions # actions agree with returns/values on the first two coordinates # meaning batch and time assert actions . shape [ 0 : 2 ] == returns . shape [ 0 : 2 ] , ( f'actions.shape was {actions.shape} and ' f'returns.shape was {returns.shape}' ) ppo_objective = rl_layers . PPOObjective ( dist_inputs , stop_gradient ( values ) , returns , dones , rewards , actions , old_log_probs , log_prob_fun = self . _policy_dist . log_prob , epsilon = self . _epsilon , normalize_advantages = self . _normalize_advantages ) # we insist that ppo_objective is a vector of shape [128,1] assert len ( ppo_objective . shape ) == 2 , ( f'ppo_objective was {ppo_objective}' ) # which agrees with returns/values/actions on the first two coordinates assert ppo_objective . shape [ 0 : 2 ] == values . shape [ 0 : 2 ] , ( f'ppo_objective.shape was {ppo_objective.shape} and ' f'values.shape was {values.shape}' ) entropy_loss = rl_layers . EntropyLoss ( dist_inputs , distribution = self . _policy_dist , coeff = self . _entropy_coeff , ) assert jnp . ndim ( entropy_loss ) == 0 , f'entropy_loss was {entropy_loss}' l2_value_loss = rl_layers . ValueLoss ( values , returns , value_loss_coeff = self . _value_loss_coeff ) assert jnp . ndim ( l2_value_loss ) == 0 , f'l2_value_loss was {l2_value_loss}' return - ppo_objective . mean ( ) + l2_value_loss - entropy_loss return tl . Fn ( 'PPOJointLoss' , f )