def step()

in lm_human_preferences/train_policy.py [0:0]

• 27 lines of code
• 2 McCabe index (conditional complexity)

    def step(self):
        step_started_at = time.time()

        queries = self.sample_queries()
        rollouts = self.policy.respond(queries, length=self.hparams.task.response_length)

        responses = rollouts['responses']
        logprobs = rollouts['logprobs']
        rollouts['queries'] = queries
        ref_logprobs = self.ref_policy.analyze_responses(queries, responses)['logprobs']
        scores, postprocessed_responses, score_stats = self.score_fn(queries, responses)

        rewards, non_score_reward, kl_coef = self.compute_rewards(
            scores=scores,
            logprobs=logprobs,
            ref_logprobs=ref_logprobs)
        rollouts['rewards'] = rewards

        train_stats = self.train(rollouts=rollouts)

        _, stats = self.record_step_stats(
            scores=scores, logprobs=logprobs, ref_logprobs=ref_logprobs, non_score_reward=non_score_reward,
            train_stats=train_stats, score_stats=score_stats, kl_coef=kl_coef)

        self.kl_ctl.update(stats['objective/kl'], self.hparams.ppo.batch_size)

        self.print_samples(queries=queries, responses=postprocessed_responses,
                           scores=scores, logprobs=logprobs, ref_logprobs=ref_logprobs)

        # Record profiles of the step times
        step = tf.get_default_session().run(tf.train.get_global_step())
        step_time = time.time() - step_started_at
        eps_per_second = float(self.hparams.ppo.batch_size) / step_time
        if self.comm.Get_rank() == 0:
            print(f"[ppo_step {step}] step_time={step_time:.2f}s, "
                  f"eps/s={eps_per_second:.2f}")