diff --git a/README.md b/README.md
index d16e43d479..c4ee98a0c3 100644
--- a/README.md
+++ b/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.
 
diff --git a/src/openpi/models/gemma.py b/src/openpi/models/gemma.py
index d1623d19d1..94454c2ba8 100644
--- a/src/openpi/models/gemma.py
+++ b/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,
diff --git a/src/openpi/models/model.py b/src/openpi/models/model.py
index 29618b4945..6862bd0e58 100644
--- a/src/openpi/models/model.py
+++ b/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,
     )
 
 
diff --git a/src/openpi/models/model_test.py b/src/openpi/models/model_test.py
index 495dc18b5f..8f29a5d741 100644
--- a/src/openpi/models/model_test.py
+++ b/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)
diff --git a/src/openpi/models/pi0.py b/src/openpi/models/pi0.py
index ae7c4590f3..ed011b3255 100644
--- a/src/openpi/models/pi0.py
+++ b/src/openpi/models/pi0.py
@@ -1,3 +1,4 @@
+import dataclasses
 import logging
 
 import einops
@@ -44,6 +45,28 @@ def make_attn_mask(input_mask, mask_ar):
     return jnp.logical_and(attn_mask, valid_mask)
 
 
+def _last_valid_indices(mask: at.Bool[at.Array, "b s"]) -> at.Int[at.Array, " b"]:
+    token_indices = jnp.arange(mask.shape[1])[None, :]
+    return jnp.max(jnp.where(mask, token_indices, -1), axis=-1)
+
+
+def _scatter_sequence(
+    base: at.Array,
+    positions: at.Int[at.Array, "b s"],
+    values: at.Array,
+    mask: at.Bool[at.Array, "b s"],
+) -> at.Array:
+    """Writes a masked token block into a padded batch sequence."""
+    max_len = base.shape[-1]
+    valid = mask & (positions >= 0) & (positions < max_len)
+    clipped_positions = jnp.clip(positions, 0, max_len - 1)
+    one_hot = jax.nn.one_hot(clipped_positions, max_len, dtype=jnp.int32)
+    put_mask = jnp.einsum("bs,bsn->bn", valid.astype(jnp.int32), one_hot).astype(jnp.bool_)
+    put_values = jnp.einsum("bs,bsn->bn", values.astype(jnp.int32) * valid.astype(jnp.int32), one_hot)
+    put_values = put_values.astype(jnp.bool_) if base.dtype == jnp.bool_ else put_values.astype(base.dtype)
+    return jnp.where(put_mask, put_values, base)
+
+
 @at.typecheck
 def posemb_sincos(
     pos: at.Real[at.Array, " b"], embedding_dim: int, min_period: float, max_period: float
@@ -67,6 +90,12 @@ class Pi0(_model.BaseModel):
     def __init__(self, config: pi0_config.Pi0Config, rngs: nnx.Rngs):
         super().__init__(config.action_dim, config.action_horizon, config.max_token_len)
         self.pi05 = config.pi05
+        self.train_subtask_prediction = config.train_subtask_prediction
+        self.sample_subtask_prediction = config.sample_subtask_prediction
+        self.subtask_loss_weight = config.subtask_loss_weight
+        self.max_subtask_len = config.max_subtask_len
+        self.subtask_temperature = config.subtask_temperature
+        self.subtask_eos_token = config.subtask_eos_token
         paligemma_config = _gemma.get_config(config.paligemma_variant)
         action_expert_config = _gemma.get_config(config.action_expert_variant)
         # TODO: rewrite gemma in NNX. For now, use bridge.
@@ -105,7 +134,7 @@ def __init__(self, config: pi0_config.Pi0Config, rngs: nnx.Rngs):
     @at.typecheck
     def embed_prefix(
         self, obs: _model.Observation
-    ) -> tuple[at.Float[at.Array, "b s emb"], at.Bool[at.Array, "b s"], at.Bool[at.Array, " s"]]:
+    ) -> tuple[at.Float[at.Array, "b s emb"], at.Bool[at.Array, "b s"], at.Bool[at.Array, "b s"]]:
         input_mask = []
         ar_mask = []
         tokens = []
@@ -122,18 +151,22 @@ def embed_prefix(
                 )
             )
             # image tokens attend to each other
-            ar_mask += [False] * image_tokens.shape[1]
+            ar_mask.append(jnp.zeros(image_tokens.shape[:2], dtype=jnp.bool_))
 
         # add language (aka tokenized inputs)
         if obs.tokenized_prompt is not None:
+            assert obs.tokenized_prompt_mask is not None, "Tokenized prompt mask is required"
             tokenized_inputs = self.PaliGemma.llm(obs.tokenized_prompt, method="embed")
             tokens.append(tokenized_inputs)
             input_mask.append(obs.tokenized_prompt_mask)
-            # full attention between image and language inputs
-            ar_mask += [False] * tokenized_inputs.shape[1]
+            if obs.token_ar_mask is None:
+                token_ar_mask = jnp.zeros_like(obs.tokenized_prompt_mask)
+            else:
+                token_ar_mask = obs.token_ar_mask.astype(jnp.bool_)
+            ar_mask.append(token_ar_mask)
         tokens = jnp.concatenate(tokens, axis=1)
         input_mask = jnp.concatenate(input_mask, axis=1)
-        ar_mask = jnp.array(ar_mask)
+        ar_mask = jnp.concatenate(ar_mask, axis=1)
         return tokens, input_mask, ar_mask
 
     @at.typecheck
@@ -142,19 +175,20 @@ def embed_suffix(
     ) -> tuple[
         at.Float[at.Array, "b s emb"],
         at.Bool[at.Array, "b s"],
-        at.Bool[at.Array, " s"],
+        at.Bool[at.Array, "b s"],
         at.Float[at.Array, "b emb"] | None,
     ]:
         input_mask = []
         ar_mask = []
         tokens = []
+        batch_size = obs.state.shape[0]
         if not self.pi05:
             # add a single state token
             state_token = self.state_proj(obs.state)[:, None, :]
             tokens.append(state_token)
             input_mask.append(jnp.ones((obs.state.shape[0], 1), dtype=jnp.bool_))
             # image/language inputs do not attend to state or actions
-            ar_mask += [True]
+            ar_mask.append(jnp.ones((batch_size, 1), dtype=jnp.bool_))
 
         action_tokens = self.action_in_proj(noisy_actions)
         # embed timestep using sine-cosine positional encoding with sensitivity in the range [0, 1]
@@ -179,10 +213,14 @@ def embed_suffix(
         tokens.append(action_expert_tokens)
         input_mask.append(jnp.ones(action_expert_tokens.shape[:2], dtype=jnp.bool_))
         # image/language/state inputs do not attend to action tokens
-        ar_mask += [True] + ([False] * (self.action_horizon - 1))
+        action_ar_mask = jnp.broadcast_to(
+            jnp.asarray([True] + ([False] * (self.action_horizon - 1))),
+            action_expert_tokens.shape[:2],
+        )
+        ar_mask.append(action_ar_mask)
         tokens = jnp.concatenate(tokens, axis=1)
         input_mask = jnp.concatenate(input_mask, axis=1)
-        ar_mask = jnp.array(ar_mask)
+        ar_mask = jnp.concatenate(ar_mask, axis=1)
         return tokens, input_mask, ar_mask, adarms_cond
 
     @override
@@ -203,7 +241,7 @@ def compute_loss(
         prefix_tokens, prefix_mask, prefix_ar_mask = self.embed_prefix(observation)
         suffix_tokens, suffix_mask, suffix_ar_mask, adarms_cond = self.embed_suffix(observation, x_t, time)
         input_mask = jnp.concatenate([prefix_mask, suffix_mask], axis=1)
-        ar_mask = jnp.concatenate([prefix_ar_mask, suffix_ar_mask], axis=0)
+        ar_mask = jnp.concatenate([prefix_ar_mask, suffix_ar_mask], axis=1)
         attn_mask = make_attn_mask(input_mask, ar_mask)
         positions = jnp.cumsum(input_mask, axis=1) - 1
         (prefix_out, suffix_out), _ = self.PaliGemma.llm(
@@ -211,7 +249,26 @@ def compute_loss(
         )
         v_t = self.action_out_proj(suffix_out[:, -self.action_horizon :])
 
-        return jnp.mean(jnp.square(v_t - u_t), axis=-1)
+        action_loss = jnp.mean(jnp.square(v_t - u_t), axis=-1)
+        if not self.train_subtask_prediction:
+            return action_loss
+        subtask_loss = self._compute_subtask_loss(prefix_out, observation)
+        return action_loss + self.subtask_loss_weight * subtask_loss[:, None]
+
+    def _compute_subtask_loss(
+        self, prefix_out: at.Float[at.Array, "b s emb"], observation: _model.Observation
+    ) -> at.Float[at.Array, " b"]:
+        assert observation.tokenized_prompt is not None, "Tokenized prompt is required for subtask prediction"
+        assert observation.token_loss_mask is not None, "Token loss mask is required for subtask prediction"
+
+        text_len = observation.tokenized_prompt.shape[1]
+        text_hidden = prefix_out[:, -text_len:-1]
+        logits = self.PaliGemma.llm(text_hidden, method="deembed")
+        logp = jax.nn.log_softmax(logits, axis=-1)
+        targets = observation.tokenized_prompt[:, 1:]
+        target_logp = jnp.take_along_axis(logp, targets[..., None], axis=-1).squeeze(-1)
+        loss_mask = observation.token_loss_mask[:, 1:]
+        return -jnp.sum(target_logp * loss_mask, axis=-1) / jnp.clip(jnp.sum(loss_mask, axis=-1), 1)
 
     @override
     def sample_actions(
@@ -223,6 +280,61 @@ def sample_actions(
         noise: at.Float[at.Array, "b ah ad"] | None = None,
     ) -> _model.Actions:
         observation = _model.preprocess_observation(None, observation, train=False)
+        prefix_tokens, prefix_mask, prefix_ar_mask = self.embed_prefix(observation)
+        return self._sample_actions_from_prefix(
+            rng,
+            observation,
+            prefix_tokens,
+            prefix_mask,
+            prefix_ar_mask,
+            num_steps=num_steps,
+            noise=noise,
+        )
+
+    def sample_actions_with_subtask(
+        self,
+        rng: at.KeyArrayLike,
+        observation: _model.Observation,
+        *,
+        num_steps: int | at.Int[at.Array, ""] = 10,
+        noise: at.Float[at.Array, "b ah ad"] | None = None,
+    ) -> dict[str, at.Array]:
+        if not self.sample_subtask_prediction:
+            return {"actions": self.sample_actions(rng, observation, num_steps=num_steps, noise=noise)}
+        if not self.pi05:
+            raise ValueError("Subtask prediction is only supported for pi0.5 models.")
+
+        subtask_rng, action_rng = jax.random.split(rng)
+        observation = _model.preprocess_observation(None, observation, train=False)
+        subtask_tokens, subtask_token_mask = self._sample_subtask_tokens(subtask_rng, observation)
+        action_observation = self._with_generated_subtask_prompt(observation, subtask_tokens, subtask_token_mask)
+        prefix_tokens, prefix_mask, prefix_ar_mask = self.embed_prefix(action_observation)
+        actions = self._sample_actions_from_prefix(
+            action_rng,
+            action_observation,
+            prefix_tokens,
+            prefix_mask,
+            prefix_ar_mask,
+            num_steps=num_steps,
+            noise=noise,
+        )
+        return {
+            "actions": actions,
+            "subtask_tokens": subtask_tokens,
+            "subtask_token_mask": subtask_token_mask,
+        }
+
+    def _sample_actions_from_prefix(
+        self,
+        rng: at.KeyArrayLike,
+        observation: _model.Observation,
+        prefix_tokens: at.Float[at.Array, "b s emb"],
+        prefix_mask: at.Bool[at.Array, "b s"],
+        prefix_ar_mask: at.Bool[at.Array, "b s"],
+        *,
+        num_steps: int | at.Int[at.Array, ""] = 10,
+        noise: at.Float[at.Array, "b ah ad"] | None = None,
+    ) -> _model.Actions:
         # note that we use the convention more common in diffusion literature, where t=1 is noise and t=0 is the target
         # distribution. yes, this is the opposite of the pi0 paper, and I'm sorry.
         dt = -1.0 / num_steps
@@ -231,7 +343,6 @@ def sample_actions(
             noise = jax.random.normal(rng, (batch_size, self.action_horizon, self.action_dim))
 
         # first fill KV cache with a forward pass of the prefix
-        prefix_tokens, prefix_mask, prefix_ar_mask = self.embed_prefix(observation)
         prefix_attn_mask = make_attn_mask(prefix_mask, prefix_ar_mask)
         positions = jnp.cumsum(prefix_mask, axis=1) - 1
         _, kv_cache = self.PaliGemma.llm([prefix_tokens, None], mask=prefix_attn_mask, positions=positions)
@@ -271,9 +382,120 @@ def step(carry):
             return x_t + dt * v_t, time + dt
 
         def cond(carry):
-            x_t, time = carry
+            _, time = carry
             # robust to floating-point error
             return time >= -dt / 2
 
         x_0, _ = jax.lax.while_loop(cond, step, (noise, 1.0))
         return x_0
+
+    def _sample_subtask_tokens(
+        self, rng: at.KeyArrayLike, observation: _model.Observation
+    ) -> tuple[at.Int[at.Array, "b l"], at.Bool[at.Array, "b l"]]:
+        assert observation.tokenized_prompt is not None, "Tokenized prompt is required for subtask prediction"
+        assert observation.tokenized_prompt_mask is not None, "Tokenized prompt mask is required for subtask prediction"
+        batch_size = observation.state.shape[0]
+        output_tokens = jnp.zeros((batch_size, self.max_subtask_len), dtype=jnp.int32)
+        output_mask = jnp.zeros((batch_size, self.max_subtask_len), dtype=jnp.bool_)
+        done = jnp.zeros((batch_size,), dtype=jnp.bool_)
+
+        def step(index, carry):
+            step_rng, tokens, token_mask, is_done = carry
+            step_rng, sample_rng = jax.random.split(step_rng)
+            prompt_observation = self._with_generated_subtask_prompt(
+                observation, tokens, token_mask, include_action_suffix=False
+            )
+            prefix_tokens, prefix_mask, prefix_ar_mask = self.embed_prefix(prompt_observation)
+            prefix_attn_mask = make_attn_mask(prefix_mask, prefix_ar_mask)
+            positions = jnp.cumsum(prefix_mask, axis=1) - 1
+            (prefix_out, _), _ = self.PaliGemma.llm([prefix_tokens, None], mask=prefix_attn_mask, positions=positions)
+            last_indices = _last_valid_indices(prefix_mask)
+            last_hidden = prefix_out[jnp.arange(prefix_out.shape[0]), last_indices][:, None, :]
+            logits = self.PaliGemma.llm(last_hidden, method="deembed")[:, 0]
+            if self.subtask_temperature > 0.0:
+                next_token = jax.random.categorical(sample_rng, logits / self.subtask_temperature, axis=-1)
+            else:
+                next_token = jnp.argmax(logits, axis=-1)
+            next_token = jnp.where(is_done, 0, next_token).astype(jnp.int32)
+            # Keep EOS in the conditioned prompt to match supervised training; output detokenization strips it.
+            next_is_valid = ~is_done
+            positions = jnp.broadcast_to(index, (batch_size, 1))
+            tokens = _scatter_sequence(tokens, positions, next_token[:, None], ~is_done[:, None])
+            token_mask = _scatter_sequence(token_mask, positions, next_is_valid[:, None], ~is_done[:, None])
+            is_done = is_done | (next_token == self.subtask_eos_token)
+            return step_rng, tokens, token_mask, is_done
+
+        _, output_tokens, output_mask, _ = jax.lax.fori_loop(
+            0, self.max_subtask_len, step, (rng, output_tokens, output_mask, done)
+        )
+        return output_tokens, output_mask
+
+    def _with_generated_subtask_prompt(
+        self,
+        observation: _model.Observation,
+        subtask_tokens: at.Int[at.Array, "b l"],
+        subtask_token_mask: at.Bool[at.Array, "b l"],
+        *,
+        include_action_suffix: bool = True,
+    ) -> _model.Observation:
+        assert observation.tokenized_prompt is not None, "Tokenized prompt is required for subtask prediction"
+        assert observation.tokenized_prompt_mask is not None, "Tokenized prompt mask is required for subtask prediction"
+        prompt_tokens = observation.tokenized_prompt
+        prompt_mask = observation.tokenized_prompt_mask
+        if observation.token_ar_mask is None:
+            prompt_ar_mask = jnp.zeros_like(prompt_mask)
+        else:
+            prompt_ar_mask = observation.token_ar_mask.astype(jnp.bool_)
+
+        prompt_len = jnp.sum(prompt_mask, axis=-1)
+        subtask_positions = prompt_len[:, None] + jnp.arange(subtask_tokens.shape[1])[None, :]
+        tokens = _scatter_sequence(prompt_tokens, subtask_positions, subtask_tokens, subtask_token_mask)
+        token_mask = _scatter_sequence(prompt_mask, subtask_positions, subtask_token_mask, subtask_token_mask)
+        token_ar_mask = _scatter_sequence(
+            prompt_ar_mask,
+            subtask_positions,
+            jnp.ones_like(subtask_token_mask),
+            subtask_token_mask,
+        )
+
+        if include_action_suffix:
+            assert observation.tokenized_action_suffix is not None, (
+                "Tokenized action suffix is required for subtask prediction"
+            )
+            assert observation.tokenized_action_suffix_mask is not None, (
+                "Tokenized action suffix mask is required for subtask prediction"
+            )
+            subtask_len = jnp.sum(subtask_token_mask, axis=-1)
+            suffix_positions = (
+                prompt_len[:, None]
+                + subtask_len[:, None]
+                + jnp.arange(observation.tokenized_action_suffix.shape[1])[None, :]
+            )
+            tokens = _scatter_sequence(
+                tokens,
+                suffix_positions,
+                observation.tokenized_action_suffix,
+                observation.tokenized_action_suffix_mask,
+            )
+            token_mask = _scatter_sequence(
+                token_mask,
+                suffix_positions,
+                observation.tokenized_action_suffix_mask,
+                observation.tokenized_action_suffix_mask,
+            )
+            token_ar_mask = _scatter_sequence(
+                token_ar_mask,
+                suffix_positions,
+                jnp.ones_like(observation.tokenized_action_suffix_mask),
+                observation.tokenized_action_suffix_mask,
+            )
+
+        return dataclasses.replace(
+            observation,
+            tokenized_prompt=tokens,
+            tokenized_prompt_mask=token_mask,
+            token_ar_mask=token_ar_mask.astype(jnp.int32),
+            token_loss_mask=None,
+            tokenized_action_suffix=None,
+            tokenized_action_suffix_mask=None,
+        )
diff --git a/src/openpi/models/pi0_config.py b/src/openpi/models/pi0_config.py
index 584d83f199..5bfbad31c5 100644
--- a/src/openpi/models/pi0_config.py
+++ b/src/openpi/models/pi0_config.py
@@ -32,6 +32,14 @@ class Pi0Config(_model.BaseModelConfig):
     # This config option is not used directly by the model, but it is read by the ModelTransformFactory.
     discrete_state_input: bool = None  # type: ignore
 
+    # Optional pi0.5 two-stage inference support. Disabled by default to preserve released checkpoint behavior.
+    train_subtask_prediction: bool = False
+    sample_subtask_prediction: bool = False
+    subtask_loss_weight: float = 1.0
+    max_subtask_len: int = 64
+    subtask_temperature: float = 0.0
+    subtask_eos_token: int = 1
+
     pytorch_compile_mode: str | None = "max-autotune"
 
     def __post_init__(self):
@@ -46,6 +54,18 @@ def __post_init__(self):
                 "max-autotune",
                 "max-autotune-no-cudagraphs",
             ]
+        if (self.train_subtask_prediction or self.sample_subtask_prediction) and not self.pi05:
+            raise ValueError("Subtask prediction is only supported for pi0.5 models.")
+        if self.train_subtask_prediction and self.subtask_loss_weight <= 0:
+            raise ValueError("subtask_loss_weight must be positive when train_subtask_prediction is enabled.")
+        if self.max_subtask_len <= 0:
+            raise ValueError("max_subtask_len must be positive.")
+        if (
+            self.train_subtask_prediction or self.sample_subtask_prediction
+        ) and self.max_subtask_len >= self.max_token_len:
+            raise ValueError("max_subtask_len must be smaller than max_token_len.")
+        if self.subtask_temperature < 0:
+            raise ValueError("subtask_temperature must be non-negative.")
 
     @property
     @override
@@ -56,7 +76,7 @@ def model_type(self) -> _model.ModelType:
 
     @override
     def create(self, rng: at.KeyArrayLike) -> "Pi0":
-        from openpi.models.pi0 import Pi0
+        from openpi.models.pi0 import Pi0  # noqa: PLC0415
 
         return Pi0(self, rngs=nnx.Rngs(rng))
 
@@ -65,21 +85,35 @@ def inputs_spec(self, *, batch_size: int = 1) -> tuple[_model.Observation, _mode
         image_spec = jax.ShapeDtypeStruct([batch_size, *_model.IMAGE_RESOLUTION, 3], jnp.float32)
         image_mask_spec = jax.ShapeDtypeStruct([batch_size], jnp.bool_)
 
+        observation_kwargs = {
+            "images": {
+                "base_0_rgb": image_spec,
+                "left_wrist_0_rgb": image_spec,
+                "right_wrist_0_rgb": image_spec,
+            },
+            "image_masks": {
+                "base_0_rgb": image_mask_spec,
+                "left_wrist_0_rgb": image_mask_spec,
+                "right_wrist_0_rgb": image_mask_spec,
+            },
+            "state": jax.ShapeDtypeStruct([batch_size, self.action_dim], jnp.float32),
+            "tokenized_prompt": jax.ShapeDtypeStruct([batch_size, self.max_token_len], jnp.int32),
+            "tokenized_prompt_mask": jax.ShapeDtypeStruct([batch_size, self.max_token_len], bool),
+        }
+        if self.train_subtask_prediction or self.sample_subtask_prediction:
+            observation_kwargs |= {
+                "token_ar_mask": jax.ShapeDtypeStruct([batch_size, self.max_token_len], jnp.int32),
+                "token_loss_mask": jax.ShapeDtypeStruct([batch_size, self.max_token_len], jnp.bool_),
+            }
+        if self.sample_subtask_prediction:
+            observation_kwargs |= {
+                "tokenized_action_suffix": jax.ShapeDtypeStruct([batch_size, self.max_token_len], jnp.int32),
+                "tokenized_action_suffix_mask": jax.ShapeDtypeStruct([batch_size, self.max_token_len], jnp.bool_),
+            }
+
         with at.disable_typechecking():
             observation_spec = _model.Observation(
-                images={
-                    "base_0_rgb": image_spec,
-                    "left_wrist_0_rgb": image_spec,
-                    "right_wrist_0_rgb": image_spec,
-                },
-                image_masks={
-                    "base_0_rgb": image_mask_spec,
-                    "left_wrist_0_rgb": image_mask_spec,
-                    "right_wrist_0_rgb": image_mask_spec,
-                },
-                state=jax.ShapeDtypeStruct([batch_size, self.action_dim], jnp.float32),
-                tokenized_prompt=jax.ShapeDtypeStruct([batch_size, self.max_token_len], jnp.int32),
-                tokenized_prompt_mask=jax.ShapeDtypeStruct([batch_size, self.max_token_len], bool),
+                **observation_kwargs,
             )
         action_spec = jax.ShapeDtypeStruct([batch_size, self.action_horizon, self.action_dim], jnp.float32)
 
diff --git a/src/openpi/models/tokenizer.py b/src/openpi/models/tokenizer.py
index 8a4966d629..7f5719174b 100644
--- a/src/openpi/models/tokenizer.py
+++ b/src/openpi/models/tokenizer.py
@@ -10,6 +10,8 @@
 import openpi.models.utils.fsq_tokenizer as fsq_tokenizer
 import openpi.shared.download as download
 
+PALIGEMMA_EOS_TOKEN = 1
+
 
 class PaligemmaTokenizer:
     def __init__(self, max_len: int = 48):
@@ -47,6 +49,97 @@ def tokenize(self, prompt: str, state: np.ndarray | None = None) -> tuple[np.nda
 
         return np.asarray(tokens), np.asarray(mask)
 
+    def tokenize_subtask_training(
+        self, prompt: str, subtask: str, state: np.ndarray | None = None
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+        """Tokenize pi0.5 inputs for supervised subtask prediction plus action conditioning."""
+        prefix_tokens = self._tokenizer.encode(self._subtask_prefix(prompt, state), add_bos=True)
+        subtask_tokens = self._tokenizer.encode(self._clean_text(subtask), add_eos=True)
+        action_suffix_tokens = self._tokenizer.encode("\nAction: ")
+
+        tokens = prefix_tokens + subtask_tokens + action_suffix_tokens
+        token_mask = [True] * len(tokens)
+        ar_mask = [0] * len(prefix_tokens) + [1] * (len(subtask_tokens) + len(action_suffix_tokens))
+        loss_mask = [False] * len(prefix_tokens) + [True] * len(subtask_tokens) + [False] * len(action_suffix_tokens)
+
+        return self._pad_token_arrays(tokens, token_mask, ar_mask, loss_mask)
+
+    def tokenize_subtask_inference(
+        self, prompt: str, state: np.ndarray | None = None
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+        """Tokenize pi0.5 inputs for two-stage inference.
+
+        The model first receives a `Subtask:` prefix, autoregressively predicts subtask tokens, then appends the
+        returned action cue before flow-matching action decoding.
+        """
+        prefix_tokens = self._tokenizer.encode(self._subtask_prefix(prompt, state), add_bos=True)
+        action_suffix_tokens = self._tokenizer.encode("\nAction: ")
+
+        prefix_tokens, prefix_mask, prefix_ar_mask, prefix_loss_mask = self._pad_token_arrays(
+            prefix_tokens,
+            [True] * len(prefix_tokens),
+            [0] * len(prefix_tokens),
+            [False] * len(prefix_tokens),
+        )
+        suffix_tokens, suffix_mask = (
+            self._pad_sequence(action_suffix_tokens, np.int32),
+            self._pad_sequence([True] * len(action_suffix_tokens), np.bool_),
+        )
+        return prefix_tokens, prefix_mask, prefix_ar_mask, prefix_loss_mask, suffix_tokens, suffix_mask
+
+    def detokenize(self, tokens: np.ndarray, mask: np.ndarray | None = None) -> str:
+        token_list = np.asarray(tokens).astype(np.int32).tolist()
+        if mask is not None:
+            mask_list = np.asarray(mask).astype(bool).tolist()
+            token_list = [token for token, valid in zip(token_list, mask_list, strict=True) if valid]
+        else:
+            token_list = [token for token in token_list if token != 0]
+
+        if PALIGEMMA_EOS_TOKEN in token_list:
+            token_list = token_list[: token_list.index(PALIGEMMA_EOS_TOKEN)]
+        return self._tokenizer.decode(token_list).strip()
+
+    def _clean_text(self, text: str) -> str:
+        return text.strip().replace("_", " ").replace("\n", " ")
+
+    def _state_str(self, state: np.ndarray) -> str:
+        discretized_state = np.digitize(state, bins=np.linspace(-1, 1, 256 + 1)[:-1]) - 1
+        return " ".join(map(str, discretized_state))
+
+    def _subtask_prefix(self, prompt: str, state: np.ndarray | None) -> str:
+        cleaned_text = self._clean_text(prompt)
+        if state is None:
+            return f"Task: {cleaned_text};\nSubtask: "
+        return f"Task: {cleaned_text}, State: {self._state_str(state)};\nSubtask: "
+
+    def _pad_sequence(self, values: list[int] | list[bool], dtype) -> np.ndarray:
+        values_len = len(values)
+        if values_len < self._max_len:
+            values = values + [False] * (self._max_len - values_len)
+        else:
+            values = values[: self._max_len]
+        return np.asarray(values, dtype=dtype)
+
+    def _pad_token_arrays(
+        self,
+        tokens: list[int],
+        token_mask: list[bool],
+        ar_mask: list[int],
+        loss_mask: list[bool],
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+        tokens_len = len(tokens)
+        if tokens_len > self._max_len:
+            logging.warning(
+                f"Token length ({tokens_len}) exceeds max length ({self._max_len}), truncating. "
+                "Consider increasing the `max_token_len` in your model config if this happens frequently."
+            )
+        return (
+            self._pad_sequence(tokens, np.int32),
+            self._pad_sequence(token_mask, np.bool_),
+            self._pad_sequence(ar_mask, np.int32),
+            self._pad_sequence(loss_mask, np.bool_),
+        )
+
 
 class FASTTokenizer:
     def __init__(self, max_len: int = 256, fast_tokenizer_path: str = "physical-intelligence/fast"):
diff --git a/src/openpi/models/tokenizer_test.py b/src/openpi/models/tokenizer_test.py
index 092c4ae3b2..14cb7223be 100644
--- a/src/openpi/models/tokenizer_test.py
+++ b/src/openpi/models/tokenizer_test.py
@@ -3,6 +3,26 @@
 from openpi.models import tokenizer as _tokenizer
 
 
+class _FakeSentencePiece:
+    def encode(self, text, **kwargs):
+        tokens = [sum(map(ord, piece)) % 97 + 2 for piece in text.split()]
+        if kwargs.get("add_bos", False):
+            tokens = [0, *tokens]
+        if kwargs.get("add_eos", False):
+            tokens = [*tokens, _tokenizer.PALIGEMMA_EOS_TOKEN]
+        return tokens
+
+    def decode(self, tokens):
+        return " ".join(map(str, tokens))
+
+
+def _make_fake_paligemma_tokenizer(max_len=64):
+    tokenizer = _tokenizer.PaligemmaTokenizer.__new__(_tokenizer.PaligemmaTokenizer)
+    tokenizer._max_len = max_len  # noqa: SLF001
+    tokenizer._tokenizer = _FakeSentencePiece()  # noqa: SLF001
+    return tokenizer
+
+
 def test_tokenize():
     tokenizer = _tokenizer.PaligemmaTokenizer(max_len=10)
     tokens, masks = tokenizer.tokenize("Hello, world!")
@@ -11,6 +31,36 @@ def test_tokenize():
     assert masks.shape == (10,)
 
 
+def test_subtask_tokenize_training():
+    tokenizer = _make_fake_paligemma_tokenizer(max_len=64)
+    tokens, token_masks, ar_masks, loss_masks = tokenizer.tokenize_subtask_training(
+        "Put apple on plate", "reach for the apple", np.zeros(3, dtype=np.float32)
+    )
+
+    assert tokens.shape == (64,)
+    assert token_masks.shape == (64,)
+    assert ar_masks.shape == (64,)
+    assert loss_masks.shape == (64,)
+    assert np.sum(loss_masks) > 0
+    assert np.all(ar_masks[loss_masks])
+
+
+def test_subtask_tokenize_inference():
+    tokenizer = _make_fake_paligemma_tokenizer(max_len=64)
+    tokens, token_masks, ar_masks, loss_masks, suffix, suffix_masks = tokenizer.tokenize_subtask_inference(
+        "Put apple on plate", np.zeros(3, dtype=np.float32)
+    )
+
+    assert tokens.shape == (64,)
+    assert token_masks.shape == (64,)
+    assert ar_masks.shape == (64,)
+    assert loss_masks.shape == (64,)
+    assert suffix.shape == (64,)
+    assert suffix_masks.shape == (64,)
+    assert not np.any(loss_masks)
+    assert np.sum(suffix_masks) > 0
+
+
 def test_fast_tokenizer():
     prompt = "Hello, world!"
     state = np.random.rand(5).astype(np.float32)
diff --git a/src/openpi/models_pytorch/pi0_pytorch.py b/src/openpi/models_pytorch/pi0_pytorch.py
index e68ddb7cc0..a9107869cd 100644
--- a/src/openpi/models_pytorch/pi0_pytorch.py
+++ b/src/openpi/models_pytorch/pi0_pytorch.py
@@ -86,6 +86,9 @@ def __init__(self, config):
         super().__init__()
         self.config = config
         self.pi05 = config.pi05
+        self.sample_subtask_prediction = config.sample_subtask_prediction
+        if config.train_subtask_prediction or config.sample_subtask_prediction:
+            raise NotImplementedError("pi0.5 subtask prediction is currently implemented for JAX models only.")
 
         paligemma_config = _gemma.get_config(config.paligemma_variant)
         action_expert_config = _gemma.get_config(config.action_expert_variant)
@@ -117,7 +120,7 @@ def __init__(self, config):
 
         msg = "transformers_replace is not installed correctly. Please install it with `uv pip install transformers==4.53.2` and `cp -r ./src/openpi/models_pytorch/transformers_replace/* .venv/lib/python3.11/site-packages/transformers/`."
         try:
-            from transformers.models.siglip import check
+            from transformers.models.siglip import check  # noqa: PLC0415
 
             if not check.check_whether_transformers_replace_is_installed_correctly():
                 raise ValueError(msg)
diff --git a/src/openpi/policies/policy.py b/src/openpi/policies/policy.py
index b9b708bdca..06eeb796b3 100644
--- a/src/openpi/policies/policy.py
+++ b/src/openpi/policies/policy.py
@@ -56,12 +56,20 @@ def __init__(
         self._pytorch_device = pytorch_device
 
         if self._is_pytorch_model:
+            if getattr(model, "sample_subtask_prediction", False):
+                raise NotImplementedError("pi0.5 subtask prediction is currently implemented for JAX policies only.")
             self._model = self._model.to(pytorch_device)
             self._model.eval()
             self._sample_actions = model.sample_actions
+            self._sample_actions_returns_dict = False
         else:
             # JAX model setup
-            self._sample_actions = nnx_utils.module_jit(model.sample_actions)
+            if getattr(model, "sample_subtask_prediction", False):
+                self._sample_actions = nnx_utils.module_jit(model.sample_actions_with_subtask)
+                self._sample_actions_returns_dict = True
+            else:
+                self._sample_actions = nnx_utils.module_jit(model.sample_actions)
+                self._sample_actions_returns_dict = False
             self._rng = rng or jax.random.key(0)
 
     @override
@@ -89,10 +97,14 @@ def infer(self, obs: dict, *, noise: np.ndarray | None = None) -> dict:  # type:
 
         observation = _model.Observation.from_dict(inputs)
         start_time = time.monotonic()
-        outputs = {
-            "state": inputs["state"],
-            "actions": self._sample_actions(sample_rng_or_pytorch_device, observation, **sample_kwargs),
-        }
+        model_outputs = self._sample_actions(sample_rng_or_pytorch_device, observation, **sample_kwargs)
+        if self._sample_actions_returns_dict:
+            outputs = {"state": inputs["state"], **model_outputs}
+        else:
+            outputs = {
+                "state": inputs["state"],
+                "actions": model_outputs,
+            }
         model_time = time.monotonic() - start_time
         if self._is_pytorch_model:
             outputs = jax.tree.map(lambda x: np.asarray(x[0, ...].detach().cpu()), outputs)
diff --git a/src/openpi/training/config.py b/src/openpi/training/config.py
index 4ca47e1286..af3f23951d 100644
--- a/src/openpi/training/config.py
+++ b/src/openpi/training/config.py
@@ -125,12 +125,30 @@ def __call__(self, model_config: _model.BaseModelConfig) -> _transforms.Group:
                 )
             case _model.ModelType.PI05:
                 assert isinstance(model_config, pi0_config.Pi0Config)
+                tokenizer = _tokenizer.PaligemmaTokenizer(model_config.max_token_len)
+                if model_config.train_subtask_prediction or model_config.sample_subtask_prediction:
+                    return _transforms.Group(
+                        inputs=[
+                            _transforms.InjectDefaultPrompt(self.default_prompt),
+                            _transforms.ResizeImages(224, 224),
+                            _transforms.TokenizePi05SubtaskInputs(
+                                tokenizer,
+                                discrete_state_input=model_config.discrete_state_input,
+                                train_subtask_prediction=model_config.train_subtask_prediction,
+                                sample_subtask_prediction=model_config.sample_subtask_prediction,
+                            ),
+                            _transforms.PadStatesAndActions(model_config.action_dim),
+                        ],
+                        outputs=(
+                            [_transforms.DetokenizeSubtask(tokenizer)] if model_config.sample_subtask_prediction else []
+                        ),
+                    )
                 return _transforms.Group(
                     inputs=[
                         _transforms.InjectDefaultPrompt(self.default_prompt),
                         _transforms.ResizeImages(224, 224),
                         _transforms.TokenizePrompt(
-                            _tokenizer.PaligemmaTokenizer(model_config.max_token_len),
+                            tokenizer,
                             discrete_state_input=model_config.discrete_state_input,
                         ),
                         _transforms.PadStatesAndActions(model_config.action_dim),
diff --git a/src/openpi/transforms.py b/src/openpi/transforms.py
index 272375ea95..d0337ab018 100644
--- a/src/openpi/transforms.py
+++ b/src/openpi/transforms.py
@@ -266,6 +266,65 @@ def __call__(self, data: DataDict) -> DataDict:
         return {**data, "tokenized_prompt": tokens, "tokenized_prompt_mask": token_masks}
 
 
+@dataclasses.dataclass(frozen=True)
+class TokenizePi05SubtaskInputs(DataTransformFn):
+    tokenizer: _tokenizer.PaligemmaTokenizer
+    discrete_state_input: bool = False
+    train_subtask_prediction: bool = False
+    sample_subtask_prediction: bool = False
+
+    def __call__(self, data: DataDict) -> DataDict:
+        if (prompt := data.pop("prompt", None)) is None:
+            raise ValueError("Prompt is required")
+
+        if self.discrete_state_input:
+            if (state := data.get("state", None)) is None:
+                raise ValueError("State is required.")
+        else:
+            state = None
+
+        if not isinstance(prompt, str):
+            prompt = prompt.item()
+
+        subtask = data.pop("subtask", None)
+        if subtask is not None:
+            if not isinstance(subtask, str):
+                subtask = subtask.item()
+            tokens, token_mask, ar_mask, loss_mask = self.tokenizer.tokenize_subtask_training(prompt, subtask, state)
+            return {
+                **data,
+                "tokenized_prompt": tokens,
+                "tokenized_prompt_mask": token_mask,
+                "token_ar_mask": ar_mask,
+                "token_loss_mask": loss_mask,
+            }
+
+        if self.train_subtask_prediction and "actions" in data:
+            raise ValueError("Subtask is required when train_subtask_prediction is enabled.")
+
+        if self.sample_subtask_prediction:
+            (
+                tokens,
+                token_mask,
+                ar_mask,
+                loss_mask,
+                action_suffix,
+                action_suffix_mask,
+            ) = self.tokenizer.tokenize_subtask_inference(prompt, state)
+            return {
+                **data,
+                "tokenized_prompt": tokens,
+                "tokenized_prompt_mask": token_mask,
+                "token_ar_mask": ar_mask,
+                "token_loss_mask": loss_mask,
+                "tokenized_action_suffix": action_suffix,
+                "tokenized_action_suffix_mask": action_suffix_mask,
+            }
+
+        tokens, token_mask = self.tokenizer.tokenize(prompt, state)
+        return {**data, "tokenized_prompt": tokens, "tokenized_prompt_mask": token_mask}
+
+
 @dataclasses.dataclass(frozen=True)
 class TokenizeFASTInputs(DataTransformFn):
     tokenizer: _tokenizer.FASTTokenizer
@@ -306,6 +365,18 @@ def __call__(self, data: DataDict) -> DataDict:
         }
 
 
+@dataclasses.dataclass(frozen=True)
+class DetokenizeSubtask(DataTransformFn):
+    tokenizer: _tokenizer.PaligemmaTokenizer
+
+    def __call__(self, data: DataDict) -> DataDict:
+        if "subtask_tokens" not in data:
+            return data
+        tokens = data["subtask_tokens"]
+        mask = data.get("subtask_token_mask")
+        return {**data, "subtask": self.tokenizer.detokenize(tokens.astype(np.int32), mask)}
+
+
 @dataclasses.dataclass(frozen=True)
 class PromptFromLeRobotTask(DataTransformFn):
     """Extracts a prompt from the current LeRobot dataset task."""
diff --git a/src/openpi/transforms_test.py b/src/openpi/transforms_test.py
index 54b1fe94f9..6773a83264 100644
--- a/src/openpi/transforms_test.py
+++ b/src/openpi/transforms_test.py
@@ -5,6 +5,26 @@
 import openpi.transforms as _transforms
 
 
+class _FakeSentencePiece:
+    def encode(self, text, **kwargs):
+        tokens = [sum(map(ord, piece)) % 97 + 2 for piece in text.split()]
+        if kwargs.get("add_bos", False):
+            tokens = [0, *tokens]
+        if kwargs.get("add_eos", False):
+            tokens = [*tokens, _tokenizer.PALIGEMMA_EOS_TOKEN]
+        return tokens
+
+    def decode(self, tokens):
+        return " ".join(map(str, tokens))
+
+
+def _make_fake_paligemma_tokenizer(max_len=64):
+    tokenizer = _tokenizer.PaligemmaTokenizer.__new__(_tokenizer.PaligemmaTokenizer)
+    tokenizer._max_len = max_len  # noqa: SLF001
+    tokenizer._tokenizer = _FakeSentencePiece()  # noqa: SLF001
+    return tokenizer
+
+
 def test_repack_transform():
     transform = _transforms.RepackTransform(
         structure={
@@ -78,6 +98,71 @@ def test_tokenize_prompt():
     assert np.allclose(tok_mask, data["tokenized_prompt_mask"])
 
 
+def test_tokenize_pi05_subtask_training():
+    tokenizer = _make_fake_paligemma_tokenizer(max_len=64)
+    transform = _transforms.TokenizePi05SubtaskInputs(
+        tokenizer,
+        discrete_state_input=True,
+        train_subtask_prediction=True,
+    )
+
+    data = transform(
+        {
+            "prompt": "Put apple on plate",
+            "subtask": "reach for the apple",
+            "state": np.zeros(3, dtype=np.float32),
+            "actions": np.zeros((2, 3), dtype=np.float32),
+        }
+    )
+
+    assert data["tokenized_prompt"].shape == (64,)
+    assert data["tokenized_prompt_mask"].shape == (64,)
+    assert data["token_ar_mask"].shape == (64,)
+    assert data["token_loss_mask"].shape == (64,)
+    assert "tokenized_action_suffix" not in data
+    assert np.sum(data["token_loss_mask"]) > 0
+
+
+def test_tokenize_pi05_subtask_inference():
+    tokenizer = _make_fake_paligemma_tokenizer(max_len=64)
+    transform = _transforms.TokenizePi05SubtaskInputs(
+        tokenizer,
+        discrete_state_input=True,
+        sample_subtask_prediction=True,
+    )
+
+    data = transform(
+        {
+            "prompt": "Put apple on plate",
+            "state": np.zeros(3, dtype=np.float32),
+        }
+    )
+
+    assert data["tokenized_prompt"].shape == (64,)
+    assert data["token_ar_mask"].shape == (64,)
+    assert data["token_loss_mask"].shape == (64,)
+    assert data["tokenized_action_suffix"].shape == (64,)
+    assert data["tokenized_action_suffix_mask"].shape == (64,)
+    assert not np.any(data["token_loss_mask"])
+
+
+def test_tokenize_pi05_subtask_training_requires_subtask():
+    tokenizer = _make_fake_paligemma_tokenizer(max_len=64)
+    transform = _transforms.TokenizePi05SubtaskInputs(
+        tokenizer,
+        train_subtask_prediction=True,
+    )
+
+    with pytest.raises(ValueError, match="Subtask is required"):
+        transform(
+            {
+                "prompt": "Put apple on plate",
+                "state": np.zeros(3, dtype=np.float32),
+                "actions": np.zeros((2, 3), dtype=np.float32),
+            }
+        )
+
+
 def test_tokenize_no_prompt():
     transform = _transforms.TokenizePrompt(_tokenizer.PaligemmaTokenizer())