reading notes: dissecting molmoact2

the best thing about this model is how open it is. whether it is the data collection strategy, the training process, the weights of the model and its finetunes, or how easy it is to tinker around and customise things. i finetuned the action head, and got it working on novel task in less than 30 minutes of teleop data, but we will talk about the inference, training and deployment part of the model/paper later

a highlighted version of the paper’s copy exists here if you’re interested: link

i think broadly the new things that the paper introduces are

• a new backbone, which specialises a general VLM for embodied/spatial perception
• new datasets, this includes teleoped manipulation data on three different embodiments, bimanual YAMs, franka panda, and so100/101s. they also share specifically finetuned versions of the model for each of these datasets (i used the so101 finetune for my experiments)
• a new action tokenizer, there is a whole process to this tokenisation, but it can be seen as a compact token representation of a high frequency action chunk that decodes back accurately
• architectural difference from vanila VLAs, a few changes that make the model architecture interesting. per layer KV cache connection between backbone and flow head, or training two different action heads together are some of these. will elaborate more later
• depth as a layer to reason before chunking actions, intermediate predictions on depth tokens, which act as a chain of thought to produce the final action chunks

note: a lot of these thoughts/notes won’t have structure and might appear random, this is more of a personal accountability and revision exercise for me to go through my paper notes and create a public artifact from them.

how they achieve embodied reasoning¶

there is a 3.3M corpus of data which they train an already existing VLM on top of. This data consists of image/video QA, pointing data, multi-image ego-exo pairs, and abstract reasoning based on the visual scene. there are no robot actions here. the authors have separated backbone specialisation from robot action training.

this dataset is used to finetune an existing VLM called Molmo2, which gives rise to a spatially aware, emboidment aware VLA backbone called Molmo2-ER

the dataset is a combination of multiple efforts to curate and experiment with multimodality. there is RoboPoint, which has 700K samples of image to coordinate pointing data. you get pixel (x,y) as an output label of this dataset. there is also a metric spatial question answer dataset which trains the model on how far, how big and directional awareness of objects in the scene. we also have multiimage reasoning where we fuse wrist camera with third person camera, making the model selfaware about the embodiment’s position from different viewpoints. VST-P(visual spatial tuning, perception) is the dataset used to do these things. we also train with a dataset, which has action consequences as question answer pairs on images. for example, ‘if you turned, whats on your right?’. pretty cool to bake in a level of dynamism in the model. not writing each of the dataset names here, you can find that in the first few sections of the paper. apart from these image QA pairs, we also train on videos, where we add temporality to the dataset. instead of a single frame, the model is trained to learn a spatially and sequentailly coherent group of frames on top of which questions are asked. RoboVQA is the dataset you should checkout for this. there is also a subtle addition of a dataset for training. this is more abstract. the model is also trained to learn what different prompts mean in different enviornment setting. for example, left of chair doesn’t mean left from the egocentric view, but what physcially on the left side of the chair.

i feel that these qa pairs teach spatial skills way better than any robot trajectory diversity could. robot trajectories are expensive, both in terms of time it takes to collect them, and the money it takes to procure equipment for collection. qa pairs are really cheap to manufacture. a lot of it can be generated programmatically from depth, geometry, or existing annotations instead of being hand labeled. so the intuive split is to buy perception with cheap qa and spend the scarce robot data only when we want tight binding of perception to actions. most players have realised this and are working on even larger scale methods to get perception strogner in these backbones. video training is a leading example of this

now to finetune an existing backbone with 3.3M new data points without catastrophic forgetting, they have split the finetuning in two halves. stage 1 force fits the data, whereas stage 2 brings back momlo2’s original data mixing it back into the model. we will discuss this in slightly more detail later

openfast tokenizer and general ways to represent actions in neural nets¶

won’t go into the internals of the tokenizer here but broadly it is taking a chunk of continuous actions as input. one value per joint at each timestep, which are smooth floats when collectively seen over time. it compresses the entire chunk chunks*number of joints, and turns them into a handful of discrete tokens. like llms, there is a reference vocabulary to be used for these tokens. there is a 2048 entry codebook, which can accurately decode back to original action numbers (with a small ronding error). the point is that when actions can be represented as tokens, the model can use the same next token machinery it has for text and images.

but btw, tokenizing is only one of the several ways actions show up in these models. there are different ways we bake in action as a modality into these transformer based architectures.

if actions are outputs to be generated, as discrete tokens, we need to tokenise as mentioend above. RT1 and OpenVLA bin each action value very crudely. openfast is smarter and more acurate with its compression. the tokenised action here can be reused for the next token prediction objective text and vision already use for the VLM.

actions as outputs, but as continuous values. here we don’t require any tokenizer. ACT regresses the numbers directly with an L1 loss, diffusion and flow heads generate the exact action values through iterative denoising. there is no text token prediction that happens here. we directly get the actions.

actions as input or conditioning. when an action feeds the model instead of being predicted, we always send it directly without any processing. we can send these directly as the action float tensors, or through adaptive layer norm parameters in transformer networks.

the main point here is that we only need an action tokeniser whenwe force actions into discrete next token stream to share same functionality/machinery with language. when we use continuous heads, or feed actions as conditoning input, we don’t need any tokenizer. molmoact is interesting because it keeps both a discrete openfast token head and a continuous flow expert at the same time!

how actions are used in a few of these architectures

note that the loss for discrete/autoregressive output always will be cross entropy because discrete outputs compare against a categorical vocabulary to find the next best token. continuous/denoising outputs can have MSE or L1 loss to train against. secondly, even if a model is autoregressive, it isn’t really adding actions back to a history which it feeds as context to chunk outputs. the model is markovian, where it just looks at current frame and state to act. the autoregressive part is predicting the entire output chunk together by sending back parts of the chunk to predict the new one. the history they feed back is just the partially decoded current action, not earlier timesteps.

chunking is orthogonal to discrete vs continuous. these are two independent choices while designing the action head. you can either predict in the token space, or the numbers space. apart from that orthogonally, you can either predict one action or predict an entire chunk of actions. combining discrete and chunked is fast. continuous and chunked is ACT, or the flow expert. discrete and single step is openvla. the only real difference is how the chunk comes out, a discrete chunk is decoded token by token (sequential) while a continuous chunk is predicted in one parallel shot.

this is actually the whole reason fast exists. if you chunk with naive discrete tokens, a one second chunk at 30hz across 7 joints is 210 tokens, and decoding 210 tokens one at a time for every action step is inefficient. fast compresses that down to a handful of opaque tokens, so you keep the discrete, language compatible tokens which are compressed without loss.

reasoning in VLAs and how reasoning can coexist with actions in the stream¶

pure textual chain of thought might be helpful in the initial planning and reflection part of physical ai agents but won’t realy be helpful for real time manipulation. so reasoning if to be baked, needs to be non textual. this reasoning gets added as an intermediate representation to condition the action output. from what i have been reading, these are a few of what has been introduced to models so far

general CoT traces - textual reasoning getting added after analyzing images. these are trained on CoT annotated demos

predicted goal images - predict a future subgoal frame and condition action on it. supervision is basically free and as we can sample a later frame from the same demo video without any additional annotating required. CoT-VLA and even pi0.7 have subgoal images to steer the action models.

point trajectories - predict a path of waypoints overlaid on top of the image before acting. need a point tracker to label, and acts as a guidance mechanism for the model. it is a way for the model to realise what its going to do, before it does it. similar problems we’re doing by coupling world models(next state prediction) to action policies. sometimes the trajectory sketch is also given as input, instead of something the model has to predict as intermediate tokens.

world model rollouts - this is where the cutting edge research is heading towards. we imagine an action sequence and roll them out in a dynamics model step by step. we do this multiple times and pick the best sequence which reached the goal with the least energy/cost spent. there is research on how to make these world model rollouts blazingly fast so that this chain of thought/internal planning can happen in milliseconds (closer to the operating frequency of these policies)

also to be noted that goal image CoT is not a pure world model. world models are forward dynamics where (state,action)->next state, which means that if I do a given a state s, what will be the next state s+1. goal image CoT predicts a target state which is NOT conditioned on a chosen action, an then inferse the aciton to reach it. it is closer to goal conditioning/inverse dynamics. the whole setup would become a world model when the next frame prediction becomes action conditioned and multistep

another advantage of having this intermediate reasoning is how diagnosable the models become during failure. you have an inspectable artificat which can localise failure to a stage. instead of figuring out if the rollout failed because of misplanning, wrong perception, or dexterity limitations, you can just read the chain of thought to see what went wrong. you can see if the model “thought that cup was on the right” which is a perception error, or whether the goal image showed the wrong end state which becomes a planning error, etc

molmoact2 brings another intermediate thinking concept called depth tokens. these are tokens created using Depth Anything (monecular depth) a variational autoencoder compresses to a 10x10(100 position) grid, each depth token being a dicrete code between 0 to 127. So a 320x320 depth image will be downsampled to 32 blocks of 100x100 each. these 100 tokens will be coarse per region 3d geometry emmitted as ordinary autoregressive tokens bracketed <depth_start...<depth_end>

how they get mixed and read in the stream:

the depth tokens aren’t a separate branch with their own decoder. they live in the same autoregressive stream as text, just bracketed inside <depth_start>...<depth_end>. the language model (LM) head that predicts the next text token is the exact same head predicting these depth codes, one at a time. so reasoning here is not an architectural change but just a vocabulary extension. if you wanted to add a second reasoning channel tomorrow, for example point trajectories, you wouldn’t add a new head but just define a new bracket pair plus a few more codes and let the same head emit them in sequence. the cost of these reasoning tokens is an increased context length, and the time it takes for the model to generate these reasoning tokens. now this does hamper the frequency but there are ways to optimise that as well which I will mention later.

the backbone feeds two action heads. the backbone (Molmo2-ER) eats the images, the instruction, the proprioceptive state, and then autoregressively decodes the depth reasoning tokens into the stream. up to here everything is the discrete LM head doing ordinary next token prediction. once that whole prefix exists (vision + text + state + depth reasoning), the continuous flow expert takes over for the actual action chunk. it doesn’t predict tokens, but denoisies the action numbers directly from the backbone context. the backbone context is provided through KV conenctions. we will discuss this too.

regarding the two action heads, the LM head does the reasoning part autoregressively and the context + reasoning is fed to the flow matching head to produce actions. though it is to be noted that during pretraining, the LM head is expected to chunk out actions as the flow matching head comes later during posttraining.

training and inference: the flow matching action expert¶

the action chunk is a fixed sized tensor (30,32) where horizon is upto 30 action timesteps and total proprio dimensions are 32. We assume the max buffer to be 1 second of action chunk at 30HZ leading to 30 timesteps, and assume that maximum robot proprioception data would be included within 32 variables. so101 only fills a handful of those 32 slots (its joints plus gripper), bimanual YAM fills many more, and the rest are masked. lower control frequency is fine too, you just get fewer real steps sitting inside the same padded horizon for example 20hz would just fill 20/30 timesteps. nothing about the model changes here.

actions are normalised per dimension to roughly [-1, 1], but using the 1st and 99th percentile of the training data, not the mean and std. the reason is robustness to outliers. teleop data is full of tiny jerky correction spikes, and if you normalise by standard deviation a few of those stretch the scale and squash all the real motion into a thin band near zero. clipping to q01/q99 throws the spikes away and lets the bulk of the motion use the full range. z-score normalisation, the more common default of subtract mean divide std, doesn’t protect from this. earlier i was just assuming z score normalisation as the norm, this feels more intuitive for action streams.

the flow matching process can be imagined as drawing a straight line from a noise sample to the real action, and teaching a network to point/direct along that line from any point on it. the network knows the slop of that line. now iamgine infinite many points on that line, but the network will always konw where exactly to push that point so that it reaches the target.

concretely, the action label is a and we place it at t=1. at t=0 we put pure gaussian noise ε. the path between them is just the linear interpolation

x_t = (1-t)·ε + t·a

at t=0 this is ε, at t=1 this is a, and in between it slides linearly. because the path is a straight line, its velocity (how x_t moves as t changes) is constant, it doesn’t depend on t at all. if you differentiate that interpolation, the (1-t) term contributes -ε and the t term contributes +a, so

dx_t/dt = a - ε

that a - ε becomes the target. it is just endpoint minus startpoint divided by unit time, which is essentially the constant slope of the line. training is then one forward pass: pick a random t, sample a noise ε, build x_t, and regress the network’s predicted velocity v_θ(x_t, t, prefix) against a - ε with an MSE loss. the network never sees the whole line, it sees one random point on it and learns to name the direction home.

inference reverses this. we don’t know a, so we start at t=0 from a fresh noise sample x_0 = ε and walk towards t=1 by following the velocity field the network learned. this is euler integration of the ODE dx/dt = v_θ. with 10 steps the step size is dt = 0.1, and each step is

x ← x + 0.1 · v_θ(x, t, prefix)

ten of those and we land at x ≈ a, the predicted action chunk. one thing which we need to note is that those 10 flow steps are not the 30 action timesteps. the flow steps are the discretisation of the noise to action ODE, and the entire 30 step (padded 32) horizon gets denoised together, in parallel, at every flow step. flow steps and action steps are different axes.

now another denoising process to get to the target output is called diffusion but that takes 4-5x more steps, making flow matching much faster. another way to make it more fast is reusing repeated variables, instead of computing them again and again. the flow expert’s action tokens cross attend into the backbone’s prefix (the vision/text/state/depth KV). but that prefix does not change across the 10 flow steps, the observation is frozen for the whole control cycle. so the keys and values from the backbone are computed once, before the loop, and reused. inside the loop the only thing that changes is x_t itself, so the only thing you recompute is the small self attention among the evolving action tokens. the heavy cross attention into the backbone is identical every step and is cached.

pseudo code to make it more intuitve:

# heavy, runs ONCE per control cycle
kv = backbone(images, instruction, state)   # also emits depth reasoning into the prefix

x = sample_noise()              # x0 ~ N(0, I), shape (32 horizon, 32 action dim)
dt = 1.0 / 10
for i in range(10):             # euler integration, t: 0 -> 1
    t = i * dt
    v = flow_expert(x, t, kv)   # cross-attends into the SAME frozen kv every step
    x = x + dt * v
chunk = x                       # x at t=1, the action chunk

the backbone line of code is expensive and only needs be run once. the tiny loop can be run iteratively

and because that loop is completely static (same shapes, same 10 iterations, no data dependent branching) we can capture it as a CUDA graph. every GPU op has to be launched by the CPU, and that launch costs the same fixed overhead no matter how small the op is. at batch size 1 the flow loop is hundreds of tiny kernels, so the GPU sits idle most of the time waiting for the CPU to hand it the next one. so it is very much dependent on how these kernels get launched by CPUs. a CUDA graph records the whole loop once and replays it as one unit/function, so you pay the launch cost once instead of per kernel. the paper reports ~2x from this, but in my own setup it was negligible, the heavy backbone forward dominated the cycle and not the little flow loop, so the graph only pays off when the loop is actually your bottleneck.

so the full control cycle is: observe once (one heavy backbone pass that also produces the depth reasoning), then run the 10 step flow loop on top of the frozen prefix to get all 30 action steps, execute them and then reobserve.

btw, there is no runtime verifier anywhere in this loop. nothing checks that the predicted chunk is good before the arm executes it. the whole thing trusts that the backbone’s KV is a strong enough representation, which measn the backbone weights are accurate enough as context to produce accurate action steps. a world model/simulator might help create a “more closed” loop here.

one more efficiency note on the reasoning side is that depth doesn’t get repredicted from scratch every cycle. most of a tabletop scene is static between control steps as only the arm and the object move. so depth patches are compared across frames by cosine similarity and patches that stay similar to the previous frame are treated as unchanged and reused instead of reemitted. so new depth tokens are only created for the parts of the observation space that moved/geometrically changed.

the VLM connector: multi-scale features and vision-language fusion¶

again this is not sequential in terms of the paper, but something I wanted to note down around how vision and language get fused together and how molmo is doing subtle changes in the process.

the basic process is the same as any vanilla VLM: the vision encoder output and the text both get projected into the language model’s embedding space and interleaved as one sequence. what’s interesting in molmoact2 is how the image tokens get pooled and which layers they come from.

pooling first. a ViT hands you a grid of patch features, and you want fewer tokens before they hit the LM. molmoact2 pools each 2x2 window of 4 patches into one token, but with attention instead of a plain average. the query for the window is the mean of its 4 patches, the scores are the dot product of that mean against each patch, and the output is the softmax weighted sum of the 4. so instead of averaging blindly, it upweights whichever patch in the window carries the most signal. the small twist vs other VLMs is exactly that: molmo uses the window mean as the query, where others learn a query parameter. this 2x2 pooling takes a 576 patch image (a 24x24 grid) down to 144 tokens (12x12, since 2x2 pooling divides the count by 4).

concretely, for one window of 4 patch vectors:

window = [p1, p2, p3, p4]              # 4 patches, each a d-dim vector
q      = mean(window)                  # query = the average patch (no learned params)
scores = [dot(q, p) for p in window]   # how aligned each patch is with the consensus
w      = softmax(scores)               # weights, sum to 1
token  = sum(w[i] * window[i])         # one pooled token out, still d-dim

plain average pooling is just this with w = [1/4, 1/4, 1/4, 1/4] hard-coded. the attention version makes the weights data dependent: a patch with a strong feature (big dot with the mean) gets upweighted, a bland or outlier patch gets damped. tile it over every 2x2 window of the 24x24 grid and 576 patches collapse to 144 tokens.

then the layer choice. instead of taking features from the ViT’s last layer like the default, they concatenate layer 3 and layer 9, an early layer and a late one. early layers carry local, spatial, low level detail (edges, texture, where things sit), late layers carry semantic abstraction (what things are). concatenating gives the LM both. the token count stays 144 but the feature dimension doubles, because you’ve stacked two layers’ features per token. an MLP at the end then fuses the two streams non linearly and projects back down into the LM’s embedding width.

this is similar to the idea of having skip connections in resnets, or a u net concatenating encoder features into the decoder where the idea is to not throw away intermediate representations, as the last layer might not be strictly more informative than the middle. there’s an interpretability flavour to it too. a lot of the meaningful, readable structure in a deep network lives in the middle layers and not the final one, so pulling features from the middle and propagating them forward is buying that structure directly.

also, how you connect a video to a language model is similar to images. the only difference being that you sample a few frames, tokenise them in the same way as images, and dump all of the frames in the same self attention together. temporal reasoning across the frames emerges on its own because attetion can relate any frame token to any other with no rigid per frame slot or fixed layout imposed. the only real constraint is compute, so frames are sparsely sampled, 8 frames at 2 fps, and pooled harder, a 3x3 window instead of 2x2, so each video frame costs ~64 tokens instead of 144. that keeps everything inside the total context budget of around 4200 tokens.

these neat subtleties are good to learn if you want to build your own architectures and training pipelines

a bit more on the dataset¶

i broadly touched what kind of data they use for pretraning but will expand a bit more here. a lot of visual data which creates a good VLM doesn’t necessarily make a good VLM which can be tuned towards robotic manipulation. manipulation requires spatial reasoning, some konwledge about world dynamics, knowledge of self/embodiment, etc. these are some of the data themes along with what gaps they fill

the abstract ones are CLEVR, which is synthetic scenes with compositionally structured questions (“what is left of the red cube behind the blue sphere”), which forces chained attribute and relation reasoning with no photographic shortcut to cheat on. The other one is GRiD-3D which teaches object intrinsic reference frames, “left of the chair” meaning the chair’s own left and not the camera’s, a distinction web data almost never labels but instruction following depends on (“put the cup to the left of the toaster” is the toaster’s left).

the design insight is that no single pillar is enough. pointing alone gives you no metric depth, metric QA alone gives you no multi view correspondence, CLEVR alone gives you no real scene grounding. the mix is pretty copmlimentary and I could see this mix come into action zero shot when I prompted the model to pick up an object it hadn’t seen before. It had the understading of what the object is in the scene, and where its placed. grasp was an issue zero shot, but I fixed it with a bit of finetuning. will tell about that in detail later.

the training pipeline: three checkpoints and a finetune¶

the model isn’t trained in one shot, it’s a tiered training process where Molmo2-ER (the VLM, already covered) flows to MolmoAct2-Pretrain (adds discrete action tokens) which eventually becomes MolmoAct2-Post (adds the flow expert) and then for deployment is an embodiment-specific fine-tune (the so101 checkpoint i actually deployed). all stages have different loss functions/different things that get supervised

regarding compute, the posttraining run is 64 H100s for around 2,300 GPU hours. the parallelism is plain data parallel: every GPU holds a full copy of the model, gets a different slice of the batch, computes its own gradients, and then all the GPUs average their gradients together before everyone applies the identical update, so the replicas never drift apart and the wall clock per step is set by the slowest GPU plus the cost of that gradient exchange. the other efficiency trick they used is sequence packing where instead of padding every example out to the 2100 or 4200 token context and wasting compute on pad tokens, you concatenate several short examples into one full length sequence and use the attention mask to stop them seeing each other, so every token in the batch is real gpu work. but these heavy training systems and strategies are not really needed to deploy this model and get it working on your robot arm. I finetuned the model’s action expert on a new task for 20$ of cloud compute using an L4 GPU. the generalisation and zero shot nature of this model allows, light finetuning to be good enough for reliable deployment.

pretraining the key idea is that everything is one token stream through one forward pass, and the LM just predicts the next token without caring whether the thing it’s predicting is a word, a pointing coordinate, a depth code, or an action. a VQA example is <img> ... question answer, a robot example is <img> <setup> <control> state ... action, and both go through the identical forward pass. what changes per example is only the supervision: VQA puts cross entropy on the answer tokens, a robot example puts cross entropy on the discrete action tokens (the openfast ones). the data mix is roughly 90% robot trajectories across the three embodiments (plus a little open-x) and 10% the multimodal stuff (the Molmo2-ER spatial corpus, Molmo2’s old general VQA, a bit of text SFT), and that 10% is there to stop the VLM forgetting its base capabilities while its representation shifts toward action relevant features. an attention mask keeps each packed example attending only to its own tokens.

these are broadly the different tokens that are used in the stream. state tokens are proprioception (joint angles) discretised into 256 bins per dimension. action tokens are the openfast codes from a 2048 entry codebook. and there are special marker tokens: <setup> describes the embodiment (which robot, the camera layout), <control> tags the convention (absolute joint pose vs delta end-effector), and <depth_start>...<depth_end> brackets the reasoning. the setup and control tags are how one model serves many robots and conventions at once, the tag is basically a language flag telling a multilingual model which dialect of “action” to speak. so101 is trained with absolute joint pose

posttraining is where the continuous flow expert gets attached and supervised. the objective becomes L_post = L_lm + L_flow: the LM head keeps its next token loss (text and discrete action tokens), and the flow expert gets its velocity MSE on the continuous chunk. the subtle, important part is knowledge insulation where the expert reads the backbone’s KV cache, but that cache is detached before it enters the expert, so the flow loss updates the expert and its KV projections without back-propagating into the VLM. the VLM only ever learns from the LM loss. the reasoning is that continuous action gradients are noisy at this scale and you don’t want them dragging the backbone’s perception around, so cross entropy (which doubles as a language and vision anchor) owns the backbone, and flow owns the expert. the learning rates carry the same caution: 5e-6 on the vision encoder and connector, 1e-5 on the LM, and a larger 5e-5 on the action expert, since it’s the newest, least trained piece. one more guard that the authors put is that the discrete action span is masked out of the expert’s conditioning, so the expert can’t cheat by reading the very action tokens it’s supposed to predict. the action tokens generated by the LM head get masked from the context stream that flows into the action expert.

the embodiment finetune (the so101 checkpoint) starts from post and differs in a few deliberate ways. it goes robot only (no VQA mixture), bumps the flow samples from K=4 to K=8 for denser supervision on a smaller dataset, and drops knowledge insulation, because at this scale the flow gradients are informative rather than noisy, so letting them touch the backbone actually helps it adapt to the specific robot. same horizon-30, 30 Hz, absolute-joint recipe, with the input camera order randomised per episode because the so100/101 data is a grab-bag of community rigs with no shared camera convention (more on that in the deployment notes).

expanding on the action expert¶

we did the flow math already (straight line from noise to action, predict the constant velocity a - ε, euler-integrate at inference). the expert is actually a small difussion transformer. a few thigns which are slightly changed in the configuration of the architecture.

RMS norm. the blocks normalise with RMSNorm rather than LayerNorm, the same choice qwen and most modern LLMs make where we divide by the root mean square of the activations, instead of subtracting mean from all, and adding a bias. its equally stable.

rms(x) = sqrt( mean(x_i²) )        # root mean square over the d features
layernorm(x) = γ · (x − μ) / σ + β  # center by mean μ, scale by std σ, then learned gain + bias
rmsnorm(x)   = γ · x / rms(x)       # no μ, no β, just rescale the magnitude

timestep conditioning through a gate (adaLN). the flow timestep t has to modulate every block, because the expert should behave differently near pure noise than near the clean action. so the gate gradually opens with t and the timestep is embedded and pushed through a small MLP that emits per block a scale, a shift, and a gate value, and each sublayer becomes

x = x + gate · sublayer(scale · norm(x) + shift)

with the gate initialised to zero. that means the expert starts life as an identity function and learns to inject signal gradually, which is the trick (adaLN-zero) that keeps DiT training stable as a fresh block can’t destabilise the residual stream until it has something useful to add.

KV mapping as a new way of sending in backbone context to action head instead of conditioning on a single final hidden vector from the backbone like earlier VLAs, the expert cross attends into the backbone’s per-layer keys and values. each backbone layer’s K and V get projected through a P_K, P_V map into the expert’s attention space, and the action tokens attend into them. so the expert isn’t handed a compressed summary of the scene, it gets to retrieve from the full depth wise stack of the backbone’s representation and the ablation says this beats the usual final hidden state conditioning

h  = embed(prefix)                          # vision + text + state + depth tokens
kv = []
for layer in backbone.layers:               # L layers deep
    K, V = layer.k_proj(h), layer.v_proj(h)
    kv.append((K, V))                        # stash this layer's keys/values, don't discard
    h = layer(h)                             # backbone forward continues as normal

# the expert mirrors that depth: expert layer l reads backbone layer l
x = noisy_action_chunk
for l, layer in enumerate(expert.layers):
    K_l, V_l = kv[l]
    K_l, V_l = P_K[l](K_l), P_V[l](V_l)            # project into the expert's attention space
    x = x + layer.self_attn(x)                      # action tokens attend to each other
    x = x + layer.cross_attn(q=x, k=K_l, v=V_l)     # + retrieve from backbone layer l
    x = x + layer.mlp(x)

so a block’s forward pass is: RMSnorm -> self-attention over the action tokens -> cross-attention into the (projected, frozen) backbone KV -> RMSnorm -> MLP, all of it adaLN-gated on t. the output projects to a velocity the same shape as the action chunk, and the loss is the masked MSE against a - ε from earlier. post training evaluates that loss at K=4 noise levels per chunk (K=8 when fine tuning regime), reusing the same backbone context for several cheap supervision points along the same trajectory. and the inference time caching (run the backbone once, freeze its KV, only recompute the evolving action tokens across the 10 flow steps) is the same cuda graph optimisation we discussed earlier along with static elements being cached. the cross attention into a fixed KV is waht makes it cacheable.

depth tokens, concretely (Depth Anything + VQ-VAE)¶

depth tokens first need to be generated, and then converted into model readable tokens that can be appended in the stream.

the generator is Depth Anything V2, an existing monocular depth estimator with frozen weights. it takes an RGB frame and outputs a dense depth map without using any depth sensor.

the discretiser is a VQ-VAE, and the “VQ” was the new part for me. a plain VAE would compress the depth map to a continuous latent, but the LM head can’t emit a continuous vector as a next token prediction as it discretely predicts over a vocabulary. so for each spatial cell of the image, the encoder generates a continuous embedding which we quantise to the nearest entry in a learned codebook, turning the depth map into a grid of discrete codes (here a 10x10 grid, 100 tokens, each an integer code 0–127). discrete is exactly what cross entropy and the autoregressive stream need, so the quantisation is required and also in turn helps as a bottleneck that forces a coarse, committed geometry instead of passing pixel level detail through. in short, just remember that the image gets converted to a depth map which further gets converted to discrete tokens that can be sent in the VLA stream and also decoded back through the integer codes that we track

as mentioend before, they also make sure that redundant depth tokens aren’t regenerated again and again. static scenes in environments get cached and only updated if cosine simlarity between two timesteps is low. when they’re training the model on depth token generation, they also subsequently do action training with and without depth tokens. they also add noise to depth tokens so that the model doesn’t get too dependent on correct depth tokens to output correct actions

benchmarks¶

fine-tuning a general VLM on the 3.3M embodied corpus lifts the embodied-QA average from ~47 to ~64, with the metric spatial benchmark (distances, sizes) roughly tripling, and the result beats GPT-5 and Gemini 2.5 Pro on that embodied QA average while being open weight and smaller. the cheap QA training unlocked the perception knowledge that these VLMs needed. the one clear weakness is episodic memory (OpenEQA, ~45 against GPT-5’s ~64): with a 4200-token budget and only 8 sparsely-sampled frames, the model can’t hold a whole episode in context, and there’s no separate memory, every inference is stateless. we can try increasing video frame length during training to solve this problem though. apart from this issue, Molmo2-ER is actually one of the best VLMs so far for embodied reasoning.

the backbone helps the policy but libero as a benchmark is unfortunately saturated as everyone strong sits at 97-98%. though MolmoAct2 is marginally the best, the noticeable SOTA is on the long horizon suite where MolmoAct2-Think performs wayy better than any other model. Groot N1.7 is a close second there

the quality of the policy shows when the eval goes real and out-of-distribution (random camera poses, unseen objects, OOD scenes). there MolmoAct2 roughly doubles the best baseline. on the DROID real world suite it averages 87% against π0.5’s 45% for manipulation tasks. on the SO-100, the platform i deployed, it clears the field too, 57% avg vs π0-SO100’s 45 and SmolVLA’s 2. and on RoboEval it wins on trajectory quality, smoothness, jerk, path efficiency, not just raw success. so on real world eval, the model is actually one of the best open source VLAs available in the market.

on really fine-grained, highly dexterous tasks, its a bit weak. the pattern is strong on reach / pick / place and on motion quality, softer on articulation and tight insertion. i faced a similar gap while deploying the policy on my so101 arm.

also we need better benchmarks/evaluations now.

some ablations¶

these can be used as architectural decisions while training/building your own VLA.

the vlm to expert connection: conditioning the flow expert on the backbone’s per layer KV beats conditioning on the final hidden state (though just by 2 percentage points), and also beats a per-head variant that rigidly ties expert heads to backbone heads. the expert wants raw retrieval across all depths, and buliding a connection for that for policies will be helpful

the number of flow samples K more is better on average (K=8 wins overall), but long-horizon prefers fewer (K=2), because long tasks chain many chunks and compound error, so they care more about low variance gradients than dense coverage of any one chunk’s flow trajectory. a clean bias variance trade living on the sampling axis, and they split it: K=4 to pre-train, K=8 to fine-tune. to think about flow sampling rate in terms of how long of a task needs to be executed?

what to finetune on top of the model: you can either just finetune the action head, finetune the head and lightly nudge the backbone using a lora adapter, or fully finetune the entire architecture. the paper suggests the isolated action head finetuning to be the worsst because the backbone doesnt know how to handle out of distribution tasks and environments. there is going to be better success passing the gradients to the VLM as well. i finetuned just the action head because it was a single task with similar taxonomy of a pick and place task, so I didn’t train the backbone. might run a test on finetuning success rate using all these three methods later

a few open threads¶

the paper helps a lot of us researchers establish what best practices might be in terms of data preperation, system optimisations, architectural nuances, training workflows and even deployment. thankfully the team has also opensourced everything for us to build on top of. in my opinion there is still a lot of scope on iterating and building better versions of these generative action policies.

some immediate things which were on my mind while reading the paper were around if depth the best intermediate reasoning element to have compared to other chain of thought modalities we discussed at the start? i dont realy have a detailed comparative study on this. also i fine-tuned the so101 checkpoint on 30 teleop demos of one pick place task, action expert-only, for about twenty dollars, and it went from reaching but never-grasping to picking the object up most of the times. the ablation says full model adaptation should beat freezing the backbone, but my own controlled probes (a diagnose.py that POSTs isolated inputs at the model) showed perception was intact for my task and the gap was the grasp-close, which lives in the flow head, so freezing the 4.9B-param backbone and training only the smaller expert was cheap and could work. the open question is whether the residual misses are motor (then more demos fix it) or perceptual grasp point localisation. another question to think about while we think to deploy these policies in production. also when can I have a LLM analogous robot policy which can zero shot any task in the world? how to build it? what are the blockers?

deploying it on a real arm¶

the mac had no usable GPU to run MolmoAct2 at anywhere close to 30Hz(33ms), so i split inference off to a cloud L4 on modal over a tiny json-numpy HTTP endpoint(could have used a raw TCP port but the overhead trimmed would have been neglible compared to larger bottlenecks), and the mac just ran the robot and the loop. compared to the ACT I trained earlier on the same task, the VLA could understand a vague object and task really well even though I’m sure that object would be out of distribution for the model. it was also able to reach the object but wouldn’t close the grasp which is a weakness even the evals identified around the model being a strong approacher but a weak grasper. i lightly finetuned it with 30 demos which partly solved the problem.

to be really honest, most of my time went into stitching different interfaces and hardware together. the config had a stale prompt from an old task, so i was asking the model to grab something that wasn’t on the table. the scene camera fed black for the first few seconds because macOS reshuffles OpenCV indices when an iphone shows up as a continuity camera. the arm sat folded at its joint limits, which after the frame conversion lands off the training range and flattens the flow head. some more ugly work on figuring out how lerobot version mismatching lead to errors in how joint action normalisation happens. thankfully molmo team had built a backward compatible converter which i used post teleop.

a structural blocker was that cloud inference can’t run at 30Hz. one chunk was ~1s of compute for me and it was latency bound, a bigger GPU doesn’t help (i measured an A100 slower than the L4, clock beats FLOPs on a sequential loop), and the network round trip stacks on top. so instead of matching the control rate over the wire i chunked: predict a second of actions and play them open loop while the next prediction runs. to mix the two chunks, i tried temporal ensembling (overlap chunks, average the shared timesteps) and it didn’t work for me, my round-trip was slower than one chunk of motion so the next chunk landed after the overlap window was already gone. RTC (inpaint the new chunk onto the actions still in flight) is the better fix but assumes the same thing, latency under one chunk, which my ~1.5s round trip blew past lol. so i reverted both and shipped a synchronous loop which is basically predict, run the whole chunk, stop, predict again. start-stop, but it moved, and there’s no point smoothing a policy that can’t do the task yet. the fix here would be cutting latency first probably through a GPU closer to home, then adding overlap back.

lerobot doesn’t support molmoact2 on its async path, or i couldn’t make it. the async_inference server hardcodes a policy list molmoact2 isn’t in, loads checkpoints in a way that chokes on the repo, and has no RTC, and the RTC reference script needs a local GPU :D. none of it fit robot on a mac plus model on a remote GPU, so i built my own and it’s small: a FastAPI /act server that calls the model’s predict_action on the cloud GPU, and a client that uses lerobot only for hardware (read joint state, drive servos, grab frames) plus my own capture-send-execute loop. json-numpy over HTTPS with a bearer token. the server’s stateless since the policy is markovian, so all episode state lives on the mac and the GPU side is just frames in, action chunk out. a couple hundred lines, and it skips the whole lerobot-async / gRPC / checkpoint-conversion stack

just finetuning the action head actually got it working on the task. this finetuning was more sample efficient than the number of episodes i had to collect to get ACT to work on the task. had more task successes, (i haven’t quantified anything was just evaluating on vibes tbh), and also was able to think about the task and objects in the scene zero shot, right out of the box, which was a nice aha moment for me.

anyway good read, learned a lot, thanks AI2!