Add optional pi0.5 subtask prediction stage

README.md

@@ -5,7 +5,7 @@ openpi holds open-source models and packages for robotics, published by the [Phy
 Currently, this repo contains three types of models:
 - the [π₀ model](https://www.physicalintelligence.company/blog/pi0), a flow-based vision-language-action model (VLA).
 - the [π₀-FAST model](https://www.physicalintelligence.company/research/fast), an autoregressive VLA, based on the FAST action tokenizer.
-- the [π₀.₅ model](https://www.physicalintelligence.company/blog/pi05), an upgraded version of π₀ with better open-world generalization trained with [knowledge insulation](https://www.physicalintelligence.company/research/knowledge_insulation). Note that, in this repository, we currently only support the flow matching head for both $\pi_{0.5}$ training and inference.
+- the [π₀.₅ model](https://www.physicalintelligence.company/blog/pi05), an upgraded version of π₀ with better open-world generalization trained with [knowledge insulation](https://www.physicalintelligence.company/research/knowledge_insulation). By default, this repository uses the flow matching head for both $\pi_{0.5}$ training and inference; optional subtask prediction can be enabled with the `Pi0Config` `train_subtask_prediction` and `sample_subtask_prediction` flags for research workflows that provide or want to generate intermediate subtask text.
 
 For all models, we provide _base model_ checkpoints, pre-trained on 10k+ hours of robot data, and examples for using them out of the box or fine-tuning them to your own datasets.
 

src/openpi/models/gemma.py

@@ -385,6 +385,10 @@ def setup(self):
     def embed(self, tokens: at.Int[at.Array, "b t"]) -> at.Float[at.Array, "b t d"]:
         return self.embedder.encode(tokens).astype(self.embed_dtype)
 
+    @at.typecheck
+    def deembed(self, hidden: at.Float[at.Array, "b t d"]) -> at.Float[at.Array, "b t v"]:
+        return self.embedder.decode(hidden)
+
     @at.typecheck
     def __call__(
         self,

src/openpi/models/model.py

@@ -66,8 +66,10 @@ class ModelType(enum.Enum):
 #     "state": float32[*b, s],  # Low-dimensional robot state
 #     "tokenized_prompt": int32[*b, l],  # Optional, tokenized language prompt
 #     "tokenized_prompt_mask": bool[*b, l],  # Optional, mask for tokenized prompt
-#     "token_ar_mask": int32[*b, l],  # Optional, autoregressive mask for FAST model
-#     "token_loss_mask": bool[*b, l],  # Optional, loss mask for FAST model
+#     "token_ar_mask": int32[*b, l],  # Optional, autoregressive token mask
+#     "token_loss_mask": bool[*b, l],  # Optional, token loss mask
+#     "tokenized_action_suffix": int32[*b, l],  # Optional, tokenized action cue for staged pi0.5 inference
+#     "tokenized_action_suffix_mask": bool[*b, l],  # Optional, mask for tokenized action cue
 #
 #      # Actions data.
 #      "actions": float32[*b ah ad]
@@ -99,19 +101,25 @@ class Observation(Generic[ArrayT]):
     # Tokenized prompt mask.
     tokenized_prompt_mask: at.Bool[ArrayT, "*b l"] | None = None
 
-    # pi0-fast model specific fields.
+    # Autoregressive token fields used by pi0-fast and optional staged pi0.5 subtask prediction.
 
-    # Token auto-regressive mask (for FAST autoregressive model).
+    # Token auto-regressive mask.
     token_ar_mask: at.Int[ArrayT, "*b l"] | None = None
-    # Token loss mask (for FAST autoregressive model).
+    # Token loss mask.
     token_loss_mask: at.Bool[ArrayT, "*b l"] | None = None
 
+    # Optional tokenized action cue used when pi0.5 first predicts a subtask and then conditions actions on it.
+    tokenized_action_suffix: at.Int[ArrayT, "*b l"] | None = None
+    tokenized_action_suffix_mask: at.Bool[ArrayT, "*b l"] | None = None
+
     @classmethod
     def from_dict(cls, data: at.PyTree[ArrayT]) -> "Observation[ArrayT]":
         """This method defines the mapping between unstructured data (i.e., nested dict) to the structured Observation format."""
         # Ensure that tokenized_prompt and tokenized_prompt_mask are provided together.
         if ("tokenized_prompt" in data) != ("tokenized_prompt_mask" in data):
             raise ValueError("tokenized_prompt and tokenized_prompt_mask must be provided together.")
+        if ("tokenized_action_suffix" in data) != ("tokenized_action_suffix_mask" in data):
+            raise ValueError("tokenized_action_suffix and tokenized_action_suffix_mask must be provided together.")
         # If images are uint8, convert them to [-1, 1] float32.
         for key in data["image"]:
             if data["image"][key].dtype == np.uint8:
@@ -126,6 +134,8 @@ def from_dict(cls, data: at.PyTree[ArrayT]) -> "Observation[ArrayT]":
             tokenized_prompt_mask=data.get("tokenized_prompt_mask"),
             token_ar_mask=data.get("token_ar_mask"),
             token_loss_mask=data.get("token_loss_mask"),
+            tokenized_action_suffix=data.get("tokenized_action_suffix"),
+            tokenized_action_suffix_mask=data.get("tokenized_action_suffix_mask"),
         )
 
     def to_dict(self) -> at.PyTree[ArrayT]:
@@ -205,6 +215,8 @@ def preprocess_observation(
         tokenized_prompt_mask=observation.tokenized_prompt_mask,
         token_ar_mask=observation.token_ar_mask,
         token_loss_mask=observation.token_loss_mask,
+        tokenized_action_suffix=observation.tokenized_action_suffix,
+        tokenized_action_suffix_mask=observation.tokenized_action_suffix_mask,
     )
 
 

src/openpi/models/model_test.py

@@ -75,6 +75,33 @@ def test_pi0_fast_lora_model():
     assert len(lora_state_elems) > 0
 
 
+def test_pi05_subtask_model():
+    key = jax.random.key(0)
+    config = pi0_config.Pi0Config(
+        pi05=True,
+        paligemma_variant="dummy",
+        action_expert_variant="dummy",
+        action_dim=4,
+        action_horizon=2,
+        max_token_len=16,
+        train_subtask_prediction=True,
+        sample_subtask_prediction=True,
+        max_subtask_len=3,
+    )
+    model = config.create(key)
+
+    batch_size = 2
+    obs, act = config.fake_obs(batch_size), config.fake_act(batch_size)
+
+    loss = nnx_utils.module_jit(model.compute_loss)(key, obs, act)
+    assert loss.shape == (batch_size, config.action_horizon)
+
+    outputs = nnx_utils.module_jit(model.sample_actions_with_subtask)(key, obs, num_steps=2)
+    assert outputs["actions"].shape == (batch_size, model.action_horizon, model.action_dim)
+    assert outputs["subtask_tokens"].shape == (batch_size, config.max_subtask_len)
+    assert outputs["subtask_token_mask"].shape == (batch_size, config.max_subtask_len)
+
+
 @pytest.mark.manual
 def test_model_restore():
     key = jax.random.key(0)