---
license: apache-2.0
library_name: pytorch
tags:
- robotics
- vla
- gr00t
- quantization
- knowledge-distillation
- robocasa
pipeline_tag: robotics
---
# 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`](https://huggingface.co/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.

| variant | computed flags | repo |
|---|---|---|
| binary | 0/1 (original) | `prehj/quantgate-student-robocasa-phase5-binary` |
| softA (ratio) | all four continuous, = fraction of K=2 merge pairs actually harmed | `prehj/quantgate-student-robocasa-phase5-softA` |
| softB (event-preserving) | cumulative flags continuous; event-type flags (grip transition, direction reversal) held near 1 | `prehj/quantgate-student-robocasa-phase5-softB` |

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.

| | |
|---|---|
| script | `vlm_gate/scripts/train_gate_module.py` |
| dataset | `kimtaey/robocasa_mg_gr00t_300` (LeRobot), uint8 memmap frame cache |
| input | 3 views (left, right, wrist) at 128x128, concatenated to 9 channels |
| conditioning | 384-d MiniLM embedding of the task instruction |
| loss | BCE on the soft teacher score `p_yes` |
| split | episode-wise, last 25% of episodes held out |
| optimiser | Adam, lr 3e-4 |
| hardware | 1 GPU |

## 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.

| configuration | success rate | steps (successful episodes) |
|---|---|---|
| uncompressed (no gate, no compression) | 0.657 | 327.0 |
| blanket K=2 (compress everything) | 0.598 | 214.0 |
| phase5 **binary** flags student | 0.638 | 274.5 |
| phase5 **softA** (ratio) student | 0.635 | 266.6 |
| phase5 **softB** (event-preserving) student | 0.627 | 252.0 |
| phase5 **DINOv3 ViT-S/16** student | 0.642 | 276.6 |

**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.

```python
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.