Quantizability Gate — architecture C (gate token inside GR00T-N1.5)
Full GR00T-N1.5 policy finetuned on RoboCasa Kitchen with a gate token added to
the flow-matching DiT. Unlike the distilled A′ student (a separate CNN), C
predicts quantizability from inside the policy, sharing the backbone that
produces the action chunk — one forward pass gives both the 16-step chunk and
the confidence that the chunk is safe to temporally compress (K2).
Variant here: v3_lam03 — one-way attention mask (the gate token reads the
action stream, the action stream does not read the gate token, so gating cannot
perturb the actions) with gate-loss weight λ=0.3, trained 60k steps on
Cosmos3-Nano labels.
Closed-loop (RoboCasa, 24 tasks × 50 episodes each)
| policy | success | avg steps |
|---|
| uncompressed baseline | 0.657 | 327 |
| naive K2 (compress everything) | 0.598 | 221 |
| C (this model), τ=0.5 | 0.638–0.647 | 289 |
| A′ distilled student (gemma4 teacher), τ=0.5 | 0.667 | 258 |
Benchmark noise is ±1.5 pp. A′ currently edges out C on this benchmark; C is
the architecturally cleaner option (no second model, no separate image encode)
and is the one to build on if you want the gate to co-adapt with the policy.
Usage
Load with the fork that defines the gate token — stock GR00T-N1.5 will not have
the extra parameters:
1git clone -b action-quantization-gate-v2 \
2 https://github.com/rakybond007/GR00T-action-quantization
1from gr00t.model.policy import Gr00tPolicy
2
3policy = Gr00tPolicy(model_path="prehj/groot-n15-quantizability-gate-C-robocasa",
4 embodiment_tag="new_embodiment", denoising_steps=4)
5out = policy.get_action(obs)
6chunk = out["action_pred"] # (16, 12)
7conf = float(out["_gate_prob"]) # P(safe to compress), float in [0,1]
Then merge when the gate allows it (deltas sum, gripper takes the last value):
1from merge_k2 import merge_k2 # see the A' repo
2if conf >= tau:
3 chunk = merge_k2(chunk, delta_dims=range(5, 11), state_dims=[11])
τ is a per-model operating point, not a constant — judges and architectures
calibrate differently, so τ=0.5 means different things across checkpoints. It is
meant to be searched, or learned by an RL head on top of _gate_prob.
Action space (RoboCasa, 12 dims)
0–4 unused · 5–7 EE delta xyz · 8–10 rotation · 11 gripper (0/1).
Summed deltas are clipped to the controller limit; raising that limit alone does
not recover success — the gate does.
Running the benchmark
Environment setup, the three-process runtime layout, the full task list, how to
sweep tau, and the pitfalls that have actually cost us runs are written up here:
RoboCasa evaluation setup guideRelated
- A′ distilled students + confidence API:
prehj/groot-n15-quantizability-gate-A-robocasa
- Code, prompt-evolution loop, label elicitation:
https://github.com/rakybond007/GR00T-action-quantization/tree/action-quantization-gate-v2
LIBERO is not covered: it was only used for policy-side K2/varK experiments, so
no gate was trained there.
