opentau.policies.pi07.high_level_planner.modeling_pi07_high_level

π07 High-Level Planner: A Vision-Language Model for Memory and Subtask Prediction.

This module implements the high-level planner for π07, built on top of the Gemma 3 VLM backbone (with a Gemma-v1 action expert; see opentau.policies.pi07.gemma3_with_expert). Given images, language instructions, robot state, and past memory, the planner autoregressively predicts updated memory and a subtask string.

Functions

`make_att_2d_masks`(pad_masks, att_masks[, ...])	Creates a 2-D attention mask given padding and 1-D attention masks.
`resize_with_pad`(img, width, height[, pad_value])	Resizes an image to fit within the specified dimensions while maintaining aspect ratio, and pads the remaining area with the specified value.

Classes

`PI07HighLevelPlannerModel`(config[, ...])	π07 High-Level Planner inner model.
`PI07HighLevelPlannerPolicy`(config[, ...])	Policy wrapper for the π07 high-level planner.

class opentau.policies.pi07.high_level_planner.modeling_pi07_high_level.PI07HighLevelPlannerModel(config: PI07HighLevelPlannerConfig, discrete_action_vocab_size: int | None = None)[source]

Bases: Module

π07 High-Level Planner inner model.

Uses the Gemma 3 VLM backbone to encode images and a composite language prompt (task + past context) and optional episode metadata, with fixed tokenizer spans ";\n ", "Updated Memory: ", and (in full training runs) "Subtask: " before the predicted text, then autoregressively predicts updated memory and subtask text:

Updated memory — next-token CE over memory_max_length slots after the "Updated Memory: " span.
Subtask (response) — next-token CE over response_max_length slots after the "Subtask: " span (training).

Inference mirrors training by inserting the live "Subtask: " token IDs into the KV cache after memory decoding and before response decoding.

Architecture (rough dataflow):

┌───────────────────────────────────────────┐
│     response content (subtask text)       │
│                   ▲                       │
│  memory, ``Subtask: ``, lang, ``";\n "``, images, … │
│     ┌───────────────────────┐             │
│     │        Gemma 3         │            │
│     │  (autoregressive LM)   │            │
│     └────────────────────────┘            │
└───────────────────────────────────────────┘

Parameters:

config – High-level planner configuration.
discrete_action_vocab_size – Vocabulary size for the discrete action tokenizer (passed through to Gemma3WithExpertModel).

__init__(config: PI07HighLevelPlannerConfig, discrete_action_vocab_size: int | None = None)[source]

Initializes the PI07HighLevelPlannerModel.

Parameters:

config – High-level planner configuration.
discrete_action_vocab_size – Vocabulary size for the discrete action tokenizer (passed through to Gemma3WithExpertModel).

embed_prefix(images: list[Tensor], img_masks: list[Tensor], lang_tokens: Tensor, lang_masks: Tensor, response_tokens: Tensor | None = None, response_masks: Tensor | None = None, memory_tokens: Tensor | None = None, memory_masks: Tensor | None = None, metadata_tokens: Tensor | None = None, metadata_masks: Tensor | None = None) → tuple[Tensor, Tensor, Tensor][source]

Embeds and concatenates all prefix modalities for the transformer.

Embeds images with SigLIP and language/metadata/memory/response spans with the Gemma 3 embedding layer. Concatenation order (training when memory and response are provided):

";\n " is gated on real metadata content — it serves as the metadata → "Updated Memory:" separator, so when no metadata is provided (or every sample’s metadata is fully padded) there is nothing to terminate and emitting it would dangle spurious tokens. The gate fires when metadata_tokens is None or every entry of metadata_masks is False (matching the low-level component’s metadata_masks.any() semantics so that training paths with all-padded metadata cleanly drop the metadata + prefix-end blocks). Like the low-level component, the decision is batch-wide: any sample with real metadata keeps both blocks present for the whole batch. The "Updated Memory: " anchor itself is unconditional because inference relies on it as the autoregressive starting point for memory decoding (memory_tokens is None at inference by design).

When memory_tokens / response_tokens are omitted (inference), only the fixed spans before those segments are present; memory and subtask text are filled in via KV-cache decoding plus an explicit "Subtask: " injection before response AR.

Attention pattern (via att_masks cumsums) – paper §VI.B says “observation tokens use bidirectional attention within themselves … the following text tokens use causal attention”:

Image patches: one bidirectional block shared across all cameras ([0] * N).

All text spans (language, metadata, ";\n ", "Updated Memory: ", "Subtask: ", memory content, response content): causal – one block per token ([1] * N).

Parameters:

images – List of image tensors, one per camera.
img_masks – List of boolean masks indicating real vs. padded images.
lang_tokens – Language token IDs of shape (B, prompt_max_length).
lang_masks – Boolean attention mask for language tokens.
response_tokens – Optional subtask response token IDs of shape (B, response_max_length). Provided during training.
response_masks – Optional boolean mask for response tokens.
memory_tokens – Optional updated memory token IDs of shape (B, memory_max_length). Provided during training.
memory_masks – Optional boolean mask for memory tokens.
metadata_tokens – Optional metadata token IDs of shape (B, metadata_max_length).
metadata_masks – Optional boolean mask for metadata tokens.

Returns:

embs: Concatenated embeddings (B, total_seq_len, D).
pad_masks: Boolean padding mask (B, total_seq_len).
att_masks: 1-D attention pattern (B, total_seq_len) used by make_att_2d_masks().

Return type:

A tuple (embs, pad_masks, att_masks) where

forward(images: list[Tensor], img_masks: list[Tensor], lang_tokens: Tensor, lang_masks: Tensor, response_tokens: Tensor | None = None, response_masks: Tensor | None = None, memory_tokens: Tensor | None = None, memory_masks: Tensor | None = None, metadata_tokens: Tensor | None = None, metadata_masks: Tensor | None = None) → dict[str, Tensor][source]

Training forward pass: embeds all modalities and computes CE losses.

The prefix matches embed_prefix() when memory and response tensors are set: fixed separators ";\n ", "Updated Memory: ", and "Subtask: " appear in addition to metadata, memory_tokens, and response_tokens. CE slices use negative offsets from the sequence tail, relying on config.subtask_indicator_max_length so memory logits align with memory contents even though "Subtask: " sits between memory and response text.

Parameters:

images – List of image tensors, one per camera.
img_masks – List of boolean masks for real vs. padded images.
lang_tokens – Language token IDs (B, prompt_max_length).
lang_masks – Boolean attention mask for language tokens.
response_tokens – Subtask response token IDs (B, response_max_length).
response_masks – Boolean mask for response tokens.
memory_tokens – Updated memory token IDs (B, memory_max_length).
memory_masks – Boolean mask for memory tokens.
metadata_tokens – Optional metadata token IDs (B, metadata_max_length).
metadata_masks – Optional boolean mask for metadata tokens.

Returns:

A dict with "MSE" (zero tensor, for interface compatibility) and "CE" (sum of memory and response cross-entropy losses).

infer_autoregressive(prefix_out: Tensor, prefix_embs: Tensor, prefix_pad_masks: Tensor, prefix_att_masks: Tensor, past_key_values: list[dict[str, Tensor]], prefix_offsets: Tensor, tokens: Tensor, auto_step: int, bsize: int, device: device) → tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, list[dict[str, Tensor]]][source]

Performs one autoregressive generation step.

At auto_step == 0 a <bos> token seeds the generation; on subsequent steps the most-recent logits are argmax-decoded into the next token. Once an <eos> or <pad> token appears in the accumulated sequence the remaining positions are filled with padding.

The method updates the KV-cache, prefix embeddings, and masks so that the next call can attend to all previously generated tokens.

Parameters:

prefix_out – Transformer output from the previous step (B, 1, D) or (B, seq, D) on the first call.
prefix_embs – Running concatenation of all embeddings fed to the transformer so far (B, current_seq, D).
prefix_pad_masks – Boolean padding mask (B, current_seq).
prefix_att_masks – 1-D attention pattern (B, current_seq).
past_key_values – KV-cache list from previous transformer calls.
prefix_offsets – Position ID offsets (B, 1) tracking the current absolute position for each batch element.
tokens – Accumulated generated token IDs (B, steps_so_far).
auto_step – Current step index (0-based).
bsize – Batch size.
device – Torch device for tensor creation.

Returns:

(prefix_out, prefix_embs, prefix_pad_masks, prefix_att_masks, prefix_offsets, tokens, past_key_values).

Return type:

A tuple of updated state tensors for the next step

sample_actions(images: list[Tensor], img_masks: list[Tensor], lang_tokens: Tensor, lang_masks: Tensor, metadata_tokens: Tensor | None = None, metadata_masks: Tensor | None = None) → tuple[Tensor, Tensor][source]

Inference forward: autoregressively generates memory and subtask tokens.

Runs memory_max_length infer_autoregressive steps, then feeds the same "Subtask: " token IDs used in training (tokenizer-dependent length subtask_indicator_max_length) through the cache, then runs response_max_length response steps. Each step conditions on prior KV-cache entries.

Parameters:

images – List of image tensors, one per camera.
img_masks – List of boolean masks for real vs. padded images.
lang_tokens – Language token IDs (B, prompt_max_length).
lang_masks – Boolean attention mask for language tokens.
metadata_tokens – Optional metadata token IDs (B, metadata_max_length).
metadata_masks – Optional boolean mask for metadata tokens.

Returns:

A tuple (memory_tokens, response_tokens) where each is a Tensor of generated token IDs with shape (B, memory_max_length) and (B, response_max_length) respectively.

class opentau.policies.pi07.high_level_planner.modeling_pi07_high_level.PI07HighLevelPlannerPolicy(config: PI07HighLevelPlannerConfig, per_dataset_stats: list[dict[str, dict[str, Tensor]]] | None = None, dataset_names: list[str] | None = None)[source]

Bases: PreTrainedPolicy

Policy wrapper for the π07 high-level planner.

Handles input normalisation, tokenisation of language/memory/response, and delegates to PI07HighLevelPlannerModel for autoregressive prediction of updated memory and subtask strings.

__init__(config: PI07HighLevelPlannerConfig, per_dataset_stats: list[dict[str, dict[str, Tensor]]] | None = None, dataset_names: list[str] | None = None)[source]

Initializes the PI07HighLevelPlannerPolicy.

Parameters:

config – Policy configuration instance.
per_dataset_stats – Ordered list of per-dataset stat dicts used to fill the stacked Normalize input-buffer. May be None when constructing for a checkpoint load.
dataset_names – Ordered list parallel to per_dataset_stats.

config_class: alias of PI07HighLevelPlannerConfig

forward(batch: dict[str, Tensor]) → dict[str, Tensor][source]

Runs a full training forward pass and computes the loss.

Tokenizes images, language (with state and past memory), target memory, and target response, then computes cross-entropy losses for both the memory and response token predictions.

Parameters:: batch – Batch of training data. Expected keys include images, "prompt", "state", "past_memory", "response", and "next_memory".
Returns:: A dict with "MSE" (always zero, kept for interface compatibility) and "CE" (sum of memory and response cross-entropy losses).

Override the from_pretrained method to handle key remapping.

Parameters:

pretrained_name_or_path – Path to the pretrained model or its name on the Hub.
config – Configuration object.
force_download – Whether to force download the model weights.
resume_download – Whether to resume download.
proxies – Proxy configuration.
token – Authentication token.
cache_dir – Directory to cache downloaded files.
local_files_only – Whether to only look for files locally.
revision – Specific model revision.
strict – Whether to strictly enforce state dict matching.
**kwargs – Additional keyword arguments.

Returns:

The loaded model instance.

Raises:

ValueError – If pretrained_name_or_path is None.

get_optim_params() → dict[source]

Returns the parameters to be optimized.

Returns:: A generator over the model parameters.

name: None = 'pi07_high_level': The name of the policy. Must be defined in subclasses.

predict_action_chunk(batch: dict[str, Tensor]) → Tensor[source]

Not implemented for the high-level planner.

Parameters:: batch – Batch of data containing environment observations.
Raises:: NotImplementedError – Always, since the high-level planner predicts memory and subtask strings, not action chunks.

prepare_discrete_state(batch: dict[str, Tensor]) → list[str][source]

Discretizes the state into bins and converts it to a string representation.

Each dimension of the state vector is discretized into 256 bins. The values of each dimension of the state are expected to be in the range [-1, 1]. The discretization bins are linearly spaced between -1 and 1. The index of the bin for each dimension is then concatenated into a space-separated string.

Parameters:: batch – Batch of data containing the “state” tensor.
Returns:: A list of strings, where each string is a space-separated list of discretized state values.
Raises:: ValueError – If the state values are not normalized between -1 and 1.

prepare_images(batch: dict[str, Tensor]) → tuple[list[Tensor], list[Tensor]][source]

Apply preprocessing to the images.

Resizes to 224x224 and padding to keep aspect ratio, and converts pixel range from [0.0, 1.0] to [-1.0, 1.0] as requested by SigLIP.

Parameters:

batch – Batch of data containing image tensors.

Returns:

images: A list of processed image tensors.
img_masks: A list of image mask tensors.

Return type:

A tuple containing

Raises:

ValueError – If no image features are present in the batch.

prepare_language(batch: dict[str, Tensor]) → tuple[Tensor, Tensor][source]

Tokenizes the composite language prompt.

Builds a prompt string from the task instruction, past memory, and discretized robot state separated by <eos> tokens, then tokenizes and pads to prompt_max_length.

Parameters:

batch – Batch containing "prompt" (task strings), "state" (state tensor), and "past_memory" (list of past memory strings).

Returns:

lang_tokens: Token IDs of shape (batch_size, prompt_max_length).
lang_masks: Boolean attention mask of the same shape.

Return type:

A tuple (lang_tokens, lang_masks) where

prepare_metadata(batch: dict[str, Tensor]) → tuple[Tensor, Tensor][source]

Tokenize episode metadata into Gemma 3 token IDs.

Wraps non-empty per-sample metadata segments into a single "Metadata: {seg1} {seg2} ..." string, then pads/truncates to metadata_max_length. Samples with no active segments emit an empty string.

Parameters:: batch – Batch dict that may contain any of: "speed", "quality", "mistake", "fps" (numeric tensors with a corresponding _is_pad bool tensor — entries marked as pad are dropped), and "robot_type", "control_mode" (lists of strings — empty string is the pad signal, no separate _is_pad flag). Missing keys are treated as fully padded.
Returns:: A tuple (metadata_tokens, metadata_masks) with shapes (B, metadata_max_length).

prepare_next_memory(batch: dict[str, Tensor]) → tuple[Tensor, Tensor][source]

Tokenizes the target updated memory for training.

Wraps each memory string with an <eos> suffix, then tokenizes and pads to memory_max_length.

Parameters:

batch – Batch containing "next_memory" (list of target memory strings) and "state" (used only to determine the device).

Returns:

memory_tokens: Token IDs of shape (batch_size, memory_max_length).
memory_masks: Boolean attention mask of the same shape (True for real tokens, False for padding).

Return type:

A tuple (memory_tokens, memory_masks) where

prepare_response(batch: dict[str, Tensor]) → tuple[Tensor, Tensor][source]

Tokenizes the target subtask response for training.

Wraps each response string with an <eos>Actions: suffix, then tokenizes and pads to response_max_length.

Parameters:

batch – Batch containing "response" (list of subtask strings) and "state" (used only to determine the device).

Returns:

response_tokens: Token IDs of shape (batch_size, response_max_length).
response_masks: Boolean attention mask of the same shape (True for real tokens, False for padding).

Return type:

A tuple (response_tokens, response_masks) where

reset() → None[source]: Resets any internal state. Call when the environment resets.

sample_actions(batch: dict[str, Tensor]) → tuple[Tensor, Tensor][source]

Run inference to predict updated memory and subtask tokens.

Normalizes inputs, prepares image and language embeddings, then delegates to the inner model for autoregressive generation.

Parameters:: batch – Batch of observations. Expected keys include images, "prompt", "state", and "past_memory".
Returns:: A tuple (memory_tokens, response_tokens) where each is a Tensor of token IDs with shape (batch_size, seq_len).

select_action(batch: dict[str, Tensor], noise: Tensor | None = None) → Tensor[source]

Not implemented for the high-level planner.

Parameters:: batch – Batch of data containing environment observations.
Raises:: NotImplementedError – Always, since the high-level planner predicts memory and subtask strings, not action chunks.

opentau.policies.pi07.high_level_planner.modeling_pi07_high_level.make_att_2d_masks(pad_masks: Tensor, att_masks: Tensor, n_cross_att_tokens: int | None = None, cross_att_pad_masks: Tensor | None = None) → Tensor[source]

Creates a 2-D attention mask given padding and 1-D attention masks.

Tokens can attend to valid inputs tokens which have a cumulative att_masks smaller or equal to theirs. This way att_masks int[B, N] can be used to setup several types of attention, for example:

[[1 1 1 1 1 1]]: pure causal attention.

[[0 0 0 1 1 1]]: prefix-lm attention. The first 3 tokens can attend between
themselves and the last 3 tokens have a causal attention. The first entry could also be a 1 without changing behaviour.

[[1 0 1 0 1 0 0 1 0 0]]: causal attention between 4 blocks. Tokens of a
block can attend all previous blocks and all tokens on the same block.

Parameters:

pad_masks – bool[B, N] true if its part of the input, false if padding.
att_masks – int32[B, N] mask that’s 1 where previous tokens cannot depend on it and 0 where it shares the same attention mask as the previous token.
n_cross_att_tokens – Add attention mask for cross-attention tokens if n_cross_att_tokens is provided.
cross_att_pad_masks – Padding masks for cross attention tokens. Required if n_cross_att_tokens is provided.

Returns:

A 2D attention mask tensor of shape (B, N + n_cross_att_tokens, N + n_cross_att_tokens) if n_cross_att_tokens is provided, else (B, N, N).

Raises:

ValueError – If att_masks or pad_masks are not 2D (including batch dimension).
AssertionError – If cross_att_pad_masks is missing when n_cross_att_tokens is set, or if its shape is incorrect.

opentau.policies.pi07.high_level_planner.modeling_pi07_high_level.resize_with_pad(img: Tensor, width: int, height: int, pad_value: int = -1) → Tensor[source]

Resizes an image to fit within the specified dimensions while maintaining aspect ratio, and pads the remaining area with the specified value.

Parameters:

img – Input image tensor of shape (batch_size, channels, current_height, current_width).
width – Target width.
height – Target height.
pad_value – Value to use for padding. Defaults to -1.

Returns:

The resized and padded image tensor of shape (batch_size, channels, height, width).

Raises:

ValueError – If the input image tensor does not have 4 dimensions.