seeker
124 lines
4.5 KiB
Python
124 lines
4.5 KiB
Python
import torch
|
|
|
|
from pi06mini.model import Pi06Config, Pi06Mini
|
|
from pi06mini.tokenizer import MiniTokenizer
|
|
|
|
|
|
def tiny_model() -> Pi06Mini:
|
|
tokenizer = MiniTokenizer()
|
|
return Pi06Mini(
|
|
Pi06Config(
|
|
vocab_size=tokenizer.vocab_size,
|
|
image_size=16,
|
|
patch_size=16,
|
|
d_model=32,
|
|
n_heads=4,
|
|
n_layers=1,
|
|
max_text_tokens=32,
|
|
max_action_tokens=4,
|
|
action_start_token_id=tokenizer.ACTION_START,
|
|
)
|
|
)
|
|
|
|
|
|
def test_flow_reads_subtask_and_advantage_but_not_fast_targets() -> None:
|
|
tokenizer = MiniTokenizer()
|
|
sequence = tokenizer.build_policy_sequence(
|
|
torch.zeros(2, 2), "steer_left", torch.zeros(2, 2), "positive"
|
|
)
|
|
model = tiny_model()
|
|
input_ids = sequence.input_ids[None]
|
|
derived = model._derive_action_prefix_mask(input_ids)
|
|
torch.testing.assert_close(derived[0], sequence.action_prefix_mask)
|
|
|
|
text_offset = 1
|
|
prefix, flow = model._masks(
|
|
batch=1,
|
|
prefix_length=1 + input_ids.shape[1],
|
|
action_length=2,
|
|
text_offset=text_offset,
|
|
attention_mask=torch.ones_like(input_ids, dtype=torch.bool),
|
|
action_prefix_mask=derived,
|
|
generated_token_start=sequence.generated_token_start,
|
|
device=torch.device("cpu"),
|
|
)
|
|
generated_start = text_offset + sequence.generated_token_start
|
|
assert prefix[0, 0, :generated_start].all()
|
|
assert not prefix[0, 0, generated_start:].any()
|
|
# Flow sees image plus all text strictly before ACTION_START.
|
|
readable_prefix = text_offset + sequence.action_start_index
|
|
assert flow[0, :, :readable_prefix].all()
|
|
assert not flow[0, :, readable_prefix : 1 + input_ids.shape[1]].any()
|
|
assert flow[0, :, 1 + input_ids.shape[1] :].all()
|
|
|
|
|
|
def test_flow_loss_is_knowledge_insulated_from_every_vlm_parameter() -> None:
|
|
torch.manual_seed(4)
|
|
tokenizer = MiniTokenizer()
|
|
sequence = tokenizer.build_policy_sequence(
|
|
torch.zeros(2, 2), "steer_left", torch.zeros(2, 2), "positive"
|
|
)
|
|
model = tiny_model()
|
|
output = model(
|
|
pixel_values=torch.randn(1, 3, 16, 16),
|
|
input_ids=sequence.input_ids[None],
|
|
noisy_actions=torch.randn(1, 2, 2),
|
|
timesteps=torch.rand(1),
|
|
action_prefix_mask=sequence.action_prefix_mask[None],
|
|
generated_token_start=sequence.generated_token_start,
|
|
target_velocity=torch.randn(1, 2, 2),
|
|
)
|
|
assert output.flow_loss is not None
|
|
output.flow_loss.backward()
|
|
|
|
vlm_gradients = [parameter.grad for parameter in model.vlm_parameters()]
|
|
assert vlm_gradients
|
|
assert all(gradient is None or torch.count_nonzero(gradient) == 0 for gradient in vlm_gradients)
|
|
action_gradients = [parameter.grad for parameter in model.action_parameters()]
|
|
assert any(gradient is not None and torch.count_nonzero(gradient) > 0 for gradient in action_gradients)
|
|
|
|
|
|
def test_knowledge_insulation_stops_gradients_not_conditioning() -> None:
|
|
torch.manual_seed(8)
|
|
tokenizer = MiniTokenizer()
|
|
model = tiny_model().eval()
|
|
positive = tokenizer.build_policy_sequence(
|
|
torch.zeros(2, 2), "keep_lane", torch.zeros(2, 2), "positive"
|
|
)
|
|
negative = tokenizer.build_policy_sequence(
|
|
torch.zeros(2, 2), "keep_lane", torch.zeros(2, 2), "negative"
|
|
)
|
|
image = torch.randn(1, 3, 16, 16)
|
|
noise = torch.randn(1, 2, 2)
|
|
timestep = torch.full((1,), 0.5)
|
|
with torch.no_grad():
|
|
pos = model(
|
|
image, positive.input_ids[None], noise, timestep,
|
|
generated_token_start=positive.generated_token_start,
|
|
).vector_field
|
|
neg = model(
|
|
image, negative.input_ids[None], noise, timestep,
|
|
generated_token_start=negative.generated_token_start,
|
|
).vector_field
|
|
assert not torch.allclose(pos, neg)
|
|
|
|
|
|
def test_joint_loss_still_trains_vlm_through_token_objective() -> None:
|
|
tokenizer = MiniTokenizer()
|
|
sequence = tokenizer.build_policy_sequence(
|
|
torch.zeros(2, 2), "stop", torch.zeros(2, 2), "negative"
|
|
)
|
|
model = tiny_model()
|
|
output = model(
|
|
torch.randn(1, 3, 16, 16),
|
|
sequence.input_ids[None],
|
|
torch.randn(1, 2, 2),
|
|
torch.rand(1),
|
|
generated_token_start=sequence.generated_token_start,
|
|
text_labels=sequence.labels[None],
|
|
target_velocity=torch.randn(1, 2, 2),
|
|
)
|
|
output.loss.backward()
|
|
assert model.token_embedding.weight.grad is not None
|
|
assert torch.count_nonzero(model.token_embedding.weight.grad) > 0
|