Buckets:
| import math | |
| import unittest | |
| import torch | |
| from rl.advantages import normalize_group_rewards | |
| from rl.objective import TrajectoryPolicyInputs, multimodal_grpo_loss | |
| class RLObjectiveTest(unittest.TestCase): | |
| def test_group_advantage_uses_population_std_and_exact_constant_zero(self): | |
| output = normalize_group_rewards( | |
| torch.tensor([[1.0, 3.0], [2.0, 2.0]]), | |
| sigma_min=0.0, | |
| epsilon=1e-8, | |
| ) | |
| torch.testing.assert_close(output.means, torch.tensor([2.0, 2.0])) | |
| torch.testing.assert_close(output.stds, torch.tensor([1.0, 0.0])) | |
| torch.testing.assert_close( | |
| output.advantages[0], torch.tensor([-1.0, 1.0]) | |
| ) | |
| torch.testing.assert_close(output.advantages[1], torch.zeros(2)) | |
| self.assertEqual( | |
| output.zero_advantage_reasons, | |
| (None, "identical_total_reward"), | |
| ) | |
| self.assertFalse(output.advantages.requires_grad) | |
| def test_factorwise_clip_then_uniform_text_visual_step_mean(self): | |
| text_current = torch.tensor([math.log(1.5)], requires_grad=True) | |
| visual_current = torch.tensor( | |
| [[math.log(1.0), math.log(2.0)]], requires_grad=True | |
| ) | |
| result = multimodal_grpo_loss( | |
| [ | |
| TrajectoryPolicyInputs( | |
| advantage=torch.tensor(1.0), | |
| text_current_log_probs=text_current, | |
| text_old_log_probs=torch.zeros(1), | |
| visual_current_factor_log_probs=visual_current, | |
| visual_old_factor_log_probs=torch.zeros(1, 2), | |
| ) | |
| ], | |
| clip_low=0.2, | |
| clip_high=0.2, | |
| reference_beta=0.0, | |
| ) | |
| self.assertAlmostEqual(float(result.surrogate), 1.15, places=6) | |
| self.assertAlmostEqual(float(result.loss), -1.15, places=6) | |
| self.assertEqual(float(result.statistics["text_clip_fraction"]), 1.0) | |
| self.assertEqual(float(result.statistics["visual_factor_clip_fraction"]), 0.5) | |
| result.loss.backward() | |
| self.assertIsNotNone(text_current.grad) | |
| self.assertIsNotNone(visual_current.grad) | |
| def test_negative_advantage_uses_the_lower_clipped_objective(self): | |
| result = multimodal_grpo_loss( | |
| [ | |
| TrajectoryPolicyInputs( | |
| advantage=torch.tensor(-1.0), | |
| text_current_log_probs=torch.tensor( | |
| [math.log(0.5)], requires_grad=True | |
| ), | |
| text_old_log_probs=torch.zeros(1), | |
| visual_current_factor_log_probs=torch.empty((0, 2)), | |
| visual_old_factor_log_probs=torch.empty((0, 2)), | |
| ) | |
| ], | |
| clip_low=0.2, | |
| clip_high=0.2, | |
| reference_beta=0.0, | |
| ) | |
| self.assertAlmostEqual(float(result.surrogate), -0.8, places=6) | |
| def test_reference_jsd_is_full_distribution_and_only_current_gets_gradient(self): | |
| current = torch.tensor([[2.0, -1.0]], requires_grad=True) | |
| reference = torch.tensor([[-1.0, 2.0]], requires_grad=True) | |
| result = multimodal_grpo_loss( | |
| [ | |
| TrajectoryPolicyInputs( | |
| advantage=torch.tensor(0.0), | |
| text_current_log_probs=torch.tensor([0.0]), | |
| text_old_log_probs=torch.tensor([0.0]), | |
| visual_current_factor_log_probs=torch.empty((0, 2)), | |
| visual_old_factor_log_probs=torch.empty((0, 2)), | |
| text_current_logits=current, | |
| text_reference_logits=reference, | |
| ) | |
| ], | |
| clip_low=0.2, | |
| clip_high=0.2, | |
| reference_beta=0.5, | |
| ) | |
| self.assertGreater(float(result.reference_jsd), 0.0) | |
| result.loss.backward() | |
| self.assertIsNotNone(current.grad) | |
| self.assertIsNone(reference.grad) | |
| def test_zero_autonomous_action_trajectory_has_backward_safe_zero_objective(self): | |
| text = torch.empty(0, requires_grad=True) | |
| visual = torch.empty((0, 2), requires_grad=True) | |
| result = multimodal_grpo_loss( | |
| [ | |
| TrajectoryPolicyInputs( | |
| advantage=torch.tensor(1.0), | |
| text_current_log_probs=text, | |
| text_old_log_probs=torch.empty(0), | |
| visual_current_factor_log_probs=visual, | |
| visual_old_factor_log_probs=torch.empty((0, 2)), | |
| ) | |
| ], | |
| clip_low=0.2, | |
| clip_high=0.2, | |
| reference_beta=0.0, | |
| ) | |
| self.assertEqual(float(result.loss), 0.0) | |
| result.loss.backward() | |
| self.assertIsNotNone(text.grad) | |
| self.assertIsNotNone(visual.grad) | |
| if __name__ == "__main__": | |
| unittest.main() | |
Xet Storage Details
- Size:
- 4.87 kB
- Xet hash:
- ac839161ddb182f0192d3ce80f2d4a1c0ad22c72b3d9831b13148132bb490ad2
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.