why need to calculate reward_stat? I see llvm_trainer.train use reward from sequence_example.reward

[18liumin]

why calculate reward_stat， I see llvm_trainer.train use reward from sequence_example.reward

"""sequence_example = _overwrite_trajectory_reward( sequence_example=sequence_example, reward=_calculate_reward( policy=policy_reward, baseline=moving_average_reward)) """

https://github.com/google/ml-compiler-opt/blob/0f895e325491793efbc45a1675f9a71f7c8c4f9e/compiler_opt/rl/compilation_runner.py#L446

for k, v in policy_result.items(): sequence_example = v[0] policy_reward = v[1] if k not in reward_stat: raise ValueError( (f'Example {k} does not exist under default policy for ' f'cmd line: {final_cmd_line}')) default_reward = reward_stat[k].default_reward moving_average_reward = reward_stat[k].moving_average_reward sequence_example = _overwrite_trajectory_reward( sequence_example=sequence_example, reward=_calculate_reward( policy=policy_reward, baseline=moving_average_reward)) sequence_example_list.append(sequence_example) reward_stat[k].moving_average_reward = ( moving_average_reward self._moving_average_decay_rate + policy_reward (1 - self._moving_average_decay_rate)) rewards.append( _calculate_reward(policy=policy_reward, baseline=default_reward)) policy_rewards.append(policy_reward) keys.append(k)

result = CompilationResult(
    sequence_examples=sequence_example_list,
    reward_stats=reward_stat,
    rewards=rewards,
    policy_rewards=policy_rewards,
    keys=keys,
    model_id=model_id)

for observer in self._observers:
  observer.observe(result)

return result

[mtrofin]

See here and then uses of self._reward_stat_map. Then note how it's used in CompilationRunner.collect_data.