| 11 | |
| 12 | class VerlPolicyLossTest(unittest.TestCase): |
| 13 | def setUp(self): |
| 14 | seed = 42 |
| 15 | torch.manual_seed(seed) |
| 16 | torch.cuda.manual_seed(seed) |
| 17 | torch.cuda.manual_seed_all(seed) |
| 18 | torch.backends.cudnn.deterministic = True |
| 19 | torch.backends.cudnn.benchmark = False |
| 20 | |
| 21 | shape = (5, 20) |
| 22 | self.logprob = 2 * torch.rand(shape) - 1 |
| 23 | self.input_data = DataProto.from_dict( |
| 24 | { |
| 25 | "old_log_probs": 2 * torch.rand(shape) - 1, |
| 26 | "ref_log_prob": 2 * torch.rand(shape) - 1, |
| 27 | "response_mask": torch.rand(shape) > 0.5, |
| 28 | "advantages": 2 * torch.rand(shape) - 1, |
| 29 | "expert_mask": torch.rand(shape[0]) > 0.5, |
| 30 | } |
| 31 | ) |
| 32 | |
| 33 | def test_ppo_policy_loss(self): |
| 34 | policy_loss_fn_cls = POLICY_LOSS_FN.get("ppo") |