admin/pi05-minimal-validation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add task-conditioned pi0.5 hierarchy

Browse Source
This commit is contained in:
seeker
2026-07-12 19:16:27 +08:00
parent ffbcca8d2b
commit b57516b57b
10 changed files with 325 additions and 174 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/pi05mini/infer.py
+14 -7
View File
@@ -40,13 +40,20 @@ def evaluate(checkpoint_path: str, manifest: str | None = None) -> Path:
generated: list[int] = []
expected_tokens = 1 + 2 + dataset.action_horizon * 2
current_tokens = tokens
for _ in range(expected_tokens):
for generation_step in range(expected_tokens):
semantic = model(
image, current_tokens, zero_action,
torch.zeros(1, device=device),
generated_token_start=tokens.shape[1]
)
next_token = int(semantic.text_logits[0, -1].argmax().cpu())
if generation_step == 0:
candidates = torch.tensor(
dataset.tokenizer.subtasks.token_ids, device=device
)
candidate_logits = semantic.text_logits[0, -1, candidates]
next_token = int(candidates[candidate_logits.argmax()].cpu())
else:
next_token = int(semantic.text_logits[0, -1].argmax().cpu())
generated.append(next_token)
current_tokens = torch.cat(
(current_tokens, torch.tensor([[next_token]], device=device)), dim=1
@@ -62,13 +69,12 @@ def evaluate(checkpoint_path: str, manifest: str | None = None) -> Path:
except ValueError:
vlm_xy = np.full_like(sample["future_xy"].numpy(), np.nan)
vlm_valid = False
conditioned_tokens = torch.cat(
(tokens, torch.tensor([[predicted_token]], device=device)), dim=1
)
conditioned_tokens = dataset.tokenizer.build_low_level_prefix(
sample["history_xy"], predicted_token
)[None].to(device)
def velocity(x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
return model(
image, conditioned_tokens, x, t,
generated_token_start=tokens.shape[1]
image, conditioned_tokens, x, t
).vector_field
normalized = flow.sample(
velocity, (1, dataset.action_horizon, 2), device,
@@ -88,6 +94,7 @@ def evaluate(checkpoint_path: str, manifest: str | None = None) -> Path:
predicted_label = f"token_{predicted_token}"
rows.append({
"sample_id": sample["sample_id"], "target_subtask": target_label,
"task_instruction": sample["task_instruction"],
"predicted_subtask": predicted_label,
"vlm_action_tokens_valid": vlm_valid,
"vlm_trajectory_mse": vlm_mse,