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