Buckets:
| import unittest | |
| from types import SimpleNamespace | |
| import torch | |
| from rl.config import RLConfig | |
| from rl.engine import execute_actor_update | |
| from rl.objective import TrajectoryPolicyInputs | |
| from rl.policy import policy_replay_mismatch, PolicyScore | |
| class _Scorer: | |
| def __init__(self, actor): | |
| self.actor = actor | |
| def score(self, trajectory, *, advantage): | |
| current = self.actor.weight.reshape(1) | |
| if trajectory == "b": | |
| current = -current | |
| inputs = TrajectoryPolicyInputs( | |
| advantage=advantage, | |
| text_current_log_probs=current, | |
| text_old_log_probs=torch.zeros(1), | |
| visual_current_factor_log_probs=current.new_empty((0, 0)), | |
| visual_old_factor_log_probs=current.new_empty((0, 0)), | |
| ) | |
| return PolicyScore( | |
| inputs=inputs, | |
| text_reference_log_probs=torch.zeros(1), | |
| visual_reference_factor_log_probs=torch.empty((0, 0)), | |
| text_replay_log_probs=current.detach(), | |
| visual_replay_factor_log_probs=torch.empty((0, 0)), | |
| behavior_log_prob_max_abs_error=0.0, | |
| text_behavior_log_prob_max_abs_error=0.0, | |
| visual_behavior_log_prob_max_abs_error=0.0, | |
| text_behavior_log_prob_mean_abs_error=0.0, | |
| text_behavior_log_prob_p99_abs_error=0.0, | |
| visual_behavior_log_prob_mean_abs_error=0.0, | |
| visual_behavior_log_prob_p99_abs_error=0.0, | |
| visual_behavior_log_prob_step_signed_mean_abs_p99=0.0, | |
| text_actions=1, | |
| visual_steps=0, | |
| ) | |
| class _RawTextMismatchScorer: | |
| def __init__(self, actor, deltas): | |
| self.actor = actor | |
| self.deltas = deltas | |
| def score(self, trajectory, *, advantage): | |
| replay = self.deltas[trajectory] | |
| old = torch.zeros_like(replay) | |
| current = self.actor.weight.reshape(1).expand_as(replay) | |
| mismatch = policy_replay_mismatch(replay, old, factorized=False) | |
| visual = current.new_empty((0, 0)) | |
| return PolicyScore( | |
| inputs=TrajectoryPolicyInputs( | |
| advantage=advantage, | |
| text_current_log_probs=current, | |
| text_old_log_probs=old, | |
| visual_current_factor_log_probs=visual, | |
| visual_old_factor_log_probs=visual, | |
| ), | |
| text_reference_log_probs=torch.zeros_like(replay), | |
| visual_reference_factor_log_probs=visual, | |
| text_replay_log_probs=replay, | |
| visual_replay_factor_log_probs=visual, | |
| behavior_log_prob_max_abs_error=mismatch["max_abs"], | |
| text_behavior_log_prob_max_abs_error=mismatch["max_abs"], | |
| visual_behavior_log_prob_max_abs_error=0.0, | |
| text_behavior_log_prob_mean_abs_error=mismatch["mean_abs"], | |
| text_behavior_log_prob_p99_abs_error=mismatch["p99_abs"], | |
| visual_behavior_log_prob_mean_abs_error=0.0, | |
| visual_behavior_log_prob_p99_abs_error=0.0, | |
| visual_behavior_log_prob_step_signed_mean_abs_p99=0.0, | |
| text_actions=replay.numel(), | |
| visual_steps=0, | |
| ) | |
| class _ObservedA5000MismatchScorer(_RawTextMismatchScorer): | |
| def __init__(self, actor): | |
| observed = torch.cat( | |
| ( | |
| torch.full((989,), 0.01), | |
| torch.full((10,), 0.272165), | |
| torch.tensor([2.45469]), | |
| ) | |
| ) | |
| super().__init__(actor, {"a": observed, "b": observed}) | |
| class RLEngineTest(unittest.TestCase): | |
| def test_one_fresh_group_produces_one_finite_policy_update(self): | |
| actor = torch.nn.Linear(1, 1, bias=False) | |
| actor.weight.data.zero_() | |
| optimizer = torch.optim.SGD(actor.parameters(), lr=0.1) | |
| scheduler = torch.optim.lr_scheduler.LambdaLR( | |
| optimizer, lambda _step: 1.0 | |
| ) | |
| config = RLConfig() | |
| config.rollout.group_size = 2 | |
| config.reference.beta = 0.0 | |
| group = SimpleNamespace( | |
| trajectories=( | |
| SimpleNamespace(trajectory="a"), | |
| SimpleNamespace(trajectory="b"), | |
| ), | |
| advantages=torch.tensor([1.0, -1.0]), | |
| ) | |
| result = execute_actor_update( | |
| groups=[group], | |
| policy_scorer=_Scorer(actor), | |
| config=config, | |
| parameters=list(actor.parameters()), | |
| actor=actor, | |
| optimizer=optimizer, | |
| scheduler=scheduler, | |
| ) | |
| self.assertTrue(result.successful) | |
| self.assertEqual(result.trajectories, 2) | |
| self.assertEqual(scheduler.last_epoch, 1) | |
| self.assertTrue(torch.isfinite(actor.weight).all()) | |
| self.assertNotEqual(float(actor.weight), 0.0) | |
| self.assertEqual(len(result.policy_records), 2) | |
| def test_observed_a5000_tail_replay_drift_does_not_reject_clean_means(self): | |
| actor = torch.nn.Linear(1, 1, bias=False) | |
| actor.weight.data.zero_() | |
| optimizer = torch.optim.SGD(actor.parameters(), lr=0.1) | |
| scheduler = torch.optim.lr_scheduler.LambdaLR( | |
| optimizer, lambda _step: 1.0 | |
| ) | |
| config = RLConfig() | |
| config.rollout.group_size = 2 | |
| config.reference.beta = 0.0 | |
| group = SimpleNamespace( | |
| trajectories=( | |
| SimpleNamespace(trajectory="a"), | |
| SimpleNamespace(trajectory="b"), | |
| ), | |
| advantages=torch.tensor([1.0, -1.0]), | |
| ) | |
| result = execute_actor_update( | |
| groups=[group], | |
| policy_scorer=_ObservedA5000MismatchScorer(actor), | |
| config=config, | |
| parameters=list(actor.parameters()), | |
| actor=actor, | |
| optimizer=optimizer, | |
| scheduler=scheduler, | |
| ) | |
| self.assertTrue(result.successful) | |
| self.assertAlmostEqual( | |
| result.statistics["behavior_log_prob_max_abs_error"], 2.45469 | |
| ) | |
| def test_sparse_trajectory_tail_uses_logical_batch_action_population(self): | |
| actor = torch.nn.Linear(1, 1, bias=False) | |
| actor.weight.data.zero_() | |
| optimizer = torch.optim.SGD(actor.parameters(), lr=0.1) | |
| scheduler = torch.optim.lr_scheduler.LambdaLR( | |
| optimizer, lambda _step: 1.0 | |
| ) | |
| config = RLConfig() | |
| config.rollout.group_size = 2 | |
| config.reference.beta = 0.0 | |
| names = ("sparse-tail", "long-clean") | |
| deltas = { | |
| names[0]: torch.full((100,), 0.01), | |
| names[1]: torch.full((10_000,), 0.01), | |
| } | |
| deltas[names[0]][-2:] = 0.320204 | |
| group = SimpleNamespace( | |
| trajectories=tuple( | |
| SimpleNamespace(trajectory=name) for name in names | |
| ), | |
| advantages=torch.ones(2), | |
| ) | |
| result = execute_actor_update( | |
| groups=[group], | |
| policy_scorer=_RawTextMismatchScorer(actor, deltas), | |
| config=config, | |
| parameters=list(actor.parameters()), | |
| actor=actor, | |
| optimizer=optimizer, | |
| scheduler=scheduler, | |
| ) | |
| self.assertTrue(result.successful) | |
| self.assertGreater( | |
| result.policy_records[0].text_behavior_log_prob_p99_abs_error, | |
| config.behavior.replay_text_p99_abs_limit, | |
| ) | |
| self.assertAlmostEqual( | |
| result.statistics["text_behavior_log_prob_p99_abs_error"], | |
| 0.01, | |
| ) | |
| def test_systematic_logical_batch_tail_still_fails_replay_guard(self): | |
| actor = torch.nn.Linear(1, 1, bias=False) | |
| actor.weight.data.zero_() | |
| optimizer = torch.optim.SGD(actor.parameters(), lr=0.1) | |
| scheduler = torch.optim.lr_scheduler.LambdaLR( | |
| optimizer, lambda _step: 1.0 | |
| ) | |
| config = RLConfig() | |
| config.rollout.group_size = 2 | |
| config.reference.beta = 0.0 | |
| tail = torch.cat((torch.zeros(89), torch.full((11,), 0.31))) | |
| group = SimpleNamespace( | |
| trajectories=( | |
| SimpleNamespace(trajectory="a"), | |
| SimpleNamespace(trajectory="b"), | |
| ), | |
| advantages=torch.tensor([1.0, -1.0]), | |
| ) | |
| with self.assertRaisesRegex( | |
| RuntimeError, | |
| "logical-batch.*text_p99_abs=0.31>0.3", | |
| ): | |
| execute_actor_update( | |
| groups=[group], | |
| policy_scorer=_RawTextMismatchScorer( | |
| actor, {"a": tail, "b": tail} | |
| ), | |
| config=config, | |
| parameters=list(actor.parameters()), | |
| actor=actor, | |
| optimizer=optimizer, | |
| scheduler=scheduler, | |
| ) | |
| if __name__ == "__main__": | |
| unittest.main() | |
Xet Storage Details
- Size:
- 8.74 kB
- Xet hash:
- 1275a39703a590477f035162653c26292038766a0b575a9996bcdeb1779d79d8
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.