quantgate-student-robocasa-phase5-softB

quantgate-student-robocasa-phase5-softB

The recommended student of the three CNN variants. 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 ~0.6 ms.
• This variant ("event-preserving" aggregation): identical architecture and identical VLM labels as the binary and softA repos. The difference is that only the cumulative risk flags are softened into fractions; the two event-type flags — gripper-state transition and direction reversal — are kept near-saturating. The reasoning: a cumulative cost genuinely scales with how many merge pairs are affected, but an event either happens inside the chunk or it does not, and a single mis-merged gripper transition is enough to drop the object. Softening those two the way softA does dilutes a real binary hazard into a small number.

Checkpoint metadata: best epoch 6, held-out AUC against the teacher 0.906.

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

Why this one is preferred despite the lowest raw success rate of the three: it holds success while cutting the most steps. It sits furthest above the blanket-compression / no-compression trade-off line (+0.0111 vs +0.0078 for binary and +0.0076 for softA on the project's tabulated closed-loop subset). Read the note under the results table before comparing these models by success rate.

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-softB"

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.