udpated multi turn reward (50907b8f) · Commits · 张泽凯 / VAGEN

vagen/trainer/ppo/ray_trainer.py

+5 −2

Original line number	Diff line number	Diff line
		@@ -175,9 +175,12 @@ def compute_advantage(data: DataProto, adv_estimator, gamma=1.0, lam=1.0, num_re
		attention_mask = data.batch['attention_mask']
		response_mask = attention_mask[:, -response_length:]
		assert "multi_turn_token_level_rewards" in data.batch.keys()
		assert "loss_mask" in data.batch.keys()
		# assert "loss_mask" in data.batch.keys()
		# loss_mask = data.batch['loss_mask'][:, -response_length:]
		if "loss_mask" in in data.batch.keys():
		loss_mask = data.batch['loss_mask'][:, -response_length:]
		print(f"[DEBUG] high_level_gamma={high_level_gamma}")
		else:
		loss_mask=data.batch['attention_mask'][:, -response_length:]
		advantages, returns = core_algos.compute_multi_turn_gae_advantage_return(token_level_rewards=data.batch['token_level_rewards'],
		values=values,
		loss_mask=loss_mask,