quantgate-student-robocasa-phase5-binary

quantgate-student-robocasa-phase5-binary

A quantizability gate: a small CNN that 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: a 5-block strided CNN over 9 input channels plus a 384-d instruction embedding, ~320k parameters, 1.3 MB. Inference cost measured at ~0.6 ms.
• This variant: teacher labels aggregated with the original binary (0/1) computed risk flags. This is the reference variant; the two soft* repos are the same student trained on continuous re-aggregations of the same labels.

Checkpoint metadata: best epoch 5, held-out AUC against the teacher 0.674. (Training ran 30 epochs; AUC peaked at epoch 5 and had decayed to 0.554 by epoch 30, which is why the later variants were trained for 10 epochs.)

Per-variant recipe: --epochs 30 --bs 256 --lr 3e-4, labels v6b_phase5_1call_full.parquet.

AUC across students trained on different label sets is not comparable and should be read only as a distillation-collapse detector. Only the closed-loop table below ranks these models.

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:

{
  "model":         state_dict of the gate network,
  "res":           128,                      # input resolution per view
  "views":         ["observation.images.left_view",
                    "observation.images.right_view",
                    "observation.images.wrist_view"],
  "temb_dim":      384,                      # instruction-embedding width
  "act_cols":      [],                       # no action inputs, by design
  "task_emb_file": path to robocasa_task_embeddings.npz,
  "epoch":         epoch this checkpoint came from,
  "val_auc":       held-out AUC against the teacher at that epoch
}

gate_module_best.pt is the best-validation-AUC epoch and is the file to serve. gate_module.pt is the final epoch, kept for reproducibility only.

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

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

class SmallGate(nn.Module):
    """9ch (3 RGB views concatenated) [+ instruction embedding] -> logit P(quantize)."""
    def __init__(self, ch=32, temb_dim=384):
        super().__init__()
        blk = lambda i, o: nn.Sequential(nn.Conv2d(i, o, 3, 2, 1), nn.BatchNorm2d(o), nn.ReLU())
        self.net  = nn.Sequential(blk(9, ch), blk(ch, ch*2), blk(ch*2, ch*4),
                                  blk(ch*4, ch*4), nn.AdaptiveAvgPool2d(1))
        self.head = nn.Sequential(nn.Linear(ch*4 + temb_dim, 128), nn.ReLU(),
                                  nn.Linear(128, 64), nn.ReLU(), nn.Linear(64, 1))
    def forward(self, x, t=None):
        f = self.net(x).flatten(1)
        if t is not None and t.numel():
            f = torch.cat([f, t], dim=1)
        return self.head(f)

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

# instruction embeddings shipped in this repo (all-MiniLM-L6-v2, 384-d, keyed by task string)
z = np.load(hf_hub_download(REPO, "robocasa_task_embeddings.npz"), allow_pickle=True)
TEMB = dict(zip(z["tasks"], z["emb"]))

def prep(img, res=128):                      # img: HxWx3 uint8 RGB
    import cv2
    img = cv2.resize(img, (res, res), interpolation=cv2.INTER_AREA)
    return (img.astype(np.float32) / 255.0).transpose(2, 0, 1)

@torch.no_grad()
def quantizable(left, right, wrist, task, tau=0.5):
    x = torch.from_numpy(np.concatenate([prep(left), prep(right), prep(wrist)], 0))[None]
    t = torch.from_numpy(TEMB[task].astype(np.float32))[None]
    p = torch.sigmoid(gate(x, t)).item()
    return p >= tau, p                       # True -> this chunk may be compressed

Serving contract. Call the gate once per action chunk, on the three RGB views at the chunk's first frame plus the task instruction. If it returns True, run the policy's action chunk through K=2 compression; if False, run the chunk uncompressed. tau is not a universal constant — it is tuned per (teacher x architecture) cell on a held-out split; 0.5 is only a starting point.

The gate never sees actions. This is a deliberate constraint, not an oversight: the gate is meant to run concurrently with the action head's denoising so its latency is hidden. Feeding it the planned action chunk would force it to wait for denoising to finish, converting a hidden cost into pure added latency. The teacher may look at the actions; the student learns E[label | image, instruction]. A lower validation AUC than an action-fed variant is the expected price of that, not a defect.