File size: 1,178 Bytes
88e15cd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 | import torch
from fugu_lite.train_rl import contextual_bandit_loss
from fugu_lite.train_sft import soft_label_loss
def test_sft_loss_prefers_high_reward_worker():
rewards = torch.tensor([[1.0, 0.0]])
good_logits = torch.tensor([[4.0, -4.0]])
bad_logits = torch.tensor([[-4.0, 4.0]])
assert soft_label_loss(good_logits, rewards, 0.1) < soft_label_loss(
bad_logits, rewards, 0.1
)
def test_expected_reward_has_useful_gradient():
logits = torch.zeros((1, 2), requires_grad=True)
rewards = torch.tensor([[1.0, 0.0]])
loss, metrics = contextual_bandit_loss(
logits,
rewards,
estimator="expected_reward",
entropy_coefficient=0.0,
)
loss.backward()
assert logits.grad[0, 0] < 0
assert logits.grad[0, 1] > 0
assert metrics["zero_spread_fraction"] == 0.0
def test_reinforce_loss_is_finite_with_equal_rewards():
torch.manual_seed(0)
logits = torch.zeros((2, 3), requires_grad=True)
rewards = torch.ones((2, 3))
loss, metrics = contextual_bandit_loss(logits, rewards, estimator="reinforce")
assert torch.isfinite(loss)
assert metrics["zero_spread_fraction"] == 1.0
|