Align pi0.7 minimal structure with paper

src/pi07mini/tokenizer.py

@@ -2,7 +2,7 @@
 
 pi0.7 conditions the low-level policy on a richer prompt than earlier pi
 models: task and subtask language, episode speed/quality/mistake metadata,
-control mode, and data provenance.  This tokenizer represents those concepts
+and control mode. This tokenizer represents those concepts
 with explicit marker/value pairs so the architecture can be validated without
 shipping a full language tokenizer.
 """
@@ -26,7 +26,6 @@ class EpisodeMetadata:
     quality: int | None = None
     mistake: bool | None = None
     control_mode: ControlMode = "joint"
-    source_type: str | None = None
 
     def __post_init__(self) -> None:
         if self.speed_steps is not None and self.speed_steps <= 0:
@@ -84,7 +83,7 @@ class MiniTokenizer:
     QUALITY_MARKER = 260
     MISTAKE_MARKER = 261
     CONTROL_MARKER = 262
-    SOURCE_MARKER = 263
+    NEXT_SUBTASK_MARKER = 263
 
     TASK_LABELS = (
         "drive",
@@ -123,18 +122,7 @@ class MiniTokenizer:
     CONTROL_JOINT = MISTAKE_TRUE + 1
     CONTROL_EE = CONTROL_JOINT + 1
 
-    SOURCE_LABELS = (
-        "demonstration",
-        "autonomous",
-        "intervention",
-        "human_video",
-        "web",
-        "unknown",
-    )
-    SOURCE_TOKEN_START = CONTROL_EE + 1
-    SOURCE_TOKEN_END = SOURCE_TOKEN_START + len(SOURCE_LABELS) - 1
-
-    ACTION_START = SOURCE_TOKEN_END + 1
+    ACTION_START = CONTROL_EE + 1
     ACTION_END = ACTION_START + 1
     ACTION_TOKEN_START = ACTION_END + 1
     ACTION_BINS = 64
@@ -145,7 +133,6 @@ class MiniTokenizer:
     def __init__(self) -> None:
         self.tasks = _LabelVocabulary(self.TASK_LABELS, self.TASK_TOKEN_START)
         self.subtasks = _LabelVocabulary(self.SUBTASK_LABELS, self.SUBTASK_TOKEN_START)
-        self.sources = _LabelVocabulary(self.SOURCE_LABELS, self.SOURCE_TOKEN_START)
 
     @property
     def vocab_size(self) -> int:
@@ -197,23 +184,6 @@ class MiniTokenizer:
             return "pick_place"
         return "unknown"
 
-    @staticmethod
-    def _classify_source(source: str) -> str:
-        if source in MiniTokenizer.SOURCE_LABELS:
-            return source
-        text = source.casefold()
-        if any(word in text for word in ("demo", "teleop")):
-            return "demonstration"
-        if any(word in text for word in ("autonomous", "rollout", "policy")):
-            return "autonomous"
-        if any(word in text for word in ("intervention", "correction")):
-            return "intervention"
-        if any(word in text for word in ("human", "ego", "video")):
-            return "human_video"
-        if any(word in text for word in ("web", "internet")):
-            return "web"
-        return "unknown"
-
     def encode_prompt(
         self,
         task: str,
@@ -257,14 +227,40 @@ class MiniTokenizer:
                     self.MISTAKE_MARKER,
                     self.MISTAKE_TRUE if metadata.mistake else self.MISTAKE_FALSE,
                 )
-            if metadata.source_type is not None:
-                source = self._classify_source(metadata.source_type)
-                add_metadata(self.SOURCE_MARKER, self.sources.encode(source))
-
         control = self.CONTROL_JOINT if metadata.control_mode == "joint" else self.CONTROL_EE
         ids.extend((self.CONTROL_MARKER, control))
         return torch.tensor(ids, dtype=torch.long), tuple(metadata_indices)
 
+    def encode_high_level_prompt(
+        self,
+        task: str,
+        past_subtasks: Iterable[str],
+        metadata: EpisodeMetadata,
+    ) -> Tensor:
+        """Encode task, past semantic decisions, metadata, and a next-subtask query."""
+
+        prompt, _ = self.encode_prompt(
+            task,
+            subtask=None,
+            metadata=metadata,
+            include_subtask=False,
+        )
+        control_pair = prompt[-2:]
+        body = prompt[:-2]
+        history: list[int] = []
+        for subtask in past_subtasks:
+            history.extend(
+                (self.SUBTASK_MARKER, self.subtasks.encode(self._classify_subtask(subtask)))
+            )
+        return torch.cat(
+            (
+                body,
+                torch.tensor(history, dtype=torch.long),
+                control_pair,
+                torch.tensor([self.NEXT_SUBTASK_MARKER], dtype=torch.long),
+            )
+        )
+
     @staticmethod
     def _orthonormal_dct_matrix(
         horizon: int, *, device: torch.device, dtype: torch.dtype
@@ -390,4 +386,3 @@ class MiniTokenizer:
             include_episode_metadata=False,
         )
         return conditioned, unconditional
-