Patch Policy: Efficient Embodied Control via Dense Visual Representations

Patch Policy: Efficient Embodied Control via Dense Visual Representations§

Core Idea§

Patch Policy modifies the attention mechanism of a transformer-based policy to directly consume dense patch tokens from a pretrained Vision Transformer (ViT) instead of global pooling. A block-causal attention mask maintains temporal causality while allowing full attention across patch tokens within each observation and across time steps, avoiding the computational burden of large vision-language models (VLMs).

Methodology§

Architecture: The policy consists of a pretrained ViT (e.g., DINOv2) that extracts patch tokens from each observation. These tokens are concatenated with a state token (robot proprioception) and an action token (to be predicted). All tokens are fed into a small transformer decoder with a custom attention mask.

Block-Causal Mask: For a sequence of $T$ time steps, each with $N = P + S + A$ tokens ($P$ patches, $S$ state tokens, $A$ action tokens; typically $S=A=1$), the attention mask $M \in \mathbb{R}^{TN \times TN}$ is defined as:

• $M_{i,j} = 0$ if token $i$ can attend to token $j$, else $M_{i,j} = -\infty$.
• All tokens can attend to all tokens from previous time steps (full temporal causality).
• Within the same time step, all tokens can attend to each other (block attention). This allows the action token to leverage all patch tokens from the current observation.

The attention computation is then: $$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{Q K^\top}{\sqrt{d}} + M\right) V $$

Training: The policy is trained via behavior cloning or reinforcement learning. The ViT backbone is frozen or fine-tuned with low-rank adaptation (LoRA) to preserve pretrained features while adapting to the domain.

PyTorch Code Snippet§

import torch
import torch.nn.functional as F

def create_block_causal_mask(T, P, S=1, A=1):
    N_per_step = P + S + A
    total_tokens = T * N_per_step
    mask = torch.full((total_tokens, total_tokens), float('-inf'))
    for t in range(T):
        start = t * N_per_step
        end = (t + 1) * N_per_step
        # attend to all previous time steps
        mask[start:end, :start] = 0
        # attend to all tokens in current time step (block)
        mask[start:end, start:end] = 0
    return mask

# Example usage in a custom transformer layer
def forward(self, x):
    T, P, S, A = ...
    mask = create_block_causal_mask(T, P, S, A).to(x.device)
    attn_output = self.self_attn(x, x, x, attn_mask=mask)
    return attn_output

Key Results§

• Simulation: Across four simulated environments, Patch Policy achieves a 40% relative improvement over policies using global-pooled representations (e.g., CLS token).
• Real-World: On three real-world robot manipulation tasks, it surpasses fine-tuned OpenVLA-OFT by 18% while using only 0.7% of the parameters.
• Efficiency: The policy runs at high frequency (e.g., 10-50 Hz) suitable for reactive control, unlike heavy VLA models.

Significance§

Patch Policy provides a simple, effective way to leverage dense pretrained visual features for robot learning without sacrificing speed or requiring large VLMs. This opens the door to integrating future advances in visual representation learning into embodied agents.