twanghcmut/efficient-temporal / docs /ARCHITECTURE_robomme16.md
twanghcmut's picture
|
download
raw
38.8 kB

System architecture: RoboMME 16-task, planner plus memory-conditioned controller

This document describes the system that produced the RoboMME 16-task eval-50 result: 505/800 = 63.1% with the planner in the loop, and 232/320 = 72.5% with oracle subgoals (20 episodes per task). The 63.1% run still passed the simulator's subgoal words to the controller (Section 12, words leak), so it is not a clean non-oracle number. A clean rerun is in progress. Numbers come from the run directories and configs under triquang/graphvla. A value that could not be confirmed from code, config, or logs is marked TODO(verify).

1. Overview

The system has two learned parts that run in a closed loop on top of one frozen vision-language backbone.

  1. A planner reads the demo video and the execution history and emits a subgoal at every node: an event kind, an ordinal, and a 2D point in the front image.
  2. A controller (the GR00T N1.6 flow-matching DiT action head, extended with subgoal conditioning) turns the current observation and that subgoal into a 40-step joint-position chunk.

The planner's subgoal replaces the language-only memory of the original VLA. The controller never sees the demo video directly. Everything it knows about the past reaches it through the subgoal.

                    demo video + execution frames (front 256x256, wrist 256x256)
                                      |
                       +--------------v---------------+
                       | Eagle-2B VLM, frozen         |  LM truncated to layer 16
                       | + 4 TCL moment tokens (end)  |  cached last hidden state
                       +------+---------------+-------+
                              | 9x9 front patches (+ mean wrist), text hidden states
              +---------------v----+          |
              | PLANNER (7.66M)    |          |   2 views x 81 image tokens + text + 4 moment
              |  node tokens 82x256|          |   = 209..228 tokens x 2048
              |  4 GraphBlocks     |          |
              |  EventCounter      |          |
 CoTracker -->|  EntityStream (192 |          |
 (frozen)     |  tracked points)   |          |
              +----+---------------+          |
                   | kind, ordinal, point     |
                   v                          v
        +----------------------+    +-------------------------------+
        | subgoal -> OracleAdapter |  | GR00T DiT action head, 32 blocks |
        |  point Fourier token  |-->|  cross-attn: image / text+moment |
        |  step token (kind,    |   |  self-attn over [state, actions] |
        |   ordinal, words)     |   |  FFN input FiLM by memory tokens |
        |  fovea goal token     |   |  (3 memory tokens, dim 256)      |
        |  mark on 81 patches   |   +----------------+----------------+
        |  fovea 64x64 crop ->  |                    |
        |   81 extra img tokens |                    v
        +----------------------+      40 x 8 absolute joint chunk,
                                     execute <= 32 steps, re-plan

Training is staged. The TCL moment tokens are trained first. The Eagle features are then cached once. The planner and the controller are trained separately from those caches. Eagle is never updated after TCL.

2. Inputs, observations, action space

Item Value
Environment RoboMME, Panda arm, 10 fps
Cameras front_view, wrist_view, RGB, 256 px shortest edge, crop fraction 0.95
Train-time image augmentation color jitter: brightness 0.3, contrast 0.4, saturation 0.5, hue 0.08
Language the task instruction string
State 8 dims: 7 joint positions, 1 gripper position. No sin/cos encoding
Action 8 dims: 7 absolute joint positions, 1 gripper close command
Action normalisation min-max to [-1, 1], statistics fit on the training demos (processor_official/statistics.json), relative_action: []
Action chunk 40 steps decoded; 50 action tokens in the DiT, loss mask on the first 40
Padding state and action padded to 128 dims inside the model
Dataset data/robomme16: 1600 episodes, 768,897 frames, 16 tasks, 100 episodes per task

Each episode starts with a demo segment (is_demo) followed by the execution segment that begins at exec_start_idx. Tasks that carry no demo video have an empty demo segment.

3. Frozen backbone

The backbone is GR00T N1.6-3B with the Eagle-Block2A-2B-v2 VLM. The VLM is frozen in every stage after TCL (tune_llm=false, tune_visual=false).

  • The language model is truncated to select_layer = 16. The cached feature is the last hidden state of that truncated model, dimension 2048. reproject_vision = false.
  • Each frame yields two views of 9x9 patches each, 162 image tokens, plus text and special tokens and 4 moment tokens. Totals observed in the cache range from 209 to 228 tokens (sampled episodes: 209, 212, 213, 214, 216, 218, 220, 223, 226, 228), always with 162 image tokens.
  • The 4 moment tokens (Section 7) are appended after the text. Eagle is causal, so they cannot change the image or text features, which equal those of the plain Eagle pass. Their 4 output vectors stay in the cached sequence as non-image tokens (image_mask = False). The DiT therefore reads them in its non-image cross-attention blocks.
  • A vlln LayerNorm(2048) is applied to the backbone features inside the action head.
  • The Eagle weights are not part of any trained checkpoint.

4. Subgoal representation

A subgoal is a triple (kind, ordinal, point), plus a short text and the previous subgoal's text as context. It is parsed from the dataset (training, oracle mode) or predicted by the planner (non-oracle test).

Field Form Size
kind pick, place, table, press, move, static, insert (regex on the subgoal string), or none 7 + none = 8 entries
ordinal first to fifth, parsed from the instruction or the subgoal ("for the second time"), or 0 5 + 0 = 6 entries
point <row, col> in the 256 px image, converted to crop coordinates in [0, 1] as (pixel - offset) / size, with size = 256 x 0.95; -1 means absent 2 floats
words lowercase [a-z]+ words of the subgoal text with the point removed, truncated or padded to 20; same for the previous subgoal vocabulary 72 words, 74 embedding rows with pad and unk

The kind regexes are ^pick up, ^(drop .* onto target|place |put down), onto table, ^press, ^move to, ^remain static, ^insert.

Labels are parsed per frame from the dataset column grounded_subgoal_online by qgraph.memory.oracle (planner: oracle_labels_16k7.pt; controller, which also needs the word ids: oracle_labels_16k7wc.pt). At test time in oracle mode the same string comes from the simulator's step info.

5. Planner: a query-conditioned spatio-temporal graph

Run runs/online/planner16_s0, code qgraph/memory/online.py (OnlineQueryGraph, kind: parallel) and qgraph/memory/entities.py (EntityStream). 7.664M trainable parameters. The planner answers one question at every time step: given the instruction q and everything seen so far (demo video and own execution), which subgoal is due now? The answer is the triple (kind, ordinal, point) of Section 4.

5.1 The graph

The planner keeps a graph G_t = (V_t, E_t) that grows by one node every 4 frames. There are two kinds of vertices.

Frame tokens. Every 4th frame f (on the grid f ≡ exec_start mod 4, so demo and execution share the grid) adds a node n. Node n holds 82 token vertices v_{n,p}: p = 1..81 are the 9x9 front-camera patches of the frozen Eagle encoder (the front camera is fixed, so slot p is always the same place on the table), and p = 82 is the mean of the 81 wrist-camera patches. Each node also carries the gripper opening g_n and a flag e_n that says whether the frame belongs to the demo (e_n = 0) or to the robot's own execution (e_n = 1).

Entity tokens. CoTracker (frozen, online variant) follows K = 192 points over every front frame. The points are seeded on the first frame by farthest-point sampling over foreground pixels (pixels whose RGB differs from an empty-table image by more than 60 in summed absolute value). At node n, entity k is the vertex u_{n,k} with its tracked position x_{n,k} in crop coordinates, a visibility bit, and two 12x12 RGB patches sampled around it (the current one and the one at the first node). An entity is a piece of an object (a cube corner, a target disc, a button), so the entity set follows objects when they move, are swapped, or are covered.

The query. The instruction q is encoded once by the frozen Eagle language model (last hidden state, one vector per token, 19 to 24 tokens on RoboMME) and projected to the graph width. Its token vectors w_1..w_L are vertices that every other vertex can read.

Edges. Each layer passes messages along four edge types. All attention is multi-head, pre-norm and residual.

Edge type From → to Rule What it carries
query w_l → v_{n,p} and w_l → u_{n,k} every token reads every instruction token (cross-attention) what to look for: "the third target", "the container hiding the red cube"
temporal v_{m,p} → v_{n,p}, m ≤ n; u_{m,k} → u_{n,k}, m ≤ n same slot or same entity, all earlier nodes (causal), learned bias per head on the log2 node distance (10 buckets) the history of one place on the table, or of one tracked object
spatial (frame) v_{n,p'} → v_{n,p} all 82 tokens of the same node the scene at one moment
spatial (entity) u_{n,k'} → u_{n,k} all entities of the same node, bias −‖x_{n,k} − x_{n,k'}‖² / 2σ_h², σ_h learned per head (initialised from 0.03 to 0.5 of the image) which objects are near each other (cube on target, cube under container)

The graph is causal: no vertex of node n reads anything later than frame n, so the answer at node n can be computed while the robot acts. The edges are fixed by the rule above. What is learned is how much each edge is used, through attention weights.

Query dependence. The instruction enters through the query edges in every layer, before the temporal and spatial edges of that layer. The same video therefore produces different node states under different instructions: for "the first target" the history of the first place event is what the later layers need, for "the third target" it is the third. Training enforces this (Section 5.5, re-query): the same VideoPlaceOrder episode is shown with its instruction rewritten to another ordinal, and the supervised answer moves to the matching demo drop location.

5.2 Layers

Frame-token stream (4 GraphBlocks, width 256, 8 heads, FFN 1024, dropout 0.1).

  1. Input: h⁰_{n,p} = LN(W_in · Eagle_{n,p}) + P_p + MLP_g(25 · g_n) + E_phase(e_n), with P a learned embedding of the 82 slots.
  2. Each block: h ← h + Attn_query(h, w); h ← h + Attn_temporal(h); h ← h + Attn_spatial(h); h ← h + FFN(h).
  3. Event counter, after the first block. For every node and every kind k (7 kinds) a score s_{n,k} = σ(Linear(LN(mean_p h_{n,p}))) is predicted. The running count is C_{n,k} = Σ_{m ≤ n} s_{m,k}, an exact prefix sum. It is encoded as a triangular one-hot over the integers 0..8 (a count of 2.9 reads mostly as 3) and its embedding is added to all 82 tokens of node n. Later blocks can then match "the third place" against the node where the place count reached 3 instead of counting by attention over hundreds of nodes.
  4. Node summary: one learned query token attends the 82 final tokens of node n, giving z_n (256-d).

Entity stream (EntityStream, 2 layers, width 128, 4 heads, FFN 512, dropout 0.1, gradient checkpointing).

  1. Input: a_{n,k} = MLP_pix([patch_{n,k}, patch_{1,k}]) + MLP_mot([10 (x_{n,k} − x_{n−1,k}), 10 (x_{n,k} − x_{1,k}), vis_{n,k}]) + c_n, where the node context c_n = Linear(LN([z_n, C_n])) brings the frame stream's summary and the 7 x 9 count code to every entity.
  2. Each layer: query edge to the instruction, then FiLM by c_n (zero-initialised), then temporal edge, then distance-biased spatial edge, then FFN.
  3. Heads: per entity a pointer logit and a 2-d shift 0.06 · tanh(·) (crop units, at most 14.6 px); per node a null logit from mean_k a_{n,k} + c_n.

5.3 Output

At node n the planner outputs:

  • kind_n = argmax Linear(z_n) over 8 classes (7 kinds + none);
  • ordinal_n = argmax Linear(z_n) over 6 classes (none, first to fifth);
  • point_n = x_{n,k*} + shift_{n,k*} with k* = argmax over the 192 pointer logits and the null logit. If the null logit wins, the point is absent (−1).

At rollout only the newest node is read. A heat head on the 81 front tokens exists in the code with loss weight 0 and is not used.

5.4 Training data and labels

The planner is trained only on the official RoboMME demonstrations: 1600 episodes (16 tasks x 100), split 1280 train / 320 val (20 per task). An episode is the demo video segment (is_demo = 1, empty for tasks without a video) followed by the scripted execution. No extra data is generated.

Labels come from the per-frame grounded_subgoal_online string of the dataset, the same text the simulator reports during rollout. qgraph.memory.oracle parses it into kind (first matching regex of Section 4, otherwise none), ordinal (regex for the (\w+) time), and point (the first <row, col> in the 256 px image, mapped to crop coordinates). Each node takes the label of its frame. Both demo nodes and execution nodes are supervised: on a demo node the label is what the demonstrator is doing at that frame, on an execution node it is what the robot must do now.

Inputs are cached once: Eagle node tokens (cache/online16_s4, qgraph.frames) and CoTracker tracks with patches (cache/entities16_p384_s4, qgraph.entities, a pool of 384 points per episode).

5.5 Losses and training procedure

Per node, normalised over the batch; execution nodes weigh 2.0 and demo nodes 1.0.

Term Weight Definition
kind 1.0 cross-entropy over 8 classes
ordinal 1.0 cross-entropy over 6 classes
events 1.0 each labelled step (a run of nodes with the same kind) must hold event mass 1 for its own kind and 0 for the others: Σ_k (Σ_{n ∈ step} s_{n,k} − [k = kind(step)])², averaged over steps
pointer 1.0 −log Σ_{k ∈ P} softmax(logits)_k, where P is the set of entities within 8 px of the label point, or {null} when the label has no point
locate 0.1 smooth-L1 (β = 1, in pixels) between x_{n,k} + shift_{n,k} and the label point, over entities in P

The counter has no count labels. Its targets come from the step boundaries in the kind labels.

Item Value
Optimiser AdamW, lr 3e-4, weight decay 0.01, grad clip 1.0, bf16 autocast
Schedule 200 warmup steps, cosine to 0 over 26,000 steps
Batch 8 episodes from one task; the task rotates every step, so all 16 tasks weigh the same
Time warp execution nodes are dropped with p 0.1 or repeated 2 to 4 times with p 0.15 (the newest node is kept), so progress cannot be read from elapsed time and a stalled robot looks like one
Symmetry one of the 8 rotations and flips of the square, applied to the 9x9 patch grid, the entity coordinates, the entity patches and the label point
Entity subsets per episode 8 fixed subsets of 192 out of the 384 tracked points (farthest-point orders from 8 starts); one is drawn per sample
Re-query VideoPlaceOrder, p 0.5: the instruction is replaced by the instruction of another training episode that differs only in the ordinal word ("first", "second", ...). The pointer target of the execution place step becomes the entities that lie within 8 px of the matching demo drop location both at the first node and at the end of that drop step
Hardware 1 A100 80GB, token store in host memory, about 3.6 h

5.6 Worked example

Val episode 214, VideoPlaceOrder. Instruction: "watch the video carefully, then place the green cube on the third target it was previously placed on" (19 Eagle tokens). The episode has 1271 frames: a 1069-frame demo, then execution from frame 1069.

The demo, as recorded in grounded_subgoal_online:

Frame Subgoal text Kind
0 pick up the cube at <80, 114> pick
98 drop the cube onto target at <125, 89> place (1st)
189 pick up the cube at <125, 87> pick
261 drop the cube onto target at <93, 91> place (2nd)
340 press the button at <121, 128> press
402 pick up the cube at <99, 83> pick
505 drop the cube onto target at <114, 166> place (3rd)
588 pick up the cube at <113, 168> pick
663 drop the cube onto target at <142, 168> place (4th)
750 pick up the cube at <144, 169> pick
826 drop the cube onto table table
939 to 1068 static, no record none
1069 pick up the cube at <81, 126> pick (execution)
1172 place the cube onto the correct target at <114, 166> place (execution)

Graph size. Nodes sit on frames 1, 5, 9, ..., so the demo gives 267 nodes before the robot moves. Execution adds about 50 more, 318 in total. Each node has 82 frame tokens and 192 entity tokens.

What the layers do on this episode:

  1. Event counter. The demo contains four place steps, so the place count C_{n,place} climbs from 0 to 4 across the demo nodes, reaching 3 at the nodes around frame 505.
  2. Query and temporal edges. The tokens read "third target" from the instruction. Along the temporal edges, the entity tokens on the target disc at <114, 166> can see the cube arrive at the node where the place count reached 3. Those on the targets at <125, 89> and <93, 91> see it arrive at counts 1 and 2.
  3. Spatial edges. These link the cube entities to the target entities they sit on.
  4. Output at an execution node during the place step (for example frame 1173). The pointer selects an entity on the target at <114, 166> (13 to 22 of the 192 tracked points lie within 8 px of it in this episode, any of them is a correct answer). The kind head says place. The ordinal head says none, since the subgoal text carries no "for the n-th time". The controller receives (place, none, point (0.442, 0.656)).

With the instruction "... on the first target ..." the same video and the same tracks must give the point near <125, 89>. Re-query training teaches exactly this.

A counting example, PickXtimes episode 100: "pick up the red cube and place it on the target, repeating this action two times, then press the button to stop". There is no demo. The labels are pick (first time), place, pick (second time), place, press. Here the ordinal head carries the count ("for the second time"), and the event counter's place count decides when "press the button" is due.

5.7 Inference loop

  1. Before the robot moves, the demo frames are replayed through the planner. CoTracker sees every demo frame. Every 4th frame on the grid becomes a node with e_n = 0.
  2. During execution, CoTracker sees every environment frame. A node with e_n = 1 is added every 4 environment steps and at every controller query.
  3. After each new node the whole graph is re-run over all nodes so far. The newest node gives (kind, ordinal, point).
  4. Watch rule (eval.watch=true). If the new (kind, ordinal) differs from the pair the running action chunk was generated with, the chunk is cut and the controller is queried at once with the new subgoal. A change of point alone does not cut the chunk.
  5. One entity subset is used (the first 192 seeds, QGRAPH_DRAWS = 1). There is no voting, no hold rule and no hand-written task logic.

5.8 Parameters per module

Module Params (M)
slot position embedding (82 x 256) 0.021
node projection (2048 → 256, LN) 0.525
gripper MLP 0.066
phase embedding 0.001
instruction projection (2048 → 256, LN) 0.525
4 GraphBlocks 5.273
final norm 0.001
heat head (unused) 0.0003
node summary attention 0.264
kind and ordinal heads 0.004
event counter 0.019
entity stream and heads 0.965
Total 7.664

6. Controller

The controller is the GR00T N1.6 action head Gr00tN1d6ActionHead with an AlternateVLDiT, plus three add-on modules (adapter, conditioning, fovea).

6.1 DiT flow-matching head

Item Value
Blocks 32, alternating: 16 cross-attention (even index), 16 self-attention (odd index)
Heads x head dim 32 x 48, inner dim 1536
Input embedding dim / hidden / output 1536 / 1024 / 1024
Norm, activation, dropout AdaLayerNorm, GELU (tanh), 0.2
Cross-attention routing attend_text_every_n_blocks = 2: blocks with idx % 4 == 0 attend the non-image tokens (text and moment tokens); idx % 4 == 2 attend the image tokens, which include the fovea tokens
Self-attention over [state token, 50 action tokens]
Flow time t = (1 - Beta(1.5, 1)) x 0.999, 1000 buckets
Interpolant x_t = (1 - t) x noise + t x action; target velocity = action - noise
Inference 4 Euler steps, t = 0, 0.25, 0.5, 0.75

Flow loss:

L_flow = mean over valid action steps of || v_theta(x_t, t, c) - (a - eps) ||^2

where c is the conditioning (backbone features, state, memory tokens).

6.2 Memory tokens

Three tokens of dimension 256 form the memory:

  1. Target point token. Fourier encoding with 6 bands (frequencies 2^k pi, sin and cos, 24 inputs) then Linear(24, 256), GELU, Linear(256, 256). A learned absent vector replaces it when the point is -1.
  2. Step token. Sum of a kind embedding, an ordinal embedding, and the mean-pooled output of a 2-layer Transformer encoder (dim 256, 4 heads, FFN 1024, GELU, pre-norm, no dropout) over word embeddings of the subgoal text (20 words) and the previous subgoal text (20 words). The context words share the embedding table and carry a learned context tag; word positions are learned.
  3. Fovea goal token (Section 6.5), appended by Fovea.summarize.

6.3 MemoryConditioning (FiLM)

Each of the 32 DiT blocks has a MemoryModulator attached as a forward pre-hook on block.ff, so it acts on the FFN input, which is the post-attention AdaLN output x.

  • Memory LayerNorm(256).
  • Query = Linear(1536 to 256) of x; key and value = Linear(256 to 256) of the memory. Multi-head attention with 4 heads (head dim 64) through SDPA.
  • FiLM = Linear(256, 256), SiLU, Linear(256, 3072), zero-initialised, giving (shift, scale).
  • Output x' = (1 + scale) x + shift.

Memory keys and values are computed once per bind and reused across the 4 denoising steps.

6.4 Mark, token drop, prefix, proprio dropout

  • Mark. A learned 2048-d mark vector is added to the 81 front-view patch tokens with weight exp(-d^2 / (2 x 0.06^2)), where d is the distance between the patch centre ((i + 0.5) / 9) and the subgoal point. Initial scale 0.5. It is applied before token drop.
  • Token drop. Each front or wrist image token is dropped with probability 0.6 during training (fovea tokens are not dropped).
  • Prefix. The prefix length d is drawn from randint(0, 17). The first d action tokens enter clean at time bucket 999 and are removed from the loss. This lets the head continue a partly executed chunk.
  • Proprio. The state token is used (proprio = true) and zeroed for a whole sample with probability 0.5 (proprio_drop). State noise is 0.
  • Point dropout for CFG. With probability 0.15 the point is set to -1 during training. This gives an unconditional branch at test time.

6.5 Fovea

For a subgoal point the model crops a 64x64 patch from the front 256 px image (stored at 96x96 with Gaussian jitter of 4 px at train time).

  • PixelReader: CNN with channels 3, 32, 32, pool, 64, 64, pool, 128, 128 (width 128), adaptive-average-pooled to 9x9 = 81 tokens, projected to 2048, plus a Fourier position encoding. The 81 tokens are appended to the image tokens.
  • Goal reader (goal, relative, attend, head_place, full_turn): an MLP over [pooled CNN feature, Fourier(centre), mean of the memory tokens, cross-attention readout over the Eagle scene features (hidden 256, 4 heads, with a null key)]. It predicts 10-d pose labels and feeds a zero-initialised goal_token appended to the memory.
  • Auxiliary labels (cache subgoal_end5r.pt), computed by forward kinematics of the demo: sin and cos of the full-turn yaw of the end pose; 10 x the end TCP position; sin and cos of the remaining yaw; 10 x the displacement from gripper to end; and a valid flag in the last column.
  • The reader is pre-trained (subgoal_reader5.safetensors, 24,000 steps, batch 128, lr 1e-3) and loaded with strict=False.
  • Training covers VideoUnmaskSwap, PickXtimes, VideoPlaceOrder, StopCube, and InsertPeg only (500 episodes in the fovea cache). See Section 12.

6.6 Losses

L = L_flow + 1.0 x L_kin + 0.1 x L_goal

  • L_kin (kinematic loss). The clean chunk estimate a_hat = x_t + (1 - t) v is un-normalised with the processor min/max, passed through a differentiable Panda forward-kinematics model, and the TCP position (pose[:3, 3] + 0.1034 x z-axis) is compared with the label chunk. The squared error is divided by a unit of 0.2 and averaged over valid steps.
  • L_goal. Squared error between the goal reader output and the 10-d pose label, masked by the valid flag. It is logged as grasp.

6.7 Classifier-free guidance

At test time v = v_uncond + 2.0 x (v_cond - v_uncond). The unconditional branch replaces the point by the learned absent token and removes the mark. Kind, ordinal, words, and the fovea crop stay. The two branches run as one doubled batch.

6.8 Parameters per module (checkpoint-72000)

Module Params (M)
DiT (action_head.model) 1091.722
action encoder 232.931
state encoder 54.608
action decoder 37.786
position embedding 1.573
vlln 0.004
subtotal, N1.6 action head about 1418.6
conditioning (32 MemoryModulators, 1.381 each) 44.188
adapter (text encoder 1.58, point MLP 0.072, word embedding 0.019, kind, ordinal, mark about 0.002 each) about 1.68
fovea (reader 0.287, project 0.264, scene 0.529, attend 0.263, scene query 0.105, goal 0.236, goal token 0.066, position 0.051, goal out 0.003, grasp 0.001) about 1.805
Total in checkpoint 1466.30

Parameter counts come from tensor shapes in checkpoint-72000. Eagle is not stored.

7. TCL moment-token encoder

Run runs/tcl_k4, 5000 steps, global batch 32, lr 1e-4 (cosine, warmup ratio 0.05, AdamW, weight decay 1e-5), 4 GPUs, bf16.

  • n_moment_tokens = 4 learned vectors of dimension 2048 (moment_tokens) are appended at the end of the Eagle input.
  • Only the 4 moment tokens and the head moment_to_repr (Linear 2048 to 2048, SiLU, Linear 2048 to 2048, about 8.4M parameters) are trained. Eagle is frozen.
  • The moment outputs are mean-pooled over the 4 tokens, projected, and L2-normalised to z.
  • Loss: InfoNCE with temperature tau = 0.07 over one positive and one negative, L = CE([z_a . z_p, z_a . z_n] / tau, label 0). The anchor, an augmented view of the same frame (positive), and a negative frame go through the backbone as three streams.
  • The checkpoint (checkpoint-5000) is the cache encoder at training time and at test time, so the 4 moment outputs are present in the cached features that the DiT sees. The normalisation statistics are re-fit on the training demos (runs/tcl_k4/processor_official).
  • The TCL checkpoint was trained on data/robomme3_train (240 training episodes of PickXtimes, VideoPlaceOrder and VideoUnmaskSwap, 162,595 frames; runs/tcl_k4/experiment_cfg/conf.yaml). It is applied unchanged to all 16 tasks. Because the moment tokens come last in a causal sequence, this choice affects only the 4 moment outputs and leaves the image and text features unchanged.

8. Controller training

Item Value
Cache cache/tcl_k4_16d: 1589 episodes (11 PatternLock episodes skipped for lack of anchors), stride 16, horizon 40, feature dim 2048, 4 moment tokens
Anchors frames with (frame - exec_start) mod 16 in the cached phases; the cache holds phases 0 and 8, giving an anchor every 8 frames
Split qgraph_split5.json: 1500 train, 100 val (20 each of PickXtimes, StopCube, VideoPlaceOrder, VideoUnmaskSwap, InsertPeg). After skipped episodes: 1489 train episodes, 31,543 train anchors
Sampling uniform over training anchors (no task balancing). Share of anchors by task ranges from 1.4% (PatternLock) to 16.3% (PickXtimes)
Steps 72,000 at global batch 32 (micro batch 32, 1 GPU), seed 0
Optimiser AdamW (fused), weight decay 1e-5, grad clip 1.0
Learning rates 5e-5 for action head parameters; 3e-4 for all memory parameters (adapter, conditioning, fovea)
Schedule linear warmup over 5% of steps (3600), cosine decay to 0
Precision bf16 autocast, TF32 on
Checkpoints saved at 24,000, 48,000, 72,000. Val loss is computed offline at saves
Final log step 72,000, loss 0.005156, 18,406 s (5.11 h)
Initialisation from a controller pre-trained on the same data (resume_policy, strict=False; fovea keys are re-initialised when shapes differ)

Each sampled batch picks anchors, reads their cached Eagle features, applies token drop and the mark, builds the memory tokens from the labels, and runs the DiT.

9. Inference

Config configs/qgraph/ctrl5wngml.yaml, overridden on the command line (Appendix).

  1. Before execution the demo frames are added to the planner as nodes. For tasks without a demo this is empty.
  2. The controller receives the current frames and state and the current subgoal, and decodes a 40-step chunk (4 Euler steps, CFG scale 2.0). At most 32 steps are executed (eval.action_steps = 32).
  3. Watch. While a chunk runs, every 4 environment steps (node_stride) a node is added and the planner is re-run over the whole history, demo included. If the new (kind, ordinal) differs from the one used for the current chunk, the chunk is cut and the controller replans. eval.replan=false in the non-oracle run, eval.watch=true.
  4. The planner point is converted to crop coordinates; the fovea crop and the mark are built from it.
  5. The step words come from the planner as well. For the predicted (task, kind, ordinal), the policy looks up the most frequent word sequence of that step in the training labels (cache/words_table16.pt, 59 entries). The context words are the looked-up words of the previous predicted step. The planner has no word output, so object colour and the order of objects that the text names ("the first green cube") come from the most frequent training text, not from the episode.
  6. Episode horizon: max_episode_steps = 1300. Inference seed 0 (per-session generator, eval.inference_seed). eval.state_sigma = 0.
  7. eval.guide = prefix is set, but eval.commit is 0, so commit() returns None and model.sample runs without a prefix. Each chunk is sampled from noise with CFG only.

Oracle mode. The planner is replaced by the simulator's grounded_subgoal_online. With eval.replan=true the chunk is cut whenever the simulator subgoal changes. The controller and the CFG settings are the same.

Latency (from runs/eval_t50_nonoracle16, wall clock with 3 clients sharing one GPU and one policy server): 69 s per episode on average, 450 env steps per episode, 16.8 controller queries per episode (26.7 env steps per query), 0.153 s per env step. These include planner re-runs, CoTracker, and server contention, so they are not single-process latencies.

10. Evaluation protocol

  • Official RoboMME test split (env_split=test), 16 tasks.
  • Non-oracle: 50 episodes per task, 800 in total, chosen stratified by difficulty from env_metadata (subset seed 0). One run, one seed.
  • Oracle: 20 episodes per task, 320 in total. Eleven tasks (220 episodes, 154 successes) come from runs/eval_t20_oracle11_wngml; the five fovea tasks (100 episodes, 78 successes) come from runs/eval_t20_ctrl5wngml_72000. Together 232/320.
  • Each GPU hosts one policy server (ZeroMQ) and 3 clients (clients_per_gpu 3, episodes_per_job 2), on 4 GPUs.
  • Reported metric: success rate per task, averaged over tasks.

11. Compute and results

11.1 Compute (A-class GPU hours, from REPORT_robomme16.md)

Stage GPU-h
TCL moment tokens 5.6
Controller feature cache 1.5
Planner node cache 2.8
CoTracker cache 3.6
Graph pre-training 0.5
Controller co-train (40,000 x 32) 2.7
Controller fine-tune chain (about 6000 to 8000 steps per link) 6.0
Controller final (72,000 x 32) 5.1
Planner (26,000 x 8) 3.6
Total about 31

11.2 Per-task success rate (%)

Task Non-oracle (50 ep) Oracle (20 ep)
BinFill 44 70
PickXtimes 60 80
SwingXtimes 92 85
StopCube 36 95
VideoUnmask 94 100
ButtonUnmask 88 85
VideoUnmaskSwap 82 80
ButtonUnmaskSwap 50 85
PickHighlight 86 70
VideoRepick 76 85
VideoPlaceButton 96 95
VideoPlaceOrder 94 100
MoveCube 46 35
InsertPeg 22 35
PatternLock 20 25
RouteStick 24 35
Average 63.1 (505/800) 72.5 (232/320)

The non-oracle column comes from the run with the words leak described in Section 12. In the clean rerun, the nine tasks finished so far score BinFill 28, ButtonUnmask 74, ButtonUnmaskSwap 48, InsertPeg 8, MoveCube 48, PatternLock 6, PickXtimes 60, VideoPlaceOrder 94 and VideoUnmaskSwap 80. The other seven tasks are still running. This table will be replaced when the rerun finishes.

Reference averages reported in the same table: SimpleMemVLA 88.25, GroundSG 84.08, FrameSamp+Modul 44.51, MemER 42.38, pi0.5 17.93.

12. Known mismatches between code, config, and cache

  • Words leak in the 63.1% run. The policy first filled every oracle_* input from the simulator's subgoal text, and the planner then overwrote only oracle_kind, oracle_ordinal and oracle_point. oracle_words and oracle_context (the current and previous subgoal text, coordinates removed) stayed true, and they feed the step token. qgraph/policy.py (planned_words, QGRAPH_WORDS) now builds both from the planner's step (Section 9, step 5). Setting QGRAPH_WORDS=truth reproduces the old behaviour.
  • Anchor phases. The config lists schedule.phases [0, 4, 8, 12]; the cache holds phases 0 and 8 only. Training sees an anchor every 8 frames.
  • Training episodes. The controller trains on 1489 episodes, not 1600: 100 are held out for the offline val loss, and 11 PatternLock episodes have no anchors.
  • Task sampler. The code samples uniformly over training anchors, not over tasks. Tasks with more anchors per episode (PickXtimes, BinFill) are seen more often.
  • Fovea scope. The fovea cache, reader pre-training, and goal labels cover five tasks. At test time policy.fovea() is active whenever the point is valid, for all 16 tasks. The other 11 tasks therefore see fovea tokens at test time that were masked to zero in the fovea-training anchors. TODO(verify) the effect.
  • Prefix guide. Training uses prefix 16. The eval command sets eval.guide=prefix, but with eval.commit unset, commit() returns no prefix, so no prefix is injected at test time (checked in qgraph/policy.py, act and commit). The prefix training term only makes the head robust to a clean prefix; it is unused at inference.
  • Action steps. schedule.action_steps = 16; eval.action_steps = 32 overrides it.
  • Config comments. Some config comments say the head predicts joint deltas; the targets are absolute joint positions.
  • TCL data. See Section 7.
  • Uniform sampler wording. The report says the controller is trained evenly over the 16 tasks; the code is uniform over anchors.

Appendix: commands and paths

Working directory triquang/graphvla/HAMLET-Isaac-GR00T, interpreter .venv/bin/python, PYTHONPATH=$PWD.

Item Path
Controller config configs/qgraph/ctrl5wngml.yaml (chain: ctrl5wngm, wngk, wngd, wngx, wngr, wngp, wng, ctrl5, oracle_official)
Planner config configs/qgraph/planner16.yaml (inherits planner5, online, entities)
Controller run ../runs/stage_b/ctrl5wngml_k32_alr5e-05_mlr0.0003_s0 (config.yaml, run.json, train_log.jsonl, checkpoint-72000)
Planner run ../runs/online/planner16_s0 (config.yaml, log.jsonl, planner.pt)
TCL run ../runs/tcl_k4 (checkpoint-5000, processor_official)
Controller cache ../cache/tcl_k4_16d
Subgoal vocab ../cache/oracle_labels_16k7wc.vocab.json (72 words)
Split manifests ../data/robomme16/meta/qgraph_split5.json, qgraph_split_planner16.json
Non-oracle eval ../runs/eval_t50_nonoracle16/ctrl5wngml_k32_alr5e-05_mlr0.0003_s0/checkpoint-72000/summary.json
Oracle eval ../runs/eval_t20_oracle11_wngml, ../runs/eval_t20_ctrl5wngml_72000
Scripts ../logs/planner/chain_eval16.sh, chain_oracle11.sh, chain_wngml.sh, chain_planner16.sh

Non-oracle evaluation (logs/planner/chain_eval16.sh):

QGRAPH_PLANNER_LOG=... PYTHONPATH=$PWD .venv/bin/python -m qgraph.evaluate \
  --config configs/qgraph/ctrl5wngml.yaml \
  --checkpoint ../runs/stage_b/ctrl5wngml_k32_alr5e-05_mlr0.0003_s0/checkpoint-72000 \
  eval.replan=false eval.watch=true eval.guide=prefix eval.cfg=2.0 eval.state_sigma=0.0 \
  eval.fill=true "eval.gpus=[0,1,2,3]" eval.port_base=25000 \
  eval.online=../runs/online/planner16_s0 "eval.tasks=[BinFill,...,VideoUnmaskSwap]" \
  eval.subset.per_task=50 eval.runs=../runs/eval_t50_nonoracle16

Oracle evaluation, 11 tasks (logs/planner/chain_oracle11.sh):

PYTHONPATH=$PWD .venv/bin/python -m qgraph.evaluate --config configs/qgraph/ctrl5wngml.yaml \
  --checkpoint ../runs/stage_b/ctrl5wngml_k32_alr5e-05_mlr0.0003_s0/checkpoint-72000 \
  eval.replan=true eval.guide=prefix eval.cfg=2.0 eval.state_sigma=0.0 eval.fill=true \
  "eval.gpus=[1,2,3]" eval.port_base=23000 \
  "eval.tasks=[BinFill,ButtonUnmask,ButtonUnmaskSwap,MoveCube,PatternLock,PickHighlight,RouteStick,SwingXtimes,VideoPlaceButton,VideoRepick,VideoUnmask]" \
  eval.runs=../runs/eval_t20_oracle11_wngml

The five fovea tasks were evaluated by logs/planner/chain_wngml.sh with the same flags and eval.runs=../runs/eval_t20_ctrl5wngml_72000. That run covers PickXtimes, VideoPlaceOrder, VideoUnmaskSwap, StopCube and InsertPeg (10 shards of 2 episodes per task).

Training launch commands for the controller, planner, and TCL are in the corresponding run.json files and logs/planner/chain_*.sh.

Xet Storage Details

Size:
38.8 kB
·
Xet hash:
67c763716b97b337864655a23ceb27c7990d68c7513d9b69bc139ce0da7a0c8d

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.