From 25acf9e60f6d82efe291c74bb574976e45d4d872 Mon Sep 17 00:00:00 2001 From: flyingdoubleG <53119724+flyingdoubleG@users.noreply.github.com> Date: Tue, 20 Oct 2026 21:43:41 +0800 Subject: [PATCH] fix(gigpo): match anchor states within each task GiGPO's step-level advantage compares discounted returns of steps that reach the same anchor state within one group of trajectories of the same task (build_step_group in the reference implementation groups by episode group first, then by observation). GiGPOAdvantageFn bucketed steps by env_state_hash over the whole batch, so identical observations from different tasks, e.g. ALFWorld's "Nothing happens." (the hash is taken over the formatted observation), were normalized together although their returns refer to different goals. Key the anchor buckets by (task id, state hash) and add a test with the same observation in two tasks. Co-Authored-By: Claude Opus 5.5 --- tests/algorithm/test_gigpo_advantage.py | 26 +++++++++++++++++++ .../algorithm/advantage_fn/gigpo_advantage.py | 8 +++--- 2 files changed, 31 insertions(+), 3 deletions(-) diff --git a/tests/algorithm/test_gigpo_advantage.py b/tests/algorithm/test_gigpo_advantage.py index 1f1e90dd279..d5ce7212635 100644 --- a/tests/algorithm/test_gigpo_advantage.py +++ b/tests/algorithm/test_gigpo_advantage.py @@ -108,6 +108,32 @@ def test_anchor_state_grouping(self): ) self.assertGreater(metrics["gigpo/anchor_groups_size_gt1"], 0) + def test_anchor_states_are_not_shared_across_tasks(self): + """The same observation in two different tasks is not one anchor group.""" + fn_full = GiGPOAdvantageFn(omega=1.0, gamma=1.0, fnorm="none") + fn_episode = GiGPOAdvantageFn(omega=0.0, gamma=1.0, fnorm="none") + shared_hash = "Observation: Nothing happens." + exps = [ + # task 0: the run that saw the shared observation goes on to succeed + _make_exp(0, 0, 0, step_reward=1.0, env_state_hash=shared_hash), + _make_exp(0, 1, 0, step_reward=0.0, env_state_hash="task0_other"), + # task 1: the run that saw the same observation fails + _make_exp(1, 0, 0, step_reward=0.0, env_state_hash=shared_hash), + _make_exp(1, 1, 0, step_reward=1.0, env_state_hash="task1_other"), + ] + for e in exps: + e.reward = e.info["step_reward"] + p_full, metrics = fn_full.process(list(exps)) + p_episode, _ = fn_episode.process(list(exps)) + + self.assertEqual(metrics["gigpo/anchor_groups_size_gt1"], 0) + + def by_run(exp): + return (exp.eid.task, exp.eid.run) + + for a, b in zip(sorted(p_full, key=by_run), sorted(p_episode, key=by_run)): + self.assertTrue(torch.allclose(a.advantages, b.advantages)) + def test_singleton_anchor_zero_step_adv(self): """Unique anchor hashes -> A^S = 0; only episode term remains.""" fn_none = GiGPOAdvantageFn(omega=1.0, fnorm="none") diff --git a/trinity/algorithm/advantage_fn/gigpo_advantage.py b/trinity/algorithm/advantage_fn/gigpo_advantage.py index a08ef472025..3652fef03aa 100644 --- a/trinity/algorithm/advantage_fn/gigpo_advantage.py +++ b/trinity/algorithm/advantage_fn/gigpo_advantage.py @@ -167,8 +167,8 @@ def process(self, exps: List[Experience]) -> Tuple[List[Experience], Dict]: """Compute GiGPO advantages for a batch of multi-step experiences. Episode-level: group by task, compare trajectory returns R(tau) across - runs (GRPO-style). Step-level: group by ``env_state_hash`` across the - batch, compare discounted returns R_t; singleton anchors get A_S = 0. + runs (GRPO-style). Step-level: group by ``env_state_hash`` within each + task, compare discounted returns R_t; singleton anchors get A_S = 0. Args: exps: Multi-step experiences with ``eid.task``, ``eid.run``, @@ -223,7 +223,9 @@ def process(self, exps: List[Experience]) -> Tuple[List[Experience], Dict]: if state_hash is None: continue r_t = exp_to_discounted_return[id(exp)] - anchor_buckets.setdefault(str(state_hash), []).append((exp, r_t)) + # Anchor states are matched within the task group (as in GiGPO), so identical observations + # from different tasks are not compared. + anchor_buckets.setdefault(f"{exp.eid.tid}/{state_hash}", []).append((exp, r_t)) exp_to_a_s: Dict[int, float] = {id(exp): 0.0 for exp in exps} for members in anchor_buckets.values():