def loss()

in lm_human_preferences/train_policy.py [0:0]

• 45 lines of code
• 4 McCabe index (conditional complexity)

    def loss(self, rollouts):
        values = rollouts['values']
        old_logprob = rollouts['logprobs']
        rewards = rollouts['rewards']
        with tf.name_scope('ppo_loss'):
            if self.hparams.ppo.whiten_rewards:
                rewards = utils.whiten(rewards, shift_mean=False)

            lastgaelam = 0
            advantages_reversed = []
            gen_length = self.hparams.task.response_length
            for t in reversed(range(gen_length)):
                nextvalues = values[:, t + 1] if t < gen_length - 1 else 0.0
                delta = rewards[:, t] + self.hparams.ppo.gamma * nextvalues - values[:, t]
                lastgaelam = delta + self.hparams.ppo.gamma * self.hparams.ppo.lam * lastgaelam
                advantages_reversed.append(lastgaelam)
            advantages = tf.stack(advantages_reversed[::-1], axis=1)
            returns = advantages + values

            advantages = utils.whiten(advantages)
            advantages = tf.stop_gradient(advantages)  # Shouldn't do anything, but better not to think about it

            outputs = self.policy.analyze_responses_op(rollouts['queries'], rollouts['responses'])

            vpred = outputs['values']
            vpredclipped = tf.clip_by_value(vpred, values - self.hparams.ppo.cliprange_value, values + self.hparams.ppo.cliprange_value)
            vf_losses1 = tf.square(vpred - returns)
            vf_losses2 = tf.square(vpredclipped - returns)
            vf_loss = .5 * tf.reduce_mean(tf.maximum(vf_losses1, vf_losses2))
            vf_clipfrac = tf.reduce_mean(tf.cast(tf.greater(vf_losses2, vf_losses1), tf.float32))

            logprob = outputs['logprobs']
            ratio = tf.exp(logprob - old_logprob)
            pg_losses = -advantages * ratio
            pg_losses2 = -advantages * tf.clip_by_value(ratio, 1.0 - self.hparams.ppo.cliprange, 1.0 + self.hparams.ppo.cliprange)
            pg_loss = tf.reduce_mean(tf.maximum(pg_losses, pg_losses2))
            pg_clipfrac = tf.reduce_mean(tf.cast(tf.greater(pg_losses2, pg_losses), tf.float32))

            loss = pg_loss + self.hparams.ppo.vf_coef * vf_loss

            entropy = tf.reduce_mean(outputs['entropies'])
            approxkl = .5 * tf.reduce_mean(tf.square(logprob - old_logprob))

            return_mean, return_var = tf.nn.moments(returns, axes=list(range(returns.shape.ndims)))
            value_mean, value_var = tf.nn.moments(values, axes=list(range(values.shape.ndims)))

            stats = dict(
                loss=dict(policy=pg_loss, value=vf_loss, total=loss),
                policy=dict(entropy=entropy, approxkl=approxkl, clipfrac=pg_clipfrac),
                returns=dict(mean=return_mean, var=return_var),
                val=dict(vpred=tf.reduce_mean(vpred), error=tf.reduce_mean((vpred - returns) ** 2),
                         clipfrac=vf_clipfrac, mean=value_mean, var=value_var)
            )
            return loss, utils.flatten_dict(stats, sep='/')