add task-conditioned pi0.5 hierarchy

src/pi05mini/data.py

@@ -19,6 +19,7 @@ from .tokenizer import MiniTokenizer
 class SampleRecord:
     clip_id: str
     t0_us: int
+    task_instruction: str
     reasoning: str
     image_path: Path
     trajectory_path: Path
@@ -111,6 +112,9 @@ def load_manifest(path: str | Path) -> list[SampleRecord]:
             SampleRecord(
                 clip_id=str(row["clip_id"]),
                 t0_us=int(row["t0_us"]),
+                task_instruction=str(
+                    row.get("task_instruction", MiniTokenizer.DEFAULT_TASK_INSTRUCTION)
+                ),
                 reasoning=str(row.get("reasoning", "")),
                 image_path=_resolve_data_path(manifest_path, row["image_path"]),
                 trajectory_path=_resolve_data_path(manifest_path, row["trajectory_path"]),
@@ -156,6 +160,7 @@ def materialize_subset(
             {
                 "clip_id": record.clip_id,
                 "t0_us": record.t0_us,
+                "task_instruction": record.task_instruction,
                 "reasoning": record.reasoning,
                 "image_path": str(image_path.relative_to(destination)),
                 "trajectory_path": str(trajectory_path.relative_to(destination)),
@@ -303,7 +308,13 @@ class Pi05OverfitDataset(Dataset[dict[str, Any]]):
         actions = self.action_normalizer.encode(history_xy, future_xy)
         return {
             "image": self._load_image(base_index, variant),
-            "prefix_tokens": self.tokenizer.build_prefix(history_xy),
+            "prefix_tokens": self.tokenizer.build_prefix(
+                history_xy, record.task_instruction
+            ),
+            "low_level_prefix_tokens": self.tokenizer.build_low_level_prefix(
+                history_xy, target
+            ),
+            "task_instruction": record.task_instruction,
             "subtask_target": target,
             "subtask_label": self.tokenizer.subtasks.decode(target),
             "actions": actions,

src/pi05mini/infer.py

@@ -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,

src/pi05mini/tokenizer.py

@@ -52,6 +52,61 @@ class CoordinateTokenizer:
         return low + encoded.float() / (self.bins - 1) * (high - low)
 
 
+class TaskVocabulary:
+    """Small overall-task vocabulary for the high-level pi0.5 inference pass."""
+
+    labels = (
+        "follow_road",
+        "clean_kitchen",
+        "clean_bedroom",
+        "pick_place",
+        "fold_laundry",
+        "open_close",
+        "unknown",
+    )
+
+    def __init__(self, first_token_id: int = 260) -> None:
+        self.first_token_id = int(first_token_id)
+        self._label_to_token = {
+            label: self.first_token_id + index for index, label in enumerate(self.labels)
+        }
+        self._token_to_label = {token: label for label, token in self._label_to_token.items()}
+
+    @property
+    def token_ids(self) -> tuple[int, ...]:
+        return tuple(self._label_to_token[label] for label in self.labels)
+
+    def classify(self, instruction: str) -> str:
+        text = instruction.casefold()
+        if any(term in text for term in ("road", "lane", "drive", "waypoint", "steer")):
+            return "follow_road"
+        if any(term in text for term in ("fold", "folding")):
+            return "fold_laundry"
+        if any(term in text for term in ("kitchen", "dish", "sink", "counter")):
+            return "clean_kitchen"
+        if any(term in text for term in ("bedroom", "bed", "pillow", "blanket")):
+            return "clean_bedroom"
+        if any(term in text for term in ("open", "close", "drawer", "cabinet", "fridge")):
+            return "open_close"
+        if any(term in text for term in ("pick", "place", "put", "move")):
+            return "pick_place"
+        return instruction if instruction in self._label_to_token else "unknown"
+
+    def encode(self, label_or_instruction: str) -> int:
+        label = (
+            label_or_instruction
+            if label_or_instruction in self._label_to_token
+            else self.classify(label_or_instruction)
+        )
+        return self._label_to_token[label]
+
+    def decode(self, token_id: int | torch.Tensor) -> str:
+        try:
+            return self._token_to_label[int(token_id)]
+        except KeyError as error:
+            raise ValueError(f"unknown task token: {int(token_id)}") from error
+
+
 class SubtaskVocabulary:
     """Small semantic vocabulary standing in for pi0.5's generated language subtask."""
 
@@ -114,7 +169,13 @@ class SubtaskVocabulary:
 
 
 class MiniTokenizer:
-    """Compose image/state, semantic, and FAST-lite action tokens.
+    """Compose hierarchical task/state, subtask, and FAST-lite action tokens.
+
+    pi0.5 performs two semantically different calls with the same model.  The
+    high-level call reads the overall task and predicts a semantic subtask.  The
+    low-level call deliberately replaces the overall task with that predicted
+    subtask, matching ``pi(a | observation, predicted_subtask)`` rather than
+    allowing continuous actions to depend directly on the overall task.
 
     Actions are z-score-normalized by the dataset before reaching this
     tokenizer.  A complete ``[horizon, 2]`` chunk is transformed along time
@@ -130,11 +191,19 @@ class MiniTokenizer:
     high-frequency coefficients.
     """
 
+    DEFAULT_TASK_INSTRUCTION = "Follow the visible road."
+
     BOS = 256
     IMAGE = 257
     STATE = 258
-    SUBTASK_QUERY = 259
-    ACTION_START = SUBTASK_QUERY + 1 + len(SubtaskVocabulary.labels)
+    TASK = 259
+    TASK_TOKEN_START = TASK + 1
+    TASK_TOKEN_END = TASK_TOKEN_START + len(TaskVocabulary.labels) - 1
+    SUBTASK_QUERY = TASK_TOKEN_END + 1
+    SUBTASK_COMMAND = SUBTASK_QUERY + 1
+    SUBTASK_TOKEN_START = SUBTASK_COMMAND + 1
+    SUBTASK_TOKEN_END = SUBTASK_TOKEN_START + len(SubtaskVocabulary.labels) - 1
+    ACTION_START = SUBTASK_TOKEN_END + 1
     ACTION_END = ACTION_START + 1
     ACTION_TOKEN_START = ACTION_END + 1
     ACTION_BINS = 64
@@ -145,14 +214,16 @@ class MiniTokenizer:
     def __init__(
         self,
         coordinate_tokenizer: CoordinateTokenizer | None = None,
+        tasks: TaskVocabulary | None = None,
         subtasks: SubtaskVocabulary | None = None,
     ) -> None:
         self.coordinates = coordinate_tokenizer or CoordinateTokenizer()
-        self.subtasks = subtasks or SubtaskVocabulary(first_token_id=self.SUBTASK_QUERY + 1)
+        self.tasks = tasks or TaskVocabulary(first_token_id=self.TASK_TOKEN_START)
+        self.subtasks = subtasks or SubtaskVocabulary(first_token_id=self.SUBTASK_TOKEN_START)
 
     @property
     def vocab_size(self) -> int:
-        return max(max(self.subtasks.token_ids), self.ACTION_TOKEN_END) + 1
+        return max(max(self.tasks.token_ids), max(self.subtasks.token_ids), self.ACTION_TOKEN_END) + 1
 
     @property
     def action_token_ids(self) -> range:
@@ -160,11 +231,37 @@ class MiniTokenizer:
 
         return range(self.ACTION_TOKEN_START, self.ACTION_TOKEN_END + 1)
 
-    def build_prefix(self, history_xy: torch.Tensor) -> torch.Tensor:
+    def build_observation(self, history_xy: torch.Tensor) -> torch.Tensor:
         coordinate_tokens = self.coordinates.encode(history_xy).reshape(-1)
         specials = torch.tensor([self.BOS, self.IMAGE, self.STATE], dtype=torch.long)
-        query = torch.tensor([self.SUBTASK_QUERY], dtype=torch.long)
-        return torch.cat((specials, coordinate_tokens.cpu(), query), dim=0)
+        return torch.cat((specials, coordinate_tokens.cpu()), dim=0)
+
+    def build_prefix(
+        self,
+        history_xy: torch.Tensor,
+        task_instruction: str = DEFAULT_TASK_INSTRUCTION,
+    ) -> torch.Tensor:
+        """Build the high-level ``observation + overall task -> subtask`` prefix."""
+
+        observation = self.build_observation(history_xy)
+        task = torch.tensor(
+            [self.TASK, self.tasks.encode(task_instruction), self.SUBTASK_QUERY],
+            dtype=torch.long,
+        )
+        return torch.cat((observation, task), dim=0)
+
+    def build_low_level_prefix(
+        self,
+        history_xy: torch.Tensor,
+        subtask: str | int | torch.Tensor,
+    ) -> torch.Tensor:
+        """Build ``observation + predicted subtask`` for FAST/flow action inference."""
+
+        token = self.subtasks.encode(subtask) if isinstance(subtask, str) else int(subtask)
+        if token not in self.subtasks.token_ids:
+            raise ValueError(f"low-level command is not a subtask token: {token}")
+        command = torch.tensor([self.SUBTASK_COMMAND, token], dtype=torch.long)
+        return torch.cat((self.build_observation(history_xy), command), dim=0)
 
     def encode_subtask(self, reasoning: str) -> torch.Tensor:
         return torch.tensor(self.subtasks.encode(reasoning), dtype=torch.long)

src/pi05mini/train.py

@@ -15,7 +15,8 @@ from .runtime import choose_device, load_config, seed_everything, write_json
 
 def _collate(rows: list[dict]) -> dict:
     tensor_keys = (
-        "image", "prefix_tokens", "subtask_target", "actions", "action_tokens"
+        "image", "prefix_tokens", "low_level_prefix_tokens", "subtask_target",
+        "actions", "action_tokens"
     )
     result = {key: torch.stack([row[key] for row in rows]) for key in tensor_keys}
     result["sample_id"] = [row["sample_id"] for row in rows]
@@ -99,8 +100,8 @@ def train(config_path: str) -> Path:
         model_labels = torch.full_like(model_tokens, -100)
         query_index = tokens.shape[1] - 1
         model_labels[:, query_index:] = targets
-        # Flow tokens may read observation/state and the generated semantic
-        # subtask, but not teacher-forced discrete action tokens.
+        # The action expert is present in the module but receives no loss in
+        # discrete pre-training, and its parameters are excluded from optimizer.
         zero_action = torch.zeros_like(actions)
         output = model(
             image, model_tokens, zero_action,
@@ -136,6 +137,7 @@ def train(config_path: str) -> Path:
         batch, iterator = next_batch(iterator)
         image = batch["image"].to(device)
         tokens = batch["prefix_tokens"].to(device)
+        low_level_tokens = batch["low_level_prefix_tokens"].to(device)
         actions = batch["actions"].to(device)
         subtask = batch["subtask_target"].to(device)[:, None]
         action_tokens = batch["action_tokens"].to(device)
@@ -143,27 +145,35 @@ def train(config_path: str) -> Path:
         model_tokens = torch.cat((tokens, targets[:, :-1]), dim=1)
         model_labels = torch.full_like(model_tokens, -100)
         model_labels[:, tokens.shape[1] - 1:] = targets
-        action_prefix_mask = torch.zeros_like(model_tokens, dtype=torch.bool)
-        action_prefix_mask[:, : tokens.shape[1] + 1] = True
         noisy, timestep, velocity = flow.training_batch(actions)
-        output = model(
-            image, model_tokens, noisy, timestep,
-            action_prefix_mask=action_prefix_mask,
+        # Text/FAST training uses the overall task and predicts the subtask.
+        # Continuous low-level actions use a second call whose prompt contains
+        # only the observation and teacher-forced subtask. This preserves the
+        # paper's factorization pi(a | observation, predicted_subtask).
+        token_output = model(
+            image, model_tokens, torch.zeros_like(actions),
+            torch.zeros(image.shape[0], device=device),
             generated_token_start=tokens.shape[1],
-            text_labels=model_labels, target_velocity=velocity,
-            text_loss_weight=cfg["train"]["semantic_loss_weight"],
-            flow_loss_weight=cfg["train"]["flow_loss_weight"],
+            text_labels=model_labels,
+        )
+        flow_output = model(
+            image, low_level_tokens, noisy, timestep,
+            target_velocity=velocity,
+        )
+        loss = (
+            cfg["train"]["semantic_loss_weight"] * token_output.text_loss
+            + cfg["train"]["flow_loss_weight"] * flow_output.flow_loss
         )
         optimizer.zero_grad(set_to_none=True)
-        output.loss.backward()
+        loss.backward()
         torch.nn.utils.clip_grad_norm_(model.parameters(), cfg["train"]["grad_clip_norm"])
         optimizer.step()
         if step == 1 or step % cfg["train"]["log_every"] == 0:
             record = {
                 "stage": "posttrain", "step": step,
-                "loss": float(output.loss.detach().cpu()),
-                "text_loss": float(output.text_loss.detach().cpu()),
-                "flow_loss": float(output.flow_loss.detach().cpu()),
+                "loss": float(loss.detach().cpu()),
+                "text_loss": float(token_output.text_loss.detach().cpu()),
+                "flow_loss": float(flow_output.flow_loss.detach().cpu()),
             }
             history.append(record)
             print(json.dumps(record))