diff --git a/src/art/tinker_native/data.py b/src/art/tinker_native/data.py index a0e7827fa..0af86f902 100644 --- a/src/art/tinker_native/data.py +++ b/src/art/tinker_native/data.py @@ -34,11 +34,13 @@ def compute_advantages( ) -> list[float]: if not rewards: return [] + if all(abs(reward - rewards[0]) <= 1e-12 for reward in rewards[1:]): + return [0.0] * len(rewards) rewards_tensor = torch.tensor(rewards, dtype=torch.float32) centered = rewards_tensor - rewards_tensor.mean() if not normalize_advantages: return centered.tolist() - std_reward = rewards_tensor.std() + std_reward = rewards_tensor.std(unbiased=False) if std_reward > 1e-8: return (centered / std_reward).tolist() return [0.0] * len(rewards) diff --git a/tests/unit/test_tinker_native_advantages.py b/tests/unit/test_tinker_native_advantages.py new file mode 100644 index 000000000..89c3a711e --- /dev/null +++ b/tests/unit/test_tinker_native_advantages.py @@ -0,0 +1,40 @@ +import math +import warnings + +import pytest + +from art.tinker_native.data import compute_advantages + + +@pytest.mark.parametrize("rewards", [[0.0, 1.0], [0.0, 1.0, 2.0, 3.0]]) +def test_advantages_use_population_standard_deviation(rewards): + mean = sum(rewards) / len(rewards) + std = math.sqrt(sum((reward - mean) ** 2 for reward in rewards) / len(rewards)) + expected = [(reward - mean) / std for reward in rewards] + advantages = compute_advantages(rewards) + assert advantages == pytest.approx(expected) + assert sum(advantage**2 for advantage in advantages) / len(rewards) == ( + pytest.approx(1.0) + ) + + +@pytest.mark.parametrize( + "rewards", + [[], [0.9], [0.9] * 7, [0.1] * 7, [0.0, 1e-12], [0.0, 1e-8], [0.0, 2e-8]], +) +def test_degenerate_advantages_are_zero_without_warnings(rewards): + with warnings.catch_warnings(): + warnings.simplefilter("error") + assert compute_advantages(rewards) == [0.0] * len(rewards) + + +def test_advantages_above_std_guard_are_normalized(): + assert compute_advantages([0.0, 4e-8]) == pytest.approx([-1.0, 1.0]) + + +def test_advantage_normalization_can_be_disabled(): + assert compute_advantages([0.0, 1.0, 2.0, 3.0], False) == [-1.5, -0.5, 0.5, 1.5] + + +def test_equal_rewards_without_normalization_are_exactly_zero(): + assert compute_advantages([0.3] * 8, False) == [0.0] * 8