add task-conditioned pi0.5 hierarchy

2026-07-12 19:16:27 +08:00
parent ffbcca8d2b
commit b57516b57b
10 changed files with 325 additions and 174 deletions
@@ -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)