| import torch | |
| from local_train.rl import gold_continuation, group_advantages | |
| def test_gold_continuation_strips_duplicated_think(): | |
| prompt = "<|im_start|>assistant\n<think>\n" | |
| completion = "<think>\nedit now\n</think>\n\n```bash\nsed -i 's/a/b/' foo.py\n```\n" | |
| assert gold_continuation(prompt, completion).startswith("edit now") | |
| assert gold_continuation("plain", completion) == completion | |
| def test_group_advantages_prefers_higher_reward_in_pair(): | |
| rewards = torch.tensor([0.2, 0.8, 0.1, 0.9]) | |
| adv = group_advantages(rewards, num_generations=2) | |
| assert adv[1] > adv[0] | |
| assert adv[3] > adv[2] | |
| def test_group_advantages_falls_back_when_groups_tie(): | |
| # Each pair is tied, but the two groups differ — batch baseline should fire. | |
| rewards = torch.tensor([0.1, 0.1, 0.9, 0.9]) | |
| adv = group_advantages(rewards, num_generations=2) | |
| assert adv[0] < 0 | |
| assert adv[2] > 0 | |
| def test_group_advantages_all_equal_is_zero(): | |
| rewards = torch.tensor([0.0, 0.0, 0.0, 0.0]) | |
| adv = group_advantages(rewards, num_generations=2) | |
| assert torch.allclose(adv, torch.zeros_like(adv)) | |