reward_fn in accelerate_ppo_trainer.py

[Jerrrrykun]

Hi, may I know why these two reward_fn function differently while they seem to be the same one passed to the PPO trainer as input? In my understanding of PPO, the reward function should output rewards for each sample instead of a sequence of (sequence_length, ) rewards.

trlx/trlx/trainer/accelerate_ppo_trainer.py

Lines 309 to 310 in 3340c2f

	# reward_fn should return list of rewards at each token per sample
	# NOTE: all_scores[0][i] is the reward due to token (action) i in prompt + response (b/c of how kl is computed)

trlx/trlx/trlx.py

Lines 38 to 40 in 3340c2f

	reward_fn (`Optional[Callable[[List[str], List[str], List[str]], List[float]]]`):
	A function to rate batches of generated samples. Its required arguments are
	(`samples`, `prompts`, `outputs`) and the return is a list of scalar rewards per each sample in batch