quantgate-student-robocasa-phase5-dinov3s

quantgate-student-robocasa-phase5-dinov3s

A quantizability gate built on a frozen DINOv3 ViT-S/16 encoder: it looks at a robot's camera views and its task instruction and predicts, per action chunk, whether that chunk can be run through K=2 action-chunk compression without breaking the task.

• Benchmark: RoboCasa Kitchen (24 tasks, MimicGen demos)
• Embodiment: RoboCasa mobile-manipulator, 3 views (left / right / wrist) at 20 fps, 32-d action space, action horizon 16
• Policy it gates: GR00T-N1.5 finetuned on RoboCasa — prehj/GR00T-N1.5-robocasa-baseline
• What the weights are: 86 MB. facebook/dinov3-vits16-pretrain-lvd1689m frozen (its weights are included in the checkpoint so the file is self-contained), plus the only trained parts: a per-view embedding, a single learned attention-pooling query with a 6-head MultiheadAttention, an instruction projection, and a 3-layer MLP head. Inference cost measured at ~7 ms, versus ~0.6 ms for the CNN students.
• Why it exists: the CNN student mixes the three views together in its first convolution and then throws away all spatial structure with a global average pool. This variant encodes each view separately, keeps the patch tokens, and pools them with one learned query, so view identity and spatial layout survive to the head.

Checkpoint metadata: best epoch 9, held-out AUC against the teacher 0.639.

Per-variant recipe: --encoder dinov3s --epochs 10 --bs 128 --lr 3e-4, labels v6b_phase5_1call_full.parquet (the binary-flag phase5 labels — the same labels as prehj/quantgate-student-robocasa-phase5-binary, so those two are the one architecture comparison in this collection that is like-for-like).

A caveat worth recording. A cheap linear probe on pre-pooled 4x4 features predicted this architecture would win comfortably (0.829 vs 0.734 for the CNN). At full training the ordering reversed on AUC (0.639 vs 0.674). Closed loop, it is the strongest single row in the table below. The probe was not predictive; the closed loop is what settled it.

Teacher labels this student was distilled from

One VLM labelling pass, shared by every phase5 student in this collection:

• Teacher: nvidia/Cosmos3-Nano, run as a local reasoner server through HuggingFace transformers (not vLLM), scored by reading the next-token distribution over the answer slots rather than by free-form generation.
• Labelled set: 247,887 action chunks from the RoboCasa MimicGen dataset kimtaey/robocasa_mg_gr00t_300 (LeRobot format), stride-8 sampling, 16 parallel judge shards.
• Prompting generation "phase5": phase-based guidance (the guidance describes each risk axis in its own paragraph) plus four grasp/hold-axis questions q_A..q_D, asked in a single VLM call per chunk.
• Aggregation: the VLM answers are combined with four computed action descriptors (grip_transition, reversal, precise_hold, infeasible_merge) under a noisy-OR, then rank-normalised to p_yes in [0, 1]. Action-derived quantities are computed and stated as facts, never asked — asking them measured AUC 0.520, i.e. chance.

The three phase5 students differ only in how those four computed flags are turned into numbers before aggregation. The VLM was called exactly once; no variant required re-labelling.

Why the continuous variants exist: under binary flags, a single flag firing saturates the noisy-OR, so 29.51% of all chunks collapsed to an identical score of 0 and the VLM's judgement was discarded outright across that whole region. Continuous flags fix that from the action numbers alone (~101 s of CPU), with no additional VLM calls.

Training recipe

Recovered from the submission script and the checkpoint metadata.

Measured closed-loop behaviour

RoboCasa Kitchen, 24 tasks x 50 episodes, GR00T-N1.5 policy (prehj/GR00T-N1.5-robocasa-baseline), K=2 action-chunk compression applied only where the gate says the chunk is quantizable.

Success rate alone cannot rank these models. Compression buys steps by spending success: a gate that fires almost never approaches the uncompressed row (high success, no speedup) and a gate that fires almost always approaches the blanket row (low success, large speedup). Both extremes are trivially reachable and neither is a good gate. The meaningful quantity is where a gate sits relative to the straight line joining the blanket-K=2 point to the uncompressed point: that line is what you get for free by simply choosing a random fraction of chunks to compress. Only distance above that line is evidence that the gate is selecting the right chunks rather than merely selecting fewer of them.

On the project's own closed-loop record (23-task completed subset, the run in which excess-over-line was tabulated), the excesses were: binary +0.0078, softA/ratio +0.0076, softB/event-preserving +0.0111 — i.e. the event-preserving aggregation is the best of the three despite having the lowest raw success rate of the three, which is exactly the point above.

How to load and serve

The checkpoint is a plain torch.save dict with model (state_dict, includes the frozen encoder weights so the repo is self-contained), res, views, encoder: "dinov3s", temb_dim, task_emb_file, epoch, val_auc.

gate_module_best.pt is the best-validation-AUC epoch and is the file to serve; gate_module.pt is the final epoch.

Requires transformers >= 4.56 for facebook/dinov3-vits16-pretrain-lvd1689m.

import numpy as np, torch, torch.nn as nn
from transformers import AutoModel
from huggingface_hub import hf_hub_download

REPO = "prehj/quantgate-student-robocasa-phase5-dinov3s"

class DinoGate(nn.Module):
    DINO = "facebook/dinov3-vits16-pretrain-lvd1689m"
    def __init__(self, temb_dim=384, res=224, nheads=6):
        super().__init__()
        self.enc = AutoModel.from_pretrained(self.DINO)
        for p in self.enc.parameters(): p.requires_grad = False
        self.enc.eval()
        d = self.enc.config.hidden_size; self.res, self.d = res, d
        self.view_emb = nn.Parameter(torch.randn(3, 1, d) * 0.02)   # per-view marker
        self.q   = nn.Parameter(torch.randn(1, 1, d) * 0.02)        # pooling query
        self.att = nn.MultiheadAttention(d, nheads, batch_first=True)
        self.temb_proj = nn.Linear(temb_dim, d) if temb_dim else None
        self.head = nn.Sequential(nn.Linear(d * (2 if temb_dim else 1), 256), nn.ReLU(),
                                  nn.Linear(256, 64), nn.ReLU(), nn.Linear(64, 1))
        self.register_buffer("mean", torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
        self.register_buffer("std",  torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
    def train(self, mode=True):
        super().train(mode); self.enc.eval(); return self
    def forward(self, x, t=None):
        B = x.shape[0]
        v = x.reshape(B * 3, 3, x.shape[-2], x.shape[-1])
        v = torch.nn.functional.interpolate(v, size=(self.res, self.res),
                                            mode="bilinear", align_corners=False)
        v = (v - self.mean) / self.std
        with torch.no_grad():
            h = self.enc(pixel_values=v).last_hidden_state
        P = (self.res // self.enc.config.patch_size) ** 2
        tok = h[:, -P:, :].reshape(B, 3, P, self.d)                  # patch tokens only
        tok = (tok + self.view_emb.unsqueeze(0)).reshape(B, 3 * P, self.d).float()
        f = self.att(self.q.expand(B, -1, -1).float(), tok, tok)[0][:, 0]
        if self.temb_proj is not None and t is not None and t.numel():
            f = torch.cat([f, self.temb_proj(t)], dim=1)
        return self.head(f)

ck = torch.load(hf_hub_download(REPO, "gate_module_best.pt"), map_location="cpu")
gate = DinoGate(temb_dim=ck["temb_dim"]); gate.load_state_dict(ck["model"]); gate.eval()

z = np.load(hf_hub_download(REPO, "robocasa_task_embeddings.npz"), allow_pickle=True)
TEMB = dict(zip(z["tasks"], z["emb"]))

Input is the same as the CNN students: the three RGB views at the chunk's first frame, each resized to 128x128, concatenated to 9 channels (the module upsamples internally to 224 for DINOv3), plus the 384-d instruction embedding. Output is a logit; sigmoid(logit) >= tau means the chunk may be compressed. tau is tuned per (teacher x architecture) cell on a held-out split — it is not portable from the CNN students.

The gate never sees actions, by design: it is meant to run concurrently with the action head's denoising so its latency is hidden, and taking the planned action chunk as input would serialise the two.