From f95b94eb7e9f95926f660ba26146b00fd1793f5d Mon Sep 17 00:00:00 2001 From: flyingdoubleG <53119724+flyingdoubleG@users.noreply.github.com> Date: Tue, 20 Oct 2026 21:46:39 +0800 Subject: [PATCH] fix(rec): handle groups with a single experience RECGroupedAdvantage.calculate_group_advantage only builds `rewards` when the group has more than one experience, but uses it right after for both branches, so a group of one (repeat_times=1, or a partial group) raises UnboundLocalError. Build `rewards` before the branch; a single experience keeps the existing no-baseline convention (mean 0, std 1). Co-Authored-By: Claude Opus 5.5 --- tests/algorithm/advantage_fn_test.py | 17 +++++++++++++++++ trinity/algorithm/advantage_fn/rec_advantage.py | 2 +- 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/tests/algorithm/advantage_fn_test.py b/tests/algorithm/advantage_fn_test.py index b9f8dab858a..34d9f7d02dc 100644 --- a/tests/algorithm/advantage_fn_test.py +++ b/tests/algorithm/advantage_fn_test.py @@ -394,3 +394,20 @@ def test_step_wise_grpo_with_std_threshold(self): # Verify that all remaining experiences are from task 2 for exp in processed_exps: self.assertEqual(exp.eid.task, 2) + + def test_rec_advantage_single_experience_group(self): + advantage_fn_cls = ADVANTAGE_FN.get("rec") + for args in [{}, {"std_normalize": True}, {"drop": "balance"}]: + advantage_fn = advantage_fn_cls(**{**advantage_fn_cls.default_args(), **args}) + exps = [ + Experience( + eid=EID(batch=0, task=0, run=0), + tokens=torch.zeros(5), + prompt_length=2, + reward=0.5, + ) + ] + processed_exps, _ = advantage_fn(exps) + self.assertEqual(len(processed_exps), 1) + # a group of one has no baseline (mean 0, std 1), as in GRPO + self.assertTrue(torch.allclose(processed_exps[0].advantages, torch.full((3,), 0.5))) diff --git a/trinity/algorithm/advantage_fn/rec_advantage.py b/trinity/algorithm/advantage_fn/rec_advantage.py index 0a99137c00a..82d0dd06648 100644 --- a/trinity/algorithm/advantage_fn/rec_advantage.py +++ b/trinity/algorithm/advantage_fn/rec_advantage.py @@ -40,11 +40,11 @@ def calculate_group_advantage( N = len(exps) metrics = {} with torch.no_grad(): + rewards = torch.tensor([exp.reward for exp in exps], dtype=torch.float32) if len(exps) == 1: group_reward_mean = torch.tensor(0.0) group_reward_std = torch.tensor(1.0) # set to 1.0 to avoid division by zero else: - rewards = torch.tensor([exp.reward for exp in exps], dtype=torch.float32) group_reward_mean = torch.mean(rewards) group_reward_std = torch.std(rewards)