Ashenzari/vla-mamba-success-monitor

VLA MAMBA Progress Monitor — v2.1.0

This revision contains a five-fold MAMBA ensemble trained to estimate causal robot-task progress from frozen VLA/VLM backbone features. It is the progress-only successor to the v1 success/failure monitor.

Version contract

• Release: v2.1.0
• Objective: progress-only
• Training recipe: original progress loss + freeze/reverse temporal augmentation
• Supported in-domain tasks: coffeecup, saucer, waterbottle, wine
• Input: 10 Hz mean-pooled 4096-dimensional VLM features
• Ensemble: mean progress from five collection-disjoint fold checkpoints
• VLM feature source: new_vla_checkpoint_vlm_30k

Cached features from v2.0.0 are not compatible with these temporal heads. Re-extract features with the 30k-step VLA backbone before inference.

The model architecture still contains an outcome head for checkpoint compatibility, but that head was not trained for this release. Do not interpret its logits as success/failure probabilities. Only head_progress output is valid.

Loading a fold

Install the source repository and the Mamba dependency when loading the Mamba release. The rank normalizer stored in each checkpoint must be applied before inference.

from huggingface_hub import snapshot_download
import json
import torch

from monitor.feature_normalization import TaskRankNormalizer
from monitor.model import Monitor

root = snapshot_download("Ashenzari/vla-mamba-success-monitor", revision="v2.1.0")
config = json.load(open(f"{root}/config.json"))
entry = config["checkpoints"][0]
checkpoint = torch.load(
    f"{root}/{entry['file']}", map_location="cpu", weights_only=True
)

model = Monitor(temporal="mamba", dropout=checkpoint["args"]["dropout"])
model.load_state_dict(checkpoint["model"], strict=True)
model.eval()
normalizer = TaskRankNormalizer.from_state_dict(checkpoint["rank_normalizer"])

# pooled_features: [T, 4096]
x = normalizer.transform(pooled_features.float(), task_id="wine")
with torch.inference_mode():
    _, progress = model(x.unsqueeze(0))
progress = progress[0]  # [T], range [0, 1]

For ensemble inference, repeat the same operation for all five checkpoint entries and average their progress curves.

Feature contract

Features were extracted from the new_vla_checkpoint_vlm_30k VLA checkpoint. Each prediction step uses the context offsets [-0.3, -0.2, -0.1, 0.0] seconds, 64 visual tokens, mean pooling, and a 4096-dimensional feature vector. Changing the VLA/VLM checkpoint can shift the feature distribution and requires revalidation.

Evaluation

See EVALUATION.md and config.json for the clean and freeze/reverse robustness metrics. The evaluation is in-domain and does not establish zero-shot performance on unseen tasks.