Workspace Models: Lightweight Robotic Memory via Saliency-Driven Supervision

Workspace Models: Lightweight Robotic Memory via Saliency-Driven Supervision§

The paper introduces workspace tokens, a train-time distilled latent memory for robotic manipulation policies. The key insight is that expensive Vision-Language Model (VLM) queries, which are typically used at deployment to extract task-salient history, can be moved to training time. This allows the deployment policy to use a lightweight, fixed-size memory representation.

Methodology§

1. Train-time saliency extraction: For each timestep, a VLM processes the full history of observations and actions to identify which past elements are necessary for the current task.
2. Distillation via set-reconstruction: The workspace token $w_t$ is trained to reconstruct the set of salient history elements using a set-reconstruction decoder loss. Formally, given a set of salient elements $S_t = \{s_1, \dots, s_k\}$, the decoder $D$ minimizes:

$$\mathcal{L}_{\text{rec}} = \sum_{s \in S_t} \ell(D(w_t), s) + \lambda \mathcal{L}_{\text{reg}}$$ where $\ell$ is a reconstruction loss (e.g., MSE for continuous states, cross-entropy for discrete actions) and $\mathcal{L}_{\text{reg}}$ optionally regularizes the token.

1. Policy conditioning: The workspace token $w_t$ is concatenated with the current observation and fed to the policy network, replacing the full history. This avoids spurious correlations from irrelevant history.

Implementation Details§

• The VLM (e.g., CLIP or a fine-tuned multimodal model) is used offline to label salient elements. The saliency can be based on relevance scores or human annotations.
• The workspace token is a learned embedding, updated recurrently or via a transformer that processes the history of observations and actions.
• The set-reconstruction decoder is a small MLP or transformer that maps $w_t$ to the salient set.
• At deployment, only the trained workspace token model and policy are used; no VLM inference is performed.
• Experiments in simulation (e.g., memory-intensive manipulation tasks) and on real hardware demonstrate that policies with workspace tokens outperform those conditioned on full histories and are more efficient than in-the-loop VLM approaches.

Code Snippet (PyTorch)§

import torch
import torch.nn as nn

class WorkspaceToken(nn.Module):
    def __init__(self, obs_dim, act_dim, token_dim, hist_len):
        super().__init__()
        self.encoder = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(d_model=obs_dim+act_dim, nhead=4),
            num_layers=2
        )
        self.token_proj = nn.Linear(obs_dim+act_dim, token_dim)
        self.decoder = nn.Sequential(
            nn.Linear(token_dim, 128),
            nn.ReLU(),
            nn.Linear(128, obs_dim+act_dim)  # reconstruct salient elements
        )
    
    def forward(self, obs_hist, act_hist):
        # obs_hist: (B, T, obs_dim), act_hist: (B, T, act_dim)
        x = torch.cat([obs_hist, act_hist], dim=-1)  # (B, T, obs_dim+act_dim)
        x = x.permute(1, 0, 2)  # (T, B, D)
        encoded = self.encoder(x)  # (T, B, D)
        # Use last token as workspace token
        w = self.token_proj(encoded[-1])  # (B, token_dim)
        return w
    
    def reconstruct(self, w):
        # Decode to salient set (simplified: reconstruct all elements, supervised on salient subset)
        return self.decoder(w)  # (B, obs_dim+act_dim)

Results§

• Workspace tokens enable policies to solve tasks requiring memory of up to 100 steps.
• Compared to full-history conditioning, they reduce spurious correlations and improve success rates by 15-30% in simulation.
• On hardware, they achieve comparable or better performance than in-the-loop VLM policies while being 10x faster at inference.

Conclusion§

Workspace tokens offer a practical solution for memory-intensive robotic tasks by leveraging train-time VLM supervision to create efficient, deployable memory representations. This approach bridges the gap between rich VLM reasoning and efficient policy execution.