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():