add task-conditioned pi0.5 hierarchy

tests/test_data.py

@@ -8,7 +8,12 @@ import torch
 from PIL import Image
 
 from pi05mini.data import Pi05OverfitDataset, load_manifest, materialize_subset
-from pi05mini.tokenizer import CoordinateTokenizer, MiniTokenizer, SubtaskVocabulary
+from pi05mini.tokenizer import (
+    CoordinateTokenizer,
+    MiniTokenizer,
+    SubtaskVocabulary,
+    TaskVocabulary,
+)
 
 
 def _make_source_project(root: Path) -> Path:
@@ -35,6 +40,7 @@ def _make_source_project(root: Path) -> Path:
     row = {
         "clip_id": "clip-a",
         "t0_us": 1234,
+        "task_instruction": "Follow the road through the construction zone.",
         "reasoning": "Steer left through the construction lane.",
         "image_path": "data/processed/images/sample.jpg",
         "trajectory_path": "data/processed/trajectories/sample.npz",
@@ -56,17 +62,28 @@ def test_coordinate_tokens_are_literal_bytes_and_round_trip() -> None:
     assert torch.allclose(decoded, coordinates, atol=10.1 / 255.0)
 
 
-def test_subtask_vocabulary_and_prefix_contract() -> None:
+def test_task_subtask_vocabulary_and_hierarchical_prefix_contract() -> None:
     tokenizer = MiniTokenizer()
     history = torch.zeros(8, 2)
-    prefix = tokenizer.build_prefix(history)
+    prefix = tokenizer.build_prefix(history, "Follow the road.")
+    other_task = tokenizer.build_prefix(history, "Clean the kitchen.")
+    subtask = tokenizer.encode_subtask("Steer left around the cones")
+    low_level = tokenizer.build_low_level_prefix(history, subtask)
 
-    assert prefix.shape == (20,)
+    assert prefix.shape == (22,)
     assert prefix[:3].tolist() == [tokenizer.BOS, tokenizer.IMAGE, tokenizer.STATE]
+    assert prefix[-3].item() == tokenizer.TASK
+    assert tokenizer.tasks.decode(prefix[-2]) == "follow_road"
     assert prefix[-1].item() == tokenizer.SUBTASK_QUERY
-    assert prefix[3:-1].min() >= 0 and prefix[3:-1].max() <= 255
-    target = tokenizer.encode_subtask("Steer left around the cones")
-    assert tokenizer.subtasks.decode(target) == "steer_left"
+    assert prefix[3:19].min() >= 0 and prefix[3:19].max() <= 255
+    assert other_task[-2] != prefix[-2]
+    assert low_level.shape == (21,)
+    torch.testing.assert_close(low_level[:19], prefix[:19])
+    assert low_level[-2].item() == tokenizer.SUBTASK_COMMAND
+    assert low_level[-1].item() == subtask
+    assert tokenizer.TASK not in low_level
+    assert tokenizer.subtasks.decode(subtask) == "steer_left"
+    assert TaskVocabulary().classify("Clean all dishes in the kitchen") == "clean_kitchen"
     assert SubtaskVocabulary().classify("Slow down for a pedestrian") == "slow_down"
 
 
@@ -98,7 +115,9 @@ def test_materialize_and_dataset_are_self_contained_and_deterministic(tmp_path:
     augmented = dataset[1]
     repeated = same_dataset[1]
     assert original["image"].shape == (3, 32, 48)
-    assert original["prefix_tokens"].shape == (20,)
+    assert original["prefix_tokens"].shape == (22,)
+    assert original["low_level_prefix_tokens"].shape == (21,)
+    assert original["task_instruction"].startswith("Follow the road")
     assert original["actions"].shape == (16, 2)
     assert original["subtask_target"].shape == ()
     assert original["subtask_label"] == "steer_left"