eff70fcc79aee2af0f4528ef16499c5dce3662e0,tf_agents/experimental/examples/ppo/ppo_learner.py,PPOLearner,_update_normalizers,#PPOLearner#Any#,216

Before Change

  @common.function(autograph=True)
  def _update_normalizers(self, dataset):
    iterator = iter(dataset)
    traj, _ = next(iterator)
    self._agent.update_observation_normalizer(traj.observation)
    self._agent.update_reward_normalizer(traj.reward)

    for traj, _ in iterator:
      self._agent.update_observation_normalizer(traj.observation)
      self._agent.update_reward_normalizer(traj.reward)

  @common.function(autograph=True)
  def _update_advantage_normalizer(self, dataset):
    self._agent._reset_advantage_normalizer()  // pylint: disable=protected-access

After Change

    num_frames += _update(traj)

    for _ in tf.range(1, self._num_batches):
      traj, _ = next(iterator)
      num_frames += _update(traj)

    return num_frames

Instances