seeker
74 lines
2.4 KiB
Python
74 lines
2.4 KiB
Python
import torch
|
|
|
|
from pi06mini.model import Pi06Config, Pi06Mini
|
|
from pi06mini.tokenizer import MiniTokenizer
|
|
from pi06mini.value import DistributionalValueModel, ValueConfig
|
|
|
|
|
|
def tiny_value_model() -> DistributionalValueModel:
|
|
tokenizer = MiniTokenizer()
|
|
return DistributionalValueModel(
|
|
ValueConfig(
|
|
vocab_size=tokenizer.vocab_size,
|
|
image_size=16,
|
|
patch_size=16,
|
|
d_model=32,
|
|
n_heads=4,
|
|
n_layers=1,
|
|
max_text_tokens=16,
|
|
)
|
|
)
|
|
|
|
|
|
def test_distributional_value_has_201_bins_and_expected_value() -> None:
|
|
tokenizer = MiniTokenizer()
|
|
model = tiny_value_model()
|
|
state = tokenizer.build_value_input(torch.zeros(2, 2))[None].repeat(2, 1)
|
|
output = model(
|
|
torch.randn(2, 3, 16, 16),
|
|
state,
|
|
target_values=torch.tensor([-1.0, 0.0]),
|
|
)
|
|
|
|
assert output.logits.shape == (2, 201)
|
|
assert output.probabilities.shape == (2, 201)
|
|
torch.testing.assert_close(output.probabilities.sum(-1), torch.ones(2))
|
|
assert torch.all((-1.0 <= output.expectation) & (output.expectation <= 0.0))
|
|
assert output.target_bins.tolist() == [0, 200]
|
|
assert output.loss is not None and output.loss.ndim == 0
|
|
|
|
|
|
def test_value_discretization_and_expectation_use_configured_support() -> None:
|
|
model = tiny_value_model()
|
|
bins = model.values_to_bins(torch.tensor([-2.0, -0.5, 1.0]))
|
|
assert bins.tolist() == [0, 100, 200]
|
|
torch.testing.assert_close(
|
|
model.bins_to_values(torch.tensor([0, 100, 200])),
|
|
torch.tensor([-1.0, -0.5, 0.0]),
|
|
)
|
|
logits = torch.full((3, 201), -100.0)
|
|
logits[0, 0] = logits[1, 100] = logits[2, 200] = 100.0
|
|
torch.testing.assert_close(
|
|
model.expectation_from_logits(logits),
|
|
torch.tensor([-1.0, -0.5, 0.0]),
|
|
)
|
|
|
|
|
|
def test_value_model_and_policy_share_no_parameters() -> None:
|
|
tokenizer = MiniTokenizer()
|
|
policy = Pi06Mini(
|
|
Pi06Config(
|
|
vocab_size=tokenizer.vocab_size,
|
|
image_size=16,
|
|
patch_size=16,
|
|
d_model=32,
|
|
n_heads=4,
|
|
n_layers=1,
|
|
action_start_token_id=tokenizer.ACTION_START,
|
|
)
|
|
)
|
|
value = tiny_value_model()
|
|
policy_ids = {id(parameter) for parameter in policy.parameters()}
|
|
value_ids = {id(parameter) for parameter in value.parameters()}
|
|
assert policy_ids.isdisjoint(value_ids)
|